Fix: Import unwrap_model from megatron.core.utils in modelopt examples

examples/post_training/modelopt/convert_model.py

@@ -17,18 +17,16 @@
 from megatron.core import mpu
 from megatron.core.enums import ModelType
 from megatron.core.parallel_state import destroy_model_parallel
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
-from megatron.post_training.utils import (
-    report_current_memory_info,
-    to_empty_if_meta,
-)
+from megatron.post_training.utils import report_current_memory_info, to_empty_if_meta
 from megatron.training import get_args
+from megatron.training.arguments import parse_and_validate_args
 from megatron.training.checkpointing import save_checkpoint
 from megatron.training.initialize import initialize_megatron
-from megatron.training.arguments import parse_and_validate_args
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 ALGO_TO_CONFIG = {

examples/post_training/modelopt/export.py

@@ -13,13 +13,13 @@
 import modelopt.torch.export as mtex
 import torch
 
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.training import get_args, get_model
 from megatron.training.arguments import parse_and_validate_args
 from megatron.training.initialize import initialize_megatron
-from megatron.training.utils import unwrap_model
 from model_provider import model_provider
 
 warnings.filterwarnings('ignore')

examples/post_training/modelopt/generate.py

@@ -14,13 +14,14 @@
 from modelopt.torch.utils.plugins import megatron_generate
 from utils import get_hf_tokenizer
 
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.utils import report_current_memory_info, to_empty_if_meta
 from megatron.training import get_args, get_model, initialize_megatron
 from megatron.training.arguments import parse_and_validate_args
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 warnings.filterwarnings('once')

examples/post_training/modelopt/mmlu.py

@@ -21,13 +21,14 @@
 from modelopt.torch.utils.plugins import megatron_mmlu
 from utils import get_hf_tokenizer
 
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.utils import report_current_memory_info
 from megatron.training import get_args, get_model, initialize_megatron
 from megatron.training.arguments import parse_and_validate_args
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 warnings.filterwarnings("ignore")

examples/post_training/modelopt/offline_feature_extract.py

@@ -12,11 +12,12 @@
 
 from examples.post_training.modelopt.finetune import SFTDataset
 from megatron.core import mpu
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.training import get_args, get_model, get_tokenizer, initialize_megatron
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 

examples/post_training/modelopt/prune.py

@@ -40,14 +40,15 @@
     get_pipeline_model_parallel_group,
     get_tensor_model_parallel_group,
 )
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.utils import report_current_memory_info
 from megatron.training import get_args, get_model, initialize_megatron
 from megatron.training.arguments import parse_and_validate_args
 from megatron.training.checkpointing import save_checkpoint
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 warnings.filterwarnings("ignore")

examples/post_training/modelopt/quantize.py

@@ -52,15 +52,15 @@
 from utils import get_hf_tokenizer
 
 from megatron.core.parallel_state import get_context_parallel_group
-from megatron.core.utils import get_batch_on_this_cp_rank
+from megatron.core.utils import get_batch_on_this_cp_rank, unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.utils import print_distributed_quant_summary, report_current_memory_info
 from megatron.training import get_args, get_model, initialize_megatron
 from megatron.training.arguments import parse_and_validate_args
 from megatron.training.checkpointing import save_checkpoint
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 warnings.filterwarnings("ignore")

examples/post_training/modelopt/validate.py

@@ -11,15 +11,16 @@
 
 import torch
 from modelopt.torch.speculative.plugins.megatron_eagle import MegatronARValidation
+from utils import get_hf_tokenizer
 
+from megatron.core.utils import unwrap_model
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.utils import get_mtbench_chat_data
 from megatron.training import get_args, get_model, initialize_megatron
 from megatron.training.arguments import parse_and_validate_args
-from utils import get_hf_tokenizer
-from megatron.training.utils import print_rank_0, unwrap_model
+from megatron.training.utils import print_rank_0
 from model_provider import model_provider
 
 warnings.filterwarnings('ignore')