Fix gather when collecting 'num_input_tokens_seen'

[CodeCreator]

What does this PR do?

Following from #29099, the allgather for num_input_tokens_seen is still getting stuck with distributed training, since the tensors are still on CPU and have not been moved to device in _prepare_inputs yet. This still results in torch.distributed errors such as:

File ".../site-packages/torch/distributed/distributed_c10d.py", line 2948, in all_gather_into_tensor
     work = group._allgather_base(output_tensor, input_tensor, opts)
ValueError: Tensors must be CUDA and dense

This PR simply moves the token count to self.args.device and then moves them back to cpu after gathering.
This problem was also mentioned in the discussion of issue #28791, but the issue was closed when the padding was fixed.

Who can review?

@pacman100
@muellerzr

[muellerzr]

Note that pacman no longer works at HF, so you can ping me in the future :)

This makes sense since we need to eventually gather. cc @SunMarc

[muellerzr]

can you run pip install -e .[quality]; make style; make quality?

[SunMarc]

Nice ! Thanks for fixing !

[CodeCreator]

@muellerzr Thanks for the pointer! I ran the code formatter as suggested.

[HuggingFaceDocBuilderDev]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[amyeroberts]

Thanks for fixing!