[BUG] Error: Preparing DeepSpeed ZeRO stage 2 optimizer HIP error on AMD MI200 GPUs

Describe the bug
Error: Preparing DeepSpeed ZeRO stage 2 optimizer HIP error on AMD MI200 GPUs

ds_report output
ds_report
[2024-02-08 10:19:24,716] [INFO] [real_accelerator.py:191:get_accelerator] Setting ds_accelerator to cuda (auto detect)

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

[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]
fused_adam ............. [NO] ....... [OKAY]
cpu_adam ............... [NO] ....... [OKAY]
cpu_adagrad ............ [NO] ....... [OKAY]
cpu_lion ............... [NO] ....... [OKAY]
[WARNING] Please specify the CUTLASS repo directory as environment variable $CUTLASS_PATH
evoformer_attn ......... [NO] ....... [NO]
fused_lamb ............. [NO] ....... [OKAY]
fused_lion ............. [NO] ....... [OKAY]
inference_core_ops ..... [NO] ....... [OKAY]
cutlass_ops ............ [NO] ....... [OKAY]
quantizer .............. [NO] ....... [OKAY]
ragged_device_ops ...... [NO] ....... [OKAY]
ragged_ops ............. [NO] ....... [OKAY]
random_ltd ............. [NO] ....... [OKAY]
[WARNING] sparse_attn is not compatible with ROCM
sparse_attn ............ [NO] ....... [NO]
spatial_inference ...... [NO] ....... [OKAY]
transformer ............ [NO] ....... [OKAY]
stochastic_transformer . [NO] ....... [OKAY]
transformer_inference .. [NO] ....... [OKAY]

DeepSpeed general environment info:
torch install path ............... ['/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch']
torch version .................... 2.1.1+git011de5c
deepspeed install path ........... ['/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed']
deepspeed info ................... 0.13.1, unknown, unknown
torch cuda version ............... None
torch hip version ................ 6.0.32830-d62f6a171
nvcc version ..................... None
deepspeed wheel compiled w. ...... torch 2.1, hip 6.0
shared memory (/dev/shm) size .... 427.71 GB

System info:

OS: Ubuntu 20.04
Python version: 3.9
DeepSpeed version: 0.13.1
PyTorch version: 2.1.1
HIP version: 6.0.32830-d62f6a171

Hardware:
=====ROCm System Management Interface =====
===== Concise Info =====
Device [Model : Revision] Temp Power Partitions SCLK MCLK Fan Perf PwrCap VRAM% GPU%
Name (20 chars) (Edge) (Avg) (Mem, Compute)

0 [0x0b0c : 0x00] 44.0°C 94.0W N/A, N/A 800Mhz 1600Mhz 0% auto 500.0W 0% 0%
AMD INSTINCT MI200 (
1 [0x0b0c : 0x00] 48.0°C N/A N/A, N/A 800Mhz 1600Mhz 0% auto 0.0W 0% 0%
AMD INSTINCT MI200 (
2 [0x0b0c : 0x00] 41.0°C 87.0W N/A, N/A 800Mhz 1600Mhz 0% auto 500.0W 0% 0%
AMD INSTINCT MI200 (
3 [0x0b0c : 0x00] 37.0°C N/A N/A, N/A 800Mhz 1600Mhz 0% auto 0.0W 0% 0%
AMD INSTINCT MI200 (
4 [0x0b0c : 0x00] 51.0°C 85.0W N/A, N/A 800Mhz 1600Mhz 0% auto 500.0W 0% 0%
AMD INSTINCT MI200 (
5 [0x0b0c : 0x00] 41.0°C N/A N/A, N/A 800Mhz 1600Mhz 0% auto 0.0W 0% 0%
AMD INSTINCT MI200 (
6 [0x0b0c : 0x00] 38.0°C 98.0W N/A, N/A 800Mhz 1600Mhz 0% auto 500.0W 0% 0%
AMD INSTINCT MI200 (
7 [0x0b0c : 0x00] 40.0°C N/A N/A, N/A 800Mhz 1600Mhz 0% auto 0.0W 0% 0%
AMD INSTINCT MI200 (

Docker context

pytorch_rocm6.0_ubuntu20.04_py3.9_pytorch_2.1.1.sif

Additional context

Terminal log:
2024-02-08 10:24:00,126] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed info: version=0.13.1, git-hash=unknown, git-branch=unknown
[2024-02-08 10:24:13,567] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Flops Profiler Enabled: False
[2024-02-08 10:24:13,569] [INFO] [logging.py:96:log_dist] [Rank 0] Using client Optimizer as basic optimizer
[2024-02-08 10:24:13,569] [INFO] [logging.py:96:log_dist] [Rank 0] Removing param_group that has no 'params' in the basic Optimizer
[2024-02-08 10:24:13,580] [INFO] [logging.py:96:log_dist] [Rank 0] DeepSpeed Basic Optimizer = AdamW
[2024-02-08 10:24:13,596] [INFO] [utils.py:56:is_zero_supported_optimizer] Checking ZeRO support for optimizer=AdamW type=<class 'torch.optim.adamw.AdamW'>
[2024-02-08 10:24:13,596] [INFO] [logging.py:96:log_dist] [Rank 0] Creating torch.bfloat16 ZeRO stage 2 optimizer
[2024-02-08 10:24:13,597] [INFO] [stage_1_and_2.py:143:init] Reduce bucket size 500,000,000
[2024-02-08 10:24:13,597] [INFO] [stage_1_and_2.py:144:init] Allgather bucket size 500,000,000
[2024-02-08 10:24:13,597] [INFO] [stage_1_and_2.py:145:init] CPU Offload: False
[2024-02-08 10:24:13,597] [INFO] [stage_1_and_2.py:146:init] Round robin gradient partitioning: False
[2024-02-08 10:24:29,532] [INFO] [utils.py:791:see_memory_usage] Before initializing optimizer states
[2024-02-08 10:24:29,535] [INFO] [utils.py:792:see_memory_usage] MA 37.72 GB Max_MA 50.27 GB CA 50.27 GB Max_CA 50 GB
[2024-02-08 10:24:29,535] [INFO] [utils.py:799:see_memory_usage] CPU Virtual Memory: used = 71.89 GB, percent = 14.3%
Traceback (most recent call last):
File "/projappl/project_xxx/code/Train_RL_MXL/train/fintune_generator.py", line 847, in
main(args)
File "/projappl/project_xxx/code/Train_RL_MXL/train/fintune_generator.py", line 816, in main
train(
File "/projappl/project_xxx/code/Train_RL_MXL/train/fintune_generator.py", line 456, in train
model, optimizer, training_dataloader, validation_dataloader, scheduler = setup_accelerator(
File "/projappl/project_xxx/code/Train_RL_MXL/train/fintune_generator.py", line 330, in setup_accelerator
model, optimizer, training_dataloader, val_dataloader, scheduler = accelerator.prepare(
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/accelerate/accelerator.py", line 1219, in prepare
result = self._prepare_deepspeed(*args)
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/accelerate/accelerator.py", line 1604, in _prepare_deepspeed
engine, optimizer, _, lr_scheduler = deepspeed.initialize(**kwargs)
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/init.py", line 171, in initialize
engine = DeepSpeedEngine(args=args,
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 308, in init
self._configure_optimizer(optimizer, model_parameters)
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1247, in _configure_optimizer
self.optimizer = self._configure_zero_optimizer(basic_optimizer)
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/runtime/engine.py", line 1503, in _configure_zero_optimizer
optimizer = DeepSpeedZeroOptimizer(
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 529, in init
self.initialize_optimizer_states()
File "/projappl/project_xxx/code/Train_RL_MXL/train-model/lib/python3.9/site-packages/deepspeed/runtime/zero/stage_1_and_2.py", line 668, in initialize_optimizer_states
self.optimizer.step()
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/optim/lr_scheduler.py", line 68, in wrapper
return wrapped(*args, **kwargs)
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/optim/optimizer.py", line 373, in wrapper
out = func(*args, **kwargs)
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/optim/optimizer.py", line 76, in _use_grad
ret = func(self, *args, **kwargs)
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/optim/adamw.py", line 173, in step
self._init_group(
File "/opt/conda/envs/py_3.9/lib/python3.9/site-packages/torch/optim/adamw.py", line 121, in _init_group
state["exp_avg"] = torch.zeros_like(
RuntimeError: HIP error: invalid argument
HIP kernel errors might be asynchronously reported at some other API call, so the stacktrace below might be incorrect.
For debugging consider passing HIP_LAUNCH_BLOCKING=1.
Compile with TORCH_USE_HIP_DSA to enable device-side assertions.

Start with deepspeed/runtime/zero/stage_1_and_2.py at initialize_optimizer_states, then inspect the torch AdamW step in the traceback. Run ds_report and reproduce with HIP_LAUNCH_BLOCKING=1 on the listed ROCm/MI200 setup; done means ZeRO stage 2 initialization completes without the reported HIP invalid-argument error.