Skip to content

Conversation

yfw
Copy link
Contributor

@yfw yfw commented Apr 30, 2025

What does this PR do ?

Add support for gemma-3 models

Issues

Closes #236

Usage

Since Gemma-3 models have tied weights for all models, the TP and SP strategies are configured to not parallelize the embed_tokens and lm_head layers. To run, use the NRL_SKIP_TIED_WEIGHT_CHECK=1 flag. Additionally, use ++policy.generation.vllm_cfg.load_format=auto because vllm has a bug using dummy weights with gemma-3 models.

Before your PR is "Ready for review"

Pre checks:

  • Make sure you read and followed Contributor guidelines
  • Did you write any new necessary tests?
  • Did you run the unit tests and functional tests locally? Visit our Testing Guide for how to run tests
  • Did you add or update any necessary documentation? Visit our Document Development Guide for how to write, build and test the docs.

Additional Information

Sample runs for 1b and 4b: https://wandb.ai/nvidia/grpo-dev-yifu/workspace?nw=7ii089mxlc5
Screenshot 2025-04-30 at 2 34 34 PM

Sample runs using TP/SP: https://wandb.ai/nvidia/grpo-dev-yifu?nw=serntftinh

yfw added 2 commits April 30, 2025 01:14
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
yfw added 5 commits April 30, 2025 12:21
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw yfw marked this pull request as ready for review April 30, 2025 21:33
@yfw yfw changed the title feat: Add support gemma-3-1b and gemma-3-4b feat: Add support gemma-3 May 6, 2025
yfw added 2 commits May 6, 2025 10:16
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw yfw changed the title feat: Add support gemma-3 feat: Add support for gemma-3 May 6, 2025
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw yfw requested review from SahilJain314 and gshennvm May 6, 2025 17:52
yfw added 2 commits May 6, 2025 17:51
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@terrykong terrykong added the CI:L1 Run doctests, unit tests, and functional tests label May 7, 2025
@terrykong terrykong mentioned this pull request May 7, 2025
@yfw yfw added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels May 7, 2025
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw yfw added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels May 7, 2025
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
yfw added 2 commits May 7, 2025 15:54
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
@yfw yfw added CI:L1 Run doctests, unit tests, and functional tests and removed CI:L1 Run doctests, unit tests, and functional tests labels May 8, 2025
@terrykong terrykong added this pull request to the merge queue May 8, 2025
@github-merge-queue github-merge-queue bot removed this pull request from the merge queue due to a conflict with the base branch May 8, 2025
@terrykong terrykong added this pull request to the merge queue May 8, 2025
Merged via the queue into main with commit f963d9b May 8, 2025
21 checks passed
@terrykong terrykong deleted the yifu/gemma branch May 8, 2025 02:56
YzjiaoNvd pushed a commit to YzjiaoNvd/NeMo-RL that referenced this pull request Jun 10, 2025
Signed-off-by: Yi-Fu Wu <yifu.wu@gmail.com>
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment
Labels
CI:L1 Run doctests, unit tests, and functional tests
Projects
None yet
Development

Successfully merging this pull request may close these issues.

gemma-3
2 participants