Drop optional unused aliases on hybrid architectures - #1579
Conversation
KYinXu
changed the title
|
|
||
|
|
||
| def _prune_hook_aliases_for_skipped(component: GeneralizedComponent, skipped: list[str]) -> None: | ||
| """Drop hook aliases whose target path starts at a skipped optional submodule. |
There was a problem hiding this comment.
The prune fires for every optional skip across 12 architectures, including native SSM aliases on nemotron_h/zamba2. Can we update the docstring to be more general, and add one non-attn test case?
| to_drop.append(alias_name) | ||
| elif len(kept) != len(target): | ||
| aliases[alias_name] = kept | ||
| elif _first_segment(target) in skipped_set: |
There was a problem hiding this comment.
On OLMo-Hybrid, BlockBridge retargets hook_mlp_out onto the optional ln2_post (block.py:78), so this branch drops it on linear-attention layers that do have an mlp, and removes the warning that flagged it. Can the prune keep aliases with a live fallback target?
| assert "hook_resid_pre" in blocks[3].hook_aliases | ||
| assert "hook_mlp_out" in blocks[3].hook_aliases | ||
|
|
||
| def test_pruned_aliases_do_not_warn_on_register(self): |
There was a problem hiding this comment.
Both new tests call _register_aliases() directly, so nothing asserts a hybrid bridge booting warning-free through TransformerBridge.__init__. Can you add a boot-level assertion using the config-only tiny-bridge fixture in test_qwen3_5_adapter.py:562?
| with warnings.catch_warnings(record=True) as caught: | ||
| warnings.simplefilter("always") | ||
| blocks[3]._register_aliases() | ||
| unresolved = [ |
There was a problem hiding this comment.
AttnAdapter declares no ln2/mlp, so block 3 still emits two legitimate unresolved warnings. This forces the substring filter instead of asserting the list is empty. Can the fixture carry those submodules so the assertion can be unconditional?
2 participants
Description
When hybrid layers skip an optional submodule (e.g.
attnon GDN / Mamba-2 layers), the sharedBlockBridgetemplate still declared HT aliases likehook_attn_out → attn.hook_out._register_aliasesthen failed to resolve them and emitted a harmless warning per hybrid layer (notably when loading Qwen3.5; see also #1547).This PR prunes those inapplicable aliases in
setup_submoduleswhen an optional submodule is skipped—per layer, so full-attention siblings keep their HT aliases and real unresolved-alias bugs still warn. Unit tests cover prune behavior and no warning on register. The Jacobian Lens Qwen3.5 demo cell is updated now that boot no longer expects that stderr.Fixes #1578
Type of change
Screenshots
Warnings no longer appear
Checklist: