stash

Author: JacobHelwig

examples/on_policy_distillation_trainer/run_qwen_gsmk8k.sh

@@ -5,6 +5,7 @@ export PATH=$CONDA_PREFIX/bin:$PATH
 export NCCL_P2P_DISABLE=1
 export CUDA_DEVICE_ORDER=PCI_BUS_ID
 export CUDA_VISIBLE_DEVICES=5,6,7,8
+# export CUDA_VISIBLE_DEVICES=7,8
 export DATA_PATH=$PWD/../verlData
 export HF_HOME=$DATA_PATH
 export VLLM_CACHE_DIR=$DATA_PATH/vllm_cache
@@ -17,14 +18,14 @@ ROLLOUT_NAME="vllm" # sglang or vllm
 
 FAMILY="Qwen"
 STUDENT_MODEL=Qwen2.5-0.5B
-TEACHER_MODEL=Qwen2.5-3B-Instruct
+TEACHER_MODEL=Qwen2.5-0.5B-Instruct
 
 USE_POLICY_GRADIENT=False
-DISTILLATION_LOSS_MODE="k3"
+# DISTILLATION_LOSS_MODE="k3"
 DISTILLATION_LOSS_MODE="forward_kl_topk"
 
-DISTILLATION_LOSS_MODE="k1"
-USE_POLICY_GRADIENT=True
+# USE_POLICY_GRADIENT=True
+# DISTILLATION_LOSS_MODE="k1"
 
 DISTILLATION_LOSS_MAX_CLAMP=10.0
 DISTILLATION_LOG_PROB_MIN_CLAMP=null
@@ -34,7 +35,7 @@ EXP_NAME="${FAMILY}/student-${STUDENT_MODEL}/teacher-${TEACHER_MODEL}/loss-${DIS
 
 MAX_PROMPT=256
 MAX_RESPONSE_LENGTH=512
-TRAIN_PROMPT_BSZ=128
+TRAIN_PROMPT_BSZ=8
 STUDENT_MICRO_BATCH_SIZE_PER_GPU=2
 STUDENT_MAX_TOKEN_LEN_PER_GPU=$(( STUDENT_MICRO_BATCH_SIZE_PER_GPU * (MAX_PROMPT + MAX_RESPONSE_LENGTH) ))
 USE_DYNAMIC_BSZ=False
@@ -44,7 +45,7 @@ STUDENT_WORLD_SIZE=2
 TEACHER_RESOURCE_POOL=True
 TEACHER_WORLD_SIZE=2
 
-ENFORCE_EAGER=False # true for faster debugging
+ENFORCE_EAGER=True # true for faster debugging
 
 ############################ Paths ############################
 
@@ -77,14 +78,14 @@ MODEL=(
 
 DISTILLATION=(
     distillation.enabled=True
-    distillation.num_workers=8
+    distillation.num_workers=1
     distillation.teacher_model.enable_resource_pool=$TEACHER_RESOURCE_POOL
     distillation.teacher_model.n_gpus_per_node=$TEACHER_WORLD_SIZE
     distillation.teacher_model.nnodes=1
     distillation.teacher_model.model_path="${FAMILY}/${TEACHER_MODEL}"
     distillation.teacher_model.inference.tensor_model_parallel_size=1
     distillation.teacher_model.inference.name=$ROLLOUT_NAME
-    distillation.teacher_model.inference.gpu_memory_utilization=0.3
+    distillation.teacher_model.inference.gpu_memory_utilization=0.6
     distillation.teacher_model.inference.enforce_eager=$ENFORCE_EAGER
     distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE
     distillation.distillation_loss.topk=64

verl/trainer/distillation/utils.py

@@ -39,6 +39,7 @@ def prepare_student_distillation_inputs(
 ) -> dict[str, torch.Tensor]:
     """Prepare student distillation inputs."""
     stage = batch.get("stage", None)
+    breakpoint()
     if not is_distillation_enabled(config) or stage != Stage.ACTOR_UPDATE:
         return {}
     loss_config: DistillationLossConfig = config.distillation_loss