[BUG] Client Optimizer is not initiated while running DeepSpeed Megatron-LM GPT2

[rraminen]

Dataset: Wikipedia
Docker Image: pytorch/pytorch:1.8.1-cuda11.1-cudnn8-devel
GPU: V100
MP_SIZE=8
NUM_WORKERS=1
NUM_GPUS_PER_WORKER=8

Steps to run:
cd DeepSpeed/DeepSpeedExamples/Megatron-LM
bash scripts/ds_zero2_pretrain_gpt2_model_parallel.sh

Following error is observed while running DeepSpeed Megatron-LM GPT2 withwikipedia dataset

Traceback (most recent call last):
File "pretrain_gpt2.py", line 713, in
main()
File "pretrain_gpt2.py", line 661, in main
model, optimizer, lr_scheduler = setup_model_and_optimizer(args)
File "pretrain_gpt2.py", line 167, in setup_model_and_optimizer
model, optimizer, _, lr_scheduler = deepspeed.initialize(
File "/opt/conda/lib/python3.8/site-packages/deepspeed/init.py", line 131, in initialize
engine = DeepSpeedEngine(args=args,
File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 166, in init
self._do_sanity_check()
File "/opt/conda/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 681, in _do_sanity_check
assert isinstance(self.client_optimizer, Optimizer),
AssertionError: Client Optimizer (type = <class 'NoneType'> is not instantiated but Client LR Scheduler is instantiated

ds_report output:

DeepSpeed C++/CUDA extension op report

NOTE: Ops not installed will be just-in-time (JIT) compiled at
runtime if needed. Op compatibility means that your system
meet the required dependencies to JIT install the op.

JIT compiled ops requires ninja
ninja .................. [OKAY]

op name ................ installed .. compatible

cpu_adam ............... [YES] ...... [OKAY]
fused_adam ............. [YES] ...... [OKAY]
fused_lamb ............. [YES] ...... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [YES] ...... [OKAY]
stochastic_transformer . [YES] ...... [OKAY]
[WARNING] async_io requires the dev libaio .so object and headers but these were not found.
[WARNING] async_io: please install the libaio-dev package with apt
[WARNING] If libaio is already installed (perhaps from source), try setting the CFLAGS and LDFLAGS environment variables to where it can be found.
async_io ............... [NO] ....... [NO]
transformer_inference .. [NO] ....... [OKAY]
utils .................. [YES] ...... [OKAY]
quantizer .............. [NO] ....... [OKAY]

DeepSpeed general environment info:
torch install path ............... ['/opt/conda/lib/python3.8/site-packages/torch']
torch version .................... 1.8.1
torch cuda version ............... 11.1
nvcc version ..................... 11.1
deepspeed install path ........... ['/opt/conda/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.5.3+a708b18, a708b18, master
deepspeed wheel compiled w. ...... torch 1.8, cuda 11.1

[tjruwase]

@rraminen. Thanks for reporting this. A brief suggest this is a regression from #1316. I will take a closer look.

[tjruwase]

@rraminen, this failure is actually due to a bug in the DeepSpeed Megatron-LM that is now exposed by tighter correctness checking in deepspeed.initialize(). However, that fixing that bug is not worth it because that example code is old and no longer compatible with latest Megatron or DeepSpeed. So I suggest using the example codes here. Do not be deterred by the ZeRO3 in the folder name, the code should work well with ZeRO stage 2. I just pushed a ZeRO2 script to help get started. One big difference is that new Megatron uses a different dataset format, but hopefully the README will bring you up to speed.

[rraminen]

Hi @tjruwase, thank you for the information. I shall try it out.