Repository navigation
fix(zamba): add use_associative_scan config flag to avoid torch.compile slowdown - #48331
Conversation
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Could we have a small fast test to check this in zamba + can you fix the linter (make style)
48840ec to
67e232a
Compare
Sure, thank you for your prompt review. |
67e232a to
3546841
Compare
After looking into it, it seems that |
|
Ah no for test I meant like running associative vs non associative scan and check that the logits are close enough |
Got it, I see what you mean now — mamba has |
3546841 to
44f520a
Compare
Test passed locally (27min15secs, |
44f520a to
5ad96df
Compare
|
@vasqu, when you have a chance, could you please review the follow-up for this PR? I'd appreciate your guidance on whether the proposed changes align with the expectations of the Hugging Face community. |
vasqu
left a comment
vasqu
left a comment
There was a problem hiding this comment.
Running the slow tests in a sec but would avoid the del model before merge
| output = model.generate(input_ids, do_sample=False, use_cache=False, max_new_tokens=10) | ||
| output_text = self.tokenizer.decode(output[0].tolist()) | ||
|
|
||
| self.assertEqual(output_text, expected_text) |
There was a problem hiding this comment.
I guess this mimics the mamba1 test. SGTM
There was a problem hiding this comment.
Yes. I compared the implementation when I drafted this PR.
|
run-slow: zamba |
Nvidia CIThis comment contains models: ["models/zamba"] |
AMD CIThis comment contains models: ["models/zamba"] |
|
The docs for this PR live here. All of your documentation changes will be reflected on that endpoint. The docs are available until 30 days after the last update. |
CI Results (AMD)Commit Info
Model CI Report❌ 1 new failed tests from this PR 😭
|
|
The AMD failures are irrelevant if you can address my last comments, I think we can merge then 🤗 |
avoid the del model before merge Co-authored-by: Anton Vlasjuk <73884904+vasqu@users.noreply.github.com>
|
[For maintainers] Suggested jobs to run (before merge) run-slow: zamba |
CI recapDashboard: View test results in Grafana |
Thank you so much for your prompt reply. I have addressed the issue you mentioned. Thanks again for your time and efforts in reviewing this. |
ZambaMambaMixer.forward() passes use_associative_scan=False directly to mamba_selective_scan, which falls back to a sequential Python loop that makes torch.compile slow.
The other Mamba-family models (mamba, jamba, falcon_mamba) instead read this setting from their config and forward it to the scan function. This PR adds the same plumbing to ZambaConfig and ZambaMambaMixer.
Changes:
@vasqu