Cartridges: Storing long contexts in tiny caches with self-study

What is this? This repository provides code for training a cartridge, a small KV cache representing a large dump of text, using a test-time training recipe called self-study. The code is based on our paper Cartridges: Lightweight and general-purpose long context representations via self-study.

tl;dr When we put lots of text (e.g. a whole code repo) into a language model's context, generation cost soars because of the KV cache's size. What if we trained a smaller KV cache for our documents offline? Using a test-time training recipe called self-study, we show that this simple idea can improve throughput by 26× while maintaining quality. (See our blogpost for more.)

Table of contents

• Setup
• Synthesizing Training Data with Self-Study
• Training a Cartridge
• Serving Cartridges
• Acknowledgments and Citation

Step 1: Clone the repository and install the Python package.

Step 2: Set some environment variables

The codebase relies on your setting the following variables. We recommend adding them to your ~/.bashrc, ~/.zshrc, DockerFile, etc.

What is self-study? Self-study is a test-time training approach where we generate synthetic conversations about a corpus of text. The process simulates two AI agents: one asks questions or makes requests about the content, and another responds using the provided context. This creates training data that teaches the model to efficiently compress and retrieve information from long contexts.

Quickstart: Take a look at the script at scripts/longhealth_synthesize.py for an example of how to generate training data with self-study. To actually run the script, you will need to spin up an inference server (either Tokasaurus or SGLang) and set the client variable to point to it.

Below we walk through the process of generating synthetic training data for a corpus of text in more detail. As a running example, we'll be training a cartridge on our paper on Cartridges. How meta! Here are the steps:

1. Create a StructuredContext object that contains the data you want to store in the cartridge
2. Ensure you have an inference server running (either Tokasaurus or SGLang) and configure your client to point to it
3. Instantiate a SynthesizeConfig object that contains the parameters for the self-study process
4. Put it all together in one script and run it!

Note: For configuration, we use Pydantic models. Pydantic models are useful for defining the schema of the config and quickly ensuring that the config is valid at the beginning of the script. We also rely on pydrantic, which provides a few utilities for working with configs.

A StructuredContext represents your corpus in a format that the self-study process can work with. We provide several built-in context types. For our example, we'll use the TexDocument context type.

We provide a few other context types including HTMLDocument, TexDocument. Can also use an arbitrary JSON object as a context.

Self-study requires an inference server to generate the synthetic conversations. We support two options:

• Tokasaurus (recommended) - We ran all of our experiments with Tokasaurus, which provides higher throughput generation and is easier to modify.
• SGLang - We're also providing support for SGLang, but we have not tested it extensively.

We found it easy to run data generation with Modal's serverless horizontal scaling.

For cloud deployment, you can deploy on Modal:

Then configure with the modal URL:

If you have access to GPUs, you can run also run a local Tokasaurus server:

1. Clone and install Tokasaurus:

2. Start the server:

3. Configure your client:

We found it easiest to run data generation with Modal because it provides serverless horizontal scaling.

For cloud deployment, you can deploy on Modal:

Then configure with the modal URL:

1. Install and launch a SGLang server following the instructions here.
2. Configure your client:

We are now going to put all of the pieces together in a SynthesizeConfig object that configures the entire self-study process.

Core Settings:

• num_samples: Total number of training examples to generate
• batch_size: Number of training examples to generate per call to the inference server.
• max_num_batches_in_parallel: Number of batches to process concurrently. When using Modal, high values

Here's a complete example script:

Once you've created the file, run it with:

Once the run is complete, it will save the results to a pickle file and print the path:

You can enhance the self-study process with tools that allow agents to dynamically retrieve additional context:

Quickstart: Take a look at the script at scripts/longhealth_train.py for an example of how to generate training data with self-study.

See cartridges.train.TrainConfig for the schema of the main config we use for training.

Below we provide an example of a config file prefaced with notes describing each part of the config:

• dataset Th
• *`

To launch a data parallel training run, you can run:

We describe two ways to serve and chat with a trained Cartridge: a simple, but slow way that just uses a pure PyTorch generation loop, and a faster one that uses a Tokasaurus server.

We've implemented (h/t @geoffreyangus) an integration with Tokasaurus, a simple LLM inference server optimized for high throughput.

To run the Tokasaurus server, you will need to (install Tokasaurus from source)[], switch to the branch geoff/cartridges, and then follow the instructions here to make API calls to the server.

We

There are tons of people and organizations who have supported this project. Below we shout out a few, but check out the the paper for a full list.

The compute for this project was provided by Modal — who made it super easy to scale out horizontally when running the synthetic data generation for self-study — and Together — who provided the compute for training the Cartridges on the synthetic data. Prime Intellect, Voltage Park, and Azure through the HAI Grants program also contributed compute towards this project.