From de7a9d36d069f80cc8d0a21cceac1e378cad483e Mon Sep 17 00:00:00 2001 From: Kion Date: Fri, 6 Mar 2026 18:31:29 -0800 Subject: [PATCH 1/3] Add top-K GJS divergence mode to Tinker SDPO engine MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit Adds a `use_topk_divergence` toggle to TrainingConfig that switches the Tinker engine from scalar KL advantages (teacher_lp - student_lp) to distributional GJS advantages computed over top-K token distributions. When enabled, the engine fetches top-K teacher/student distributions via `sample_async(topk_prompt_logprobs=K)` and computes per-token Generalized Jensen-Shannon Divergence — the same formulation used by the local/Modal engine's torch-based `sdpo_loss.py`, but implemented in pure Python for the non-GPU Tinker path. Co-Authored-By: Claude Opus 4.6 --- claas/core/types.py | 1 + claas/eval/configs/base.yaml | 1 + claas/eval/types.py | 1 + claas/training/engine/tinker/engine.py | 275 +++++++++++++++++++++---- claas/training/gjs.py | 198 ++++++++++++++++++ tests/test_gjs.py | 168 +++++++++++++++ tests/test_tinker_engine.py | 188 +++++++++++++++++ 7 files changed, 788 insertions(+), 44 deletions(-) create mode 100644 claas/training/gjs.py create mode 100644 tests/test_gjs.py diff --git a/claas/core/types.py b/claas/core/types.py index 2e5ecfc..07ab79a 100644 --- a/claas/core/types.py +++ b/claas/core/types.py @@ -32,6 +32,7 @@ class TrainingConfig: max_grad_norm: float = 1.0 kl_reg_weight: float = 0.0 teacher_top_k: int = 100 + use_topk_divergence: bool = False steps_per_batch: int = 4 feedback_repetitions: int = 1 diff --git a/claas/eval/configs/base.yaml b/claas/eval/configs/base.yaml index 26e658b..d5f213a 100644 --- a/claas/eval/configs/base.yaml +++ b/claas/eval/configs/base.yaml @@ -35,5 +35,6 @@ training: max_grad_norm: 1.0 kl_reg_weight: 0.0 teacher_top_k: 100 + use_topk_divergence: false steps_per_batch: 4 feedback_repetitions: 1 diff --git a/claas/eval/types.py b/claas/eval/types.py index a43fce5..b2c685c 100644 --- a/claas/eval/types.py +++ b/claas/eval/types.py @@ -131,6 +131,7 @@ class TinkerDistillMetrics: completion_len: int = 0 batch_size: int = 0 steps_per_batch_applied: int = 1 + divergence_mode: str = "scalar_kl" @dataclass diff --git a/claas/training/engine/tinker/engine.py b/claas/training/engine/tinker/engine.py index ed0305e..e226c51 100644 --- a/claas/training/engine/tinker/engine.py +++ b/claas/training/engine/tinker/engine.py @@ -13,8 +13,9 @@ import asyncio import logging import os +from dataclasses import dataclass from datetime import datetime, timezone -from typing import Any, TypedDict +from typing import Any from urllib.parse import quote import httpx @@ -46,6 +47,12 @@ lora_exists as state_lora_exists, set_tinker_path, ) +from claas.training.gjs import ( + SequenceTopK, + compute_topk_gjs, + extract_token_logprobs, + slice_completion_topk, +) from claas.training.teacher_helpers import build_teacher_messages, teacher_messages_to_chat_template logger = logging.getLogger(__name__) @@ -55,6 +62,59 @@ _MAX_KL_GAIN = 4.0 +# ------------------------------------------------------------------ +# Discriminated-union sample/behavior types +# ------------------------------------------------------------------ + + +@dataclass(frozen=True) +class SampleCore: + """Fields shared across both scalar and top-K prepared samples.""" + + full_tokens: list[int] + input_tokens: list[int] + target_tokens: list[int] + prompt_len: int + completion_len: int + teacher_scored_text: str + + +@dataclass(frozen=True) +class ScalarPreparedSample(SampleCore): + """Sample with scalar teacher logprobs (current approach).""" + + teacher_logprobs: list[float] + + +@dataclass(frozen=True) +class TopKPreparedSample(SampleCore): + """Sample with top-K teacher distributions (GJS approach).""" + + teacher_logprobs: list[float] + teacher_topk: SequenceTopK + + +PreparedSample = ScalarPreparedSample | TopKPreparedSample + + +@dataclass(frozen=True) +class ScalarBehavior: + """Scalar student logprobs for IS correction.""" + + logprobs: list[float] + + +@dataclass(frozen=True) +class TopKBehavior: + """Top-K student distributions + scalar logprobs for IS correction.""" + + logprobs: list[float] + topk: SequenceTopK + + +BehaviorSignal = ScalarBehavior | TopKBehavior + + class TinkerTrainingEngine(TrainingEngine): """Executes training and LoRA management through the Tinker Python SDK.""" @@ -187,6 +247,9 @@ async def distill( kl_coef = payload.training.alpha steps_per_batch = payload.training.steps_per_batch feedback_repetitions = payload.training.feedback_repetitions + use_topk = payload.training.use_topk_divergence + teacher_top_k = payload.training.teacher_top_k + alpha = payload.training.alpha # ── Phase 1: Setup (once per batch) ── training_client = await self.service.create_training_client_from_state_async( @@ -205,12 +268,14 @@ async def distill( tokenizer=tokenizer, teacher_sampling=teacher_sampling, feedback_repetitions=feedback_repetitions, + use_topk=use_topk, + teacher_top_k=teacher_top_k, ) for sample in payload.samples ] results = await asyncio.gather(*tasks) - prepared_samples = [r[0] for r in results] - behavior_logprobs = [r[1] for r in results] + prepared_samples: list[PreparedSample] = [r[0] for r in results] + behavior_signals: list[BehaviorSignal] = [r[1] for r in results] # ── Phase 3: Multi-step training ── step_metrics: list[dict[str, float | int]] = [] @@ -219,8 +284,9 @@ async def distill( datum_metrics = [ _build_sample_datum( prepared=prepared, - student_logprobs=behavior_logprobs[sample_idx], + behavior=behavior_signals[sample_idx], kl_coef=kl_coef, + alpha=alpha, ) for sample_idx, prepared in enumerate(prepared_samples) ] @@ -256,10 +322,18 @@ async def distill( if step_idx < steps_per_batch - 1: student_sampling = await training_client.save_weights_and_get_sampling_client_async() - behavior_logprobs = await _compute_student_logprobs_for_batch( - student_sampling=student_sampling, - prepared_samples=prepared_samples, - ) + if use_topk: + behavior_signals = await _compute_student_topk_for_batch( + student_sampling=student_sampling, + prepared_samples=prepared_samples, + top_k=teacher_top_k, + ) + else: + scalar_logprobs = await _compute_student_logprobs_for_batch( + student_sampling=student_sampling, + prepared_samples=prepared_samples, + ) + behavior_signals = [ScalarBehavior(logprobs=lps) for lps in scalar_logprobs] # ── Phase 4: Save & return (once per request) ── final_step = step_metrics[-1] @@ -297,9 +371,10 @@ async def distill( "lr": lr, "loss_fn": "importance_sampling", "timestamp": datetime.now(timezone.utc).isoformat(), - "teacher_scored_texts": [p["teacher_scored_text"] for p in prepared_samples], + "teacher_scored_texts": [p.teacher_scored_text for p in prepared_samples], "steps_per_batch_applied": steps_per_batch, "per_step_metrics": step_metrics, + "divergence_mode": "topk_gjs" if use_topk else "scalar_kl", } if final_fwd_metrics is not None: @@ -308,14 +383,9 @@ async def distill( return DistillResponse(lora_id=payload.lora_id, metadata=metadata) -class PreparedSample(TypedDict): - full_tokens: list[int] - input_tokens: list[int] - target_tokens: list[int] - prompt_len: int - completion_len: int - teacher_logprobs: list[float] - teacher_scored_text: str +# ------------------------------------------------------------------ +# Sample preparation +# ------------------------------------------------------------------ async def _prepare_sample_inputs( @@ -324,8 +394,10 @@ async def _prepare_sample_inputs( tokenizer: Any, teacher_sampling: Any, feedback_repetitions: int, -) -> tuple[PreparedSample, list[float]]: - """Prepare sample-invariant tensors and initial behavior logprobs.""" + use_topk: bool = False, + teacher_top_k: int = 100, +) -> tuple[PreparedSample, BehaviorSignal]: + """Prepare sample-invariant tensors and initial behavior signal.""" prompt_tokens = list(sample.prompt_token_ids) response_tokens = list(sample.response_token_ids) completion_len = len(response_tokens) @@ -361,61 +433,124 @@ async def _prepare_sample_inputs( teacher_full = T.ModelInput.from_ints(teacher_full_tokens) teacher_scored_text = tokenizer.decode(teacher_full_tokens, skip_special_tokens=False) - teacher_logprobs_full = await teacher_sampling.compute_logprobs_async(teacher_full) - teacher_logprobs = _slice_completion_logprobs( - teacher_logprobs_full, - teacher_prompt_len, - completion_len, - ) - - prepared = PreparedSample( + core_kwargs = dict( full_tokens=full_tokens, input_tokens=input_tokens, target_tokens=target_tokens, prompt_len=prompt_len, completion_len=completion_len, - teacher_logprobs=teacher_logprobs, teacher_scored_text=teacher_scored_text, ) - return prepared, list(sample.response_logprobs) + + initial_behavior = ScalarBehavior(logprobs=list(sample.response_logprobs)) + + if use_topk: + # Top-K path: fetch teacher top-K distributions via sample_async + response = await teacher_sampling.sample_async( + prompt=teacher_full, + num_samples=1, + sampling_params=T.SamplingParams(max_tokens=1), + include_prompt_logprobs=True, + topk_prompt_logprobs=teacher_top_k, + ) + topk_full = response.topk_prompt_logprobs + teacher_topk = slice_completion_topk(topk_full, teacher_prompt_len, completion_len) + teacher_logprobs = extract_token_logprobs( + teacher_topk, response_tokens, + ) + + prepared: PreparedSample = TopKPreparedSample( + **core_kwargs, + teacher_logprobs=teacher_logprobs, + teacher_topk=teacher_topk, + ) + else: + # Scalar path: fetch teacher logprobs via compute_logprobs_async + teacher_logprobs_full = await teacher_sampling.compute_logprobs_async(teacher_full) + teacher_logprobs = _slice_completion_logprobs( + teacher_logprobs_full, + teacher_prompt_len, + completion_len, + ) + + prepared = ScalarPreparedSample( + **core_kwargs, + teacher_logprobs=teacher_logprobs, + ) + + return prepared, initial_behavior + + +# ------------------------------------------------------------------ +# Datum construction +# ------------------------------------------------------------------ def _build_sample_datum( *, prepared: PreparedSample, - student_logprobs: list[float], + behavior: BehaviorSignal, kl_coef: float, + alpha: float, ) -> tuple[T.Datum, dict[str, float]]: """Build a Tinker datum from prepared teacher signals + current behavior policy.""" - completion_len = prepared["completion_len"] + completion_len = prepared.completion_len + student_logprobs = behavior.logprobs if len(student_logprobs) != completion_len: raise ValueError( f"student_logprobs length ({len(student_logprobs)}) != " f"completion_len ({completion_len})" ) - teacher_logprobs = prepared["teacher_logprobs"] - raw_kl_deltas = [t - s for s, t in zip(student_logprobs, teacher_logprobs, strict=True)] - adv_abs_mean_raw = sum(abs(d) for d in raw_kl_deltas) / max(len(raw_kl_deltas), 1) + if isinstance(prepared, TopKPreparedSample) and isinstance(behavior, TopKBehavior): + # Top-K GJS path + raw_advantages = compute_topk_gjs( + prepared.teacher_topk, + behavior.topk, + alpha, + ) + elif isinstance(prepared, ScalarPreparedSample) and isinstance(behavior, ScalarBehavior): + # Scalar KL path (current) + raw_advantages = [ + t - s + for s, t in zip(student_logprobs, prepared.teacher_logprobs, strict=True) + ] + elif isinstance(prepared, TopKPreparedSample) and isinstance(behavior, ScalarBehavior): + # Initial step in top-K mode: behavior is still scalar from rollout cache + # Fall back to scalar KL using teacher_logprobs derived from top-K + raw_advantages = [ + t - s + for s, t in zip(student_logprobs, prepared.teacher_logprobs, strict=True) + ] + else: + raise TypeError( + f"Incompatible prepared/behavior types: " + f"{type(prepared).__name__} + {type(behavior).__name__}" + ) + + adv_abs_mean_raw = sum(abs(a) for a in raw_advantages) / max(len(raw_advantages), 1) gain = 1.0 if adv_abs_mean_raw > 0: gain = min(max(_TARGET_ADV_ABS_MEAN / adv_abs_mean_raw, 1.0), _MAX_KL_GAIN) effective_kl_coef = kl_coef * gain - advantages = [ - effective_kl_coef * (t - s) - for s, t in zip(student_logprobs, teacher_logprobs, strict=True) + advantages = [effective_kl_coef * a for a in raw_advantages] + + # Also compute scalar KL for metrics (always available via teacher_logprobs) + raw_kl_deltas = [ + t - s + for s, t in zip(student_logprobs, prepared.teacher_logprobs, strict=True) ] - full_logprobs = [0.0] * prepared["prompt_len"] + student_logprobs - full_advantages = [0.0] * prepared["prompt_len"] + advantages + full_logprobs = [0.0] * prepared.prompt_len + student_logprobs + full_advantages = [0.0] * prepared.prompt_len + advantages shifted_logprobs = full_logprobs[1:] shifted_advantages = full_advantages[1:] datum = T.Datum( - model_input=T.ModelInput.from_ints(prepared["input_tokens"]), + model_input=T.ModelInput.from_ints(prepared.input_tokens), loss_fn_inputs={ - "target_tokens": TensorData(data=prepared["target_tokens"], dtype="int64"), + "target_tokens": TensorData(data=prepared.target_tokens, dtype="int64"), "logprobs": TensorData(data=shifted_logprobs, dtype="float32"), "advantages": TensorData(data=shifted_advantages, dtype="float32"), }, @@ -433,6 +568,11 @@ def _build_sample_datum( return datum, metrics +# ------------------------------------------------------------------ +# Student behavior recomputation +# ------------------------------------------------------------------ + + async def _compute_student_logprobs_for_batch( *, student_sampling: Any, @@ -455,13 +595,60 @@ async def _compute_student_logprobs_for_sample( prepared: PreparedSample, ) -> list[float]: """Compute completion logprobs for one sample under current student weights.""" - student_full = T.ModelInput.from_ints(prepared["full_tokens"]) + student_full = T.ModelInput.from_ints(prepared.full_tokens) student_logprobs_full = await student_sampling.compute_logprobs_async(student_full) return _slice_completion_logprobs( student_logprobs_full, - prepared["prompt_len"], - prepared["completion_len"], + prepared.prompt_len, + prepared.completion_len, + ) + + +async def _compute_student_topk_for_batch( + *, + student_sampling: Any, + prepared_samples: list[PreparedSample], + top_k: int, +) -> list[TopKBehavior]: + """Recompute student top-K behavior under the updated student policy.""" + tasks = [ + _compute_student_topk_for_sample( + student_sampling=student_sampling, + prepared=prepared, + top_k=top_k, + ) + for prepared in prepared_samples + ] + return await asyncio.gather(*tasks) + + +async def _compute_student_topk_for_sample( + *, + student_sampling: Any, + prepared: PreparedSample, + top_k: int, +) -> TopKBehavior: + """Compute student top-K distributions for one sample.""" + student_full = T.ModelInput.from_ints(prepared.full_tokens) + response = await student_sampling.sample_async( + prompt=student_full, + num_samples=1, + sampling_params=T.SamplingParams(max_tokens=1), + include_prompt_logprobs=True, + topk_prompt_logprobs=top_k, ) + topk_full = response.topk_prompt_logprobs + student_topk = slice_completion_topk(topk_full, prepared.prompt_len, prepared.completion_len) + + # Extract response token IDs from full_tokens + response_tokens = prepared.full_tokens[prepared.prompt_len:] + scalar_logprobs = extract_token_logprobs(student_topk, response_tokens) + return TopKBehavior(logprobs=scalar_logprobs, topk=student_topk) + + +# ------------------------------------------------------------------ +# Utilities +# ------------------------------------------------------------------ def _require_entry(lora_id: str, state_path: str) -> LoraEntry: diff --git a/claas/training/gjs.py b/claas/training/gjs.py new file mode 100644 index 0000000..fe882a7 --- /dev/null +++ b/claas/training/gjs.py @@ -0,0 +1,198 @@ +"""Pure-Python GJS divergence for the Tinker SDPO engine. + +Mirrors the torch-based GJS computation in ``sdpo_loss.py`` but uses only +the Python standard library (``math``). This is the non-GPU code path +used when the Tinker engine fetches top-K distributions via +``topk_prompt_logprobs`` on ``sample_async``. + +Reference: veRL ``core_algos.py:1120-1122`` for the renormalize-over-top-K +pattern and Hübotter et al. (2026) for the GJS formulation. +""" + +from __future__ import annotations + +import math + +# --------------------------------------------------------------------------- +# Type aliases +# --------------------------------------------------------------------------- + +TokenTopK = list[tuple[int, float]] +"""K (token_id, logprob) pairs for one position.""" + +SequenceTopK = list[TokenTopK] +"""One TokenTopK per token in the sequence.""" + +# --------------------------------------------------------------------------- +# Default floor log-probability (matches ``sdpo_loss._lookup_token_in_topk``) +# --------------------------------------------------------------------------- + +_DEFAULT_FLOOR_LOGPROB = -20.0 + + +# --------------------------------------------------------------------------- +# Public API +# --------------------------------------------------------------------------- + + +def compute_topk_gjs( + teacher_topk: SequenceTopK, + student_topk: SequenceTopK, + alpha: float, + floor_logprob: float = _DEFAULT_FLOOR_LOGPROB, +) -> list[float]: + """Per-token GJS divergence over top-K distributions. + + For each position the teacher's top-K token IDs define the shared + support (following the veRL reference). Student logprobs are looked + up in its own top-K set — missing entries fall back to *floor_logprob*. + + When ``alpha == 1.0`` the function degenerates to reverse-KL style + scalar advantages (``teacher_logprob - student_logprob`` at the + *first* token in the teacher's top-K set — the most-probable token). + This is kept for backward-compatibility with the scalar path. + + Args: + teacher_topk: Per-position teacher top-K distributions. + student_topk: Per-position student top-K distributions. + alpha: GJS interpolation weight (0 < alpha <= 1). + floor_logprob: Logprob to use when a token is absent from a top-K set. + + Returns: + List of per-token GJS divergence values (length = len(teacher_topk)). + """ + if len(teacher_topk) != len(student_topk): + raise ValueError( + f"teacher_topk length ({len(teacher_topk)}) != " + f"student_topk length ({len(student_topk)})" + ) + + result: list[float] = [] + for t_topk, s_topk in zip(teacher_topk, student_topk): + result.append(_gjs_at_position(t_topk, s_topk, alpha, floor_logprob)) + return result + + +def extract_token_logprobs( + topk: SequenceTopK, + token_ids: list[int], + floor_logprob: float = _DEFAULT_FLOOR_LOGPROB, +) -> list[float]: + """Look up specific token logprobs in a per-position top-K set. + + Pure-Python counterpart to ``sdpo_loss._lookup_token_in_topk``. + + Args: + topk: Per-position top-K distributions. + token_ids: Token IDs to look up (one per position). + floor_logprob: Value returned when the token is absent. + + Returns: + List of logprobs (length = len(token_ids)). + """ + if len(topk) != len(token_ids): + raise ValueError( + f"topk length ({len(topk)}) != token_ids length ({len(token_ids)})" + ) + + result: list[float] = [] + for pos_topk, tid in zip(topk, token_ids): + lp = floor_logprob + for tok_id, tok_lp in pos_topk: + if tok_id == tid: + lp = tok_lp + break + result.append(lp) + return result + + +def slice_completion_topk( + topk_full: list[TokenTopK | None], + prompt_len: int, + completion_len: int, +) -> SequenceTopK: + """Extract the completion portion of a full-sequence top-K list. + + Analogous to ``_slice_completion_logprobs`` but for top-K data. + ``None`` entries (e.g. the first unconditional position) are replaced + with empty lists. + + Args: + topk_full: Full-sequence top-K data (may contain ``None``). + prompt_len: Number of prompt tokens to skip. + completion_len: Number of completion tokens to extract. + + Returns: + SequenceTopK of length *completion_len*. + """ + raw = topk_full[prompt_len : prompt_len + completion_len] + return [entry if entry is not None else [] for entry in raw] + + +# --------------------------------------------------------------------------- +# Internal helpers +# --------------------------------------------------------------------------- + + +def _gjs_at_position( + teacher_topk: TokenTopK, + student_topk: TokenTopK, + alpha: float, + floor_logprob: float, +) -> float: + """GJS divergence at a single position over the teacher's top-K support.""" + if not teacher_topk: + return 0.0 + + # Build student lookup + student_lookup: dict[int, float] = {tid: lp for tid, lp in student_topk} + + # Shared support = teacher's top-K token IDs + teacher_ids = [tid for tid, _ in teacher_topk] + teacher_lps = [lp for _, lp in teacher_topk] + + student_lps = [ + student_lookup.get(tid, floor_logprob) for tid in teacher_ids + ] + + # Renormalize both distributions over the shared support + teacher_probs = _renormalize(teacher_lps) + student_probs = _renormalize(student_lps) + + if alpha >= 1.0: + # Degenerate to reverse KL: scalar teacher_lp - student_lp + # for the highest-probability teacher token + return teacher_lps[0] - student_lps[0] + + # Mixture: M = alpha * teacher + (1 - alpha) * student + mixture = [ + alpha * tp + (1.0 - alpha) * sp + for tp, sp in zip(teacher_probs, student_probs) + ] + + kl_teacher_m = _kl_divergence(teacher_probs, mixture) + kl_student_m = _kl_divergence(student_probs, mixture) + + return alpha * kl_teacher_m + (1.0 - alpha) * kl_student_m + + +def _renormalize(logprobs: list[float]) -> list[float]: + """Convert logprobs to normalized probabilities over the support.""" + max_lp = max(logprobs) if logprobs else 0.0 + # exp(lp - max_lp) for numerical stability + exps = [math.exp(lp - max_lp) for lp in logprobs] + total = sum(exps) + if total < 1e-30: + # Uniform fallback to avoid division by zero + n = len(exps) + return [1.0 / n] * n + return [e / total for e in exps] + + +def _kl_divergence(p: list[float], q: list[float]) -> float: + """KL(p || q) where p and q are probability vectors.""" + total = 0.0 + for pi, qi in zip(p, q): + if pi > 0 and qi > 0: + total += pi * math.log(pi / qi) + return total diff --git a/tests/test_gjs.py b/tests/test_gjs.py new file mode 100644 index 0000000..62837aa --- /dev/null +++ b/tests/test_gjs.py @@ -0,0 +1,168 @@ +"""Tests for the pure-Python GJS divergence module.""" + +from __future__ import annotations + +import math + +import pytest + +from claas.training.gjs import ( + compute_topk_gjs, + extract_token_logprobs, + slice_completion_topk, +) + +# ── compute_topk_gjs ──────────────────────────────────────────────── + + +def _make_topk(probs: list[float], start_id: int = 0) -> list[tuple[int, float]]: + """Build a TokenTopK from raw probabilities (converts to logprobs).""" + return [(start_id + i, math.log(p)) for i, p in enumerate(probs)] + + +def test_gjs_identical_distributions_is_zero(): + """GJS(p, p) = 0 for any alpha in (0, 1).""" + topk = [_make_topk([0.5, 0.3, 0.2])] + for alpha in [0.1, 0.5, 0.9]: + result = compute_topk_gjs(topk, topk, alpha=alpha) + assert len(result) == 1 + assert abs(result[0]) < 1e-10, f"GJS(p,p) should be 0, got {result[0]} for alpha={alpha}" + + +def test_gjs_symmetric_at_half(): + """GJS_0.5(p, q) == GJS_0.5(q, p) — Jensen-Shannon is symmetric at alpha=0.5.""" + p_topk = [_make_topk([0.7, 0.2, 0.1])] + q_topk = [_make_topk([0.3, 0.4, 0.3])] + + gjs_pq = compute_topk_gjs(p_topk, q_topk, alpha=0.5) + gjs_qp = compute_topk_gjs(q_topk, p_topk, alpha=0.5) + + assert abs(gjs_pq[0] - gjs_qp[0]) < 1e-10, ( + f"JSD should be symmetric: {gjs_pq[0]} != {gjs_qp[0]}" + ) + + +def test_gjs_non_negative(): + """GJS >= 0 for all inputs and alpha values.""" + p_topk = [_make_topk([0.6, 0.3, 0.1])] + q_topk = [_make_topk([0.2, 0.5, 0.3])] + + for alpha in [0.1, 0.3, 0.5, 0.7, 0.9]: + result = compute_topk_gjs(p_topk, q_topk, alpha=alpha) + assert result[0] >= -1e-10, f"GJS should be >= 0, got {result[0]} for alpha={alpha}" + + +def test_gjs_floor_logprob_used_for_missing_tokens(): + """When student's top-K doesn't contain a teacher token, floor is applied.""" + # Teacher has tokens 0, 1, 2 + teacher = [[(0, math.log(0.5)), (1, math.log(0.3)), (2, math.log(0.2))]] + # Student only has tokens 0, 3, 4 — tokens 1 and 2 will use floor + student = [[(0, math.log(0.6)), (3, math.log(0.3)), (4, math.log(0.1))]] + + result = compute_topk_gjs(teacher, student, alpha=0.5, floor_logprob=-20.0) + assert len(result) == 1 + # Should be positive (distributions differ significantly) + assert result[0] > 0.01 + + +def test_gjs_alpha_one_degenerates(): + """alpha=1.0 path returns scalar KL-like values (teacher_lp - student_lp).""" + teacher_lp = math.log(0.8) + student_lp = math.log(0.3) + # With alpha=1.0, result should be teacher_lp[0] - student_lp[0] + teacher = [[(0, teacher_lp), (1, math.log(0.2))]] + student = [[(0, student_lp), (1, math.log(0.7))]] + + result = compute_topk_gjs(teacher, student, alpha=1.0) + expected = teacher_lp - student_lp + assert abs(result[0] - expected) < 1e-10 + + +def test_gjs_multi_position(): + """Verify per-position output for a 3-position sequence.""" + teacher = [ + _make_topk([0.5, 0.3, 0.2]), + _make_topk([0.8, 0.1, 0.1]), + _make_topk([0.4, 0.4, 0.2]), + ] + student = [ + _make_topk([0.5, 0.3, 0.2]), # identical → 0 + _make_topk([0.3, 0.4, 0.3]), # different → positive + _make_topk([0.1, 0.1, 0.8]), # very different → larger positive + ] + + result = compute_topk_gjs(teacher, student, alpha=0.5) + assert len(result) == 3 + assert abs(result[0]) < 1e-10 # identical + assert result[1] > 0.01 # different + assert result[2] > result[1] # more different + + +def test_gjs_length_mismatch_raises(): + """Mismatched teacher/student lengths raise ValueError.""" + with pytest.raises(ValueError, match="length"): + compute_topk_gjs( + [_make_topk([0.5, 0.5])], + [_make_topk([0.5, 0.5]), _make_topk([0.3, 0.7])], + alpha=0.5, + ) + + +# ── extract_token_logprobs ────────────────────────────────────────── + + +def test_extract_token_logprobs_found(): + """Token in top-K returns correct logprob.""" + topk = [[(10, -1.5), (20, -2.0), (30, -3.0)]] + result = extract_token_logprobs(topk, [20]) + assert result == [-2.0] + + +def test_extract_token_logprobs_missing(): + """Token not in top-K returns floor.""" + topk = [[(10, -1.5), (20, -2.0)]] + result = extract_token_logprobs(topk, [99], floor_logprob=-20.0) + assert result == [-20.0] + + +def test_extract_token_logprobs_multi_position(): + """Multi-position extraction with mixed hits and misses.""" + topk = [ + [(10, -1.0), (20, -2.0)], + [(30, -0.5), (40, -1.5)], + [(50, -3.0)], + ] + result = extract_token_logprobs(topk, [10, 99, 50], floor_logprob=-20.0) + assert result == [-1.0, -20.0, -3.0] + + +def test_extract_token_logprobs_length_mismatch(): + """Mismatched lengths raise ValueError.""" + with pytest.raises(ValueError, match="length"): + extract_token_logprobs([[(1, -1.0)]], [1, 2]) + + +# ── slice_completion_topk ─────────────────────────────────────────── + + +def test_slice_completion_topk(): + """Correct slicing with None handling.""" + topk_full = [ + None, # position 0 (prompt, skipped) + [(1, -1.0)], # position 1 (prompt, skipped) + [(2, -2.0), (3, -2.5)], # position 2 (completion start) + None, # position 3 (None → empty list) + [(4, -3.0)], # position 4 + ] + result = slice_completion_topk(topk_full, prompt_len=2, completion_len=3) + assert len(result) == 3 + assert result[0] == [(2, -2.0), (3, -2.5)] + assert result[1] == [] # None replaced + assert result[2] == [(4, -3.0)] + + +def test_slice_completion_topk_all_none(): + """All None entries become empty lists.""" + topk_full = [None, None, None, None] + result = slice_completion_topk(topk_full, prompt_len=1, completion_len=2) + assert result == [[], []] diff --git a/tests/test_tinker_engine.py b/tests/test_tinker_engine.py index 7d43074..b91fd79 100644 --- a/tests/test_tinker_engine.py +++ b/tests/test_tinker_engine.py @@ -856,3 +856,191 @@ def test_engine_distill_uses_user_prompt_for_teacher(tinker_engine, mock_trainin tok = mock_training_client.get_tokenizer() tok.apply_chat_template.assert_called_once() assert tok.apply_chat_template.call_args[1]["add_generation_prompt"] is True + + +# ── Top-K GJS mode tests ───────────────────────────────────────────── + + +def _make_mock_topk_response(topk_data): + """Build a mock Tinker sample_async response with topk_prompt_logprobs.""" + resp = MagicMock() + resp.topk_prompt_logprobs = topk_data + return resp + + +def test_engine_distill_topk_mode(tinker_engine, mock_training_client): + """End-to-end distill with use_topk_divergence=True. + + Verifies that sample_async is called (instead of compute_logprobs_async) + for teacher signal acquisition and that GJS-derived advantages flow + through to the datum. + """ + engine, mock_service = tinker_engine + + set_tinker_path("test/lora", "tinker://init-ckpt", "gpt-oss/GPT-OSS-120B", 32, step=0) + mock_service.create_training_client_from_state_async = AsyncMock( + return_value=mock_training_client + ) + + # Teacher sampling client: returns top-K data via sample_async + teacher_sampler = MagicMock() + # Build top-K data: 100 positions (enough for slicing), each with K=3 entries + topk_entry = [(0, -0.5), (1, -1.0), (2, -2.0)] + topk_full = [topk_entry] * 100 + teacher_sampler.sample_async = AsyncMock( + return_value=_make_mock_topk_response(topk_full) + ) + # compute_logprobs_async should NOT be called in topk mode + teacher_sampler.compute_logprobs_async = AsyncMock( + side_effect=AssertionError("compute_logprobs_async should not be called in topk mode") + ) + mock_service.create_sampling_client_async = AsyncMock( + return_value=teacher_sampler + ) + + payload = DistillBatchRequestPayload( + lora_id="test/lora", + training=TrainingConfig( + steps_per_batch=1, + use_topk_divergence=True, + teacher_top_k=3, + ), + samples=[ + DistillBatchItem( + prompt="Hello", + response="World", + feedback="Good job", + response_logprobs=[-0.1, -0.2, -0.3, -0.4, -0.5], + prompt_token_ids=[0, 1, 2, 3, 4], + response_token_ids=[0, 1, 2, 3, 4], + user_prompt="Hello", + system_prompt="You are a helpful assistant.", + ) + ], + ) + + result = asyncio.run(engine.distill(payload)) + + assert isinstance(result, DistillResponse) + assert result.lora_id == "test/lora" + assert result.metadata["step"] == 1 + assert result.metadata["divergence_mode"] == "topk_gjs" + + # Teacher used sample_async, not compute_logprobs_async + teacher_sampler.sample_async.assert_called_once() + call_kwargs = teacher_sampler.sample_async.call_args[1] + assert call_kwargs["topk_prompt_logprobs"] == 3 + assert call_kwargs["include_prompt_logprobs"] is True + + # Training step happened + mock_training_client.forward_backward_async.assert_called_once() + mock_training_client.optim_step_async.assert_called_once() + + +def test_engine_distill_topk_multistep(tinker_engine, mock_training_client): + """Multi-step distill with top-K: verifies student top-K recomputation.""" + engine, mock_service = tinker_engine + + set_tinker_path("test/lora", "tinker://ckpt", "gpt-oss/GPT-OSS-120B", 32, step=0) + mock_service.create_training_client_from_state_async = AsyncMock( + return_value=mock_training_client + ) + + # Teacher sampling client with top-K + teacher_sampler = MagicMock() + topk_entry = [(0, -0.3), (1, -1.5), (2, -2.5)] + topk_full = [topk_entry] * 100 + teacher_sampler.sample_async = AsyncMock( + return_value=_make_mock_topk_response(topk_full) + ) + mock_service.create_sampling_client_async = AsyncMock( + return_value=teacher_sampler + ) + + # Student sampling client also returns top-K via sample_async + student_sampler = MagicMock() + student_topk_entry = [(0, -0.2), (1, -1.2), (2, -2.2)] + student_topk_full = [student_topk_entry] * 100 + student_sampler.sample_async = AsyncMock( + return_value=_make_mock_topk_response(student_topk_full) + ) + mock_training_client.save_weights_and_get_sampling_client_async = AsyncMock( + return_value=student_sampler + ) + + payload = DistillBatchRequestPayload( + lora_id="test/lora", + training=TrainingConfig( + steps_per_batch=2, + use_topk_divergence=True, + teacher_top_k=3, + ), + samples=[ + DistillBatchItem( + prompt="Hello", + response="World", + feedback="Nice", + response_logprobs=[-0.1, -0.2, -0.3, -0.4, -0.5], + prompt_token_ids=[0, 1, 2, 3, 4], + response_token_ids=[0, 1, 2, 3, 4], + user_prompt="Hello", + system_prompt="You are a helpful assistant.", + ) + ], + ) + + result = asyncio.run(engine.distill(payload)) + + assert isinstance(result, DistillResponse) + assert result.metadata["steps_per_batch_applied"] == 2 + assert result.metadata["step"] == 2 + assert result.metadata["divergence_mode"] == "topk_gjs" + + # Two optimizer steps + assert mock_training_client.forward_backward_async.call_count == 2 + assert mock_training_client.optim_step_async.call_count == 2 + + # Student recomputation used sample_async (not compute_logprobs_async) + mock_training_client.save_weights_and_get_sampling_client_async.assert_called_once() + student_sampler.sample_async.assert_called() + + # Step-2 datum should use different advantages (from student top-K GJS) + first_call = mock_training_client.forward_backward_async.call_args_list[0] + second_call = mock_training_client.forward_backward_async.call_args_list[1] + first_advantages = first_call.args[0][0].loss_fn_inputs["advantages"].data + second_advantages = second_call.args[0][0].loss_fn_inputs["advantages"].data + assert first_advantages != second_advantages + + +def test_engine_distill_scalar_mode_reports_divergence_mode(tinker_engine, mock_training_client): + """Scalar mode (default) reports divergence_mode='scalar_kl' in metadata.""" + engine, mock_service = tinker_engine + + set_tinker_path("test/lora", "tinker://ckpt", "gpt-oss/GPT-OSS-120B", 32, step=0) + mock_service.create_training_client_from_state_async = AsyncMock( + return_value=mock_training_client + ) + + teacher_sampler = MagicMock() + teacher_sampler.compute_logprobs_async = AsyncMock(return_value=[-0.2] * 100) + mock_service.create_sampling_client_async = AsyncMock(return_value=teacher_sampler) + + payload = DistillBatchRequestPayload( + lora_id="test/lora", + training=TrainingConfig(steps_per_batch=1), + samples=[ + DistillBatchItem( + prompt="Hello", + response="World", + feedback="Nice", + response_logprobs=[-0.1, -0.2, -0.3, -0.4, -0.5], + prompt_token_ids=[0, 1, 2, 3, 4], + response_token_ids=[0, 1, 2, 3, 4], + user_prompt="Hello", + system_prompt="You are a helpful assistant.", + ) + ], + ) + + result = asyncio.run(engine.distill(payload)) + assert result.metadata["divergence_mode"] == "scalar_kl" From db7e2e5f70b279e819d1acec9495961ab20144ec Mon Sep 17 00:00:00 2001 From: Kion Date: Fri, 6 Mar 2026 18:36:21 -0800 Subject: [PATCH 2/3] Fix ty type errors: list covariance and dict splat inference - Use explicit `list[BehaviorSignal](...)` cast for top-K behavior list assignment (ty enforces list invariance) - Replace `**core_kwargs` dict splat with explicit keyword args in dataclass constructors (ty can't narrow dict value union types) Co-Authored-By: Claude Opus 4.6 --- claas/training/engine/tinker/engine.py | 26 ++++++++++++++------------ 1 file changed, 14 insertions(+), 12 deletions(-) diff --git a/claas/training/engine/tinker/engine.py b/claas/training/engine/tinker/engine.py index e226c51..fc7d088 100644 --- a/claas/training/engine/tinker/engine.py +++ b/claas/training/engine/tinker/engine.py @@ -323,11 +323,12 @@ async def distill( if step_idx < steps_per_batch - 1: student_sampling = await training_client.save_weights_and_get_sampling_client_async() if use_topk: - behavior_signals = await _compute_student_topk_for_batch( + topk_behaviors = await _compute_student_topk_for_batch( student_sampling=student_sampling, prepared_samples=prepared_samples, top_k=teacher_top_k, ) + behavior_signals = list[BehaviorSignal](topk_behaviors) else: scalar_logprobs = await _compute_student_logprobs_for_batch( student_sampling=student_sampling, @@ -433,15 +434,6 @@ async def _prepare_sample_inputs( teacher_full = T.ModelInput.from_ints(teacher_full_tokens) teacher_scored_text = tokenizer.decode(teacher_full_tokens, skip_special_tokens=False) - core_kwargs = dict( - full_tokens=full_tokens, - input_tokens=input_tokens, - target_tokens=target_tokens, - prompt_len=prompt_len, - completion_len=completion_len, - teacher_scored_text=teacher_scored_text, - ) - initial_behavior = ScalarBehavior(logprobs=list(sample.response_logprobs)) if use_topk: @@ -460,7 +452,12 @@ async def _prepare_sample_inputs( ) prepared: PreparedSample = TopKPreparedSample( - **core_kwargs, + full_tokens=full_tokens, + input_tokens=input_tokens, + target_tokens=target_tokens, + prompt_len=prompt_len, + completion_len=completion_len, + teacher_scored_text=teacher_scored_text, teacher_logprobs=teacher_logprobs, teacher_topk=teacher_topk, ) @@ -474,7 +471,12 @@ async def _prepare_sample_inputs( ) prepared = ScalarPreparedSample( - **core_kwargs, + full_tokens=full_tokens, + input_tokens=input_tokens, + target_tokens=target_tokens, + prompt_len=prompt_len, + completion_len=completion_len, + teacher_scored_text=teacher_scored_text, teacher_logprobs=teacher_logprobs, ) From c7a8096fab553c355beb5a3266d016a6b616ece8 Mon Sep 17 00:00:00 2001 From: Kion Date: Fri, 6 Mar 2026 18:39:07 -0800 Subject: [PATCH 3/3] Move engine types to types.py, default use_topk_divergence=True MIME-Version: 1.0 Content-Type: text/plain; charset=UTF-8 Content-Transfer-Encoding: 8bit - Extract SampleCore, ScalarPreparedSample, TopKPreparedSample, ScalarBehavior, TopKBehavior into claas/training/engine/tinker/types.py - Change use_topk_divergence default from False to True in TrainingConfig and base.yaml — top-K GJS is now the default divergence mode Co-Authored-By: Claude Opus 4.6 --- claas/core/types.py | 2 +- claas/eval/configs/base.yaml | 2 +- claas/training/engine/tinker/engine.py | 63 ++++---------------------- claas/training/engine/tinker/types.py | 55 ++++++++++++++++++++++ 4 files changed, 65 insertions(+), 57 deletions(-) create mode 100644 claas/training/engine/tinker/types.py diff --git a/claas/core/types.py b/claas/core/types.py index 07ab79a..08a9873 100644 --- a/claas/core/types.py +++ b/claas/core/types.py @@ -32,7 +32,7 @@ class TrainingConfig: max_grad_norm: float = 1.0 kl_reg_weight: float = 0.0 teacher_top_k: int = 100 - use_topk_divergence: bool = False + use_topk_divergence: bool = True steps_per_batch: int = 4 feedback_repetitions: int = 1 diff --git a/claas/eval/configs/base.yaml b/claas/eval/configs/base.yaml index d5f213a..5af11c3 100644 --- a/claas/eval/configs/base.yaml +++ b/claas/eval/configs/base.yaml @@ -35,6 +35,6 @@ training: max_grad_norm: 1.0 kl_reg_weight: 0.0 teacher_top_k: 100 - use_topk_divergence: false + use_topk_divergence: true steps_per_batch: 4 feedback_repetitions: 1 diff --git a/claas/training/engine/tinker/engine.py b/claas/training/engine/tinker/engine.py index fc7d088..5c71092 100644 --- a/claas/training/engine/tinker/engine.py +++ b/claas/training/engine/tinker/engine.py @@ -13,7 +13,6 @@ import asyncio import logging import os -from dataclasses import dataclass from datetime import datetime, timezone from typing import Any from urllib.parse import quote @@ -47,8 +46,15 @@ lora_exists as state_lora_exists, set_tinker_path, ) +from claas.training.engine.tinker.types import ( + BehaviorSignal, + PreparedSample, + ScalarBehavior, + ScalarPreparedSample, + TopKBehavior, + TopKPreparedSample, +) from claas.training.gjs import ( - SequenceTopK, compute_topk_gjs, extract_token_logprobs, slice_completion_topk, @@ -62,59 +68,6 @@ _MAX_KL_GAIN = 4.0 -# ------------------------------------------------------------------ -# Discriminated-union sample/behavior types -# ------------------------------------------------------------------ - - -@dataclass(frozen=True) -class SampleCore: - """Fields shared across both scalar and top-K prepared samples.""" - - full_tokens: list[int] - input_tokens: list[int] - target_tokens: list[int] - prompt_len: int - completion_len: int - teacher_scored_text: str - - -@dataclass(frozen=True) -class ScalarPreparedSample(SampleCore): - """Sample with scalar teacher logprobs (current approach).""" - - teacher_logprobs: list[float] - - -@dataclass(frozen=True) -class TopKPreparedSample(SampleCore): - """Sample with top-K teacher distributions (GJS approach).""" - - teacher_logprobs: list[float] - teacher_topk: SequenceTopK - - -PreparedSample = ScalarPreparedSample | TopKPreparedSample - - -@dataclass(frozen=True) -class ScalarBehavior: - """Scalar student logprobs for IS correction.""" - - logprobs: list[float] - - -@dataclass(frozen=True) -class TopKBehavior: - """Top-K student distributions + scalar logprobs for IS correction.""" - - logprobs: list[float] - topk: SequenceTopK - - -BehaviorSignal = ScalarBehavior | TopKBehavior - - class TinkerTrainingEngine(TrainingEngine): """Executes training and LoRA management through the Tinker Python SDK.""" diff --git a/claas/training/engine/tinker/types.py b/claas/training/engine/tinker/types.py new file mode 100644 index 0000000..ff24f88 --- /dev/null +++ b/claas/training/engine/tinker/types.py @@ -0,0 +1,55 @@ +"""Tinker engine data types for SDPO distillation.""" + +from __future__ import annotations + +from dataclasses import dataclass + +from claas.training.gjs import SequenceTopK + + +@dataclass(frozen=True) +class SampleCore: + """Fields shared across both scalar and top-K prepared samples.""" + + full_tokens: list[int] + input_tokens: list[int] + target_tokens: list[int] + prompt_len: int + completion_len: int + teacher_scored_text: str + + +@dataclass(frozen=True) +class ScalarPreparedSample(SampleCore): + """Sample with scalar teacher logprobs (current approach).""" + + teacher_logprobs: list[float] + + +@dataclass(frozen=True) +class TopKPreparedSample(SampleCore): + """Sample with top-K teacher distributions (GJS approach).""" + + teacher_logprobs: list[float] + teacher_topk: SequenceTopK + + +PreparedSample = ScalarPreparedSample | TopKPreparedSample + + +@dataclass(frozen=True) +class ScalarBehavior: + """Scalar student logprobs for IS correction.""" + + logprobs: list[float] + + +@dataclass(frozen=True) +class TopKBehavior: + """Top-K student distributions + scalar logprobs for IS correction.""" + + logprobs: list[float] + topk: SequenceTopK + + +BehaviorSignal = ScalarBehavior | TopKBehavior