No HF format checkpoint been saved in GRPO training, only saw pytorch checkpoint

[wukaixingxp]

🐛 Describe the bug

I successfully run coding-grpo on 1.7B and get step-2500 checkpoint.. but this checkpoint is still pytorch checkpoint __0_0.distcp, even if I set last_save_in_hf: true.. looking closer, I can not find any code that uses this argument last_save_in_hf. Given that we have a way to update to vllm weights (which should be HF format).. i think we are just missing a conversion step in the end.

ls forge/1_7B_checkpoint/step-2500/
__0_0.distcp

Versions

I forked the main on my branch.

[DNXie]

I just ran it, I can see hf ckpt saved:

model-00001-of-00002.safetensors  model-00002-of-00002.safetensors

In your config, the training.step is set to 2000, so the hf ckpt should be saved to step-2000. Could you also check this folder?

[wukaixingxp]

Thanks for your help @DNXie, so it seems that only one HF checkpoint will be saved at the end of the training.. but I hope we can save a HF checkpoint every interval : 500, as there is a pytorch checkpoint saved every 500 steps. I think it will be easier for me to load the checkpoints to vllm to do some manual testing.

[DNXie]

@wukaixingxp Unfortunately, the current implementation doesn't support saving hf format in between. We should be able to support that in the future.

[DNXie]

To support this feature we need to touch the code in titan. Currently the switch on ckpt format is in dcp_save.
But when we call save, it doesn't support such parameter to be passed in and would call dcp_save with to_hf=False

cc @ebsmothers

[JenniferWang]

@wukaixingxp I don't understand why HF format checkpointing is needed along side with the PyTorch checkpointing? What are they used for?

[wukaixingxp]

@JenniferWang I am trying to do a coding grpo and would like to evaluate the intermediate checkpoints with livecodebench so I need to load them back to vllm. The final checkpoint does not always get the highest score on livecodebench as sometime the model collapse during the training. I only saw the Pytorch weights and did not find a way to convert it back to HF format.

[JenniferWang]

I don't fully understand the use case but I wonder if this what you need? @wukaixingxp

https://github.com/meta-pytorch/forge/blob/main/src/forge/actors/trainer.py#L343

[DNXie]

@JenniferWang Yes, it is the same that titan checkpointer would save at the last step.