DeepSpeed hangs with iGPT
🐛 Bug
iGPT has caused issues with FairScale Sharded DDP before, so it's not a surprise DeepSpeed has some issues with running this model.
When training with ZeRO Optimization, DeepSpeed crashes:
RuntimeError: expected scalar type Float but found Half
    result = self.forward(*input, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/modules/normalization.py", line 170, in forward
    return F.layer_norm(
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/functional.py", line 2202, in layer_norm
    result = self.forward(*input, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/pl_bolts/models/vision/image_gpt/gpt2.py", line 23, in forward
    x = self.ln_1(x)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/modules/module.py", line 889, in _call_impl
    result = self.forward(*input, **kwargs)
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/modules/normalization.py", line 170, in forward
    return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
RuntimeError: expected scalar type Float but found Half
    return F.layer_norm(
  File "/home/sean/miniconda3/lib/python3.8/site-packages/torch/nn/functional.py", line 2202, in layer_norm
    return torch.layer_norm(input, normalized_shape, weight, bias, eps, torch.backends.cudnn.enabled)
RuntimeError: expected scalar type Float but found Half

This is due to this line here:
https://github.com/PyTorchLightning/pytorch-lightning-bolts/blob/master/pl_bolts/models/vision/image_gpt/gpt2.py#L97
That doesn't take in the dtype or do a type_as conversion. We've managed to deal with this for now since when precision=16 we use PyTorch Native AMP which converts arguments automatically, but in DeepSpeed they use a custom FP16_Optimizer for ZeRO which does not handle these conversions automatically.
One fix for this would be for DeepSpeed to bring Native AMP support, and I'll make a ticket for this, but short term we can add a type_as conversion for lightning-bolts:
        sos = torch.ones(1, batch, self.hparams.embed_dim).type_as(x) * self.sos

However after remedying the above issue, the code hangs after a few iterations (I'm going to assume that this iteration may be where an optimizer step happens/takes place).
...
Epoch 0:   0%|▏                                                        | 14/3750 [00:03<15:59,  3.89it/s, loss=3.11, v_num=54]
[2021-02-18 13:31:07,760] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 262144.0, reducing to 131072.0
[2021-02-18 13:31:07,760] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 1 Skipping step. Attempted loss scale: 262144.0, reducing to 131072.0
[2021-02-18 13:31:07,760] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 2 Skipping step. Attempted loss scale: 262144.0, reducing to 131072.0
[2021-02-18 13:31:07,760] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 3 Skipping step. Attempted loss scale: 262144.0, reducing to 131072.0
Epoch 0:   0%|▏                                                        | 15/3750 [00:03<15:22,  4.05it/s, loss=3.11, v_num=54]
[2021-02-18 13:31:07,868] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 1 Skipping step. Attempted loss scale: 131072.0, reducing to 65536.0
[2021-02-18 13:31:07,868] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 2 Skipping step. Attempted loss scale: 131072.0, reducing to 65536.0
[2021-02-18 13:31:07,868] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 131072.0, reducing to 65536.0
[2021-02-18 13:31:07,868] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 3 Skipping step. Attempted loss scale: 131072.0, reducing to 65536.0
Epoch 0:   0%|▏                                                        | 16/3750 [00:03<14:50,  4.19it/s, loss=3.11, v_num=54]
[2021-02-18 13:31:07,976] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 65536.0, reducing to 32768.0
[2021-02-18 13:31:07,976] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 1 Skipping step. Attempted loss scale: 65536.0, reducing to 32768.0
[2021-02-18 13:31:07,977] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 3 Skipping step. Attempted loss scale: 65536.0, reducing to 32768.0
[2021-02-18 13:31:07,977] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 2 Skipping step. Attempted loss scale: 65536.0, reducing to 32768.0
Epoch 0:   0%|▎                                                        | 17/3750 [00:03<14:21,  4.33it/s, loss=3.11, v_num=54]
[2021-02-18 13:31:08,088] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 1 Skipping step. Attempted loss scale: 32768.0, reducing to 16384.0
[2021-02-18 13:31:08,088] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 0 Skipping step. Attempted loss scale: 32768.0, reducing to 16384.0
[2021-02-18 13:31:08,088] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 3 Skipping step. Attempted loss scale: 32768.0, reducing to 16384.0
[2021-02-18 13:31:08,088] [INFO] [stage2.py:1357:step] [deepscale] OVERFLOW! Rank 2 Skipping step. Attempted loss scale: 32768.0, reducing to 16384.0
Epoch 0:   0%|▎                                                        | 18/3750 [00:04<13:56,  4.46it/s, loss=3.11, v_num=54] # stalls here

This could also be due to waiting for buckets that do not have gradients. For DDP, we have to set find_unused_parameters=True in order to run, so this might highlight an issue in the DeepSpeed distributed code.
Install PyTorch Lightning from master.
pip install lightning-bolts deepspeed
import logging

from pl_bolts.datamodules import MNISTDataModule
from pl_bolts.models import ImageGPT

import pytorch_lightning as pl
from pytorch_lightning.plugins import DeepSpeedPlugin, DDPPlugin

dm = MNISTDataModule('.', batch_size=4)

model = ImageGPT(embed_dim=2048, layers=4, heads=4, vocab_size=32, num_pixels=28)

deepspeed_plugin = DeepSpeedPlugin(logging_level=logging.INFO)
ddp_plugin = DDPPlugin(find_unused_parameters=True)
trainer = pl.Trainer(
    max_epochs=1,
    gpus=4,
    precision=16,
    plugins=deepspeed_plugin,
)
trainer.fit(model, dm)
Environment

CUDA:
- GPU:
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- A100-SXM4-40GB
- available:         True
- version:           11.2
Packages:
- numpy:             1.20.1
- pyTorch_debug:     False
- pyTorch_version:   1.8.0+cu112
- pytorch-lightning: 1.2.0rc1
- tqdm:              4.56.0
System:
- OS:                Linux
- architecture:
- 64bit
- ELF
- processor:         x86_64
- python:            3.8.5