fix(kimi_k25): map lm_head to model.language_model.lm_head

[akx]

What does this PR do ?

Fixes the Kimi-K2.5 state-dict adapter mapping language_model.lm_head.* to a top-level lm_head.* key that no parameter has, which left the output head at random init on every load through NeMoAutoModelForCausalLM / KimiK25VLForConditionalGeneration.

Changelog

• kimi_k25_vl/state_dict_adapter.py: from_hf maps language_model.lm_head.* to model.language_model.lm_head.* (the model's actual FQN), using the same "model." + key rule the adapter already applies to the other language_model.*` keys.

Background: loading moonshotai/Kimi-K2.5 logged Checkpoint key mismatch ... missing=['model.language_model.lm_head.weight'] unexpected=['lm_head.weight'] and continued, because the base-model load uses set_model_state_dict(strict=False). Training then runs against an uninitialised head with a plausible-looking loss curve. With this fix the warning is gone and step-0 SFT loss on chat data is 2.18 (a random head gives ~ln(163840) ≈ 12).

Before your PR is "Ready for review"

Pre checks:

• Make sure you read and followed Contributor guidelines
• Did you write any new necessary tests?
  • updated the test that asserted the buggy key; added a round-trip test
• Did you add or update any necessary documentation?
  • none needed — no user-facing option changes

Additional Information

• This seems to have been a thing ever since Kimi K2.5 support was introduced in feat: add support for kimi K2.5 VL #1132. Cc @HuiyingLi as the author of that PR.

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

[HuiyingLi]

I think the FQN correction itself is valid, but the current root-cause statement ("the head was left at random init on every load since #1132") does not appear to be accurate.

I tested this with the actual Kimi-K2.5 checkpoint:

• On the PR base with 8xH100 / EP8, every local lm_head shard is byte-for-byte equal to the corresponding checkpoint rows, despite the missing/unexpected-key warning.
• The PR base and this branch also produce identical loss with the same model/input.

The reason is that the multi-rank DCP path builds its HF state dict from tensors that alias the final model storage. DCP writes the checkpoint into the real lm_head before from_hf() returns the incorrect lm_head.weight FQN. The subsequent warning is real, but in this path it does not mean the parameter remained at initialization.

I can reproduce an actual regression at 875637f3 / #3574, but only for the world_size == 1 non-write-through custom-model path:

• Parent 7a6483ba: loaded head exactly matches the checkpoint.
• 875637f3: loaded head does not match the checkpoint.
• This PR: loaded head matches again.

That commit routes the single-rank case through the full-state path, where checkpoint tensors no longer alias model storage and the incorrect FQN is silently ignored by strict=False.

Could you clarify the topology used for the reported step-0 loss and whether 2.18 was true same-input base/head A/B? If it was a normal multi-rank run, this mapping change should not affect the loaded head or loss.

I suggest reframing this as fixing the Kimi adapter contract and the single-rank full-state regression introduced by #3574. It would also be good to add a regression test through the single-device Checkpointer full-state path; the current mapping/round-trip unit tests do not exercise the path that actually fails.