Migrate torch.cuda.amp.autocast to torch.amp

[xyf5432]

Fixes #592

Summary

torch.cuda.amp.autocast has been deprecated since torch 2.4 and is scheduled for removal. This PR migrates the seven uses in the core flagai/ package to torch.amp without raising the torch floor, since the repo documents PyTorch >= 1.8.0 (README "Requirements and Installation") and pins no torch version:

• Imports (6 deprecated from torch.cuda.amp import autocast lines across 5 files — AltDiffusionM18.py carries it twice): removed. Nothing is imported from torch.amp at import time any more; each of these modules defines the version-guarded helper below instead.
  • flagai/model/predictor/predictor.py:21 (public Predictor API)
  • flagai/model/mm/AltDiffusion.py:17, AltDiffusionM18.py:17, 19, Unets/Unet.py:19, modules/diffusionmodules/openaimodel.py:10
• Call sites (5, in the same files): with autocast(): → with _autocast():, where a version-guarded helper keeps both call signatures correct:

  def _autocast(**kwargs):
      """torch.amp.autocast("cuda") on torch >= 2.0, else torch.cuda.amp.autocast."""
      if hasattr(torch, "amp"):
          return torch.amp.autocast("cuda", **kwargs)
      return torch.cuda.amp.autocast(**kwargs)

  Unets/Unet.py and modules/diffusionmodules/openaimodel.py import torch as th, so their copies of this helper reference th (see the review note at the end).
  Why not a plain autocast("cuda")? The two APIs take different call forms: torch.amp.autocast (>= 2.0) requires the device type as the first positional argument, while torch.cuda.amp.autocast (< 2.0) has enabled as its first positional parameter — a bare autocast("cuda") would silently feed "cuda" into the old API's enabled slot. That happens to behave equivalently (truthy check), but it relies on unspecified behavior, so the helper dispatches to the exact signature each version expects.
• modules/diffusionmodules/util.py:136 — same helper pattern for the kwargs form inside the autograd function (torch.cuda.amp.autocast(**ctx.gpu_autocast_kwargs) takes keyword args only):

  def _autocast_ctx(**kwargs):
      """torch.amp.autocast("cuda") on torch >= 2.0, else torch.cuda.amp.autocast."""
      if hasattr(torch, "amp"):
          return torch.amp.autocast("cuda", **kwargs)
      return torch.cuda.amp.autocast(**kwargs)

The vendored examples/ca-lora/src/section-4.1/opendelta/ copy is third-party embedded code and intentionally left untouched.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Validation

Verified on torch 2.13.0+cpu, executing each helper in a namespace that contains only the names its own module binds at top level, so a helper referencing an unbound name fails the check instead of passing silently:

• Both helpers dispatch to torch.amp.autocast("cuda", ...) on torch >= 2.0 — in the bare and the enabled=... kwargs form — with no FutureWarning under warnings.simplefilter("error", FutureWarning).
• Fallback branch: with torch.amp removed from the module (simulating torch < 2.0), both helpers dispatch to torch.cuda.amp.autocast(**kwargs) — the signature the old API actually takes (the old API emits no warning before torch 2.4).
• Control group: torch.cuda.amp.autocast() raises FutureWarning under the same filter — the fix is effective and the filter is sensitive.
• py_compile passes on all six touched files; the only torch.cuda.amp references left in the core flagai/ package are the guarded fallback branches (grep-verified).

User impact

No behavior change on any supported torch version (the documented floor of 1.8 is preserved); the FutureWarning emitted by any Predictor-based inference and AltDiffusion inference/training on torch 2.4+ is eliminated, as is the breakage risk once torch.cuda.amp is removed.

Notes for reviewers

• All five call sites previously used the bare with autocast(): form (default enabled, float16), and _autocast() preserves exactly that on both torch paths — no semantics change.
• Same migration as huggingface/lerobot#3167.

---

This replaces #593, which was closed when its branch was force-pushed to a commit with unrelated history while adding the DCO Signed-off-by line. The code change is identical apart from the review fix noted above (the two th-bound helpers); the commit now carries the sign-off.

🤖 Generated with Claude Code

[Copilot]

🟡 Changes recommended

Fix the undefined torch references in the U-Net and OpenAI model autocast helpers.

Get a fresh assessment by requesting another Copilot review.

Pull request overview

Migrates core AMP usage to torch.amp while retaining compatibility with older PyTorch versions.

Changes:

• Adds version-compatible autocast helpers and updates inference/training call sites.
• Preserves the torch.cuda.amp fallback for PyTorch <2.0.
• Two helpers currently reference undefined torch despite importing it as th, causing critical NameErrors.

File summaries

File	Reviewed change
flagai/model/predictor/predictor.py	Updates Predictor autocast usage.
flagai/model/mm/Unets/Unet.py	Adds autocast helper; helper references undefined torch.
flagai/model/mm/modules/diffusionmodules/util.py	Updates checkpoint backward autocast handling.
flagai/model/mm/modules/diffusionmodules/openaimodel.py	Adds autocast helper; helper references undefined torch.
flagai/model/mm/AltDiffusionM18.py	Updates AltDiffusion M18 autocast usage.
flagai/model/mm/AltDiffusion.py	Updates AltDiffusion autocast usage.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 2
• Review effort level: Lite

---

💡 Add a code-review agent skill or configure MCP servers for context-aware, tailored reviews. Learn more in the docs.

[Copilot]

🟢 Approval recommended

No unresolved issues were identified, and all reviewed changes are approval-ready.

Review details

• Files reviewed: 6/6 changed files
• Comments generated: 0 new
• Review effort level: Lite