Adding RoPE to TransformerEncoder

nemo/collections/asr/modules/transformer_encoder.py

@@ -24,6 +24,7 @@
     PositionalEncoding,
     RelPositionalEncoding,
     RelPositionMultiHeadAttention,
+    RotaryPositionalEncoding,
 )
 from nemo.collections.asr.parts.submodules.subsampling import FeatureStacking, StackingSubsampling
 from nemo.core.classes.module import freeze, unfreeze
@@ -75,6 +76,13 @@ class TransformerEncoderConfig:
             - ``"no_pos"`` (or ``None``): no positional encoding at all. The pre-encoder output
               is consumed directly by the Transformer blocks. ``xscaling``, ``pos_emb_max_len``,
               ``dropout_pre_encoder`` and ``dropout_emb`` are unused in this mode.
+            - ``"rope"``: rotary position embedding applied to Q and K inside attention. No
+              additive positional embedding is added to the embeddings; the standard scaled
+              dot-product attention runs through FlexAttention with ``score_mod=None``.
+        rope_base: Theta base for the rotary position embedding. Only used when
+            ``self_attention_model='rope'``.
+        rotary_fraction: Fraction of the per-head dim to rotate. Only used when
+            ``self_attention_model='rope'``.
     """
 
     feat_in: int = 128
@@ -91,6 +99,8 @@ class TransformerEncoderConfig:
     # Future: "lookahead", "local", "sliding_window".
     attn_mode: str = "full"
     self_attention_model: str = "rel_pos"
+    rope_base: float = 10000.0
+    rotary_fraction: float = 1.0
 
 
 def _make_padding_mod(lengths):
@@ -112,7 +122,7 @@ def causal(b, h, q_idx, kv_idx):
 
 
 _SUPPORTED_ATTENTION_MODES = ("full", "causal")
-_SUPPORTED_SELF_ATTENTION_MODELS = ("abs_pos", "rel_pos", "no_pos")
+_SUPPORTED_SELF_ATTENTION_MODELS = ("abs_pos", "rel_pos", "no_pos", "rope")
 
 
 class FeedForward(nn.Module):
@@ -132,13 +142,14 @@ def forward(self, x):
 
 
 class MultiHeadAttention(nn.Module):
-    def __init__(self, cfg: TransformerEncoderConfig):
+    def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None):
         super().__init__()
         self.n_heads = cfg.n_heads
         self.head_dim = cfg.d_model // cfg.n_heads
         self.d_model = cfg.d_model
         self.self_attention_model = cfg.self_attention_model
         self._uses_rel_pos = self.self_attention_model == "rel_pos"
+        self._uses_rope = self.self_attention_model == "rope"
         if self.self_attention_model not in _SUPPORTED_SELF_ATTENTION_MODELS:
             raise ValueError(
                 f"self_attention_model='{self.self_attention_model}' is not supported. "
@@ -150,6 +161,15 @@ def __init__(self, cfg: TransformerEncoderConfig):
                 f"but got head_dim={self.head_dim} from d_model={self.d_model}, n_heads={self.n_heads}."
             )
 
+        # Rotary position embedding shared across layers; rotates Q/K before the attention
+        # kernel. The shared module owns the cos/sin buffers (see ``TransformerEncoder``).
+        if self._uses_rope:
+            if pos_enc is None:
+                raise ValueError("'rope' attention requires a RotaryPositionalEncoding via pos_enc.")
+            self.rope = pos_enc
+        else:
+            self.rope = None
+
         self.w_qkv = nn.Linear(cfg.d_model, 3 * cfg.d_model, bias=cfg.qkv_bias)
         self.out_proj = nn.Linear(cfg.d_model, cfg.d_model)
 
@@ -212,9 +232,8 @@ def _build_rel_pos_score_mod(self, q, pos_emb):
         p = self.linear_pos(pos_emb).view(pos_emb.size(0), -1, H, D).transpose(1, 2)
         # pos_bias_{u,v}: (H, D) -> (1, H, 1, D) so they broadcast over the (B, H, T, D)
         # Q tensor against the head/depth axes rather than (incorrectly) against time.
-        # Match dtype under AMP so fp32 bias params do not upcast q before FlexAttention.
-        bias_u = self.pos_bias_u.view(1, H, 1, D).to(dtype=q.dtype)
-        bias_v = self.pos_bias_v.view(1, H, 1, D).to(dtype=q.dtype)
+        bias_u = self.pos_bias_u.view(1, H, 1, D).to(q.dtype)
+        bias_v = self.pos_bias_v.view(1, H, 1, D).to(q.dtype)
         # Matrix b + d: ((Q + v) @ R^T) shifted into (q_idx, kv_idx) space, then scaled
         # by 1/sqrt(D) so it can be added directly to FlexAttention's already-scaled scores.
         q_with_bias_v = q + bias_v
@@ -240,6 +259,10 @@ def forward(self, x, block_mask=None, pos_emb=None):
             q = self.q_norm(q).to(v.dtype)
             k = self.k_norm(k).to(v.dtype)
 
+        if self._uses_rope:
+            # RoPE rotates Q/K in place; it is orthogonal to FlexAttention's score_mod.
+            q, k = self.rope(q, k)
+
         score_mod = None
         if self._uses_rel_pos:
             score_mod, q = self._build_rel_pos_score_mod(q, pos_emb)
@@ -251,10 +274,10 @@ def forward(self, x, block_mask=None, pos_emb=None):
 
 
 class TransformerBlock(nn.Module):
-    def __init__(self, cfg: TransformerEncoderConfig):
+    def __init__(self, cfg: TransformerEncoderConfig, pos_enc=None):
         super().__init__()
         self.norm1 = nn.LayerNorm(cfg.d_model)
-        self.attn = MultiHeadAttention(cfg)
+        self.attn = MultiHeadAttention(cfg, pos_enc=pos_enc)
         self.drop = nn.Dropout(cfg.drop_rate)
         self.norm2 = nn.LayerNorm(cfg.d_model)
         self.ffn = FeedForward(cfg)
@@ -324,8 +347,16 @@ class TransformerEncoder(nn.Module):
               ``pos_emb_max_len``, ``dropout_pre_encoder`` and ``dropout_emb`` have no effect
               in this mode. ``None`` is accepted as a YAML-friendly alias for ``"no_pos"``
               (an unset field in a config maps to ``None``).
+            - ``"rope"``: rotary position embedding applied to Q/K inside attention. No additive
+              positional embedding is added to the embeddings; ``dropout_pre_encoder`` (and
+              ``xscaling`` if set) are still applied to the pre-encoder output, and ``dropout_emb``
+              has no effect in this mode.
 
             ``"rel_pos_local_attn"`` is not implemented yet.
+        rope_base: Theta base for the rotary position embedding. Only used when
+            ``self_attention_model='rope'``. Defaults to 10000.0.
+        rotary_fraction: Fraction of the per-head dim to rotate. Only used when
+            ``self_attention_model='rope'``. Defaults to 1.0.
         pos_emb_max_len: Initial maximum length for sinusoidal positional embeddings.
         xscaling: If True, scale embeddings by ``sqrt(d_model)`` before adding positional encodings,
             following "Attention Is All You Need" article. Originally intended to balance the magnitude
@@ -355,6 +386,8 @@ def __init__(
         ff_expansion: float = 4.0,
         pre_block_norm: bool = True,
         self_attention_model: Optional[str] = "rel_pos",
+        rope_base: float = 10000.0,
+        rotary_fraction: float = 1.0,
         pos_emb_max_len: int = 5000,
         xscaling: bool = False,
         attn_mode: str = "full",
@@ -375,7 +408,7 @@ def __init__(
         if self_attention_model not in _SUPPORTED_SELF_ATTENTION_MODELS:
             raise ValueError(
                 f"self_attention_model='{self_attention_model}' is not supported. "
-                "Currently only 'abs_pos', 'rel_pos', and 'no_pos' (or None) are available."
+                "Currently only 'abs_pos', 'rel_pos', 'rope', and 'no_pos' (or None) are available."
             )
         if dropout_pre_encoder is None:
             dropout_pre_encoder = drop_rate
@@ -393,6 +426,8 @@ def __init__(
             subsampling_factor=subsampling_factor,
             attn_mode=attn_mode,
             self_attention_model=self_attention_model,
+            rope_base=rope_base,
+            rotary_fraction=rotary_fraction,
         )
         self.d_model = d_model
         self.n_layers = n_layers
@@ -443,10 +478,23 @@ def __init__(
                 xscale=self.xscale,
                 dropout_rate_emb=dropout_emb,
             )
+        elif self_attention_model == "rope":
+            # RoPE has no additive pos-emb step to host dropout, so pre-encoder
+            # dropout is applied separately in forward_internal.
+            self.dropout_pre_encoder = nn.Dropout(dropout_pre_encoder)
+            self.pos_enc = RotaryPositionalEncoding(
+                d_k=d_model // n_heads,
+                rotary_fraction=rotary_fraction,
+                rope_base=rope_base,
+                max_len=pos_emb_max_len,
+            )
         else:  # "no_pos"
             self.pos_enc = None
         self.embed_norm = nn.LayerNorm(d_model) if pre_block_norm else nn.Identity()
-        self.layers = nn.ModuleList([TransformerBlock(cfg) for _ in range(n_layers)])
+        # For 'rope', the shared RotaryPositionalEncoding is passed into each block so the
+        # cos/sin buffers are computed once and reused across all attention modules.
+        layer_pos_enc = self.pos_enc if self_attention_model == "rope" else None
+        self.layers = nn.ModuleList([TransformerBlock(cfg, pos_enc=layer_pos_enc) for _ in range(n_layers)])
         self.final_norm = nn.LayerNorm(d_model)
         if feat_out > 0 and feat_out != self._feat_out:
             self.out_proj = nn.Linear(self._feat_out, feat_out)
@@ -510,7 +558,13 @@ def forward_internal(self, audio_signal, length, bypass_pre_encode=False):
             x = audio_signal
             length = length.to(torch.int64)
 
-        if self.pos_enc is not None:
+        if self.self_attention_model == "rope":
+            # RoPE: no pos emb added; just apply xscale (if set) + pre-encoder dropout here.
+            if self.xscale:
+                x = x * self.xscale
+            x = self.dropout_pre_encoder(x)
+            pos_emb = None
+        elif self.pos_enc is not None:
             x, pos_emb = self.pos_enc(x=x)
         else:  # "no_pos": pre-encoder output flows in unchanged
             pos_emb = None

tests/collections/asr/test_transformer_encoder.py

@@ -21,6 +21,7 @@
     TransformerEncoder,
     TransformerEncoderConfig,
 )
+from nemo.collections.asr.parts.submodules.multi_head_attention import RotaryPositionalEncoding
 
 
 class TestTransformerEncoderConfig:
@@ -39,6 +40,8 @@ def test_default_config(self):
         assert cfg.subsampling_factor == 4
         assert cfg.attn_mode == "full"
         assert cfg.self_attention_model == "rel_pos"
+        assert cfg.rope_base == 10000.0
+        assert cfg.rotary_fraction == 1.0
 
     @pytest.mark.unit
     def test_custom_config(self):
@@ -496,7 +499,7 @@ def test_default_is_rel_pos(self):
         assert model.self_attention_model == "rel_pos"
 
     @pytest.mark.unit
-    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos"])
+    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos", "rope"])
     def test_valid_modes_are_accepted(self, mode):
         model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model=mode)
         assert model.self_attention_model == mode
@@ -533,9 +536,9 @@ def test_rel_pos_attention_params_allocated(self):
             assert attn.pos_bias_v.shape == (n_heads, head_dim)
 
     @pytest.mark.unit
-    @pytest.mark.parametrize("mode", ["abs_pos", "no_pos"])
+    @pytest.mark.parametrize("mode", ["abs_pos", "no_pos", "rope"])
     def test_non_rel_pos_modes_have_no_rel_params(self, mode):
-        """abs_pos and no_pos modes must not allocate the rel-pos parameters."""
+        """abs_pos, no_pos and rope modes must not allocate the rel-pos parameters."""
         model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=2, self_attention_model=mode)
         for layer in model.layers:
             attn = layer.attn
@@ -552,7 +555,7 @@ def test_no_pos_has_no_positional_encoding_module(self):
         assert model.max_audio_length == model.pos_emb_max_len
 
     @pytest.mark.unit
-    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos", None])
+    @pytest.mark.parametrize("mode", ["abs_pos", "rel_pos", "no_pos", "rope", None])
     def test_forward_each_mode_cpu(self, mode):
         """Each ``self_attention_model`` choice (including ``None``) must produce a valid forward."""
         model = TransformerEncoder(
@@ -603,3 +606,42 @@ def test_rel_pos_broadcasts_when_T_differs_from_n_heads(self):
 
         assert out.shape == (B, 64, T // 4)
         assert not torch.isnan(out).any()
+
+    @pytest.mark.unit
+    def test_rope_uses_shared_rotary_pos_enc(self):
+        """rope mode builds a single ``RotaryPositionalEncoding`` reused by every attention layer.
+
+        The cos/sin buffers are computed once on the shared module (see ``TransformerEncoder``),
+        so each layer's ``attn.rope`` must be the *same* object as ``model.pos_enc``.
+        """
+        model = TransformerEncoder(feat_in=128, d_model=64, n_heads=4, n_layers=3, self_attention_model="rope")
+        assert isinstance(model.pos_enc, RotaryPositionalEncoding)
+        for layer in model.layers:
+            attn = layer.attn
+            assert attn._uses_rope is True
+            assert attn.rope is model.pos_enc
+
+    @pytest.mark.unit
+    def test_rope_partial_rotation_forward_cpu(self):
+        """``rotary_fraction`` < 1.0 rotates only part of each head dim (exercises the pass-through split)."""
+        model = TransformerEncoder(
+            feat_in=128,
+            d_model=64,
+            n_heads=4,
+            n_layers=2,
+            drop_rate=0.0,
+            subsampling_factor=4,
+            self_attention_model="rope",
+            rotary_fraction=0.5,
+        )
+        model.eval()
+
+        B, C, T = 2, 128, 200
+        x = torch.randn(B, C, T)
+        lengths = torch.tensor([T, 160])
+
+        with torch.no_grad():
+            out, _ = model(audio_signal=x, length=lengths)
+
+        assert out.shape == (B, 64, T // 4)
+        assert not torch.isnan(out).any()