Skip to content

[WIP] support value head - #2099

Open
PengchengShi00 wants to merge 2 commits into
InternLM:mainfrom
PengchengShi00:feat/critic
Open

PengchengShi00 wants to merge 2 commits into
InternLM:mainfrom
PengchengShi00:feat/critic

Conversation

@PengchengShi00

Copy link
Copy Markdown
Collaborator

Summary

  • Add a critic ValueHead (nn.Linear(hidden, 1)) selected by TransformerConfig.head_type="value_head". Runtime module name stays lm_head; HF save/load uses value_head.weight. Compose/VL only needs text_config.head_type.
  • Add CriticLossConfig: SFT MSE when old_values is omitted, OpenAI-style clipped value loss when it is provided. value_head rejects MTP and does not tie embeddings.
  • Rename FSDPConfig.fp32_lm_headfp32_head (covers LM and value heads). RolloutConfig.fp32_lm_head is unchanged. Actor checkpoints load into a critic with strict=False; missing value-head weights are randomly initialized on all ranks.

Out of scope

  • Trainer default lm_loss_cfg is still CELossConfig; critic is not wired into TrainEngine / Trainer._prepare_model_input.
  • TrainingWorker.fit remains the actor (advantages / logprobs) path. WorkerConfig.loss_cfg accepts CriticLossConfig for typing only.
  • No CriticWorker, GAE, or vision-side value head.

Test plan

  • tests/engine/test_critic_sft.py with QWEN3_5_MOE_PATH (Qwen3.5-35B-A3B): MTP reject, HF key remap, SFT MSE, clipped value loss, actor→critic load / save_hf / reload
  • Existing RL configs still construct FSDPConfig(fp32_head=...)

@PengchengShi00 PengchengShi00 changed the title support value head [Feature] support value head Sep 20, 2026
@PengchengShi00 PengchengShi00 changed the title [Feature] support value head [WIP] support value head Sep 20, 2026
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

1 participant