Skip to content

Fix rcps test and mamba imports#65

Open
eric-czech wants to merge 1 commit intokuleshov-group:mainfrom
eric-czech:testskip
Open

Fix rcps test and mamba imports#65
eric-czech wants to merge 1 commit intokuleshov-group:mainfrom
eric-czech:testskip

Conversation

@eric-czech
Copy link
Collaborator

The test_rcps.py#test_rcps_add_norm_wrapper test was failing for me for two reasons:

  1. It was only looking for RMSNorm using the Mamba v1 package structure: mamba_ssm.ops.triton.layernorm rather than also checking mamba_ssm.ops.triton.layer_norm (like modeling_caduceus.py does).
  2. The logic for flipping RC results wasn't working (maybe it changed in v2?).

I cleaned up the way version-dependent mamba modules and classes are loaded in this PR and fixed that test.

@yair-schiff
Copy link
Contributor

This is really clean and nice, thanks! Did you test that the change is non-breaking for running pre-training, e.g.?

If so, happy to merge this in.

@eric-czech
Copy link
Collaborator Author

Did you test that the change is non-breaking for running pre-training, e.g.?

Non-breaking in what sense? I'd argue that the bar for testing should be what's in the unit tests or CI workflows (or some other repeatable process for everyone), so here are some options that come to mind:

  1. Testing training orchestration over Hugging Face and Composer APIs would be included as a part of Add support for gradient checkpointing #60. We could wait for that to get merged and rebase these changes on it, though I would be shocked if that had any interactions with what's here.
  2. Make a CI process and figure out how to test multiple environments, e.g. with multiple versions of Mamba. That would require figuring out how to do CI w/ GPUs though so it likely wouldn't be free (cf. Open-Athena/openathena-cornell-dna/issues/3).
  3. Use tox or something like it to create unit tests that build multiple environments locally and then run all the tests within them (possibly as a pre-commit hook). We still don't have a documented way to create one environment for this repo though (e.g. caduceus_env.yml doesn't currently resolve), so it would be a significant lift to get there and I don't see much value in it yet.

My suggestion would be to make a best-effort to carry support for Mamba 1.x without any real guarantees (for a time), and test Mamba 2.x+ well. Would you agree, or is there a good argument for not building on Mamba 2 instead?

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants