On-policy Distillation (OPD) Across Model Families

In the 003 tutorial, we learned the basics of OPD: minimize the reverse-KL of a teacher and student model. Our toy environment used Qwen3-8B to distill basic math answering capabilities to Qwen3-4B. What if we want our teacher model to be from a different, more powerful model family? We can’t compute reverse-KL divergence betweeen teacher and token logprobs because their tokenizer vocabularies are different!

This tutorial uses SimCT, a novel technique for aligning tokenizers, to train a Qwen3.6-35B-A3B student model using a DeepSeek-V4 Flash teacher model. Our target capability is multi-turn tool calling on Berkeley Function-Calling Leaderboard (https://gorilla.cs.berkeley.edu/blogs/13_bfcl_v3_multi_turn.html). We’ll first measure the accuracy of the teacher on our baseline eval and compare its accuracy to the base student model. After the teacher proves it can outperform the student, training will begin with a reverse-K learning cirriculum. We want to initialize the student model’s context up to the Kth tool call in the conversation, where K is initialized to N calls - 1. As the student model completes a task, K is decremented until the student is given the starter prompt or rollouts have completed.

Steps

1. Deploy DeepSeek V4 Flash as the teacher.
2. Load BFCL multi_turn_base, carve a train/eval split, and define reverse-K curriculum + shaped reward.
3. Evaluate the base student and teacher models on the held-out eval.
4. Define SimCT alignment, reward function, and OPD-adjustment.
5. Train with reverse-K curriculum + GRPO + cross-tokenizer OPD.
6. Evaluate the trained student with the same eval function.

import asyncio
import json
import re

from modal_training_gym import (
    DeploymentConfig,
    EvalConfig,
    EvalRowResult,
    ModelDeployment,
    Qwen3_6_35B,
    TrainConfig,
    list_checkpoints,
)
from modal_training_gym.common.models.base import HFModelConfiguration, ToolCall

from modal_training_gym.common.environments import (
    BfclMultiTurnConfig,
    BfclMultiTurnDataset,
    build_bfcl_env as build_env,
    build_bfcl_prefix_messages as build_prefix_messages,
    bfcl_prefix_turn_index,
    bfcl_tool_schemas_to_openai as tool_schemas_to_openai,
    run_bfcl_episode,
    to_json_schema,
)
from modal_training_gym.deploy_recipes.sglang_recipe import (
    DeepSeek_V4_Flash_SglangRecipe,
    Qwen3_6_35b_SglangRecipe,
)
from modal_training_gym.train_recipes.slime_recipe import Qwen3_6_35b_Recipe

Deploy DeepSeek-V4 Flash Teacher Model

Modal training gym provides a production SGLang recipe for the 284B-A13B FP4 checkpoint on 4×B200. The teacher model assigns its own logprobs to student token outputs, which is a prefill-bound process. Training-gym ships MegaMoE DeepGEMM kernels and DP attention (dp=4) for DSV4 to increase prefill speed on lengthy trajectories.

TEACHER_READY_TIMEOUT = 30 * 60
# OPD fires one /generate prefill per trajectory. With 16×8=128 traj/step, the
# default max_running_requests=16 saturates and returns 503s for minutes — raise
# the teacher queue and throttle client-side (see TEACHER_RM_CONCURRENCY below).
teacher_deployment = DeploymentConfig(
    model=HFModelConfiguration(model_name="deepseek-ai/DeepSeek-V4-Flash"),
    recipe=DeepSeek_V4_Flash_SglangRecipe(
        context_length=16384,
        startup_timeout=TEACHER_READY_TIMEOUT,
        max_running_requests=64,
    ),
    app_name="dsv4-teacher-model",
    served_model_name="deepseek-v4-flash",
).serve()
print(f"Teacher URL: {teacher_deployment.url}")

teacher_deployment.wait_until_ready(timeout=TEACHER_READY_TIMEOUT)

TEACHER_GENERATE_URL = f"{teacher_deployment.url}/generate"
TEACHER_RM_CONCURRENCY = 24

base_model = Qwen3_6_35B()

Loading BFCL V3 and Defining Train/Eval Split

The Berkeley Function Calling Leaderboard contains human-verified, multi-turn tasks with sufficient context to complete tasks. Each tool-call gets processed locally, and the final trajectory is verified against BFCL’s terminal evaluation. We save 30 out of the 200 multi-turn base tasks for the held-out evaluation split.

BFCL_EVAL_TAIL = 30
MAX_TURNS = 16
CURRICULUM_TAIL_MIN = 1
EVAL_TAIL_STEPS = CURRICULUM_TAIL_MIN
STUDENT_ENABLE_THINKING = False
OPD_SKIP_ON_TEACHER_FAILURE = True

dataset_config = BfclMultiTurnConfig(eval_tail=BFCL_EVAL_TAIL)
dataset = BfclMultiTurnDataset(split="train", config=dataset_config)
eval_dataset = BfclMultiTurnDataset(split="eval", config=dataset_config)

Custom Student Training Curriculum

Some multi-turn conversations run a dozen calls deep across several user turns, which adds unnecessary difficulty early in the training run. A well-known technique in machine learning is gradually increasing the difficulty of the training task, which has been found to speed up convergence and improve the quality of local optima (https://dl.acm.org/doi/epdf/10.1145/1553374.1553380). To increase training efficiency and accuracy, we start the student model with the previous context up to the Kth tool call, where K is initialized to N calls - 1 (curriculum tail T = 1, so K = N - T). After each rollout we lengthen the tail by one (T ← T + 1), so later rounds ask the student to finish more of the conversation from scratch.

def curriculum_rollout(args, rollout_id, data_source, evaluation=False):
    from slime.rollout.sglang_rollout import generate_rollout

    if evaluation:
        return generate_rollout(args, rollout_id, data_source, evaluation=True)

    if not hasattr(args, "curriculum_tail") or rollout_id == 0:
        args.curriculum_tail = CURRICULUM_TAIL_MIN

    T = args.curriculum_tail
    result = generate_rollout(args, rollout_id, data_source, evaluation=False)

    # Fixed schedule: lengthen the remaining horizon by one after every rollout.
    args.curriculum_tail = T + 1
    print(
        f"[curriculum] iter={rollout_id} tail={T} -> {args.curriculum_tail}",
        flush=True,
    )
    return result

Reward Function

Each trajectory gets a shaped score in [0, 1]:

0.25 * mean_partial + 0.20 * first_call + 0.10 * exec_score + 0.45 * terminal_pass

• mean_partial: average partial credit over all student tool calls.
• first_call: partial credit on the first tool call, providing extra credit for getting the first step right.
• exec_score: fraction of tool calls that executed without error.
• terminal_pass: 1 if the task is passed by BFCL’s terminal evaluation, 0 otherwise.

Per-call partial credit in [0, 1]:

1. +0.20 if the call parses (has a name)
2. +0.15 if that name is in the task’s tool catalog
3. +0.15 if args validate against the tool’s JSON schema
4. +0.50 * structural_match:
   • 0.4 for matching toolcall name
   • 0.3 for exactly matching toolcall keys else 0.15 for partial overlap
   • 0.3 * number of toolcall keys with matching values / total number of toolcall keys.

The ground-truth tool-call for structural match is derived by comparing each tool-call from the student’s response with the possible_answer field at the same sequence position as the student’s output.

_UUID_RE = re.compile(r"[0-9a-f]{8}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{4}-[0-9a-f]{12}")
_TS_RE = re.compile(r"\d{4}-\d{2}-\d{2}[T ]\d{2}:\d{2}:\d{2}")
_TMP_RE = re.compile(r"/tmp/[^\s/]+")

def _normalize_value(v) -> str:
    s = str(v).strip().lower()
    s = _UUID_RE.sub("<uuid>", s)
    s = _TS_RE.sub("<timestamp>", s)
    s = _TMP_RE.sub("/tmp/<tmp>", s)
    return s

def _coerce_args(call: dict | None) -> dict:
    if not call:
        return {}
    args = call.get("arguments", {})
    if isinstance(args, dict):
        return args
    if isinstance(args, str):
        try:
            parsed = json.loads(args)
            return parsed if isinstance(parsed, dict) else {}
        except (json.JSONDecodeError, TypeError):
            return {}
    return {}

def _validates(call: dict, schema: dict) -> bool:
    try:
        import jsonschema
        # BFCL's own func docs use `dict`/`list` where JSON Schema expects `object`/`array`.
        jsonschema.validate(_coerce_args(call), to_json_schema(schema))
        return True
    except Exception:
        return False

def _score_structural_match(student_call: dict | None, expert_call: dict) -> float:
    if not student_call:
        return 0.0
    if student_call.get("name") != expert_call.get("name"):
        return 0.0
    score = 0.4

    s_vals = _coerce_args(student_call)
    e_vals = _coerce_args(expert_call)
    s_args = set(s_vals.keys())
    e_args = set(e_vals.keys())
    if s_args == e_args:
        score += 0.3
    elif s_args & e_args:
        score += 0.15
    else:
        return score

    shared = s_args & e_args
    if not shared:
        score += 0.3
    else:
        matches = sum(
            1 for k in shared
            if _normalize_value(s_vals.get(k)) == _normalize_value(e_vals.get(k))
        )
        score += 0.3 * (matches / len(shared))
    return score

def _partial_credit(student_call: dict | None, expert_call: dict, tool_schemas: dict) -> float:
    if not student_call or "name" not in student_call:
        return 0.0
    score = 0.20
    spec = tool_schemas.get(student_call.get("name"))
    if spec is not None:
        score += 0.15
        schema = spec.get("parameters", spec) if isinstance(spec, dict) else spec
        if _validates(student_call, schema):
            score += 0.15
    score += 0.50 * _score_structural_match(student_call, expert_call)
    return min(1.0, score)

def trajectory_reward(
    student_calls: list,
    exec_successes: list,
    expert_calls: list,
    tool_schemas: dict,
    task_passed: bool | None,
    tail_len: int,
) -> float:
    T = max(int(tail_len), 1)
    graded = min(len(student_calls), len(expert_calls), T)
    partial_sum = sum(
        _partial_credit(student_calls[j], expert_calls[j], tool_schemas)
        for j in range(graded)
    )
    mean_partial = partial_sum / T
    first_call_score = (
        _partial_credit(student_calls[0], expert_calls[0], tool_schemas)
        if graded and expert_calls else 0.0
    )
    successful_execs = sum(1.0 for ok in exec_successes if ok)
    exec_score = min(successful_execs, T) / T
    if task_passed is not None:
        verify_score = 1.0 if task_passed else 0.0
        return (
            0.25 * mean_partial
            + 0.20 * first_call_score
            + 0.10 * exec_score
            + 0.45 * verify_score
        )
    return mean_partial

Baseline eval

Before training, we score both DeepSeek-V4 Flash (teacher) and Qwen3.6-35B-A3B (student) on the held-out BFCL split. Using our K curriculum, we initialize the agent context and the task’s class instances to the Kth ground-truth call, where K is set to N calls - 1. The teacher returns OpenAI-style structured tool_calls (SGLang --tool-call-parser deepseekv4); the student uses Qwen’s <tool_call> wire format via base_model.parse_response.

SERVED_CONTEXT_LEN = 16384
RESPONSE_TOKEN_CAP = 8192
CONTEXT_SAFETY_MARGIN = 512

EVAL_MAX_TURNS = EVAL_TAIL_STEPS * 2
MAX_CONSECUTIVE_TOOL_ERRORS = 3
DEPLOYMENT_READY_TIMEOUT = 1200

def _prompt_token_count(messages, tools=None) -> int:
    try:
        tok = _get_student_tokenizer()
        text = tok.apply_chat_template(
            messages, tools=tools, tokenize=False, add_generation_prompt=True
        )
        return len(tok(text, add_special_tokens=False)["input_ids"])
    except Exception:
        return sum(len(str(m.get("content", ""))) for m in messages) // 3

def _chat(deployment, messages, tools=None, max_tokens=None, max_attempts=12, *, qwen_thinking=False):
    if max_tokens is None:
        max_tokens = RESPONSE_TOKEN_CAP
    remaining = SERVED_CONTEXT_LEN - _prompt_token_count(messages, tools) - CONTEXT_SAFETY_MARGIN
    capped = min(max_tokens, remaining)
    if capped <= 0:
        return {"content": "", "tool_calls": []}

    kwargs = {
        "temperature": 0.0,
        "max_tokens": capped,
    }
    if tools is not None:
        kwargs["tools"] = tools
    if qwen_thinking is not None:
        kwargs["chat_template_kwargs"] = {"enable_thinking": qwen_thinking}
    return deployment.chat(
        messages,
        ensure_ready=False,
        max_attempts=max_attempts,
        **kwargs,
    )

def _actions_from_message(msg: dict) -> tuple[str, list[ToolCall]]:
    """Prefer structured API tool_calls (DeepSeek); else parse Qwen wire text."""
    content = msg.get("content") or msg.get("reasoning_content") or ""
    structured = msg.get("tool_calls") or []
    if structured:
        actions: list[ToolCall] = []
        for tc in structured:
            fn = tc.get("function") or {}
            raw_args = fn.get("arguments", {})
            if isinstance(raw_args, str):
                try:
                    raw_args = json.loads(raw_args) if raw_args else {}
                except json.JSONDecodeError:
                    raw_args = {}
            if not isinstance(raw_args, dict):
                raw_args = {}
            name = fn.get("name") or ""
            if name:
                actions.append(ToolCall(name=name, arguments=raw_args))
        return content, actions
    parsed = base_model.parse_response(content)
    return parsed.content, parsed.tool_calls

def bfcl_eval_fn(deployment: ModelDeployment, example: dict) -> EvalRowResult:
    label = json.loads(example.get("label", "{}"))
    task_id = label.get("task_id", "")
    N = label.get("total_steps", 1)
    flattened_calls = label.get("flattened_calls", [])
    K = max(0, N - EVAL_TAIL_STEPS)
    expert_call = flattened_calls[K] if K < len(flattened_calls) else {}

    served = deployment.deployment_config.served_model_name
    is_student = served != "deepseek-v4-flash"

    deployment.wait_until_ready(timeout=DEPLOYMENT_READY_TIMEOUT)

    episode = run_bfcl_episode(
        label,
        start_step=K,
        generate=lambda messages, tools: _chat(
            deployment,
            messages,
            tools=tools,
            qwen_thinking=STUDENT_ENABLE_THINKING if is_student else None,
        ),
        parse_response=_actions_from_message,
        max_turns=EVAL_MAX_TURNS,
        max_consecutive_errors=MAX_CONSECUTIVE_TOOL_ERRORS,
    )
    first_call = episode.first_call
    model_calls = episode.calls
    exec_successes = episode.execution_successes
    task_passed = episode.verdict.passed
    expert_calls = flattened_calls[K:K + len(model_calls)]
    shaped_score = trajectory_reward(
        model_calls,
        exec_successes,
        expert_calls,
        label.get("tool_schemas", {}),
        bool(task_passed),
        EVAL_TAIL_STEPS,
    )

    return EvalRowResult(
        score=shaped_score,
        response=episode.final_response,
        metadata={
            "task": task_id,
            "step_K": K,
            "eval_tail": EVAL_TAIL_STEPS,
            "task_passed": bool(task_passed),
            "terminal_score": 1.0 if task_passed else 0.0,
            "shaped_reward": shaped_score,
            "exec_successes": sum(exec_successes),
            "exec_calls": len(exec_successes),
            "parsed_call": first_call is not None,
            "tool_match": bool(first_call and first_call.get("name") == expert_call.get("name")),
        },
    )

def _print_eval_summary(name: str, result) -> None:
    def _frac(rows, key):
        if not rows:
            return 0.0
        return sum(1 for r in rows if r.metadata.get(key)) / len(rows)

    n = len(result.rows)
    print(f"{'Metric':<25} {name:>10}")
    print("-" * 37)
    print(f"{'Eval rows':<25} {n:>10d}")
    print(f"{'Shaped reward':<25} {result.mean:>10.3f}")
    print(f"{'Terminal pass rate':<25} {_frac(result.rows, 'task_passed'):>10.1%}")
    if n:
        print(f"{'Parsed tool call':<25} {_frac(result.rows, 'parsed_call'):>10.1%}")
        print(f"{'First-call tool match':<25} {_frac(result.rows, 'tool_match'):>10.1%}")

student_recipe = Qwen3_6_35b_SglangRecipe(context_length=SERVED_CONTEXT_LEN)
base_deployment = DeploymentConfig(model=base_model, recipe=student_recipe).serve()
print(f"Student URL: {base_deployment.url}")

eval_config = EvalConfig(dataset=eval_dataset, eval_fn=bfcl_eval_fn)

teacher_eval = None
print("--- Evaluating teacher (DeepSeek V4 Flash)... ---")
try:
    teacher_eval = eval_config.evaluate(teacher_deployment, max_concurrency=4)
    print(f"Teacher shaped reward: {teacher_eval.mean:.3f}")
    _print_eval_summary("Teacher", teacher_eval)
except Exception as e:
    print(
        f"[teacher-eval] FAILED ({e!r}) — continuing with student baseline",
        flush=True,
    )

print("--- Evaluating base student (shaped live reward + terminal verdict metadata)... ---")
try:
    base_eval = eval_config.evaluate(base_deployment, max_concurrency=4)
    print(f"Base shaped reward: {base_eval.mean:.3f}")
    _print_eval_summary("Base", base_eval)
except Exception as e:
    print(
        f"[base-eval] FAILED ({e!r}) — skipping baseline, proceeding to training",
        flush=True,
    )
    base_eval = None

if teacher_eval is not None and base_eval is not None:
    def _frac(rows, key):
        if not rows:
            return 0.0
        return sum(1 for r in rows if r.metadata.get(key)) / len(rows)

    t_pass = _frac(teacher_eval.rows, "task_passed")
    b_pass = _frac(base_eval.rows, "task_passed")
    print(
        f"[baseline] teacher pass={t_pass:.1%} student pass={b_pass:.1%} "
        f"(gap={t_pass - b_pass:+.1%})",
        flush=True,
    )

Cross-tokenizer alignment via SimCT Minimally Aligned Units (MTUs)

To solve the problem of cross-tokenizer misalignment, where two tokenizers may not share the same token vocabularies or may merge tokens differently, we employ an algorithm called SimCT (https://arxiv.org/abs/2605.07711) that constructs Minimally Aligned Units (MAUs) between boundaries of character-aligned tokens.

Here’s an example where SimCT would create a MAU for normalizing both model’s different vocabularies: For the sentence “I am happy today”, the student model tokenizes it as [“I”, “am”, “ha”, “pp”, “y”, “today”] and the teacher model tokenizes it as [“I”, “am”, “hap”, “py”, “today”]. The “I” and “am” tokens match up perfectly, but the word “happy” is split across multiple tokens.

SimCT greedily finds shared character boundaries between both tokenizations. Where boundaries don’t align (the word “happy”), it groups the tokens into a MAU spanning [“ha”, “pp”, “y”] for the student and [“hap”, “py”] for the teacher. The MAU log-probability is the joint probability in log space, which is the sum of its constituent token logprobs (or the log of the product of token probabilities). To integrate with slime’s per-token KL, we distribute the teacher’s MAU logprob sum equally across student tokens in the MAU, so that slime’s per-token summation reconstructs the correct MAU-level reverse KL.

Special tokens are excluded from alignment via skip_special_tokens=True during decoding, ensuring the character-level alignment only operates on actual content where the teacher’s logprobs are meaningful.

def align_cross_tokenizer(
    teacher_token_texts: list[str],
    teacher_logprobs: list[float],
    student_token_texts: list[str],
    return_coverage: bool = False,
):
    import torch

    def _offsets(texts):
        out, pos = [], 0
        for t in texts:
            out.append((pos, pos + len(t)))
            pos += len(t)
        return out

    t_off = _offsets(teacher_token_texts)
    s_off = _offsets(student_token_texts)
    t_bounds = {s for s, e in t_off} | {e for s, e in t_off}
    s_bounds = {s for s, e in s_off} | {e for s, e in s_off}
    shared = sorted(t_bounds & s_bounds)

    result = torch.zeros(len(student_token_texts), dtype=torch.float32)
    covered = torch.zeros(len(student_token_texts), dtype=torch.bool)
    for i in range(len(shared) - 1):
        lo, hi = shared[i], shared[i + 1]
        t_idx = [j for j, (s, e) in enumerate(t_off) if s >= lo and e <= hi]
        s_idx = [j for j, (s, e) in enumerate(s_off) if s >= lo and e <= hi]
        if t_idx and s_idx:
            mau_lp_sum = sum(teacher_logprobs[j] for j in t_idx)
            per_student_token = mau_lp_sum / len(s_idx)
            for j in s_idx:
                result[j] = per_student_token
                covered[j] = True
    if return_coverage:
        return result, covered
    return result

OPD-adjusted GRPO Advantage

Slime combines the reverse-KL loss from OPD into the GRPO advantage:

$$ A_t = A_t^{\mathrm{GRPO}} + \lambda \left(\log \pi_{\mathrm{student}}(y_t) - \log \pi_{\mathrm{teacher}}(y_t)\right) $$

where $\lambda$ is --opd-kl-coef, set to 0.3 in this tutorial.

_student_tokenizer_cache = {}

def _get_student_tokenizer():
    if "tok" not in _student_tokenizer_cache:
        from transformers import AutoTokenizer
        _student_tokenizer_cache["tok"] = AutoTokenizer.from_pretrained(
            "Qwen/Qwen3.6-35B-A3B", trust_remote_code=True
        )
    return _student_tokenizer_cache["tok"]

_teacher_tokenizer_cache = {}

def _get_teacher_tokenizer():
    if "tok" not in _teacher_tokenizer_cache:
        from transformers import AutoTokenizer
        _teacher_tokenizer_cache["tok"] = AutoTokenizer.from_pretrained(
            "deepseek-ai/DeepSeek-V4-Flash", trust_remote_code=True
        )
    return _teacher_tokenizer_cache["tok"]

def _teacher_response_boundary(teacher_tok, prefix_text, full_text):
    try:
        offsets = teacher_tok(
            full_text, add_special_tokens=False, return_offsets_mapping=True
        )["offset_mapping"]
        return sum(1 for start, _end in offsets if start < len(prefix_text))
    except (TypeError, KeyError, NotImplementedError, ValueError):
        return len(teacher_tok(prefix_text, add_special_tokens=False)["input_ids"])

_teacher_rm_sem: asyncio.Semaphore | None = None

def _get_teacher_rm_sem(limit: int) -> asyncio.Semaphore:
    global _teacher_rm_sem
    if _teacher_rm_sem is None:
        _teacher_rm_sem = asyncio.Semaphore(max(1, int(limit)))
    return _teacher_rm_sem

async def cross_tokenizer_reward(args, sample, **kwargs):
    """Collect teacher log-probs over the student's response (same context, no privileged info).

    Returns the teacher /generate JSON (logprobs for OPD). The BFCL *task* reward is
    computed later in ``cross_tokenizer_post_process`` — slime's rollout dump of
    ``reward: {'text': '', ...}`` is this payload, not a zero task score.
    """
    import aiohttp
    import random

    from modal_training_gym.common.deployment import _modal_proxy_auth_headers

    tokenizer = _get_student_tokenizer()

    resp_len = max(1, sample.response_length)
    split = max(0, len(sample.tokens) - resp_len)
    prefix_text = tokenizer.decode(sample.tokens[:split], skip_special_tokens=True)
    response_text = tokenizer.decode(sample.tokens[split:], skip_special_tokens=True)
    full_text = prefix_text + response_text

    teacher_tok = _get_teacher_tokenizer()
    prompt_length = _teacher_response_boundary(teacher_tok, prefix_text, full_text)

    payload = {
        "text": full_text,
        "sampling_params": {"temperature": 0, "max_new_tokens": 0, "skip_special_tokens": False},
        "return_logprob": True,
        "logprob_start_len": prompt_length,
        "return_text_in_logprobs": True,
    }
    response_token_count = resp_len

    skip_opd = {"meta_info": {"input_token_logprobs": []}}

    # Prefill-bound; long BFCL prefixes need more than a flat 60s.
    request_timeout = max(180, 60 + response_token_count // 20)
    max_attempts = 8
    rm_limit = int(getattr(args, "teacher_rm_concurrency", 24) or 24)
    sem = _get_teacher_rm_sem(rm_limit)
    async with sem:
        for attempt in range(max_attempts):
            try:
                async with aiohttp.ClientSession() as session:
                    async with session.post(
                        args.rm_url, json=payload,
                        headers=_modal_proxy_auth_headers(),
                        timeout=aiohttp.ClientTimeout(total=request_timeout),
                    ) as resp:
                        resp.raise_for_status()
                        return await resp.json()
            except (aiohttp.ClientError, asyncio.TimeoutError) as e:
                status = getattr(e, "status", None)
                if status == 400:
                    return skip_opd
                is_retryable = status in (503, 502, 504, 429, None)
                if attempt == max_attempts - 1 or not is_retryable:
                    return skip_opd
                base = min(30.0, 2 ** min(attempt, 5))
                wait = base + random.uniform(0, base * 0.25)
                await asyncio.sleep(wait)

    return skip_opd

async def tool_step_generate(args, sample, sampling_params):
    from slime.rollout.sglang_rollout import GenerateState
    from slime.utils.http_utils import post
    from slime.utils.types import Sample

    state = GenerateState(args)
    url = f"http://{args.sglang_router_ip}:{args.sglang_router_port}/generate"
    label = json.loads(getattr(sample, "label", "{}"))

    task_id = label["task_id"]
    N = label.get("total_steps", 1)
    flattened_calls = label.get("flattened_calls", [])
    tool_schemas = label.get("tool_schemas", {})

    T = int(getattr(args, "curriculum_tail", EVAL_TAIL_STEPS))
    K = max(0, N - T)
    max_turns = min(int(getattr(args, "max_turns", MAX_TURNS)), max(T * 2, 4))
    expert_call = flattened_calls[K] if K < len(flattened_calls) else {}

    prefix_msgs = build_prefix_messages(label, K)
    tools_list = tool_schemas_to_openai(tool_schemas)
    prompt_text = state.tokenizer.apply_chat_template(
        prefix_msgs, tools=tools_list, tokenize=False,
        add_generation_prompt=True, enable_thinking=STUDENT_ENABLE_THINKING,
    )
    prompt_ids = state.tokenizer(prompt_text, add_special_tokens=False)["input_ids"]

    base_render = state.tokenizer.apply_chat_template(
        prefix_msgs, tools=tools_list, tokenize=False,
        add_generation_prompt=False, enable_thinking=STUDENT_ENABLE_THINKING,
    )
    gen_suffix = prompt_text[len(base_render):]
    probe = state.tokenizer.apply_chat_template(
        prefix_msgs
        + [
            {"role": "assistant", "content": "\x01A\x01"},
            {"role": "tool", "content": "\x00OBS\x00"},
        ],
        tools=tools_list, tokenize=False,
        add_generation_prompt=True, enable_thinking=STUDENT_ENABLE_THINKING,
    )
    after_assistant = probe.split("\x01A\x01", 1)[1]
    obs_open, _rest = after_assistant.split("\x00OBS\x00", 1)
    obs_close = _rest[: len(_rest) - len(gen_suffix)]
    user_probe = state.tokenizer.apply_chat_template(
        prefix_msgs
        + [
            {"role": "assistant", "content": "\x01A\x01"},
            {"role": "user", "content": "\x00USER\x00"},
        ],
        tools=tools_list, tokenize=False,
        add_generation_prompt=True, enable_thinking=STUDENT_ENABLE_THINKING,
    )
    after_assistant = user_probe.split("\x01A\x01", 1)[1]
    user_open, _rest = after_assistant.split("\x00USER\x00", 1)
    user_close = _rest[: len(_rest) - len(gen_suffix)]
    stop_tok = obs_open.split("\n", 1)[0]

    def _log(msg):
        print(f"[rollout:{task_id} T={T} K={K}] {msg}", flush=True)

    def _abort_sample():
        pad_id = (
            state.tokenizer.pad_token_id
            if state.tokenizer.pad_token_id is not None
            else state.tokenizer.eos_token_id
            if state.tokenizer.eos_token_id is not None
            else prompt_ids[-1]
        )
        sample.status = Sample.Status.ABORTED
        sample.tokens = prompt_ids + [pad_id]
        sample.response_length = 1
        sample.response = ""
        sample.loss_mask = [0]
        return sample

    try:
        env = await asyncio.to_thread(build_env, label, K)
    except Exception as e:
        _log(f"env build failed: {e!r} — skipping rollout")
        return _abort_sample()
    trajectory_text = ""
    response_segments: list[tuple[str, int]] = []

    student_calls: list = []
    exec_successes: list = []
    finish_type = "stop"
    current_turn = bfcl_prefix_turn_index(label, K)
    for turn in range(max_turns):
        output = await post(url, {
            "text": prompt_text + trajectory_text,
            "sampling_params": sampling_params,
        })
        finish_type = output["meta_info"]["finish_reason"]["type"]
        if finish_type == "abort":
            return _abort_sample()

        model_text = output["text"]
        trajectory_text += model_text
        response_segments.append((model_text, 1))

        actions = base_model.parse_response(model_text).tool_calls
        action = actions[0] if actions else None

        if action is None:
            next_turn = current_turn + 1
            turns = label.get("turns", [])
            if finish_type != "length" and next_turn < len(turns):
                seg_open = (
                    user_open[len(stop_tok):]
                    if model_text.endswith(stop_tok)
                    else user_open
                )
                user_segment = (
                    seg_open + turns[next_turn]["user"] + user_close + gen_suffix
                )
                trajectory_text += user_segment
                response_segments.append((user_segment, 0))
                current_turn = next_turn
                continue
            break

        student_calls.append(
            {"name": action.name, "arguments": action.arguments}
        )
        try:
            result = await asyncio.to_thread(env.step, action)
            obs_text, is_error = result.observation.text, result.observation.is_error
        except Exception as e:
            _log(f"turn {turn} execution error: {e!r} — ending episode")
            exec_successes.append(False)
            finish_type = "stop"
            break
        exec_successes.append(not is_error)
        seg_open = (
            obs_open[len(stop_tok):] if model_text.endswith(stop_tok) else obs_open
        )
        obs_segment = seg_open + obs_text[:2000] + obs_close + gen_suffix
        trajectory_text += obs_segment
        response_segments.append((obs_segment, 0))

        if finish_type == "length":
            break

    try:
        verdict = env.evaluate()
        task_passed = verdict.passed
    except Exception as e:
        _log(f"evaluate() failed: {e!r} — marking failed")
        task_passed = False

    expert_calls = flattened_calls[K:K + len(student_calls)]
    shaped = trajectory_reward(
        student_calls,
        exec_successes,
        expert_calls,
        tool_schemas,
        bool(task_passed),
        T,
    )

    response_token_ids: list[int] = []
    loss_masks: list[int] = []
    for seg, trainable in response_segments:
        tids = state.tokenizer(seg, add_special_tokens=False)["input_ids"]
        response_token_ids += tids
        loss_masks += [trainable] * len(tids)

    sample.tokens = prompt_ids + response_token_ids
    sample.response_length = len(response_token_ids)
    sample.response = trajectory_text
    sample.loss_mask = loss_masks
    sample.status = (
        Sample.Status.TRUNCATED if finish_type == "length"
        else Sample.Status.COMPLETED
    )
    sample.metadata = {
        "student_calls": student_calls,
        "exec_successes": exec_successes,
        "expert_calls": expert_calls,
        "tool_schemas": tool_schemas,
        "task_passed": bool(task_passed),
        "tail_len": T,
        "curriculum_tail": T,
        "step_K": K,
        # Stash before custom_rm overwrites sample.reward with the teacher
        # /generate dict — dashboard rollout logging reads this when reward
        # is still non-scalar (OPD).
        "shaped_reward": shaped,
        "task_reward": shaped,
    }

    sample.bfcl_student_calls = student_calls
    sample.bfcl_exec_successes = exec_successes
    sample.bfcl_expert_calls = expert_calls
    sample.bfcl_tool_schemas = tool_schemas
    sample.bfcl_task_passed = bool(task_passed)
    sample.bfcl_tail_len = T
    return sample

Cross-tokenizer post-process

Upon successful completion of a response from the student trainer node, we take the raw prompt prefix + response token IDs and decode these (Qwen tokenizer) into the raw text. This raw text is posted to the teacher /generate endpoint for logprob computation. Upon successful logprob computation, the student receives a response from the teacher that contains the logprob, token ID, and decoded text (DeepSeek tokenizer).

To align the text, we create two arrays that contain the start and end character boundaries for teacher and student tokens. The intersection of start and end boundaries between the teacher and student arrays is the aligned text, and any tokens that are in between these two boundaries are used for the SimCT MAU calculation described above.

def cross_tokenizer_post_process(args, samples, **kwargs):
    """Compute step rewards + align teacher logprobs via SimCT MAUs."""
    import torch

    tokenizer = _get_student_tokenizer()
    raw_rewards = [s.get_reward_value(args) for s in samples]

    rewards = []
    first_call_hits = 0
    for sample in samples:
        meta = getattr(sample, "metadata", {}) or {}
        sc_list = (
            getattr(sample, "bfcl_student_calls", None)
            or meta.get("student_calls", [])
            or []
        )
        ec_list = (
            getattr(sample, "bfcl_expert_calls", None)
            or meta.get("expert_calls", [])
            or []
        )
        sc0 = sc_list[0] if sc_list else None
        ec0 = ec_list[0] if ec_list else {}
        if sc0 and ec0 and sc0.get("name") == ec0.get("name"):
            first_call_hits += 1
        task_passed_from_attr = (
            getattr(sample, "bfcl_task_passed", None)
            if hasattr(sample, "bfcl_task_passed")
            else None
        )
        task_passed_value = (
            task_passed_from_attr if task_passed_from_attr is not None
            else meta.get("task_passed")
        )
        r = trajectory_reward(
            sc_list,
            getattr(sample, "bfcl_exec_successes", None)
            or meta.get("exec_successes", [])
            or [],
            ec_list,
            getattr(sample, "bfcl_tool_schemas", None)
            or meta.get("tool_schemas", {}),
            task_passed_value,
            getattr(sample, "bfcl_tail_len", None)
            or meta.get("tail_len", 1),
        )
        rewards.append(r)
        try:
            sample.reward = r
        except Exception:
            pass
        if isinstance(meta, dict):
            meta["shaped_reward"] = r
            meta["task_reward"] = r
            try:
                sample.metadata = meta
            except Exception:
                pass

    cur_tail = getattr(args, "curriculum_tail", None)
    hit_rate = (first_call_hits / len(samples)) if samples else 0.0

    if rewards:
        n = len(samples)
        n_pass = sum(1 for s in samples if (getattr(s, "metadata", {}) or {}).get("task_passed"))
        all_exec = [
            ok for s in samples
            for ok in ((getattr(s, "metadata", {}) or {}).get("exec_successes") or [])
        ]
        exec_ok = (sum(1 for ok in all_exec if ok) / len(all_exec)) if all_exec else 0.0
        print(
            f"[bfcl] tail={cur_tail} "
            f"pass={n_pass}/{n} first_call={first_call_hits}/{n} exec_ok={exec_ok:.2f}",
            flush=True,
        )

    unaligned = total_resp = opd_dropped = 0
    for sample, reward in zip(samples, raw_rewards):
        r_meta = (reward or {}).get("meta_info", {})
        raw_logprobs = r_meta.get("input_token_logprobs", [])
        entries = raw_logprobs

        t_lps = [e[0] if e[0] is not None else 0.0 for e in entries if e is not None]
        t_texts = [e[2] if len(e) > 2 else "" for e in entries if e is not None]

        resp_tokens = sample.tokens[-sample.response_length:]
        s_texts = [tokenizer.decode([tid], skip_special_tokens=True) for tid in resp_tokens]
        aligned, covered = align_cross_tokenizer(t_texts, t_lps, s_texts, return_coverage=True)

        if not raw_logprobs and OPD_SKIP_ON_TEACHER_FAILURE:
            sample.teacher_log_probs = torch.full(
                (len(aligned),), float("nan"), dtype=torch.float32
            )
            opd_dropped += 1
            continue

        sample.teacher_log_probs = aligned
        resp_covered = covered
        unaligned += int((~resp_covered).sum().item())
        total_resp += int(resp_covered.numel())

    if total_resp:
        print(
            f"[simct] aligned={1 - (unaligned / total_resp):.1%}",
            flush=True,
        )
    if opd_dropped:
        print(
            f"[opd] dropped={opd_dropped}/{len(samples)} "
            "(teacher unavailable; GRPO retained)",
            flush=True,
        )

    return rewards, rewards

Training

Each rollout draws rollout_batch_size × n_samples_per_prompt trajectories (16 × 8 = 128 with the settings below). We run num_rollout=5 steps; after each step the reverse-K curriculum lengthens the remaining horizon by one (T ← T + 1). Tune rollout_temperature and num_rollout if you want more exploration or a longer run.

Training uses 2×8 H100 actor nodes.

training_run = TrainConfig(
    model=base_model,
    dataset=dataset,
    recipe=Qwen3_6_35b_Recipe(
        custom_rm_function=cross_tokenizer_reward,
        custom_generate_function=tool_step_generate,
        custom_reward_post_process_function=cross_tokenizer_post_process,
        rollout_function=curriculum_rollout,
        image_overlay=lambda img: img.pip_install(
            "modal~=1.5.2",
            "huggingface_hub~=1.12.0",
            "aiohttp~=3.13.0",
            "jsonschema~=4.23.0",
            "bfcl-eval==2026.3.23",
        ),

        gpu_type="H100",
        colocate=False,
        actor_num_nodes=2,
        actor_num_gpus_per_node=8,
        rollout_num_gpus=8,
        tensor_model_parallel_size=2,
        sequence_parallel=True,
        pipeline_model_parallel_size=2,
        context_parallel_size=2,
        expert_model_parallel_size=4,
        rollout_num_gpus_per_engine=8,
        sglang_dp_size=8,
        sglang_ep_size=8,
        sglang_cuda_graph_bs=[1, 2, 4, 8, 16, 24, 32, 48],
        sglang_max_running_requests=48,

        num_rollout=5,
        rollout_batch_size=16,
        n_samples_per_prompt=8,
        rollout_max_response_len=4000,
        rollout_temperature=1,
        sglang_mem_fraction_static=0.75,

        global_batch_size=16,
        lr=1e-6,
        kl_loss_coef=0.02,
        # The demo has five rollouts, so save one model-only checkpoint at the end.
        save_interval=5,
        no_save_optim=True,

        environment={
            "PYTHONPATH": "/root/Megatron-LM/:/root",
            "CUDA_DEVICE_MAX_CONNECTIONS": "1",
            "NCCL_NVLS_ENABLE": "1",
            "PYTORCH_CUDA_ALLOC_CONF": "expandable_segments:True",
        },

        extra_config={
            "use_opd": True,
            "opd_type": "sglang",
            "opd_kl_coef": 0.3,
            "rm_url": TEACHER_GENERATE_URL,
            "teacher_rm_concurrency": TEACHER_RM_CONCURRENCY,
            "max_turns": MAX_TURNS,
        },
    ),
)

print("--- Starting GRPO + cross-tokenizer OPD training... ---")
print(f"  Teacher: DeepSeek V4 Flash")
print(f"  Student: Qwen3.6-35B-A3B")
print(f"  Dataset: BFCL multi_turn_base, prefix-conditioned (task, K) rows")
print(f"  Reward: schema + live exec + structural match + terminal state/response verdict")
train_result = training_run.train()
print(f"Training run id: {train_result.training_run_id}")
print("--- Training complete ---")

Evaluate the trained student

Deploy the last checkpoint and re-run the held-out BFCL ids with the same evaluator from our earlier baseline.

checkpoint = list_checkpoints(train_result.training_run_id)[-1]
print(f"Checkpoint: {checkpoint.path}")

trained_deployment = DeploymentConfig(
    model=Qwen3_6_35B(),
    recipe=student_recipe,
    checkpoint=checkpoint,
    app_name="qwen3-6-35b-bfcl-trained",
    served_model_name="qwen3-6-35b-bfcl-trained",
).serve()
print(f"Trained student URL: {trained_deployment.url}")

print("--- Evaluating trained student (shaped live reward + terminal verdict metadata)... ---")
trained_eval = eval_config.evaluate(trained_deployment, max_concurrency=4)
print(f"Trained shaped reward: {trained_eval.mean:.3f}")

def _frac(rows, key):
    if not rows:
        return 0.0
    return sum(1 for r in rows if r.metadata.get(key)) / len(rows)

if base_eval is None:
    n_trained = len(trained_eval.rows)
    print(f"{'Metric':<25} {'Trained':>10}")
    print("-" * 37)
    print(f"{'Eval rows':<25} {n_trained:>10d}")
    print(f"{'Shaped reward':<25} {trained_eval.mean:>10.3f}")
    print(f"{'Terminal pass rate':<25} {_frac(trained_eval.rows, 'task_passed'):>10.1%}")
    print("(baseline eval was skipped — trained metrics only)")
    return

n_base = len(base_eval.rows)
n_trained = len(trained_eval.rows)

print(f"{'Metric':<25} {'Base':>10} {'Trained':>10} {'Delta':>10}")
print("-" * 57)
print(f"{'Eval rows':<25} {n_base:>10d} {n_trained:>10d} {'':>10}")
print(f"{'Shaped reward':<25} {base_eval.mean:>10.3f} {trained_eval.mean:>10.3f} {trained_eval.mean - base_eval.mean:>+10.3f}")

base_pass_rate = _frac(base_eval.rows, "task_passed")
trained_pass_rate = _frac(trained_eval.rows, "task_passed")
print(f"{'Terminal pass rate':<25} {base_pass_rate:>10.1%} {trained_pass_rate:>10.1%} {trained_pass_rate - base_pass_rate:>+10.1%}")

if n_base and n_trained:
    base_parse = _frac(base_eval.rows, "parsed_call")
    trained_parse = _frac(trained_eval.rows, "parsed_call")
    print(f"{'Parsed tool call':<25} {base_parse:>10.1%} {trained_parse:>10.1%} {trained_parse - base_parse:>+10.1%}")

    base_match = _frac(base_eval.rows, "tool_match")
    trained_match = _frac(trained_eval.rows, "tool_match")
    print(f"{'First-call tool match':<25} {base_match:>10.1%} {trained_match:>10.1%} {trained_match - base_match:>+10.1%}")
else:
    print("(no eval rows — check the eval dataset / deployment)")

Example results

Metric	Base	Trained	Delta
Eval rows	30	30	—
Shaped reward	0.707	0.789	+0.082
Terminal pass rate	53.3%	63.3%	+10.0 pp
Parsed tool call	90.0%	96.7%	+6.7 pp
First-call tool match	86.7%	93.3%	+6.6 pp

After only 5 rollouts, shaped reward rose from 0.707 → 0.789 and terminal pass rate from 53.3% → 63.3%, with parsed tool calls and first-call match also up. Still well below the teacher’s 83.3% pass rate. More rollouts would likely see further improvement.

Teacher reference

Metric	Teacher
Eval rows	30
Shaped reward	0.751
Terminal pass rate	83.3%
Parsed tool call	100.0%
First-call tool match	93.3%

The teacher’s terminal pass rate is 30 percentage points above the base student.

Future Possibilities

Some possible next steps for the tutorial:

1. Augment the training dataset with more long-context tasks from BFCL’s Long-Context Multi-Turn category.
2. Experiment with different student and teacher model configurations.
3. Add in privledged information to the teacher model for an even stronger distillation signal.
4. Extend Slime’s OPD loss to include top-k logprobs from the teacher model.
5. Try a different cross-tokenizer alignment strategy than averaging logprobs for reverse-KL such as f-divergence (https://neurips.cc/virtual/2025/loc/san-diego/poster/119176).
6. Increase the —opd-kl-coef value from 0.3 to see if a stronger reverse-KL signal improves training.
7. Ablate the OPD term completely from the GRPO advantage by setting —opd-kl-coef to 0.0.

Cross-tokenizer distillation is a novel OPD technique, and the misalignment of tokenizers across model families makes it an interesting research problem!

---

Related API Reference

• Qwen3_6_35B
• DeploymentConfig
• EvalConfig
• EvalRowResult
• Qwen3_6_35b_Recipe
• TrainConfig

Source: tutorials/rl/009_cross_tokenizer_distillation/009_cross_tokenizer_distillation.py | Open in Modal Notebook