[Sagemaker] sagemaker distributed features in Trainer broken since Transformers 4.29
Author: JingyaHuang
Pull request: no

### System Info

* `transformers: 4.29.1`
* `datasets: 2.12.0`
* `evaluate: 0.4.0`
* `accelerate: 0.19.0`
* `torch: 2.0.0`
* `diffusers: 0.16.1`

### Who can help?

@philschmid @sgugger

### Information

- [X] The official example scripts
- [ ] 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

The issue is found while I update AWS Sagemaker deep learning container for Transformers-PyTorch. Sagemaker's distributed data parallel and model parallel features are broken since Transformers 4.29.* .

* Test scripts: [`test_smmp.py`](https://github.com/aws/deep-learning-containers/blob/master/test/sagemaker_tests/huggingface_pytorch/training/integration/sagemaker/test_smmp.py) and [`test_smdp.py`](https://github.com/aws/deep-learning-containers/blob/master/test/sagemaker_tests/huggingface_pytorch/training/integration/sagemaker/test_smdp.py).
* Related PR in the DLC repo: https://github.com/aws/deep-learning-containers/pull/2993
* Error log:
```
AttributeError: 'TrainingArguments' object has no attribute 'distributed_state'
```

<details close>
<summary>More detailed error log</summary>
<br>
```
_____________________________ test_smmp_gpu[gloo] ______________________________

ecr_image = '669063966089.dkr.ecr.us-west-2.amazonaws.com/pr-huggingface-pytorch-training:2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48'
sagemaker_regions = ['us-west-2', 'us-east-1', 'eu-west-1']
instance_type = 'ml.p3.8xlarge', framework_version = '2.0.0', py_version = 'py3'
dist_gpu_backend = 'gloo'

    @pytest.mark.processor("gpu")
    @pytest.mark.integration("smmp")
    @pytest.mark.model("hf_qa_smmp")
    @pytest.mark.skip_cpu
    @pytest.mark.skip_py2_containers
    @pytest.mark.skip_trcomp_containers
    def test_smmp_gpu(
        ecr_image, sagemaker_regions, instance_type, framework_version, py_version, dist_gpu_backend
    ):
>       invoke_sm_helper_function(ecr_image, sagemaker_regions, _test_smmp_gpu_function, py_version, 1)

integration/sagemaker/test_smmp.py:76:
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _
../../__init__.py:119: in invoke_sm_helper_function
    raise e
../../__init__.py:113: in invoke_sm_helper_function
    test_function(tested_ecr_image, sagemaker_session, *test_function_args)
integration/sagemaker/test_smmp.py:117: in _test_smmp_gpu_function
    huggingface_estimator.fit(job_name=sagemaker.utils.unique_name_from_base("test-hf-pt-qa-smmp"))
2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48/lib/python3.8/site-packages/sagemaker/workflow/pipeline_context.py:272: in wrapper
    return run_func(*args, **kwargs)
2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48/lib/python3.8/site-packages/sagemaker/estimator.py:1156: in fit
    self.latest_training_job.wait(logs=logs)
2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48/lib/python3.8/site-packages/sagemaker/estimator.py:2297: in wait
    self.sagemaker_session.logs_for_job(self.job_name, wait=True, log_type=logs)
2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48/lib/python3.8/site-packages/sagemaker/session.py:4216: in logs_for_job
    self._check_job_status(job_name, description, "TrainingJobStatus")
_ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _ _

self = <sagemaker.session.Session object at 0x7f466ef3fc40>
job = 'test-hf-pt-qa-smmp-1684001017-f81e'
desc = {'AlgorithmSpecification': {'EnableSageMakerMetricsTimeSeries': True, 'TrainingImage': '669063966089.dkr.ecr.us-west-2...)), 'DebugHookConfig': {'CollectionConfigurations': [], 'S3OutputPath': 's3://sagemaker-us-west-2-669063966089/'}, ...}
status_key_name = 'TrainingJobStatus'
```
    def _check_job_status(self, job, desc, status_key_name):
        """Check to see if the job completed successfully.

        If not, construct and raise a exceptions. (UnexpectedStatusException).

        Args:
            job (str): The name of the job to check.
            desc (dict[str, str]): The result of ``describe_training_job()``.
            status_key_name (str): Status key name to check for.

        Raises:
            exceptions.CapacityError: If the training job fails with CapacityError.
            exceptions.UnexpectedStatusException: If the training job fails.
        """
        status = desc[status_key_name]
        # If the status is capital case, then convert it to Camel case
        status = _STATUS_CODE_TABLE.get(status, status)

        if status == "Stopped":
            LOGGER.warning(
                "Job ended with status 'Stopped' rather than 'Completed'. "
                "This could mean the job timed out or stopped early for some other reason: "
                "Consider checking whether it completed as you expect."
            )
        elif status != "Completed":
            reason = desc.get("FailureReason", "(No reason provided)")
            job_type = status_key_name.replace("JobStatus", " job")
            message = "Error for {job_type} {job_name}: {status}. Reason: {reason}".format(
                job_type=job_type, job_name=job, status=status, reason=reason
            )
            if "CapacityError" in str(reason):
                raise exceptions.CapacityError(
                    message=message,
                    allowed_statuses=["Completed", "Stopped"],
                    actual_status=status,
                )
>           raise exceptions.UnexpectedStatusException(
                message=message,
                allowed_statuses=["Completed", "Stopped"],
                actual_status=status,
            )
E           sagemaker.exceptions.UnexpectedStatusException: Error for Training job test-hf-pt-qa-smmp-1684001017-f81e: Failed. Reason: AlgorithmError: ExecuteUserScriptError:
E           ExitCode 1
E           ErrorMessage "AttributeError: 'TrainingArguments' object has no attribute 'distributed_state'
E            ╭───────────────────── Traceback (most recent call last) ──────────────────────╮
E            │ /opt/conda/lib/python3.10/runpy.py:196 in _run_module_as_main                │
E            │                                                                              │
E            │   193 │   main_globals = sys.modules["__main__"].__dict__                    │
E            │   194 │   if alter_argv:                                                     │
E            │   195 │   │   sys.argv[0] = mod_spec.origin                                  │
E            │ ❱ 196 │   return _run_code(code, main_globals, None,                         │
E            │   197 │   │   │   │   │    "__main__", mod_spec)                             │
E            │   198                                                                        │
E            │   199 de, exit code: 1

2.0.0-transformers4.29.1-gpu-py310-cu118-ubuntu20.04-pr-2993-2023-05-13-17-25-48/lib/python3.8/site-packages/sagemaker/session.py:3749: UnexpectedStatusException
```

</details>


### Expected behavior

Find either what to adapt from the sagemaker side or from our side to make sure distributed features work, so that we would be update transformers to higher versions the next time. 

Comments:
cc @pacman100  

cc @muellerzr as he has full insights on the related PR.

@JingyaHuang could you show the full stack trace so I can see where this error is specifically being raised from/its context?

I need any context of a trace that relates to something in transformers to be able to know where/how this is stemming from. SageMaker should not be handled by the current accelerate implementation, so it's critical to know just where that logic fault is :) 

@muellerzr Sure, here is the complete log: https://gist.github.com/JingyaHuang/9327adb701d9989da8cf4c33bfeb043e
(sorry if it looks messy, it's populated by sagemaker. Related tests are test_smmp and test_smdp.py)


@JingyaHuang could you try installing transformers via `pip install git+https://github.com/huggingface/transformers@muellerzr-fix-sagemaker` and verify this fixes it? And if not, what other errors arise? Thanks!

Trying! It will take a while for the image build, will update asap.

@philschmid I thought SageMaker MP was broken since many releases ago?

> @philschmid I thought SageMaker MP was broken since many releases ago?

It was not broken, just not using all the features by default, e.g. Tensor Parallelism. 

Hey @muellerzr, here is the log that I got with your patch: https://gist.github.com/JingyaHuang/3b60725d0a6f22f377b27694d22c18ca

There seems to be another issue now:
```
TypeError: init_process_group() got multiple values for argument 'backend'
```  

@JingyaHuang what version of `accelerate` are you running with? That should be fixed in v0.19.0

@JingyaHuang there was a fix on main for accelerate that may be related to this, are we trying to spawn on cpu/do distributed CPU?

Thanks @muellerzr, yes the tests were running in a docker container with accelerate 0.19.0 installed
```
Name: accelerate
Version: 0.19.0
Summary: Accelerate
Home-page: https://github.com/huggingface/accelerate
Author: The HuggingFace team
Author-email: sylvain@huggingface.co
License: Apache
Location: /opt/conda/lib/python3.10/site-packages
Requires: numpy, packaging, psutil, pyyaml, torch
Required-by: 
``` 
The smmp test ran sagemaker distributed on GPU, I am not familiar with how sagemaker's distributed model parallel works. @philschmid might have a better answer, does it spawns CPU processes? 

@JingyaHuang looked at the trace again, yes it does make sense that the fix to main should have changed it actually based on what's happening here. Can you try one more time with `pip install git+https://github.com/huggingface/accelerate` to verify? Thanks 🙏 

(Auto-closed due to PR merging, will keep open until we know for sure w/ Accelerate fix :) )

Hey @muellerzr, here is the log that I got by the test with accelerate from source last week: https://gist.github.com/JingyaHuang/0026e8801e99d0df522fb2bcb2b2334c
(I did not configure accelerate, not sure if I should do that?)

Thanks @JingyaHuang (and appreciate your patience).

Let's try via the following:

```bash
pip install git+https://github.com/huggingface/transformers@muellerzr-sagemaker-dp git+https://github.com/huggingface/accelerate@sagemakerdp
```
Thanks!

@muellerzr No worries!

Another error occurs, we don't have the luck :( . Here is the log: https://gist.github.com/JingyaHuang/01de393f9da716ff094c248f12ec1465 

@JingyaHuang I'd actually say we're making progress! New errors, with easier solutions :) Let's try again, same branches etc for everything

Boom!
🚀 
![image](https://github.com/huggingface/transformers/assets/44135271/b0d3ed73-55fa-48ad-ba71-bcb49d54dbf2)


@muellerzr I just ran the sagemaker data parallel test but with question answering task(`run_qa.py` ❌ ) instead of text generation task(`run_glue.py` ✅ ) this time, and it failed during the evaluation (while doing post-processing of the predictions).

```
ValueError: Got 676 predictions and 10784 features.
``` 

Full tracing log here: https://gist.github.com/JingyaHuang/b824d3abd17c6db23e68968dec0cee13 

Do you think it is related to the issue, or just an update need to be done for the examples? 

@JingyaHuang to know for sure, try running it with `transformers==4.28.1`

@muellerzr I just did two tests:

* Run __qa__ example on __smmp__ test with __patched transformers & accelerate__ -> tests passed ✅ 
![image](https://github.com/huggingface/transformers/assets/44135271/023e2eed-9215-4040-9e70-447cd69b8f34)

(so with the patch, smmp test passed for both text classification and qa)

p.s. `run_qa` and `run_glue` are fetched from the main branch

* Run __qa__ example on __smdp__ test with __trfrs 4.28.1 & accelerate 0.19.0__ -> tests  passed ✅  
![image](https://github.com/huggingface/transformers/assets/44135271/479df65c-d891-4c42-a66b-9890e4d23356)

p.s. `run_qa` and `run_glue` are fetched from the v4.28.1 branch

And the previous error log on the qa task was for __smdp__ test, so it seems smmp is good but smdp is still broken for 4.29.*, is there anything that needs to be done for smdp during the prediction step maybe?

@muellerzr The number of features for evaluation is 10784. Since the smdp test was run on two `ml.p3.16xlarge`(8 gpus) instances and 10784 / 676 = 15.9526627, intuitively I doubt when using smdp only the predictions on one worker are kept (676). 

Ref smdp test: https://github.com/aws/deep-learning-containers/blob/master/test/sagemaker_tests/huggingface_pytorch/training/integration/sagemaker/test_smdp.py#L108

just a thought (ㆆᴗㆆ)

Closing now via https://github.com/huggingface/transformers/pull/23681, as all tests pass