fix: correctly resume moe lora ckpt

[ZhiyuLi-Nvidia]

Problem

Resuming training of MoE LoRA or QLoRA models from a checkpoint causes a loss spike (back to random-init levels). The LoRA adapter weights are silently not loaded.

Root Cause

PyTorch DCP (get_model_state_dict / set_model_state_dict) loading cannot handle:

• EP (Expert Parallelism): MoE expert modules use custom FQNs (e.g. gate_up_linear.weight0) that DCP's FQN resolution cannot traverse → KeyError on save, silent skip on load, i.e. the model weight is from scratch.

Solution

Bypass DCP entirely for PEFT models with EP or quantization. All changes in nemo_automodel/components/checkpoint/stateful_wrappers.py:

• Load: _set_peft_state_dict() — matches saved tensors by name, redistributes full tensors back into EP DTensor shards.

Test

• loss curve matched now: https://wandb.ai/nvidia/automodel-dev-zhiyul/workspace?nw=9bbvcvbsr3k

If you haven't finished some of the above items you can still open "Draft" PR.

Additional Information

• Related to # (issue)

[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.

[ZhiyuLi-Nvidia]

/ok to test 0751476