72 lines
3.0 KiB
Python
72 lines
3.0 KiB
Python
# SPDX-FileCopyrightText: Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
|
|
# SPDX-License-Identifier: Apache-2.0
|
|
|
|
import aiohttp
|
|
import torch
|
|
|
|
from slime.utils.processing_utils import encode_image_for_rollout_engine
|
|
from slime.utils.types import Sample
|
|
|
|
|
|
async def reward_func(args, sample, **kwargs):
|
|
# For Lightning OPD: teacher log-probs are pre-computed in metadata,
|
|
# no teacher server call needed. Return a sentinel so post_process_rewards knows.
|
|
metadata = sample.metadata or {}
|
|
if metadata.get("is_lightning_opd", False) or metadata.get("is_offline_opd", False):
|
|
return {"lightning_opd": True}
|
|
|
|
payload = {
|
|
"input_ids": sample.tokens,
|
|
"sampling_params": {
|
|
"temperature": 0,
|
|
"max_new_tokens": 0,
|
|
"skip_special_tokens": False,
|
|
},
|
|
"return_logprob": True,
|
|
"logprob_start_len": 0,
|
|
}
|
|
|
|
if sample.multimodal_inputs and sample.multimodal_inputs.get("images"):
|
|
image_data = sample.multimodal_inputs["images"]
|
|
payload["image_data"] = [encode_image_for_rollout_engine(image) for image in image_data]
|
|
|
|
session_kwargs = {}
|
|
async with aiohttp.ClientSession(**session_kwargs) as session:
|
|
async with session.post(args.rm_url, json=payload) as resp:
|
|
resp.raise_for_status()
|
|
return await resp.json()
|
|
|
|
|
|
def post_process_rewards(args, samples: list[Sample], **kwargs):
|
|
"""Process rewards from teacher model and extract teacher log probabilities.
|
|
|
|
This function:
|
|
1. Extracts teacher log-probs from the reward response (which contains sglang's logprob output)
|
|
2. Trims them to match the response length
|
|
3. Stores them in sample.teacher_log_probs for OPD KL penalty computation
|
|
4. Returns scalar rewards (0.0 for pure distillation) compatible with GRPO/PPO
|
|
|
|
For Lightning OPD, teacher log-probs are pre-computed in the parquet
|
|
metadata instead of being fetched from a teacher server at runtime.
|
|
"""
|
|
raw_rewards = [sample.get_reward_value(args) for sample in samples]
|
|
response_lengths = [sample.response_length for sample in samples]
|
|
|
|
for i, (sample, reward) in enumerate(zip(samples, raw_rewards)):
|
|
metadata = sample.metadata or {}
|
|
if isinstance(reward, dict) and reward.get("lightning_opd"):
|
|
# Lightning OPD: teacher log-probs are pre-computed in metadata
|
|
pre_teacher_lp = metadata.get("teacher_log_probs", [])
|
|
sample.teacher_log_probs = torch.tensor(
|
|
[float(x) for x in pre_teacher_lp], dtype=torch.float32
|
|
)
|
|
else:
|
|
# Online OPD: extract teacher log-probs from sglang response
|
|
t_log_probs = torch.tensor(
|
|
[item[0] for item in reward["meta_info"]["input_token_logprobs"][1:]],
|
|
dtype=torch.float32,
|
|
)
|
|
sample.teacher_log_probs = t_log_probs[-response_lengths[i]:]
|
|
|
|
scalar_rewards = [0.0] * len(samples)
|
|
return scalar_rewards, scalar_rewards |