Add config tests for mamba and mamba2 architecture adapter#1532
Conversation
|
hey @jlarson4, please review and merge this |
jlarson4
left a comment
jlarson4
left a comment
There was a problem hiding this comment.
Hi @MdSadiqMd, thank you for putting these tests together, sorry for the delay on the review. These tests are great, I just want to trim them down a little bit before merging. Notes below, nothing crazy. Let me know if you have any questions!
| """Mamba2 only applies to Phase 4 (generation + text quality).""" | ||
|
|
||
| def test_applicable_phases_is_phase_4_only(self, adapter: Mamba2ArchitectureAdapter) -> None: | ||
| assert adapter.applicable_phases == [4] |
There was a problem hiding this comment.
Something I completely missed in my first review. Mamba and Mamba2 are now able to run through phase 1,2, and 3 as of PR #1481. Because you originally based your PR on main, this detail was lost as that has not been released to main yet. Since this runs the default phases now, it is probably safe to drop this test.
Thank you for your patience on this one, I am aiming to get the release out tomorrow which will get main up to date with these changes
There was a problem hiding this comment.
np, will do the fixes, and thank you very much for maintaining this repo
| """Mamba only applies to Phase 4 (generation + text quality).""" | ||
|
|
||
| def test_applicable_phases_is_phase_4_only(self, adapter: MambaArchitectureAdapter) -> None: | ||
| assert adapter.applicable_phases == [4] |
|
Hey @jlarson4 made the changes, please review it once again, let me know if anything else is needed from my side |
|
Looks great @MdSadiqMd, merging now! |
2 participants
Description
This PR adds unit tests for two SSM architecture adapters currently lacking focused testing:
MambaArchitectureAdapter (39 tests): Config flags (RMS norm, no positional embeddings, stateful), SSMBlockBridge structure with norm and SSMMixerBridge, mixer submodules (in_proj, conv1d, x_proj, dt_proj, out_proj), factory registration, and applicable_phases.
Mamba2ArchitectureAdapter (39 tests): Config flags plus computed fields (intermediate_size, conv_dim, expected_in_proj_out_features), SSM2MixerBridge structure (in_proj, conv1d, inner_norm, out_proj — no x_proj/dt_proj), factory registration, and architecture guards distinguishing Mamba2 from Mamba1.
Behavioural coverage (forward pass, hook firing) already lives in the integration tests:
tests/integration/model_bridge/test_mamba_adapter.pytests/integration/model_bridge/test_mamba2_adapter.pyThese unit tests complement the integration tier, following the pattern of existing adapter tests like
test_opt_adapter.pyandtest_mixtral_adapter.py.black,mypy, andpytest(unit tests) all pass locally.Parts of #1302
Type of change
Screenshots
N/A
Checklist: