DGXAI/driftcall
0
1"""Custom trainer + dataset adapter (docs/modules/training.md §2.2, §3.2.3).2 3Two public types:4 5- :class:`EpisodeDatasetAdapter` — stateless iterable feeding6 ``GRPOTrainer.train_dataset``. Each ``__iter__`` tick yields7 ``{"prompt": str, "_meta": {...}}`` where ``_meta`` carries the8 ``GoalSpec``, the monotonically-derived ``episode_seed``, the curriculum9 ``stage``, and the ``language_weights``. One call to10 ``task_generator.generate`` per step; one call to11 ``tokenizer.apply_chat_template(messages, tokenize=False,12 add_generation_prompt=True)`` to render the prompt.13 14- :class:`DriftCallGRPOTrainer` — ``GRPOTrainer`` subclass whose15 ``_generate_and_score_completions`` override runs G multi-turn episodes16 via a caller-provided ``RolloutGroupFn`` and plumbs the resulting17 frozen ``Episode`` tuple into ``reward_fn`` (step_13) before handing the18 G reward scalars + padded completions back to the inherited GRPO19 advantage / KL / optimizer step path. **The inherited code path is20 untouched** (training.md §3.2.3).21 22``trl`` and ``torch`` are imported lazily. Pure-Python fallbacks for23``_generate_and_score_completions`` are provided so the class shape24can be verified on CPU-only CI.25"""26 27from __future__ import annotations28 29import math30from dataclasses import dataclass31from typing import TYPE_CHECKING, Any, Literal, Protocol, cast32 33if TYPE_CHECKING: # pragma: no cover - typing only34 from collections.abc import Callable, Iterator35 36from cells.step_13_grpo_config import BETA_KL37 38PINNED_SYSTEM_PROMPT: str = (39 "You are a concierge assistant. Use the provided tools. "40 "Respond in the caller's language. Submit with calibrated confidence."41)42 43LanguageCode = Literal["hi", "ta", "kn", "en", "hinglish"]44 45 46class EpisodeSampler(Protocol):47 """Draws a ``GoalSpec`` for one prompt slot (training.md §2.2)."""48 49 def __call__(self, step: int) -> Any: ...50 51 52class EnvFactory(Protocol):53 """Returns a fresh ``DriftCallEnv`` per rollout (training.md §3.2)."""54 55 def __call__(self) -> Any: ...56 57 58class RolloutGroupFn(Protocol):59 """Runs G multi-turn rollouts sharing one goal.60 61 Returns a tuple ``(episodes, completions)`` of length G each.62 """63 64 def __call__(65 self,66 *,67 model: Any,68 tokenizer: Any,69 goal: Any,70 episode_seed: int,71 num_generations: int,72 env_factory: EnvFactory,73 ) -> tuple[tuple[Any, ...], tuple[str, ...]]: ...74 75 76@dataclass(frozen=True)77class AdapterRecord:78 """Frozen view of one :class:`EpisodeDatasetAdapter` yield.79 80 Tests consume this view rather than dict-typing ``_meta`` inline.81 """82 83 prompt: str84 goal: Any85 episode_seed: int86 stage: Literal[1, 2, 3]87 language_weights: dict[LanguageCode, float]88 89 90def render_initial_prompt(tokenizer: Any, goal: Any) -> str:91 """Render the turn-0 chat template (training.md §3.2.1).92 93 Messages: pinned system prompt + ``goal.seed_utterance`` as the user94 turn. ``add_generation_prompt=True`` tells Gemma to emit an assistant95 turn. Tool schemas live in later turns so only these two messages96 appear at ``step == 0``.97 """98 seed_utterance = getattr(goal, "seed_utterance", "")99 messages: list[dict[str, str]] = [100 {"role": "system", "content": PINNED_SYSTEM_PROMPT},101 {"role": "user", "content": seed_utterance},102 ]103 result = tokenizer.apply_chat_template(104 messages,105 tokenize=False,106 add_generation_prompt=True,107 )108 return str(result)109 110 111class EpisodeDatasetAdapter:112 """Stateless streaming dataset (training.md §2.2).113 114 Constructor signature matches training.md §2.2: a ``task_gen`` callable115 accepting ``(seed, stage, language_weights)``, an ``env_factory``116 producing fresh envs, the curriculum ``stage``, a ``stage_base_seed``117 used to derive per-step ``episode_seed``, the per-language sampling118 ``language_weights``, and the ``tokenizer`` used to render prompts.119 120 Iteration is infinite — exactly one record per GRPO training step.121 Step counter is local to ``__iter__`` so resume simply restarts from122 whatever step TRL's ``resume_from_checkpoint`` restores.123 """124 125 def __init__(126 self,127 *,128 task_gen: Callable[..., Any],129 env_factory: EnvFactory,130 stage: Literal[1, 2, 3],131 stage_base_seed: int,132 language_weights: dict[LanguageCode, float],133 tokenizer: Any,134 ) -> None:135 self.task_gen = task_gen136 self.env_factory = env_factory137 self.stage: Literal[1, 2, 3] = stage138 self.stage_base_seed = stage_base_seed139 self.language_weights = dict(language_weights)140 self.tokenizer = tokenizer141 142 def _build_record(self, step: int) -> dict[str, Any]:143 episode_seed = self.stage_base_seed + step144 goal = self.task_gen(145 seed=episode_seed,146 stage=self.stage,147 language_weights=self.language_weights,148 )149 prompt = render_initial_prompt(self.tokenizer, goal)150 return {151 "prompt": prompt,152 "_meta": {153 "goal": goal,154 "episode_seed": episode_seed,155 "stage": self.stage,156 "language_weights": dict(self.language_weights),157 },158 }159 160 def __iter__(self) -> Iterator[dict[str, Any]]:161 step = 0162 while True:163 yield self._build_record(step)164 step += 1165 166 def __len__(self) -> int:167 """Length sentinel for TRL 0.24+ ``RepeatSampler``.168 169 The dataset is logically infinite (one record per GRPO step), but170 TRL 0.24's ``RepeatSampler`` calls ``len(data_source)`` to size the171 sampler. Returning a large finite number lets training proceed; the172 actual step count is bounded by ``GRPOConfig.max_steps``.173 """174 return 1_000_000175 176 def __getitem__(self, idx: int) -> dict[str, Any]:177 """Map-style indexing for TRL 0.24+ DataLoader.178 179 TRL 0.24 treats the train_dataset as a Map-style dataset and looks180 records up by integer index. We honour the contract by deriving the181 record purely from ``idx`` — the adapter is stateless so any index182 produces a deterministic ``(prompt, _meta)`` pair for that step.183 """184 return self._build_record(int(idx))185 186 def peek(self, step: int) -> AdapterRecord:187 """Materialize the record at ``step`` without advancing iteration.188 189 Used by tests (§1.2 U14–U18) to assert record shape at arbitrary190 steps without consuming a generator.191 """192 rec = self._build_record(step)193 meta = rec["_meta"]194 return AdapterRecord(195 prompt=rec["prompt"],196 goal=meta["goal"],197 episode_seed=meta["episode_seed"],198 stage=meta["stage"],199 language_weights=meta["language_weights"],200 )201 202 203def _import_grpo_trainer() -> type[Any]:204 """Lazy import of ``trl.GRPOTrainer``; isolated for mocking in tests."""205 from trl import GRPOTrainer206 207 return cast("type[Any]", GRPOTrainer)208 209 210def _make_driftcall_init(211 base_cls: type[Any],212) -> Callable[..., None]:213 """Build an ``__init__`` bound to ``base_cls``; avoids super() recursion214 when the returned class is itself further subclassed.215 216 DriftCall-specific kwargs added on top of ``GRPOTrainer.__init__``:217 218 - ``rollout_group_fn``, ``env_factory``, ``reward_fn_driftcall`` — the219 multi-turn rollout override surface (see class docstring).220 - ``enable_adaptive_kl`` (default ``True``) — auto-attach an221 :class:`AdaptiveKLCallback` so β retargets to the measured KL each222 logging tick (training.md §3.3.1). Set ``False`` to disable.223 - ``adaptive_kl_target`` — override the default ``target_kl=BETA_KL``.224 - ``adaptive_kl_kp`` — override the proportional gain.225 - ``adaptive_kl_beta_min`` / ``adaptive_kl_beta_max`` — override clamp226 bounds.227 """228 229 def _init(230 self: Any,231 *args: Any,232 rollout_group_fn: RolloutGroupFn,233 env_factory: EnvFactory,234 reward_fn_driftcall: Callable[..., list[float]],235 enable_adaptive_kl: bool = True,236 adaptive_kl_target: float | None = None,237 adaptive_kl_kp: float = DEFAULT_KP,238 adaptive_kl_beta_min: float = DEFAULT_BETA_MIN,239 adaptive_kl_beta_max: float = DEFAULT_BETA_MAX,240 **kwargs: Any,241 ) -> None:242 # TRL 0.24 made ``reward_funcs`` a required arg on GRPOTrainer.243 # Our custom ``_generate_and_score_completions`` short-circuits the244 # base reward path entirely (calls ``reward_fn_driftcall`` directly),245 # so the parent's ``reward_funcs`` value is never invoked. Pass a246 # placeholder identity reward to satisfy the signature on TRL>=0.24.247 if "reward_funcs" not in kwargs:248 def _placeholder_reward(249 completions: Any = None,250 **_unused: Any,251 ) -> list[float]:252 n = len(completions) if completions is not None else 0253 return [0.0] * n254 255 kwargs["reward_funcs"] = [_placeholder_reward]256 base_cls.__init__(self, *args, **kwargs)257 self.rollout_group_fn = rollout_group_fn258 self.env_factory = env_factory259 self.reward_fn_driftcall = reward_fn_driftcall260 261 if enable_adaptive_kl:262 target = (263 adaptive_kl_target if adaptive_kl_target is not None else BETA_KL264 )265 callback = AdaptiveKLCallback(266 target_kl=target,267 kp=adaptive_kl_kp,268 beta_min=adaptive_kl_beta_min,269 beta_max=adaptive_kl_beta_max,270 )271 self.adaptive_kl_callback = callback272 add_callback = getattr(base_cls, "add_callback", None)273 if callable(add_callback):274 # Production path (TRL ≥ 0.23): register through the TRL275 # callback handler so ``on_log`` fires alongside default276 # loggers with the correct ``args``/``state``/``control``.277 self.add_callback(callback)278 else:279 # Fallback: minimal bases in tests lack ``add_callback``.280 # Keep a private list so callers can still invoke the hook.281 if not hasattr(self, "_driftcall_callbacks"):282 self._driftcall_callbacks = []283 self._driftcall_callbacks.append(callback)284 else:285 self.adaptive_kl_callback = None286 287 return _init288 289 290def _driftcall_generate_and_score_completions(291 self: Any, inputs: list[dict[str, Any]]292) -> dict[str, Any]:293 """Run the multi-turn rollout, then call ``reward_fn``.294 295 Expects ``inputs`` to carry one row per prompt slot with the296 ``_meta`` dict produced by :class:`EpisodeDatasetAdapter`.297 Returns a dict with keys ``episodes``, ``completions``, ``rewards``,298 ``prompts`` — each length G (num_generations).299 """300 if not inputs:301 raise ValueError("inputs must be a non-empty list")302 303 row = inputs[0]304 meta = row["_meta"]305 prompt = row["prompt"]306 goal = meta["goal"]307 episode_seed = meta["episode_seed"]308 309 num_generations = int(getattr(self.args, "num_generations", 8))310 episodes, completions = self.rollout_group_fn(311 model=self.model,312 tokenizer=self.processing_class,313 goal=goal,314 episode_seed=episode_seed,315 num_generations=num_generations,316 env_factory=self.env_factory,317 )318 319 if len(episodes) != num_generations or len(completions) != num_generations:320 raise ValueError(321 f"rollout_group_fn produced {len(episodes)} episodes and "322 f"{len(completions)} completions; expected {num_generations} each"323 )324 325 prompts = [prompt] * num_generations326 metas = [dict(meta) for _ in range(num_generations)]327 rewards = self.reward_fn_driftcall(328 prompts=prompts,329 completions=list(completions),330 _meta=metas,331 episodes=list(episodes),332 )333 334 return {335 "episodes": episodes,336 "completions": completions,337 "rewards": rewards,338 "prompts": prompts,339 }340 341 342def make_driftcall_grpo_trainer_cls(base_cls: type[Any] | None = None) -> type[Any]:343 """Build the :class:`DriftCallGRPOTrainer` class bound to ``base_cls``.344 345 Default ``base_cls`` is ``trl.GRPOTrainer`` (imported lazily). Tests346 pass a stub base class so they can exercise the override path without347 TRL installed.348 349 GRPOTrainer subclass with multi-turn rollout override350 (training.md §3.2.3). Construction adds three DriftCall-specific351 kwargs over the standard ``GRPOTrainer.__init__``:352 353 - ``rollout_group_fn``: :class:`RolloutGroupFn` running G multi-turn354 episodes and returning ``(episodes, completions)``.355 - ``env_factory``: :class:`EnvFactory` producing a fresh356 ``DriftCallEnv`` per rollout.357 - ``reward_fn_driftcall``: the step_13 ``reward_fn`` — called358 directly with the frozen ``Episode`` tuple after rollout.359 360 ``_generate_and_score_completions`` replaces the TRL default.361 Advantage + KL + optimizer step paths are inherited unchanged.362 """363 resolved_base: type[Any] = (364 base_cls if base_cls is not None else _import_grpo_trainer()365 )366 return type(367 "DriftCallGRPOTrainer",368 (resolved_base,),369 {370 "__init__": _make_driftcall_init(resolved_base),371 "_generate_and_score_completions": _driftcall_generate_and_score_completions,372 "__doc__": "GRPOTrainer subclass with multi-turn rollout override.",373 },374 )375 376 377def driftcall_grpo_trainer_methods() -> tuple[str, ...]:378 """Return the method names the subclass overrides (introspection helper).379 380 Used by the shape test (U in §1.x) to verify the override surface.381 """382 return ("__init__", "_generate_and_score_completions")383 384 385# ---------------------------------------------------------------------------386# Adaptive KL controller (training.md §3.3 — retarget β from measured KL)387# ---------------------------------------------------------------------------388 389 390DEFAULT_BETA_MIN: float = 0.001391DEFAULT_BETA_MAX: float = 1.0392DEFAULT_KP: float = 2.0393 394 395def _trainer_callback_base() -> type:396 """Return ``transformers.TrainerCallback`` if importable, else ``object``.397 398 Importing transformers lazily keeps step_14 importable on CPU-only CI399 runners that don't have transformers installed.400 """401 try:402 from transformers.trainer_callback import TrainerCallback403 return TrainerCallback404 except Exception:405 return object406 407 408class AdaptiveKLCallback(_trainer_callback_base()): # type: ignore[misc]409 """Retarget β each step based on the ratio of measured KL to ``target_kl``.410 411 Proportional controller with symmetric log-space update:412 413 err = (kl - target_kl) / target_kl414 new_beta = beta * exp(kp * err)415 new_beta = clamp(new_beta, beta_min, beta_max)416 417 When ``kl`` matches ``target_kl``, ``err == 0`` and β is left unchanged.418 Safe on missing / NaN / non-numeric KL signals (no-op, no exception).419 420 Inherits from :class:`transformers.trainer_callback.TrainerCallback` when421 available (production path) so all the no-op callback events422 (``on_train_begin``, ``on_step_begin``, etc.) come for free; falls back423 to ``object`` on CPU-only CI when transformers is not installed.424 """425 426 def __init__(427 self,428 target_kl: float = BETA_KL,429 *,430 kp: float = DEFAULT_KP,431 beta_min: float = DEFAULT_BETA_MIN,432 beta_max: float = DEFAULT_BETA_MAX,433 ) -> None:434 if target_kl <= 0.0:435 raise ValueError(f"target_kl must be > 0; got {target_kl}")436 if beta_min <= 0.0 or beta_max <= 0.0:437 raise ValueError(438 f"beta bounds must be > 0; got min={beta_min}, max={beta_max}"439 )440 if beta_min > beta_max:441 raise ValueError(442 f"beta_min ({beta_min}) must be <= beta_max ({beta_max})"443 )444 self.target_kl = float(target_kl)445 self.kp = float(kp)446 self.beta_min = float(beta_min)447 self.beta_max = float(beta_max)448 449 def _coerce_kl(self, raw: Any) -> float | None:450 """Return a finite float or ``None`` — propagates no-op on bad input."""451 try:452 value = float(raw)453 except (TypeError, ValueError):454 return None455 if math.isnan(value) or math.isinf(value):456 return None457 return value458 459 def _next_beta(self, beta: float, kl: float) -> tuple[float, bool, bool]:460 """Return ``(new_beta, clamped_to_min, clamped_to_max)``."""461 err = (kl - self.target_kl) / self.target_kl462 # Clamp the exponent so extreme KL spikes don't overflow math.exp;463 # the result is clamped anyway and exp(±50) easily saturates either bound.464 exponent = max(-50.0, min(50.0, self.kp * err))465 scaled = beta * math.exp(exponent)466 if scaled <= self.beta_min:467 return self.beta_min, True, False468 if scaled >= self.beta_max:469 return self.beta_max, False, True470 return scaled, False, False471 472 def on_log(473 self,474 args: Any,475 state: Any,476 control: Any,477 *,478 logs: dict[str, Any] | None = None,479 **_kwargs: Any,480 ) -> Any:481 """TRL hook — called with every ``trainer.log(...)`` dict.482 483 On a well-formed KL signal: mutates ``args.beta`` with the new484 coefficient and writes five diagnostic fields back into ``logs``485 so TRL's default reporter forwards them to wandb / CSV / etc.:486 487 - ``train/beta_adaptive`` current KL coefficient488 - ``train/kl_measured`` sanitised KL input489 - ``train/kl_target`` constant — aids chart-by-reference490 - ``train/beta_clamped_to_min`` 0/1 — fires on collapse491 - ``train/beta_clamped_to_max`` 0/1 — fires on runaway divergence492 """493 if logs is None:494 return control495 if "kl" not in logs:496 return control497 kl = self._coerce_kl(logs["kl"])498 if kl is None:499 return control500 beta = float(getattr(args, "beta", BETA_KL))501 new_beta, clamped_lo, clamped_hi = self._next_beta(beta, kl)502 args.beta = new_beta503 logs["train/beta_adaptive"] = new_beta504 logs["train/kl_measured"] = kl505 logs["train/kl_target"] = self.target_kl506 logs["train/beta_clamped_to_min"] = 1 if clamped_lo else 0507 logs["train/beta_clamped_to_max"] = 1 if clamped_hi else 0508 return control509 510 511__all__ = [512 "AdapterRecord",513 "AdaptiveKLCallback",514 "DEFAULT_BETA_MAX",515 "DEFAULT_BETA_MIN",516 "DEFAULT_KP",517 "EnvFactory",518 "EpisodeDatasetAdapter",519 "EpisodeSampler",520 "LanguageCode",521 "PINNED_SYSTEM_PROMPT",522 "RolloutGroupFn",523 "driftcall_grpo_trainer_methods",524 "make_driftcall_grpo_trainer_cls",525 "render_initial_prompt",526]527