RuntimeError: Missing key in checkpoint when training with Nemotron Nano 3 and Qwen3 MoE 30b

Describe the bug

I tried to run training example with Nemotron Nano 3 and Qwen3 MoE 30b and I got the error below

[rank1]:   File "/usr/local/lib/python3.12/dist-packages/torch/distributed/checkpoint/default_planner.py", line 471, in create_default_local_load_plan                                                
[rank1]:     raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")                                                                                                                      
[rank1]: RuntimeError: Missing key in checkpoint state_dict: backbone.embeddings.weight.

# The above is for Nemotron Nano3, for Qwen3 MoE 30b, it is:
RuntimeError: Missing key in checkpoint state_dict: lm_head.weight.

Steps/Code to reproduce bug

1- Pull repo
git clone https://github.com/NVIDIA-NeMo/Automodel.git (commit 7261e73a826a3b9f41b4f392c0cf1166ce49c428)
2- Run container
cd Automodel
docker run --gpus all --network=host -it --rm -v $(pwd):/workspace/Automodel -v $(pwd):/opt/Automodel --shm-size=32g nvcr.io/nvidia/nemo-automodel:25.11.00 /bin/bash

3- Run training
export HF_TOKEN=...
cd /workspace/Automodel
pip install -e .

TORCHDYNAMO_VERBOSE=1 torchrun --nproc-per-node=8 examples/llm_finetune/finetune.py --config examples/llm_finetune/nemotron/nemotron_nano_v3_squad.yaml

Expected behavior

Training should work

Additional context

I didn't use uv this time because I have another issue with uv (basically, torch and torchvision somehow use different CUDA version 12.9 and 13.0 when using uv and the training fails with another error)

The above error happens to nemotron_nano_v3_squad.yaml and qwen3_moe_30b_lora.yaml.
nemotron_nano_v3_squad_peft.yaml works fine