From d9a8404dd52f17700cc720d266eb1299a658227b Mon Sep 17 00:00:00 2001 From: Shiyang Chen Date: Wed, 28 Oct 2026 18:27:27 -0700 Subject: [PATCH] Add context parallelism to HF QAT and QAD examples/llm_qat gets --cp_size, which turns on accelerate's context parallelism (FSDP2, SDPA attention). Under context parallelism the Trainer adds shift_labels, shifted before the sequence is split across ranks. KDTrainer and the Liger loss in ModelOptHFTrainer passed them to the model, which crashed with Liger, and the KD losses shifted the labels again on each shard. They now keep shift_labels out of the model inputs and score every position of the shard against them. KDTrainer refuses a teacher that does not use SDPA, and QATTrainer refuses enabled p_bmm quantizers, since either would attend only within its own shard. Co-Authored-By: Claude Opus 5.5 Signed-off-by: Shiyang Chen --- CHANGELOG.rst | 1 + examples/llm_qat/ARGUMENTS.md | 1 + examples/llm_qat/README.md | 12 ++++ examples/llm_qat/arguments.py | 10 +++ examples/llm_qat/train.py | 31 ++++++++++ modelopt/torch/distill/plugins/huggingface.py | 62 ++++++++++++++----- modelopt/torch/opt/plugins/transformers.py | 30 +++++++-- .../plugins/transformers_trainer.py | 24 +++++++ tests/examples/llm_qat/test_llm_qat.py | 35 ++++++++++- .../distill/plugins/test_huggingface_kd.py | 31 ++++++++++ 10 files changed, 215 insertions(+), 22 deletions(-) diff --git a/CHANGELOG.rst b/CHANGELOG.rst index 90aa9a89b4c..037859713ee 100755 --- a/CHANGELOG.rst +++ b/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* diff --git a/examples/llm_qat/ARGUMENTS.md b/examples/llm_qat/ARGUMENTS.md index 8e1022d49dc..ac5bdd5c3e4 100644 --- a/examples/llm_qat/ARGUMENTS.md +++ b/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. | diff --git a/examples/llm_qat/README.md b/examples/llm_qat/README.md index ccdc2facd09..252a36808f3 100644 --- a/examples/llm_qat/README.md +++ b/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()`. diff --git a/examples/llm_qat/arguments.py b/examples/llm_qat/arguments.py index 52a3bda8983..9c635af2c45 100644 --- a/examples/llm_qat/arguments.py +++ b/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={ diff --git a/examples/llm_qat/train.py b/examples/llm_qat/train.py index 53b94e0a284..ac81287b4a8 100644 --- a/examples/llm_qat/train.py +++ b/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}") diff --git a/modelopt/torch/distill/plugins/huggingface.py b/modelopt/torch/distill/plugins/huggingface.py index 80675a3f4aa..140c28e8c6e 100644 --- a/modelopt/torch/distill/plugins/huggingface.py +++ b/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, diff --git a/modelopt/torch/opt/plugins/transformers.py b/modelopt/torch/opt/plugins/transformers.py index 3ebef890f56..f5238d3ad85 100644 --- a/modelopt/torch/opt/plugins/transformers.py +++ b/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, diff --git a/modelopt/torch/quantization/plugins/transformers_trainer.py b/modelopt/torch/quantization/plugins/transformers_trainer.py index b3d5619beb2..5bdb8a261f5 100644 --- a/modelopt/torch/quantization/plugins/transformers_trainer.py +++ b/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): diff --git a/tests/examples/llm_qat/test_llm_qat.py b/tests/examples/llm_qat/test_llm_qat.py index b9c7ac09624..4a17197603b 100644 --- a/tests/examples/llm_qat/test_llm_qat.py +++ b/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" diff --git a/tests/unit/torch/distill/plugins/test_huggingface_kd.py b/tests/unit/torch/distill/plugins/test_huggingface_kd.py index ab8c9b3e9b5..8cec84e6295 100644 --- a/tests/unit/torch/distill/plugins/test_huggingface_kd.py +++ b/tests/unit/torch/distill/plugins/test_huggingface_kd.py @@ -231,6 +231,37 @@ def test_standard_kd_loss_without_labels_uses_mean(tmp_path): assert loss.item() == pytest.approx(expected.item()) +def test_context_parallel_shards_add_up_to_the_full_kd_loss(tmp_path): + """Under context parallelism each rank gets a slice of the sequence and ``shift_labels``. + + The Trainer shifts the labels before splitting, so the shards' KD losses, each scaled by the + global token count, add up to the unsplit loss. Neither model may receive ``shift_labels``. + """ + student, teacher = _make_models() + batch = _make_batch() + trainer = _make_trainer(tmp_path, student, teacher) + trainer.model.train() + labels = batch["labels"] + shift_labels = F.pad(labels, (0, 1), value=IGNORE_INDEX)[..., 1:] + num_items = (shift_labels != IGNORE_INDEX).sum().item() + + full = trainer.compute_loss(trainer.model, batch.copy(), num_items_in_batch=num_items) + shards = [ + trainer.compute_loss( + trainer.model, + { + "input_ids": batch["input_ids"][:, part], + "labels": labels[:, part], + "shift_labels": shift_labels[:, part], + }, + num_items_in_batch=num_items, + ) + for part in (slice(0, 2), slice(2, 4)) + ] + + assert sum(shards).item() == pytest.approx(full.item()) + + def test_kd_loss_ddp_scaling(tmp_path): """KD loss should be scaled by world_size when average_tokens_across_devices is True.""" student, teacher = _make_models()