Fine-tune GPT-J with SageMaker Model Parallelism 
## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->

- `transformers` version: 4.11.0
- Platform: SageMaker
- Python version: 3.8
- PyTorch version (GPU?): 1.9.0
- Using GPU in script?: Yes
- Using distributed or parallel set-up in script?: No

### Who can help
@philschmid and @patil-suraj

Models: GPT-J


## Information

I am trying to fine-tune GPT-J on SageMaker using Model Parallelism and I am getting CUDA OOM errors when loading the model. **I have confirmed that my code works with GPT-2**.

I'm on a relatively new account and am in the process of increasing my instance allocations; however, I have attempted to fine-tune GPT-J on one `ml.p3.16xlarge` (8 V100 16GB GPUs) instance and then on two `ml.p3.16xlarge` instances. 

In both cases, I set the per device batch size to 1 and `fp16` to `True`. I am also only using a single partition of the model. 

I expected GPT-J to exceed 96GB, but I did not think that loading the model alone would exceed 128GB (i.e. I expected to at least be able to fit the model into memory on a single `ml.p3.16xlarge`). And the model should definitely fit one two `ml.p3.16xlarge`s. That said, when I tried to train on two instances, my logs indicated that the model was loaded on each instance -- in contrast, I expected a single initialization of the model to be spread across 16 GPUs (I interpret this as a sign of user error...).

I've scoured the Hugging Face and SageMaker docs and haven't found a solution to my problem, so I am hoping for a little guidance. I assume the problem is related to one or more of the following possibilities: 

* I am misusing SageMaker's Model Parallelism
* I simply need more VRAM. 
* SageMaker's Model Parallelism is doing something inefficient with GPT-J

The problem arises when using:
* [x] my own modified scripts: I am using a very lightly modified version of a SageMaker mod of [run_clm.py](https://github.com/aws/amazon-sagemaker-examples/blob/master/sagemaker-training-compiler/huggingface/pytorch_multiple_gpu_single_node/scripts/run_clm.py), which is associated with a tutorial on SageMaker's Training Compiler. The only additional modification I made was to import `SageMakerTrainer` and `SageMakerTrainingArguments` as follows:

```python
from transformers.sagemaker import SageMakerTrainingArguments as TrainingArguments
from transformers.sagemaker import SageMakerTrainer as Trainer
```

The tasks I am working on is:
* [x ] an official GLUE/SQUaD task: sst2

## To reproduce

Steps to reproduce the behavior:
I am starting my SageMaker Training job with the HuggingFace SDK.

1. **Specify hyperparameters**

```python
from sagemaker.huggingface import HuggingFace, TrainingCompilerConfig

INSTANCE_TYPE = "ml.p3.16xlarge"  # ml.p3.8xlarge is easily available. However, p3.16xlarge provides better performance.

per_device_batch_size = 1
num_gpus_per_instance = 8

learning_rate = (
    float("5e-5") / 32 * per_device_batch_size * num_gpus_per_instance
)

hyperparameters = {
    "tokenizer_name": "EleutherAI/gpt-j-6B",
    "model_name_or_path": "EleutherAI/gpt-j-6B",
    "dataset_name": "glue",
    "dataset_config_name": "sst2",
    "do_train": True,
    "do_eval": True,
    "fp16": True,
    "per_device_train_batch_size": per_device_batch_size,
    'per_device_eval_batch_size': per_device_batch_size,
    "learning_rate": learning_rate,
    "num_train_epochs": 1,
    "block_size": 700,
    "overwrite_output_dir": True,
    "save_strategy": "no",
    "logging_strategy": "epoch",
    "output_dir": "/opt/ml/model",
    "max_train_samples": 100,
    "cache_dir":"/tmp",
    'max_grad_norm': 0,
    'preprocessing_num_workers': 1,
}
```
2. **Specify MP distribution**

```python
# configuration for running training on smdistributed model parallel
mpi_options = {
    "enabled" : True,
    "processes_per_host" : num_gpus_per_instance
}

smp_options = {
    "enabled":True,
    "parameters": {
        "microbatches": 1,
        "placement_strategy": "spread",
        "pipeline": "interleaved",
        "optimize": "memory",
        "partitions": 1,
        "ddp": True,
    }
}

distribution={
    "smdistributed": {"modelparallel": smp_options},
    "mpi": mpi_options
}
```

3. **Initialize the estimator**

```python


# configure the training job
optimized_estimator = HuggingFace(
    entry_point="run_clm_orig_mp.py",  # Wrapper around training script that enables multi GPU training
    source_dir="../scripts",
    instance_type=INSTANCE_TYPE,
    instance_count=1,
    role=role,
    volume_size=500,
    py_version="py38",
    transformers_version="4.11.0",
    pytorch_version="1.9.0",
    hyperparameters=hyperparameters,
    distribution = distribution,
    debugger_hook_config=False,  # Disabling SageMaker Debugger to avoid overheads during benchmarking
)
```

4. **Start Training job**

```python
# start the training job
optimized_estimator.fit(wait=False)
optimized_estimator.latest_training_job.name
```


## Expected behavior

I expected to GPT-J to be distributed across N devices and then fine-tuned.

Hi @joehoover -

Thanks for raising this issue! Couple of things based on the above configuration setup -

1. Can you try with single instance training job with SMP configuration with partitions = number of devices, disable ddp by setting to False and set processes_per_host = <number of GPU devices>  for the first experiment and see if the partition fits to the single GPU Memory. Please note - Single GPU memory for p3.16xlarge is 16GB/GPU. With above setup you still have single partition and given the size of FP16 quantized GPT-J model, it might not fit to single GPU memory and thus OOM. You might need to tune the number of partitions value once OOM error goes away.
2. Please attach the training script that you are using so that I can guide you further.

I will also try to reproduce the issue at my end and let you know the results.
Looking forward to your response!<|||||>Hi @dhawalkp, thanks for looking into this! I think you've already cleared up some confusion on my end; I somehow misunderstood the definition of partitions in SageMaker's model parallelism library. For some reason, I thought partitions referred to replications of the model, e.g. I thought `partitions = 2, ddp=True` would implement 2 way data parallelism. 

I tried rerunning with your suggested hyperparameters and I still hit an OOM error; however, the discrepancy between the memory allocated and the memory available decreased substantially: 

```bash
2021-12-21T10:20:21.096-05:00

Copy
[1,0]<stderr>:RuntimeError: CUDA out of memory. Tried to allocate 64.00 MiB (GPU 0; 15.78 GiB total capacity; 14.34 GiB already allocated; 22.00 MiB free; 14.35 GiB reserved in total by PyTorch)
```

I'm trying again with a much smaller block size (changed from `700` to `250`) and I'll update my response with the results of that test. 

I've also attached my complete training script, which I converted to .txt so I could upload to github. 


[run_clm_orig_mp.txt](https://github.com/huggingface/transformers/files/7756450/run_clm_orig_mp.txt)

If you don't mind another question, perhaps you could clarify how I would specify 2-way data parallelism and 4 way model parallelism with the library. I understand that that would involve 8 processes and therefore 8 devices. And I now understand (I think) that setting `partitions=8` and `ddp = False` specifies 8-way model parallelism. But, I'm not sure how I would switch to 2-way data parallelism and 4-way model parallelism. 

Thanks!

### Update

Ran into the same OOM error with `block=250`. I'm now trying `instances=2` and `partitions=16`.

### Update 2

Running into resource issues when requesting 2 `ml.p3.16xlarge` instances. I'll report back when the request goes through.<|||||>Hi Joe -
I can reproduce the CUDA out of memory even after trying with partitions= number of GPUs with 2 machines in the cluster. I am checking this issue further and working with sagemaker product team on this. Will keep you posted.
Thanks<|||||>Hey @dhawalkp, thanks so much for the update! <|||||>Hi Joe -

I analyzed the stack trace of the thread that errored with CUDA OOM. Based on the stack trace, the  initial Model tracing part is erroring out because of the size of the model is 24GB (FP32).The initial model tracing is using FP32 GPT-J model and is using GPU device's memory by default. As the GPU device's Memory is limited to 16GB, its erroring out with OOM.  The SageMaker Model Parallel library has added a parameter called as trace_device which if set to CPU, the initial tracing will happen on the host memory instead of GPU's memory. I will perform another test with  HF Estimator with trace_device=cpu and see if everything runs good. 

Another option is to adopt latest Prebuilt SageMaker container image for PyTorch (https://github.com/aws/deep-learning-containers/releases/tag/v1.12-pt-1.8.1-tr-gpu-py36)  and extend the below example of training the HF Models with Model parallelism using SageMaker PyTorch prebuilt container. You will have to modify your training script as per the example shown here - https://github.com/aws/amazon-sagemaker-examples/blob/8b4789a6a91c8fd3b342a8749619cfde53875666/training/distributed_training/pytorch/model_parallel/gpt2/train_gpt_simple.py

Example notebook: https://github.com/aws/amazon-sagemaker-examples/blob/8b4789a6a91c8fd3b342a8749619cfde53875666/training/distributed_training/pytorch/model_parallel/gpt2/submit-train-gpt-simple.ipynb
The above example has been demonstrated to train 100B parameters Model size using SMP library with the latest container image.

I will keep you posted on the results of the first option I suggested.
Thanks
<|||||>Hey Dhawalkumar, 

I'm trying to figure out how to set `trace_device` to `cpu` but I've not had any luck so far. I took a wild guess and specified `        "trace_device": 'cpu'` in my `modelparallel` specification (i.e. it added it to my distribution specification). However, my logs still indicate that tracing was done on GPU:

```
[1,0]<stdout>:[2022-01-03 19:58:03.507: I smdistributed/modelparallel/torch/worker.py:280] Tracing on GPU. If the model parameters do not fit in a single GPU, you can set trace_device to `cpu`. | [1,0]<stdout>:[2022-01-03 19:58:03.507: I smdistributed/modelparallel/torch/worker.py:280] Tracing on GPU. If the model parameters do not fit in a single GPU, you can set trace_device to `cpu`.
```

I've since tried to identify where `trace_device` should be specified. This led me to SageMaker [docs](https://sagemaker.readthedocs.io/en/v2.21.0/api/training/smd_model_parallel_pytorch.html) for `DistributedModel` and it appears that `trace_device` should be specified when a `DistributedModel` is instantiated. 

However, the Transformer's `Trainer` class doesn't seem to pass `trace_device` or kwargs [when it instantiates](https://github.com/huggingface/transformers/blob/f2ab21833f3c1fcc8dea76988bdacd368fca779e/src/transformers/trainer.py#L958) a `DistributedModel` instance. 

I very well might be misunderstanding something, but I'm nonetheless still not sure how to specify `train_device` when using the HuggingFace estimator. 

Thanks!
<|||||>Hi Joe - 
HuggingFace SageMaker container might not have yet support for this parameter - trace_device. We are working on upgrading the container. Meanwhile, please use SageMaker PyTorch container option mentioned in my previous post. That should definitely work. 
Thanks<|||||>@dhawalkp in which version did `trace_device` get introduced? <|||||>@philschmid - Its not present in the current version of HF DLC yet. It's in the roadmap.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.<|||||>@josephevans - Checking if there is a separate PR for Latest SageMaker Model Parallel library version support for SageMaker hugging face DLC.<|||||>Hi @joehoover - I have published the notebook for training GPT-J using SageMaker Model Parallel here - https://github.com/aws/amazon-sagemaker-examples/tree/main/training/distributed_training/pytorch/model_parallel/gpt-j. Thanks for collaborating on this! Let me know if this issue is resolved.<|||||>This issue has been automatically marked as stale because it has not had recent activity. If you think this still needs to be addressed please comment on this thread.

Please note that issues that do not follow the [contributing guidelines](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md) are likely to be ignored.