bigscience-workshop / bigscience-workshop/Megatron-DeepSpeed

User Warnings for accessing grad attribute of non-leaf Tensors thrown with TP=1 and PP>1

Open
#356 3 comments 0 reactions 0 assignees View on GitHub
Dominant language
Python
Stars
1.4k
Forks
226
PR merge metrics
No merged PRs in 30d

Description

# Problem
On pretaining GPT like models using this [script](https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/7b5f175b73a12d602cdadd3ee205ed666f6d4234/pretrain_gpt.py) , with **_*tensor_parallelsim(TP)=1 and pipeline_paralleism(PP)>1 for model of any size and batch size*_** I get the following user warning multiple times
```
[default3]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default3]: return self._grad
[default2]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default2]: return self._grad
[default1]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default1]: return self._grad
[default0]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default0]: return self._grad
[default3]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default3]: return self._grad
[default2]:/p/software/juwelsbooster/stages/2022/software/PyTorch/1.11-gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/lib/python3.9/site-packages/torch/_tensor.py:1104: UserWarning: The .grad attribute of a Tensor that is not a leaf Tensor is being accessed. Its .grad attribute won't be populated during autograd.backward(). If you indeed want the .grad field to be populated for a non-leaf Tensor, use .retain_grad() on the non-leaf Tensor. If you access the non-leaf Tensor by mistake, make sure you access the leaf Tensor instead. See github.com/pytorch/pytorch/pull/30531 for more informations. (Triggered internally at /dev/shm/strube1/juwelsbooster/PyTorch/1.11/gcccoremkl-11.2.0-2021.4.0-CUDA-11.5/pytorch/build/aten/src/ATen/core/TensorBody.h:470.)
[default2]: return self._grad

```
The training does not halt and the logs looks fine but the warning talks about the gradients not being back propagated is concerning. It seems that TP=1, creates non leaf tensors and TP>1 creates leaf tensors which is rather confusing to me. From [here](https://pytorch.org/docs/stable/generated/torch.Tensor.is_leaf.html), the tensors that are result of an operation are not leaf tensors, but again why only for TP=1 ?

From my observations, the warning happens only in the beginning of the training and the number of times the error appears is equal to the number of models trained(Data Parallel,DP) times the number of pipeline passes. For example for a 6.7B parameter model:
1. Nodes=4; GPUs=16;TP=1; PP=2 => DP = 8 ; number of pipleline pass = PP-1=1; error appears 8*1= 8 times
2. Nodes=8; GPUSs=32;TP=1; PP=8 => DP=4; number of pipeline passes = PP-1=7 ; error appears 4*7=28 times

Also looking at the pipeline communication [here](https://github.com/bigscience-workshop/Megatron-DeepSpeed/blob/7b5f175b73a12d602cdadd3ee205ed666f6d4234/megatron/p2p_communication.py#L24); the tensors are communicated with `required_grad=True` flag.

# System and Repo specifics:
- PyTorch Version : 1.11
- PyTorch Cuda Version : 11.5
- NVCC Version : 11.5
- NCCL : 2.12.7+cuda11.5
- Deepspeed : https://github.com/microsoft/DeepSpeed/tree/0f5c2012ce19c36936be562094ef3d73b230e364 (Deepspeed wheel compiled with torch 1.12, cuda 11.6)
- Megatron-DeepSpeed : https://github.com/bigscience-workshop/Megatron-DeepSpeed/tree/7b5f175b73a12d602cdadd3ee205ed666f6d4234
- Apex : https://github.com/NVIDIA/apex/tree/21e415479b134309a2ba1af95b2319b7bf068f7a
- GPU model: NVIDIA A100 Tensor Core GPU with 40 GB; Each node contains 4 GPUs connected via NVLink3 to each other. More info on system [here](https://apps.fz-juelich.de/jsc/hps/juwels/booster-overview.html)

# Example Launch Command
```
>>Megatron-DeepSpeed/pretrain_gpt.py --tensor-model-parallel-size 1 --pipeline-model-parallel-size 2 --num-layers 32 --hidden-size 4096 --num-attention-heads 32 --seq-length 2048 --max-position-embeddings 2048 --micro-batch-size 2 --global-batch-size 2048 --train-samples 69_335_938 --vocab-file vocab.json --merge-file merges.txt --loss-scale 12 --fp16 --seed 42 --checkpoint-activations --train-tokens 142_000_000_000 --optimizer adam --adam-beta1 0.9 --adam-beta2 0.95 --adam-eps 1e-8 --lr 1.2e-4 --min-lr 1.2e-5 --lr-decay-style cosine --lr-decay-samples 126_953_125 --lr-warmup-samples 183_105 --clip-grad 1.0 --weight-decay 1e-1 --log-interval 1 --save-interval 300 --eval-interval 300 --tensorboard-dir tensorboard --tensorboard-queue-size 5 --log-timers-to-tensorboard --log-batch-size-to-tensorboard --log-validation-ppl-to-tensorboard --save checkpoints --data-path merged_german_only --split 949,50,1 --data-impl mmap --distributed-backend nccl --deepspeed --deepspeed_config ds_config.6416902.json --zero-stage 1 --deepspeed-activation-checkpointing
```

Contributor guide

No contributing guide indexed for this repository

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.