Hacktoberfest 2026：维护者为十月标记出来的 issue，仍然开放、适合新手。 浏览 Hacktoberfest issue

Support bounded-memory CPU export of FusedAdam checkpoint states

未关闭
  • 哪里不对？
#3,668 0 条评论 0 个 reaction 已指派 0 人 在 GitHub 查看

维护者通常 2 天内回复

还没有人认领这个 Issue。

评估

难度
4/5
预计耗时
3-5 天
新手友好度
48/100
Issue 类型
功能
描述清晰度
基本清楚
活跃度
活跃
技术栈
python, pytorch
领域
machine-learning

调研方向

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.

由索引模型根据 Issue 内容生成。

描述

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?

主要语言
Python
星标
3.6k
派生
851
平均合并
4 天 15 小时
30 天内合并 PR
54

环境准备

从这里开始

  1. 先读完整个 Issue，再读项目的贡献指南。
  2. 在 Issue 下留言说明你要接手 —— 这能避免两个人做同样的事。
  3. Fork 仓库，在一个分支上完成修改。
  4. 提交 Pull Request，并在描述里引用这个 Issue 编号。

NVIDIA/TransformerEngine 的其他 Issue

查看 NVIDIA/TransformerEngine 的全部 Issue

相似的 Issue

更多 Python Issue

把新 issue 发到你的邮箱

精选适合新手参与的 GitHub issue 摘要。