[ZenAlexa]

description

add opt-in state_dict_cpu_offload=False to FusedAdam, allowing checkpoint callers that use the no-argument state_dict() API to move each converted state to CPU before exporting the next one. this addresses the GPU expansion peak reported in Automodel #4209, where checkpoint CPU offload runs after the complete optimizer dictionary has already been materialized.

Fixes #3668.

type of change

• new feature, with existing behavior retained by default
• fix for the bounded-memory checkpoint export path

changes

keep the existing state conversion and DTensor wrapping, then copy each export state to CPU when the option is enabled. placing the copy after wrapping avoids CUDA meshes moving a CPU local tensor back to the GPU. exported values, dtypes, keys and distributed metadata remain unchanged, while live optimizer state and parameter-group metadata retain their existing ownership and placement.

add a standalone lifecycle scenario to the existing L0 runner, covering repeated exports, restoration and continued updates, including sharded DCP with the flatten_optimizer_state_dict=True configuration used by Automodel. the option is a constructor setting so the existing optimizer factory can forward it without changing DCP's call signature.

verification

merged main 63451d6b868a through a normal merge at d080a53d1b75 on November 1. the FusedAdam owner retains main's te_device_type import migration, and the CPU export implementation and both QA entry points are unchanged. current-head scoped pre-commit, repository license checks, Python syntax, QA shell syntax and contribution-diff checks pass. the merged revision has not been rebuilt or rerun on a GPU.

earlier native validation

the native build recorded for the 7a3c25482cd0 contribution revision succeeded on H200 with PyTorch 2.12.0.dev20260408+cu128, CUDA 12.8.61, GCC 13.3, cuDNN 9.20.0.48 and the source-pinned CUTLASS revision. the compiled common library, PyTorch extension and Python package are installed together as TE 2.22.0.dev0, without a source overlay or platform-getter replacement. the earlier TE 2.20.2 compatibility run remains historical evidence.

the distributed QA job now invokes the scenario with two ranks, closing the review finding about skipped DTensor/DCP checks. its failure status goes through the job's existing test_fail path.

• 14 ordinary-tensor lifecycle combinations pass, covering FP32/FP16/BF16/FP8 moments, FP32/FP16 master weights, int16 remainders and capturable settings, with unchanged live-state values, devices and storage pointers after export.
• two-rank CUDA-mesh DTensor export, restore and continued native optimizer updates pass for FP32, FP16 and BF16 moments, including uneven shards.
• two-rank sharded DCP passes all four flatten=False/True × sync/async combinations. full-state DCP also passes default/CPU export with synchronous and asynchronous saves, followed by load and continued updates.
• a two-rank checkpoint also loads into one rank and produces bitwise-equal parameters after the next native Adam update.
• that native package passed 27 selected FusedAdam regressions with 12 deselected, using -k "TestFusedAdam and not MXFP8 and not Float8Block". the earlier main/candidate Python comparison passed the same selection for both arms.

the GPU allocation increments measured in that earlier native run were:

parameters	default export peak	CPU export peak
1 × 1,048,576 elements	8 MiB	4 MiB
8 × 1,048,576 elements	64 MiB	4 MiB
32 × 1,048,576 elements	256 MiB	4 MiB
8 × 4,194,304 elements	256 MiB	16 MiB

these measurements use two BF16 moments and report temporary PyTorch allocated memory above the pre-export level. comparing with a baseline that subsequently copies the complete dictionary to CPU gives the same peak differences and the same final CPU payload, with zero additional GPU export allocation retained at the common endpoint. total process memory and checkpoint storage are separate quantities.

a full-state gather can move exported states back to the mesh device, as documented. the complete L0 matrix, minimum supported PyTorch version, larger cross-world-size configurations and original 64-GB200 recipe remain unrun. the FP8 scenarios emit the release's Float8Tensor _to_copy warning and still pass their value and continued-update checks. please run the project's full build and distributed optimizer matrix before merging.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has insufficient remaining quota.

[greptile-apps]

Confidence Score: 5/5

[Medium impact] The PR appears safe to merge; no new actionable issues were found.

Summary

Adds opt-in state_dict_cpu_offload to FusedAdam. Each exported state moves to CPU before the next state is converted.

• FusedAdam can export optimizer states to CPU one at a time.

No new actionable issues were found. GPU tests were not run during this review.

_{Reviews (3) · Last reviewed commit: "chore: merge current main for cpu optimi..." · Reviewed by Greptile}