Hacktoberfest 2026: the issues maintainers tagged for October, open and beginner-friendly. Browse Hacktoberfest issues

Support bounded-memory CPU export of FusedAdam checkpoint states

Open
  • Is something wrong?
#3,668 0 comments 0 reactions 0 assignees View on GitHub

Maintainers usually reply within 2 days

Nobody has claimed this yet.

Assessment

Difficulty
4/5
Estimated time
3-5 days
Newbie friendliness
48/100
Issue type
Feature
Clarity
Mostly clear
Activity status
Active
Tech stack
python, pytorch
Domain
machine-learning

Research direction

Start by reading FusedAdam.state_dict() and its get_unscaled_state() conversion, then trace how exported state is wrapped as a DTensor and consumed by sharded DCP. Add coverage for ordinary tensors and CUDA-mesh DTensors, including save/load and continued training, synchronous and asynchronous staging, and the unchanged default full-state path. Done means opt-in CPU export preserves values, dtypes, keys, and DTensor placement metadata while reducing GPU export temporaries; validate the memory and round-trip behavior described in the issue.

Written by the indexing model from the issue text.

Description

Is your feature request related to a problem?

Automodel #4209 reports that the shipped DeepSeek V4.1 Flash recipe completes training on 64 GB200 GPUs, then runs out of memory in FusedAdam.state_dict() during BF16-to-FP32 moment conversion. The recipe already enables checkpoint CPU offload. Its sharded DCP path calls optimizer.state_dict() before offloading the returned dictionary, so all FP32 moment copies coexist on the GPU with the live BF16 states.

For N local parameter elements with two BF16 moments, that export retains an additional 8N bytes of FP32 GPU buffers. Model state, master weights, allocator reserve, and staging buffers contribute separately.

Describe the solution you'd like

Add an opt-in constructor setting, provisionally state_dict_cpu_offload=False, that moves each per-parameter export state to CPU immediately after the existing get_unscaled_state() conversion and DTensor wrapping. The default retains current behavior. The option preserves exported values, dtypes, keys, and DTensor placement metadata; live optimizer state and parameter-group metadata retain their existing behavior.

Proposed usage:

optimizer = FusedAdam(
    model.parameters(),
    exp_avg_dtype=torch.bfloat16,
    exp_avg_sq_dtype=torch.bfloat16,
    state_dict_cpu_offload=True,
)
state = get_optimizer_state_dict(
    model,
    optimizer,
    options=StateDictOptions(
        flatten_optimizer_state_dict=True,
        cpu_offload=True,
    ),
)

Automodel's existing optimizer _target_ configuration forwards constructor arguments to FusedAdam, allowing this policy to reach DCP's no-argument state_dict() call. Applying CPU offload after DTensor wrapping also avoids DTensor.from_local() moving a CPU local tensor back to a CUDA mesh.

Releasing each converted GPU state before exporting the next would reduce moment-export GPU temporaries from 8N to approximately 4 * max_local_state_numel bytes. The same FP32 checkpoint representation and CPU memory requirement remain. Explicit CPU export followed by a full-state gather can add a transfer back to the mesh device; the default full-state path remains unchanged.

Describe alternatives you've considered

  • Offloading after the complete state_dict() returns retains the GPU expansion peak.
  • A state_dict(cpu_offload=True) argument alone does not reach the existing DCP consumer.
  • Exporting raw BF16 moments changes the checkpoint dtype and ownership contract. Reusing the current TE conversion keeps that contract intact.

Additional context

The original precision-aware optimizer review identified CPU checkpoint buffers as a memory-improvement direction. The DTensor checkpoint fix establishes the need to retain distributed metadata and validate resharding.

Validation would cover ordinary tensors and CUDA-mesh DTensors, sharded DCP save/load followed by continued training, synchronous and asynchronous staging, the default full-state path, and peak GPU/CPU memory across independently varied total state size and largest parameter size. This is a placement-API proposal; native memory and round-trip measurements are the next step.

Would this constructor-level policy fit the optimizer API, or is a checkpoint-scoped placement mechanism preferred?

Dominant language
Python
Stars
3.6k
Forks
851
Avg merge
4d 15h
Merged PRs (30d)
54

Getting set up

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

More from NVIDIA/TransformerEngine

All issues in NVIDIA/TransformerEngine

Similar issues

More Python issues

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.