Fix model correctness and Hugging Face compatibility

README.md

@@ -152,8 +152,7 @@ even at inference, consistent with the original training code.
 ## Credits
 
 Original model and code by Penić et al. Source: [GitHub lbcb-sci/RiNALMo](https://github.com/lbcb-sci/RiNALMo).
-The HF conversion code was authored primarily by [Claude Code](https://claude.ai/code)
-and reviewed manually by Taykhoom Dalal.
+Hugging Face port maintained by Taykhoom Dalal.
 
 ## License
 

modeling_rinalmo.py

@@ -18,25 +18,51 @@ def _rotate_half(x):
 
 
 def _apply_rotary_pos_emb(q, k, cos, sin):
-    cos = cos.to(q.dtype)
-    sin = sin.to(q.dtype)
+    cos = cos.to(device=q.device, dtype=q.dtype)
+    sin = sin.to(device=q.device, dtype=q.dtype)
     return (q * cos) + (_rotate_half(q) * sin), (k * cos) + (_rotate_half(k) * sin)
 
 
 class RotaryPositionEmbedding(nn.Module):
     def __init__(self, dim: int, base: int = 10000):
         super().__init__()
+        self.dim = dim
+        self.base = base
         inv_freq = 1.0 / (base ** (torch.arange(0, dim, 2).float() / dim))
         self.register_buffer("inv_freq", inv_freq)
         self._seq_len_cached = None
+        self._device_cached = None
         self._cos_cached = None
         self._sin_cached = None
 
     def _update_cache(self, seq_len: int, device, dtype):
-        if seq_len != self._seq_len_cached:
+        if (
+            seq_len != self._seq_len_cached
+            or device != self._device_cached
+        ):
             self._seq_len_cached = seq_len
-            t = torch.arange(seq_len, device=device).type_as(self.inv_freq)
-            freqs = torch.einsum("i,j->ij", t, self.inv_freq)
+            self._device_cached = device
+            if (
+                self.inv_freq.device == device
+                and self.inv_freq.dtype == torch.float32
+            ):
+                inv_freq = self.inv_freq
+            else:
+                inv_freq = 1.0 / (
+                    self.base
+                    ** (
+                        torch.arange(
+                            0,
+                            self.dim,
+                            2,
+                            dtype=torch.float32,
+                            device=device,
+                        )
+                        / self.dim
+                    )
+                )
+            t = torch.arange(seq_len, device=device, dtype=torch.float32)
+            freqs = torch.einsum("i,j->ij", t, inv_freq)
             emb = torch.cat((freqs, freqs), dim=-1)
             self._cos_cached = emb.cos()[None, None, :, :]
             self._sin_cached = emb.sin()[None, None, :, :]
@@ -81,11 +107,11 @@ class RiNALMoAttention(nn.Module):
         if key_padding_mask is not None:
             attn = attn.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"))
 
-        attn = attn.softmax(dim=-1)
+        attn = torch.nan_to_num(attn.float().softmax(dim=-1), nan=0.0)
         attn_weights = attn if output_attentions else None
         attn = self.attn_dropout(attn)
 
-        out = torch.matmul(attn, v)
+        out = torch.matmul(attn.to(v.dtype), v)
         out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
         out = self.out_proj(out)
         return out, attn_weights
@@ -113,7 +139,13 @@ class RiNALMoSdpaAttention(RiNALMoAttention):
             attn_mask = torch.zeros(B, 1, 1, T, dtype=q.dtype, device=q.device)
             attn_mask = attn_mask.masked_fill(key_padding_mask.unsqueeze(1).unsqueeze(2), float("-inf"))
 
-        out = F.scaled_dot_product_attention(q, k, v, attn_mask=attn_mask, dropout_p=0.0)
+        out = F.scaled_dot_product_attention(
+            q,
+            k,
+            v,
+            attn_mask=attn_mask,
+            dropout_p=self.attn_dropout.p if self.training else 0.0,
+        )
         out = out.transpose(1, 2).contiguous().view(B, T, self.embed_dim)
         out = self.out_proj(out)
         return out, None
@@ -149,11 +181,11 @@ class RiNALMoFlashAttention2(RiNALMoAttention):
             q = q_t.transpose(1, 2)
             k = k_t.transpose(1, 2)
 
-        orig_dtype = q.dtype
         if q.dtype not in (torch.float16, torch.bfloat16):
-            q = q.to(torch.bfloat16)
-            k = k.to(torch.bfloat16)
-            v = v.to(torch.bfloat16)
+            raise ValueError(
+                "flash_attention_2 requires float16 or bfloat16 weights. "
+                f"Received {q.dtype}."
+            )
 
         if key_padding_mask is not None and key_padding_mask.any():
             attend_mask = ~key_padding_mask
@@ -165,14 +197,20 @@ class RiNALMoFlashAttention2(RiNALMoAttention):
                 q_unpad, k_unpad, v_unpad,
                 cu_seqlens_q=cu_seqlens, cu_seqlens_k=cu_seqlens,
                 max_seqlen_q=max_seqlen, max_seqlen_k=max_seqlen,
+                dropout_p=self.attn_dropout.p if self.training else 0.0,
                 causal=False,
             )
             out = pad_input(out_unpad.view(-1, self.embed_dim), indices, B, T)
         else:
-            out = flash_attn_func(q, k, v, causal=False)
+            out = flash_attn_func(
+                q,
+                k,
+                v,
+                dropout_p=self.attn_dropout.p if self.training else 0.0,
+                causal=False,
+            )
             out = out.view(B, T, self.embed_dim)
 
-        out = out.to(orig_dtype)
         out = self.out_proj(out)
         return out, None
 
@@ -312,6 +350,14 @@ class RiNALMoModel(RiNALMoPreTrainedModel):
 
         x = self.final_layer_norm(x)
 
+        if not return_dict:
+            output = (x,)
+            if output_hidden_states:
+                output += (tuple(all_hidden_states),)
+            if output_attentions:
+                output += (tuple(all_attentions),)
+            return output
+
         return BaseModelOutput(
             last_hidden_state=x,
             hidden_states=tuple(all_hidden_states) if output_hidden_states else None,
@@ -338,15 +384,25 @@ class RiNALMoForMaskedLM(RiNALMoPreTrainedModel):
         return_dict = return_dict if return_dict is not None else self.config.use_return_dict
         out = self.model(input_ids, attention_mask=attention_mask,
                          output_hidden_states=output_hidden_states,
-                         output_attentions=output_attentions, return_dict=return_dict)
+                         output_attentions=output_attentions, return_dict=True)
         logits = self.lm_head(out.last_hidden_state)
         loss = None
         if labels is not None:
             loss = F.cross_entropy(logits.view(-1, self.config.vocab_size),
                                    labels.view(-1), ignore_index=-100)
-        return MaskedLMOutput(loss=loss, logits=logits,
-                              hidden_states=out.hidden_states,
-                              attentions=out.attentions)
+        if not return_dict:
+            output = (logits,)
+            if output_hidden_states:
+                output += (out.hidden_states,)
+            if output_attentions:
+                output += (out.attentions,)
+            return (loss,) + output if loss is not None else output
+        return MaskedLMOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=out.hidden_states,
+            attentions=out.attentions,
+        )
 
 
 class RiNALMoLMHead(nn.Module):