Potentially incorrect calculation of `total_updates` on >=4.46.0 since #34198 affecting multi gpu training
Author: chiragjn
Pull request: no
Labels: name: trainer

name: bug

### System Info

Okay I have been pulling my hair out for a few hours and turns out this bug only happens when `average_tokens_across_devices` is True and epochs > 1

Simplest case to reproduce

DDP
**world size 2**
dataset length = 4
**epochs = 2**
micro batch size  = 1 (aka per gpu batch size)
gradient accumulation = 1
**average_tokens_across_devices = True**

So every epoch, total 2 steps on both devices

but as first epoch finishes, we get

```
[rank1]:   File "/home/jovyan/axolotl/src/axolotl/train.py", line 191, in train
[rank1]:     trainer.train(resume_from_checkpoint=resume_from_checkpoint)
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/transformers/trainer.py", line 2164, in train
[rank1]:     return inner_training_loop(
[rank1]:            ^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/transformers/trainer.py", line 2473, in _inner_training_loop
[rank1]:     batch_samples, num_items_in_batch = self.get_batch_samples(epoch_iterator, num_batches)
[rank1]:                                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/transformers/trainer.py", line 5142, in get_batch_samples
[rank1]:     num_items_in_batch = self.accelerator.gather(num_items_in_batch).sum().item()
[rank1]:                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/accelerate/accelerator.py", line 2459, in gather
[rank1]:     return gather(tensor)
[rank1]:            ^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/accelerate/utils/operations.py", line 376, in wrapper
[rank1]:     return function(*args, **kwargs)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/accelerate/utils/operations.py", line 437, in gather
[rank1]:     return _gpu_gather(tensor)
[rank1]:            ^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/accelerate/utils/operations.py", line 356, in _gpu_gather
[rank1]:     return recursively_apply(_gpu_gather_one, tensor, error_on_other_type=True)
[rank1]:            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
[rank1]:   File "/home/jovyan/.conda/envs/jupyter-base/lib/python3.11/site-packages/accelerate/utils/operations.py", line 129, in recursively_apply
[rank1]:     raise TypeError(
[rank1]: TypeError: Unsupported types (<class 'NoneType'>) passed to `_gpu_gather_one`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.
```

The main culprit here is 
https://github.com/huggingface/transformers/blob/8f38f58f3de5a35f9b8505e9b48985dce5470985/src/transformers/trainer.py#L2468-L2473

`steps_in_epoch` per rank is correctly calculated as 2 but total updates is 3
Normally that is harmless because dataloader would be exhausted and would result in empty batch and it won't enter the loop on 2473.
However, when using the recently added option `average_tokens_across_devices`, it will try to gather number of total items in batches across all ranks and gather doesn't like broadcasting `None`

https://github.com/huggingface/transformers/blob/8f38f58f3de5a35f9b8505e9b48985dce5470985/src/transformers/trainer.py#L5139-L5156

This problem does not surface with 1 gpu because `average_tokens_across_devices` is auto set to `False` and neither under epoch = 1 because `DefaultFlowCallback` stops the training process considering global step and expected max steps

### Who can help?

@muellerzr

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

Train any LM on more than one gpu, set at least average_tokens_across_devices = True and epochs  > 1

### Expected behavior

Either we fix `total_updates` count or we handle `None` for gather

Comments:
Thanks for the detailed report @chiragjn ! I think the simplest solution would be to set num_items_in_batch as a tensor of 0 when the value is None since it means that `batch_samples` doesn't have any element. Can you check if this solves the issue ? 
```python
if self.args.average_tokens_across_devices:
    if num_items_in_batch is None:
        num_items_in_batch = torch.tensor(0)
```

Fixed by #35102