Using apex leeads to a CUDA out of memory on an A100

username: Juffin_Hally
cooked: <p>I’m using apex in a conda environment, it was installed from <a href="https://anaconda.org/conda-forge/nvidia-apex/0.1/download/linux-64/nvidia-apex-0.1-py37h519209e_4.tar.bz2" rel="noopener nofollow ugc">https://anaconda.org/conda-forge/nvidia-apex/0.1/download/linux-64/nvidia-apex-0.1-py37h519209e_4.tar.bz2</a>.<br>
When I run my training script:</p>
<pre><code class="lang-auto">#!/usr/bin/bash
LANG=python                       # set python for py150
DATADIR=/home/ubuntu/CodeXGLUE/Code-Code/CodeCompletion-token/dataset/py150/token_completion
LITFILE=/home/ubuntu/CodeXGLUE/Code-Code/CodeCompletion-token/dataset/py150/literals.json
OUTPUTDIR=../save/py150
#PRETRAINDIR="sberbank-ai/rugpt3large_based_on_gpt2"       # directory of your saved model
#PRETRAINDIR="gpt2-xl"
#PRETRAINDIR="gpt2-large"
PRETRAINDIR="gpt2-medium"
LOGFILE=completion_py150_eval.log


GPT2MED='gpt2-medium'
GPT2LARGE='gpt2-large'
GPT3LARGE='rugpt3large_based_on_gpt2'
GPT2XL='gpt2-xl'
GPT27B='gpt-neo-2.7B'
GPT6B='gpt-j-6B'

PER_NODE_GPU=1


python3 -u -m torch.distributed.launch --nproc_per_node=$PER_NODE_GPU  code/run_lm.py \
        --data_dir=$DATADIR \
        --lit_file=$LITFILE \
        --langs=$LANG \
        --output_dir=$OUTPUTDIR \
        --pretrain_dir=$PRETRAINDIR \
        --log_file=$LOGFILE \
        --model_type=$GPT2MED \
        --block_size=1024 \
        --do_train \
        --evaluate_during_training \
        --do_eval \
        --per_gpu_train_batch_size 1\
        --per_gpu_eval_batch_size=1 \
        --logging_steps=100 \
        --save_steps=500 \
        --save_total_limit 4 \
        --seed=42 \
        --fp16 \
        --gpu_per_node $PER_NODE_GPU \
        --learning_rate=8e-5 \
        --weight_decay=0.01 \
        --evaluate_during_training \
        --per_gpu_train_batch_size=2 \
        --per_gpu_eval_batch_size=4 \
        --gradient_accumulation_steps=4 \
        --num_train_epochs=5 \
        --overwrite_output_dir \
        --not_pretrain \
        --tensorboard_dir ./tensorboard_logs

</code></pre>
<p>I run into an error:</p>
<pre><code class="lang-auto">/usr/local/lib/python3.8/dist-packages/torch/distributed/launch.py:178: FutureWarning: The module torch.distributed.launch is deprecated
and will be removed in future. Use torchrun.
Note that --use_env is set by default in torchrun.
If your script expects `--local_rank` argument to be set, please
change it to read from `os.environ['LOCAL_RANK']` instead. See 
https://pytorch.org/docs/stable/distributed.html#launch-utility for 
further instructions

  warnings.warn(
11/25/2021 16:11:40 - WARNING - __main__ -   Process rank: -1, device: cuda:0, n_gpu: 1, distributed training: False, 16-bits training: True, world size: 1
Special tokens have been added in the vocabulary, make sure the associated word embeddings are fine-tuned or trained.
/usr/local/lib/python3.8/dist-packages/transformers/models/auto/modeling_auto.py:708: FutureWarning: The class `AutoModelWithLMHead` is deprecated and will be removed in a future version. Please use `AutoModelForCausalLM` for causal language models, `AutoModelForMaskedLM` for masked language models and `AutoModelForSeq2SeqLM` for encoder-decoder models.
  warnings.warn(
11/25/2021 16:11:44 - INFO - __main__ -   Model has a total of 355066880 trainable parameters
11/25/2021 16:11:44 - INFO - __main__ -   Training/evaluation parameters Namespace(adam_epsilon=1e-08, block_size=1024, cache_dir='', config_dir=None, data_dir='/home/ubuntu/CodeXGLUE/Code-Code/CodeCompletion-token/dataset/py150/token_completion', device=device(type='cuda', index=0), do_eval=True, do_lower_case=False, do_train=True, eval_all_checkpoints=False, evaluate_during_training=True, fp16=True, fp16_opt_level='O1', gpu_per_node=1, gradient_accumulation_steps=4, langs='python', learning_rate=8e-05, lit_file='/home/ubuntu/CodeXGLUE/Code-Code/CodeCompletion-token/dataset/py150/literals.json', load_name='pretrained', local_rank=-1, log_file='completion_py150_eval.log', logging_steps=100, max_grad_norm=1.0, max_steps=-1, mlm=False, mlm_probability=0.15, model_type='gpt2-medium', n_gpu=1, no_cuda=False, node_index=-1, not_pretrain=True, num_train_epochs=5.0, output_dir='../save/py150', overwrite_cache=False, overwrite_output_dir=True, per_gpu_eval_batch_size=4, per_gpu_train_batch_size=2, pretrain_dir='gpt2-medium', save_steps=500, save_total_limit=4, seed=42, server_ip='', server_port='', start_epoch=0, start_step=0, tensorboard_dir='./tensorboard_logs', tokenizer_dir=None, warmup_steps=0, weight_decay=0.01)
11/25/2021 16:11:44 - WARNING - __main__ -   Loading features from cached file ../save/py150/train_blocksize_1024_wordsize_1_rank_0
Selected optimization level O1:  Insert automatic casts around Pytorch functions and Tensor methods.

Defaults for this optimization level are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
Processing user overrides (additional kwargs that are not None)...
After processing overrides, optimization options are:
enabled                : True
opt_level              : O1
cast_model_type        : None
patch_torch_functions  : True
keep_batchnorm_fp32    : None
master_weights         : None
loss_scale             : dynamic
11/25/2021 16:11:53 - INFO - __main__ -   ***** Running training *****
11/25/2021 16:11:53 - INFO - __main__ -     Num examples = 126276
11/25/2021 16:11:53 - INFO - __main__ -     Num epoch = 4
11/25/2021 16:11:53 - INFO - __main__ -     Instantaneous batch size per GPU = 2
11/25/2021 16:11:53 - INFO - __main__ -     Total train batch size (w. parallel, distributed &amp; accumulation) = 8
11/25/2021 16:11:53 - INFO - __main__ -     Gradient Accumulation steps = 4
11/25/2021 16:11:53 - INFO - __main__ -     Total optimization steps = 78920
Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 4096.0
/usr/local/lib/python3.8/dist-packages/torch/optim/lr_scheduler.py:122: UserWarning: Seems like `optimizer.step()` has been overridden after learning rate scheduler initialization. Please, make sure to call `optimizer.step()` before `lr_scheduler.step()`. See more details at https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate
  warnings.warn("Seems like `optimizer.step()` has been overridden after learning rate scheduler "
Gradient overflow.  Skipping step, loss scaler 0 reducing loss scale to 512.0
Traceback (most recent call last):
  File "code/run_lm.py", line 725, in &lt;module&gt;
    main()
  File "code/run_lm.py", line 713, in main
    global_step, tr_loss = train(args, train_dataset, model, tokenizer, fh, pool)
  File "code/run_lm.py", line 184, in train
    outputs = model(inputs, labels=labels)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/transformers/models/gpt2/modeling_gpt2.py", line 1073, in forward
    loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/module.py", line 1102, in _call_impl
    return forward_call(*input, **kwargs)
  File "/usr/local/lib/python3.8/dist-packages/torch/nn/modules/loss.py", line 1150, in forward
    return F.cross_entropy(input, target, weight=self.weight,
  File "/usr/local/lib/python3.8/dist-packages/apex/amp/wrap.py", line 25, in wrapper
    new_args = utils.casted_args(cast_fn,
  File "/usr/local/lib/python3.8/dist-packages/apex/amp/utils.py", line 81, in casted_args
    new_args.append(cast_fn(x))
  File "/usr/local/lib/python3.8/dist-packages/apex/amp/utils.py", line 74, in maybe_float
    return x.float()
RuntimeError: CUDA out of memory. Tried to allocate 396.00 MiB (GPU 0; 22.20 GiB total capacity; 19.82 GiB already allocated; 170.12 MiB free; 19.88 GiB reserved in total by PyTorch) If reserved memory is &gt;&gt; allocated memory try setting max_split_size_mb to avoid fragmentation.  See documentation for Memory Management and PYTORCH_CUDA_ALLOC_CONF
ERROR:torch.distributed.elastic.multiprocessing.api:failed (exitcode: 1) local_rank: 0 (pid: 1403) of binary: /usr/bin/python3
Traceback (most recent call last):
  File "/usr/lib/python3.8/runpy.py", line 194, in _run_module_as_main
    return _run_code(code, main_globals, None,
  File "/usr/lib/python3.8/runpy.py", line 87, in _run_code
    exec(code, run_globals)
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/launch.py", line 193, in &lt;module&gt;
    main()
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/launch.py", line 189, in main
    launch(args)
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/launch.py", line 174, in launch
    run(args)
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/run.py", line 710, in run
    elastic_launch(
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/launcher/api.py", line 131, in __call__
    return launch_agent(self._config, self._entrypoint, list(args))
  File "/usr/local/lib/python3.8/dist-packages/torch/distributed/launcher/api.py", line 259, in launch_agent
    raise ChildFailedError(
torch.distributed.elastic.multiprocessing.errors.ChildFailedError: 
============================================================
code/run_lm.py FAILED
------------------------------------------------------------
Failures:
  &lt;NO_OTHER_FAILURES&gt;
------------------------------------------------------------
Root Cause (first observed failure):
[0]:
  time      : 2021-11-25_16:12:04
  host      : ip-172-31-4-91.ec2.internal
  rank      : 0 (local_rank: 0)
  exitcode  : 1 (pid: 1403)
  error_file: &lt;N/A&gt;
  traceback : To enable traceback see: https://pytorch.org/docs/stable/elastic/errors.html
============================================================
</code></pre>
<p>Which seems ridiculous since A100 is full of memory.</p>
<p>How can this be fixed? Perhaps there’s some working set of pytorch, cuda, transformers and apex that doesn’t run into the error</p>

username: ptrblck
cooked: <p><code>apex.amp</code> is depredcated in favor of the native mixed-precision utility used via <code>torch.cuda.amp</code>. Take a look at the <a href="https://pytorch.org/docs/stable/notes/amp_examples.html#amp-examples">examples</a> to see its usage.</p>
<p>The linked package is also not official and thus also doesn’t come with pre-built C++/CUDA kernels.</p>