CP support fused CE loss

unsloth-cli.py

@@ -45,6 +45,13 @@ def run(args):
 
     logging.getLogger("hf-to-gguf").setLevel(logging.WARNING)
 
+    # Convert gradient checkpointing arg to expected type
+    use_gradient_checkpointing = args.use_gradient_checkpointing
+    if use_gradient_checkpointing.lower() == "true":
+        use_gradient_checkpointing = True
+    elif use_gradient_checkpointing.lower() == "false":
+        use_gradient_checkpointing = False
+
     # Load model and tokenizer
     device_map, distributed = prepare_device_map()
     model, tokenizer = FastLanguageModel.from_pretrained(
@@ -53,6 +60,7 @@ def run(args):
         dtype = args.dtype,
         load_in_4bit = args.load_in_4bit,
         device_map = device_map,
+        use_gradient_checkpointing = use_gradient_checkpointing,
     )
 
     # Configure PEFT model
@@ -71,7 +79,7 @@ def run(args):
         lora_alpha = args.lora_alpha,
         lora_dropout = args.lora_dropout,
         bias = args.bias,
-        use_gradient_checkpointing = args.use_gradient_checkpointing,
+        use_gradient_checkpointing = use_gradient_checkpointing,
         random_state = args.random_state,
         use_rslora = args.use_rslora,
         loftq_config = args.loftq_config,
@@ -133,6 +141,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 +165,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 +375,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(

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

unsloth/context_parallel.py

@@ -0,0 +1,527 @@
+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
+    for name, module in model.named_modules():
+        if name.endswith("self_attn"):
+            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 get_global_tokens(self, inputs) -> torch.Tensor:
+        """Compute global token count across CP group."""
+        shift_labels = inputs.get("shift_labels")
+        if shift_labels is None:
+            return None
+
+        local_tokens = shift_labels.ne(-100).sum().float()
+        global_tokens = local_tokens.clone()
+        dist.all_reduce(global_tokens, op = dist.ReduceOp.SUM, group = self._cp_group)
+        return global_tokens
+
+    def reduce_loss(self, loss):
+        """Reduce loss across CP group for reporting."""
+        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, _ = loss[0], loss[1:]
+        else:
+            tensor = loss
+
+        # Sum across ranks and divide by cp_size for correct mean.
+        global_loss = tensor.detach().clone()
+        dist.all_reduce(global_loss, op = dist.ReduceOp.SUM, group = self._cp_group)
+        global_loss = global_loss / self.settings.size
+        self._set_report_loss(global_loss)
+
+        return loss
+
+
+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_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
+            elif not dist.is_initialized():
+                raise RuntimeError(
+                    f"Context parallelism requires torch.distributed to be initialized. "
+                    f"Use torchrun or accelerate launch with {settings.size} processes."
+                )
+            elif dist.get_world_size() < settings.size:
+                raise RuntimeError(
+                    f"Context parallelism size ({settings.size}) exceeds world size "
+                    f"({dist.get_world_size()}). Launch with at least {settings.size} processes."
+                )
+            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)
+
+        # Enable sync_each_batch when using CP to ensure consistent computation graph
+        # for DDP + static_graph mode. This is needed because ring attention changes
+        # the graph structure, and sync_each_batch keeps it constant.
+        if (
+            manager
+            and accelerator is not None
+            and hasattr(accelerator, "gradient_state")
+        ):
+            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)
+
+        # Check if we're in CP mode with pre-shifted labels
+        shift_labels = inputs.get("shift_labels")
+        use_cp = manager and isinstance(shift_labels, torch.Tensor)
+
+        if use_cp:
+            global_tokens = manager.get_global_tokens(inputs)
+
+            # For CP, pass both labels and shift_labels to the model.
+            # The model's forward (e.g., llama.py) detects shift_labels and uses
+            # unsloth_fused_ce_loss with the pre-shifted labels.
+            # We keep labels so the fused CE loss path is triggered (requires labels is not None).
+
+            # Scale by GA and cp_size to match gradient behavior of CP=1.
+            ga_steps = getattr(self.args, "gradient_accumulation_steps", 1)
+            cp_size = manager.settings.size
+            inputs["num_items_in_batch"] = global_tokens * ga_steps / cp_size
+
+            outputs = model(**inputs)
+            loss = outputs.loss if hasattr(outputs, "loss") else outputs[0]
+
+            # Clean up temporary input
+            inputs.pop("num_items_in_batch", None)
+
+            if return_outputs:
+                loss = (loss, outputs)
+        else:
+            # No CP - use original path
+            loss = original_compute_loss(
+                self,
+                model,
+                inputs,
+                return_outputs = return_outputs,
+                **kwargs,
+            )
+
+        if manager:
+            loss = manager.reduce_loss(loss)
+        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
+
+    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.__unsloth_context_parallel__ = True
+    if original_get_train_sampler is not None:
+        trainer_cls._get_train_sampler = _patch_train_sampler(
+            original_get_train_sampler
+        )

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

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)

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