DeepSpeed ZeRO 3 issue

I followed all the instructions for training, ran the pre-training script on 1 image for debugging and was getting a lot of GPU memory usage and the following error:

"[rank0]: Traceback (most recent call last):
[rank0]: File "videollama2/train_flash_attn.py", line 12, in
[rank0]: train(attn_implementation="flash_attention_2")
[rank0]: File "/home/ayaan/VideoLLaMA2/./videollama2/train.py", line 964, in train
[rank0]: safe_save_model_for_hf_trainer(trainer=trainer, output_dir=training_args.output_dir)
[rank0]: File "/home/ayaan/VideoLLaMA2/./videollama2/train.py", line 226, in safe_save_model_for_hf_trainer
[rank0]: weight_to_save = get_mm_adapter_state_maybe_zero_3(trainer.model.named_parameters(), keys_to_match)
[rank0]: File "/home/ayaan/VideoLLaMA2/./videollama2/train.py", line 196, in get_mm_adapter_state_maybe_zero_3
[rank0]: to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
[rank0]: File "/home/ayaan/VideoLLaMA2/./videollama2/train.py", line 196, in
[rank0]: to_return = {k: maybe_zero_3(v, ignore_status=True).cpu() for k, v in to_return.items()}
[rank0]: File "/home/ayaan/VideoLLaMA2/./videollama2/train.py", line 154, in maybe_zero_3
[rank0]: param = param.data.detach().cpu().clone()
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 2230, in __exit__
[rank0]: self.params[0].partition(param_list=self.params, has_been_updated=False)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1375, in partition
[rank0]: self._partition(param_list, has_been_updated=has_been_updated)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1524, in _partition
[rank0]: self._partition_param(param, has_been_updated=has_been_updated)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 1557, in _partition_param
[rank0]: free_param(param)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/utils/nvtx.py", line 15, in wrapped_fn
[rank0]: ret_val = func(*args, **kwargs)
[rank0]: File "/home/ayaan/miniconda3/envs/myenv/lib/python3.8/site-packages/deepspeed/runtime/zero/partition_parameters.py", line 281, in free_param
[rank0]: assert not param.ds_active_sub_modules, param.ds_summary()
[rank0]: AssertionError: {'id': 682, 'status': 'AVAILABLE', 'numel': 4194304, 'ds_numel': 4194304, 'shape': (4096, 1024, 1, 1), 'ds_shape': (4096, 1024, 1, 1), 'requires_grad': True, 'grad_shape': None, 'persist': False, 'active_sub_modules': {755}, 'ds_tensor.shape': torch.Size([524288])}"

Did you guys encounter this and how can I fix this? I only get this error in DeepSpeed ZeRO 3 and not DeepSpeed ZeRO 2.

Start with videollama2/train_flash_attn.py and the maybe_zero_3 and safe_save_model_for_hf_trainer functions in videollama2/train.py, especially the ZeRO 3 save path around lines 154-226. Reproduce the one-image pre-training run with ZeRO 3, compare it with ZeRO 2, and trace the active submodule assertion. Done means the ZeRO 3 run completes saving without the reported assertion and the behavior is documented or covered by a regression test.

Written by the indexing model from the issue text.