Share equivalent RoPE cache storage

torchtitan/distributed/pipeline_parallel.py

@@ -488,7 +488,11 @@ def _split_module(
     model = copy.deepcopy(whole_model)
     # Create a set of modules to keep for faster lookup
     modules_to_keep = set(module_names)
+    if "rope_modules" in model._modules:
+        modules_to_keep.add("rope_modules")
     for module_name, module_value in model.named_children():
+        if module_name in modules_to_keep:
+            continue
         # Handle layer-like structures (e.g., "layers.0", "layers.1")
         if isinstance(
             module_value, (nn.ModuleDict, nn.ModuleList, ModuleDict, ModuleList)

torchtitan/models/common/__init__.py

@@ -36,7 +36,11 @@
     RMSNorm,
     SiLU,
 )
-from .rope import ComplexRoPE, CosSinRoPE, RoPE
+from .rope import (
+    ComplexRoPE,
+    CosSinRoPE,
+    RoPE,
+)
 
 __all__ = [
     "Conv1d",

torchtitan/models/common/attention.py

@@ -13,7 +13,7 @@
 
 from collections.abc import Callable, Mapping
 from dataclasses import dataclass, field
-from typing import Any, ClassVar, NamedTuple
+from typing import Any, ClassVar, NamedTuple, cast
 
 import spmd_types as spmd
 import torch
@@ -43,7 +43,7 @@
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.rope import RoPE
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 from torchtitan.tools.utils import round_up
 
 
@@ -667,6 +667,10 @@ def __post_init__(self):
             assert self.n_heads > 0, "n_heads must be > 0"
 
 
+def _resolve_rope(rope_config: RoPE.Config, rope_modules: ModuleDict) -> RoPE:
+    return cast(RoPE, rope_modules[rope_config.rope_key()])
+
+
 class BaseQKVLinear(Module):
     """Base class for Q/K/V projection strategies.
 
@@ -892,6 +896,8 @@ class GQAttention(BaseAttention):
     :class:`FusedQKVLinear` for a single fused projection.
     """
 
+    rope: RoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int
@@ -919,7 +925,7 @@ def __post_init__(self) -> None:
                     f"n_kv_heads ({n_kv_heads})"
                 )
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.n_heads = config.n_heads
         self.n_kv_heads = (
@@ -931,7 +937,8 @@ def __init__(self, config: Config):
             else config.dim // config.n_heads
         )
         self.enable_gqa = self.n_heads > self.n_kv_heads
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
         # Pluggable QKV projection
         self.qkv_linear = config.qkv_linear.build()

torchtitan/models/common/decoder.py

@@ -9,19 +9,19 @@
 from dataclasses import dataclass
 
 import torch
-from torch.nn.attention.flex_attention import _mask_mod_signature, and_masks, BlockMask
+from torch.nn.attention.flex_attention import BlockMask, _mask_mod_signature, and_masks
 
 from torchtitan.distributed.utils import is_in_batch_invariant_mode
 from torchtitan.models.common.attention import (
     AttentionMasksType,
     BaseAttention,
+    FlexAttention,
+    ScaledDotProductAttention,
+    VarlenAttention,
     create_attention_mask,
     create_varlen_metadata_for_document,
-    FlexAttention,
     get_causal_mask_mod,
     get_efficient_causal_mask_mod_for_packed_document,
-    ScaledDotProductAttention,
-    VarlenAttention,
 )
 from torchtitan.models.common.embedding import Embedding
 from torchtitan.models.common.feed_forward import FeedForward
@@ -35,6 +35,22 @@
 __all__ = ["Decoder", "TransformerBlock"]
 
 
+def _rope_config(layer_config):
+    attention_config = getattr(layer_config, "attention", None)
+    return getattr(attention_config, "rope", None)
+
+
+def _register_rope_modules(layer_configs, rope_modules: ModuleDict) -> None:
+    """Build one canonical RoPE module for each layer configuration key."""
+    for layer_config in layer_configs:
+        rope_config = _rope_config(layer_config)
+        if rope_config is None:
+            continue
+        key = rope_config.rope_key()
+        if key not in rope_modules:
+            rope_modules[key] = rope_config.build()
+
+
 # TODO: we can unify the TransformerBlock impl across all models when
 # there is no special logic for each model, including
 # ffn vs. moe naming and creation, etc.
@@ -248,12 +264,19 @@ def update_from_config(
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
-
         self.tok_embeddings = config.tok_embeddings.build()
 
+        self.rope_modules = ModuleDict()
+        _register_rope_modules(config.layers, self.rope_modules)
+
         self.layers = ModuleDict()
         for i, layer_config in enumerate(config.layers):
-            self.layers[str(i)] = layer_config.build()
+            rope_config = _rope_config(layer_config)
+            if rope_config is None:
+                layer = layer_config.build()
+            else:
+                layer = layer_config.build(rope_modules=self.rope_modules)
+            self.layers[str(i)] = layer
 
         self.norm = config.norm.build()
         self.lm_head = config.lm_head.build()

torchtitan/models/common/rope.py

@@ -5,7 +5,8 @@
 # LICENSE file in the root directory of this source tree.
 
 import math
-from dataclasses import dataclass
+import re
+from dataclasses import dataclass, fields
 from typing import Literal
 
 import spmd_types as spmd
@@ -21,6 +22,21 @@
 ]
 
 
+def _format_rope_key_value(value: object) -> str:
+    """Format config values into safe, readable ``ModuleDict`` key fragments."""
+    if value is None:
+        text = "none"
+    elif isinstance(value, bool):
+        text = "true" if value else "false"
+    elif isinstance(value, float):
+        text = format(value, ".17g")
+    elif isinstance(value, (list, tuple)):
+        text = "-".join(_format_rope_key_value(item) for item in value)
+    else:
+        text = str(value)
+    return re.sub(r"[^A-Za-z0-9_-]", "p", text)
+
+
 # pyrefly: ignore [not-callable]
 @spmd.no_typecheck()
 def _maybe_check_max_pos(positions: torch.Tensor, *, max_valid_pos: int) -> None:
@@ -89,6 +105,8 @@ class RoPE(Module):
     cosine/sine caches.
     """
 
+    cache: torch.Tensor
+
     @dataclass(kw_only=True, slots=True)
     class Config(Module.Config):
         dim: int
@@ -107,6 +125,66 @@ class Config(Module.Config):
         original_seq_len: int = 4096
         truncate: bool = True
 
+        def rope_key(self) -> str:
+            """Return the stable, descriptive key for this RoPE implementation."""
+            owner = self._owner
+            owner_name = owner.__name__ if owner is not None else type(self).__name__
+            parts = [owner_name]
+            common_fields = {"dim", "max_context_length", "theta", "scaling"}
+            scaling_fields = {
+                "llama": {
+                    "scaling_factor",
+                    "low_freq_factor",
+                    "high_freq_factor",
+                    "original_max_position_embeddings",
+                },
+                "yarn": {
+                    "rope_factor",
+                    "beta_fast",
+                    "beta_slow",
+                    "original_seq_len",
+                    "truncate",
+                },
+            }.get(self.scaling, set())
+            key_names = {
+                "dim": "d",
+                "max_context_length": "ctx",
+                "scaling_factor": "sf",
+                "low_freq_factor": "lf",
+                "high_freq_factor": "hf",
+                "original_max_position_embeddings": "orig",
+                "rope_factor": "rf",
+                "beta_fast": "bf",
+                "beta_slow": "bs",
+                "original_seq_len": "orig",
+                "truncate": "trunc",
+            }
+            for config_field in fields(self):
+                if config_field.name in {"param_init", "sharding_config"}:
+                    continue
+                if (
+                    config_field.name not in common_fields
+                    and config_field.name not in scaling_fields
+                    and config_field.name
+                    in {
+                        "scaling_factor",
+                        "low_freq_factor",
+                        "high_freq_factor",
+                        "original_max_position_embeddings",
+                        "rope_factor",
+                        "beta_fast",
+                        "beta_slow",
+                        "original_seq_len",
+                        "truncate",
+                    }
+                ):
+                    continue
+                value = _format_rope_key_value(getattr(self, config_field.name))
+                name = key_names.get(config_field.name, config_field.name)
+                separator = "" if config_field.name in common_fields - {"scaling"} else "_"
+                parts.append(f"{name}{separator}{value}")
+            return "_".join(parts)
+
     def __init__(self, config: Config):
         super().__init__()
         self.config = config
@@ -167,13 +245,13 @@ def forward(
         return self.apply_rotary_emb(query, key, reshaped_cache)
 
     def _init_self_buffers(self, *, buffer_device: torch.device | None = None) -> None:
-        # TODO: In long-term we need to have buffer abstraction in `Module`` class to infer the buffer_device
         if buffer_device is None:
             # After ``to_empty()``, the existing cache records the target device.
             # Recompute there when the caller does not pass an explicit buffer device.
             buffer_device = self.cache.device
         with torch.device(buffer_device):
-            self.cache = self._precompute_cache()
+            cache = self._precompute_cache()
+            self.register_buffer("cache", cache, persistent=False)
 
 
 class ComplexRoPE(RoPE):

torchtitan/models/deepseek_v3/model.py

@@ -15,14 +15,15 @@
     AttentionMasksType,
     BaseAttention,
     FlexAttention,
+    _resolve_rope,
 )
 from torchtitan.models.common.decoder import TransformerBlock
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
 from torchtitan.models.common.rope import RoPE
 from torchtitan.models.deepseek_v3.mtp import MTPDecoder
 from torchtitan.models.utils import get_moe_model_nparams_and_flops
-from torchtitan.protocols.module import Module
+from torchtitan.protocols.module import Module, ModuleDict
 
 
 class Attention(BaseAttention):
@@ -32,6 +33,8 @@ class Attention(BaseAttention):
     This is DeepSeek V3-specific and NOT shared with other models.
     """
 
+    rope: RoPE
+
     @dataclass(kw_only=True, slots=True)
     class Config(BaseAttention.Config):
         n_heads: int
@@ -53,7 +56,7 @@ class Config(BaseAttention.Config):
         inner_attention: Module.Config = field(default_factory=FlexAttention.Config)
         mscale: float = 1.0
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
         self.dim = config.dim
         self.n_heads = config.n_heads
@@ -88,7 +91,8 @@ def __init__(self, config: Config):
             self.softmax_scale = self.softmax_scale * mscale * mscale
 
         self.inner_attention = config.inner_attention.build()
-        self.rope = config.rope.build()
+        # Keep the canonical module registered only under Decoder.rope_modules.
+        object.__setattr__(self, "rope", _resolve_rope(config.rope, rope_modules))
 
     def forward(
         self,
@@ -160,9 +164,9 @@ class DeepSeekV3TransformerBlock(TransformerBlock):
     class Config(TransformerBlock.Config):
         pass
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         self.attention_norm = config.attention_norm.build()
         self.ffn_norm = config.ffn_norm.build()
 

torchtitan/models/deepseek_v3/mtp.py

@@ -23,10 +23,14 @@
 from torchtitan.config import CompileConfig
 from torchtitan.distributed.fsdp import apply_fsdp_to_decoder
 from torchtitan.models.common.attention import AttentionMasksType
-from torchtitan.models.common.decoder import Decoder, TransformerBlock
+from torchtitan.models.common.decoder import (
+    Decoder,
+    TransformerBlock,
+    _register_rope_modules,
+)
 from torchtitan.models.common.linear import Linear
 from torchtitan.models.common.nn_modules import RMSNorm
-from torchtitan.protocols.module import ModuleList
+from torchtitan.protocols.module import ModuleDict, ModuleList
 
 
 def roll_mtp_sequence(
@@ -120,9 +124,9 @@ class Config(TransformerBlock.Config):
         eh_proj: Linear.Config
         mtp_norm: RMSNorm.Config
 
-    def __init__(self, config: Config):
+    def __init__(self, config: Config, *, rope_modules: ModuleDict):
         super().__init__()
-        self.attention = config.attention.build()
+        self.attention = config.attention.build(rope_modules=rope_modules)
         self.attention_norm = config.attention_norm.build()
         self.ffn_norm = config.ffn_norm.build()
         self.enorm = config.enorm.build()
@@ -211,14 +215,17 @@ def __init__(self, config: Config):
             self.mtp_layers = None
             return
 
+        _register_rope_modules(config.mtp_layers, self.rope_modules)
         self.mtp_layers = ModuleList()
         for layer_config in config.mtp_layers:
             if not isinstance(layer_config, MTPTransformerBlock.Config):
                 raise ValueError(
                     "MTPDecoder requires Config.mtp_layers to contain "
                     "MTPTransformerBlock.Config instances."
                 )
-            self.mtp_layers.append(layer_config.build())
+            self.mtp_layers.append(
+                layer_config.build(rope_modules=self.rope_modules)
+            )
 
     def forward(
         self,