esmfold2: keep distogram_head in fp32 as well

src/transformers/models/esmfold2/modeling_esmfold2.py

@@ -1864,6 +1864,7 @@ class EsmFold2PreTrainedModel(PreTrainedModel):
         "norm_start",
         "norm_single",
         "boundaries",
+        "distogram_head",
     ]
     _supports_sdpa = True
 

src/transformers/models/esmfold2/modular_esmfold2.py

@@ -1796,6 +1796,7 @@ class EsmFold2PreTrainedModel(PreTrainedModel):
         "norm_start",
         "norm_single",
         "boundaries",
+        "distogram_head",
     ]
     _supports_sdpa = True
 

tests/models/esmfold2/test_modeling_esmfold2.py

@@ -25,6 +25,7 @@
 
 from transformers import EsmFold2Config, is_torch_available
 from transformers.testing_utils import (
+    Expectations,
     TestCasePlus,
     require_torch,
     require_torch_accelerator,
@@ -208,6 +209,8 @@ def test_bf16_load_pins_norms_but_not_adaln_linears(self):
         self.assertEqual(adaln.cond_norm.weight.dtype, torch.float32)
         tri_mul = reloaded.msa_encoder.layers[0].tri_mul_in
         self.assertEqual(tri_mul.norm_start.weight.dtype, torch.float32)  # prefix-named norm stays pinned
+        self.assertEqual(reloaded.distogram_head.weight.dtype, torch.float32)
+        self.assertEqual(reloaded.distogram_head.bias.dtype, torch.float32)
         with torch.no_grad():
             out = reloaded.infer_protein(self.seq, num_loops=1, num_diffusion_samples=1, num_sampling_steps=2)
         self.assertTrue(torch.isfinite(out["sample_atom_coords"].float()).all())
@@ -543,7 +546,14 @@ def test_inference_deterministic_bf16(self):
             with torch.no_grad():
                 output = model.infer_protein(seq, num_loops=4, num_diffusion_samples=2, num_sampling_steps=32)
 
-            expected_distogram = torch.tensor([6.4062, 7.7500, 9.5625, 9.5000, 16.2500, 18.7500, 19.7500, 22.7500])
+            # fmt: off
+            expected_distogram = Expectations(
+                {
+                    ("cuda", 8): torch.tensor([6.3493, 7.7382, 9.4400, 9.4147, 16.2251, 18.6971, 19.6784, 22.7508]),
+                    ("xpu", 5): torch.tensor([6.1462, 7.4562, 9.1028, 9.0691, 15.9278, 18.4071, 19.3449, 22.4337]),
+                }
+            ).get_expectation()
+            # fmt: on
             torch.testing.assert_close(
                 output["distogram_logits"][0, 0, 1, :8].float().cpu(), expected_distogram, rtol=0, atol=0.2
             )