Add context parallelism to HF QAT and QAD

CHANGELOG.rst

@@ -22,6 +22,7 @@ Changelog
 - Add ``examples/hf_ptq/convert_ggml_to_int8.py`` to expand a checkpoint exported with a GGML recipe (IQ formats, ``q8_0``) into an INT8 weight-only compressed-tensors checkpoint. Stock vLLM serves it with its integer kernels at about half the BF16 weight memory, reproducing the GGML weights up to rounding their scales to the model dtype.
 - Add a fake-quantized weight cache for QAT and QAD. Enroll the model with ``mtq.utils.enable_weight_cache(model)`` and wrap the forward and backward passes of each optimizer step in ``mtq.utils.weight_cache_scope(model)``: each weight is then fake-quantized once per step instead of once per micro-batch and activation recomputation, at the cost of one extra copy of each cached weight. It needs a straight-through weight gradient (``pass_through_bwd=True``, the default).
 - Speed up NVFP4 QAT and QAD of Hugging Face fused MoE experts: all experts of a projection are now fake-quantized in one kernel launch instead of one per expert, about 2x faster per MoE layer on Qwen3-30B-A3B with FSDP2.
+- Add context parallelism to QAT and QAD in ``examples/llm_qat``: ``--cp_size N`` splits each sequence across ``N`` GPUs with accelerate (FSDP2, ``--attn_implementation sdpa``, models whose layers all use full causal attention). ``QATTrainer``, ``QADTrainer`` and ``KDTrainer`` now train correctly under accelerate's context parallelism.
 
 *Speculative Decoding*
 

examples/llm_qat/ARGUMENTS.md

@@ -67,4 +67,5 @@ Extends [HuggingFace TrainingArguments](https://huggingface.co/docs/transformers
 | `--lr_config` | `str` | `None` | Path to a YAML file mapping fnmatch patterns to optimizer kwargs (e.g. lr, weight_decay). First matching pattern wins per parameter. See examples/llm_qat/configs/train/lr/lr_config_example.yaml. |
 | `--manual_gc` | `bool` | `False` | Run `gc.collect()` before each training/prediction step to work around GPU memory leaks during QAT/distillation. |
 | `--liger_ce_label_smoothing` | `float` | `0.0` | Label smoothing for Liger fused CE loss. Only used when --use_liger_kernel is enabled. |
+| `--cp_size` | `int` | `1` | Context parallel size: split each sequence across this many GPUs with accelerate's context parallelism. Needs the FSDP2 accelerate config, attn_implementation 'sdpa' and a model whose layers all use full causal attention. |
 | `--lora` | `bool` | `False` | Whether to add LoRA (Low-Rank Adaptation) adapter before training. When using real quantization, the LoRA adapter must be set, as quantized weights will be frozen during training. |

examples/llm_qat/README.md

@@ -150,6 +150,18 @@ Built-in recipes support full-model, partial-layer, and mixed-precision quantiza
 
 Replace `--config-file configs/accelerate/fsdp2.yaml` with the desired backend config in any of the commands above.
 
+### Context Parallelism
+
+For sequences too long for one GPU, `--cp_size N` splits each sequence across `N` GPUs with accelerate's context parallelism. FSDP2 still shards the model across all GPUs, so the GPU count must be a multiple of `N`:
+
+```bash
+accelerate launch --config-file configs/accelerate/fsdp2.yaml train.py \
+    --config configs/train/qad_nvfp4.yaml \
+    --attn_implementation sdpa --cp_size 2 ...
+```
+
+Context parallelism needs the FSDP2 config, SDPA attention for the student and the teacher, and `--model_max_length` a multiple of `2 * cp_size`. It supports models whose layers all use full causal attention, such as Llama, Qwen3 and Qwen3-MoE, or DeepSeek-V3-style MLA; not sliding-window, linear-attention or sparse (DSA) layers. Softmax (`p_bmm`) quantization is not supported with it.
+
 ## QLoRA (Real Quantization)
 
 [QLoRA](https://arxiv.org/pdf/2305.14314) reduces training memory by quantizing LoRA backbone weights with real quantization via `mtq.compress()`.

examples/llm_qat/arguments.py

@@ -98,6 +98,16 @@ class TrainingArguments(ModelOptTrainerArguments, transformers.TrainingArguments
         default=True,
         metadata={"help": "Use Liger kernel for fused loss computation. Reduces memory usage."},
     )
+    cp_size: int = field(
+        default=1,
+        metadata={
+            "help": (
+                "Context parallel size: split each sequence across this many GPUs with accelerate's "
+                "context parallelism. Needs the FSDP2 accelerate config, attn_implementation 'sdpa' "
+                "and a model whose layers all use full causal attention."
+            )
+        },
+    )
     lora: bool = field(
         default=False,
         metadata={

examples/llm_qat/train.py

@@ -56,8 +56,39 @@
 mto.enable_huggingface_checkpointing()
 
 
+def enable_context_parallelism(model_args, training_args):
+    """Split each sequence across ``cp_size`` GPUs with accelerate's context parallelism."""
+    from accelerate import ParallelismConfig
+
+    # Other backends would give the ranks of a context-parallel group different samples.
+    if os.environ.get("ACCELERATE_USE_FSDP") != "true" or os.environ.get("FSDP_VERSION") != "2":
+        raise ValueError("--cp_size > 1 needs the FSDP2 accelerate config, e.g. fsdp2.yaml.")
+    if model_args.attn_implementation != "sdpa":
+        raise ValueError(
+            "--cp_size > 1 needs --attn_implementation sdpa: context parallelism splits the "
+            "sequence inside SDPA only."
+        )
+    seq_chunks = 2 * training_args.cp_size
+    if model_args.model_max_length % seq_chunks:
+        raise ValueError(
+            f"--model_max_length must be a multiple of 2 * --cp_size ({seq_chunks}) so each "
+            "sequence splits evenly across the context-parallel group."
+        )
+    world_size = int(os.environ.get("WORLD_SIZE", 1))
+    if world_size % training_args.cp_size:
+        raise ValueError(
+            f"The number of GPUs ({world_size}) must be a multiple of --cp_size "
+            f"({training_args.cp_size})."
+        )
+    training_args.parallelism_config = ParallelismConfig(
+        cp_size=training_args.cp_size, dp_shard_size=world_size // training_args.cp_size
+    )
+
+
 def train():
     model_args, training_args, data_args, distill_args = get_training_args()
+    if training_args.cp_size > 1:
+        enable_context_parallelism(model_args, training_args)
 
     print_rank_0(f"arguments: {model_args}, {training_args}, {data_args}, {distill_args}")
 

modelopt/torch/distill/plugins/huggingface.py

@@ -183,6 +183,18 @@ def _ensure_teacher_prepared(self):
         """Prepare teacher model via accelerator (handles FSDP2, DeepSpeed, DDP)."""
         if self._teacher_prepared:
             return
+        # The Trainer checks only the student; another attention would see only this rank's slice.
+        parallelism_config = getattr(self.accelerator, "parallelism_config", None)
+        teacher_attn = getattr(self._teacher_model.config, "_attn_implementation", None)
+        if (
+            parallelism_config is not None
+            and parallelism_config.cp_enabled
+            and teacher_attn != "sdpa"
+        ):
+            raise ValueError(
+                "Context parallelism splits the sequence inside SDPA only, but the teacher uses "
+                f"{teacher_attn!r} attention. Load it with attn_implementation='sdpa'."
+            )
         self._teacher_prepared = True
         self._teacher_model = self._prepare_model(self._teacher_model)
         print_rank_0("Teacher model prepared for distillation.")
@@ -224,8 +236,13 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         self._ensure_teacher_prepared()
         # ``skip_logits`` (added for Liger eval by transformers>=5.15) would make the teacher's
         # Liger forward demand labels; the fused KD loss needs its hidden states instead.
-        kd_inputs = {k: v for k, v in inputs.items() if k not in ("labels", "skip_logits")}
+        # Under context parallelism the Trainer adds ``shift_labels``, shifted before the sequence
+        # was split so they stay aligned across shards; given to a Liger forward, they skip logits.
+        kd_inputs = {
+            k: v for k, v in inputs.items() if k not in ("labels", "skip_logits", "shift_labels")
+        }
         labels = inputs.get("labels")
+        shift_labels = inputs.get("shift_labels")
         is_training = model.training
 
         if is_training:
@@ -238,13 +255,17 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
                 else:
                     outputs = model(**kd_inputs)
                     if self._kd_loss_alpha < 1.0:
-                        ce_loss = self._liger_loss_func(outputs, labels, **kwargs)
+                        ce_loss = self._liger_loss_func(
+                            outputs, labels, shift_labels=shift_labels, **kwargs
+                        )
         else:
             ce_loss, outputs = super().compute_loss(model, inputs, return_outputs=True, **kwargs)
             batch_size = find_batch_size(inputs)
             self._record_eval_ce_loss(ce_loss, batch_size)
 
-        kd_loss = self._compute_kd_loss(outputs, labels, kd_inputs, **kwargs)
+        kd_loss = self._compute_kd_loss(
+            outputs, labels, kd_inputs, shift_labels=shift_labels, **kwargs
+        )
 
         if self._kd_loss_alpha < 1.0:
             loss = self._kd_loss_alpha * kd_loss + (1.0 - self._kd_loss_alpha) * ce_loss
@@ -276,16 +297,24 @@ def _compute_teacher_outputs(self, inputs):
             self._teacher_model.eval()
             return self._teacher_model(**inputs)
 
-    def _standard_kd_loss(self, outputs, labels, num_items_in_batch=None, **kwargs):
-        """KD loss with causal shift and ignore-index masking."""
+    def _standard_kd_loss(
+        self, outputs, labels, num_items_in_batch=None, shift_labels=None, **kwargs
+    ):
+        """KD loss with causal shift and ignore-index masking.
+
+        ``shift_labels`` (context parallelism) are already aligned with the positions, so every
+        position is scored; otherwise the last one, which has no next label, is dropped.
+        """
         # Match causal LM CE: logits at position t are scored against label t+1.
-        student_logits = outputs.logits[..., :-1, :].float()
-        teacher_logits = self._last_teacher_outputs.logits[..., :-1, :].float()
+        end = None if shift_labels is not None else -1
+        student_logits = outputs.logits[..., :end, :].float()
+        teacher_logits = self._last_teacher_outputs.logits[..., :end, :].float()
         self._last_teacher_outputs = None
         per_token_loss = self._kd_criterion(student_logits, teacher_logits)
-        if labels is None:
-            return per_token_loss.mean()
-        shift_labels = labels[..., 1:]
+        if shift_labels is None:
+            if labels is None:
+                return per_token_loss.mean()
+            shift_labels = labels[..., 1:]
         mask = shift_labels != IGNORE_INDEX
         num_tokens = mask.sum().clamp(min=1)
         loss = (per_token_loss * mask).sum() / num_tokens
@@ -317,7 +346,7 @@ def _teacher_liger_identity_lm_head(self):
         finally:
             teacher_lm_head.forward = teacher_orig
 
-    def _liger_kd_loss(self, outputs, labels, num_items_in_batch=None, **kwargs):
+    def _liger_kd_loss(self, outputs, labels, num_items_in_batch=None, shift_labels=None, **kwargs):
         """Fused lm_head + JSD for KD."""
         from liger_kernel.transformers import LigerFusedLinearJSD
 
@@ -331,10 +360,13 @@ def _liger_kd_loss(self, outputs, labels, num_items_in_batch=None, **kwargs):
         teacher_hs = self._last_teacher_outputs.logits.to(teacher_lm_head.weight.dtype)
         self._last_teacher_outputs = None
 
-        # Causal LM shift
-        student_hs = student_hs[..., :-1, :].contiguous().view(-1, student_hs.size(-1))
-        teacher_hs = teacher_hs[..., :-1, :].contiguous().view(-1, teacher_hs.size(-1))
-        shift_labels = labels[..., 1:].contiguous().view(-1)
+        # Causal LM shift, unless context parallelism already shifted the labels
+        end = None if shift_labels is not None else -1
+        student_hs = student_hs[..., :end, :].contiguous().view(-1, student_hs.size(-1))
+        teacher_hs = teacher_hs[..., :end, :].contiguous().view(-1, teacher_hs.size(-1))
+        if shift_labels is None:
+            shift_labels = labels[..., 1:]
+        shift_labels = shift_labels.contiguous().view(-1)
 
         jsd = LigerFusedLinearJSD(
             jsd_beta=self._liger_jsd_beta,

modelopt/torch/opt/plugins/transformers.py

@@ -17,6 +17,7 @@
 
 import dataclasses
 import fnmatch
+import functools
 import gc
 import json
 import os
@@ -758,14 +759,30 @@ def compute_loss(self, model, inputs, return_outputs=False, **kwargs):
         """Compute loss, patching lm_head to identity when using liger fused loss."""
         if self.use_liger_kernel:
             # The fused loss is computed here from the hidden states, so drop the ``skip_logits``
-            # that transformers>=5.15 adds for eval: Liger would then need the popped labels.
-            inputs = {k: v for k, v in inputs.items() if k != "skip_logits"}
-            with self._liger_identity_lm_head():
-                return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs)
+            # that transformers>=5.15 adds for eval: Liger would then need the popped labels. Under
+            # context parallelism the Trainer adds ``shift_labels``, shifted before the sequence was
+            # split; a Liger forward given them skips its logits, so the fused loss takes them.
+            shift_labels = inputs.get("shift_labels")
+            inputs = {k: v for k, v in inputs.items() if k not in ("skip_logits", "shift_labels")}
+            loss_func = self.compute_loss_func
+            if shift_labels is not None:
+                self.compute_loss_func = functools.partial(loss_func, shift_labels=shift_labels)
+            try:
+                with self._liger_identity_lm_head():
+                    return super().compute_loss(
+                        model, inputs, return_outputs=return_outputs, **kwargs
+                    )
+            finally:
+                self.compute_loss_func = loss_func
         return super().compute_loss(model, inputs, return_outputs=return_outputs, **kwargs)
 
-    def _liger_loss_func(self, outputs, labels, num_items_in_batch=None, **kwargs):
-        """Fused lm_head + CE loss via liger kernel."""
+    def _liger_loss_func(
+        self, outputs, labels, num_items_in_batch=None, shift_labels=None, **kwargs
+    ):
+        """Fused lm_head + CE loss via liger kernel.
+
+        ``shift_labels`` are labels the Trainer shifted before context parallelism split the sequence.
+        """
         from liger_kernel.transformers.model.loss_utils import LigerForCausalLMLoss
 
         model = self.accelerator.unwrap_model(self.model)
@@ -777,6 +794,7 @@ def _compute():
                 hidden_states=hidden_states,
                 lm_head_weight=lm_head.weight,
                 labels=labels,
+                shift_labels=shift_labels,
                 hidden_size=hidden_states.size(-1),
                 num_items_in_batch=num_items_in_batch,
                 ignore_index=IGNORE_INDEX,

modelopt/torch/quantization/plugins/transformers_trainer.py

@@ -169,6 +169,7 @@ def __init__(
 
         self.quant_args = quant_args
         self.quant_cfg = resolve_quant_cfg_from_args(quant_args)
+        self._context_parallel_checked = False
 
         # Add lora adapter before quantizing the model
         if getattr(self.args, "lora_config", None) is not None and not hasattr(
@@ -259,13 +260,36 @@ def training_step(self, *args, **kwargs):
         gc.collect()
         if self.quant_cfg is not None and not is_quantized(self.model):
             self._quantize_model()
+        if not self._context_parallel_checked:
+            self._check_context_parallel()
         return super().training_step(*args, **kwargs)
 
+    def _check_context_parallel(self):
+        """Refuse quantized attention that context parallelism would split wrongly."""
+        self._context_parallel_checked = True
+        parallelism_config = getattr(self.accelerator, "parallelism_config", None)
+        if parallelism_config is None or not parallelism_config.cp_enabled:
+            return
+        # Softmax (P) quantization runs ModelOpt's own attention kernel instead of SDPA, so each
+        # rank would attend only within its own slice of the sequence.
+        names = [
+            name
+            for name, module in self.model.named_modules()
+            if name.endswith("p_bmm_quantizer") and getattr(module, "is_enabled", False)
+        ]
+        if names:
+            raise ValueError(
+                "Context parallelism does not support softmax (p_bmm) quantization, which runs "
+                f"its own attention kernel: {', '.join(names[:3])}. Disable it or use cp_size=1."
+            )
+
     def prediction_step(self, *args, **kwargs):
         """Prediction step."""
         gc.collect()
         if self.quant_cfg is not None and not is_quantized(self.model):
             self._quantize_model()
+        if not self._context_parallel_checked:
+            self._check_context_parallel()
         return super().prediction_step(*args, **kwargs)
 
     def evaluate(self, *args, **kwargs):

tests/examples/llm_qat/test_llm_qat.py

@@ -15,6 +15,7 @@
 
 
 import json
+import os
 
 import pytest
 import torch
@@ -74,7 +75,13 @@ def _run_quantize(config: str, extra_cmd_args: list[str], cache_dir: str = ""):
     )
 
 
-def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", cache_dir: str = ""):
+def _run_train(
+    config: str,
+    extra_cmd_args: list[str],
+    backend: str = "fsdp2",
+    cache_dir: str = "",
+    env: dict[str, str] | None = None,
+):
     config_file = BACKEND_CONFIGS[backend]
     gradient_args = (
         ["--gradient_checkpointing", "True"]
@@ -94,6 +101,7 @@ def _run_train(config: str, extra_cmd_args: list[str], backend: str = "fsdp2", c
         ],
         "llm_qat",
         setup_free_port=True,
+        env=env,
     )
 
 
@@ -266,6 +274,31 @@ def test_qwen3_qad_nvfp4(tied_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_
     )
 
 
+@pytest.mark.skipif(torch.cuda.device_count() < 2, reason="context parallelism needs 2 GPUs")
+@pytest.mark.parametrize("distill", [False, True])
+def test_qwen3_nvfp4_context_parallel(
+    tied_qwen3_path, nvfp4_ptq_dir, dataset_cache_dir, tmp_path, distill
+):
+    """QAT and QAD, eval included, with each sequence split across 2 GPUs by accelerate."""
+    distill_args = ["--distill", "True", "--teacher_model", tied_qwen3_path] if distill else []
+    visible = os.environ.get(
+        "CUDA_VISIBLE_DEVICES", ",".join(map(str, range(torch.cuda.device_count())))
+    )
+    _run_train(
+        f"configs/train/{'qad' if distill else 'qat'}_nvfp4.yaml",
+        [
+            "--model_name_or_path", nvfp4_ptq_dir,
+            "--do_train", "True",
+            "--output_dir", str(tmp_path / "cp"),
+            "--attn_implementation", "sdpa",
+            "--cp_size", "2",
+            *distill_args,
+        ],
+        cache_dir=dataset_cache_dir,
+        env={**os.environ, "CUDA_VISIBLE_DEVICES": ",".join(visible.split(",")[:2])},
+    )
+
+
 def test_qwen3_qlora_nvfp4(tiny_qwen3_path, dataset_cache_dir, tmp_path):
     ptq_output_dir = tmp_path / "ptq"