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:
- Clone the Github Repo https://github.com/microsoft/Megatron-DeepSpeed
- Use Wikipedia Dump as Dataset
- Run the DeepSpeed Example Script run_deepspeed_example.sh after adjusting the environment.
- 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
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.
Any suggestion on how to debug?
To Reproduce
Steps to reproduce the behavior:
Expected behavior
I expect better performance than DeepSpeed 0.6.0 in ZeRO stage 3.
ds_report output
Screenshots
Not Applicable
System info (please complete the following information):
Launcher context
DeepSpeed Launcher as suggested by the sample script.
Docker context
Bare-metal run
Additional context
None