Skip to content

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

Open
tivris wants to merge 1 commit into
TransformerLensOrg:devfrom
tivris:fix-compat-created-bias-dtype
Open

tivris wants to merge 1 commit into
TransformerLensOrg:devfrom
tivris:fix-compat-created-bias-dtype

Conversation

@tivris

@tivris tivris commented Oct 16, 2026

Copy link
Copy Markdown

Description

Fixes #1911.

process_weights upcasts non-float32 floating tensors to float32 for the folds and restores the recorded dtypes afterwards. Biases created during processing (q/k/v biases from LN folding, the MLP input and gate biases, b_O from fold_value_biases) were not in that record, so they stayed float32, and the next forward of a bf16/fp16 bridge after enable_compatibility_mode() failed with self and mat2 must have the same dtype. On MPS the process aborts.

Each created key now takes the original dtype of its .weight sibling before the existing downcast runs. Keys that were already in the state dict, fp32 models and the float32 arithmetic during the folds are unchanged.

Besides GPT-Neo, GPT-J, Qwen2 and StableLM from the issue, this fixes GLM, GLM-4, Qwen2.5-VL, Falcon, CodeGen, Qwen2-MoE and Seed-OSS, which fail the same way on tiny from_config models in bf16 and fp16.

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

@koriyoshi2041

Copy link
Copy Markdown
Contributor

Validated exact head 300fe85b9b43cf522ae3a6f3acfca1dd0f6b20af locally.

  • The added focused test matrix passes: 8 passed, 28 deselected.
  • I also exercised three architectures named in the description but not present in BIAS_CREATING_CONFIGS: CodeGenForCausalLM, FalconForCausalLM, and Qwen2MoeForCausalLM, each in both torch.float16 and torch.bfloat16 (6 additional cases). For each case I built a tiny model from config, enabled compatibility mode, asserted every model parameter retained the requested dtype, and ran a forward pass. All six passed.

That gives coverage across the Conv1D/parallel CodeGen path, Falcon's packed attention projection, and the Qwen2 MoE/shared-expert path. I did not find a dtype regression in the added architecture controls.

This branch has not been deployed

No deployments
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants