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()