Fix ESM3 bf16 forward

esm/layers/geom_attention.py

@@ -63,7 +63,8 @@ def forward(self, s, affine, affine_mask, sequence_id, chain_id):
         )
 
         ns = self.s_norm(s)
-        vec_rot, vec_dist = self.proj(ns).split(
+        vec = self.proj(ns).float()
+        vec_rot, vec_dist = vec.split(
             [
                 self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages,
                 self.v_heads * 2 * 3,
@@ -144,6 +145,6 @@ def forward(self, s, affine, affine_mask, sequence_id, chain_id):
         )
         if self.mask_and_zero_frameless:
             attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0)
-        s = self.out_proj(attn_out)
+        s = self.out_proj(attn_out.to(s.dtype))
 
         return s

esm/models/esm3.py

@@ -106,9 +106,10 @@ def forward(
         rbf_16_fn = partial(rbf, v_min=0.0, v_max=1.0, n_bins=16)
         # the `masked_fill(padding_mask.unsqueeze(2), 0)` for the two below is unnecessary
         # as pad tokens never even interact with the "real" tokens (due to sequence_id)
-        plddt_embed = self.plddt_projection(rbf_16_fn(average_plddt))
+        dtype = self.plddt_projection.weight.dtype
+        plddt_embed = self.plddt_projection(rbf_16_fn(average_plddt).to(dtype))
         structure_per_res_plddt = self.structure_per_res_plddt_projection(
-            rbf_16_fn(per_res_plddt)
+            rbf_16_fn(per_res_plddt).to(dtype)
         )
 
         # Structure + "structural features" embeds

tests/layers/geom_attention_test.py

@@ -0,0 +1,34 @@
+"""Geometric attention in reduced precision."""
+
+import torch
+
+from esm.layers.geom_attention import GeometricReasoningOriginalImpl
+from esm.utils.structure.affine3d import build_affine3d_from_coordinates
+
+C_S = 64
+V_HEADS = 4
+B, L = 2, 10
+
+
+def test_geometric_attention_runs_in_bf16_with_fp32_frames():
+    """Issue #178: a bf16 ESM3 still builds its frames from fp32 coordinates.
+
+    ``ESM3.forward`` passes ``structure_coords`` straight to
+    ``build_affine3d_from_coordinates``, so geometric attention has to mix fp32
+    rotations with its own bf16 projections and hand ``out_proj`` bf16 input.
+    """
+    torch.manual_seed(0)
+    attention = GeometricReasoningOriginalImpl(c_s=C_S, v_heads=V_HEADS)
+    attention = attention.to(torch.bfloat16).eval()
+
+    affine, affine_mask = build_affine3d_from_coordinates(torch.randn(B, L, 3, 3))
+    s = torch.randn(B, L, C_S, dtype=torch.bfloat16)
+    sequence_id = torch.zeros(B, L, dtype=torch.long)
+    chain_id = torch.ones(B, L, dtype=torch.long)
+
+    with torch.no_grad():
+        out = attention(s, affine, affine_mask, sequence_id, chain_id)
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == (B, L, C_S)
+    assert torch.isfinite(out.float()).all()

tests/models/esm3_test.py

@@ -0,0 +1,45 @@
+"""ESM3 in reduced precision."""
+
+import torch
+
+from esm.models.esm3 import EncodeInputs
+from esm.utils.constants import esm3 as C
+
+D_MODEL = 32
+B, L = 2, 10
+
+
+def test_encode_inputs_runs_in_bf16_with_the_default_plddt():
+    """Issue #178: ``ESM3.forward`` fills pLDDT in with ``.float()``.
+
+    Those fp32 values go through ``rbf`` into the bf16 pLDDT projections, so the
+    first forward of a bf16 model fails with a dtype mismatch.
+    """
+    torch.manual_seed(0)
+    encoder = EncodeInputs(D_MODEL).to(torch.bfloat16).eval()
+
+    # The inputs ESM3.forward builds for a sequence-only call.
+    sequence_tokens = torch.randint(4, 24, (B, L))
+    structure_tokens = torch.full((B, L), C.STRUCTURE_MASK_TOKEN)
+    average_plddt = torch.ones(B, L)
+    per_res_plddt = torch.zeros(B, L)
+    ss8_tokens = torch.full((B, L), C.SS8_PAD_TOKEN)
+    sasa_tokens = torch.full((B, L), C.SASA_PAD_TOKEN)
+    function_tokens = torch.full((B, L, 8), C.INTERPRO_PAD_TOKEN)
+    residue_annotation_tokens = torch.full((B, L, 16), C.RESIDUE_PAD_TOKEN)
+
+    with torch.no_grad():
+        out = encoder(
+            sequence_tokens,
+            structure_tokens,
+            average_plddt,
+            per_res_plddt,
+            ss8_tokens,
+            sasa_tokens,
+            function_tokens,
+            residue_annotation_tokens,
+        )
+
+    assert out.dtype == torch.bfloat16
+    assert out.shape == (B, L, D_MODEL)
+    assert torch.isfinite(out.float()).all()