[NeMo-RL-DSV4] Review missing optimizer states for MoE gate and DSV4 indexer during checkpoint resume

During DeepSeek V4 bring-up in NeMo-RL, we observed an optimizer checkpoint resume failure in a small Automodel smoke run using a 4-layer DeepSeek V4 variant. We have not reproduced this issue in a standalone Automodel training run or on the full DeepSeek V4 model, so this report is mainly to share the integration-time observation and ask for Automodel maintainers' review.

Issue

The smoke run used a 4-layer DeepSeek V4 variant with BF16 training, EP=8, and checkpoint save/resume enabled. When resuming from a checkpoint saved after step 1, loading the optimizer state failed inside PyTorch DCP with:

TypeError: cannot pickle code objects

After adding temporary debug prints around Checkpointer.load_optimizer(), the underlying issue appeared to be a mismatch between the current optimizer state dict and the saved DCP metadata:

metadata key diff: expected=1627 saved=1597 missing=30 extra=0

The missing optimizer state entries were for:

optim.state.model.layers.{0,1,2,3}.mlp.gate.weight.{step,exp_avg,exp_avg_sq}
optim.state.model.layers.2.self_attn.compressor.indexer.{ape,kv_norm,wgate,wkv,wq_b,weights_proj}.{step,exp_avg,exp_avg_sq}

So the surface DCP error seems to be a secondary wrapper failure; the actionable issue is that some trainable parameters are present in the optimizer param groups at resume time, but their Adam states were not materialized in the saved checkpoint.

Local workaround in fork

In my fork, I added two local mitigations:

https://github.com/zpqiu/Automodel/commit/da506f3df75432f5d77cb2857ba8249657fc0e26

This changes the EP all-gather path in nemo_automodel/components/moe/experts.py to preserve gradients for routing weights when needed:

weights = weights.float()
weights = _all_gather_dim0_var(weights, differentiable=weights.requires_grad)

Previously this path used differentiable=False for weights, which can detach the routing-weight path under EP. In the 4-layer DSV4 smoke run, that appeared to leave mlp.gate.weight in the optimizer param groups without materialized Adam state.

https://github.com/zpqiu/Automodel/commit/26b97d5f97864826b3098aa25c6f931226ff40c4

This freezes DeepseekV4Indexer parameters during training. In the current DSV4 training path, the indexer output is consumed as hard topk(...).indices, so there does not appear to be a continuous autograd path from the loss back to the indexer parameters. Freezing those parameters prevents them from entering the optimizer expected state.

Request

Could you review whether these two patterns are expected in Automodel?

In particular:

• Should EP gathering of MoE routing weights preserve gradients when weights.requires_grad=True?
• Should DeepseekV4Indexer parameters be trainable in the current training path if only hard top-k indices are consumed?
• Is there a better upstream-friendly way to handle optimizer-state gaps for trainable parameters that do not receive gradients before the first checkpoint?

These workarounds have not been rigorously validated for standalone Automodel correctness or performance.

Start with Checkpointer.load_optimizer() and nemo_automodel/components/moe/experts.py, then trace the DeepseekV4Indexer parameters through the checkpoint and optimizer setup. Reproduce the 4-layer BF16 EP=8 smoke resume if possible, and compare the current optimizer metadata with the saved DCP metadata. Done means determining whether the two workarounds are correct and documenting an upstream-safe resolution for missing optimizer states.