Training a Million-Token Context Window Is Only the Beginning

From brittle workflows to longer context

An enterprise introduces a new workflow, and its machine-learning platform needs rebuilding. An incremental improvement becomes a project lasting quarters or years. Mark Huang describes that experience as a founding motivation for Gradient: build systems that survive changing workflows, handle unfamiliar inputs and improve alongside their users.

Huang left quantitative finance seeking more collaboration and experience with large-scale data and machine learning. He worked on product analytics and go-to-market at Box, then streaming analytics, search and deep learning at Splunk. In this May 2024 conversation, he describes Gradient as a full-stack AI platform for moving codified enterprise automation toward more autonomous workflows. His test for an agent is practical: additional steps and nondeterministic execution should increase the probability of success.

Long context grew out of the same ambition. Gradient had worked extensively on fine-tuning and retrieval-augmented generation—retrieving relevant material to supply to a model. Huang treats those techniques and long-context prompting as complementary ways to adapt a model to a workload. A larger context window lets the model receive more information and demonstrations in one input.

Why extend Llama 3?

Gemini’s million-token window encouraged exploration beyond prompts containing ten or a hundred demonstrations. Gradient was preparing to extend other open-source models when Llama 3 arrived. Its initial 8,000-token window offered an obvious opportunity. Its training on 15 trillion tokens also suggested to Huang that it contained substantial capabilities that might adapt well to longer inputs.

One million tokens was an exploratory target. The team would extend the window, evaluate, and continue while results looked promising. Reaching it required substantial preparation and compute: Crusoe supplied L40S GPU instances and helped arrange a dedicated cluster.

Gradient used a context-length curriculum: progressively increase the training-sequence length of an already trained model. The intuition resembles studying a textbook chapter by chapter, doing exercises along the way, before attempting problems that integrate the whole book. Whether this approach generally outperforms training at long context throughout remained unresolved.

Make longer positions usable

The starting model must already predict tokens well. The discussion suggests that an ordinary 4K-context model could be extended progressively, without establishing a minimum starting window. Context extension cannot repair weak next-token prediction. The model’s positional encoding—its representation of where tokens occur—also matters.

For Llama 3’s rotational positional encoding, discussed as RoPE, the parameter theta controls how frequently positional patterns repeat. In the conversation’s qualitative explanation, increasing theta changes those patterns so that longer sequences have positional relationships more like those encountered during training. Positional interpolation compresses positions toward the original range; the speaker favors it over extrapolating beyond familiar positions. Gradient’s write-up on scaling rotational embeddings accompanies this discussion.

The team evaluated between extensions, using a theta-scaling formula from prior work that held over the lengths it tried. Empirical success did not guarantee continued scaling. Theta scaling was simple and familiar to the team; it had not directly compared its approach with the alternative ALiBi method. The speaker also notes that YaRN can be used with RoPE.

Make training work on the cluster

Position scaling addressed how the model would handle longer sequences. Processing those sequences also required changes to the training system. Gradient combined FlashAttention and Ring Attention, configuring the GPU network topology to maximize memory bandwidth.

The original Berkeley JAX implementation offered mature topology configurations but was difficult to run on the team’s GPUs out of the box. Other PyTorch implementations failed in its setup too, possibly partly because of configuration mistakes. Active breaking changes added maintenance risk during the experiment.

EasyContext was the first PyTorch implementation using native libraries that worked well for Gradient. The team adapted it to its cluster topology, credited Zhang Peiyuan’s contribution, and recommended it as an accessible starting point for engineers comfortable with PyTorch.

Long examples must require distant information

Gradient describes two training stages: 1. Continual pretraining: filter and concatenate the dataset referred to as SlimPajamas into sequences matching the target context length. 2. Chat training: filter and reformat UltraChat, or a curated derivative, for conversational use. Truncation boundaries and data diversity needed care to preserve existing capabilities. The speaker suggests embeddings as a possible prefilter for assessing diversity relative to the original training corpus.

Length alone does not create a useful training example. If the final 200 tokens always contain everything needed to answer, the model has no reason to use earlier information. A useful long-context learning signal requires connections across the sequence. The practical dataset question is: could the model answer from the final fragment alone?

Gradient used full fine-tuning. The speaker cautions that even a billion additional tokens is small beside pretraining measured in tens of trillions, making it unwise to assume that such an update teaches large amounts of new knowledge. The challenge is gaining capabilities while retaining old ones; replaying older data is described as a common defense against forgetting.

Code Llama’s specialization from Llama 2 serves as a cautionary example: the speaker judges it to have degraded general language ability, though no benchmark results accompany that assessment. Another participant proposes gradually changing the data mixture while retaining original-source material, or modifying the loss to constrain drift. The response distinguishes useful heuristics from a proven solution. DoReMi is cited as interesting data-mixture research, but success on simple tasks leaves open how to preserve complex abilities that are themselves difficult to evaluate.

Gradient also used GPT-4 to rephrase and reformat chat examples and generate more varied language and data, including attempts to introduce out-of-domain examples. The speaker locates much of a model developer’s competitive advantage in this data pipeline: its construction shapes what the extension teaches and which existing behaviors it preserves.

Can capabilities travel through model merges?

A related experiment explores model alchemy: combining model weights or adapters to mix existing skills without another training run. Gradient subtracted a Llama base model from its instruction-tuned counterpart and applied the resulting difference to another model through an adapter, described as LoRA. The changes were sometimes surprisingly effective, according to the speaker, but complex abilities transferred less reliably than stylistic behavior.

The proposed explanation—that complex abilities depend on deeper network pathways disrupted by brute-force merging—remains a hypothesis. Merging also creates an evaluation hazard: selecting models for benchmark strengths can produce targeted leaderboard gains. The speaker finds merge, then fine-tune more promising, possibly adding a little noise to help subsequent learning, but reports no measured acceleration.

Finding a needle is only a prerequisite

Needle in a Haystack asks whether a model can retrieve an isolated fact from a long input. Gradient treated it as a prerequisite for useful long-context behavior. Passing does not establish understanding or instruction following: a model can still ignore the request or mistake supplied source material for instructions.

The team therefore used RULER, which it found more demanding. Tasks discussed include: - Retrieving multiple key-value pairs, with multi-value and multi-query variants. - Tracking variables through chains of assignments. - Counting common words distributed throughout the input. These probe whether a model can connect or aggregate information across a sequence, beyond locating one fact.

That distinction matters when information is scattered across documents. Gradient describes carefully tuned retrieval that worked for one collection but stopped finding the necessary context when other documents were added. A model that generalizes across those changes would reduce this brittleness. Yet benchmarks cover only part of the capability space, and richer assessments can be expensive: teacher-reviewed lesson plans in Gemini 1.5 evaluations exemplify studies the startup would struggle to reproduce.

Asked about benchmark degradation when extending from one million to four million tokens, the respondent suggests that a very large theta and floating-point precision might cause numerical problems. Clamping activations during training, and possibly inference, is proposed as a remedy. This remains a tentative diagnosis, without degradation scores, a specified precision format or a tested fix.

The harder workload is changing state

One million tokens generated excitement; four million felt like another checkpoint. The practical question became what could fill the window usefully. An entire code repository is one candidate, but iterative, agentic techniques might produce better modules or suggestions than a single large-context request.

Gradient was especially interested in evolving conversational state: the requirements and concepts that change over successive turns. Its proposed evaluations introduce nuanced changes and test whether the model produces a known target output. Coding sessions and investment-management discussions are potential applications. Static documents can sometimes use retrieval or successive layers of summarization; rapidly changing session state is harder to manage without explicitly codifying the workflow and its transitions.

Many-shot prompting supplies hundreds or thousands of demonstrations in the input. Perhaps a thousand or two thousand examples could reduce dependence on retrieving precisely the right ones. A participant cites an unnamed CMU study reporting reduced example sensitivity on a few extraction and classification tasks with many possible classes. Whether enough demonstrations make random sampling adequate remains open; no quantitative result or model-specific recipe is given.

A participant’s coding experience captures the other side of remembering: a model repeatedly reapplies earlier CSS styles after being told they may be inappropriate. Useful context handling requires retaining relevant history while suppressing obsolete information.

Multimodality changes what fills the window

Video frames and large image collections provide another reason for longer context: they consume substantial token budgets. Financial charts paired with explanations and medical images paired with text are proposed applications. Connecting visual and verbal information could give models better grounding.

The discussion turns to early fusion, integrating images and text within multimodal training. A speaker cites Meta’s Chameleon work as showing greater sample efficiency, without specifying comparison conditions. Another participant asks whether an existing Llama 3 model could be retrofitted this way. A related idea was on Gradient’s research backlog, and the response invites collaboration; no implemented technique is described.

This raises a product-strategy problem: progress in reasoning, modality integration and context length can follow different paths. A generic expectation of tenfold improvement gives little guidance about which capability users need. The principle offered is to pursue improvements that materially increase user value, rather than treating technical difficulty as a reason to build.

Watch the training curve—and choose the research

The research routine described combines AI News aggregation, selected researchers on Twitter, implementation and dataset discussions on Discord, and AI-assisted paper searches. Trying products is another form of study: they embody accumulated research choices, and reasoning backward from their behavior can help reveal those choices. Perplexity, the product, is offered as an example.

Returning to the training metric, perplexity indicates how well the model predicts tokens. Gradient watched its early trajectory: a sharp decline suggested that context length and theta scaling worked together, while substantial oscillation prompted the team to stop and retry another theta scale. A value near four is mentioned as a rough sign of progress for Llama extension, but without the dataset, model size or evaluation protocol needed to make it a portable threshold.

Research selection starts with three to five areas that could help the company, then excludes work that is irrelevant or infeasible. Algorithms for training on roughly 2,000 H100s were beyond the team’s available scale; evaluations, post-training methods and synthetic data construction were useful priorities. Accumulated topic knowledge helps identify novelty, while expert conversations test whether an apparently new result is credible and worth pursuing.

The closing invitation is to collaborate on long-context evaluations and paired multimodal datasets, such as video connected to text. Both address unfinished work: teaching models to connect information across large inputs, and determining whether those connections produce useful behavior.