[Trainer] accelerate contextparallel support in trainer

src/transformers/trainer.py

@@ -778,6 +778,7 @@ def __init__(
             stateful_callbacks=[
                 cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
             ],
+            parallelism_config=getattr(self, "parallelism_config", None),
         )
         # Internal variable to count flos in each process, will be accumulated in `self.state.total_flos` then
         # returned to 0 every time flos need to be logged
@@ -2424,7 +2425,8 @@ def _inner_training_loop(
         self.state = TrainerState(
             stateful_callbacks=[
                 cb for cb in self.callback_handler.callbacks + [self.control] if isinstance(cb, ExportableState)
-            ]
+            ],
+            parallelism_config=self.parallelism_config,
         )
         self.state.is_hyper_param_search = trial is not None
         self.state.train_batch_size = self._train_batch_size
@@ -3870,7 +3872,34 @@ def training_step(
             return loss_mb.reduce_mean().detach().to(self.args.device)
 
         with self.compute_loss_context_manager():
-            loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
+            # Wrap forward pass with context parallelism if enabled
+            if getattr(self, "is_cp_enabled", False):
+                # Prepare buffers for context parallelism
+                buffers = []
+                buffer_seq_dims = []
+
+                # position_ids are already provided by the data collator, no need to create them manually
+
+                if "input_ids" in inputs:
+                    buffers.append(inputs["input_ids"])
+                    buffer_seq_dims.append(1)  # Sequence dimension
+                if "labels" in inputs:
+                    buffers.append(inputs["labels"])
+                    buffer_seq_dims.append(1)
+                if "attention_mask" in inputs:
+                    buffers.append(inputs["attention_mask"])
+                    buffer_seq_dims.append(1)
+                # Include position_ids in context parallelism splitting (key for Flash Attention compatibility)
+                if "position_ids" in inputs and inputs["position_ids"] is not None:
+                    buffers.append(inputs["position_ids"])
+                    buffer_seq_dims.append(1)
+
+                with self.accelerator.maybe_context_parallel(
+                    buffers=buffers, buffer_seq_dims=buffer_seq_dims, no_restore_buffers=set(buffers)
+                ):
+                    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
+            else:
+                loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
 
         del inputs
         if (
@@ -4150,6 +4179,11 @@ def _save(self, output_dir: Optional[str] = None, state_dict=None):
         os.makedirs(output_dir, exist_ok=True)
         logger.info(f"Saving model checkpoint to {output_dir}")
 
+        # Handle context parallelism state dict properly
+        if state_dict is None and self.is_cp_enabled:
+            # Use accelerator.get_state_dict() for proper gathering when context parallelism is enabled
+            state_dict = self.accelerator.get_state_dict(self.model)
+
         supported_classes = (PreTrainedModel,) if not is_peft_available() else (PreTrainedModel, PeftModel)
         # Save a trained model and configuration using `save_pretrained()`.
         # They can then be reloaded using `from_pretrained()`
@@ -5324,6 +5358,14 @@ def create_accelerator_and_postprocess(self):
         args = {
             "deepspeed_plugin": self.args.deepspeed_plugin,
         }
+
+        # Add parallelism_config if context parallelism is enabled
+        if getattr(self.args, "parallelism_config", None) is not None:
+            args["parallelism_config"] = self.args.parallelism_config
+
+        # Add fsdp_plugin if it was created for context parallelism
+        if getattr(self.args, "fsdp_plugin", None) is not None:
+            args["fsdp_plugin"] = self.args.fsdp_plugin
         if is_accelerate_available("0.28.0"):
             args["dataloader_config"] = dataloader_config
         else:
@@ -5351,6 +5393,11 @@ def create_accelerator_and_postprocess(self):
         self.is_deepspeed_enabled = getattr(self.accelerator.state, "deepspeed_plugin", None) is not None
         self.is_fsdp_enabled = getattr(self.accelerator.state, "fsdp_plugin", None) is not None
         self.is_tp_enabled = getattr(self.accelerator.state, "torch_tp_plugin", None) is not None
+        self.parallelism_config = getattr(self.accelerator.state, "parallelism_config", None)
+        self.is_cp_enabled = self.parallelism_config and self.parallelism_config.cp_enabled
+        # Update state with parallelism config for callbacks and logging (if state exists)
+        if hasattr(self, "state") and self.state is not None:
+            self.state.parallelism_config = self.parallelism_config
         # post accelerator creation setup
         if self.is_fsdp_enabled:
             fsdp_plugin = self.accelerator.state.fsdp_plugin
@@ -5457,7 +5504,22 @@ def get_batch_samples(
                 pass
 
         if num_items_in_batch is not None:
-            if self.args.average_tokens_across_devices:
+            # Handle context parallelism token counting
+            if self.is_cp_enabled:
+                # For context parallelism, use AVG instead of SUM to avoid over-accounting tokens
+                import torch.distributed as dist
+
+                if dist.is_initialized():
+                    # Convert to tensor for all_reduce
+                    if not torch.is_tensor(num_items_in_batch):
+                        num_items_in_batch = torch.tensor(num_items_in_batch, dtype=torch.long)
+                    num_items_in_batch = num_items_in_batch.to(device)
+
+                    # All-reduce with AVG across context parallel group
+                    # Note: We use all processes here as we don't have easy access to CP group
+                    dist.all_reduce(num_items_in_batch, op=dist.ReduceOp.AVG)
+                    num_items_in_batch = int(num_items_in_batch.item())
+            elif self.args.average_tokens_across_devices:
                 num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum()
 
             if torch.is_tensor(num_items_in_batch):

src/transformers/trainer_callback.py

@@ -19,9 +19,13 @@
 import json
 import math
 from dataclasses import dataclass
-from typing import Optional, Union
+from typing import TYPE_CHECKING, Optional, Union
 
 import numpy as np
+
+
+if TYPE_CHECKING:
+    from accelerate import ParallelismConfig
 from tqdm.auto import tqdm
 
 from .trainer_utils import HPSearchBackend, IntervalStrategy, SaveStrategy, has_length
@@ -91,6 +95,8 @@ class TrainerState:
         stateful_callbacks (`list[StatefulTrainerCallback]`, *optional*):
             Callbacks attached to the `Trainer` that should have their states be saved or restored.
             Relevant callbacks should implement a `state` and `from_state` function.
+        parallelism_config (`ParallelismConfig`, *optional*):
+            Configuration for parallel training modes including context parallelism.
     """
 
     epoch: Optional[float] = None
@@ -113,6 +119,7 @@ class TrainerState:
     trial_name: Optional[str] = None
     trial_params: dict[str, Union[str, float, int, bool]] = None
     stateful_callbacks: list["TrainerCallback"] = None
+    parallelism_config: Optional["ParallelismConfig"] = None
 
     def __post_init__(self):
         if self.log_history is None:

src/transformers/training_args.py

@@ -1282,6 +1282,16 @@ class TrainingArguments:
     label_smoothing_factor: float = field(
         default=0.0, metadata={"help": "The label smoothing epsilon to apply (zero means no label smoothing)."}
     )
+    context_parallel_size: int = field(
+        default=1,
+        metadata={
+            "help": (
+                "Context Parallelism size for splitting long sequences across devices. "
+                "Requires FSDP v2 and accelerate >= 1.10.0. Default is 1 (disabled). "
+                "When enabled, sequences are split across multiple GPUs to handle longer contexts."
+            )
+        },
+    )
 
     default_optim = "adamw_torch"
     if is_torch_available():
@@ -2070,6 +2080,48 @@ def __post_init__(self):
             self.deepspeed_plugin.set_mixed_precision(mixed_precision)
             self.deepspeed_plugin.set_deepspeed_weakref()
 
+        # Initialize parallelism_config and fsdp_plugin for context parallelism
+        self.parallelism_config = None
+        self.fsdp_plugin = None
+
+        # Handle context parallelism setup
+        if self.context_parallel_size > 1:
+            from accelerate import __version__ as accelerate_version
+            from packaging import version
+
+            if version.parse(accelerate_version) < version.parse("1.10.0"):
+                raise ValueError(f"Context parallelism requires accelerate >= 1.10.0, but found {accelerate_version}")
+
+            from accelerate.parallelism_config import ParallelismConfig
+
+            # Create ParallelismConfig
+            self.parallelism_config = ParallelismConfig(cp_size=self.context_parallel_size)
+
+            # Context parallelism requires FSDP v2
+            # Only create if not already using FSDP
+            if not self.fsdp:
+                from accelerate.utils import FullyShardedDataParallelPlugin
+
+                self.fsdp_plugin = FullyShardedDataParallelPlugin(
+                    fsdp_version=2,
+                    auto_wrap_policy="transformer_based_wrap",
+                    state_dict_type="FULL_STATE_DICT",
+                )
+            else:
+                # Ensure FSDP v2 is used when context parallelism is enabled
+                if self.fsdp_config.get("version", 1) != 2:
+                    logger.warning("Context parallelism requires FSDP v2. Updating FSDP config to use version 2.")
+                    self.fsdp_config["version"] = 2
+
+            # Validation for context parallelism
+            if self.per_device_train_batch_size != 1:
+                logger.warning(
+                    f"Context parallelism typically requires batch_size=1, "
+                    f"but got {self.per_device_train_batch_size}. "
+                    f"Setting per_device_train_batch_size=1."
+                )
+                self.per_device_train_batch_size = 1
+
         if self.use_cpu:
             self.dataloader_pin_memory = False