Hacktoberfest 2026: los issues que los mantenedores marcaron para octubre, abiertos y aptos para principiantes. Explorar issues de Hacktoberfest

Support bounded-memory CPU export of FusedAdam checkpoint states

Abierto
  • ¿Algo no está bien?
#3,668 0 comentarios 0 reacciones 0 asignados Ver en GitHub

Los mantenedores suelen responder en 2 días

Nadie ha tomado este issue todavía.

Evaluación

Dificultad
4/5
Tiempo estimado
3-5 días
Aptitud para principiantes
48/100
Tipo de issue
Nueva funcionalidad
Claridad
Bastante claro
Estado de actividad
Activo
Stack tecnológico
python, pytorch
Área
machine-learning

Línea de trabajo

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.

Escrito por el modelo de indexación a partir del texto del issue.

Descripción

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?

Lenguaje dominante
Python
Estrellas
3.6k
Forks
851
Merge medio
4 d 15 h
PR fusionados (30 d)
54

Preparar el entorno

Primeros pasos

  1. Lee el issue completo y luego la guía de contribución del proyecto.
  2. Comenta en el issue que vas a ocuparte — evita que dos personas hagan lo mismo.
  3. Haz un fork del repositorio y trabaja en una rama.
  4. Abre un pull request que haga referencia al número del issue.

Más de NVIDIA/TransformerEngine

Todos los issues de NVIDIA/TransformerEngine

Issues similares

Más issues de Python

Recibe los nuevos issues en tu correo

Un resumen breve de issues de GitHub para principiantes.