Skip to content

Add lora layer tests #981

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom
Merged
merged 14 commits into from
Nov 8, 2025
Merged

Add lora layer tests #981

merged 14 commits into from
Nov 8, 2025

Conversation

@vanbasten23
Copy link
Collaborator

@vanbasten23 vanbasten23 commented Oct 30, 2025 •
edited
Loading

Description

This PR adds unit tests for the MergedColumnParallelLinearWithLoRA lora wrapper.

Tests

pytest -s -vv tests/lora/test_layers.py

Checklist

Before submitting this PR, please make sure:

  • I have performed a self-review of my code.
  • I have necessary comments in my code, particularly in hard-to-understand areas.
  • I have made or will make corresponding changes to any relevant documentation.

@vanbasten23 vanbasten23 force-pushed the xiowei/add_lora_layer_tests branch from 9369e12 to c17bace Compare October 30, 2025 23:20
@vanbasten23 vanbasten23 marked this pull request as ready for review October 31, 2025 18:39
@vanbasten23
Copy link
Collaborator Author

vanbasten23 commented Oct 31, 2025

cc: @qihqi @kyuyeunk @hfan @aman2930

kyuyeunk
kyuyeunk reviewed Nov 5, 2025
View reviewed changes
@vanbasten23 vanbasten23 force-pushed the xiowei/add_lora_layer_tests branch from 82d8a25 to 7ea2939 Compare November 5, 2025 21:27
@vanbasten23 vanbasten23 requested a review from kyuyeunk November 6, 2025 17:06
kyuyeunk
kyuyeunk reviewed Nov 7, 2025
View reviewed changes
@vanbasten23
Save an initial version. Not capable of running at all.
e73674a 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
now the test can run to completion.
da64f46 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
ok. The case without lora passed.
823fbc1 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
ok, the test passed. Need to make it simpler next.
e8718d8 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
also check if the correct and the sharding is correct.
539c639 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
cleaned up
90888cb 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
ok, fixed the torchax.view.item() issue.
223cf15 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
add multi-chip test case
c1f5841 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
fix the format
70e1ac2 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
clean up
bdb47e9 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
Add lora unit tests to the CI
3022e6a 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
extract the function _create_column_parallel_packed_layer out of the …
3f3355b 
…test.

Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
reduce the size of the test
fde81d3 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23
in the middle of refactoring
e9a705b 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
@vanbasten23 vanbasten23 force-pushed the xiowei/add_lora_layer_tests branch from 0afa56b to e9a705b Compare November 7, 2025 19:22
@vanbasten23 vanbasten23 requested a review from kyuyeunk November 7, 2025 19:22
kyuyeunk
kyuyeunk approved these changes Nov 7, 2025
View reviewed changes
Copy link
Collaborator

@kyuyeunk kyuyeunk left a comment

Choose a reason for hiding this comment

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

LGTM. Can you run a ci and verify it's all good before submitting it?

@vanbasten23
Copy link
Collaborator Author

vanbasten23 commented Nov 8, 2025

The CI https://buildkite.com/tpu-commons/tpu-inference-ci/builds/5099#_ is green. The faililng tests are irrelevant to this PR.

Thanks for the review!

@vanbasten23 vanbasten23 merged commit d86986b into main Nov 8, 2025
3 checks passed
sierraisland pushed a commit that referenced this pull request Nov 8, 2025
@vanbasten23 @sierraisland
Add lora layer tests (#981)
37279aa 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
sierraisland pushed a commit that referenced this pull request Nov 8, 2025
@vanbasten23 @sierraisland
Add lora layer tests (#981)
63cccf6 
Signed-off-by: Xiongfei Wei <isaacwxf23@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@kyuyeunk kyuyeunk kyuyeunk approved these changes

Assignees

No one assigned

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

3 participants

@vanbasten23 @kyuyeunk