diff --git a/esm/layers/geom_attention.py b/esm/layers/geom_attention.py index 69656d0f..64834948 100644 --- a/esm/layers/geom_attention.py +++ b/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 diff --git a/esm/models/esm3.py b/esm/models/esm3.py index 05141f33..26298f94 100644 --- a/esm/models/esm3.py +++ b/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 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() diff --git a/tests/models/esm3_test.py b/tests/models/esm3_test.py new file mode 100644 index 00000000..33fce8bd --- /dev/null +++ b/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()