Support Megatron ckpt inputs for distill.py + doc updates

Signed-off-by: Keval Morabia <28916987+kevalmorabia97@users.noreply.github.com>

Author: kevalmorabia97

examples/megatron_bridge/README.md

@@ -86,7 +86,88 @@ torchrun --nproc_per_node 1 prune_minitron.py --help
 
 ## Distillation
 
-TODO - Add info!
+This section shows how to distill a student model from a teacher model in the Megatron-Bridge framework.
+The student and teacher can each be loaded from a HuggingFace model (hub name or local dir) or a Megatron checkpoint (`iter_*` directory) - auto-detected based on the path provided.
+
+This can be used stand-alone or after pruning (see [Pruning](#pruning)) / quantization (see [Quantization](#quantization)) to recover accuracy of the model by distilling from the original model (teacher).
+
+The [distill.py](distill.py) script loads student and teacher models from HuggingFace or Megatron checkpoints and saves the distilled model to `<output_dir>/checkpoints` in Megatron distributed checkpoint format.
+
+Supported input model path formats for student and teacher (auto-detected by the script):
+    - HuggingFace: model hub name (e.g. `Qwen/Qwen3-8B`) or local HF dir
+    - Megatron: `iter_*` directory inside a Megatron checkpoint (e.g. `/path/to/ckpt/iter_0000000`)
+
+### Data Preparation
+
+The distillation script expects pre-tokenized data in Megatron's binary format (`.bin` / `.idx` files).
+You can tokenize your JSONL dataset using the following function. If you have multiple JSONL files, you can tokenize them one by one and pass all the paths to the `--data_paths` argument.
+
+```python
+from modelopt.torch.utils.plugins import megatron_preprocess_data
+
+megatron_preprocess_data(
+    input_path="/path/to/your/data.jsonl",
+    output_dir="/path/to/tokenized/data",
+    tokenizer_name_or_path="Qwen/Qwen3-0.6B",
+    json_keys=["text"],  # change to your JSON key if needed
+    workers=32,
+    log_interval=100000,
+    max_sequence_length=256000,  # To avoid rare OOM errors if text is too long
+)
+```
+
+### Distillation with Real Data
+
+Example usage to distill a 4B student (HF) from an 8B teacher (HF) on 8 GPUs:
+
+```bash
+torchrun --nproc_per_node 8 distill.py \
+    --teacher_path Qwen/Qwen3-8B \
+    --student_path Qwen/Qwen3-4B \
+    --tp_size 8 \
+    --data_paths 1.0 /path/to/tokenized/data \
+    --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \
+    --seq_length 8192 \
+    --mbs 1 \
+    --gbs 768 \
+    --train_iters 15000 \
+    --lr 1e-4 \
+    --min_lr 1e-5 \
+    --lr_warmup_iters 50 \
+    --eval_interval 100 \
+    --eval_iters 32 \
+    --log_interval 10 \
+    --output_dir /output/qwen3_8b_to_4b_distill
+```
+
+Tensorboard logging is enabled by default and logs are saved to `<output_dir>/tensorboard` directory.
+To use Weights & Biases for logging, set the `WANDB_API_KEY` environment variable and pass the `--wandb_project` argument.
+Optionally, you can also pass `--wandb_entity` and `--wandb_exp_name` arguments to group runs under a project and experiment name.
+
+To see all available arguments:
+
+```bash
+torchrun --nproc_per_node 1 distill.py --help
+```
+
+### Quick Test with Mock Data
+
+Example usage with mock data for quick testing (no pre-tokenized data needed):
+
+```bash
+torchrun --nproc_per_node 8 distill.py \
+    --teacher_path Qwen/Qwen3-0.6B \
+    --student_path Qwen/Qwen3-0.6B \
+    --tp_size 8 \
+    --use_mock_data \
+    --seq_length 512 \
+    --mbs 1 \
+    --gbs 8 \
+    --train_iters 100 \
+    --eval_interval 10 \
+    --eval_iters 4 \
+    --output_dir /tmp/test_distill
+```
 
 ## Quantization
 

examples/megatron_bridge/distill.py

@@ -14,66 +14,18 @@
 # limitations under the License.
 """Distillation script for Megatron-Bridge.
 
-Loads student and teacher models directly from HuggingFace checkpoints (local or remote) and saves the distilled model
-to <output_dir>/checkpoints in megatron distributed checkpoint format.
-
-Example usage to distill a 4B student from an 8B teacher on 8 GPUs:
-
-.. code-block:: bash
-
-    torchrun --nproc_per_node 8 distill.py \
-        --teacher_hf_path Qwen/Qwen3-8B \
-        --student_hf_path Qwen/Qwen3-4B \
-        --tp_size 8 \
-        --data_paths 1.0 /path/to/tokenized/data \
-        --data_path_to_cache /path/to/cache/dataset_indices_qwen3 \
-        --seq_length 8192 \
-        --mbs 1 \
-        --gbs 768 \
-        --train_iters 15000 \
-        --lr 1e-4 \
-        --min_lr 1e-5 \
-        --lr_warmup_iters 50 \
-        --eval_interval 100 \
-        --eval_iters 32 \
-        --log_interval 10 \
-        --output_dir /output/qwen3_8b_to_4b_distill
-
-Example usage to use mock data for quick testing:
-
-.. code-block:: bash
-
-    torchrun --nproc_per_node 8 distill.py \
-        --teacher_hf_path Qwen/Qwen3-0.6B \
-        --student_hf_path Qwen/Qwen3-0.6B \
-        --tp_size 8 \
-        --use_mock_data \
-        --seq_length 512 \
-        --mbs 1 \
-        --gbs 8 \
-        --train_iters 100 \
-        --eval_interval 10 \
-        --eval_iters 4 \
-        --output_dir /tmp/test_distill
-
-If you want to tokenize your own data for a specific tokenizer, you can use the following command:
-
-.. code-block:: python
-
-    from modelopt.torch.utils.plugins import megatron_preprocess_data
-
-    megatron_preprocess_data(
-        input_path="/path/to/your/data.jsonl",
-        output_dir="/path/to/tokenized/data",
-        tokenizer_name_or_path="Qwen/Qwen3-0.6B",
-        json_keys=["text"],
-        workers=32,
-        log_interval=100000,
-        max_sequence_length=256000,
-    )
+Loads student and teacher models from HuggingFace or Megatron checkpoints and saves the distilled model
+to `<output_dir>/checkpoints` in Megatron distributed checkpoint format.
+
+Supported input model path formats for student and teacher (auto-detected):
+    - HuggingFace: model hub name (e.g. `Qwen/Qwen3-8B`) or local HF dir
+    - Megatron: `iter_*` directory inside a Megatron checkpoint (e.g. `/path/to/ckpt/iter_0000000`)
+
+See `README.md` in this directory for example usage and data preparation instructions.
 """
 
 import argparse
+import contextlib
 import os
 
 import torch
@@ -82,6 +34,7 @@
 from megatron.bridge.recipes.utils.optimizer_utils import (
     distributed_fused_adam_with_cosine_annealing,
 )
+from megatron.bridge.training.checkpointing import _load_model_weights_from_checkpoint
 from megatron.bridge.training.config import (
     CheckpointConfig,
     ConfigContainer,
@@ -93,6 +46,7 @@
     TrainingConfig,
 )
 from megatron.bridge.training.distill import distill
+from megatron.bridge.training.model_load_save import load_model_config
 from megatron.bridge.training.post_training.distillation import ModelOptDistillConfig
 from megatron.core.datasets.utils import get_blend_from_list
 from megatron.core.distributed import DistributedDataParallelConfig
@@ -106,18 +60,18 @@
 def get_args():
     """Parse command-line arguments."""
     parser = argparse.ArgumentParser(description="Distillation for Megatron-Bridge.")
-    # Model arguments
+    # Model arguments (accepts HuggingFace paths or Megatron checkpoint directories)
     parser.add_argument(
-        "--student_hf_path",
+        "--student_path",
         type=str,
         required=True,
-        help="HuggingFace model name or path for the student (e.g. Qwen/Qwen3-0.6B)",
+        help="Student model: HuggingFace name/path (e.g. Qwen/Qwen3-0.6B) or Megatron dir (e.g. /path/iter_0000000)",
     )
     parser.add_argument(
-        "--teacher_hf_path",
+        "--teacher_path",
         type=str,
         required=True,
-        help="HuggingFace model name or path for the teacher (e.g. Qwen/Qwen3-8B)",
+        help="Teacher model: HuggingFace name/path (e.g. Qwen/Qwen3-8B) or Megatron dir (e.g. /path/iter_0000000)",
     )
     # Parallelism arguments
     parser.add_argument("--tp_size", type=int, default=1, help="Tensor parallel size")
@@ -179,26 +133,93 @@ def get_args():
     return args
 
 
+def _build_provider_from_hf(hf_path):
+    """Build a model provider from a HuggingFace model path."""
+    bridge = AutoBridge.from_hf_pretrained(hf_path)
+    provider = bridge.to_megatron_provider(load_weights=True)
+    return provider
+
+
+def _build_provider_from_megatron(iter_path):
+    """Build a model provider from a Megatron checkpoint iter directory.
+
+    Uses load_model_config() to reconstruct the provider (architecture) from the
+    checkpoint's saved run config, then registers a pre_wrap_hook that loads the
+    checkpoint weights after model creation.
+    """
+    print_rank_0(f"Loading model config from Megatron checkpoint: {iter_path}")
+    provider, _ = load_model_config(iter_path)
+    provider.perform_initialization = False  # Will load weights from ckpt
+
+    def _load_megatron_weights_hook(model):
+        """Pre-wrap hook: load weights from a Megatron checkpoint into the model.
+
+        For distillation kd_models, hide teacher / loss modules so only the
+        student weights are loaded (same pattern as HF weight loading).
+        """
+        cms = []
+        for m in model:
+            if hasattr(m, "hide_teacher_model"):
+                cms.append(m.hide_teacher_model())
+            if hasattr(m, "hide_loss_modules"):
+                cms.append(m.hide_loss_modules())
+        with contextlib.ExitStack() as stack:
+            for cm in cms:
+                stack.enter_context(cm)
+            _load_model_weights_from_checkpoint(iter_path, model)
+        return model
+
+    provider.register_pre_wrap_hook(_load_megatron_weights_hook)
+    return provider
+
+
+def _build_model_provider(path, *, tp_size, pp_size, seq_length):
+    """Build a model provider, auto-detecting HuggingFace vs Megatron checkpoint.
+
+    Megatron checkpoints are detected by the `iter_` prefix in the directory name
+    (e.g. `/path/to/checkpoints/iter_0000000`).
+    """
+    if os.path.basename(os.path.normpath(path)).startswith("iter_"):
+        print_rank_0(f"Detected Megatron checkpoint: {path}")
+        provider = _build_provider_from_megatron(path)
+    else:
+        print_rank_0(f"Loading from HuggingFace: {path}")
+        provider = _build_provider_from_hf(path)
+
+    # Override parallelism / training settings
+    provider.tensor_model_parallel_size = tp_size
+    provider.pipeline_model_parallel_size = pp_size
+    provider.context_parallel_size = 1
+    provider.sequence_parallel = tp_size > 1
+    provider.seq_length = seq_length
+    provider.pipeline_dtype = torch.bfloat16
+    provider.cross_entropy_fusion_impl = "te"
+
+    # Run deferred __post_init__ to compute derived fields (e.g. init_method from init_method_std).
+    # Must be called after all overrides are set.
+    # Safe to call on HF-sourced providers too (idempotent).
+    provider.finalize()
+
+    return provider
+
+
 def main(args: argparse.Namespace):
     checkpoint_dir = os.path.join(args.output_dir, "checkpoints")
     tensorboard_dir = os.path.join(args.output_dir, "tb_logs")
 
-    # Build student and teacher model providers
-    def _build_model_provider(hf_path):
-        bridge = AutoBridge.from_hf_pretrained(hf_path)
-        provider = bridge.to_megatron_provider(load_weights=True)
-        provider.tensor_model_parallel_size = args.tp_size
-        provider.pipeline_model_parallel_size = args.pp_size
-        provider.context_parallel_size = 1
-        provider.sequence_parallel = args.tp_size > 1
-        provider.seq_length = args.seq_length
-        provider.pipeline_dtype = torch.bfloat16
-        provider.cross_entropy_fusion_impl = "te"
-        return provider
-
-    # TODO: Support megatron-ckpt as an alternative to HF checkpoints
-    student_provider = _build_model_provider(args.student_hf_path)
-    teacher_provider = _build_model_provider(args.teacher_hf_path)
+    # Build student and teacher model providers (auto-detects HF vs Megatron ckpt)
+    student_provider = _build_model_provider(
+        args.student_path,
+        tp_size=args.tp_size,
+        pp_size=args.pp_size,
+        seq_length=args.seq_length,
+    )
+    teacher_provider = _build_model_provider(
+        args.teacher_path,
+        tp_size=args.tp_size,
+        pp_size=args.pp_size,
+        seq_length=args.seq_length,
+    )
 
     # Wrap into DistillationProvider
     kd_config = ModelOptDistillConfig()

examples/megatron_bridge/prune_minitron.py

@@ -28,6 +28,8 @@
 
 To see the full usage for advanced configurations, run:
     torchrun --nproc_per_node 1 prune_minitron.py --help
+
+See `README.md` in this directory for more details.
 """
 
 import argparse