TRL: document 1M-token training and add a context-parallelism example

[qgallouedec]

Two things, both about training on sequences far longer than the usual few thousand tokens:

• a "Training at 1M tokens and beyond" section in docs/source/distributing_training.md, with the configurations I verified and the levers in the order you hit them;
• examples/sft_long_context/, a runnable example that trains a book-length sequence per step on one 8xH100 node, plus examples/accelerate_configs/context_parallel_8gpu.yaml.

Measured

One 8xH100 node, bf16, per_device_train_batch_size=1, loss_type="chunked_nll" (the default):

Model	Sequence length	Step time	Peak GPU memory
Qwen3-0.6B	1M	137 s	27.9 GB
Qwen3-8B	1M	364 s	56.2 GB
Qwen3-30B-A3B (MoE)	1M	483 s	46.0 GB
Qwen3-32B	1M	1295 s	60.8 GB
Qwen3-0.6B	4M	2135 s	73.7 GB

Context length scales with node count (Qwen3-8B: 1M on 1 node at 364 s, 2M on 2 nodes at 696 s, 4M on
4 nodes at 1346 s), because each GPU keeps the same shard of the sequence.

It is real training, not just "it fits": Qwen3-30B-A3B with YaRN x32 on PG-19 books at 1M context goes
from loss 4.88 to 2.37 in six steps.

Assumptions

Everything here is written as if the rest of the series has landed. Each of these is a separate PR, and the text or the example changes if one of them does not merge:

assumption	where it shows up	if it does not land
huggingface/accelerate#4175 — FSDP2 activation_checkpointing_offload	context_parallel_8gpu.yaml sets fsdp_activation_checkpointing_offload: true; the guide lists it as lever 3	drop the line from the YAML; models at or above 8B no longer fit at 1M
huggingface/transformers#48163 — cuDNN SDPA preference on Hopper	the guide says Transformers selects it automatically; the example does not select a backend	the example needs an SFTTrainer subclass wrapping training_step in sdpa_kernel([SDPBackend.CUDNN_ATTENTION]), and every step time above is ~1.7x larger
huggingface/accelerate#4177 — refuse sliding-window models under CP	the constraints section says accelerate rejects them	they are silently trained with full causal attention instead, and the section has to say so
#6843 — refuse packing under CP	the constraints section says TRL raises	the "Do not use packing" bullet becomes advice rather than a guarantee
transformers cp-eval-loss-scaling — apply CP to the evaluation path	not mentioned, but any user who evaluates during a long-context run gets eval_loss multiplied by cp_size	worth a warning in the guide

Two more that are not assumptions, just constraints of what exists today, both stated in the example:

• MLP sequence tiling (lever 6) is a user-side monkeypatch, not a feature. It is what makes Qwen3-32B fit.
• Context parallelism expresses only full causal attention, so gpt-oss, Gemma 3/4, Mistral and Qwen3.5
  and later cannot be used. Qwen3.5+ is the interesting one: three quarters of its layers are linear
  attention, which is a different problem from the sliding-window case and is not covered by [New Trainer] TreeRPO: Hierarchical Credit Assignment for Deterministic-Correctness Tasks #4177's
  fix, only by its refusal.

---

Note

Low Risk
Documentation-only change; no runtime, training, or security code is modified.

Overview
Adds a Training at 1M tokens and beyond section to the Ring Attention docs, with verified step times and peak memory for Qwen3 models from 0.6B to 32B (and a 30B MoE) on 8×H100, plus multi-node scaling for Qwen3-8B at 1M/2M/4M.

Documents the memory/speed levers in the order they become necessary: chunked NLL, cuDNN SDPA, FSDP activation offload, alltoall ring rotation + bf16 merge, param CPU offload, MLP sequence tiling, and YaRN RoPE scaling. Also states CP constraints: full causal attention only (no sliding-window models) and no packing.

^{Reviewed by Cursor Bugbot for commit 851fd9a. Bugbot is set up for automated code reviews on this repo. Configure here.}

[bot-ci-comment]

The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update.

[qgallouedec]

Reopening from a branch actually based on #6820 — this one was cut from main, so the diff showed 125 unrelated files.

[chatgpt-codex-connector]

💡 Codex Review

Here are some automated review suggestions for this pull request.

Reviewed commit: 851fd9a0a6

ℹ️ About Codex in GitHub

Your team has set up Codex to review pull requests in this repo. Reviews are triggered when you

• Open a pull request for review
• Mark a draft as ready
• Comment "@codex review".

If Codex has suggestions, it will comment; otherwise it will react with 👍.

Codex can also answer questions or update the PR. Try commenting "@codex address that feedback".

[chatgpt-codex-connector]

Set packing false in the referenced million-token setup

The referenced setup explicitly sets packing=True, while the constraints below explain that context parallelism drops the block-diagonal mask. SFTTrainer currently only warns about packing with non-Flash attention; it does not raise based on context parallelism as claimed later. Following “the setup above” can therefore allow packed documents to attend to one another and silently corrupt the training results, so the million-token recipe must explicitly disable packing rather than relying on a nonexistent guard.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Disable Liger before relying on chunked NLL

Under the setup referenced here, use_liger_kernel=True; SFTConfig.__post_init__ consequently changes the default loss to "nll", and explicitly selecting "chunked_nll" raises an incompatibility error in SFTTrainer. Thus users either hit an exception or materialize the full million-token logits tensor that this recommendation is intended to avoid, likely causing OOM. The recipe needs to set use_liger_kernel=False before calling chunked NLL the default.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Gate activation offload on a supporting Accelerate version

This option is not available in the documented minimum Accelerate 1.11.0 setup; support depends on the newer FSDP2 activation-offload change. A user satisfying the stated requirements can therefore have this configuration key rejected instead of obtaining the memory savings required for the listed 8B+ runs. Specify the first released Accelerate version that provides the option and raise the section's minimum accordingly.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Provide the Transformers 4.x RoPE override

TRL still supports Transformers 4.x (transformers>=4.56.2), where Qwen3 expects rope_scaling plus the top-level rope_theta, not the v5-only rope_parameters dictionary shown here. On that supported dependency range this override does not configure YaRN as intended, so users can start a costly million-token run with the unscaled positional encoding and the high-loss behavior described immediately above. Include the 4.x form or explicitly require Transformers v5 for this recipe.

Useful? React with 👍 / 👎.

[chatgpt-codex-connector]

Require the Accelerate version that rejects local attention

The stated Accelerate 1.11.0 minimum does not provide the advertised rejection of sliding-window or chunked-attention models, and TRL performs no equivalent SFT validation. With such a supported installation, selecting one of the listed models can proceed after context parallelism drops its per-layer mask, silently training full causal attention instead of the model's intended architecture. Tie this guarantee to the Accelerate release containing the validation, or warn that older versions do not enforce it.

Useful? React with 👍 / 👎.