Fix GA loss for Deepspeed
Author: timjeffrey10
Pull request: yes

# What does this PR do?

This PR addresses an issue where the `grad_norm` differs when the `gradient_accumulation_steps` changes when using DeepSpeed. The issue does not occur when using DDP.

In the accelerate library, the `Accelerator` class ignores the `gradient_accumulation_steps` parameter when using DeepSpeed because DeepSpeed handles loss scaling internally. The relevant code in `Accelerate` is:
[Reference](https://github.com/huggingface/accelerate/blob/78b8126bff9882a66074a082cb2a80aa89026118/src/accelerate/accelerator.py#L2234-L2236)
```python
        if self.distributed_type != DistributedType.DEEPSPEED:
            # deepspeed handles loss scaling by gradient_accumulation_steps in its `backward`
            loss = loss / self.gradient_accumulation_steps
```
In DeepSpeed, loss scaling occurs in the `_scale_loss_by_gas` method when the `backward` function is called. While setting `gradient_accumulation_steps` to 1 in the `Accelerator` disables loss scaling, it cannot be set to 1 in DeepSpeed. This is because DeepSpeed relies on this parameter to manage internal operations like counting `micro_steps` in its `step` function.

In DeepSpeed `backward` function, it allow us to pass `scale_wrt_gas` to turn off the loss scaling. The `_scale_loss_by_gas` will not be called.
[Reference](https://github.com/microsoft/DeepSpeed/blob/de4596bedc61100db9385b5d99efd9025db13a7d/deepspeed/runtime/engine.py#L2028-L2029)
```python
        if do_gradient_reduction and self.gradient_accumulation_steps() > 1 and scale_wrt_gas:
            loss = self._scale_loss_by_gas(loss.float())
```

This PR passes the `kwargs["scale_wrt_gas"]=False` to the `backward` function when the DeepSpeed is enabled to prevent potential loss scaling in DeepSpeed engine.

## Before submitting
- [ ] This PR fixes a typo or improves the docs (you can dismiss the other checks if that's the case).
- [X] 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.
- [X] 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?

Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
members/contributors who may be interested in your PR.

<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @

 If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
 Please tag fewer than 3 people.

Models:

- text models: @ArthurZucker
- vision models: @amyeroberts, @qubvel
- speech models: @ylacombe, @eustlb
- graph models: @clefourrier

Library:

- flax: @sanchit-gandhi
- generate: @zucchini-nlp (visual-language models) or @gante (all others)
- pipelines: @Rocketknight1
- tensorflow: @gante and @Rocketknight1
- tokenizers: @ArthurZucker
- trainer: @muellerzr and @SunMarc
- chat templates: @Rocketknight1

Integrations:

- deepspeed: HF Trainer/Accelerate: @muellerzr
- ray/raytune: @richardliaw, @amogkam
- Big Model Inference: @SunMarc
- quantization (bitsandbytes, autogpt): @SunMarc @MekkCyber

Documentation: @stevhliu

HF projects:

- accelerate: [different repo](https://github.com/huggingface/accelerate)
- datasets: [different repo](https://github.com/huggingface/datasets)
- diffusers: [different repo](https://github.com/huggingface/diffusers)
- rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)

Maintained examples (not research project or legacy):

- Flax: @sanchit-gandhi
- PyTorch: See Models above and tag the person corresponding to the modality of the example.
- TensorFlow: @Rocketknight1

 -->


Comments:
cc @sunmarc @muellerzr

Thanks for the review.
I have committed another solution to pass `scale_wrt_gas` to the backward function, which prevents the loss scaling in DeepSpeed engine.
Another solution is to set `scale_wrt_gas` after initializing DeepSpeed engine (i.e., after calling `accelerator.prepare`). I have no idea how to set it during initilizing :(
```python
        model, self.optimizer = self.accelerator.prepare(self.model, self.optimizer)
        if self.is_deepspeed_enabled:
            model.scale_wrt_gas = False
```
This can help not to always add `scale_wrt_gas` to kwargs in `training_step`. If reviewer prefer this implementation, I can try to modify it.

I will upload the wandb logs and testing result ASAP.

# Training logs (wandb)
- DeepSpeed: ZeRO-2
- 2 GPUs, Global batch size = 512
- Compare `gradient_accumulation_steps`=1 and 256

Here is the orinal version, the `grad_norm` is extreme small when `gradient_accumulation_steps`=256
![image](https://github.com/user-attachments/assets/49a8770b-f8d5-46b8-b5e6-e709837e8e69)

Here is the revised version, the `grad_norm` is almost the same in both `gradient_accumulation_steps`=1 and 256
![image](https://github.com/user-attachments/assets/7cedd2f4-5753-460c-a265-1a0008103931)

# GA Tests Results
Below is the screenshot of the results of releted GA tests I found. Please let me know if any tests are missing.
![image](https://github.com/user-attachments/assets/4a52ec38-34fa-4e92-83ff-0e27705d3007)

@muellerzr If there are any problem, please let me know. Thank a lot.


Failing test is unrelated, cc @ArthurZucker it's batching