fix: keep created biases in the model dtype during weight processing

tests/unit/test_weight_processing.py

@@ -10,8 +10,10 @@
 import einops
 import pytest
 import torch
+from transformers import AutoConfig, AutoModelForCausalLM
 
 from transformer_lens.config.transformer_lens_config import TransformerLensConfig
+from transformer_lens.model_bridge.sources import build_bridge_from_module
 from transformer_lens.weight_processing import ProcessWeights
 
 
@@ -1101,3 +1103,46 @@ def test_resolve_state_dict_key_dense_mlp_fallback():
         ProcessWeights._resolve_state_dict_key(sd5, "blocks.0.mlp.shared_expert.in.weight", 0)
         == "blocks.0.mlp.shared_expert.in.weight"
     )
+
+
+BIAS_CREATING_CONFIGS = {
+    "GPTNeoForCausalLM": dict(
+        model_type="gpt_neo",
+        hidden_size=32,
+        num_layers=2,
+        num_heads=4,
+        attention_types=[[["global", "local"], 1]],
+    ),
+    "GPTJForCausalLM": dict(model_type="gptj", n_embd=32, n_layer=2, n_head=4, rotary_dim=4),
+    "Qwen2ForCausalLM": dict(
+        model_type="qwen2",
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+    ),
+    "StableLmForCausalLM": dict(
+        model_type="stablelm",
+        hidden_size=32,
+        intermediate_size=64,
+        num_hidden_layers=2,
+        num_attention_heads=4,
+        num_key_value_heads=2,
+    ),
+}
+
+
+@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16])
+@pytest.mark.parametrize("architecture", sorted(BIAS_CREATING_CONFIGS))
+def test_process_weights_created_biases_keep_model_dtype(architecture, dtype):
+    """Biases created by folding take the dtype of the projection they are added to."""
+    hf_config = AutoConfig.for_model(vocab_size=64, **BIAS_CREATING_CONFIGS[architecture])
+    hf_config.architectures = [architecture]
+    hf_model = AutoModelForCausalLM.from_config(hf_config).to(dtype)
+    bridge = build_bridge_from_module(hf_model, architecture, hf_config=hf_config)
+
+    bridge.enable_compatibility_mode()
+
+    assert not [name for name, param in hf_model.named_parameters() if param.dtype != dtype]
+    bridge(torch.tensor([[1, 2, 3, 4]]))

transformer_lens/weight_processing.py

@@ -1657,6 +1657,7 @@ def process_weights(
         # folding involve multiplications that accumulate rounding errors when
         # performed in low precision.
         original_dtypes: Dict[str, torch.dtype] = {}
+        original_keys = set(state_dict)
         for k, v in state_dict.items():
             if isinstance(v, torch.Tensor) and v.is_floating_point() and v.dtype != torch.float32:
                 original_dtypes[k] = v.dtype
@@ -1730,7 +1731,11 @@ def process_weights(
                 state_dict, cfg, adapter=adapter
             )
 
-        # Downcast back to original dtypes
+        # Downcast back to original dtypes; biases created by folding take their weight's dtype
+        for k in state_dict.keys() - original_keys:
+            weight_dtype = original_dtypes.get(k.removesuffix(".bias") + ".weight")
+            if weight_dtype is not None:
+                original_dtypes[k] = weight_dtype
         for k, orig_dtype in original_dtypes.items():
             if k in state_dict and isinstance(state_dict[k], torch.Tensor):
                 state_dict[k] = state_dict[k].to(orig_dtype)