# Tinker Cookbook Agent Guide

Working notes for future agents hacking on `tinker-cookbook`. Additional docs can be found in the `llms.txt` (condensed) / `llms-full.txt` (complete), `CONTRIBUTING`, and the bundled documentation.

## Mission & Scope
- `tinker-cookbook` is the client-side layer for the hosted **Tinker** service. You author training/eval loops that run on a CPU machine; Tinker executes the heavy GPU work (LoRA fine-tuning, sampling, checkpointing) on synchronized worker pools (a.k.a. clock cycles).
- The cookbook must mirror the public docs. Both `llms.txt` and `llms-full.txt` are autogenerated outside this repo—treat them as read-only and coordinate with maintainers when they need a refresh.
- Primary users: (1) researchers cloning recipes and swapping in their data/envs; (2) SDK developers extending abstractions like renderers, datasets, evaluators, completers.

## Tooling & Setup
- Python ≥3.11. Follow the onboarding instructions: join the waitlist, create a `TINKER_API_KEY` in the console, `pip install tinker`, then `pip install -e .[dev]` (or `uv pip install -e .[dev]`). Most contributors already have the env variable set; if requests fail with auth errors, re-export it.
- Optional extras (`vector-search`, `wandb`, `verifiers`, etc.) are defined in `pyproject.toml`.
- CLI utilities expect datasets, logs, and checkpoints to live under user-controlled paths (default `/tmp/tinker-examples/...`). Clean up disk usage between runs.
- Heavy examples (smoke tests, RL recipes) download Hugging Face datasets and call the hosted API; run them only when you have network+API access.

## Architecture & Patterns
- **Builder pattern (per CONTRIBUTING):**
  - Config objects are lightweight `chz` dataclasses (e.g., `SupervisedDatasetBuilder`, `RLDatasetBuilder`, `EnvGroupBuilder`, `EvaluatorBuilder`). They capture parameters, stay serializable, and usually expose a `.build()`/`__call__()` that returns heavyweight runtime objects.
  - Launch scripts define a CLI-facing `CLIConfig` (parsed by `chz`) that instantiates the richer training `Config`. This gives every recipe a consistent `python -m ... key=value` interface.
  - Env builders compose like `RLDatasetBuilder → EnvGroupBuilder → Env`. Groups let us share metadata (tags, pairwise comparisons) and center rewards across related rollouts.
- **Completers:** algorithms interact with the `TokenCompleter` interface. `TinkerTokenCompleter` (wrapping a `SamplingClient`) is the default implementation, but evaluators may accept any `TokenCompleter` or `MessageCompleter`.
- **Renderers & tokenizer utils:** pick the renderer that matches your tokenizer/model pair (e.g., `role_colon`, `llama3`, `qwen3`). `TrainOnWhat` controls which tokens get weight=1 in SFT. Tokenizers are cached via `tokenizer_utils.get_tokenizer`, with Llama-3 names remapped to `baseten/Meta-Llama-3-tokenizer` to bypass HF gating.
- **Loss plumbing:** every `tinker.Datum` bundles a `model_input` plus `loss_fn_inputs` (`TensorData`). Use helpers such as `conversation_to_datum`, `datum_from_tokens_weights`, and `_remove_mask` instead of constructing dicts manually. Built-in losses: `cross_entropy`, `importance_sampling`, `ppo`; `forward_backward_custom` covers bespoke differentiable objectives.

## Conventions & Notation (from CONTRIBUTING)
- **Subscripts:** `_P` (problems/prompts), `_G` (groups of rollouts sharing metadata), `_T` (tokens/time), `_D` (datums), with flattened forms like `_PG`. Example: `tokens_P_G_T[p][g][t]` indexes tokens for problem `p`, group member `g`, token `t`. Keep these suffixes when naming tensors/metrics so downstream tooling can interpret shapes.
- **Env lifecycle:** `Env` objects are single-use (no `reset`); create them via `EnvGroupBuilder`, which returns correlated envs (for GRPO-style centering or multi-agent comparisons). Datasets return groups, not individual envs.
- **Typing:** prefer explicit typing, avoid `Any` / `type: ignore`. Keep generics readable. Converters like `TensorData.from_numpy` and helper casting utilities already exist; use them.
- **`chz` usage:** configuration objects (`Config`, dataset builders, CLI configs) are `@chz.chz` classes so they can be serialized, logged, and hydrated from CLI key-value pairs.
- **Logging style:** training scripts rely on `ml_log` for metrics (`metrics.jsonl`, optional W&B) and `logtree` for HTML transcripts. When adding new metrics, follow the `ml_log.log_metrics` shape conventions (`str → float/int/str`).
- **Safe iteration:** functions like `safezip`, `timed`, and `scope` (tracing) are widely used; follow those patterns instead of hand-writing logging/zip logic.

## Data & Rendering
- Rendering is the bridge between chat-style data and token sequences. `renderers.py` defines `Renderer.build_supervised_example`, `build_generation_prompt`, `get_stop_sequences`, and `parse_response`. Use `TrainOnWhat` to switch between “last assistant only” vs “all assistant messages” vs “prompt distillation” setups.
- For supervised chat datasets, reuse `SupervisedDatasetFromHFDataset`, `StreamingSupervisedDatasetFromHFDataset`, or `FromConversationFileBuilder`. They expect HF rows with `messages` arrays; map them through a renderer and optional max length.
- RL data is organized by dimensions `_P` (problems), `_G` (group members / rollouts per problem), `_T` (tokens), `_D` (datums). Keep arrays ragged-aware, and document shape suffixes when introducing new tensors.

## Training Playbooks
### Supervised Learning
- **Main loop:** `tinker_cookbook/supervised/train.py`. It pipelines batches by submitting `forward_backward_async` and `optim_step_async` immediately, storing futures inside `SubmittedBatch`. Metrics/logging run through `ml_log`, stdout previews via `display.colorize_example`.
- **Configs:** include LR schedule (`linear` multiplier via `compute_schedule_lr_multiplier`), LoRA rank, checkpoint cadence (`save_every`), eval cadence (`eval_every`, `infrequent_eval_every`), and dataset builders.
- **Hyperparameters:** Call `hyperparam_utils.get_lr(model_name)`; LR is independent of LoRA rank.
- **Prompt distillation:** see `tinker_cookbook/recipes/prompt_distillation`. Renderers assign weight=0 to context instructions and weight=1 to distilled responses.
- **Sweeps:** `tinker_cookbook/recipes/sl_loop.py` doubles as a sweep harness (log-scale grid, aggregate from `metrics.jsonl`). Keep these scripts runnable; they double as docs tests.

### Reinforcement Learning
- **Main loop:** `tinker_cookbook/rl/train.py`. Steps: build dataset (`RLDatasetBuilder`), get groups of envs (`EnvGroupBuilder`), collect rollouts (`do_group_rollout`), compute advantages (`compute_advantages`), assemble datums (`assemble_training_data`), run `forward_backward_async(..., loss_fn="importance_sampling" | "ppo")`, apply `optim_step_async`.
- **Policies:** implement the `TokenCompleter` interface. Training loops usually instantiate `TinkerTokenCompleter`, but tests may stub a completer.
- **Hyperparameters:** key knobs are `batch_size` vs `group_size`, `num_substeps` (similar to PPO epochs but still single-pass), and advanced configs:
  - `StreamMinibatchConfig` overlaps sampling with training (still on-policy).
  - `AsyncConfig` enables bounded off-policy lag (“off-by-K”). Monitor KL metrics (`compute_kl_sample_train`, `compute_post_kl`) plus reward trends to make sure drift stays manageable.
- **Environments:** `Env`, `EnvGroupBuilder`, and `RLDataset` live in `tinker_cookbook/rl/types.py`. Groups make it easy to compute pairwise rewards (preference models) or multi-agent games. Example: `recipes/multiplayer_rl/twenty_questions`.
- **Recipes:** `rl_basic.py` demonstrates default metrics: reward, entropy, `ac_tokens_per_turn`, format rate, KL approximations, and progress/time tokens.

### Preferences & Distillation
- **DPO:** `tinker_cookbook/preference/train_dpo.py` (CLI in `recipes/preference/train.py`). Important knobs: dataset builder (choose whichever comparison corpus you need), `renderer_name`, `dpo_beta`, LR (often 1e-5 to 1e-6). Metrics like `dpo_loss`, `accuracy`, `margin`, `chosen/rejected_reward` come from the implicit reward model.
- **RLHF pipeline:** `recipes/preference/rlhf/rlhf_pipeline.py` walks through the standard three stages (supervised warm-start, preference model, RL self-play using pairwise comparisons).
- **Distillation:** `distillation/train_on_policy.py` handles on-policy or SFT-style distillation; combine with `renderers`, `hyperparam_utils`, and `sampling_client` utilities.

### Evaluations & Sampling
- Inline evaluators implement either `TrainingClientEvaluator` or `SamplingClientEvaluator`. Training loops accept builder lists (`evaluator_builders`, `infrequent_evaluator_builders`). Inspect AI integration is in `eval/inspect_evaluators.py` and `eval/run_inspect_evals.py`.
- Sampling clients come from `training_client.save_weights_and_get_sampling_client(name=...)`. To export weights, use `RestClient.download_checkpoint_archive_from_tinker_path`.

## Async & Performance
- Worker pools advance in ~10s clock cycles. Submit `forward_backward_async` and `optim_step_async` back-to-back, then await both futures to keep them on the same cycle.
- Pipeline batches: enqueue the next `forward_backward_async` before awaiting the previous batch’s results so there’s always work when a clock cycle begins.
- Use async everywhere performance matters (RL loops, production SFT). The synchronous helpers exist only for pedagogy (e.g., `recipes/sl_loop.py`) and small tests.

- Training CLIs (`recipes/*/*.py`) call `cli_utils.check_log_dir` at startup to decide whether to delete, resume, or prompt about an existing `log_path`. This convention is specific to training/eval entry points and is wired through `chz` CLI configs so users can choose behaviors (`delete`, `resume`, `ask`, `raise`).
- `ml_log` handles structured logging: metrics stream to stdout, `metrics.jsonl`, and optionally Weights & Biases (`wandb_project`, `wandb_name`). Use `logtree` scopes for HTML transcripts when you need qualitative review of rollouts.
- `checkpoint_utils.save_checkpoint_async` writes `{log_path}/checkpoints.jsonl` entries for state and/or sampler checkpoints. `get_last_checkpoint` filters by key (`state_path`, `sampler_path`) before resuming.
- After each optimizer step sequence, call `save_weights_for_sampler[_async]` (or `save_weights_and_get_sampling_client`) and then create a **new** sampling client. Existing `SamplingClient`s do not automatically pick up fresh weights, so evaluators must use the newly returned client handle.

## Testing & Troubleshooting
- Lightweight checks: `pytest tinker_cookbook/tests/test_renderers.py`, `pytest tinker_cookbook/tests/test_utils.py`. `tests/smoke_tests.py` spins up real training runs (needs HF + API access).
- Example data lives in `example-data/` (e.g., `conversations.jsonl`, `multilingual.txt`) and mirrors the formats documented in `training-sampling`.
- If you hit auth/network issues, double-check `TINKER_API_KEY`, ensure your environment can reach the Tinker service, and verify dependencies (`pip show tinker`).
- Resize datasets/batch sizes in recipes when debugging; `dataset_builder` objects usually accept `n_batches`, `batch_size`, and `group_size` fields so you can shrink workloads.

## Common Pitfalls
- **LoRA LR mismatch:** LoRA typically needs learning rates tens of times higher than full fine-tuning. Use `hyperparam_utils.get_lr` or the LR formula above. Rank does not change the optimal LR.
- **Renderer/tokenizer mismatch:** The renderer determines BOS/EOS tokens and stop sequences. Pair `renderer_name` with the tokenizer family your model expects (`llama3`, `qwen3`, `role_colon`, etc.). Otherwise loss weights and sampling stops will be wrong.
- **Loss inputs wrong shape:** Stick to helper functions so `loss_fn_inputs["weights"]`, `["target_tokens"]`, `["advantages"]`, etc., end up as `TensorData` with the right dtype. Custom DPO/RL objectives often fail here.
- **Async gaps:** Awaiting `forward_backward` before submitting `optim_step` wastes two extra clock cycles. Submit both first, then await results.
- **Sampler desync:** Saving weights isn’t enough; always request a new sampling client (e.g., via `save_weights_and_get_sampling_client`) before running evals so the client reflects the latest checkpoint.
- **Group semantics:** RL advantages are centered within each group.
- **DPO beta / LR:** Too large a beta or LR makes the policy collapse; start with `dpo_beta=0.1`, LR≈1e-5, and watch `accuracy` + `margin` trends.

## Quick Reference Commands
1. **Environment setup:**
   ```bash
   python -m venv .venv
   source .venv/bin/activate
   pip install tinker
   pip install -e .[dev]
   # Set once per shell if needed
   export TINKER_API_KEY=sk-...
   ```
2. **Basic SFT run (default recipe):**
   ```bash
   python -m tinker_cookbook.recipes.sl_basic \
     model_name=meta-llama/Llama-3.2-1B \
     log_path=/tmp/tinker-examples/sl_basic
   ```
3. **Custom JSONL SFT (bring your own conversations file):**
   ```bash
   python -m tinker_cookbook.recipes.sl_basic \
     dataset_path=/path/to/conversations.jsonl \
     renderer_name=role_colon \
     train_on_what=all_assistant_messages \
     log_path=/tmp/tinker-examples/sl_jsonl
   ```
4. **RL basic run (default reward):**
   ```bash
   python -m tinker_cookbook.recipes.rl_basic \
     model_name=meta-llama/Llama-3.1-8B \
     log_path=/tmp/tinker-examples/rl_basic
   ```
5. **DPO training (generic preference dataset):**
   ```bash
   python -m tinker_cookbook.recipes.preference.train \
     log_path=/tmp/dpo-run \
     model_name=meta-llama/Llama-3.2-1B \
     dataset=<preference_dataset> renderer_name=role_colon \
     learning_rate=1e-5 dpo_beta=0.1
   ```
6. **Inspect eval after training:**
   ```bash
   python -m tinker_cookbook.eval.run_inspect_evals \
     model_path=tinker://YOUR_MODEL \
     model_name=meta-llama/Llama-3.2-1B \
     tasks=<inspect_task_id> \
     renderer_name=role_colon
   ```