On Policy Distillation checkpoints Loading

[White1973]

Does anybody would like to help me solve this problem.

When I finsh the Teacher Model pretraining bash examples/algorithms/dapo/run_dapo_aime_qwen3_4b_aime.sh and obtain ckpts/dapo_qwen3_4b_base/global_step_90. It contains data.pt, trainer_state.pt and policy folder, where policy folder includes huggingface folder (extra_state_world_size_8_rank_{0-8}.pt, model_world_size_8_rank_{0-7}, optim_world_size_8_rank_{0-8}.pt and fsdp_config.json)

when I set TEACHER_MODEL as "ckpts/dapo_qwen3_4b_base/global_step_90/" and execute "bash examples/on_policy_distillation/run_on_policy_distill_math_qwen3_4b.sh", I got the following error:

raise OSError(
OSError: Error no file named pytorch_model.bin, model.safetensors, tf_model.h5, model.ckpt.index or flax_model.msgpack found in directory /personal/jiayu2026/ckpts/dapo_qwen3_4b_base/global_step_150/.

[erictang000]

Hey @White1973, sorry for the late reply - this is because teacher model expects a HF model path (i.e. with model.safetensors), while your ckpts/dapo_qwen3_4b_base/global_step_90/ is a sharded checkpoint (https://skyrl.readthedocs.io/en/latest/checkpointing-logging/checkpointing.html#fsdp-checkpointing).

To save model parameters in safetensors format throughout training, you can use the trainer.hf_save_interval flag (although this can be expensive for disk storage, which is why we do not enable it by default). To merge your checkpoint into an hf format, you can modify your examples/algorithms/dapo/run_dapo_aime_qwen3_4b_aime.sh script to load from the existing checkpoint, and then set trainer.hf_save_interval to 1, or add a call to self.dispatch.save_hf_model("policy", policy_export_dir, self.tokenizer) after this line in the trainer:

SkyRL/skyrl-train/skyrl_train/trainer.py

Line 180 in 6a9ab82

self.global_step, _ = self.load_checkpoints()

if you want to avoid taking another optimizer step.

Let me know if that works for you!