初始化项目，由ModelHub XC社区提供模型

Model: ayh015/myLightningOPD
Source: Original Platform

slime/rollout/on_policy_distillation.py

@@ -0,0 +1,72 @@
+# 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