Hacktoberfest 2026：メンテナが10月に向けて印を付けた、オープンで初心者向けの issue。 Hacktoberfest の issue を見る

Support bounded-memory CPU export of FusedAdam checkpoint states

オープン
  • 何か問題がありますか？
#3,668 コメント 0 件 リアクション 0 件 担当者 0 名 GitHub で見る

メンテナーはふだん 2 日以内に返信

まだ誰も着手していません。

評価

難易度
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
フォーク
862
平均マージ
4日 21時間
マージ済み PR（30日）
62

環境構築

はじめの一歩

  1. issue を最後まで読み、次にプロジェクトのコントリビューションガイドを読みます。
  2. 着手することを issue にコメントします — 二人が同じ作業をするのを防げます。
  3. リポジトリをフォークし、ブランチを切って変更します。
  4. issue 番号を参照したプルリクエストを送ります。

NVIDIA/TransformerEngine のほかの issue

NVIDIA/TransformerEngine の issue をすべて見る

似ている issue

Python の issue をもっと見る

新しい issue をメールで受け取る

初心者向けの GitHub issue を短くまとめたダイジェスト。