refactor(train): drive SFT/RL training through AutoModel's DeepSpeed-style Engine by HuiyingLi · Pull Request #92 · NVIDIA-NeMo/labs-molt

HuiyingLi · 2026-08-23T07:25:36Z

Summary

Molt stops maintaining its own training-execution loop and drives SFT and RL through AutoModel's DeepSpeed-style eager Engine, following the OpenRLHF↔DeepSpeed boundary:

model_output = actor(sequences, action_mask, ...)
loss = policy_loss(model_output.action_log_probs, old_log_probs, advantages, ...)
actor.model.backward(loss)
actor.model.step()

The Engine is a thin nn.Module wrapper (forward / backward(loss) / step() / zero_grad() / get_global_grad_norm()) around an already-distributed model. There is no Datum, loss callback, output envelope, or packing protocol at this boundary — AutoModel stays RL-oblivious.

Net diff: +2594 / −2936 across 49 files.

Main changes

Training execution delegated to Engine. FsdpStrategy loses its duplicate backward / gradient-sync / clip / optimizer / accumulation loop; SFT, policy, and critic trainers call engine.backward(loss) + engine.step() directly, with the RL formula staying in the trainer.
Model layer split by responsibility. models/base.py (918 → ~330 lines) keeps only the runtime forward contract; checkpoint-loading policy moves to models/loading.py, THD/VLM packing to models/packing.py, and the critic value head to models/critic.py.
Vocab-parallel primitives move to AutoModel. trainer/fsdp/packing.py is deleted; under TP the lm-head output is a vocab-sharded DTensor and log-prob/entropy come from AutoModel's token_log_probs / token_entropy, which handle dense and vocab-sharded logits without gathering the vocabulary.
HybridEP unequal-token handling deleted. Token-count equalization is fixed inside AutoModel's dispatcher (fix(moe): equalize and align per-rank token counts for HybridEP dispatch Automodel#3641, included in the pin), so molt's caller-side padding (~90 lines) is gone.
Legacy paths removed. HF-fallback MoE training and the FA2 varlen THD adapter are intentionally unsupported (fail fast); MoE requires an AutoModel-native implementation.

Ownership boundary

Molt owns:

logical batch → model inputs (padding, THD packing, CP/R3 contexts, VLM prep) and restoring dense token outputs
PPO/GSPO/CISPO, KL, entropy, value objectives; global token normalization
the critic value head: a plain fp32 module installed after AutoModel wraps the backbone, replicated across ranks (rank-0 broadcast, sync_replicated_grads before the step)
CPU optimizer offload (CpuOptimizerOffloader: AdamW step on CPU, params stay on GPU; --fsdp.offload none|optimizer)
optimizer/scheduler construction, checkpoint policy, vLLM refit, Ray orchestration, metrics

AutoModel owns:

distributed model construction and execution (FSDP2/TP/CP/EP), including HybridEP token equalization
Engine: backward, gradient finalization, clipping, optimizer update, gradient clearing
vocab-parallel token log-probability / entropy primitives; routing-replay adapter

Details: docs/automodel-engine-integration.md.

Validation

Unit: tests/unit green (234) in the 26.06 molt container; ruff + compileall clean.
SFT: 2-GPU FSDP2 parity (loss/grads/updates vs reference) and 8-GPU smoke.
RL: MoE RL smoke with CPU optimizer offload (Qwen3-30B-A3B class, 65.5 GB peak, nonzero grads on hard-math data); GAE critic smoke with the post-wrap value head (value_loss 2.79 → 2.27).
Parity on AutoModel side of the boundary: DSV4 PP2/EP2, CP+TP and dense packed THD runs match main.