Repository: huggingface/transformers
Pull request: Make it possible to evaluate when using sequence parallel in HF Trainer

# What does this PR do?

When using that sp(sequence parallel) in evaluate, the sequence length must be divisible by `world_size`, and the eval batch size must be the same as the train batch size to be able to use it, and this is a PR to modify that.
When using deepspeed sp, due to monkey patching the attention function, the following problems occur.

1. Eval batch size must be the same as train batch size.
2. During eval, the sequence must always be divisible by world_size.

In this situation, I propose to add the following two things in this PR.

1. A callback that restores the monkey-patched attention function to the original function during evaluate, and then reapplies it when training starts.
2. Check if it's training or not with self.model.training in the compute_loss section.

I think it should be made so that it can work without sequence parallel during evaluate.

However, some discussion is needed regarding these changes.
I'm not sure if when you first designed sequence parallel, you only considered it for training, or if you designed it to work together in evaluate as well.
Looking at the current code, it seems like it was made assuming it would be used in training, so I modified the code so that it can work normally in evaluate.

---

- `transformers` version: 5.0.0.dev0
- Platform: Linux-5.15.0-153-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 1.3.3
- Safetensors version: 0.7.0
- Accelerate version: 1.12.0
- Accelerate config:    not found
- DeepSpeed version: 0.18.4
- PyTorch version (accelerator?): 2.9.1+cu128 (CUDA)
- Using distributed or parallel set-up in script?: <fill in>
- Using GPU in script?: <fill in>
- GPU type: NVIDIA H100 NVL

I'm using the dev version via `pip install git+https://github.com/huggingface/transformers.git@dfe30827b8ebdd974eb7ce69c7d5d8cf8e6cf852`

## Before submitting

- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md#create-a-pull-request),
      Pull Request section?
- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
      to it if that's the case.
- [ ] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/main/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/main/docs#writing-source-documentation).
- [ ] Did you write any new necessary tests?

## Who can review?

@SunMarc


Changed files:
filename: src/transformers/trainer.py
patch: @@ -3837,7 +3837,7 @@ def compute_loss(
         make sure to overwrite `self.model_accepts_loss_kwargs` to `False`. Otherwise, the loss calculating might be slightly inaccurate when performing gradient accumulation.
         """
         pc = getattr(self.accelerator, "parallelism_config", None)
-        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled:
+        if pc is not None and pc.sp_backend == "deepspeed" and pc.sp_enabled and self.model.training:
             return self._deepspeed_sp_compute_loss(model, inputs, return_outputs, pc)
 
         if (self.label_smoother is not None or self.compute_loss_func is not None) and "labels" in inputs:

filename: tests/deepspeed/test_alst_ulysses_sp.py
patch: @@ -75,6 +75,7 @@ def test_sp_equivalence(self):
                    """,
         )
 
+        sp_yes_eval_metrics_path = sp_yes_output_dir / "sp_yes_eval_metrics.json"
         cmd_sp = f"""
             accelerate launch
             --config_file {sp_yes_accelerate_config_path}
@@ -87,7 +88,9 @@ def test_sp_equivalence(self):
             --logging_steps 1
             --remove_unused_columns False
             --seed 42
+            --per_device_eval_batch_size 1
             --loss_output_file {sp_yes_losses_path}
+            --eval_output_file {sp_yes_eval_metrics_path}
         """.split()
 
         execute_subprocess_async(cmd_sp, env=self.get_env())
@@ -147,21 +150,35 @@ def test_sp_equivalence(self):
             msg=f"SP-enabled losses {sp_yes_losses} do not match SP-disabled losses {sp_no_losses}",
         )
 
+        # Eval should succeed even though eval sequences are not divisible by sp_size,
+        # because SP is disabled in eval mode.
+        eval_metrics = read_json_file(sp_yes_eval_metrics_path)
+        assert "eval_loss" in eval_metrics, f"Expected eval_loss in metrics, got keys: {list(eval_metrics.keys())}"
+        assert torch.isfinite(torch.tensor(eval_metrics["eval_loss"])), f"Eval loss is not finite: {eval_metrics}"
+
 
 if __name__ == "__main__":
     model_name = "hf-internal-testing/tiny-random-LlamaForCausalLM"
 
     # Parse custom arguments (not TrainingArguments parameters)
     loss_output_file = None
+    eval_output_file = None
 
     if "--loss_output_file" in sys.argv:
         idx = sys.argv.index("--loss_output_file")
         loss_output_file = sys.argv[idx + 1]
         sys.argv.pop(idx)
         sys.argv.pop(idx)
 
+    if "--eval_output_file" in sys.argv:
+        idx = sys.argv.index("--eval_output_file")
+        eval_output_file = sys.argv[idx + 1]
+        sys.argv.pop(idx)
+        sys.argv.pop(idx)
+
     parser = HfArgumentParser((TrainingArguments,))
     training_args = parser.parse_args_into_dataclasses()[0]
+    training_args.do_eval = True
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     if tokenizer.pad_token is None:
@@ -183,7 +200,13 @@ def test_sp_equivalence(self):
     def tokenize_function(examples):
         return tokenizer(examples, max_length=128, truncation=True, padding="max_length")
 
+    def tokenize_eval_function(examples):
+        # Use an odd max_length to ensure it's not divisible by sp_size (2),
+        # so eval would fail if SP isn't disabled in eval mode.
+        return tokenizer(examples, max_length=127, truncation=True, padding="max_length")
+
     train_dataset = [tokenize_function(text) for text in texts]
+    eval_dataset = [tokenize_eval_function(text) for text in texts]
 
     # Use standard DataCollatorForLanguageModeling for causal LM
     # pad_to_multiple_of=4 ensures sequences are divisible by sp_size * 2 (for sp_size=2)
@@ -198,6 +221,7 @@ def tokenize_function(examples):
         model=model,
         args=training_args,
         train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
         data_collator=data_collator,
     )
 
@@ -207,6 +231,12 @@ def tokenize_function(examples):
     # Verify training completed
     assert trainer.state.global_step > 0, "Training should have completed at least one step"
 
+    if training_args.do_eval:
+        eval_metrics = trainer.evaluate()
+        if eval_output_file and training_args.process_index == 0:
+            with open(eval_output_file, "w") as f:
+                json.dump(eval_metrics, f)
+
     # Save losses to file if requested (for equivalence testing)
     if loss_output_file and training_args.process_index == 0:
         losses = [log["loss"] for log in trainer.state.log_history if "loss" in log]