diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 89fb2a8b8d7..9ff43dea159 100755 --- a/CHANGELOG.rst +++ b/CHANGELOG.rst @@ -162,7 +162,7 @@ Changelog *Megatron Framework (M-LM / M-Bridge)* - Add chat-data distillation to ``examples/megatron_bridge/distill.py``: ``--sft --sft_hf_dataset`` takes a local chat ``.jsonl`` or a Hugging Face Hub dataset, renders each conversation with the model's chat template, and ``--sft_loss_mode`` selects which turns are trained. Requires a Megatron-Bridge with ``DirectHFSFTDatasetConfig``. -- Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` in ``examples/megatron_bridge/distill.py``): the student's MTP head keeps its input weights and its loss is skipped, so distillation trains the rest of the model on the distillation loss alone. +- Add ``freeze_mtp`` to the Megatron ``DistillationConfig`` (``--freeze_mtp`` / ``--no-freeze_mtp`` in ``examples/megatron_bridge/distill.py``): ``True`` freezes the student's MTP-specific parameters and disables its MTP loss, ``False`` keeps them training. Unset, they are frozen only when ``kd_loss_alpha`` is 1.0 and the student is quantized except for its MTP head, as before. - Add ``output_format: chat_jsonl`` to ``prepare_megatron_data_blend``: it writes a weighted blend as ``train.jsonl`` / ``validation.jsonl`` chat data for ``distill.py --sft --sft_hf_dataset``, so the distillation loss can be restricted to assistant tokens. See the `Megatron data preparation guide `_. - Add mergeable LoRA QAD for fake-quantized dense Megatron/Transformer Engine linear layers. Enable it with ``--lora_rank`` in Megatron-Bridge PTQ, train through the normal QAD flow, and export a standalone quantized checkpoint without deployment adapters. - Add optional Top-P (nucleus) truncation to the Megatron top-k logits KD loss (now ``TopLogitsKLLoss``) via ``logit_kl_top_p`` and ``logit_kl_top_p_min_k`` in ``DistillationConfig``: after the global Top-K selection, only the smallest prefix whose cumulative teacher probability reaches ``top_p`` (with a floor of ``min_k`` entries) contributes to the KL, mirroring ``--logits-save-top-p`` / ``--logits-save-top-p-min-k`` in Megatron-LM's logits saver. diff --git a/examples/megatron_bridge/distill.py b/examples/megatron_bridge/distill.py index 8b233e05916..ad1ae3327f4 100644 --- a/examples/megatron_bridge/distill.py +++ b/examples/megatron_bridge/distill.py @@ -302,8 +302,11 @@ def get_args(): ) parser.add_argument( "--freeze_mtp", - action="store_true", - help="Keep the student's MTP head at its input weights: skip its loss and freeze it.", + action=argparse.BooleanOptionalAction, + default=None, + help="Freeze the student's MTP-specific parameters and disable its MTP loss " + "(--no-freeze_mtp keeps them training). Unset, they are frozen only when " + "--kd_loss_alpha is 1.0 and the student is quantized except for its MTP head.", ) parser.add_argument( "--moe_aux_loss_coeff", diff --git a/modelopt/torch/distill/plugins/megatron.py b/modelopt/torch/distill/plugins/megatron.py index 0126e3f6065..c4a5df11bf5 100644 --- a/modelopt/torch/distill/plugins/megatron.py +++ b/modelopt/torch/distill/plugins/megatron.py @@ -71,8 +71,11 @@ class DistillationConfig: Only the smallest prefix of the (sorted) Top-K whose cumulative teacher probability reaches this value contributes to the loss. Requires ``logit_kl_topk``. Must be in (0, 1]. logit_kl_top_p_min_k: Minimum number of Top-K entries kept per token when top-P is active. - freeze_mtp: Keep the student's MTP head fixed: skip its loss, which otherwise also trains the - shared layers it reads from, and freeze its parameters. Default: ``False``. + freeze_mtp: Whether to freeze the student's MTP-specific parameters and disable its auxiliary + MTP loss, which otherwise also trains the shared layers it reads from. ``True`` freezes + it and ``False`` keeps it training. ``None`` (default) freezes it only when + ``kd_loss_alpha == 1.0`` and the student is quantized except for its MTP head, since + there is then no quantization error for the MTP loss to recover. """ intermediate_layer_pairs: list[tuple[str, ...]] = field(default_factory=list) @@ -82,7 +85,7 @@ class DistillationConfig: logit_kl_topk: int | None = None logit_kl_top_p: float | None = None logit_kl_top_p_min_k: int = 1 - freeze_mtp: bool = False + freeze_mtp: bool | None = None criterion: Criterion | None = None loss_balancer: mtd.DistillationLossBalancer | None = None @@ -743,20 +746,16 @@ def _init_weights(self, module): ######################################################## -def _mtp_excluded_from_quantization(model: nn.Module) -> bool: - """Whether the model is quantized but its MTP submodules are not. - - Only then is the MTP loss pure overhead, since distillation has no quantization error to - recover there. Plain distillation (e.g. pruning recovery) still trains the MTP head. - """ - has_mtp = model_quantized = mtp_quantized = False +def _has_quantized_backbone_with_unquantized_mtp(model: nn.Module) -> bool: + """Whether ``model`` has MTP submodules and enabled quantizers, none of them inside MTP.""" + has_mtp = backbone_quantized = mtp_quantized = False for name, module in model.named_modules(): is_mtp = "mtp" in name.split(".") has_mtp |= is_mtp if isinstance(module, TensorQuantizer) and module.is_enabled: - model_quantized = True mtp_quantized |= is_mtp - return has_mtp and model_quantized and not mtp_quantized + backbone_quantized |= not is_mtp + return has_mtp and backbone_quantized and not mtp_quantized def adjust_distillation_model_for_mcore( @@ -775,13 +774,15 @@ def _sharded_state_dict(self, *args, **kwargs) -> "ShardedStateDict": # Uses a per-forward call counter so that MTP head calls (which always precede the # main LM head call in _postprocess) still receive real CE loss even when # skip_lm_loss=True — only the final main-head call is zeroed. - # An MTP head left out of quantization is exempt from that: there is no quantization - # error to recover there, and its CE materialises an fp32 [seq, vocab] tensor. - # ``freeze_mtp`` skips the MTP loss unconditionally. + # ``freeze_mtp`` decides whether the MTP head trains. Unset, it is frozen only when the LM loss + # is skipped and the student is quantized except for its MTP head: there is then no + # quantization error to recover there, and its CE materialises an fp32 [seq, vocab] tensor. skip_lm_loss = distill_cfg.skip_lm_loss - skip_mtp_loss = distill_cfg.freeze_mtp or ( - skip_lm_loss and _mtp_excluded_from_quantization(model) - ) + if distill_cfg.freeze_mtp is None: + with model.hide_teacher_model(): + skip_mtp_loss = skip_lm_loss and _has_quantized_backbone_with_unquantized_mtp(model) + else: + skip_mtp_loss = distill_cfg.freeze_mtp if skip_mtp_loss: # Freeze the untrained MTP head: DDP's overlapped grad reduce asserts on params with no grad. num_frozen = 0 diff --git a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py index 57a841adc48..ff2a1d0f3fc 100644 --- a/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py +++ b/tests/gpu_megatron/torch/distill/plugins/test_distill_megatron.py @@ -36,7 +36,7 @@ LogitsKLLoss, TopKLogitsKLLoss, TopLogitsKLLoss, - _mtp_excluded_from_quantization, + _has_quantized_backbone_with_unquantized_mtp, adjust_distillation_model_for_mcore, setup_distillation_config, ) @@ -240,7 +240,8 @@ def _test_topk_logits_kl_loss(kd_kwargs, rank, size): def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size): """Test that skip_lm_loss only zeroes the main LM head, unless MTP is left out of quantization. - ``freeze_mtp`` zeroes the MTP loss and freezes the MTP head regardless. + An explicit ``freeze_mtp`` overrides that: ``True`` zeroes the MTP loss and freezes the MTP head, + ``False`` keeps both. """ set_seed(SEED) @@ -284,7 +285,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size) mtp_num_layers=mtp_num_layers, ).cuda() if quantized: - # A quantized body with the MTP head left out of quantization. + # Quantization-state fixture: a quantized body with the MTP head left out of quantization. student_model.decoder.weight_quantizer = TensorQuantizer() distill_cfg = setup_distillation_config( @@ -304,7 +305,7 @@ def _test_skip_lm_loss_with_mtp(quantized, skip_lm_loss, freeze_mtp, rank, size) adjust_distillation_model_for_mcore(distillation_model, distill_cfg) # An untrained MTP head is frozen so DDP doesn't wait for its grads. - mtp_skipped = freeze_mtp or (quantized and skip_lm_loss) + mtp_skipped = (quantized and skip_lm_loss) if freeze_mtp is None else freeze_mtp with distillation_model.hide_teacher_model(): mtp_params = [p for n, p in distillation_model.named_parameters() if "mtp" in n.split(".")] assert mtp_params @@ -346,8 +347,8 @@ def _recording_loss(labels, logits): assert recorded_losses[-1].any() != skip_lm_loss, "Main LM head loss is zero iff skip_lm_loss" -def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): - """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" +def _test_freeze_mtp_training_step(mtp_quantized, freeze_mtp, kd_loss_alpha, rank, size): + """Through a training step, the MTP head is frozen exactly when ``freeze_mtp`` resolves to it.""" set_seed(SEED) vocab_size = 32 @@ -368,10 +369,12 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): } teacher_model = get_mcore_gpt_model(initialize_megatron=True, **model_kwargs).cuda() student_model = get_mcore_gpt_model(initialize_megatron=False, **model_kwargs).cuda() - # A recipe that quantizes the MTP head as well as the body. + # Quantization-state fixtures: a quantized body, with or without a quantized MTP head. student_model.decoder.weight_quantizer = TensorQuantizer() - student_model.mtp.weight_quantizer = TensorQuantizer() - assert not _mtp_excluded_from_quantization(student_model) + if mtp_quantized: + student_model.mtp.weight_quantizer = TensorQuantizer() + assert _has_quantized_backbone_with_unquantized_mtp(student_model) != mtp_quantized + frozen = (kd_loss_alpha == 1.0 and not mtp_quantized) if freeze_mtp is None else freeze_mtp distill_cfg = setup_distillation_config( config_or_path=DistillationConfig(kd_loss_alpha=kd_loss_alpha, freeze_mtp=freeze_mtp), @@ -389,7 +392,7 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): adjust_distillation_model_for_mcore(distillation_model, distill_cfg) if rank == 0: frozen_msg = [w for w in caught if "MTP head is frozen" in str(w.message)] - assert bool(frozen_msg) == freeze_mtp + assert bool(frozen_msg) == frozen with distillation_model.hide_teacher_model(): named_params = dict(distillation_model.named_parameters()) @@ -428,8 +431,8 @@ def _test_freeze_mtp_with_quantized_mtp(freeze_mtp, kd_loss_alpha, rank, size): assert any(p.grad is not None and p.grad.any() for p in body_params.values()) # The MTP head trains unless frozen; frozen, it gets no gradient and the step leaves it as is. mtp_trained = any(p.grad is not None and p.grad.any() for p in mtp_params.values()) - assert mtp_trained != freeze_mtp - if freeze_mtp: + assert mtp_trained != frozen + if frozen: assert all(torch.equal(p, mtp_before[n]) for n, p in mtp_params.items()) @@ -639,11 +642,12 @@ def test_distillation_config_top_p_validation(): @pytest.mark.parametrize( ("quantized", "skip_lm_loss", "freeze_mtp"), [ - (False, True, False), - (True, True, False), - (True, False, False), + (False, True, None), + (True, True, None), + (True, False, None), (False, True, True), (False, False, True), + (True, True, False), # explicit False overrides the automatic freeze ], ) def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss, freeze_mtp): @@ -651,14 +655,27 @@ def test_skip_lm_loss_with_mtp(dist_workers, quantized, skip_lm_loss, freeze_mtp dist_workers.run(partial(_test_skip_lm_loss_with_mtp, quantized, skip_lm_loss, freeze_mtp)) -@pytest.mark.parametrize(("freeze_mtp", "kd_loss_alpha"), [(True, 1.0), (True, 0.5), (False, 1.0)]) -def test_freeze_mtp_with_quantized_mtp(dist_workers, freeze_mtp, kd_loss_alpha): - """With the MTP head quantized too, only ``freeze_mtp`` keeps it out of a training step.""" - dist_workers.run(partial(_test_freeze_mtp_with_quantized_mtp, freeze_mtp, kd_loss_alpha)) +@pytest.mark.parametrize( + ("mtp_quantized", "freeze_mtp", "kd_loss_alpha"), + [ + (True, None, 1.0), # automatic: a quantized MTP head keeps training + (False, None, 1.0), # automatic: an unquantized MTP head is frozen + (False, False, 1.0), # explicit False overrides the automatic freeze + (True, True, 1.0), + (True, True, 0.5), + (False, True, 0.5), + (True, False, 1.0), + ], +) +def test_freeze_mtp_training_step(dist_workers, mtp_quantized, freeze_mtp, kd_loss_alpha): + """Through a training step, the MTP head is frozen exactly when ``freeze_mtp`` resolves to it.""" + dist_workers.run( + partial(_test_freeze_mtp_training_step, mtp_quantized, freeze_mtp, kd_loss_alpha) + ) -def test_mtp_excluded_from_quantization(): - """MTP loss is skipped only when the model is quantized and MTP is left out of it.""" +def test_has_quantized_backbone_with_unquantized_mtp(): + """True only for a model with MTP and enabled quantizers, none of them inside MTP.""" def _quantizer(enabled: bool) -> TensorQuantizer: quantizer = TensorQuantizer() @@ -676,18 +693,20 @@ def _model(*, with_mtp: bool, body_quant: bool, mtp_quant: bool) -> nn.Module: model.mtp.weight_quantizer = _quantizer(mtp_quant) return model - # Plain distillation (e.g. pruning recovery) still trains the MTP head. - assert not _mtp_excluded_from_quantization( + # Nothing quantized. + assert not _has_quantized_backbone_with_unquantized_mtp( _model(with_mtp=True, body_quant=False, mtp_quant=False) ) - # QAD with MTP excluded from the recipe: no quantization error to recover there. - assert _mtp_excluded_from_quantization(_model(with_mtp=True, body_quant=True, mtp_quant=False)) - # QAD with a quantized MTP head: keep its loss. - assert not _mtp_excluded_from_quantization( + # Quantized backbone, MTP left out of quantization. + assert _has_quantized_backbone_with_unquantized_mtp( + _model(with_mtp=True, body_quant=True, mtp_quant=False) + ) + # MTP quantized too. + assert not _has_quantized_backbone_with_unquantized_mtp( _model(with_mtp=True, body_quant=True, mtp_quant=True) ) # No MTP at all. - assert not _mtp_excluded_from_quantization( + assert not _has_quantized_backbone_with_unquantized_mtp( _model(with_mtp=False, body_quant=True, mtp_quant=False) )