kevalmorabia97 · Oct 19, 2026 · Oct 19, 2026 · Oct 19, 2026 · Oct 19, 2026 · Oct 19, 2026
@@ -16,6 +16,12 @@ Changelog
 
 - Add the **external draft** speculative-decoding mode: train a pretrained standalone draft model against a cached base model, instead of grafting a draft head onto the target. Select it with ``mtsp.convert(draft, [("external", cfg)])`` or the ``speculative_external`` recipe; the draft keeps its own embeddings and lm_head and exports as an ordinary HF checkpoint.
 
+*Megatron Framework (M-LM / M-Bridge)*
+
+- Add ``--sft_in_batch_packing``, ``--moe_aux_loss_coeff`` and ``--dsa_kernel_backend`` to ``examples/megatron_bridge/distill.py`` (``--dsa_kernel_backend`` also to ``quantize.py``): pack ``--sft_hf_dataset`` conversations per micro-batch instead of padding each to ``--seq_length``, override the MoE load-balancing loss weight (e.g. 0 for padded SFT data), and select fused DeepSeek sparse-attention kernels for long sequences.
+- Add ``--vllm_fakequant_export`` to ``examples/megatron_bridge/export_quantized_megatron_to_hf.py`` to evaluate formats vLLM has no kernels for (e.g. GGML IQ) as BF16 fake-quant checkpoints, and allow ``--ep_size`` > 1 for checkpoints built with ``--no_moe_grouped_gemm``.
+- GGML GPTQ checkpoints from ``examples/megatron_bridge`` keep their calibrated codes when the Megatron checkpoint is reloaded for QAD or export, and layerwise calibration now runs on Megatron vision-language models.
+
 *Misc*
 
 - The ``evaluation`` agent skill can stop looping ("doom loop") reasoning early: ask it to avoid doom loops and it adds vLLM's ``repetition_detection`` request field (vLLM 0.17.0 or later) to every task, ending a request once its output repeats. Use the same setting for baseline and candidate runs.

@@ -178,6 +178,15 @@ def get_args():
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
     parser.add_argument("--cp_size", type=int, default=1, help="Context parallel size")
     parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size")
+    parser.add_argument(
+        "--dsa_kernel_backend",
+        choices=["none", "tilelang", "cudnn"],
+        default=None,
+        help=(
+            "Fused kernels for DeepSeek sparse attention (e.g. GLM-5.x). The unfused fallback "
+            "builds dense [seq, seq] scores per head, which does not fit at long sequence lengths."
+        ),
+    )
 
     # Dataset arguments
     parser.add_argument(
@@ -238,6 +247,13 @@ def get_args():
         "or the whole conversation. Multi-turn records train all of their responses under "
         "'assistant'.",
     )
+    parser.add_argument(
+        "--sft_in_batch_packing",
+        action="store_true",
+        help="Pack each micro-batch's --sft_hf_dataset conversations into one THD sequence (attention "
+        "restarts per conversation) instead of padding every conversation to --seq_length. --mbs "
+        "then sets the conversations per pack.",
+    )
     # Training & Eval arguments
     parser.add_argument(
         "--output_dir", type=str, required=True, help="Folder for logging and checkpoint saving"
@@ -264,6 +280,13 @@ def get_args():
         action="store_true",
         help="Keep the student's MTP head at its input weights: skip its loss and freeze it.",
     )
+    parser.add_argument(
+        "--moe_aux_loss_coeff",
+        type=float,
+        default=None,
+        help="Override the MoE load-balancing loss weight (e.g. 0 to drop it). It is computed over every "
+        "token, so with --sft padding it also sends gradient into the padded positions.",
+    )
     parser.add_argument(
         "--no_async_save",
         action="store_true",
@@ -421,8 +444,10 @@ def get_args():
             )
         if args.eval_iters > 0 and not args.sft_hf_validation:
             raise ValueError("--sft_hf_dataset with --eval_iters > 0 needs --sft_hf_validation.")
-    elif args.sft_hf_validation or args.sft_loss_mode != "assistant":
-        raise ValueError("--sft_hf_validation / --sft_loss_mode require --sft_hf_dataset.")
+    elif args.sft_hf_validation or args.sft_loss_mode != "assistant" or args.sft_in_batch_packing:
+        raise ValueError(
+            "--sft_hf_validation / --sft_loss_mode / --sft_in_batch_packing require --sft_hf_dataset."
+        )
     if args.sft and (args.data_paths or args.use_mock_data):
         raise ValueError(
             "--sft is mutually exclusive with --data_paths / --use_mock_data: the SFT branch wins "
@@ -508,7 +533,11 @@ def _build_model_provider(hf_path, load_weights=True, moe_grouped_gemm=True):
         provider.expert_model_parallel_size = args.ep_size
         provider.expert_tensor_parallel_size = 1  # Expert tensor parallelism is not supported
         provider.seq_length = args.seq_length
+        if args.dsa_kernel_backend is not None:
+            provider.dsa_kernel_backend = args.dsa_kernel_backend
         set_moe_expert_layout(provider, moe_grouped_gemm)
+        if args.moe_aux_loss_coeff is not None:
+            provider.moe_aux_loss_coeff = args.moe_aux_loss_coeff
         if per_token_loss:
             provider.calculate_per_token_loss = True
         if args.recompute_granularity is not None:
@@ -632,9 +661,14 @@ def _restore_student_hook(model_chunks):
                 _hf_source(args.sft_hf_validation) if args.sft_hf_validation else None
             ),
             preprocessing=ChatSFTPreprocessingConfig(loss_mode=args.sft_loss_mode),
-            dataloader_type="batch",
+            # "batch" collates a rank's whole global batch at once and then splits it into
+            # micro-batches, which a packed THD batch cannot be; packing collates per micro-batch.
+            dataloader_type="single" if args.sft_in_batch_packing else "batch",
             do_validation=args.eval_iters > 0,
             do_test=False,
+            enable_in_batch_packing=args.sft_in_batch_packing,
+            # Each packed conversation must split evenly across the 2 * CP context-parallel chunks.
+            in_batch_packing_pad_to_multiple_of=2 * args.cp_size if args.cp_size > 1 else 1,
         )
     elif args.sft:
         # SFT-masked distillation via Bridge's FinetuningDatasetConfig -> NeMo-style GPTSFTDataset,

@@ -23,7 +23,7 @@
 The HuggingFace unified exporter does not gather tensor-parallel-sharded weights, so this script
 always loads the checkpoint at tensor_model_parallel_size=1 (re-sharding from whatever TP was used
 during quantization). Use --pp_size and/or --ep_size to shard a large model across GPUs for
-export; --ep_size needs grouped-GEMM experts (not --no_moe_grouped_gemm).
+export.
 
 Example usage to export an FP8 checkpoint produced by quantize.py:
 
@@ -53,6 +53,7 @@
 import modelopt.torch.quantization as mtq
 import modelopt.torch.utils.distributed as dist
 from modelopt.torch.export import export_mcore_gpt_to_hf
+from modelopt.torch.export.plugins.vllm_fakequant_megatron import export_mcore_gpt_to_hf_vllm_fq
 from modelopt.torch.utils import print_args, print_rank_0
 from modelopt.torch.utils.mlflow import Tool, masked_args
 from modelopt.torch.utils.plugins.mbridge import (
@@ -106,6 +107,15 @@ def get_args() -> argparse.Namespace:
             "HuggingFace, in which case SequentialMLP is selected automatically."
         ),
     )
+    parser.add_argument(
+        "--vllm_fakequant_export",
+        action="store_true",
+        help=(
+            "Fold the weight quantizers into BF16 weights, for accuracy evaluation in a stock vLLM "
+            "of formats it has no kernels for (e.g. GGML IQ2). Activation quantizers go to a "
+            "quantizer_state.pth sidecar."
+        ),
+    )
     parser.add_argument(
         "--export_extra_modules",
         action="store_true",
@@ -118,7 +128,7 @@ def get_args() -> argparse.Namespace:
         "--ep_size",
         type=int,
         default=1,
-        help="Expert parallel size (grouped-GEMM experts only); shards the experts across ranks",
+        help="Expert parallel size; shards the experts across ranks",
     )
     parser.add_argument(
         "--num_layers_in_first_pipeline_stage",
@@ -199,7 +209,10 @@ def main(args: argparse.Namespace):
         f"Exporting to HuggingFace (unified) checkpoint at {args.export_unified_hf_path}..."
     )
     # TODO: Gemma3-VL is not in export_mcore_gpt_to_hf's per-arch mappings yet.
-    export_mcore_gpt_to_hf(
+    export_fn = (
+        export_mcore_gpt_to_hf_vllm_fq if args.vllm_fakequant_export else export_mcore_gpt_to_hf
+    )
+    export_fn(
         unwrapped_model,
         args.hf_model_name_or_path,
         export_extra_modules=export_extra_modules,

@@ -152,6 +152,15 @@ def get_args() -> argparse.Namespace:
     parser.add_argument("--pp_size", type=int, default=1, help="Pipeline parallel size")
     parser.add_argument("--ep_size", type=int, default=1, help="Expert parallel size")
     parser.add_argument("--cp_size", type=int, default=1, help="Context parallel size")
+    parser.add_argument(
+        "--dsa_kernel_backend",
+        choices=["none", "tilelang", "cudnn"],
+        default=None,
+        help=(
+            "Fused kernels for DeepSeek sparse attention (e.g. GLM-5.x). The unfused fallback "
+            "builds dense [seq, seq] scores per head, which does not fit at long sequence lengths."
+        ),
+    )
 
     # Quantization arguments
     parser.add_argument(
@@ -345,6 +354,11 @@ def main(args: argparse.Namespace):
             "pipeline_dtype": torch.bfloat16,
             "seq_length": args.seq_length,
             "gradient_accumulation_fusion": False,  # not supported
+            **(
+                {"dsa_kernel_backend": args.dsa_kernel_backend}
+                if args.dsa_kernel_backend is not None
+                else {}
+            ),
         },
         init_model_parallel=True,
     )

@@ -394,10 +394,15 @@ def save_safetensors_by_layer_index(
         }
         for layer_index in range(total_layers):
             meta_filename = name_template.format(layer_index + 1, total_layers) + ".json"
-            with open(save_directory + "/" + meta_filename) as f:
-                shard = json.load(f)
-            safetensor_index["metadata"]["total_size"] += shard["metadata"]["total_size"]
-            safetensor_index["weight_map"].update(shard["weight_map"])
+            # Plus the expert shards that EP ranks other than 0 write beside the layer's own.
+            ep_shards = sorted(
+                Path(save_directory).glob(meta_filename[: -len(".json")] + "-ep*.json")
+            )
+            for shard_file in [Path(save_directory) / meta_filename, *ep_shards]:
+                with open(shard_file) as f:
+                    shard = json.load(f)
+                safetensor_index["metadata"]["total_size"] += shard["metadata"]["total_size"]
+                safetensor_index["weight_map"].update(shard["weight_map"])
 
         with open(save_directory + "/model.safetensors.index.json", "w") as f:
             json.dump(safetensor_index, f, indent=4)

@@ -322,7 +322,15 @@ def save_pretrained(
             index: self._get_quantizer_state(state_dict)
             for index, state_dict in enumerate([*self.layer_state_dicts.values(), self._state_dict])
         }
+
+        # The weights are written in BF16, so a source checkpoint's own quantization (e.g. FP8)
+        # no longer describes them.
+        if hasattr(self._hf_config, "quantization_config"):
+            del self._hf_config.quantization_config
+
         super().save_pretrained(save_directory, pretrained_model_name_or_path)
+        # The base exporter copies a pre-quantized source's hf_quant_config.json along with it.
+        (Path(save_directory) / "hf_quant_config.json").unlink(missing_ok=True)
 
         # Save fresh quantizer files after the base exporter copies source files.
         recipe = self._extract_quantizer_recipe_markers(quantizer_state_dicts)

@@ -93,9 +93,11 @@
     from megatron.core.models.multimodal.llava_model import LLaVAModel
     from megatron.core.parallel_state import (
         get_data_parallel_rank,
+        get_expert_data_parallel_rank,
         get_expert_model_parallel_group,
         get_expert_model_parallel_rank,
         get_expert_model_parallel_world_size,
+        get_expert_tensor_parallel_rank,
         get_pipeline_model_parallel_rank,
         get_pipeline_model_parallel_world_size,
         get_tensor_model_parallel_rank,
@@ -161,6 +163,8 @@ def __init__(
 
         self._state_dict = OrderedDict()
         self._layer_state_dicts = OrderedDict()
+        # Keys of the SequentialMLP experts held by this rank (all ranks' differ at EP>1).
+        self._local_expert_keys: set[str] = set()
         self._hf_pretrained_model_name = pretrained_model_name_or_path
         self._hf_config = transformers.AutoConfig.from_pretrained(
             pretrained_model_name_or_path, trust_remote_code=trust_remote_code
@@ -346,24 +350,24 @@ def save_pretrained(
                 raise NotImplementedError(
                     "Megatron GGML unified export currently requires pipeline model parallel size 1"
                 )
-        # SequentialMLP rules index experts by local position, so EP>1 would collide. Agreed across
-        # ranks so that stages without such experts don't block in the collectives below.
+        # Packed local-expert rules hand over all of a rank's experts as one, so EP>1 would collide.
+        # Agreed across ranks so that stages without such experts don't block in the collectives below.
         if (
             torch.distributed.is_initialized()
             and get_expert_model_parallel_world_size() > 1
+            and self.rules.get("use_packed_local_experts", False)
             and self._any_rank(any(hasattr(m, "local_experts") for m in self.model.modules()))
         ):
             raise NotImplementedError(
-                "Export at expert parallel size > 1 needs grouped-GEMM experts; "
-                "export SequentialMLP (--no_moe_grouped_gemm) checkpoints at EP=1."
+                "Export at expert parallel size > 1 does not support packed local experts; "
+                "export at EP=1."
             )
         # One layer-shard writer per pipeline stage: other TP / DP / EP ranks hold the same layers
         # (EP>1 ranks hold no gathered experts at all), and writing them too would race on the same
         # files. Unlike is_writer_rank (one global writer of save_directory metadata), this has no
         # PP term on purpose: each stage owns a disjoint slice of the layer shards.
-        is_stage_layer_writer = (
-            tp_rank == 0 and get_data_parallel_rank() == 0 and get_expert_model_parallel_rank() == 0
-        )
+        ep_rank = get_expert_model_parallel_rank()
+        is_stage_layer_writer = tp_rank == 0 and get_data_parallel_rank() == 0 and ep_rank == 0
 
         # Main export process
         layer_state_dicts = self.layer_state_dicts
@@ -486,6 +490,10 @@ def save_pretrained(
         # Bracket the writer's config.json read-modify-write with barriers so peers
         # never observe a truncated file (also ensures export_dir exists).
         torch.distributed.barrier()
+        if is_writer_rank:
+            # Expert shards left by an earlier export at a larger EP would enter the index.
+            for stale in Path(save_directory).glob("model-*-of-*-ep*.*"):
+                stale.unlink()
         config_json_file = save_directory + "/config.json"
         if is_writer_rank and self._hf_quant_config and os.path.exists(config_json_file):
             with open(config_json_file) as f:
@@ -497,11 +505,30 @@ def save_pretrained(
                 json.dump(config_dict, f, indent=4)
         torch.distributed.barrier()
 
+        # SequentialMLP experts are not gathered: each other EP rank writes its own to extra shards.
+        # Expert ownership follows the expert grid (EP x ETP x expert-DP), not the dense TP rank.
+        writes_own_experts = (
+            get_expert_tensor_parallel_rank() == 0
+            and get_expert_data_parallel_rank() == 0
+            and ep_rank > 0
+            and bool(self._local_expert_keys)
+        )
+        name_template = "model-{:05d}-of-{:05d}"
+        if writes_own_experts:
+            # Layers without local experts (e.g. dense ones) would write empty shards.
+            layer_state_dicts = {
+                i: experts
+                for i, sd in layer_state_dicts.items()
+                if (experts := {k: v for k, v in sd.items() if k in self._local_expert_keys})
+            }
+            name_template += f"-ep{ep_rank:03d}"
         save_safetensors_by_layer_index(
-            layer_state_dicts=layer_state_dicts if is_stage_layer_writer else {},
+            layer_state_dicts=layer_state_dicts
+            if is_stage_layer_writer or writes_own_experts
+            else {},
             total_layers=self.model.config.num_layers,
             save_directory=save_directory,
-            name_template="model-{:05d}-of-{:05d}",
+            name_template=name_template,
         )
 
         # Every rank has written its shards; one rank now checks nothing was dropped. The result is
@@ -785,13 +812,18 @@ def _get_transformer_layer_state_dict(self, layer, layer_id, is_mtp=False):
                         )
                 if hasattr(layer.mlp.experts, "local_experts"):
                     if not self.rules.get("use_packed_local_experts", False):
-                        for expert_id, expert in enumerate(layer.mlp.experts.local_experts):
+                        keys_before = set(self._state_dict)
+                        # Global ids: at EP>1 each rank holds a different slice of the experts.
+                        for expert_id, expert in zip(
+                            layer.mlp.local_expert_indices, layer.mlp.experts.local_experts
+                        ):
                             self.rules["local_experts.linear_fc1"](
                                 expert.linear_fc1, layer_id, expert_id, is_mtp=is_mtp
                             )
                             self.rules["local_experts.linear_fc2"](
                                 expert.linear_fc2, layer_id, expert_id, is_mtp=is_mtp
                             )
+                        self._local_expert_keys |= set(self._state_dict) - keys_before
                     else:
                         # For llama 4, in hf unified checkpoint, all local experts share one scale
                         self.rules["local_experts.linear_fc1"](

@@ -154,8 +154,12 @@ def pinned_packed_weight(quantizer, format_name: str) -> torch.Tensor | None:
     return getattr(quantizer, _pin_name(format_name), None)
 
 
+# Prefix of the buffer holding a pinned payload; checkpoint code keeps buffers with it.
+PINNED_PAYLOAD_PREFIX = "_ggml_pinned_"
+
+
 def _pin_name(format_name: str) -> str:
-    return f"_ggml_pinned_{format_name}"
+    return f"{PINNED_PAYLOAD_PREFIX}{format_name}"
 
 
 def _matching_rows(table: torch.Tensor, rows: torch.Tensor) -> torch.Tensor | None:

@@ -24,6 +24,7 @@
 
 from modelopt.torch.utils.distributed import ParallelState
 
+from ..ggml.common import PINNED_PAYLOAD_PREFIX
 from ..nn import (
     GroupedQuantizer,
     NVFP4StaticQuantizer,
@@ -132,7 +133,10 @@ def modelopt_post_restore(self, prefix: str = ""):
 
         def _check_unsupported_states(quantizer: TensorQuantizer):
             for k in quantizer.state_dict():
-                if k not in ["_amax", "_pre_quant_scale", "_global_amax"]:
+                # GGML payloads pinned by GPTQ are loaded with the weights from the sharded checkpoint.
+                if k not in ["_amax", "_pre_quant_scale", "_global_amax"] and not k.startswith(
+                    PINNED_PAYLOAD_PREFIX
+                ):
                     warnings.warn(
                         f"Restore of {k} for {prefix} is not supported. The restore of this layer might be "
                         f"incorrect. Please implement a custom restore for {k}."