Baike.dev
Log in
Back to tool/Back to issues
#2024·apex

`torch.cuda.amp` still used in 5 files (incl. private `autocast_mode._cast` / `grad_scaler.OptState`) — deprecated since torch 2.3/2.4, no guards

Author: xyf5432Created Aug 20, 2026Updated Aug 22, 2026

Summary

apex still references the deprecated torch.cuda.amp APIs in 5 files, 10 references at HEAD (9e3568a6). PR #1813 ("deprecate uses of torch.cuda.amp", merged 2024-06-29) migrated 16 files, and the open PR #1867 only covers cudnn_gbn/batch_norm.py (stale since 2024-12, currently not mergeable). The remaining sites below are still unaddressed:

File Line Usage
apex/_autocast_utils.py 26 torch.cuda.amp.autocast_mode._cast(...) — private internal API, the core AMP interop path: _cast_if_autocast_enabled is imported by 5 production modules (apex/fused_dense/fused_dense.py:5, apex/mlp/mlp.py:8, apex/normalization/fused_layer_norm.py:10, apex/contrib/layer_norm/layer_norm.py:5)
apex/contrib/cudnn_gbn/batch_norm.py 5 from torch.cuda.amp import custom_bwd, custom_fwd (covered by stale PR #1867)
apex/contrib/optimizers/distributed_fused_adam.py 2343, 2398 torch.cuda.amp.grad_scaler.OptState — private module symbol
apex/contrib/optimizers/distributed_fused_adam.py 2310, 2371 torch.cuda.amp.GradScaler type annotations
apex/contrib/test/optimizers/test_distributed_fused_lamb.py 4 from torch.cuda.amp import GradScaler
tests/L0/run_mlp/test_mlp.py 79 torch.cuda.amp.autocast_mode.autocast(...) — private internal API

Deprecation status — torch.cuda.amp.autocast / custom_fwd / custom_bwd are deprecated since torch 2.4, GradScaler since torch 2.3. Verified on torch 2.11 that even the private autocast_mode._cast is deprecated-wrapped and emits the FutureWarning:

FutureWarning: `torch.cuda.amp.autocast_mode._cast(value, dtype)` is deprecated. Please use `torch.amp.au...

torch.cuda.amp.grad_scaler.OptState is a private module symbol with no warning at all — it would break silently when the module is removed. The deprecation notice says these APIs "will be removed in a future release".

No version guard, floor-only constraints

  • None of the sites has any guard: no LooseVersion, no hasattr(torch.amp, ...), no try/except fallback around them.
  • setup.py:142-148 only enforces a floor (TORCH_MAJOR==0 and TORCH_MINOR<4 → error; anything ≥0.4 passes), and requirements.txt:8 torch>=2.6.0 is also floor-only (added by #1972, 2025-12).
  • pip metadata does not depend on torch at all: setup.py:916 install_requires=["packaging>20.6"].
  • README.md:22-24 recommends "the latest stable release … or nightly" — so users run current torch, where every site above warns (or silently depends on a private symbol).

Fix direction — precedent already exists in-repo

torch.amp is already used in fused_layer_norm.py:674-720, fused_dense.py:63-75, conv_bias_relu.py:11-89, examples/imagenet/main_amp.py:151 — the migration direction is established. Remaining work:

  • Route _autocast_utils.py through torch.amp.autocast (with device_type="cuda", which the old API defaulted to implicitly).
  • Replace custom_fwd/custom_bwd with torch.amp.custom_fwd(device_type='cuda') / custom_bwd(...).
  • Guard or replace torch.cuda.amp.grad_scaler.OptState in distributed_fused_adam.py (the only non-deprecated-wrapped but private symbol).

Reference

  • #1813 — "deprecate uses of torch.cuda.amp", merged 2024-06-29, migrated 16 files
  • #1867 — open PR for remaining FutureWarnings, stale since 2024-12, covers 1 file
  • PyTorch AMP docs — official deprecation notice
  • huggingface/lerobot#3167 — same migration, merged

Source: NVIDIA/apex

View original on GitHubView discussion on GitHub