Trainer.training_step incorrectly normalizes mean token loss when n_gpu > 1
Author: wilswu99
Pull request: no
Labels: name: Good First Issue

name: bug

### System Info

```
- `transformers` version: 4.46.0
- Platform: Linux-5.15.0-136-generic-x86_64-with-glibc2.35
- Python version: 3.10.12
- Huggingface_hub version: 0.29.2
- Safetensors version: 0.5.3
- Accelerate version: 1.4.0
- Accelerate config:    not found
- PyTorch version (GPU?): 2.4.1+cu121 (True)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using distributed or parallel set-up in script?: yes
- Using GPU in script?: yes
- GPU type: NVIDIA RTX A5000
```

### Who can help?

@zach-huggingface @SunMarc @ArthurZucker

### Information

- [ ] The official example scripts
- [x] My own modified scripts

### Tasks

- [x] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [ ] My own task or dataset (give details below)

### Reproduction

Full example setup:

```
config = AutoConfig.from_pretrained('EleutherAI/pythia-14m')
model = GPTNeoXForCausalLM(config=config).to('cuda')
tokenizer = AutoTokenizer.from_pretrained('EleutherAI/pythia-14m')
tokenizer.pad_token = tokenizer.eos_token
train_data = load_dataset("wiwu2390/minipile-100k", split="train")

def tokenize_function(sample):
    return tokenizer(sample["text"], truncation=True, max_length=512)

tokenized_dataset = train_data.map(tokenize_function, batched=True, remove_columns=["text"])

data_collator = DataCollatorForLanguageModeling(
    tokenizer=tokenizer, mlm=False
)

training_args = TrainingArguments(
    output_dir="../data/pythia-14m-minipile-100k",
    num_train_epochs=3,
    per_device_train_batch_size=16,
    per_device_eval_batch_size=16,
    evaluation_strategy="no",
    logging_steps=1,
    save_steps=100,
    learning_rate=1e-3,
    weight_decay=0.01,
    warmup_steps=100,
    fp16=True,
)

trainer = Trainer(
    model=model,
    args=training_args,
    train_dataset=tokenized_dataset,
    tokenizer=tokenizer,
    data_collator=data_collator,
)

trainer.train()
```

With 4 GPUs, the training loss at step 1 is ~2.7. However, the expected value is ~10.8. Indeed, this is what we get if we set CUDA_VISIBLE_DEVICES=0.

### Expected behavior

Since the model is being trained from initialization, the training loss at the first few steps should be around ~log(vocab_size)=10.8. However, when using 4 GPUs, the reported loss is 1/4 of that (2.7).

The reason that this is happening is that the DataParallel-wrapped model gets `num_items_in_batch` as an input kwarg in `Trainer.compute_loss`; this is equal to the number of tokens in the batch (combined across all devices). Each device gets a 1/4-size per-device batch and returns the sum of token losses divided by `num_items_in_batch` (see `transformers.loss.loss_utils.fixed_cross_entropy`). The correct way to aggregate these per-device losses is then to *sum* them. However, `Trainer.training_step` takes the mean:
https://github.com/huggingface/transformers/blob/953196a43dae6a3c474165fba7d215fcbc7b7730/src/transformers/trainer.py#L3759

A quick and dirty fix would be:
```
if self.args.n_gpu > 1:
    loss = loss.mean() if num_items_in_batch is None else loss.sum()
```
I'm not sure if this is compatible with other workflows though.

Comments:
@wiwu2390 thanks for the report. 

> The reason that this is happening is that the DataParallel-wrapped model gets num_items_in_batch as an input kwarg in Trainer.compute_loss; this is equal to the number of tokens in the batch (combined across all devices). Each device gets a 1/4-size per-device batch and returns the sum of token losses divided by num_items_in_batch (see transformers.loss.loss_utils.fixed_cross_entropy). 

`num_items_in_batch` shouldn't be equal to the number of tokens in the batch combined across all devices but only on the respective device. We only combine if you set `average_tokens_across_devices=True`. However, the default for this arg is False.

Did you test both runs with the same `per_device_train_batch_size` ? You need to divide it by 4 when running with 4 gpus to be comparable. Otherwise, you are actually using a bigger global batch size. However, I don't think this is the real issue.

@SunMarc Thanks for taking a look at this. (And apologies for the delayed reply -- somehow I didn't get a notification about your response.)

I am dividing the `per_device_train_batch_size` by 4 for 4 gpus, but as you said I don't think this is the root cause.
Also I have the default setting `average_tokens_across_devices=False`.

When I run my above example code, the behavior I'm observing is:
- `Trainer.compute_loss` is called with `num_items_in_batch=(inputs['labels'] != -100).sum()`
- Inside `compute_loss`, it runs `outputs = model(**inputs); loss = outputs["loss"]`
- In my example, `model` is a `DataParallel` object wrapping a `GPTNeoXForCausalLM`. I don't construct a `DataParallel` explicitly, so I'm guessing that `Trainer` automatically does this somewhere when it detects multiple devices?
- `DataParallel` splits  `input_ids` into 4 chunks for the 4 gpus, **but it does not modify `num_items_in_batch`**
- Thus `GPTNeoXForCausalLM.forward` calls `self.loss_function` using `input_ids` for just 1 gpu but `num_items_in_batch` equal to the sum across all gpus.

Are you able to replicate this behavior on your end? I've provided the full example code (except imports) above.

Also, maybe I'm incorrectly using `Trainer` for multi-device training? Is there a different config that I should be using?

I'll check ! Let me know if you find something on your side 

This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/main/CONTRIBUTING.md) are likely to be ignored.

Hi! I'd like to work on this, is it still available? Thanks.

This PR addresses issue #37474 by introducing a dedicated _reduce_loss method in Trainer to correctly normalize losses across single and multi-GPU setups. Unlike a one-line patch, this solution provides a clean, reusable design with full test coverage, ensuring accurate token-level averaging and long-term maintainability.

> This PR addresses issue [#37474](https://github.com/huggingface/transformers/issues/37474) by introducing a dedicated _reduce_loss method in Trainer to correctly normalize losses across single and multi-GPU setups. Unlike a one-line patch, this solution provides a clean, reusable design with full test coverage, ensuring accurate token-level averaging and long-term maintainability.

Hi @moonrunnerkc !
I will let @SunMarc judge, but in my opinion the PR you linked does not fix the issue raised by @SunMarc above:

> num_items_in_batch shouldn't be equal to the number of tokens in the batch combined across all devices but only on the respective device

Your approach is very similar to the change already proposed to wiwu in https://github.com/huggingface/transformers/pull/37475: the value of the loss is fixed and thus becomes correct, but the fundamental issue of num_items_in_batch being wrong when n_gpus > 1 remains and could cause other silence fails down the road. On the other hand I think that the PR I put together last week: https://github.com/huggingface/transformers/pull/40747/files addresses the fundamental issue as `self.training_step` now sees the correct number of tokens!



> > This PR addresses issue [#37474](https://github.com/huggingface/transformers/issues/37474) by introducing a dedicated _reduce_loss method in Trainer to correctly normalize losses across single and multi-GPU setups. Unlike a one-line patch, this solution provides a clean, reusable design with full test coverage, ensuring accurate token-level averaging and long-term maintainability.
> 
> Hi [@moonrunnerkc](https://github.com/moonrunnerkc) ! I will let [@SunMarc](https://github.com/SunMarc) judge, but in my opinion the PR you linked does not fix the issue raised by [@SunMarc](https://github.com/SunMarc) above:
> 
> > num_items_in_batch shouldn't be equal to the number of tokens in the batch combined across all devices but only on the respective device
> 
> Your approach is very similar to the change already proposed to wiwu in [#37475](https://github.com/huggingface/transformers/pull/37475): the value of the loss is fixed and thus becomes correct, but the fundamental issue of num_items_in_batch being wrong when n_gpus > 1 remains and could cause other silence fails down the road. On the other hand I think that the PR I put together last week: https://github.com/huggingface/transformers/pull/40747/files addresses the fundamental issue as `self.training_step` now sees the correct number of tokens!

That makes sense. My implementation fixed the loss output but didn’t address the underlying num_items_in_batch bug.

 I’ll review your PR and see if I can adapt my work to build on the corrected token handling instead.

I had a look at this issue and I've opened a PR to fix this ! Please test it 

Hi! I’m interested in working on this. Is it still available?

The issue is fixed, so i'm closing this !