Tinker Cookbook

repository·main·Indexed 26 days ago

https://github.com/thinking-machines-lab/tinker-cookbook

A library of abstractions and recipes for customizing language models using the Tinker training SDK. It provides tools for supervised fine-tuning (SFT), reinforcement learning (RL), preference learning (DPO/RLHF), and model evaluation, including a benchmark framework for datasets like gsm8k, mmlu_pro, and ifeval.

Tokens
81.3K
Snippets
234
Records
378
Agent score
87%

What's inside tinker-cookbook

  1. Overview of Audio fine-tuning recipes

    main

    The tinker_cookbook/recipes/audio/ directory contains recipes for fine-tuning audio-input models (Inkling) using supervised fine-tuning (SFT) and reinforcement learning (RL).

    Available recipes include:

    • Speech recognition: SFT and RL on LibriSpeech using a -WER reward.
    • Emotion + speech recognition: A two-stage SFT → RL pipeline for classifying speaking style and transcribing expressive speech.
    • Medical ASR domain adaptation: SFT for specialist vocabulary (e.g., drug names) using WER and medical-entity recall metrics.

    Each recipe directory contains its own env.py, tests, and training entrypoints (such as sl_train.py, rl_train.py, or train.py).

  2. Overview of Tinker Tutorials

    main

    The Tinker Cookbook provides a structured learning path through several categories of notebooks:

    Basics (1xx)

    Covers architecture, client hierarchy, sampling, SFT (Supervised Fine-Tuning), async patterns, and introductory RL (Reinforcement Learning).

    Core Concepts (2xx)

    Covers rendering, tokenization, loss functions (cross_entropy, IS, PPO, CISPO), completers, weights/checkpoint lifecycle, and evaluations.

    Cookbook Abstractions (3xx)

    Focuses on high-level abstractions like Env, EnvGroupBuilder, RLDataset, ProblemEnv, and configuration-based training using train.Config and train.main().

    Advanced (4xx)

    Covers hyperparameters (SL and RL), DPO (Direct Preference Optimization), sequence extension (multi-turn RL), multi-agent RL, prompt distillation, and full RLHF pipelines.

    Deployment (5xx)

    Covers exporting models to HuggingFace, building LoRA adapters for vLLM/SGLang, and publishing to the Hub.

  3. Explore post-training LLM training examples

    main

    The repository contains a wide range of specialized training environments for LLMs. Each example is located in its own subfolder and includes a README.md with implementation details, execution commands, and expected performance.

    Available examples include:

    • Chat supervised learning: SFT on conversational datasets (e.g., Tulu3).
    • Math reasoning: RL for math question accuracy.
    • Code reasoning: Training on competitive programming with sandboxed execution.
    • Preference learning: A three-stage RLHF pipeline (SFT $\rightarrow$ Reward Model $\rightarrow$ RL).
    • Tool use: Training for retrieval tool usage.
    • Prompt distillation: Internalizing complex instructions.
    • Multi-Agent: Optimizing LLMs for competitive play.
    • Model distillation: On-policy or SFT distillation from teacher models.
    • Rubric-based grading: Using LLM graders with rubrics for RL rewards.
    • Verifiers environments: Integration with Prime Intellect's Environments Hub.
    • VLM image classification: Training vision-language models.
    • Audio: SFT and RL for speech recognition and domain adaptation.
    • Harbor RL: RL on Harbor-formatted tasks with sandboxed execution.
    • Self-Distillation Fine-Tuning (SDFT): Self-distillation via top-K forward KL loss.
    • True Thinking Score (TTS): Quantifying CoT faithfulness.
  4. Data requirements for Medical ASR (EkaCare)

    main

    The recipe uses the EkaCare medical ASR evaluation dataset (MIT, en config).

    • Content: Indian-accented English dictated prescriptions and consultations containing drug names and dosages.
    • Preparation: The dataset downloads automatically from Hugging Face on the first run; no manual preparation is required.
    • Splits: The recipe creates a seeded random 80/20 train/eval split from the single available split.
  5. Understand the Rubric-based Grading Recipe Structure

    main

    The rubric-based grading recipe for LLMs is composed of several key modules:

    • data.py: Defines the datapoint class. Each datapoint contains a conversation prefix (convo) and a list of rubric_items.
    • generate_data.py: Generates example datasets (e.g., for addition tasks).
    • env.py: Defines the environment logic. It allows a policy to read a prefix, generate a response, and then uses a grader LLM to grade that response based on rubric_items. The final reward is the sum of the responses from each grader.
    • train.py: The training script used to train LLMs on datasets formatted according to data.py.
    • prometheus_experimental.py: An experimental script for training LLMs using rubrics from the prometheus-eval/Feedback-Collection dataset.
  6. Explore preference learning recipes

    main

    The preference recipe directory provides implementations and guides for learning from non-scalar rewards. Key workflows include:

    • Shorter responses: Uses the PairwisePreferenceRLDatasetBuilder abstraction to train models to generate shorter outputs.
    • RLHF (Reinforcement Learning from Human Feedback): Implements the standard three-stage pipeline: supervised fine-tuning, reward model learning, and reinforcement learning.
    • DPO (Direct Preference Optimization): Implements the DPO algorithm using a custom loss function to optimize for human preferences directly.
  7. Setup and Use Vision-Language (VL) Renderers

    main

    VL renderers require an image_processor to handle image tokens correctly. Forgetting this will raise an AssertionError when rendering image content.

    Setup

    from tinker_cookbook.image_processing_utils import get_image_processor
    from tinker_cookbook.renderers import get_renderer
    from tinker_cookbook.tokenizer_utils import get_tokenizer
    
    model_name = "Qwen/Qwen3.6-35B-A3B"
    tokenizer = get_tokenizer(model_name)
    image_processor = get_image_processor(model_name)
    
    renderer = get_renderer(
        "qwen3_5",
        tokenizer,
        image_processor=image_processor,  # Required for image content
    )

    Troubleshooting VL Issues

    • Token Mismatch: Ensure the same transformers version and image_processor configuration (e.g., max resolution) are used during both training and serving.
    • Image Token Count: If you see Expected X tokens, got Y from image, upgrade to transformers>=5.0 to fix a known Qwen2VLImageProcessor bug.
    • Weight Loading: VL models use a model.language_model.* prefix. Use weights.build_hf_model() to ensure LoRA adapters are correctly applied to the prefixed weights.
    from tinker_cookbook.image_processing_utils import get_image_processor
    from tinker_cookbook.renderers import get_renderer
    from tinker_cookbook.tokenizer_utils import get_tokenizer
    
    model_name = "Qwen/Qwen3.6-35B-A3B"
    tokenizer = get_tokenizer(model_name)
    image_processor = get_image_processor(model_name)
    
    renderer = get_renderer(
        "qwen3_5",
        tokenizer,
        image_processor=image_processor,  # Required for image content,
    )
  8. Implement stateful tools using classes

    main

    To create tools that maintain state or share configuration (like an API key), apply the @tool decorator to methods within a class. When passing these to the environment, pass the bound methods of a single class instance.

    class MyTools:
        def __init__(self, api_key: str):
            self._api_key = api_key
    
        @tool
        async def search(self, query: Annotated[str, "Query"]) -> ToolResult:
            """Search using the configured API."""
            results = await search_api(query, self._api_key)
            return simple_tool_result(json.dumps(results))
    
        @tool
        async def lookup(self, id: Annotated[str, "Document ID"]) -> ToolResult:
            """Look up a document by ID."""
            result = await lookup_api(id, self._api_key)
            return simple_tool_result(json.dumps(result))
    
    # Usage: tools share the same instance state
    tools_obj = MyTools(api_key="...")
    env = build_agent_tool_env(..., tools=[tools_obj.search, tools_obj.lookup])
  9. Follow the Tinker Research methodology

    main

    When conducting research using the Tinker Cookbook, follow this structured methodology:

    1. Understand the problem: Clearly define the task.
    2. Know your models: Understand the capabilities and limitations of the models being used.
    3. Set up evaluation FIRST: Establish evaluation protocols before starting any training process.
    4. Run baseline eval BEFORE any training: Always establish a baseline performance score before introducing training changes.
    5. Analyze failure modes: Load incorrect examples from evaluations to understand why the model is failing.
  10. Train on Terminal-Bench 2.0 tasks

    main

    To run RL training on Terminal-Bench 2.0, follow these steps:

    1. Download the tasks:
    uvx harbor datasets download terminal-bench@2.0 -o ~/.cache/harbor/tasks/terminal-bench-2.0/
    1. Launch the training script:
    uv run python tinker_cookbook/recipes/harbor_rl/scripts/train_terminal_bench.py \
        model_name=moonshotai/Kimi-K2.6 \
        max_tokens=8192 \
        group_size=4 \
        groups_per_batch=8 \
        learning_rate=1e-5 \
        lora_rank=32 \
        wandb_project=cookbook_harbor_rl
  11. Implement the `TwentyQuestionsEnv` training environment

    main

    To create a custom Twenty Questions environment, implement the TwentyQuestionsEnv class. The core logic resides in the TwentyQuestionsEnv.step method.

    TwentyQuestionsEnv.step Input

    • Action: A sequence of integer tokens from the language model that can be parsed into a natural language question (e.g., "Is it a plant?").

    TwentyQuestionsEnv.step Output (StepResult)

    • Reward: 1 if the player correctly guesses the secret keyword (e.g., "Apple"), 0 otherwise.
    • Termination: An indicator of whether the episode should end (triggered when the player guesses correctly or exceeds 20 questions).
    • next_stop_condition: Typically the stop token of the policy language model.
    • next_observation: The conversation history, the last question asked, and the answer to that question, provided in token space.