refactor(engine): training Engine API

[HuiyingLi]

What does this PR do ?

Adds training Engine and migrates the LLM/VLM finetune recipes onto it. The earlier Datum-based design (structured collation metadata, forward_backward with loss callbacks) was dropped after review; the boundary is now the same one OpenRLHF uses on DeepSpeed:

out = engine(**batch)          # ordinary nn.Module forward
loss = caller_computed_loss(out.logits, ...)
engine.backward(loss)
engine.step()                  # clip + optimizer + scheduler at the accumulation boundary

Engine (nemo_automodel/engine/_engine.py, ~300 lines) wraps an already-distributed model and owns backward, gradient finalization/clipping, optimizer update, gradient clearing, and explicit gradient-accumulation windows (set_gradient_accumulation_steps). It does not know about losses, RL, packing, or collation — callers compute the loss and prepare inputs. This is the execution boundary the molt RL integration builds on (NVIDIA-NeMo/labs-molt#92 pins this branch).

Changelog

• Add nemo_automodel/engine/_engine.py: eager forward / backward(loss) / step() / zero_grad() / get_global_grad_norm() with explicit accumulation windows.
• Add components/loss/vocab_parallel.py: token_log_probs / token_entropy unified over dense tensors and vocab-sharded DTensor logits (distributed reductions, no vocabulary gather).
• Extend components/moe/router_replay.py: routing-replay adapter covering forward, activation recomputation, and backward, including pipeline-split model parts.
• Equalize and align per-rank token counts inside the HybridEP dispatcher so dynamic (variable-length/packed) batches work; also proposed standalone as fix(moe): equalize and align per-rank token counts for HybridEP dispatch #3641.
• Migrate recipes/llm/train_ft.py and recipes/vlm/finetune.py to drive training through Engine, removing duplicated backward/clip/step logic.
• VLM data: generic media collation and variable-resolution pixel padding; preserve media axes for THD forwards (Qwen3-VL, Qwen3.5-MoE).
• Checkpoint: stream passthrough model weights; preserve a supplied config in custom loading; avoid re-enabling input gradients being lost after evaluation under PP.

Validation

• Unit: 1590 tests green in the 26.06 container, including tests/unit_tests/test_engine.py (new), loss/test_vocab_parallel.py (new), moe/test_router_replay.py (new), and moe/test_token_dispatcher.py equalization coverage.
• Functional: tests/functional_tests/moe/run_hybridep_unequal_tokens.py (2×H100, bitwise forward+backward parity for unequal/unaligned token counts); DeepSeek-V4 PP2/EP2 and CP+TP dense-packed THD parity runs match main.
• Downstream: the molt SFT/RL integration (labs-molt#92) runs SFT parity, MoE RL, and critic smokes against this exact revision.

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
• Did you add or update any necessary documentation?

Draft while full CI and reviewer feedback are pending.

Additional Information

• Related: refactor(train): drive SFT/RL training through AutoModel's DeepSpeed-style Engine labs-molt#92 (molt-side integration), fix(moe): equalize and align per-rank token counts for HybridEP dispatch #3641 (HybridEP equalization standalone).
• Multi-axis mRoPE with packed TE THD context parallelism remains intentionally unsupported.

🤖 Generated with Claude Code

[copy-pr-bot]

This pull request requires additional validation before any workflows can run on NVIDIA's runners.

Pull request vetters can view their responsibilities here.

Contributors can view more details about this message here.