Skip to content

Commit 1109004

Browse files
authored
docs: add yarn doc (#2283)
Signed-off-by: ruit <ruit@nvidia.com>
1 parent 84cdfeb commit 1109004

2 files changed

Lines changed: 109 additions & 0 deletions

File tree

docs/guides/yarn-long-context.md

Lines changed: 101 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,101 @@
1+
# YaRN Long-Context Training
2+
3+
[**YaRN** (Yet another RoPE extensioN)](https://arxiv.org/abs/2309.00071) extends a model's usable context window beyond the length it was pretrained on by rescaling RoPE frequencies. NeMo RL supports YaRN RoPE scaling for SFT, GRPO, DPO, RM, and distillation workflows, letting you fine-tune or RL-train models at sequence lengths much larger than their original pretraining context.
4+
5+
## Requirements
6+
7+
YaRN is only supported with the **Megatron backend**. The DTensor (Automodel) backend will raise an assertion error if `rope_scaling.rope_type=yarn` is set. Make sure:
8+
9+
1. Megatron submodules are initialized: `git submodule update --init --recursive`
10+
2. Megatron backend is enabled: `policy.megatron_cfg.enabled=True` and `policy.dtensor_cfg.enabled=False`
11+
12+
## Enablement
13+
14+
YaRN is configured through `policy.hf_config_overrides.rope_scaling`. All YaRN fields are required — NeMo RL validates that none are missing before training starts.
15+
16+
```yaml
17+
policy:
18+
max_total_sequence_length: 65536
19+
megatron_cfg:
20+
enabled: true
21+
dtensor_cfg:
22+
enabled: false
23+
hf_config_overrides:
24+
rope_scaling:
25+
rope_type: yarn
26+
rope_theta: 1000000
27+
factor: 1.6
28+
original_max_position_embeddings: 40960
29+
truncate: true
30+
beta_fast: 32
31+
beta_slow: 1
32+
mscale: 1
33+
mscale_all_dim: 0
34+
```
35+
36+
### Required Fields
37+
38+
| Field | Description |
39+
| --- | --- |
40+
| `rope_type` | Must be `yarn` to enable YaRN. |
41+
| `rope_theta` | RoPE base frequency used when recomputing scaled frequencies. |
42+
| `factor` | Scaling factor — typically `max_total_sequence_length / original_max_position_embeddings`. |
43+
| `original_max_position_embeddings` | The model's original pretraining context length. |
44+
| `truncate` | Whether to truncate out-of-range positions. |
45+
| `beta_fast` / `beta_slow` | YaRN interpolation thresholds on the fast and slow RoPE dimensions. |
46+
| `mscale` / `mscale_all_dim` | Attention temperature scaling terms used by YaRN. |
47+
48+
You can also compute `factor` directly from other config values using the `div:` interpolation helper:
49+
50+
```yaml
51+
factor: ${div:${policy.max_total_sequence_length},${policy.hf_config_overrides.rope_scaling.original_max_position_embeddings}}
52+
```
53+
54+
> [!NOTE]
55+
> YaRN only takes effect when `max_total_sequence_length` exceeds `original_max_position_embeddings`. If the two are equal, YaRN is a no-op.
56+
57+
## How Conversion Works
58+
59+
When `hf_config_overrides` are present, NeMo RL's Megatron setup:
60+
61+
1. Validates that every required YaRN field is specified.
62+
2. Computes a stable hash over `hf_config_overrides` and appends it to the converted-checkpoint directory name (`<model>__hfovr_<hash>`). Different override sets therefore produce separate cached checkpoints and will not collide.
63+
3. Re-imports the HF model into Megatron format if no cached checkpoint at that path exists.
64+
65+
The same `hf_config_overrides` are also propagated to vLLM during generation, so the rollout engine applies the identical YaRN settings as the trainer.
66+
67+
## Forcing a Fresh Conversion
68+
69+
If you change yarn parameters in `hf_config_overrides` after a conversion has been cached (for example, adjusting the YaRN `factor`), set:
70+
71+
```yaml
72+
policy:
73+
megatron_cfg:
74+
force_reconvert_from_hf: true # Default: false
75+
```
76+
77+
This re-runs the HF → Megatron conversion and overwrites the cached checkpoint. It is equivalent to deleting the cached directory and rerunning. See also the [Training Backends design doc](../design-docs/training-backends.md#force-reconvert) for background on the checkpoint cache.
78+
79+
## Example Recipes
80+
81+
Two end-to-end YaRN recipes ship with NeMo RL:
82+
83+
- **SFT, 64K context**: [`examples/configs/recipes/llm/sft-qwen3-0.6B-1n8g-megatron-yarn-64k.yaml`](../../examples/configs/recipes/llm/sft-qwen3-0.6B-1n8g-megatron-yarn-64k.yaml) — Qwen3-0.6B fine-tuned to a 64K sequence length on Nemotron-Cascade-2-SFT-Math using `factor: 1.6` (64K / 40960).
84+
- **GRPO, 256K context**: [`examples/configs/recipes/llm/grpo-qwen2.5-1.5B-4n8g-megatron-yarn-256k.yaml`](../../examples/configs/recipes/llm/grpo-qwen2.5-1.5B-4n8g-megatron-yarn-256k.yaml) — Qwen2.5-1.5B trained at 256K sequence length with `factor` derived from `max_total_sequence_length / original_max_position_embeddings`.
85+
86+
Launch them the same way as any other recipe:
87+
88+
```bash
89+
uv run examples/run_sft.py \
90+
--config examples/configs/recipes/llm/sft-qwen3-0.6B-1n8g-megatron-yarn-64k.yaml
91+
92+
uv run examples/run_grpo_math.py \
93+
--config examples/configs/recipes/llm/grpo-qwen2.5-1.5B-4n8g-megatron-yarn-256k.yaml
94+
```
95+
96+
## Practical Tips
97+
98+
- **Set context parallelism appropriately.** Long sequences typically require `policy.megatron_cfg.context_parallel_size > 1`. The 256K recipe uses `context_parallel_size: 32`; the 64K recipe uses `8`.
99+
- **Check `make_sequence_length_divisible_by`.** Long-context recipes usually need a larger divisor (e.g. `64` at 256K) so sequences align with CP/TP shapes.
100+
- **Keep override configs identical across trainer and generator.** Because `hf_config_overrides` flows into both Megatron and vLLM, editing only one side will cause mismatched RoPE behavior.
101+
- **Reconvert after changing overrides.** Cached checkpoints are keyed by override hash, so a new hash produces a new cache entry — but if you deliberately mutate an existing cache path, use `force_reconvert_from_hf: true`.

docs/index.md

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -121,6 +121,13 @@ Configure offline and online Eagle3 draft-model workflows to accelerate rollout
121121
Learn how to add support for new model architectures in NeMo RL.
122122
:::
123123

124+
:::{grid-item-card} {octicon}`arrow-both` YaRN Long-Context Training
125+
:link: guides/yarn-long-context
126+
:link-type: doc
127+
128+
Extend a model's context window with YaRN RoPE scaling on the Megatron backend for SFT, GRPO, and other workflows.
129+
:::
130+
124131
::::
125132

126133
## Advanced Topics
@@ -227,6 +234,7 @@ guides/deepseek.md
227234
model-quirks.md
228235
guides/async-grpo.md
229236
guides/eagle3-speculative-decoding.md
237+
guides/yarn-long-context.md
230238
guides/muon-optimizer.md
231239
guides/dtensor-tp-accuracy.md
232240
guides/ft-launcher-guide.md

0 commit comments

Comments
 (0)