Skip to content

Feat/add lora for sglangjax #826

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
copybara-service merged 1 commit into from
Jan 17, 2026
Merged

Feat/add lora for sglangjax #826

copybara-service merged 1 commit into from
Jan 17, 2026

Conversation

@aolemila
Copy link
Collaborator

@aolemila aolemila commented Dec 2, 2025 •
edited
Loading

Resolves #825.

  1. Add codes in scripts/grpo_demo_llama3_qwen2.py to run LoRA.
  2. Add the sglang_jax_lora_test.py to ensure update_params works, and put it into tpu-tests.yml. verify_update_params will be executed when VERIFY_UPDATE_PARAMS_KEY is configured.
  3. Add more fields for SGLangJax in RolloutConfig.
  4. Pass the following tests. Environment: TPU-v6e-4.

Test1: Run verification of update_params

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 tests/generate/sglang_jax_lora_test.py.
image

Test2: Run scripts/grpo_demo_llama3_qwen2.py without LoRA

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax.
Pasted Graphic 20

Test3: Run scripts/grpo_demo_llama3_qwen2.py with LoRA

JAX_COMPILATION_CACHE_DIR=/tmp/jit_cache python3 scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax --enable-lora --lora-target-modules all.
Pasted Graphic 21

Reference

Colab Notebook

Checklist

  • I have added all the necessary unit tests for my change.
  • I have verified that my change does not break existing code and all unit tests pass.
  • I have added all appropriate doc-strings/documentation.
  • My PR is based on the latest changes of the main branch (if unsure, rebase the code).
  • I have signed the Contributor License Agreement.
  • I have followed Contribution Guidelines.

@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 32a06fd to 3af36c3 Compare December 3, 2025 09:54
@aolemila aolemila changed the title [WIP] Feat/add lora for sglangjax [Feature] Feat/add lora for sglangjax Dec 3, 2025
@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 3ea01a4 to a56ef3a Compare December 4, 2025 06:46
@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 3dfcb05 to bb7bd21 Compare December 5, 2025 02:46
@aolemila aolemila mentioned this pull request Dec 5, 2025
Closed
6 tasks
@aolemila aolemila mentioned this pull request Dec 17, 2025
Open
2 tasks
@wang2yn84
Copy link
Collaborator

wang2yn84 commented Dec 19, 2025

Hi @aolemila , thank you for the PR! Can you rebase to head and resolve the conflicts? We've removed the sglang script so it should merge into our main script. And please squash the commits.

wang2yn84
wang2yn84 reviewed Dec 19, 2025
View reviewed changes
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

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

Thank you for your PR! Left some comments.

@aolemila
Copy link
Collaborator Author

aolemila commented Dec 25, 2025

Hi @wang2yn84 , thanks for your reply. I will rebase the main and modify codes according to your advice.

@aolemila
Copy link
Collaborator Author

aolemila commented Dec 25, 2025

I am rerunning scripts and fix new problems I meet.

@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from bb7bd21 to e32615c Compare December 25, 2025 12:44
@aolemila
Copy link
Collaborator Author

aolemila commented Dec 26, 2025

Hi, @wang2yn84 . I have updated codes according to your suggestions. In addition to modifications, I have passed three test cases. You can see more details in PR descriptions.

  • Test1: Run verification of update_params
  • Test2: Run scripts/grpo_demo_llama3_qwen2.py without LoRA
  • Test3: Run scripts/grpo_demo_llama3_qwen2.py with LoRA

JamesBrianD
JamesBrianD reviewed Jan 4, 2026
View reviewed changes
@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 874bfbc to fd80b54 Compare January 7, 2026 11:30
wang2yn84
wang2yn84 reviewed Jan 15, 2026
View reviewed changes
@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from fd80b54 to 401c62f Compare January 15, 2026 08:59
@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 401c62f to 48c01da Compare January 15, 2026 09:01
@aolemila
Copy link
Collaborator Author

aolemila commented Jan 15, 2026

I run python scripts/grpo_demo_llama3_qwen2.py --num-batches 2 --num-test-batches 1 --root-dir=/home/gcpuser/aolemila --rollout-engine sglang_jax --enable-lora --lora-target-modules all.

image

wang2yn84
wang2yn84 reviewed Jan 16, 2026
View reviewed changes
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

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

Much better! Thank you!

@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 48c01da to 5b87a5c Compare January 16, 2026 04:00
wang2yn84
wang2yn84 reviewed Jan 16, 2026
View reviewed changes
Copy link
Collaborator

@wang2yn84 wang2yn84 left a comment

Choose a reason for hiding this comment

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

Thank you! Can you squash the commits?

@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 5b87a5c to 62c6701 Compare January 16, 2026 04:10
@aolemila
Copy link
Collaborator Author

aolemila commented Jan 16, 2026

image

@aolemila
Copy link
Collaborator Author

aolemila commented Jan 16, 2026

Thank you! Can you squash the commits?

done

@aolemila aolemila force-pushed the feat/add-lora-for-sglangjax branch from 62c6701 to b516523 Compare January 16, 2026 04:52
@copybara-service copybara-service bot merged commit f71f2f7 into main Jan 17, 2026
9 checks passed
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Reviewers

@wang2yn84 wang2yn84 wang2yn84 left review comments

@JamesBrianD JamesBrianD JamesBrianD left review comments

@tianshub tianshub Awaiting requested review from tianshub tianshub is a code owner

@lc5211 lc5211 Awaiting requested review from lc5211 lc5211 is a code owner

@hgao327 hgao327 Awaiting requested review from hgao327 hgao327 is a code owner

@sizhit2 sizhit2 Awaiting requested review from sizhit2 sizhit2 is a code owner

@abheesht17 abheesht17 Awaiting requested review from abheesht17 abheesht17 is a code owner

@jiangyangmu jiangyangmu Awaiting requested review from jiangyangmu jiangyangmu is a code owner

Assignees

@aolemila aolemila

Labels

None yet

Projects

None yet

Milestone

No milestone

Development

Successfully merging this pull request may close these issues.

[Feature] Support LoRA For SGLangJax Rollout

3 participants

@aolemila @wang2yn84 @JamesBrianD