From a8e8a302a649f6514ded46b271fcb61f176ca5a0 Mon Sep 17 00:00:00 2001 From: Dan Saunders Date: Tue, 23 Dec 2025 14:29:44 -0500 Subject: [PATCH] Add context parallelism support --- unsloth-cli.py | 34 ++ unsloth/__init__.py | 5 + unsloth/context_parallel.py | 590 ++++++++++++++++++++++++++++ unsloth/models/llama.py | 139 +++++-- unsloth/trainer.py | 15 +- unsloth/utils/attention_dispatch.py | 25 +- 6 files changed, 772 insertions(+), 36 deletions(-) create mode 100644 unsloth/context_parallel.py diff --git a/unsloth-cli.py b/unsloth-cli.py index 612da11eb2..b1c1f21c34 100644 --- a/unsloth-cli.py +++ b/unsloth-cli.py @@ -31,6 +31,7 @@ Happy fine-tuning! import argparse import os +from typing import Optional def run(args): @@ -133,6 +134,10 @@ def run(args): print("Data is formatted and ready!") # Configure training arguments + pad_multiple = args.pad_to_multiple_of + if pad_multiple is None and args.context_parallel_size > 1: + pad_multiple = 2 * args.context_parallel_size + training_args = SFTConfig( per_device_train_batch_size = args.per_device_train_batch_size, per_device_eval_batch_size = args.per_device_eval_batch_size, @@ -153,6 +158,9 @@ def run(args): dataset_num_proc = 2, ddp_find_unused_parameters = False if distributed else None, packing = args.packing, + context_parallel_size = args.context_parallel_size, + pad_to_multiple_of = pad_multiple, + shuffle_dataset = args.shuffle_dataset, ) # Initialize trainer @@ -360,6 +368,32 @@ if __name__ == "__main__": action = "store_true", help = "Enable padding-free sample packing via TRL's bin packer.", ) + training_group.add_argument( + "--pad_to_multiple_of", + type = int, + default = None, + help = ( + "Pad every batch to a multiple of this value. " + "Defaults to `2 * context_parallel_size` when context parallelism is enabled." + ), + ) + training_group.add_argument( + "--shuffle_dataset", + action = argparse.BooleanOptionalAction, + default = True, + help = "Shuffle the dataset during training (default: True).", + ) + + context_group = parser.add_argument_group("🧩 Context Parallelism") + context_group.add_argument( + "--context_parallel_size", + type = int, + default = 1, + help = ( + "Number of distributed ranks participating in PyTorch context parallelism. " + "Set >1 only when running with torch.distributed initialized on PyTorch >= 2.7." + ), + ) report_group = parser.add_argument_group("📊 Report Options") report_group.add_argument( diff --git a/unsloth/__init__.py b/unsloth/__init__.py index dbacd551c4..5240a95764 100644 --- a/unsloth/__init__.py +++ b/unsloth/__init__.py @@ -328,3 +328,8 @@ from unsloth_zoo.rl_environments import ( # Patch TRL trainers for backwards compatibility _patch_trl_trainer() + +from .context_parallel import patch_sft_config, patch_sft_trainer + +patch_sft_config() +patch_sft_trainer() diff --git a/unsloth/context_parallel.py b/unsloth/context_parallel.py new file mode 100644 index 0000000000..30ead77ad6 --- /dev/null +++ b/unsloth/context_parallel.py @@ -0,0 +1,590 @@ +from __future__ import annotations + +import contextlib +import contextvars +import functools +import sys +import warnings +from dataclasses import dataclass, field +from typing import Iterator, Optional, Tuple + +import torch +import torch.nn.functional as F +import torch.distributed as dist +import trl + +try: + from torch.distributed.tensor.experimental import context_parallel + from torch.distributed.tensor import DeviceMesh +except (ImportError, AttributeError): + context_parallel = None + DeviceMesh = None + +from .device_type import DEVICE_TYPE_TORCH +from .utils.packing import mask_packed_sequence_boundaries + +_ACTIVE_MANAGER: contextvars.ContextVar[Optional["ContextParallelManager"]] = ( + contextvars.ContextVar("unsloth_active_cp_manager", default = None) +) + +_BUFFER_NAMES = ( + "input_ids", + "attention_mask", + "labels", + "position_ids", + "shift_labels", +) + + +def get_cp_manager() -> Optional["ContextParallelManager"]: + return _ACTIVE_MANAGER.get() + + +@dataclass +class ContextParallelSettings: + size: int = field( + default = 1, + metadata = { + "help": ( + "Number of ranks that should participate in context parallelism. " + "Set to >1 only when running under torch.distributed / accelerate." + ) + }, + ) + + @classmethod + def from_args(cls, args: Optional[object]) -> "ContextParallelSettings": + if args is None: + return cls() + size = int(getattr(args, "context_parallel_size", 1)) + return cls(size = size) + + +def _attach_context_parallel_attention_hooks(model: torch.nn.Module) -> list: + """ + Attach forward_pre_hooks to self_attn modules to ensure correct attention behavior + during context parallelism with load balancing. + + Args: + model: The model to attach hooks to + + Returns: + List of hook handles that can be used to remove the hooks later + """ + handles = [] + + def _self_attn_pre_forward_hook(_module, module_args, module_kwargs): + # Remove attention_mask and set is_causal=True + # This ensures ring attention uses causal masking correctly + if "attention_mask" in module_kwargs: + module_kwargs["attention_mask"] = None + if "is_causal" in module_kwargs or hasattr(_module, "is_causal"): + module_kwargs["is_causal"] = True + return module_args, module_kwargs + + # Find all self_attn modules - they may be nested in PEFT wrappers + attn_modules = [] + for name, module in model.named_modules(): + # Attach to modules ending with self_attn (transformers convention) + if name.endswith("self_attn"): + attn_modules.append((name, module)) + + for _, module in attn_modules: + handle = module.register_forward_pre_hook( + _self_attn_pre_forward_hook, with_kwargs = True, prepend = True + ) + handles.append(handle) + + return handles + + +class ContextParallelManager: + """Toggles PyTorch context parallelism.""" + + def __init__(self, settings: ContextParallelSettings): + self.settings = settings + self._mesh: Optional[DeviceMesh] = None + self._device_mesh: Optional[DeviceMesh] = None + self._cp_group: Optional[dist.ProcessGroup] = None + self._cp_rank_index: int = 0 + self._dp_world_size: int = 1 + self._world_size: int = dist.get_world_size() + self._report_loss: Optional[torch.Tensor] = None + self._attention_hook_handles: list = [] + self._mesh = self._build_mesh() + self._device_mesh = self._build_device_mesh() + + def attach_attention_hooks(self, model: torch.nn.Module) -> None: + """ + Attach hooks to self_attn modules to ensure correct attention behavior during + context parallelism with load balancing. + """ + if self._attention_hook_handles: + return + self._attention_hook_handles = _attach_context_parallel_attention_hooks(model) + + def _build_mesh(self) -> DeviceMesh: + rank = torch.distributed.get_rank() + group_index = rank // self.settings.size + start = group_index * self.settings.size + cp_ranks = torch.arange(start, start + self.settings.size, dtype = torch.int64) + mesh = DeviceMesh(DEVICE_TYPE_TORCH, cp_ranks) + self._cp_group = mesh.get_group() + self._cp_rank_index = int(rank - start) + return mesh + + def _build_device_mesh(self) -> DeviceMesh: + self._dp_world_size = self._world_size // self.settings.size + mesh = torch.arange(self._world_size, dtype = torch.int64).reshape( + self._dp_world_size, self.settings.size + ) + return DeviceMesh( + DEVICE_TYPE_TORCH, mesh, mesh_dim_names = ("dp_replicate", "cp") + ) + + @property + def device_mesh(self) -> Optional[DeviceMesh]: + return self._device_mesh + + @property + def data_parallel_world_size(self) -> int: + return self._dp_world_size + + @property + def cp_rank_index(self) -> int: + return self._cp_rank_index + + def data_parallel_rank(self) -> int: + return dist.get_rank() // self.settings.size + + def _collect_buffers( + self, inputs: dict[str, torch.Tensor] + ) -> Tuple[list[torch.Tensor], list[int], set[torch.Tensor]]: + buffers: list[torch.Tensor] = [] + for name in _BUFFER_NAMES: + tensor = inputs.get(name) + if tensor is None or not isinstance(tensor, torch.Tensor): + continue + if tensor.ndim <= 1: + continue + buffers.append(tensor) + return buffers, [1] * len(buffers), set(buffers) + + def _ensure_position_ids(self, inputs: dict[str, torch.Tensor]) -> None: + if "position_ids" in inputs: + return + input_ids = inputs.get("input_ids") + if input_ids is None: + return + seq_len = input_ids.size(1) + positions = torch.arange(seq_len, dtype = torch.long, device = input_ids.device) + inputs["position_ids"] = positions.unsqueeze(0).expand(input_ids.size(0), -1) + + def _ensure_shift_labels(self, inputs: dict[str, torch.Tensor]) -> None: + """Pre-shift labels globally before sharding for correct next-token prediction.""" + if "shift_labels" in inputs: + return + labels = inputs.get("labels") + if labels is None: + return + # Pad with -100, then take [1:] to get shifted labels + shift_labels = F.pad(labels, (0, 1), value = -100)[:, 1:].contiguous() + packed_seq_lengths = inputs.get("packed_seq_lengths") + if packed_seq_lengths is not None: + mask_packed_sequence_boundaries(shift_labels, packed_seq_lengths) + inputs["shift_labels"] = shift_labels + + @contextlib.contextmanager + def apply(self, inputs: dict[str, torch.Tensor]) -> Iterator[None]: + """Wrap training step to shard buffers and patch SDPA for ring attention.""" + token = _ACTIVE_MANAGER.set(self) + self._ensure_position_ids(inputs) + self._ensure_shift_labels(inputs) + buffers, seq_dims, no_restore = self._collect_buffers(inputs) + with context_parallel( + self._mesh, + buffers = buffers, + buffer_seq_dims = seq_dims, + no_restore_buffers = no_restore, + ): + yield + _ACTIVE_MANAGER.reset(token) + + def _set_report_loss(self, value: torch.Tensor) -> None: + self._report_loss = value.detach() if torch.is_tensor(value) else None + + def consume_report_loss(self) -> Optional[torch.Tensor]: + value = self._report_loss + self._report_loss = None + return value + + def reduce_loss(self, loss, inputs): + if self._cp_group is None: + return loss + + # Handle (loss, outputs) tuple from return_outputs=True + is_tuple = isinstance(loss, tuple) + if is_tuple: + tensor, rest = loss[0], loss[1:] + else: + tensor = loss + + # Count local valid tokens + shift_labels = inputs["shift_labels"] + local_tokens = ( + shift_labels.ne(-100).sum().to(dtype = tensor.dtype, device = tensor.device) + ) + + # Get global token count + global_tokens = local_tokens.clone() + dist.all_reduce(global_tokens, op = dist.ReduceOp.SUM, group = self._cp_group) + + # Weight loss by local fraction + weight = local_tokens.detach() / global_tokens.detach() + weighted_loss = tensor * weight + + # Reduce for reporting + global_loss = weighted_loss.detach().clone() + dist.all_reduce(global_loss, op = dist.ReduceOp.SUM, group = self._cp_group) + self._set_report_loss(global_loss) + + return (weighted_loss, *rest) if is_tuple else weighted_loss + + def reduce_grad_norm(self, grad_norm: float) -> float: + """ + Reduce gradient norm across CP group. + + Each rank computes a local gradient norm (L2). The global norm is: + sqrt(sum(local_norm_i^2 for all ranks)) + + This is needed because the Trainer computes grad_norm locally, but with + CP each rank only has partial gradients for its sequence shard. + """ + if self._cp_group is None: + return grad_norm + + # Square the local norm, sum across ranks, then sqrt + local_norm_sq = torch.tensor( + grad_norm**2, + dtype = torch.float32, + device = torch.device(DEVICE_TYPE_TORCH), + ) + dist.all_reduce(local_norm_sq, op = dist.ReduceOp.SUM, group = self._cp_group) + return float(local_norm_sq.sqrt().item()) + + +def patch_sft_config(): + """Patch SFTConfig to add context_parallel_size and shuffle_dataset fields.""" + base_cls = trl.SFTConfig + if hasattr(base_cls, "context_parallel_size"): + return + + @dataclass + class PatchedSFTConfig(base_cls): # type: ignore[misc, valid-type] + context_parallel_size: int = field( + default = 1, + metadata = { + "help": ( + "Number of ranks participating in context parallelism. " + "Set to 1 to disable context parallelism." + ) + }, + ) + shuffle_dataset: bool = field( + default = True, + metadata = { + "help": ( + "Whether to shuffle the training dataset before each epoch. " + "Exposed for CP = 1 vs. CP > 1 debugging purposes." + ) + }, + ) + + PatchedSFTConfig.__name__ = base_cls.__name__ + PatchedSFTConfig.__qualname__ = base_cls.__qualname__ + PatchedSFTConfig.__module__ = base_cls.__module__ + module = sys.modules.get(base_cls.__module__) + if module is not None: + setattr(module, base_cls.__name__, PatchedSFTConfig) + trl.SFTConfig = PatchedSFTConfig + if hasattr(trl, "trainer") and hasattr(trl.trainer, "sft_trainer"): + trl.trainer.sft_trainer.SFTConfig = PatchedSFTConfig + + +def patch_sft_trainer() -> None: + """Patch SFTTrainer to add context parallelism support.""" + trainer_cls = trl.SFTTrainer + if hasattr(trainer_cls, "__unsloth_context_parallel__"): + return + + original_init = trainer_cls.__init__ + original_compute_loss = trainer_cls.compute_loss + original_prediction_step = trainer_cls.prediction_step + original_training_step = trainer_cls.training_step + original_log = trainer_cls.log + original_get_train_sampler = getattr(trainer_cls, "_get_train_sampler", None) + + def _patch_train_sampler(original_fn): + @functools.wraps(original_fn) + def wrapper(self, *args, **kwargs): + sampler = original_fn(self, *args, **kwargs) + manager = getattr(self, "_context_parallel_manager", None) + dataset = args[0] if args else None + if dataset is None: + dataset = getattr(self, "train_dataset", None) + shuffle_dataset = getattr(self.args, "shuffle_dataset", True) + if ( + manager + and torch.distributed.is_available() + and torch.distributed.is_initialized() + and dataset is not None + ): + dp_world = manager.data_parallel_world_size + world_size = torch.distributed.get_world_size() + if dp_world != world_size: + try: + from torch.utils.data.distributed import DistributedSampler + except ImportError: + return sampler + dp_rank = manager.data_parallel_rank() + shuffle = shuffle_dataset and not getattr( + self.args, "group_by_length", False + ) + return DistributedSampler( + dataset, + num_replicas = dp_world, + rank = dp_rank, + shuffle = shuffle, + drop_last = getattr(self.args, "dataloader_drop_last", False), + ) + if not shuffle_dataset and dataset is not None: + try: + from torch.utils.data import SequentialSampler + except ImportError: + return sampler + return SequentialSampler(dataset) + return sampler + + return wrapper + + @functools.wraps(original_init) + def patched_init(self, *args, **kwargs): + original_init(self, *args, **kwargs) + settings = ContextParallelSettings.from_args(getattr(self, "args", None)) + if settings.size > 1: + if context_parallel is None or DeviceMesh is None: + warnings.warn( + "Context parallelism requested but PyTorch >= 2.7 is required.", + stacklevel = 2, + ) + self._context_parallel_manager = None + else: + self._context_parallel_manager = ContextParallelManager(settings) + else: + self._context_parallel_manager = None + accelerator = getattr(self, "accelerator", None) + manager = self._context_parallel_manager + if manager: + print( + f"Unsloth: Context parallelism enabled with size={manager.settings.size}" + ) + mesh = getattr(manager, "device_mesh", None) if manager else None + existing_mesh = ( + getattr(accelerator, "torch_device_mesh", None) + if accelerator is not None + else None + ) + if ( + accelerator is not None + and mesh is not None + and ( + existing_mesh is None + or "cp" not in getattr(existing_mesh, "mesh_dim_names", ()) + ) + ): + setattr(accelerator.state, "device_mesh", mesh) + + # When using pure context parallelism (dp_world_size=1), disable DDP + # to avoid gradient checkpointing compatibility issues + if manager and manager.data_parallel_world_size == 1: + try: + from accelerate.utils import DistributedType + + args = getattr(self, "args", None) + distributed_state = getattr(args, "distributed_state", None) + if ( + distributed_state is not None + and distributed_state.distributed_type == DistributedType.MULTI_GPU + ): + distributed_state.distributed_type = DistributedType.NO + except ImportError: + pass + + # Enable sync_each_batch when using CP with gradient accumulation and DDP. + # This keeps the computation graph constant for DDP + static_graph mode. + if ( + manager + and manager.data_parallel_world_size > 1 # Only needed with actual DP + and accelerator is not None + and hasattr(accelerator, "gradient_state") + ): + grad_accum_steps = getattr( + getattr(self, "args", None), "gradient_accumulation_steps", 1 + ) + if grad_accum_steps > 1: + accelerator.gradient_state.plugin_kwargs["sync_each_batch"] = True + + # Attach attention hooks for proper ring attention behavior with load balancing. + # This ensures attention_mask is removed and is_causal=True for all self_attn calls. + if manager: + model = getattr(self, "model", None) + if model is not None: + manager.attach_attention_hooks(model) + + @functools.wraps(original_compute_loss) + def patched_compute_loss(self, model, inputs, return_outputs = False, **kwargs): + manager = getattr(self, "_context_parallel_manager", None) + kwargs.pop("num_items_in_batch", None) + + # For context parallelism with shift_labels, prefer letting the model + # handle the pre-shifted targets when it advertises support. Otherwise + # fall back to an external loss that consumes the sharded tensors. + shift_labels = inputs.get("shift_labels") + use_cp_shift_labels = manager and isinstance(shift_labels, torch.Tensor) + model_supports_shift_labels = bool( + getattr( + model, + "_unsloth_supports_context_parallel_shift_labels", + False, + ) + ) + + if use_cp_shift_labels and not model_supports_shift_labels: + # Remove labels so model doesn't compute loss internally + saved_labels = inputs.pop("labels", None) + # Also remove shift_labels from inputs (model doesn't expect it) + local_shift_labels = inputs.pop("shift_labels", None) + + # Get model outputs (logits only, no loss) + outputs = model(**inputs) + logits = outputs.logits if hasattr(outputs, "logits") else outputs[0] + + # Compute loss using pre-shifted labels + from unsloth.kernels.cross_entropy_loss import fast_cross_entropy_loss + + loss = fast_cross_entropy_loss( + logits = logits, + labels = local_shift_labels, + ) + + # Restore labels for reduce_loss token counting + if saved_labels is not None: + inputs["labels"] = saved_labels + if local_shift_labels is not None: + inputs["shift_labels"] = local_shift_labels + + if return_outputs: + loss = (loss, outputs) + else: + loss = original_compute_loss( + self, + model, + inputs, + return_outputs = return_outputs, + **kwargs, + ) + + if manager: + loss = manager.reduce_loss(loss, inputs) + return loss + + @functools.wraps(original_prediction_step) + def patched_prediction_step( + self, + model, + inputs, + prediction_loss_only, + ignore_keys = None, + **kwargs, + ): + manager = getattr(self, "_context_parallel_manager", None) + context = manager.apply(inputs) if manager else contextlib.nullcontext() + with context: + return original_prediction_step( + self, + model, + inputs, + prediction_loss_only, + ignore_keys, + **kwargs, + ) + + def _maybe_enable_sync_each_batch(trainer): + """Enable sync_each_batch at runtime if gradient checkpointing is detected.""" + if getattr(trainer, "_sync_each_batch_checked", False): + return + setattr(trainer, "_sync_each_batch_checked", True) + + accelerator = getattr(trainer, "accelerator", None) + if accelerator is None or not hasattr(accelerator, "gradient_state"): + return + + # Check if already enabled + if accelerator.gradient_state.plugin_kwargs.get("sync_each_batch", False): + return + + model = getattr(trainer, "model", None) + is_checkpointing = getattr(model, "is_gradient_checkpointing", False) + grad_accum_steps = getattr(trainer.args, "gradient_accumulation_steps", 1) + + if is_checkpointing and grad_accum_steps > 1: + accelerator.gradient_state.plugin_kwargs["sync_each_batch"] = True + + @functools.wraps(original_training_step) + def patched_training_step(self, model, inputs, *args, **kwargs): + manager = getattr(self, "_context_parallel_manager", None) + original_n_gpu = getattr(self.args, "n_gpu", 1) + if manager: + setattr(self.args, "_n_gpu", manager.data_parallel_world_size) + _maybe_enable_sync_each_batch(self) + # Attach attention hooks if not already done (model may not be ready at init) + if not manager._attention_hook_handles: + m = getattr(self, "model", None) + if m is not None: + manager.attach_attention_hooks(m) + + # Wrap entire training step (forward + backward) in context_parallel + # This keeps SDPA patched and buffers sharded throughout, including + # during gradient checkpoint recomputation in backward pass. + cp_context = manager.apply(inputs) if manager else contextlib.nullcontext() + with cp_context: + loss = original_training_step(self, model, inputs, *args, **kwargs) + + if manager: + setattr(self.args, "_n_gpu", original_n_gpu) + + report_loss = manager.consume_report_loss() if manager else None + if report_loss is not None: + return report_loss + return loss + + @functools.wraps(original_log) + def patched_log(self, logs, start_time = None): + manager = getattr(self, "_context_parallel_manager", None) + # Reduce grad_norm across CP group if present + if manager and "grad_norm" in logs: + grad_norm = logs["grad_norm"] + if isinstance(grad_norm, (int, float)): + logs["grad_norm"] = manager.reduce_grad_norm(grad_norm) + return original_log(self, logs, start_time) + + trainer_cls.__init__ = patched_init + trainer_cls.compute_loss = patched_compute_loss + trainer_cls.prediction_step = patched_prediction_step + trainer_cls.training_step = patched_training_step + trainer_cls.log = patched_log + trainer_cls.__unsloth_context_parallel__ = True + if original_get_train_sampler is not None: + trainer_cls._get_train_sampler = _patch_train_sampler( + original_get_train_sampler + ) diff --git a/unsloth/models/llama.py b/unsloth/models/llama.py index 93d93e26d6..bace5c68b1 100644 --- a/unsloth/models/llama.py +++ b/unsloth/models/llama.py @@ -16,6 +16,7 @@ import torch import gc import math import functools +import os from typing import Optional, Tuple, List, Union from ._utils import * @@ -76,6 +77,8 @@ from transformers.modeling_attn_mask_utils import ( ) from ..kernels import * from ..tokenizer_utils import * +from ..context_parallel import get_cp_manager + from .vision import FastBaseModel # Final patching code @@ -148,7 +151,7 @@ from math import sqrt as math_sqrt KV_CACHE_INCREMENT = 512 # KV Cache update size torch_nn_functional_softmax = torch.nn.functional.softmax # SDPA has GQA internally -SDPA_HAS_GQA = "enable_gqa" in scaled_dot_product_attention.__doc__ +SDPA_HAS_GQA = "enable_gqa" in F.scaled_dot_product_attention.__doc__ from peft.utils.other import ModulesToSaveWrapper @@ -588,7 +591,7 @@ def LlamaAttention_fast_forward_inference( attention_mask = attention_mask.eq(0) if SDPA_HAS_GQA: - A = scaled_dot_product_attention( + A = F.scaled_dot_product_attention( Qn, Knn, Vnn, @@ -597,7 +600,7 @@ def LlamaAttention_fast_forward_inference( enable_gqa = True, ) else: - A = scaled_dot_product_attention( + A = F.scaled_dot_product_attention( Qn, Knn, Vnn, attn_mask = attention_mask, is_causal = is_causal ) A = A.transpose(1, 2) @@ -724,6 +727,10 @@ def LlamaAttention_fast_forward( head_dim = self.head_dim assert n_kv_heads * n_groups == n_heads + cp_manager = get_cp_manager() + cp_active = cp_manager is not None + cp_size = cp_manager.settings.size if cp_manager else 1 + cp_rank_index = cp_manager.cp_rank_index if cp_manager else 0 Q, K, V = self.apply_qkv(self, hidden_states) Q = Q.view(bsz, q_len, n_heads, head_dim).transpose(1, 2) K = K.view(bsz, q_len, n_kv_heads, head_dim).transpose(1, 2) @@ -734,25 +741,66 @@ def LlamaAttention_fast_forward( if past_key_value is not None: kv_seq_len += past_key_value[0].shape[-2] - if position_embeddings and kv_seq_len <= position_embeddings[0].shape[0]: + required_seq_len = kv_seq_len + if isinstance(position_ids, torch.Tensor) and position_ids.numel() > 0: + max_position = int(position_ids.max().item()) + 1 + required_seq_len = max(required_seq_len, max_position) + elif cp_active and cp_size > 1: + required_seq_len = max(required_seq_len, q_len * cp_size) + + if ( + position_embeddings + and required_seq_len <= position_embeddings[0].shape[0] + and required_seq_len <= position_embeddings[1].shape[0] + ): cos, sin = position_embeddings else: rotary_emb = self.rotary_emb - rotary_emb.extend_rope_embedding(V, seq_len = kv_seq_len) - cos, sin = rotary_emb.get_cached(kv_seq_len, Q.device.index) + rotary_emb.extend_rope_embedding(V, seq_len = required_seq_len) + cos, sin = rotary_emb.get_cached(required_seq_len, Q.device.index) cos = cos.to(device = Q.device, dtype = Q.dtype) sin = sin.to(device = Q.device, dtype = Q.dtype) + # For padding-free/packing, get position_ids from kwargs if not provided + # (TRL's collator puts them there when padding_free=True) rope_position_ids = position_ids if rope_position_ids is None and seq_info is not None: rope_position_ids = kwargs.get("position_ids") - # Q, K = ( - # fast_rope_embedding(Q, K, cos, sin) - # if rope_position_ids is None - # else inplace_rope_embedding(Q, K, cos, sin, rope_position_ids) - # ) - Q, K = fast_rope_embedding(Q, K, cos, sin, rope_position_ids) + def _slice_rope_frequencies( + cos: torch.Tensor, + sin: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + target_len = q_len + if isinstance(rope_position_ids, torch.Tensor): + ids = rope_position_ids + if ids.ndim > 1: + flat_ids = ids.view(-1, ids.shape[-1]) + if flat_ids.shape[0] > 1: + all_equal = torch.all(flat_ids == flat_ids[0]).item() + if not all_equal: + raise RuntimeError( + "fast_rope_embedding requires identical position_ids across the batch." + ) + ids = flat_ids[0] + ids = ids.to(device = cos.device, dtype = torch.long) + ids = ids[..., -target_len:] + cos_slice = cos.index_select(0, ids) + sin_slice = sin.index_select(0, ids) + return cos_slice, sin_slice + start = 0 + if cp_active and cp_size > 1: + start = cp_rank_index * target_len + if cos.shape[0] < start + target_len or sin.shape[0] < start + target_len: + raise RuntimeError( + "RoPE cache is smaller than the requested context-parallel slice." + ) + cos_slice = cos.narrow(0, start, target_len) + sin_slice = sin.narrow(0, start, target_len) + return cos_slice, sin_slice + + cos, sin = _slice_rope_frequencies(cos, sin) + Q, K = fast_rope_embedding(Q, K, cos, sin) if past_key_value is not None: K = torch.cat([past_key_value[0], K], dim = 2) @@ -784,7 +832,6 @@ def LlamaAttention_fast_forward( attention_mask = attention_mask, causal_mask = causal_mask, ) - A = run_attention(config = config, context = context, Q = Q, K = K, V = V) attn_output = A.reshape(bsz, q_len, n_heads * head_dim) attn_output = self.apply_o(self, attn_output) @@ -1433,6 +1480,7 @@ def CausalLM_fast_forward(fast_forward_inference): past_key_values: Optional[List[torch.FloatTensor]] = None, inputs_embeds: Optional[torch.FloatTensor] = None, labels: Optional[torch.LongTensor] = None, + shift_labels: Optional[torch.LongTensor] = None, use_cache: Optional[bool] = None, output_attentions: Optional[bool] = None, output_hidden_states: Optional[bool] = None, @@ -1499,6 +1547,24 @@ def CausalLM_fast_forward(fast_forward_inference): hidden_states = hidden_states.to(lm_head_device) if labels is not None: labels = labels.to(lm_head_device) + if shift_labels is not None: + shift_labels = shift_labels.to(lm_head_device) + + has_pre_shift_labels = torch.is_tensor(shift_labels) + shift_label_cache: Optional[torch.Tensor] = None + + def _get_shift_labels(): + nonlocal shift_label_cache + if shift_label_cache is not None: + return shift_label_cache + if has_pre_shift_labels: + shift_label_cache = shift_labels + elif labels is not None: + cached = torch.empty_like(labels) + cached[..., :-1] = labels[..., 1:] + cached[..., -1] = -100 + shift_label_cache = cached + return shift_label_cache # Output last hidden states without logits if asked if os.environ.get("UNSLOTH_RETURN_HIDDEN_STATES", "0") == "1": @@ -1541,18 +1607,22 @@ def CausalLM_fast_forward(fast_forward_inference): # num_items_in_batch = n_items, # logit_softcapping = logit_softcapping, # ) + effective_labels = ( + _get_shift_labels() if has_pre_shift_labels else labels + ) loss = unsloth_fused_ce_loss( trainer = None, hidden_states = hidden_states, lm_head_weight = lm_head, lm_head_bias = None, - labels = labels, + labels = effective_labels, mask = None, n_items = n_items, scaling = getattr(self, "accelerator_scaler", None), target_gb = None, torch_compile = True, logit_softcapping = logit_softcapping, + shift_labels = not has_pre_shift_labels, ) if not return_dict: output = (logits,) + outputs[1:] @@ -1582,26 +1652,25 @@ def CausalLM_fast_forward(fast_forward_inference): elif self.config.model_type == "falcon_h1": logit_scaling = self.config.lm_head_multiplier - if labels is not None: + if labels is not None or has_pre_shift_labels: shift_logits = logits - # if not hasattr(self, "extra_ignored_labels"): - # # Fixes https://github.com/unslothai/unsloth/issues/10 - # self.extra_ignored_labels = torch.full((self.max_seq_length, 1), -100, device = "cuda:0") - # pass - shift_labels = torch.empty_like(labels) - shift_labels[..., :-1] = labels[..., 1:] - shift_labels[..., -1] = -100 - mask_packed_sequence_boundaries( - shift_labels, - kwargs.get("packed_seq_lengths"), - ) - # shift_labels = torch.hstack((labels[..., 1:], self.extra_ignored_labels[:labels.shape[0]])) + loss_shift_labels = _get_shift_labels() + # Mask packed sequence boundaries for padding-free/packing modes. + # Skip if has_pre_shift_labels (CP already masked boundaries pre-sharding). + if ( + not has_pre_shift_labels + and kwargs.get("packed_seq_lengths") is not None + ): + mask_packed_sequence_boundaries( + loss_shift_labels, + kwargs.get("packed_seq_lengths"), + ) n_items = kwargs.get("num_items_in_batch", None) if n_items is None: n_items = kwargs.get("n_items", None) loss = fast_cross_entropy_loss( logits = shift_logits, - labels = shift_labels, + labels = loss_shift_labels, logit_softcapping = logit_softcapping, logit_scaling = logit_scaling, n_items = n_items, @@ -1644,6 +1713,7 @@ def PeftModel_fast_forward( attention_mask = None, inputs_embeds = None, labels = None, + shift_labels = None, output_attentions = None, output_hidden_states = None, return_dict = None, @@ -1665,6 +1735,9 @@ def PeftModel_fast_forward( **kwargs, ) else: + # Only pass shift_labels if set (for context parallelism) + if shift_labels is not None: + kwargs["shift_labels"] = shift_labels return self.base_model( input_ids = input_ids, causal_mask = causal_mask, @@ -2165,7 +2238,17 @@ class FastLlamaModel: LlamaForCausalLM.forward = CausalLM_fast_forward( LlamaModel_fast_forward_inference ) + setattr( + LlamaForCausalLM, + "_unsloth_supports_context_parallel_shift_labels", + True, + ) PeftModelForCausalLM.forward = PeftModel_fast_forward + setattr( + PeftModelForCausalLM, + "_unsloth_supports_context_parallel_shift_labels", + True, + ) fix_prepare_inputs_for_generation(LlamaForCausalLM) # Solves https://github.com/unslothai/unsloth/issues/168 diff --git a/unsloth/trainer.py b/unsloth/trainer.py index 65abe6801f..764860ead2 100644 --- a/unsloth/trainer.py +++ b/unsloth/trainer.py @@ -358,24 +358,30 @@ def _patch_sft_trainer_auto_packing(trl_module): processing_class = kwargs.get("processing_class") or kwargs.get("tokenizer") data_collator = kwargs.get("data_collator") - # We also disable vision language models for padding free collators + # Check if context parallelism is enabled + cp_size = getattr(config_arg, "context_parallel_size", 1) or 1 + is_context_parallel = cp_size > 1 + + # Block packing/padding-free for incompatible configurations blocked = ( (data_collator is not None) or isinstance(processing_class, ProcessorMixin) or is_vlm or is_unsupported_model + or is_context_parallel # CP uses ring attention which doesn't support packed masks or ( os.environ.get("UNSLOTH_RETURN_LOGITS", "0") == "1" ) # Disable padding free on forced logits ) requested_pack = bool(getattr(config_arg, "packing", False)) + padding_free_requested = getattr(config_arg, "padding_free", None) is True if blocked: if hasattr(config_arg, "packing"): setattr(config_arg, "packing", False) if hasattr(config_arg, "padding_free"): setattr(config_arg, "padding_free", False) - if blocked and requested_pack: + if blocked and (requested_pack or padding_free_requested): reason = "custom data collator" if data_collator is None and isinstance(processing_class, ProcessorMixin): reason = "processor-based model" @@ -383,7 +389,9 @@ def _patch_sft_trainer_auto_packing(trl_module): reason = "vision-language model" elif is_unsupported_model: reason = f"unsupported model type(s): {', '.join(model_types)}" - message = "Unsloth: Sample packing skipped " f"({reason} detected)." + elif is_context_parallel: + reason = "context parallelism enabled" + message = "Unsloth: Sample packing/padding-free skipped " f"({reason})." print(message) packing_active = False @@ -394,7 +402,6 @@ def _patch_sft_trainer_auto_packing(trl_module): # Resolve padding_free: None (default) = auto-enable unless env-disabled or packing auto_padding_free_active = False - padding_free_requested = getattr(config_arg, "padding_free", None) is True if not blocked: if padding_free_requested: configure_padding_free(config_arg) diff --git a/unsloth/utils/attention_dispatch.py b/unsloth/utils/attention_dispatch.py index 72d52ab376..bb00d0397e 100644 --- a/unsloth/utils/attention_dispatch.py +++ b/unsloth/utils/attention_dispatch.py @@ -22,9 +22,10 @@ from typing import Any, Optional, Tuple import torch from torch import Tensor -from torch.nn.functional import scaled_dot_product_attention +import torch.nn.functional as F from ..models._utils import * +from ..context_parallel import get_cp_manager from ..utils.packing import ( build_sdpa_packed_attention_mask, build_xformers_block_causal_mask, @@ -33,13 +34,14 @@ from ..utils.packing import ( if HAS_FLASH_ATTENTION: from flash_attn import flash_attn_func, flash_attn_varlen_func HAS_XFORMERS = xformers is not None -SDPA_HAS_GQA = "enable_gqa" in (scaled_dot_product_attention.__doc__ or "") +SDPA_HAS_GQA = "enable_gqa" in (F.scaled_dot_product_attention.__doc__ or "") FLASH_VARLEN = "flash_varlen" FLASH_DENSE = "flash_dense" XFORMERS = "xformers" SDPA = "sdpa" +_CP_SDPA_FALLBACK_LOGGED = False XFORMERS_BLOCK_DIAG_CLS = ( xformers.attn_bias.BlockDiagonalCausalMask if HAS_XFORMERS else None @@ -86,6 +88,21 @@ class AttentionContext: def select_attention_backend(use_varlen: bool = False) -> str: """Return attention backend based on availability / priority order.""" + # Context parallelism requires SDPA + # TODO(djsaunde): integrate ring-flash-attn for FA CP support + cp_manager = get_cp_manager() + if cp_manager is not None: + if use_varlen: + raise ValueError( + "Context parallelism does not support varlen/packing mode. " + "Disable packing or set context_parallel_size=1." + ) + global _CP_SDPA_FALLBACK_LOGGED + if not _CP_SDPA_FALLBACK_LOGGED: + print("Unsloth: Context parallelism requires SDPA backend).") + _CP_SDPA_FALLBACK_LOGGED = True + return SDPA + if HAS_FLASH_ATTENTION: if use_varlen: return FLASH_VARLEN @@ -321,7 +338,7 @@ def run_attention( if use_sdpa_gqa: kwargs.setdefault("enable_gqa", True) - out = scaled_dot_product_attention(Q, K, V, **kwargs) + out = F.scaled_dot_product_attention(Q, K, V, **kwargs) return out.transpose(1, 2) K_mod = K @@ -336,7 +353,7 @@ def run_attention( K_mod = K_mod.reshape(bsz, n_heads, kv_seq_len, head_dim) V_mod = V_mod.reshape(bsz, n_heads, kv_seq_len, head_dim) - out = scaled_dot_product_attention( + out = F.scaled_dot_product_attention( Q.contiguous(), K_mod.contiguous(), V_mod.contiguous(),