diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index 852720805b..e4e0e8de08 100644 --- a/slime/backends/megatron_utils/model.py +++ b/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