Fix native init: resolve initializer_range sentinel, thread gain into xavier/kaiming - #1685
Conversation
|
@cjnegao11-cmyk I just posted 6 fresh issues as well if you want to take a look! Thanks for contributing, it is appreciated. I will opt for this solution to #1568 & review it tomorrow if we don't here from @tomatotomata by then |
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for putting this together @cjnegao11-cmyk and for verifying the gain equivalence rather than just asserting it. A couple small edit requests below:
… xavier/kaiming, document residual scaling delta
c8a3423 to
81b9400
Compare
jlarson4
left a comment
There was a problem hiding this comment.
Thanks for the update & rebase @cjnegao11-cmyk! Just one additional comment
| optimizer.zero_grad() | ||
|
|
||
|
|
||
| def test_boot_native_resolves_initializer_range_sentinel(): |
There was a problem hiding this comment.
These assert weight statistics after boot_native, but dev's init.py already falls back to 0.8/math.sqrt(cfg.d_model) and applies the gain, so both tests pass on dev without your fixes.
Lets keep the end-to-end assertions and add the config-level ones back alongside them — cfg.initializer_range == pytest.approx(0.8 / math.sqrt(cfg.d_model)) for init_mode="gpt2", plus a case for a non-gpt2 mode resolving to 1.0, which is currently untested. Both read -1.0 without this PR.
|
Great work! Merging now |
|
Thanks for the thorough review and for merging! Happy to keep contributing. |
Addresses #1568
initializer_range = -1.0sentinel inTransformerBridgeConfig.__post_init__, mirroringHookedTransformerConfigdevand keeps the existingdevxavier/kaiming gain handling instead of duplicating it1/sqrt(2*n_layers)residual output scaling in gpt2-mode as an intentional delta fromHookedTransformer._init_weights_gpt2boot_native()is exercised end-to-end and verifies the resulting embedding stdValidation:
python -m pytest tests/unit/model_bridge/test_boot_native.py -q→ 31 passedgit diff --check→ passedNote: this issue was previously assigned to @tomatotomata; @jlarson4 asked for a status check before this PR was opened.