Sequence Parallelism Disabled for Evaluation
Author: amanichopra
Pull request: no

Hi, I am training a model with sequence parllelism as per these [docs](qwen 3.6 27b). I am run evals evals every 50 training steps, and encountered an OOM during my first eval step:

```

2026-06-24T08:21:11.619Z | 2%\|▏ \| 50/2064 [1:12:45<49:41:22, 88.82s/it]
-- | --
  | 2026-06-24T08:21:11.619Z | {'loss': '0.6005', 'grad_norm': '0.2946', 'learning_rate': '4.712e-06', 'epoch': '0.02422'}
  | 2026-06-24T08:21:11.619Z | 2%\|▏ \| 50/2064 [1:12:45<49:41:22, 88.82s/it]
  | 2026-06-24T08:21:51.626Z | [rank4]: Traceback (most recent call last):
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/ml/code/cpt.py", line 561, in <module>
  | 2026-06-24T08:21:51.626Z | [rank4]: main()
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/ml/code/cpt.py", line 555, in main
  | 2026-06-24T08:21:51.626Z | [rank4]: train_function(model_args, script_args, training_args)
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/ml/code/cpt.py", line 357, in train_function
  | 2026-06-24T08:21:51.626Z | [rank4]: train_result = trainer.train(resume_from_checkpoint=last_checkpoint)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 1427, in train
  | 2026-06-24T08:21:51.626Z | [rank4]: return inner_training_loop(
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 1509, in _inner_training_loop
  | 2026-06-24T08:21:51.626Z | [rank4]: self._run_epoch(
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 1778, in _run_epoch
  | 2026-06-24T08:21:51.626Z | [rank4]: self._maybe_log_save_evaluate(
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2084, in _maybe_log_save_evaluate
  | 2026-06-24T08:21:51.626Z | [rank4]: metrics = self._evaluate(trial, ignore_keys_for_eval)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2995, in _evaluate
  | 2026-06-24T08:21:51.626Z | [rank4]: metrics = self.evaluate(ignore_keys=ignore_keys_for_eval)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2565, in evaluate
  | 2026-06-24T08:21:51.626Z | [rank4]: dataset_metrics = self.evaluate(
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2582, in evaluate
  | 2026-06-24T08:21:51.626Z | [rank4]: output = self.evaluation_loop(
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2707, in evaluation_loop
  | 2026-06-24T08:21:51.626Z | [rank4]: losses, logits, labels = self.prediction_step(model, inputs, prediction_loss_only, ignore_keys=ignore_keys)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 2961, in prediction_step
  | 2026-06-24T08:21:51.626Z | [rank4]: loss, outputs = self.compute_loss(
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/trainer.py", line 1981, in compute_loss
  | 2026-06-24T08:21:51.626Z | [rank4]: outputs = model(**inputs)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1775, in _wrapped_call_impl
  | 2026-06-24T08:21:51.626Z | [rank4]: return self._call_impl(*args, **kwargs)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1881, in _call_impl
  | 2026-06-24T08:21:51.626Z | [rank4]: return inner()
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/usr/local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1829, in inner
  | 2026-06-24T08:21:51.626Z | [rank4]: result = forward_call(*args, **kwargs)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/utils/generic.py", line 900, in wrapper
  | 2026-06-24T08:21:51.626Z | [rank4]: output = func(self, *args, **kwargs)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/models/qwen3_5/modeling_qwen3_5.py", line 1770, in forward
  | 2026-06-24T08:21:51.626Z | [rank4]: loss = self.loss_function(logits=logits, labels=labels, vocab_size=self.config.vocab_size, **kwargs)
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: File "/opt/venv/lib/python3.12/site-packages/transformers/loss/loss_utils.py", line 56, in ForCausalLMLoss
  | 2026-06-24T08:21:51.626Z | [rank4]: logits = logits.float()
  | 2026-06-24T08:21:51.626Z | [rank4]: ^^^^^^^^^^^^^^
  | 2026-06-24T08:21:51.626Z | [rank4]: torch.OutOfMemoryError: CUDA out of memory. Tried to allocate 74.00 GiB. GPU 4 has a total capacity of 79.19 GiB of which 22.50 GiB is free. Including non-PyTorch memory, this process has 56.68 GiB memory in use. Of the allocated memory 49.90 GiB is allocated by PyTorch, and 157.45 MiB is reserved by PyTorch but unallocated. If reserved but unallocated memory is large try setting PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True to avoid fragmentation. See documentation for Memory Management (https://pytorch.org/docs/stable/notes/cuda.html#environment-variables)
  | 2026-06-24T08:21:52.627Z | [rank6]: Traceback
```

Given that this only happened during evals and not training, it makes me think something is strange with how SP is wired. I was looking at the trainer source code and came across [this](https://github.com/huggingface/transformers/blame/50b4cc1882b95f2a238042bcf251de11e04421e3/src/transformers/trainer.py#L1975. ) line which suggests that SP is disabled for evals. In deesppeed, I also see they have a comment about disabling SP during evals specifically for HF transformers: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L82. 

Not exactly sure how to reconcile all this ifnormation, but I want to understand: **Is SP activate during eval?** If not, then there is no sharding over the input sequence which ofc will result in an OOM if we bump up the sequence length.. So this seems to defeat the purpose of SP in the first place. I understand that we can probably disable evals entirely to bypass this issue.. but that doesn't seem like the right solution. Would appreciate any advance/insights. Thanks!

Comments:
@amanichopra  hey can you provide repro for this

@Prachi-kushwaha sorry missed your comment earlier. 

I'm  training like:

```
# training config
max_seq_length: 100000  
sp_size: 8              # must equal number of GPUs per node
per_device_eval_batch_size: 1
eval_strategy: steps
eval_steps: 50
```

```
# training script
from accelerate.utils import ParallelismConfig, DeepSpeedSequenceParallelConfig

parallelism_config = ParallelismConfig(
    sp_backend="deepspeed",
    sp_size=8,
    dp_shard_size=1,
    sp_handler=DeepSpeedSequenceParallelConfig(
        sp_seq_length_is_variable=True,
    ),
)
training_args.parallelism_config = parallelism_config

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=train_dataset,
    eval_dataset=eval_dataset,  # required to trigger the bug
    ...
)
trainer.train()
```

I'm getting OOM at the first eval step regardless of how long the sequences actually are because it seems like SP is not applied during eval. In [trainer.py](https://github.com/huggingface/transformers/blob/main/src/transformers/trainer.py#L1981), `compute_loss` dispatches to `deepspeed_sp_compute_loss` only when `self.model.training` is `True`. Also, as I said above [deepspeed/runtime/sequence_parallel/ulysses_sp.py](https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/runtime/sequence_parallel/ulysses_sp.py#L262) has a `disable_in_eval` flag that is meant to bypass SP at eval time...  In fact, there's a comment stating: 

```
# Skip SP operations during eval if disable_in_eval is True
# This avoids issues with frameworks like HF Trainer that may run eval with different data distribution
```

Not exactly sure what they mean by data distribution but I imagine they might have said this because huggingface decide to not to apply sp to the eval dataloader? And if they still applied SP, then if rank 0 and rank 1 have completely different sequences from the eval set, then the all-to-all would still run. This means the two sequences would b concatenated together and attention would be computed as if there was just 1 long seq... 

So, I have a solution just to create a subclass of `Trainer` and override the `get_eval_dataloader`. With this fix, I'm not getting OOM during evals anymore... 

```
from transformers.integrations.deepspeed import deepspeed_sp_compute_loss

class SPAwareTrainer(Trainer):
  def get_eval_dataloader(self, eval_dataset=None):
      dataloader = super().get_eval_dataloader(eval_dataset)
      pc = getattr(self.accelerator, "parallelism_config", None)
      if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
          dataloader = self.accelerator.deepspeed_ulysses_dl_adapter(dataloader, self.model_wrapped)
      return dataloader

  def compute_loss(self, model, inputs, return_outputs=False, num_items_in_batch=None):
      pc = getattr(self.accelerator, "parallelism_config", None)
      if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
          return deepspeed_sp_compute_loss(self.accelerator, model, inputs, return_outputs, pc)
      return super().compute_loss(model, inputs, return_outputs=return_outputs, num_items_in_batch=num_items_in_batch)
```

I just don't understanding why HF didn't do this to begin with??