cp: ModelOpt Knowledge-Distillation support (880) into r0.2.0

.github/workflows/cicd-main.yml

@@ -111,6 +111,7 @@ jobs:
             timeout: 40
           - script: L2_Launch_recipes_llama_1b
           - script: L2_Launch_recipes_llama_3b
+          - script: L2_Launch_recipes_llama_distill
           - script: L2_Launch_recipes_mamba
           - script: L2_Launch_recipes_qwen
           - script: L2_Launch_data

docs/index.md

@@ -49,6 +49,7 @@ training/activation-recomputation.md
 training/cpu-offloading.md
 training/peft.md
 training/packed-sequences.md
+training/distillation.md
 ```
 
 ```{toctree}

docs/training/distillation.md

@@ -0,0 +1,121 @@
+# Knowledge Distillation
+
+Megatron Bridge provides a streamlined setup for Knowledge Distillation (KD) training, making it easy to enable and integrate into your workflow. This section explains how to use this feature effectively.
+
+Knowledge Distillation is a technique where a pre-trained model (the "teacher") transfers its learned knowledge to a second model (the "student"), which is typically smaller and faster. This process helps the student model learn more efficiently by mimicking the behavior of the teacher. KD offers two key advantages over traditional training: faster convergence and higher final accuracy.
+
+In Megatron Bridge, KD is enabled by NVIDIA TensorRT Model Optimizer (ModelOpt) — a library to optimize deep-learning models for inference on GPUs.
+
+## Knowledge Distillation Process
+
+The KD process involves these steps:
+
+1. **Loads Checkpoints**: Loads both the student and teacher model checkpoints.
+2. **Replaces Loss Function**: Replaces the standard loss function with the KL-Divergence between the output logits (and potentially additional losses between pairs of intermediate model states).
+3. **Trains Models**: Runs forward passes on both models, but executes the backward pass only on the student model.
+4. **Saves Checkpoints**: Saves only the student model checkpoint, allowing it to be used later in the same manner as before.
+
+## Limitations
+
+* Only GPT-based checkpoints are currently supported.
+* Student and teacher models must support the same parallelism strategy.
+* If Pipeline Parallelism is enabled, intermediate-state based KD losses are only supported on the final pipeline stage.
+
+## Configuration
+
+### Knowledge Distillation Config
+
+You can configure the KD process via the `ModelOptDistillConfig` class or a YAML file. The configuration includes:
+
+* `logit_layers`: The layer names of student and teacher model logit layers. These names correspond to the PyTorch submodule attributes of the Megatron Core model. (For GPT-based models, this is `"output_layer"`). Default: `["output_layer", "output_layer"]`
+* `intermediate_layer_pairs`: A list of pairs of intermediate layer names. These pairs will by default have a Cosine-Similarity loss between them, and if tensor-parallelism is enabled, these layers must have sequence parallel outputs (i.e. LayerNorms), as Cosine loss cannot have a split hidden dimension. Default: `[["decoder.final_layernorm", "decoder.final_layernorm"]]`
+* `skip_lm_loss`: Whether to skip the default language modeling (LM) loss. If `false`, it will be added to the distillation loss. (Note it consumes more memory). Default: `true`
+* `kd_loss_scale`: Relative scale factor for the distillation loss. The cumulative logits-and-intermediate loss gets scaled to `kd_loss_scale` times the magnitude of the LM loss. Not used if `skip_lm_loss` is `true`. Default: `1.0`
+* `logit_kl_temperature`: Temperature variable for KL Divergence loss calculation. Default: `1.0`
+
+Example YAML configuration:
+
+```yaml
+logit_layers: ["output_layer", "output_layer"]
+intermediate_layer_pairs:
+  - ["decoder.final_layernorm", "decoder.final_layernorm"]
+logit_kl_temperature: 2.0
+```
+
+## Usage
+
+### Basic Usage with Default Configuration
+
+The simplest way to run knowledge distillation is to use or adapt one of the provided recipe scripts. Here's an example for distilling Llama3.2-3B into Llama3.2-1B:
+
+```bash
+torchrun --nproc_per_node=1 examples/recipes/llama/distill_llama32_3b-1b.py
+```
+
+### Using a Custom YAML Config File
+
+You can provide a custom YAML configuration file to override default settings:
+
+```bash
+torchrun --nproc_per_node=1 examples/recipes/llama/distill_llama32_3b-1b.py \
+    --config-file my_custom_config.yaml
+```
+
+### Using CLI Overrides
+
+Megatron Bridge supports Hydra-style CLI overrides for flexible configuration:
+
+```bash
+torchrun --nproc_per_node=2 examples/recipes/llama/distill_llama32_3b-1b.py \
+    model.tensor_model_parallel_size=2 \
+    model.teacher.tensor_model_parallel_size=2
+```
+
+### Combining YAML and CLI Overrides
+
+CLI overrides take precedence over YAML configuration:
+
+```bash
+torchrun --nproc_per_node=2 examples/recipes/llama/distill_llama32_3b-1b.py \
+    --config-file conf/my_config.yaml \
+    train.global_batch_size=512
+```
+
+## Model Support
+
+Currently, distillation is supported for GPT and Mamba-based models
+
+To enable distillation for a model:
+
+1. Use `GPTDistillationProvider` instead of `GPTModelProvider`
+2. Set the `teacher` attribute to the teacher model configuration
+3. Configure `kd_config` with desired distillation settings
+
+## Checkpointing
+
+During distillation training:
+
+* Only the **student model** checkpoints are saved
+* Teacher model remains frozen and is not modified
+* Checkpoints can be used for inference or further training like any standard checkpoint
+
+## Best Practices
+
+1. **Match Parallelism**: Ensure student and teacher use compatible parallelism configurations
+2. **Monitor Loss**: Track both distillation loss and (if enabled) language modeling loss
+3. **Batch Size**: Use larger batch sizes for better stability during distillation
+4. **Learning Rate**: Start with a smaller LR than pretraining
+5. **Data Quality**: Use high-quality, diverse training data for best distillation results
+
+## Troubleshooting
+
+### Out of Memory Errors
+
+* Reduce `train.micro_batch_size`
+* Increase parallelism sizes
+* Set `model.kd_config.skip_lm_loss = True` to save memory
+
+## References
+
+For more information on the underlying implementation, see:
+* [NVIDIA TensorRT Model Optimizer](https://github.com/NVIDIA/TensorRT-Model-Optimizer)

examples/recipes/llama/conf/llama32_3b-1b_distill_override_example.yaml

@@ -0,0 +1,82 @@
+# Copyright (c) 2025, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+# Example override file
+
+# To override a parameter, ensure the structure matches the ConfigContainer
+# and its sub-configurations (e.g., model, train, etc.)
+# Top-level ConfigContainer fields are dataclasses themselves
+
+model:
+  seq_length: 4096
+  tensor_model_parallel_size: 2
+  pipeline_model_parallel_size: 1
+  context_parallel_size: 1
+  sequence_parallel: true
+  teacher:
+    seq_length: 4096
+    tensor_model_parallel_size: 2
+    pipeline_model_parallel_size: 1
+    context_parallel_size: 1
+    sequence_parallel: true
+  kd_config:
+    logit_layers: ["output_layer", "output_layer"]
+    intermediate_layer_pairs: []
+    skip_lm_loss: true
+    kd_loss_scale: 1.0
+    logit_kl_temperature: 1.0
+
+train:
+  train_iters: 10
+  global_batch_size: 8
+  micro_batch_size: 1
+  eval_iters: 8
+
+optimizer:
+  lr: 1e-4
+  min_lr: 1e-5
+
+scheduler:
+  lr_warmup_iters: 3
+
+checkpoint:
+  # Directory to save to. If null, no checkpoint will be saved.
+  save: "./distill_llama32_3b-1b"
+
+dist:
+  use_megatron_fsdp: false
+  use_torch_fsdp2: false
+
+logger:
+  log_interval: 1
+
+dataset:
+  sequence_length: 4096
+
+rng:
+  seed: 42
+
+ddp:
+  grad_reduce_in_fp32: true
+
+profiling:
+  # For optional fields in the config, specify the target to instantiate the object.
+  _target_: megatron.bridge.training.config.ProfilingConfig
+  use_nsys_profiler: false
+  profile_step_start: 5
+  profile_step_end: 10
+  use_pytorch_profiler: true
+  profile_ranks: [0, 1]
+  record_shapes: true
+