Replace deprecated torch.cuda.amp API with torch.amp equivalents - #2421
Open
xyf5432 wants to merge 1 commit into
Open
Replace deprecated torch.cuda.amp API with torch.amp equivalents#2421xyf5432 wants to merge 1 commit into
xyf5432 wants to merge 1 commit into
Conversation
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>
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Fixes #2420
Summary
torch.cuda.amp.GradScalerhas been deprecated since torch 2.3 andtorch.cuda.amp.autocastsince torch 2.4, both scheduled for removal. This PR migrates the eight__main__self-test usages totorch.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 inlibrary/ipex/hijacks.pyalready targetstorch.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, keepingdtype=torch.bfloat16networks/control_net_lllite.py(2) — same migrationnetworks/control_net_lllite_for_train.py(2) — same migration, keepingdtype=torch.bfloat16library/ipex/hijacks.py(0) — comment only: thetorch.cuda.ampaliasing inipex_hijacks()(IPEX-only path) replacestorch.cuda.ampwith thetorch.amp/torch.xpu.ampimplementations, so deprecated names are never invoked; the comment documents why the references stayAll 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.pyhave zerotorch.cuda.ampreferences).Type of change
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 thedtype=torch.bfloat16variant enter cleanly (autocast is disabled when CUDA is unavailable, by design).torch.cuda.amp.GradScaler(enabled=True)raisesFutureWarningunder the same filter — the migration is effective and the filter is sensitive.torch.cuda.amp.custom_fwd(attribute access, not a call) emits no warning, and aftertorch.cuda.amp = torch.ampthe alias resolves to thetorch.ampimplementation.py_compilepasses on all five touched files.User impact
No behavior change; the
FutureWarningemitted 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 removestorch.cuda.amp.Notes for reviewers
torch.amp.autocastexists since torch 2.0 andtorch.amp.GradScalersince torch 2.3, so no version branching is needed on the declared floor (README: "PyTorch 2.6.0 or later is required").