Unify CosSinRoPE and ComplexRoPE YaRN computation (#3787)

### Background

`ComplexRoPE` and `CosSinRoPE` are the two RoPE variants in torchtitan:
`ComplexRoPE` uses an interleaved complex cache (DeepSeek V3, LLaMA 3,
Qwen3.5); `CosSinRoPE` uses a concatenated cos/sin cache with
`rotate_half` application (GPT-OSS, Qwen3). Both implement YaRN
("NTK-by-parts") for context-length extension, but they drifted apart in
following ways:

1. **`beta_fast`/`beta_slow` assignment** — `CosSinRoPE` mapped `low ←
beta_slow`, `high ← beta_fast`, **the inverse of** the YaRN paper
convention and of `ComplexRoPE`. GPT-OSS's config compensated with
swapped values, hiding the inconsistency from users.

2. **Embedded rope `mscale`** — `CosSinRoPE` multiplied the YaRN
attention-magnitude factor (`0.1·log(factor)+1`) into its cos/sin cache,
coupling the rope to the attention. GPT-OSS needed this; other
`rotate_half` consumers may not. The rope should be a pure rotation,
with mscale applied where each model's attention expects it.

This PR makes the two classes share one YaRN implementation, agreeing
exactly when both enter their respective YaRN (or non-YaRN) branches.

### Commits

**1. Align `CosSinRoPE` `beta_fast`/`beta_slow` to the YaRN paper/HF
convention**
Swap `low ← beta_slow` to `low ← beta_fast` (extrapolation), `high ←
beta_fast` to `high ← beta_slow` (interpolation). Un-swap GPT-OSS config
to the official `beta_fast=32.0, beta_slow=1.0`. Numerically neutral
(the two un-swaps cancel).

**2. Extract a shared `_yarn_inv_freq` helper + `truncate` flag**
A single source-of-truth function called by both classes. Add a
`truncate` config flag (floor/ceil for DeepSeek; fractional for GPT-OSS,
which uses `truncate=False`). Always clamps to `[0, dim-1]`.

**3. Make `CosSinRoPE` a pure rotation; let GPT-OSS handle mscale in its
attention**
Remove `mscale` from `CosSinRoPE._precompute_cache` (and the unused
`mscale` field from `RoPE.Config`). GPT-OSS folds `mscale²` into its
`softmax_scale` — mathematically equivalent. The rope is now modular:
each model's attention owns its own mscale.

### References

- YaRN paper: [arXiv:2309.00071](https://arxiv.org/abs/2309.00071) (§3.2
NTK-by-parts, §3.3 attention temperature)
- DeepSeek-V3.2:
[`precompute_freqs_cis`](https://huggingface.co/deepseek-ai/DeepSeek-V3.2/blob/c69397e/inference/model.py#L324-L402)
— `low = floor(fcd(beta_fast))`, `high = ceil(fcd(beta_slow))`
- GPT-OSS official: [`ntk_alpha=1,
ntk_beta=32`](https://github.com/openai/gpt-oss/blob/d21f34e/gpt_oss/torch/model.py#L28-L29);
[`low ← ntk_beta` in
rope](https://github.com/openai/gpt-oss/blob/d21f34e/gpt_oss/torch/model.py#L85-L131)

Author: sdmyzlp

torchtitan/models/common/rope.py

@@ -39,6 +39,49 @@ def _maybe_check_max_pos(positions: torch.Tensor, *, max_valid_pos: int) -> None
     )
 
 
+def _yarn_inv_freq(
+    dim: int,
+    base: float,
+    rope_factor: float,
+    beta_fast: float,
+    beta_slow: float,
+    original_seq_len: int,
+    truncate: bool,
+) -> torch.Tensor:
+    """Shared YaRN ("NTK-by-parts") inverse-frequency computation.
+
+    Single source of truth for both ``ComplexRoPE`` and ``CosSinRoPE`` so the
+    two cache formats are guaranteed to agree. Follows the YaRN paper / HF
+    convention: ``low <- beta_fast`` (extrapolation boundary), ``high <-
+    beta_slow`` (interpolation boundary). ``truncate`` floors/ceils the cutoffs
+    (DeepSeek style); ``truncate=False`` keeps fractional cutoffs (gpt-oss
+    style). The range is always clamped to ``[0, dim/2 - 1]``. The YaRN
+    attention "mscale" is intentionally NOT applied here -- the rope stays a
+    pure rotation and the model folds mscale into its softmax scale.
+    """
+    inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim))
+
+    def find_correction_dim(num_rotations: float) -> float:
+        return (dim * math.log(original_seq_len / (num_rotations * 2 * math.pi))) / (
+            2 * math.log(base)
+        )
+
+    low = find_correction_dim(beta_fast)
+    high = find_correction_dim(beta_slow)
+    if truncate:
+        low = math.floor(low)
+        high = math.ceil(high)
+    low, high = max(low, 0), min(high, dim - 1)
+    assert (
+        0 < low < high < dim - 1
+    ), f"Invalid YaRN params: 0 < {low} < {high} < {dim - 1}"
+
+    ramp = ((torch.arange(dim // 2, dtype=torch.float32) - low) / (high - low)).clamp(
+        0, 1
+    )
+    return inv_freq / rope_factor * ramp + inv_freq * (1 - ramp)
+
+
 class RoPE(Module):
     """Shared Rotary Position Embedding module.
 
@@ -63,7 +106,7 @@ class Config(Module.Config):
         beta_fast: float = 32.0
         beta_slow: float = 1.0
         original_seq_len: int = 4096
-        mscale: float = 0.0
+        truncate: bool = True
 
     def __init__(self, config: Config):
         super().__init__()
@@ -175,47 +218,15 @@ def _precompute_cache(self) -> torch.Tensor:
             freqs = torch.where(is_medium_freqs, smoothed_freqs, freqs)
         elif cfg.scaling == "yarn" and end > cfg.original_seq_len:
             # YaRN (DeepSeek V3 style)
-            beta_fast = cfg.beta_fast
-            beta_slow = cfg.beta_slow
-            base = theta
-            original_seq_len = cfg.original_seq_len
-            factor = cfg.rope_factor
-
-            def find_correction_dim(
-                num_rotations: float, dim: int, base: float, max_seq_len: int
-            ) -> float:
-                return (
-                    dim
-                    * math.log(max_seq_len / (num_rotations * 2 * math.pi))
-                    / (2 * math.log(base))
-                )
-
-            def find_correction_range(
-                low_rot: float,
-                high_rot: float,
-                dim: int,
-                base: float,
-                max_seq_len: int,
-            ) -> tuple[int, int]:
-                low = math.floor(find_correction_dim(low_rot, dim, base, max_seq_len))
-                high = math.ceil(find_correction_dim(high_rot, dim, base, max_seq_len))
-                return max(low, 0), min(high, dim - 1)
-
-            def linear_ramp_factor(
-                min_val: float, max_val: float, dim: int
-            ) -> torch.Tensor:
-                if min_val == max_val:
-                    max_val += 0.001
-                linear_func = (torch.arange(dim, dtype=torch.float32) - min_val) / (
-                    max_val - min_val
-                )
-                return torch.clamp(linear_func, 0, 1)
-
-            low, high = find_correction_range(
-                beta_fast, beta_slow, dim, base, original_seq_len
+            freqs = _yarn_inv_freq(
+                dim,
+                theta,
+                cfg.rope_factor,
+                cfg.beta_fast,
+                cfg.beta_slow,
+                cfg.original_seq_len,
+                cfg.truncate,
             )
-            smooth = 1 - linear_ramp_factor(low, high, dim // 2)
-            freqs = freqs / factor * (1 - smooth) + freqs * smooth
 
         t = torch.arange(end, device=freqs.device)
         freqs = torch.outer(t, freqs).float()
@@ -270,48 +281,30 @@ def _precompute_cache(self) -> torch.Tensor:
         max_seq_len = cfg.max_seq_len
         base = cfg.theta
 
-        freq = base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
-        mscale = 1.0
-
         if cfg.scaling == "llama":
             raise NotImplementedError("Cos/sin RoPE does not support Llama scaling.")
 
         if cfg.scaling == "yarn" and cfg.rope_factor > 1.0:
-            rope_factor = cfg.rope_factor
-            # YaRN mscale for attention magnitude preservation
-            mscale = 0.1 * math.log(rope_factor) + 1.0
-
-            # Compute correction range (NTK by parts)
-            d_half = dim / 2
-            low = (
-                d_half
-                * math.log(cfg.original_seq_len / (cfg.beta_slow * 2 * math.pi))
-                / math.log(base)
-            )
-            high = (
-                d_half
-                * math.log(cfg.original_seq_len / (cfg.beta_fast * 2 * math.pi))
-                / math.log(base)
+            inv_freq = _yarn_inv_freq(
+                dim,
+                base,
+                cfg.rope_factor,
+                cfg.beta_fast,
+                cfg.beta_slow,
+                cfg.original_seq_len,
+                cfg.truncate,
             )
-            assert (
-                0 < low < high < d_half - 1
-            ), f"Invalid YaRN params: 0 < {low} < {high} < {d_half - 1}"
-
-            ramp = (torch.arange(d_half, dtype=torch.float32) - low) / (high - low)
-            mask = 1 - ramp.clamp(0, 1)
-
-            interpolation = 1.0 / (rope_factor * freq)
-            extrapolation = 1.0 / freq
-            inv_freq = interpolation * (1 - mask) + extrapolation * mask
         else:
-            inv_freq = 1.0 / freq
+            inv_freq = 1.0 / (
+                base ** (torch.arange(0, dim, 2)[: (dim // 2)].float() / dim)
+            )
 
         t = torch.arange(max_seq_len, dtype=inv_freq.dtype, device=inv_freq.device)
         freqs = torch.outer(t, inv_freq).float()
         theta = torch.cat([freqs, freqs], dim=-1)
 
-        cos = theta.cos() * mscale
-        sin = theta.sin() * mscale
+        cos = theta.cos()
+        sin = theta.sin()
         return torch.cat([cos, sin], dim=-1)
 
     def _reshape_cache(

torchtitan/models/gpt_oss/__init__.py

@@ -283,8 +283,9 @@ def _debugmodel(
                 theta=150000.0,
                 scaling="yarn",
                 rope_factor=32,
-                beta_slow=32.0,
-                beta_fast=1.0,
+                beta_fast=32.0,
+                beta_slow=1.0,
+                truncate=False,
                 original_seq_len=4096,
             ),
         ),
@@ -326,8 +327,9 @@ def _20b(
                 theta=150000.0,
                 scaling="yarn",
                 rope_factor=32,
-                beta_slow=32.0,
-                beta_fast=1.0,
+                beta_fast=32.0,
+                beta_slow=1.0,
+                truncate=False,
                 original_seq_len=4096,
             ),
         ),
@@ -369,8 +371,9 @@ def _120b(
                 theta=150000.0,
                 scaling="yarn",
                 rope_factor=32,
-                beta_slow=32.0,
-                beta_fast=1.0,
+                beta_fast=32.0,
+                beta_slow=1.0,
+                truncate=False,
                 original_seq_len=4096,
             ),
         ),

torchtitan/models/gpt_oss/model.py

@@ -104,6 +104,11 @@ def __init__(self, config: Config):
 
         # Standard attention softmax scale (1/sqrt(head_dim))
         self.softmax_scale = 1.0 / math.sqrt(self.head_dim)
+        if config.rope.scaling == "yarn" and config.rope.rope_factor > 1.0:
+            mscale = 0.1 * math.log(config.rope.rope_factor) + 1.0
+            # Merge YaRN attention mscale into softmax_scale, with
+            # m**2 being equivalent to scaling q / k each by mscale.
+            self.softmax_scale *= mscale * mscale
 
         self.qkv_linear = config.qkv_linear.build()
         self.wo = config.wo.build()