Skip to content

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 #2024

Description

@xyf5432

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 Require torch>=2.6 #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

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    No labels
    No labels

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions