Hacktoberfest 2026: le issue che i maintainer hanno segnato per ottobre, aperte e adatte ai principianti. Sfoglia le issue Hacktoberfest

Support bounded-memory CPU export of FusedAdam checkpoint states

Aperta
  • Qualcosa non va?
#3,668 0 commenti 0 reazioni 0 assegnatari Vedi su GitHub

I maintainer di solito rispondono entro 2 giorni

Nessuno ha ancora preso questa issue.

Valutazione

Difficoltà
4/5
Tempo stimato
3-5 giorni
Idoneità per principianti
48/100
Tipo di issue
Funzionalità
Chiarezza
Abbastanza chiara
Stato di attività
Attiva
Stack tecnologico
python, pytorch
Ambito
machine-learning

Direzione di ricerca

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.

Scritto dal modello di indicizzazione a partire dal testo della issue.

Descrizione

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?

Lingua principale
Python
Stelle
3.6k
Fork
862
Merge medio
4g 18h
PR unite (30g)
64

Preparare l'ambiente

Come iniziare

  1. Leggi tutta la issue e poi la guida ai contributi del progetto.
  2. Commenta sulla issue per dire che te ne occupi tu — evita che due persone facciano lo stesso lavoro.
  3. Fai un fork del repository e lavora su un branch.
  4. Apri una pull request che faccia riferimento al numero della issue.

Altre issue di NVIDIA/TransformerEngine

Tutte le issue di NVIDIA/TransformerEngine

Issue simili

Altre issue su Python

Ricevi le nuove issue nella tua casella

Un breve riepilogo di issue GitHub adatte ai principianti.