[fsdp,megatron,vllm,trainer,algo] feat: On-Policy Distillation (#5041)

# What does this PR do?

Adds on-policy distillation support across FSDP and Megatron backends.

Collaboration with @wuxibin89, including design guidance, restructuring,
and added support for parallelism.

Supports: 

- FSDP and Megatron engines 
- top-k distillation loss and KL estimator distillation losses 
- Supervised and policy-gradient-style updates
- Teacher logprobs computation using a vLLM teacher server
- LLM and VLM distillation
- FSDP sequence parallel
- Megatron context parallel and tensor parallel

## Losses 

1. top-k distillation loss: **forward** KL estimated using top-k logits
**from teacher**.
2. KL estimator distillation losses: **reverse** KL estimated using only
the log prob for the sampled token via the same estimators used by the
reference model (e.g., k1, k3)

## Updates

1. Supervised: distillation loss is directly backpropagated, as in
https://arxiv.org/abs/2306.13649
2. Policy gradient: negative distillation loss is used as a reward, as
in https://thinkingmachines.ai/blog/on-policy-distillation/

# Test

- LLM distillation with FSDP:
`examples/on_policy_distillation_trainer/run_qwen_gsmk8k.sh`
- VLM distillation with FSDP:
`examples/on_policy_distillation_trainer/run_qwen3_vl_geo3k.sh`
- LLM distillation with megatron:
`examples/on_policy_distillation_trainer/run_qwen_gsmk8k_megatron.sh`.


## Main results 

### LLM Distillation 

These experiments compare 3 training runs with student model
Qwen2.5-0.5B using
`examples/on_policy_distillation_trainer/run_qwen_gsmk8k.sh`:

1. Forward top-k KL with Qwen2.5-3B-Instruct teacher (gold)
2. Forward top-k KL with Qwen2.5-7B-Instruct teacher (green)
3. k3 estimator KL with Qwen2.5-7B-Instruct teacher (red)

#### GSM8K eval acc

<img width="969" height="641" alt="image"
src="https://github.com/user-attachments/assets/19c1bee7-b688-4d24-a41e-4426761a26f1"
/>

#### GSM8K train acc

<img width="972" height="639" alt="image"
src="https://github.com/user-attachments/assets/8f932649-5d45-4964-9d59-18a3706004d5"
/>

#### Distillation loss

<img width="963" height="633" alt="image"
src="https://github.com/user-attachments/assets/609f817b-2247-42df-86a8-5e07d637ea7c"
/>

### VLM Distillation 

- Data: Geometry3K
- Student: Qwen3-VL-2B-Instruct
- Teacher: Qwen3-VL-4B-Instruct
- OPD algo: k1 KL estimator as reward with policy gradient loss

#### Geo3K eval acc

<img width="967" height="640" alt="image"
src="https://github.com/user-attachments/assets/e511fedb-8cf1-4576-b214-e992984c7550"
/>

#### Geo3K train acc

<img width="963" height="630" alt="image"
src="https://github.com/user-attachments/assets/0eb0e281-218d-4283-825e-2d0b1aa095d4"
/>

#### Distillation loss

<img width="964" height="633" alt="image"
src="https://github.com/user-attachments/assets/7b539c60-2af8-4d30-a9d3-3cb4cf485847"
/>

## LLM Distillation: Top-k training stability

Clamping the top-k forward KL loss was needed for training stability.
These experiments compare 3 types of clamping:

1. No clamping (grey)
2. Clamping the distillation loss to a maximum value of 10 (blue)
3. Clamping the student and teacher log probs to a minimum value of -10
(gold)

### Distillation loss

<img width="947" height="633" alt="image"
src="https://github.com/user-attachments/assets/e90815f5-f745-4a07-bc41-dbb3eb16b1dc"
/>

### GSM8K eval acc

<img width="971" height="639" alt="image"
src="https://github.com/user-attachments/assets/d5cd1833-7e66-40f3-b0a3-fae6a9bc0d7b"
/>

### GSM8K train acc

<img width="958" height="636" alt="image"
src="https://github.com/user-attachments/assets/53742e8d-cd41-4483-943e-cea2b5f0c27b"
/>

## LLM Distillation: Policy-gradient results

While the VLM results in this PR use the k1 KL estimator with policy
gradient updates, all LLM distillation results outside of this section
rely on supervised updates. LLM distillation with policy gradient
updates are validated in this section:

1. Forward top-k KL with supervised update (green)
2. k1 estimator KL with policy gradient update (purple)
3. k3 estimator KL with supervised update (red)

While purple seems best, it also is generating responses that exceed the
maximum response length of 512.

### Distillation loss

<img width="971" height="627" alt="image"
src="https://github.com/user-attachments/assets/208b8625-23d7-46d5-982d-6ab1b5049b21"
/>


### GSM8K eval acc

<img width="976" height="632" alt="image"
src="https://github.com/user-attachments/assets/dd46c6d3-651a-45cb-8bc5-75a765a9a38e"
/>


### GSM8K train acc

<img width="969" height="634" alt="image"
src="https://github.com/user-attachments/assets/f2db7e16-5426-4fa2-b017-fb78b58b8dd4"
/>

### Response length

<img width="979" height="640" alt="image"
src="https://github.com/user-attachments/assets/f61fadd7-d1a8-4110-ba26-a5a2735f8107"
/>

## LLM Distillation: Megatron 

To verify parity of megatron engine with FSDP, these experiments compare
3 training runs with student model Qwen2.5-0.5B:

1. Forward top-k KL with Qwen2.5-7B-Instruct teacher + clamping log
probs to minimum value of -10.0 (teal)
2. Forward top-k KL with Qwen2.5-3B-Instruct teacher + clamping log
probs to minimum value of -10.0 (red)
3. Forward top-k KL with Qwen2.5-3B-Instruct teacher + clamping loss to
maximum value of 10.0 (blue)
4. k3 estimator reverse KL with Qwen2.5-7B-Instruct teacher + clamping
loss to maximum value of 10.0 (green)

The solid line uses megatron engine with TP=2, the dotted line uses
FSDP.

### GSM8K Eval Acc.

<img width="1656" height="839" alt="image"
src="https://github.com/user-attachments/assets/e43a31a9-f012-452f-b606-99ede49a5fce"
/>

### GSM8K Train Acc.

<img width="1653" height="838" alt="image"
src="https://github.com/user-attachments/assets/609b5917-dbb4-4a72-bbb8-acd4a7f4224a"
/>

### Distillation Loss

<img width="1651" height="846" alt="image"
src="https://github.com/user-attachments/assets/6842295a-aa9a-4b11-9153-b9f83130e352"
/>


### Grad Norm

<img width="988" height="636" alt="image"
src="https://github.com/user-attachments/assets/5883e67c-71e3-4e65-87e3-79e0a2527757"
/>

## LLM Distillation: Note on reverse KL

Initially, this PR included top-k reverse KL and top-k Jensen-Shannon
divergences (JSD interpolates between forward and reverse KL). For the
student distribution $q$ and teacher distribution $p$, the top-k reverse
KL is given by

$$
KL_{\text{top-}k}(q||p) = \sum_i \bf{1}(q_i\in
\text{top-}k)q_i\log\frac{q_i}{p_i}.
$$ 

Unfortunately, this was unstable. The reason is because one way to make
this loss small is to make $q_i$ as small as possible for all $q_i \in
\text{top}-k$. This can be seen from the logs tracking the amount of
mass captured in the top-$k$ probabilities:

<img width="1118" height="1286" alt="image"
src="https://github.com/user-attachments/assets/6e335263-7f4d-48d1-96f4-181f89b24e21"
/>

## LLM Distillation: Ablation: performance with more lenient parser

Note that the only loss used is the distillation loss (no rewards for
correctness on GSM8K). Any increase in the logged rewards=GSM8k accuracy
are an indirect result of minimizing the distillation loss. The reason
that the base model has Pass@1~=0 is because the default GSM8k answer
formatting (`#### 42`) is OOD for the model. The base model is answering
the questions correctly, but using incorrect formatting, so none of the
answers can be parsed. The base model can be evaluated using a reward
function that is more lenient on formatting by adding the following to
the script:

```bash
...
    reward_model.reward_manager=remote \
    custom_reward_function.path=tests/experimental/reward_loop/reward_fn.py \
    custom_reward_function.name=compute_score_math_verify \
    trainer.val_only=True
``` 

The results are:

```bash
(TaskRunner pid=904198) ("Initial validation metrics: {'val-aux/openai/gsm8k/reward/mean@1': "
(TaskRunner pid=904198)  "np.float64(0.31766489764973466), 'val-core/openai/gsm8k/acc/mean@1': "
(TaskRunner pid=904198)  "np.float64(0.31766489764973466), 'val-aux/num_turns/min': np.int32(2), "
(TaskRunner pid=904198)  "'val-aux/num_turns/max': np.int32(2), 'val-aux/num_turns/mean': "
(TaskRunner pid=904198)  'np.float64(2.0)}')
```

# Design & Code Changes

- Teacher workers are used in the agent loop, similar to the generative
reward model: after a student worker finishes its rollout, the teacher
worker obtains logprobs
- In the initial version of this PR (#4897), requests were submitted to
the `vLLMHttpServer` via the `v1/completions` endpoint, which does not
support multi-modal data. While `v1/chat/completions` does support
multi-modal inputs, text must be passed as raw text instead of token
IDs, preventing exact scoring of student generations since `student gen
IDs -> student gen text -> teacher input IDs via v1/chat/completions
tokenization` will not always give `student gen IDs == teacher input
IDs` (https://vllm.ai/blog/agent-lightning).
- This PR instead follows a path similar to how rollout replicas
directly call the `generate` method on the `vLLMHttpServer`. This
enables multi-modal inputs while representing text as token IDs.
Requests to the teacher server now call the newly-added
`compute_logprobs` method of `vLLMHttpServer`.

---------

Co-authored-by: wuxibin <wuxibin@bytedance.com>

Author: JacobHelwig

.github/workflows/gpu_unit_tests.yml

@@ -117,6 +117,9 @@ jobs:
       - name: Testing LinearCrossEntropyTP Correctness, Computation Time and Memory Consumption
         run: |
           LOW_MEMORY=True torchrun --standalone --nnodes=1 --nproc-per-node=8 tests/utils/test_special_linear_cross_entropy_tp.py
+      - name: Testing Megatron KL Loss TP Correctness
+        run: |
+          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/utils/test_special_megatron_kl_loss_tp.py
       - name: Testing FSDP2 actor functionality
         run: |
           torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/actor/test_special_dp_actor.py

examples/on_policy_distillation_trainer/run_qwen3_vl_geo3k.sh

@@ -0,0 +1,159 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+############################ Quick Config ############################
+
+ROLLOUT_NAME="vllm" # sglang or vllm
+
+FAMILY="Qwen"
+STUDENT_MODEL=Qwen3-VL-2B-Instruct
+TEACHER_MODEL=Qwen3-VL-4B-Instruct
+
+# USE_POLICY_GRADIENT=False
+# DISTILLATION_LOSS_MODE="k3"
+# DISTILLATION_LOSS_MODE="forward_kl_topk"
+# USE_FUSED_KERNELS=False
+
+USE_POLICY_GRADIENT=True
+DISTILLATION_LOSS_MODE="k1"
+USE_FUSED_KERNELS=True
+
+DISTILLATION_LOSS_MAX_CLAMP=10.0
+DISTILLATION_LOG_PROB_MIN_CLAMP=-10.0
+
+PROJECT_NAME='verl_on_policy_distillation_example_geo3k'
+
+MAX_PROMPT=1024
+MAX_RESPONSE_LENGTH=2048
+MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 ))
+TRAIN_PROMPT_BSZ=128
+STUDENT_MICRO_BATCH_SIZE_PER_GPU=1
+STUDENT_MAX_TOKEN_LEN_PER_GPU=$(( STUDENT_MICRO_BATCH_SIZE_PER_GPU * (MAX_PROMPT + MAX_RESPONSE_LENGTH) ))
+USE_DYNAMIC_BSZ=False
+
+STUDENT_WORLD_SIZE=4
+
+TEACHER_RESOURCE_POOL=True
+TEACHER_WORLD_SIZE=4
+
+SP=1
+
+EXP_NAME="fsdp/student-${STUDENT_MODEL}/teacher-${TEACHER_MODEL}/loss-${DISTILLATION_LOSS_MODE}/pg-${USE_POLICY_GRADIENT}"
+
+ENFORCE_EAGER=False # true for faster debugging
+
+############################ Paths ############################
+
+geo3k_train_path=$DATA_PATH/geo3k/train.parquet
+geo3k_test_path=$DATA_PATH/geo3k/test.parquet
+
+TRAIN_FILES="['$geo3k_train_path']"
+TEST_FILES="['$geo3k_test_path']"
+
+############################ Parameter Groups ############################
+
+DATA=(
+    data.train_files="$TRAIN_FILES"
+    data.val_files="$TEST_FILES"
+    data.max_prompt_length=$MAX_PROMPT
+    data.max_response_length=$MAX_RESPONSE_LENGTH
+    data.train_batch_size=$TRAIN_PROMPT_BSZ
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=False
+    data.image_key=images
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="${FAMILY}/${STUDENT_MODEL}"
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.rollout.enforce_eager=$ENFORCE_EAGER
+)
+
+DISTILLATION=(
+    distillation.enabled=True
+    distillation.num_workers=8
+    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.8
+    distillation.teacher_model.inference.enforce_eager=$ENFORCE_EAGER
+    distillation.teacher_model.inference.max_model_len=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_batched_tokens=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_seqs=$MAX_NUM_TOKENS
+    +distillation.teacher_model.inference.engine_kwargs.vllm.disable_mm_preprocessor_cache=True
+    distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE
+    distillation.distillation_loss.topk=64
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT
+    distillation.distillation_loss.loss_max_clamp=$DISTILLATION_LOSS_MAX_CLAMP
+    distillation.distillation_loss.log_prob_min_clamp=$DISTILLATION_LOG_PROB_MIN_CLAMP
+)
+
+STUDENT=(
+    actor_rollout_ref.actor.optim.lr=1e-6
+    actor_rollout_ref.actor.ppo_mini_batch_size=$TRAIN_PROMPT_BSZ
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.actor.use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1
+    actor_rollout_ref.rollout.name=$ROLLOUT_NAME
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.8
+    actor_rollout_ref.rollout.calculate_log_probs=False
+    actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_seqs=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.n=1
+    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True
+)
+
+ALGORITHM=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+)
+
+TRAINER=(
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=$PROJECT_NAME
+    trainer.experiment_name=$EXP_NAME
+    trainer.n_gpus_per_node=$STUDENT_WORLD_SIZE
+    trainer.nnodes=1
+    trainer.save_freq=200
+    trainer.test_freq=5
+    trainer.total_epochs=15
+    trainer.val_before_train=True
+    trainer.use_legacy_worker_impl=disable
+    trainer.resume_mode=disable
+    trainer.log_val_generations=5
+)
+
+
+
+############################ Launch ############################
+
+python3 -m verl.trainer.main_ppo \
+    --config-path=config \
+    --config-name='ppo_trainer.yaml' \
+    "${DATA[@]}" \
+    "${ALGORITHM[@]}" \
+    "${MODEL[@]}" \
+    "${DISTILLATION[@]}" \
+    "${ROLLOUT[@]}" \
+    "${STUDENT[@]}" \
+    "${TRAINER[@]}" \
+    "$@"

examples/on_policy_distillation_trainer/run_qwen_gsm8k.sh

@@ -0,0 +1,156 @@
+#!/usr/bin/env bash
+set -xeuo pipefail
+
+############################ Quick Config ############################
+
+ROLLOUT_NAME="vllm" # sglang or vllm
+
+FAMILY="Qwen"
+STUDENT_MODEL=Qwen2.5-0.5B
+TEACHER_MODEL=Qwen2.5-0.5B
+
+# USE_POLICY_GRADIENT=False
+# DISTILLATION_LOSS_MODE="k3"
+# DISTILLATION_LOSS_MODE="forward_kl_topk"
+# USE_FUSED_KERNELS=False
+
+USE_POLICY_GRADIENT=True
+DISTILLATION_LOSS_MODE="k1"
+USE_FUSED_KERNELS=False
+
+DISTILLATION_LOSS_MAX_CLAMP=10.0
+DISTILLATION_LOG_PROB_MIN_CLAMP=-10.0
+
+PROJECT_NAME='verl_on_policy_distillation_example_gsm8k'
+
+MAX_PROMPT=256
+MAX_RESPONSE_LENGTH=512
+MAX_NUM_TOKENS=$(( MAX_PROMPT + MAX_RESPONSE_LENGTH + 1 ))
+TRAIN_PROMPT_BSZ=128
+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=True
+
+STUDENT_WORLD_SIZE=2
+
+TEACHER_RESOURCE_POOL=False
+TEACHER_WORLD_SIZE=4
+
+SP=1
+
+EXP_NAME="fsdp/student-${STUDENT_MODEL}/teacher-${TEACHER_MODEL}/loss-${DISTILLATION_LOSS_MODE}/pg-${USE_POLICY_GRADIENT}"
+
+ENFORCE_EAGER=True # true for faster debugging
+
+############################ Paths ############################
+
+gsm8k_train_path=$DATA_PATH/gsm8k/train.parquet
+gsm8k_test_path=$DATA_PATH/gsm8k/test.parquet
+
+TRAIN_FILES="['$gsm8k_train_path']"
+TEST_FILES="['$gsm8k_test_path']"
+
+############################ Parameter Groups ############################
+
+DATA=(
+    data.train_files="$TRAIN_FILES"
+    data.val_files="$TEST_FILES"
+    data.max_prompt_length=$MAX_PROMPT
+    data.max_response_length=$MAX_RESPONSE_LENGTH
+    data.train_batch_size=$TRAIN_PROMPT_BSZ
+    data.filter_overlong_prompts=True
+    data.truncation='error'
+    data.shuffle=False
+)
+
+MODEL=(
+    actor_rollout_ref.model.path="${FAMILY}/${STUDENT_MODEL}"
+    actor_rollout_ref.model.enable_gradient_checkpointing=True
+    actor_rollout_ref.model.use_remove_padding=True
+    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS
+    actor_rollout_ref.actor.use_torch_compile=True
+    actor_rollout_ref.rollout.enforce_eager=$ENFORCE_EAGER
+)
+
+DISTILLATION=(
+    distillation.enabled=True
+    distillation.num_workers=8
+    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.enforce_eager=$ENFORCE_EAGER
+    distillation.teacher_model.inference.max_model_len=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_batched_tokens=$MAX_NUM_TOKENS
+    distillation.teacher_model.inference.max_num_seqs=$MAX_NUM_TOKENS
+    distillation.distillation_loss.loss_mode=$DISTILLATION_LOSS_MODE
+    distillation.distillation_loss.topk=64
+    distillation.distillation_loss.use_task_rewards=False
+    distillation.distillation_loss.use_policy_gradient=$USE_POLICY_GRADIENT
+    distillation.distillation_loss.loss_max_clamp=$DISTILLATION_LOSS_MAX_CLAMP
+    distillation.distillation_loss.log_prob_min_clamp=$DISTILLATION_LOG_PROB_MIN_CLAMP
+)
+
+STUDENT=(
+    actor_rollout_ref.actor.optim.lr=1e-6
+    actor_rollout_ref.actor.ppo_mini_batch_size=$TRAIN_PROMPT_BSZ
+    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.actor.use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.actor.fsdp_config.param_offload=True
+    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True
+    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$SP
+)
+
+ROLLOUT=(
+    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$STUDENT_MICRO_BATCH_SIZE_PER_GPU
+    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=$STUDENT_MAX_TOKEN_LEN_PER_GPU
+    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=$USE_DYNAMIC_BSZ
+    actor_rollout_ref.rollout.tensor_model_parallel_size=1
+    actor_rollout_ref.rollout.name=$ROLLOUT_NAME
+    actor_rollout_ref.rollout.gpu_memory_utilization=0.3
+    actor_rollout_ref.rollout.calculate_log_probs=False
+    actor_rollout_ref.rollout.max_model_len=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_batched_tokens=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.max_num_seqs=$MAX_NUM_TOKENS
+    actor_rollout_ref.rollout.n=1
+)
+
+ALGORITHM=(
+    algorithm.adv_estimator=grpo
+    algorithm.use_kl_in_reward=False
+)
+
+TRAINER=(
+    trainer.logger='["console","wandb"]'
+    trainer.project_name=$PROJECT_NAME
+    trainer.experiment_name=$EXP_NAME
+    trainer.n_gpus_per_node=$STUDENT_WORLD_SIZE
+    trainer.nnodes=1
+    trainer.save_freq=200
+    trainer.test_freq=5
+    trainer.total_epochs=15
+    trainer.val_before_train=False
+    trainer.use_legacy_worker_impl=disable
+    trainer.resume_mode=disable
+    trainer.log_val_generations=5
+)
+
+
+
+############################ Launch ############################
+
+python3 -m verl.trainer.main_ppo \
+    --config-path=config \
+    --config-name='ppo_trainer.yaml' \
+    "${DATA[@]}" \
+    "${ALGORITHM[@]}" \
+    "${MODEL[@]}" \
+    "${DISTILLATION[@]}" \
+    "${ROLLOUT[@]}" \
+    "${STUDENT[@]}" \
+    "${TRAINER[@]}" \
+    "$@"