GPT-OSS on-policy distillation fails during checkpoint load (Missing key: model.layers.0.mlp.experts.down_proj)

[uygarmv]

Summary

Trying to use GPT-OSS with on-policy distillation (per #1679: “should work by changing model names”) fails during model initialization when loading the base checkpoint via torch.distributed.checkpoint (DCP). After working around a Transformers attention-impl error by forcing attn_implementation="eager", the run still crashes with a missing key in the checkpoint state_dict:

RuntimeError: Missing key in checkpoint state_dict: model.layers.0.mlp.experts.down_proj.

---

Expected behavior

Running examples/run_distillation_math.py with GPT-OSS teacher/student should initialize successfully and start distillation training (or at least reach the first training step) after swapping model names as described in #1679.

---

Actual behavior

The Ray teacher actors die during initialization (DTensorPolicyWorkerV2.__init__) while loading the base model checkpoint. All ranks report a CheckpointException, with the root cause being a missing key in the checkpoint.

---

Exact command

uv run python examples/run_distillation_math.py \
  policy.model_name="openai/gpt-oss-20b" \
  teacher.model_name="openai/gpt-oss-20b" \
  cluster.gpus_per_node=8

---

Changes made to try GPT-OSS

1) Config change

In distillation_math.yaml:
policy:
dtensor_cfg:
automodel_kwargs:
force_hf: true

2. Workaround for Transformers attention error

Original error:

GptOssForCausalLM does not support an attention implementation through torch.nn.functional.scaled_dot_product_attention yet...
... load your model with attn_implementation="eager"

I applied the suggested workaround by forcing eager attention in nemo_rl/models/automodel/setup.py:
• set attn_impl = "eager"
• and passed attn_implementation="eager" into config load:

model_config = AutoConfig.from_pretrained(
model_name,
torch_dtype=torch.float32, # Always load in float32 for master weights
trust_remote_code=True,
attn_implementation="eager",
**hf_config_overrides,
)

⸻

Error log (root cause)

This is the first actionable failure I see (repeated across ranks):

RuntimeError: Missing key in checkpoint state_dict: model.layers.0.mlp.experts.down_proj.

Full trace excerpt (trimmed to the relevant portion):

ray::teacher-0-7:DTensorPolicyWorkerV2.init()
File ".../dtensor_policy_worker_v2.py", line 228, in init
model_and_optimizer_state = setup_model_and_optimizer(...)
File ".../nemo_rl/models/automodel/setup.py", line 484, in setup_model_and_optimizer
checkpoint_manager.load_base_model(...)
File ".../nemo_rl/utils/automodel_checkpoint.py", line 237, in load_base_model
self.checkpointer.load_base_model(...)
File ".../nemo_automodel/components/checkpoint/checkpointing.py", line 369, in load_base_model
self.load_model(...)
File ".../checkpointing.py", line 474, in _do_load
dcp.load(state_dict, checkpoint_id=path, storage_reader=storage_reader)
...
File ".../torch/distributed/checkpoint/default_planner.py", line 471, in create_default_local_load_plan
raise RuntimeError(f"Missing key in checkpoint state_dict: {fqn}.")
RuntimeError: Missing key in checkpoint state_dict: model.layers.0.mlp.experts.down_proj.

[zpqiu]

Hi @uygarmv We may need some required changes to use DTensor v2 path for GPT-OSS models:

1. Configuration Changes

The key is to properly configure the dtensor_cfg for both policy and teacher models, following the pattern from examples/configs/recipes/llm/sft-gpt-oss-20b-1n8g-fsdp8ep8-automodel.yaml:

policy:
    model_name: "path/to/gpt-oss-20b"
    dequantize_base_checkpoint: true  # if checkpoint is quantized
    dtensor_cfg:
      tensor_parallel_size: 1
      context_parallel_size: 1
      expert_parallel_size: 8
      automodel_kwargs:
        backend:
          _target_: nemo_automodel.components.moe.utils.BackendConfig
          attn: flex
          linear: te
          rms_norm: te
          enable_deepep: true
          fake_balanced_gate: false
          enable_hf_state_dict_adapter: true  # Critical for GPT-OSS checkpoint loading

teacher:
    model_name: "path/to/gpt-oss-20b"
    dequantize_base_checkpoint: true
    dtensor_cfg:
      # Same configuration as policy
      ...

2. Code Changes Required for Custom MoE Models

There may be a bug in the current main branch. (FYI. @terrykong) Custom MoE models loaded via AutoModel can output logits directly as tensors (Tensor/DTensor) rather than wrapped in an object with a .logits attribute.

While the main branch has added handling for this in the train method, it's missing in get_logprobs and get_topk_logits methods. You'll need to manually patch nemo_rl/models/policy/workers/dtensor_policy_worker_v2.py as a temporary workaround:

In get_topk_logits method (around line 1226-1230):
Replace:

if not hasattr(outputs, "logits"):
    logits = self.model.lm_head(outputs.last_hidden_state)
else:
    logits = outputs.logits
del outputs

With:

# Get logits
if isinstance(outputs, (torch.Tensor, DTensor)):
    # custom models (e.g., those coming from AutoModel) can output logits directly
    logits = outputs
elif not hasattr(outputs, "logits"):
    logits = self.model.lm_head(outputs.last_hidden_state)
else:
    logits = outputs.logits

[guyueh1]

@zpqiu is it ok to assign you to fix this ticket?

[zpqiu]

Hi @guyueh1, to use DTensor v2 for distillation, we need to configure the teacher and student models properly, similar to the SFT example. I’ve posted the reference above.

Regarding the necessary code changes mentioned above, I have submitted a PR #1815 to address them. Could you let me know which colleagues I should tag for a review?

[hemildesai]

Hi @zpqiu, code changes in #1815 look good. But we are also doing a refactor for Dtensor policy v2 in #1757 which also has this fix. Its in the extract_logits helper function, can you take a look at that?

[zpqiu]

Hi @hemildesai, I just noticed this PR and it looks great – it should resolve the issue. I'll close #1815.

[riturajj-cerebras]

@uygarmv were you able to solve the issue? if yes, can you please share your final config with version of nemo-rl you're using?