[feat] support opd rl#9641
Conversation
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request implements Megatron On-Policy Distillation as RL (OPD-RL) by integrating teacher KL as a GRPO advantage across local and Ray-based GKD and GRPO trainers. It also introduces OpenEnvScheduler and OpenEnvWrapper to support multi-turn rollouts in OpenEnv environments. The review feedback highlights several critical issues, including missing imports and potential AttributeErrors in gkd_helpers.py, a rank-guarding mismatch in teacher_mixin.py that could cause runtime failures, and synchronous blocking calls in OpenEnvScheduler that should be run in separate threads to avoid blocking the asyncio event loop. Additionally, defensive checks are recommended to prevent potential IndexError, StopIteration, and type promotion issues.
Important
The consumer version of Gemini Code Assist on GitHub is being sunset. Starting June 18, 2026, new organization installations will be blocked, and all code review activity will officially cease on July 17, 2026.
For more details on the timeline and next steps, please review the Help Documentation.
| # Teacher returned logprobs for the full sequence (prompt + response + end tokens). | ||
| # Locate the response portion by matching token IDs from the end. | ||
| rti = response_token_ids[i] | ||
| if isinstance(rti[0], list): |
There was a problem hiding this comment.
If rti is empty (e.g., if the model generated an empty response), accessing rti[0] will raise an IndexError. We should check if rti is non-empty before accessing its first element.
| if isinstance(rti[0], list): | |
| if rti and isinstance(rti[0], list): |
| tid, info = next(iter(pos_lp.items())) | ||
| lps.append([info['logprob']]) | ||
| ixs.append([int(tid)]) |
There was a problem hiding this comment.
If pos_lp is empty, next(iter(pos_lp.items())) will raise a StopIteration exception. We should handle the empty case defensively to avoid crashes.
| tid, info = next(iter(pos_lp.items())) | |
| lps.append([info['logprob']]) | |
| ixs.append([int(tid)]) | |
| if pos_lp: | |
| tid, info = next(iter(pos_lp.items())) | |
| lps.append([info['logprob']]) | |
| ixs.append([int(tid)]) | |
| else: | |
| lps.append([0.0]) | |
| ixs.append([0]) |
| """ | ||
| d = teacher_per_token_logps - policy_per_token_logps | ||
| per_token = torch.exp(d) - d - 1 | ||
| return per_token * completion_mask |
There was a problem hiding this comment.
To prevent potential type promotion or precision mismatch errors (especially when using bfloat16 or float16 for logprobs and bool/long for the completion mask), we should explicitly cast completion_mask to the dtype of d before multiplication.
| return per_token * completion_mask | |
| return per_token * completion_mask.to(d.dtype) |
1 participant
No description provided.