Will it run?
Models

TRL 1.13 releases guide for million-token context training on single node

By Nadia Ksiazek Clawpit staff
TRL 1.13 releases guide for million-token context training on single node

The open-source RL library TRL 1.13 centers long-context training with a practical guide that demonstrates fine-tuning on sequences of one million tokens and beyond on a single 8-GPU H100 node. The release includes general speed and memory improvements, but the immediate value lies in the technical breakdown of the real memory bottleneck — not model weights, but loss computation.

The guide fine-tunes Qwen3-8B on PG-19 books concatenated to roughly 1.049 million tokens per example. A single training step takes about 380 seconds, a little over six minutes, and the initial loss sits at 4.31, a healthy value confirming correct configuration; a misconfigured long-context run typically starts around 10. Naively, the sequence demands 288 GB per GPU; after the guide's optimizations that drops to 56 GB.

Profiling shows the memory spike appears between the forward and backward passes, exactly where the loss is computed. The final decoder layer emits a hidden state of shape (sequence_length × hidden_size), which is multiplied by the language-modeling head to produce a logits matrix of shape (sequence_length × vocab_size) — tens of thousands of columns per token. At a million tokens that matrix no longer fits in memory.

The fix: chunked loss computation. Instead of materializing the full logits matrix at once, TRL splits the calculation into chunks of 256 rows. Cross-entropy sums over rows, and softmax runs across the vocabulary for a single row, so each chunk contains everything needed for its portion of the loss. One large multiplication becomes a series of small ones, and the memory spike disappears; loss is no longer the bottleneck.

Caveats apply. The guide requires transformers from main branch, not a released version, because it relies on gradient checkpointing with offload that has not yet landed in an official release. The code runs immediately with accelerate and the supplied configuration, but it is a technical demonstration, not a production pipeline. No downstream benchmarks exist yet for models trained this way, and the guide itself notes the run was tested on a single node only.