Describe the bug
Git commit: Repro on 44af130
Crashes when check-pointing when saving OPTIMIZER_OFFLOAD=1, PR=fp8, when it hits for the save interval from the DeepSeek configs from Megatron Model Zoo.
A clear and concise description of what the bug is
Steps/Code to reproduce bug
24: [rank24]: Traceback (most recent call last):
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/pretrain_gpt.py", line 233, in <module>
24: [rank24]: pretrain(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 737, in pretrain
24: [rank24]: iteration, num_floating_point_operations_so_far = train(
24: [rank24]: ^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 2405, in train
24: [rank24]: should_exit = checkpoint_and_decide_exit(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1850, in checkpoint_and_decide_exit
24: [rank24]: save_checkpoint_and_time(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/training.py", line 1721, in save_checkpoint_and_time
24: [rank24]: save_checkpoint(
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 509, in save_checkpoint
24: [rank24]: state_dict = generate_state_dict(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/training/checkpointing.py", line 817, in generate_state_dict
24: [rank24]: optimizer_sd = optimizer.sharded_state_dict(state_dict, **(optim_sd_kwargs or {}))
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/optimizer.py", line 1204, in sharded_state_dict
24: [rank24]: optim_state_dict = optimizer.sharded_state_dict(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1252, in sharded_state_dict
24: [rank24]: param_state = self.sharded_param_state_dp_reshardable(
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: File "/mnt/sharefs/users/runner/joshcopy/a2aoverlap/Megatron-MoE-ModelZoo/Megatron-LM/megatron/core/optimizer/distrib_optimizer.py", line 1646, in sharded_param_state_dp_reshardable
24: [rank24]: assert tensors[key].shape == (gbuf_local_end - gbuf_local_start,), (
24: [rank24]: ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
24: [rank24]: AssertionError: (torch.Size([]), 0, 1050820)
25: [rank25]: Traceback (most recent call last):Not sure if this interacts with the singleton save. Turning OPTIMIZER_OFFLOAD=0 seems to mitigate this but need to double check.
A clear and concise description of what you expected to happen.
Checkpointing saving works properly when optimizer offload is enabled.
Additional context
Add any other context about the problem here.
Describe the bug
Git commit: Repro on 44af130
Crashes when check-pointing when saving OPTIMIZER_OFFLOAD=1, PR=fp8, when it hits for the save interval from the DeepSeek configs from Megatron Model Zoo.
A clear and concise description of what the bug is
Steps/Code to reproduce bug
Not sure if this interacts with the singleton save. Turning OPTIMIZER_OFFLOAD=0 seems to mitigate this but need to double check.
A clear and concise description of what you expected to happen.
Checkpointing saving works properly when optimizer offload is enabled.
Additional context
Add any other context about the problem here.