Use shared ModelOpt calibration loop on 0.45+ with 0.44 fallback fix

examples/post_training/modelopt/finetune.py

@@ -2,11 +2,10 @@
 
 """Supervised Finetuning GPT."""
 import itertools
-import json
 import os
 import sys
 from functools import partial
-from typing import Any, Dict, Optional
+from typing import Any, Dict
 
 sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), "../../../")))
 
@@ -17,16 +16,13 @@
 from megatron.core import mpu, tensor_parallel
 from megatron.core.enums import ModelType
 from megatron.core.models.gpt import GPTModel
+from megatron.core.utils import get_batch_on_this_cp_rank
 from megatron.post_training.arguments import add_modelopt_args
 from megatron.post_training.loss_func import loss_func
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
 from megatron.post_training.non_loss_data_func import report_draft_acceptance_length
 from megatron.training import get_args, get_timers, pretrain
-from megatron.training.utils import (
-    get_batch_on_this_cp_rank,
-    get_ltor_masks_and_position_ids,
-    print_rank_0,
-)
+from megatron.training.utils import get_ltor_masks_and_position_ids, print_rank_0
 from utils import get_hf_tokenizer
 from model_provider import model_provider
 from megatron.core.parallel_state import get_context_parallel_group

examples/post_training/modelopt/prune.py

@@ -20,8 +20,20 @@
 import modelopt.torch.prune as mtp
 from modelopt.torch.export import import_mcore_gpt_from_hf
 from modelopt.torch.prune.plugins.mcore_minitron import SUPPORTED_HPARAMS
-from modelopt.torch.utils.dataset_utils import get_dataset_dataloader, get_supported_datasets
+from modelopt.torch.utils import get_dataset_samples
+from modelopt.torch.utils.dataset_utils import get_supported_datasets
 from modelopt.torch.utils.plugins import megatron_generate, megatron_prefill
+
+# modelopt 0.45+ exposes a shared Megatron calibration forward loop. Fall back to an
+# inline pack=True implementation on 0.44 so this script works on both releases.
+try:
+    from modelopt.torch.utils.plugins.megatron_calibration import (
+        get_megatron_calibration_forward_loop,
+    )
+
+    _HAS_SHARED_CALIB = True
+except ImportError:
+    _HAS_SHARED_CALIB = False
 from utils import get_hf_tokenizer
 
 from megatron.core.parallel_state import (
@@ -63,7 +75,7 @@ def add_prune_args(parser):
     group.add_argument(
         "--calib-max-sequence-length",
         type=int,
-        default=512,
+        default=4096,
         help="Maximum sequence length for calibration samples.",
     )
     group.add_argument(
@@ -224,24 +236,44 @@ def _custom_prompt_forward_loop_func(model):
             if all_references[idx] is not None:
                 assert all_references[idx] == generated_texts[0], all_references[idx]
 
-    def _hf_dataset_forward_loop_func(model):
-        if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
-            tokenizer.pad_token = tokenizer.eos_token
-
-        dataloader = get_dataset_dataloader(
+    if _HAS_SHARED_CALIB:
+        forward_loop = get_megatron_calibration_forward_loop(
+            tokenizer,
             dataset_name=args.calib_dataset,
-            tokenizer=tokenizer,
             num_samples=args.calib_size,
-            max_sample_length=args.calib_max_sequence_length,
+            seq_length=args.calib_max_sequence_length,
             batch_size=1,
-            device="cuda",
+            # pack=True uses Megatron pretraining-style global-stream document packing
             pack=True,
         )
-        for sample in tqdm(dataloader, disable=torch.distributed.get_rank()):
-            megatron_prefill(model, sample["input_ids"], skip_return_logits=True)
+    else:
+        # modelopt 0.44 fallback: inline pack=True (concatenate raw samples into a single
+        # EOS-separated token stream, slice into fixed-length chunks). Equivalent in
+        # behavior to get_megatron_calibration_forward_loop at batch_size=1.
+        def forward_loop(model):
+            if not hasattr(tokenizer, "pad_token") or tokenizer.pad_token is None:
+                tokenizer.pad_token = tokenizer.eos_token
+            seq_len = args.calib_max_sequence_length
+            samples = get_dataset_samples(args.calib_dataset, num_samples=args.calib_size * 2)
+            sep_id = tokenizer.eos_token_id
+            token_stream: list[int] = []
+            for s in samples:
+                token_stream.extend(tokenizer.encode(s, add_special_tokens=False))
+                token_stream.append(sep_id)
+                if len(token_stream) >= args.calib_size * seq_len:
+                    break
+            n_chunks = min(args.calib_size, len(token_stream) // seq_len)
+            print_rank_0(
+                f"Calibration packing: {len(samples)} raw samples -> {len(token_stream)} tokens "
+                f"-> {n_chunks} chunks of {seq_len} tokens."
+            )
+            for i in tqdm(range(n_chunks), disable=torch.distributed.get_rank()):
+                chunk = token_stream[i * seq_len : (i + 1) * seq_len]
+                input_ids = torch.tensor([chunk], dtype=torch.long, device="cuda")
+                megatron_prefill(model, input_ids, skip_return_logits=True)
 
     print_rank_0(f"Pruning model with export_config: {args.prune_export_config}")
-    config = {"forward_loop": _hf_dataset_forward_loop_func}
+    config = {"forward_loop": forward_loop}
     if args.prune_intermediate_ckpt is not None:
         config["checkpoint"] = args.prune_intermediate_ckpt
     mtp.prune(

examples/post_training/modelopt/quantize.py

@@ -24,6 +24,18 @@
 from modelopt.torch.utils.dataset_utils import get_dataset_dataloader
 from modelopt.torch.utils.plugins import megatron_generate, megatron_prefill
 
+# modelopt 0.45+ exposes a shared Megatron calibration forward loop. Fall back to the
+# legacy local-JSONL + HF-dataset calibration path on 0.44 so this script works on both
+# releases.
+try:
+    from modelopt.torch.utils.plugins.megatron_calibration import (
+        get_megatron_calibration_forward_loop,
+    )
+
+    _HAS_SHARED_CALIB = True
+except ImportError:
+    _HAS_SHARED_CALIB = False
+
 try:
     import modelopt.torch.quantization.plugins.psx_formats as mtq_psx
 except ImportError:
@@ -37,9 +49,9 @@
     mtq_luts = None
     warnings.warn("luts is not installed. LUTs quantization configs will not be available.")
 
-from megatron.core.parallel_state import get_context_parallel_group
 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.post_training.arguments import add_modelopt_args
 from megatron.post_training.checkpointing import load_modelopt_checkpoint
@@ -79,7 +91,10 @@ def add_text_generate_ptq_args(parser):
     """Add additional arguments for ModelOpt text generation PTQ."""
     group = parser.add_argument_group(title="ModelOpt text generation ptq")
     group.add_argument(
-        "--calib-size", type=int, default=512, help="Number of samples to use for ptq calibration."
+        "--calib-size",
+        type=int,
+        default=1024,
+        help="Number of samples to use for ptq calibration.",
     )
     group.add_argument(
         "--calib-dataset-path-or-name",
@@ -90,7 +105,7 @@ def add_text_generate_ptq_args(parser):
     group.add_argument(
         "--calib-max-sequence-length",
         type=int,
-        default=512,
+        default=4096,
         help="Maximum sequence length for calibration.",
     )
     group.add_argument(
@@ -355,20 +370,32 @@ def _custom_prompt_forward_loop_func(model):
             if all_references[idx] is not None:
                 assert all_references[idx] == generated_texts[0], all_references[idx]
 
-    def _dataset_forward_loop_func(model):
-        dataloader = get_calib_dataloader(
-            dataset_path_or_name=args.calib_dataset_path_or_name,
-            tokenizer=tokenizer,
-            calib_size=args.calib_size,
-            max_sequence_length=args.calib_max_sequence_length,
-            use_random_offset=args.calib_use_random_offset,
+    if _HAS_SHARED_CALIB:
+        _dataset_forward_loop_func = get_megatron_calibration_forward_loop(
+            tokenizer,
+            dataset_name=args.calib_dataset_path_or_name,
+            num_samples=args.calib_size,
+            seq_length=args.calib_max_sequence_length,
             batch_size=args.calib_batch_size,
+            # pack=True uses Megatron pretraining-style global-stream document packing
+            # Leave to False for backward compatibility
+            pack=False,
         )
-        for sample in tqdm(dataloader, disable=torch.distributed.get_rank()):
-            sample = get_batch_on_this_cp_rank(
-                sample, is_hybrid_cp=False, cp_group=get_context_parallel_group()
+    else:
+        def _dataset_forward_loop_func(model):
+            dataloader = get_calib_dataloader(
+                dataset_path_or_name=args.calib_dataset_path_or_name,
+                tokenizer=tokenizer,
+                calib_size=args.calib_size,
+                max_sequence_length=args.calib_max_sequence_length,
+                use_random_offset=args.calib_use_random_offset,
+                batch_size=args.calib_batch_size,
             )
-            megatron_prefill(model, sample["input_ids"], skip_return_logits=True)
+            for sample in tqdm(dataloader, disable=torch.distributed.get_rank()):
+                sample = get_batch_on_this_cp_rank(
+                    sample, is_hybrid_cp=False, cp_group=get_context_parallel_group()
+                )
+                megatron_prefill(model, sample["input_ids"], skip_return_logits=True)
 
     unwrapped_model = unwrap_model(model)[0]