Skip to content

Replace deprecated torch.cuda.amp API with torch.amp equivalents - #2421

Open
xyf5432 wants to merge 1 commit into
kohya-ss:mainfrom
xyf5432:fix/amp-deprecation
Open

Replace deprecated torch.cuda.amp API with torch.amp equivalents#2421
xyf5432 wants to merge 1 commit into
kohya-ss:mainfrom
xyf5432:fix/amp-deprecation

Conversation

@xyf5432

@xyf5432 xyf5432 commented Aug 20, 2026

Copy link
Copy Markdown

Fixes #2420

Summary

torch.cuda.amp.GradScaler has been deprecated since torch 2.3 and torch.cuda.amp.autocast since torch 2.4, both scheduled for removal. This PR migrates the eight __main__ self-test usages to torch.amp, which provides the same APIs (GradScaler since torch 2.3, autocast since torch 2.0) and is compatible with the torch ≥ 2.6.0 floor the README already declares. The IPEX hijack aliasing in library/ipex/hijacks.py already targets torch.amp, so it only gets a clarifying comment.

Changes:

  • library/sdxl_original_unet.py (2) — torch.cuda.amp.GradScaler(enabled=True)torch.amp.GradScaler("cuda", enabled=True); with torch.cuda.amp.autocast(enabled=True):with torch.amp.autocast("cuda", enabled=True):
  • library/sdxl_original_control_net.py (2) — same migration, keeping dtype=torch.bfloat16
  • networks/control_net_lllite.py (2) — same migration
  • networks/control_net_lllite_for_train.py (2) — same migration, keeping dtype=torch.bfloat16
  • library/ipex/hijacks.py (0) — comment only: the torch.cuda.amp aliasing in ipex_hijacks() (IPEX-only path) replaces torch.cuda.amp with the torch.amp / torch.xpu.amp implementations, so deprecated names are never invoked; the comment documents why the references stay

All eight changes are inside if __name__ == "__main__": self-test blocks (developer smoke tests), none on the main training path (train_network.py / sdxl_train.py / train_util.py have zero torch.cuda.amp references).

Type of change

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

Validation

Verified on torch 2.11.0 (CPU build) with warnings.simplefilter("error", FutureWarning):

  • torch.amp.GradScaler("cuda", enabled=True) constructs and runs warning-free.
  • torch.amp.autocast("cuda", enabled=True) and the dtype=torch.bfloat16 variant enter cleanly (autocast is disabled when CUDA is unavailable, by design).
  • Control group: torch.cuda.amp.GradScaler(enabled=True) raises FutureWarning under the same filter — the migration is effective and the filter is sensitive.
  • IPEX hijack pattern verified: reading torch.cuda.amp.custom_fwd (attribute access, not a call) emits no warning, and after torch.cuda.amp = torch.amp the alias resolves to the torch.amp implementation.
  • py_compile passes on all five touched files.

User impact

No behavior change; the FutureWarning emitted when running the self-test blocks on the documented torch ≥ 2.6.0 install path is eliminated, as is the risk of breakage when torch removes torch.cuda.amp.

Notes for reviewers

  • torch.amp.autocast exists since torch 2.0 and torch.amp.GradScaler since torch 2.3, so no version branching is needed on the declared floor (README: "PyTorch 2.6.0 or later is required").
  • Same migration as ultralytics/yolov5#13244 and huggingface/lerobot#3167.

torch.cuda.amp.GradScaler is deprecated since torch 2.3 and autocast since
torch 2.4, both scheduled for removal. The README already requires torch
2.6.0 or later, where torch.amp provides the same APIs (GradScaler since
2.3, autocast since 2.0). Swap the eight __main__ self-test usages to
torch.amp and document the IPEX hijack aliasing, which already targets
torch.amp.

Co-Authored-By: Claude <noreply@anthropic.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

torch.cuda.amp in __main__ self-test blocks and the IPEX hijack — deprecated since torch 2.3/2.4, unguarded on the main install path

1 participant