diff --git a/.github/workflows/pr-test.yml b/.github/workflows/pr-test.yml index 8ec019285b..f9924d0d7d 100644 --- a/.github/workflows/pr-test.yml +++ b/.github/workflows/pr-test.yml @@ -123,7 +123,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] + info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_opsd.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] defaults: run: working-directory: ${{ github.workspace }} @@ -506,7 +506,7 @@ jobs: strategy: fail-fast: false matrix: - info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] + info: [{"num_gpus": 4, "test_file": "test_full_disk_weight_update.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_quick_start_glm4_9B.py"}, {"num_gpus": 8, "test_file": "test_glm4.7_30B_A3B_pd_mooncake.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"num_gpus": 8, "test_file": "test_qwen3.6_35B_A3B_pd_mooncake.py", "use_deepep": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_30B_A3B_r3.py", "use_deepep": "1", "use_fp8_rollout": "1"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_disaggregate.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_qwen3_4B_ppo_train_critic_only.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B.py"}, {"enable_eval": "0", "num_gpus": 8, "test_file": "test_moonlight_16B_A3B_r3.py"}, {"num_gpus": 8, "test_file": "test_mimo_7B_mtp_only_grad.py"}, {"num_gpus": 8, "test_file": "test_qwen3_0.6B_parallel_check.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_debug_rollout_then_train.py"}, {"num_gpus": 8, "test_file": "test_qwen2.5_0.5B_opd_sglang.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_opsd.py"}, {"num_gpus": 6, "test_file": "test_qwen3_4B_external_pd.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fully_async_short.py"}, {"num_gpus": 4, "test_file": "test_qwen2.5_0.5B_fanout_short.py"}, {"num_gpus": 8, "test_file": "test_qwen3_4B_streaming_partial_rollout.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_short.py"}, {"num_gpus": 4, "test_file": "test_qwen3.5_0.8B_gsm8k_async_short.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer gpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer cpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--save-optimizer cpu --load-optimizer gpu", "test_file": "test_qwen3_4B_ckpt.py"}, {"num_gpus": 8, "test_args": "--async-save", "test_file": "test_qwen3_4B_ckpt.py"}] defaults: run: working-directory: ${{ github.workspace }} diff --git a/.github/workflows/pr-test.yml.j2 b/.github/workflows/pr-test.yml.j2 index 1e2a656c9f..5e61c6b137 100644 --- a/.github/workflows/pr-test.yml.j2 +++ b/.github/workflows/pr-test.yml.j2 @@ -14,6 +14,7 @@ {'test_file': 'test_qwen3_0.6B_parallel_check.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_debug_rollout_then_train.py', 'num_gpus': 8}, {'test_file': 'test_qwen2.5_0.5B_opd_sglang.py', 'num_gpus': 8}, + {'test_file': 'test_qwen2.5_0.5B_opsd.py', 'num_gpus': 4}, {'test_file': 'test_qwen3_4B_external_pd.py', 'num_gpus': 6}, {'test_file': 'test_qwen2.5_0.5B_fully_async_short.py', 'num_gpus': 4}, {'test_file': 'test_qwen2.5_0.5B_fanout_short.py', 'num_gpus': 4}, diff --git a/docs/en/advanced/on-policy-distillation.md b/docs/en/advanced/on-policy-distillation.md index 1930ff8e47..61a5e14deb 100644 --- a/docs/en/advanced/on-policy-distillation.md +++ b/docs/en/advanced/on-policy-distillation.md @@ -69,6 +69,49 @@ The teacher model is loaded directly into Megatron via `--opd-teacher-load`. Tea > **Note**: The teacher checkpoint must be in Megatron format (`torch_dist` or `torch`). You can convert from HuggingFace format using `tools/convert_hf_to_torch_dist.py`. +### Self Mode (`--opd-type self`) — On-Policy Self-Distillation (OPSD) + +This mode implements **On-Policy Self-Distillation** ([Self-Distilled Reasoner](https://github.com/siyan-zhao/OPSD)). A single model acts as both student and teacher, differing only by *context*: + +- **Student**: conditioned on the problem only; generates the on-policy rollout (the trainable current policy). +- **Teacher**: the **same model, frozen at the initial-policy checkpoint** (`--opd-teacher-load`), conditioned on the problem **plus privileged information** (the ground-truth solution). The teacher does not generate — it scores the student's response tokens in a single forward pass. + +Unlike the other modes (which fold a sampled-token reverse-KL into the advantage), OPSD uses a **direct, full-vocab token-level Jensen-Shannon divergence (JSD)** as the loss, with **no task reward** (pure distillation): + +$$ +\mathcal{L}_{\text{OPSD}} = \mathbb{E}_{t}\Big[\min\big(\text{JSD}_\beta\big(P_{\text{teacher}}(\cdot|x, s, y_{ **Limitations (current MVP)**: OPSD requires `--context-parallel-size 1`. The teacher response logits are held over the full vocabulary between the teacher and student forwards, so memory scales with `response_length × vocab_size`; pass `--opsd-offload-teacher-logits` to offload them to CPU (trading GPU memory for host↔device transfer), and prefer smaller models / shorter responses for now. A future chunked JSD would avoid materializing them. The teacher's micro-batches are re-packed by the (longer) privileged sequence lengths so the teacher forward respects `--max-tokens-per-gpu`. The example script is `examples/on_policy_distillation/run-qwen3-8B-opsd.sh`. + ## Running the Examples Complete example scripts are provided in `examples/on_policy_distillation/`: diff --git a/examples/on_policy_distillation/run-qwen3-8B-opsd.sh b/examples/on_policy_distillation/run-qwen3-8B-opsd.sh new file mode 100644 index 0000000000..3ea334abcd --- /dev/null +++ b/examples/on_policy_distillation/run-qwen3-8B-opsd.sh @@ -0,0 +1,172 @@ +#!/bin/bash + +# On-Policy Self-Distillation (OPSD), following "Self-Distilled Reasoner". +# +# A single model acts as both student and teacher: +# - Student: conditioned on the problem only; generates the on-policy rollout. +# - Teacher: the SAME model, frozen at the init checkpoint (--opd-teacher-load), +# conditioned on the problem + privileged information (the ground-truth solution). +# The teacher scores the student's response tokens in one forward pass, and the loss +# is a token-level full-vocab JSD between teacher and student over the student's tokens. +# +# Requirements: +# - --opd-type self (Megatron-only; context-parallel-size must be 1) +# - --opsd-privileged-info-key names the dataset field holding the ground-truth solution +# - --opd-teacher-load points at the initial-policy checkpoint (frozen teacher) +# - No task reward is used (pure distillation). + +set -ex + +export PYTHONUNBUFFERED=1 + +NVLINK_COUNT=$(nvidia-smi topo -m 2>/dev/null | grep -o 'NV[0-9][0-9]*' | wc -l) +if [ "$NVLINK_COUNT" -gt 0 ]; then + HAS_NVLINK=1 +else + HAS_NVLINK=0 +fi +echo "HAS_NVLINK: $HAS_NVLINK (detected $NVLINK_COUNT NVLink references)" + +source "/root/slime/scripts/models/qwen3-8B.sh" + + +CKPT_ARGS=( + --hf-checkpoint /root/Qwen3-8B + --ref-load /root/Qwen3-8B_torch_dist + --load /root/Qwen3-8B_slime/ + --save /root/Qwen3-8B_slime/ + --save-interval 20 +) + +# Data: the OPSD release uses siyanzhao/Openthoughts_math_30k_opsd, with fields +# "problem" -> the question (student + teacher context) +# "solution" -> the ground-truth solution (privileged info, teacher-only) +# Download and convert to jsonl, e.g.: +# hf download --repo-type dataset siyanzhao/Openthoughts_math_30k_opsd --local-dir /root/opsd-math-30k +ROLLOUT_ARGS=( + --prompt-data /root/opsd-math-30k/train.jsonl + --input-key problem + --apply-chat-template + --rollout-shuffle + --num-rollout 300 + --rollout-batch-size 16 + --n-samples-per-prompt 4 + --rollout-max-response-len 16384 + --rollout-temperature 1 + + --global-batch-size 64 + --balance-data +) + +# OPSD is pure distillation; the reward is computed but unused (no advantages). +RM_ARGS=( + --rm-type math +) + +EVAL_ARGS=( + # --eval-interval 20 + # --eval-prompt-data aime ${DATA_DIR}/aime-2024/aime-2024.jsonl + # --n-samples-per-eval-prompt 16 + # --eval-max-response-len 16384 + # --eval-top-p 1 +) + +PERF_ARGS=( + --tensor-model-parallel-size 2 + --sequence-parallel + --pipeline-model-parallel-size 1 + --context-parallel-size 1 # OPSD requires context-parallel-size 1 + --expert-model-parallel-size 1 + --expert-tensor-parallel-size 1 + + --recompute-granularity full + --recompute-method uniform + --recompute-num-layers 1 + + # --micro-batch-size 1 + --use-dynamic-batch-size + --max-tokens-per-gpu 16384 +) + +OPSD_ARGS=( + # On-Policy Self-Distillation + --use-opd # Enable on-policy distillation + --opd-type self # OPSD: single model is both student and teacher + --opd-teacher-load /root/Qwen3-8B_torch_dist # Frozen teacher = initial-policy checkpoint + --opsd-beta 0.5 # JSD interpolation (0=fwd KL, 1=rev KL, 0.5=symmetric) + --opsd-jsd-clip 0.05 # Per-token JSD clamp (stops style tokens dominating) + --opsd-privileged-info-key solution # Dataset field with the ground-truth solution (teacher-only) +) + +OPTIMIZER_ARGS=( + --optimizer adam + --lr 1e-6 + --lr-decay-style constant + --weight-decay 0.1 + --adam-beta1 0.9 + --adam-beta2 0.98 +) + +WANDB_ARGS=( + #--use-wandb + # --wandb-project slime-dev + # --wandb-group qwen3-8B-opsd + # --wandb-key ${WANDB_KEY} +) + +SGLANG_ARGS=( + --rollout-num-gpus-per-engine 1 + --sglang-mem-fraction-static 0.4 +) + + +MISC_ARGS=( + --attention-dropout 0.0 + --hidden-dropout 0.0 + --accumulate-allreduce-grads-in-fp32 + --attention-softmax-in-fp32 + --attention-backend flash +) + + + + +# launch the master node of ray in container +export MASTER_ADDR=${MASTER_ADDR:-"127.0.0.1"} +ray start --head --node-ip-address ${MASTER_ADDR} --num-gpus 8 --disable-usage-stats --dashboard-host=0.0.0.0 --dashboard-port=8265 + + +ray job submit --address="http://127.0.0.1:8265" \ + --runtime-env-json='{ + "env_vars": { + "PYTHONPATH": "/root/Megatron-LM/", + "CUDA_DEVICE_MAX_CONNECTIONS": "1" + } + }' \ + -- python3 train.py \ + --actor-num-nodes 1 \ + --actor-num-gpus-per-node 2 \ + --rollout-num-gpus 4 \ + ${MODEL_ARGS[@]} \ + ${CKPT_ARGS[@]} \ + ${ROLLOUT_ARGS[@]} \ + ${OPTIMIZER_ARGS[@]} \ + ${OPSD_ARGS[@]} \ + ${WANDB_ARGS[@]} \ + ${PERF_ARGS[@]} \ + ${EVAL_ARGS[@]} \ + ${SGLANG_ARGS[@]} \ + ${MISC_ARGS[@]} \ + ${RM_ARGS[@]} + + + +####clear after training +pkill -9 sglang +sleep 3 +ray stop --force +pkill -9 ray +pkill -9 python +sleep 3 +pkill -9 ray +pkill -9 python diff --git a/slime/backends/megatron_utils/actor.py b/slime/backends/megatron_utils/actor.py index 9830cc19cf..3467ff26b4 100644 --- a/slime/backends/megatron_utils/actor.py +++ b/slime/backends/megatron_utils/actor.py @@ -17,6 +17,7 @@ from slime.utils import train_dump_utils from slime.utils.data import process_rollout_data from slime.utils.distributed_utils import get_gloo_group +from slime.utils.dp_schedule import repack_micro_batches_by_length from slime.utils.logging_utils import init_tracking from slime.utils.memory_utils import clear_memory, print_memory from slime.utils.misc import Box @@ -32,7 +33,7 @@ from .data import DataIterator, get_data_iterator, log_perf_data, log_rollout_data from .hf_checkpoint_saver import save_hf_model_to_path from .initialize import init, is_megatron_main_rank -from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_values +from .loss import compute_advantages_and_returns, get_log_probs_and_entropy, get_response_logits, get_values from .model import forward_only, initialize_model_and_optimizer, save, train from .update_weight.common import named_params_and_buffers from .update_weight.update_weight_from_disk import UpdateWeightFromDisk @@ -235,6 +236,13 @@ def _get_rollout_data(self, rollout_data_ref: Box) -> RolloutBatch: rollout_data["loss_masks"] = [ torch.tensor(t, dtype=torch.int, device=torch.cuda.current_device()) for t in rollout_data["loss_masks"] ] + # OPSD: move the teacher's privileged token sequence to GPU and record its lengths. + if "teacher_tokens" in rollout_data: + rollout_data["teacher_tokens"] = [ + torch.tensor(t, dtype=torch.long, device=torch.cuda.current_device()) + for t in rollout_data["teacher_tokens"] + ] + rollout_data["teacher_total_lengths"] = [t.numel() for t in rollout_data["teacher_tokens"]] if "rollout_mask_sums" in rollout_data: # Promote precomputed per-rollout mask totals to GPU tensors here # (matching loss_masks) so the loss reducer can just divide. @@ -400,6 +408,67 @@ def compute_log_prob( store_prefix=store_prefix, ) + def compute_teacher_response_logits( + self, + rollout_data: RolloutBatch, + num_microbatches: list[int], + ) -> dict[str, list[torch.Tensor]]: + """Run the OPSD privileged teacher forward and return response-position logits. + + The currently-active model must be the teacher. The teacher is forwarded on + the privileged token sequence ``[prompt + privileged_info + response]`` (stored + in ``rollout_data["teacher_tokens"]``); per-sample response logits ``[R, V_local]`` + are returned under the key ``teacher_response_logits``, aligned 1:1 with the + student's response positions (response tokens are identical). + """ + assert mpu.get_context_parallel_world_size() == 1, "OPSD currently supports only context-parallel-size 1." + # Build a rollout_data view whose token stream is the privileged teacher sequence. + # The sample-to-rank assignment is unchanged (so teacher and student response + # positions stay aligned); only the micro-batch *groupings* are re-packed by the + # teacher (longer) lengths so the teacher forward respects max_tokens_per_gpu. + teacher_view = dict(rollout_data) + teacher_view["tokens"] = rollout_data["teacher_tokens"] + teacher_view["total_lengths"] = rollout_data["teacher_total_lengths"] + teacher_num_microbatches = num_microbatches + if self.args.use_dynamic_batch_size: + # Re-packing changes the per-step micro-batch count, which is not aligned to the + # VPP micro-batch group; OPSD with dynamic batching therefore requires no VPP. + assert (mpu.get_virtual_pipeline_model_parallel_world_size() or 1) == 1, ( + "OPSD with --use-dynamic-batch-size does not support virtual pipeline parallelism." + ) + max_tokens_per_bin = self.args.max_tokens_per_gpu * mpu.get_context_parallel_world_size() + teacher_micro_batch_indices, teacher_num_microbatches = repack_micro_batches_by_length( + rollout_data["micro_batch_indices"], + num_microbatches, + teacher_view["total_lengths"], + max_tokens_per_bin, + ) + teacher_view["micro_batch_indices"] = teacher_micro_batch_indices + if self.args.qkv_format == "bshd": + pad_size = mpu.get_tensor_model_parallel_world_size() * self.args.data_pad_size_multiplier + max_seq_len = max(teacher_view["total_lengths"]) + max_seq_len = (max_seq_len + pad_size - 1) // pad_size * pad_size + teacher_view["max_seq_lens"] = [max_seq_len] * len(teacher_view["tokens"]) + + data_iterator = get_data_iterator(teacher_view) + with timer("opsd_teacher_response_logits"): + res = forward_only( + get_response_logits, + self.args, + self.model, + data_iterator, + teacher_num_microbatches, + store_prefix="teacher_", + ) + + # Optionally offload the full-vocab teacher logits to CPU to cut peak GPU memory + # (moved back to the device per micro-batch in opsd_loss_function). + if self.args.opsd_offload_teacher_logits and "teacher_response_logits" in res: + res["teacher_response_logits"] = [ + t.cpu() if t is not None else None for t in res["teacher_response_logits"] + ] + return res + def train(self, rollout_id: int, rollout_data_ref: Box, external_data=None): if self.args.debug_rollout_only: return None @@ -461,75 +530,92 @@ def train_actor(self, rollout_id: int, rollout_data: RolloutBatch, external_data with inverse_timer("train_wait"), timer("train"): if self.args.compute_advantages_and_returns: - if "ref" in self.weights_backuper.backup_tags: - if self.args.use_routing_replay: - os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" - self._switch_model("ref") - rollout_data.update( - self.compute_log_prob( - data_iterator, - num_microbatches, - store_prefix="ref_", - ) - ) - - # Forward teacher model to get teacher_log_probs for Megatron-based OPD - if "teacher" in self.weights_backuper.backup_tags: - if self.args.use_routing_replay: - os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" + opsd = self.args.use_opd and self.args.opd_type == "self" + if opsd: + # OPSD: only the privileged teacher forward is needed. The JSD loss is + # computed during the student training forward (loss_type="opsd"), so no + # ref / old-logprob / advantage passes are required (pure distillation). + assert ( + "teacher" in self.weights_backuper.backup_tags + ), "OPSD requires the teacher model to be loaded (--opd-teacher-load)." self._switch_model("teacher") - rollout_data.update( - self.compute_log_prob( - data_iterator, - num_microbatches, - store_prefix="teacher_", + try: + rollout_data.update(self.compute_teacher_response_logits(rollout_data, num_microbatches)) + finally: + # Always restore the actor as the live model, even if the teacher + # forward raises (e.g. OOM), so the post-train backup("actor") never + # snapshots teacher weights. + self._switch_model("actor") + else: + if "ref" in self.weights_backuper.backup_tags: + if self.args.use_routing_replay: + os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" + self._switch_model("ref") + rollout_data.update( + self.compute_log_prob( + data_iterator, + num_microbatches, + store_prefix="ref_", + ) ) - ) - self._switch_model("old_actor" if self.args.keep_old_actor else "actor") - can_reuse_log_probs_in_loss = ( - len(num_microbatches) == 1 - and self.args.loss_type == "policy_loss" - and self.args.kl_coef == 0 - and not self.args.use_rollout_logprobs - and not self.args.get_mismatch_metrics - and not self.args.use_critic - and not self.args.keep_old_actor - and not self.args.use_opd - and not self.args.use_routing_replay - and self.args.advantage_estimator != "gspo" - ) - if ( - not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics - ) and not can_reuse_log_probs_in_loss: - if self.args.use_routing_replay: - if self.args.use_rollout_routing_replay: - os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" - else: - os.environ["ROUTING_REPLAY_STAGE"] = "record" - rollout_data.update( - self.compute_log_prob( - data_iterator, - num_microbatches, - store_prefix="", + # Forward teacher model to get teacher_log_probs for Megatron-based OPD + if "teacher" in self.weights_backuper.backup_tags: + if self.args.use_routing_replay: + os.environ["ROUTING_REPLAY_STAGE"] = "fallthrough" + self._switch_model("teacher") + rollout_data.update( + self.compute_log_prob( + data_iterator, + num_microbatches, + store_prefix="teacher_", + ) ) + + self._switch_model("old_actor" if self.args.keep_old_actor else "actor") + can_reuse_log_probs_in_loss = ( + len(num_microbatches) == 1 + and self.args.loss_type == "policy_loss" + and self.args.kl_coef == 0 + and not self.args.use_rollout_logprobs + and not self.args.get_mismatch_metrics + and not self.args.use_critic + and not self.args.keep_old_actor + and not self.args.use_opd + and not self.args.use_routing_replay + and self.args.advantage_estimator != "gspo" ) - if self.args.use_rollout_routing_replay: - RoutingReplay.clear_all_forward() - - if self.args.use_critic: - if external_data is not None and mpu.is_pipeline_last_stage(): - values = external_data.get("values") - if values is not None: - from slime.backends.megatron_utils.data import tensors_to_gpu - - rollout_data["values"] = tensors_to_gpu(values) - if self._active_model_tag != "actor": - self._switch_model("actor") - - # Calculate adv and returns. Need to performed before training (instead of on the fly), - # because we may need normalize the whole rollout. - compute_advantages_and_returns(self.args, rollout_data) + if ( + not self.args.use_rollout_logprobs or self.args.get_mismatch_metrics + ) and not can_reuse_log_probs_in_loss: + if self.args.use_routing_replay: + if self.args.use_rollout_routing_replay: + os.environ["ROUTING_REPLAY_STAGE"] = "replay_forward" + else: + os.environ["ROUTING_REPLAY_STAGE"] = "record" + rollout_data.update( + self.compute_log_prob( + data_iterator, + num_microbatches, + store_prefix="", + ) + ) + if self.args.use_rollout_routing_replay: + RoutingReplay.clear_all_forward() + + if self.args.use_critic: + if external_data is not None and mpu.is_pipeline_last_stage(): + values = external_data.get("values") + if values is not None: + from slime.backends.megatron_utils.data import tensors_to_gpu + + rollout_data["values"] = tensors_to_gpu(values) + if self._active_model_tag != "actor": + self._switch_model("actor") + + # Calculate adv and returns. Need to performed before training (instead of on the fly), + # because we may need normalize the whole rollout. + compute_advantages_and_returns(self.args, rollout_data) if self.rollout_data_postprocess is not None: self.rollout_data_postprocess(self.args, rollout_id, rollout_data) diff --git a/slime/backends/megatron_utils/data.py b/slime/backends/megatron_utils/data.py index 7c12a5a778..1aaddb9e7d 100644 --- a/slime/backends/megatron_utils/data.py +++ b/slime/backends/megatron_utils/data.py @@ -303,6 +303,11 @@ def log_rollout_data( "global_batch_sizes", "num_microbatches", "micro_batch_indices", + # OPSD: privileged teacher token sequence and its per-token response logits + # are not scalar/token-level metrics and must not be reduced here. + "teacher_tokens", + "teacher_total_lengths", + "teacher_response_logits", ]: continue # Emit (sum, count) so gather_log_data can do a weighted average across diff --git a/slime/backends/megatron_utils/loss.py b/slime/backends/megatron_utils/loss.py index a63fa159f0..4419aca70d 100644 --- a/slime/backends/megatron_utils/loss.py +++ b/slime/backends/megatron_utils/loss.py @@ -16,6 +16,7 @@ compute_gspo_kl, compute_opsm_mask, compute_policy_loss, + compute_vocab_parallel_jsd, get_advantages_and_returns_batch, get_grpo_returns, get_reinforce_plus_plus_baseline_advantages, @@ -530,6 +531,49 @@ def get_values( return torch.empty((0,), device=logits.device), res +def get_response_logits( + logits: torch.Tensor, + *, + args: Namespace, + unconcat_tokens: list[torch.Tensor], + total_lengths: list[int], + response_lengths: list[int], + with_entropy: bool = False, + non_loss_data: bool = True, + max_seq_lens: list[int] | None = None, +) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]: + """Extract per-sample response-aligned logits (full local vocab shard). + + Used by OPSD to capture the teacher's response-position logits during a + privileged forward pass. Each returned tensor has shape `[R, V_local]` where + `V_local` is this rank's vocab shard. The response tokens scored by the + teacher are a verbatim copy of the student's response, so position `r` of the + returned chunk aligns 1:1 with response position `r` of the student forward. + + Returns a dict with key "response_logits" so the caller (forward_only) stores + it under the chosen prefix (e.g. "teacher_response_logits"). + """ + logits_list = [] + for logits_chunk, _ in get_responses( + logits.float(), + args=args, + unconcat_tokens=unconcat_tokens, + total_lengths=total_lengths, + response_lengths=response_lengths, + max_seq_lens=max_seq_lens, + apply_temperature=False, + ): + # .detach().clone() the response slice: ``logits_chunk`` is a view into the full + # ``[1, T_padded, V_local]`` microbatch tensor, so keeping the view would pin the + # entire padded backing buffer in GPU memory for every microbatch until the student + # forward runs. Cloning retains only the ``[R, V_local]`` response data. (Teacher is + # frozen -> detach.) For further savings, --opsd-offload-teacher-logits moves these + # to CPU. + logits_list.append(logits_chunk.detach().clone()) + + return torch.empty((0,), device=logits.device), {"response_logits": logits_list} + + def apply_opd_kl_to_advantages( args: Namespace, rollout_data: RolloutBatch, @@ -1028,6 +1072,84 @@ def policy_loss_function( return loss, reported_loss +def opsd_loss_function( + args: Namespace, + batch: RolloutBatch, + logits: torch.Tensor, + sum_of_sample_mean: Callable[[torch.Tensor], torch.Tensor], +) -> tuple[torch.Tensor, dict[str, torch.Tensor]]: + """Compute the on-policy self-distillation (OPSD) loss. + + For each sample, computes the per-token full-vocab generalized JSD between the + student's next-token distribution (from `logits`) and the teacher's + distribution (precomputed during the privileged teacher forward and stored in + `batch["teacher_response_logits"]`). The per-token JSD is clamped to + `args.opsd_jsd_clip` and reduced with the standard per-sample mean reducer. + + Args: + args: Configuration containing `opsd_beta` and `opsd_jsd_clip`. + batch: Mini-batch with "teacher_response_logits" (per-sample `[R, V_local]`), + "unconcat_tokens", "total_lengths", "response_lengths", "loss_masks". + logits: Student policy logits with shape `[1, T, V]` (float32). + sum_of_sample_mean: Reduction function that averages per-sample values. + + Returns: + Tuple of `(loss, metrics)` where `metrics` contains detached "loss" and + "opsd_jsd" scalars. + """ + teacher_response_logits = batch.get("teacher_response_logits") + if teacher_response_logits is None: + raise ValueError("OPSD requires 'teacher_response_logits' in the batch, but it is missing.") + + tp_group = mpu.get_tensor_model_parallel_group() + + student_response_logits = [chunk for chunk, _ in get_responses( + logits.float(), + args=args, + unconcat_tokens=batch["unconcat_tokens"], + total_lengths=batch["total_lengths"], + response_lengths=batch["response_lengths"], + max_seq_lens=batch.get("max_seq_lens", None), + apply_temperature=False, + )] + assert len(student_response_logits) == len(teacher_response_logits), ( + f"OPSD sample count mismatch: {len(student_response_logits)} student vs " + f"{len(teacher_response_logits)} teacher" + ) + + jsd_list = [] + for student_logits_chunk, teacher_logits_chunk in zip( + student_response_logits, teacher_response_logits, strict=True + ): + assert student_logits_chunk.shape == teacher_logits_chunk.shape, ( + f"OPSD student/teacher logits shape mismatch: " + f"{student_logits_chunk.shape} vs {teacher_logits_chunk.shape}" + ) + jsd = compute_vocab_parallel_jsd( + student_logits_chunk.float(), + # teacher logits may be offloaded to CPU (--opsd-offload-teacher-logits). + teacher_logits_chunk.to(student_logits_chunk.device).float(), + beta=args.opsd_beta, + process_group=tp_group, + temperature=args.opsd_temperature, + ) + jsd = jsd.clamp(max=args.opsd_jsd_clip) + jsd_list.append(jsd) + + jsd_cat = torch.cat(jsd_list, dim=0) + loss = sum_of_sample_mean(jsd_cat) + + # make sure the gradient could backprop correctly when this rank has no tokens. + if jsd_cat.numel() == 0: + loss = loss + 0 * logits.sum() + + reported_loss = { + "loss": loss.clone().detach(), + "opsd_jsd": loss.clone().detach(), + } + return loss, reported_loss + + def value_loss_function( args: Namespace, batch: RolloutBatch, @@ -1190,6 +1312,8 @@ def loss_function( func = value_loss_function case "sft_loss": func = sft_loss_function + case "opsd": + func = opsd_loss_function case "custom_loss": func = load_function(args.custom_loss_function_path) case _: diff --git a/slime/backends/megatron_utils/model.py b/slime/backends/megatron_utils/model.py index db6020a94d..3e6710ce42 100644 --- a/slime/backends/megatron_utils/model.py +++ b/slime/backends/megatron_utils/model.py @@ -501,6 +501,7 @@ def forward_step(data_iterator: DataIterator, model: GPTModel, return_schedule_p "rollout_log_probs", "max_seq_lens", "teacher_log_probs", + "teacher_response_logits", "rollout_mask_sums", ], args.data_pad_size_multiplier, diff --git a/slime/ray/placement_group.py b/slime/ray/placement_group.py index 182481c433..e65e223b21 100644 --- a/slime/ray/placement_group.py +++ b/slime/ray/placement_group.py @@ -187,8 +187,10 @@ def create_training_models(args, pgs, rollout_manager): actor_model.async_init( actor_args, role="actor", - with_ref=actor_args.kl_coef != 0 or actor_args.use_kl_loss, - with_opd_teacher=actor_args.use_opd and actor_args.opd_type == "megatron", + # OPSD (opd_type=self) never forwards the ref model, so skip loading it. + with_ref=(actor_args.kl_coef != 0 or actor_args.use_kl_loss) + and not (actor_args.use_opd and actor_args.opd_type == "self"), + with_opd_teacher=actor_args.use_opd and actor_args.opd_type in ("megatron", "self"), ) ) # TODO how to decide rollout start id when critic is involved? For now we just require user to specify it via args. diff --git a/slime/ray/rollout.py b/slime/ray/rollout.py index 3cbdb7e2cb..1db0a78718 100644 --- a/slime/ray/rollout.py +++ b/slime/ray/rollout.py @@ -399,6 +399,12 @@ def __init__(self, args, pg): self.custom_convert_samples_to_train_data_func = load_function( self.args.custom_convert_samples_to_train_data_path ) + # Tokenizer for OPSD: used to tokenize the teacher's privileged information. + self.tokenizer = None + if self.args.use_opd and self.args.opd_type == "self": + from slime.utils.processing_utils import load_tokenizer + + self.tokenizer = load_tokenizer(self.args.hf_checkpoint, trust_remote_code=True) logger.info(f"import {self.args.rollout_function_path} as generate_rollout function.") logger.info(f"import {self.args.eval_function_path} as eval_generate_rollout function.") @@ -776,6 +782,30 @@ def _convert_samples_to_train_data(self, samples: list[Sample] | list[list[Sampl if samples[0].teacher_log_probs is not None: train_data["teacher_log_probs"] = [sample.teacher_log_probs for sample in samples] + # OPSD: build the teacher's privileged token sequence [prompt + privileged_info + response]. + # The response segment is a verbatim copy of the student's response, so its R teacher + # response positions align 1:1 with the student's response positions during training. + if self.args.use_opd and self.args.opd_type == "self": + teacher_tokens = [] + num_missing_privileged_info = 0 + for sample in samples: + prompt_length = len(sample.tokens) - sample.response_length + prompt_tokens = list(sample.tokens[:prompt_length]) + response_tokens = list(sample.tokens[prompt_length:]) + privileged_info = sample.privileged_info or "" + if not privileged_info: + num_missing_privileged_info += 1 + privileged_tokens = self.tokenizer.encode(privileged_info, add_special_tokens=False) + teacher_tokens.append(prompt_tokens + privileged_tokens + response_tokens) + train_data["teacher_tokens"] = teacher_tokens + if num_missing_privileged_info: + logger.warning( + f"OPSD: {num_missing_privileged_info}/{len(samples)} samples have empty/None " + f"privileged_info (key '{self.args.opsd_privileged_info_key}'). For those samples the " + f"teacher context equals the student's, yielding ~0 distillation signal. Check the dataset " + f"field name and that rows are populated." + ) + return train_data def set_train_parallel_config(self, config: dict): @@ -825,6 +855,7 @@ def _split_train_data_by_dp(self, data): "rollout_routed_experts", "prompt", "teacher_log_probs", + "teacher_tokens", ]: if key not in data: continue diff --git a/slime/rollout/data_source.py b/slime/rollout/data_source.py index ca7171ecae..2eb8a14db8 100644 --- a/slime/rollout/data_source.py +++ b/slime/rollout/data_source.py @@ -76,6 +76,7 @@ def __init__(self, args): prompt_key=args.input_key, multimodal_keys=args.multimodal_keys, label_key=args.label_key, + privileged_info_key=args.opsd_privileged_info_key, metadata_key=args.metadata_key, tool_key=args.tool_key, apply_chat_template=args.apply_chat_template, diff --git a/slime/utils/arguments.py b/slime/utils/arguments.py index 494632f255..e0ee765efa 100644 --- a/slime/utils/arguments.py +++ b/slime/utils/arguments.py @@ -1083,12 +1083,16 @@ def add_on_policy_distillation_arguments(parser): parser.add_argument( "--opd-type", type=str, - choices=["sglang", "megatron"], + choices=["sglang", "megatron", "self"], default=None, help=( "Type of on-policy distillation. " "'sglang': Teacher log-probs are obtained from external SGLang server during rollout. " - "'megatron': Teacher model is loaded via --opd-teacher-load and forwarded during training." + "'megatron': Teacher model is loaded via --opd-teacher-load and forwarded during training. " + "'self': On-policy self-distillation (OPSD). A single model is both student and teacher; " + "the teacher (frozen at the init checkpoint, loaded via --opd-teacher-load) is conditioned " + "on privileged information (the ground-truth solution) and the loss is a token-level " + "full-vocab JSD over the student's on-policy tokens." ), ) parser.add_argument( @@ -1102,13 +1106,58 @@ def add_on_policy_distillation_arguments(parser): type=str, default=None, help=( - "The checkpoint for OPD teacher model. Required when --opd-type=megatron. " - "The teacher model should have the same architecture as policy/ref model." + "The checkpoint for OPD teacher model. Required when --opd-type=megatron or --opd-type=self. " + "The teacher model should have the same architecture as policy/ref model. For OPSD " + "(--opd-type=self) this should point at the initial-policy checkpoint." ), ) parser.add_argument( "--opd-teacher-ckpt-step", type=int, default=None, help="The checkpoint step for OPD teacher model." ) + # --- On-policy self-distillation (OPSD, --opd-type=self) --- + parser.add_argument( + "--opsd-beta", + type=float, + default=0.5, + help=( + "Generalized JSD interpolation weight for OPSD. beta=0 -> forward KL (teacher||student), " + "beta=1 -> reverse KL (student||teacher), 0.5 -> symmetric JSD. Default 0.5." + ), + ) + parser.add_argument( + "--opsd-jsd-clip", + type=float, + default=0.05, + help=( + "Per-token JSD clamp (max) for OPSD. Caps each token's divergence so high-divergence " + "style tokens do not dominate the gradient. Default 0.05. Set to a large value to disable." + ), + ) + parser.add_argument( + "--opsd-temperature", + type=float, + default=1.0, + help="Softmax temperature applied to both student and teacher logits in the OPSD JSD. Default 1.0.", + ) + parser.add_argument( + "--opsd-offload-teacher-logits", + action="store_true", + default=False, + help=( + "Offload the teacher response logits to CPU between the teacher and student forwards " + "to reduce peak GPU memory (moved back per micro-batch during the loss). Trades GPU " + "memory for host<->device transfer. Future work: chunked JSD to avoid storing them." + ), + ) + parser.add_argument( + "--opsd-privileged-info-key", + type=str, + default=None, + help=( + "JSON dataset key holding the privileged information (e.g. the ground-truth solution) " + "shown only to the OPSD teacher. Required when --opd-type=self." + ), + ) return parser def add_router_arguments(parser): @@ -1766,12 +1815,14 @@ def slime_validate_args(args): # Validate on-policy distillation (OPD) arguments if args.use_opd: if args.opd_type is None: - raise ValueError("--opd-type must be specified when --use-opd is enabled. Choose 'sglang' or 'megatron'.") + raise ValueError( + "--opd-type must be specified when --use-opd is enabled. Choose 'sglang', 'megatron' or 'self'." + ) - if args.opd_type == "megatron": + if args.opd_type in ("megatron", "self"): if args.opd_teacher_load is None: raise ValueError( - "--opd-teacher-load is required when --opd-type=megatron. " + f"--opd-teacher-load is required when --opd-type={args.opd_type}. " "Please provide the path to the teacher model checkpoint." ) if not os.path.exists(args.opd_teacher_load): @@ -1784,6 +1835,36 @@ def slime_validate_args(args): "please make sure it is a valid megatron checkpoint directory." ) + if args.opd_type == "self": + # OPSD is a pure-distillation, full-vocab JSD loss computed during the training + # forward pass; it requires the JSD loss type and the privileged-info data key. + if args.opsd_privileged_info_key is None: + raise ValueError( + "--opsd-privileged-info-key is required when --opd-type=self. It names the dataset " + "field holding the privileged information (ground-truth solution) shown to the teacher." + ) + if args.loss_type not in ("policy_loss", "opsd"): + raise ValueError( + f"--opd-type=self (OPSD) requires --loss-type=opsd (got {args.loss_type!r})." + ) + if args.loss_type != "opsd": + logger.info("OPSD: setting --loss-type=opsd (was %r).", args.loss_type) + args.loss_type = "opsd" + if not (0.0 <= args.opsd_beta <= 1.0): + raise ValueError(f"--opsd-beta must be in [0, 1], got {args.opsd_beta}.") + if args.kl_coef != 0 or args.use_kl_loss: + raise ValueError( + "--opd-type=self (OPSD) is pure distillation and does not use a reference-model KL. " + "Set --kl-coef 0 and do not pass --use-kl-loss (the ref model is intentionally not " + "loaded under OPSD)." + ) + if not args.compute_advantages_and_returns: + raise ValueError( + "--opd-type=self (OPSD) requires compute_advantages_and_returns to be enabled " + "(do not pass --disable-compute-advantages-and-returns): the privileged teacher " + "forward runs in that phase." + ) + elif args.opd_type == "sglang": if args.opd_teacher_load is not None: raise ValueError( diff --git a/slime/utils/data.py b/slime/utils/data.py index 0d26b6dda5..a0b8867afc 100644 --- a/slime/utils/data.py +++ b/slime/utils/data.py @@ -210,6 +210,7 @@ def __init__( prompt_key="text", multimodal_keys=None, label_key=None, + privileged_info_key=None, tool_key=None, metadata_key="metadata", seed=42, @@ -254,10 +255,16 @@ def __init__( else: multimodal_inputs = None + if privileged_info_key is not None and privileged_info_key not in data: + raise KeyError( + f"Dataset row is missing the privileged-info field {privileged_info_key!r} " + f"required by --opsd-privileged-info-key. Available keys: {sorted(data.keys())}." + ) origin_samples.append( Sample( prompt=output_prompt, label=data[label_key] if label_key is not None else None, + privileged_info=data[privileged_info_key] if privileged_info_key is not None else None, metadata=metadata, multimodal_inputs=multimodal_inputs, ) diff --git a/slime/utils/dp_schedule.py b/slime/utils/dp_schedule.py index 1735fad3b4..9706b7656f 100644 --- a/slime/utils/dp_schedule.py +++ b/slime/utils/dp_schedule.py @@ -79,6 +79,51 @@ def _pack_step_into_mbs( return [list(range(i, min(i + micro_batch_size, n))) for i in range(0, n, micro_batch_size)] +def repack_micro_batches_by_length( + micro_batch_indices: list[list[int]], + num_microbatches: list[int], + lengths: list[int], + max_tokens_per_bin: int, +) -> tuple[list[list[int]], list[int]]: + """Re-pack one DP rank's micro-batches for a different per-sample length budget. + + Used by OPSD: the teacher is forwarded on the longer ``[prompt + privileged_info + + response]`` sequences, so its micro-batches must be packed by the teacher token + lengths even though the **sample-to-rank assignment is unchanged** (so teacher and + student response positions stay aligned). This re-groups *only within each training + step* (preserving which samples belong to which step), so every local sample is + still covered exactly once and the per-step structure matches the student schedule. + + Because the OPSD teacher pass is forward-only and per-rank independent (no + cross-DP collective), ranks may end up with different ``num_microbatches`` here. + + Args: + micro_batch_indices: this rank's mbs (flat across steps); each mb is a list of + local sample indices. (i.e. ``build_dp_schedule(...)[1][rank]``) + num_microbatches: number of mbs per step for this rank (same layout as the + value returned by ``build_dp_schedule``). + lengths: per-local-sample token length to pack by (teacher lengths). + max_tokens_per_bin: token cap per micro-batch (``max_tokens_per_gpu * cp_size``). + + Returns: + ``(new_micro_batch_indices, new_num_microbatches)`` with the same per-sample + coverage and step boundaries, re-binned by ``lengths``. + """ + new_micro_batch_indices: list[list[int]] = [] + new_num_microbatches: list[int] = [] + mb_cursor = 0 + for step_num_mb in num_microbatches: + step_mbs = micro_batch_indices[mb_cursor : mb_cursor + step_num_mb] + mb_cursor += step_num_mb + step_local_indices = [idx for mb in step_mbs for idx in mb] + step_lengths = [lengths[idx] for idx in step_local_indices] + bins = first_fit_pack(step_lengths, max_tokens_per_bin) + step_new_mbs = [[step_local_indices[p] for p in bin_] for bin_ in bins] + new_micro_batch_indices.extend(step_new_mbs) + new_num_microbatches.append(len(step_new_mbs)) + return new_micro_batch_indices, new_num_microbatches + + def build_dp_schedule( args: Any, train_parallel_config: dict, diff --git a/slime/utils/ppo_utils.py b/slime/utils/ppo_utils.py index 2a858e7a3f..744e8f9099 100644 --- a/slime/utils/ppo_utils.py +++ b/slime/utils/ppo_utils.py @@ -1,6 +1,7 @@ # Adapt from https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/models/utils.py # and https://github.com/OpenRLHF/OpenRLHF/blob/10c733694ed9fbb78a0a2ff6a05efc7401584d46/openrlhf/trainer/ppo_utils/experience_maker.py +import math from argparse import Namespace import torch @@ -198,6 +199,142 @@ def compute_entropy_from_logits(logits: torch.Tensor, process_group) -> torch.Te return _VocabParallelEntropy.apply(logits, process_group) +class _VocabParallelAllReduceSum(torch.autograd.Function): + """All-reduce (SUM) across the vocab-parallel (TP) group with **identity** backward. + + Correct only when the all-reduced value feeds a *replicated* downstream (same on + every rank), so the cotangent ``dL/dy`` computed locally already equals the global + one. Used for the final per-token JSD reduction (its downstream — clamp + loss + reduction — is replicated across TP ranks). This is Megatron's "g" (reduce) op. + + Do NOT use this for an all-reduced value consumed with a *rank-local* shard (e.g. + a softmax normalizer); that needs :class:`_VocabParallelAllReduceSumGradAllReduce`. + """ + + @staticmethod + def forward(ctx, x: torch.Tensor, process_group: dist.ProcessGroup | None) -> torch.Tensor: + if process_group is not None and dist.get_world_size(group=process_group) > 1: + x = x.clone() + dist.all_reduce(x, group=process_group) + return x + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + return grad_output, None + + +class _VocabParallelAllReduceSumGradAllReduce(torch.autograd.Function): + """All-reduce (SUM) across the vocab-parallel (TP) group with **all-reduce** backward. + + Correct when the all-reduced value ``y = sum_r x_r`` is replicated but then consumed + with a *rank-local* shard (so ``dL/dy`` differs per rank and only sums to the true + global cotangent across ranks). Used for the softmax normalizer ``sum_exp`` in + vocab-parallel log-softmax: every rank's log-probs depend on the same global + normalizer, so its gradient must be all-reduced before flowing back to the local + ``x_r``. Without this, the normalization gradient is under-counted by ~1/TP. + """ + + @staticmethod + def forward(ctx, x: torch.Tensor, process_group: dist.ProcessGroup | None) -> torch.Tensor: + ctx.process_group = process_group + if process_group is not None and dist.get_world_size(group=process_group) > 1: + x = x.clone() + dist.all_reduce(x, group=process_group) + return x + + @staticmethod + def backward(ctx, grad_output: torch.Tensor): + process_group = ctx.process_group + if process_group is not None and dist.get_world_size(group=process_group) > 1: + grad_output = grad_output.clone() + dist.all_reduce(grad_output, group=process_group) + return grad_output, None + + +def _vocab_parallel_log_softmax(logits: torch.Tensor, process_group: dist.ProcessGroup | None) -> torch.Tensor: + """Numerically-stable log-softmax over a vocab dimension sharded across `process_group`. + + `logits` has shape `[N, V_local]`; the returned tensor has the same shape and + holds the log-probabilities of the local vocab shard, normalized over the + *global* vocabulary. Differentiable w.r.t. `logits` (including the cross-rank + coupling through the global normalizer). + """ + # The max is a constant shift for softmax; detach so no gradient flows through it. + logits_max = logits.detach().max(dim=-1, keepdim=True).values + if process_group is not None and dist.get_world_size(group=process_group) > 1: + dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=process_group) + shifted = logits - logits_max + sum_exp = shifted.exp().sum(dim=-1, keepdim=True) + # The normalizer couples all ranks' log-probs -> its gradient must be all-reduced. + sum_exp = _VocabParallelAllReduceSumGradAllReduce.apply(sum_exp, process_group) + return shifted - sum_exp.log() + + +def compute_vocab_parallel_jsd( + student_logits: torch.Tensor, + teacher_logits: torch.Tensor, + beta: float, + process_group: dist.ProcessGroup | None, + temperature: float = 1.0, +) -> torch.Tensor: + """Per-token generalized Jensen-Shannon divergence over a vocab-parallel logits shard. + + Follows the generalized JSD used by OPSD ("Self-Distilled Reasoner") / TRL's + GOLD trainer: + + M = (1 - beta) * P_student + beta * P_teacher + JSD_beta = beta * KL(P_teacher || M) + (1 - beta) * KL(P_student || M) + + with the limiting cases ``beta == 0`` -> ``KL(P_teacher || P_student)`` (forward KL) + and ``beta == 1`` -> ``KL(P_student || P_teacher)`` (reverse KL). At ``beta == 0.5`` + this is the symmetric Jensen-Shannon divergence. + + Gradients flow only through ``student_logits``; the teacher is treated as a + constant target (detached). `student_logits` / `teacher_logits` have shape + `[N, V_local]` (the local vocab shard); the returned tensor has shape `[N]`. + + Args: + student_logits: Trainable student logits for the local vocab shard. + teacher_logits: Frozen teacher logits for the local vocab shard. + beta: Interpolation weight in [0, 1]. + process_group: Tensor/vocab-parallel group, or None for single-rank. + temperature: Softmax temperature applied to both logits before the JSD + (matches OPSD; default 1.0 is a no-op). + """ + if temperature != 1.0: + student_logits = student_logits / temperature + teacher_logits = teacher_logits / temperature + + student_log_probs = _vocab_parallel_log_softmax(student_logits, process_group) + teacher_log_probs = _vocab_parallel_log_softmax(teacher_logits.detach(), process_group).detach() + + if beta == 0.0: + # KL(teacher || student) = sum_v P_teacher * (log P_teacher - log P_student) + per_shard = teacher_log_probs.exp() * (teacher_log_probs - student_log_probs) + elif beta == 1.0: + # KL(student || teacher) = sum_v P_student * (log P_student - log P_teacher) + per_shard = student_log_probs.exp() * (student_log_probs - teacher_log_probs) + else: + # Mixture log-probs computed stably in log space: log((1-beta)*q + beta*p). + mixture_log_probs = torch.logsumexp( + torch.stack( + [ + student_log_probs + math.log(1.0 - beta), + teacher_log_probs + math.log(beta), + ], + dim=0, + ), + dim=0, + ) + kl_teacher = teacher_log_probs.exp() * (teacher_log_probs - mixture_log_probs) + kl_student = student_log_probs.exp() * (student_log_probs - mixture_log_probs) + per_shard = beta * kl_teacher + (1.0 - beta) * kl_student + + jsd = per_shard.sum(dim=-1) + jsd = _VocabParallelAllReduceSum.apply(jsd, process_group) + return jsd + + def get_grpo_returns( rewards: torch.Tensor, kl: list[torch.Tensor], diff --git a/slime/utils/types.py b/slime/utils/types.py index 54092680e2..c06a8af658 100644 --- a/slime/utils/types.py +++ b/slime/utils/types.py @@ -35,6 +35,9 @@ class Sample: rollout_routed_experts: list[list[int]] | None = None # Routed experts from rollout engine remove_sample: bool = False teacher_log_probs: list[float] | None = None # Log probabilities from teacher model for OPD + # Privileged information (e.g. the ground-truth solution) shown only to the OPSD teacher, + # used to build the teacher's privileged context for on-policy self-distillation. + privileged_info: str | None = None class Status(Enum): PENDING = "pending" diff --git a/tests/test_opsd_jsd.py b/tests/test_opsd_jsd.py new file mode 100644 index 0000000000..38ddc414b4 --- /dev/null +++ b/tests/test_opsd_jsd.py @@ -0,0 +1,109 @@ +"""Unit tests for the OPSD (on-policy self-distillation) full-vocab JSD core. + +These run on CPU with a single (no-op) process group, exercising the same code +path used under tensor/vocab parallelism (where the all-reduces become real). +""" + +import math + +import pytest +import torch +import torch.nn.functional as F + +from slime.utils.dp_schedule import repack_micro_batches_by_length +from slime.utils.ppo_utils import compute_vocab_parallel_jsd + + +def test_repack_micro_batches_preserves_coverage_and_respects_budget(): + # One rank, two steps. Student schedule (packed by student lengths) below; + # teacher lengths are larger so the teacher must use more, smaller micro-batches. + micro_batch_indices = [[0, 1], [2, 3]] # step 0: 1 mb (samples 0,1); step 1: 1 mb (2,3) + num_microbatches = [1, 1] + teacher_lengths = [600, 600, 600, 600] + new_mbs, new_num_mb = repack_micro_batches_by_length( + micro_batch_indices, num_microbatches, teacher_lengths, max_tokens_per_bin=1000 + ) + # Each step's samples must be covered exactly once, step boundaries preserved. + assert sum(new_num_mb) == len(new_mbs) + assert len(new_num_mb) == 2 + assert sorted(i for mb in new_mbs for i in mb) == [0, 1, 2, 3] + # 600+600 > 1000 -> each sample lands in its own bin: 2 mbs per step. + assert new_num_mb == [2, 2] + # Every micro-batch respects the token budget. + for mb in new_mbs: + assert sum(teacher_lengths[i] for i in mb) <= 1000 + + +def test_repack_micro_batches_oversized_sample_alone(): + new_mbs, new_num_mb = repack_micro_batches_by_length( + [[0, 1, 2]], [1], lengths=[100, 5000, 100], max_tokens_per_bin=1000 + ) + assert sorted(i for mb in new_mbs for i in mb) == [0, 1, 2] + # The 5000-token sample must occupy its own bin. + assert any(mb == [1] for mb in new_mbs) + + +def _reference_jsd(student_logits: torch.Tensor, teacher_logits: torch.Tensor, beta: float) -> torch.Tensor: + """Independent dense reference using F.kl_div, mirroring OPSD's GOLD trainer. + + M = (1 - beta) * student + beta * teacher + JSD = beta * KL(teacher || M) + (1 - beta) * KL(student || M) + Endpoints: beta=0 -> KL(teacher || student); beta=1 -> KL(student || teacher). + """ + s = F.log_softmax(student_logits, dim=-1) + t = F.log_softmax(teacher_logits, dim=-1) + if beta == 0.0: + # F.kl_div(input, target, log_target) = sum target * (log target - input) = KL(target || exp(input)) + return F.kl_div(s, t, reduction="none", log_target=True).sum(-1) + if beta == 1.0: + return F.kl_div(t, s, reduction="none", log_target=True).sum(-1) + m = torch.logsumexp(torch.stack([s + math.log(1 - beta), t + math.log(beta)]), dim=0) + kl_teacher = F.kl_div(m, t, reduction="none", log_target=True) + kl_student = F.kl_div(m, s, reduction="none", log_target=True) + return (beta * kl_teacher + (1 - beta) * kl_student).sum(-1) + + +@pytest.mark.parametrize("beta", [0.0, 0.3, 0.5, 1.0]) +def test_jsd_value_matches_reference(beta): + torch.manual_seed(0) + student = torch.randn(7, 23, dtype=torch.float64) + teacher = torch.randn(7, 23, dtype=torch.float64) + got = compute_vocab_parallel_jsd(student, teacher, beta, process_group=None) + expected = _reference_jsd(student, teacher, beta) + assert torch.allclose(got, expected, atol=1e-9) + + +@pytest.mark.parametrize("beta", [0.0, 0.3, 0.5, 1.0]) +def test_jsd_gradient_only_through_student(beta): + torch.manual_seed(1) + teacher = torch.randn(5, 17, dtype=torch.float64) + + student = torch.randn(5, 17, dtype=torch.float64, requires_grad=True) + compute_vocab_parallel_jsd(student, teacher, beta, process_group=None).sum().backward() + + student_ref = student.detach().clone().requires_grad_(True) + _reference_jsd(student_ref, teacher, beta).sum().backward() + + assert torch.allclose(student.grad, student_ref.grad, atol=1e-9) + + +@pytest.mark.parametrize("temperature", [0.5, 2.0]) +def test_jsd_temperature_matches_reference(temperature): + torch.manual_seed(3) + student = torch.randn(6, 19, dtype=torch.float64) + teacher = torch.randn(6, 19, dtype=torch.float64) + got = compute_vocab_parallel_jsd(student, teacher, 0.5, process_group=None, temperature=temperature) + expected = _reference_jsd(student / temperature, teacher / temperature, 0.5) + assert torch.allclose(got, expected, atol=1e-9) + + +@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0]) +def test_jsd_non_negative_and_zero_for_identical(beta): + torch.manual_seed(2) + logits = torch.randn(4, 11) + assert (compute_vocab_parallel_jsd(logits, logits.clone(), beta, process_group=None) >= -1e-6).all() + assert torch.allclose( + compute_vocab_parallel_jsd(logits, logits.clone(), beta, process_group=None), + torch.zeros(4), + atol=1e-6, + ) diff --git a/tests/test_opsd_jsd_dist.py b/tests/test_opsd_jsd_dist.py new file mode 100644 index 0000000000..beb93cb62f --- /dev/null +++ b/tests/test_opsd_jsd_dist.py @@ -0,0 +1,81 @@ +"""Distributed (tensor-parallel) test for the OPSD vocab-parallel JSD. + +Spawns 2 gloo workers, shards the vocabulary across them, and checks that the +per-token JSD value *and the student-logit gradient* match a single-process dense +computation. This is the configuration (TP > 1) that exercises the cross-rank +gradient coupling through the softmax normalizer — a TP=1 test cannot catch a +wrong all-reduce backward. +""" + +from __future__ import annotations + +import os +import socket + +import pytest +import torch +import torch.multiprocessing as mp + +from slime.utils.ppo_utils import compute_vocab_parallel_jsd + +N, V = 4, 8 + + +def _free_port() -> int: + s = socket.socket() + s.bind(("", 0)) + port = s.getsockname()[1] + s.close() + return port + + +def _full_logits(): + g = torch.Generator().manual_seed(1234) + student = torch.randn(N, V, dtype=torch.float64, generator=g) + teacher = torch.randn(N, V, dtype=torch.float64, generator=g) + return student, teacher + + +def _worker(rank: int, world_size: int, master_port: int, beta: float, result_dir: str) -> None: + import torch.distributed as dist + + os.environ["MASTER_ADDR"] = "127.0.0.1" + os.environ["MASTER_PORT"] = str(master_port) + dist.init_process_group(backend="gloo", rank=rank, world_size=world_size) + group = dist.group.WORLD + + student, teacher = _full_logits() + shard = V // world_size + sl = slice(rank * shard, (rank + 1) * shard) + s = student[:, sl].clone().requires_grad_(True) + t = teacher[:, sl].clone() + + jsd = compute_vocab_parallel_jsd(s, t, beta, group) # [N], replicated across ranks + jsd.sum().backward() + + torch.save({"jsd": jsd.detach(), "grad": s.grad}, os.path.join(result_dir, f"rank{rank}.pt")) + dist.barrier() + dist.destroy_process_group() + + +@pytest.mark.parametrize("beta", [0.0, 0.5, 1.0]) +def test_vocab_parallel_jsd_matches_dense_under_tp(beta, tmp_path): + world_size = 2 + port = _free_port() + mp.spawn(_worker, args=(world_size, port, beta, str(tmp_path)), nprocs=world_size, join=True) + + # Distributed results: jsd is replicated; reassemble the grad from vocab shards. + rank_outs = [torch.load(os.path.join(str(tmp_path), f"rank{r}.pt")) for r in range(world_size)] + jsd_dist = rank_outs[0]["jsd"] + grad_dist = torch.cat([rank_outs[r]["grad"] for r in range(world_size)], dim=1) + + # Dense single-process reference (group=None) on the full vocab. + student, teacher = _full_logits() + s_full = student.clone().requires_grad_(True) + jsd_dense = compute_vocab_parallel_jsd(s_full, teacher.clone(), beta, None) + jsd_dense.sum().backward() + + assert torch.allclose(jsd_dist, jsd_dense, atol=1e-9), (jsd_dist - jsd_dense).abs().max() + # The student-logit gradient is the key check: a wrong (identity) normalizer + # backward under TP would corrupt this even though the forward value is correct. + assert torch.allclose(grad_dist, s_full.grad, atol=1e-9), (grad_dist - s_full.grad).abs().max() diff --git a/tests/test_qwen2.5_0.5B_opsd.py b/tests/test_qwen2.5_0.5B_opsd.py new file mode 100644 index 0000000000..5d71ff8cd9 --- /dev/null +++ b/tests/test_qwen2.5_0.5B_opsd.py @@ -0,0 +1,121 @@ +import os + +import slime.utils.external_utils.command_utils as U + +MODEL_NAME = "Qwen2.5-0.5B-Instruct" +MODEL_TYPE = "qwen2.5-0.5B" +NUM_GPUS = 4 + + +def prepare(): + U.exec_command("mkdir -p /root/models /root/datasets") + U.exec_command(f"hf download Qwen/{MODEL_NAME} --local-dir /root/models/{MODEL_NAME}") + U.hf_download_dataset("zhuzilin/dapo-math-17k") + # OPSD loads the teacher (frozen init policy) directly into Megatron, so we need a + # Megatron (torch_dist) checkpoint for both the student load and the teacher. + U.convert_checkpoint(MODEL_NAME, MODEL_TYPE, NUM_GPUS) + + +def execute(): + torch_dist_ckpt = f"/root/{MODEL_NAME}_torch_dist" + + # Teacher (frozen init policy) and student both start from the converted checkpoint. + ckpt_args = ( + f"--hf-checkpoint /root/models/{MODEL_NAME}/ " + f"--ref-load {torch_dist_ckpt}/ " + f"--load {torch_dist_ckpt}/ " + ) + + # Smoke test: the dapo-math dataset has no full reference-solution field, so we use the + # `label` field as the privileged information just to exercise the OPSD data/forward path. + rollout_args = ( + "--prompt-data /root/datasets/dapo-math-17k/dapo-math-17k.jsonl " + "--input-key prompt " + "--label-key label " + "--apply-chat-template " + "--rollout-shuffle " + "--rm-type deepscaler " + "--num-rollout 2 " + "--rollout-batch-size 4 " + "--n-samples-per-prompt 4 " + "--rollout-max-response-len 1024 " + "--rollout-temperature 0.8 " + "--global-batch-size 16 " + "--balance-data " + ) + + perf_args = ( + "--tensor-model-parallel-size 1 " + "--sequence-parallel " + "--pipeline-model-parallel-size 1 " + "--context-parallel-size 1 " # OPSD requires context-parallel-size 1 + "--expert-model-parallel-size 1 " + "--expert-tensor-parallel-size 1 " + "--use-dynamic-batch-size " + "--max-tokens-per-gpu 9216 " + ) + + # On-Policy Self-Distillation (OPSD): single model is both student and teacher; the + # frozen teacher is conditioned on privileged info and the loss is a full-vocab JSD. + opsd_args = ( + "--use-opd " + "--opd-type self " + f"--opd-teacher-load {torch_dist_ckpt}/ " + "--opsd-beta 0.5 " + "--opsd-jsd-clip 0.05 " + "--opsd-privileged-info-key label " + ) + + optimizer_args = ( + "--optimizer adam " + "--lr 1e-6 " + "--lr-decay-style constant " + "--weight-decay 0.1 " + "--adam-beta1 0.9 " + "--adam-beta2 0.98 " + ) + + sglang_args = ( + "--rollout-num-gpus-per-engine 1 " + "--sglang-mem-fraction-static 0.7 " + "--sglang-cuda-graph-max-bs 16 " + "--sglang-enable-metrics " + ) + + ci_args = "--ci-test " + + misc_args = ( + "--attention-dropout 0.0 " + "--hidden-dropout 0.0 " + "--accumulate-allreduce-grads-in-fp32 " + "--attention-softmax-in-fp32 " + "--attention-backend flash " + "--actor-num-nodes 1 " + f"--actor-num-gpus-per-node {NUM_GPUS} " + "--colocate " + ) + + train_args = ( + f"{ckpt_args} " + f"{rollout_args} " + f"{optimizer_args} " + f"{opsd_args} " + f"{U.get_default_wandb_args(__file__)} " + f"{perf_args} " + f"{sglang_args} " + f"{ci_args} " + f"{misc_args} " + ) + + U.execute_train( + train_args=train_args, + num_gpus_per_node=NUM_GPUS, + megatron_model_type=MODEL_TYPE, + ) + + +if __name__ == "__main__": + prepare() + for proxy_var in ("http_proxy", "https_proxy", "HTTP_PROXY", "HTTPS_PROXY"): + os.environ.pop(proxy_var, None) + execute()