Qwen3.5 GRPO training fails with ZeRO-3 + gradient checkpointing enabled

[lpenghao]

I encountered an error when training Qwen3.5 with GRPO.

My setup is:

DeepSpeed ZeRO-3

gradient_checkpointing: true

bf16 training

The training crashes during backward with:

[rank0]: torch.utils.checkpoint.CheckpointError: torch.utils.checkpoint: Recomputed values for the following tensors have different metadata than during the forward pass.
[rank0]: tensor at position 6:
[rank0]: saved metadata: {'shape': torch.Size([12288, 5120]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 13:
[rank0]: saved metadata: {'shape': torch.Size([1024, 5120]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 20:
[rank0]: saved metadata: {'shape': torch.Size([1024, 5120]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 36:
[rank0]: saved metadata: {'shape': torch.Size([5120, 6144]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 43:
[rank0]: saved metadata: {'shape': torch.Size([17408, 5120]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: tensor at position 46:
[rank0]: saved metadata: {'shape': torch.Size([17408, 5120]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}
[rank0]: recomputed metadata: {'shape': torch.Size([0]), 'dtype': torch.bfloat16, 'device': device(type='cuda', index=0)}

It looks like some parameters become empty tensors during checkpoint recomputation under ZeRO-3.

Is ZeRO-3 + gradient_checkpointing=true currently supported for GRPO training? How can I solve this error？

[2U1]

Were you using --lora_enable True --vision_lora True?

[lpenghao]

When finetune_grpo+zero3:
--freeze_vision_tower False
--freeze_llm True
--freeze_merger False
When finetune_lora+zero3:
--lora_enable True
--vision_lora True
--freeze_vision_tower True
--freeze_llm True
--freeze_merger False
I encountered this error in both of the above situations.