Skip to content

fix: apply dtype and device changes to gated MLPs in compatibility mode - #1929

Merged
jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
tivris:fix-gated-mlp-processed-weights
Oct 21, 2026
Merged

jlarson4 merged 1 commit into
TransformerLensOrg:devfrom
tivris:fix-gated-mlp-processed-weights

Conversation

@tivris

@tivris tivris commented Oct 21, 2026

Copy link
Copy Markdown

Description

Fixes #1927.

GatedMLPBridge.set_processed_weights stored the processed weights and biases twice: in the gate, in and out projections, and in six plain attributes that the processed forward read. Conversions and copy.deepcopy update only the projections' parameters, so the forward read stale tensors.

The processed forward now reads each projection's current weight and bias, detached as before, and the six attributes are gone. JointGateUpMLPBridge now checks only the processed-weights flag. Outputs, hook order and gradients are unchanged when nothing is converted. The hook test's fixture now installs its weights through set_processed_weights instead of setting the attributes.

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 jlarson4 linked an issue Oct 21, 2026 that may be closed by this pull request
1 task done

@jlarson4 jlarson4 left a comment

Copy link
Copy Markdown
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This PR is approved, just waiting for CI to complete before I can merge. Going to rerun BackwardLens once the job is done, I have a fix for that pending but this PR should be ready before that is merged.

@jlarson4
jlarson4 merged commit e7feb73 into TransformerLensOrg:dev Oct 21, 2026
55 of 56 checks passed
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.

[Bug Report] Gated MLPs fail after a dtype or device change in compatibility mode

2 participants