[recipe] feat: add Online Policy Distillation (OPD) recipe

[narutolhy]

Summary

• Add a lightweight Online Policy Distillation (OPD) recipe that extends verl's RayPPOTrainer with external teacher knowledge distillation
• Student trains with GRPO while minimizing KL divergence against a teacher model's per-token distribution
• Teacher is queried via ZMQ (reusing GKD recipe's teacher server)
• Zero modification to verl source code — pure inheritance + method override

Architecture

Student (FSDP + vLLM async rollout)     Teacher Server (ZMQ)
  ├─ Rollout: generate responses          ├─ Qwen3-32B (TP=4)
  ├─ Reward: task-specific scoring        └─ Returns top-k logprobs
  ├─ Teacher query: get ref_log_prob ◄────┘
  ├─ Compute GRPO advantages
  └─ Update: L_GRPO + kl_coef * KL(student || teacher)

Key Design

OPDTrainer inherits RayPPOTrainer and overrides a single method:

class OPDTrainer(RayPPOTrainer):
    def _compute_ref_log_prob(self, batch):
        # Query external teacher instead of reference model
        teacher_log_probs = self.teacher_client.get_teacher_log_probs(...)
        return DataProto.from_dict({"ref_log_prob": teacher_log_probs})

This plugs into verl's existing use_kl_loss mechanism — no changes needed to actor loss computation.

Files

File	Description
opd/opd_trainer.py	OPDTrainer + TeacherClient (ZMQ)
opd/main_opd.py	Entry point with custom OPDTaskRunner
opd/reward_gsm8k.py	GSM8K answer extraction and scoring
opd/run_opd.sh	Launch script with configurable KL coefficient
opd/README.md	Documentation + experimental results

Experimental Results

Qwen3-8B student + Qwen3-32B teacher on GSM8K, 500 steps, n=8:

Config	Avg Score	Non-zero Steps	Behavior
KL=0.0 (GRPO baseline)	0.0178	169/500 (34%)	Higher peaks, fewer consistent hits
KL=0.001 (OPD)	0.0108	213/500 (43%)	More consistent, lower peaks

Test plan

• Verified with Qwen3-0.6B student + 0.6B teacher (pipeline test)
• Verified with Qwen3-8B student + 32B teacher (real OPD)
• Ablation: KL=0.0, 0.001, 0.01, 0.1 coefficients
• 200-step and 500-step training runs

🤖 Generated with Claude Code

[gemini-code-assist]

Code Review

This pull request introduces Online Policy Distillation (OPD), allowing a student model to match a teacher model's distribution while optimizing task rewards. Key additions include the OPDTrainer, which overrides reference log probability computation to query an external teacher via ZMQ, and a GSM8K-specific reward function. Feedback focuses on critical improvements to the TeacherClient: addressing a security risk by using weights_only=True in torch.load, batching ZMQ requests for efficiency, fixing an off-by-one error in log probability indexing, and implementing socket timeouts to prevent indefinite blocking.

[gemini-code-assist]

Using weights_only=False with torch.load is a security risk as it can lead to arbitrary code execution when loading untrusted data. Since the data exchanged here consists of standard Python types and tensors, weights_only=True should be used instead.

Suggested change

	return torch.load(buf, weights_only=False)
	return torch.load(buf, weights_only=True)

[gemini-code-assist]

There are several issues in this block:

1. Efficiency: The teacher server is queried sequentially for each sample in the batch. This should be batched to avoid massive overhead during training.
2. Correctness (Off-by-one): The current logic src_idx = actual_prompt_len + t gets the log probability of the next token (ids[src_idx + 1]). For the first response token (t=0), it calculates log_prob(ids[actual_prompt_len + 1] | ...ids[actual_prompt_len]), effectively skipping the first response token and shifting all log probabilities by one.
3. Robustness: Initializing with zeros (logprob 0.0 => prob 1.0) is dangerous because if a query fails or a sample is skipped, it defaults to a perfect-match signal, which will provide incorrect training gradients.

I suggest batching the requests and fixing the indexing logic.

        # Initialize with a small value to avoid 0.0 (logprob) which implies prob=1.0
        teacher_log_probs = torch.full((batch_size, response_length), -100.0, dtype=torch.float32)

        all_ids = []
        valid_indices = []
        for i in range(batch_size):
            mask = attention_mask[i].bool()
            ids = input_ids[i][mask].tolist()
            if len(ids) >= 2:
                all_ids.append(ids)
                valid_indices.append(i)

        if not all_ids:
            return teacher_log_probs

        # Query teacher for the entire batch at once
        request = _serialize({
            "prompt_token_ids": all_ids,
            "temperature": 0.0,
            "max_tokens": 1,
            "only_response": False,
        })
        self.socket.send(request)
        response = _deserialize(self.socket.recv())

        if response.get("status") != "ok":
            logger.warning(f"Teacher error: {response.get('reason', 'unknown')}")
            return teacher_log_probs

        all_topk_logps = response["teacher_topk_logprobs"]
        all_topk_indices = response["teacher_topk_indices"]

        for idx_in_batch, i in enumerate(valid_indices):
            topk_logps = all_topk_logps[idx_in_batch]
            topk_indices = all_topk_indices[idx_in_batch]
            ids = all_ids[idx_in_batch]
            actual_prompt_len = int(attention_mask[i][:prompt_length].sum().item())

            for t in range(response_length):
                # To get logprob of token at ids[j], we look at teacher distribution after ids[j-1]
                src_idx = actual_prompt_len + t - 1
                if src_idx < 0 or src_idx >= topk_logps.shape[0] or src_idx + 1 >= len(ids):
                    continue
                next_token = ids[src_idx + 1]
                match = (topk_indices[src_idx] == next_token).nonzero(as_tuple=True)
                if len(match[0]) > 0:
                    teacher_log_probs[i, t] = topk_logps[src_idx, match[0][0]].item()

[gemini-code-assist]

The ZMQ REQ socket is synchronous and lacks timeouts. If the teacher server hangs or the network fails, the entire training process will block indefinitely. It is recommended to set RCVTIMEO and SNDTIMEO options.

Suggested change

	self.socket = self.context.socket(zmq.REQ)
	self.socket.connect(f"tcp://{server_ip}:{server_port}")
	self.socket = self.context.socket(zmq.REQ)
	self.socket.setsockopt(zmq.RCVTIMEO, 60000) # 60s timeout
	self.socket.setsockopt(zmq.SNDTIMEO, 60000)
	self.socket.connect(f"tcp://{server_ip}:{server_port}")