Repository navigation
Conversation
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Looks good, just one comment, let me know if you have any questions!
| return state_dict | ||
|
|
||
| @staticmethod | ||
| def _unembed_model_axis(weight: torch.Tensor, d_model: int, adapter=None) -> int: |
There was a problem hiding this comment.
center_unembed (around line 1431) still picks the vocab axis by comparing shapes against d_vocab, so a square readout in HF layout gets centered along d_model and the log-probs shift, the same ambiguity this helper resolves for the fold path. Since center_unembed=True is the compatibility-mode default, can we route that path through the helper as well?
There was a problem hiding this comment.
Fixed in 8bed6ad. center_unembed now uses the same layout helper as folding and centers along the complementary vocabulary axis, including when cfg is omitted. I added centering coverage for square and padded readouts, plus native-model compatibility-mode regressions with folding both enabled and disabled.
Description
Final LayerNorm folding can change predictions when the unembedding is square (
d_vocab == d_model). After an adapter converts the readout to TL layout, the old shape-first checks still select the HF axis for gamma and beta. The centering path has the opposite ambiguity for square readouts left in HF layout.Resolve the hidden axis from the adapter conversion contract before using unambiguous dimensions as a legacy fallback. Adapter-free square weights follow the TL layout; square Bridge weights without a tensor conversion retain their HF layout. Use the same axis resolution for final-norm scaling, bias folding and hidden-axis centering.
Regression tests compare the folded parameters against analytical formulas and verify logits and log-softmax parity on config-built tiny GPT-2 models with square and non-square vocabularies. Additional tensor-only cases cover TL, HF, explicit conversion and legacy mappings, ordinary/offset RMSNorm, optional centering and disabled bias folding. No model downloads are needed.
Type of change
Checklist
No public API or documentation contract changes.
Validation: the five-file regression surface passed with 222 passed, 1 skipped. The skip is the existing cross-device test requiring CUDA/MPS. Full mypy passed for 389 source files; pinned pycln/isort/Black and
git diff --checkpassed for the changed files. The complete unit suite was not run locally. Existing SWIG deprecations and the Phi3 unsupported-fold warning remain unchanged. GPU/MPS and pretrained-model validation were not run.