Unable to Resume Training from Checkpoints – Possible Issue in Checkpoint Loading Logic

[xytsakura]

Hi, thanks for your amazing work on Search-R1!

I am facing an issue when trying to resume training from saved checkpoints. During training, the checkpoints are successfully generated under the checkpoint-* directories. However, when I restart the training script after interruption, the model does not properly load the previous checkpoint and instead initializes some weights from scratch.

Here is the warning message I consistently get:
(WorkerDict pid=3488384) Some weights of Qwen2ForTokenClassification were not initialized from the model checkpoint at /data/USTC_Kevin/yantai/search_r1_project/models/Qwen2.5-1.5B-Instruct and are newly initialized: ['score.bias'] [repeated 2x across cluster]

🙏 My Question

Could you please help confirm:

1.Which file/function in the repository handles loading existing checkpoints when training restarts?

2.Is there any additional flag or configuration required to ensure resume-from-checkpoint works properly?

3.Why would some parameters like score.bias be initialized instead of loaded—does this indicate a mismatch between model head and saved state dict?

[xytsakura]

my train_ppo.sh:

export CUDA_VISIBLE_DEVICES=4,5
export DATA_DIR='data/nq_search'

WAND_PROJECT='Search-R1'

export BASE_MODEL='meta-llama/Llama-3.2-3B'

export EXPERIMENT_NAME=nq-search-r1-ppo-llama3.2-3b-em

export BASE_MODEL='/data/USTC_Kevin/yantai/search_r1_project/models/Qwen2.5-1.5B-Instruct'
export EXPERIMENT_NAME=nq-search-r1-ppo-qwen2.5-1.5b-it-em

export BASE_MODEL='meta-llama/Llama-3.2-3B-Instruct'

export EXPERIMENT_NAME=nq-search-r1-ppo-llama3.2-3b-it-em

export BASE_MODEL='meta-llama/Llama-3.2-3B-Instruct'

export EXPERIMENT_NAME=nq-search-r1-ppo-llama3.2-3b-it-em

export BASE_MODEL='meta-llama/Llama-3.1-8B'

export EXPERIMENT_NAME=nq-search-r1-ppo-llama3.1-8b-em

export BASE_MODEL='meta-llama/Llama-3.1-8B-Instruct'

export EXPERIMENT_NAME=nq-search-r1-ppo-llama3.1-8b-it-em

export BASE_MODEL='Qwen/Qwen2.5-3B'

export EXPERIMENT_NAME=nq-search-r1-ppo-qwen2.5-3b-em

export BASE_MODEL='Qwen/Qwen2.5-3B-Instruct'

export EXPERIMENT_NAME=nq-search-r1-ppo-qwen2.5-3b-it-em

export BASE_MODEL='Qwen/Qwen2.5-7B'

export EXPERIMENT_NAME=nq-search-r1-ppo-qwen2.5-7b-em

export BASE_MODEL='Qwen/Qwen2.5-7B-Instruct'

export EXPERIMENT_NAME=nq-search-r1-ppo-qwen2.5-7b-it-em

set -x

export VLLM_ATTENTION_BACKEND=XFORMERS # vllm + qwen2-7b with flash_attn has some issues

max_prompt_length = (config['training']['max_start_length'] + config['training']['max_response_length'] * (config['training']['max_turns'] - 1) + config['training']['max_obs_length'] * config['training']['max_turns'])

PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo
data.train_files=$DATA_DIR/train.parquet
data.val_files=$DATA_DIR/test.parquet
data.train_data_num=null
data.val_data_num=128
data.train_batch_size=32
data.val_batch_size=32
data.max_prompt_length=3072
data.max_response_length=500
data.max_start_length=2048
data.max_obs_length=580
data.shuffle_train_dataloader=True
algorithm.adv_estimator=gae
actor_rollout_ref.model.path=$BASE_MODEL
actor_rollout_ref.actor.optim.lr=1e-6
actor_rollout_ref.model.enable_gradient_checkpointing=true
actor_rollout_ref.model.use_remove_padding=False
actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.285
actor_rollout_ref.actor.ppo_mini_batch_size=16
actor_rollout_ref.actor.ppo_micro_batch_size=2
actor_rollout_ref.actor.fsdp_config.param_offload=true
actor_rollout_ref.actor.fsdp_config.grad_offload=false
actor_rollout_ref.actor.fsdp_config.optimizer_offload=false
actor_rollout_ref.rollout.log_prob_micro_batch_size=2
actor_rollout_ref.rollout.tensor_model_parallel_size=1
actor_rollout_ref.rollout.name=hf
actor_rollout_ref.rollout.gpu_memory_utilization=0.6
actor_rollout_ref.ref.log_prob_micro_batch_size=2
actor_rollout_ref.ref.fsdp_config.param_offload=True
actor_rollout_ref.rollout.n_agent=1
actor_rollout_ref.rollout.temperature=1
actor_rollout_ref.actor.state_masking=true
critic.optim.lr=1e-5
critic.model.use_remove_padding=False
critic.optim.lr_warmup_steps_ratio=0.015
critic.model.path=$BASE_MODEL
critic.model.enable_gradient_checkpointing=true
critic.ppo_micro_batch_size=4
critic.model.fsdp_config.param_offload=true
critic.model.fsdp_config.grad_offload=false
critic.model.fsdp_config.optimizer_offload=false
algorithm.kl_ctrl.kl_coef=0.001
algorithm.no_think_rl=false
trainer.critic_warmup=1
trainer.logger=['wandb']
+trainer.val_only=false
+trainer.val_before_train=false
trainer.default_hdfs_dir=null
trainer.n_gpus_per_node=2
trainer.nnodes=1
trainer.save_freq=2
trainer.test_freq=50
trainer.project_name=$WAND_PROJECT
trainer.experiment_name=$EXPERIMENT_NAME
trainer.total_epochs=2
trainer.total_training_steps=200
trainer.default_hdfs_dir=null
trainer.default_local_dir=verl_checkpoints/$EXPERIMENT_NAME
max_turns=2
retriever.url="http://127.0.0.1:8000/retrieve"
retriever.topk=3
2>&1 | tee $EXPERIMENT_NAME.log

from flash_attn.bert_padding import pad_input, unpad_input