2026-07-11 nightly release (51c197c)

Author: pytorchbot

tests/integration_tests/models.py

@@ -18,22 +18,30 @@ def _enable_spmd_backend(t: OverrideDefinitions, backend: str) -> OverrideDefini
         variant = tuple(
             arg.replace(f"{t.test_name}/", f"{test_name}/") for arg in variant
         )
-        prefix = [f"--parallelism.spmd_backend {backend}"]
+        is_qwen3_5 = any("--module qwen3_5" in arg for arg in variant)
+        prefix = []
+        if backend != "spmd_types" or not is_qwen3_5:
+            prefix.append(f"--parallelism.spmd_backend {backend}")
         suffix = []
         # Compile, PP, and explicit AC modes are not compatible with SPMD
         # typechecking yet; keep those as backend-only coverage.
-        if backend == "spmd_types" and not any(
-            token in arg
-            for arg in variant
-            for token in (
-                "compile.enable",
-                "pipeline_parallel_degree",
-                "activation-checkpoint:",
+        if (
+            prefix
+            and backend == "spmd_types"
+            and not any(
+                token in arg
+                for arg in variant
+                for token in (
+                    "compile.enable",
+                    "pipeline_parallel_degree",
+                    "activation-checkpoint:",
+                )
             )
         ):
             prefix.append("--debug.spmd_typechecking")
             suffix.append("activation-checkpoint:none")
         new_args.append(tuple(prefix) + tuple(variant) + tuple(suffix))
+
     return dataclasses.replace(
         t,
         override_args=tuple(new_args),

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/common/token_dispatcher.py

@@ -269,7 +269,7 @@ def _token_count_exchange(
         assert self.ep_mesh is not None
         if (
             torch.compiler.is_compiling() or torch.compiler._is_non_strict_tracing()
-        ) and get_spmd_backend() != "spmd_types":
+        ) or get_spmd_backend() != "spmd_types":
             return all_to_all_single(
                 num_local_tokens_per_expert_E.view(ep_size, -1),
                 None,
@@ -336,7 +336,7 @@ def _dispatch_token_exchange(
         assert self.ep_mesh is not None
         if (
             torch.compiler.is_compiling() or torch.compiler._is_non_strict_tracing()
-        ) and get_spmd_backend() != "spmd_types":
+        ) or get_spmd_backend() != "spmd_types":
             return all_to_all_single(
                 routed_input_ND,
                 output_splits,
@@ -364,7 +364,7 @@ def _combine_token_exchange(
         assert self.ep_mesh is not None
         if (
             torch.compiler.is_compiling() or torch.compiler._is_non_strict_tracing()
-        ) and get_spmd_backend() != "spmd_types":
+        ) or get_spmd_backend() != "spmd_types":
             return all_to_all_single(
                 routed_output_RD,
                 input_splits,

torchtitan/models/deepseek_v3/model.py

@@ -132,7 +132,7 @@ def forward(
                 kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1
             )
             k = torch.cat([k_nope, k_pe.expand(-1, -1, k_nope.size(2), -1)], dim=-1)
-            if get_spmd_backend() == "spmd_types":
+            if get_spmd_backend() == "spmd_types" and not torch.compiler.is_compiling():
                 for t in [k, v]:
                     spmd.assert_type(
                         t,

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()