getting out of memory error with distillgpt as well as gptneo-125 M in ml.p3.2xlarge AWS Sagemaker instance
code:
import sagemaker
from sagemaker.huggingface import HuggingFace
# gets role for executing training job
role = sagemaker.get_execution_role()
hyperparameters = {
    'per_device_train_batch_size': 16,
    'train_file': '/opt/ml/input/data/training/domain-gen-training.csv',
	'model_name_or_path':'distilgpt2',
	'output_dir':'/opt/ml/model',
    'do_train': True,
    'do_eval': False,
    #'data_dir':'/opt/ml/input/data/training'
}

# git configuration to download our fine-tuning script
git_config = {'repo': 'https://github.com/huggingface/transformers.git','branch': 'v4.6.1'}

# creates Hugging Face estimator
huggingface_estimator = HuggingFace(
	entry_point='run_clm.py',
	source_dir='./examples/pytorch/language-modeling',
	instance_type='ml.p3.2xlarge',
	instance_count=1,
	role=role,
	git_config=git_config,
	transformers_version='4.6.1',
	pytorch_version='1.7.1',
	py_version='py36',
	hyperparameters = hyperparameters,
)

# starting the train job
huggingface_estimator.fit({'training': 's3://domain-gen-data/domain-gen-training.csv'})

error:
[INFO|trainer.py:1156] 2021-08-31 10:43:00,451 >> ***** Running training *****
[INFO|trainer.py:1157] 2021-08-31 10:43:00,452 >>   Num examples = 13083
[INFO|trainer.py:1158] 2021-08-31 10:43:00,452 >>   Num Epochs = 3
[INFO|trainer.py:1159] 2021-08-31 10:43:00,452 >>   Instantaneous batch size per device = 32
[INFO|trainer.py:1160] 2021-08-31 10:43:00,452 >>   Total train batch size (w. parallel, distributed & accumulation) = 32
[INFO|trainer.py:1161] 2021-08-31 10:43:00,452 >>   Gradient Accumulation steps = 1
[INFO|trainer.py:1162] 2021-08-31 10:43:00,453 >>   Total optimization steps = 1227
#015  0%|          | 0/1227 [00:00<?, ?it/s]Traceback (most recent call last):
  File "run_clm.py", line 468, in <module>
    main()
  File "run_clm.py", line 421, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1272, in train
    tr_loss += self.training_step(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1734, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1766, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 954, in forward
    return_dict=return_dict,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 797, in forward
    output_attentions=output_attentions,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 354, in forward
    feed_forward_hidden_states = self.mlp(hidden_states)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 282, in forward
    hidden_states = self.act(hidden_states)
  File "/opt/conda/lib/python3.6/site-packages/transformers/activations.py", line 42, in gelu_new
[INFO|trainer.py:1156] 2021-08-31 10:43:00,451 >> ***** Running training *****
[INFO|trainer.py:1157] 2021-08-31 10:43:00,452 >>   Num examples = 13083
[INFO|trainer.py:1158] 2021-08-31 10:43:00,452 >>   Num Epochs = 3
[INFO|trainer.py:1159] 2021-08-31 10:43:00,452 >>   Instantaneous batch size per device = 32
[INFO|trainer.py:1160] 2021-08-31 10:43:00,452 >>   Total train batch size (w. parallel, distributed & accumulation) = 32
[INFO|trainer.py:1161] 2021-08-31 10:43:00,452 >>   Gradient Accumulation steps = 1
[INFO|trainer.py:1162] 2021-08-31 10:43:00,453 >>   Total optimization steps = 1227
#015  0%|          | 0/1227 [00:00<?, ?it/s]Traceback (most recent call last):
  File "run_clm.py", line 468, in <module>
    main()
  File "run_clm.py", line 421, in main
    train_result = trainer.train(resume_from_checkpoint=checkpoint)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1272, in train
    tr_loss += self.training_step(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1734, in training_step
    loss = self.compute_loss(model, inputs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/trainer.py", line 1766, in compute_loss
    outputs = model(**inputs)
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 954, in forward
    return_dict=return_dict,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 797, in forward
    output_attentions=output_attentions,
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 354, in forward
    feed_forward_hidden_states = self.mlp(hidden_states)
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
RuntimeError: CUDA out of memory. Tried to allocate 384.00 MiB (GPU 0; 15.78 GiB total capacity; 14.66 GiB already allocated; 156.75 MiB free; 14.70 GiB reserved in total by PyTorch)
#015  0%|          | 0/1227 [00:02<?, ?it/s]
  File "/opt/conda/lib/python3.6/site-packages/torch/nn/modules/module.py", line 756, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/opt/conda/lib/python3.6/site-packages/transformers/models/gpt2/modeling_gpt2.py", line 282, in forward
    hidden_states = self.act(hidden_states)
  File "/opt/conda/lib/python3.6/site-packages/transformers/activations.py", line 42, in gelu_new
    return 0.5 * x * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
RuntimeError: CUDA out of memory. Tried to allocate 384.00 MiB (GPU 0; 15.78 GiB total capacity; 14.66 GiB already allocated; 156.75 MiB free; 14.70 GiB reserved in total by PyTorch)
#015  0%|          | 0/1227 [00:02<?, ?it/s]


2021-08-31 10:43:11 Uploading - Uploading generated training model
2021-08-31 10:43:11 Failed - Training job failed
---------------------------------------------------------------------------
UnexpectedStatusException                 Traceback (most recent call last)
<ipython-input-4-80ef9c507021> in <module>
     31 
     32 # starting the train job
---> 33 huggingface_estimator.fit({'training': 's3://domain-gen-data/domain-gen-training.csv'})

~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/estimator.py in fit(self, inputs, wait, logs, job_name, experiment_config)
    680         self.jobs.append(self.latest_training_job)
    681         if wait:
--> 682             self.latest_training_job.wait(logs=logs)
    683 
    684     def _compilation_job_name(self):

~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/estimator.py in wait(self, logs)
   1625         # If logs are requested, call logs_for_jobs.
   1626         if logs != "None":
-> 1627             self.sagemaker_session.logs_for_job(self.job_name, wait=True, log_type=logs)
   1628         else:
   1629             self.sagemaker_session.wait_for_job(self.job_name)

~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/session.py in logs_for_job(self, job_name, wait, poll, log_type)
   3731 
   3732         if wait:
-> 3733             self._check_job_status(job_name, description, "TrainingJobStatus")
   3734             if dot:
   3735                 print()

~/anaconda3/envs/pytorch_p36/lib/python3.6/site-packages/sagemaker/session.py in _check_job_status(self, job, desc, status_key_name)
   3291                 ),
   3292                 allowed_statuses=["Completed", "Stopped"],
-> 3293                 actual_status=status,
   3294             )
   3295 

UnexpectedStatusException: Error for Training job huggingface-pytorch-training-2021-08-31-10-34-40-424: Failed. Reason: AlgorithmError: ExecuteUserScriptError:
Command "/opt/conda/bin/python3.6 run_clm.py --do_eval False --do_train True --model_name_or_path distilgpt2 --output_dir /opt/ml/model --per_device_train_batch_size 32 --train_file /opt/ml/input/data/training/domain-gen-training.csv"

0 tables [00:00, ? tables/s]
6 tables [00:00, 53.43 tables/s]
15 tables [00:00, 59.52 tables/s]
24 tables [00:00, 64.73 tables/s]
33 tables [00:00, 68.73 tables/s]
42 tables [00:00, 72.00 tables/s]
50 tables [00:00, 73.93 tables/s]
59 tables [00:00, 75.86 tables/s]
                                 
[INFO|file_utils.py:1532] 2021-08-31 10:41:10,579 >> https://huggingface.co/distilgpt2/resolve/main/config.json not found in cache or force_download set to True, downloading to /root/.cache/huggingface/transformers/tmpx0ysz91j

Downloading:   0%|          | 0.00/762 [00:00<?, ?B/s]
Downloading: 100%|ââââââââââ| 762/762 [00:00<00:00, 677kB/s]
[INFO|file_utils.py:1536] 2021-08-31 10:41:10,600 >> storing https://huggingface.co/