adding option to save/reload scaler
Author: hsilva664
Pull request: yes

# What does this PR do?

Attempts to solve inconsistencies when saving and reloading models discussed in #34053 .


@pacman100 @yitongh @SunMarc 


Comments:
@hsilva664 Great, thank you! I quickly tested the code, and the results of resumed training are now consistent with the initial training.
Regarding lines 288 and 291, the `SCALER_NAME` is defined twice.

@SunMarc  I have added a test that runs the same dummy language model training example discussed in [34053](https://github.com/huggingface/transformers/issues/34053). It compares the interrupted run with the from-scratch one and checks to see if the result is exactly the same, ensuring reproducibility. If, say, the scaler---or anything really---was to affect the training progress in any way, this test would catch it. The test also checks for the presence of the file saving the scaler. The test is run with:

```bash
python -m pytest -n auto --dist=loadfile -s -v ./tests/trainer/test_trainer.py::TrainerIntegrationTest::test_can_resume_training
```

There were already some tests checking for this functionality, but they were mostly done on a simpler linear regression setting, so maybe this is why they did not catch the scaler issue. 

A problem that I am running into is that, to ensure determinism and proper checking of the saving/reloading, this example requires setting `os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"`. When I run the test in isolation, it passes. Still, when running it together with the other tests using

```bash
python -m pytest -n auto --dist=loadfile -s -v ./tests/trainer/test_trainer.py
```

It seems that setting the environment variable like I am doing is interfering with other tests and causing errors. I tried setting it in a way that this would not happen, but it seems to not have been enough. The reason this variable needs to be set is because, otherwise, if we call `torch.use_deterministic_algorithms(True)` and then an operation using CuBLAS (with CUDA >= 10.2), the following error occurs:

> RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)` or `at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic because it uses CuBLAS and you have CUDA >= 10.2. To enable deterministic behavior in this case, you must set an environment variable before running your PyTorch application: CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. For more information, go to https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility


@SunMarc @muellerzr  Can you give me a quick pointer so that I can change the code and make it pass the tests? I think the problem I'm facing might be familiar.

I implemented a new test that reproduces the problem in the original issue. To try to keep the code organized, I added the test in the class `TrainerIntegrationTest`, under `test_trainer.py`. My test, however, only works if an environment variable is set right at the beginning of the script (before CuBLAS is loaded, to be precise), otherwise, tests running in parallel will start CuBLAS before the variable is set and the test will break. I am hesitant to set it in the beginning of the file, as I believe it might interfere with pre-existing tests. Setting the variable later does not work, its effects are ignored. I could create a new file just for this test, but that might make things disorganized needlessly. Do you know a cleaner way to address this? 

Using the RNGStateManager you coded doesn't work ? Also could you do make style and rebase for PR to pass the CI ? cc @muellerzr @ydshieh  for the questions about the global variables

The docs for this PR live [here](https://moon-ci-docs.huggingface.co/docs/transformers/pr_34932). All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

@SunMarc I made the desired changes. For the isolation of the new test, I ended up simply creating another file. Please see if the PR is ready to be merged now.

cc @ydshieh for the test file 

One quick question first:

> cc @ydshieh for the questions about the global variables

@SunMarc  is this still relevant and require me to check? I don't follow which `global variables` (is it about `SCALER_NAME` or some env. variables ..?).

> @SunMarc is this still relevant and require me to check?

Not needed anymore for global variables @ydshieh  ! OP created a new test file because of the following: 

> Creating this file was necessary, instead of adding this function to test_trainer.py. The reason is that, to run
> these tests with determinism, we would need to set `os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8"` before importing
> pytorch, otherwise the call to any function using CuBLAS after using `torch.use_deterministic_algorithms(True)`
> would throw an error. If this is run together with other workers, these workers can start CuBLAS before the variable is
> set. If, instead, we had set this variable at the beginning of test_trainer.py, this might interfere with the other
> tests. Details at https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility.

Just wanted to know if this was fine and just to let you know that there was such behavior.
He set os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":16:8" at the beginning of the file



Got it.

@hsilva664: For me to better understand the situation I have a few questions


Q1: The `determinism` is required for

> self._inner.assertTrue(torch.allclose(model_params[indices], model_params_sample))

or for 

> self._inner.check_trainer_state_are_the_same(state1, state2)

, both of them, or even something else I miss?


Q2:

> tests running in parallel will start CuBLAS before the variable is set and the test will break

>  If this is run together with other workers, these workers can start CuBLAS before the variable is
set

Each worker is running tests in a separate python process, so I don't think the error is caused by parallel works. However, the tests collected to be run in a single python process may running a test that will call `use_deterministic_algorithms` then later call the new tests you added in this PR. I believe this is likely the situation where we get the error.

I might be wrong however.

I am a bit surprised that we don't have such issue so far (on `main`). Do you know which operation throws the error (the one requiring `CuBLAS `). Could you provide the log.

Q3:

If my gess it correct, we probably could use `run_test_using_subprocess`, see #34213 , and avoing using the new test file + no need to set the environment variable?





@ydshieh 

**Q1 - What is the determinism for?**

The training needs to be deterministic for both lines that you wrote. The test was based on an [issue mentioned above](https://github.com/huggingface/transformers/issues/34053) where the NN after interrupting the training and then resuming from a checkpoint was very different than training without interruption. There, training was breaking because the accelerator had a scaling factor that was not being saved/loaded. The existing tests verified saving/loading behaviour for a very simple linear regression case, so I added a short test training with GPT-2 based on the problem presented in the issue.

**Q2 - What operation throws the determinism error?**

I agree with your assessment that probably a shared process among the pool runs one test that starts  CuBLAS and then runs mine afterwards, resulting in the error.

The error is thrown when I call `train`. It seems to simply be occurring on the last layer, the one that generates the logits. Here is a log of a call that reproduces the error:

```
Traceback (most recent call last):
  File "transformers/temp.py", line 100, in <module>
    trainer.train(
  File "transformers/src/transformers/trainer.py", line 2163, in train
    return inner_training_loop(
           ^^^^^^^^^^^^^^^^^^^^
  File "transformers/src/transformers/trainer.py", line 2524, in _inner_training_loop
    tr_loss_step = self.training_step(model, inputs, num_items_in_batch)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/src/transformers/trainer.py", line 3711, in training_step
    loss = self.compute_loss(model, inputs, num_items_in_batch=num_items_in_batch)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/src/transformers/trainer.py", line 3767, in compute_loss
    outputs = model(**inputs)
              ^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/accelerate/utils/operations.py", line 823, in forward
    return model_forward(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/accelerate/utils/operations.py", line 811, in __call__
    return convert_to_fp32(self.model_forward(*args, **kwargs))
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/amp/autocast_mode.py", line 44, in decorate_autocast
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "transformers/src/transformers/models/gpt2/modeling_gpt2.py", line 1083, in forward
    lm_logits = self.lm_head(hidden_states)
                ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "transformers/hugging-face/lib/python3.12/site-packages/torch/nn/modules/linear.py", line 125, in forward
    return F.linear(input, self.weight, self.bias)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: Deterministic behavior was enabled with either `torch.use_deterministic_algorithms(True)` or `at::Context::setDeterministicAlgorithms(true)`, but this operation is not deterministic because it uses CuBLAS and you have CUDA >= 10.2. To enable deterministic behavior in this case, you must set an environment variable before running your PyTorch application: CUBLAS_WORKSPACE_CONFIG=:4096:8 or CUBLAS_WORKSPACE_CONFIG=:16:8. For more information, go to https://docs.nvidia.com/cuda/cublas/index.html#results-reproducibility
```

**Q3 - New suggestion using subprocess**

I will try to run it using subprocess to see if I can move the test back to the other file. Thanks for the suggestion.

@ydshieh @SunMarc I updated the code to run using the test subprocess. It worked correctly on my machine. After pushing, however, one of the tests seem to have failed. Inspecting CircleCI log above yields:

> FAILED tests/models/qwen2_vl/test_modeling_qwen2_vl.py::Qwen2VLModelTest::test_assisted_decoding_matches_greedy_search_1_same - ValueError: Image features and image tokens do not match: tokens: 2, features 1

I did not do anything to this file though, and I don't believe it shares environment with my subprocess either.

that is irrelevant to your PR, no worry :-)

Glad subprocess works ! Will take a final look tomorrow

BTW, how long it takes to run this single test? On a T4, it takes more than 10 minutes.

OK, it blocks when there are multiple GPU. On a single GPU runner, it takes 3 minutess.

Let's decorate the test with `require_torch_non_multi_accelerator` for now. 
(add a comment about the reason)

Also let's change the steps from 20 to 3, checkpoint to resume to 1, batch per device from 4 to 2 to speedup the test. (180 seconds to 30 seconds), possibly incrase lr.

WDYT?

@ydshieh On my tests it takes very little (less than 1 minute), but good to know it is not as fast on T4. I will add the decorator and comment.

About why I used 20 and 11, it was to reproduce [this behaviour](https://github.com/huggingface/transformers/issues/34053#issuecomment-2501830389). If I simplify the tests too much, I lose it. I've just tried your suggestion and the test without the PR---which should fail---passes anyway, since the grad norm never stops being infinity and the buggy behaviour does not kick in. I will try to tweak parameters a little to get that behaviour with less compute.

@ydshieh I spent some time trying to write the test a faster way without making major changes (like changing the model, dataset, or stop using accelerator). As I've said before, the 20-step test has the behaviour of [my linked comment](https://github.com/huggingface/transformers/issues/34053#issuecomment-2501830389) and your suggested 3-step does not. This is because of the accelerator starting with scale `2**16` by default, which causes these `inf` grad norms initially. On the 20-step, it takes some time for it to get to more mild scales, on the 3-step, it cannot get there because the training ends so fast. This causes the average gradient in the 3-step to always have `inf` norm and the saving-reload behaviour makes no difference. One way of making the 3-step version work as it should is to do `trainer.accelerator.scaler._init_scale = 1.0` by hand before calling `trainer.train()`, however I think this is likely to generate unpredictable results more easily in the future with slight changes to the current codebase, I don't think accelerator was meant to be tweaked with like this. The current interfaces do not seem to provide a way of setting this scale by parameter-passing.

Maybe another idea would be to simply do:

- `per_device_train_batch_size = 1`
- `max_steps=10`
- resume from step 8 (`checkpoint-9`)

This is faster to run while also keeping the `inf` behaviour that caused the issue in the first place (i.e. `inf` grad causes optimizer and lr scheduler to be skipped, which caused this vastly different behaviour between save-reload and uninterrupted trainings, which was easy to detect and report as an issue). It avoids having to tweak accelerator or having to change the model. Is that reasonable with you? It works as expected, the test fails without the PR changes and passes with them.

OK, sounds great, thanks a lot for the detailed investigation and explanation. I will  check how long the test will take on T4 runner tomorrow and come back to you, but I think it should be fine.

(BTW, if you don't mind spend a bit of more time to explain one more thing to me: is it normal to have these `inf` values at the beginning of training ...? From what you described, it seems to me it is that cases and after a few steps, it won't be `inf` anymore. But I am a bit surprised how this is possible 🤔 .)

There is a `trainer.accelerator` attribute that runs the forward pass. This accelerator has a `scaler` attribute, which multiplies the loss by a number before backpropagating. This number is gradually adjusted throughout training, but it starts at `2**16=65536`. It is because of this multiplication that the grads start at infinity. I'm not sure about the workings of the accelerator or why the initial scale is chosen to be so high, as I learned about it by inspecting the stack trace of this issue, but its purpose is to enable distributed training and seamless operation across different types of devices, as well as training with mixed precision. I don't know if having these initial `-inf` is intentional (since it will go away shortly) or not.

> Maybe another idea would be to simply do:

37 seconds on T4 runners. Let's go with it 🙏 

Will leave @ArthurZucker for a final look (file `src/transformers/trainer.py`)

Okay, let's go with this then! 🤗