|
| 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`. |
0 commit comments