tinkder-cookbook

Imported from yanzhen4/search_uncertainty (tinkder-cookbook/AGENTS.md). Install upstream with npx skills add yanzhen4/search_uncertainty --skill tinkder-cookbook. Copyright stays with the author.

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 SamplingClients 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:

   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):

   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):

   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):

   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):

   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:

   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