From ab1d17c98b92ddb9527b8305ef858444b3eb8a96 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Tue, 19 May 2026 16:14:23 -0700
Subject: [PATCH 1/4] Make ModelOpt calibration loop dual-compatible with 0.44
 and 0.45

The new modelopt 0.45 shared util get_megatron_calibration_forward_loop
unifies prune/quantize calibration with pack=True. Wrap both prune.py and
quantize.py with try-import + _HAS_SHARED_CALIB so they continue to work
on modelopt 0.44 (inline pack=True for prune, legacy local-JSONL /
HF-dataset pad+truncate for quantize) and use the shared util on 0.45+.

Co-Authored-By: Claude Opus 4.7 <noreply@anthropic.com>
Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/post_training/modelopt/prune.py    |  58 +++++--
 examples/post_training/modelopt/quantize.py | 164 +++++++++-----------
 2 files changed, 121 insertions(+), 101 deletions(-)

diff --git a/examples/post_training/modelopt/prune.py b/examples/post_training/modelopt/prune.py
index ebc474254b0..17a5876553a 100644
--- a/examples/post_training/modelopt/prune.py
+++ b/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,42 @@ 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,
         )
-        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(
diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index 9ed5de36313..f0eaefc5b9f 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -7,7 +7,6 @@
 import inspect
 import json
 import os
-import random
 import sys
 import warnings
 
@@ -24,6 +23,20 @@
 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. The 0.45 path uses pack=True (no padding/truncation loss); the 0.44 fallback
+# uses padding+truncation, which is slightly worse for long-document corpora but
+# functional.
+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:
@@ -39,7 +52,6 @@
 
 from utils import get_hf_tokenizer
 
-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
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
@@ -78,7 +90,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",
@@ -89,7 +104,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(
@@ -223,79 +238,6 @@ def get_modelopt_torch_quantization_config():
     return mtq_config
 
 
-def get_calib_dataloader(
-    dataset_path_or_name,
-    tokenizer,
-    calib_size=512,
-    max_sequence_length=512,
-    use_random_offset=False,
-    batch_size=1,
-):
-    """Return a dataloader/iterator for calibration using SFT or HF datasets.
-
-    Supports either a local path (.jsonl) or a HuggingFace dataset name.
-    """
-    if os.path.isfile(dataset_path_or_name):
-        # Local file
-        print_rank_0(f"Loading calibration dataset from local file: {dataset_path_or_name}")
-        all_texts = []
-        with open(dataset_path_or_name) as f:
-            for i, line in enumerate(f):
-                if len(all_texts) == calib_size:
-                    break
-                if not line.strip():
-                    continue
-                sample = json.loads(line)
-
-                # Extract text field from various possible keys
-                if isinstance(sample, dict) and "text" in sample:
-                    if not sample["text"]:
-                        warnings.warn(f"Sample {i} has empty text, skipping")
-                        continue
-                    full_text = sample["text"]
-                elif isinstance(sample, dict) and "messages" in sample:
-                    conversations = sample["messages"]
-                    assert "role" in conversations[0] and "content" in conversations[0]
-                    full_text = "".join([f"{msg['role']}: {msg['content']}" for msg in conversations])
-                elif isinstance(sample, list) and isinstance(sample[0], dict):
-                    assert "role" in sample[0] and "content" in sample[0]
-                    full_text = "".join([f"{msg['role']}: {msg['content']}" for msg in sample])
-                else:
-                    raise ValueError(f"Sample {i} has unexpected format")
-
-                # Slice text
-                max_text_length = int(max_sequence_length / 0.75)  # tokenized text is roughtly ~75% length of original
-                start_idx = 0
-                if use_random_offset and len(full_text) > max_text_length:
-                    start_idx = random.randint(0, len(full_text) - max_text_length)
-                text = full_text[start_idx : start_idx + max_text_length]
-                all_texts.append(text)
-
-        print_rank_0(f"Loaded calibration dataset ({dataset_path_or_name}) with {len(all_texts)} samples")
-        print_rank_0(f"Actual num samples: {len(all_texts)}, max seq length: {max_sequence_length}")
-        print_rank_0(f"Sampling Strategy: {'Random Index' if use_random_offset else 'From Beginning'}")
-
-        # Tokenize all texts at once and move to device
-        tokens = tokenizer(
-            all_texts, return_tensors="pt", padding="max_length", max_length=max_sequence_length, truncation=True
-        )
-        all_input_ids = tokens.input_ids.cuda()
-        return [{"input_ids": all_input_ids[i:i+batch_size]} for i in range(0, len(all_input_ids), batch_size)]
-    else:
-        # HuggingFace dataset
-        if use_random_offset:
-            warnings.warn("Random offset is not supported for HuggingFace datasets.")
-        print_rank_0(f"Loading calibration dataset from HuggingFace: {dataset_path_or_name}")
-        return get_dataset_dataloader(
-            dataset_name=dataset_path_or_name,
-            tokenizer=tokenizer,
-            num_samples=calib_size,
-            max_sample_length=max_sequence_length,
-            batch_size=batch_size,
-            device="cuda",
-        )
-
-
 if __name__ == "__main__":
     parse_and_validate_args(extra_args_provider=add_text_generate_ptq_args, args_defaults={
             "tokenizer_type": "HuggingFaceTokenizer",
@@ -354,18 +296,66 @@ 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,
         )
-        for sample in tqdm(dataloader, disable=torch.distributed.get_rank()):
-            sample = get_batch_on_this_cp_rank(sample)
-            megatron_prefill(model, sample["input_ids"], skip_return_logits=True)
+    else:
+        # modelopt 0.44 fallback: pad+truncate path via get_dataset_dataloader. Uses the
+        # tokenizer's `padding="max_length", truncation=True` defaults — long documents get
+        # cut at calib_max_sequence_length and short ones pad to it (some activation
+        # statistics noise, less ideal than pack=True but functional).
+        if os.path.isfile(args.calib_dataset_path_or_name):
+            # Local JSONL: load via HF builder, tokenize batched with padding.
+            all_texts = []
+            with open(args.calib_dataset_path_or_name) as f:
+                for line in f:
+                    if len(all_texts) == args.calib_size:
+                        break
+                    if not line.strip():
+                        continue
+                    sample = json.loads(line)
+                    if isinstance(sample, dict) and "text" in sample:
+                        if sample["text"]:
+                            all_texts.append(sample["text"])
+                    elif isinstance(sample, dict) and "messages" in sample:
+                        all_texts.append("".join(
+                            f"{m['role']}: {m['content']}" for m in sample["messages"]
+                        ))
+                    elif isinstance(sample, list):
+                        all_texts.append("".join(
+                            f"{m['role']}: {m['content']}" for m in sample
+                        ))
+            tokens = tokenizer(
+                all_texts,
+                return_tensors="pt",
+                padding="max_length",
+                max_length=args.calib_max_sequence_length,
+                truncation=True,
+            )
+            all_input_ids = tokens.input_ids.cuda()
+            _calib_loader = [
+                {"input_ids": all_input_ids[i : i + args.calib_batch_size]}
+                for i in range(0, len(all_input_ids), args.calib_batch_size)
+            ]
+        else:
+            # HF dataset name: route through modelopt's dataset_utils.
+            _calib_loader = get_dataset_dataloader(
+                dataset_name=args.calib_dataset_path_or_name,
+                tokenizer=tokenizer,
+                num_samples=args.calib_size,
+                max_sample_length=args.calib_max_sequence_length,
+                batch_size=args.calib_batch_size,
+                device="cuda",
+            )
+
+        def _dataset_forward_loop_func(model):
+            for sample in tqdm(_calib_loader, disable=torch.distributed.get_rank()):
+                megatron_prefill(model, sample["input_ids"], skip_return_logits=True)
 
     unwrapped_model = unwrap_model(model)[0]
 

From 3f9279660180a5449a1305373393b275a019dd10 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Thu, 21 May 2026 13:24:53 -0700
Subject: [PATCH 2/4] cleanup

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/post_training/modelopt/quantize.py | 135 ++++++++++++--------
 1 file changed, 84 insertions(+), 51 deletions(-)

diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index 70abb185d7c..fac54325b6b 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -7,6 +7,7 @@
 import inspect
 import json
 import os
+import random
 import sys
 import warnings
 
@@ -25,9 +26,7 @@
 
 # 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. The 0.45 path uses pack=True (no padding/truncation loss); the 0.44 fallback
-# uses padding+truncation, which is slightly worse for long-document corpora but
-# functional.
+# releases.
 try:
     from modelopt.torch.utils.plugins.megatron_calibration import (
         get_megatron_calibration_forward_loop,
@@ -50,9 +49,10 @@
     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
 from megatron.post_training.model_builder import modelopt_gpt_hybrid_builder
@@ -239,6 +239,79 @@ def get_modelopt_torch_quantization_config():
     return mtq_config
 
 
+def get_calib_dataloader(
+    dataset_path_or_name,
+    tokenizer,
+    calib_size=512,
+    max_sequence_length=512,
+    use_random_offset=False,
+    batch_size=1,
+):
+    """Return a dataloader/iterator for calibration using SFT or HF datasets.
+
+    Supports either a local path (.jsonl) or a HuggingFace dataset name.
+    """
+    if os.path.isfile(dataset_path_or_name):
+        # Local file
+        print_rank_0(f"Loading calibration dataset from local file: {dataset_path_or_name}")
+        all_texts = []
+        with open(dataset_path_or_name) as f:
+            for i, line in enumerate(f):
+                if len(all_texts) == calib_size:
+                    break
+                if not line.strip():
+                    continue
+                sample = json.loads(line)
+
+                # Extract text field from various possible keys
+                if isinstance(sample, dict) and "text" in sample:
+                    if not sample["text"]:
+                        warnings.warn(f"Sample {i} has empty text, skipping")
+                        continue
+                    full_text = sample["text"]
+                elif isinstance(sample, dict) and "messages" in sample:
+                    conversations = sample["messages"]
+                    assert "role" in conversations[0] and "content" in conversations[0]
+                    full_text = "".join([f"{msg['role']}: {msg['content']}" for msg in conversations])
+                elif isinstance(sample, list) and isinstance(sample[0], dict):
+                    assert "role" in sample[0] and "content" in sample[0]
+                    full_text = "".join([f"{msg['role']}: {msg['content']}" for msg in sample])
+                else:
+                    raise ValueError(f"Sample {i} has unexpected format")
+
+                # Slice text
+                max_text_length = int(max_sequence_length / 0.75)  # tokenized text is roughtly ~75% length of original
+                start_idx = 0
+                if use_random_offset and len(full_text) > max_text_length:
+                    start_idx = random.randint(0, len(full_text) - max_text_length)
+                text = full_text[start_idx : start_idx + max_text_length]
+                all_texts.append(text)
+
+        print_rank_0(f"Loaded calibration dataset ({dataset_path_or_name}) with {len(all_texts)} samples")
+        print_rank_0(f"Actual num samples: {len(all_texts)}, max seq length: {max_sequence_length}")
+        print_rank_0(f"Sampling Strategy: {'Random Index' if use_random_offset else 'From Beginning'}")
+
+        # Tokenize all texts at once and move to device
+        tokens = tokenizer(
+            all_texts, return_tensors="pt", padding="max_length", max_length=max_sequence_length, truncation=True
+        )
+        all_input_ids = tokens.input_ids.cuda()
+        return [{"input_ids": all_input_ids[i:i+batch_size]} for i in range(0, len(all_input_ids), batch_size)]
+    else:
+        # HuggingFace dataset
+        if use_random_offset:
+            warnings.warn("Random offset is not supported for HuggingFace datasets.")
+        print_rank_0(f"Loading calibration dataset from HuggingFace: {dataset_path_or_name}")
+        return get_dataset_dataloader(
+            dataset_name=dataset_path_or_name,
+            tokenizer=tokenizer,
+            num_samples=calib_size,
+            max_sample_length=max_sequence_length,
+            batch_size=batch_size,
+            device="cuda",
+        )
+
+
 if __name__ == "__main__":
     parse_and_validate_args(extra_args_provider=add_text_generate_ptq_args, args_defaults={
             "tokenizer_type": "HuggingFaceTokenizer",
@@ -306,56 +379,16 @@ def _custom_prompt_forward_loop_func(model):
             batch_size=args.calib_batch_size,
         )
     else:
-        # modelopt 0.44 fallback: pad+truncate path via get_dataset_dataloader. Uses the
-        # tokenizer's `padding="max_length", truncation=True` defaults — long documents get
-        # cut at calib_max_sequence_length and short ones pad to it (some activation
-        # statistics noise, less ideal than pack=True but functional).
-        if os.path.isfile(args.calib_dataset_path_or_name):
-            # Local JSONL: load via HF builder, tokenize batched with padding.
-            all_texts = []
-            with open(args.calib_dataset_path_or_name) as f:
-                for line in f:
-                    if len(all_texts) == args.calib_size:
-                        break
-                    if not line.strip():
-                        continue
-                    sample = json.loads(line)
-                    if isinstance(sample, dict) and "text" in sample:
-                        if sample["text"]:
-                            all_texts.append(sample["text"])
-                    elif isinstance(sample, dict) and "messages" in sample:
-                        all_texts.append("".join(
-                            f"{m['role']}: {m['content']}" for m in sample["messages"]
-                        ))
-                    elif isinstance(sample, list):
-                        all_texts.append("".join(
-                            f"{m['role']}: {m['content']}" for m in sample
-                        ))
-            tokens = tokenizer(
-                all_texts,
-                return_tensors="pt",
-                padding="max_length",
-                max_length=args.calib_max_sequence_length,
-                truncation=True,
-            )
-            all_input_ids = tokens.input_ids.cuda()
-            _calib_loader = [
-                {"input_ids": all_input_ids[i : i + args.calib_batch_size]}
-                for i in range(0, len(all_input_ids), args.calib_batch_size)
-            ]
-        else:
-            # HF dataset name: route through modelopt's dataset_utils.
-            _calib_loader = get_dataset_dataloader(
-                dataset_name=args.calib_dataset_path_or_name,
+        def _dataset_forward_loop_func(model):
+            dataloader = get_calib_dataloader(
+                dataset_path_or_name=args.calib_dataset_path_or_name,
                 tokenizer=tokenizer,
-                num_samples=args.calib_size,
-                max_sample_length=args.calib_max_sequence_length,
+                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,
-                device="cuda",
             )
-
-        def _dataset_forward_loop_func(model):
-            for sample in tqdm(_calib_loader, disable=torch.distributed.get_rank()):
+            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()
                 )

From 53bb6c536e99272432fbde6bf049b9db20efa423 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Fri, 22 May 2026 14:50:18 -0700
Subject: [PATCH 3/4] minor

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/post_training/modelopt/prune.py    | 2 ++
 examples/post_training/modelopt/quantize.py | 2 ++
 2 files changed, 4 insertions(+)

diff --git a/examples/post_training/modelopt/prune.py b/examples/post_training/modelopt/prune.py
index 17a5876553a..a6a2e62ad0c 100644
--- a/examples/post_training/modelopt/prune.py
+++ b/examples/post_training/modelopt/prune.py
@@ -243,6 +243,8 @@ def _custom_prompt_forward_loop_func(model):
             num_samples=args.calib_size,
             seq_length=args.calib_max_sequence_length,
             batch_size=1,
+            # pack=True uses Megatron pretraining-style global-stream document packing
+            pack=True,
         )
     else:
         # modelopt 0.44 fallback: inline pack=True (concatenate raw samples into a single
diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index fac54325b6b..cbaebb1b3f8 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -377,6 +377,8 @@ def _custom_prompt_forward_loop_func(model):
             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
+            pack=True,
         )
     else:
         def _dataset_forward_loop_func(model):

From bb9b6b596998bd5eafde32b0791f57c76f7011f3 Mon Sep 17 00:00:00 2001
From: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
Date: Sat, 23 May 2026 10:32:26 -0700
Subject: [PATCH 4/4] fix finetune.py import

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>
---
 examples/post_training/modelopt/finetune.py | 10 +++-------
 examples/post_training/modelopt/quantize.py |  3 ++-
 2 files changed, 5 insertions(+), 8 deletions(-)

diff --git a/examples/post_training/modelopt/finetune.py b/examples/post_training/modelopt/finetune.py
index 9e19c98a3c0..006a559aa71 100755
--- a/examples/post_training/modelopt/finetune.py
+++ b/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
diff --git a/examples/post_training/modelopt/quantize.py b/examples/post_training/modelopt/quantize.py
index cbaebb1b3f8..c4c24e1bfa0 100644
--- a/examples/post_training/modelopt/quantize.py
+++ b/examples/post_training/modelopt/quantize.py
@@ -378,7 +378,8 @@ def _custom_prompt_forward_loop_func(model):
             seq_length=args.calib_max_sequence_length,
             batch_size=args.calib_batch_size,
             # pack=True uses Megatron pretraining-style global-stream document packing
-            pack=True,
+            # Leave to False for backward compatibility
+            pack=False,
         )
     else:
         def _dataset_forward_loop_func(model):