Skip to content

[BUG] Strange Behaviour on Megatron-DeepSpeed for training GPT2 model using ZeRO-3 on DeepSpeed>0.6.0 #2185

Description

@hibagus

Describe the bug
I was able to run the provided example on Megatron-DeepSpeed repo to train the GPT2 model using ZeRO stage 3 and DeepSpeed version 0.6.0.

The same script output repeated Warning Message for DeepSpeed version 0.6.3, 0.6.4, 0.6.5, 0.6.6, and 0.6.7.

[WARNING] [stage3.py:106:_apply_to_tensors_only] A module has unknown inputs or outputs type (<class 'torch.nn.parameter.Parameter'>) and the tensors embedded in it cannot be detected. The ZeRO-3 hooks designed to trigger before or after the backward pass of the module relies on knowing the input and output tensors and therefore may not get triggered properly.

I see that this warning is added to stage3.py on 0.6.3 and is moved to the separate file parameter_offload.py on 0.6.6.

If I comment out that warning, the training can be run. Is it okay to ignore this problem? Will it affect training performance?

If I run the same script using the newest DeepSpeed version 0.7.0, the following error message is produced and training cannot continue.

Traceback (most recent call last):
  File "/home/bagus/DeepSpeed/Megatron-DeepSpeed/pretrain_gpt.py", line 276, in <module>
    pretrain(train_valid_test_datasets_provider, model_provider, forward_step,
  File "/home/bagus/DeepSpeed/Megatron-DeepSpeed/megatron/training.py", line 130, in pretrain
    model, optimizer, lr_scheduler = setup_model_and_optimizer(model_provider, teacher=False)
  File "/home/bagus/DeepSpeed/Megatron-DeepSpeed/megatron/training.py", line 420, in setup_model_and_optimizer
    model, optimizer, _, lr_scheduler = deepspeed.initialize(
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/__init__.py", line 124, in initialize
    engine = DeepSpeedEngine(args=args,
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 319, in __init__
    self._configure_optimizer(optimizer, model_parameters)
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1133, in _configure_optimizer
    self.optimizer = self._configure_zero_optimizer(basic_optimizer)
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/engine.py", line 1436, in _configure_zero_optimizer
    optimizer = DeepSpeedZeroOptimizer_Stage3(
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 308, in __init__
    self._setup_for_real_optimizer()
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 363, in _setup_for_real_optimizer
    self.initialize_optimizer_states()
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 923, in initialize_optimizer_states
    self._optimizer_step(i)
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed/runtime/zero/stage3.py", line 843, in _optimizer_step
    self.optimizer.step()
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/torch/optim/optimizer.py", line 109, in wrapper
    return func(*args, **kwargs)
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/apex/optimizers/fused_adam.py", line 180, in step
    multi_tensor_applier(self.multi_tensor_adam,
  File "/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/apex/multi_tensor_apply/multi_tensor_apply.py", line 27, in __call__
    return op(self.chunk_size,
RuntimeError: expected input to be on cuda

Any suggestion on how to debug?

To Reproduce
Steps to reproduce the behavior:

  1. Clone the Github Repo https://github.com/microsoft/Megatron-DeepSpeed
  2. Use Wikipedia Dump as Dataset
  3. Run the DeepSpeed Example Script run_deepspeed_example.sh after adjusting the environment.
  4. Try with different DeepSpeed versions.

Expected behavior
I expect better performance than DeepSpeed 0.6.0 in ZeRO stage 3.

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 ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
fused_adam ............. [NO] ....... [OKAY]
fused_lamb ............. [NO] ....... [OKAY]
sparse_attn ............ [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
async_io ............... [NO] ....... [OKAY]
utils .................. [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]
--------------------------------------------------
DeepSpeed general environment info:
torch install path ............... ['/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/torch']
torch version .................... 1.12.0+cu113
torch cuda version ............... 11.3
torch hip version ................ None
nvcc version ..................... 11.6
deepspeed install path ........... ['/home/bagus/anaconda3/envs/MegatronDeepSpeed/lib/python3.8/site-packages/deepspeed']
deepspeed info ................... 0.7.0, unknown, unknown
deepspeed wheel compiled w. ...... torch 1.12, cuda 11.3

Screenshots
Not Applicable

System info (please complete the following information):

  • OS: Ubuntu 20.04
  • GPU: 4 x NVIDIA A100 40GB SXM4
  • Interconnects: Mellanox 200GbE SN3700 switch, 2 x ConnectX-6 in each node.
  • Python 3.8

Launcher context
DeepSpeed Launcher as suggested by the sample script.

Docker context
Bare-metal run

Additional context
None

Activity

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Metadata

Metadata

Assignees

No one assigned

    Labels

    bugSomething isn't working

    Type

    No type

    Projects

    No projects

      Milestone

      No milestone

      Relationships

      None yet

      Development

      No branches or pull requests

      Issue actions