Skip to content

fix: load critic from policy checkpoints without value head - #2259

Open
Dodojordi wants to merge 1 commit into
THUDM:mainfrom
Dodojordi:fix/critic-policy-checkpoint-value-head
Open

fix: load critic from policy checkpoints without value head#2259
Dodojordi wants to merge 1 commit into
THUDM:mainfrom
Dodojordi:fix/critic-policy-checkpoint-value-head

Conversation

@Dodojordi

Copy link
Copy Markdown

Summary

Loading a critic from a policy checkpoint can fail before slime gets a chance
to initialize the critic value head. This is especially visible with converted
Qwen3.5 checkpoints whose tied policy output head is absent from the distributed
checkpoint:

KeyError: 'output_layer.bias'

slime already detects a missing or shape-incompatible critic output layer and
plans to reinitialize it after checkpoint loading. However, Megatron's strict
distributed-checkpoint loader tries to restore that output layer first, so the
load fails before the existing reinitialization code can run.

This PR excludes only the incompatible critic value head from the checkpoint
load request. All shared transformer parameters remain strictly loaded. The
original value-head Parameter objects are then restored and initialized by
the existing critic initialization path.

Motivation

A policy language-model head and a critic value head are not interchangeable:

  • A policy head projects hidden states to the vocabulary dimension.
  • A critic head projects hidden states to one scalar value.
  • With tied word embeddings, a policy checkpoint may not store a separate
    output-layer weight at all.

The runtime critic nevertheless owns an independent scalar output_layer.
When its weight or bias is missing from the checkpoint, or has the policy
head's incompatible shape, strict distributed-checkpoint loading treats this
as an error.

The existing metadata check correctly recognizes that the value head must be
reinitialized, but reinitializing it after a failed strict load is too late.
The incompatible tensors must be omitted from the load request itself.

Changes

  • Detect the existing reinit_critic_output_layer condition before loading.
  • Temporarily unregister the critic output_layer.weight and
    output_layer.bias while constructing/loading the checkpoint state.
  • Restore the exact same Parameter objects in a finally block, including
    when checkpoint loading raises an exception.
  • Keep strict checkpoint loading for every compatible/shared model parameter.
  • Reuse the existing critic value-head initialization:
    • weight: normal initialization using init_method_std
    • bias: zero initialization
  • Preserve optimizer references to the original parameters and refresh the
    low-precision optimizer's main parameters through the existing
    reload_model_params() path.

Normal actor loading and critic checkpoint resume are unchanged when the
checkpoint value head already has the expected shape.

Behavior

Load scenario Result
Actor loads a policy checkpoint Unchanged
Critic checkpoint contains a compatible scalar value head Loaded normally
Policy checkpoint omits the output head because embeddings are tied Shared tensors load; critic value head is initialized
Policy checkpoint contains a vocabulary-sized output head Shared tensors load; incompatible head is skipped and critic value head is initialized
Checkpoint loading raises another error Original value-head parameters are still restored

Validation

Added focused CPU unit tests covering:

  • both value-head weight and bias are absent from the checkpoint load request;
  • the exact original Parameter objects are restored afterward;
  • optimizer parameter references remain valid;
  • parameters are restored even if checkpoint loading raises an exception.

Local checks:

tests/test_critic_value_head_load.py: 2 passed
python -m py_compile slime/backends/megatron_utils/model.py tests/test_critic_value_head_load.py
git diff --check

The same loading strategy was also exercised in a ProRL Qwen3.5 TP=8 critic
smoke run (pt-q35-vhead-smk-0807a): the converted distributed checkpoint
loaded successfully, the critic initialized, and a finite critic update was
completed.

Scope

This change does not relax checkpoint loading globally. It skips only a critic
value head that slime has already identified as missing or shape-incompatible;
unexpected mismatches in the rest of the model continue to fail strictly.

@Dodojordi
Dodojordi force-pushed the fix/critic-policy-checkpoint-value-head branch from 8a88537 to 9ce0b5c Compare August 9, 2026 09:01
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