Skip to content

rope tests#134

Merged
ncfrey merged 2 commits intomainfrom
n/rope-tests
Jul 3, 2025
Merged

rope tests#134
ncfrey merged 2 commits intomainfrom
n/rope-tests

Conversation

@ncfrey
Copy link
Contributor

@ncfrey ncfrey commented Jul 2, 2025

This pull request introduces enhancements to rotary embedding functionality, including a CPU fallback implementation, improved handling of flash attention availability, and comprehensive tests for the rotary embedding module. The most important changes are the addition of CPU-based rotary embedding methods, conditional fallback for environments without CUDA or flash attention, and new test cases to validate functionality and positional encoding behavior.

Enhancements to rotary embedding functionality:

  • CPU fallback implementation: Added apply_rotary_emb_cpu and apply_rotary_emb_unpad_cpu methods to enable rotary embedding on CPU for unpadded sequences. These methods ensure compatibility in environments without CUDA or flash attention. (src/lobster/model/modern_bert/_rotary.py, [1] [2]
  • Conditional fallback: Modified apply_rotary_emb_unpad to use the CPU fallback implementation when CUDA or flash attention is unavailable. (src/lobster/model/modern_bert/_rotary.py, src/lobster/model/modern_bert/_rotary.pyR219-R222)

Improvements to testing:

  • Unit tests for rotary embedding: Added tests to verify the forward pass of UnpaddedRotaryEmbedding and ensure positional encoding behavior. These tests check output correctness, device compatibility, and positional invariance. (tests/lobster/model/modern_bert/test_rotary.py, tests/lobster/model/modern_bert/test_rotary.pyR1-R144)
  • Flash attention handling in tests: Incorporated logic to gracefully handle environments without flash attention during test setup, ensuring compatibility across different configurations. (tests/lobster/model/modern_bert/test_rotary.py, tests/lobster/model/modern_bert/test_rotary.pyR1-R144)

Documentation update:

  • Installation instructions: Updated README.md to include a new command for resolving flash attention installation issues during setup. (README.md, README.mdR87-R89)

@ncfrey ncfrey requested a review from karinazad July 2, 2025 21:17
@ncfrey ncfrey temporarily deployed to test.pypi.org July 2, 2025 21:30 — with GitHub Actions Inactive

assert outputs[0].device == model.device

@pytest.mark.skip(
Copy link
Collaborator

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

+1

@ncfrey ncfrey merged commit c5a8720 into main Jul 3, 2025
5 checks passed
@ncfrey ncfrey deleted the n/rope-tests branch July 3, 2025 20:51
ncfrey added a commit that referenced this pull request Jul 7, 2025
* rope tests

* skip test

---------

Co-authored-by: freyn6 <freyn6@gene.com>
ncfrey added a commit that referenced this pull request Jul 7, 2025
* rope tests

* skip test

---------

Co-authored-by: freyn6 <freyn6@gene.com>
ncfrey added a commit that referenced this pull request Jul 7, 2025
* rope tests (#134)

* rope tests

* skip test

---------

Co-authored-by: freyn6 <freyn6@gene.com>

* update tokenizer

* updates

* save results

* real layers

* fix layers

* Update uv.lock after merge

---------

Co-authored-by: freyn6 <freyn6@gene.com>
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