Support LSQ (scale learning) in MCore

examples/post_training/modelopt/finetune.py

@@ -261,7 +261,16 @@ def _process_example(self, example: Dict[str, Any]):
                 return None
 
         # We always add eos between samples for training purpose.
-        input_ids = self.tokenizer.apply_chat_template(example)
+        # Newer transformers (>=5) return a BatchEncoding/dict from apply_chat_template instead of a
+        # bare list of token ids; normalize to a flat list either way.
+        _enc = self.tokenizer.apply_chat_template(example)
+        if isinstance(_enc, list):
+            input_ids = _enc
+        else:
+            input_ids = _enc["input_ids"]
+            if len(input_ids) and isinstance(input_ids[0], (list, tuple)):
+                input_ids = input_ids[0]
+        input_ids = list(input_ids)
         current_loss_mask = [1] * len(input_ids)
         input_ids = input_ids + [get_eos_token_id(self.tokenizer)]
         current_loss_mask += [0]

megatron/post_training/arguments.py

@@ -114,6 +114,15 @@ def add_modelopt_args(parser):
         'parameters and only train MTP heads.',
     )
 
+    # Scale-learning QAD (LSQ): the checkpoint is assumed to already have LSQ enabled
+    group.add_argument(
+        '--lsq-scale-lr',
+        type=float,
+        default=None,
+        help='Optional separate learning rate for LSQ scale (amax) parameters. If unset, they '
+        'use the base learning rate.',
+    )
+
     # Special model architecture option
     group.add_argument(
         '--export-qk-l2-norm',

megatron/post_training/model_builder.py

@@ -5,6 +5,8 @@
 import logging
 import os
 from argparse import Namespace
+
+import torch
 from typing import Any, Dict
 
 import modelopt.torch.distill as mtd
@@ -218,6 +220,34 @@ def _freeze_base_for_mtp(model):
     _freeze_for_qad(model, "mtp")
 
 
+def _propagate_expert_allreduce_to_lsq_params(model):
+    """Copy each weight's MCore ``allreduce`` flag onto its LSQ ``_amax_{pre,post}`` params.
+
+    Without the flag, expert amax params get bucketed into the global data-parallel group
+    instead of the expert-data-parallel group (wrong grad reduction under expert parallelism).
+    """
+    module_by_name = dict(model.named_modules())
+    for name, param in model.named_parameters():
+        if not (name.endswith("_amax_pre") or name.endswith("_amax_post")):
+            continue
+        marker = ".weight_quantizer"
+        if marker not in name:
+            continue
+        linear = module_by_name.get(name.split(marker, 1)[0])
+        if linear is None:
+            continue
+        allreduce = next(
+            (
+                w.allreduce
+                for wname, w in linear.named_parameters(recurse=False)
+                if "weight" in wname and hasattr(w, "allreduce")
+            ),
+            None,
+        )
+        if allreduce is not None:
+            param.allreduce = allreduce
+
+
 def modelopt_gpt_hybrid_builder(
     args,
     pre_process,
@@ -427,6 +457,26 @@ def modelopt_gpt_hybrid_builder(
     if qad_train_target is not None:
         _freeze_for_qad(model, qad_train_target)
 
+    # LSQ _amax_* params lack MCore's ``allreduce`` attribute; propagate it from each weight
+    # before get_model() builds DDP and the distributed optimizer. No-op without LSQ params.
+    _propagate_expert_allreduce_to_lsq_params(model)
+
+    # LSQ scale-only diagnostic: freeze every non-amax parameter so ONLY the learnable NVFP4 scales
+    # (_amax_pre / _amax_post) train. Env-gated; runs before get_model() wraps DDP and builds the
+    # optimizer, so frozen weights are excluded from the grad buffer and optimizer state entirely.
+    if os.environ.get("LSQ_SCALE_ONLY") == "1":
+        n_frozen = n_train = 0
+        for pname, p in model.named_parameters():
+            if pname.endswith("_amax_pre") or pname.endswith("_amax_post"):
+                p.requires_grad = True
+                n_train += 1
+            else:
+                p.requires_grad = False
+                n_frozen += 1
+        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:
+            print(f"[LSQ_SCALE_ONLY] froze {n_frozen} non-amax params, training {n_train} amax params",
+                  flush=True)
+
     _add_load_convert_hooks(model)
 
     # Distillation mode.

megatron/training/training.py

@@ -1955,6 +1955,17 @@ def get_megatron_optimizer_config(args: Any) -> OptimizerConfig:
     #  can be added to as needed by the user, or replaced entirely with a custom override.
     config_overrides = get_standard_config_overrides(config=config)
 
+    # Scale-learning QAD (LSQ): route the learnable scale (amax) parameters into their own
+    # optimizer param group with a (typically higher) learning rate. 
+    lsq_scale_lr = getattr(args, "lsq_scale_lr", None)
+    if lsq_scale_lr is not None:
+        from megatron.core.optimizer.optimizer_config import ParamKey
+        from megatron.core.optimizer_param_scheduler import ParamGroupOverride
+
+        config_overrides[ParamKey(name=("*_amax_pre", "*_amax_post"))] = ParamGroupOverride(
+            max_lr=lsq_scale_lr
+        )
+
     return config, config_overrides
 
 def get_megatron_ddp_config(args: argparse.Namespace) -> DistributedDataParallelConfig: