Leo Pekelis on Extending Llama 3 to a Million-Token Context

Matthew Berman introduces Leo Pekelis, Gradient’s chief scientist, to discuss the team’s extension of Llama 3 to a million-token context window. At the time of this June 2024 interview, the original Llama 3 supported an 8K context. Berman asks what a context window means and why making it longer matters.

Pekelis explains that autoregressive language models generate one token after another. Instructions, documents and conversation history seed that sequence, steering what the model generates next. The context window is the amount of text available to the model in that sequence; Berman notes that input and output together count toward the total.

Pekelis compares context to working memory and training to knowledge held deeper in memory. A longer window lets a model attend over more supplied material, potentially reducing the need to preprocess, select or summarize it outside the model. The analogy describes a useful distinction, not a literal claim that models have human memory.

Berman observes that a longer window can accommodate books or multimodal material that previously had to be split into chunks. Pekelis says the advantage is not just convenience: when the information fits, the model can attempt to choose which pieces matter instead of relying entirely on the user’s advance selection. The practical quality and speed still depend on the model and task.

For an enterprise coding use case, Pekelis contrasts completing a small function with implementing a feature or workflow that requires information from an entire codebase or several repositories. An agent could fetch pieces sequentially, he says, but that can be brittle, slow and costly. A larger window makes it possible to supply more material in one request and potentially reduce the number of calls.

Berman describes the limitations of assistants that mainly see the open file or a partial map of a project. Pekelis says Gradient has experimented with placing disparate codebases in context to ask how they might be integrated. These are examples of the team’s experiments, not a guarantee that a model understands or safely modifies every repository.

Asked about his own workflow, Pekelis says he uses VS Code and has tried Cursor. He also experiments with Gradient’s own long-context models, although they are less integrated into the editor and require more back-and-forth. The team runs them on internal clusters and accepts some waiting time while it studies how best to serve them.

The team is experimenting with caching the initial portion of a query, such as a codebase that remains largely unchanged across questions. That allows subsequent requests to reuse computation associated with the shared prefix. Pekelis describes this as an approximation: the implementation is more involved than simply storing all pairwise attention results.

Pekelis explains the compute bottleneck through attention: in a straightforward formulation, interactions among tokens grow roughly with the square of sequence length. Caching can avoid repeating some work when a prefix is reused, but processing a long input still consumes substantial resources. Efficient serving is therefore a separate challenge from training a model to handle long context.

Berman asks how Gradient extends Meta’s base Llama 3. Pekelis identifies two challenges: the compute needed to process long sequences, and teaching a model trained for shorter contexts to work with much longer ones. Supplying extra tokens alone is not evidence that the model can use them reliably.

Pekelis describes the training as closer to continued pretraining than a simple conventional fine-tune, with an important role for positional encoding. The model needs a way to represent the relative positions of tokens that may be separated by hundreds or far more. He presents positional encoding as central to extending the range over which the model can operate.

Long-context training is costly, Pekelis says, and Gradient had been working on its efficiency before Llama 3 appeared. He credits open research and resources for influencing the team’s stack. Compared with a published Berkeley Large World Model paper, he estimates roughly 30 times greater compute-time efficiency and perhaps 100 times greater sample efficiency; these are his team’s comparisons, not an independently established universal ratio.

Berman asks whether the extended model is more demanding on an end user’s hardware. Pekelis says using a short 8K context is not more intensive simply because the model supports a larger maximum, but providing a million tokens requires much more computation. The distinction is between the model’s maximum capacity and the length actually used.

Pekelis adds that Gradient tried to preserve Llama 3’s short-context performance while extending its range. For the team, a longer context would not be a worthwhile improvement if it severely degraded the base model. His claim describes the design goal and the team’s monitoring, rather than proof that every short-context capability is unchanged.

To explain evaluation, Pekelis describes a needle-in-a-haystack test: place one piece of information within a large body of text and ask the model to retrieve it. Gradient sometimes uses War and Peace as part of the haystack, which he estimates at roughly 750,000 tokens. He sees associative recall as a useful basic test and a precursor to richer in-context learning.

Berman asks whether retrieving one isolated fact is enough when an application must connect information across a codebase. Pekelis agrees that there is a substantial gap between finding a single item and reasoning about a whole project. The test is a first step, and many long-context models were already close to perfect on it at the time of the interview.

Berman describes Gradient’s reported retrieval heatmap as all green across the tested positions; Pekelis says Gemini also scored all green. That result shows success on the particular retrieval test, not universal reasoning ability across a million tokens. More difficult evaluations are needed to distinguish systems that all perform well on the simple task.

Pekelis points to NVIDIA’s RULER benchmark, which extends needle retrieval into 13 task variations. These include multiple needles, different facts sharing a key, and selecting one relevant item among many competing items. Its variable-tracking task requires following references such as X2 = X1 and X3 = X2 to recover a value.

Berman compares variable tracking to a scavenger hunt in which one clue leads to the next. Pekelis says Gradient’s model was fourth on the RULER leaderboard at the time, behind Gemini, GPT-4 and Cohere Command R+. Berman questions the comparison with models that have different maximum context sizes.

Pekelis acknowledges the key limitation: the RULER comparisons he describes ran only up to 128K tokens because the tasks are more complex. The leaderboard position is therefore not evidence of equivalent performance at one million tokens. Gradient was looking into more detailed testing beyond 128K.

Looking ahead, Pekelis is interested in ways to serve long-context models with less memory and computation. At the time, users largely depended on providers such as Google or Gradient to host demanding models. He suggests approaches that compress portions of memory and selectively expand them when needed, borrowing an analogy from how people recall relevant chapters rather than hold an entire textbook in mind.

Pekelis invites researchers and users to connect with Gradient through its social channels and Discord, and says the team benefits from collaboration with the open-source community. Berman closes by saying he wants to keep testing the million-token Llama 3 model and hopes the team will eventually reach two million tokens. Pekelis treats that as a future goal, not a release announcement.