[Sonica-B]

Description

Anyone who loads a bridge in float64 and calls enable_compatibility_mode() gets weights rounded to float32. This hits the fp64 step of Debugging Numerical Divergence ("run fp64 ... If the diff stays the same magnitude, it's a bug. If it drops by ~8 orders of magnitude, it was noise"), the same check in the contributing guide's numerical-work section, and any analysis that uses float64 for tight tolerances. On gpt2 the documented check drops by one order instead of eight, so plain float32 noise reads as a bug.

The cause is the upcast at the top of ProcessWeights.process_weights. It is meant for reduced-precision dtypes ("bfloat16, float16"), but its guard is v.dtype != torch.float32, which also matches float64. So float64 tensors are cast down to float32, folded, and cast back, and every processed weight of a float64 model ends up exactly representable in float32. It goes back to #1176; #1912 changed the same block for created biases and kept the condition.

The change records every floating tensor's dtype, casts everything except float64 to float32, and restores every key afterwards:

• float64 is folded in float64.
• float32, bf16 and fp16 are processed exactly as before.
• Every key, including biases created during folding, leaves with the dtype it came in with. That also holds in a state dict that mixes float32 and float64.

flowchart LR
    subgraph Before
        A1[float64 state dict] --> B1["v.dtype != float32<br/>→ v.float()"]
        B1 --> C1[fold_ln / center_* / fold_value_biases<br/>in float32]
        C1 --> D1[cast back to float64<br/>values float32-rounded]
    end
    subgraph After
        A2[float64 state dict] --> B2["dtype recorded<br/>float64 kept as is"]
        B2 --> C2[fold_ln / center_* / fold_value_biases<br/>in float64]
        C2 --> D2[every key restored<br/>to its input dtype]
    end

Loading

Before and after, gpt2 on CPU, compatibility mode against AutoModelForCausalLM with attn_implementation="eager", prompt "The Eiffel Tower is in the city of":

for dtype in (torch.float32, torch.float64):
    hf = AutoModelForCausalLM.from_pretrained("gpt2", dtype=dtype, attn_implementation="eager").eval()
    bridge = TransformerBridge.boot_transformers("gpt2", device="cpu", dtype=dtype)
    bridge.enable_compatibility_mode()
    tokens = bridge.to_tokens(prompt)
    # max |log_softmax(bridge(tokens)) - log_softmax(hf(tokens).logits)|, and whether W_U == W_U.float().to(W_U.dtype)

dev @ a7b3213
torch.float32  max |log-prob diff| vs HF = 7.82e-05   W_U dtype torch.float32, every entry float32-representable: True
torch.float64  max |log-prob diff| vs HF = 7.82e-06   W_U dtype torch.float64, every entry float32-representable: True

this branch
torch.float32  max |log-prob diff| vs HF = 7.82e-05   W_U dtype torch.float32, every entry float32-representable: True
torch.float64  max |log-prob diff| vs HF = 1.78e-13   W_U dtype torch.float64, every entry float32-representable: False

On tiny from_config models in float64, compat-mode log-probs against HF go from 2.3e-08 to 8.9e-16 (GPT-2), from 2.9e-08 to 8.9e-16 (GPT-J) and from 2.2e-08 to 8.9e-16 (Llama).

Tests, in tests/unit/test_weight_processing.py next to the #1912 dtype tests:

• test_process_weights_keeps_float64_precision (tiny GPT-2 and GPT-J) builds the HF model in float64, enables compatibility mode, and compares log-probs at atol=1e-12. On dev it fails at 2.6e-08 and 2.8e-08.
• test_process_weights_created_biases_keep_model_dtype now also runs in torch.float64, so created biases are checked to stay float64.

Run on Windows 11, Python 3.11.9, torch 2.11.0+cpu, transformers 5.13.0:

• tests/unit/test_weight_processing.py and tests/unit/test_unembed_folding_layout.py: 167 passed, 1 skipped. On dev the 2 new precision tests fail and the rest pass.
• black, isort 5.8.0, pycln and mypy are clean on the changed files.

Type of change

• Bug fix (non-breaking change which fixes an issue)

Checklist:

• I have commented my code, particularly in hard-to-understand areas
• I have made corresponding changes to the documentation
• My changes generate no new warnings
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have not rewritten tests relating to key interfaces which would affect backward compatibility

[jlarson4]

Good catch and good fix @Sonica-B, thank you for putting this together. I am merging as is!