failed to start On-policy distillation #4390

Reproduction

Failed to start training with the given script in section Building it for yourself Distillation Recipe.

accelerate launch \
  --config_file examples/accelerate_configs/multi_gpu.yaml trl/experimental/gold/gold.py \
  --model_name_or_path <sft-model> \
  --dtype auto \
  --attn_implementation kernels-community/flash-attn \
  --dataset_name allenai/tulu-3-sft-mixture \
  --dataset_train_split train \
  --bf16 \
  --learning_rate 1e-7 \
  --gradient_checkpointing \
  --per_device_train_batch_size 1 \
  --gradient_accumulation_steps 64 \
  --num_train_epochs 1 \
  --eval_strategy steps \
  --eval_steps 100 \
  --temperature 1.0 \
  --top_p 0.95 \
  --top_k 0 \
  --max_new_tokens 2048 \
  --max_prompt_length 512 \
  --lmbda 0.25 \
  --beta 0.0 \
  --use_uld_loss \
  --use_extended_uld \
  --uld_use_hybrid_loss \
  --uld_crossentropy_weight 0.0 \
  --uld_distillation_weight 1.0 \
  --uld_student_temperature 1.0 \
  --uld_teacher_temperature 1.0 \
  --uld_hybrid_unmatched_weight 1.0 \
  --uld_hybrid_matched_weight 1.0 \
  --teacher_model_name_or_path Qwen/Qwen3-4B-Instruct-2507 \
  --logging_steps 1 \
  --push_to_hub \
  --hub_model_id <your-username>/Qwen3-4B-GKD-Tulu \
  --report_to trackio \
  --trackio_project Qwen3-4B-GKD-Tulu \
  --seed 42 \
  --warmup_ratio 0.05 \
  --lr_scheduler_type cosine_with_min_lr

System Info

• Platform: Linux-5.4.250-2-velinux1u1-amd64-x86_64-with-glibc2.35
• Python version: 3.12.12
• TRL version: 0.25.0.dev0+709866b
• PyTorch version: 2.8.0
• accelerator(s): NVIDIA A100-SXM4-80GB
• Transformers version: 4.57.1
• Accelerate version: 1.11.0
• Accelerate config:
  • compute_environment: LOCAL_MACHINE
  • distributed_type: DEEPSPEED
  • mixed_precision: bf16
  • use_cpu: False
  • debug: False
  • num_processes: 8
  • machine_rank: 0
  • num_machines: 1
  • main_process_ip: localhost
  • main_process_port: 9090
  • rdzv_backend: static
  • same_network: True
  • main_training_function: main
  • enable_cpu_affinity: False
  • deepspeed_config: {'deepspeed_multinode_launcher': 'standard', 'gradient_accumulation_steps': 8, 'gradient_clipping': 1.0, 'offload_optimizer_device': 'none', 'offload_param_device': 'none', 'zero3_init_flag': False, 'zero_stage': 2}
  • downcast_bf16: no
  • tpu_use_cluster: False
  • tpu_use_sudo: False
  • tpu_env: []
• Datasets version: 4.3.0
• HF Hub version: 0.36.0
• bitsandbytes version: not installed
• DeepSpeed version: 0.18.1
• Liger-Kernel version: 0.6.3
• LLM-Blender version: not installed
• OpenAI version: 2.6.1
• PEFT version: 0.17.1
• vLLM version: 0.11.0

Checklist

• I have checked that my issue isn't already filed (see open issues)
• I have included my system information
• Any code provided is minimal, complete, and reproducible (more on MREs)
• Any code provided is properly formatted in code blocks, (no screenshot, more on code blocks)
• Any traceback provided is complete