Pin SDPA math backend (bit-reproducible fp32 oracle); device-aware rotary cache; drop dead dropout

modeling_mdlm.py

@@ -7,7 +7,9 @@ that removes the hard dependency on ``flash-attn``. The only substantive changes
   * Attention uses ``torch.nn.functional.scaled_dot_product_attention`` (non-causal,
     full bidirectional) instead of ``flash_attn_varlen_qkvpacked_func``. The two are
     mathematically equivalent here: a single sequence per batch element, no causal
-    mask, default softmax scale ``1/sqrt(head_dim)``.
+    mask, default softmax scale ``1/sqrt(head_dim)``. The call is pinned to the **math**
+    SDPA backend (``sdpa_kernel(SDPBackend.MATH)``) so the fp32 oracle is bit-reproducible
+    across GPUs, rather than silently dispatching to a fused kernel with different rounding.
   * Rotary position embeddings are applied with an explicit ``rotate_half`` instead of
     ``flash_attn.layers.rotary``. This reproduces flash-attn's non-interleaved
     convention exactly (verified: ``q*cos + rotate_half(q)*sin`` with full-dim cos/sin).
@@ -36,6 +38,7 @@ import typing
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+from torch.nn.attention import SDPBackend, sdpa_kernel
 import transformers
 from transformers import modeling_outputs
 
@@ -86,7 +89,10 @@ class Rotary(nn.Module):
 
   def forward(self, x, seq_dim=1):
     seq_len = x.shape[seq_dim]
-    if seq_len != self.seq_len_cached:
+    # Rebuild on a sequence-length OR device change: cos/sin are plain attributes (not
+    # buffers), so `.to(device)` does not move them -- key the cache on device too.
+    if (seq_len != self.seq_len_cached or self.cos_cached is None
+        or self.cos_cached.device != x.device):
       self.seq_len_cached = seq_len
       t = torch.arange(seq_len, device=x.device).type_as(self.inv_freq)
       freqs = torch.einsum("i,j->ij", t, self.inv_freq)
@@ -197,16 +203,15 @@ class DDiTBlock(nn.Module):
     self.norm1 = LayerNorm(dim)
     self.attn_qkv = nn.Linear(dim, 3 * dim, bias=False)
     self.attn_out = nn.Linear(dim, dim, bias=False)
-    self.dropout1 = nn.Dropout(dropout)
 
     self.norm2 = LayerNorm(dim)
     self.mlp = nn.Sequential(
       nn.Linear(dim, mlp_ratio * dim, bias=True),
       nn.GELU(approximate='tanh'),
       nn.Linear(mlp_ratio * dim, dim, bias=True))
-    self.dropout2 = nn.Dropout(dropout)
-    self.dropout = dropout
 
+    # Inference-only oracle: the upstream training-time residual dropout is intentionally
+    # omitted (it is identity at eval). To *resume training*, use the official checkpoint.
     self.adaLN_modulation = nn.Linear(cond_dim, 6 * dim, bias=True)
 
   def forward(self, x, rotary_cos_sin, c, seqlens=None):
@@ -227,7 +232,9 @@ class DDiTBlock(nn.Module):
     k = apply_rotary(k, cos, sin)
     # SDPA wants [b, h, s, d]; attention is full bidirectional (no causal mask).
     q, k, v = q.transpose(1, 2), k.transpose(1, 2), v.transpose(1, 2)
-    attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
+    # Pin the math backend so the fp32 oracle is bit-reproducible across GPUs.
+    with sdpa_kernel(SDPBackend.MATH):
+      attn = F.scaled_dot_product_attention(q, k, v, is_causal=False)
     attn = attn.transpose(1, 2).reshape(batch_size, seq_len, -1)
 
     x = x_skip + gate_msa * self.attn_out(attn)