diff --git a/megatron/post_training/model_builder.py b/megatron/post_training/model_builder.py index 48f783367da..5bd6ca5942e 100644 --- a/megatron/post_training/model_builder.py +++ b/megatron/post_training/model_builder.py @@ -21,12 +21,11 @@ from megatron.core.post_training.modelopt.gpt.state_dict_hooks import ( mcore_gpt_load_te_state_dict_pre_hook, ) -from megatron.post_training.checkpointing import load_modelopt_checkpoint, load_modelopt_state +from megatron.post_training.checkpointing import load_modelopt_state +from megatron.post_training.utils import print_distributed_quant_summary from megatron.training import get_args, print_rank_0 from megatron.training.arguments import core_transformer_config_from_args -from megatron.post_training.utils import print_distributed_quant_summary - def count_parameters_in_layer(model, layer_name): num_params = 0 @@ -121,7 +120,7 @@ def _load_teacher_model_config(checkpoint_path: str) -> Namespace: return Namespace(**args_dict) -def _load_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str, Any]) -> MCoreGPTModel: +def _build_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str, Any]) -> MCoreGPTModel: """Teacher model creator.""" args = get_args() @@ -148,19 +147,10 @@ def _load_teacher_model(config, config_raw: Namespace, model_kwargs: Dict[str, A use_arbitrary_attention_mask=False, ) teacher = MCoreGPTModel(config=config, **model_kwargs) + _add_load_convert_hooks(teacher) - print_rank_0(f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ...") - # [WAR]: load checkpoint will check checkpoint's saved args and rng state if not finetune. - # To avoid error out on loading teacher's checkpoint, we temporarily set args.finetune to - # True while loading the teacher checkpoint. - original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format - args.finetune = True - if args.export_kd_teacher_ckpt_format is not None: - args.ckpt_format = args.export_kd_teacher_ckpt_format - load_modelopt_checkpoint([teacher], load_arg='export_kd_teacher_load') - args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format - print_rank_0("...teacher loaded successfully.") + # NOTE: Checkpoint loading now handled in `megatron/training/checkpointing.py`. return teacher @@ -354,7 +344,7 @@ def modelopt_gpt_hybrid_builder( args.export_kd_cfg, student_cfg=config, teacher_cfg=teacher_config ) kd_config = { - "teacher_model": _load_teacher_model(teacher_config, teacher_config_raw, model_kwargs), + "teacher_model": _build_teacher_model(teacher_config, teacher_config_raw, model_kwargs), "criterion": distill_cfg.criterion, "loss_balancer": distill_cfg.loss_balancer, } @@ -365,7 +355,7 @@ def modelopt_gpt_hybrid_builder( mtd_mcore.adjust_distillation_model_for_mcore(model, distill_cfg) # Also remove KD mode state to prevent issues with re-conversion after restore. mto.ModeloptStateManager(model).state_dict().pop() # TODO(aanoosheh): remove once fixed in ModelOpt - + print_distributed_quant_summary(model) return model diff --git a/megatron/training/checkpointing.py b/megatron/training/checkpointing.py index 389cb7ee7e7..d4dae645e76 100644 --- a/megatron/training/checkpointing.py +++ b/megatron/training/checkpointing.py @@ -39,7 +39,7 @@ from megatron.core.num_microbatches_calculator import update_num_microbatches from megatron.core.optimizer import DistributedOptimizer from megatron.core.rerun_state_machine import get_rerun_state_machine -from megatron.core.utils import get_pg_rank, get_pg_size, get_torch_version, is_torch_min_version +from megatron.core.utils import get_pg_rank, get_pg_size from ..core.dist_checkpointing.utils import _clean_metadata_for_serialization from . import ft_integration, wandb_utils @@ -2033,10 +2033,7 @@ def load_model_state_dict(module, state_dict, strict: bool): f'[ t {mpu.get_tensor_model_parallel_rank() + 1}/{mpu.get_tensor_model_parallel_world_size()}, ' f'p {mpu.get_pipeline_model_parallel_rank() + 1}/{mpu.get_pipeline_model_parallel_world_size()} ] ' f'at iteration {iteration}') - - if has_nvidia_modelopt: - print_distributed_quant_summary(model, msg="After loading checkpoint") - + # Additional callback for wandb (last rank) if not torch.distributed.is_initialized() \ or is_last_rank(): @@ -2059,6 +2056,28 @@ def load_model_state_dict(module, state_dict, strict: bool): print_rank_0(">>> Inserting 'default_config' field into optimizer.param_groups...") log_printed = True + if has_nvidia_modelopt: + print_distributed_quant_summary(model, msg="After loading checkpoint") + + # Load teacher model in Distillation mode. + if getattr(args, "export_kd_teacher_load", None): + from megatron.post_training.checkpointing import load_modelopt_checkpoint + + unwrapped_model = unwrap_model(model)[0] + # Note: load_modelopt_checkpoint may call this function so we prevent infinite recursion. + if hasattr(unwrapped_model, 'teacher_model'): + teacher = unwrapped_model.teacher_model + print_rank_0(f"Loading teacher as {type(teacher).__name__} from {args.export_kd_teacher_load} ...") + # [WAR]: To avoid error out on loading teacher's checkpoint, we temporarily + # set args.finetune to True while loading the teacher checkpoint. + original_args_finetune, original_ckpt_format = args.finetune, args.ckpt_format + args.finetune = True + if args.export_kd_teacher_ckpt_format is not None: + args.ckpt_format = args.export_kd_teacher_ckpt_format + load_modelopt_checkpoint([teacher], load_arg='export_kd_teacher_load') + args.finetune, args.ckpt_format = original_args_finetune, original_ckpt_format + print_rank_0("... teacher loaded successfully.") + return iteration, num_floating_point_operations_so_far