fix ppo value head load bugs

slime/backends/megatron_utils/model.py

@@ -18,7 +18,7 @@
 from megatron.core.optimizer.optimizer import MegatronOptimizer
 from megatron.core.optimizer_param_scheduler import OptimizerParamScheduler
 from megatron.core.pipeline_parallel import get_forward_backward_func
-from megatron.core.utils import get_model_config
+from megatron.core.utils import get_model_config, unwrap_model
 from megatron.training.global_vars import get_args
 from megatron.training.training import get_model
 
@@ -33,6 +33,68 @@
 logger = logging.getLogger(__name__)
 
 
+def _iter_critic_output_layers(model: Sequence[DDP]):
+    for chunk_id, module in enumerate(unwrap_model(model)):
+        output_layer = getattr(module, "output_layer", None)
+        if output_layer is not None:
+            yield chunk_id, output_layer
+
+
+def _critic_output_layer_needs_reinit(args: Namespace, model: Sequence[DDP], role: str) -> bool:
+    if role != "critic" or args.load is None:
+        return False
+
+    from megatron.core.dist_checkpointing.serialization import load_tensors_metadata
+    from megatron.training.checkpointing import get_load_checkpoint_path_by_args
+
+    checkpoint_path = Path(get_load_checkpoint_path_by_args(args))
+    if not (checkpoint_path / ".metadata").is_file():
+        return False
+
+    checkpoint_metadata = load_tensors_metadata(str(checkpoint_path))
+    for _chunk_id, output_layer in _iter_critic_output_layers(model):
+        for name in ("weight", "bias"):
+            param = getattr(output_layer, name, None)
+            if param is None:
+                continue
+
+            param_name = f"output_layer.{name}"
+            ckpt_tensor_metadata = next(
+                (
+                    tensor_metadata
+                    for key, tensor_metadata in checkpoint_metadata.items()
+                    if key == param_name or key.endswith(f".{param_name}")
+                ),
+                None,
+            )
+            expected_shape = tuple(param.shape)
+            checkpoint_shape = tuple(ckpt_tensor_metadata.global_shape) if ckpt_tensor_metadata is not None else None
+            if checkpoint_shape == expected_shape:
+                continue
+
+            reason = (
+                "missing from checkpoint metadata"
+                if checkpoint_shape is None
+                else f"shape mismatch checkpoint={checkpoint_shape} runtime={expected_shape}"
+            )
+            logger.warning(
+                "Will reinitialize critic %s after checkpoint load because it is %s",
+                param_name,
+                reason,
+            )
+            return True
+
+    return False
+
+
+@torch.no_grad()
+def _reinitialize_critic_output_layer(model: Sequence[DDP]) -> None:
+    for _chunk_id, output_layer in _iter_critic_output_layers(model):
+        output_layer.weight.data.normal_(mean=0.0, std=0.02)
+        if output_layer.bias is not None:
+            output_layer.bias.data.zero_()
+
+
 def get_optimizer_param_scheduler(args: Namespace, optimizer: MegatronOptimizer) -> OptimizerParamScheduler:
     """Create and configure the optimizer learning-rate/weight-decay scheduler.
 
@@ -776,6 +838,7 @@ def initialize_model_and_optimizer(
 
     model, optimizer, opt_param_scheduler = setup_model_and_optimizer(args, role)
     model[0].role = role
+    reinit_critic_output_layer = _critic_output_layer_needs_reinit(args, model, role)
     clear_memory()
     iteration, _ = load_checkpoint(
         model,
@@ -784,6 +847,10 @@ def initialize_model_and_optimizer(
         checkpointing_context={},
         skip_load_to_model_and_opt=False,
     )
+    if reinit_critic_output_layer:
+        _reinitialize_critic_output_layer(model)
+        if (args.fp16 or args.bf16) and optimizer is not None:
+            optimizer.reload_model_params()
     clear_memory()
 
     return model, optimizer, opt_param_scheduler, iteration