Refactor spec modification/introspection to make references to Submodules typed - #2834

Closed

nschank wants to merge 8 commits into

NVIDIA:mainfrom

nschank:submodules

Commits on Jan 30, 2026

Refactor to ensure TransformerLayerSubmodules references can easily be identified
nschank
committed
9894b88
Look for .self_attention modules
nschank
committed
e971c5b
Look for .mlp modules
nschank
committed
ada5417
Look for MambaStackSubmodules members
nschank
committed
0c6222d
Look for MLA Submodules members
nschank
committed
52b9a65
Clean up imports
nschank
committed
459d783
Switch from cast to assert isinstance
nschank
committed
6800e72
Preserve the signature of the wrappers in gpt_layer_specs
nschank
committed
592747b