From 8a554b1d853f627e1da45e9f04917730265933f7 Mon Sep 17 00:00:00 2001 From: Ben Viggiano Date: Wed, 30 Sep 2026 10:21:49 -0700 Subject: [PATCH 1/4] Add test for ESM3 pLDDT encoding in bf16 (#178) --- tests/models/esm3_test.py | 49 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 49 insertions(+) create mode 100644 tests/models/esm3_test.py diff --git a/tests/models/esm3_test.py b/tests/models/esm3_test.py new file mode 100644 index 00000000..3575dc4a --- /dev/null +++ b/tests/models/esm3_test.py @@ -0,0 +1,49 @@ +"""ESM3 in reduced precision. + +The published weights load in bf16, so every tensor ESM3 builds internally has to +reach the bf16 layers in their own dtype. +""" + +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() From 35d362e5326f32eac44c9a93f0b7a0cab8e98ae6 Mon Sep 17 00:00:00 2001 From: Ben Viggiano Date: Wed, 30 Sep 2026 10:21:49 -0700 Subject: [PATCH 2/4] Add test for geometric attention in bf16 (#178) --- tests/layers/__init__.py | 0 tests/layers/geom_attention_test.py | 34 +++++++++++++++++++++++++++++ 2 files changed, 34 insertions(+) create mode 100644 tests/layers/__init__.py create mode 100644 tests/layers/geom_attention_test.py diff --git a/tests/layers/__init__.py b/tests/layers/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/tests/layers/geom_attention_test.py b/tests/layers/geom_attention_test.py new file mode 100644 index 00000000..aa8387af --- /dev/null +++ b/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() From 475bc0d440e122ae40b28652ae45cbac1866f18f Mon Sep 17 00:00:00 2001 From: Ben Viggiano Date: Wed, 30 Sep 2026 11:49:18 -0700 Subject: [PATCH 3/4] Fix ESM3 bf16 forward (#178) --- esm/layers/geom_attention.py | 6 ++++-- esm/models/esm3.py | 6 ++++-- 2 files changed, 8 insertions(+), 4 deletions(-) diff --git a/esm/layers/geom_attention.py b/esm/layers/geom_attention.py index 69656d0f..7999cc9e 100644 --- a/esm/layers/geom_attention.py +++ b/esm/layers/geom_attention.py @@ -63,7 +63,9 @@ def forward(self, s, affine, affine_mask, sequence_id, chain_id): ) ns = self.s_norm(s) - vec_rot, vec_dist = self.proj(ns).split( + # The frames are fp32 by design, so the geometric terms are computed in fp32. + 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 +146,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 diff --git a/esm/models/esm3.py b/esm/models/esm3.py index 05141f33..e2b5aca4 100644 --- a/esm/models/esm3.py +++ b/esm/models/esm3.py @@ -106,9 +106,11 @@ 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)) + # rbf is computed in the pLDDT's own dtype; cast at the boundary so bf16 weights work + 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 From 730cf9eaa84bb03fc7b2dd42c053780e1a384d8b Mon Sep 17 00:00:00 2001 From: Ben Viggiano Date: Wed, 30 Sep 2026 13:17:09 -0700 Subject: [PATCH 4/4] Remove comments --- esm/layers/geom_attention.py | 1 - esm/models/esm3.py | 1 - tests/models/esm3_test.py | 6 +----- 3 files changed, 1 insertion(+), 7 deletions(-) diff --git a/esm/layers/geom_attention.py b/esm/layers/geom_attention.py index 7999cc9e..64834948 100644 --- a/esm/layers/geom_attention.py +++ b/esm/layers/geom_attention.py @@ -63,7 +63,6 @@ def forward(self, s, affine, affine_mask, sequence_id, chain_id): ) ns = self.s_norm(s) - # The frames are fp32 by design, so the geometric terms are computed in fp32. vec = self.proj(ns).float() vec_rot, vec_dist = vec.split( [ diff --git a/esm/models/esm3.py b/esm/models/esm3.py index e2b5aca4..26298f94 100644 --- a/esm/models/esm3.py +++ b/esm/models/esm3.py @@ -106,7 +106,6 @@ 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) - # rbf is computed in the pLDDT's own dtype; cast at the boundary so bf16 weights work 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( diff --git a/tests/models/esm3_test.py b/tests/models/esm3_test.py index 3575dc4a..33fce8bd 100644 --- a/tests/models/esm3_test.py +++ b/tests/models/esm3_test.py @@ -1,8 +1,4 @@ -"""ESM3 in reduced precision. - -The published weights load in bf16, so every tensor ESM3 builds internally has to -reach the bf16 layers in their own dtype. -""" +"""ESM3 in reduced precision.""" import torch