[
  {
    "path": ".gemini/config.yaml",
    "content": "have_fun: false\ncode_review:\n  disable: false\n  comment_severity_threshold: HIGH\n  max_review_comments: -1\n  pull_request_opened:\n    help: false\n    summary: false\n    code_review: true\nignore_patterns: []\n"
  },
  {
    "path": ".git-blame-ignore-revs",
    "content": "# Local uasge: git config blame.ignoreRevsFile .git-blame-ignore-revs\n\n# [dev] feat: immigrate from yapf & pylint to ruff based on pre-commit\n# Changed 268 files, +10k/-9k lines. This is the biggest formatter change.\nb00f77d8559b48d57a33c0132a5ba1c81891a536\n\n# [ci] refactor: reduce ruff line-length from 300 to 120\n# Changed 238 files, +6k/-1k lines. Global formatting change.\n00a10a8ef389556f957a2f36132b2358fd6a109f\n\n# [Lint] fix: linting errors in all files\n# Changed 179 files, +1k/-3k lines. Global lint fix.\n8e5ad4688a13de81727c014a3c2e2fb26324bc20\n"
  },
  {
    "path": ".github/CODEOWNERS",
    "content": "/docs @eric-haibin-lin @zhaochenyang20 @hongpeng-guo\n/docs/amd_tutorial @yushengsu-thu\n/docs/slang_multiturn @zhaochenyang20 @SwordFaith\n/docs/ascend_tutorial @FightingZhen\n\n/third_party/sglang @zhaochenyang20 @SwordFaith\n/third_party/vllm @PeterSH6 @wuxibin89\n\n/examples/grpo_trainer @vermouth1992 @PeterSH6 @tardis-key @FightingZhen @ji-huazhong\n\n/verl/single_controller @zw0610 @wuxibin89 @hongpeng-guo\n/verl/trainer @eric-haibin-lin @vermouth1992 @tongyx361 @PeterSH6\n/verl/models/mcore @ISEEKYAN @vermouth1992\n/verl/models/transformers @vermouth1992 @PeterSH6 @tardis-key @FightingZhen @ji-huazhong\n/verl/workers/engine @eric-haibin-lin @vermouth1992 @ZihengJiang\n/verl/workers/roles @eric-haibin-lin @vermouth1992 @ZihengJiang\n/verl/workers/engine/fsdp @eric-haibin-lin @vermouth1992 @ZihengJiang\n/verl/workers/rollout/vllm_rollout @wuxibin89 @PeterSH6 @chenhaiq\n/verl/workers/rollout/sglang_rollout @zhaochenyang20 @SwordFaith @chenhaiq\n/verl/workers/actor/megatron_actor.py @ISEEKYAN @vermouth1992\n/verl/workers/critic/megatron_critic.py @ISEEKYAN @vermouth1992\n/verl/workers/megatron_workers.py @ISEEKYAN @vermouth1992\n/verl/experimental @wuxibin89 @ArronHZG\n\n/tests/single_controller @zw0610 @wuxibin89\n/tests/trainer @eric-haibin-lin @vermouth1992 @tongyx361 @PeterSH6\n/tests/workers/rollout/vllm_rollout @wuxibin89 @PeterSH6 @chenhaiq\n"
  },
  {
    "path": ".github/ISSUE_TEMPLATE/bug-report.yml",
    "content": "# modified from https://github.com/huggingface/transformers/blob/main/.github/ISSUE_TEMPLATE/bug-report.yml?plain=1\nname: \"\\U0001F41B Bug Report\"\ndescription: Submit a bug report to help us improve verl\nlabels: [ \"bug\" ]\nbody:\n  - type: markdown\n    attributes:\n      value: |\n        Thanks for taking the time to fill out this bug report! 🤗\n\n  - type: textarea\n    id: system-info\n    attributes:\n      label: System Info\n      description: Please share your system info with us. You can run the command `python scripts/diagnose.py` and copy-paste its output below.\n      placeholder: verl version, platform, python version, ...\n    validations:\n      required: true\n\n  - type: checkboxes\n    id: information-scripts-examples\n    attributes:\n      label: Information\n      description: 'The problem arises when using:'\n      options:\n        - label: \"The official example scripts\"\n        - label: \"My own modified scripts\"\n\n  - type: checkboxes\n    id: information-tasks\n    attributes:\n      label: Tasks\n      description: \"The tasks I am working on are:\"\n      options:\n        - label: \"An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)\"\n        - label: \"My own task or dataset (give details below)\"\n\n  - type: textarea\n    id: reproduction\n    validations:\n      required: true\n    attributes:\n      label: Reproduction\n      description: |\n        Please provide a code sample that reproduces the problem you ran into. It can be a Colab link or just a code snippet.\n        Please include relevant config information with your code.\n        If you have code snippets, error messages, stack traces please provide them here as well.\n        Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting\n        Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.\n\n      placeholder: |\n        Steps to reproduce the behavior:\n\n          1.\n          2.\n          3.\n\n\n  - type: textarea\n    id: expected-behavior\n    validations:\n      required: true\n    attributes:\n      label: Expected behavior\n      description: \"A clear and concise description of what you would expect to happen.\""
  },
  {
    "path": ".github/ISSUE_TEMPLATE/config.yml",
    "content": "blank_issues_enabled: true\nversion: 0.1\n"
  },
  {
    "path": ".github/ISSUE_TEMPLATE/feature-request.yml",
    "content": "# modified from https://github.com/huggingface/transformers/blob/main/.github/ISSUE_TEMPLATE/feature-request.yml?plain=1\nname: \"\\U0001F680 Feature request\"\ndescription: Submit a proposal/request for a new verl feature\nlabels: [ \"Feature request\" ]\nbody:\n  - type: textarea\n    id: feature-request\n    validations:\n      required: true\n    attributes:\n      label: Feature request\n      description: |\n        A clear and concise description of the feature proposal. Please provide a link to the paper and code in case they exist.\n\n  - type: textarea\n    id: motivation\n    validations:\n      required: true\n    attributes:\n      label: Motivation\n      description: |\n        Please outline the motivation for the proposal. Is your feature request related to a problem? e.g., I'm always frustrated when [...]. If this is related to another GitHub issue, please link here too.\n\n\n  - type: textarea\n    id: contribution\n    validations:\n      required: true\n    attributes:\n      label: Your contribution\n      description: |\n        Is there any way that you could help, e.g. by submitting a PR? Make sure to read the CONTRIBUTING.MD [readme](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md)"
  },
  {
    "path": ".github/PULL_REQUEST_TEMPLATE.md",
    "content": "### What does this PR do?\n\n> Add **concise** overview of what this PR aims to achieve or accomplish. Reference related GitHub issues and PRs that help with the review.\n\n### Checklist Before Starting\n\n- [ ] Search for similar PRs. Paste at least one query link here: ...\n- [ ] Format the PR title as `[{modules}] {type}: {description}` (This will be checked by the CI)\n  - `{modules}` include `fsdp`, `megatron`, `veomni`, `sglang`, `vllm`, `rollout`, `trainer`, `ci`, `training_utils`, `recipe`, `hardware`, `deployment`, `ray`, `worker`, `single_controller`, `misc`, `perf`, `model`, `algo`, `env`, `tool`, `ckpt`, `doc`, `data`, `cfg`, `reward`, `fully_async`, `one_step_off`\n  - If this PR involves multiple modules, separate them with `,` like `[megatron, fsdp, doc]`\n  - `{type}` is in `feat`, `fix`, `refactor`, `chore`, `test`\n  - If this PR breaks any API (CLI arguments, config, function signature, etc.), add `[BREAKING]` to the beginning of the title.\n  - Example: `[BREAKING][fsdp, megatron] feat: dynamic batching`\n\n### Test\n\n> For changes that can not be tested by CI (e.g., algorithm implementation, new model support), validate by experiment(s) and show results like training curve plots, evaluation results, etc.\n\n### API and Usage Example\n\n> Demonstrate how the API changes if any, and provide usage example(s) if possible.\n\n```python\n# Add code snippet or script demonstrating how to use this\n```\n\n### Design & Code Changes\n\n> Demonstrate the high-level design if this PR is complex, and list the specific changes.\n\n### Checklist Before Submitting\n\n> [!IMPORTANT]\n> Please check all the following items before requesting a review, otherwise the reviewer might deprioritize this PR for review.\n\n- [ ] Read the [Contribute Guide](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md).\n- [ ] Apply [pre-commit checks](https://github.com/volcengine/verl/blob/main/CONTRIBUTING.md#code-linting-and-formatting): `pre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always`\n- [ ] Add / Update [the documentation](https://github.com/volcengine/verl/tree/main/docs).\n- [ ] Add unit or end-to-end test(s) to [the CI workflow](https://github.com/volcengine/verl/tree/main/.github/workflows) to cover all the code. If not feasible, explain why: ...\n- [ ] Once your PR is ready for CI, send a message in [the `ci-request` channel](https://verl-project.slack.com/archives/C091TCESWB1) in [the `verl` Slack workspace](https://join.slack.com/t/verl-project/shared_invite/zt-3855yhg8g-CTkqXu~hKojPCmo7k_yXTQ). (If not accessible, please try [the Feishu group (飞书群)](https://applink.larkoffice.com/client/chat/chatter/add_by_link?link_token=772jd4f1-cd91-441e-a820-498c6614126a).)\n- [ ] If your PR is related to the `recipe` submodule, please also update the reference to the submodule commit via `git submodule update --remote` or `cd recipe && git pull origin main`.\n"
  },
  {
    "path": ".github/dependabot.yml",
    "content": "## Enabled the dependabot to check the dependencies of the project\n## Dependabot will open pull requests to update dependencies automatically\n\nversion: 2\nupdates:\n  - package-ecosystem: pip\n    directory: \"/\"\n    schedule:\n      interval: weekly"
  },
  {
    "path": ".github/workflows/README.md",
    "content": "### Adding a New Workflow\n\nWhen adding a new workflow for continuous integration (CI), you have two runner options: a fixed runner or a machine from the vemlp.\n\n- **Fixed Runner**: To use a fixed runner, specify it in your workflow using the `runs-on` keyword, like `runs-on: [L20x8]`. \n- **Vemlp Runner**: Opting for a Vemlp machine allows you to launch tasks elastically. \n\nHere is a template to assist you. This template is designed for using Vemlp machines. Currently, for each workflow, you need to create a `setup` and a `cleanup` job. When using this template, the main parts you need to modify are the `IMAGE` environment variable and the specific `job steps`.\n\n```yaml\nname: Your Default Workflow\n\non:\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \".github/workflows/template.yml\"\n\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"your vemlp image\" # e.g. \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_URL: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\" # public veFaas api\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      task-id: ${{ steps.create-runner.outputs.task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1 \n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_URL }}\"\n          image: \"${{ env.DEFAULT_IMAGE }}\"\n\n  your_job:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'default-runner' }}\"]\n    steps:\n      xxxx # your jobs\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, your_job]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_URL }}\"\n          task-id: \"${{ needs.setup.outputs.task-id }}\"\n```\n\n### Model and Dataset\nTo avoid CI relies on network, we pre-download dataset on a NFS on the CI machine. The path for models are \\${HOME}/models and the path for dataset is \\${HOME}/models/hf_data."
  },
  {
    "path": ".github/workflows/check-pr-title.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests \n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\n\non:\n  pull_request:\n    types: [opened, edited, synchronize]\n\njobs:\n  check-title:\n    runs-on: ubuntu-latest\n    steps:\n      - name: Checkout code\n        uses: actions/checkout@v4\n\n      - name: Set up Python\n        uses: actions/setup-python@v5\n        with:\n          python-version: '3.11'\n\n      - name: Run PR title checker\n        run: python3 tests/special_sanity/check_pr_title.py\n        env:\n          PR_TITLE: ${{ github.event.pull_request.title }}\n\n      - name: Run PR description checker\n        run: python3 tests/special_sanity/check_pr_description.py\n        env:\n          PR_TITLE: ${{ github.event.pull_request.title }}\n          GITHUB_EVENT_PATH: ${{ github.event_path }}\n"
  },
  {
    "path": ".github/workflows/cpu_unit_tests.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: cpu_unit_tests\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - .github/workflows/cpu_unit_tests.yml\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  cpu_unit_tests:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 20 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      TORCH_COMPILE_DISABLE: 1\n      TORCHINDUCTOR_DISABLE: 1\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install --upgrade \"transformers>=5.0.0\"\n      - name: Download datasets\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n          python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k\n      - name: Running CPU unit tests\n        run: |\n          echo '[pytest]' > pytest.ini\n          echo 'python_files = *_on_cpu.py' >> pytest.ini\n          pytest -s -x --asyncio-mode=auto tests/\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, cpu_unit_tests]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/doc.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests \n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\n\nname: doc_test\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"docs/**\"\n      - .github/workflows/doc.yml\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read      # for checkout\n  pages: write        # for deploy-pages\n  id-token: write     # for deploy-pages\n\njobs:\n  doc_test:\n    runs-on: ubuntu-latest\n    timeout-minutes: 5 # Increase this timeout value as needed\n    strategy:\n      matrix:\n        python-version: [\"3.10\"]\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n      - name: Set up Python ${{ matrix.python-version }}\n        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n        with:\n          python-version: ${{ matrix.python-version }}\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip install -r docs/requirements-docs.txt\n\n      - name: Run doc make html\n        run: |\n          cd docs \n          make clean\n          make html SPHINXOPTS=\"--keep-going -w _build/sphinx.log\"\n          if grep -q \": ERROR:\" _build/sphinx.log; then\n            echo \"🚨 Sphinx doc build contained ERRORs - see _build/sphinx.log\"\n            exit 1\n          fi\n          if grep -q \"WARNING: document isn't included in any toctree\" _build/sphinx.log; then\n            echo \"🚨 Sphinx doc build contained WARNING. Please include newly added docs in index.rst. See _build/sphinx.log for details\"\n            exit 1\n          fi\n          if grep -q \"WARNING: Inline emphasis\" _build/sphinx.log; then\n            echo \"🚨 Sphinx doc build contained WARNING. Please check inline emphasis is correct. See _build/sphinx.log for details\"\n            exit 1\n          fi\n          if grep -q \"WARNING: Definition list ends without a blank line\" _build/sphinx.log; then\n            echo \"🚨 Sphinx doc build contained WARNING. Please check if the indentation is correct. See _build/sphinx.log for details\"\n            exit 1\n          fi\n"
  },
  {
    "path": ".github/workflows/docker-build-ascend-a2.yml",
    "content": "name: docker-build-ascend-a2\n\non:\n  workflow_dispatch:\n  push:\n    branches: [\"main\"]\n    paths:\n      - \"docker/ascend/Dockerfile.ascend_8.5.0_a2\"\n      - \".github/workflows/docker-build-ascend-a2.yml\"\n  release:\n    types: [published]\n  schedule:\n    - cron: \"0 16 * * *\"\n\njobs:\n  build-ascend-image-a2:\n    if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }}\n    runs-on: ubuntu-latest\n    concurrency:\n      group: ${{ github.workflow }}-${{ github.ref }}-build-ascend-image-a2\n      cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n    steps:\n      - name: Remove unnecessary parts in github actions runners to free up disk space\n        uses: jlumbroso/free-disk-space@v1.3.1\n        with:\n          tool-cache: true\n\n      - name: Checkout code\n        uses: actions/checkout@v4\n\n      - name: Set up Python\n        uses: actions/setup-python@v5\n        with:\n          python-version: \"3.11\"\n\n      - name: Get base image name and tag\n        id: base_image\n        run: |\n          BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend_8.5.0_a2 | head -1 | cut -d' ' -f2)\n          echo \"Base image full: $BASE_IMAGE_FULL\" \n          BASE_IMAGE_TAG=$(echo \"$BASE_IMAGE_FULL\" | cut -d':' -f2)\n          echo \"Base image tag: $BASE_IMAGE_TAG\"\n          NEW_IMAGE_NAME=\"verl-$BASE_IMAGE_TAG\"\n          echo \"New image name: $NEW_IMAGE_NAME\"  \n          echo \"base_image_tag=$BASE_IMAGE_TAG\" >> \"$GITHUB_OUTPUT\"\n          echo \"new_image_name=$NEW_IMAGE_NAME\" >> \"$GITHUB_OUTPUT\"\n\n      - name: Get image tag\n        id: version\n        run: |\n          BRANCH_NAME=$(echo \"${{ github.ref }}\" | sed 's/refs\\/heads\\///g' | sed 's/[^a-zA-Z0-9._-]/_/g')\n          if [ \"${{ github.event_name }}\" = \"release\" ]; then\n            echo \"tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}\" >> \"$GITHUB_OUTPUT\"\n          elif [ \"$BRANCH_NAME\" = \"main\" ]; then\n            echo \"tag=${{ steps.base_image.outputs.new_image_name }}-latest\" >> \"$GITHUB_OUTPUT\"\n          fi\n\n      - name: Set up Docker Buildx\n        uses: docker/setup-buildx-action@v3\n\n      - name: Login to Quay.io\n        uses: docker/login-action@v3\n        with:\n          registry: quay.io\n          username: ${{ secrets.QUAY_USERNAME }}\n          password: ${{ secrets.QUAY_PASSWORD }}\n\n      - name: Clean Docker cache before build\n        run: |\n          docker system prune -a -f --volumes || true\n\n      - name: Build and push images Quay\n        uses: docker/build-push-action@v6\n        with:\n          context: .\n          platforms: linux/amd64,linux/arm64\n          file: ./docker/ascend/Dockerfile.ascend_8.5.0_a2\n          push: true\n          tags: |\n            quay.io/ascend/verl:${{ steps.version.outputs.tag }}\n          cache-from: type=gha\n          cache-to: type=gha,mode=max\n          build-args: |\n            BUILDKIT_INLINE_CACHE=1\n"
  },
  {
    "path": ".github/workflows/docker-build-ascend-a3.yml",
    "content": "name: docker-build-ascend-a3\n\non:\n  workflow_dispatch:\n  push:\n    branches: [\"main\"]\n    paths:\n      - \"docker/ascend/Dockerfile.ascend_8.5.0_a3\"\n      - \".github/workflows/docker-build-ascend-a3.yml\"\n  release:\n    types: [published]\n  schedule:\n    - cron: \"0 19 * * *\"\n\njobs:\n  build-ascend-image-a3:\n    if: ${{ github.event_name != 'pull_request' && github.repository_owner == 'verl-project' }}\n    runs-on: ubuntu-latest\n    concurrency:\n      group: ${{ github.workflow }}-${{ github.ref }}-build-ascend-image-a3\n      cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n    steps:\n      - name: Remove unnecessary parts in github actions runners to free up disk space\n        uses: jlumbroso/free-disk-space@v1.3.1\n        with:\n          tool-cache: true\n\n      - name: Checkout code\n        uses: actions/checkout@v4\n\n      - name: Set up Python\n        uses: actions/setup-python@v5\n        with:\n          python-version: \"3.11\"\n\n      - name: Get base image name and tag\n        id: base_image\n        run: |\n          BASE_IMAGE_FULL=$(grep '^FROM' ./docker/ascend/Dockerfile.ascend_8.5.0_a3 | head -1 | cut -d' ' -f2)\n          echo \"Base image full: $BASE_IMAGE_FULL\" \n          BASE_IMAGE_TAG=$(echo \"$BASE_IMAGE_FULL\" | cut -d':' -f2)\n          echo \"Base image tag: $BASE_IMAGE_TAG\"\n          NEW_IMAGE_NAME=\"verl-$BASE_IMAGE_TAG\"\n          echo \"New image name: $NEW_IMAGE_NAME\"  \n          echo \"base_image_tag=$BASE_IMAGE_TAG\" >> \"$GITHUB_OUTPUT\"\n          echo \"new_image_name=$NEW_IMAGE_NAME\" >> \"$GITHUB_OUTPUT\"\n\n      - name: Get image tag\n        id: version\n        run: |\n          BRANCH_NAME=$(echo \"${{ github.ref }}\" | sed 's/refs\\/heads\\///g' | sed 's/[^a-zA-Z0-9._-]/_/g')\n          if [ \"${{ github.event_name }}\" = \"release\" ]; then\n            echo \"tag=${{ steps.base_image.outputs.new_image_name }}-${{ github.event.release.tag_name }}\" >> \"$GITHUB_OUTPUT\"\n          elif [ \"$BRANCH_NAME\" = \"main\" ]; then\n            echo \"tag=${{ steps.base_image.outputs.new_image_name }}-latest\" >> \"$GITHUB_OUTPUT\"\n          fi\n\n      - name: Set up Docker Buildx\n        uses: docker/setup-buildx-action@v3\n\n      - name: Login to Quay.io\n        uses: docker/login-action@v3\n        with:\n          registry: quay.io\n          username: ${{ secrets.QUAY_USERNAME }}\n          password: ${{ secrets.QUAY_PASSWORD }}\n\n      - name: Clean Docker cache before build\n        run: |\n          docker system prune -a -f --volumes || true\n\n      - name: Build and push images Quay\n        uses: docker/build-push-action@v6\n        with:\n          context: .\n          platforms: linux/amd64,linux/arm64\n          file: ./docker/ascend/Dockerfile.ascend_8.5.0_a3\n          push: true\n          tags: |\n            quay.io/ascend/verl:${{ steps.version.outputs.tag }}\n          cache-from: type=gha\n          cache-to: type=gha,mode=max\n          build-args: |\n            BUILDKIT_INLINE_CACHE=1\n"
  },
  {
    "path": ".github/workflows/e2e_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n    paths:\n      - \".github/workflows/e2e_ascend.yml\"\n      - \"examples/data_preprocess/**\"\n      - \"examples/grpo_trainer/**\"\n      - \"examples/ppo_trainer/**\"\n      - \"examples/sft/**\"\n      - \"verl/experimental/one_step_off_policy/**\"\n      - \"tests/special_npu/**\"\n      - \"tests/special_sanity/check_device_api_usage.py\"\n      - \"verl/**\"\n      - \"pyproject.toml\"\n      - \"requirements-npu.txt\"\n      - \"setup.py\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\npermissions:\n  contents: read\n\njobs:\n  llm_rl_job:\n    if: github.repository_owner == 'verl-project'\n    name: E2E Ascend testing for RL training scenarios of LLM models\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 120\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout volcengine/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Preprocess gsm8k dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running gsm8k e2e training tests with PPO on ASCEND NPU (FSDP backend)\n        run: |\n          ray stop --force\n          bash tests/special_npu/run_qwen3_06b_ppo.sh\n          rm -rf $HOME/ckpts\n      - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (FSDP backend)\n        run: |\n          ray stop --force\n          bash tests/special_npu/run_qwen2_5_05b_grpo.sh\n          rm -rf $HOME/ckpts\n      - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeed backend)\n        run: |\n          ray stop --force\n          USE_DIST_CKPT=True bash tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh\n          rm -rf $HOME/dist_ckpt/qwen2_5_05b_grpo_mindspeed\n          rm -rf $HOME/ckpts\n      - name: Running gsm8k e2e training tests with GRPO on ASCEND NPU (MindSpeed backend, MoE Model)\n        run: |\n          ray stop --force\n          USE_DIST_CKPT=True USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json DUMMY_MODEL_PATH=$HOME/dist_ckpt/qwen3_30b_grpo_mindspeed bash tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh\n\n  vlm_rl_job:\n    if: github.repository_owner == 'verl-project'\n    name: E2E Ascend testing for RL training scenarios of VLM models\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 120\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout volcengine/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Preprocess geo3k dataset\n        run: |\n          python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k\n      - name: Running geo3k e2e training tests with GRPO on ASCEND NPU\n        run: |\n          ray stop --force\n          bash tests/special_npu/run_qwen2_5_vl_3b_npu.sh\n          rm -rf $HOME/ckpts\n"
  },
  {
    "path": ".github/workflows/e2e_fully_async_policy.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_fully_async_policy\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/*trainer*\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/experimental/fully_async_policy\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Home\n      - \"verl/experimental/fully_async_policy\"\n      # Entrypoints\n      - \".github/workflows/e2e_fully_async_policy.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/run_fully_async_policy.sh\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  # Test FSDP2 strategy\n  e2e_fully_async_policy_fsdp2:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 10 # Increase timeout for async training\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"fsdp2\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install cupy-cuda12x==13.6.0\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running the E2E test with fully_async_policy algorithm (FSDP2)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_fully_async_policy.sh\n\n  # Test Megatron strategy\n  e2e_fully_async_policy_megatron:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 10 # Increase timeout for async training\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"megatron\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install cupy-cuda12x==13.6.0\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running the E2E test with fully_async_policy algorithm (Megatron)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_fully_async_policy.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, e2e_fully_async_policy_fsdp2]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_fully_async_policy_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_fully_async_policy_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/*trainer*\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/experimental/fully_async_policy\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Home\n      - \"verl/experimental/fully_async_policy\"\n      # Entrypoints\n      - \".github/workflows/e2e_fully_async_policy_ascend.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/run_fully_async_policy.sh\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  # Test FSDP2 strategy\n  e2e_fully_async_policy_fsdp2_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"fsdp2\"\n      device_name: \"npu\"\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running the E2E test with fully_async_policy algorithm (FSDP2)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_fully_async_policy.sh\n\n  # Test Megatron strategy\n  e2e_fully_async_policy_megatron_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"megatron\"\n      device_name: \"npu\"\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running the E2E test with fully_async_policy algorithm (Megatron)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_fully_async_policy.sh\n"
  },
  {
    "path": ".github/workflows/e2e_one_step_off_policy.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_one_step_off_policy\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/*trainer*\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/experimental/one_step_off_policy\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Home\n      - \"verl/experimental/one_step_off_policy\"\n      # Entrypoints\n      - \".github/workflows/e2e_one_step_off_policy.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/run_one_step_off_policy.sh\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  # Test FSDP2 strategy\n  e2e_one_step_off_policy_fsdp2:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 10 # Increase timeout for async training\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"fsdp2\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install cupy-cuda12x==13.6.0\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running the E2E test with one_step_off_policy algorithm (FSDP2)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_one_step_off_policy.sh\n\n  # Test Megatron strategy\n  e2e_one_step_off_policy_megatron:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 10 # Increase timeout for async training\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"megatron\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install cupy-cuda12x==13.6.0\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running the E2E test with one_step_off_policy algorithm (Megatron)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_one_step_off_policy.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [setup, e2e_one_step_off_policy_fsdp2, e2e_one_step_off_policy_megatron]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_one_step_off_policy_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_one_step_off_policy_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/*trainer*\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/experimental/one_step_off_policy\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - \"!**/*.md\"\n      - \"!**/*.sh\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Home\n      - \"verl/experimental/one_step_off_policy\"\n      # Entrypoints\n      - \".github/workflows/e2e_one_step_off_policy_ascend.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/run_one_step_off_policy.sh\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  # Test FSDP2 strategy\n  e2e_one_step_off_policy_fsdp2_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"fsdp2\"\n      device_name: \"npu\"\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running the E2E test with one_step_off_policy algorithm (FSDP2)\n        run: |\n          ray stop --force\n          bash tests/special_e2e/run_one_step_off_policy.sh\n\n  # Test Megatron strategy\n  e2e_one_step_off_policy_megatron_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ACTOR_STRATEGY: \"megatron\"\n      device_name: \"npu\"\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running the E2E test with one_step_off_policy algorithm (Megatron)\n        run: |\n          ray stop --force\n          export PYTHONPATH=$PYTHONPATH:/Megatron-LM\n          bash tests/special_e2e/run_one_step_off_policy.sh\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_grpo_trainer_trtllm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_trtllm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Recipes\n      - \"!recipe/**\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Recipes\n      - \"!recipe/**\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      # Entrypoints\n      - \"verl/workers/rollout/trtllm_rollout/**\"\n      - \"tests/workers/rollout/rollout_trtllm/**\"\n      - \".github/workflows/e2e_ppo_grpo_trainer_trtllm.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"examples/data_preprocess/dapo_multiturn_w_tool.py\"\n      - \"examples/data_preprocess/aime2024_multiturn_w_tool.py\"\n      - \"examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh\"\n      - \"examples/grpo_trainer/run_qwen2-7b_math_megatron_trtllm.sh\"\n      - \"examples/grpo_trainer/run_qwen3-30b_dapo_megatron_fp8_trtllm.sh\"\n      # add back when ppo flow is ready\n      # - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      # - \"verl/trainer/main_ppo.py\"\n      # - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:trtllm1.3.0rc4\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  trtllm_unit_tests:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install pytest-asyncio\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Run TRTLLM unit tests\n        run: |\n          export TRTLLM_TEST_MODEL_PATH_ROOT=\"${HOME}/models\"\n          ray stop --force\n          pytest -v -s \\\n            tests/workers/rollout/rollout_trtllm/test_adapter.py \\\n            tests/workers/rollout/rollout_trtllm/test_async_server.py \\\n            tests/workers/rollout/rollout_trtllm/test_trtllm_rollout_utils.py\n\n  e2e_grpo_trainer_fsdp-qwen2:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_save_dir ${PWD}/data/gsm8k\n      - name: Running GSM8K E2E training tests with FSDP on 8 L20 GPUs (Qwen)\n        run: |\n          ray stop --force\n          DATADIR=${HOME}/data \\\n            bash examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh 2 \\\n            trainer.total_training_steps=1 \\\n            data.train_files=\"['${PWD}/data/gsm8k/train.parquet']\" \\\n            data.val_files=\"['${PWD}/data/gsm8k/test.parquet']\" \\\n            trainer.logger='[\"console\"]' \\\n            actor_rollout_ref.model.path=\"${HOME}/models/Qwen/Qwen2.5-0.5B-Instruct\"\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  e2e_grpo_trainer_megatron-qwen2:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_save_dir ${PWD}/data/gsm8k\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen)\n        run: |\n          ray stop --force\n          DATADIR=${HOME}/data \\\n          ACTOR_TP=2 \\\n            bash examples/grpo_trainer/run_qwen2-7b_math_megatron_trtllm.sh 2 \\\n            trainer.total_training_steps=1 \\\n            data.train_files=\"['${PWD}/data/gsm8k/train.parquet']\" \\\n            data.val_files=\"['${PWD}/data/gsm8k/test.parquet']\" \\\n            trainer.logger='[\"console\"]' \\\n            actor_rollout_ref.model.path=\"${HOME}/models/Qwen/Qwen2.5-0.5B-Instruct\"\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n  e2e_grpo_trainer_fsdp-vlm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install qwen_vl_utils \n          pip3 install mathruler\n      - name: Prepare GEO3K dataset\n        run: |\n          python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k --local_save_dir ${PWD}/data/geo3k\n      - name: Running GEO3K E2E training tests with FSDP on 8 L20 GPUs (VLM)\n        run: |\n          ray stop --force\n          DATADIR=${HOME}/data \\\n            bash examples/grpo_trainer/run_qwen2_5_vl_3b_trtllm.sh 2 \\\n            trainer.total_training_steps=1 \\\n            data.train_files=\"['${PWD}/data/geo3k/train.parquet']\" \\\n            data.val_files=\"['${PWD}/data/geo3k/test.parquet']\" \\\n            trainer.logger='[\"console\"]' \\\n            actor_rollout_ref.model.path=\"${HOME}/models/Qwen/Qwen3-VL-2B-Instruct\"\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n      - name: Prepare DAPO-Math-17k and AIME-2024 datasets (data_preprocess)\n        run: |\n          python3 examples/data_preprocess/dapo_multiturn_w_tool.py --local_save_dir ${PWD}/data/dapo-math-17k\n          python3 examples/data_preprocess/aime2024_multiturn_w_tool.py --local_save_dir ${PWD}/data/aime-2024\n      - name: Running DAPO E2E with FP8 TRT-LLM rollout (Qwen3-0.6B)\n        run: |\n          ray stop --force\n          export INFER_TP=2 ACTOR_TP=2 ACTOR_PP=2 ACTOR_VPP=2 ACTOR_EP=1 ACTOR_CP=2 REF_TP=2 REF_PP=2 REF_VPP=2 REF_EP=1 REF_CP=2 GEN_MOE_TP=null GEN_MOE_EP=null\n          export NNODES=1 GPUS_PER_NODE=8 TRTLLM_MOE_BACKEND=CUTLASS\n          export DATA_DIR=${PWD} DAPO_MATH_TRAIN=${PWD}/data/dapo-math-17k/train.parquet AIME_VAL=${PWD}/data/aime-2024/train.parquet MODEL_PATH=${HOME}/models/Qwen/Qwen3-0.6B\n          bash examples/grpo_trainer/run_qwen3-30b_dapo_megatron_fp8_trtllm.sh \\\n            reward_model.reward_kwargs.overlong_buffer_cfg.len=258 \\\n            reward_model.reward_kwargs.max_resp_len=512 \\\n            data.max_prompt_length=512 \\\n            data.max_response_length=512 \\\n            data.train_batch_size=32 \\\n            actor_rollout_ref.rollout.n=4 \\\n            actor_rollout_ref.rollout.max_num_seqs=16 \\\n            actor_rollout_ref.rollout.max_num_batched_tokens=1024 \\\n            actor_rollout_ref.rollout.max_model_len=1024 \\\n            actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=False \\\n            actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=False \\\n            trainer.total_training_steps=1 \\\n            trainer.logger='[\"console\"]'\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, trtllm_unit_tests, e2e_grpo_trainer_fsdp-qwen2, e2e_grpo_trainer_megatron-qwen2, e2e_grpo_trainer_fsdp-vlm]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer.yml",
    "content": "name: e2e_ppo_trainer\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!**/*.md\"\n      - \"!docker/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Docs\n      - \"!docs/**\"\n\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_ppo_trainer.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/ppo_trainer\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  pre_commit_for_ppo:\n    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n        python-version: [\"3.12\"]\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n      - name: Set up Python ${{ matrix.python-version }}\n        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n        with:\n          python-version: ${{ matrix.python-version }}\n      - name: Install the current repository\n        run: |\n          pip install pre-commit hydra-core\n          pip3 install --no-deps -e .\n      - name: Set ruff --output-format=github\n        run: |\n          sed -i 's/--output-format=full/--output-format=github/' .pre-commit-config.yaml\n          git add .pre-commit-config.yaml\n      - uses: pre-commit/action@v3.0.1\n        with:\n          extra_args: \"\" # Overriding default \"--all-files\"\n\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_megatron_sglang.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_sglang\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\" # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\" # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n      # Entrypoints\n      - \"verl/worksers/rollout/sglang_rollout/*\"\n      - \".github/workflows/e2e_ppo_trainer_megatron_sglang.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  e2e_ppo_trainer_megatron-deepseek:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ENGINE: sglang\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation\n          pip3 install --no-deps -e .\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (DeepSeek)\n        run: |\n          ray stop --force\n          OPTIM_MEMORY_EFFICIENT=True ENGINE=sglang SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (DeepSeek)\n        run: |\n          ray stop --force\n          export VLLM_USE_V1=1\n          ray start --head\n          ENGINE=sglang MODE=async RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Profiling GRPO GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Deepseek)\n        run: |\n          ray stop --force\n          PROFILE_ENABLE=True ENGINE=sglang ADV_ESTIMATOR=grpo USE_DYNAMIC_BSZ=False MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct bash tests/special_e2e/run_ppo_trainer_megatron.sh\n          if [ -z \"$( ls -A '/tmp/ray/session_latest/logs/nsight/' )\" ]; then\n            echo \"[ERROR] not found any profiling files\"\n            exit 1\n          else\n            echo \"[SUCCESS] profile success\"\n          fi\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  # Qwen3-0.6B: dense, tie_word_embeddings=True\n  e2e_ppo_trainer_megatron-qwen3:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      ENGINE: sglang\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen3) testing learning rate scheduler\n        run: |\n          ray stop --force\n          ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout\n        run: |\n          ray stop --force\n          export VLLM_USE_V1=1\n          ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [setup, e2e_ppo_trainer_megatron-deepseek, e2e_ppo_trainer_megatron-qwen3]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_megatron_sglang_2.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_sglang_2\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\" # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\" # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n      # Entrypoints\n      - \"verl/worksers/rollout/sglang_rollout/*\"\n      - \".github/workflows/e2e_ppo_trainer_megatron_sglang.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  e2e_ppo_trainer_fsdp_sglang:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 40 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm and save ckpt\n        run: |\n          ray stop --force\n          ENGINE=sglang bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n  e2e_ppo_trainer_fsdp-qwen2_5vl-3b:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      # Geo3k\n      - name: Prepare GEO3K dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/\n      - name: Running GEO3K VLM E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \\\n            ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GEO3K VLM E2E with rmpad using torch fused kernel (Qwen2.5-VL)\n        run: |\n          ray stop --force\n          FUSED_KERNELS=True TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \\\n            ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GEO3K VLM E2E with rmpad using triton fused kernel (Qwen2.5-VL)\n        run: |\n          ray stop --force\n          FUSED_KERNELS=True FUSED_KERNEL_BACKEND=triton \\\n            TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            ENGINE=sglang ROLLOUT_MODE=async GPU_MEMORY_UTILIZATION=0.6 ACTOR_FSDP_PARAM_OFFLOAD=True \\\n            ACTOR_FSDP_OPTIMIZER_OFFLOAD=True REF_FSDP_PARAM_OFFLOAD=True \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [setup, e2e_ppo_trainer_fsdp-qwen2_5vl-3b, e2e_ppo_trainer_fsdp_sglang]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_megatron_vllm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_vllm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_ppo_trainer_megatron_vllm.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  # deepseek-ai/deepseek-coder-1.3b-instruct: dense, tie_word_embeddings=False\n  e2e_ppo_trainer_megatron-deepseek:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps --force-reinstall .\n          pip3 install git+https://github.com/ISEEKYAN/mbridge.git@main --no-deps --no-build-isolation\n          pip3 install math-verify\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      # Full training save&load\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use mbridge e2e to pre-load and save (Deepseek)\n        run: |\n          ray stop --force\n          ALL_OFFLOAD=True SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct COMMON_PP=4 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True USE_DIST_CKPT=False \\\n          bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use mbridge e2e to pre-load and save (Deepseek)\n        run: |\n          ray stop --force\n          RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 SAVE_FREQ=1 COMMON_PP=4 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True USE_DIST_CKPT=False \\\n          bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      # LoRA training save&load\n      - name: clean up and install Megatron-Bridge\n        run: |\n          rm -rf checkpoints\n          pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@83a7c11 --no-deps --no-build-isolation\n          pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@5455f0a --no-deps --no-build-isolation\n          pip3 install \"nvidia-modelopt[torch]>=0.37.0\" transformers==4.57.1\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use Megatron-Bridge LoRA e2e to pre-load and save (Deepseek)\n        run: |\n          ray stop --force\n          ALL_OFFLOAD=True SAVE_FREQ=1 MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct COMMON_PP=4 LORA_RANK=8 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False USE_DIST_CKPT=False \\\n          bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron, use Megatron-Bridge LoRA e2e to pre-load and save (Deepseek)\n        run: |\n          ray stop --force\n          RESUME_MODE=auto MODEL_ID=deepseek-ai/deepseek-coder-1.3b-instruct TOTAL_TRAIN_STEPS=2 SAVE_FREQ=1 COMMON_PP=4 LORA_RANK=8 COMMON_VPP=null COMMON_CP=1 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False USE_DIST_CKPT=False \\\n          bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  # Qwen3-0.6B: dense, tie_word_embeddings=True\n  e2e_ppo_trainer_megatron-qwen3:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install math-verify\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron (Qwen3) testing learning rate scheduler\n        run: |\n          ray stop --force\n          ALL_OFFLOAD=True VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 LR_WARMUP_STEPS=1 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with FP8 rollout\n        run: |\n          ray stop --force\n          export VLLM_USE_V1=1\n          ROLLOUT_QUANTIZATION=fp8 TOTAL_TRAIN_STEPS=2 MODEL_ID=Qwen/Qwen3-0.6B bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [setup, e2e_ppo_trainer_megatron-deepseek, e2e_ppo_trainer_megatron-qwen3]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_vllm_2\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_ppo_trainer_megatron_vllm_2.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  e2e_ppo_trainer_megatron-moe-expert-parallel:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps --force-reinstall .\n          pip3 install git+https://github.com/NVIDIA-NeMo/Megatron-Bridge.git@83a7c11 --no-deps --no-build-isolation\n          pip3 install git+https://github.com/NVIDIA/Megatron-LM.git@5455f0a --no-deps --no-build-isolation\n          pip3 install \"nvidia-modelopt[torch]>=0.37.0\" transformers==4.57.1\n      - name: Prepare GSM8K dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron-Bridge (Qwen3-30B-A3B-Instruct-2507)\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \\\n          PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \\\n          MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \\\n          MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \\\n          COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=8 \\\n          USE_DIST_CKPT=True ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: Running GSM8K E2E training tests with 3D parallelism with FP8 rollout on 8 L20 GPUs with Megatron-Bridge (Qwen3-30B-A3B-Instruct-2507)\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \\\n          PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \\\n          MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 \\\n          MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \\\n          COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=4 COMMON_ETP=1 INFER_TP=2 \\\n          USE_DIST_CKPT=True ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 ROLLOUT_QUANTIZATION=fp8 bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n      - name: Running GSM8K E2E training tests with 3D parallelism on 8 L20 GPUs with Megatron-Bridge LoRA (Qwen3-30B-A3B-Instruct-2507)\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_DUMMY_MODEL=True DUMMY_MODEL_CONFIG_PATH=tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json \\\n          PPO_MAX_TOKEN_LEN=1024 FWD_MAX_TOKEN_LEN=1024 \\\n          MAX_PROMPT_LENGTH=512 MAX_RESPONSE_LENGTH=512 LORA_RANK=8 CRITIC_LORA_RANK=8 \\\n          MODEL_ID=Qwen/Qwen3-30B-A3B-Instruct-2507 USE_MBRIDGE=True VANILLA_MBRIDGE=False VALUE_VANILLA_MBRIDGE=False \\\n          COMMON_PP=2 COMMON_VPP=null COMMON_CP=1 COMMON_TP=4 COMMON_EP=2 COMMON_ETP=1 INFER_TP=8 \\\n          USE_DIST_CKPT=False LORA_MERGE=True ALL_OFFLOAD=True SKIP_SAVE_HF_MODEL=1 bash tests/special_e2e/run_ppo_trainer_megatron.sh\n      - name: clean up\n        run: |\n          rm -rf checkpoints\n\n  e2e_ppo_trainer_fsdp_vllm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare GSM8K dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      # Function RM\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP_SIZE=8)\n        run: |\n          ray stop --force\n          VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-fsdp-size8\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm after resuming\n        run: |\n          ray stop --force\n          RESUME_MODE=auto VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-fsdp-size8\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test merging FSDP checkpoints (Qwen Actor)\n        run: |\n          exp_name=\"qwen2.5-0.5b-function-reward-minimal-fsdp-size8\"\n          python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (DDP_SIZE=2, FSDP_SIZE=4)\n        run: |\n          ray stop --force\n          VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True FSDP_SIZE=4 USE_KL=True VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test merging DDP+FSDP checkpoints (Qwen Actor)\n        run: |\n          exp_name=\"qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4\"\n          python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP2)\n        run: |\n          ray stop --force\n          VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-fsdp2-size8\" STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test merging FSDP2 checkpoints (Qwen Actor)\n        run: |\n          exp_name=\"qwen2.5-0.5b-function-reward-minimal-fsdp2-size8\"\n          python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E without rmpad using function rm\n        run: |\n          ray stop --force\n          RM_PAD=False bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (GRPO)\n        run: |\n          ray stop --force\n          CUSTOM_REWARD_FN=True ADV_ESTIMATOR=grpo USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      # - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (ReMax)\n      #   run: |\n      #     ray stop --force\n      #     ADV_ESTIMATOR=remax USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      # LoRA tests\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True TOTAL_TRAIN_STEPS=1 SAVE_FREQ=1 FSDP_SIZE=4 VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test GRPO LoRA checkpoints merging function\n        run: |\n          export EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal\"\n          ls checkpoints/verl-test/${EXP_NAME}/global_step_1/actor\n          cat checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface/config.json\n          python3 -m verl.model_merger merge --backend fsdp --local_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/ --target_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon with fsdp2\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n  e2e_ppo_trainer_fsdp-qwen2_5vl-3b:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 40 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      # Geo3k\n      - name: Prepare GEO3K dataset\n        run: |\n          python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/\n      - name: Running GEO3K VLM GRPO E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n      - name: Running GEO3K VLM PPO E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=gae RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GEO3K VLM GRPO E2E lora training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            LORA_RANK=32 LORA_EXCLUDE=\".*visual.*\" \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [\n        setup,\n        e2e_ppo_trainer_megatron-moe-expert-parallel,\n        e2e_ppo_trainer_fsdp-qwen2_5vl-3b,\n        e2e_ppo_trainer_fsdp_vllm,\n      ]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_megatron_vllm_2_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      - \"!verl/utils/fsdp_utils.py\"\n      - \"!verl/utils/checkpoint/fsdp_checkpoint_manager.py\"\n      - \"!verl/model_merger/fsdp_model_merger.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_ppo_trainer_megatron_vllm_2_ascend.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_megatron.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_megatron_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  e2e_ppo_trainer_fsdp_vllm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 90 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      # Function RM\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (DDP_SIZE=2, FSDP_SIZE=4)\n        run: |\n          ray stop --force\n          VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True FSDP_SIZE=4 USE_KL=True VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test merging DDP+FSDP checkpoints (Qwen Actor)\n        run: |\n          exp_name=\"qwen2.5-0.5b-function-reward-minimal-ddp-size2-fsdp-size4\"\n          python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm with validation and saving (FSDP2)\n        run: |\n          ray stop --force\n          VAL_BEFORE_TRAIN=True TEST_FREQ=1 SAVE_FREQ=1 SAVE_HF_MODEL=True VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal-fsdp2-size8\" STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test merging FSDP2 checkpoints (Qwen Actor)\n        run: |\n          exp_name=\"qwen2.5-0.5b-function-reward-minimal-fsdp2-size8\"\n          python -m verl.model_merger test --backend fsdp --local_dir checkpoints/verl-test/${exp_name}/global_step_1/actor --test_hf_dir checkpoints/verl-test/${exp_name}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E without rmpad using function rm\n        run: |\n          ray stop --force\n          RM_PAD=False bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm (GRPO)\n        run: |\n          ray stop --force\n          CUSTOM_REWARD_FN=True ADV_ESTIMATOR=grpo USE_KL=True bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True TOTAL_TRAIN_STEPS=1 SAVE_FREQ=1 FSDP_SIZE=4 VERL_EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal\" bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Test GRPO LoRA checkpoints merging function\n        run: |\n          export EXP_NAME=\"qwen2.5-0.5b-function-reward-minimal\"\n          ls checkpoints/verl-test/${EXP_NAME}/global_step_1/actor\n          cat checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface/config.json\n          python3 -m verl.model_merger merge --backend fsdp --local_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/ --target_dir checkpoints/verl-test/${EXP_NAME}/global_step_1/actor/huggingface\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with grpo lora using function rm with use_shm and layered_summon with fsdp2\n        run: |\n          ray stop --force\n          ADV_ESTIMATOR=grpo USE_SHM=True LORA_RANK=32 LOAD_FORMAT=safetensors LAYERED_SUMMON=True STRATEGY=fsdp2 bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n\n  e2e_ppo_trainer_fsdp-qwen2_5vl-3b_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n          pip install trl==0.26.0\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      # Geo3k\n      - name: Prepare GEO3K dataset\n        run: |\n          python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k\n      - name: Running GEO3K VLM GRPO E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GEO3K VLM PPO E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=gae RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n      - name: Running GEO3K VLM GRPO E2E lora training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          TRAIN_FILES=$HOME/data/geo3k/train.parquet VAL_FILES=$HOME/data/geo3k/test.parquet \\\n            MAX_PROMPT_LEN=1536 MAX_RESPONSE_LEN=1536 \\\n            MODEL_ID=Qwen/Qwen2.5-VL-3B-Instruct \\\n            ADV_ESTIMATOR=grpo RM_PAD=True USE_KL=True ENABLE_CHUNKED_PREFILL=False \\\n            SP_SIZE=2 \\\n            LORA_RANK=32 LORA_EXCLUDE=\".*visual.*\" \\\n            bash tests/special_e2e/ppo_trainer/run_function_reward.sh\n"
  },
  {
    "path": ".github/workflows/e2e_ppo_trainer_veomni_vllm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_ppo_trainer_veomni_vllm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch.\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!docker/**\"\n      # Docs\n      - \"!**/*.md\"\n      - \"!docs/**\"\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_ppo_trainer_veomni_vllm.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"examples/data_preprocess/geo3k.py\"\n      - \"tests/special_e2e/run_ppo_trainer_veomni.sh\"\n      - \"verl/trainer/main_ppo.py\"\n      - \"verl/trainer/config/ppo_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  e2e_ppo_trainer_veomni_vllm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4\n      - name: Prepare GSM8K dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Prepare GEO3K dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/models/hf_data/hiyouga/geometry3k/\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with veomni engine (FSDP_SIZE=4, USP=2)\n        run: |\n          ray stop --force\n          FSDP_SIZE=4 SP_SIZE=2 bash tests/special_e2e/run_ppo_trainer_veomni.sh\n      - name: Running GEO3K E2E training tests on 8 L20 GPUs with veomni engine (FSDP_SIZE=8, USP=1)\n        run: |\n          ray stop --force\n          MODEL_ID=Qwen/Qwen3-VL-2B-Instruct TRAIN_FILES=${HOME}/data/geo3k/train.parquet VAL_FILES=${HOME}/data/gsm8k/test.parquet FSDP_SIZE=8 SP_SIZE=1 bash tests/special_e2e/run_ppo_trainer_veomni.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs:\n      [\n        setup,\n        e2e_ppo_trainer_veomni_vllm,\n      ]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_sft_llm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_sft_llm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_sft_llm.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/sft\"\n      - \"verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/trainer/config/sft_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n  e2e_sft_llm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install peft\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs w/o rmpad using function rm\n        run: |\n          ray stop --force\n          RM_PAD=False bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with sequence parallism\n        run: |\n          ray stop --force\n          SP_SIZE=2 bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests on 8 L20 GPUs with sequence parallism and liger\n        run: |\n          ray stop --force\n          SP_SIZE=2 LIGER=True bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests with LoRA\n        run: |\n          ray stop --force\n          LORA_RANK=32 bash tests/special_e2e/sft/run_sft.sh\n      - name: Run GSM8K E2E training and resume tests resuming from the checkpoint manager\n        run: |\n          ray stop --force\n          LORA_RANK=32 RESUME_MODE=auto TOTAL_TRAIN_STEP=2 bash tests/special_e2e/sft/run_sft.sh\n      # TODO: multiturn\n      - name: Running GSM8K E2E training tests with multiturn and various configs and compare results\n        run: |\n          bash tests/special_e2e/sft/test_sft_engine_all.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, e2e_sft_llm]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/e2e_sft_llm_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_sft_llm_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_sft_llm_ascend.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/sft\"\n      - \"verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/trainer/config/sft_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions: \n  contents: read\n\njobs:\n  e2e_sft_llm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 90 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install -e .\n          pip install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4\n          pip install pandas==2.3.3\n          pip uninstall -y mbridge\n          pip install git+https://github.com/ISEEKYAN/mbridge.git@89eb10\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare gsm8k dataset\n        run: |\n          python3 examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running GSM8K E2E training tests on 8 NPUs with rmpad using function rm\n        run: |\n          ray stop --force\n          bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests on 8 NPUs w/o rmpad using function rm\n        run: |\n          ray stop --force\n          RM_PAD=False bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests on 8 NPUs with sequence parallism\n        run: |\n          ray stop --force\n          SP_SIZE=2 bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests with LoRA\n        run: |\n          ray stop --force\n          LORA_RANK=32 bash tests/special_e2e/sft/run_sft.sh\n      - name: Run GSM8K E2E training and resume tests resuming from the checkpoint manager\n        run: |\n          ray stop --force\n          LORA_RANK=32 RESUME_MODE=auto TOTAL_TRAIN_STEP=2 bash tests/special_e2e/sft/run_sft.sh\n      - name: Running GSM8K E2E training tests with multiturn and various configs and compare results\n        run: |\n          ray stop --force\n          rm -rf ~/verl/test/log\n          mkdir -p ~/verl/test/log\n          export VERL_FILE_LOGGER_ROOT=~/verl/test/log\n          # test with single gpu as golden\n          echo \"run with single gpu as golden\"\n          BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp VERL_FILE_LOGGER_PATH=~/verl/test/log/golden.jsonl bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with fsdp 1\n          echo \"run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode no_padding\"\n          BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with fsdp 1 use_remove_padding and pad_mode no_padding\n          echo \"run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding use_remove_padding False\"\n          BACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with fsdp 2\n          echo \"run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp2\"\n          BACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with veomni\n          echo \"run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2\"\n          BACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with megatron\n          echo \"run with tp2 pp2 vpp2 cp2 num_gpus8\"\n          BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=NULL CP_SIZE=2 NUM_GPUS=8 bash tests/special_e2e/sft/run_sft_engine.sh\n          # test with cp in ray\n          echo \"run with tp2 pp2 vpp2 cp2 num_gpus8 mode=ray\"\n          BACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=NULL CP_SIZE=2 NUM_GPUS=8 mode=ray bash tests/special_e2e/sft/run_sft_engine.sh\n          rm -rf ~/verl/test/log\n"
  },
  {
    "path": ".github/workflows/e2e_sft_vlm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: e2e_sft_vlm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # Entrypoints\n      - \".github/workflows/e2e_sft_vlm.yml\"\n      - \"examples/data_preprocess/gsm8k.py\"\n      - \"tests/special_e2e/sft\"\n      - \"verl/trainer/fsdp_sft_trainer.py\"\n      - \"verl/trainer/config/sft_trainer.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n  e2e_sft_vlm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install peft\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install git+https://github.com/ByteDance-Seed/VeOmni.git@v0.1.4\n      - name: Prepare pokemon-gpt4o-captions dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/pokemon.py --local_dataset_path ${HOME}/models/hf_data/pokemon-gpt4o-captions\n      - name: Running Pokemon E2E training tests with multiturn and various configs and compare results\n        run: |\n          MODEL_ID=Qwen/Qwen3-VL-2B-Instruct DATASET_DIR=~/data/pokemon-gpt4o-captions VPP_SIZE=null bash tests/special_e2e/sft/test_sft_engine_all.sh\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, e2e_sft_vlm]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/gpu_unit_tests.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: GPU unit tests\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.4.x\n    paths:\n      - \"**/*.py\"\n      - .github/workflows/gpu_unit_tests.yml\n  pull_request:\n    branches:\n      - main\n      - v0.4.x\n    paths:\n      # The order that you define paths patterns matters:\n      # A matching negative pattern (prefixed with !) after a positive match will exclude the path.\n      # A matching positive pattern after a negative match will include the path again.\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # Entrypoints\n      - .github/workflows/gpu_unit_tests.yml\n      - \"tests/**test_*.py\"\n      # Ignore CPU tests\n      - \"!tests/*_on_cpu.py\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  gpu_unit_tests:\n    if: github.repository_owner == 'verl-project'\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 60 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1\"\n      HF_HUB_ENABLE_HF_TRANSFER: 1\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install hf_transfer\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install cupy-cuda12x==13.6.0 pytest-asyncio\n          pip3 install --ignore-installed blinker\n          pip3 install --ignore-installed mlflow \"numpy<2.0\"\n      - name: Run all GPU unit tests\n        run: |\n          pytest -s -x --ignore-glob=\"*on_npu.py\" --ignore-glob=\"*test_special_*.py\" --ignore-glob='*on_cpu.py' --ignore-glob=\"*test_vllm*\" --ignore-glob=\"*_sglang*\" --ignore-glob=\"*_hf_rollout*\" --ignore-glob=\"tests/models/\" --ignore-glob='tests/special*' --ignore-glob=\"tests/experimental\" --ignore-glob=\"tests/workers/reward_model\" --ignore-glob=\"*test_shared_memory*\" --ignore-glob=\"tests/workers/rollout/rollout_trtllm\" --ignore-glob=\"*test_bucketed_weight_transfer*\" tests/\n      - name: Testing LinearCrossEntropyTP Correctness, Computation Time and Memory Consumption\n        run: |\n          LOW_MEMORY=True torchrun --standalone --nnodes=1 --nproc-per-node=8 tests/utils/test_special_linear_cross_entropy_tp.py\n      - name: Testing FSDP2 actor functionality\n        run: |\n          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/actor/test_special_dp_actor.py\n      - name: Testing FSDP2 critic functionality\n        run: |\n          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/critic/test_special_dp_critic.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, gpu_unit_tests]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/model.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: model\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"verl/**/*.py\"\n      # Entrypoints\n      - \".github/workflows/model.yml\"\n      - \"tests/special_distributed/test_fsdp_ckpt.py\"\n      - \"tests/special_distributed/test_tensor_dict.py\"\n      - \"tests/models/**\"\n      - \"tests/special_distributed/run_all.sh\"\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  model_rmpad:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 20 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository and upgrade to latest transformers(4.54.0)/flash_attn, transformers 4.55.0 has strange behavior with model backward\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install --upgrade \"transformers<5.0.0\"\n      - name: Running rmpad model tests on 8 L20 GPUs + flash_attn 2.5.8\n        run: |\n          pytest -s tests/models/test_transformer.py\n      - name: Running rmpad model tests on 8 L20 GPUs + latest flash_attn\n        run: |\n          pytest -s tests/models/test_transformer.py\n      - name: Running FSDP rmpad model tests on 8 L20 GPUs + latest flash_attn\n        run: |\n          STRATEGY=fsdp torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py\n      - name: Running transformers ulysses tests on 8 L20 GPUs + latest transformers\n        run: |\n          torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py\n      - name: Running transformers ulysses tests on 8 L20 GPUs + transformers 4.54.1\n        run: |\n          pip3 install transformers==4.54.1\n          torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py\n      - name: Run distributed test\n        run: |\n          bash tests/special_distributed/run_all.sh\n\n  # TODO: Move this back to model_rmpad once FSDP2 is stable.\n  # NOTE: List as an independent job to make rerun easier.\n  model_rmpad_fsdp2_unstable:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 20 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository and upgrade to latest transformers/flash_attn\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Running FSDP2 rmpad model tests on 8 L20 GPUs + latest flash_attn\n        run: |\n          STRATEGY=fsdp2 torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py\n\n  model_engine:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 20 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Download model config files\n        run: |\n          hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir $HOME/models/Qwen/Qwen2.5-0.5B-Instruct\n\n      - name: Running mcore engine tests on 8 L20 GPUs\n        run: |\n          ray stop --force\n          pytest -s -x tests/models/test_engine.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, model_rmpad, model_rmpad_fsdp2_unstable, model_engine]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/model_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests \n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: model_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"verl/**/*.py\"\n      # Entrypoints\n      - \".github/workflows/model_ascend.yml\"\n      - \"tests/special_distributed/test_fsdp_ckpt.py\"\n      - \"tests/special_distributed/test_tensor_dict.py\"\n      - \"tests/models/**\"\n      - \"tests/special_distributed/run_all.sh\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\npermissions:\n  contents: read\n\njobs:\n  model_rmpad_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .[test]\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Running rmpad model tests on 8 NPUs\n        run: |\n          pytest -s tests/models/test_transformer.py\n      - name: Running FSDP rmpad model tests on 8 NPUs\n        run: |\n          STRATEGY=fsdp torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py\n      - name: Running transformers ulysses tests on 8 NPUs\n        run: |\n          torchrun --nproc_per_node=8 -m pytest tests/models/test_transformers_ulysses.py\n      - name: Run distributed test\n        run: |\n          bash tests/special_distributed/run_all.sh\n\n  # TODO: Move this back to model_rmpad once FSDP2 is stable.\n  # NOTE: List as an independent job to make rerun easier.\n  model_rmpad_fsdp2_unstable_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .[test]\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Running FSDP2 rmpad model tests on 8 NPUs\n        run: |\n          STRATEGY=fsdp2 torchrun --nproc_per_node=8 tests/special_distributed/test_fsdp_ckpt.py\n"
  },
  {
    "path": ".github/workflows/nightly_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: nightly_ci_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  # For push, for now only anti-patterns are specified so it is more conservative\n  # and achieves higher coverage.\n  schedule:\n    - cron: \"0 17 * * *\"\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  # Test ppo qwen3-8b fsdp+vllm\n  nightlyCI_ppo-qwen3-8b-fsdp-vllm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 180 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running nightlyCI_ppo-qwen3-8b-fsdp-vllm_ascend\n        run: |\n          ray stop --force\n          bash tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh\n\n  # Test grpo qwen25-7b-Instruct fsdp+vllm\n  nightlyCI_grpo-qwen25-7b-Instruct-fsdp-vllm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 180 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare GSM8K dataset\n        run: |\n          python examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k\n      - name: Running nightlyCI_grpo-qwen25-7b-Instruct-fsdp-vllm_ascend\n        run: |\n          ray stop --force\n          bash tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh\n\n  # Test grpo qwen25-vl-3b-Instruct fsdp+vllm\n  nightlyCI_grpo-qwen25-vl-3b-Instruct-fsdp-vllm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 180 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Preprocess geo3k dataset\n        run: |\n          python examples/data_preprocess/geo3k.py --local_dataset_path ${HOME}/.cache/datasets/hiyouga/geometry3k\n      - name: Running nightlyCI_grpo-qwen25-vl-3b-Instruct-fsdp-vllm_ascend\n        run: |\n          ray stop --force\n          bash tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh\n"
  },
  {
    "path": ".github/workflows/npu_unit_tests.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - `npu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix on ascend device.\n#   - Since cpu/gpu/npu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: NPU unit tests\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - .github/workflows/npu_unit_tests.yml\n  pull_request:\n    branches:\n      - main\n    paths:\n      # The order that you define paths patterns matters:\n      # A matching negative pattern (prefixed with !) after a positive match will exclude the path.\n      # A matching positive pattern after a negative match will include the path again.\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      - \"!recipe/**\"\n      # Entrypoints\n      - .github/workflows/npu_unit_tests.yml\n      - \"tests/**test_*.py\"\n      # Ignore CPU tests\n      - \"!tests/*_on_cpu.py\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  npu_unit_tests:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout volcengine/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .[test]\n          pip install mlflow pytest-asyncio\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Run all NPU unit tests\n        run: |\n          pytest -s -x --ignore-glob=\"*test_special_*.py\" --ignore-glob=\"*on_cpu.py\" --ignore-glob=\"*test_vllm*\" --ignore-glob=\"*_sglang*\" --ignore-glob=\"*_hf_rollout*\" --ignore-glob=\"tests/models/\" --ignore-glob=\"tests/special*\" --ignore-glob=\"tests/experimental\" --ignore-glob=\"tests/workers/reward_model\" --ignore-glob=\"*test_rvdz*\" --ignore-glob=\"*test_ray_collectives*\" --ignore-glob=\"*test_nvtx_profile*\" --ignore-glob=\"tests/checkpoint_engine\" --ignore-glob=\"*test_shared_memory*\" --ignore-glob=\"tests/workers/rollout/rollout_trtllm\" --ignore-glob=\"*test_fsdp_lora_merge*\" --ignore-glob=\"*test_activation_offload*\" --ignore-glob=\"*test_normalize_peft_param_name.py*\" tests/\n      - name: Testing activation offload\n        run: |\n          pytest -s -x tests/utils/test_activation_offload.py\n      - name: Testing normalize peft param name\n        run: |\n          pytest -s -x tests/utils/test_normalize_peft_param_name.py\n      - name: Testing FSDP2 actor functionality\n        run: |\n          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/actor/test_special_dp_actor.py\n      - name: Testing FSDP2 critic functionality\n        run: |\n          torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/workers/critic/test_special_dp_critic.py\n      - name: Running NPU profiling unit tests\n        run: |\n          pytest -s -x tests/utils/test_special_mstx_profile.py\n"
  },
  {
    "path": ".github/workflows/pre-commit.yml",
    "content": "# c.f. https://github.com/pre-commit/action?tab=readme-ov-file#using-this-action\nname: pre-commit\n\n# No need to avoid / cancel lightweight pre-commit jobs\non:\n  schedule:\n    - cron: \"0 0 * * 0\"\n  pull_request:\n  push:\n    branches:\n      - main\n      - v0.*\n  # Allow manual triggering\n  workflow_dispatch:\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  pre-commit:\n    runs-on: ubuntu-latest\n    strategy:\n      matrix:\n        python-version: [\"3.12\"]\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n      - name: Set up Python ${{ matrix.python-version }}\n        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n        with:\n          python-version: ${{ matrix.python-version }}\n      - name: Install the current repository\n        run: |\n          pip install pre-commit hydra-core\n          pip install --no-deps -e .\n      - name: Set ruff --output-format=github\n        run: |\n          sed -i 's/--output-format=full/--output-format=github/' .pre-commit-config.yaml\n          git add .pre-commit-config.yaml\n      # Check \"--all-files\" by default\n      - uses: pre-commit/action@v3.0.1\n"
  },
  {
    "path": ".github/workflows/precommit-autofix.yml",
    "content": "name: scheduled pre-commit autofix\n\non:\n  schedule:\n    # Every hour\n    - cron: \"0 * * * *\"\n  workflow_dispatch:\n\npermissions:\n  contents: write\n  pull-requests: write\n\njobs:\n  precommit:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n\n    steps:\n      - name: Checkout repository\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n\n      - name: Set up Python\n        uses: actions/setup-python@v5\n        with:\n          python-version: \"3.10\"\n\n      - name: Install pre-commit\n        run: |\n          python -m pip install --upgrade pip\n          pip install pre-commit hydra-core\n\n      - name: Run pre-commit\n        run: |\n          pre-commit run --all-files || true\n\n      - name: Create or update PR\n        uses: peter-evans/create-pull-request@v6\n        with:\n          branch: bot/precommit-autofix\n          delete-branch: true\n          title: \"[ci] chore: scheduled pre-commit autofix\"\n          commit-message: \"chore: auto-fix pre-commit issues\"\n          body: |\n            This PR was created automatically by a scheduled GitHub Action.\n\n            - Runs `pre-commit run --all-files`\n            - Triggered hourly\n          labels: |\n            automated\n            pre-commit\n"
  },
  {
    "path": ".github/workflows/reward_model_sglang.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: reward_model_sglang\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"verl/**/*.py\"\n      # Entrypoints\n      - \".github/workflows/reward_model_sglang.yml\"\n      - \"tests/experimental/reward_loop/**\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  reward_model_sglang:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: \"True\"\n      NCCL_SHM_DISABLE: \"1\"\n      NCCL_P2P_DISABLE: \"1\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install sglang-router==0.2.2\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_dir ${HOME}/data/gsm8k\n      - name: Running sglang generative reward model tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py\n      - name: Running sglang discriminative reward model tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py\n      - name: Running sglang agent loop with reward manager tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py\n      - name: Running sglang agent loop with reward model colocate tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=sglang pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, reward_model_sglang]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/reward_model_vllm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: reward_model_vllm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"verl/**/*.py\"\n      # Entrypoints\n      - \".github/workflows/reward_model_vllm.yml\"\n      - \"tests/experimental/reward_loop/**\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  reward_model_vllm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 30 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n      SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: \"True\"\n      NCCL_SHM_DISABLE: \"1\"\n      NCCL_P2P_DISABLE: \"1\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k --local_dir ${HOME}/data/gsm8k\n      - name: Running vllm generative reward model tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py\n      - name: Running vllm discriminative reward model tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py\n\n      - name: Running vllm agent loop with reward manager tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py\n      - name: Running vllm agent loop with reward model colocate tests on 8 L20 GPUs\n        run: |\n          unset http_proxy https_proxy HTTP_PROXY HTTPS_PROXY\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, reward_model_vllm]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/reward_model_vllm_ascend.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests \n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: reward_model_vllm_ascend\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"verl/**/*.py\"\n      # Entrypoints\n      - \".github/workflows/reward_model_vllm_ascend.yml\"\n      - \"tests/experimental/reward_loop/**\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions: \n  contents: read\n\njobs:\n  reward_model_vllm_ascend:\n    if: github.repository_owner == 'verl-project'\n    runs-on: linux-aarch64-a2b3-8\n    timeout-minutes: 60 # Increase this timeout value as needed\n    container:\n      image: swr.cn-southwest-2.myhuaweicloud.com/modelfoundry/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n      options: >-\n        --shm-size 16g\n    env:\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - name: Check npu and CANN info\n        run: |\n          cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n          npu-smi info\n      - name: Check initial pip list from image\n        run: |\n          pip list\n      - name: Checkout verl-project/verl repo\n        uses: actions/checkout@v4\n        with:\n          fetch-depth: 0\n          clean: true\n      - name: Install the current repository\n        run: |\n          pip install -r requirements-npu.txt\n          pip install --no-deps -e .[test]\n      - name: Check final pip list\n        run: |\n          pip list\n      - name: Prepare weights\n        run: |\n          ln -s /root/.cache/models ~/models\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k --local_dir ${HOME}/data/gsm8k\n      - name: Running vllm generative reward model tests on 8 NPUs\n        run: |\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_genrm.py\n      - name: Running vllm discriminative reward model tests on 8 NPUs\n        run: |\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_reward_model_disrm.py\n      - name: Running vllm agent loop with reward manager tests on 8 NPUs\n        run: |\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_standalone.py\n      - name: Running vllm agent loop with reward model colocate tests on 8 NPUs\n        run: |\n          export HCCL_HOST_SOCKET_PORT_RANGE=auto\n          export HCCL_NPU_SOCKET_PORT_RANGE=auto\n          ROLLOUT_NAME=vllm pytest -s -x tests/experimental/reward_loop/test_agent_reward_loop_colocate.py"
  },
  {
    "path": ".github/workflows/sanity.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n# name: Check PR Title\n\nname: sanity\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - .github/workflows/sanity.yml\n      - \"tests/special_sanity/**\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\njobs:\n  sanity:\n    runs-on: ubuntu-latest\n    timeout-minutes: 5 # Increase this timeout value as needed\n    strategy:\n      matrix:\n        python-version: [\"3.10\"]\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n      - name: Set up Python ${{ matrix.python-version }}\n        uses: actions/setup-python@0b93645e9fea7318ecaed2b359559ac225c90a2b # v5.3.0\n        with:\n          python-version: ${{ matrix.python-version }}\n      - name: Install the current repository\n        run: |\n          pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu\n          pip3 install -r requirements.txt\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Run sanity test\n        run: |\n          pytest -s -x tests/special_sanity\n      - name: Run license test\n        run: |\n          python3 tests/special_sanity/check_license.py --directories .\n      - name: Assert naming convention\n        run: |\n          if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=__pycache__ 'veRL' .; then\n            echo \"Please use verl instead of veRL in the codebase\"\n            exit 1\n          fi\n      - name: Assert SGLang naming convention\n        run: |\n          if grep -rIn --exclude-dir=.git --exclude-dir=.github --exclude-dir=venv --exclude-dir=__pycache__ --exclude=ascend_sglang_best_practices.rst -E 'Sglang|sgLang|sglAng|sglaNg|sglanG' .; then\n            echo \"Please use SGLang or sglang as the formal name of SGLang rollout engine\"\n            exit 1\n          fi\n      - name: Validate test folder structure\n        run: python3 tests/special_sanity/validate_structure.py\n      - name: Assert documentation requirement for functions\n        run: python3 tests/special_sanity/validate_imported_docs.py\n      - name: Assert device api usage in verl/verl\n        run: python3 tests/special_sanity/check_device_api_usage.py --directory ./verl\n      - name: Assert documentation time info\n        run: python3 tests/special_sanity/check_docs_time_info.py\n      - name: Check docstrings for specified files\n        run: python3 tests/special_sanity/check_docstrings.py\n      - name: Check DataProto for specified folders\n        run: python3 tests/special_sanity/check_dataproto_usage.py -d ./verl/workers/engine\n"
  },
  {
    "path": ".github/workflows/scorecard.yml",
    "content": "# This workflow uses actions that are not certified by GitHub. They are provided\n# by a third-party and are governed by separate terms of service, privacy\n# policy, and support documentation.\n\nname: Scorecard supply-chain security\non:\n  # For Branch-Protection check. Only the default branch is supported. See\n  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#branch-protection\n  branch_protection_rule:\n  # To guarantee Maintained check is occasionally updated. See\n  # https://github.com/ossf/scorecard/blob/main/docs/checks.md#maintained\n  schedule:\n    - cron: \"27 7 * * 1\"\n  push:\n    branches:\n      - main\n      - v0.*\n\n# Declare default permissions as read only.\npermissions: read-all\n\njobs:\n  analysis:\n    name: Scorecard analysis\n    runs-on: ubuntu-latest\n    permissions:\n      # Needed to upload the results to code-scanning dashboard.\n      security-events: write\n      # Needed to publish results and get a badge (see publish_results below).\n      id-token: write\n      # Uncomment the permissions below if installing in a private repository.\n      # contents: read\n      # actions: read\n\n    steps:\n      - name: \"Checkout code\"\n        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1\n        with:\n          persist-credentials: false\n\n      - name: \"Run analysis\"\n        uses: ossf/scorecard-action@0864cf19026789058feabb7e87baa5f140aac736 # v2.3.1\n        with:\n          results_file: results.sarif\n          results_format: sarif\n          # (Optional) \"write\" PAT token. Uncomment the `repo_token` line below if:\n          # - you want to enable the Branch-Protection check on a *public* repository, or\n          # - you are installing Scorecard on a *private* repository\n          # To create the PAT, follow the steps in https://github.com/ossf/scorecard-action?tab=readme-ov-file#authentication-with-fine-grained-pat-optional.\n          # repo_token: ${{ secrets.SCORECARD_TOKEN }}\n\n          # Public repositories:\n          #   - Publish results to OpenSSF REST API for easy access by consumers\n          #   - Allows the repository to include the Scorecard badge.\n          #   - See https://github.com/ossf/scorecard-action#publishing-results.\n          # For private repositories:\n          #   - `publish_results` will always be set to `false`, regardless\n          #     of the value entered here.\n          publish_results: true\n\n      # Upload the results to GitHub's code scanning dashboard (optional).\n      # Commenting out will disable upload of results to your repo's Code Scanning dashboard\n      - name: \"Upload to code-scanning\"\n        uses: github/codeql-action/upload-sarif@9e8d0789d4a0fa9ceb6b1738f7e269594bdd67f0 #v3.28.9\n        with:\n          sarif_file: results.sarif\n"
  },
  {
    "path": ".github/workflows/secrets_scan.yml",
    "content": "on:\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n\npermissions:\n  contents: read\n\njobs:\n  test:\n    runs-on: ubuntu-latest\n    steps:\n      - name: Checkout code\n        uses: actions/checkout@b4ffde65f46336ab88eb53be808477a3936bae11 # v4.1.1\n        with:\n          fetch-depth: 0\n      - name: Secret Scanning\n        uses: trufflesecurity/trufflehog@7dc056a193116ba8d82154bf0549381c8fb8545c # v3.88.14\n        with:\n          extra_args: --results=verified,unknown\n"
  },
  {
    "path": ".github/workflows/sgl.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: sgl\n\non:\n  #  workflow_dispatch: # Manual\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      - .github/workflows/sgl.yml\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\" # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # vLLM\n      - \"!**/*vllm*\"\n\n      # Entrypoints\n      - \".github/workflows/sgl.yml\"\n      - \"tests/rollout/*sglang*\"\n      - \"tests/rollout/async_rollout_utils.py\"\n      - \"tests/workers/rollout/*interaction*\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:sgl059.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  sgl:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 35 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: 1\n      SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: \"True\"\n      NCCL_SHM_DISABLE: \"1\"\n      NCCL_P2P_DISABLE: \"1\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install cupy-cuda12x==13.6.0 pytest-asyncio\n          pip3 install hf_transfer fastmcp pytest-asyncio\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Test the latest SGLang Rollout async with agent loop\n        run: |\n          ROLLOUT_NAME=sglang pytest -svvv tests/experimental/agent_loop\n\n  sgl_checkpoint_engine:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 35 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: 1\n      SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK: \"True\"\n      NCCL_SHM_DISABLE: \"1\"\n      NCCL_P2P_DISABLE: \"1\"\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install cupy-cuda12x==13.6.0 pytest-asyncio\n          pip3 install hf_transfer fastmcp pytest-asyncio\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n      - name: Test SGLang ServerAdapter with Checkpoint Engine (NCCL)\n        run: |\n          ROLLOUT_NAME=sglang pytest -svvv tests/checkpoint_engine/test_special_server_adapter.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, sgl, sgl_checkpoint_engine]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".github/workflows/type-coverage-check.yml",
    "content": "name: Type Annotation and Docstring Coverage\n\non:\n  pull_request:\n    paths:\n      - '**/*.py'\n      - '.github/workflows/type-coverage-check.yml'\n\njobs:\n  type-coverage-check:\n    runs-on: ubuntu-latest\n    steps:\n      - uses: actions/checkout@v4\n        with:\n          fetch-depth: 0  # 🚨 Important: fetch full history so `origin/main` is available\n      - name: Set up Python\n        uses: actions/setup-python@v5\n        with:\n          python-version: '3.10'\n\n      - name: Install dependencies\n        run: |\n          pip3 install torch torchvision --index-url https://download.pytorch.org/whl/cpu\n          pip3 install -r requirements.txt\n          pip3 install --no-deps -e .\n      - name: Run type annotation coverage check\n        run: |\n          python3 tests/special_sanity/type_coverage_check.py\n      - name: Run docstring coverage check\n        run: |\n          python3 tests/special_sanity/check_api_docs.py verl\n"
  },
  {
    "path": ".github/workflows/vllm.yml",
    "content": "# # Tests layout\n\n# Each folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n# - `tests/trainer` for testing functionality related to `verl/trainer`\n# - `tests/models` for testing functionality related to `verl/models`\n# - ...\n\n# There are a few folders with `special_` prefix, created for special purposes:\n# - `special_distributed`: unit tests that must run with multiple GPUs\n# - `special_e2e`: end-to-end tests with training/generation scripts\n# - `special_npu`: tests for NPUs\n# - `special_sanity`: a suite of quick sanity tests\n# - `special_standalone`: a set of test that are designed to run in dedicated environments\n\n# Accelerators for tests\n# - By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n# - For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# # Workflow layout\n\n# All CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n# 1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n# 2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n# 3. End-to-end tests: `e2e_*.yml`\n# 4. Unit tests\n#   - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n#   - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n#   - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n#     - new workflow yaml is added to `.github/workflows`\n#     - new tests are added to workflow mentioned in 2.\n\nname: vllm\n\non:\n  # Trigger the workflow on push or pull request,\n  # but only for the main branch\n  push:\n    branches:\n      - main\n      - v0.*\n  pull_request:\n    branches:\n      - main\n      - v0.*\n    paths:\n      - \"**/*.py\"\n      # Other entrypoints\n      - \"!examples/**\"\n      - \"!tests/**\"\n      - \"!verl/trainer/main_*.py\"\n      - \"!verl/trainer/fsdp_sft_trainer.py\"\n      # FSDP\n      - \"!verl/workers/**/*dp_*.py\"\n      # Megatron\n      - \"!verl/workers/**/megatron_*.py\"\n      # SGLang\n      - \"!**/*sglang*\"\n      # Entrypoints\n      - \".github/workflows/vllm.yml\"\n      - \"tests/special_e2e/generation\"\n      - \"tests/workers/rollout\"\n      - \"verl/trainer/main_generation.py\"\n      - \"verl/trainer/config/generation.yaml\"\n\n# Cancel jobs on the same ref if a new one is triggered\nconcurrency:\n  group: ${{ github.workflow }}-${{ github.ref }}\n  cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}\n\n# Declare permissions just read content.\npermissions:\n  contents: read\n\nenv:\n  IMAGE: \"verl-ci-cn-beijing.cr.volces.com/verlai/verl:vllm017.dev2\"\n  DYNAMIC_RUNNER_ENDPOINT: \"https://sd10g3clalm04ug7alq90.apigateway-cn-beijing.volceapi.com/runner\"\n\njobs:\n  setup:\n    if: github.repository_owner == 'verl-project'\n    runs-on: ubuntu-latest\n    outputs:\n      runner-label: ${{ steps.create-runner.outputs.runner-label }}\n      mlp-task-id: ${{ steps.create-runner.outputs.mlp-task-id }}\n    steps:\n      - uses: actions/checkout@v4\n      - id: create-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"create\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-image: \"${{ env.IMAGE }}\"\n\n  vllm:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 35 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install --upgrade \"transformers<5.0\"\n      #      - name: Download Model to Use\n      #        run: |\n      #          hf download Qwen/Qwen2.5-0.5B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-0.5B-Instruct\n      #          hf download Qwen/Qwen2.5-1.5B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-1.5B-Instruct\n      #          hf download Qwen/Qwen2.5-VL-3B-Instruct --local-dir ${HOME}/models/Qwen/Qwen2.5-VL-3B-Instruct\n      #          hf download OldKingMeister/Qwen2.5-1.5B-Instruct-YaRN --local-dir ${HOME}/models/OldKingMeister/Qwen2.5-1.5B-Instruct-YaRN\n      #          export HF_HUB_OFFLINE=1\n      - name: Prepare gsm8k dataset\n        run: |\n          ray stop --force\n          python3 examples/data_preprocess/gsm8k.py --local_dataset_path ${HOME}/models/hf_data/gsm8k\n      - name: Test the latest vLLM Rollout async with agent loop\n        run: |\n          ROLLOUT_NAME=vllm pytest -svvv tests/experimental/agent_loop\n      - name: Test vllm server abort functionality\n        run: |\n          pytest tests/workers/rollout/rollout_vllm/test_vllm_abort.py -v -s\n\n  vllm_checkpoint_engine:\n    needs: setup\n    runs-on: [\"${{ needs.setup.outputs.runner-label || 'L20x8' }}\"]\n    timeout-minutes: 35 # Increase this timeout value as needed\n    env:\n      HTTP_PROXY: ${{ secrets.PROXY_HTTP }}\n      HTTPS_PROXY: ${{ secrets.PROXY_HTTPS }}\n      NO_PROXY: \"localhost,127.0.0.1,hf-mirror.com\"\n      HF_ENDPOINT: \"https://hf-mirror.com\"\n      HF_HUB_ENABLE_HF_TRANSFER: \"0\" # This is more stable\n    steps:\n      - uses: actions/checkout@11bd71901bbe5b1630ceea73d27597364c9af683 # v4.2.2\n        with:\n          fetch-depth: 0\n      - name: Install the current repository\n        run: |\n          pip3 install pytest-asyncio\n          pip3 install -r requirements-test.txt\n          pip3 install --no-deps -e .\n          pip3 install --upgrade \"transformers<5.0\"\n          pip3 install cupy-cuda12x==13.6.0\n      - name: Test vLLM ServerAdapter with Checkpoint Engine (NCCL)\n        run: |\n          ROLLOUT_NAME=vllm pytest -svvv tests/checkpoint_engine/test_special_server_adapter.py\n      - name: Test bucketed weight transfer\n        run: |\n          pytest -svvv tests/utils/test_bucketed_weight_transfer.py\n\n  cleanup:\n    runs-on: ubuntu-latest\n    needs: [setup, vllm, vllm_checkpoint_engine]\n    if: always()\n    steps:\n      - id: destroy-runner\n        uses: volcengine/vemlp-github-runner@v1\n        with:\n          mode: \"destroy\"\n          faas-url: \"${{ env.DYNAMIC_RUNNER_ENDPOINT }}\"\n          mlp-task-id: \"${{ needs.setup.outputs.mlp-task-id }}\"\n"
  },
  {
    "path": ".gitignore",
    "content": "**/*.pt\n**/checkpoints\n**/wget-log\n**/_build/\n**/*.ckpt\n**/outputs\n**/*.tar.gz\n**/playground\n**/wandb\n\n/pyrightconfig.json\n\n# Byte-compiled / optimized / DLL files\n__pycache__/\n*.py[cod]\n*$py.class\ndataset/*\ntensorflow/my_graph/*\n.idea/\n# C extensions\n*.so\n\n# Distribution / packaging\n.Python\n# env/\nbuild/\ndevelop-eggs/\ndist/\ndownloads/\neggs/\n.eggs/\nlib/\nlib64/\nparts/\nsdist/\nvar/\ntmp/\n*.egg-info/\n.installed.cfg\n*.egg\n\n# PyInstaller\n#  Usually these files are written by a python script from a template\n#  before PyInstaller builds the exe, so as to inject date/other infos into it.\n*.manifest\n*.spec\n\n# Installer logs\npip-log.txt\npip-delete-this-directory.txt\n\n# Unit test / coverage reports\nhtmlcov/\n.tox/\n.coverage\n.coverage.*\n.cache\nnosetests.xml\ncoverage.xml\n*,cover\n.hypothesis/\npytest.ini\noutput.txt\n\n# Translations\n*.mo\n*.pot\n\n# Django stuff:\n*.log\nlocal_settings.py\n\n# Flask stuff:\ninstance/\n.webassets-cache\n\n# Scrapy stuff:\n.scrapy\n\n# Sphinx documentation\ndocs/_build/\n\n# PyBuilder\ntarget/\n\n# IPython Notebook\n.ipynb_checkpoints\n\n# pyenv\n.python-version\n\n# celery beat schedule file\ncelerybeat-schedule\n\n# dotenv\n.env\n\n# virtualenv\nvenv/\n.venv/\nENV/\n\n# Spyder project settings\n.spyderproject\n\n# Rope project settings\n.ropeproject\n\n# vscode\n.vscode\n\n# Mac\n.DS_Store\n\n# vim\n*.swp\n\n# emacs\n*~\n\n# ckpt\n*.lock\n\n# data\n*.parquet\n\n\n# local logs\nlogs\nlog\noutputs\n.history"
  },
  {
    "path": ".gitmodules",
    "content": "[submodule \"recipe\"]\n\tpath = recipe\n\turl = https://github.com/verl-project/verl-recipe.git\n"
  },
  {
    "path": ".pre-commit-config.yaml",
    "content": "repos:\n  - repo: https://github.com/astral-sh/ruff-pre-commit\n    rev: \"v0.12.2\"\n    hooks:\n      - id: ruff\n        args: [\"--fix\", \"--show-fixes\", \"--output-format=full\"]\n        exclude: ^.*\\.(ipynb)$\n      - id: ruff-format\n\n  - repo: https://github.com/pre-commit/mirrors-mypy\n    rev: \"v1.17.0\"\n    hooks:\n      - id: mypy\n\n  - repo: local\n    hooks:\n      - id: autogen-trainer-cfg\n        name: Generate and verify verl/trainer/config/_generated_*.yaml\n        entry: scripts/generate_trainer_config.sh\n        language: script\n        pass_filenames: false\n\n  - repo: local\n    hooks:\n      - id: check-docstrings\n        name: Check doc string coverage\n        entry: python3 tests/special_sanity/check_docstrings.py\n        language: python\n        pass_filenames: false\n\n  - repo: local\n    hooks:\n      - id: check-license\n        name: Check license\n        entry: python3 tests/special_sanity/check_license.py --directories examples scripts tests verl setup.py\n        language: python\n        pass_filenames: false\n\n  - repo: local\n    hooks:\n      - id: compileall\n        name: Compile all python files\n        entry: sh -c 'PYTHONWARNINGS=error python3 -m compileall -q . -x \"(^|[\\\\/])(\\.venv|venv|\\.git)([\\\\/]|$)\"'\n        language: python\n        pass_filenames: false\n"
  },
  {
    "path": ".readthedocs.yaml",
    "content": "# Read the Docs configuration file\n# See https://docs.readthedocs.io/en/stable/config-file/v2.html for details\n\nversion: 2\n\nbuild:\n  os: ubuntu-22.04\n  tools:\n    python: \"3.11\"\n    rust: \"1.70\"\n\nsphinx:\n  configuration: docs/conf.py\n\npython:\n  install:\n    - requirements: docs/requirements-docs.txt\n    - method: pip\n      path: .\n"
  },
  {
    "path": "CONTRIBUTING.md",
    "content": "# Contributing to verl\n\nThank you for considering a contribution to verl! We welcome contributions of any kind - bug fixes, enhancements, documentation improvements, or even just feedback. Whether you're an experienced developer or this is your first open-source project, your help is invaluable.\n\nYour support can take many forms:\n- Report issues or unexpected behaviors.\n- Suggest or implement new features.\n- Improve or expand documentation.\n- Review pull requests and assist other contributors.\n- Spread the word: share verl in blog posts, social media, or give the repo a ⭐.\n\n## Finding Issues to Contribute\n\nLooking for ways to dive in? Check out these issues:\n- [Good first issues](https://github.com/volcengine/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22good%20first%20issue%22)\n- [Call for contribution](https://github.com/volcengine/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22call%20for%20contribution%22)\nFurthermore, you can learn the development plan and roadmap via [RFC](https://github.com/volcengine/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3ARFC) and [Roadmap](https://github.com/volcengine/verl/issues?q=state%3Aopen%20label%3A%22roadmap%22).\n\n\n## Developing\n\n- **Python-only**: install verl via `pip install -e .[test,vllm]` or `pip install -e .[test,sglang]` and iterate quickly. For full dependency setup, check out the verl [installation doc](https://verl.readthedocs.io/en/latest/start/install.html).\n\n## Code Linting and Formatting\n\nWe rely on pre-commit to keep our code consistent. To set it up:\n\n```bash\npip install pre-commit\npre-commit install\n# for staged changes\npre-commit run\n# for all files in the repo\npre-commit run --all-files\n# run a specific hook with pre-commit\n# pre-commit run --all-files --show-diff-on-failure --color=always <hood-id>\npre-commit run --all-files --show-diff-on-failure --color=always ruff\npre-commit run --all-files --show-diff-on-failure --color=always autogen-trainer-cfg\n```\n\n## Testing\n\nOur test suites run on GitHub Actions. Check these workflows for details:\n- [GPU unit tests](https://github.com/volcengine/verl/blob/main/.github/workflows/gpu_unit_tests.yml)\n- [CPU unit tests](https://github.com/volcengine/verl/blob/main/.github/workflows/cpu_unit_tests.yml)\n- [vLLM tests](https://github.com/volcengine/verl/blob/main/.github/workflows/vllm.yml)\n- [SGLang tests](https://github.com/volcengine/verl/blob/main/.github/workflows/sgl.yml)\n\n### Adding CI tests\n\nIf possible, please add CI test(s) for your new feature:\n\n1. Find the most relevant workflow yml file, which usually corresponds to a `hydra` default config (e.g. `ppo_trainer`, `ppo_megatron_trainer`, `sft_trainer`, etc).\n2. Add related path patterns to the `paths` section if not already included.\n3. Minimize the workload of the test script(s) (see existing scripts for examples).\n\n## Building the Docs\n```\n# Ensure verl is on your PYTHONPATH, e.g.:\npip install -e .[test]\n\n# Install documentation dependencies\ncd docs\npip install -r requirements-docs.txt\n\n# Generate HTML docs\nmake clean\nmake html\n\n# Preview locally\npython -m http.server -d _build/html/\n```\nOpen your browser at http://localhost:8000 to explore the docs.\n\n## Pull Requests & Code Reviews\n\nThanks for submitting a PR! To streamline reviews:\n- Follow our Pull Request Template for title format and checklist.\n- Adhere to our pre-commit lint rules and ensure all checks pass.\n- Update docs for any user-facing changes.\n- Add or update tests in the CI workflows, or explain why tests aren't applicable.\n\n## License\n\nSee the [LICENSE](https://github.com/volcengine/verl/blob/main/LICENSE) file for full details.\n\n## Thank You\n\nWe appreciate your contributions to verl. Your efforts help make the project stronger and more user-friendly. Happy coding!\n\n"
  },
  {
    "path": "LICENSE",
    "content": "\n                                 Apache License\n                           Version 2.0, January 2004\n                        http://www.apache.org/licenses/\n\n   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION\n\n   1. Definitions.\n\n      \"License\" shall mean the terms and conditions for use, reproduction,\n      and distribution as defined by Sections 1 through 9 of this document.\n\n      \"Licensor\" shall mean the copyright owner or entity authorized by\n      the copyright owner that is granting the License.\n\n      \"Legal Entity\" shall mean the union of the acting entity and all\n      other entities that control, are controlled by, or are under common\n      control with that entity. For the purposes of this definition,\n      \"control\" means (i) the power, direct or indirect, to cause the\n      direction or management of such entity, whether by contract or\n      otherwise, or (ii) ownership of fifty percent (50%) or more of the\n      outstanding shares, or (iii) beneficial ownership of such entity.\n\n      \"You\" (or \"Your\") shall mean an individual or Legal Entity\n      exercising permissions granted by this License.\n\n      \"Source\" form shall mean the preferred form for making modifications,\n      including but not limited to software source code, documentation\n      source, and configuration files.\n\n      \"Object\" form shall mean any form resulting from mechanical\n      transformation or translation of a Source form, including but\n      not limited to compiled object code, generated documentation,\n      and conversions to other media types.\n\n      \"Work\" shall mean the work of authorship, whether in Source or\n      Object form, made available under the License, as indicated by a\n      copyright notice that is included in or attached to the work\n      (an example is provided in the Appendix below).\n\n      \"Derivative Works\" shall mean any work, whether in Source or Object\n      form, that is based on (or derived from) the Work and for which the\n      editorial revisions, annotations, elaborations, or other modifications\n      represent, as a whole, an original work of authorship. For the purposes\n      of this License, Derivative Works shall not include works that remain\n      separable from, or merely link (or bind by name) to the interfaces of,\n      the Work and Derivative Works thereof.\n\n      \"Contribution\" shall mean any work of authorship, including\n      the original version of the Work and any modifications or additions\n      to that Work or Derivative Works thereof, that is intentionally\n      submitted to Licensor for inclusion in the Work by the copyright owner\n      or by an individual or Legal Entity authorized to submit on behalf of\n      the copyright owner. For the purposes of this definition, \"submitted\"\n      means any form of electronic, verbal, or written communication sent\n      to the Licensor or its representatives, including but not limited to\n      communication on electronic mailing lists, source code control systems,\n      and issue tracking systems that are managed by, or on behalf of, the\n      Licensor for the purpose of discussing and improving the Work, but\n      excluding communication that is conspicuously marked or otherwise\n      designated in writing by the copyright owner as \"Not a Contribution.\"\n\n      \"Contributor\" shall mean Licensor and any individual or Legal Entity\n      on behalf of whom a Contribution has been received by Licensor and\n      subsequently incorporated within the Work.\n\n   2. Grant of Copyright License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      copyright license to reproduce, prepare Derivative Works of,\n      publicly display, publicly perform, sublicense, and distribute the\n      Work and such Derivative Works in Source or Object form.\n\n   3. Grant of Patent License. Subject to the terms and conditions of\n      this License, each Contributor hereby grants to You a perpetual,\n      worldwide, non-exclusive, no-charge, royalty-free, irrevocable\n      (except as stated in this section) patent license to make, have made,\n      use, offer to sell, sell, import, and otherwise transfer the Work,\n      where such license applies only to those patent claims licensable\n      by such Contributor that are necessarily infringed by their\n      Contribution(s) alone or by combination of their Contribution(s)\n      with the Work to which such Contribution(s) was submitted. If You\n      institute patent litigation against any entity (including a\n      cross-claim or counterclaim in a lawsuit) alleging that the Work\n      or a Contribution incorporated within the Work constitutes direct\n      or contributory patent infringement, then any patent licenses\n      granted to You under this License for that Work shall terminate\n      as of the date such litigation is filed.\n\n   4. Redistribution. You may reproduce and distribute copies of the\n      Work or Derivative Works thereof in any medium, with or without\n      modifications, and in Source or Object form, provided that You\n      meet the following conditions:\n\n      (a) You must give any other recipients of the Work or\n          Derivative Works a copy of this License; and\n\n      (b) You must cause any modified files to carry prominent notices\n          stating that You changed the files; and\n\n      (c) You must retain, in the Source form of any Derivative Works\n          that You distribute, all copyright, patent, trademark, and\n          attribution notices from the Source form of the Work,\n          excluding those notices that do not pertain to any part of\n          the Derivative Works; and\n\n      (d) If the Work includes a \"NOTICE\" text file as part of its\n          distribution, then any Derivative Works that You distribute must\n          include a readable copy of the attribution notices contained\n          within such NOTICE file, excluding those notices that do not\n          pertain to any part of the Derivative Works, in at least one\n          of the following places: within a NOTICE text file distributed\n          as part of the Derivative Works; within the Source form or\n          documentation, if provided along with the Derivative Works; or,\n          within a display generated by the Derivative Works, if and\n          wherever such third-party notices normally appear. The contents\n          of the NOTICE file are for informational purposes only and\n          do not modify the License. You may add Your own attribution\n          notices within Derivative Works that You distribute, alongside\n          or as an addendum to the NOTICE text from the Work, provided\n          that such additional attribution notices cannot be construed\n          as modifying the License.\n\n      You may add Your own copyright statement to Your modifications and\n      may provide additional or different license terms and conditions\n      for use, reproduction, or distribution of Your modifications, or\n      for any such Derivative Works as a whole, provided Your use,\n      reproduction, and distribution of the Work otherwise complies with\n      the conditions stated in this License.\n\n   5. Submission of Contributions. Unless You explicitly state otherwise,\n      any Contribution intentionally submitted for inclusion in the Work\n      by You to the Licensor shall be under the terms and conditions of\n      this License, without any additional terms or conditions.\n      Notwithstanding the above, nothing herein shall supersede or modify\n      the terms of any separate license agreement you may have executed\n      with Licensor regarding such Contributions.\n\n   6. Trademarks. This License does not grant permission to use the trade\n      names, trademarks, service marks, or product names of the Licensor,\n      except as required for reasonable and customary use in describing the\n      origin of the Work and reproducing the content of the NOTICE file.\n\n   7. Disclaimer of Warranty. Unless required by applicable law or\n      agreed to in writing, Licensor provides the Work (and each\n      Contributor provides its Contributions) on an \"AS IS\" BASIS,\n      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or\n      implied, including, without limitation, any warranties or conditions\n      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A\n      PARTICULAR PURPOSE. You are solely responsible for determining the\n      appropriateness of using or redistributing the Work and assume any\n      risks associated with Your exercise of permissions under this License.\n\n   8. Limitation of Liability. In no event and under no legal theory,\n      whether in tort (including negligence), contract, or otherwise,\n      unless required by applicable law (such as deliberate and grossly\n      negligent acts) or agreed to in writing, shall any Contributor be\n      liable to You for damages, including any direct, indirect, special,\n      incidental, or consequential damages of any character arising as a\n      result of this License or out of the use or inability to use the\n      Work (including but not limited to damages for loss of goodwill,\n      work stoppage, computer failure or malfunction, or any and all\n      other commercial damages or losses), even if such Contributor\n      has been advised of the possibility of such damages.\n\n   9. Accepting Warranty or Additional Liability. While redistributing\n      the Work or Derivative Works thereof, You may choose to offer,\n      and charge a fee for, acceptance of support, warranty, indemnity,\n      or other liability obligations and/or rights consistent with this\n      License. However, in accepting such obligations, You may act only\n      on Your own behalf and on Your sole responsibility, not on behalf\n      of any other Contributor, and only if You agree to indemnify,\n      defend, and hold each Contributor harmless for any liability\n      incurred by, or claims asserted against, such Contributor by reason\n      of your accepting any such warranty or additional liability.\n\n   END OF TERMS AND CONDITIONS\n\n   APPENDIX: How to apply the Apache License to your work.\n\n      To apply the Apache License to your work, attach the following\n      boilerplate notice, with the fields enclosed by brackets \"[]\"\n      replaced with your own identifying information. (Don't include\n      the brackets!)  The text should be enclosed in the appropriate\n      comment syntax for the file format. We also recommend that a\n      file or class name and description of purpose be included on the\n      same \"printed page\" as the copyright notice for easier\n      identification within third-party archives.\n\n   Copyright [yyyy] [name of copyright owner]\n\n   Licensed under the Apache License, Version 2.0 (the \"License\");\n   you may not use this file except in compliance with the License.\n   You may obtain a copy of the License at\n\n       http://www.apache.org/licenses/LICENSE-2.0\n\n   Unless required by applicable law or agreed to in writing, software\n   distributed under the License is distributed on an \"AS IS\" BASIS,\n   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n   See the License for the specific language governing permissions and\n   limitations under the License.\n"
  },
  {
    "path": "Notice.txt",
    "content": "Copyright 2023-2024 Bytedance Ltd. and/or its affiliates "
  },
  {
    "path": "README.md",
    "content": "<div align=\"center\">\n 👋 Hi, everyone!\n    verl is a RL training library initiated by <b>ByteDance Seed team</b> and maintained by the verl community.\n    <br>\n    <br>\n</div>\n\n<div align=\"center\">\n\n<a href=\"https://deepwiki.com/volcengine/verl\"><img src=\"https://devin.ai/assets/deepwiki-badge.png\" alt=\"Ask DeepWiki.com\" style=\"height:20px;\"></a>\n[![GitHub Repo stars](https://img.shields.io/github/stars/volcengine/verl)](https://github.com/volcengine/verl/stargazers)\n[![Twitter](https://img.shields.io/twitter/follow/verl_project)](https://twitter.com/verl_project)\n<a href=\"https://join.slack.com/t/verl-project/shared_invite/zt-3c6mc2khw-v0lo6NfDPuFP6OnkrZwfqw\"><img src=\"https://img.shields.io/badge/Slack-verl-blueviolet?logo=slack&amp\"></a>\n<a href=\"https://arxiv.org/pdf/2409.19256\"><img src=\"https://img.shields.io/static/v1?label=EuroSys&message=Paper&color=red\"></a>\n[![Documentation](https://img.shields.io/badge/documentation-blue)](https://verl.readthedocs.io/en/latest/)\n<a href=\"https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/WeChat.JPG\"><img src=\"https://img.shields.io/badge/微信-green?logo=wechat&amp\"></a>\n\n</div>\n\n![seed logo](https://github.com/user-attachments/assets/c42e675e-497c-4508-8bb9-093ad4d1f216)\n\n<h1 style=\"text-align: center;\">verl: Volcano Engine Reinforcement Learning for LLMs</h1>\n\nverl is a flexible, efficient and production-ready RL training library for large language models (LLMs).\n\nverl is the open-source version of **[HybridFlow: A Flexible and Efficient RLHF Framework](https://arxiv.org/abs/2409.19256v2)** paper.\n\nverl is flexible and easy to use with:\n\n- **Easy extension of diverse RL algorithms**: The hybrid-controller programming model enables flexible representation and efficient execution of complex post-training dataflows. Build RL dataflows such as GRPO, PPO in a few lines of code.\n\n- **Seamless integration of existing LLM infra with modular APIs**: Decouples computation and data dependencies, enabling seamless integration with existing LLM frameworks, such as FSDP, Megatron-LM, vLLM, SGLang, etc\n\n- **Flexible device mapping**: Supports various placement of models onto different sets of GPUs for efficient resource utilization and scalability across different cluster sizes.\n\n- Ready integration with popular HuggingFace models\n\nverl is fast with:\n\n- **State-of-the-art throughput**: SOTA LLM training and inference engine integrations and SOTA RL throughput.\n\n- **Efficient actor model resharding with 3D-HybridEngine**: Eliminates memory redundancy and significantly reduces communication overhead during transitions between training and generation phases.\n\n<div align=\"center\">\n <img src=\"https://github.com/verl-project/verl-data/blob/main/images/verl-arch.png?raw=true\" width=\"400\" alt=\"verl-arch.png\">\n</div>\n\n</p>\n\n## News\n\n- [2026/01] verl has been migrated to the [verl-project](https://github.com/verl-project)\n- [2026/01] verl first meetup was successfully held in Shanghai on 01/10, hosted by Volcengine and NVIDIA, the slides has been uploaded to [verl-data](https://github.com/verl-project/verl-data).\n- [2026/01] The `recipe` directory has been migrated to a dedicated repository: [verl-recipe](https://github.com/verl-project/verl-recipe) and added as a submodule. See https://github.com/volcengine/verl/pull/4795. It can be used as it was after `git submodule update --init --recursive recipe`. Note that [`transfer_queue`](verl/experimental/transfer_queue), [`fully_async_policy`](verl/experimental/fully_async_policy), [`one_step_off_policy`](verl/experimental/one_step_off_policy) and [`vla`](verl/experimental/vla) are kept under [`verl/experimental`](verl/experimental) since they are planned to be merged into the main library. Use them through `verl.experimental.{module}`.\n- [2025/12] [Mind Lab](https://macaron.im/mindlab) successfully used [verl](https://github.com/volcengine/verl) and [Megatron-bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) to train GRPO Lora for Trillion-parameter model on 64 H800 - See their [techblog](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus).\n- [2025/10] verl is presented in the [PyTorch Conference 2025](https://pytorch.org/event/pytorch-conference-2025/).\n- [2025/08] verl is presented in the [PyTorch Expert Exchange Webinar](https://www.youtube.com/watch?v=Vd79NmmqY3Q&t=2s). [Slides](https://github.com/eric-haibin-lin/verl-community/blob/main/slides/verl_talk_pytorch_2025_08.pdf) available.\n- [2025/07] The [ReTool](https://arxiv.org/pdf/2504.11536) recipe is fully open sourced. [Blog](https://www.notion.so/verl-reTool-recipe-Using-multi-round-conversations-and-code-sandboxing-to-improve-the-math-of-large-23a8b5b7feba80b386b2e5b5e3c1cde0)\n- [2025/07] The first verl meetup will be held at ICML Vancouver on July 16th! Please [join us](https://lu.ma/0ek2nyao) if you are at ICML! (onsite only)\n- [2025/06] verl with Megatron backend enables large MoE models such as [DeepSeek-671B and Qwen3-235B](https://verl.readthedocs.io/en/latest/perf/dpsk.html).\n- [2025/03] [DAPO](https://dapo-sia.github.io/) is the open-sourced SOTA RL algorithm that achieves 50 points on AIME 2024 based on the Qwen2.5-32B pre-trained model, surpassing the previous SOTA achieved by DeepSeek's GRPO (DeepSeek-R1-Zero-Qwen-32B). DAPO's training is fully powered by verl and the reproduction code is available in `recipe/dapo` now.\n<details><summary> more... </summary>\n<ul>\n  <li>[2025/04] [Seed-Thinking-v1.5](https://github.com/ByteDance-Seed/Seed-Thinking-v1.5/blob/main/seed-thinking-v1.5.pdf) tech report is released! Trained with verl, Seed-Thinking-v1.5 achieves 86.7 on AIME 2024, 55.0 on Codeforces and 77.3 on GPQA, demonstrating excellent reasoning abilities in STEM and coding. Beyond reasoning tasks, the method demonstrates notable generalization across diverse domains.</li>\n  <li>[2025/07] verl keynote at [AWS AI Hours Singapore](https://pages.awscloud.com/aws-ai-hours-sg.html#agenda) on 7/8, verl & verl-agent project updates at [Agent for SWE meetup](https://lu.ma/e498qhsi) by LF AI & Data Singapore on 7/11.</li>\n  <li>[2025/06] verl team will provide latest project updates at [PyTorch Day China](https://www.lfasiallc.com/pytorch-day-china/) on June 7th. Meet our dev team in Beijing!</li>\n  <li> [2025/04] [VAPO](https://arxiv.org/pdf/2504.05118) (value-based augmented PPO) paper covers our latest RL method for reasoning models. Trained from Qwen-32B-base model, VAPO achieves 60.4 on AIME 2024, outperforming DAPO-32B.</li>\n  <li>[2025/05] [PF-PPO](https://arxiv.org/abs/2409.06957), accepted to ICML 2025, is now supported in verl! PF-PPO enhances policy learning efficiency and robustness by filtering potentially noisy reward signals and reusing high-quality experiences via a replay buffer.</li>\n  <li>[2025/04] We will give a tutorial about latest post-training techniques and programming guide for verl at [ICLR 2025 Expo](https://iclr.cc/virtual/2025/calendar?filter_events=Expo+Talk+Panel&filter_rooms=), [SCI-FM workshop](https://open-foundation-model.github.io/) and [LMSys afterparty](https://lu.ma/d23nyynm). Talk materials available [here](https://github.com/eric-haibin-lin/verl-community/tree/main/iclr25). </li>\n  <li>[2025/03] verl v0.3.0.post1 is released! See [release note](https://github.com/volcengine/verl/releases/) for details. It achieves [~1.4x speedup](https://tongyx361.github.io/blogs/posts/verl-intro/#/verl-flexible-and-efficient-rl-for-llms) compared to prev versions.</li>\n  <li>[2025/05] verl will be presented at [A2M Shanghai](https://a2m.msup.com.cn/home/?aid=4488&city=shanghai) on 5/16 - 5/17.</li>\n  <li>[2025/05] verl will be presented at [GOSIM x PyTorch Day 2025](https://paris2025.gosim.org/). See you in Paris! </li>\n  <li>[2025/03] We introduced the programming model of verl at the [vLLM Beijing Meetup](https://mp.weixin.qq.com/s/n77GibL2corAtQHtVEAzfg) and [verl intro and updates](https://github.com/eric-haibin-lin/verl-community/blob/main/slides/verl-lmsys-meetup.pdf) at the [SGLang-LMSYS Org Meetup](https://lu.ma/ntjrr7ig) in Sunnyvale mid-March.</li>\n  <li>[2025/03] We will present verl(HybridFlow) at EuroSys 2025. See you in Rotterdam!</li>\n  <li>[2025/02] verl v0.2.0.post2 is released!</li>\n  <li>[2025/02] We presented verl in the <a href=\"https://lu.ma/ji7atxux\">Bytedance/NVIDIA/Anyscale Ray Meetup</a>. See you in San Jose!</li>\n  <li>[2025/01] [Doubao-1.5-pro](https://team.doubao.com/zh/special/doubao_1_5_pro) is released with SOTA-level performance on LLM & VLM. The RL scaling preview model is trained using verl, reaching OpenAI O1-level performance on math benchmarks (70.0 pass@1 on AIME).</li>\n  <li>[2024/12] verl is presented at Ray Forward 2024. Slides available <a href=\"https://github.com/eric-haibin-lin/verl-community/blob/main/slides/Ray_Forward_2024_%E5%B7%AB%E9%94%A1%E6%96%8C.pdf\">here</a></li>\n  <li>[2024/12] The team presented <a href=\"https://neurips.cc/Expo/Conferences/2024/workshop/100677\">Post-training LLMs: From Algorithms to Infrastructure</a> at NeurIPS 2024. <a href=\"https://github.com/eric-haibin-lin/verl-data/tree/neurips\">Slides</a> and <a href=\"https://neurips.cc/Expo/Conferences/2024/workshop/100677\">video</a> available.</li>\n  <li>[2024/10] verl is presented at Ray Summit. <a href=\"https://www.youtube.com/watch?v=MrhMcXkXvJU&list=PLzTswPQNepXntmT8jr9WaNfqQ60QwW7-U&index=37\">Youtube video</a> available.</li>\n  <li>[2024/08] HybridFlow (verl) is accepted to EuroSys 2025.</li>\n</ul>\n</details>\n\n## Key Features\n\n- **FSDP**, **FSDP2** and **Megatron-LM** for training.\n- **vLLM**, **SGLang** and **HF Transformers** for rollout generation.\n- Compatible with Hugging Face Transformers and Modelscope Hub: [Qwen-3](https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3-8b.sh), Qwen-2.5, Llama3.1, Gemma2, DeepSeek-LLM, etc\n- Supervised fine-tuning.\n- Reinforcement learning with [PPO](examples/ppo_trainer/), [GRPO](examples/grpo_trainer/), [GSPO](https://github.com/verl-project/verl-recipe/tree/main/gspo/), [ReMax](examples/remax_trainer/), [REINFORCE++](https://verl.readthedocs.io/en/latest/examples/config.html#algorithm), [RLOO](examples/rloo_trainer/), [PRIME](https://github.com/verl-project/verl-recipe/tree/main/prime/), [DAPO](https://github.com/verl-project/verl-recipe/tree/main/dapo/), [DrGRPO](https://github.com/verl-project/verl-recipe/tree/main/drgrpo), [KL_Cov & Clip_Cov](https://github.com/verl-project/verl-recipe/tree/main/entropy) etc.\n  - Support model-based reward and function-based reward (verifiable reward) for math, [coding](https://github.com/verl-project/verl-recipe/tree/main/dapo), etc\n  - Support vision-language models (VLMs) and [multi-modal RL](examples/grpo_trainer/run_qwen2_5_vl-7b.sh) with Qwen2.5-vl, Kimi-VL\n  - [Multi-turn with tool calling](https://github.com/volcengine/verl/tree/main/examples/sglang_multiturn)\n- LLM alignment recipes such as [Self-play preference optimization (SPPO)](https://github.com/verl-project/verl-recipe/tree/main/sppo)\n- Flash attention 2, [sequence packing](examples/ppo_trainer/run_qwen2-7b_seq_balance.sh), [sequence parallelism](examples/ppo_trainer/run_deepseek7b_llm_sp2.sh) support via DeepSpeed Ulysses, [LoRA](examples/sft/gsm8k/run_qwen_05_peft.sh), [Liger-kernel](examples/sft/gsm8k/run_qwen_05_sp2_liger.sh).\n- Scales up to 671B models and hundreds of GPUs with [expert parallelism](https://github.com/volcengine/verl/pull/1467)\n- Multi-gpu [LoRA RL](https://verl.readthedocs.io/en/latest/advance/ppo_lora.html) support to save memory.\n- Experiment tracking with wandb, swanlab, mlflow and tensorboard.\n- Hardware Support: Supports NVIDIA, AMD, [Ascend](https://github.com/volcengine/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_quick_start.rst)\n\n## Upcoming Features and Changes\n\n- Q3 Roadmap https://github.com/volcengine/verl/issues/2388\n- DeepSeek 671b optimizations with Megatron https://github.com/volcengine/verl/issues/1033\n- Multi-turn rollout and tools using optimizations https://github.com/volcengine/verl/issues/1882\n- [Agent integration](https://github.com/volcengine/verl/tree/main/verl/experimental/agent_loop)\n- Async and off-policy architecture https://github.com/volcengine/verl/pull/2231\n- List of breaking changes since v0.4 https://github.com/volcengine/verl/discussions/2270\n\n## Getting Started\n\n<a href=\"https://verl.readthedocs.io/en/latest/index.html\"><b>Documentation</b></a>\n\n**Quickstart:**\n\n- [Installation](https://verl.readthedocs.io/en/latest/start/install.html)\n- [Quickstart](https://verl.readthedocs.io/en/latest/start/quickstart.html)\n- [Programming Guide](https://verl.readthedocs.io/en/latest/hybrid_flow.html) & [Tech Talk](https://hcqnc.xetlk.com/sl/3vACOK) (in Chinese)\n- [PPO in verl](https://verl.readthedocs.io/en/latest/algo/ppo.html)\n- [GRPO in verl](https://verl.readthedocs.io/en/latest/algo/grpo.html)\n\n**Running a PPO example step-by-step:**\n\n- [Prepare Data for Post-Training](https://verl.readthedocs.io/en/latest/preparation/prepare_data.html)\n- [Implement Reward Function for Dataset](https://verl.readthedocs.io/en/latest/preparation/reward_function.html)\n- [PPO Example Architecture](https://verl.readthedocs.io/en/latest/examples/ppo_code_architecture.html)\n- [Config Explanation](https://verl.readthedocs.io/en/latest/examples/config.html)\n\n**Reproducible algorithm baselines:**\n\n- [RL performance on coding, math](https://verl.readthedocs.io/en/latest/algo/baseline.html)\n\n**For code explanation and advance usage (extension):**\n\n- PPO Trainer and Workers\n\n  - [PPO Ray Trainer](https://verl.readthedocs.io/en/latest/workers/ray_trainer.html)\n  - [PyTorch FSDP Backend](https://verl.readthedocs.io/en/latest/workers/fsdp_workers.html)\n  - [Megatron-LM Backend](https://verl.readthedocs.io/en/latest/index.html)\n\n- Advanced Usage and Extension\n  - [Add Models with the FSDP Backend](https://verl.readthedocs.io/en/latest/advance/fsdp_extension.html)\n  - [Add Models with the Megatron-LM Backend](https://verl.readthedocs.io/en/latest/advance/megatron_extension.html)\n  - [Multi-turn Rollout Support](https://verl.readthedocs.io/en/latest/sglang_multiturn/multiturn.html)\n  - [Search Tool Integration](https://verl.readthedocs.io/en/latest/sglang_multiturn/search_tool_example.html)\n  - [Sandbox Fusion Integration](https://verl.readthedocs.io/en/latest/examples/sandbox_fusion_example.html)\n  - [Deployment using Separate GPU Resources](https://github.com/volcengine/verl/tree/main/examples/split_placement)\n  - [Extend to Other RL(HF) algorithms](https://verl.readthedocs.io/en/latest/advance/dpo_extension.html)\n  - [Ray API design tutorial](https://verl.readthedocs.io/en/latest/advance/placement.html)\n\n**Blogs from the community**\n\n- [When Reasoning Models Break Tokenization: The Hidden Complexity of Multiturn Training](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/fast_tokenization/multiturn_tokenization_and_masking.md)\n- [verl deployment on AWS SageMaker](https://medium.com/@kaige.yang0110/run-verl-on-sagemaker-using-4x8-l40s-gpus-8e6d5c3c61d3)\n- [verl x SGLang Multi-turn Code Walkthrough](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/code-walk-through/readme_EN.md)\n- [Optimizing SGLang Memory Usage in verl](https://hebiao064.github.io/rl-memory-management)\n- [SGLang, verl, OpenBMB and Tsinghua University: Pioneering End-to-End Multi-Turn RLHF](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/verl-multiturn-rollout-Release.md)\n- [Reinforcement Learning from Human Feedback on AMD GPUs with verl and ROCm Integration](https://rocm.blogs.amd.com/artificial-intelligence/verl-large-scale/README.html)\n- [veMLP x verl ：玩转强化学习训练](https://mp.weixin.qq.com/s/7nbqxk4knMGd-hQE9ls2tA)\n- [使用 verl 进行 GRPO 分布式强化学习训练最佳实践](https://www.volcengine.com/docs/6459/1463942)\n- [HybridFlow verl 原文浅析](https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/readme.md)\n- [最高提升 20 倍吞吐量！豆包大模型团队发布全新 RLHF 框架，现已开源！](https://team.doubao.com/en/blog/%E6%9C%80%E9%AB%98%E6%8F%90%E5%8D%8720%E5%80%8D%E5%90%9E%E5%90%90%E9%87%8F-%E8%B1%86%E5%8C%85%E5%A4%A7%E6%A8%A1%E5%9E%8B%E5%9B%A2%E9%98%9F%E5%8F%91%E5%B8%83%E5%85%A8%E6%96%B0-rlhf-%E6%A1%86%E6%9E%B6-%E7%8E%B0%E5%B7%B2%E5%BC%80%E6%BA%90)\n\n## Performance Tuning Guide\n\nThe performance is essential for on-policy RL algorithm. We have written a detailed [performance tuning guide](https://verl.readthedocs.io/en/latest/perf/perf_tuning.html) to help you optimize performance.\n\n## Upgrade to vLLM >= v0.8.2\n\nverl now supports vLLM>=0.8.2 when using FSDP as the training backend. Please refer to [this document](https://github.com/volcengine/verl/blob/main/docs/README_vllm0.8.md) for the installation guide and more information. Please avoid vllm 0.7.x, which contains bugs that may lead to OOMs and unexpected errors.\n\n## Use Latest SGLang\n\nSGLang is fully supported with verl, and SGLang RL Group is working extensively on building unique features, including multi-turn agentic RL, VLM RLHF, server-based RL, and partial rollout. Please refer to [this document](https://verl.readthedocs.io/en/latest/workers/sglang_worker.html) for the installation guide and more information.\n\n## Upgrade to FSDP2\n\nverl is fully embracing FSDP2! FSDP2 is recommended by torch distributed team, providing better throughput and memory usage, and is composible with other features (e.g. torch.compile). To enable FSDP2, simply use verl main and set the following options:\n\n```\nactor_rollout_ref.ref.strategy=fsdp2\nactor_rollout_ref.actor.strategy=fsdp2\ncritic.strategy=fsdp2\n```\n\nFurthermore, FSDP2 cpu offloading is compatible with gradient accumulation. You can turn it on to save memory with `actor_rollout_ref.actor.fsdp_config.offload_policy=True`. For more details, see https://github.com/volcengine/verl/pull/1026\n\n## AMD Support (ROCm Kernel)\n\nverl now supports FSDP as the training engine (Megatron support coming soon) and both integrates with vLLM and SGLang as inference engines. Please refer to [this document](https://github.com/volcengine/verl/blob/main/docs/amd_tutorial/amd_build_dockerfile_page.rst) for the installation guide and more information, and [this document](https://github.com/volcengine/verl/blob/main/docs/amd_tutorial/amd_vllm_page.rst) for the vLLM performance tuning for ROCm.\n\n## Citation and acknowledgement\n\nIf you find the project helpful, please cite:\n\n- [HybridFlow: A Flexible and Efficient RLHF Framework](https://arxiv.org/abs/2409.19256v2)\n- [A Framework for Training Large Language Models for Code Generation via Proximal Policy Optimization](https://i.cs.hku.hk/~cwu/papers/gmsheng-NL2Code24.pdf)\n\n```bibtex\n@article{sheng2024hybridflow,\n  title   = {HybridFlow: A Flexible and Efficient RLHF Framework},\n  author  = {Guangming Sheng and Chi Zhang and Zilingfeng Ye and Xibin Wu and Wang Zhang and Ru Zhang and Yanghua Peng and Haibin Lin and Chuan Wu},\n  year    = {2024},\n  journal = {arXiv preprint arXiv: 2409.19256}\n}\n```\n\nverl is inspired by the design of Nemo-Aligner, Deepspeed-chat and OpenRLHF. The project is adopted and contributed by Bytedance, Anyscale, LMSys.org, [Alibaba Qwen team](https://github.com/QwenLM/), Shanghai AI Lab, Tsinghua University, UC Berkeley, UCLA, UIUC, University of Hong Kong, ke.com, [All Hands AI](https://www.all-hands.dev/), [ModelBest](http://modelbest.cn/), JD AI Lab, Microsoft Research, [StepFun](https://www.stepfun.com/), Amazon, LinkedIn, Meituan, [Camel-AI](https://www.camel-ai.org/), [OpenManus](https://github.com/OpenManus), Xiaomi, NVIDIA research, [Baichuan](https://www.baichuan-ai.com/home), [RedNote](https://www.xiaohongshu.com/), [SwissAI](https://www.swiss-ai.org/), [Moonshot AI (Kimi)](https://www.moonshot-ai.com/), Baidu, Snowflake, Skywork.ai, JetBrains, [IceSword Lab](https://www.iceswordlab.com), and many more.\n\n## Awesome Projects Built with `verl`\n\nWelcome to register your awesome project build with `verl` for other developers' reference!\n\n- [TinyZero](https://github.com/Jiayi-Pan/TinyZero): a reproduction of **DeepSeek R1 Zero** recipe for reasoning tasks ![GitHub Repo stars](https://img.shields.io/github/stars/Jiayi-Pan/TinyZero)\n- [SkyThought](https://github.com/NovaSky-AI/SkyThought): RL training for Sky-T1-7B by NovaSky AI team. ![GitHub Repo stars](https://img.shields.io/github/stars/NovaSky-AI/SkyThought)\n- [simpleRL-reason](https://github.com/hkust-nlp/simpleRL-reason): SimpleRL-Zoo: Investigating and Taming Zero Reinforcement Learning for Open Base Models in the Wild ![GitHub Repo stars](https://img.shields.io/github/stars/hkust-nlp/simpleRL-reason)\n- [Easy-R1](https://github.com/hiyouga/EasyR1): **Multi-modal** RL training framework ![GitHub Repo stars](https://img.shields.io/github/stars/hiyouga/EasyR1)\n- [OpenManus-RL](https://github.com/OpenManus/OpenManus-RL): LLM Agents RL tuning framework for multiple agent environments. ![GitHub Repo stars](https://img.shields.io/github/stars/OpenManus/OpenManus-RL)\n- [rllm](https://github.com/agentica-project/rllm): async RL training with [verl-pipeline](https://github.com/agentica-project/verl-pipeline) ![GitHub Repo stars](https://img.shields.io/github/stars/agentica-project/rllm)\n- [RAGEN](https://github.com/ZihanWang314/ragen): a general-purpose reasoning **agent** training framework ![GitHub Repo stars](https://img.shields.io/github/stars/ZihanWang314/ragen)\n- [Search-R1](https://github.com/PeterGriffinJin/Search-R1): RL with reasoning and **searching (tool-call)** interleaved LLMs ![GitHub Repo stars](https://img.shields.io/github/stars/PeterGriffinJin/Search-R1)\n- [ReSearch](https://github.com/Agent-RL/ReSearch): Learning to **Re**ason with **Search** for LLMs via Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/Agent-RL/ReSearch)\n- [Skywork-OR1](https://github.com/SkyworkAI/Skywork-OR1): Skywork open reaonser series ![GitHub Repo stars](https://img.shields.io/github/stars/SkyworkAI/Skywork-OR1)\n- [ToRL](https://github.com/GAIR-NLP/ToRL): Scaling tool-integrated RL ![GitHub Repo stars](https://img.shields.io/github/stars/GAIR-NLP/ToRL)\n- [Absolute Zero Reasoner](https://github.com/LeapLabTHU/Absolute-Zero-Reasoner): [A no human curated data self-play framework for reasoning](https://arxiv.org/abs/2505.03335) ![GitHub Repo stars](https://img.shields.io/github/stars/LeapLabTHU/Absolute-Zero-Reasoner)\n- [verl-agent](https://github.com/langfengQ/verl-agent): A scalable training framework for **long-horizon LLM/VLM agents**, along with a new algorithm **GiGPO** ![GitHub Repo stars](https://img.shields.io/github/stars/langfengQ/verl-agent)\n- [RL-Factory](https://github.com/Simple-Efficient/RL-Factory): An easy and efficient RL post-training framework for Agentic Learning ![GitHub Repo stars](https://img.shields.io/github/stars/Simple-Efficient/RL-Factory)\n- [ReTool](https://retool-rl.github.io/): ReTool: reinforcement learning for strategic tool use in LLMs. Code release is in progress...\n- [verl-tool](https://github.com/TIGER-AI-Lab/verl-tool): An unified and easy-to-extend tool-agent training framework based on verl![GitHub Repo stars](https://img.shields.io/github/stars/TIGER-AI-Lab/verl-tool)\n- [PRIME](https://github.com/PRIME-RL/PRIME): Process reinforcement through implicit rewards ![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/PRIME)\n- [MemAgent](https://github.com/BytedTsinghua-SIA/MemAgent): MemAgent: Reshaping Long-Context LLM with Multi-Conv RL based Memory Agent ![GitHub Repo stars](https://img.shields.io/github/stars/BytedTsinghua-SIA/MemAgent)\n- [POLARIS](https://github.com/ChenxinAn-fdu/POLARIS): A Post-training recipe for scaling RL on Advanced Reasoning models ![GitHub Repo stars](https://img.shields.io/github/stars/ChenxinAn-fdu/POLARIS)\n- [GUI-R1](https://github.com/ritzz-ai/GUI-R1): **GUI-R1**: A Generalist R1-style Vision-Language Action Model For **GUI Agents** ![GitHub Repo stars](https://img.shields.io/github/stars/ritzz-ai/GUI-R1)\n- [DeepRetrieval](https://github.com/pat-jj/DeepRetrieval): RL Training of **Search Agent** with **Search/Retrieval Outcome** ![GitHub Repo stars](https://img.shields.io/github/stars/pat-jj/DeepRetrieval)\n- [Code-R1](https://github.com/ganler/code-r1): Reproducing R1 for **Code** with Reliable Rewards ![GitHub Repo stars](https://img.shields.io/github/stars/ganler/code-r1)\n- [DeepResearcher](https://github.com/GAIR-NLP/DeepResearcher): Scaling deep research via reinforcement learning in real-world environments ![GitHub Repo stars](https://img.shields.io/github/stars/GAIR-NLP/DeepResearcher)\n- [VAGEN](https://github.com/RAGEN-AI/VAGEN): Training VLM agents with multi-turn reinforcement learning ![GitHub Repo stars](https://img.shields.io/github/stars/RAGEN-AI/VAGEN)\n- [RM-R1](https://arxiv.org/abs/2505.02387): RL training of reasoning reward models ![GitHub Repo stars](https://img.shields.io/github/stars/RM-R1-UIUC/RM-R1)\n- [Dr. MAS](https://arxiv.org/pdf/2602.08847): Stable **end-to-end RL** post-training for **multi-agent LLM systems** ![GitHub Repo stars](https://img.shields.io/github/stars/langfengQ/DrMAS)\n- [LUFFY](https://arxiv.org/pdf/2504.14945): Learning to Reason under Off-Policy Guidance![GitHub Repo stars](https://img.shields.io/github/stars/ElliottYan/LUFFY)\n- [DeepMath](https://github.com/zwhe99/DeepMath): DeepMath-103K data and series models for math reasoning![GitHub Repo stars](https://img.shields.io/github/stars/zwhe99/DeepMath)\n- [PACS](https://github.com/ritzz-ai/PACS): Implicit Actor Critic Coupling via a Supervised Learning Framework for RLVR ![GitHub Repo stars](https://img.shields.io/github/stars/ritzz-ai/PACS)\n- [Entropy Mechanism of RL](https://github.com/PRIME-RL/Entropy-Mechanism-of-RL): The Entropy Mechanism of Reinforcement Learning for Large Language Model Reasoning![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/Entropy-Mechanism-of-RL)\n- [LLaSA-TTS-GRPO](https://github.com/channel-io/ch-tts-llasa-rl-grpo): TTS fine-tuning with GRPO optimization based on LLASA models ![GitHub Repo stars](https://img.shields.io/github/stars/channel-io/ch-tts-llasa-rl-grpo)\n- [PF-PPO](https://arxiv.org/abs/2409.06957): Policy Filtration for PPO based on the reliability of reward signals for more efficient and robust RLHF.\n- [RACRO](https://github.com/gyhdog99/RACRO2): Build multi-modal reasoning models via decoupling it into query-conditioned captioning and text-only reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/gyhdog99/RACRO2)\n- [Agent Lightning](https://github.com/microsoft/agent-lightning): A flexible and extensible framework that enables seamless agent optimization for any existing agent framework. ![GitHub Repo stars](https://img.shields.io/github/stars/microsoft/agent-lightning)\n- [VTool-R1](https://github.com/VTOOL-R1/vtool-r1): VLMs Learn to Think with Images via Reinforcement Learning on Multimodal Tool Use. ![GitHub Repo stars](https://img.shields.io/github/stars/VTOOL-R1/vtool-r1)\n- [Kimina-Prover-RL](https://github.com/project-numina/kimina-prover-rl/tree/main/recipe/kimina_prover_rl): Training pipeline for formal theorem proving, based on a paradigm inspired by DeepSeek-R1.\n- [RL-PLUS](https://github.com/YihongDong/RL-PLUS): Countering Capability Boundary Collapse of LLMs in Reinforcement Learning with Hybrid-policy Optimization.\n- [rStar2-Agent](https://github.com/microsoft/rStar): Using reinforcement learning with multi-step tool-calling for math tasks, rStar2-Agent-14B reaches frontier-level math reasoning in just 510 RL training steps ![GitHub Repo stars](https://img.shields.io/github/stars/microsoft/rStar)\n- [Vision-SR1](https://github.com/zli12321/Vision-SR1): Self-Rewarding Vision-Language Model via Reasoning Decomposition ![GitHub Repo stars](https://img.shields.io/github/stars/zli12321/Vision-SR1)\n- [SimpleVLA-RL](https://github.com/PRIME-RL/SimpleVLA-RL): SimpleVLA-RL: A Simple yet Effective Vision-Language Action Model for Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/PRIME-RL/SimpleVLA-RL)\n- [Table-R1](https://github.com/Table-R1/Table-R1): Table-R1: Inference-Time Scaling for Table Reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/Table-R1/Table-R1)\n- [Revisual-R1](https://github.com/CSfufu/Revisual-R1): Revisual-R1: Advancing Multimodal Reasoning From Optimized Cold Start to Staged Reinforcement Learning ![GitHub Repo stars](https://img.shields.io/github/stars/CSfufu/Revisual-R1)\n- [ARES](https://github.com/shawn0728/ARES): ARES: Multimodal Adaptive Reasoning via Difficulty-Aware Token-Level Entropy Shaping ![GitHub Repo stars](https://img.shields.io/github/stars/shawn0728/ARES)\n- [Meta-Bandit-LLM](https://github.com/sanxing-chen/meta-bandit-llm): Meta-Bandit-LLM: Long-horizon multiturn interactive training for meta-bandit agents ![GitHub Repo stars](https://img.shields.io/github/stars/sanxing-chen/meta-bandit-llm)\n- [PokeeResearch](https://github.com/Pokee-AI/PokeeResearchOSS): PokeeResearch: State-of-the-art 7B DeepResearch Agent that leverages web search and content reading capabilities to answer complex questions using the most up-to-date information available online. ![Github Repo Stars](https://img.shields.io/github/stars/Pokee-AI/PokeeResearchOSS)\n- [Search Self-play](https://github.com/Alibaba-Quark/SSP): Pushing the Frontier of Agent Capability without Supervision ![GitHub Repo stars](https://img.shields.io/github/stars/Alibaba-Quark/SSP)\n- [OneThinker](https://github.com/tulerfeng/OneThinker): All-in-one Reasoning Model for Image and Video ![GitHub Repo stars](https://img.shields.io/github/stars/tulerfeng/OneThinker)\n- [OpenTinker](https://github.com/open-tinker/OpenTinker): Democratizing Agentic Reinforcement Learning as a Service ![GitHub Repo stars](https://img.shields.io/github/stars/open-tinker/OpenTinker)\n- [FlowRL](https://github.com/Xuekai-Zhu/FlowRL): Matching reward distributions via **flow balance** for diverse exploration and generalizable reasoning ![GitHub Repo stars](https://img.shields.io/github/stars/Xuekai-Zhu/FlowRL)\n- [Logic-RL](https://github.com/Unakar/Logic-RL): a reproduction of DeepSeek R1 Zero on 2K Tiny Logic Puzzle Dataset. ![GitHub Repo stars](https://img.shields.io/github/stars/Unakar/Logic-RL)\n- [Seed-Coder](https://github.com/ByteDance-Seed/Seed-Coder): RL training of Seed-Coder boosts performance on competitive programming ![GitHub Repo stars](https://img.shields.io/github/stars/ByteDance-Seed/Seed-Coder)\n- [all-hands/openhands-lm-32b-v0.1](https://www.all-hands.dev/blog/introducing-openhands-lm-32b----a-strong-open-coding-agent-model): A strong, open coding agent model, trained with [multi-turn fine-tuning](https://github.com/volcengine/verl/pull/195)\n- [s3](https://github.com/pat-jj/s3) **Efficient Yet Effective** Search Agent Training via RL ![GitHub Repo stars](https://img.shields.io/github/stars/pat-jj/s3)\n- [Rec-R1](https://arxiv.org/pdf/2503.24289): Bridging Generative Large Language Models and Recommendation Systems via Reinforcement Learning\n- [Explore RL Data Scaling](https://arxiv.org/abs/2503.22230): Exploring Data Scaling Trends and Effects in Reinforcement Learning from Human Feedback\n- [FIRE](https://arxiv.org/abs/2410.21236): Flaming-hot initiation with regular execution sampling for large language models\n- [DQO](https://arxiv.org/abs/2410.09302): Enhancing multi-Step reasoning abilities of language models through direct Q-function optimization\n- [ProRL](https://arxiv.org/abs/2505.24864): Prolonged Reinforcement Learning Expands Reasoning Boundaries in Large Language Models\n- [cognition-engineering](https://github.com/gair-nlp/cognition-engineering): Test time scaling drives cognition engineering. ![GitHub Repo stars](https://img.shields.io/github/stars/gair-nlp/cognition-engineering)\n- [Trust Region Preference Approximation](https://github.com/XueruiSu/Trust-Region-Preference-Approximation): A simple and stable **reinforcement learning algorithm** for LLM reasoning. ![GitHub Repo stars](https://img.shields.io/github/stars/XueruiSu/Trust-Region-Preference-Approximation)\n- [AdaRFT](https://github.com/uscnlp-lime/verl): Efficient Reinforcement Finetuning via **Adaptive Curriculum Learning** ![GitHub Repo stars](https://img.shields.io/github/stars/uscnlp-lime/verl)\n- [critic-rl](https://github.com/HKUNLP/critic-rl): LLM critics for code generation ![GitHub Repo stars](https://img.shields.io/github/stars/HKUNLP/critic-rl)\n- [self-rewarding-reasoning-LLM](https://arxiv.org/pdf/2502.19613): self-rewarding and correction with **generative reward models** ![GitHub Repo stars](https://img.shields.io/github/stars/RLHFlow/Self-rewarding-reasoning-LLM)\n- [DeepEnlighten](https://github.com/DolbyUUU/DeepEnlighten): Reproduce R1 with **social reasoning** tasks and analyze key findings ![GitHub Repo stars](https://img.shields.io/github/stars/DolbyUUU/DeepEnlighten)\n- [MetaSpatial](https://github.com/PzySeere/MetaSpatial): Reinforcing **3D Spatial Reasoning** in **VLMs** for the **Metaverse** ![GitHub Repo stars](https://img.shields.io/github/stars/PzySeere/MetaSpatial)\n- [PURE](https://github.com/CJReinforce/PURE): **Credit assignment** is the key to successful reinforcement fine-tuning using **process reward model** ![GitHub Repo stars](https://img.shields.io/github/stars/CJReinforce/PURE)\n- [cognitive-behaviors](https://github.com/kanishkg/cognitive-behaviors): Cognitive Behaviors that Enable Self-Improving Reasoners, or, Four Habits of Highly Effective STaRs ![GitHub Repo stars](https://img.shields.io/github/stars/kanishkg/cognitive-behaviors)\n- [deepscaler](https://github.com/agentica-project/rllm/tree/deepscaler): iterative context scaling with GRPO ![GitHub Repo stars](https://img.shields.io/github/stars/agentica-project/deepscaler)\n- [DAPO](https://dapo-sia.github.io/): the fully open source SOTA RL algorithm that beats DeepSeek-R1-zero-32B ![GitHub Repo stars](https://img.shields.io/github/stars/volcengine/verl)\n- [NoisyRollout](https://github.com/NUS-TRAIL/NoisyRollout): Reinforcing Visual Reasoning with Data Augmentation ![GitHub Repo stars](https://img.shields.io/github/stars/NUS-TRAIL/NoisyRollout)\n- [SPEAR](https://github.com/TencentYoutuResearch/SPEAR): **Self-imitation** with **Progressive Exploration** for Agentic Reinforcement Learning (ICLR 2026) ![GitHub Repo stars](https://img.shields.io/github/stars/TencentYoutuResearch/SPEAR)\n- [RuleReasoner](https://github.com/bigai-nlco/RuleReasoner): **RuleReasoner:** Reinforced Rule-based Reasoning via **Domain-aware Dynamic Sampling** (ICLR 2026) ![GitHub Repo stars](https://img.shields.io/github/stars/bigai-nlco/RuleReasoner)\n- [MetaphorStar](https://metaphorstar.github.io/): **Image Metaphor** Understanding and Reasoning with End-to-End **Visual Reinforcement Learning** ![GitHub Repo stars](https://img.shields.io/github/stars/MING-ZCH/MetaphorStar)\n\n## Contribution Guide\n\nSee [contributions guide](CONTRIBUTING.md)\n\n## About [ByteDance Seed Team](https://team.doubao.com/)\n\nFounded in 2023, ByteDance Seed Team is dedicated to crafting the industry's most advanced AI foundation models. The team aspires to become a world-class research team and make significant contributions to the advancement of science and society. You can get to know Bytedance Seed better through the following channels👇\n\n<div>\n  <a href=\"https://team.doubao.com/\">\n    <img src=\"https://img.shields.io/badge/Website-%231e37ff?style=for-the-badge&logo=bytedance&logoColor=white\"></a>\n  <a href=\"https://github.com/user-attachments/assets/469535a8-42f2-4797-acdf-4f7a1d4a0c3e\">\n    <img src=\"https://img.shields.io/badge/WeChat-07C160?style=for-the-badge&logo=wechat&logoColor=white\"></a>\n <a href=\"https://www.xiaohongshu.com/user/profile/668e7e15000000000303157d?xsec_token=ABl2-aqekpytY6A8TuxjrwnZskU-6BsMRE_ufQQaSAvjc%3D&xsec_source=pc_search\">\n    <img src=\"https://img.shields.io/badge/Xiaohongshu-%23FF2442?style=for-the-badge&logo=xiaohongshu&logoColor=white\"></a>\n  <a href=\"https://www.zhihu.com/org/dou-bao-da-mo-xing-tuan-dui/\">\n    <img src=\"https://img.shields.io/badge/zhihu-%230084FF?style=for-the-badge&logo=zhihu&logoColor=white\"></a>\n</div>\n\nWe are HIRING! Send us an [email](mailto:the.verl.project@gmail.com) if you are interested in internship/FTE opportunities in RL for agents.\n"
  },
  {
    "path": "docker/Dockerfile.isaaclab230",
    "content": "\n#FROM nvcr.nju.edu.cn/nvidia/isaac-lab:2.3.0\nFROM isaac-lab-base:latest\n\nENV ACCEPT_EULA=Y\nENTRYPOINT []\n\n# desktop\nRUN --mount=type=cache,target=/var/cache/apt \\\n    sed -i 's/archive.ubuntu.com/mirrors.ivolces.com/g' /etc/apt/sources.list && \\\n    sed -i 's/security.ubuntu.com/mirrors.ivolces.com/g' /etc/apt/sources.list && \\\n    apt-get update && \\\n    DEBIAN_FRONTEND=noninteractive apt-get install -y locales && \\\n    locale-gen en_US.UTF-8 && \\\n    update-locale LANG=en_US.UTF-8 LC_CTYPE=en_US.UTF-8 && \\\n    apt-get install -y wget curl \\\n        xfce4 \\\n        xfce4-goodies \\\n        xorg \\\n        dbus-x11 \\\n        x11-xserver-utils \\\n        tigervnc-standalone-server \\\n        tigervnc-common \\\n        tigervnc-tools \\\n        fonts-dejavu \\\n        fonts-liberation\n# cuda 12.2\nRUN --mount=type=cache,target=/var/cache/apt \\\n    cd /tmp && \\\n    wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/3bf863cc.pub && \\\n    apt-key add 3bf863cc.pub && \\\n    echo \"deb https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64 /\" > /etc/apt/sources.list.d/cuda.list && \\\n    apt-get update && \\\n    apt-get install -y libcusparselt0 libnccl2=2.27.3-1+cuda12.2 libglfw3 libgl1-mesa-glx libosmesa6 && \\\n    rm -f 3bf863cc.pub\n\n# libero\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install easydict==1.9 robosuite==1.4.0 bddl==1.0.1 future==0.18.2 cloudpickle==2.1.0\n\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install transformers[hf_xet]\n\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install --upgrade numpy==1.26.4 ray[default] \\\n    accelerate codetiming datasets dill hydra-core pandas peft pyarrow>=19.0.0 pybind11 pylatexenc\n\n# openvla-oft\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install pre-commit torchdata packaging>=20.0 uvicorn fastapi latex2sympy2_extended math_verify tensorboard\n\n\n# flash_attn\nRUN cd /tmp && \\\n    wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl && \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install /tmp/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl && \\\n    rm -f /tmp/flash_attn-2.8.0.post2+cu12torch2.7cxx11abiFALSE-cp311-cp311-linux_x86_64.whl\n\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install --upgrade protobuf==3.20.3 timm==0.9.16\n\nRUN --mount=type=cache,target=/root/.cache/pip \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install orjson==3.11.3 pyvers==0.1.0 tensordict==0.10.0 --force --no-deps\n\n\nRUN mkdir -p /root/.vnc && \\\n    cat <<'EOP' > /root/.vnc/xstartup \n#!/bin/sh\nunset SESSION_MANAGER\nunset DBUS_SESSION_BUS_ADDRESS\n[ -r \\$HOME/.Xresources ] && xrdb \\$HOME/.Xresources\nxsetroot -solid grey\nexec startxfce4\nEOP\n\nRUN cat <<'EOP' > /root/.vnc/config\ngeometry=1920x1080\ndepth=24\ndesktop=Isaac-Sim-Desktop\ndpi=96\nlocalhost=no\nEOP\n\nRUN cat <<'EOP' > /root/start_isaac_vnc.sh\n#!/bin/bash\n# 设置显示变量\nexport DISPLAY=:1\n\n# 检查VNC是否运行\nif ! pgrep -f \"Xvnc.*:1\" > /dev/null; then\n    echo \"Starting VNC server...\"\n    vncserver :1 -localhost no -geometry 1920x1080 -depth 24 -desktop \"Isaac-Sim-Desktop\"\n    sleep 3\nfi\n\n# 启动Isaac Sim\necho \"Starting Isaac Sim...\"\n/workspace/isaaclab/_isaac_sim/isaac-sim.sh --allow-root\nEOP\n\nRUN chmod +x /root/.vnc/xstartup && \\\n    chmod +x /root/start_isaac_vnc.sh\n\nRUN /workspace/isaaclab/_isaac_sim/isaac-sim.sh --allow-root --ext-precache-mode\n\nRUN cd /root && \\\n    git clone https://github.com/Lifelong-Robot-Learning/LIBERO.git && \\\n    cd LIBERO && \\\n    git apply <<'EOP'\ndiff --git a/setup.py b/setup.py\nindex 59d4900..dbe9811 100644\n--- a/setup.py\n+++ b/setup.py\n@@ -13,7 +13,8 @@ long_description = \"\".join(lines)\n\n setup(\n     name=\"libero\",\n-    packages=[package for package in find_packages() if package.startswith(\"libero\")],\n+    #packages=[package for package in find_packages() if package.startswith(\"libero\")],\n+    packages=[\"libero\"],\n     install_requires=[],\n     eager_resources=[\"*\"],\n     include_package_data=True,\nEOP\n\nRUN cd /root/LIBERO && \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install -e .\n\n# libero config\nRUN mkdir -p /root/.libero && \\\ncat <<'EOP' > /root/.libero/config.yaml\nassets: /root/LIBERO/libero/libero/./assets\nbddl_files: /root/LIBERO/libero/libero/./bddl_files\nbenchmark_root: /root/LIBERO/libero/libero\ndatasets: /root/LIBERO/libero/libero/../datasets\ninit_states: /root/LIBERO/libero/libero/./init_files\nEOP\n\n# from https://github.com/nvidia-china-sae/RobotLearningLab\nCOPY RobotLearningLab/ /root/RobotLearningLab/\n\nRUN cd /workspace/isaaclab/ && \\\n    rm -rf source && \\\n    ln -s /root/RobotLearningLab/source source && \\\n    /workspace/isaaclab/_isaac_sim/python.sh -m pip install -e ./source/isaaclab\n# Ray cmd\nRUN /workspace/isaaclab/_isaac_sim/python.sh -m pip install colorama && \\\ncat <<'EOP' >> /root/.bashrc\nalias ray='/workspace/isaaclab/_isaac_sim/python.sh /workspace/isaaclab/_isaac_sim/kit/python/lib/python3.11/site-packages/ray/scripts/scripts.py'\nEOP"
  },
  {
    "path": "docker/Dockerfile.stable.sglang",
    "content": "# sgl059\n\nFROM lmsysorg/sglang:v0.5.9\n\nARG PIP_NO_CACHE_DIR=1\n\nRUN pip install pybind11\n\nRUN pip install nvidia-mathdx\n\nRUN MAX_JOBS=128 pip install -v --disable-pip-version-check --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" git+https://github.com/NVIDIA/apex.git\n\nRUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.12\n\n# RUN pip install --upgrade transformers tokenizers\n\nRUN pip install codetiming mathruler pylatexenc qwen_vl_utils cachetools pytest-asyncio\n\nRUN pip install --no-build-isolation flash_attn==2.8.3\n\nRUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ \"$(uname -m)\" = \"aarch64\" ]; then echo \"arm64\"; else echo \"amd64\"; fi) && \\\n    wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0 && \\\n    apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-${NSIGHT_VERSION}.deb\n\n# sglang image has already installed DeepEP\n\nRUN pip3 install --no-deps trl==0.27.0\n\nRUN pip3 install nvtx matplotlib liger_kernel\n\nRUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n\nRUN pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.16.0\n\nRUN pip install git+https://github.com/volcengine/verl.git@v0.7.0 && \\\n    pip uninstall -y verl\n\nRUN sed -i '/nvidia-cudnn-cu12/d' /usr/local/lib/python3.12/dist-packages/torch-2.9.1+cu129.dist-info/METADATA && \\\n    pip install --no-deps --force-reinstall nvidia-cudnn-cu12==9.16.0.29\n\n# for packages compiled from source code\nRUN apt-get update && \\\n    apt-get install -y --allow-downgrades --allow-change-held-packages \\\n    libcudnn9-cuda-12=9.16.0.29-1 \\\n    libcudnn9-dev-cuda-12=9.16.0.29-1 \\\n    libcudnn9-headers-cuda-12=9.16.0.29-1 && \\\n    rm -rf /var/lib/apt/lists/*\n"
  },
  {
    "path": "docker/Dockerfile.stable.trtllm",
    "content": "# Base image from NGC TensorRT-LLM, which includes a pre-installed TensorRT-LLM.\n# For available images, visit: https://nvidia.github.io/TensorRT-LLM/installation/containers.html\n# Use TRTLLM_BASE_IMAGE to specify the base image (default: release:1.2.0rc6)\nARG TRTLLM_BASE_IMAGE=nvcr.io/nvidia/tensorrt-llm/release:1.3.0rc4\nFROM ${TRTLLM_BASE_IMAGE}\n\n\n# ==============================================================================\n# Install Megatron dependencies\n# ==============================================================================\n# DeepEP is required for IBGDA support.\n# Clone and build gdrcopy and deepep-nvshmem dependencies.\nWORKDIR /home/dpsk_a2a\nRUN git clone -b v2.5.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    pushd gdrcopy && \\\n    make prefix=/usr/local lib_install && \\\n    popd && rm -rf gdrcopy && \\\n    pip install nvidia-nvshmem-cu13==3.3.20 && \\\n    export NVSHMEM_DIR=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem && \\\n    export LD_LIBRARY_PATH=\"${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH\" && \\\n    export PATH=\"${NVSHMEM_DIR}/bin:$PATH\" && \\\n    pushd ${NVSHMEM_DIR}/lib && \\\n    ln -s libnvshmem_host.so.3 libnvshmem_host.so && \\\n    popd && \\\n    git clone -b v1.2.1 https://github.com/deepseek-ai/DeepEP.git && \\\n    pushd DeepEP && \\\n    wget https://raw.githubusercontent.com/NVIDIA/Megatron-LM/refs/tags/core_v0.15.0/docker/patches/deepep.patch && \\\n    patch -p1 < deepep.patch && \\\n    TORCH_CUDA_ARCH_LIST=\"9.0 10.0 12.0\" python setup.py install && \\\n    popd && rm -rf deepep\n\n# Install Python dependencies\nRUN pip3 install --no-cache-dir --no-deps trl && \\\n    pip3 install --no-cache-dir nvtx matplotlib liger_kernel cachetools && \\\n    pip install --no-cache-dir -U git+https://github.com/ISEEKYAN/mbridge.git && \\\n    pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.15.0\n\n\n# ==============================================================================\n# Install verl dependencies\n# ==============================================================================\nRUN pip install git+https://github.com/volcengine/verl.git@v0.7.0\nRUN pip uninstall -y verl\nRUN pip install \"verl[mcore] @ git+https://github.com/volcengine/verl.git@v0.7.0\"\nRUN pip uninstall -y verl\n\n\n# ==============================================================================\n# Install a specific TensorRT-LLM on demand\n# ==============================================================================\n# Note: The NGC image already includes a pre-installed TensorRT-LLM, but you can install a specific version if needed.\n# Refer to https://nvidia.github.io/TensorRT-LLM/installation/index.html for more details.\n"
  },
  {
    "path": "docker/Dockerfile.stable.vllm",
    "content": "# vllm017\n\nFROM nvidia/cuda:12.9.1-devel-ubuntu22.04\n\nARG DEBIAN_FRONTEND=noninteractive\nARG PIP_NO_CACHE_DIR=1\n\nRUN apt-get update && apt-get install -y \\\n    git \\\n    wget \\\n    cmake \\\n    build-essential \\\n    libibverbs-dev \\\n    libnuma-dev \\\n    librdmacm-dev \\\n    numactl \\\n    software-properties-common \\\n    vim && \\\n    add-apt-repository ppa:deadsnakes/ppa -y && \\\n    apt-get update && \\\n    apt-get install -y \\\n    python3.12 \\\n    python3.12-dev \\\n    && rm -rf /var/lib/apt/lists/*\n\nRUN wget https://bootstrap.pypa.io/get-pip.py && \\\n    python3.12 get-pip.py && \\\n    rm get-pip.py\n\nRUN ln -sf /usr/bin/python3.12 /usr/bin/python3 && \\\n    ln -sf /usr/bin/python3.12 /usr/bin/python\n\nRUN pip install torch==2.10.0 torchvision torchaudio --index-url https://download.pytorch.org/whl/cu129\n\nRUN pip install vllm==0.17.0\n\nRUN pip install pybind11\n\nRUN wget https://developer.download.nvidia.com/compute/cuda/repos/debian12/x86_64/cuda-keyring_1.1-1_all.deb && \\\n    dpkg -i cuda-keyring_1.1-1_all.deb && \\\n    apt-get update && \\\n    apt-get -y install cudnn && \\\n    rm -rf /var/lib/apt/lists/*\n\nRUN pip install nvidia-mathdx\n\nRUN MAX_JOBS=128 pip install -v --disable-pip-version-check --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" git+https://github.com/NVIDIA/apex.git\n\nRUN export NVTE_FRAMEWORK=pytorch && \\\n    MAX_JOBS=128 \\\n    NVTE_BUILD_THREADS_PER_JOB=4 \\\n    pip3 install --resume-retries 999 --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.12\n\n# RUN pip install --upgrade transformers tokenizers\n\nRUN pip install codetiming mathruler pylatexenc qwen_vl_utils cachetools pytest-asyncio\n\nRUN export FLASH_ATTENTION_FORCE_BUILD=\"TRUE\" && MAX_JOBS=16 pip install --no-build-isolation flash_attn==2.8.3\n\nRUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ \"$(uname -m)\" = \"aarch64\" ]; then echo \"arm64\"; else echo \"amd64\"; fi) && \\\n    wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0 && \\\n    apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-${NSIGHT_VERSION}.deb && \\\n    rm -rf /var/lib/apt/lists/*\n\n# =========================\n# Install DeepEP\n# =========================\n# Clone and build deepep and deepep-nvshmem\nWORKDIR /home/dpsk_a2a\nRUN git clone -b v2.5.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    cd gdrcopy && \\\n    make prefix=/usr/local lib_install && \\\n    cd .. && rm -rf gdrcopy\n\nENV GDRCOPY_HOME=/usr/local\n\nRUN git clone -b hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \\\n    export NVSHMEM_DIR=/usr/local/lib/python3.12/dist-packages/nvidia/nvshmem && \\\n    export LD_LIBRARY_PATH=\"${NVSHMEM_DIR}/lib:$LD_LIBRARY_PATH\" && \\\n    export PATH=\"${NVSHMEM_DIR}/bin:$PATH\" && \\\n    cd ${NVSHMEM_DIR}/lib && \\\n    ln -sf libnvshmem_host.so.3 libnvshmem_host.so && \\\n    cd /home/dpsk_a2a/DeepEP && \\\n    export CPATH=/usr/local/cuda/targets/x86_64-linux/include/cccl:$CPATH && \\\n    python setup.py install\n\nRUN pip3 install --no-deps trl==0.27.0\n\nRUN pip3 install nvtx matplotlib liger_kernel\n\nRUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n\nRUN pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.16.0\n\nRUN pip install git+https://github.com/volcengine/verl.git@v0.7.0 && \\\n    pip uninstall -y verl\n\nRUN apt-get update && apt-get install -y curl \\\n    && rm -rf /var/lib/apt/lists/*\n\nRUN apt-get update && \\\n    apt-get install -y --allow-downgrades --allow-change-held-packages \\\n    libcudnn9-cuda-12=9.16.0.29-1 \\\n    libcudnn9-dev-cuda-12=9.16.0.29-1 \\\n    libcudnn9-headers-cuda-12=9.16.0.29-1 && \\\n    rm -rf /var/lib/apt/lists/*\n"
  },
  {
    "path": "docker/README.md",
    "content": "# Dockerfiles of verl\n\nWe provide pre-built Docker images for quick setup. And from this version, we utilize a new image release hierarchy for productivity and stability.\n\nStart from v0.6.0, we use vllm and sglang release image as our base image.\n\nStart from v0.7.0, since vllm/vllm-openai:v0.12.0 is a minimal image without some essential libraries, we use nvidia/cuda:12.9.1-devel-ubuntu22.04 as our base image for vllm.\n\n## Base Image\n\n- vLLM: https://hub.docker.com/r/nvidia/cuda\n- SGLang: https://hub.docker.com/r/lmsysorg/sglang\n\n## Application Image\n\nUpon base image, the following packages are added:\n- flash_attn\n- Megatron-LM\n- Apex\n- TransformerEngine\n- DeepEP\n\nLatest docker file:\n- [Dockerfile.stable.vllm](https://github.com/volcengine/verl/blob/main/docker/Dockerfile.stable.vllm)\n- [Dockerfile.stable.sglang](https://github.com/volcengine/verl/blob/main/docker/Dockerfile.stable.sglang)\n\nAll pre-built images are available in dockerhub: https://hub.docker.com/r/verlai/verl. For example, `verlai/verl:sgl059.latest`, `verlai/verl:vllm017.latest`.\n\nYou can find the latest images used for development and ci in our github workflows:\n- [.github/workflows/vllm.yml](https://github.com/volcengine/verl/blob/main/.github/workflows/vllm.yml)\n- [.github/workflows/sgl.yml](https://github.com/volcengine/verl/blob/main/.github/workflows/sgl.yml)\n\n\n## Installation from Docker\n\nAfter pulling the desired Docker image and installing desired inference and training frameworks, you can run it with the following steps:\n\n1. Launch the desired Docker image and attach into it:\n\n```sh\ndocker create --runtime=nvidia --gpus all --net=host --shm-size=\"10g\" --cap-add=SYS_ADMIN -v .:/workspace/verl --name verl <image:tag> sleep infinity\ndocker start verl\ndocker exec -it verl bash\n```\n\n2. If you use the images provided, you only need to install verl itself without dependencies:\n\n```sh\n# install the nightly version (recommended)\ngit clone https://github.com/volcengine/verl && cd verl\npip3 install --no-deps -e .\n```\n\n[Optional] If you hope to switch between different frameworks, you can install verl with the following command:\n\n```sh\n# install the nightly version (recommended)\ngit clone https://github.com/volcengine/verl && cd verl\npip3 install -e .[vllm]\npip3 install -e .[sglang]\n```\n\n## Release History\n\n- 2026/03/10: update vllm stable image to vllm==0.17.0; update sglang stable image to sglang==0.5.9\n- 2026/01/17: update vllm stable image to torch==2.9.1, cudnn==9.16, deepep==1.2.1\n- 2025/12/23: update vllm stable image to vllm==0.12.0; update sglang stable image to sglang==0.5.6\n- 2025/11/18: update vllm stable image to vllm==0.11.1; update sglang stable image to sglang==0.5.5\n\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a2",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11\n\nARG ASCEND_CANN_PATH=\"/usr/local/Ascend\"\nARG PIP_INDEX_URL=\"https://mirrors.aliyun.com/pypi/simple\"\nARG PTA_BASE_VERSION=\"torch_npu-2.7.1.post2-cp311-cp311-manylinux_2_28\"\nARG PTA_URL=\"https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.7.1\"\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip config set global.index-url ${PIP_INDEX_URL} && \\\n    pip config set install.trusted-host mirrors.aliyun.com && \\\n    pip install --upgrade pip setuptools packaging && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    # Set extra pip index for x86_64 platform\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.5.8 https://github.com/sgl-project/sglang.git && \\\n    git clone https://github.com/sgl-project/sgl-kernel-npu.git && cd sgl-kernel-npu && git checkout 46b73de && cd .. && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. \n\n# Install repositories with low update frequency\nRUN cd sglang && \\\n    # Install sglang\n    mv python/pyproject.toml python/pyproject.toml.backup && \\\n    mv python/pyproject_other.toml python/pyproject.toml && \\\n    pip install -e \"python[srt_npu]\" && \\\n    pip install torch==2.7.1 torchvision==0.22.1 && \\\n    # Install torch_npu\n    ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \\\n    echo \"[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl\" && \\\n    cd .. \n\n# Install sgl-kernel-npu\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/ascend-toolkit/8.3.RC1/${ARCH}-linux/devlib/linux/${ARCH}:$LD_LIBRARY_PATH && \\\n    source ${ASCEND_CANN_PATH}/ascend-toolkit/set_env.sh && \\\n    source ${ASCEND_CANN_PATH}/nnal/atb/set_env.sh && \\\n    pip install pybind11 && \\\n    cd sgl-kernel-npu && \\\n    bash build.sh  && \\\n    pip install output/torch_memory_saver*.whl && \\\n    pip install output/sgl_kernel_npu*.whl && \\\n    # Deep_ep package is compiled for A3 by default; Recompile in deepep2 mode for A2, following https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md.\n    bash build.sh -a deepep2 && \\\n    pip install output/deep_ep*.whl && \\\n    cd \"$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')\" && ln -s deep_ep/deep_ep_cpp*.so && cd - && \\\n    cd ..\n\n# Install MindSpeed & Megatron\nRUN pip install -e MindSpeed && \\\n    pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \\\n    # Remove existing triton or triton-ascend installed by some third-party packages\n    pip uninstall -y triton timm && \\\n    # Install mbridge\n    pip install mbridge && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Prepare and install verl (update frequently)\nRUN git clone --recursive https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    pip install ray==2.46.0 click==8.2.1 cachetools && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a3",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-ubuntu22.04-py3.11\n\nARG ASCEND_CANN_PATH=\"/usr/local/Ascend\"\nARG PIP_INDEX_URL=\"https://mirrors.aliyun.com/pypi/simple\"\nARG PTA_BASE_VERSION=\"torch_npu-2.7.1.post2-cp311-cp311-manylinux_2_28\"\nARG PTA_URL=\"https://gitcode.com/Ascend/pytorch/releases/download/v7.3.0-pytorch2.7.1\"\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential net-tools iputils-ping && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip config set global.index-url ${PIP_INDEX_URL} && \\\n    pip config set install.trusted-host mirrors.aliyun.com && \\\n    pip install --upgrade pip setuptools packaging && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    # Set extra pip index for x86_64 platform\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.5.8 https://github.com/sgl-project/sglang.git && \\\n    git clone https://github.com/sgl-project/sgl-kernel-npu.git && cd sgl-kernel-npu && git checkout 46b73de && cd .. && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. \n\n# Install repositories with low update frequency\nRUN cd sglang && \\\n    # Install sglang\n    mv python/pyproject.toml python/pyproject.toml.backup && \\\n    mv python/pyproject_other.toml python/pyproject.toml && \\\n    pip install -e \"python[srt_npu]\" && \\\n    pip install torch==2.7.1 torchvision==0.22.1 && \\\n    # Install torch_npu\n    ARCH=$(uname -m) && wget ${PTA_URL}/${PTA_BASE_VERSION}_${ARCH}.whl && pip install ${PTA_BASE_VERSION}_${ARCH}.whl && \\\n    echo \"[LOG INFO] Torch_npu version is: ${PTA_BASE_VERSION}_${ARCH}.whl\" && \\\n    cd .. \n\n# Install sgl-kernel-npu\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    export LD_LIBRARY_PATH=${ASCEND_CANN_PATH}/ascend-toolkit/8.3.RC1/${ARCH}-linux/devlib/linux/${ARCH}:$LD_LIBRARY_PATH && \\\n    source ${ASCEND_CANN_PATH}/ascend-toolkit/set_env.sh && \\\n    source ${ASCEND_CANN_PATH}/nnal/atb/set_env.sh && \\\n    pip install pybind11 && \\\n    cd sgl-kernel-npu && \\\n    bash build.sh  && \\\n    pip install output/torch_memory_saver*.whl && \\\n    pip install output/sgl_kernel_npu*.whl && \\\n    pip install output/deep_ep*.whl && \\\n    cd \"$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')\" && ln -s deep_ep/deep_ep_cpp*.so && cd - && \\\n    cd ..\n\n# Install MindSpeed & Megatron\nRUN pip install -e MindSpeed && \\\n    pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 && \\\n    # Remove existing triton or triton-ascend installed by some third-party packages\n    pip uninstall -y triton timm && \\\n    # Install mbridge\n    pip install mbridge && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Prepare and install verl (update frequently)\nRUN git clone --recursive https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    pip install ray==2.46.0 click==8.2.1 cachetools && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.2.rc1_a2",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.2.rc1-910b-ubuntu22.04-py3.11\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip setuptools packaging && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm && \\\n    git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install torch & torch_npu & torchvision\n    pip install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 && \\\n    # Install vllm\n    cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\nENV PYTHONPATH=\"/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}\"\n\n# Prepare and install verl (update frequently)\nRUN git clone --depth 1 https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.2.rc1_a3",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.2.rc1-a3-ubuntu22.04-py3.11\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip setuptools packaging && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm && \\\n    git clone --depth 1 --branch v0.9.1 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.2.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install torch & torch_npu & torchvision\n    pip install torch==2.5.1 torch_npu==2.5.1 torchvision==0.20.1 && \\\n    # Install vllm\n    cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\nENV PYTHONPATH=\"/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}\"\n\n# Prepare and install verl (update frequently)\nRUN git clone --depth 1 https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.3.rc1_a2",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-910b-ubuntu22.04-py3.11\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip packaging setuptools==80.10.2 && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm.git && \\\n    git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install torch & torch_npu & torchvision\n    pip install torch==2.7.1 torch_npu==2.7.1 torchvision==0.22.1 transformers==4.57.6 && \\\n    # Install vllm\n    cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    # Remove existing triton or triton-ascend installed by some third-party packages\n    pip uninstall -y triton triton-ascend && \\\n    # Install mbridge\n    pip install mbridge && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\nENV PYTHONPATH=\"/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}\"\n\n# Prepare and install verl (update frequently)\nRUN git clone --depth 1 https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.3.rc1_a3",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.3.rc1-a3-ubuntu22.04-py3.11\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip packaging setuptools==80.10.2 && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm.git && \\\n    git clone --depth 1 --branch v0.11.0 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout f2b0977e && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/ascend-toolkit/8.3.RC1/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install torch & torch_npu & torchvision\n    pip install torch==2.7.1 torch_npu==2.7.1 torchvision==0.22.1 transformers==4.57.6 && \\\n    # Install vllm\n    cd vllm && VLLM_TARGET_DEVICE=empty pip install -v -e . && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    # Remove existing triton or triton-ascend installed by some third-party packages\n    pip uninstall -y triton triton-ascend && \\\n    # Install mbridge\n    pip install mbridge && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\nENV PYTHONPATH=\"/Megatron-LM${PYTHONPATH:+:${PYTHONPATH}}\"\n\n# Prepare and install verl (update frequently)\nRUN git clone --depth 1 https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.5.0_a2",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-910b-ubuntu22.04-py3.11\n\nARG SOC_VERSION=\"ascend910b1\"\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip packaging setuptools==80.10.2 && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.13.0 https://github.com/vllm-project/vllm.git && \\\n    git clone -b releases/v0.13.0 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install transformers\n    pip install transformers==4.57.6 && \\\n    # Install vllm\n    cd vllm && pip install -r requirements/build.txt && \\\n    VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -r requirements.txt && \\\n    export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    pip install -e Megatron-LM && \\\n    # Remove existing triton installed by some third-party packages\n    pip uninstall -y triton && \\\n    # Install mbridge\n    pip install mbridge torch_npu==2.8.0 && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n\n# Prepare and install verl (update frequently)\nRUN git clone --recursive https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]"
  },
  {
    "path": "docker/ascend/Dockerfile.ascend_8.5.0_a3",
    "content": "# Pull base image\nFROM swr.cn-south-1.myhuaweicloud.com/ascendhub/cann:8.5.0-a3-ubuntu22.04-py3.11\n\nARG SOC_VERSION=\"ascend910_9392\"\n\n# Prepare required system dependencies\nRUN apt-get update -y && \\\n    apt-get install -y --no-install-recommends gcc g++ cmake libnuma-dev wget git curl jq vim build-essential && \\\n    apt-get clean && \\\n    rm -rf /var/lib/apt/lists/* /tmp/* /var/tmp/* && \\\n    pip install --upgrade pip packaging setuptools==80.10.2 && \\\n    pip cache purge\n\n# Prepare repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Set extra pip index for x86_64 platform\n    echo \"[LOG INFO] Detected architecture: $ARCH\" && \\\n    if [ \"$ARCH\" = \"x86_64\" ]; then \\\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"; \\\n    fi && \\\n    # Clone libs\n    git clone --depth 1 --branch v0.13.0 https://github.com/vllm-project/vllm.git && \\\n    git clone -b releases/v0.13.0 https://github.com/vllm-project/vllm-ascend.git && \\\n    git clone https://gitcode.com/Ascend/MindSpeed.git && \\\n    cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd .. && \\\n    git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n\n# Install repositories with low update frequency\nRUN ARCH=$(uname -m) && \\\n    # Export and source env\n    if [ \"$ARCH\" = \"aarch64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/aarch64-linux/devlib/linux/aarch64:$LD_LIBRARY_PATH; \\\n    elif [ \"$ARCH\" = \"x86_64\" ]; then \\\n        export LD_LIBRARY_PATH=/usr/local/Ascend/cann-8.5.0/x86_64-linux/devlib/linux/x86_64/:$LD_LIBRARY_PATH; \\\n    fi && \\\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh && \\\n    source /usr/local/Ascend/nnal/atb/set_env.sh && \\\n    # Install transformers\n    pip install transformers==4.57.6 && \\\n    # Install vllm\n    cd vllm && pip install -r requirements/build.txt && \\\n    VLLM_TARGET_DEVICE=empty pip install -v -e. && cd .. && \\\n    # Install vllm-ascend\n    cd vllm-ascend && pip install -r requirements.txt && \\\n    export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd .. && \\\n    # Install MindSpeed & Megatron\n    pip install -e MindSpeed && \\\n    pip install -e Megatron-LM && \\\n    # Remove existing triton installed by some third-party packages\n    pip uninstall -y triton && \\\n    # Install mbridge\n    pip install mbridge torch_npu==2.8.0 && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n\n# Prepare and install verl (update frequently)\nRUN git clone --recursive https://github.com/volcengine/verl.git && \\\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd .. && \\\n    # Clear extra files\n    rm -rf /tmp/* /var/tmp/* && \\\n    pip cache purge\n\n# Show install results\nRUN pip list\n\n# Setting Default Commands\nCMD [\"/bin/bash\"]\n"
  },
  {
    "path": "docker/aws/Dockerfile.extention.awsefa",
    "content": "# Base Image support aws EFA\n# Build Image with frameworks based on this\nFROM verlai/verl:app-verl0.6-transformers4.56.1-sglang0.5.2-mcore0.13.0-te2.2\n\n# For aws instances with EFA net interface (Sagemaker AI Pod)\n#     install EFA driver:\n######## AWS EFA ############\nENV NCCL_VERSION=2.25.1-1\nENV DEBIAN_FRONTEND=noninteractive\nENV EFA_INSTALLER_VERSION=1.40.0\nENV AWS_OFI_NCCL_VERSION=1.14.2\nENV FI_EFA_SET_CUDA_SYNC_MEMOPS=0\nENV FI_PROVIDER=efa\n\nRUN apt update && apt install -y linux-image-generic libhwloc-dev\n\nRUN cd /tmp && \\\n    curl -O https://efa-installer.amazonaws.com/aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz  && \\\n    tar -xf aws-efa-installer-${EFA_INSTALLER_VERSION}.tar.gz && \\\n    cd aws-efa-installer && \\\n    ./efa_installer.sh -y -g --skip-kmod --skip-limit-conf --no-verify && \\\n    ldconfig && \\\n    rm -rf /tmp/aws-efa-installer /var/lib/apt/lists/*\n\n# NCCL EFA Plugin\nRUN cd /tmp && \\\n    curl -LO https://github.com/aws/aws-ofi-nccl/archive/refs/tags/v${AWS_OFI_NCCL_VERSION}.tar.gz && \\\n    tar -xzf /tmp/v${AWS_OFI_NCCL_VERSION}.tar.gz && \\\n    rm /tmp/v${AWS_OFI_NCCL_VERSION}.tar.gz && \\\n    mv aws-ofi-nccl-${AWS_OFI_NCCL_VERSION} aws-ofi-nccl && \\\n    cd /tmp/aws-ofi-nccl && \\\n    ./autogen.sh && \\\n    ./configure --prefix=/opt/amazon/efa \\\n    --with-libfabric=/opt/amazon/efa \\\n    --with-cuda=/usr/local/cuda \\\n    --enable-platform-aws \\\n    --with-mpi=/opt/amazon/openmpi && \\\n    make -j$(nproc) install && \\\n    rm -rf /tmp/aws-ofi/nccl\n\n# NCCL\nRUN echo \"/usr/local/lib\"      >> /etc/ld.so.conf.d/local.conf && \\\n    echo \"/opt/amazon/openmpi/lib\" >> /etc/ld.so.conf.d/efa.conf && \\\n    ldconfig\n\nENV OMPI_MCA_pml=^cm,ucx            \\\n    OMPI_MCA_btl=tcp,self           \\\n    OMPI_MCA_btl_tcp_if_exclude=lo,docker0,veth_def_agent \\\n    OPAL_PREFIX=/opt/amazon/openmpi \\\n    NCCL_SOCKET_IFNAME=^docker,lo,veth_def_agent  \\\n    FI_EFA_USE_HUGE_PAGE=0\n\n# docker build -t verl:awsefa --label \"commit=$(git rev-parse --short HEAD)\" .\n# on aws:\n# docker run --ipc=host --privileged --name verldev --gpus all --network=host --shm-size=1800gb -itd verl:awsefa\n"
  },
  {
    "path": "docker/aws/Dockerfile.ngc.vllm0.8.sagemaker",
    "content": "# Using a pre-built image from AWS DLC which contains the current version of python (3.10) and supported cuda version (12.1)\nFROM 763104351884.dkr.ecr.us-east-1.amazonaws.com/huggingface-pytorch-training:2.1.0-transformers4.36.0-gpu-py310-cu121-ubuntu20.04\n\n# uninstall nv-pytorch fork\nRUN pip3 uninstall -y pytorch-quantization \\\n    pytorch-triton torch torch-tensorrt torchvision \\\n    xgboost transformer_engine flash_attn apex megatron-core\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install tini\nRUN apt-get update && \\\n    apt-get install -y tini && \\\n    apt-get clean\n\n# Install torch-2.6.0 + vllm-0.8.2\nRUN pip install --no-cache-dir vllm==0.8.2 torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0 tensordict torchdata==0.11.0 \\\n    transformers>=4.49.0 accelerate datasets peft hf-transfer \\\n    ray[default] codetiming hydra-core pandas pyarrow>=15.0.0 pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \\\n    pytest pre-commit py-spy pyext ruff tensorboard\n\n# Install flash_attn-2.7.4.post1\nRUN pip uninstall -y transformer-engine flash-attn && \\\n    pip install flash-attn==2.7.4.post1 --no-build-isolation\n\n# Fix cv2\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --no-cache-dir nvidia-ml-py>=12.560.30 opencv-python-headless==4.8.0.74 fastapi==0.115.6 && \\\n    pip install --no-cache-dir --upgrade optree>=0.13.0\n\n# Install verl\nRUN pip install --no-cache-dir verl[vllm] -U\n\n# Reset pip config\nRUN pip config unset global.index-url && \\\n    pip config unset global.extra-index-url\n"
  },
  {
    "path": "docker/rocm/Apptainerfile.rocm",
    "content": "Bootstrap: docker\n\n# Support - Traing: fsdp; Inference: vllm\n# FROM: rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n# Support - Traing: fsdp; Inference: vllm, sglang\nFROM lmsysorg/sglang:v0.4.5-rocm630\n\n%environment\n    export PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n\n    export HIPCC_COMPILE_FLAGS_APPEND=\"--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__\"\n    export CFLAGS=\"-D__HIP_PLATFORM_AMD__\"\n    export CXXFLAGS=\"-D__HIP_PLATFORM_AMD__\"\n\n%post\n    # Create source directory\n    mkdir -p /opt/src\n\n    # Uninstall and reinstall vllm\n    pip uninstall -y vllm\n    cd /opt/src\n    git clone -b v0.6.3 https://github.com/vllm-project/vllm.git\n    cd vllm\n    MAX_JOBS=$(nproc) python3 setup.py install\n    cd /opt\n    rm -rf /opt/src/vllm\n\n    # Install dependencies\n    pip install \"tensordict<0.6\" --no-deps\n    pip install accelerate \\\n        codetiming \\\n        datasets \\\n        dill \\\n        hydra-core \\\n        liger-kernel \\\n        numpy \\\n        pandas \\\n        peft \\\n        \"pyarrow>=15.0.0\" \\\n        pylatexenc \\\n        \"ray[data,train,tune,serve]\" \\\n        torchdata \\\n        transformers \\\n        wandb \\\n        orjson \\\n        pybind11\n\n    # Clone and install verl from GitHub\n    cd /opt\n    git clone https://github.com/volcengine/verl.git\n    cd verl\n    # Uncomment to use a specific version\n    # git checkout v0.3.0.post0\n    pip install -e . --no-deps\n\n    # Install torch_memory_saver\n    pip install git+https://github.com/ExtremeViscent/torch_memory_saver.git --no-deps"
  },
  {
    "path": "docker/rocm/Dockerfile.rocm",
    "content": "# FROM \"compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compute-rocm-rel-6.4:94_ubuntu22.04_py3.10_pytorch_release-2.7_575e247\"\n# FROM \"rlfoundation.azurecr.io/rocm6.3.4:vllm-0.8.5-numa-patch-ubuntu-22.04\"\nFROM \"rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04\"\n\nSHELL [\"/bin/bash\", \"-ceuxo\", \"pipefail\"]\n\nENV MAX_JOBS=512\n\nENV PATH=\"/usr/local/python3.12/bin:$PATH\"\nRUN ln -sf /usr/bin/python3.12 /usr/bin/python && \\\n    ln -sf /usr/bin/pip3.12 /usr/bin/pip\n\n############################################\n############################################\nRUN apt-get update\nRUN apt-get install -y pkg-config liblzma-dev\n############################################\n############################################\n\n\n###########################################\n##########Install TransformerEngine########\n###########################################\nWORKDIR /workspace/\n# transformer-engine install\n# https://github.com/ROCm/TransformerEngine\n\nRUN rm -rf TransformerEngine \nRUN git clone --recursive https://github.com/ROCm/TransformerEngine.git\nWORKDIR /workspace/TransformerEngine\nRUN git checkout 236178e5\n# git checkout bb061ade\n# git checkout 864405c\n\nENV NVTE_FRAMEWORK=pytorch \nENV NVTE_ROCM_ARCH=gfx942 \nENV NVTE_USE_HIPBLASLT=1\nENV NVTE_USE_ROCM=1  \n\n# export CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}\"\nENV CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr\"\n\n\n# ENV NVTE_BUILD_MAX_JOBS=$(MAX_JOBS)\n\nRUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv \n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n\n####################################################################################\n################Install vllm - sglang require vllm 0.6.7 dependency#################\n####################################################################################\n#### Require vllm 0.6.7 - checkout 113274a0\nWORKDIR /workspace/\nRUN rm -rf vllm\nRUN pip uninstall -y vllm\n# Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html\nRUN git clone https://github.com/ROCm/vllm.git\n# git clone https://github.com/vllm-project/vllm.git\nWORKDIR /workspace/vllm\nRUN git checkout 113274a0\nENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n#ENV MAX_JOBS=512\nENV MAX_JOBS=${MAX_JOBS}\nRUN pip install \"boto3>=1.26.0\"\nRUN pip install setuptools_scm\n# will add src into py. You can delete the repo\nRUN python3 setup.py install\nWORKDIR /workspace/\n####################################################################################\n####################################################################################\n####################################################################################\n\n\n\n###########################################\n############For hack docker################\n###########################################\nRUN pip install setuptools==75.8.0\n###########################################\n###########################################\n###########################################\n\n\n\n###########################################\n############build sgalng###################\n###########################################\n# Set environment variables\nENV BASE_DIR=/sgl-workspace\nENV BUILD_TYPE=all\nENV SGL_REPO=https://github.com/sgl-project/sglang\nENV SGL_BRANCH=v0.4.6.post5\nENV TRITON_REPO=https://github.com/ROCm/triton.git\nENV TRITON_COMMIT=improve_fa_decode_3.0.0\nENV AITER_REPO=https://github.com/ROCm/aiter.git\nENV AITER_COMMIT=v0.1.2\n# v0.1.2 version - commit id: 9d11f47\n# ENV AITER_COMMIT=9d11f47\n\nENV HIP_FORCE_DEV_KERNARG=1\nENV HSA_NO_SCRATCH_RECLAIM=1\nENV SGLANG_SET_CPU_AFFINITY=1\nENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1\nENV NCCL_MIN_NCHANNELS=112\nENV MOE_PADDING=1\nENV VLLM_FP8_PADDING=1\nENV VLLM_FP8_ACT_PADDING=1\nENV VLLM_FP8_WEIGHT_PADDING=1\nENV VLLM_FP8_REDUCE_CONV=1\nENV TORCHINDUCTOR_MAX_AUTOTUNE=1\nENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1\nENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\nENV AMDGPU_TARGETS=gfx942\nENV ROCM_ARCH=gfx942\nENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n\n# Switch to working directory\nWORKDIR /sgl-workspace\n\n# Clean and create directory\nRUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace\n\n# Clone and build sglang\nRUN git clone ${SGL_REPO} \\\n    && cd sglang \\\n    && git checkout ${SGL_BRANCH} || echo \"Using default branch\" \\\n    && cd sgl-kernel \\\n    && rm -f pyproject.toml \\\n    && mv pyproject_rocm.toml pyproject.toml \\\n    && python setup_rocm.py install \\\n    && cd .. \\\n    && if [ \"$BUILD_TYPE\" = \"srt\" ]; then \\\n         python -m pip --no-cache-dir install -e \"python[srt_hip]\"; \\\n       else \\\n         python -m pip --no-cache-dir install -e \"python[all_hip]\"; \\\n       fi \\\n    && cd /sgl-workspace \\\n    && cp -r /sgl-workspace/sglang /sglang \\\n    && python -m pip cache purge\n\n# Install common Python packages\nRUN pip install IPython orjson python-multipart torchao pybind11\n\n# Rebuild Triton\nRUN pip uninstall -y triton || true \\\n    && git clone ${TRITON_REPO} \\\n    && cd triton \\\n    && git checkout ${TRITON_COMMIT} \\\n    && cd python \\\n    && python3 setup.py install \\\n    && cd /sgl-workspace\n\n\n# ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1\"\n# ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\n\n# Build aiter\n#version: Commit 9d11f47\n    # && git checkout ${AITER_COMMIT} \\\nRUN pip uninstall -y aiter || true\nRUN git clone ${AITER_REPO} \\\n    && cd aiter \\\n    && git checkout ${AITER_COMMIT} \\\n    && git submodule sync \\\n    && git submodule update --init --recursive \\\n    && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \\\n    && cd /sgl-workspace\n    # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \\\n    # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \\\n\n# Copy MI300X config \nRUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \\\n         /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \\\n         -type f -name '*MI300X*' | \\\n         xargs -I {} sh -c 'vf_config=$(echo \"$1\" | sed \"s/MI300X/MI300X_VF/\"); cp \"$1\" \"$vf_config\"' -- {}\n\n# Environment setup complete.\nRUN echo \"Environment setup complete.\"\n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n\n\n###########################################\n###############vllm v0.8.5#################\n###########################################\n# ENV GITHUB_USERNAME=yushengsu-thu\n# ENV GITHUB_MAIL=yushengsu@gmail.com\n\n# RUN git config --global user.name \"${GITHUB_USERNAME}\" \\\n#     && git config --global user.email \"${GITHUB_MAIL}\" \n\nWORKDIR /workspace/\n\nENV VLLM_TARGET_DEVICE=rocm \nENV ROCM_PATH=/opt/rocm \nENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev\n\n# Find the repo path in: DockerFile/Dockerfile.rocm_yang\n# RUN git clone https://github.com/RLFoundation/vllm-patch.git\nRUN pip uninstall -y vllm || true\nRUN rm -rf vllm-patch\nRUN git clone https://github.com/RLFoundation/vllm-patch.git \\\n    && cd vllm-patch \\\n    && git checkout v0.8.5-sleep-numa \\\n    && rm -rf build/ dist/ *.egg-info \\\n    && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \\\n    && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py install\n    # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py develop\n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n#########################################\n#### Install megatron-core###############\n#########################################\nRUN pip uninstall -y megatron-core && \\\n    git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \\\n    cd Megatron-LM-amd_version && \\\n    pip install -vvv -e . && \\\n    cd /workspace/\n#########################################\n#########################################\n#########################################\n\n\n\n\n#######################################\n################apex###################\n#######################################\nWORKDIR /workspace/\nRUN pip uninstall -y apex && \\\n    git clone https://github.com/ROCm/apex.git && \\\n    cd apex && \\\n    python setup.py install && \\\n    cd /workspace/ \n#######################################\n#######################################\n#######################################\n\n\n\n\n################################################################################\n###########################Add torch_memory_saver###############################\n################################################################################\n# Set environment variables\nENV HIPCC_COMPILE_FLAGS_APPEND=\"--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__\"\nENV CFLAGS=\"-D__HIP_PLATFORM_AMD__\"\nENV CXXFLAGS=\"-D__HIP_PLATFORM_AMD__\"\nRUN pip install \"git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa\"\n################################################################################\n################################################################################\n################################################################################\n\n\n\n########################################\n######Install ray#######################\n########################################\n# need to add this patch: https://github.com/ray-project/ray/pull/53531/files\nRUN pip uninstall ray -y\nRUN pip install \"ray[data,train,tune,serve]>=2.47.0\" \n########################################\n########################################\n########################################\n\n\n\n##########################################\n#######Install other dependencies#########\n##########################################\nRUN pip install \"tensordict==0.6.2\" --no-deps && \\\n    pip install accelerate \\\n    codetiming \\\n    datasets \\\n    dill \\\n    hydra-core \\\n    liger-kernel \\\n    numpy \\\n    pandas \\\n    peft \\\n    \"pyarrow>=15.0.0\" \\\n    pylatexenc \\\n    torchdata \\\n    wandb \\\n    orjson \\\n    pybind11\n    \nWORKDIR /workspace/\nRUN git clone https://github.com/volcengine/verl.git && \\\n    cd verl && \\\n    pip install -e . \n##########################################\n##########################################\n##########################################\n\n\n\nWORKDIR /workspace/\n\nCMD [\"/usr/bin/bash\"]\n"
  },
  {
    "path": "docker/rocm/Dockerfile.rocm7",
    "content": "# default base image\nARG REMOTE_VLLM=\"1\"\nARG COMMON_WORKDIR=/app\nARG BASE_IMAGE=rocm/vllm-dev:base\n\nFROM ${BASE_IMAGE} AS base\n\nARG ARG_PYTORCH_ROCM_ARCH\nENV PYTORCH_ROCM_ARCH=${ARG_PYTORCH_ROCM_ARCH:-${PYTORCH_ROCM_ARCH}}\n\n# Install some basic utilities\nRUN apt-get update -q -y && apt-get install -q -y \\\n    sqlite3 libsqlite3-dev libfmt-dev libmsgpack-dev libsuitesparse-dev \\\n    apt-transport-https ca-certificates wget curl\n# Remove sccache\nRUN python3 -m pip install --upgrade pip\nRUN apt-get purge -y sccache; python3 -m pip uninstall -y sccache; rm -f \"$(which sccache)\"\nARG COMMON_WORKDIR\nWORKDIR ${COMMON_WORKDIR}\n\n\n# -----------------------\n# vLLM fetch stages\nFROM base AS fetch_vllm_0\nONBUILD COPY ./ vllm/\nFROM base AS fetch_vllm_1\n#ARG VLLM_REPO=\"https://github.com/ROCm/vllm.git\"\n#ARG VLLM_BRANCH=\"main\"\nARG VLLM_REPO=https://github.com/HollowMan6/vllm.git\nARG VLLM_BRANCH=\"sleep_amd\"\nONBUILD RUN git clone ${VLLM_REPO} \\\n            && cd vllm \\\n            && git checkout ${VLLM_BRANCH}\nFROM fetch_vllm_${REMOTE_VLLM} AS fetch_vllm\n\n# -----------------------\n# vLLM build stages\nFROM fetch_vllm AS build_vllm\n# Build vLLM\nRUN cd vllm \\\n    && python3 -m pip install -r requirements/rocm.txt \\\n    && python3 setup.py clean --all  \\\n    && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \\\n    && VLLM_TARGET_DEVICE=rocm ROCM_PATH=/opt/rocm/ VLLM_GPU_LANG=HIP SETUPTOOLS_SCM_PRETEND_VERSION=0.11.0.dev python3 setup.py bdist_wheel --dist-dir=dist\n    #&& python3 setup.py bdist_wheel --dist-dir=dist\nFROM scratch AS export_vllm\nARG COMMON_WORKDIR\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/dist/*.whl /\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/requirements /requirements\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/benchmarks /benchmarks\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/tests /tests\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/examples /examples\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm/.buildkite /.buildkite\n\n# -----------------------\n# Test vLLM image\nFROM base AS test\n\nRUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*\n\n# Install vLLM\n#RUN --mount=type=bind,from=export_vllm,src=/,target=/install \\\nCOPY --from=export_vllm /*.whl /install\nCOPY --from=export_vllm /requirements /install/requirements\nCOPY --from=export_vllm /benchmarks /install/benchmarks\nCOPY --from=export_vllm /tests /install/tests\nCOPY --from=export_vllm /examples /install/examples\nCOPY --from=export_vllm /.buildkite /install/.buildkite\n\nRUN cd /install \\\n    && pip install -U -r requirements/rocm.txt \\\n    && pip install -U -r requirements/rocm-test.txt \\\n    && pip uninstall -y vllm \\\n    && pip install *.whl\n\nWORKDIR /vllm-workspace\nARG COMMON_WORKDIR\nCOPY --from=build_vllm ${COMMON_WORKDIR}/vllm /vllm-workspace\n\n# install development dependencies (for testing)\nRUN cd /vllm-workspace \\\n    && rm -rf vllm \\\n    && python3 -m pip install -e tests/vllm_test_utils \\\n    && python3 -m pip install lm-eval[api]==0.4.4 \\\n    && python3 -m pip install pytest-shard\n\n# -----------------------\n# Final vLLM image\nFROM base AS final\n\nRUN python3 -m pip install --upgrade pip && rm -rf /var/lib/apt/lists/*\n# Error related to odd state for numpy 1.20.3 where there is no METADATA etc, but an extra LICENSES_bundled.txt.\n# Manually remove it so that later steps of numpy upgrade can continue\nRUN case \"$(which python3)\" in \\\n        *\"/opt/conda/envs/py_3.9\"*) \\\n            rm -rf /opt/conda/envs/py_3.9/lib/python3.9/site-packages/numpy-1.20.3.dist-info/;; \\\n        *) ;; esac\n\nRUN python3 -m pip install --upgrade huggingface-hub[cli]\n\n# Install vLLM\nRUN --mount=type=bind,from=export_vllm,src=/,target=/install \\\n    cd /install \\\n    && pip install -U -r requirements/rocm.txt \\\n    && pip uninstall -y vllm \\\n    && pip install *.whl\n\nARG COMMON_WORKDIR\n\n# Copy over the benchmark scripts as well\nCOPY --from=export_vllm /benchmarks ${COMMON_WORKDIR}/vllm/benchmarks\nCOPY --from=export_vllm /examples ${COMMON_WORKDIR}/vllm/examples\n\nENV RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1\nENV TOKENIZERS_PARALLELISM=false\n\n# ENV that can improve safe tensor loading, and end-to-end time\nENV SAFETENSORS_FAST_GPU=1\n\n# Performance environment variable.\nENV HIP_FORCE_DEV_KERNARG=1\n\n# -----------------------\n# Install verl\nARG VERL_REPO=https://github.com/volcengine/verl.git\nARG VERL_BRANCH=main\nRUN pip install \"tensordict==0.6.2\" --no-deps && \\\n    pip install accelerate \\\n    codetiming \\\n    datasets \\\n    dill \\\n    hydra-core \\\n    liger-kernel \\\n    numpy \\\n    pandas \\\n    peft \\\n    \"pyarrow>=15.0.0\" \\\n    pylatexenc \\\n    torchdata \\\n    wandb \\\n    orjson \\\n    pybind11\n\nWORKDIR /workspace/\nRUN git clone ${VERL_REPO} && \\\n    cd verl && \\\n    git checkout ${VERL_BRANCH} && \\\n    pip install -e .\n\nCMD [\"/bin/bash\"]\n\n"
  },
  {
    "path": "docker/rocm/Dockerfile.rocm_verl-0.3.0.post1",
    "content": "#  Build the docker in the repo dir:\n# docker build -f docker/Dockerfile.rocm -t verl-rocm:03.04.2015 .\n# docker images # you can find your built docker\n\n\n# Support - Traing: fsdp; Inference: vllm\n# FROM rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n# Support - Traing: fsdp; Inference: vllm, sglang\nFROM lmsysorg/sglang:v0.4.6.post5-rocm630\n\n# Set working directory\n# WORKDIR $PWD/app\n\n# Set environment variables\nENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n\nENV HIPCC_COMPILE_FLAGS_APPEND=\"--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__\"\nENV CFLAGS=\"-D__HIP_PLATFORM_AMD__\"\nENV CXXFLAGS=\"-D__HIP_PLATFORM_AMD__\"\n\n# Install vllm\nRUN pip uninstall -y vllm && \\\n    rm -rf vllm && \\\n    git clone -b v0.6.3 https://github.com/vllm-project/vllm.git && \\\n    cd vllm && \\\n    MAX_JOBS=$(nproc) python3 setup.py install && \\\n    cd .. && \\\n    rm -rf vllm\n\n# Copy the entire project directory\nCOPY . .\n\n# Install dependencies\nRUN pip install \"tensordict==0.6.2\" --no-deps && \\\n    pip install accelerate \\\n    codetiming \\\n    datasets \\\n    dill \\\n    hydra-core \\\n    liger-kernel \\\n    numpy \\\n    pandas \\\n    peft \\\n    \"pyarrow>=15.0.0\" \\\n    pylatexenc \\\n    \"ray[data,train,tune,serve]<2.45.0\" \\\n    torchdata \\\n    transformers \\\n    wandb \\\n    orjson \\\n    pybind11\n    \nRUN git clone https://github.com/volcengine/verl.git && \\\n    cd verl && \\\n    pip install -e . \n\n# Install torch_memory_saver\nRUN pip install git+https://github.com/ExtremeViscent/torch_memory_saver.git --no-deps\n"
  },
  {
    "path": "docker/rocm/Dockerfile.rocm_verl-0.4.1",
    "content": "# FROM \"compute-artifactory.amd.com:5000/rocm-plus-docker/framework/compute-rocm-rel-6.4:94_ubuntu22.04_py3.10_pytorch_release-2.7_575e247\"\n# FROM \"rlfoundation.azurecr.io/rocm6.3.4:vllm-0.8.5-numa-patch-ubuntu-22.04\"\nFROM \"rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04\"\n\nSHELL [\"/bin/bash\", \"-ceuxo\", \"pipefail\"]\n\nENV MAX_JOBS=512\n\nENV PATH=\"/usr/local/python3.12/bin:$PATH\"\nRUN ln -sf /usr/bin/python3.12 /usr/bin/python && \\\n    ln -sf /usr/bin/pip3.12 /usr/bin/pip\n\n############################################\n############################################\nRUN apt-get update\nRUN apt-get install -y pkg-config liblzma-dev\n############################################\n############################################\n\n\n###########################################\n##########Install TransformerEngine########\n###########################################\nWORKDIR /workspace/\n# transformer-engine install\n# https://github.com/ROCm/TransformerEngine\n\nRUN rm -rf TransformerEngine \nRUN git clone --recursive https://github.com/ROCm/TransformerEngine.git\nWORKDIR /workspace/TransformerEngine\nRUN git checkout 236178e5\n# git checkout bb061ade\n# git checkout 864405c\n\nENV NVTE_FRAMEWORK=pytorch \nENV NVTE_ROCM_ARCH=gfx942 \nENV NVTE_USE_HIPBLASLT=1\nENV NVTE_USE_ROCM=1  \n\n# export CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}\"\nENV CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr\"\n\n\n# ENV NVTE_BUILD_MAX_JOBS=$(MAX_JOBS)\n\nRUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv \n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n\n####################################################################################\n################Install vllm - sglang require vllm 0.6.7 dependency#################\n####################################################################################\n#### Require vllm 0.6.7 - checkout 113274a0\nWORKDIR /workspace/\nRUN rm -rf vllm\nRUN pip uninstall -y vllm\n# Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html\nRUN git clone https://github.com/ROCm/vllm.git\n# git clone https://github.com/vllm-project/vllm.git\nWORKDIR /workspace/vllm\nRUN git checkout 113274a0\nENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n#ENV MAX_JOBS=512\nENV MAX_JOBS=${MAX_JOBS}\nRUN pip install \"boto3>=1.26.0\"\nRUN pip install setuptools_scm\n# will add src into py. You can delete the repo\nRUN python3 setup.py install\nWORKDIR /workspace/\n####################################################################################\n####################################################################################\n####################################################################################\n\n\n\n###########################################\n############For hack docker################\n###########################################\nRUN pip install setuptools==75.8.0\n###########################################\n###########################################\n###########################################\n\n\n\n###########################################\n############build sgalng###################\n###########################################\n# Set environment variables\nENV BASE_DIR=/sgl-workspace\nENV BUILD_TYPE=all\nENV SGL_REPO=https://github.com/sgl-project/sglang\nENV SGL_BRANCH=v0.4.6.post5\nENV TRITON_REPO=https://github.com/ROCm/triton.git\nENV TRITON_COMMIT=improve_fa_decode_3.0.0\nENV AITER_REPO=https://github.com/ROCm/aiter.git\nENV AITER_COMMIT=v0.1.2\n# v0.1.2 version - commit id: 9d11f47\n# ENV AITER_COMMIT=9d11f47\n\nENV HIP_FORCE_DEV_KERNARG=1\nENV HSA_NO_SCRATCH_RECLAIM=1\nENV SGLANG_SET_CPU_AFFINITY=1\nENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1\nENV NCCL_MIN_NCHANNELS=112\nENV MOE_PADDING=1\nENV VLLM_FP8_PADDING=1\nENV VLLM_FP8_ACT_PADDING=1\nENV VLLM_FP8_WEIGHT_PADDING=1\nENV VLLM_FP8_REDUCE_CONV=1\nENV TORCHINDUCTOR_MAX_AUTOTUNE=1\nENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1\nENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\nENV AMDGPU_TARGETS=gfx942\nENV ROCM_ARCH=gfx942\nENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n\n# Switch to working directory\nWORKDIR /sgl-workspace\n\n# Clean and create directory\nRUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace\n\n# Clone and build sglang\nRUN git clone ${SGL_REPO} \\\n    && cd sglang \\\n    && git checkout ${SGL_BRANCH} || echo \"Using default branch\" \\\n    && cd sgl-kernel \\\n    && rm -f pyproject.toml \\\n    && mv pyproject_rocm.toml pyproject.toml \\\n    && python setup_rocm.py install \\\n    && cd .. \\\n    && if [ \"$BUILD_TYPE\" = \"srt\" ]; then \\\n         python -m pip --no-cache-dir install -e \"python[srt_hip]\"; \\\n       else \\\n         python -m pip --no-cache-dir install -e \"python[all_hip]\"; \\\n       fi \\\n    && cd /sgl-workspace \\\n    && cp -r /sgl-workspace/sglang /sglang \\\n    && python -m pip cache purge\n\n# Install common Python packages\nRUN pip install IPython orjson python-multipart torchao pybind11\n\n# Rebuild Triton\nRUN pip uninstall -y triton || true \\\n    && git clone ${TRITON_REPO} \\\n    && cd triton \\\n    && git checkout ${TRITON_COMMIT} \\\n    && cd python \\\n    && python3 setup.py install \\\n    && cd /sgl-workspace\n\n\n# ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1\"\n# ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\n\n# Build aiter\n#version: Commit 9d11f47\n    # && git checkout ${AITER_COMMIT} \\\nRUN pip uninstall -y aiter || true\nRUN git clone ${AITER_REPO} \\\n    && cd aiter \\\n    && git checkout ${AITER_COMMIT} \\\n    && git submodule sync \\\n    && git submodule update --init --recursive \\\n    && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \\\n    && cd /sgl-workspace\n    # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \\\n    # && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py develop \\\n\n# Copy MI300X config \nRUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \\\n         /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \\\n         -type f -name '*MI300X*' | \\\n         xargs -I {} sh -c 'vf_config=$(echo \"$1\" | sed \"s/MI300X/MI300X_VF/\"); cp \"$1\" \"$vf_config\"' -- {}\n\n# Environment setup complete.\nRUN echo \"Environment setup complete.\"\n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n\n\n###########################################\n###############vllm v0.8.5#################\n###########################################\n# ENV GITHUB_USERNAME=yushengsu-thu\n# ENV GITHUB_MAIL=yushengsu@gmail.com\n\n# RUN git config --global user.name \"${GITHUB_USERNAME}\" \\\n#     && git config --global user.email \"${GITHUB_MAIL}\" \n\nWORKDIR /workspace/\n\nENV VLLM_TARGET_DEVICE=rocm \nENV ROCM_PATH=/opt/rocm \nENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev\n\n# Find the repo path in: DockerFile/Dockerfile.rocm_yang\n# RUN git clone https://github.com/RLFoundation/vllm-patch.git\nRUN pip uninstall -y vllm || true\nRUN rm -rf vllm-patch\nRUN git clone https://github.com/RLFoundation/vllm-patch.git \\\n    && cd vllm-patch \\\n    && git checkout v0.8.5-sleep-numa \\\n    && rm -rf build/ dist/ *.egg-info \\\n    && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \\\n    && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py install\n    # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py develop\n\nWORKDIR /workspace/\n###########################################\n###########################################\n###########################################\n\n\n\n\n#########################################\n#### Install megatron-core###############\n#########################################\nRUN pip uninstall -y megatron-core && \\\n    git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \\\n    cd Megatron-LM-amd_version && \\\n    pip install -vvv -e . && \\\n    cd /workspace/\n#########################################\n#########################################\n#########################################\n\n\n\n\n#######################################\n################apex###################\n#######################################\nWORKDIR /workspace/\nRUN pip uninstall -y apex && \\\n    git clone https://github.com/ROCm/apex.git && \\\n    cd apex && \\\n    python setup.py install && \\\n    cd /workspace/ \n#######################################\n#######################################\n#######################################\n\n\n\n\n################################################################################\n###########################Add torch_memory_saver###############################\n################################################################################\n# Set environment variables\nENV HIPCC_COMPILE_FLAGS_APPEND=\"--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__\"\nENV CFLAGS=\"-D__HIP_PLATFORM_AMD__\"\nENV CXXFLAGS=\"-D__HIP_PLATFORM_AMD__\"\nRUN pip install \"git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa\"\n################################################################################\n################################################################################\n################################################################################\n\n\n\n########################################\n######Install ray#######################\n########################################\n# need to add this patch: https://github.com/ray-project/ray/pull/53531/files\nRUN pip uninstall ray -y\nRUN pip install \"ray[data,train,tune,serve]>=2.47.0\" \n########################################\n########################################\n########################################\n\n\n\n##########################################\n#######Install other dependencies#########\n##########################################\nRUN pip install \"tensordict==0.6.2\" --no-deps && \\\n    pip install accelerate \\\n    codetiming \\\n    datasets \\\n    dill \\\n    hydra-core \\\n    liger-kernel \\\n    numpy \\\n    pandas \\\n    peft \\\n    \"pyarrow>=15.0.0\" \\\n    pylatexenc \\\n    torchdata \\\n    wandb \\\n    orjson \\\n    pybind11\n    \nWORKDIR /workspace/\nRUN git clone https://github.com/volcengine/verl.git && \\\n    cd verl && \\\n    pip install -e . \n##########################################\n##########################################\n##########################################\n\n\n\nWORKDIR /workspace/\n\nCMD [\"/usr/bin/bash\"]\nCMD [\"/usr/bin/bash\"]\n"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.6.post5 and torch-memory-saver\nRUN pip install --resume-retries 999 \"sglang[all]==0.4.6.post5\" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir\n\n# Some sglang operations in 0.4.6.post5 require vllm\n# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Fix for transformers 4.53.0\nRUN pip3 install --no-cache-dir \"transformers[hf_xet]<4.52.0\"\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.12.deepep",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.6.post5 and torch-memory-saver\nRUN pip install --resume-retries 999 \"sglang[all]==0.4.6.post5\" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir\n\n# Some sglang operations in 0.4.6.post5 require vllm\n# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Fix for transformers 4.53.0\nRUN pip3 install --no-cache-dir \"transformers[hf_xet]<4.52.0\"\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.sglang.vllm.mcore0.13.preview",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.6.post5 and torch-memory-saver\nRUN pip install --resume-retries 999 \"sglang[all]==0.4.6.post5\" --no-cache-dir --find-links https://flashinfer.ai/whl/cu124/torch2.6/flashinfer-python && pip install torch-memory-saver --no-cache-dir\n\n# Some sglang operations in 0.4.6.post5 require vllm\n# [Warning] vllm can have some packages not compatible with sglang, for example, flashinfer\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.13.0\n\n# Fix for transformers 4.53.0\nRUN pip3 install --no-cache-dir \"transformers[hf_xet]<4.52.0\"\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install torch-2.6.0+cu124 + vllm-0.8.5.post1\n# torch-2.6.0+cu124: cxx11abi=False\n# torch-2.6.0+cu126: cxx11abi=True\n# see https://github.com/flashinfer-ai/flashinfer/issues/911\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True)\n# vllm-0.8.3 does not support flashinfer>=0.2.3\n# see https://github.com/vllm-project/vllm/pull/15777\nRUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Fix for transformers 4.53.0\nRUN pip3 install --no-cache-dir \"transformers[hf_xet]<4.52.0\"\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.12.deepep",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install torch-2.6.0+cu124 + vllm-0.8.5.post1\n# torch-2.6.0+cu124: cxx11abi=False\n# torch-2.6.0+cu126: cxx11abi=True\n# see https://github.com/flashinfer-ai/flashinfer/issues/911\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True)\n# vllm-0.8.3 does not support flashinfer>=0.2.3\n# see https://github.com/vllm-project/vllm/pull/15777\nRUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Fix for transformers 4.53.0\nRUN pip3 install --no-cache-dir \"transformers[hf_xet]<4.52.0\"\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install torch-2.6.0+cu124 + vllm-0.8.5.post1\n# torch-2.6.0+cu124: cxx11abi=False\n# torch-2.6.0+cu126: cxx11abi=True\n# see https://github.com/flashinfer-ai/flashinfer/issues/911\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.8.5.post1\n\n# Install flashinfer-0.2.2.post1+cu126 (cxx11abi=True)\n# vllm-0.8.3 does not support flashinfer>=0.2.3\n# see https://github.com/vllm-project/vllm/pull/15777\nRUN aria2c --max-tries=9999 https://github.com/flashinfer-ai/flashinfer/releases/download/v0.2.2.post1/flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    pip install --no-cache-dir flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl && \\\n    rm flashinfer_python-0.2.2.post1+cu124torch2.6-cp38-abi3-linux_x86_64.whl\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.base",
    "content": "# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks\n# Target: verlai/verl:base-v2-cu124-cudnn9.8-torch2.6-fa2.8.0-te2.3\n# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10)\n# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html\nFROM nvcr.io/nvidia/pytorch:24.08-py3\n\n# Define environments\nENV MAX_JOBS=16\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Define installation arguments\nARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/\nARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple\n\n# Set apt source\nRUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \\\n    { \\\n    echo \"deb ${APT_SOURCE} jammy main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-updates main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-backports main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-security main restricted universe multiverse\"; \\\n    } > /etc/apt/sources.list\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install tini\nRUN apt-get update && \\\n    apt-get install -y tini aria2 && \\\n    apt-get clean\n\n# Change pip source\nRUN pip config set global.index-url \"${PIP_INDEX}\" && \\\n    pip config set global.extra-index-url \"${PIP_INDEX}\" && \\\n    python -m pip install --upgrade pip\n\n# Uninstall nv-pytorch fork\nRUN pip uninstall -y torch torchvision torchaudio \\\n    pytorch-quantization pytorch-triton torch-tensorrt \\\n    xgboost transformer_engine flash_attn apex megatron-core grpcio\n\n# Reinstall CUDA 12.4\nRUN aria2c https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-ubuntu2204.pin && \\\n    mv cuda-ubuntu2204.pin /etc/apt/preferences.d/cuda-repository-pin-600\n\nRUN aria2c --always-resume=true --max-tries=99999 https://developer.download.nvidia.com/compute/cuda/12.4.1/local_installers/cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \\\n    dpkg -i cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \\\n    cp /var/cuda-repo-ubuntu2204-12-4-local/cuda-*-keyring.gpg /usr/share/keyrings/ && \\\n    apt-get update && \\\n    apt-get -y install cuda-toolkit-12-4 && \\\n    rm cuda-repo-ubuntu2204-12-4-local_12.4.1-550.54.15-1_amd64.deb && \\\n    update-alternatives --set cuda /usr/local/cuda-12.4 && \\\n    rm -rf /usr/local/cuda-12.6\n\nRUN pip install --resume-retries 999 --no-cache-dir torch==2.6.0 torchvision==0.21.0 torchaudio==2.6.0\n\nRUN pip install --resume-retries 999 --no-cache-dir \"tensordict==0.6.2\" torchdata \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\n# Install flash-attn-2.7.4.post1 (cxx11abi=False)\nRUN wget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl && \\\n    pip install --no-cache-dir flash_attn-2.7.4.post1+cu12torch2.6cxx11abiFALSE-cp310-cp310-linux_x86_64.whl\n\n# Fix packages\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\n# Install cudnn\nRUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \\\n    apt-get update && \\\n    apt-get -y install cudnn-cuda-12 && \\\n    rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb\n\n# Install Apex\nRUN git clone https://github.com/NVIDIA/apex.git && \\\n    cd apex && \\\n    pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" ./\n\n# Profiling tools\nRUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0 && \\\n    dpkg -i ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb\n\n# Fix opencv\nRUN pip install --resume-retries 999 --no-cache-dir opencv-python\n\nRUN pip install --resume-retries 999 --no-cache-dir opencv-fixer && \\\n    python -c \"from opencv_fixer import AutoFix; AutoFix()\"\n\nRUN pip install --resume-retries 999 --no-cache-dir cuda-bindings\n\n# Reset pip config\nRUN pip config unset global.index-url && \\\n    pip config unset global.extra-index-url\n\nRUN apt-get update && \\\n    apt-get install -y libfreeimage3 libfreeimage-dev zlib1g htop\n\n"
  },
  {
    "path": "docker/verl0.4-cu124-torch2.6-fa2.7.4/README.md",
    "content": "# verl image with verl v0.4.x\n\n## Important packages version\n\n```txt\ncuda==12.4\ncudnn==9.8.0\ntorch==2.6.0\nflash_attn=2.7.4\nsglang==0.4.6.post5\nvllm==0.8.5.post1\nnvidia-cudnn-cu12==9.8.0.87\ntransformer_engine==2.3\nmegatron.core==core_v0.12.2\n# Preview\ntransformer_engine==2.5\nmegatron.core==core_r0.13.0\n```\n\n## Target\n\n- Base image: \n    - `verlai/verl:base-verl0.4-cu124-cudnn9.8-torch2.6-fa2.7.4`\n- App image:\n    - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.12.2-te2.2`: SGLang requires vLLM in 0.4.6.post5 version, vLLM can have some package conflicts with SGLang\n    - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.12.2-te2.2-deepep`: Built with deepep\n    - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.12.2-te2.2`\n    - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.12.2-te2.2-deepep`: Built with deepep\n- Preview image:\n    - `verlai/verl:app-verl0.4-sglang0.4.6.post5-vllm0.8.5-mcore0.13.0-te2.2-preview`\n    - `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.13.0-te2.2-preview`"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.10.post2.mcore0.13",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=8\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.10\n# Install FlashInfer Python package\nRUN pip install --upgrade pip setuptools packaging\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.9rc1\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation \"sglang[all]==0.4.10.post2\"\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]==4.55.4\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.sglang0.4.9.post6.mcore0.13",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=8\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.10\n# Install FlashInfer Python package\nRUN pip install --upgrade pip setuptools packaging\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.9rc1\nRUN pip install --resume-retries 999  --no-cache-dir --no-build-isolation \"sglang[all]==0.4.9.post6\"\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]==4.55.4\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.13",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install torch-2.7.1+cu126 + vllm-0.10.0\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.10.0\n\n# Fix packages\n# transformers 4.54.0 still not support\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.55.4\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.0\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge\n\n# Fix qwen vl\nRUN pip3 install --no-cache-dir --no-deps trl"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.app.vllm.mcore0.15",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM iseekyan/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4-h100\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install torch-2.7.1+cu126 + vllm-0.10.0\nRUN pip install --resume-retries 999 --no-cache-dir vllm==0.10.0\n\n# Fix packages\n# transformers 4.54.0 still not support\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.55.4\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.7\nRUN pip install onnxscript\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.15.0rc4\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge==v0.15.0\n\n# Fix qwen vl\nRUN pip3 install --no-cache-dir --no-deps trl"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/Dockerfile.base.torch2.7.1",
    "content": "# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks\n# Target: verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6\n# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10)\n# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html\nFROM nvcr.io/nvidia/pytorch:24.08-py3\n\n# Define environments\nENV MAX_JOBS=16\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Define installation arguments\nARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/\nARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple\n\n# Set apt source\nRUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \\\n    { \\\n    echo \"deb ${APT_SOURCE} jammy main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-updates main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-backports main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-security main restricted universe multiverse\"; \\\n    } > /etc/apt/sources.list\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install tini\nRUN apt-get update && \\\n    apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \\\n    apt-get clean\n\n# Change pip source\nRUN pip config set global.index-url \"${PIP_INDEX}\" && \\\n    pip config set global.extra-index-url \"${PIP_INDEX}\" && \\\n    python -m pip install --upgrade pip\n\n# Uninstall nv-pytorch fork\nRUN pip uninstall -y torch torchvision torchaudio \\\n    pytorch-quantization pytorch-triton torch-tensorrt \\\n    xgboost transformer_engine flash_attn apex megatron-core grpcio\n\nRUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1\n\n# Install flash-attn-2.7.4.post1, although built with torch2.6, it is compatible with torch2.7\n# https://github.com/Dao-AILab/flash-attention/issues/1644#issuecomment-2899396361\nRUN ABI_FLAG=$(python -c \"import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')\") && \\\n    URL=\"https://github.com/Dao-AILab/flash-attention/releases/download/v2.7.4.post1/flash_attn-2.7.4.post1+cu12torch2.6cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl\" && \\\n    FILE=\"flash_attn-2.7.4.post1+cu12torch2.6cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl\" && \\\n    wget -nv \"${URL}\" && \\\n    pip install --no-cache-dir \"${FILE}\"\n\n# Fix packages\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\n# Install cudnn\nRUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \\\n    apt-get update && \\\n    apt-get -y install cudnn-cuda-12 && \\\n    rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb\n\n# Install Apex\nRUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" --resume-retries 999 git+https://github.com/NVIDIA/apex.git\n\n# Profiling tools\nRUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0\n\nRUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb\n\nRUN pip install --resume-retries 999 --no-cache-dir \"tensordict==0.6.2\" torchdata \"transformers[hf_xet]>=4.52.3\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas cuda-bindings \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install\n\n# Reset pip config\nRUN pip config unset global.index-url && \\\n    pip config unset global.extra-index-url\n\n"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7-fa2.7.4/README.md",
    "content": "# verl image with verl v0.5\n\n## Important packages version\n\n```txt\ncuda==12.6\ncudnn==9.8.0\ntorch==2.7.1\nflash_attn=2.7.4.post1\nsglang==0.4.9.post6\nvllm==0.8.5.post1\nnvidia-cudnn-cu12==9.8.0.87\ntransformer_engine==2.3\nmegatron.core==core_v0.12.2\n# Preview\ntransformer_engine==2.5\nmegatron.core==core_r0.13.0\n```\n\n## Target\n\n- Base image:\n  - `verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.7.4`: We offer a base image with deep ep built in, for vllm/sglang\n- App image:\n  - `verlai/verl:app-verl0.5-transformers4.55.4-vllm0.10.0-mcore0.13.0-te2.2`\n  - `verlai/verl:app-verl0.5-transformers4.55.4-sglang0.4.10.post2-mcore0.13.0-te2.2`\n  - `iseekyan/verl:app-verl0.5-transformers4.55.4-vllm0.10.0-mcore0.15.0-te2.7`\n"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.12",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0\n\n# Define environments\nENV MAX_JOBS=8\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.8 and torch-memory-saver\n# Install FlashInfer Python package\nRUN pip install --upgrade pip setuptools packaging\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1\nRUN pip install --resume-retries 999  --no-cache-dir \"sglang[all]==0.4.8\" && pip install torch-memory-saver --no-cache-dir\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.3\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.mcore0.13.preview",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0\n\n# Define environments\nENV MAX_JOBS=8\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.8 and torch-memory-saver\n# Install FlashInfer Python package\nRUN pip install --upgrade pip setuptools packaging\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1\nRUN pip install --resume-retries 999  --no-cache-dir \"sglang[all]==0.4.8\" && pip install torch-memory-saver --no-cache-dir\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.2\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7.1-fa2.8.0/Dockerfile.base",
    "content": "# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks\n# Target: verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6\n# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10)\n# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html\nFROM nvcr.io/nvidia/pytorch:24.08-py3\n\n# Define environments\nENV MAX_JOBS=16\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Define installation arguments\nARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/\nARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple\n\n# Set apt source\nRUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \\\n    { \\\n    echo \"deb ${APT_SOURCE} jammy main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-updates main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-backports main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-security main restricted universe multiverse\"; \\\n    } > /etc/apt/sources.list\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install tini\nRUN apt-get update && \\\n    apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \\\n    apt-get clean\n\n# Change pip source\nRUN pip config set global.index-url \"${PIP_INDEX}\" && \\\n    pip config set global.extra-index-url \"${PIP_INDEX}\" && \\\n    python -m pip install --upgrade pip\n\n# Uninstall nv-pytorch fork\nRUN pip uninstall -y torch torchvision torchaudio \\\n    pytorch-quantization pytorch-triton torch-tensorrt \\\n    xgboost transformer_engine flash_attn apex megatron-core grpcio\n\nRUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1\n\n# Install flash-attn-2.8.0.post2 (cxx11abi=True)\nRUN ABI_FLAG=$(python -c \"import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')\") && \\\n    URL=\"https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl\" && \\\n    FILE=\"flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp310-cp310-linux_x86_64.whl\" && \\\n    wget -nv \"${URL}\" && \\\n    pip install --no-cache-dir \"${FILE}\"\n\n# Fix packages\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\n# Install cudnn\nRUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \\\n    apt-get update && \\\n    apt-get -y install cudnn-cuda-12 && \\\n    rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb\n\n# Install Apex\nRUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" --resume-retries 999 git+https://github.com/NVIDIA/apex.git\n\n# Profiling tools\nRUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0\n\nRUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb\n\nRUN pip install --resume-retries 999 --no-cache-dir \"tensordict==0.6.2\" torchdata \"transformers[hf_xet]>=4.53\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas cuda-bindings \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pyext pre-commit ruff\n\n# Install DeepEP\n## the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n## Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\n\n## Build deepep-nvshmem\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install\n\n# Reset pip config\nRUN pip config unset global.index-url && \\\n    pip config unset global.extra-index-url\n\n"
  },
  {
    "path": "docker/verl0.5-cu126-torch2.7.1-fa2.8.0/README.md",
    "content": "# verl image with verl v0.5\n\n## Important packages version\n\n```txt\ncuda==12.6\ncudnn==9.8.0\ntorch==2.7.1\nflash_attn=2.8.0    ##\nsglang==0.4.8\nvllm==0.8.5.post1\nnvidia-cudnn-cu12==9.8.0.87\ntransformer_engine==2.3\nmegatron.core==core_v0.12.2\n# Preview\ntransformer_engine==2.5\nmegatron.core==core_r0.13.0\n```\n\n## Target\n\n- Base image:\n    - `verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.1-fa2.8.0`: We offer a base image with deep ep built in\n- App image:\n    - `verlai/verl:app-verl0.5-sglang0.4.9-mcore0.12.2`\n    - `verlai/verl:app-verl0.5-sglang0.4.9-mcore0.13.0-preview`\n- vllm temporarily not support latest version"
  },
  {
    "path": "docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.app.sglang.megatron",
    "content": "# Start from the verl base image\n# Dockerfile.base\nFROM verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6\n\n# Define environments\nENV MAX_JOBS=8\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Install sglang-0.4.8 and torch-memory-saver\n# Install FlashInfer Python package\nRUN pip install --resume-retries 999 --no-cache-dir --no-build-isolation flashinfer-python==0.2.6.post1\nRUN pip install --resume-retries 999  --no-cache-dir \"sglang[all]==0.4.8\" && pip install torch-memory-saver --no-cache-dir\n\n# Fix packages\nRUN pip install --no-cache-dir \"tensordict==0.6.2\" \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pre-commit ruff\n\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --resume-retries 999 --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\nRUN pip install --resume-retries 999 --no-cache-dir nvidia-cudnn-cu12==9.8.0.87\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --resume-retries 999 --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.5\n\n# Install Megatron-LM\nRUN pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/Megatron-LM.git@core_r0.13.0\n\n# Install mbridge\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/Dockerfile.base",
    "content": "# Base Docker Image of verl, with CUDA/Torch/FlashAttn/Apex/TransformerEngine, without other frameworks\n# Target: verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0-fi0.2.6\n# Start from the NVIDIA official image (ubuntu-22.04 + cuda-12.6 + python-3.10)\n# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-24-08.html\nFROM nvcr.io/nvidia/pytorch:25.02-py3\n\n# Define environments\nENV MAX_JOBS=16\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\n\n# Define installation arguments\nARG APT_SOURCE=https://mirrors.tuna.tsinghua.edu.cn/ubuntu/\nARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple\n\n# Set apt source\nRUN cp /etc/apt/sources.list /etc/apt/sources.list.bak && \\\n    { \\\n    echo \"deb ${APT_SOURCE} jammy main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-updates main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-backports main restricted universe multiverse\"; \\\n    echo \"deb ${APT_SOURCE} jammy-security main restricted universe multiverse\"; \\\n    } > /etc/apt/sources.list\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install tini\nRUN apt-get update && \\\n    apt-get install -y tini aria2 libfreeimage3 libfreeimage-dev zlib1g htop && \\\n    apt-get clean\n\n# Change pip source\nRUN pip config set global.index-url \"${PIP_INDEX}\" && \\\n    pip config set global.extra-index-url \"${PIP_INDEX}\" && \\\n    python -m pip install --upgrade pip\n\n# Uninstall nv-pytorch fork\nRUN pip uninstall -y torch torchvision torchaudio \\\n    pytorch-quantization pytorch-triton torch-tensorrt \\\n    xgboost transformer_engine flash_attn apex megatron-core grpcio\n\nRUN pip install --resume-retries 999 --no-cache-dir torch==2.7.1 torchvision==0.22.1 torchaudio==2.7.1 --index-url https://download.pytorch.org/whl/cu128\n\n# Install flash-attn-2.8.0.post2 (cxx11abi=True)\nRUN ABI_FLAG=$(python -c \"import torch; print('TRUE' if torch._C._GLIBCXX_USE_CXX11_ABI else 'FALSE')\") && \\\n    URL=\"https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.0.post2/flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp312-cp312-linux_x86_64.whl\" && \\\n    FILE=\"flash_attn-2.8.0.post2+cu12torch2.7cxx11abi${ABI_FLAG}-cp312-cp312-linux_x86_64.whl\" && \\\n    wget -nv \"${URL}\" && \\\n    pip install --no-cache-dir \"${FILE}\"\n\n# Fix packages\nRUN pip uninstall -y pynvml nvidia-ml-py && \\\n    pip install --no-cache-dir --upgrade \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\n# Install cudnn\nRUN aria2c --max-tries=9999 https://developer.download.nvidia.com/compute/cudnn/9.8.0/local_installers/cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    dpkg -i cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb && \\\n    cp /var/cudnn-local-repo-ubuntu2204-9.8.0/cudnn-*-keyring.gpg /usr/share/keyrings/ && \\\n    apt-get update && \\\n    apt-get -y install cudnn-cuda-12 && \\\n    rm cudnn-local-repo-ubuntu2204-9.8.0_1.0-1_amd64.deb\n\n# Install Apex\nRUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" --resume-retries 999 git+https://github.com/NVIDIA/apex.git\n\n# Profiling tools\nRUN aria2c --always-resume=true --max-tries=99999 https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_3/nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0\n\nRUN apt-get install -y ./nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.3.1/target-linux-x64/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-2025.3.1_2025.3.1.90-1_amd64.deb\n\nRUN pip install --resume-retries 999 --no-cache-dir \"tensordict==0.6.2\" torchdata \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas cuda-bindings \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pre-commit ruff\n\n# Reset pip config\nRUN pip config unset global.index-url && \\\n    pip config unset global.extra-index-url\n\n"
  },
  {
    "path": "docker/verl0.5-preview-cu128-torch2.7.1-fa2.8.0/README.md",
    "content": "# verl image with verl v0.5\n\n## Important packages version\n\n```txt\ncuda==12.8\ncudnn==9.8.0\ntorch==2.7.1\nflash_attn=2.8.0    ##\nsglang==0.4.8\ntransformer_engine==2.5\nmegatron.core==core_r0.13.0\nnvidia-cudnn-cu12==9.8.0.87\n```\n\n## Target\n\n- Base image:\n    - `verlai/verl:base-verl0.5-preview-cu128-cudnn9.8-torch2.7.1-fa2.8.0`: We offer a base image with flash infer 0.2.6.post1 built in\n- App image:\n    - `verlai/verl:app-verl0.5-preview-sglang0.4.8-mcore0.13.0-preview`\n- vllm temporarily not support latest version\n\n## !!!Notice!!!\n\n- pyext is lack of maintainace and cannot work with python 3.12, consider using replacement and deprecating this package."
  },
  {
    "path": "docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.app.sglang",
    "content": "FROM verlai/verl:base-verl0.6-cu128-cudnn9.8-torch2.8.0-fa2.7.4\n\nRUN pip install --no-cache-dir \"sglang[all]==0.5.2\"\nRUN pip install --no-cache-dir \"torch-memory-saver==0.0.9rc1\"\n"
  },
  {
    "path": "docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.base",
    "content": "# Start from the NVIDIA official image (ubuntu-24.04 + cuda-12.8 + python-3.12)\n# https://docs.nvidia.com/deeplearning/frameworks/pytorch-release-notes/rel-25-03.html\nFROM nvcr.io/nvidia/pytorch:25.03-py3\n\n# Define environments\nENV MAX_JOBS=32\nENV VLLM_WORKER_MULTIPROC_METHOD=spawn\nENV DEBIAN_FRONTEND=noninteractive\nENV NODE_OPTIONS=\"\"\nENV PIP_ROOT_USER_ACTION=ignore\nENV HF_HUB_ENABLE_HF_TRANSFER=\"1\"\nENV PIP_CONSTRAINT=\"\"\n\nARG PIP_INDEX=https://mirrors.tuna.tsinghua.edu.cn/pypi/web/simple\n\n# Change pip source\nRUN pip config set global.index-url \"${PIP_INDEX}\" && \\\n    pip config set global.extra-index-url \"${PIP_INDEX}\" && \\\n    pip config set global.no-cache-dir \"true\" && \\\n    python -m pip install --upgrade pip\n\n# Install systemctl\nRUN apt-get update && \\\n    apt-get install -y -o Dpkg::Options::=\"--force-confdef\" systemd && \\\n    apt-get clean\n\n# Install libxml2\nRUN apt-get update && \\\n    apt-get install -y libxml2 aria2 && \\\n    apt-get clean\n\n# Uninstall nv-pytorch fork\nRUN pip uninstall -y torch torchvision torchaudio \\\n    pytorch-quantization pytorch-triton torch-tensorrt \\\n    transformer_engine flash_attn apex megatron-core \\\n    xgboost opencv grpcio\n\n# Fix packages\nRUN pip install --no-cache-dir tensordict torchdata \"transformers[hf_xet]==4.55.4\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=19.0.1\" pandas \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler blobfile xgrammar \\\n    pytest py-spy pre-commit ruff\n\n# Fix cv2\nRUN rm -rf /usr/local/lib/python3.11/dist-packages/cv2\n\n# Install torch\nRUN pip install --no-cache-dir torch==2.8.0 --index-url https://download.pytorch.org/whl/cu128\n\n# Install flash-attn\nRUN pip install --no-cache-dir --no-build-isolation flash_attn==2.7.4.post1\n\n# Install DeepEP\n# the dependency of IBGDA\nRUN ln -s /usr/lib/x86_64-linux-gnu/libmlx5.so.1 /usr/lib/x86_64-linux-gnu/libmlx5.so\n\n# Clone and build deepep and deepep-nvshmem\nRUN git clone -b v2.3.1 https://github.com/NVIDIA/gdrcopy.git && \\\n    git clone https://github.com/deepseek-ai/DeepEP.git  && \\\n    cd DeepEP && git checkout a84a248\n\n# Prepare nvshmem\nRUN wget https://developer.nvidia.com/downloads/assets/secure/nvshmem/nvshmem_src_3.2.5-1.txz && \\\n    tar -xvf nvshmem_src_3.2.5-1.txz && mv nvshmem_src deepep-nvshmem && \\\n    cd deepep-nvshmem && git apply ../DeepEP/third-party/nvshmem.patch\n\n## Build deepep-nvshmem\nRUN apt-get install -y ninja-build cmake\n\nENV CUDA_HOME=/usr/local/cuda\n### Set MPI environment variables. Having errors when not set.\nENV CPATH=/usr/local/mpi/include:$CPATH\nENV LD_LIBRARY_PATH=/usr/local/mpi/lib:$LD_LIBRARY_PATH\nENV LD_LIBRARY_PATH=/usr/local/x86_64-linux-gnu:$LD_LIBRARY_PATH\nENV GDRCOPY_HOME=/workspace/gdrcopy\nENV GDRCOPY_INCLUDE=/workspace/gdrcopy/include\n\nRUN cd deepep-nvshmem && \\\n    NVSHMEM_SHMEM_SUPPORT=0 \\\n    NVSHMEM_UCX_SUPPORT=0 \\\n    NVSHMEM_USE_NCCL=0 \\\n    NVSHMEM_MPI_SUPPORT=0 \\\n    NVSHMEM_IBGDA_SUPPORT=1 \\\n    NVSHMEM_PMIX_SUPPORT=0 \\\n    NVSHMEM_TIMEOUT_DEVICE_POLLING=0 \\\n    NVSHMEM_USE_GDRCOPY=1 \\\n    cmake -G Ninja -S . -B build/ -DCMAKE_INSTALL_PREFIX=/workspace/deepep-nvshmem/install && cmake --build build/ --target install\n\nENV NVSHMEM_DIR=/workspace/deepep-nvshmem/install\nENV LD_LIBRARY_PATH=$NVSHMEM_DIR/lib:$LD_LIBRARY_PATH\nENV PATH=$NVSHMEM_DIR/bin:$PATH\n\n## Build deepep\nRUN cd DeepEP && \\\n    python setup.py install\n\n# Install Apex\nRUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" git+https://github.com/NVIDIA/apex.git\n\n# Install TransformerEngine\nRUN export NVTE_FRAMEWORK=pytorch && pip3 install --no-deps --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@v2.2.1\n\n# Install Megatron-LM\nRUN git clone -b core_v0.13.0 https://github.com/NVIDIA/Megatron-LM.git && \\\n    cd Megatron-LM && pip3 install --no-deps -e .\n\n# Install mbridge\nRUN pip3 install --no-cache-dir git+https://github.com/ISEEKYAN/mbridge.git\n"
  },
  {
    "path": "docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.vllm011.mcore_gpt-oss",
    "content": "FROM nvcr.io/nvidia/nemo:25.07.gpt_oss\n\nRUN git clone -b v0.11.0 --depth 1 https://github.com/vllm-project/vllm.git /opt/vllm\n\nRUN pip install setuptools_scm\n\nRUN cd /opt/vllm && pip install --no-deps --no-build-isolation --no-cache-dir -e .\n\nRUN pip install cbor2 setproctitle blake3 openai_harmony pybase64 msgspec partial_json_parser py-cpuinfo diskcache gguf\n\nRUN pip install --upgrade transformers tokenizers\n\nRUN pip install codetiming tensordict mathruler pylatexenc\n\nRUN pip3 install --no-cache-dir mbridge"
  },
  {
    "path": "docker/verl0.6.1-experimental/Dockerfile.sglang056exp",
    "content": "# Dockerfile for verlai/verl:sgl056.exp\nFROM lmsysorg/sglang:v0.5.6.post1\n\nRUN pip install pybind11\n\nRUN pip install nvidia-mathdx\n\nRUN pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" git+https://github.com/NVIDIA/apex.git\n\nRUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.11\n\nRUN pip install --upgrade --no-cache-dir transformers tokenizers\n\nRUN pip install codetiming tensordict mathruler pylatexenc qwen_vl_utils\n\nRUN pip install --no-cache-dir --no-build-isolation flash_attn==2.8.1\n\nRUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ \"$(uname -m)\" = \"aarch64\" ]; then echo \"arm64\"; else echo \"amd64\"; fi) && \\\n    wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0 && \\\n    apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-${NSIGHT_VERSION}.deb\n\n\n# =========================\n# Install HybridEP\n# =========================\nWORKDIR /home/\nRUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \\\n    cd DeepEP && git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \\\n    TORCH_CUDA_ARCH_LIST=\"9.0;10.0\" pip install --no-build-isolation .\n\n# =========================\n# Install Qwen3-Next dependencies\n# =========================\nWORKDIR /home/\n# Install causal-conv1d and flash-linear-attention\nRUN cd /tmp && \\\n    git clone https://github.com/Dao-AILab/causal-conv1d.git && \\\n    cd causal-conv1d && \\\n    unset PIP_CONSTRAINT && \\\n    CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install --no-build-isolation . && \\\n    cd .. && \\\n    rm -rf causal-conv1d && \\\n    pip install flash-linear-attention\n\nRUN pip install --no-cache-dir torch-memory-saver\n\nRUN pip3 install --no-cache-dir --no-deps trl\n\nRUN pip3 install nvtx matplotlib liger_kernel\n\nRUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n\nRUN pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@1d462bd37dac21cfa14177405d4921eedb987052 # latest dev branch on 20251209\n\nRUN pip install git+https://github.com/volcengine/verl.git@v0.6.1\n\nRUN pip uninstall -y verl"
  },
  {
    "path": "docker/verl0.6.1-experimental/Dockerfile.vllm012exp",
    "content": "# dockerfile for verlai/verl:vll012.exp\nFROM nvcr.io/nvidia/pytorch:25.11-py3\n\nRUN git clone -b v0.12.0 --depth 1 https://github.com/vllm-project/vllm.git /opt/vllm\n\nRUN pip install setuptools_scm\n\nRUN cd /opt/vllm && pip install --no-deps --no-build-isolation --no-cache-dir -e .\n\nRUN pip install -r /opt/vllm/requirements/common.txt\n\n\nRUN pip install pybind11\n\nRUN export NVTE_FRAMEWORK=pytorch && MAX_JOBS=128 NVTE_BUILD_THREADS_PER_JOB=4 pip3 install --resume-retries 999 --no-cache-dir --no-build-isolation git+https://github.com/NVIDIA/TransformerEngine.git@release_v2.11\n\nRUN pip install --upgrade --no-cache-dir transformers tokenizers\n\nRUN pip install codetiming tensordict mathruler pylatexenc qwen_vl_utils\n\nRUN pip install flash_attn\n#==2.8.1\n\nRUN apt update && apt install numactl\n\nRUN NSIGHT_VERSION=2025.6.1_2025.6.1.190-1_$(if [ \"$(uname -m)\" = \"aarch64\" ]; then echo \"arm64\"; else echo \"amd64\"; fi) && \\\n    wget https://developer.nvidia.com/downloads/assets/tools/secure/nsight-systems/2025_6/nsight-systems-${NSIGHT_VERSION}.deb && \\\n    apt-get update && apt-get install -y libxcb-cursor0 && \\\n    apt-get install -y ./nsight-systems-${NSIGHT_VERSION}.deb && \\\n    rm -rf /usr/local/cuda/bin/nsys && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys  /usr/local/cuda/bin/nsys && \\\n    rm -rf /usr/local/cuda/bin/nsys-ui && \\\n    ln -s /opt/nvidia/nsight-systems/2025.6.1/nsys-ui /usr/local/cuda/bin/nsys-ui && \\\n    rm nsight-systems-${NSIGHT_VERSION}.deb\n\n\n# =========================\n# Install HybridEP\n# =========================\nWORKDIR /home/\nRUN git clone --branch hybrid-ep https://github.com/deepseek-ai/DeepEP.git && \\\n    cd DeepEP && git checkout 3f601f7ac1c062c46502646ff04c535013bfca00 && \\\n    TORCH_CUDA_ARCH_LIST=\"9.0;10.0\" pip install --no-build-isolation .\n\n# =========================\n# Install Qwen3-Next dependencies\n# =========================\nWORKDIR /home/\n# Install causal-conv1d and flash-linear-attention\nRUN cd /tmp && \\\n    git clone https://github.com/Dao-AILab/causal-conv1d.git && \\\n    cd causal-conv1d && \\\n    unset PIP_CONSTRAINT && \\\n    CAUSAL_CONV1D_FORCE_BUILD=TRUE pip install --no-build-isolation . && \\\n    cd .. && \\\n    rm -rf causal-conv1d && \\\n    pip install flash-linear-attention\n\nRUN pip3 install --no-cache-dir --no-deps trl\n\nRUN pip3 install nvtx matplotlib liger_kernel\n\nRUN pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n\nRUN pip install --no-deps --no-cache-dir git+https://github.com/NVIDIA/Megatron-LM.git@1d462bd37dac21cfa14177405d4921eedb987052 # latest dev branch on 20251209\n\nRUN pip install git+https://github.com/volcengine/verl.git@v0.6.1\n\nRUN pip uninstall -y verl"
  },
  {
    "path": "docs/Makefile",
    "content": "# Minimal makefile for Sphinx documentation\n#\n\n# You can set these variables from the command line.\nSPHINXOPTS    =\nSPHINXBUILD   = sphinx-build\nSPHINXPROJ    = verl\nSOURCEDIR     = .\nBUILDDIR      = _build\n\n# Put it first so that \"make\" without argument is like \"make help\".\nhelp:\n\t@$(SPHINXBUILD) -M help \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n\n.PHONY: help Makefile\n\n# Catch-all target: route all unknown targets to Sphinx using the new\n# \"make mode\" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).\n%: Makefile\n\t@$(SPHINXBUILD) -M $@ \"$(SOURCEDIR)\" \"$(BUILDDIR)\" $(SPHINXOPTS) $(O)\n"
  },
  {
    "path": "docs/README.md",
    "content": "# verl documentations\n\n## Build the docs\n\n```bash\n# If you want to view auto-generated API docstring, please make sure verl is available in python path. For instance, install verl via:\n# pip install .. -e[test]\n\n# Install dependencies needed for building docs.\npip install -r requirements-docs.txt\n\n# Build the docs.\nmake clean\nmake html\n```\n\n## Open the docs with your browser\n\n```bash\npython -m http.server -d _build/html/\n```\nLaunch your browser and navigate to http://localhost:8000 to view the documentation. Alternatively you could drag the file `_build/html/index.html` to your local browser and view directly.\n"
  },
  {
    "path": "docs/README_vllm0.7.md",
    "content": "# Upgrading to vllm >= 0.7\n\nNote: verl+vllm 0.8.3 is now stable. Please see ``docs/README_vllm0.8.md`` for upgrade guide.\n\n## Installation\n\nNote: At time of writing, verl+vllm 0.7.x supports **FSDP** for training and **vLLM** for rollout.\n\n```\n# Create the conda environment\nconda create -n verl python==3.10\nconda activate verl\n\n# Install verl\ngit clone https://github.com/volcengine/verl.git\ncd verl\npip3 install -e .\n\n# Install the latest stable version of vLLM\npip3 install vllm==0.7.3 \n\n# Install flash-attn\npip3 install flash-attn --no-build-isolation\n\n```\n\nNote that if you are installing lower versions of vLLM (0.7.0, 0.7.1, 0.7.2), you need to make some tiny patches manually on vllm (/path/to/site-packages/vllm after installation) after the above steps:\n\n- vllm/distributed/parallel_state.py: Remove the assertion below:\n\n```\nif (world_size\n        != tensor_model_parallel_size * pipeline_model_parallel_size):\n    raise RuntimeError(\n        f\"world_size ({world_size}) is not equal to \"\n        f\"tensor_model_parallel_size ({tensor_model_parallel_size}) x \"\n        f\"pipeline_model_parallel_size ({pipeline_model_parallel_size})\")\n\n```\n\n- vllm/executor/uniproc_executor.py: change `local_rank = rank` to `local_rank = int(os.environ[\"LOCAL_RANK\"])`\n- vllm/model_executor/model_loader/weight_utils.py: remove the `torch.cuda.empty_cache()` in `pt_weights_iterator`\n\n## Features\n\n### Use cuda graph\n\nAfter installation, examples using FSDP as training backends can be used. By default, the `enforce_eager` is set to True, which disables the cuda graph. To enjoy cuda graphs and the sleep mode of vLLM>=0.7, add the following lines to the bash script:\n\n```\nactor_rollout_ref.rollout.enforce_eager=False \\\nactor_rollout_ref.rollout.free_cache_engine=True \\\n\n```\n\nFor a typical job like examples/ppo_trainer/run_qwen2-7b_seq_balance.sh, the rollout generation time is 85 seconds with vLLM0.7.0. By enabling the cudagraph, the generation duration is further reduced to 62 seconds.\n\n**Note:** Currently, if the `n` is greater than 1 in `SamplingParams` in vLLM>=0.7, there is a potential performance issue on the stability of rollout generation time (Some iterations would see generation time bursts) using vLLM's V0 Engine.\n\n### Use vLLM V1 Engine\n\nUsing the vLLM V1 engine can avoid instability issues and achieve additional performance improvements. To use the V1 engine, you can first uninstall the previously installed vLLM and then follow the steps below to install the newer version.\n\n```\ngit clone https://github.com/vllm-project/vllm.git\ncd vllm\ngit checkout 2275784\nsed -i \"903a\\    data_parallel_size = world_size // pipeline_model_parallel_size // tensor_model_parallel_size\" ./vllm/distributed/parallel_state.py\nVLLM_USE_PRECOMPILED=1 pip install --editable .\n```\n\nThen you can enable the V1 engine by setting `export VLLM_USE_V1=1`. In some benchmark tests, the V1 engine demonstrates a 1.5x speed improvement over the vLLM V0 engine.\nThe stable support of the vLLM V1 engine is available on verl main.\n"
  },
  {
    "path": "docs/README_vllm0.8.md",
    "content": "# Upgrading to vLLM >= 0.8\n\nLast updated: 05/04/2025.\n\n## Installation\n\nNote: This version of verl+vLLM 0.8+ supports **FSDP** for training and **vLLM** for rollout.\n\n```bash\n# Create the conda environment\nconda create -n verl python==3.10\nconda activate verl\n\n# Install verl\ngit clone https://github.com/volcengine/verl.git\ncd verl\npip3 install -e .\n\n# Install the latest stable version of vLLM\npip3 install vllm==0.8.3\n\n# Install flash-attn\npip3 install flash-attn --no-build-isolation\n\n```\n\nWe have a pre-built docker image for verl+vLLM 0.8.3. You can direct import it with the following command:\n\n```bash\ndocker pull hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0\n```\n\n## Features\n\nvLLM 0.8+ supports cuda graph and V1 engine by default in verl. To enable these features, remember to add the following lines to the bash script:\n\n```bash\nactor_rollout_ref.rollout.enforce_eager=False \\\nactor_rollout_ref.rollout.free_cache_engine=True \\\n```\n\nand also **remove** the environment variable if it exists:\n\n## Notes\n\nWhen you just directly upgrade vllm>=0.8, some dependency packages may undergo version changes. If you encounter the following problems:\n\n```bash\nin <module> from torch.multiprocessing.reductions import ForkingPickler ImportError: cannot import name 'ForkingPickler' from 'torch.multiprocessing.reductions' (/opt/conda/lib/python3.11/site-packages/torch/multiprocessing/reductions.py)\n```\n\nYou need to upgrade `tensordict` to version 0.6.2 using the command `pip install tensordict==0.6.2`.\n"
  },
  {
    "path": "docs/_static/custom.css",
    "content": "/* Make the documentation use full screen width */\n.wy-nav-content {\n    max-width: none !important;\n    width: 100% !important;\n    padding: 1.618em 3.236em !important;\n}\n\n/* Adjust the content wrapper - will be set by JavaScript */\n.wy-nav-content-wrap {\n    margin-left: 300px;\n    transition: margin-left 0.2s ease;\n    width: auto !important;\n    position: relative !important;\n    background: white !important;\n    min-height: 100vh !important;\n}\n\n/* Make the main content area responsive */\n.rst-content {\n    max-width: none !important;\n    width: 100% !important;\n}\n\n/* Optional: Adjust table widths to prevent overflow */\n.rst-content table.docutils {\n    width: 100% !important;\n    table-layout: auto !important;\n}\n\n/* Optional: Better code block width handling */\n.rst-content .highlight {\n    width: 100% !important;\n}\n\n/* Content area positioning already handled above */\n\n/* Optional: Improve readability with some margin on very wide screens */\n@media (min-width: 1400px) {\n    .wy-nav-content {\n        max-width: none !important;\n        margin: 0 auto !important;\n    }\n}\n\n/* Resizable sidebar styles */\n.wy-nav-side {\n    position: fixed !important;\n    top: 0 !important;\n    bottom: 0 !important;\n    left: 0 !important;\n    width: 300px;\n    min-width: 200px;\n    max-width: 600px;\n    display: flex;\n    flex-direction: column;\n    z-index: 200 !important;\n}\n\n/* Ensure sidebar header (logo, search) adapts to width */\n.wy-side-nav-search {\n    width: 100% !important;\n    box-sizing: border-box !important;\n    padding: 0.809em 0.809em !important;\n}\n\n.wy-side-nav-search input[type=\"text\"] {\n    width: 100% !important;\n    box-sizing: border-box !important;\n}\n\n/* Make logo/title area responsive */\n.wy-side-nav-search > div.version {\n    width: 100% !important;\n}\n\n.wy-side-nav-search > a {\n    width: 100% !important;\n    display: block !important;\n    white-space: nowrap !important;\n    overflow: hidden !important;\n    text-overflow: ellipsis !important;\n}\n\n/* Responsive adjustments for narrow sidebar */\n@media (max-width: 300px) {\n    .wy-side-nav-search > a {\n        font-size: 0.9em !important;\n    }\n    \n    .wy-side-nav-search input[type=\"text\"] {\n        font-size: 0.8em !important;\n    }\n}\n\n/* Ensure search input doesn't overflow */\n.wy-side-nav-search form {\n    width: 100% !important;\n    margin: 0 !important;\n}\n\n/* Make search icon responsive */\n.wy-side-nav-search .wy-dropdown {\n    width: 100% !important;\n}\n\n/* Adjust search results dropdown width */\n.wy-side-nav-search .wy-dropdown-menu {\n    width: 100% !important;\n    max-width: none !important;\n    left: 0 !important;\n    right: 0 !important;\n}\n\n/* Resize handle is created by JavaScript */\n\n/* Make sure the sidebar content doesn't overflow */\n.wy-side-scroll {\n    width: 100% !important;\n    flex: 1 !important;\n    overflow-y: auto !important;\n    overflow-x: hidden !important;\n    padding-right: 10px !important;\n    box-sizing: border-box !important;\n    scroll-behavior: auto !important; /* Prevent smooth scrolling on sidebar itself */\n}\n\n/* Ensure proper scroll behavior for main content area */\nhtml {\n    scroll-behavior: smooth !important;\n}\n\n/* Ensure anchor links work properly in main content */\n.wy-nav-content-wrap {\n    scroll-behavior: smooth !important;\n}\n\n/* Fix scroll to target for anchor links */\n.rst-content {\n    scroll-behavior: smooth !important;\n}\n\n/* Fix anchor scroll offset to account for fixed header */\n.rst-content .section {\n    scroll-margin-top: 60px;\n}\n\n/* Fix anchor scroll offset for headers */\n.rst-content h1, .rst-content h2, .rst-content h3, .rst-content h4, .rst-content h5, .rst-content h6 {\n    scroll-margin-top: 60px;\n}\n\n/* Fix anchor scroll offset for specific scroll targets */\n.rst-content .headerlink {\n    scroll-margin-top: 60px;\n}\n\n/* Fix sidebar navigation styling */\n.wy-menu-vertical {\n    width: 100% !important;\n}\n\n.wy-menu-vertical li {\n    width: 100% !important;\n}\n\n.wy-menu-vertical a {\n    width: 100% !important;\n    word-wrap: break-word !important;\n    white-space: normal !important;\n}\n\n/* Content area margin is handled by JavaScript */\n\n/* Custom drag handle (more visible) */\n.resize-handle {\n    position: absolute;\n    top: 0;\n    right: 0;\n    width: 8px;\n    height: 100%;\n    background: #ccc;\n    cursor: col-resize;\n    z-index: 1001;\n    opacity: 0.3;\n    transition: opacity 0.2s ease;\n}\n\n.resize-handle:hover {\n    opacity: 0.8;\n    background: #999;\n}\n\n.resize-handle::before {\n    content: '';\n    position: absolute;\n    top: 50%;\n    left: 50%;\n    width: 2px;\n    height: 20px;\n    background: #666;\n    transform: translate(-50%, -50%);\n    border-radius: 1px;\n}\n\n.resize-handle:hover::before {\n    background: #333;\n}\n\n/* Ensure smooth resizing */\n.wy-nav-side.resizing {\n    user-select: none;\n    pointer-events: none;\n}\n\n.wy-nav-side.resizing .wy-side-scroll {\n    overflow: hidden;\n}"
  },
  {
    "path": "docs/_static/js/resizable-sidebar.js",
    "content": "// Resizable sidebar functionality\ndocument.addEventListener('DOMContentLoaded', function() {\n    const sidebar = document.querySelector('.wy-nav-side');\n    const content = document.querySelector('.wy-nav-content-wrap');\n    \n    if (!sidebar || !content) return;\n    \n    // Create resize handle\n    const resizeHandle = document.createElement('div');\n    resizeHandle.className = 'resize-handle';\n    sidebar.appendChild(resizeHandle);\n    \n    let isResizing = false;\n    let startX = 0;\n    let startWidth = 0;\n    \n    // Get initial width\n    const getInitialWidth = () => {\n        return 300; // Default width\n    };\n    \n    // Save width to localStorage\n    const saveWidth = (width) => {\n        localStorage.setItem('sidebar-width', width);\n    };\n    \n    // Load width from localStorage\n    const loadWidth = () => {\n        const savedWidth = localStorage.getItem('sidebar-width');\n        if (savedWidth) {\n            const width = parseInt(savedWidth, 10);\n            if (width >= 200 && width <= 600) {\n                return width;\n            }\n        }\n        return getInitialWidth();\n    };\n    \n    // Apply width to sidebar and content\n    const applyWidth = (width) => {\n        // Update sidebar width\n        sidebar.style.width = width + 'px';\n        \n        // Update content margin with !important to override any CSS\n        content.style.setProperty('margin-left', width + 'px', 'important');\n        \n        // Also update any other content wrapper that might exist\n        const contentInner = document.querySelector('.wy-nav-content');\n        if (contentInner) {\n            contentInner.style.setProperty('margin-left', '0px', 'important');\n        }\n        \n        // Force reflow and repaint\n        sidebar.offsetHeight;\n        content.offsetHeight;\n        \n        // Trigger window resize event to notify other components\n        window.dispatchEvent(new Event('resize'));\n    };\n    \n    // Initialize with saved width\n    const initialWidth = loadWidth();\n    applyWidth(initialWidth);\n    \n    // Mouse down on resize handle\n    resizeHandle.addEventListener('mousedown', (e) => {\n        isResizing = true;\n        startX = e.clientX;\n        startWidth = parseInt(window.getComputedStyle(sidebar).width, 10);\n        \n        sidebar.classList.add('resizing');\n        document.body.style.cursor = 'col-resize';\n        document.body.style.userSelect = 'none';\n        \n        // Add overlay to prevent iframe issues\n        const overlay = document.createElement('div');\n        overlay.style.cssText = `\n            position: fixed;\n            top: 0;\n            left: 0;\n            width: 100%;\n            height: 100%;\n            z-index: 9999;\n            cursor: col-resize;\n        `;\n        overlay.id = 'resize-overlay';\n        document.body.appendChild(overlay);\n        \n        e.preventDefault();\n    });\n    \n    // Mouse move\n    document.addEventListener('mousemove', (e) => {\n        if (!isResizing) return;\n        \n        const width = startWidth + e.clientX - startX;\n        const clampedWidth = Math.max(200, Math.min(600, width));\n        applyWidth(clampedWidth);\n    });\n    \n    // Mouse up\n    document.addEventListener('mouseup', () => {\n        if (!isResizing) return;\n        \n        isResizing = false;\n        sidebar.classList.remove('resizing');\n        document.body.style.cursor = '';\n        document.body.style.userSelect = '';\n        \n        // Remove overlay\n        const overlay = document.getElementById('resize-overlay');\n        if (overlay) {\n            overlay.remove();\n        }\n        \n        // Save the current width\n        const currentWidth = parseInt(window.getComputedStyle(sidebar).width, 10);\n        saveWidth(currentWidth);\n    });\n    \n    // Handle window resize - removed to prevent infinite loop\n    // The sidebar width is fixed and managed by drag functionality, no need to recalculate on window resize\n    \n    // Double-click to reset to default width\n    resizeHandle.addEventListener('dblclick', () => {\n        const defaultWidth = 300;\n        applyWidth(defaultWidth);\n        saveWidth(defaultWidth);\n    });\n});\n\n// Fix navigation issues - Using MutationObserver for reliable initialization\ndocument.addEventListener('DOMContentLoaded', function() {\n    let navigationFixed = false;\n    \n    function setupNavigationFix() {\n        if (navigationFixed) return;\n        \n        // Find all links in the sidebar\n        const sidebarLinks = document.querySelectorAll('.wy-menu-vertical a');\n        \n        // Only proceed if we have sidebar links\n        if (sidebarLinks.length === 0) return;\n        \n        console.log('Setting up navigation fix...');\n        \n        sidebarLinks.forEach(function(link) {\n            const href = link.getAttribute('href');\n            \n            // Clone the link to remove all existing event listeners\n            const newLink = link.cloneNode(true);\n            \n            // Add our own click handler\n            newLink.addEventListener('click', function(e) {\n                console.log('Link clicked:', href);\n                \n                // If it's an anchor link within the same page\n                if (href && href.startsWith('#') && href !== '#') {\n                    e.preventDefault();\n                    e.stopPropagation();\n                    \n                    const targetId = href.substring(1);\n                    const targetElement = document.getElementById(targetId);\n                    \n                    if (targetElement) {\n                        // Calculate offset for fixed header\n                        const headerHeight = 60;\n                        const elementPosition = targetElement.getBoundingClientRect().top;\n                        const offsetPosition = elementPosition + window.pageYOffset - headerHeight;\n                        \n                        window.scrollTo({\n                            top: offsetPosition,\n                            behavior: 'smooth'\n                        });\n                        \n                        // Update URL hash\n                        if (history.pushState) {\n                            history.pushState(null, null, '#' + targetId);\n                        } else {\n                            location.hash = '#' + targetId;\n                        }\n                    }\n                }\n                // For external links, navigate normally\n                else if (href && !href.startsWith('#') && !href.startsWith('javascript:')) {\n                    console.log('Navigating to external link:', href);\n                    window.location.href = href;\n                }\n            });\n            \n            // Replace the old link with the new one\n            link.parentNode.replaceChild(newLink, link);\n        });\n        \n        navigationFixed = true;\n        \n        // Handle initial page load with hash\n        if (window.location.hash) {\n            // Use requestAnimationFrame for better timing\n            requestAnimationFrame(() => {\n                const targetId = window.location.hash.substring(1);\n                const targetElement = document.getElementById(targetId);\n                if (targetElement) {\n                    const headerHeight = 60;\n                    const elementPosition = targetElement.getBoundingClientRect().top;\n                    const offsetPosition = elementPosition + window.pageYOffset - headerHeight;\n                    \n                    window.scrollTo({\n                        top: offsetPosition,\n                        behavior: 'smooth'\n                    });\n                }\n            });\n        }\n    }\n    \n    // Try to set up navigation fix immediately\n    setupNavigationFix();\n    \n    // If it didn't work, use MutationObserver to watch for when sidebar links are added\n    if (!navigationFixed) {\n        const observer = new MutationObserver(function(mutations) {\n            mutations.forEach(function(mutation) {\n                if (mutation.type === 'childList' && mutation.addedNodes.length > 0) {\n                    // Check if sidebar links were added\n                    const sidebarLinks = document.querySelectorAll('.wy-menu-vertical a');\n                    if (sidebarLinks.length > 0) {\n                        setupNavigationFix();\n                        if (navigationFixed) {\n                            observer.disconnect();\n                        }\n                    }\n                }\n            });\n        });\n        \n        // Start observing the document for changes\n        observer.observe(document.body, {\n            childList: true,\n            subtree: true\n        });\n        \n        // Fallback timeout in case MutationObserver doesn't work\n        setTimeout(function() {\n            if (!navigationFixed) {\n                setupNavigationFix();\n            }\n            observer.disconnect();\n        }, 5000);\n    }\n});"
  },
  {
    "path": "docs/_static/js/runllm-widget.js",
    "content": "document.addEventListener(\"DOMContentLoaded\", function () {\n    var script = document.createElement(\"script\");\n    script.type = \"module\";\n    script.id = \"runllm-widget-script\";\n    script.src = \"https://widget.runllm.com\";\n    script.setAttribute(\"version\", \"stable\");\n    script.setAttribute(\"crossorigin\", \"true\");\n    script.setAttribute(\"runllm-keyboard-shortcut\", \"Mod+j\");\n    script.setAttribute(\"runllm-name\", \"verl Chatbot\");\n    script.setAttribute(\"runllm-position\", \"TOP_RIGHT\");\n    script.setAttribute(\"runllm-assistant-id\", \"679\");\n    script.async = true;\n    document.head.appendChild(script);\n  });"
  },
  {
    "path": "docs/advance/agent_loop.rst",
    "content": "Agent Loop\n==========\n\nLast updated: 07/17/2025.\n\n.. versionadded:: 0.4.2\n   [status: alpha]\n\n.. warning::\n   Agent Loop is ready for use, but the API may change in future releaes.\n\nAgent Loop is designed as general interface for multi-turn rollout and agentic reinforcement learning.\n\n**Design goal**:\n\n- Plugable user defined agent loop\n- Provide standard request generate api with different inference frameworks\n- Provide request level load balance between multiple inference servers\n\n**Non-goal**:\n\n- How tool is defined and how to call tool\n\nIn high level overview, agent loop is given a prompt, run user defined loop: call LLM generate api, call tools, ...\nand return the final output. The final output is then calculated reward and used as trajectory for RL training.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_overview.svg?raw=true\n\n\nAPI Design\n----------\n\n``AgentLoopBase`` class is the abstraction of agent loop, and ``run`` method is the only interface that user need to implement.\nThe run method, given prompt messages in format: [{\"role\": \"user\"}, {\"content\": \"...\"}], and additional sampling params,\ncould do whatever user wants, such as\n\n- call LLM generate api\n- call tools: web search, database query, code sandbox, ...\n- environment interaction\n- reflection\n- ...\n\n.. code:: python\n\n   class AgentLoopBase(ABC):\n       @abstractmethod\n       async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:\n           \"\"\"Run agent loop to interact with LLM server and environment.\n\n           Args:\n               sampling_params (Dict[str, Any]): LLM sampling params.\n               **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`.\n\n           Returns:\n               AgentLoopOutput: Agent loop output.\n           \"\"\"\n           raise NotImplementedError\n\nAfter running user defined loop, run method should return ``AgentLoopOutput``, including prompt token ids,\nresponse token ids, and response mask.\n\n.. code:: python\n\n   class AgentLoopOutput(BaseModel):\n       \"\"\"Agent loop output.\"\"\"\n\n       prompt_ids: list[int]\n       \"\"\"Prompt token ids.\"\"\"\n       response_ids: list[int]\n       \"\"\"Response token ids including LLM generated token, tool response token.\"\"\"\n       response_mask: list[int]\n       \"\"\"Response mask, 1 for LLM generated token, 0 for tool response token.\"\"\"\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_output.svg?raw=true\n\n.. note:: AgentLoopOutput only output one trajectory for a given prompt, multiple trajectories output is still under discussion.\n\nArchitecture Design\n-------------------\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop_architecture.png?raw=true\n\nA single PPO step contain two phase: rollout and train. In rollout phase:\n\n1. PPOTrainer sample a batch from dataset and call ``AgentLoopManager.generate_sequences``.\n2. AgentLoopManager ``wake_up`` all async LLM server instances, which will sync weights between inference engine(vLLM/SGLang) and training engine(FSDP/Megatron-LM).\n3. AgentLoopManager split batch into chunks and send each chunk to ``AgentLoopWorker``.\n4. AgentLoopWorker receive chunk and for each prompt, spawn a user defined ``AgentLoopBase`` instance, run ``run`` coroutine until end and get ``AgentLoopOutput``.\n\n.. tip::\n   AgentLoopWorker schedules multiple coroutines concurrently. If number of AgentLoopWorker equals batch_size, then each worker is response for one prompt.\n\nIn agent loop, when user need LLM generate response:\n\n5. Call ``AsyncLLMServerManager.generate`` with prompt_ids.\n6. AsyncLLMServerManager select a server instance with least request in first turn and send request to it. (In following turns, the request will be sent to the same server instance).\n7. AsyncLLMServer receive a request, issue ipc/rpc with model_runner, and generate response. (There's slight differences between vLLM and SGLang, see below).\n\nWhen all prompts in all AgentLoopWorker finish, AgentLoopManager gather results and return to PPOTrainer.\n\n8. AgentLoopManager ``sleep`` all server instances, which will free kv cache and offload weights to CPU memory.\n\nAsyncLLMServer\n~~~~~~~~~~~~~~\n\nAsyncLLMServer is the abstraction of LLM server with two types of generation api:\n\n- `OpenAI chat completion <https://platform.openai.com/docs/api-reference/chat>`_: generate response for the given chat conversation.\n- Token in token out: generate response ids for the given token ids.\n\nWe have officially supported vLLM and SGLang AsyncLLMServer, both of them implement the two api and are well tested.\nOther inference engine should be easy to plug-in by implement the ``AsyncServerBase`` class.\n\n.. code:: python\n\n   class AsyncServerBase(ABC):\n       @abstractmethod\n       async def chat_completion(self, raw_request: Request) -> JSONResponse:\n           \"\"\"OpenAI chat completion API.\n\n           Args:\n               raw_request (Request): raw json request\n           \n           Returns:\n               JSONResponse: json response\n\n           API reference: https://platform.openai.com/docs/api-reference/chat/create\n           \"\"\"\n           raise NotImplementedError\n\n       @abstractmethod\n       async def generate(self, prompt_ids: list[int], sampling_params: dict[str, Any], request_id: str) -> list[int]:\n           \"\"\"Generate response ids given prompt ids.\n\n           Args:\n               prompt_ids (List[int]): prompt ids\n               sampling_params (Dict[str, Any]): sampling params\n               request_id (str): request id\n\n           Returns:\n               List[int]: response ids\n           \"\"\"\n           raise NotImplementedError\n\n\nChat completion vs Token in token out\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. warning::\n   The following conclusion is based on our recent experience and is still open to investigation and discussion.\n\nAlmost all agent frameworks (LangGraph, CrewAI, LlamaIndex, etc) call LLM with OpenAI chat completion api, and \nkeep chat history as messages. So user may expect that we should use the chat completion api in multi-turn rollout.\n\nBut based on our recent experience on single-turn training on DAPO and multi-turn training on `retool <https://github.com/verl-project/verl-recipe/tree/main/retool>`_,\nwe found the token_ids from apply the final messages may not equal to the token_ids by concat prompt_ids and response_ids in each turn.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/multi_turn.png?raw=true\n\n**Where does this inconsistency happened?**\n\nFirst, the tool parser may alter the content. For example\n\n.. code:: json\n\n   {\"role\": \"assistant\", \"content\": \"Let me call a <tool_call>...</tool_call> and get the result\"}\n\nAfter tool_calls extraction, the messages is like this:\n\n.. code:: json\n\n   {\"role\": \"assistant\", \"content\": \"Let me call a and get the result\", \"tool_calls\": [{\"name\": \"foo\", \"arguments\": \"{}\"}]}\n\nEncode the extracted message back is not equal to the original LLM generated response_ids.\n\nSecond,  the `decode-encode` may also lead to inconsistency: `Agent-R1 issue#30 <https://github.com/0russwest0/Agent-R1/issues/30#issuecomment-2826155367>`_.\n\n**What is the impact of this inconsistency?**\n\nThis inconsistency is not a big problem for serving/agent system, but is critical to RL training.\nIt causes the trajectory deviate from the policy model distribution. We have observed that apply_chat_template\nto the final chat history messages make PPO training not even converged in single-turn.\n\nvLLM\n^^^^\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/async_vllm.png?raw=true\n\nFor vLLM, the Async LLM Engine is running in same process as the server, and ModelRunner is running in same process as FSDP/Megatron-LM workers.\nAsync LLM Engine communicate with ModelRunner through ZeroMQ. When server receive a request, it directly call engine to generate response_ids.\n\nSGLang\n^^^^^^\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/async_sglang.png?raw=true\n\nFor SGLang, the Async LLM Engine is running in same process as FSDP/Megatron-LM worker-0, and it spawn multiple subprocesses as ModelRunner.\nAlso, Async LLM Engine communicate with ModelRunner through ZeroMQ. When server receive a request, it remote call the worker-0 and get response_ids.\n\nAsyncLLMServerManager\n~~~~~~~~~~~~~~~~~~~~~\n\nAsyncLLMServerManager serve as proxy to multiple AsyncLLMServer instances, provides:\n\n- load balance: select a server instance with least request in first turn and send request to it.\n- sticky session: bind request_id to server instance, so that the same request_id will be sent to the same server instance in following turns.\n\nAsyncLLMServerManager is passed to ``AgentLoopBase.__init__``, whenever user want to interact with LLM in agent loop,\nthey can call ``AsyncLLMServerManager.generate`` to generate response_ids.\n\n.. code:: python\n\n   class AsyncLLMServerManager:\n       async def generate(\n           self,\n           request_id,\n           *,\n           prompt_ids: list[int],\n           sampling_params: dict[str, Any],\n       ) -> list[int]:\n           \"\"\"Generate tokens from prompt ids.\n\n           Args:\n               request_id (str): request id for sticky session.\n               prompt_ids (List[int]): List of prompt token ids.\n               sampling_params (Dict[str, Any]): Sampling parameters for the chat completion.\n\n           Returns:\n               List[int]: List of generated token ids.\n           \"\"\"\n           ...\n\nNext\n----\n\n- :doc:`Agentic RL Training<../start/agentic_rl>`: Quick start agentic RL training with gsm8k dataset.\n- `LangGraph MathExpression <https://github.com/verl-project/verl-recipe/tree/main/langgraph_agent/example>`_: Demonstrate how to use LangGraph to build agent loop.\n- `Retool <https://github.com/verl-project/verl-recipe/tree/main/retool>`_: End-to-end retool paper reproduction using tool agent.\n"
  },
  {
    "path": "docs/advance/async-on-policy-distill.md",
    "content": "# Recipe: Async On-Policy Knowledge Distillation Trainer\n\n**Authors:** Brilliant Hanabi, furunding\n\n**Last updated:** 2025-11-08\n\n## 1. Background\n\nOn-policy knowledge distillation (KD) trains a student policy to imitate a stronger teacher using samples drawn from the student's current policy. For each on-policy rollout the teacher returns soft, top-k token distributions and the student is optimized with a token-wise sparse KL objective that focuses learning on the teacher's high-probability modes. Because training examples come from the student's own state distribution, KD reduces distributional mismatch relative to off-policy distillation or supervised fine-tuning (SFT), improving stability and sample efficiency. Compared with reinforcement learning, KD avoids high-variance reward-based optimization and complex reward design by providing dense, informative per-token targets, which typically yields faster convergence and simpler scaling. Recent empirical and implementation-focused writeups (e.g., [ThinkingMachines' blog on on-policy distillation](https://thinkingmachines.ai/blog/on-policy-distillation/)) also demonstrate that on-policy distillation can deliver high-quality behavior with substantially lower compute and data requirements than many alternative approaches.\n\nBuilt on verl’s Ray-based single-controller components, we initially assembled a strictly on-policy KD pipeline where rollout generation, teacher knowledge acquisition, and policy optimization ran in lockstep. In practice, this synchronous design proved highly inefficient: the three stages had to wait for one another, creating pipeline bubbles and underutilized GPUs. To address this, we extend the asynchronous schedulers introduced by the One-Step-Off Policy pipeline to overlap these phases. This overlap preserves the same distillation objective while trading some strict on-policy guarantees for substantial gains in end-to-end throughput and hardware utilization.\n\n## 2. Distillation Overview and Objective\n\nThis recipe centers on on-policy knowledge distillation: the student policy learns from a stronger teacher on samples generated by the current policy (on-policy). For each input prompt, the student (actor) generates responses; the teacher provides top-k token distributions, and the student is trained to match them token-wise.\n\nCore components:\n\n1. Teacher signal: top-k log-probabilities and token indices per valid token position.\n2. Student objective: sparse, token-level KL divergence between student logits and teacher top-k distribution.\n\nObjective: encourage student probabilities $Q$ to cover teacher modes $P$ using token-wise $\\mathrm{KL}(P\\,\\|\\,Q)$ computed on the teacher's top-k support.\n\n## 3. Efficient System Design\n\n### 3.1 Schedulers (One-Step / Two-Step Off-Policy)\n\nThe native (serial) on-policy distillation process is shown in the figure below.\n\n![Zero-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/zero-step-off-distill.png)\n\nThis recipe supports optional schedulers that overlap generation, teacher querying, and updates to improve throughput without changing the distillation objective.\n\n#### 3.1.1 One-Step-Off-Policy\n\n![One-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one-step-off-distill.png)\n\n- Warm-up: 2 steps.\n- Overlap pattern: rollout while actor update; weight sync while teacher retrieving.\n- Timing keys: `sync_rollout_weights`, `wait_prev_gen`, `wait_prev_teacher`.\n\n#### 3.1.2 Two-Step-Off-Policy\n\n![Two-Step-Off Scheduler](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/two-step-off-distill.png)\n\n- Warm-up: 3 steps.\n- Overlap pattern: rollout, actor update while teacher retrieving; interleave weight sync.\n- Timing keys: `sync_rollout_weights`, `max(wait_prev_gen, wait_prev_prev_teacher)`.\n\nTip: Use `two_step_off` when teacher takes much more time than sync; `one_step_off` for simpler overlapping.\n\nPractical details:\n\n- Inputs per batch: `teacher_topk_logps`, `teacher_topk_indices`, `attention_mask` (to select valid token positions).\n- Loss injection: last pipeline stage computes KL via a logits processor; earlier stages remain unchanged.\n- Optional dynamic micro-batching groups sequences by density to reduce padding overhead.\n\nThe pipeline:\n\n1. Actor parameters are synchronized to a rollout worker group (nccl broadcast) with a little bit latency.\n2. Rollout workers (vLLM-backed) generate sequences asynchronously (`async_generate_sequences`).\n3. Teacher client service (ZeroMQ based) returns top-k log-probabilities + token indices for each sequence (batched micro-requests), enabling KL-based guidance.\n4. Megatron actor performs a KL divergence computation between student logits and teacher top-k distributions (custom TP-aware kernel in `megatron_kl_loss.py`).\n5. Scheduling strategies (`one_step_off_scheduler`, `two_step_off_scheduler`) can overlap phases (optional for throughput):\n\n### 3.2 Weights sync between actor and rollout\n\nWe initially followed the weight synchronization path from the One-Step-Off-Policy recipe (Ray collective broadcast across all actor and rollout ranks, plus Megatron-side allgather of parameter shards). In practice this became the dominant bottleneck, so we made three changes:\n\n1. Batch-and-bulk load on the rollout side: instead of streaming tensors one-by-one (in one-step-off-policy recipe), we stage a bundle of parameter tensors and issue a single batched load into the rollout engine. In our setup this reduced the weight-loading time by roughly 3×.\n2. Batch-and-bulk broadcast between the actor and rollout: instead of streaming tensors one-by-one (in one-step-off-policy recipe), we stage a bundle of parameter tensors and issue a single batched broadcast between the actor and rollout workers.\n3. Replace allgather with gather-to-root in Megatron: parameter shards are gathered to actor rank 0 (rather than allgathered to everyone), and that root then serves as the single source for broadcasting to rollout ranks. On top of the previous change, 2 and 3 changes delivered an additional ~4× speedup in the synchronization phase.\n\n## 4. High-Level Data & Control Flow\n\n```\nDriver (TaskRunner)\n  ├─ Initialize Ray, tokenizer, datasets, worker groups\n  ├─ Build ResourcePoolManager (actor vs rollout GPU layouts)\n  ├─ Trainer.fit()\n      ├─ init_workers(): build actor + rollout groups, broadcast weight metadata, create nccl collective group\n      ├─ continuous_iterator(): epochs → batches\n      ├─ scheduler (see Section 6)\n        • _async_gen_next_batch(): optional weight sync + non-blocking rollout\n        • _async_get_teacher_knowledge(): submit teacher requests, store future\n        ├─ For each step:\n            • Sync rollout weights\n            • Retrieve (batch, gen_output, teacher_output) from futures\n            • Merge gen + teacher outputs → DataProto\n            • Compute metrics (response length stats, timing, throughput)\n            • Update actor (forward_backward_batch + KL loss + optimizer step)\n            • (Optional) save checkpoint\n```\n\n> Note: Schedulers are optional and explained later; the distillation objective is independent of how phases are overlapped.\n\n## 5. Key Components\n\n### 5.1 `OnPolicyDistillTrainer` (`ray_trainer.py`)\n- Creates `GenerationBatchFuture` objects holding rollout and (later) teacher futures.\n- Adds scheduling + teacher integration + modified metric emission (KL, timing, MFU).\n\n### 5.2 Actor Worker (Megatron)\n- `OnPolicyDistillActor.update_policy()` orchestrates micro-batch forward/backward.\n- KL Loss injection via `logits_processor` during forward on pipeline last stage.\n\n### 5.3 Rollout Worker (vLLM / SGLang)\n- Pure inference mode (`init_model` builds model; no optimizer). \n- `async_generate_sequences` returns a Ray future for overlapping.\n\n### 5.4 Teacher Service (`teacher/`)\n- Proxy + worker architecture (ZMQ REQ/REP) for batched top-k retrieval.\n- `TeacherClient.submit()` returns a `Future`; aggregator composes micro-batches.\n- Configurable temperature, max tokens, only-response mode.\n\n### 5.5 KL Loss (`megatron_kl_loss.py`)\n- Performs normalization & stable per-token probability construction across TP shards.\n- Gradient is (student_probs - teacher_sparse_probs) scaled by upstream grad.\n\n## 6. Configuration Highlights (`on_policy_distill_trainer.yaml`)\n\n| Section | Purpose | Notable Keys |\n|---------|---------|-------------|\n| actor_rollout_ref.teacher | Teacher server | server_ip, server_port, n_server_workers |\n| trainer | Global training control | total_epochs, save_freq, scheduler (one_step_off | two_step_off), n_gpus_per_node, nnodes |\n| rollout | Resource split for rollout | n_gpus_per_node, nnodes |\n\n**Remember to set `trainer.n_gpus_per_node`, `trainer.nnodes`, `rollout.n_gpus_per_node` and `rollout.nnodes` to allocate GPU resources.**\n\n### Dynamic Batch Size\n\nEnable by: \n\n```\nactor_rollout_ref.actor.use_dynamic_bsz=True\nactor_rollout_ref.actor.max_token_len=6000  # cap post-group token length\n```\n\nImproves utilization under variable sequence lengths.\n\n### Resource Guidelines\n\n- Actor pool: `trainer.nnodes * trainer.n_gpus_per_node` GPUs.\n- Rollout pool: `rollout.nnodes * rollout.n_gpus_per_node` GPUs.\n- Ensure teacher server capacity ≈ `n_server_workers` to avoid stalls (monitor `wait_prev_teacher`).\n\n## 7. Usage Examples\n\n### 7.1 Launch Teacher Server\n\nBefore training process, you should have a teacher server to provide logp information.\n\nWe provide a toy teacher server example with vLLM. It needs `telnet` to check proxy status, and `python` command to run. So if you have not installed `telnet`, you can just delete these code in `start_server.sh`. And some OS use `python3` rather than `python`, so you also need to modify it. Also you can change the port of teacher if you meet port conflict.\n\nThere are 3 arguments can be set for vllm backend `--tp-size`, `--n-logprobs` and `--ckpt-path` in `start_server.sh` / `worker.py`. You should set before you start server.\n\nWe also provide a toy multi-node teacher server. You can start the main node using `start_server.sh` and start the slave nodes using `join_server.sh`. Still remember to set args in `join_server.sh`, especially the `$PROXY_IP` and `$PROXY_BACKEND_PORT` of main node.\n\nWhen training, student will automatically use the teacher's topk (n-logprobs) to set its own topk argument at line 83 of `recipe/gkd/megatron_kl_loss.py`, so you don't need to set student's topk argument.\n\n```bash\ncd recipe/gkd/teacher\nbash start_server.sh\n# Exports ports and launches proxy + worker (default vLLM backend)\n```\n\nVerify with:\n\n```bash\ntelnet localhost 15555\n```\n\n### 7.2 Minimal Local (Megatron + vLLM) Run\n\n```bash\npython3 -m recipe.gkd.main_gkd \\\n  --config-path=recipe/gkd/config \\\n  --config-name=on_policy_distill_trainer \\\n  actor_rollout_ref.model.path=/path/to/MODEL \\\n  data.train_files=/path/to/train.parquet \\\n  trainer.total_epochs=2 \\\n  trainer.n_gpus_per_node=4 rollout.n_gpus_per_node=2 \\\n  actor_rollout_ref.teacher.server_ip=127.0.0.1 \\\n  actor_rollout_ref.teacher.server_port=15555 \\\n  trainer.scheduler=one_step_off\n```\n\n(Requires a running teacher server).\n\n### 7.3 Ray Job Submission (Distilled 16B Example)\n\nSee `run_moonlight_dsv3_training.sh` for a full script including:\n\n- Dist ckpt path setup (`dist_checkpointing_path`)\n- Expert parallel sizing (EP / ETP)\n- Dynamic batch sizing\n- Two-step-off scheduling for deeper overlap.\n\nSubmit (after adjusting paths):\n\n```bash\nbash recipe/gkd/run_moonlight_dsv3_training.sh\n```\n\n## 8. Metrics & Monitoring\n\nEmitted metrics include (prefixes may vary):\n\n- Timing: `timing/wait_prev_gen`, `timing/sync_rollout_weights`, `timing/get_teacher_knowledge`, `timing/update_actor`.\n- Sequence stats: `response_seq_len/*` (avg, max, min, counts).\n- Performance: `perf/mfu/actor`, `perf/max_memory_allocated_gb`, `perf/cpu_memory_used_gb`.\n- Distillation: `actor/kl_loss`, `actor/grad_norm`, `actor/lr`.\n\nInterpretation Tips:\n\n- High `wait_prev_teacher` → scale `n_server_workers` and allocate more teacher GPUs or reduce per-request batch size, or just use `two_step_off`.\n- High `wait_prev_gen` with uniform lengths → allocate more rollout GPUs.\n- High `sync_rollout_weights` → check NCCL env / network congestion and try to modify `actor_rollout_ref.rollout.update_weights_bucket_megabytes`.\n\n## 9. Extensibility Notes\n\n- Add new schedulers by following interface returning `(epoch, batch, gen_output, teacher_output, timing_dict)`.\n- Integrate different distillation signals (e.g., hidden states, intermediate reasoning tokens) by extending `teacher_utils.get_teacher_knowledge` and modifying `logits_processor`.\n\n## 10. Functional Support Summary\n\n| Category | Supported |\n|----------|-----------|\n| Train engine | Megatron |\n| Rollout engine | vLLM |\n| Distillation signal | Teacher top-k logprobs & indices |\n| Scheduling | one_step_off, two_step_off |\n\n## 11. Quick Checklist Before Running\n\n- Teacher server reachable (`telnet <ip> <port>`).\n- `actor_rollout_ref.model.path` contains the correct Megatron/HF config artifacts.\n- `train_files` points to a parquet dataset compatible with this recipe's dataset loader.\n- NCCL environment vars set (see `config/runtime_env.yaml`).\n\n---\nFeel free to open issues or PRs to extend scheduler variants, add new distillation objectives, or broaden engine support, and more improvement.\n"
  },
  {
    "path": "docs/advance/attention_implementation.rst",
    "content": ".. _attention-implementation-override:\n\nAttention Implementation Override\n==================================\n\nLast updated: 10/31/2025.\n\nBy default, VERL's FSDP workers use ``flash_attention_2`` as the attention implementation for improved performance. \nHowever, you can now override this setting to use different attention implementations based on your needs.\n\nSupported Attention Implementations\n-----------------------------------\n\nThe following attention implementations are supported (subject to model and hardware compatibility):\n\n- ``flash_attention_2``: High-performance attention implementation (default)\n- ``eager``: Standard PyTorch attention implementation\n- ``sdpa``: Scaled Dot-Product Attention (PyTorch native)\n\nWhen to Override\n----------------\n\nYou might want to override the attention implementation in the following scenarios:\n\n- **Debugging**: Use ``eager`` for easier debugging and better error messages\n- **Compatibility**: Some models or hardware configurations may not support ``flash_attention_2``\n- **Memory constraints**: Different implementations have different memory characteristics\n- **Performance tuning**: Testing different implementations for optimal performance\n\nConfiguration Examples\n-----------------------\n\nPPO Training with Eager Attention\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nTo override the attention implementation for the actor, rollout, and reference models:\n\n.. code:: bash\n\n    python3 ppo_trainer.py \\\n        +actor_rollout_ref.model.override_config.attn_implementation=eager \\\n        [other parameters...]\n\nPPO Training with SDPA Attention\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: bash\n\n    python3 ppo_trainer.py \\\n        +actor_rollout_ref.model.override_config.attn_implementation=sdpa \\\n        [other parameters...]\n\nCritic Model Override\n~~~~~~~~~~~~~~~~~~~~~\n\nFor training configurations that include a critic model, you can also override its attention implementation:\n\n.. code:: bash\n\n    python3 ppo_trainer.py \\\n        +actor_rollout_ref.model.override_config.attn_implementation=eager \\\n        +critic.model.override_config.attn_implementation=eager \\\n        [other parameters...]\n\nYAML Configuration\n~~~~~~~~~~~~~~~~~~\n\nYou can also specify the attention implementation in your YAML configuration file:\n\n.. code:: yaml\n\n    actor_rollout_ref:\n      model:\n        override_config:\n          attn_implementation: eager\n          # other overrides...\n\n    critic:  # if using a critic model\n      model:\n        override_config:\n          attn_implementation: eager\n          # other overrides...\n\nImportant Notes\n---------------\n\n**Backward Compatibility**: If you don't specify ``attn_implementation`` in the override config, \nVERL will continue to use ``flash_attention_2`` by default, ensuring backward compatibility with existing configurations.\n\n**Model Support**: Not all models support all attention implementations. Ensure your model is compatible \nwith the chosen attention implementation before training.\n\n**Performance Impact**: Different attention implementations have varying performance characteristics. \n``flash_attention_2`` typically offers the best performance, while ``eager`` provides better debugging capabilities.\n\n**Hardware Dependencies**: Some attention implementations (like ``flash_attention_2``) may require \nspecific hardware or CUDA versions. If you encounter compatibility issues, try using ``eager`` or ``sdpa``.\n\nTroubleshooting\n---------------\n\nIf you encounter errors when using a specific attention implementation:\n\n1. **Check model compatibility**: Verify that your model supports the chosen attention implementation\n2. **Try eager attention**: Use ``attn_implementation=eager`` as a fallback for debugging\n3. **Check hardware requirements**: Ensure your hardware supports the attention implementation\n4. **Review error messages**: Attention implementation errors often provide clear guidance on supported options\n\nExample Error Resolution\n~~~~~~~~~~~~~~~~~~~~~~~~\n\nIf you see an error like \"flash_attention_2 is not supported\", you can resolve it by switching to eager attention:\n\n.. code:: bash\n\n    # Instead of the default flash_attention_2\n    python3 ppo_trainer.py +actor_rollout_ref.model.override_config.attn_implementation=eager\n\nThis override ensures your training can proceed while you investigate the flash attention compatibility issue.\n"
  },
  {
    "path": "docs/advance/checkpoint.rst",
    "content": ".. _checkpoint-page:\n\nUsing Checkpoints to Support Fault Tolerance Training\n=====================================================\n\nLast updated: 06/25/2025.\n\nThere could be training errors or machine failure during the whole RLHF training process, \nso it is recommended to enable checkpoints to minimize your loss.\n\nThe API Interface has already been listed in :ref:`config-explain-page`,\nand we will not repeat them. But there are still some technique details\nwe hope to clarify.\n\n.. note:: \n\n    Notice that the ``checkpoint.contents`` field has no effect to FSDP checkpoint except ``hf_model``, \n    the other 3 fields are binded together to save and load. We recommend to include ``model``, ``optimizer`` and ``extra`` all.\n\nCheckpoint Saving Directory Structure\n-------------------------------------\n\nCommonly, we use the ``default_local_dir`` declared in ``ppo_trainer.yaml`` or ``ppo_megatron_trainer.yml``\nto work as preffix when saving checkpoints, which is ``checkpoints/${trainer.project_name}/${trainer.experiment_name}``.\n\nSo the inner checkpoint structure of **FSDP** is like:\n\n.. code::\n\n    checkpoints/${trainer.project_name}/${trainer.experiment_name}\n    ├── global_steps_${i}\n    │   ├── actor\n    │   │   ├── huggingface      # default save config and tokenizer, save huggingface model if include ``hf_model`` in checkpoint.contents\n    │   │   └── fsdp_config.json # FSDP config file, including world_size and fsdp version\n    │   │   ├── model_world_size_{self.world_size}_rank_{self.rank}.pt\n    │   │   ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt\n    │   │   └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt\n    │   ├── critic\n    │   │   ├── huggingface\n    │   │   └── fsdp_config.json\n    │   │   ├── model_world_size_{self.world_size}_rank_{self.rank}.pt\n    │   │   ├── optim_world_size_{self.world_size}_rank_{self.rank}.pt\n    │   │   └── extra_state_world_size_{self.world_size}_rank_{self.rank}.pt\n    └── latest_checkpointed_iteration.txt\n\nAll model shards, optimizers and extra states are stored together, in a sharded and distributed way.\n\nWhile **Megatron** current checkpoint structure is:\n\n.. code::\n\n    checkpoints/${trainer.project_name}/${trainer.experiment_name}\n    ├── global_steps_${i}\n    │   ├── actor\n    │   │   ├── huggingface     # default save config and tokenizer, save huggingface model if include ``hf_mode`` in checkpoint.contents\n    │   │   └── dist_ckpt       # save sharded model/optimizer/rng_states, naming the same as Megatron\n    │   └── critic\n    │   │   ├── huggingface\n    │   │   └── dist_ckpt\n    └── latest_checkpointed_iteration.txt\n\nConvert FSDP and Megatron Checkpoints to HuggingFace Format Model\n-----------------------------------------------------------------\n\nWe provide a tool to convert the FSDP and Megatron checkpoints to HuggingFace format model.\nThe tool is located in ``verl/model_merger``. For older versions of verl that don't include fsdp_config.json in checkpoints, you can use the legacy model merger located at ``verl/scripts/legacy_model_merger.py``.\n\nThe script supports two main sub-commands: `merge` (to convert and save checkpoints) and `test` (to validate merged checkpoints against a reference model).\nThe arguments for the `merge` sub-command are as follows:\n\n.. code:: bash\n\n    usage: python -m verl.model_merger merge [-h] --backend {fsdp,megatron} [--local_dir LOCAL_DIR] [--tie-word-embedding] [--is-value-model] [--use_cpu_initialization] [--target_dir TARGET_DIR]\n                         [--hf_upload_path HF_UPLOAD_PATH] [--private]\n\n    options:\n    -h, --help            show this help message and exit\n    --backend {fsdp,megatron}\n                            The backend of the model\n    --local_dir LOCAL_DIR\n                            Path to the saved model checkpoints\n    --tie-word-embedding  Whether to tie word embedding weights (currently only Megatron supported)\n    --is-value-model      Whether the model is a value model (currently only Megatron supported)\n    --use_cpu_initialization\n                            Whether to use CPU initialization for the model. This is useful for large models that cannot fit into GPU memory during initialization.\n    --target_dir TARGET_DIR\n                            Directory to save the merged huggingface model\n    --hf_upload_path HF_UPLOAD_PATH\n                            Hugging Face repository ID to upload the model\n    --private             Whether to upload the model to a private Hugging Face repository\n\nExample usage for merging Megatron checkpoints:\n\n.. code:: bash\n\n    python -m verl.model_merger merge \\\n        --backend megatron \\\n        --tie-word-embedding \\\n        --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \\\n        --target_dir /path/to/merged_hf_model\n\nExample usage for distributed merging Megatron checkpoints:\n\n.. code:: bash\n\n    torchrun --nproc_per_node 1 --nnodes 8 --node_rank ${RANK} -m verl.model_merger merge \\\n        --backend megatron \\\n        --tie-word-embedding \\\n        --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \\\n        --target_dir /path/to/merged_hf_model\n\nExample usage for merging FSDP checkpoints:\n\n.. code:: bash\n\n    python -m verl.model_merger merge \\\n        --backend fsdp \\\n        --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \\\n        --target_dir /path/to/merged_hf_model\n\n\nMegatron Merger details\n-----------------------\n\nCurrent implement of decoder layers uses ``nn.ModuleList`` to store the layers, \nand thus the model layers on every PP rank and VPP rank starts their index from 0.\n\nThere are 3 ways to correct this behavior:\n\n1. Modify the decoder layer's state_dict, add ``offset`` to each layer's index, thus rewrite ``nn.ModuleList`` implementation.\n2. Modify the layer index when saving checkpoint and recover them when loading checkpoint.\n3. The Checkpoint merger do this work, calculate the actual ``offset`` from ``state_dict`` only, a little complex.\n\nCurrent implementation use solution 2.\n\n\nHuggingFace to Megatron DistCheckpoint details\n----------------------------------------------\n\nThrough ``mbridge``, we can directly save the mcore model to huggingface format during training.\nNo need to convert the model to Megatron dist-checkpoint format.\n\nOriginal Checkpoint Utils\n-------------------------\n\nOriginal Checkpoint Utils refer to original checkpoint implementation in ``verl/models/[model]/megatron/checkpoint_utils``.\n\nWe only need ``[model]_loader.py`` in original checkpoint utils now, since we get rid of storing ``hf_model`` every time (which is not recommended for large model training, try only saving sharded models if you can).\n\n.. note:: \n\n    Note that ``[model]_loader`` only support environments where **storage clusters are able to connect with every calculation nodes**. \n    Because it utilizes **sharded load way to minimize the loading checkpoint overhead**. \n    Every rank loads its own data from ``state_dict`` which can be accessed by all of them.\n    While there is also no need to broadcast among DP ranks, since the saved state_dict is only produced by DP rank 0.\n\n    For users who can **only place the huggingface model on one device**, we keep the original costly implementation in ``[model]_loader_deprecated``. In this implementation, rank 0 broadcast all weights to each tp and pp rank, and then dp rank 0 broadcast to all dp ranks. There may be at risks of OOM.\n\n    To use deprecated loader, change the import package of ``load_state_dict_to_megatron_llama``.\n"
  },
  {
    "path": "docs/advance/dpo_extension.rst",
    "content": "Extend to other RL(HF) algorithms\n=================================\n\nLast updated: 02/25/2025.\n\nWe already implemented the complete training pipeline of the PPO\nalgorithms. To extend to other algorithms, we analyze the high-level\nprinciple to use verl and provide a tutorial to implement the DPO\nalgorithm. Users can follow the similar paradigm to extend to other RL algorithms.\n\n.. note:: **Key ideas**: Single process drives multi-process computation and data communication.\n\nOverall Approach\n----------------\n\nStep 1: Consider what multi-machine multi-GPU computations are needed\nfor each model, such as ``generate_sequence`` , ``compute_log_prob`` and\n``update_policy`` in the actor_rollout model. Implement distributed\nsingle-process-multiple-data (SPMD) computation and encapsulate them\ninto APIs\n\nStep 2: Based on different distributed scenarios, including FSDP and 3D\nparallelism in Megatron-LM, implement single-process control of data\ninteraction among multi-process computations.\n\nStep 3: Utilize the encapsulated APIs to implement the control flow\n\nExample: Online DPO\n-------------------\n\nWe use verl to implement a simple online DPO algorithm. The algorithm\nflow of Online DPO is as follows:\n\n1. There is a prompt (rollout) generator which has the same weight as\n   the actor model. After a batch of prompts are fed into the generator,\n   it generates N responses for each prompt.\n2. Send all the prompts + responses to a verifier for scoring, which can\n   be reward model or a rule-based function. Then sort them in pairs to\n   form a training batch.\n3. Use this training batch to train the actor model using DPO. During\n   the process, a reference policy is needed.\n\nStep 1: What are the multi-machine multi-GPU computations\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n**Sample Generator**\n\nImplementation details:\n\n.. code:: python\n\n   from verl.single_controller.base import Worker\n   from verl.single_controller.ray import RayWorkerGroup, RayClassWithInitArgs, RayResourcePool\n   import ray\n\n   @ray.remote\n   class SampleGenerator(Worker):\n       def __init__(self, config):\n           super().__init__()\n           self.config = config\n           \n       def generate_sequences(self, data):\n           pass\n\nHere, ``SampleGenerator`` can be viewed as a multi-process pulled up by\n``torchrun``, with each process running the same code (SPMD).\n``SampleGenerator`` needs to implement a ``generate_sequences`` API for\nthe control flow to call. The implementation details inside can use any\ninference engine including vllm, sglang and huggingface. Users can\nlargely reuse the code in\nverl/verl/workers/rollout/vllm_rollout/vllm_rollout.py and we won't\ngo into details here.\n\n**ReferencePolicy inference**\n\nAPI: compute reference log probability\n\n.. code:: python\n\n   from verl.single_controller.base import Worker\n   import ray\n\n   @ray.remote\n   class ReferencePolicy(Worker):\n       def __init__(self):\n           super().__init__()\n           self.model = Model()\n           \n       def infer(self, data):\n           return self.model(data)\n\n**Actor update**\n\nAPI: Update actor model parameters\n\n.. code:: python\n\n   from verl.single_controller.base import Worker\n   import ray\n\n   @ray.remote\n   class DPOActor(Worker):\n       def __init__(self):\n           super().__init__()\n           self.model = Model()\n           self.model = FSDP(self.model)  # or other distributed strategy\n           self.optimizer = optim.Adam(self.model.parameters(), lr=1e-3)\n           self.loss_fn = xxx\n           \n       def update(self, data):\n           self.optimizer.zero_grad()\n           logits = self.model(data)\n           loss = self.loss_fn(logits)\n           loss.backward()\n           self.optimizer.step()\n\n**Notes: How to distinguish between control processes and distributed computation processes**\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n- Control processes are generally functions directly decorated with\n  ``@ray.remote``\n- Computation processes are all wrapped into a ``RayWorkerGroup``.\n\nUsers can reuse most of the distribtued computation logics implemented\nin PPO algorithm, including FSDP and Megatron-LM backend in\nverl/verl/trainer/ppo.\n\nStep 2: Based on different distributed scenarios, implement single-process control of multi-process data interaction\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n**The core problem to solve here is how a single process sends data to\nmultiple processes, drives multi-process computation, and how the\ncontrol process obtains the results of multi-process computation.**\nFirst, we initialize the multi-process ``WorkerGroup`` in the control\nprocess.\n\n.. code:: python\n\n   @ray.remote(num_cpus=1)\n   def main_task(config):\n       # construct SampleGenerator\n       resource_pool = RayResourcePool(process_on_nodes=[8] * 2)  # 16 GPUs\n       ray_cls = RayClassWithInitArgs(SampleGenerator, config=config)\n       # put SampleGenerator onto resource pool\n       worker_group = RayWorkerGroup(resource_pool, ray_cls)\n       \n       # construct reference policy\n\nAs we can see, in the control process, multiple processes are wrapped\ninto a ``RayWorkerGroup``. Inside this ``WorkerGroup``, there is a\n``self._workers`` member, where each worker is a RayActor\n(https://docs.ray.io/en/latest/ray-core/actors.html) of SampleGenerator.\nray_trainer.md also provide an implementation of\n``MegatronRayWorkerGroup``.\n\nAssuming the model is distributed using FSDP, and there is a batch of\ndata on the control process, for data parallelism, the underlying\ncalling process is:\n\n.. code:: python\n\n   data = xxx\n   data_list = data.chunk(dp_size)\n\n   output = []\n   for d in data_list:\n       # worker_group._workers[i] is a SampleGenerator\n       output.append(worker_group._workers[i].generate_sequences.remote(d))\n\n   output = ray.get(output)\n   output = torch.cat(output)\n\nSingle process calling multiple processes involves the following 3\nsteps:\n\n1. Split the data into DP parts on the control process.\n2. Send the data to remote, call the remote computation through RPC, and\n   utilize multi-process computation.\n3. Obtain the computation results of each worker on the control process\n   and merge them.\n\nFrequently calling these 3 steps on the controller process greatly hurts\ncode readability. **In verl, we have abstracted and encapsulated these 3\nsteps, so that the worker's method + dispatch + collect can be\nregistered into the worker_group**\n\n.. code:: python\n\n   from verl.single_controller.base.decorator import register\n\n   def dispatch_data(worker_group, data):\n       return data.chunk(worker_group.world_size)\n       \n   def collect_data(worker_group, data):\n       return torch.cat(data)\n\n   dispatch_mode = {\n       'dispatch_fn': dispatch_data,\n       'collect_fn': collect_data\n   }\n\n   @register(dispatch_mode=dispatch_mode)\n   def generate_sequences(self, data):\n       pass\n\nIn this way, we can directly call the method inside the worker through\nthe ``worker_group`` on the control (driver) process (which is a single\nprocess):\n\n.. code:: python\n\n   output = worker_group.generate_sequences(data)\n\nThis single line includes data splitting, data distribution and\ncomputation, and data collection.\n\nFurthermore, the model parallelism size of each model is usually fixed,\nincluding dp, tp, pp. So for these common distributed scenarios, we have\npre-implemented specific dispatch and collect methods,in `decorator.py <https://github.com/volcengine/verl/blob/main/verl/single_controller/base/decorator.py>`_, which can be directly used to wrap the computations.\n\n.. code:: python\n\n   from verl.single_controller.base.decorator import register, Dispatch\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def generate_sequences(self, data: DataProto) -> DataProto:\n       pass\n\nHere it requires the data interface to be ``DataProto``. Definition of\n``DataProto`` is in `protocol.py <https://github.com/volcengine/verl/blob/main/verl/protocol.py>`_.\n\nStep 3: Main training loop\n~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nWith the above training flows, we can implement the algorithm's control\nflow. It is recommended that ``main_task`` is also a ray remote process.\n\n.. code:: python\n\n   @ray.remote(num_cpus=1)\n   def main_task(config):\n       # construct SampleGenerator\n       resource_pool = RayResourcePool(process_on_nodes=[8] * 2)  # 16 GPUs\n       ray_cls = RayClassWithInitArgs(SampleGenerator, config=config) \n       # put SampleGenerator onto resource pool\n       sample_gen = RayWorkerGroup(resource_pool, ray_cls)\n       \n       # construct reference policy\n       ray_cls = RayClassWithInitArgs(ReferencePolicy)\n       ref_policy = RayWorkerGroup(resource_pool, ray_cls)\n       \n       # construct actor\n       ray_cls = RayClassWithInitArgs(DPOActor)  \n       dpo_policy = RayWorkerGroup(resource_pool, ray_cls)\n       \n       dataloader = DataLoader()\n       \n       for data in dataloader:\n           # generate data\n           data = sample_gen.generate_sequences(data)\n           # generate scores for each data \n           data = generate_scores(data)\n           # generate pairwise data using scores\n           data = generate_pairwise_data(data)\n           # generate ref_log_prob\n           data.batch['ref_log_prob'] = ref_policy.infer(data)\n           # update using dpo\n           dpo_policy.update(data)\n           # logging\n\nHere, different ``WorkerGroups`` can be placed in the same resource pool or\nin different resource pools using ``create_colocated_worker_cls``\nsimilar as in `ray_trainer.py <https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/ray_trainer.py>`_.\n"
  },
  {
    "path": "docs/advance/fp8.md",
    "content": "# FP8 RL in verl\n\nLast updated: 03/05/2026\n\nverl supports two FP8 modes for accelerating RL training:\n\n| Mode | Training Precision | Rollout Precision |\n|------|-------------------|-------------------|\n| **FP8 Rollout Only** | BF16 | FP8 |\n| **FP8 End-to-End** | FP8 (Megatron) | FP8 (vLLM) |\n\n> [!TIP]\n> For ready-to-run scripts, see the [low-precision recipe directory](https://github.com/verl-project/verl-recipe/low_precision).\n\n---\n\n## FP8 Rollout Only\n\nFP8 rollout-only mode keeps training in BF16 and quantizes rollout inference to FP8. This reduces GPU memory during generation and speeds up rollout without affecting training precision.\n\n### Implementation\n\nWe monkey patch several vLLM functions to enable FP8 rollout for reinforcement learning:\n\n1. **Quantize weights**: Quantize model weights on-the-fly from higher-precision formats to FP8.\n2. **Process weights after loading**: For vLLM, we replace the `vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.process_weights_after_loading` function to handle weight processing after quantization. For SGLang, this patch is not needed as it natively supports loading quantized weights.\n\n### Support Matrix\n\n- FP8 blockwise quantization for rollout\n  - Used in Deepseek, which is 1x128 quantization for activations and 128x128 quantization for model weights\n- Dense models and MoE models\n- Async rollout interfaces\n- vLLM 0.10.x & vLLM 0.11 & vLLM 0.12 & SGLang 0.5.5\n- FSDP and Megatron training backends\n\n### Usage\n\nEnable in config file:\n\n```yaml\nrollout:\n  quantization: \"fp8\"\n```\n\nOr via command line:\n\n```bash\nactor_rollout_ref.rollout.quantization=fp8\n```\n\n### Experiments and Outcomes\n\n#### Qwen3-8B-Base Dense Model\n\n**Configuration**\n- DAPO recipe. AIME24 online validation.\n- vLLM(FP8 spmd rollout) + FSDP\n  - Note that SPMD rollout has been deprecated, so we removed the FP8 SPMD rollout.\n- Prompt batch size 32, n=16.\n- Rollout batch size: 32\\*3*16\n- Train_batch_size & ppo_mini_batch_size 32\n- Max response length 20K\n- Token-level TIS, C=2\n- 8*H100\n- vLLM 0.10.0+CUDA 12.6 vs vLLM 0.11.0+CUDA 12.9\n\n**Accuracy**\n![Qwen3-8b-base_fp8_acc](\nhttps://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-8b-base_fp8_acc.png?raw=true)\n*dark green: BF16, orange: FP8 rollout + token-level TIS, light green: FP8 rollout without TIS*\n\nResults and observations:\n- With TIS, FP8 rollout aligns with BF16\n- Obvious accuracy drop when TIS is not enabled\n- Higher mismatch kl but within acceptable range throughout the training\n\n\n**Performance**\n\n![Qwen3-8b-base_fp8_rollout_perf](\nhttps://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-8b-base_fp8_rollout_perf.png?raw=true)\n*green: BF16, orange: FP8 rollout + CUDA12.6 + DeepGemm, purple: FP8 rollout + CUDA 12.9 + DeepGemm*\n\nResults and observations:\n- FP8 rollout leads to around ~12% rollout speedup with CUDA 12.6 + DeepGemm\n- When upgrading to CUDA 12.9, speedup can be up to ~18%\n\n#### Qwen3-30B-A3B-Base MoE Model\n\n**Configuration**\n- DAPO recipe. AIME24 online validation.\n- FP8 async rollout, vLLM+FSDP\n- Prompt batch size 32\n- Rollout batch size: 32\\*3*16\n- Train_batch_size & ppo_mini_batch_size 32\n- Max response length 20K\n- Token-level TIS, C=2\n- 2\\*8*H100\n- vLLM 0.10.0+CUDA 12.6\n\n**Accuracy**\n![Qwen3-30b-a3b_fp8_acc](\nhttps://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-30b-a3b_fp8_acc.png?raw=true)\n*grey: BF16 + token-level TIS, red: FP8 rollout + token-level TIS*\n\nResults and observations:\n- Rollout & training distribution mismatch is in general higher for MoE\n- Rollout correction required even for BF16\n- FP8 rollout with token-level TIS aligns with BF16\n\n\n**Performance**\n\n![Qwen3-30b-a3b_fp8_perf](\nhttps://github.com/Agoniii/verl/blob/xueh/fp8_pr_images/docs/advance/images/Qwen3-30b-a3b_fp8_perf.png?raw=true)\n*grey: BF16 + token-level TIS, red: FP8 rollout + token-level TIS​*\n\nResults and observations:\n- FP8 rollout : over 35% rollout speedup\n- Expecting more perf gain with CUDA 12.9\n\n---\n\n## FP8 End-to-End (Training + Rollout)\n\nFP8 E2E applies FP8 to the entire RL pipeline: forward/backward passes via Transformer Engine, FP8 optimizer states, and FP8 rollout inference via vLLM. This maximizes memory savings and throughput.\n\n### Requirements\n\n- **CUDA 12.9+** (required for block-wise FP8 scaling)\n- **Transformer Engine** with block-wise FP8 support\n- Environment variable: `NVTE_FP8_BLOCK_SCALING_FP32_SCALES=1`\n\n### Key Configuration\n\n```yaml\n# FP8 training via Transformer Engine\nactor_rollout_ref.actor.megatron.override_transformer_config:\n  fp8: \"hybrid\"              # FP8 forward + backward; also supports \"e4m3\"\n  fp8_recipe: \"blockwise\"    # block-wise scaling\n\n# FP8 optimizer\nactor_rollout_ref.actor.optim.override_optimizer_config:\n  fp8_recipe: \"blockwise\"\n\n# FP8 rollout inference (vLLM)\nactor_rollout_ref.rollout:\n  quantization: fp8\n```\n\n### Support Matrix\n\n- Megatron training backend (via Megatron-Bridge)\n- Verified on Qwen3-30B-A3B and Qwen3-8B\n- Block-wise FP8 scaling (`fp8_recipe: \"blockwise\"`)\n\n### Experiments and Results\n\n#### Qwen3-30B-A3B MoE Model\n\n**Configuration**\n- DAPO recipe. AIME24 online validation.\n- Megatron + Megatron-Bridge, FP8 async rollout with vLLM\n- MoE router in BF16 for both vLLM and Megatron-Core\n- Prompt batch size 128, n=16\n- Max response length 20K\n- Token-level TIS, C=2\n- 2\\*8*H100, CUDA 12.9\n\n![Qwen3-30b-a3b_fp8_e2e](https://github.com/user-attachments/assets/70fb1396-ec73-40d7-9a43-1d48553c0ad9)\n*Orange: BF16, Green: FP8 E2E, Red: FP8 rollout + BF16 training*\n\nResults and observations:\n- FP8 E2E achieves comparable accuracy to the BF16 baseline, with the two curves closely aligned throughout training.\n- The training/inference precision mismatch (measured by KL divergence) follows the ordering: FP8 rollout-only > FP8 E2E > BF16 E2E. This is expected, as FP8 E2E maintains consistent precision across both training and inference, resulting in lower distribution mismatch than the FP8 rollout-only setting where training remains in BF16.\n\n---\n\n## Citation\n\nFor more extensive experiments, ablation studies, and analysis on FP8 reinforcement learning, please refer to our technical report:\n\n```bibtex\n@article{qiu2026fp8rl,\n  title={FP8-RL: A Practical and Stable Low-Precision Stack for LLM Reinforcement Learning},\n  author={Qiu, Zhaopeng and Yu, Shuang and Zhang, Jingqi and Zhang, Shuai and Huang, Xue and Yang, Jingyi and Lai, Junjie},\n  journal={arXiv preprint arXiv:2601.18150},\n  year={2026},\n  url={https://arxiv.org/abs/2601.18150}\n}\n```\n"
  },
  {
    "path": "docs/advance/fsdp_extension.rst",
    "content": "\nAdd models with the FSDP backend\n==================================\n\nLast updated: 02/09/2025.\n\nModel\n--------------------------\n\nIn principle, our FSDP backend can support any HF model and we can\nsychronoize the actor model weight with vLLM using `hf_weight_loader.py` under `third_party/vllm`.\nHowever, ``hf_weight_loader`` is will gather the full state_dict of a\nmodel during synchronization, which may cause OOM. We suggest using\n``dtensor_weight_loader`` which gather the full model parameter layer by\nlayer to reduce the peak memory usage. We already support dtensor weight\nloader for the models below in `dtensor_weight_loader.py` under `third_party/vllm`:\n\n- ``GPT2LMHeadModel``\n- ``LlamaForCausalLM``\n- ``LLaMAForCausalLM``\n- ``MistralForCausalLM``\n- ``InternLMForCausalLM``\n- ``AquilaModel``\n- ``AquilaForCausalLM``\n- ``Phi3ForCausalLM``\n- ``GemmaForCausalLM``\n- ``Gemma2ForCausalLM``\n- ``GPTBigCodeForCausalLM``\n- ``Starcoder2ForCausalLM``\n- ``Qwen2ForCausalLM``\n- ``DeepseekV2ForCausalLM``\n\nTo implement ``dtensor_weight_loader`` of a model that's supported in\nvLLM, follow the guide of gemma model below:\n\n1. Copy the\n   ``load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]])`` from the vllm model class\n   to ``dtensor_weight_loaders.py``\n2. Modify the arguments to\n   ``(actor_weights: Dict, vllm_model: nn.Module)``\n3. Replace the ``self`` to ``vllm_model``\n4. Add the\n   ``local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight)``\n   before each ``param = params_dict[name]`` and modify the following\n   weight loading using ``local_loaded_weight``.\n5. Register the implemented dtensor weight loader to ``__MODEL_DTENSOR_WEIGHT_LOADER_REGISTRY__``.\n\n.. code-block:: diff\n\n    - def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]):\n    + def gemma_dtensor_weight_loader(actor_weights: Dict, vllm_model: nn.Module) -> nn.Module:\n        stacked_params_mapping = [\n            # (param_name, shard_name, shard_id)\n            (\"qkv_proj\", \"q_proj\", \"q\"),\n            (\"qkv_proj\", \"k_proj\", \"k\"),\n            (\"qkv_proj\", \"v_proj\", \"v\"),\n            (\"gate_up_proj\", \"gate_proj\", 0),\n            (\"gate_up_proj\", \"up_proj\", 1),\n        ]\n    -   params_dict = dict(self.named_parameters())\n    +   params_dict = dict(vllm_model.named_parameters())\n        loaded_params = set()\n    -   for name, loaded_weight in weights:\n    +   for name, loaded_weight in actor_weights.items():\n            for (param_name, shard_name, shard_id) in stacked_params_mapping:\n                if shard_name not in name:\n                    continue\n                name = name.replace(shard_name, param_name)\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n    +           local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight)\n                param = params_dict[name]\n                weight_loader = param.weight_loader\n    -           weight_loader(param, loaded_weight, shard_id)\n    +           weight_loader(param, local_loaded_weight.to(dtype=param.dtype), shard_id)\n                break\n            else:\n                # lm_head is not used in vllm as it is tied with embed_token.\n                # To prevent errors, skip loading lm_head.weight.\n                if \"lm_head.weight\" in name:\n                    continue\n                # Skip loading extra bias for GPTQ models.\n                if name.endswith(\".bias\") and name not in params_dict:\n                    continue\n    +           local_loaded_weight = redistribute_dtensor(param_name=name, loaded_weights=loaded_weight)\n                param = params_dict[name]\n                weight_loader = getattr(param, \"weight_loader\",\n                                        default_weight_loader)\n    -           weight_loader(param, loaded_weight)\n    +           weight_loader(param, local_loaded_weight.to(dtype=param.dtype))\n            loaded_params.add(name)\n        unloaded_params = params_dict.keys() - loaded_params\n        if unloaded_params:\n            raise RuntimeError(\n                \"Some weights are not initialized from checkpoints: \"\n                f\"{unloaded_params}\")"
  },
  {
    "path": "docs/advance/fully_async.md",
    "content": "# Recipe: Fully Async Policy Trainer\n\n**Author:** `https://github.com/meituan-search`\n\nLast updated: 02/05/2026.\n\nThis document introduces a fully asynchronous PPO training system that completely decouples the Trainer and Rollouter,\nsupporting asynchronous sample generation and training.\nUnder this system, we achieved a 2.35x-2.67x performance improvement when training the Qwen2.5-7B model with 128 GPUs,\nwithout significantly affecting the results.\n\n## Introduction\n\n### Background\n\nThe separated rollout and train architecture, compared to the colocate architecture, can allocate resources more\nflexibly and design more flexible training logic, thereby addressing issues such as low GPU utilization and training\nefficiency caused by long-tail problems.\nThe one_step_off_policy alleviates the problem of long rollout times and achieves some gains in training efficiency by\ndesigning a separated architecture and performing asynchronous training between rollout and train for one round.\nHowever, it forcibly uses data from one round of asynchronous training, which is not flexible enough and cannot\ncompletely eliminate the impact of long-tail on training efficiency.\nIn other frameworks such as AReaL, Magistral, StreamRL, and AsyncFlow, asynchronous training and streaming training have\nbeen implemented based on the separated architecture and have achieved gains.\nWe borrow from their methods and implemented them in VERL. The fully_async_policy supports asynchronous, streaming, and\npartial\nrollout training.\nBy reasonably setting parameters such as resource allocation and parameter synchronization frequency, fully_async_policy\ncan significantly improve training efficiency.\n\n> Magistral https://arxiv.org/abs/2506.10910\n>\n> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language\n> Reasoning https://arxiv.org/abs/2505.24298\n>\n> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream\n> Generation https://arxiv.org/abs/2504.15930\n>\n> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663\n>\n\n### Core Contributions\n\n* **Resource Isolation**: Unlike using hybrid_engine, Rollouter and Trainer use separate computing resources and need to\n  specify the resources they occupy separately.\n* **Parallel Generation and Training**: While the Trainer is training, the Rollouter is generating new samples.\n* **Multi-step Asynchronous**: Compared to one step off policy, it supports asynchronous settings from 0.x steps to\n  multiple steps, making the asynchronous solution more flexible.\n* **NCCL Parameter Synchronization**: Based on the nccl communication primitive, refer\n  to [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine) to\n  achieve efficient parameter synchronization between Rollouter and Trainer.\n* **Stream Inference and Training**: Rollouter generates data sample by sample, and data transmission uses a single\n  sample as the minimum transmission unit.\n* **Asynchronous Training and Freshness Control**: By setting the parameter async_training.staleness_threshold, it\n  supports training with samples generated by old parameters.\n* **PartialRollout**: The Rollouter's inference process supports partial rollout logic. During parameter\n  synchronization, by adding `sleep() and resume()` logic, it\n  saves samples from ongoing rollouts and continues using them in the next rollout, reducing the time spent waiting for\n  ongoing tasks to finish during parameter synchronization.\n\nCurrently, the supported usage mode is megatron/fsdp+vllm. vllm must use the server mode based on AgentLoop.\n\n## Design\n\nThe overall architecture of fully_async_policy is shown in the figure below. fully_async_policy mainly consists of four\nparts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer.\n\n![fully_async_policy_structure](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)\n\n1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the\n   production speed controlled by freshness.\n2. MessageQueue is used to temporarily store samples generated by Rollouter.\n3. Trainer fetches samples from MessageQueue sample by sample. After fetching `require_batches*ppo_mini_batch_size`\n   samples, it will perform training. After training for async_training.trigger_parameter_sync_step rounds, it triggers\n   a parameter synchronization with Rollouter.\n4. ParameterSynchronizer implements the NCCL synchronous parameter synchronization capability.\n\nThe source of benefits compared to the base scheme lies in the fact that in the colocate case, using more resources for\nrollout cannot solve the idleness caused by long-tail samples.\nAfter we perform resource isolation, the time for rollout and train may be longer than before (because fewer resources\nare used),\nbut the overlap in their time consumption reduces the end-to-end time consumption.\n\n![fully_async_policy_revenue](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)\n\n## Usage\n\n### Parameter Description\n\n| super params                                                     | implication                                                                                    |\n|------------------------------------------------------------------|------------------------------------------------------------------------------------------------|\n| `trainer.nnodes`                                                 | Number of nodes for Trainer                                                                    |\n| `trainer.n_gpus_per_node`                                        | Number of GPUs per node for Trainer                                                            |\n| `rollout.nnodes`                                                 | Number of nodes for Rollouter                                                                  |\n| `rollout.n_gpus_per_node`                                        | Number of GPUs per node for Rollouter                                                          |\n| `data.train_batch_size`                                          | In the fully async strategy, this value is not effective (default is 0)                        |\n| `data.gen_batch_size`                                            | In the fully async strategy, uses streaming sample production logic (default is 1)             |\n| `rollout.total_rollout_steps`                                    | Total number of rollout samples                                                                |\n| `rollout.test_freq`                                              | How many times Rollouter updates parameters before performing a validation                     |\n| `actor_rollout_ref.actor.ppo_mini_batch_size`                    | The ppo_mini_batch_size is a global num across all workers/gpus                                |\n| `actor_rollout_ref.actor.use_rollout_log_probs=True`             | Use log_probs generated by rollout                                                             |\n| `algorithm.rollout_correction.bypass_mode`                       | Whether to compute log_prob using the training model's parameters during the training phase.   |\n| `async_training.require_batches`                                 | Number of ppo_mini_batch_size that FullyAsyncTrainer fetches at once                           |\n| `async_training.trigger_parameter_sync_step`                     | Indicates how many local updates FullyAsyncTrainer performs before a parameter synchronization |\n| `async_training.staleness_threshold`                             | Freshness control                                                                              |\n| `async_training.partial_rollout`                                 | Whether to perform partial_rollout                                                             |\n| `async_training.use_trainer_do_validate`                         | Whether use trainer node to do validate process, default `False`                               |\n\n**Further Explanation:**\n\n* `rollout.total_rollout_steps`\n\n  Compared to colocate, the quantity can be aligned by multiplying train_batch_size and step:\n  `rollout.total_rollout_steps = data.train_batch_size * step`.\n\n* `async_training.trigger_parameter_sync_step`\n\n  In the fully async strategy, it indicates how many local updates the Trainer performs (i.e., how many times it fetches\n  `require_batches * ppo_mini_batch_size` samples) before a parameter synchronization with Rollouter.\n  Between every two parameter synchronizations between Rollouter and Trainer, the Trainer will process\n  `trigger_parameter_sync_step* require_batches*ppo_mini_batch_size` samples.\n  To fairly compare speed with colocate, `trigger_parameter_sync_step` should be set to\n  `data.train_batch_size / (require_batches * ppo_mini_batch_size)`.\n\n* `async_training.staleness_threshold`\n\n  In the fully async strategy, it indicates the maximum proportion of stale samples allowed to be used.\n\n    * `staleness_threshold`=0, indicates synchronous training.\n      Rollouter will generate a fixed number of samples between two parameter updates, the sample count is:\n\n      `rollout_num = (trigger_parameter_sync_step*require_batches*ppo_mini_batch_size)`\n    * `staleness_threshold`>0, indicates asynchronous training, can be set to a decimal for more flexible asynchronous\n      calls.\n      Rollouter will generate at most the following number of samples between two parameter updates:\n\n      `rollout_num = (1+staleness_threshold)*(trigger_parameter_sync_step*require_batches*ppo_mini_batch_size) - num_staleness_sample`\n\n  `num_staleness_sample` represents the number of stale samples generated in excess during the last rollout.\n\n  Since it's a streaming system, rollout continues to generate and trainer continues to consume. If rollouter is slower,\n  trainer will trigger parameter synchronization earlier, and rollouter will not actually produce rollout_num samples.\n  When rollout is fast enough, setting `staleness_threshold` to 1 is basically equivalent to one_step_off policy.\n  To avoid too many expired samples affecting training accuracy, it is recommended to set this value to less than 1.\n\n* `async_training.partial_rollout`\n\n  partial_rollout only actually takes effect when staleness_threshold>0.\n\n* `async_training.require_batches`\n\n  In streaming training, require_batches should be set to 1, indicating that training is performed after producing\n  enough ppo_mini_batch_size samples.\n  In actual testing, we found that if fewer samples are issued at once, due to the order of data distribution, it can\n  cause training instability and longer response lengths.\n  Here, we additionally provide require_batches for streaming distribution and control the number of samples\n  participating in training at once.\n\n* `actor_rollout_ref.actor.use_rollout_log_probs=True`\n\n  In reinforcement learning algorithms, log_probs have implicit correlations with parameter versions and tokens. Due to\n  the settings of algorithms like PPO/GRPO/DAPO, when calculating importance sampling,\n  old_log_prob must use the log_probs corresponding to the rollout parameters and tokens to ensure algorithm\n  correctness. In the fully\n  async strategy, we default to old_log_prob being calculated by rollout rather than by trainer.\n\n* `algorithm.rollout_correction.bypass_mode`\n\n  > algorithm.rollout_correction.bypass_mode default is True, using rollout log prob.\n\n  During the training process, we observed that metrics and response lengths may become unstable in the later\n  stages of training. To mitigate this issue, we can use\n  the [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html)\n  technique for importance sampling. To utilize Rollout Importance Sampling, we need to compute log_prob using\n  the training engine, which requires enabling this switch.\n  Additionally, when `algorithm.rollout_correction.bypass_mode=False` and Rollout Importance Sampling are enabled under\n  mode d\n  (async stream pipeline with partial rollout), our implementation approximates `Areal's Decoupled PPO`.\n\n* `async_training.use_trainer_do_validate`\n\n  It controls whether to use the trainer's `do_validate` method for validation.\n  If set to True, the trainer will perform validation after each parameter update. It can reduce the validation time\n  overhead and trainer node idle time.\n  If set to False, the trainer will not perform validation.\n\n### Supported Modes\n\n1. on policy pipeline:\n    1. **trigger_parameter_sync_step=1, staleness_threshold=0**\n    2. Rollouter produces `require_batches*ppo_mini_batch_size` samples at once, Trainer fetches these samples for\n       training, and after training completes, Trainer and Rollouter perform a parameter synchronization;\n    3. During the rollout phase, if there are long-tail samples but few rollout samples, shorter samples cannot fill\n       idle resources, causing some resource waste.\n    4. As shown in figure a;\n\n2. stream off policy pipeline:\n    1. **trigger_parameter_sync_step>1, staleness_threshold=0**\n    2. Synchronous streaming training will be performed. Rollouter produces\n       `require_batches*ppo_mini_batch_size*trigger_parameter_sync_step` samples at once, Trainer performs a local\n       training every time it fetches `require_batches*ppo_mini_batch_size` samples, and after training\n       trigger_parameter_sync_step times, Trainer and Rollouter perform a parameter synchronization;\n    3. Compared to a, since more samples are generated at once, resource idleness will be lower.\n    4. In one step training, there will be two periods of resource idleness: when fetching the first batch of samples,\n       train waits for `require_batches*ppo_mini_batch_size` samples to be produced, and during the last parameter\n       update, rollout waits for training to complete.\n    5. As shown in figure b;\n\n3. async stream pipeline with stale samples:\n    1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=False**\n    2. After each parameter update, Rollouter will plan to produce at most rollout_num samples (in practice, the number\n       of samples generated may be less than this value depending on rollout speed).\n    3. If the rollout process is relatively fast, Rollouter will generate some additional samples num_stale_samples\n       before parameter synchronization for immediate use by Trainer after synchronization.\n       When triggering parameter synchronization, if Rollouter has ongoing tasks, it will wait for the tasks to complete\n       and not add new tasks;\n    4. Compared to b, except for the first step training, subsequent training will not have the time to wait for the\n       first batch rollout to finish, but will have the time to wait for active tasks to finish.\n    5. As shown in figure c;\n\n4. async stream pipeline with partial rollout:\n    1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=True**\n    2. Compared to c, when triggering parameter synchronization, if Rollouter has samples being produced, it will\n       interrupt the rollout process and perform parameter synchronization. The interrupted samples will continue to be\n       generated after synchronization. This reduces the time to wait for active tasks to finish.\n    3. As shown in figure d;\n\n![fully_async_policy_mode](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)\n\n### Key Metrics\n\n| metrics                                        | implication                                                                                            |\n|------------------------------------------------|--------------------------------------------------------------------------------------------------------|\n| `trainer/idle_ratio`                           | Trainer idle rate                                                                                      |\n| `rollouter/idle_ratio`                         | Rollouter idle rate                                                                                    |\n| `fully_async/count/stale_samples_processed`    | Total number of old samples used in training                                                           |\n| `fully_async/count/stale_trajectory_processed` | Total number of old trajectories used in training (one sample produces rollout.n trajectories)         |\n| `fully_async/partial/total_partial_num`        | Number of partial samples processed by Trainer between two trigger_parameter_sync_step                 |\n| `fully_async/partial/partial_ratio`            | Ratio of partial samples processed by Trainer between two trigger_parameter_sync_step                  |\n| `fully_async/partial/max_partial_span`         | Maximum parameter span of partial samples processed by Trainer between two trigger_parameter_sync_step |\n\n### Parameter Tuning Recommendations\n\n* Resource Allocation and Adjustment:\n    * Reasonable resource allocation is the prerequisite for achieving good training efficiency. The ideal resource\n      allocation should make the rollout time and train time close, thereby minimizing pipeline bubbles in the entire\n      training process,\n      avoiding resource idleness, and ensuring Trainer does not use old samples. In real training scenarios, resource\n      allocation can be adjusted based on the idle time of rollout and train during actual training,\n      which can be obtained from rollouter/idle_ratio and trainer/idle_ratio. If rollouter/idle_ratio is high and\n      trainer/idle_ratio is low,\n      Trainer resources should be increased and Rollouter resources should be reduced, and vice versa.\n\n* Key Parameters:\n    * staleness_threshold: Setting it too high will cause more old samples to be used, affecting model performance. It\n      is recommended to set it to less than 1.\n    * require_batches: The closer to 1, the closer to a pure streaming process, the smaller the training bubbles, and\n      the faster the acceleration effect that can be achieved in terms of speed, but it will affect the order of sample\n      processing;\n    * trigger_parameter_sync_step: The smaller the setting, the closer to on policy, but it will cause frequent\n      parameter synchronization. Long-tail samples waste resources that cannot be filled by short samples, resulting in\n      low resource utilization.\n      The larger the setting, the higher the computational efficiency, but the accuracy will be affected by off policy.\n    * rollout.test_freq: It will occupy Rollouter resources and is not recommended to be set too small.\n\n* Mode Selection: By adjusting different parameters, the Fully Async architecture supports optimization acceleration at\n  different levels, suitable for tasks in different scenarios.\n    * For small-scale tasks that need to ensure training stability and on-policy nature, and have low speed\n      requirements, the on policy pipeline mode (Mode 1) can be tried.\n    * For scenarios that need to improve training throughput but are sensitive to staleness, the stream off policy\n      pipeline mode can be tried. That is, by\n      setting trigger_parameter_sync_step>1 to improve training efficiency, but still maintaining the synchronization\n      mechanism (staleness_threshold=0) (Mode 2).\n    * For large-scale tasks with high training speed requirements and can tolerate a certain degree of off-policy and\n      staleness, setting staleness_threshold>\n      0 and partial_rollout=True can improve training efficiency, using the async stream pipeline mode (Mode 3 or 4).\n\n### Quick Start\n\n```shell\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=10\nstaleness_threshold=0\ntrigger_parameter_sync_step=16\npartial_rollout=False\n\n\npython -m recipe.fully_async_policy.fully_async_main \\\n\ttrain_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    rollout.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n```\n\n## Experiments\n\n### Asynchronous Training on 7B Model\n\nWe used Qwen2.5-Math-7B to verify the benefits of the fully async strategy under long candidates and multiple resources.\nUsing the `async stream pipeline with stale samples` strategy, we achieved about 2x performance improvement on 32 cards,\n64 cards, and 128 cards without significantly affecting experimental results.\n\n* Machine: H20\n* Model: Qwen2.5-Math-7B\n* Rollout length: max_response_length FSDP2: 28K tokens;\n* Algorithm: DAPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 20\n\n* colocate sync:\n    * step: 400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|  training mode   \t   | resource allocation \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t  | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1          \t      |\n|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:---------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:|\n| colocate sync      \t | 32                  \t | 790.10 \t | 357.41 \t | 107.71       \t | 269.80        \t | 13h 44m                \t | 1d 3h 43m              \t | 2d 9h 22m              \t | 3d 17h 5m              \t | max: 0.3313<br>last: 0.2448  \t  |\n| fully_async_policy \t | 16:16               \t |  294.77  |  21.26   | \\            \t |     313.81      |    7h 58m<br>(1.72x)     |    16h 21m<br>(1.70x)    |   1d 0h 53m<br>(2.31x)   |   1d 9h 26m<br>(2.66x)   | max: 0.3302<br>last: 0.2333   \t |\n| colocate sync      \t | 64                  \t | 365.28 \t | 150.72 \t | 70.26        \t | 133.41       \t  | 10h 22m                \t | 20h 45m                \t | 1d 7h 6m               \t | 1d 17h 32m             \t | max: 0.3365<br>last:  0.2333 \t  |\n| fully_async_policy \t | 32:32               \t | 189.26 \t | 28.46  \t | \\            \t | 156.98       \t  | 4h 57m<br>(2.09x)      \t | 10h 14m<br>(2.03x)     \t | 16h 58m<br>(1.83x)     \t | 21h 40m<br>(1.92x)     \t | max: 0.3677<br>last: 0.3406  \t  |\n| colocate sync      \t | 128                 \t | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t  | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t  |\n| fully_async_policy \t | 64:64               \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t  | 3h 13m<br>(2.67x)      \t | 6h 46m<br>(2.65x)      \t | 10h 53m<br>(2.67x)     \t | 17h 22m<br>(2.35x)     \t | max: 0.3521<br>last: 0.3094  \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg\n\n### 128-card 7B Asynchronous Mode Experiment\n\nWe used Qwen2.5-Math-7B to verify the effects of various modes supported by fully async.\nWe can see that the benefit brought by streaming is approximately 1.6x, and after combining staleness and\npartial_rollout, the benefit reaches 2.35x.\n\n|                             mode                                         \t                              | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1         \t      |\n|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:|\n|                                          colocate sync      \t                                           | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t |\n| `stream off policy pipeline`<br>(+fully async: trigger_parameter_sync_step= 4,<br>require_batches= 4) \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t  |\n|          `async stream pipeline with stale samples`<br>(+staleness_threshold=0.5)            \t          |    \t     |    \t     |       \t        |       \t        |            \t             |            \t             |            \t             |            \t             |               \t                |\n|        `async stream pipeline with partial rollout`<br>(+partial_rollout=True)                 \t        | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg\n\n### 128-card Stale Ablation Experiment\n\nUnder the `async stream pipeline with partial rollout` mode, we verified the impact of staleness settings on training\nefficiency.\nWe found that the larger the staleness, the more obvious the final gains.\nWe also noticed that the times for staleness values of 0.3 and 0.5 are quite close, because as the training steps\nincrease, the response length changes significantly, causing training instability.\nFurther analysis and optimization are needed for this issue.\n\n| staleness_threshold \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |     acc/mean@1         \t      |\n|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 0                   \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t |\n| 0.1                 \t | 171.30 \t | 58.17  \t | \\            \t | 109.12       \t | 3h 53m                 \t | 8h 37m                 \t | 14h 25m                \t | 19h 59m                \t | max: 0.3542<br>last: 0.2979 \t |\n| 0.3                 \t | 146.11 \t | 38.88  \t | \\            \t | 103.22       \t | 3h 18m                 \t | 6h 49m                 \t | 11h 40m                \t | 17h 20m                \t | max: 0.3469<br>last: 0.2865 \t |\n| 0.5                 \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg\n\n### 128-card 7B require_batches Ablation Experiment\n\nIn multiple tests, we found that the number of samples issued each time in streaming affects the response length during\ntraining, which in turn affects training time. We verified the impact on results by modifying\n`async_training.require_batches`.\n\n| require_batches \t | step  \t  | gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t |     acc/mean@1         \t      |\n|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 1               \t | 203.47 \t | 30.88 \t | \\            \t | 181.08       \t | 3h 31m                 \t | 8h 29m                 \t | 17h 36m                \t | max: 0.349<br>last: 0.326   \t |\n| 2               \t | 158.72 \t | 26.32 \t | \\            \t | 128.08       \t | 3h 35m                 \t | 7h 38m                 \t | 13h 57m                \t | max: 0.351<br>last: 0.3406  \t |\n| 4               \t | 124.64 \t | 25.62 \t | \\            \t | 95.06        \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | max: 0.3521<br>last: 0.3521 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg\n\n### 30B Model Mode Experiment\n\nWe achieved a 1.7x performance improvement with `async stream pipeline with staleness samples` strategy on the\nQwen3-30B-A3B-Base model compared to the colocate setup. It is worth noting that this is far from the upper limit of\nperformance gains achievable through asynchrony. Firstly, the comparative experiments used a maximum response length of\nonly 8k, which is much shorter than the 20k sequence length in previous experiments, resulting in a less pronounced\nrollout tail effect. Secondly, we adopted a highly skewed resource allocation, with rollout using 96 GPUs and trainer\nusing 32 GPUs, which is not an optimal configuration. During the experiments, we observed that the current verl\nimplementation imposes certain constraints, such as requiring data to be evenly divisible by the number of GPUs, making\nresource adjustment less flexible. Additionally, as asynchronous training and deployment accelerate, the performance gap\nis gradually narrowing. Therefore, enabling more flexible resource allocation and dynamic resource adjustment in the\nfuture will be our next focus.\n\n* Machine: H20\n* Model: Qwen3-30B-A3B-Base\n* Rollout length: max_response_length : 8K tokens;\n* Algorithm: GRPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+Megatron\n* rollout.n: 16\n* ppo_mini_batch_size: 128\n* test_freq: 20\n\n* colocate sync:\n    * step:400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * trigger_parameter_sync_step: 512/128 = 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n| Training Mode      | Resource Allocation | Step   | Gen    | Old Log Prob | Ref   | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1                  |\n|--------------------|---------------------|--------|--------|--------------|-------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------|\n| Colocate Sync      | 128                 | 497.89 | 348.05 | 28.73        | 20.86 | 86.27        | 13h 36m             | 1d 3h 48m           | 1d 19h 4m           | 2d 11h 39m          | max: 0.3500<br>last: 0.3208 |\n| Fully Async Policy | 96:32               | 282.75 | 22.06  | \\            | 50.05 | 206.63       | 6h 45m (2.01x)      | 14h 48m (1.88x)     | 1d 0h 9m (1.78x)    | 1d 10h 41m (1.72x)  | max: 0.3813<br>last: 0.3448 |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg           | | |\n\n### checkpoint-engine Ablation Experiment\nWe tested the single-step parameter synchronization time of the checkpoint-engine on three models: Qwen2.5-Math-7B, Qwen3-30B-A3B, and Qwen3-235B-A22B, using default checkpoint-engine configurations. All experiments were performed on H20 machines, and the Megatron engine was used for training.\n\n|      model      | trainer rank \t | rollout rank\t | checkpoint-engine \t | total sync time \t |\n|:---------------:|:--------------:|:-------------:|:-------------------:|:-----------------:|\n| Qwen2.5-Math-7B |       4        |       4       |        False        |       0.12s       |\n| Qwen2.5-Math-7B |       4        |       4       |        True         |       0.02s       |\n|  Qwen3-30B-A3B  |       16       |      16       |        False        |      15.76s       |\n|  Qwen3-30B-A3B  |       16       |      16       |        True         |       4.38s       |\n| Qwen3-235B-A22B |       64       |      64       |        False        |      58.57s       |\n| Qwen3-235B-A22B |       64       |      64       |        True         |      23.70s       |\n\n\n### use_trainer_do_validate Experiment\nWe tested the effect of setting `use_trainer_do_validate=True` on the training process. The results show that setting\nthis parameter to True can reduce the validation time overhead and trainer node idle time.\nWe used Qwen2.5-Math-7B to verify the benefits of `use_trainer_do_validate=True` on the training process, we achieved about 2x performance improvement on validation time, and the trainer node idle time is reduced by about 40%.\n\n* Machine: H20\n* Model: Qwen2.5-Math-7B\n* Rollout length: max_response_length FSDP2: 10K tokens;\n* Algorithm: DAPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 10\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|   training mode    | resource allocation |  step   |   gen   | old_log_prob | update_actor | validate time | total time<br>50 step | acc/mean@2 |\n|:------------------:|:-------------------:|:-------:|:-------:|:------------:|:------------:|:-------------:|:---------------------:|:----------:|\n|   colocate sync    |         16          | 484.623 | 52.939\t |      0\t      |   430.263    |  205.080  \t   |        7h9m  \t        |    22.6    |\n| fully_async_policy |         8:8         | 489.953 | 52.622\t |      0\t      |   435.874    |   95.699  \t   |        7h2m  \t        |    21.0    |\n\n\n## Multi-Turn Tool Calling\n\nReferencing **recipe/retool** and **ToolAgentLoop**, we implemented **AsyncPartialToolAgentLoop**, a multi-turn\ntool-calling loop that supports partial_rollout for **fully_async_policy**.\n\n### Core Design\n\n`AsyncPartialToolAgentLoop` inherits from `ToolAgentLoop` and is adapted for the asynchronous training mode of\n`fully_async_policy`. When `partial_rollout=True`, the Rollouter interrupts ongoing generation tasks before\nsynchronizing parameters with the Trainer. `AsyncPartialToolAgentLoop` is capable of:\n\n1. **Interrupting Tasks**: Responding to an interrupt signal to save the current state. Currently, interruptions occur\n   during the `GENERATING` process or after other states have completed.\n2. **Resuming Tasks**: Resuming execution from the saved state after parameter synchronization is complete, rather than\n   starting over.\n\n### How to Use\n\nRL training with multi-turn tool calling in `fully_async_policy` is similar to `recipe/retool`. It is enabled by\nspecifying `multi_turn` configurations in the config file.\n\n1. **SFT Stage**: First, the model should undergo SFT to learn how to follow tool-calling format instructions.\n2. **Multi-turn Configuration**: In the `fully_async_policy` training configuration, set the following parameters:\n   ```yaml\n   actor_rollout_ref:\n     rollout:\n       multi_turn:\n         enable: True # AsyncPartialToolAgentLoop will be used by default in fully_async_policy mode\n         # Other multi_turn related configurations\n   ```\n3. **Async Parameters**: To improve efficiency, enable `partial_rollout` and `staleness_threshold` when using multi-turn\n   tool calling:\n   ```yaml\n   async_training:\n     partial_rollout: True\n     staleness_threshold: 0.5\n     # Other async parameters\n   ```\n4. **Example**: See `recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`.\n\n### Experimental Results\n\nTo validate the performance of `fully_async_policy` on multi-turn tool-calling tasks, we compared it with the standard\n`colocate` synchronous mode. Key parameter settings are as follows.\n\n* **SFT Model**: Based on `Qwen2.5-7B-Instruct`, trained for 6 epochs on the `ReTool-SFT` dataset\n* **RL Algorithm**: DAPO\n* **Dataset**:\n    * Train: `DAPO-Math-17k`\n    * Test: `aime_2025`\n* **Resource and Mode Comparison**:\n    * `colocate sync`: 32 H20 gpus\n    * `fully_async_policy`: 16 gpus for Trainer + 16 gpus for Rollouter\n* **Key Configurations**:\n    1. **Tool Calling Configuration**:\n        * `multi_turn.enable: True`\n        * `multi_turn.max_user_turns: 16`\n        * `multi_turn.max_assistant_turns: 16`\n        * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml`\n    2. **`colocate sync` Configuration**:\n        * `ppo_mini_batch_size: 16`\n        * `train_batch_size: 64`\n    3. **`fully_async_policy` Configuration**:\n        * `ppo_mini_batch_size: 16`\n        * `trigger_parameter_sync_step: 4`\n        * `require_batches: 1`\n        * `staleness_threshold: 1`\n        * `partial_rollout: True`\n\n|  training mode   \t   | Resource allocation \t |  step  \t  |  gen   \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t |   aime_2025<br>acc/mean@30  \t   |\n|:--------------------:|:---------------------:|:---------:|:---------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:-------------------------------:|\n| colocate           \t | 32                  \t | 375.47  \t | 228.03 \t  | 35.19        \t | 111.84       \t | 9h 46m                 \t | 22h 28m                \t | start:0.1078<br>last:0.2056   \t |\n| fully_async_policy \t | 16: 16              \t | 221.36 \t  | 40.59   \t | \\            \t | 179.58       \t | 6h 19m<br>(1.55x)      \t | 14h 4m<br>(1.60x)      \t |   start:0.11<br>last:0.2044 \t   |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg\n\n## Future Plans\n\n* GRPO experiments\n* Megatron adaptation\n* SGLang integration\n* Transfer queue integration\n* Asynchronous parameter synchronization\n* AReaL asynchronous algorithm implementation\n* TPPO algorithm implementation\n* Multi-turn and Tool support\n"
  },
  {
    "path": "docs/advance/grafana_prometheus.md",
    "content": "# Use Prometheus and Grafana to Monitor Rollout\n\n**Author:** `https://github.com/meituan-search`\n\nLast updated: 12/05/2025.\n\nMonitor the rollout computation process using Prometheus and Grafana when using verl to enhance system observability and facilitate further performance optimization.\n\nWe provide an additional training monitoring capability, leveraging Prometheus and Grafana to display rollout information during training and enhance system observability to facilitate further performance optimization.\n\nThe system automatically configures Prometheus to scrape metrics from rollout servers, eliminating manual configuration steps.\n\n## Overview\n\nThe figures below show the performance of Qwen235B on the AIME2024 dataset with a response length of 20k, where the emergence of a long-tail problem is clearly observable.\n\n![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana_validate.png?raw=true)\n\nThe following figure presents the fully asynchronous training of the Qwen235B model. Here, resource idleness is distinctly noticeable, indicating that rollout resources can be reduced.\n\n![fully_async_policy_structure](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana_fully_async_train.png?raw=true)\n\nThrough the above two examples, we also illustrate the necessity of system observability.\n\n## Architecture Overview\n\nThe overall workflow consists of the following steps:\n\n1. **Multi-node Ray Cluster Setup**: Start Ray cluster across multiple nodes with Grafana and Prometheus information configured in environment variables on the master node\n2. **Start Grafana Service**: Launch Grafana on the master node for visualization of monitoring dashboards\n3. **Start Prometheus Service**: Launch Prometheus on the master node for metrics collection and storage\n4. **verl Async Rollout Mode**: verl uses async rollout mode to obtain rollout server ports and IP addresses\n5. **Automatic Prometheus Configuration**: verl automatically rewrites the Prometheus configuration to add monitoring for rollout servers and notifies Prometheus to reload the configuration\n6. **Metrics Collection**: After program execution, metrics can be viewed in Prometheus\n7. **Dashboard Visualization**: Upload and view monitoring metrics in Grafana dashboards\n\n## Detailed Setup Steps\n\n### Step 1: Environment Variables and Start Ray Cluster\n\nFirst, set the necessary environment variables and start the Ray service.\n\n> Reference: [configure-manage-dashboard](https://docs.ray.io/en/latest/cluster/configure-manage-dashboard.html)\n\n```bash\n# Master node environment variables\nexport GF_SERVER_HTTP_PORT=3000                     # Grafana service default port (customizable)\nexport PROMETHEUS_PORT=9090                         # Prometheus service default port (customizable)\nexport RAY_HEAD_PORT=6379                           # Ray master node port (customizable)\nexport RAY_DASHBOARD_PORT=8265                      # Ray dashboard default port (customizable)\nexport GRAFANA_PATHS_DATA=/tmp/grafana              # Grafana data storage directory (customizable)\nexport RAY_GRAFANA_HOST=\"http://${master_ip}:${GF_SERVER_HTTP_PORT}\"        # Ray-associated Grafana address\nexport RAY_PROMETHEUS_HOST=\"http://${master_ip}:${PROMETHEUS_PORT}\"         # Ray-associated Prometheus address\n\n# Start Ray on master node\nray start --head --port=${RAY_HEAD_PORT} --dashboard-port=${RAY_DASHBOARD_PORT}\n\n# Start Ray on worker nodes\nray start --address={master_addr}:${RAY_HEAD_PORT}\n```\n\n**Verification:** Visit `http://master_ip:8265` to confirm Ray has started successfully.\n\n### Step 2: Start Grafana (Visualization Dashboard)\n\nGrafana is used to display metrics collected by Prometheus (such as cache hit rate, throughput, etc.):\n\n```bash\n# Master node\nnohup grafana-server \\\n  --config /tmp/ray/session_latest/metrics/grafana/grafana.ini \\\n  --homepath /usr/share/grafana \\\n  web > grafana.log 2>&1 &\n```\n\n**Verification:** Visit `http://master_ip:3000` to confirm Grafana has started successfully (default credentials: `admin/admin`).\n\nIf you need to change the port, modify the `GF_SERVER_HTTP_PORT` environment variable, and grafana-server will automatically recognize it.\n\n### Step 3: Start Prometheus (Metrics Collection)\n\nPrometheus is responsible for scraping metrics from vLLM services and storing them as time-series data:\n\n```bash\n# Master node\nnohup prometheus \\\n  --config.file /tmp/ray/session_latest/metrics/prometheus/prometheus.yml \\\n  --web.enable-lifecycle \\\n  --web.listen-address=:${PROMETHEUS_PORT} \\\n  > prometheus.log 2>&1 &\n```\n\n**Verification:** Visit `http://master_ip:9090` to confirm Prometheus service has started successfully.\n\n### Step 4 & 5: Start verl Training\n\nStart verl training with the following parameters configured:\n\n**Required Configuration:**\n\n- `actor_rollout_ref.rollout.mode=\"async\"`\n- `actor_rollout_ref.rollout.disable_log_stats=False`\n- `actor_rollout_ref.rollout.prometheus.enable=True`\n\nIf use default port, this parameter can be omitted.\n\n- `actor_rollout_ref.rollout.prometheus.port=9090`\n\nIf use default path, this parameter can be omitted.\n\n- `actor_rollout_ref.rollout.prometheus.file=\"/tmp/ray/session_latest/metrics/prometheus/prometheus.yml\"`\n\nserved_model_name uses `model_path.split(\"/\")[-1]` for data statistics by default.\nUsers can also customize other aliases:\n\n- `actor_rollout_ref.rollout.prometheus.served_model_name=\"Qwen3-235B\"`\n\n**Shell Script Example:**\n\n```bash\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\"  # Options: sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Synchronous training\nray job submit --no-wait --runtime-env=\"${RUNTIME_ENV}\" \\\n    --working-dir \"${WORKING_DIR}\" \\\n    -- python3 -m verl.trainer.main_ppo \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.disable_log_stats=False \\\n    actor_rollout_ref.rollout.prometheus.enable=True\n    ...\n\n# Asynchronous training\nray job submit --no-wait --runtime-env=\"${RUNTIME_ENV}\" \\\n    --working-dir \"${WORKING_DIR}\" \\\n    -- python3 verl.experimental.fully_async_policy.fully_async_main \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.disable_log_stats=False \\\n    actor_rollout_ref.rollout.prometheus.enable=True\n    ...\n```\n\n### Step 6: View Metrics in Prometheus\n\nAfter task execution, verify that Prometheus is correctly collecting metrics.\n\n**Verification:** Visit the Prometheus interface at `http://master_ip:9090` and search for `vllm:` or `sglang:` to\nconfirm metrics are being reported correctly.\n\n**Troubleshooting:**\n\nIf no metrics appear:\n\n1. Check logs for `AgentLoopManager` to find the server port\n2. Visit `http://master_ip:server_port/metrics` to verify server metrics are available\n3. Confirm that `actor_rollout_ref.rollout.disable_log_stats=False` is set\n\n### Step 7: View Metrics in Grafana\n\nAfter task execution, log in to Grafana to view and customize monitoring dashboards.\n\n**Login:** Visit `http://master_ip:3000` (default credentials: `admin/admin`)\n\n**Import Dashboard:**\n\n1. Select `Dashboards` → `New` → `Import` → `Upload dashboard JSON file`\n2. Upload a pre-built dashboard JSON file\n\n**Available Dashboards:**\n\n- [vLLM Grafana Dashboard style 1](https://github.com/ArronHZG/verl-community/blob/main/docs/grafana/vllm_grafana.json)\n- [vLLM Grafana Dashboard style 2](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/dashboards/grafana/performance_statistics.json)\n- [vLLM Grafana Dashboard style 2](https://github.com/vllm-project/vllm/blob/main/examples/online_serving/dashboards/grafana/query_statistics.json)\n- [SGLang Grafana Dashboard](https://github.com/sgl-project/sglang/blob/main/examples/monitoring/grafana/dashboards/json/sglang-dashboard.json)\n\n## Additional Resources\n\n- [Ray Monitoring Documentation](https://docs.ray.io/en/latest/cluster/configure-manage-dashboard.html)\n- [Prometheus Documentation](https://prometheus.io/docs/)\n- [Grafana Documentation](https://grafana.com/docs/)\n- [vLLM GitHub Repository](https://github.com/vllm-project/vllm)\n- [SGLang GitHub Repository](https://github.com/sgl-project/sglang)\n"
  },
  {
    "path": "docs/advance/megatron_extension.rst",
    "content": "Add models with the Megatron-LM backend\n=========================================\n\nLast updated: 04/25/2025.\n\nModel\n-----------\n\n\nIf use latest verl, we have direct support of ``GPTModel`` for Megatron backend. \nYou can use the similar way of using Megatron to pretrain custom models. \nWe list the steps here:\n\n1. Find `model_initializer.py <https://github.com/volcengine/verl/blob/main/verl/models/mcore/model_initializer.py>`_\n2. If your model is configurable by ``TransformerLayerSpec`` , you can\n   directly use ``GPTModel``. Otherwise, Please implement a new\n   ``ModelLayerSpec`` and ``ModelLayer`` here.\n3. Use the right ``LayerSpec`` , ``TransformerConfig`` and ``HuggingfaceConfig`` \n   as arguments to initialize the GPTModel.\n4. Return the model at last.\n"
  },
  {
    "path": "docs/advance/mtp.md",
    "content": "# Guide to Using MTP in SFT/RL Training and Inference\n\n**Author**: `https://github.com/meituan-search`\n\nLast updated: 02/15/2026\n\n# 1. Scope of Support\n\nCurrently, RL training can be performed on mimo-7B-RL, Qwen-next, and Deepseek series models based on the MTP architecture. The support rules for training and inference engines are as follows:\n\n- **Training Engine**: Only supports the `mbridge/Megatron-Bridge + megatron` combination; other training engines are not compatible at this time;\n\n- **Inference Engine**: Compatible with all engines, but the model must be in the corresponding engine's compatibility list;\n\n- **Dependency Versions**:\n\n    - mbridge: Apply the patches and review suggestions from PR: [#62](https://github.com/ISEEKYAN/mbridge/pull/62) (Already merged into the main branch);\n\n    - Megatron-Bridge: Apply the patches and review suggestions from PR if you want to try out mimo-7B-RL: [#2387](https://github.com/NVIDIA-NeMo/Megatron-Bridge/pull/2387) (will be merged into the main branch in the future);\n\n    - megatron: Use the latest dev version (commit: [23e092f41ec8bc659020e401ddac9576c1cfed7e](https://github.com/NVIDIA/Megatron-LM/tree/23e092f41ec8bc659020e401ddac9576c1cfed7e)), which supports MTP + CP training methods.\n    \n    - sglang: Use the specified branch: [https://github.com/ArronHZG/sglang/tree/fix_mtp_update_weights_from_tensor](https://github.com/ArronHZG/sglang/tree/fix_mtp_update_weights_from_tensor), [PR](https://github.com/sgl-project/sglang/pull/17870) , which fix the MTP update weights from tensor OOM issue.\n\n# 2. MTP Training Configuration (Core Parameters)\n\nThe MTP training process can be flexibly controlled through the following configurations. All configurations are based on the `actor_rollout_ref.model.mtp` prefix:\n\n| Configuration Scenario | Core Parameters                                                                                                                                                                                                                                                                                               | Description                                             |\n|------------------------|---------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------|---------------------------------------------------------|\n| Load MTP Parameters Only | `enable=True`                                                                                                                                                                                                                                                                                              | VRAM usage will increase, but the exported parameters include the MTP module and can be directly used for online deployment              |\n| Full-Parameter MTP Training | `enable=True`<br>`enable_train=True`<br>`mtp_loss_scaling_factor=0.1`                                                                                                                                                                                                                              | MTP Loss will apply to all model parameters                            |\n| MTP Parameter-Only Training | `enable=True`<br>`enable_train=True`<br>`detach_encoder=True`                                                                                                                                                                                                                                      | Freeze the Encoder layer, update only MTP module parameters, MTP Loss applies only to MTP parameters |\n| MTP Accelerated Rollout | 1. vLLM configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`method=\"mtp\"`<br>`num_speculative_tokens=1`<br>2. SGLang configuration:<br>`enable=True`<br>`enable_rollout=True`<br>`speculative_algorithm=\"EAGLE\"`<br>`speculative_num_steps=2`<br>`speculative_eagle_topk=2`<br>`speculative_num_draft_tokens=4` | Achieve inference acceleration during the Rollout phase based on MTP                      |\n\n# 3. Experimental Results\n\nThe experiment was conducted as follows:\n\n* model = mimo-7B-math\n* max_response_length = 8k\n\nExperiment chart:\n\n![fully_async_policy_revenue](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/mimo-7b-mtp.png?raw=true)\n\nThe wandb link for the graph: [wandb](https://wandb.ai/hou-zg-meituan/mimo-7b-sft-mtp?nw=nwuserhouzg)\n\n**Scenarios with No Significant Effect**\n\nThe following configurations will not have a noticeable impact on training results:\n\n1. The base model does not carry MTP parameters;\n\n2. The base model carries MTP parameters, but the MTP module is not trained;\n\n3. The base model carries MTP parameters and trains MTP, with `mtp_loss_scaling_factor=0`;\n\n4. The base model carries MTP parameters, trains MTP and detaches the encoder, with `mtp_loss_scaling_factor=0.1`.\n\n**Scenarios with Significant Effect**\n\nOnly the following configuration will have a noticeable impact on training results:\n\n- The base model carries MTP parameters, MTP Loss applies to all model parameters, and `mtp_loss_scaling_factor=0.1`.\n\n**Recommended Training Method**\n\nIt is recommended to adopt the `detach_encoder=True` approach for MTP training.\n\n# 4. Performance Notes for MTP in Rollout Inference\n\nEnabling MTP improves the rollout acceptance rate by around 14%. However, on H20 GPUs, overall throughput does not increase and even decreases slightly.\n\n![spec_log](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/spec_log.png?raw=true)\n\nThe effectiveness of MTP-accelerated Rollout is significantly affected by **model size** and **inference hardware**. Key reference information is as follows:\n\n**Hardware Tensor Core Performance**\n\n| Hardware Model | FP16 Performance (TFLOPS) |\n|----------------|---------------------------|\n| H20  | 148            |\n| H800 | 1,671          |\n| H200 | 1,979          |\n\n**Measured Performance and Recommendations**\n\nTaking the mimo-7B model deployed separately on H20 hardware using SGLang as an example: After enabling MTP speculative decoding, the Rollout throughput decreases by approximately 50%.\n\n- Current priority recommendation: Do not enable MTP acceleration during the inference phase for now;\n\n- Future planning: Further optimization of the speculative logic in the Rollout phase will be conducted to improve throughput performance.\n\n# 5. SFT training\n\nThe SFT training with MTP is supported, using the same MTP training configuration as RL training.\n\nAn example configuration for running SFT can be found in `examples/sft/gsm8k/run_mimo_megatron_mtp.sh`\n\n**SFT result**\n\nThe experiment was conducted using following data:\n- model = mimo-7B-math\n- dataset = gsm8k\n\nThe result: [wandb link](https://wandb.ai/hou-zg-meituan/mimo-7b-sft-mtp?nw=nwuserhouzg)\n\nThe presence of mtp layer has limited effect on main loss. However, when MTP layer is detached, the mtp_loss converges to a higher value.\n\n"
  },
  {
    "path": "docs/advance/one_step_off.md",
    "content": "# Recipe: One Step Off Policy Async Trainer\n\n**Author:** `https://github.com/meituan-search`\n\nLast updated: 07/17/2025.\n\n## Introduction\n\n### Background\n\nThe current reinforcement learning training process implemented by verl is synchronous, adhering to the algorithmic\nworkflows of established methods like PPO, GRPO, and DAPO. In each step, training samples are generated by the latest\nmodel, and the model is updated after training completes. While this approach aligns with off-policy reinforcement\nlearning and stabilizes RL training, but it suffers from severe efficiency issues.\nModel updates must wait for the longest output in the generation phase to complete.\nDuring the generation of long-tail samples, GPUs remain idle, resulting in significant underutilization.\nThe more severe the long-tail problem in sample generation, the lower the overall training efficiency.\nFor example, in DAPO 32B training, the Rollout phase accounts for approximately 70% of the total time,\nand increasing resources does not reduce the Rollout duration.\n\n![DAPO 32B Math Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/dapo_32b_math.png)\n\n> source data: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=nwusertongyuxuan361\n\n### Solution\n\nWe have implemented the **One Step Off Async Trainer** to help alleviate this issue. This approach parallelizes the\ngeneration and training processes, utilizing samples generated in the previous step for current training.\nIt also involves appropriately partitioning resources, allocating dedicated resources for generation while automatically\nassigning the remainder to training. By reducing resources allocated to the generation phase, we mitigate GPU idle time\nduring long-tail sample generation. Throughout this process, generation and training parameters maintain a one-step off\npolicy.\n\n![One Step Off Policy Diagram](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_policy.png)\n\n> reference: [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://arxiv.org/abs/2505.24298)\n\nOur core contributions include:\n\n1. **Parallel Generation and Training**:\n   Samples for the next batch are asynchronously generated while the current batch is being trained.\n\n2. **Resource Isolation**:\n   Unlike `hybrid_engine`, this method requires explicit resource allocation for rollout, with remaining resources\n   automatically assigned to training.\n\n3. **NCCL Parameter Synchronization**:\n   Employs NCCL communication primitives for seamless parameter transfer between generation and training modules.\n\n### Experimental Results\n\n- **Machine Configuration**: 2 nodes with 16 H20 GPUs each\n  - Generation: 4 GPUs\n  - Training: 12 GPUs\n- **Model**: Qwen2.5-Math-7B\n- **Rollout Configuration**:\n- **Max Response Length**: FSDP2: 20,480 tokens; Megatron: 8,192 tokens\n- **Algorithm**: DAPO\n- **Rollout Engine**: vLLM\n\n| training mode          | engine        | step | gen | wait_prev_gen | generate_sequences | old_log_prob | update_actor | total time     | acc/best@32/mean | acc/maj@32/mean |\n| ---------------------- | ------------- | ---- | --- | ------------- | ------------------ | ------------ | ------------ | -------------- | ---------------- | --------------- |\n| colocate sync          | VLLM+FSDP2    | 749  | 321 | -             | 247                | 88           | 286          | 19h18m         | 0.5948           | 0.417           |\n| one-step-overlap async | VLLM+FSDP2    | 520  | -   | 45            | 458                | 108          | 337          | 15h34m（+23%） | 0.6165           | 0.494           |\n| colocate sync          | VLLM+Megatron | 699  | 207 | -             | 162                | 119          | 344          | 18h21m         | 0.605            | 0.4217          |\n| one-step-overlap async | VLLM+Megatron | 566  | -   | 59            | 501                | 120          | 347          | 13h06m (+40%)  | 0.6569           | 0.4038          |\n\n- colocate sync: step ≈ gen + old_log_prob + update_actor\n- one-step-overlap async: step ≈ wait_prev_gen + old_log_prob + update_actor\n\n![One Step Off Megatron Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_megatron.png)\n\n> source data: https://wandb.ai/hou-zg-meituan/one-step-off-policy?nw=nwuserhouzg\n\n## Implementation\n\n### One Step Off Policy Async Pipeline\n\nOur implemented **One Step Off Policy Async Pipeline** integrates seamlessly into existing training logic at minimal\ncost,\neliminating the need for additional sample storage management. The core mechanism uses `async_gen_next_batch`\nfor asynchronous rollout generation while maintaining continuous operation during epoch transitions\nvia `create_continuous_iterator`.\n\n```python\n# iterator generator, simplify one-step integration of the training process\ndef _create_continuous_iterator(self):\n   for epoch in range(self.config.trainer.total_epochs):\n      iterator = iter(self.train_dataloader)\n      for batch_dict in iterator:\n         yield epoch, batch_dict\n\n\n# read next batch samples, parameters sync and launch asyn gen_seq\ndef _async_gen_next_batch(self, continuous_iterator):\n   # read train_data\n   try:\n      epoch, batch_dict = next(continuous_iterator)\n   except StopIteration:\n      return None\n   batch = DataProto.from_single_dict(batch_dict)\n   gen_batch = batch_pocess(batch)\n   # sync weights from actor to rollout\n   self.sync_rollout_weights()\n   # async generation\n   gen_batch_output = self.rollout_wg.async_generate_sequences(gen_batch)\n   # future encapsulated\n   return GenerationBatchFuture(epoch, batch, gen_batch_output)\n\n\ncontinuous_iterator = self._create_continuous_iterator()\n# run rollout first to achieve one-step-off\nbatch_data_future = self._async_gen_next_batch(continuous_iterator)\n\nwhile batch_data_future is not None:\n   # wait for the gen_seq result from the previous step\n   batch = batch_data_future.get()\n   # launch the next async call to generate sequences\n   batch_data_future = self._async_gen_next_batch(continuous_iterator)\n\n   # compute advantages\n   batch = critic.compute_values(batch)\n   batch = reference.compute_log_prob(batch)\n   batch = reward.compute_reward(batch)\n   batch = compute_advantages(batch)\n\n   # model update\n   critic_metrics = critic.update_critic(batch)\n   actor_metrics = actor.update_actor(batch)\n```\n\n### Parameter Synchronization\n\nThe exciting point is that our nccl based weights updating for rollout model has great performance.\nAt most of time, the latency is under 300ms, which is negligible for RLHF.\n\n> **sync_rollout_weights**：The time for synchronizing parameters from actor to rollout is extremely fast and can almost\n> be ignored because it is implemented with nccl.\n\n```python\nclass ActorRolloutRefWorker:\n   # actor acquires the meta-info of model parameters for parameter sync\n   @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n   def get_actor_weights_info(self):\n      params = self._get_actor_params()\n      ret = []\n      for key, tensor in params.items():\n         ret.append((key, tensor.size(), tensor.dtype))\n      self._weights_info = ret\n      return ret\n\n   # rollout sets the meta-info of model parameters for parameter sync\n   @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n   def set_actor_weights_info(self, weights_info):\n      self._weights_info = weights_info\n\n\nclass AsyncRayPPOTrainer(RayPPOTrainer):\n   def init_workers(self):\n      ...\n      # rollout obtains the meta-info of model parameters from the actor for parameter sync\n      weights_info = self.actor_wg.get_actor_weights_info()[0]\n      self.rollout_wg.set_actor_weights_info(weights_info)\n\n      # Create an actor-rollout communication group for parameter sync\n      self.create_weight_sync_group\n```\n\n```python\n# The driving process invokes the actor and rollout respectively to create a weight synchronization group based on nccl/hccl.\ndef create_weight_sync_group(self):\n   master_address = ray.get(self.actor_wg.workers[0]._get_node_ip.remote())\n   master_port = ray.get(self.actor_wg.workers[0]._get_free_port.remote())\n   world_size = len(self.actor_wg.workers + self.rollout_wg.workers)\n   self.actor_wg.create_weight_sync_group(\n      master_address,\n      master_port,\n      0,\n      world_size,\n   )\n   ray.get(\n      self.rollout_wg.create_weight_sync_group(\n            master_address,\n            master_port,\n            len(self.actor_wg.workers),\n            world_size,\n      )\n   )\n\n# drive process call the actor and rollout respectively to sync parameters by nccl\ndef sync_rollout_weights(self):\n   self.actor_wg.sync_rollout_weights()\n   ray.get(self.rollout_wg.sync_rollout_weights())\n\n\n# fsdp model parameter sync\n@register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\ndef sync_rollout_weights(self):\n   params = self._get_actor_params() if self._is_actor else None\n   if self._is_rollout:\n      inference_model = (\n         self.rollout.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model\n      )\n      from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader\n      patch_vllm_moe_model_weight_loader(inference_model)\n   # Model parameters are broadcast tensor-by-tensor from actor to rollout\n   for key, shape, dtype in self._weights_info:\n      tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())\n      if self._is_actor:\n         assert key in params\n         origin_data = params[key]\n         if hasattr(origin_data, \"full_tensor\"):\n            origin_data = origin_data.full_tensor()\n         if torch.distributed.get_rank() == 0:\n            tensor.copy_(origin_data)\n      from ray.util.collective import collective\n\n      collective.broadcast(tensor, src_rank=0, group_name=\"actor_rollout\")\n      if self._is_rollout:\n         inference_model.load_weights([(key, tensor)])\n```\n\n### PPO Correctness\n\nTo ensure the correctness of the PPO algorithm, we use rollout log_probs for PPO importance sampling.\nFor the related algorithm details, please refer to: https://verl.readthedocs.io/en/latest/algo/rollout_corr_math.html\nThe default mode is `bypass_ppo_clip`, but other modification strategies can also be explored.\n\n### AgentLoop\n\nIn the current implementation, we no longer provide SPMD model rollout mode.\nInstead, we have switched to AgentLoop mode, which also supports multi-turn tool calling.\n\n## Usage\n\n### FSDP2 Configuration Example\n\n```shell\npython3 -m verl.experimental.one_step_off_policy.async_main_ppo \\\n    --config-path=config \\\n    --config-name='one_step_off_ppo_trainer.yaml' \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    # actor and rollout are placed separately\n    actor_rollout_ref.hybrid_engine=False \\\n    # actor and rollout resource\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=6 \\\n    rollout.nnodes=1 \\\n    rollout.n_gpus_per_node=2\n```\n\n### Megatron Configuration Example\n\n```shell\npython3 -m verl.experimental.one_step_off_policy.async_main_ppo \\\n    --config-path=config \\\n    --config-name='one_step_off_ppo_megatron_trainer.yaml' \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    # actor and rollout are placed separately\n    actor_rollout_ref.hybrid_engine=False \\\n    # actor and rollout resource\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=6 \\\n    rollout.nnodes=1 \\\n    rollout.n_gpus_per_node=2\n```\n\n### Configuration Guidelines\n\n1. **Card Number Relationships**\n   Maintain either of these relationships for optimal batch distribution:\n\n   - `actor_rollout_ref.rollout.n` should be an integer divisor of:\n     `trainer.n_gpus_per_node * trainer.nnodes`\n   - `actor_rollout_ref.rollout.n * data.train_batch_size` should be evenly divisible by:\n     `trainer.n_gpus_per_node * trainer.nnodes`\n\n   > Rationale: Ensures training samples can be evenly distributed across training GPUs when using partial resources for\n   > generation.\n\n2. **Dynamic Resource Tuning**\n   Adjust `trainer.nnodes` `trainer.n_gpus_per_node` `rollout.nnodes` `rollout.n_gpus_per_node` based on phase\n   durations:\n   - **Ideal state**: Rollout and training phases have comparable durations\n   - **Diagnostic metrics**:\n     - Monitor `wait_prev_gen` duration\n     - Analyze `sequence_length` distribution\n   - **Adjustment strategy**:\n     - High `wait_prev_gen` + uniform sequence lengths → Increase rollout resources\n     - High `wait_prev_gen` + long-tail sequences → Optimize stopping criteria (resource increase won't help)\n       > **wait_prev_gen**：The time consumed waiting for the previous rollout to end (the part that is not fully\n       > overlapped).\n       > **Resource Configuration Strategies:**\n   - **Resource-constrained scenario**: Optimize resource utilization by adjusting GPU allocation ratios,\n     keeping the number of nodes equal to allow training and rollout to share nodes;\n     - Configure `trainer.nnodes = rollout.nnodes` with\n       `trainer.n_gpus_per_node + rollout.n_gpus_per_node = physical_gpus_per_node`. Control rollout resource\n       allocation by adjusting `n_gpus_per_node`.\n   - **Resource-abundant scenario**: Optimize performance by adjusting the number of nodes,\n     keeping the number of GPUs per node equal to enable independent scaling of training and rollout\n     parallelism.\n     - Configure `trainer.n_gpus_per_node = rollout.n_gpus_per_node` and control rollout resource allocation by\n       adjusting `trainer.nnodes` and `rollout.nnodes`to achieve optimal performance.\n       > **Note**: The total number of nodes required by the system is not simply `trainer.nnodes + rollout.nnodes`. The\n       > actual calculation depends on GPU capacity:\n       >\n       > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node <= physical_gpus_per_node`,\n       >   the required node count is `max(trainer.nnodes, rollout.nnodes)`\n       > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`,\n       >   the required node count is `trainer.nnodes + rollout.nnodes`\n\n## Functional Support\n\n| Category           | Support Situation                                                                                               |\n| ------------------ | --------------------------------------------------------------------------------------------------------------- |\n| train engine       | FSDP2 <br/> Megatron                                                                                            |\n| rollout engine     | vLLM                                                                                                            |\n| AdvantageEstimator | GRPO <br/> GRPO_PASSK <br/> REINFORCE_PLUS_PLUS <br/> RLOO <br/> OPO <br/> REINFORCE_PLUS_PLUS_BASELINE<br/>GPG |\n| Reward             | all                                                                                                             |\n"
  },
  {
    "path": "docs/advance/placement.rst",
    "content": "Ray API Design Tutorial\n=======================================\n\nLast updated: 10/30/2024.\n\nWe provide a tutorial for our Ray API design, including:\n\n- Ray basic concepts\n- Resource Pool and RayWorkerGroup\n- Data Dispatch, Execution and Collection\n- Initialize the RayWorkerGroup and execute the distributed computation in the given Resource Pool\n\nSee details in `tutorial.ipynb <https://github.com/volcengine/verl/blob/main/examples/ray/tutorial.ipynb>`_."
  },
  {
    "path": "docs/advance/ppo_lora.rst",
    "content": "RL(HF) algorithms with LoRA Support\n===========================================\n\nLast updated: 02/03/2026.\n\nWe support LoRA (Low-Rank Adaptation) for reinforcement learning algorithms such as PPO, GRPO, and others.\n\nLoRA is a parameter-efficient fine-tuning technique that injects trainable low-rank matrices into pre-trained weights (typically linear layers). This reduces memory footprint and compute cost, making it possible to fine-tune large models with limited hardware.\n\nThe benefits this brings include:\n\n- reinforcement learning with very large models (e.g. 70B+) with modest hardware (e.g. 8x80G GPUs),\n- enable larger batch sizes due to reduced memory usage,\n- simplify model transfer and deployment, as only LoRA adapters need to be saved,\n- Combine with techniques like `SLoRA <https://arxiv.org/abs/2311.03285>`_ or `CCoE <https://arxiv.org/abs/2407.11686>`_ to serve multiple LoRA adapters efficiently\n\nThis guide explains how to enable LoRA in RL training and configure related parameters.\n\nFSDP Backend Usage Guide\n------------------------\n\n.. note::\n\n   This section applies to **FSDP/FSDP2 backend only**. For Megatron backend, see the :ref:`megatron-lora` section below.\n\n1. Lora is available in the `verl.trainer.ppo.ray_trainer.RayPPOTrainer`. Examples are provided via the `verl.trainer.main_ppo` entry point.\n\n2. Currently, LoRA is supported via huggingface peft, only with fsdp/fsdp2 and vllm backend (sglang support coming soon).\n\n- `strategy=fsdp` or `strategy=fsdp2`\n- `rollout.name=vllm`\n\n3. Required configurations for LoRA:\n\n- `actor_rollout_ref.model.lora_rank`: int, set to a reasonable value greater than 0 (e.g., 8, 16, 32, 64)\n- `actor_rollout_ref.model.lora_alpha`: float, the alpha term in LoRA\n- `actor_rollout_ref.rollout.load_format=\"safetensors\"`: required. This enables vLLM to load the base model.\n- `actor_rollout_ref.model.target_modules`: the target modules for LoRA. Typically set to \"all-linear\".\n\n4. Optional configurations for LoRA:\n\n- `actor_rollout_ref.model.lora_adapter_path`: string, path to a pretrained LoRA adapter directory. \n   If provided, loads existing adapter instead of creating new one. Enables multi-stage training from previously saved adapters.\n   Directory need contain `adapter_model.safetensors` and `adapter_config.json`.\n- `actor_rollout_ref.model.lora.merge`: bool, whether to merge LoRA adapters into the base model weights before transferring to vLLM. \n   If True, it will merge LoRA adapters into the base model weights before transferring to vLLM. If False, it will transfer only adapters to vLLM. This option is currently supported **only for engine-based rollout workers** (i.e. vLLM engine workers using the new worker implementation with ``trainer.use_legacy_worker_impl`` disabled) and is not available when using the legacy worker implementation.\n\n5. Recommend options:\n\n- `actor_rollout_ref.model.use_shm=True`: preload the model into `/dev/shm` to improve model loading speed.\n- `actor_rollout_ref.rollout.layered_summon=True`: this enables the actor-model to gather the FSDP shards per layers when synchronizing the LoRA Adapter to vLLM, thereby reducing GPU peak memory. Recommended if the model is very large (70B+) or the GPU memory is limited (< 48GB)\n\n.. _megatron-lora:\n\nMegatron Backend Usage Guide\n----------------------------\n\n.. warning::\n\n   The FSDP-specific config options are **NOT applicable** to Megatron backend, and they will be ignored if set. Only options listed under ``lora`` key are applicable:\n\n   - ``actor_rollout_ref.model.lora.*``\n   - ``critic.model.lora.*``\n\nYou need to install and enable Megatron-Bridge for Megatron LoRA support.\n\nMake sure you use Megatron-Bridge later than 0.2.0, and we recommended using `this commit <https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/83a7c1134c562d8c6decd10a1f0a6e6a7a8a3a44>`_ or later for proper support, and use the following settings to enable Megatron-Bridge:\n\n- ``actor_rollout_ref.actor.megatron.use_mbridge=True``\n- ``actor_rollout_ref.actor.megatron.vanilla_mbridge=False``\n\n**Key Differences from FSDP LoRA:**\n\n1. **LoRA Implementation**: Verl Megatron backend uses Megatron-Bridge's native LoRA implementation, which differs from HuggingFace PEFT.\n\n2. **Weight Sync / Refit Mechanism**: Currently, Megatron-Bridge can support syncing weights by either merging LoRA adapters into the base model weights before transferring to vLLM (for better inference speed but more refit time and potential precision loss), as well as loading separate adapters.\n\n**Configuration for Megatron LoRA:**\n\n.. code-block:: yaml\n\n   actor_rollout_ref:\n     model:\n      lora:\n        # LoRA type: \"lora\", \"vlm_lora\", \"canonical_lora\", or \"dora\"\n        type: lora\n\n        # whether to sync weights / refit by either merging LoRA adapters into the base model weights before transferring to vLLM (for better inference speed but more refit time and potential precision loss). If this is False, it will load separate adapters.\n        merge: False\n\n        # LoRA rank (Dimension of the low-rank projection space.). Set to 0 to disable LoRA\n        rank: 0\n        \n        #  Weighting factor for the low-rank projection. Defaults to 32\n        alpha: 32\n        \n        # Dropout rate for the low-rank projection. Defaults to 0.0\n        dropout: 0.0\n        \n        # A list of module names to apply LoRA to.\n        # For fused LoRA, Defaults to all linear layers ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2'].\n        # For canonical LoRA: [\"linear_q\", \"linear_k\", \"linear_v\", \"linear_proj\", \"linear_fc1_up\", \"linear_fc1_gate\", \"linear_fc2\"]\n        # - 'linear_qkv': Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention\n        # - 'linear_proj': Apply LoRA to the linear layer used for projecting the output of self-attention\n        # - 'linear_fc1': Apply LoRA to the first fully-connected layer in MLP\n        # - 'linear_fc2': Apply LoRA to the second fully-connected layer in MLP\n        # Target modules can also contain wildcards. For example, you can specify\n        # target_modules=['*.layers.0.*.linear_qkv', '*.layers.1.*.linear_qkv'] to add LoRA to only linear_qkv on the first two layers\n        # \n        # Note:\n        # For MLA (e.g., DeepSeek), you should use [\"linear_kv_down_proj\",\"linear_kv_up_proj\",\"linear_q_down_proj\",\"linear_q_up_proj\",\"linear_q_proj\"]\n        # Instead of \"linear_qkv\" or [\"linear_q\",\"linear_k\",\"linear_v\"]\n        # By default, MoE routers are excluded from LoRA adaptation, and you will need to specify \"router\" in target_modules to include them.\n        target_modules:\n          - linear_qkv\n          - linear_proj\n          - linear_fc1\n          - linear_fc2\n        \n        # A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name\n        # does not match any string in exclude_modules. If used, will require target_modules to be empty list or None\n        exclude_modules: []\n\n        # Position for applying dropout, can be 'pre' (before the low-rank projection) or 'post' (after). Defaults to 'pre'\n        dropout_position: pre\n\n        # Initialization method for the low-rank matrix A. Defaults to \"xavier\".\n        lora_A_init_method: xavier\n\n        # Initialization method for the low-rank matrix B. Defaults to \"zero\".\n        lora_B_init_method: zero\n\n        # Enables the experimental All-to-All (A2A) communication strategy. Defaults to False\n        a2a_experimental: False\n\n        # Parameter data type for LoRA weights. Default to null, which will use model's dtype.\n        dtype: null\n\n        # Path to pre-trained LoRA adapter weights (null to train from scratch)\n        adapter_path: null\n\n        # Whether to fully shard LoRA adapters. Defaults to False\n        # https://docs.vllm.ai/en/latest/api/vllm/config/lora/#vllm.config.lora.LoRAConfig.fully_sharded_loras\n        fully_sharded_loras: bool\n\n        # VLMLoRA additionally allows the user to specify whether the language or vision models should be frozen.\n        # For example, a common finetuning workload for multimodal models is to apply adapters to language model and fully\n        # finetune the vision model.\n        freeze_vision_model: True\n        freeze_vision_projection: True\n        freeze_language_model: True\n\nLoRA training experiment with Qwen3-8B on 8 * H200 single node comparing FSDP and Megatron backend (script adapted from examples/grpo_trainer/run_qwen2-7b_math_megatron_lora.sh):\n\n.. image:: https://github.com/user-attachments/assets/0482f423-01a3-4e52-a7ee-8b9cd79b7b1a\n.. image:: https://github.com/user-attachments/assets/6ce10400-8164-47d8-90a6-c1bf002fb9e8\n.. image:: https://github.com/user-attachments/assets/092d3a43-4eba-425e-a584-8d83c1f02de4\n\n\nBest Practices and Notes\n-------------------------\n\n1. **Learning rate**: it is recommended to increase the value of learning rate by an order of magnitude.\n\n2. **LoRA Rank**:\n\n- Too small a rank can hurt convergence.\n- LoRA rank recommendation from @thelongestusernameofall:\n\n  - A very small lora_rank can lead to slower convergence or worse training performance. It is recommended to set lora_rank to be>=32. Tests have shown that for a 0.5B model, with lora_rank=32,the training convergence speed and final performance are almost identical to non-LoRA training\n  - For a 32B model,with lora_rank=128,the training convergence speed and final performance are also almost identical to non-LoRA training.\n  - More comprehensive reference results are coming soon.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/f2b80b8b26829124dd393b7a795a0640eff11644/docs/lora.jpg?raw=true\n\n3. **FSDP-Specific:** Reference configuration for RL training with the Qwen2.5-72B model using 8 x 80GB GPUs (increase lora_rank if needed):\n\n.. code-block::\n\n    data.train_batch_size=64 \\\n    actor_rollout_ref.model.use_shm=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=8 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=64 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n\nExample Scripts\n-------------------\n\nFor end-to-end examples, refer to the scripts below:\n\n**FSDP Examples:**\n\n- LoRA training from scratch: examples/grpo_trainer/run_qwen2_5-3b_gsm8k_grpo_lora.sh\n- LoRA training from adapter path: examples/grpo_trainer/run_qwen2_5-3b_gsm8k_grpo_lora_from_adapter.sh\n\n**Megatron Examples:**\n\n- LoRA training with Dense: examples/grpo_trainer/run_qwen2-7b_math_megatron_lora.sh\n- LoRA training with MoE: examples/grpo_trainer/run_qwen3moe-30b_megatron_lora.sh\n"
  },
  {
    "path": "docs/advance/reward_loop.rst",
    "content": "Reward Loop\n===========\n\n.. _yyding: https://yyding1.github.io\n\nAuthor: `Yuyang Ding <https://yyding1.github.io>`_\n\nLast updated: 2/10/2026.\n\nIntroduction\n------------\n\nReward Loop is the default reward computation implementation in verl.\nIt is designed to support efficient, flexible, and easy-to-use reward computation.\n\nThis document introduces the usage and architectural design.\n\nKey features include:\n\n1. **Distributed reward manager**, enabling scalable and efficient reward computation.\n2. **Support for hybrid reward settings**, including both generative and discriminative reward models, as well as more complex reward scenarios.\n3. **Simple and extensible interface**, for easily defining customized reward functions.\n\nDistributed Reward manager\n--------------------------\n\n.. image:: https://github.com/yyDing1/verl-materials/blob/main/distributed_reward_manager.svg?raw=true\n\nHow distributed\n~~~~~~~~~~~~~~~\n\nUnder the single_controller setup, actor rollout and reward computation can be abstracted as:\n\n.. code:: python\n\n   # initalize rollout manager and async reward loop manager\n   async_rollout_manager = AgentLoopManager(config)\n   async_reward_manager = RewardLoopManager(config)\n   # actor rollout using `async_rollout_manager`\n   gen_batch = async_rollout_manager.generate_sequences(batch)\n   # compute reward using `async_reward_manager`\n   reward_batch = async_reward_manager.compute_rm_score(gen_batch)\n\nWithin the ``RewardLoopManager``, multiple ``RewardWorker`` are launched across all nodes to enable distributed reward computation. \nThe number of parallel workers can be configured via ``config.reward.num_workers``.\n\nUpon receiving a batch reward request, the batch is partitioned into smaller chunks and distributed to each reward worker for parallel execution.\nUser only need to invoke ``compute_rm_score``.\n\n.. code:: python\n\n   class RewardLoopManager:\n      \"\"\"\n      RewardLoopManager run in single controller.\n      This class will create reward loop workers and manage them.\n      \"\"\"\n      def _init_reward_loop_workers(self):\n         self.reward_loop_workers = [...]\n\n      def compute_rm_score(self, data):\n         chunks = data.chunk(len(self.reward_loop_workers))\n         outputs = ray.get(\n            [\n               worker.compute_score_batch.remote(chunk)\n               for worker, chunk in zip(self.reward_loop_workers, chunks, strict=True)\n            ]\n         )\n         outputs_flat = [item for sublist in outputs for item in sublist]\n         ...\n\nThis is how the reward manager is parallelized and distributed across all nodes.\n\nStreaming Reward with Rollout\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nFurthermore, we check whether actor rollout and reward computation can be performed in a streaming manner,\nwhere the reward is calculated as soon as each sample is rolled out.\n\n.. code:: python\n\n   # agent_reward_loop: streaming reward computation with actor rollout\n   # two conditions satisfied: (1) rule-based reward, or (2) reward model with extra resource pool\n   enable_agent_reward_loop = not use_rm or config.reward.reward_model.enable_resource_pool\n\n   # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager\n   # to stream reward computation with actor rollout\n   reward_loop_worker_handles = async_reward_manager.reward_loop_workers if enable_agent_reward_loop else None\n   async_rollout_manager = AgentLoopManager(\n      config=config,\n      worker_group=actor_rollout_wg,\n      rollout_resource_pool=actor_rollout_resource_pool,\n      reward_loop_worker_handles=reward_loop_worker_handles,\n   )\n\nHybrid Reward Scenarios Usage\n-----------------------------\n\nAs described above, each ``reward_loop_worker`` is responsible for handling reward requests.\nThe rewards can be categorized as follows:\n\n- **Rule-based Reward**: The reward is determined by predefined rules, e.g., checking whether the predicted answer matches the ground truth via string matching.\n- **Discriminative Reward Model (DisRM)**: The reward is produced by a specified discriminative reward model, such as ``Skywork/Skywork-Reward-Llama-3.1-8B-v0.2``.\n- **Generative Reward Model (GenRM)**: The reward is obtained using a generative reward model, for example ``dyyyyyyyy/FAPO-GenRM-4B``.\n- **Hybrid Reward Scenarios**: A combination of the above reward types, e.g., rule + GenRM.\n\n.. code:: python\n\n   class RewardLoopWorker:\n\n      async def compute_score_batch(self, data: DataProto) -> list[dict]:\n         tasks = []\n         for i in range(len(data)):\n            tasks.append(asyncio.create_task(self.compute_score(data[i : i + 1])))\n         outputs = await asyncio.gather(*tasks)\n         return outputs\n\n      async def compute_score(self, data: DataProto) -> dict:\n         assert len(data) == 1, \"RewardLoopWorker only support single data item\"\n         if self.config.reward.custom_reward_function.path is not None:\n            # directly use user-customized reward function\n            return await self.reward_manager.run_single(data)\n         else:\n            if self.config.reward.reward_model.enable:\n               # we assume the rm is disrm\n               # genrm must set custom_reward_function\n               return await self.compute_score_disrm(data)\n            else:\n               return await self.reward_manager.run_single(data)\n\nEach ``RewardLoopWorker`` will initalize one ``RewardManager``, splits the batch into individual data items and processes them in parallel using asynchronous tasks.\n\nReward Manager\n~~~~~~~~~~~~~~\n\nThe ``RewardManager`` maintains a reward function and defines its computation logic, including:\n\n- **naive**: The simplest implementation.\n- **dapo**: DAPO implementation with an overlong reward penalty.\n- **limit**: Restricts the concurrency of the reward function, useful when external API calls are rate-limited.\n- **remote**: Runs in a separate process, effective for CPU-intensive tasks such as ``Math-Verify``.\n\nUsers can also customize their own ``RewardManager``, inheriting from ``RewardManagerBase``, and implementing the ``run_single`` function.\n\n.. code:: python\n\n   @register(\"user_costomized\")\n   class UserCostomizedRewardManager(RewardManagerBase):\n      async def run_single(self, data: DataProto) -> dict:\n         assert len(data) == 1, \"Only support single data item\"\n         # your own reward manager\n         ...\n\nAfter defining it, users can specify their custom reward manager by setting ``reward.reward_manager.name=user_costomized``.\n\nRule-Based Reward\n~~~~~~~~~~~~~~~~~\n\nIf ``reward.custom_reward_function`` is provided, the user-defined reward function will be used. Otherwise, it falls back to the default reward function.\n\nNote that The custom function can be either synchronous or asynchronous; the system automatically detects its type and loads it accordingly.\n\nWe recommend **using asynchronous functions** when reward computation need to involve external model API calls or sandboxed execution, as they are significantly more efficient.\n\n.. code:: python\n\n   async def compute_score(data_source, solution_str, ground_truth, extra_info):\n      \"\"\"Compute a score by sending an async request to a remote service.\"\"\"\n      \n      # prepare request payload\n      payload = {\"messages\": [{\"role\": \"user\", \"content\": \"check the correcness of the question and response ...\"}], ...}\n\n      # send async HTTP request\n      async with aiohttp.ClientSession() as session:\n         async with session.post(\"https://api.openai.com/v1/chat/completions\", json=payload) as resp:\n               result = await resp.json()\n\n      # parse and return score\n      score = int(result[\"choices\"][0][\"message\"][\"content\"].strip().split(\"\\n\")[-1])\n      return {\"score\": score}\n\nModel-Base Reward\n~~~~~~~~~~~~~~~~~\n\n**For discriminative reward model (DisRM)**, we provide a simple implementation:\n\n.. code:: python\n\n   class RewardLoopWorker:\n      async def compute_score_disrm(self, data) -> dict:\n         disrm_prompt = await self._preprocess_reward_inputs(data)\n         payloads = {\n            \"model\": model_name,\n            \"input\": disrm_prompt,\n            \"activation\": False,\n         }\n         output = await self._post_request(payloads, \"classify\")\n         rm_score = output[\"data\"][-1][\"probs\"][-1]\n         return {\"reward_score\": rm_score}\n\npass the question and the model rollout as inputs to the reward model and obtain a reward score. This is also the standard practice for most DisRM.\n\nUsers should provide ``reward.reward_model.model_path`` to specify the reward model.\n\n**For generative reward model (GenRM)**\n\nFor generative reward model scenarios, users need to specify both ``reward.reward_model.model_path`` and ``reward.custom_reward_function``.\n\nThe custom reward function should implement the following components:\n\n- Convert the question and the model rollout into a GenRM input prompt using a custom prompt template.\n- Invoke the GenRM to perform generation with custom sampling parameters. For this purpose, the Reward Loop provides an HTTP interface (i.e., ``reward_router_address``) for interacting with GenRM.\n- Parse the GenRM output using a custom parser and extract the reward score.\n\nAs these steps are highly customizable and task-dependent, we offer this flexibility entirely to the user-defined reward function.\n\nBelow we provide an example of a custom reward function using GenRM.\n\n.. code:: python\n\n   async def compute_score_gsm8k(\n      data_source: str,\n      solution_str: str,\n      ground_truth: str,\n      extra_info: dict,\n      reward_router_address: str,  # an HTTP router endpoint provided by Reward Loop\n      reward_model_tokenizer: PreTrainedTokenizer,\n   ):\n      \"\"\"Compute the reward score.\"\"\"\n\n      # Step 1: Prepare prompt and request payload\n      grm_prompt = GRM_PROMPT_TEMPLATE.format(problem=extra_info[\"question\"], solution=solution_str)\n      messages = [{\"role\": \"user\", \"content\": grm_prompt}]\n      sampling_params = {\"temperature\": 0.7, \"top_p\": 0.8, \"max_tokens\": 4096}\n      chat_complete_request = {\"messages\": messages, **sampling_params}\n\n      # Step 2: Send async request to the reward model\n      # here, chat_complete sends async http request to the router address\n      result = await chat_complete(\n         router_address=reward_router_address,\n         chat_complete_request=chat_complete_request,\n      )\n\n      # Step 3: Parse model response and extract score\n      grm_response = result.choices[0].message.content.strip()\n      try:\n         score_str = grm_response.split(\"\\n\\n\")[-1].strip()\n         score = int(score_str)\n      except Exception:\n         score = 0\n\n      return {\"score\": score}\n\n**For hybrid reward scenarios**, such as combining rule-based rewards with GenRM similarly as above,\n\n.. _recipe/fapo: https://github.com/verl-project/verl-recipe/tree/main/fapo\n\nA runnable and reproducible example that demonstrates how to use a rule-based reward function together with a GenRM is provided in the `recipe/fapo`_ directory for reference. Welcome to use and cite.\n\nReward Model Arch Design\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nWe support multiple execution modes for reward models during:\n\n- **Colocate Mode**: The reward model shares the same resource pool as the actor/rollout/reference models. In this setup, all rollouts must complete first, after which the reward model is awakened to perform inference.\n- **Standalone Mode**: The reward model runs on a separate resource pool, independent from the actor/rollout/reference models. In this setup, each sample is evaluated by the reward model immediately after its rollout finishes.\n\nThe standalone mode can enable the streaming manner stated above.\n\nBy default, the system runs in colocate mode. Users can enable standalone mode by setting ``reward.reward_model.enable_resource_pool=True`` and allocating the corresponding resources via ``reward.reward_model.nnodes`` and ``reward.reward_model.n_gpus_per_node``.\n\n.. image:: https://github.com/yyDing1/verl-materials/blob/main/reward_loop.svg?raw=true\n\n\nTo support flexible and scalable reward model computation, we implement a reward router that coordinates requests among multiple reward model servers.\n\nEach reward model runs as an independent server and is registered with the router.\nThis router will forward the requests to the registered reward servers with load balancing and return the results.\nThis design allows us to expose a single unified router address to user-defined reward functions, enabling them to access various reward models seamlessly through the same interface.\n\n.. image:: https://github.com/yyDing1/verl-materials/blob/main/reward_loop_full.svg?raw=true\n\n.. code:: python\n\n   class RewardModelManager:\n      \"\"\"Reward model manager.\"\"\"\n\n      def __init__(\n         self,\n         config: RewardModelConfig,\n         resource_pool: RayResourcePool = None,\n      ):\n         \"\"\"\n         Initialize the reward model manager.\n\n         Args:\n            config (RewardModelConfig): Reward model configuration.\n            resource_pool (RayResourcePool, optional): Resource pool. Defaults to None.\n         \"\"\"\n         self.config = config\n         self.resource_pool = resource_pool\n         self._initialize_llm_servers()\n         self._initialize_router()\n\n"
  },
  {
    "path": "docs/advance/rollout_skip.rst",
    "content": "RolloutSkip Function Usage Documentation\n========================================\n\nLast updated: 08/01/2025.\n\nApplicable Scenarios\n--------------------\n\nThe RolloutSkip functionality is designed to accelerate the rollout process in reinforcement learning training by caching and reusing previously generated sequences. This feature is particularly useful when:\n\n1. You need to repeatedly run experiments with the same configuration\n\n2. You want to save time by avoiding redundant sequence generation to come close to the optimal policy\n\n\nAPI and Usage Example\n----------------------\n\n2.1 Trainer Adaptation\n~~~~~~~~~~~~~~~~~~~~~~\n\nBoth`RayDAPOTrainer()` (in `verl/recipe/dapo/dapo_ray_trainer.py`) and `RayPPOTrainer()`(in `verl/trainer/ppo/ray_trainer.py``) have already been adapted.\n\nThis is an example of how to patch rollout_skip in RayPPOTrainer.\n\n.. code-block:: python\n\n    #* Import the RolloutSkip class\n    from verl.utils.rollout_skip import RolloutSkip\n\n    ...\n    class RayPPOTrainer:\n        ...\n        def fit(self):\n            ...\n\n            #* Add code as follow:\n            rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)\n            rollout_skip.wrap_generate_sequences()\n\n            ...\n\n            for epoch in range(self.config.trainer.total_epochs):\n                for batch_dict in self.train_dataloader:\n                    ...\n\n2.2 Basic Configuration\n~~~~~~~~~~~~~~~~~~~~~~~\n\nThen, you should add the following parameters to your config to enable the RolloutSkip feature:\n\n.. code-block:: bash\n\n    actor_rollout_ref.rollout.skip_rollout=True \\\n    actor_rollout_ref.rollout.skip_dump_dir=\"/tmp/rollout_dump\" \\\n\n\nNote:\n\n1. The `skip_dump_dir` is the directory where the cached sequences will be stored. Ensure that this directory is writable and accessible by your training process. And make sure that `skip_dump_dir` is not relative path because ray will store the data in `/tmp/ray/session_<session_id>/` and the relative path will not be found in the worker.\n2. The dumped data path follows this naming pattern `{experiment_name}_{project_name}_TrainGBS{train_gbs}__InferGBS{gen_gbs}__N{n}`, once you change the `experiment_name`, `project_name`, `train_gbs`, `gen_gbs`, or `n`, the cached data will be stored in a new directory.\n"
  },
  {
    "path": "docs/advance/rollout_trace.rst",
    "content": "Trace Function Usage Instructions\n========================================\n\nLast updated: 07/10/2025.\n\nApplicable Scenarios\n--------------------\n\nAgentic RL involves multiple turns of conversations, tool invocations, and user interactions during the rollout process. During the Model Training process, it is necessary to track function calls, inputs, and outputs to understand the flow path of data within the application. The Trace feature helps, in complex multi-round conversations, to view the transformation of data during each interaction and the entire process leading to the final output by recording the inputs, outputs, and corresponding timestamps of functions, which is conducive to understanding the details of how the model processes data and optimizing the training results.\n\nThe Trace feature integrates commonly used Agent trace tools, including wandb weave and mlflow, which are already supported. Users can choose the appropriate trace tool according to their own needs and preferences. Here, we introduce the usage of each tool.\n\n\nTrace Parameter Configuration\n-----------------------------\n\n- ``actor_rollout_ref.rollout.trace.backend=mlflow|weave`` # the trace backend type\n- ``actor_rollout_ref.rollout.trace.token2text=True`` # To show decoded text in trace view\n- ``actor_rollout_ref.rollout.trace.max_samples_per_step_per_worker=N`` # Limit traces per worker (optional)\n\nLimiting Trace Volume\n~~~~~~~~~~~~~~~~~~~~~~\n\nBy default, all samples are traced, which can generate large amounts of data and incur significant costs with trace backends like Weave or MLflow. To limit trace volume while maintaining representative coverage, use ``max_samples_per_step_per_worker``.\n\nExample configuration:\n\n.. code-block:: yaml\n\n   actor_rollout_ref:\n     rollout:\n       trace:\n         backend: weave\n         token2text: False\n         max_samples_per_step_per_worker: 5  # Each worker traces 5 random samples\n\nEach agent loop worker independently selects up to N unique samples to trace per training step. For GRPO (``n > 1``), all rollouts for selected samples are traced. Total traces per step = max_samples_per_step_per_worker * num_workers * n.\n\nExample: With 4 workers, max_samples_per_step_per_worker=5, and GRPO n=4, you get 4 * 5 * 4 = 80 traces per step instead of tracing all samples. Set to null (default) to trace all samples.\n\n\nGlossary\n--------\n\n+----------------+------------------------------------------------------------------------------------------------------+\n| Object         | Explaination                                                                                         |\n+================+======================================================================================================+\n| trajectory     | A complete multi-turn conversation includes:                                                         |\n|                | 1. LLM output at least once                                                                          |\n|                | 2. Tool Call                                                                                         |\n+----------------+------------------------------------------------------------------------------------------------------+\n| step           | The training step corresponds to the global_steps variable in the trainer                            |\n+----------------+------------------------------------------------------------------------------------------------------+\n| sample_index   | The identifier of the sample, defined in the extra_info.index of the dataset. It is usually a number,|\n|                | but may also be a uuid in some cases.                                                                |\n+----------------+------------------------------------------------------------------------------------------------------+\n| rollout_n      | In the GROP algorithm, each sample is rolled out n times. rollout_n represents the serial number of  |\n|                | the rollout.                                                                                         |\n+----------------+------------------------------------------------------------------------------------------------------+\n| validate       | Whether the test dataset is used for evaluation?                                                     |\n+----------------+------------------------------------------------------------------------------------------------------+\n\nRollout trace functions\n-----------------------\n\nThere are 2 functions used for tracing:\n\n1. ``rollout_trace_op``: This is a decorator function used to mark the functions to trace. In default, only few method has it, you can add it to more functions to trace more infor.\n2. ``rollout_trace_attr``: This function is used to mark the entry of a trajectory and input some info to trace. If you add new type of agent, you may need to add it to enable trace.\n\n\nUsage of wandb weave\n--------------------\n\n1.1 Basic Configuration\n~~~~~~~~~~~~~~~~~~~~~~~\n\n1. Set the ``WANDB_API_KEY`` environment variable\n2. Configuration Parameters\n\n   1. ``actor_rollout_ref.rollout.trace.backend=weave``\n   2. ``trainer.logger=['console', 'wandb']``: This item is optional. Trace and logger are independent functions. When using Weave, it is recommended to also enable the wandb logger to implement both functions in one system.\n   3. ``trainer.project_name=$project_name``\n   4. ``trainer.experiment_name=$experiment_name``\n   5. ``actor_rollout_ref.rollout.mode=async``: Since trace is mainly used for agentic RL, need to enable agent toop using async mode for either vllm or sglang.\n\nNote:\nThe Weave Free Plan comes with a default monthly network traffic allowance of 1GB. During the training process, the amount of trace data generated is substantial, reaching dozens of gigabytes per day, so it is necessary to select an appropriate wandb plan.\n\n\n1.2 View Trace Logs\n~~~~~~~~~~~~~~~~~~~\n\nAfter executing the training, on the project page, you can see the WEAVE sidebar. Click Traces to view it.\n\nEach Trace project corresponds to a trajectory. You can filter and select the trajectories you need to view by step, sample_index, rollout_n, and experiment_name.\n\nAfter enabling token2text, prompt_text and response_text will be automatically added to the output of ToolAgentLoop.run, making it convenient to view the input and output content.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/weave_trace_list.png?raw=true\n\n1.3 Compare Trace Logs\n~~~~~~~~~~~~~~~~~~~~~~\n\nWeave can select multiple trace items and then compare the differences among them.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/weave_trace_compare.png?raw=true\n\nUsage of mlflow\n---------------\n\n1. Basic Configuration\n~~~~~~~~~~~~~~~~~~~~~~\n\n1. Set the ``MLFLOW_TRACKING_URI`` environment variable, which can be:\n\n   1. Http and https URLs corresponding to online services\n   2. Local files or directories, such as ``sqlite:////tmp/mlruns.db``, indicate that data is stored in ``/tmp/mlruns.db``. When using local files, it is necessary to initialize the file first (e.g., start the UI: ``mlflow ui --backend-store-uri sqlite:////tmp/mlruns.db``) to avoid conflicts when multiple workers create files simultaneously.\n\n2. Configuration Parameters\n\n   1. ``actor_rollout_ref.rollout.trace.backend=mlflow``\n   2. ``trainer.logger=['console', 'mlflow']``. This item is optional. Trace and logger are independent functions. When using mlflow, it is recommended to also enable the mlflow logger to implement both functions in one system.\n   3. ``trainer.project_name=$project_name``\n   4. ``trainer.experiment_name=$experiment_name``\n\n\n2. View Log\n~~~~~~~~~~~\n\nSince ``trainer.project_name`` corresponds to Experiments in mlflow, in the mlflow view, you need to select the corresponding project name, then click the \"Traces\" tab to view traces. Among them, ``trainer.experiment_name`` corresponds to the experiment_name of tags, and tags corresponding to step, sample_index, rollout_n, etc., are used for filtering and viewing.\n\nFor example, searching for ``\"tags.step = '1'\"`` can display all trajectories of step 1.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/mlflow_trace_list.png?raw=true\n\nOpening one of the trajectories allows you to view each function call process within it.\n\nAfter enabling token2text, prompt_text and response_text will be automatically added to the output of ToolAgentLoop.run, making it convenient to view the content.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/mlflow_trace_view.png?raw=true\n\nNote:\n\n1. mlflow does not support comparing multiple traces\n2. rollout_trace can not associate the mlflow trace with the run, so the trace content cannot be seen in the mlflow run logs.\n"
  },
  {
    "path": "docs/advance/rope.rst",
    "content": "RoPE Scaling override\n=======================================\n\nLast updated: 05/14/2025.\n\nSome models such as `Qwen/Qwen2.5-7B-Instruct <https://huggingface.co/Qwen/Qwen2.5-7B-Instruct#processing-long-texts>`_ support RoPE Scaling but don't have it defined in their config.json file.\nFor example, this model supports this configuration:\n\n.. code:: python\n\n    {\n        ...,\n        \"rope_scaling\": {\n            \"factor\": 4.0,\n            \"original_max_position_embeddings\": 32768,\n            \"type\": \"yarn\"\n        }\n    }\n\n\n\nIn order to support a longer context for such models, you must override the model configs when starting the trainer.\n\nPPO example:\n\n.. code:: bash\n\n    +actor_rollout_ref.model.override_config.rope_scaling.type=yarn \\\n    +actor_rollout_ref.model.override_config.rope_scaling.factor=4.0 \\\n    +actor_rollout_ref.model.override_config.rope_scaling.original_max_position_embeddings=32768 \\\n\n\nAnd for the critic model\n\n.. code:: bash\n\n    +critic.model.override_config.rope_scaling.type=yarn \\\n    +critic.model.override_config.rope_scaling.factor=4.0 \\\n    +critic.model.override_config.rope_scaling.original_max_position_embeddings=32768 \\\n"
  },
  {
    "path": "docs/algo/baseline.md",
    "content": "# Algorithm Baselines\n\nLast updated: 06/18/2025.\n\n## Math related datasets\n\n### GSM8k\n\nAssuming GSM8k/math dataset is preprocessed via:\n\n```bash\npython3 examples/data_preprocess/*.py\n```\n\nRefer to the table below to reproduce RL training from different pre-trained checkpoints. Below is the performance on the GSM8k dataset if not specified otherwise. More comprehensive benchmark results areavailable in the recipe folder.\n\n| Hardware   | Model                            | Method          | Test score   | Details                                                                                                                                                                                                                       |\n| ---------- | -------------------------------- | --------------- | ------------ | ----------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| NVIDIA GPU | google/gemma-2-2b-it             | hf checkpoint   | 23.9         | [Huggingface](https://huggingface.co/google/gemma-2-2b-it#benchmark-results)                                                                                                                                                  |\n| NVIDIA GPU | google/gemma-2-2b-it             | SFT             | 52.06        | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/gemma-2-2b-it-sft-0.411.log)                                                                                                           |\n| NVIDIA GPU | google/gemma-2-2b-it             | SFT + PPO       | 64.02        | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/gemma-2-2b-it-ppo-bsz512_4-prompt1024-resp-512-0.640.log), [wandb](https://api.wandb.ai/links/verl-team/h7ux8602)                      |\n| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct       | hf checkpoint   | 49.6         | [Qwen blog](https://qwen.ai/blog?id=qwen2.5-llm)                                                                                                                                                                              |\n| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct       | PPO             | 56.7         | [command and log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log)                                                                                     |\n| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct       | PRIME           | 58.7         | [script](https://github.com/verl-project/verl-recipe/blob/main//prime/run_prime_qwen.sh), [wandb](https://api.wandb.ai/links/zefan-wang-thu-tsinghua-university/rxd1btvb)                                                     |\n| NVIDIA GPU | Qwen/Qwen2.5-0.5B-Instruct       | GRPO-LoRA       | 54.3         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz64_2-prompt512-resp1024-lorarank32-score0.543.log)                                                                     |\n| NVIDIA GPU | Qwen/Qwen2.5-1.5B-Instruct       | GRPO-LoRA       | 77.9         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-1.5B-bsz64_2-prompt512-resp1024-lorarank32-score0.779.log)                                                                     |\n| NVIDIA GPU | Qwen/Qwen2.5-3B-Instruct         | GRPO-LoRA       | 86.1         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-3B-bsz64_2-prompt512-resp1024-lorarank32-score0.861.log)                                                                       |\n| NVIDIA GPU | deepseek-ai/deepseek-llm-7b-chat | PPO (Megatron)  | 69.5 [1]     | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/deepseek-llm-7b-chat-megatron-bsz256_4-prompt512-resp512-0.695.log), [wandb](https://wandb.ai/verl-team/verl_megatron_gsm8k_examples/runs/10fetyr3) |\n| NVIDIA GPU | Qwen/Qwen2-7B-Instruct           | GRPO            | 89           | [script](https://github.com/volcengine/verl/blob/a65c9157bc0b85b64cd753de19f94e80a11bd871/examples/grpo_trainer/run_qwen2-7b_seq_balance.sh)                                                                                  |\n| NVIDIA GPU | Qwen/Qwen2-7B-Instruct           | GRPO (FSDP2)    | 89.8         | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log)                                                                                                                                 |\n| NVIDIA GPU | Qwen/Qwen2-7B-Instruct           | GRPO (Megatron) | 89.6         | [log](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b_math_megatron.log)                                                                                                                         |\n| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct         | ReMax           | 97           | [script](https://github.com/eric-haibin-lin/verl/blob/main/examples/remax_trainer/run_qwen2.5-3b_seq_balance.sh), [wandb](https://wandb.ai/liziniu1997/verl_remax_example_gsm8k/runs/vxl10pln)                                |\n| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct         | SPPO            | 65.6 (MATH)  | [SPPO script](https://github.com/verl-project/verl-recipe/tree/main/sppo/README.md)                                                                                                                                             |\n| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct         | GRPO-LoRA       | 93.4         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-7B-bsz64_8-prompt512-resp1024-lorarank32-score0.934.log)                                                                       |\n| NVIDIA GPU | Mixtral-8x22B-Instruct-v0.1      | Instruct model  | 83.7         | [Qwen Blog](https://qwen.ai/blog?id=qwen2.5-llm)                                                                                                                                                                              |\n| NVIDIA GPU | Mixtral-8x22B-Instruct-v0.1      | RLOO (Megatron) | 92.3         | [wandb](https://api.wandb.ai/links/ppo_dev/sbuiuf2d)                                                                                                                                                                          |\n| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct         | SPIN            | 92           | [script](https://github.com/verl-project/verl-recipe/tree/main/spin/README.md)                                                                                                                                                  |\n| NVIDIA GPU | Qwen/Qwen2-7B-Instruct           | GPG             | 88           | [log](https://github.com/diqiuzhuanzhuan/verldata/blob/main/run_logs/qwen2-7b_math.log), [wandb](https://wandb.ai/diqiuzhuanzhuan/verl_gpg_example_gsm8k_math/runs/ab86c4va)                                                  |\n| NVIDIA GPU | Qwen/Qwen2-7B-Instruct           | GPG (Megatron)  | 88           | [log](https://github.com/diqiuzhuanzhuan/verldata/blob/main/run_logs/qwen2-7b_math_megatron.log), [wandb](https://wandb.ai/diqiuzhuanzhuan/verl_gpg_example_gsm8k_math/runs/yy8bheu8)                                         |\n| NVIDIA GPU | Qwen/Qwen2.5-VL-7B-Instruct      | GRPO (Megatron) | 65.4 (GEO3k) | [script](https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl-7b-megatron.sh), [wandb](https://api.wandb.ai/links/megatron-core-moe-dev/1yngvkek)                                                |\n| AMD MI300  | deepseek-ai/deepseek-llm-7b-chat | PPO             | 70.5 [1]     | [log](https://github.com/yushengsu-thu/verl_training_log/blob/main/gsm8k/ppo_run_deepseek7b_llm.log)                                                                                                                          |\n| AMD MI300  | deepseek-ai/deepseek-llm-7b-chat | GRPO            | 71.4 [1]     | [log](https://github.com/yushengsu-thu/verl_training_log/blob/main/gsm8k/grpo_run_deepseek7b_llm.log)                                                                                                                         |\n| NVIDIA GPU | Qwen/Qwen2.5-14B-Instruct        | GRPO-LoRA       | 94.6         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-14B-bsz64_8-prompt512-resp1024-lorarank32-score0.946.log)                                                                      |\n| NVIDIA GPU | Qwen/Qwen2.5-32B-Instruct        | GRPO-LoRA       | 95.8         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-32B-bsz64_8-prompt512-resp1024-lorarank32-score0.958.log)                                                                      |\n| NVIDIA GPU | Qwen/Qwen2.5-72B-Instruct        | GRPO-LoRA       | 96.0         | [command and logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-72B-bs64_8-prompt512-resp1024-lorarank32-score0.960.log)                                                                       |\n\n### DAPO math-17k\n\n- Training DAPO math-17k dataset: https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k\n- Testing: AIME'24: https://huggingface.co/datasets/BytedTsinghua-SIA/AIME-2024\n\nNote:\n\n- For Qwen/Qwen2.5-Math-7B, we directly modify the max_position_embeddings to 32768 without observing performance degradation in order to train longer response length.\n\n| Hardware   | Model                      | Method                  | Test score | Details                                                                                                                                                                                             |\n| ---------- | -------------------------- | ----------------------- | ---------- | --------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| NVIDIA GPU | Qwen/Qwen2.5-Math-7B (32k) | DAPO                    | 36.3       | [command](https://github.com/verl-project/verl-recipe/blob/main//dapo/test_dapo_7b_math.sh), [logs](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361) |\n| NVIDIA GPU | Qwen/Qwen2.5-7B-Instruct   | DAPO + Code Interpreter | 40.0       | [command](https://github.com/verl-project/verl-recipe/blob/main//retool/run_qwen2_7b_dapo.sh)                                                                                                       |\n\n## Coding related datasets\n\nBelow is the result on leetcode if not specified otherwise.\n\n| Hardware   | Model                   | Method | Test score | Details                                                                                                                                                                                |\n| ---------- | ----------------------- | ------ | ---------- | -------------------------------------------------------------------------------------------------------------------------------------------------------------------------------------- |\n| NVIDIA GPU | PRIME-RL/Eurus-2-7B-SFT | RPIME  | 36.1       | [script](https://github.com/verl-project/verl-recipe/blob/main//prime/run_prime_qwen_code.sh), [swanlab](https://swanlab.cn/@wangzefan/prime_example/runs/7f541qhspgmy8nmhdlx35/chart) |\n\n### Notes\n\n[1] During evaluation, we have only extracted answers following the format `\"####\"`. A more flexible answer extraction, longer response length, and better prompt engineering may lead to a higher score.\n\n[2] The default value of `actor_rollout_ref.actor.entropy_coeff` is set to `0.0` since verl 0.3.x on 2025-05-30, which is different from previous versions.\n"
  },
  {
    "path": "docs/algo/collabllm.md",
    "content": "# Recipe: CollabLLM \n\nLast updated: 09/22/2025.\n\n> Open-Source Algorithm Implementation & Expriement Running: [Haiquan Chen](https://github.com/chenhaiq), [Shirley Wu](https://github.com/Wuyxin)\n\n🏠 [Homepage](https://aka.ms/CollabLLM) | 📝 [Paper](https://arxiv.org/pdf/2502.00640) | 🤗 [Datasets & Models](https://huggingface.co/collabllm) | ⭐️ [Original Implementation](https://github.com/Wuyxin/collabllm)\n\n`verl` provides a recipe for the Outstanding Paper at ICML 2025, **\"CollabLLM: From Passive Responders to Active Collaborators\"**. [CollabLLM](https://aka.ms/CollabLLM) is a unified fine-tuning framework that optimizes LLMs for effective and efficient multiturn collaboration with users.\n\n**Core Idea:** Models are rewarded based on how well their responses enable effective *future* collaboration with users.\n\nPaper Authors: [Shirley Wu](https://cs.stanford.edu/~shirwu/), [Michel Galley](https://www.microsoft.com/en-us/research/people/mgalley/), Baolin Peng, Hao Cheng, Gavin Li, Yao Dou, Weixin Cai, [James Zou](https://www.james-zou.com/), [Jure Leskovec](https://cs.stanford.edu/people/jure/), [Jianfeng Gao](https://www.microsoft.com/en-us/research/people/jfgao/)\n\n\n---\n## Quick Start\n\n### 0. Environment\nMake sure the required packages for `verl` are installed. Additionally, install `litellm` and export the required API keys. The API model will be used for user simulators and, optionally, LLM Judges (see the Configuration section below).\n\n### 1. Prepare Your Dataset\n\nFirst, process your dataset using the provided script (see example commands and usage in `process_dataset.py`):\n\n```bash\npython process_dataset.py --dataset <> ... --dataset_type <sft or rl>\n```\n\n\n**Requirements:**\n- Input: A Hugging Face multiturn dataset. Existing datasets: `collabllm/collabllm-multiturn-$DATASET`, with `DATASET` in one of [`math-hard(-large)`, `medium(-large)`, `bigcodebench(-large)`] (*-large are the datasets used in the CollabLLM paper)\n- Example format: See [collabllm-multiturn-math-hard](https://huggingface.co/datasets/collabllm/collabllm-multiturn-math-hard)\n- To generate your own dataset: Use [build_dataset.py](https://github.com/Wuyxin/collabllm/blob/main/scripts/engine/build_dataset.py) from the original CollabLLM repository\n\n\n### 2. Train Your Model\n\n**(Optional) For Supervised Fine-Tuning (SFT):**\n```bash\nbash train_sft_collabllm.sh\n```\n\n**For Reinforcement Learning (RL):**\n\n```bash\nbash train_rl_collabllm.sh\n```\n\nThe RL script shows an example to train CollabLLM on `math-hard-large`. \n\n- The config to sample future conversations are in `recipe/collabllm/config/collabllm_interaction_config.yaml`. \n- The Multiturn-aware Reward is aggregated from these three conversational-level rewards:\n\n    ```\n    +reward_model.reward_kwargs.metric_weights.accuracy=1 \\\n    +reward_model.reward_kwargs.metric_weights.interactivity=1 \\\n    +reward_model.reward_kwargs.metric_weights.token_amount=-0.0001 \\\n    ```\n\n    You can remove, add, or modify the weights depending on your task. A list of implemented metrics you can already add are under `recipe/collabllm/metrics`. For example, on `medium-large`, you can replace `accuracy` with `bleu_score` via\n    ```\n    +reward_model.reward_kwargs.metric_weights.bleu_score=1 \n    ```\n    which will instead apply bleu score on the sampled future conversations. \n\n## Algorithm\n\n| Step | Name                          | Description                                                                 |\n|------|-------------------------------|-----------------------------------------------------------------------------|\n| 1    | Model response generation     | The model generates multiple responses for each prompt in a batch.          |\n| 2    | Collaborative simulation      | A user simulator (e.g., GPT or Claude) samples `num_repeat_rollouts` conversations for up to `max_user_turns` additional turns. |\n| 3    | Compute Multiturn-aware Reward | Customized conversational reward functions are applied to the sampled conversations. Rewards are aggregated, then averaged across rollouts. |\n| 4    | Update model                  | The model weights are updated using the computed multiturn-aware rewards.  |\n\n---\n\n## Configuration\n\nThe primary configuration is managed through the launch script `train_rl_collabllm.sh` and the YAML file `recipe/collabllm/config/collabllm_interaction_config.yaml`. Key configuration sections:\n\n| Section              | Key Parameters / Notes                                                                 |\n|----------------------|-----------------------------------------------------------------------------------------|\n| `data`               | Paths to training/validation files, batch sizes, sequence lengths.                      |\n| `actor_rollout_ref` (common) | Base model path (used for actor + initial reference), FSDP settings, optimization (LR, scheduler). |\n| `actor_rollout_ref` (CollabLLM-specific) | Hyperparameters under `actor_rollout_ref.rollout.multi_turn`: `max_user_turns`, `max_assistant_turns`, `num_repeat_rollouts`. |\n| `interaction`        | Defined in `collabllm_interaction_config.yaml`. Specifies user simulator and hyperparameters. Requires exported API keys. |\n| `reward_model`       | Manager set to `collabllm` by default. Modify `reward_model.reward_kwargs.metric_weights` for conversational rewards and weights. LLM Judge hyperparameters (e.g., `model`, `temperature`) go under `reward_model.reward_kwargs.llm_judge_kwargs`. |\n| `algorithm`          | GRPO-specific hyperparameters such as `actor_rollout_ref.rollout.n`.                    |\n| `trainer`            | Distributed training (nodes, GPUs per node), logging (WandB), checkpointing frequency.  |\n\n---\n\n## Key Files\n\n| File Path | Purpose |\n|-----------|---------|\n| `recipe/collabllm/collabllm_agent_loop.py` | Main logic to sample future conversations, using `CollabLLMInteraction` from `verl/interactions/collabllm_interaction.py`. |\n| `verl/workers/reward_manager/collabllm.py` | Computes rewards for future conversations, leveraging `recipe/collabllm/reward_function.py` to apply each metric. |\n\n---\n\n## Acknowledgement\n\nWe sincerely thank the `verl` community and advisors for their contributions and guidance!\n"
  },
  {
    "path": "docs/algo/dapo.md",
    "content": "# Recipe: Decoupled Clip and Dynamic Sampling Policy Optimization (DAPO)\n\nLast updated: 06/19/2025.\n\n> Open-Source Algorithm Implementation & Expriement Running: [Yuxuan Tong](https://tongyx361.github.io/), [Guangming Sheng](https://hk.linkedin.com/in/guangming-sheng-b50640211)\n\n🏠 [Homepage](https://dapo-sia.github.io/) | 📝 [Paper@arXiv](https://arxiv.org/abs/2503.14476) | 🤗 [Datasets&Models@HF](https://huggingface.co/collections/BytedTsinghua-SIA/dapo-67d7f1517ee33c8aed059da0) | 🐱 [Code@GitHub](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo) | 🐱 [Repo@GitHub](https://github.com/BytedTsinghua-SIA/DAPO)\n\n> We propose the **D**ecoupled Clip and Dynamic s**A**mpling **P**olicy **O**ptimization (DAPO) algorithm. By making our work publicly available, we provide the broader research community and society with practical access to scalable reinforcement learning, enabling all to benefit from these advancements. Our system is based on the awesome [verl](https://github.com/volcengine/verl) framework. Thanks for their great work! Applying DAPO training to Qwen2.5-32B base model proves to outperform the previous state-of-the-art DeepSeek-R1-Zero-Qwen-32B on AIME 2024, achieving **50%** accuracy with **50%** less training steps.\n>\n> ![dapo-main-result](https://dapo-sia.github.io/static/images/score.png)\n\n## Quickstart\n\n1. Prepare the datasets **on the Ray cluster**:\n\n```bash\nbash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default\n```\n\n2. Submit the job to the Ray cluster **from any machine**:\n\n```bash\ncd verl # Repo root\nexport RAY_ADDRESS=\"http://${RAY_IP:-localhost}:8265\" # The Ray cluster address to connect to\nexport WORKING_DIR=\"${PWD}\" # The local directory to package to the Ray cluster\n# Set the runtime environment like env vars and pip packages for the Ray cluster in yaml\nexport RUNTIME_ENV=\"./recipe/dapo/runtime_env.yaml\" # This sets environment variables for the Ray cluster\nbash recipe/dapo/run_dapo_qwen2.5_32b.sh # or other scripts\n```\n\n## Reproduction Runs\n\n| Setup                                        | AIME 2024 Acc. | Hardware  | Image                                                                | Commit                                                                                       | Environment Variables                                                                                                             | Training Script                                                                                                                                             | Training Record                                                                           |\n| -------------------------------------------- | -------------- | --------- | -------------------------------------------------------------------- | -------------------------------------------------------------------------------------------- | --------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------------------------------------------------------------------------- | ----------------------------------------------------------------------------------------- |\n| DAPO                                         | 52%            | 16x8xH800 | `hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0` | [`4f80e4`](https://github.com/volcengine/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_qwen2.5_32b.sh](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_qwen2.5_32b.sh)             | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) |\n| DAPO w/o Dynamic Sampling                    | 50%            | 16x8xH800 | `hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.3-flashinfer0.2.2-cxx11abi0` | [`4f80e4`](https://github.com/volcengine/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_wo_ds_qwen2.5_32b.sh](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_wo_ds_qwen2.5_32b.sh) | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) |\n| DAPO w/o Token-level Loss & Dynamic Sampling | 44%            | 16x8xH20  | `hiyouga/verl:ngc-th2.5.1-cu120-vllm0.7.4-hotfix`                    | [`4f80e4`](https://github.com/volcengine/verl/tree/4f80e465c2ec79ab9c3c30ec74b9745de61d0490) | [runtime_env.yaml](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/runtime_env.yaml) | [run_dapo_early_qwen2.5_32b.sh](https://github.com/volcengine/verl/blob/4f80e465c2ec79ab9c3c30ec74b9745de61d0490/recipe/dapo/run_dapo_early_qwen2.5_32b.sh) | [W&B](https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=wmb4qxfht0n) |\n\n> [!IMPORTANT]\n>\n> **📢 Call for Contribution!**\n>\n> Welcome to submit your reproduction runs and setups!\n\n## Configuration\n\n### Separated Clip Epsilons (-> Clip-Higher)\n\nAn example configuration:\n\n```yaml\nactor_rollout_ref:\n  actor:\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.28\n```\n\n`clip_ratio_low` and `clip_ratio_high` specify the $\\varepsilon_{\\text {low }}$ and $\\varepsilon_{\\text {high }}$ in the DAPO objective.\n\nCore relevant code:\n\n```python\npg_losses1 = -advantages * ratio\npg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high)\npg_losses = torch.maximum(pg_losses1, pg_losses2)\n```\n\n### Dynamic Sampling (with Group Filtering)\n\nAn example configuration:\n\n```yaml\ndata:\n  gen_batch_size: 1536\n  train_batch_size: 512\nalgorithm:\n  filter_groups:\n    enable: True\n    metric: acc # score / seq_reward / seq_final_reward / ...\n    max_num_gen_batches: 10 # Non-positive values mean no upper limit\n```\n\nSetting `filter_groups.enable` to `True` will filter out groups whose outputs' `metric` are all the same, e.g., for `acc`, groups whose outputs' accuracies are all 1 or 0.\n\nThe trainer will repeat sampling with `gen_batch_size` until there are enough qualified groups for `train_batch_size` or reaching the upper limit specified by `max_num_gen_batches`.\n\nCore relevant code:\n\n```python\nprompt_bsz = self.config.data.train_batch_size\nif num_prompt_in_batch < prompt_bsz:\n    print(f'{num_prompt_in_batch=} < {prompt_bsz=}')\n    num_gen_batches += 1\n    max_num_gen_batches = self.config.algorithm.filter_groups.max_num_gen_batches\n    if max_num_gen_batches <= 0 or num_gen_batches < max_num_gen_batches:\n        print(f'{num_gen_batches=} < {max_num_gen_batches=}. Keep generating...')\n        continue\n    else:\n        raise ValueError(\n            f'{num_gen_batches=} >= {max_num_gen_batches=}. Generated too many. Please check your data.'\n        )\nelse:\n    # Align the batch\n    traj_bsz = self.config.data.train_batch_size * self.config.actor_rollout_ref.rollout.n\n    batch = batch[:traj_bsz]\n```\n\n### Flexible Loss Aggregation Mode (-> Token-level Loss)\n\nAn example configuration:\n\n```yaml\nactor_rollout_ref:\n  actor:\n    loss_agg_mode: \"token-mean\" # / \"seq-mean-token-sum\" / \"seq-mean-token-mean\"\n    # NOTE: \"token-mean\" is the default behavior\n```\n\nSetting `loss_agg_mode` to `token-mean` will mean the (policy gradient) loss across all the tokens in all the sequences in a mini-batch.\n\nCore relevant code:\n\n```python\nif loss_agg_mode == \"token-mean\":\n    loss = verl_F.masked_mean(loss_mat, loss_mask)\nelif loss_agg_mode == \"seq-mean-token-sum\":\n    seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)  # token-sum\n    loss = torch.mean(seq_losses)  # seq-mean\nelif loss_agg_mode == \"seq-mean-token-mean\":\n    seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / torch.sum(loss_mask, dim=-1)  # token-mean\n    loss = torch.mean(seq_losses)  # seq-mean\nelse:\n    raise ValueError(f\"Invalid loss_agg_mode: {loss_agg_mode}\")\n```\n\n### Overlong Reward Shaping\n\nAn example configuration:\n\n```yaml\ndata:\n  max_response_length: 20480 # 16384 + 4096\nreward_model:\n  overlong_buffer:\n    enable: True\n    len: 4096\n    penalty_factor: 1.0\n```\n\nSetting `overlong_buffer.enable` to `True` will penalize the outputs whose lengths are overlong but still within the hard context limit.\n\nSpecifically, the penalty increases linearly from `0` to `overlong_buffer.penalty_factor` when the length of the output exceeds the `max_response_length - overlong_buffer.len` by `0` to `overlong_buffer.len` tokens.\n\nCore relevant code:\n\n```python\nif self.overlong_buffer_cfg.enable:\n    overlong_buffer_len = self.overlong_buffer_cfg.len\n    expected_len = self.max_resp_len - overlong_buffer_len\n    exceed_len = valid_response_length - expected_len\n    overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor\n    overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)\n    reward += overlong_reward\n```\n\n## FAQ\n\n### Where is the \"Overlong Filtering\" in the paper?\n\nMost experiments in the paper, including the best-performant one, are run without Overlong Filtering because it's somehow overlapping with Overlong Reward Shaping in terms of properly learning from the longest outputs. So we don't implement it here.\n\n### What's the difference between [the `recipe/dapo` directory in the `main` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo) and the [`recipe/dapo` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo)?\n\n[The `recipe/dapo` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo/recipe/dapo) is for **as-is reproduction** and thus won't be updated with new features.\n\n[The `recipe/dapo` directory in the `main` branch](https://github.com/verl-project/verl-recipe/tree/main/dapo) works as an example of how to extend the latest `verl` to implement an algorithm recipe, which will be maintained with new features.\n\n### Why can't I produce similar results after modifications?\n\nRL infrastructures nowadays still have inherent unrobustness, on which we are still working hard to improve.\n\nWe strongly recommend to only modify one thing at a time.\n\nWe also list some known problems here:\n\n1. Enabling CUDA graph (`enforce_eager=False`) might cause model performance degradation, whose cause is still under investigation.\n"
  },
  {
    "path": "docs/algo/dppo.md",
    "content": "# Divergence Proximal Policy Optimization (DPPO)\n\nLast updated: 02/25/2026.\n\n\n<div align=\"center\">\n\n## Rethinking the Trust Region in LLM Reinforcement Learning\n\n[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white )](https://arxiv.org/pdf/2602.04879)\n[![Github](https://img.shields.io/badge/Stable_RL-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/sail-sg/Stable-RL)\n[![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/QPHutu/status/2019435642539897303)\n\n</div>\n\n\n## ✨Getting started\n\n1. Prepare the datasets by running [prepare_dapo_data.sh](https://github.com/verl-project/verl-recipe/blob/3490a22a0a3adeb7e4787fe70b1060b642efbae4/dapo/prepare_dapo_data.sh):\n\n```bash\nbash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default\n```\n\n2. Prepare the model:\n\n```bash\nhf download Qwen/Qwen3-30B-A3B-Base --local-dir ${HOME}/verl/models/Qwen3-30B-A3B-Base\n```\n\n3. Run the script:\n\n```bash\n# run DPPO-Binary-KL\nLOSS_MODE=dppo_kl bash examples/dppo_trainer/run_qwen30b_dppo.sh\n\n# run DPPO-Binary-TV\nLOSS_MODE=dppo_tv bash examples/dppo_trainer/run_qwen30b_dppo.sh\n\n# run GRPO baseline\nLOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.2 bash examples/dppo_trainer/run_qwen30b_dppo.sh\n# or GRPO with clip higher\nLOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.28 bash examples/dppo_trainer/run_qwen30b_dppo.sh\n```\n\n## 📖Introduction\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/ppo_vs_dppo.jpg?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nComparison of **PPO** and the proposed **DPPO** (the Binary-TV variant). **(Left)** The surrogate objective and corresponding masks for PPO and DPPO. PPO (and variants like GRPO) employs a heuristic mask based on the probability ratio. In contrast, DPPO utilizes a more principled mask based on a direct approximation of policy divergence (e.g., Total Variation), ensuring updates stay within a theoretically grounded trust region. **(Right)** Experimental results on the AIME24 using Qwen3-30B-A3B-Base. DPPO significantly outperforms GRPO baselines, achieving superior training stability and final performance even without rollout routing replay (R3).\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/sanity_test.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nDPPO variants achieve stable training while controlling the training-inference mismatch at a low level. In contrast, methods without a trust region (PG-IS, CISPO) or with a misspecified one (MiniRL) suffer from growing mismatch and eventual collapse.\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/moe_prob_ratio_tv.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nThe plots show numerical differences between a training and an inference engine for Qwen3-30B-A3B-Base with identical parameters. **(Left)** The probability ratio (used in PPO) is highly volatile for low-probability tokens. **(Right)** In contrast, the TV divergence is more stable. This highlights a key flaw of PPO's clipping mechanism: it **over-penalizes low-probability tokens**, which can slow down learning; and **under-penalizes high-probability tokens**, which can permit large, destabilizing updates.\n\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/clipped_tokens.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nThe most frequently clipped tokens (by GRPO) are important to the reasoning task! \nThey are dominated by:\n- numbers, like 1, 4\n- mathematical symbols, like +, -, =\n- reasoning and structural Words: Wait, Thus, Next\n\n## Top-K divergence approximation\n\nWe only implement the DPPO-Binary-TV/DPPO-Binary-KL here due to their simplicity.\n\nFor the TopK divergence approximation, please refer to the [the original repo](https://github.com/sail-sg/Stable-RL) for a complete implementation.\n\n## Citation\nIf you find our works useful for your research, please consider citing:\n\n```bibtex\n@article{qi2026dppo,\n  title={Rethinking the Trust Region in LLM Reinforcement Learning},\n  author={Qi, Penghui and Zhou, Xiangxin and Liu, Zichen and Pang, Tianyu and Du, Chao and Lin, Min and Lee, Wee Sun},\n  journal={arXiv preprint arXiv:2602.04879},\n  year={2026}\n}\n```\n\n## 🌻Acknowledgement\nWe implement our reinforcement learning algorithm extending from [verl](https://github.com/volcengine/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) and [sglang](https://github.com/sgl-project/sglang) for inference. Our models are trained primarily on [Qwen3 family](https://huggingface.co/collections/Qwen/qwen3). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions!\n"
  },
  {
    "path": "docs/algo/entropy.md",
    "content": "# Recipe: Entropy Mechanism\n\nLast updated: 06/27/2025.\n\n\n<div align=\"center\">\n\n  The Entropy Mechanism of Reinforcement Learning for Large Language Model Reasoning.\n\n[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white)](https://arxiv.org/pdf/2505.22617)  [![Github](https://img.shields.io/badge/PRIME-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/PRIME-RL/Entropy-Mechanism-of-RL) [![alphaXiv](https://img.shields.io/badge/discussion-A42C25?style=for-the-badge&logo=arxiv&logoColor=white&color=blue\n)](https://www.alphaxiv.org/abs/2505.22617) [![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/stingning/status/1928088554166505667) [![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/charlesfornlp/status/1928089451080585283) [![Twitter-ak](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/_akhaliq/status/1928077929105268861)\n\n\n<div align=\"center\" style=\"font-family: Arial, sans-serif;\">\n  <p>\n    <a href=\"#🎉news\" style=\"text-decoration: none; font-weight: bold;\">🎉 News</a> •\n    <a href=\"#✨getting-started\" style=\"text-decoration: none; font-weight: bold;\">✨ Getting Started</a> •\n    <a href=\"#📖introduction\" style=\"text-decoration: none; font-weight: bold;\">📖 Introduction</a>\n  </p>\n  <p>\n    <a href=\"#🎈citation\" style=\"text-decoration: none; font-weight: bold;\">🎈 Citation</a> •\n    <a href=\"#🌻acknowledgement\" style=\"text-decoration: none; font-weight: bold;\">🌻 Acknowledgement</a> •\n    <a href=\"#📬Contact\" style=\"text-decoration: none; font-weight: bold;\">📬 Contact</a> •\n    <a href=\"#📈star-history\" style=\"text-decoration: none; font-weight: bold;\">📈 Star History</a>\n  </p>\n</div>\n\n</div>\n\n\n## 🎉News\n\n- **[2025/05/29]** 🎉 Ranked **#1** of the day on [Huggingface Daily Papers](https://huggingface.co/papers?date=2025-05-29).\n- **[2025/05/29]** Released our Paper on arXiv. See [here](https://arxiv.org/pdf/2505.22617). We provide insights into the entropy mechanism of RL for LLMs and propose two simple yet effective strategies to alleviate the entropy collapse. \n\n\n\n## ✨Getting started\n\nAfter preparing the training data, for training Qwen2.5-7B on a single node, taking the KL-Cov approach as an example, you can simply run:\n\n```\ncd verl\nconda activate your_env\nbash recipe/dapo/7b_kl_cov.sh\n```\n\nWhile for training Qwen2.5-32B on multi nodes, you can run the following commands:\n\n```\ncd verl\nconda activate your_env\nbash recipe/dapo/32b_kl_cov.sh\n```\n\n## 📖Introduction\n\n<div align=\"left\">\n  <img src=\"https://github.com/PRIME-RL/Entropy-Mechanism-of-RL/blob/main/figures/e2a.jpg?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nThis paper addresses the entropy collapse issue in scaling reinforcement learning (RL) for large language models (LLMs), where policy entropy drops sharply during training, leading to overconfidence and performance saturation. We empirically establish a relationship between entropy ($H$) and performance ($R$): $R=−aexp(H)+b$, showing performance is bottlenecked by entropy exhaustion. \n\n<div align=\"left\">\n  <img src=\"https://github.com/PRIME-RL/Entropy-Mechanism-of-RL/blob/main/figures/cov.jpg?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nTheoretically, we find entropy changes are driven by the covariance between action probability and logit updates, which correlates with advantage in Policy Gradient methods. High-probability, high-advantage actions reduce entropy, while rare, high-advantage actions increase it. Empirically, the covariance term remains positive, explaining entropy’s monotonic decline. To mitigate this, we propose ​​Clip-Cov​​ and ​​KL-Cov​​, which restrict updates for high-covariance tokens. These methods effectively prevent entropy collapse, and improve performance. \n\n## 📃Evaluation\n\n<div align=\"left\">\n  <img src=\"https://github.com/PRIME-RL/Entropy-Mechanism-of-RL/blob/main/figures/performance_fig.jpg?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\n\nOur method is able to maintain a considerably higher level of entropy throughout training. For example, when the baseline's entropy reaches a plateau and can no longer be consumed, the KL-Cov method still sustains an entropy level over 10 times higher. Meanwhile, the response length of the policy model steadily increases, and its performance on the test set consistently surpasses that of the baseline. This indicates that our model is able to explore more freely during training, learning better policy through RL. \n| **Method**        | **AIME24** | **AIME25** |  **AMC** | **MATH-500** | **OMNI-MATH** | **OlympiadBench** | **Minerva** | **Avg.** |\n| ----------------- | ---------: | ---------: | -------: | -----------: | ------------: | ----------------: | ----------: | -------: |\n| *Qwen2.5-7B*      |            |            |          |              |               |                   |             |          |\n| GRPO              |       21.2 |        9.6 |     58.7 |         78.8 |          27.9 |              40.7 |        36.7 |     38.6 |\n| w. Clip-higher    |       18.1 |       11.5 |     56.6 |         79.2 |          29.8 |              43.3 |        40.4 |     38.8 |\n| w. **`CLIP-Cov`** |       22.1 |   **15.8** |     58.2 |         80.4 |      **30.5** |          **44.1** |    **41.1** |     40.4 |\n| w. **`KL-Cov`**   |   **22.6** |       12.9 | **61.4** |     **80.8** |          29.1 |              42.6 |        38.2 | **40.6** |\n| *Qwen2.5-32B*     |            |            |          |              |               |                   |             |          |\n| GRPO              |       21.8 |       16.2 |     69.7 |         84.2 |          35.2 |              43.6 |        45.5 |     45.8 |\n| w. Clip-higher    |       35.6 |       22.3 |     69.5 |         77.2 |          35.1 |              42.5 |        43.0 |     47.2 |\n| w. **`CLIP-Cov`** |       32.3 |       22.7 |     67.2 |     **87.0** |      **42.0** |          **57.2** |        46.0 |     50.3 |\n| w. **`KL-Cov`**   |   **36.8** |   **30.8** | **74.5** |         84.6 |          39.1 |              49.0 |    **46.3** | **52.2** |\n\nOur two approaches both achieve non-trivial improvements across all benchmarks. Compared to GRPO, our method outperforms it by 2.0% on average for the 7B model and by 6.4% for the 32B model. Moreover, we observe that our method yields more substantial gains on the larger Qwen2.5-32B. Specifically, our method achieves improvements of 15.0% and 14.6% compared to GRPO on the most challenging benchmarks, AIME24 and AIME25, respectively.\n\n\n## 🎈Citation\nIf you find this paper or repo helpful, please cite us.\n\n```bibtex\n@article{cui2025entropy,\n  title={The Entropy Mechanism of Reinforcement Learning for Reasoning Language Models},\n  author={Cui, Ganqu and Zhang, Yuchen and Chen, Jiacheng and Yuan, Lifan and Wang, Zhi and Zuo, Yuxin and Li, Haozhan and Fan, Yuchen and Chen, Huayu and Chen, Weize and others},\n  journal={arXiv preprint arXiv:2505.22617},\n  year={2025}\n}\n```\n## 🌻Acknowledgement\nWe implement our reinforcement learning algorithm extending from [verl](https://github.com/volcengine/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) for inference. Our models are trained primarily on [Qwen2.5 family](https://github.com/QwenLM/Qwen2.5). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions!\n\n## 📬 Contact\n\nFor questions, discussion, or collaboration opportunities, feel free to contact:\n- Ganqu Cui: cuiganqu@pjlab.org.cn\n- Yuchen Zhang: yuchen.zhang2003@gmail.com\n- Jiacheng Chen: jackchan9345@gmail.com\n- Ning Ding: ningding.cs@gmail.com\n\n"
  },
  {
    "path": "docs/algo/gpg.md",
    "content": "# GPG: Group Policy Gradient\n\nLast updated: 07/03/2025.\n\nGroup Policy Gradient (GPG) is a minimalist reinforcement learning (RL) method that enhances the reasoning ability of large language models without relying on supervised fine-tuning or complex tricks. GPG revisits traditional policy gradients and directly optimizes the RL objective—no surrogate losses, no KL penalties, no critic, and no reference model. Compared to GRPO, GPG is simpler, more efficient, and achieves better results on many tasks. For more details, please refer to the original paper [GPG: A Simple and Strong Reinforcement Learning Baseline for Model Reasoning\n](https://arxiv.org/abs/2504.02546).\n\n## Key Components\n- Use a corrected advantage function to improve policy gradient accuracy and training efficiency.\n- By eliminating the critic and reference models, avoiding KL divergence constraints, significantly simplifies the training process compared to Group Relative Policy Optimization (GRPO)\n\n## Configuration\nTo configure GPG within the framework, use the following YAML settings.\n\n```yaml\nalgorithm:\n  adv_estimator: gpg \nactor_rollout_ref:\n  actor:\n    policy_loss:\n      loss_mode: \"gpg\"\n```\n\n## Advanced Extensions\nGPG is a simple and strong baseline for model reasoning. Although it avoids using KL loss in its original form, you can still use KL loss to further improve the performance.\n\n```yaml\nalgorithm:\n  adv_estimator: gpg\nactor_rollout_ref:\n  actor:\n    use_kl_loss: True # enable kl regularization\n    kl_loss_coef: 0.01\n    policy_loss:\n      loss_mode: \"gpg\"\n```"
  },
  {
    "path": "docs/algo/grpo.md",
    "content": "# Group Relative Policy Optimization (GRPO)\n\nLast updated: 05/31/2025.\n\nIn reinforcement learning, classic algorithms like PPO rely on a \"critic\" model to estimate the value of actions, guiding the learning process. However, training this critic model can be resource-intensive. \n\nGRPO simplifies this process by eliminating the need for a separate critic model. Instead, it operates as follows:\n- Group Sampling: For a given problem, the model generates multiple possible solutions, forming a \"group\" of outputs.\n- Reward Assignment: Each solution is evaluated and assigned a reward based on its correctness or quality.\n- Baseline Calculation: The average reward of the group serves as a baseline. \n- Policy Update: The model updates its parameters by comparing each solution's reward to the group baseline, reinforcing better-than-average solutions and discouraging worse-than-average ones.\n\nThis approach reduces computational overhead by avoiding the training of a separate value estimation model, making the learning process more efficient. For more details, refer to the original paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://arxiv.org/pdf/2402.03300)\n\n## Key Components\n\n- No Value Function (Critic-less): unlike PPO, GRPO does not train a separate value network (critic)\n- Group Sampling (Grouped Rollouts): instead of evaluating one rollout per input, GRPO generates multiple completions (responses) from the current policy for each prompt. This set of completions is referred to as a group.\n- Relative Rewards: within each group, completions are scored (e.g., based on correctness), and rewards are normalized relative to the group.\n\n## Configuration\n\nNote that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior.\n\nDespite that many configurations start with the `ppo_` prefix, they work across different RL algorithms in verl, as the GRPO training loop is similar to that of PPO (without critic).\n\n![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d)\n\n- `actor_rollout.ref.rollout.n`: For each prompt, sample n times. Default to 1. For GRPO, please set it to a value larger than 1 for group sampling.\n\n- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n`\n\n- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers.\n\n- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for GRPO updates on one set of sampled trajectories for actor\n\n- `actor_rollout_ref.actor.clip_ratio`: The GRPO clip range. Default to 0.2\n\n- `algorithm.adv_estimator`: Default is gae. Please set it to grpo instead\n\n- `actor_rollout_ref.actor.loss_agg_mode`: Default is \"token-mean\". Options include \"token-mean\", \"seq-mean-token-sum\", \"seq-mean-token-mean\". The original GRPO paper takes the sample-level loss (seq-mean-token-mean), which may be unstable in long-CoT scenarios. All GRPO example scripts provided in verl uses the default configuration \"token-mean\" for loss aggregation instead.\n\nInstead of adding KL penalty in the reward, GRPO regularizes by directly adding the KL divergence between the trained policy and the reference policy to the loss:\n\n- `actor_rollout_ref.actor.use_kl_loss`: To use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False. Please set it to True for GRPO.\n\n- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001.\n\n- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending \"+\" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/volcengine/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\n## Advanced Extensions\n\n### DrGRPO\n\n[Understanding R1-Zero-Like Training: A Critical Perspective](https://arxiv.org/pdf/2503.20783) claims there's optimization bias in GRPO, which leads to artificially longer responses, especially for incorrect outputs. This inefficiency stems from the way GRPO calculates advantages using group-based reward normalization. Instead, DrGRPO aggregates token-level losses by normalizing with a global constant to eliminate length bias.\n\nConfigure the following to enable DrGRPO, with all other parameters the same as GRPO's:\n\n- `actor_rollout_ref.actor.loss_agg_mode`: \"seq-mean-token-sum-norm\", which turns off seq-dim averaging\n- `actor_rollout_ref.actor.loss_scale_factor`: (Optional) Set to a constant integer (e.g., max response length) to ensure consistent normalization throughout training. If not set, uses the current batch's response length.\n- `actor_rollout_ref.actor.use_kl_loss`: Please set it to False for DrGRPO\n- `algorithm.norm_adv_by_std_in_grpo`: False, which turns off standard deviation norm\n\n## Reference Example\n\nQwen2.5 GRPO training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log)\n\n```bash\nbash examples/grpo_trainer/run_qwen3-8b.sh\n```\n\nFor more reference performance, please see https://verl.readthedocs.io/en/latest/algo/baseline.html\n"
  },
  {
    "path": "docs/algo/opo.md",
    "content": "# On-Policy RL with Optimal Reward Baseline (OPO)\n\nLast updated: 06/02/2025.\n\nLoose on-policy constraints and suboptimal baselines in reinforcement learning often lead to training instability such as large policy shifts and entropy collapse. OPO addresses these challenges by using exact on-policy training with the theretically optimal reward baseline for advantage estimation. It achieves lower policy shifts and higher output entropy, encouraging more diverse and less repetitive responses.\n\nOPO uses group sampling to generate multiple outputs for each input like GRPO. Unlike group-based algorithms which typically use the mean reward of a group as its baseline, OPO employs a theoretically optimal baseline: the length-weighted reward of the group. It also  omits the standard deviation normalization. By adopting these two key components, OPO enables the training of a single policy model with the objective of maximizing only the expected reward. For more detailes, refer to the original paper [On-Policy RL with Optimal Reward Baseline](https://arxiv.org/pdf/2505.23585).\n\n## Key Components\n\n- Exact On-Policy Training: always generates responses from the current policy, without using any pre-generated data or off-policy data.\n- Optimal Reward Baseline: uses a length-weighted reward of the group as the baseline for normalizing the rewards.\n\n## Configuration\n\nTo configure OPO within the framework, use the following YAML settings. These parameters are crucial for enabling exact on-policy training and activating the optimal reward baseline.\n\n```yaml\nalgorithm:\n  adv_estimator: opo  # Use OPO for optimal reward baseline \ndata:\n  train_batch_size: 1024\nactor_rollout_ref:\n  actor:\n    ppo_mini_batch_size: 1024 # ppo_mini_batch_size should equal to train_batch_size to enable exact on-policy training\n    entropy_coeff: 0 # disable entropy regularization\n    use_kl_loss: False # disable kl regularization\n    kl_loss_coef: 0 \n```\n\n## Advanced Extensions\n\nOPO can also be extended to other algorithms like RLOO and Reinforce++. It just needs to adjust their configurations to enable exact on-policy training and incorporate the optimal length-weighted reward baseline with minimal modifications to their advantage estimation functions.\n"
  },
  {
    "path": "docs/algo/otb.md",
    "content": "# Optimal Token Baseline (OTB)\n\nLast updated: 02/23/2026.\n\n📝 [ArXiv](https://www.arxiv.org/abs/2602.07078) | 📒 [Blog](https://richardli.xyz/optimal-token-baseline) | 🤗 [Datasets](https://huggingface.co/datasets/Jiawei415/DPAO_filter) \n\nOptimal Token Baseline (OTB) is a dynamic token-level baseline for gradient variance reduction in policy-gradient reinforcement learning. It weights updates with the \"Realized Energy\" statistic that tracks how much uncertainty has accumulated up to each token, so noisy regions get downweighted while confident regions carry more weight.\n\n## Key properties\n\n- _Token-level baselines:_ OTB adapts per token by tracking realized energy, avoiding the padding artifacts that appear when group means dilute the signal with `EOS` tokens.\n- _Forward-only overhead:_ The realized-energy statistic is computed via the **Logit-Gradient Proxy**, so OTB requires no extra backward passes or gradient-norm kernels.\n\n## Logit-Gradient Proxy\n\nComputing true uncertainty per token would normally mandate per-token backward passes. OTB sidesteps this by estimating realized energy entirely from forward probabilities, so it introduces negligible runtime overhead in practice.\n\n## Mechanics at a glance\n\nFor each prompt group of size `N`, OTB computes rewards-to-go `G_t` and cumulative variance weights `W_t`. The optimal baseline per token is\n\n```\nB*_t = (Σ_i G_t^{(i)} · W_t^{(i)}) / (Σ_i W_t^{(i)} + ε),\nW_t = Σ_{j=1}^t (1 - 2π_j + Σπ_j²),\nΣπ_j² = exp(logsumexp(2·logits_j) - 2·logsumexp(logits_j)).\n```\n\nThe final advantage is `(G_t - B*_t) · mask_t`, so padding tokens stay at zero.\n\n## Integration in VERL\n\n- `AdvantageEstimator.OPTIMAL_TOKEN_BASELINE` registers `compute_optimal_token_baseline_advantage`, invoked whenever `algorithm.adv_estimator` is set to `optimal_token_baseline`.\n- `ActorRolloutRefWorker.compute_log_prob` emits an additional tensor `sum_pi_squared` (Σπ² per token) when `actor.calculate_sum_pi_squared=True`. This requires disabling fused log-prob kernels, because they do not surface logits.\n- Trainers assert `sum_pi_squared` exists, regroup trajectories by `non_tensor_batch[\"uid\"]`, and run the OTB calculation. If rollout IS is active, they rescale the weights by `rollout_is_weights**2` before aggregating.\n- In Ulysses sequence-parallel setups, the actor gathers, unpads, and returns Σπ² in the same way it handles log-probabilities, so OTB supports sharded sequence-parallel models out of the box.\n- `sum_pi_squared_checkpointing` is available to trade compute for memory when Σπ² tensors become large (e.g., lengthy chain-of-thought reasoning).\n\n## Configuration checklist\n\n- `actor_rollout_ref.actor.calculate_sum_pi_squared: true` (mandatory).\n- `actor_rollout_ref.model.use_fused_kernels: false` (required until fused kernels emit logits).\n- `algorithm.adv_estimator: optimal_token_baseline` for single-turn RL and `tir_optimal_token_baseline` for multi-turn RL.\n- Group sampling (`actor_rollout_ref.rollout.n > 1`) to unlock OTB’s variance reduction; with `n=1` the baseline collapses to returns.\n\nExample OmegaConf overlay:\n\n```yaml\nalgorithm:\n  adv_estimator: optimal_token_baseline\n\nactor_rollout_ref:\n  actor:\n    calculate_sum_pi_squared: true\n    sum_pi_squared_checkpointing: false # optional memory saver\n  rollout:\n    n: 8\n```\n\n## Example script\n\nSee `examples/otb_trainer/run_qwen2_5-7b.sh` for a reference training loop.\n\n## Gradient Variance Proxy Metrics\n\nAll gradient-variance analysis in the Optimal Token Baseline work starts from the variance identity\n\n```\nVar(ĝ) = E[||ĝ||²] - ||E[ĝ]||²,\n```\n\nwhich states that the variance of any stochastic gradient equals the mean squared magnitude minus the squared norm of its expectation.\n\nFor a trajectory `τ`, the policy-gradient estimator is\n\n```\nĝ(τ) = ∇ log π_θ(τ) · A(τ),        A(τ) = R(τ) - B.\n```\n\nThe logit-gradient proxy approximates the squared gradient norm without an extra backward pass:\n\n```\n||ĝ(τ)||² ≈ Ŵ(τ) · A(τ)²,\n```\n\nwhere `Ŵ(τ)` is the realized energy built. Given a mini-batch `{τ_i}` of size `N`, we decompose its statistics into three diagnostics:\n\n- **Signal strength (squared norm of the mean gradient)**\n  ```\n  S = || (1/N) · Σ ĝ(τ_i) ||²\n  ```\n- **Total power (signal + noise)**\n  ```\n  P_total = (1/N) · Σ Ŵ(τ_i) · A(τ_i)²\n  ```\n- **Pure noise (estimated variance of the batch mean)**\n  ```\n  Var_proxy = (1/(N-1)) · (P_total - S)\n  ```\n\n`verl/trainer/ppo/metric_utils.py#L306` implements these diagnostics via `compute_variance_proxy_metrics`, emitting `variance_proxy/proxy1_signal_strength`, `variance_proxy/proxy2_total_power`, and `variance_proxy/proxy3_pure_noise`.\n\nTracking these metrics provides a forward-only, low-overhead view of gradient health for any advantage estimator that supplies `sum_pi_squared`.\n"
  },
  {
    "path": "docs/algo/ppo.md",
    "content": "# Proximal Policy Optimization (PPO)\n\nLast updated: 06/19/2025.\n\nProximal Policy Optimization (PPO) is a family of policy gradient methods for reinforcement learning, proposed by OpenAI in 2017. PPO strikes a balance between simplicity, stability, and performance, making it one of the most widely used algorithms in modern RL applications, including large-scale language model fine-tuning.\n\nTraditional policy gradient methods like REINFORCE or Vanilla Policy Gradient suffer from:\n\n- High variance and sample inefficiency.\n- Instability due to large policy updates.\n\nPPO addresses this problem using a clipped surrogate objective that avoids overly large updates without requiring second-order derivatives.\n\nFor more technical details regarding PPO, we suggest reading the introduction in the [OpenAI spinning up tutorial](https://spinningup.openai.com/en/latest/algorithms/ppo.html), and the paper [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347).\n\n## Key Components\n\n- Actor-Critic Architecture: PPO requires both an actor model (policy) and a critic model (value function). This differs from other algorithms like GRPO and RLOO that don't require a critic model.\n\n- Generalized Advantage Estimation (GAE): PPO uses GAE for computing advantage values, which helps reduce variance in policy gradient estimates while maintaining low bias.\n\n- Clipped Surrogate Objective: The core of PPO is implemented through the clipped surrogate objective function that limits policy updates.\n\n## Configuration\n\nNote that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior.\n\nMost critic configs are similar to those of actors. Note that the critic model is omitted from the figure below.\n\n![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d)\n\n- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n`\n\n- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers\n\n- `critic.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO critic updates. The ppo_mini_batch_size is a global size across all workers\n\n- `actor_rollout_ref.actor.clip_ratio`: The PPO clip range. Default to 0.2\n\n- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for actor\n\n- `critic.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for critic. Defaults to `actor_rollout_ref.actor.ppo_epochs`\n\n- `algorithm.gemma`: discount factor\n\n- `algorithm.lam`: The lambda term that trades off between bias and variance in the GAE estimator\n\n- `algorithm.adv_estimator`: Support gae, grpo, reinforce_plus_plus, reinforce_plus_plus_baseline, rloo\n\n## Advanced Extensions\n\n### KL Divergence Control\n\nOptions to prevent the policy from diverging too far from a reference policy. Two mechanisms are available: KL reward penalty and KL loss. For more technical details, see [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155)\n\nOptions to use KL loss for KL divergence control: \n\n- `actor_rollout_ref.actor.use_kl_loss`: to use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False\n\n- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001.\n\n- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending \"+\" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/volcengine/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\nOptions to use KL penalty in the reward:\n\n- `algorithm.use_kl_in_reward`: Whether to enable in-reward kl penalty. Default is False.\n\n- `algorithm.kl_penalty`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. This defines the way to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty` in core_algos.py. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\n- `algorithm.kl_ctrl.kl_coef`: The (initial) coefficient of in-reward kl_penalty. Default is 0.001.\n- `algorithm.kl_ctrl.type`: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController.\n- `algorithm.kl_ctrl.horizon`: See source code of AdaptiveKLController for details.\n- `algorithm.kl_ctrl.target_kl`: See source code of AdaptiveKLController for details.\n\n### Dual-clip PPO\n\nThe Dual-Clip PPO introduces a approach by applying a lower bound to the policy ratio when the advantage is less than zero, when multiplied by a large raito, does not exceed a specified lower bound.\n\n![image](https://github.com/user-attachments/assets/fc232181-d8b0-4307-8dd2-4dc0a4c1c139)\n\n- `actor_rollout_ref.actor.clip_ratio_c`: lower bound of the value for Dual-clip PPO, defaults to 3.0\n\n## Reference Example\n\nQwen2.5 training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log)\n\n```bash\nbash run_gemma.sh\n  trainer.n_gpus_per_node=1 \\\n  actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n  trainer.logger=console \\\n  critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  data.train_batch_size=256 \\\n  actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n  actor_rollout_ref.actor.ppo_micro_batch_size=2 \\\n  critic.ppo_micro_batch_size=2\n```\n\nReference performance with verl v0.2:\n\n| Model                          | Method          | Score | Link                                                                                           |\n|-------------------------------|------------------|-------|------------------------------------------------------------------------------------------------|\n| Qwen/Qwen2.5-0.5B-Instruct     | pretrained model | 36.4  | [Qwen Blog](https://qwenlm.github.io/blog/qwen2.5-llm/)                                        |\n| Qwen/Qwen2.5-0.5B-Instruct     | PPO              | 56.7  | [PPO Command and Logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) |\n"
  },
  {
    "path": "docs/algo/rollout_corr.md",
    "content": "# Rollout Correction\n\n**Author:** [Yingru Li](https://richardli.xyz/)\n\nLast updated: 10/30/2025.\n\n---\n\n> **📖 Documentation Structure**\n>\n> - **This document** - Practical usage guide: configurations, presets, troubleshooting\n> - **[Mathematical Formulations](rollout_corr_math.md)** - Theoretical foundations, derivations, and algorithmic details\n>\n> Start here for implementation, refer to the math doc for theory and design rationale.\n\n---\n\nThis document provides a comprehensive overview of the Rollout Correction implementation in verl.\n\n**Note on Naming**: This feature is called \"Rollout Correction\" to reflect the complete functionality: importance sampling (IS) weights and rejection sampling (RS). The internal variable `rollout_is_weights` retains its name as it specifically refers to the IS weights component.\n\n### BibTeX Citation\n\n```bibtex\n@online{liu-li-2025-rl-collapse,\n  title = {When Speed Kills Stability: Demystifying {RL} Collapse from the Training-Inference Mismatch},\n  author = {Liu, Jiacai and Li, Yingru and Fu, Yuqian and Wang, Jiawei and Liu, Qian and Shen, Yu},\n  year = {2025},\n  month = sep,\n  url = {https://richardli.xyz/rl-collapse}\n}\n\n@article{li2025trust,\n  title={Trust Region Masking for Long-Horizon LLM Reinforcement Learning},\n  author={Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Liu, Qian and Wang, Baoxiang},\n  journal={arXiv preprint arXiv:2512.23075},\n  year={2025}\n}\n```\n\n### Blog Series\n\n- Main blog post: https://richardli.xyz/rl-collapse\n- [Part 1: Why Mismatch Breaks LLM-RL](https://richardli.xyz/rl-collapse-1) (analytical framework using TV distance for bias and χ²-divergence for variance)\n- [Part 2: The Gradient Estimator Trials](https://richardli.xyz/rl-collapse-2) (token-level vs sequence-level correction bias-variance tradeoff)\n- [Part 3: When Math Meets Reality—Toxic Tails and Length Traps](https://richardli.xyz/rl-collapse-3) (why rejection over clipping, and geometric-level RS)\n- Latest Paper: https://arxiv.org/abs/2512.23075\n\n## Overview\n\nRollout Correction provides a unified framework to handle **general off-policy problems** in RL training. Any scenario where the data collection distribution differs from the training distribution can benefit from these methods.\n\n**Common off-policy scenarios:**\n\n1. **Policy Mismatch** (Implementation Differences)\n\n   - Different precision: FP8 vs FP16 vs BF16 vs FP32\n   - Different backends: vLLM vs SGLang vs FSDP vs Megatron\n   - Different implementations even with identical weights\n\n2. **Temporal Lag** (Model Staleness)\n\n   - Rollout uses older checkpoint while training has progressed\n   - Asynchronous rollout workers with stale parameters\n   - Common in distributed/async RL systems\n\n3. **Replay Buffers**\n\n   - Training on historical trajectories from earlier iterations\n   - Experience replay from different policy versions\n   - Data augmentation or resampling strategies\n\n4. **Off-Policy Algorithms**\n\n   - Behavioral cloning from expert demonstrations\n   - DAPO (data from auxiliary policies)\n   - Any algorithm using trajectories from a different policy\n\n5. **Data Quality Filtering**\n   - Reweighting or filtering collected data\n   - Preference learning with modified distributions\n   - Curriculum learning with distribution shifts\n\nThese off-policy gaps can cause training instability and policy collapse. Rollout Correction uses importance sampling (IS) weights and rejection sampling (RS) to correct for any distribution shift between data collection and training.\n\n**Important Note on Common Implementation Mistakes:**\n\nMany LLM-RL implementations incorrectly apply PPO by **ignoring the actual rollout policy** π_rollout and assuming the training reference policy π_old is the behavior policy. This is mathematically incorrect when π_rollout ≠ π_old (which is typical in LLM-RL due to precision/backend differences between rollout and training).\n\n**This is not PPO's fault** - PPO itself is mathematically correct. The issue is the incorrect assumption that π_old = π_rollout in naive implementations.\n\nThis critical implementation mistake that leads to RL training collapse was identified in the blog post [\"When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch\"](https://richardli.xyz/rl-collapse) and motivated the development of this rollout correction framework.\n\n**Mathematically correct approaches:**\n\n- **Decoupled mode**: Three policies (π_rollout, π_old, π_θ) with IS correction from π_rollout to π_old\n- **Bypass mode**: Two policies (π_rollout = π_old, π_θ) using actual rollout policy as PPO anchor\n- **Bypass + Policy Gradient mode**: Two policies (π_rollout, π_θ) with IS/RS correction and no PPO clipping\n\nSee [Mathematical Formulations](rollout_corr_math.md#37-common-implementation-mistake) for detailed explanation.\n\n### Key Design Principle: Separation of IS Weights and Rejection Sampling\n\nThe implementation cleanly separates two orthogonal mechanisms:\n\n1. **IS Weights** (`rollout_is_weights`): Continuous reweighting for gradient correction\n\n   - Policy ratio: π_old/π_rollout (decoupled) or π_θ/π_rollout (bypass)\n   - **Safety-bounded**: Clamped to [exp(-20), exp(20)] ≈ [2e-9, 5e8] to prevent overflow\n     - Token level: Bounds per-token ratios\n     - Sequence level: Bounds product of ratios (broadcast to all tokens)\n   - **Truncated**: Upper clamped via `.clamp(max=rollout_is_threshold)` (TIS: Truncated Importance Sampling)\n   - **Zeroed at padding**: Multiplied by response_mask to zero out padding positions\n   - Used to weight policy gradients (variance reduction)\n\n2. **Rejection Sampling** (`modified_response_mask`): Binary filtering for outlier exclusion\n   - Creates binary mask: 1 = keep, 0 = reject\n   - Rejects tokens/sequences with IS ratios outside [lower_threshold, upper_threshold]\n   - Modifies response_mask to exclude rejected samples from training\n\nThis separation ensures:\n\n- ✅ IS weights provide continuous reweighting (reduce variance)\n- ✅ Rejection sampling provides hard filtering (remove extreme outliers)\n- ✅ Both mechanisms can be enabled independently or together\n- ✅ Safety bounds prevent numerical overflow in all cases\n\n## Quick Start: Using Verified Presets\n\n**NEW**: We now provide typed configuration with verified presets for common scenarios. These presets have been validated with tens of thousands of GPU hours across various models and training scenarios.\n\n### Python API\n\n```python\nfrom verl.trainer.config.algorithm import RolloutCorrectionConfig\n\n# === Decoupled PPO mode (3 policies: π_rollout, π_old, π_θ) ===\n# IS weights correct for gap between π_old and π_rollout\nconfig = RolloutCorrectionConfig.decoupled_token_is()           # Token-TIS\nconfig = RolloutCorrectionConfig.decoupled_seq_is()             # Seq-TIS\nconfig = RolloutCorrectionConfig.decoupled_seq_is_rs()          # Seq-MIS\nconfig = RolloutCorrectionConfig.decoupled_geo_rs()             # Geo-RS (ratio mode)\nconfig = RolloutCorrectionConfig.decoupled_geo_rs_token_tis()   # Geo-RS + Token-TIS\n\n# === K3 KL Estimator presets (more stable for small KL) ===\nconfig = RolloutCorrectionConfig.decoupled_k3_rs()              # K3-RS only\nconfig = RolloutCorrectionConfig.decoupled_k3_rs_token_tis()    # K3-RS + Token-TIS\n\n# === Bypass PPO mode (2 policies: π_rollout = π_old, π_θ) - fast ===\n# PPO ratio handles IS, so no explicit IS weights needed\nconfig = RolloutCorrectionConfig.bypass_ppo_clip()              # PPO-clip only\nconfig = RolloutCorrectionConfig.bypass_ppo_clip_geo_rs()       # PPO-clip + Geo-RS\nconfig = RolloutCorrectionConfig.bypass_ppo_clip_k3_rs()        # PPO-clip + K3-RS\n\n# === Bypass PG mode (2 policies, no PPO clipping) - fast ===\n# IS weights computed on-the-fly as π_θ / π_rollout\nconfig = RolloutCorrectionConfig.bypass_pg_is()                 # Seq-TIS + PG\nconfig = RolloutCorrectionConfig.bypass_pg_geo_rs()             # Geo-RS + PG\nconfig = RolloutCorrectionConfig.bypass_pg_geo_rs_token_tis()   # Geo-RS + Token-TIS + PG\n\n# === Other ===\nconfig = RolloutCorrectionConfig.disabled()             # Metrics only (no correction)\n```\n\n### YAML Configuration (Advanced)\n\nFor advanced customization or YAML-based configs:\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token # IS weights: \"token\", \"sequence\", or null\n    rollout_is_threshold: 2.0 # Upper threshold for IS weights\n    rollout_is_batch_normalize: false # Batch normalize IS weights to mean=1.0\n    rollout_rs: null # Rejection sampling: comma-separated canonical options (e.g. \"token_k1,seq_max_k2\")\n    rollout_rs_threshold: null # Threshold spec: float(s) or \"lower_upper\" string(s)\n    bypass_mode: false # Skip old_log_prob computation (sets π_old = π_rollout)\n    loss_type: ppo_clip # Loss type in bypass mode: \"ppo_clip\" (default) or \"reinforce\"\n\n# REQUIRED: Enable log prob calculation\nactor_rollout_ref:\n  rollout:\n    calculate_log_probs: true\n```\n\n## Files\n\n### **Core Implementation**\n\n- `verl/trainer/ppo/rollout_corr_helper.py` - Contains `compute_rollout_correction_and_rejection_mask()` and `compute_offpolicy_metrics()`\n- `verl/trainer/ppo/core_algos.py` - Rollout Correction integration with PPO and REINFORCE modes (`compute_policy_loss_bypass_mode()`, `compute_policy_loss_reinforce()`)\n- `verl/trainer/ppo/ray_trainer.py` - Bypass mode implementation (skips `old_log_prob` computation)\n- `verl/workers/actor/dp_actor.py` - Mode selection logic and metrics collection\n\n### **Configuration Files**\n\n- `verl/trainer/config/algorithm.py` - Rollout Correction parameters in `RolloutCorrectionConfig`\n- `verl/workers/config/actor.py` - Rollout Correction parameters in `PolicyLossConfig`\n- `verl/trainer/config/actor/actor.yaml` - Rollout Correction configuration section\n- `verl/trainer/config/ppo_trainer.yaml` - Algorithm config with Rollout Correction\n\n### **Documentation**\n\n- `docs/examples/config.rst` - Configuration parameter descriptions\n\n### **Example Scripts**\n\n- `recipe/dapo/run_dapo_qwen2.5_32b_rollout_corr.sh` - DAPO example with Rollout Correction\n- `examples/rollout_correction/run_with_rollout_corr.sh` - Basic example\n- `examples/rollout_correction/run_with_rollout_corr_multi_rs.sh` - Multi-RS example\n\n### **Tests**\n\n- `tests/trainer/ppo/test_rollout_corr.py` - Unit tests for IS/RS mechanisms\n- `tests/trainer/ppo/test_rollout_corr_integration.py` - Integration tests\n\n## Configuration Parameters\n\nAll parameters are under `algorithm.rollout_correction`:\n\n### `rollout_is` (str or null)\n\nImportance sampling weights aggregation level:\n\n- `null` = No IS weights computed (metrics-only mode)\n- `\"token\"`: Per-token IS weights\n  - **Decoupled mode**: ρ_t = π_old(t)/π_rollout(t)\n  - **Bypass/Pure IS mode**: ρ_t = π_θ(t)/π_rollout(t)\n  - Independent truncation per token\n  - Typical threshold: 1.5 - 5.0\n- `\"sequence\"`: Per-sequence weight ρ_seq = ∏_t ρ_t\n  - Multiplicative aggregation across sequence\n  - Typical threshold: 2.0 - 10.0\n\nAll IS weights are safety-bounded to [exp(-20), exp(20)] ≈ [2e-9, 5e8]\n\n### `rollout_is_threshold` (float)\n\nUpper threshold for IS weight truncation. Default: `2.0`\n\n- Truncates IS weights via `.clamp(max=rollout_is_threshold)` (TIS: Truncated Importance Sampling)\n- Applied to IS weights for variance reduction\n- Separate from rejection sampling (controlled by `rollout_rs` parameters)\n\n### `rollout_is_batch_normalize` (bool)\n\nApply batch normalization to IS weights. Default: `False`\n\n- `True`: Normalize IS weights to have mean=1.0 within each batch\n  - **Token-level IS**: Normalizes over all token weights\n  - **Sequence-level IS**: Normalizes over sequence means (one weight per sequence)\n- `False`: Use raw (truncated) IS weights\n- Reduces variance by ensuring average weight is 1.0 per batch\n- Applied AFTER truncation to preserve truncation semantics\n- Only affects IS weight values, not rejection sampling\n\n### `rollout_rs` (str or null)\n\nRejection sampling aggregation modes. Supply a comma-separated string (spaces optional) using the canonical options implemented in `rollout_corr_helper`:\n\n- `token_k1`: Token-level rejection with `-log r` bounds (ratio thresholds supplied as `lower_upper`). Example: `\"0.6_1.4\"`\n- `token_k2`: Token-level rejection with `0.5 * (log r)^2` (upper bound only)\n- `token_k3`: Token-level rejection with `exp(log r) - 1 - log r` (upper bound only)\n- `seq_sum_k1`: Sequence-level rejection with sum of `-log r` (ratio bounds)\n- `seq_sum_k2`: Sequence-level rejection with sum of `0.5 * (log r)^2` (upper bound only)\n- `seq_sum_k3`: Sequence-level rejection with sum of `exp(log r) - 1 - log r` (upper bound only)\n- `seq_mean_k1`: Sequence-level rejection with mean of `-log r` (ratio bounds)\n- `seq_mean_k2`: Sequence-level rejection with mean of `0.5 * (log r)^2` (upper bound only)\n- `seq_mean_k3`: Sequence-level rejection with mean of `exp(log r) - 1 - log r` (upper bound only)\n- `seq_max_k2`: Sequence-level rejection with max of `0.5 * (log r)^2` (upper bound only)\n- `seq_max_k3`: Sequence-level rejection with max of `exp(log r) - 1 - log r` (upper bound only)\n\n### `rollout_rs_threshold` (str, float, or null)\n\nThreshold specification for rejection sampling.\n\n- Provide **one entry per option**, separated by commas. A single entry is broadcast to every option.\n- **K1 KL modes (`*k1`)**: Use `\"lower_upper\"` strings (e.g. `\"0.7_1.3\"`). Supplying a float implies only the upper bound; the lower bound defaults to its reciprocal.\n- **K2/K3 KL modes (`*k2`/`*k3`)**: Supply positive upper bounds (float or numeric string).\n- Set to `null` to disable thresholds entirely (only valid when `rollout_rs` is null).\n\n## Understanding the Framework: Components and Combinations\n\nThe rollout correction framework is built from **orthogonal components** that can be combined flexibly. Understanding these components helps you choose the right configuration for your scenario.\n\n### Key Components\n\n1. **Operating Mode** (Section: [Operation Modes](#operation-modes))\n\n   - **Decoupled**: Three policies (π_rollout, π_old, π_θ) with separate π_old computation\n   - **Bypass**: Two policies (π_rollout = π_old, π_θ), skips π_old computation\n\n2. **Loss Function** (in bypass mode, controlled by `loss_type`)\n\n   - **PPO-clip** (`loss_type=\"ppo_clip\"`, default): PPO clipped objective (IS handled by ratio)\n   - **REINFORCE** (`loss_type=\"reinforce\"`): Policy gradient with explicit IS weights (no clipping)\n\n3. **IS/RS Aggregation Level**\n   - **Token**: Per-token IS weights/rejection\n   - **Sequence**: Sequence-level IS weights/rejection\n\nSee [Mathematical Formulations](rollout_corr_math.md#3-algorithmic-components-and-combinations) for detailed theory.\n\n---\n\n## Preset Configuration Guide\n\nThis section provides detailed guidance on choosing and using the verified presets. Each preset is a specific combination of components optimized for common scenarios.\n\n### Understanding the Presets\n\n#### Available Preset Methods\n\n| Preset Method                                                                  | Estimator        | Mode               | IS Level | RS Level | Properties                              |\n| ------------------------------------------------------------------------------ | ---------------- | ------------------ | -------- | -------- | --------------------------------------- |\n| **Decoupled PPO Mode** (3 policies: π_rollout, π_old, π_θ)                     |\n| `decoupled_token_is()`                                                         | Token-TIS        | Decoupled          | token    | -        | Token-level IS weights                    |\n| `decoupled_seq_is()`                                                           | Seq-TIS          | Decoupled          | sequence | -        | Sequence-level IS weights               |\n| `decoupled_seq_is_rs()`                                                        | Seq-MIS          | Decoupled          | sequence | sequence | Sequence IS + seq_sum_k1 RS               |\n| `decoupled_geo_rs()`                                                           | Geo-RS           | Decoupled          | -        | sequence | Geometric RS (seq_mean_k1)               |\n| `decoupled_geo_rs_token_tis()`                                                 | Geo-RS-Token-TIS | Decoupled          | token    | sequence | Geometric RS + token IS |\n| **K3 KL Estimator** (more stable for small KL values)                          |\n| `decoupled_k3_rs()`                                                            | K3-RS            | Decoupled          | -        | sequence       | seq_mean_k3 RS             |\n| `decoupled_k3_rs_token_tis()`                                                  | K3-RS-Token-TIS  | Decoupled          | token    | sequence       | seq_mean_k3 RS  + token IS        |\n| **Bypass Mode (PPO-clip)** (2 policies; ratio handles IS, RS masks outliers)   |\n| `bypass_ppo_clip()`                                                            | -                | Bypass (PPO-clip)  | -        | -        | PPO-clip only                           |\n| `bypass_ppo_clip_geo_rs()`                                                     | Geo-RS           | Bypass (PPO-clip)  | -        | sequence | PPO-clip + Geo-RS                       |\n| `bypass_ppo_clip_k3_rs()`                                                      | K3-RS            | Bypass (PPO-clip)  | -        | sequence       | PPO-clip + K3-RS                       |\n| **Bypass Mode (REINFORCE)** (2 policies; explicit IS weights, no PPO clipping) |\n| `bypass_pg_is()`                                                               | Seq-TIS          | Bypass (REINFORCE) | sequence | -        | REINFORCE with explicit IS              |\n| `bypass_pg_geo_rs()`                                                           | Geo-RS           | Bypass (REINFORCE) | -        | sequence | REINFORCE with Geo-RS                   |\n| `bypass_pg_geo_rs_token_tis()`                                                 | Geo-RS-Token-TIS | Bypass (REINFORCE) | token    | sequence | REINFORCE + Geo-RS + token IS       |\n| **Other**                                                                      |\n| `disabled()`                                                                   | -                | -                  | -        | -        | Metrics only, no correction             |\n\n**Note:**\n\n- **Bypass mode** sets π_old = π_rollout and uses `loss_type` to select the loss function:\n  - `\"ppo_clip\"` (default): PPO clipped objective where ratio = π_θ/π_rollout already handles IS\n  - `\"reinforce\"`: REINFORCE with explicit IS weights as π_θ/π_rollout\n- Both loss types benefit from rejection sampling (RS) which masks out-of-distribution samples.\n- All estimators (Token-TIS, Seq-TIS, Seq-MIS, Geo-RS, ...) are compatible with Decoupled and Bypass modes.\n\n#### Other Supported Combinations (Manual Configuration Required)\n\n**Other supported combinations without preset methods:**\n\n- Token IS + Token RS: Token-level IS weights + Token-level RS mask\n- Pure token RS: Token-level RS only, no IS weights\n- Pure sequence RS: Sequence-level RS only, no IS weights\n\nSee [detailed configuration examples below](#additional-useful-configurations-not-exposed-as-presets) for manual configurations.\n\n**Key properties:**\n\n- Any aggregation level (token/sequence) works in either decoupled or bypass mode\n- All combinations are fully supported by the implementation\n- Rejection sampling is independent of IS weighting\n- Pure RS (`bypass_pg_rs`) uses bypass + geometric RS with `loss_type=\"reinforce\"` (no IS weights)\n\n---\n\n### 1. Decoupled Mode with Token-level Importance Sampling (`decoupled_token_is`)\n\n**Configuration:**\n\n```python\nconfig = RolloutCorrectionConfig.decoupled_token_is(threshold=2.0)\n```\n\n**Components:**\n\n- **Operating Mode**: Decoupled (3 policies)\n- **Loss**: PPO with clipping (only for the second drift correction)\n- **IS Aggregation**: Token-level\n- **RS**: None (can be added separately)\n\n**Equivalent YAML:**\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    bypass_mode: false # Decoupled mode\n```\n\n**Properties:**\n\n- Independent truncation per token\n- Lower variance than sequence-level (product of ratios bounded individually)\n- Typical threshold: 1.5 - 5.0\n\n**Theory:** See [rollout_corr_math.md §3.3.1](rollout_corr_math.md#331-token-level-aggregation)\n\n---\n\n### 2. Decoupled Mode with Sequence-level Importance Sampling (`decoupled_seq_is`)\n\n**Also known as: Seq-TIS (Sequence-Level Truncated IS)**\n\n**Configuration:**\n\n```python\nconfig = RolloutCorrectionConfig.decoupled_seq_is(threshold=2.0)\n```\n\n**Components:**\n\n- **Operating Mode**: Decoupled (3 policies)\n- **Loss**: PPO with clipping (only for the second drift correction)\n- **IS Aggregation**: Sequence-level (Seq-TIS)\n- **RS**: None (can be added separately)\n\n**Equivalent YAML:**\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: sequence\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    bypass_mode: false # Decoupled mode\n```\n\n**Properties:**\n\n- Multiplicative aggregation across sequence\n- More sensitive to outliers than token-level\n- Typical threshold: 2.0 - 10.0 (higher than token-level)\n\n**Theory:** See [rollout_corr_math.md §3.3.2](rollout_corr_math.md#332-sequence-level-aggregation)\n\n---\n\n### 3. Decoupled Mode with Sequence-level IS + Rejection Sampling (`decoupled_seq_is_rs`)\n\n**Also known as: Seq-MIS (Sequence-Level Masked IS)**\n\n**Configuration:**\n\n```python\nconfig = RolloutCorrectionConfig.decoupled_seq_is_rs(is_threshold=2.0, rs_threshold=\"0.5_2.0\")\n```\n\n**Components:**\n\n- **Operating Mode**: Decoupled (3 policies)\n- **Loss**: PPO with clipping (only for the second drift correction)\n- **IS Aggregation**: Sequence-level (Seq-TIS)\n- **RS**: Sequence-level rejection (Seq-MIS)\n\n**Equivalent YAML:**\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: sequence\n    rollout_is_threshold: 2.0\n    rollout_rs: seq_sum_k1\n    rollout_rs_threshold: 0.5_2.0\n    bypass_mode: false # Decoupled mode\n```\n\n**Properties:**\n\n- Double mechanism: IS reweighting (Seq-TIS) + rejection filtering (Seq-MIS)\n- Lower effective sample size (rejects outliers)\n- For severe off-policy gaps or when the distribution tail is \"toxic\" (garbage/adversarial samples)\n\n**When to use Seq-MIS over Seq-TIS:**\n\n- **Seq-TIS (clipping only)**: Maximizes information efficiency; extracts signal from all samples. Use when data is clean and mismatch is moderate.\n- **Seq-MIS (rejection)**: Maximizes safety; acts as a hard trust region filter. Use when mismatch is severe or when high-weight samples are likely garbage rather than signal.\n\n**Theory:** See [rollout_corr_math.md §3.5](rollout_corr_math.md#35-rejection-sampling-rs)\n\n---\n\n### 6. Bypass Mode with PPO-clip (`bypass_ppo_clip`)\n\n**Configuration:**\n\n```python\nconfig = RolloutCorrectionConfig.bypass_ppo_clip()\n```\n\n**Components:**\n\n- **Operating Mode**: Bypass (2 policies: π_rollout = π_old, π_θ)\n- **Loss**: PPO-clip (IS handled by ratio, no explicit IS weights)\n- **IS Aggregation**: None (PPO ratio handles it)\n- **RS**: None\n\n**Equivalent YAML:**\n\n```yaml\nrollout_correction:\n  rollout_is: null\n  rollout_rs: null\n  bypass_mode: true\n  loss_type: ppo_clip\n```\n\n**Properties:**\n\n- PPO clipped objective in bypass mode\n- The PPO ratio = π_θ/π_rollout already handles IS (no explicit IS weights needed)\n- Skips `actor.compute_log_prob()` forward pass (2 policies instead of 3)\n- No rejection sampling - use `bypass_ppo_clip_geo_rs()` for RS\n\n**Configuration requirement:**\n\n- Set `actor_rollout_ref.rollout.calculate_log_probs: true`\n\n**Additional requirements for bypass mode:**\n\n- Set `actor_rollout_ref.actor.use_rollout_log_probs: true`\n- Set `actor_rollout_ref.actor.policy_loss.loss_mode: bypass_mode`\n- Set rollout correction config via `actor_rollout_ref.actor.policy_loss.rollout_correction`\n\n**Theory:** See [rollout_corr_math.md §3.1.2](rollout_corr_math.md#312-bypass-mode-two-policies)\n\n---\n\n### 7. REINFORCE with IS (`bypass_pg_is`)\n\n**Configuration:**\n\n```python\nconfig = RolloutCorrectionConfig.bypass_pg_is(threshold=2.0)\n```\n\n**Components:**\n\n- **Operating Mode**: Bypass (2 policies: π_rollout, π_θ)\n- **Loss**: REINFORCE (policy gradient with explicit IS weights, no PPO clipping)\n- **IS Aggregation**: Sequence-level\n- **RS**: None\n\n**Equivalent YAML:**\n\n```yaml\nrollout_correction:\n  rollout_is: sequence\n  rollout_is_threshold: 2.0\n  rollout_rs: null\n  bypass_mode: true\n  loss_type: reinforce # REINFORCE with explicit IS weights\n```\n\n**Properties:**\n\n- REINFORCE loss with explicit IS weights (no PPO clipping)\n- Single forward pass (skips old_log_prob computation)\n- IS weights computed on-the-fly in loss function\n\n**Theory:** See [rollout_corr_math.md §3.2.2](rollout_corr_math.md#322-policy-gradient-loss-with-isrs-correction)\n\n---\n\n## Additional Useful Configurations (Not Exposed as Presets)\n\nThese configurations are **fully supported** but don't have convenience preset methods yet.\n\n### 1. Token IS + Token RS (`token_is_rs`)\n\nToken-level IS weights with token-level RS mask.\n\n**Python:**\n\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=\"token\",\n    rollout_is_threshold=2.0,\n    rollout_rs=\"token_k1\",\n    rollout_rs_threshold=2.0,\n)\n```\n\n**Properties:** Per-token IS weights + per-token RS mask.\n\n### 2. Pure Token RS (`token_rs`)\n\nToken-level RS only, no IS weights.\n\n**Python:**\n\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=None,\n    rollout_rs=\"token_k1\",\n    rollout_rs_threshold=2.0,\n)\n```\n\n**Properties:** Token-level RS mask, no IS reweighting.\n\n### 3. Pure Sequence RS (`seq_rs`)\n\nSequence-level RS only, no IS weights.\n\n**Python:**\n\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=None,\n    rollout_rs=\"seq_sum_k1\",\n    rollout_rs_threshold=\"0.5_2.0\",\n)\n```\n\n**Properties:** Sequence-level RS mask, no IS reweighting.\n\n---\n\n### Summary: How IS Weights are Processed\n\nIS weights (`rollout_is_weights`) go through a fixed processing pipeline:\n\n**Stage 1: Safety Bound (Prevent Overflow)**\n\n- Token level: `exp(clamp(log_ratio, -20, 20))` per token → bounds each token to [2e-9, 5e8]\n- Sequence level: `exp(clamp(sum(log_ratio), -20, 20))` → bounds product to [2e-9, 5e8], broadcast to all tokens\n\n**Stage 2: Truncation (Reduce Variance)**\n\n- `.clamp(max=rollout_is_threshold)` → caps weights at upper threshold (TIS: Truncated Importance Sampling)\n- No lower truncation (preserves unbiasedness for small weights)\n\n**Stage 3: Padding Zeroing (Correct Aggregation)**\n\n- `weights * response_mask` → zeros out padding positions\n\n**Stage 4: Optional Batch Normalization**\n\n- If `rollout_is_batch_normalize=True`: Normalize weights to mean=1.0 within batch\n- Applied after truncation to preserve truncation semantics\n\n**Rejection Sampling (Separate Mechanism)**\n\nRejection sampling modifies `response_mask` (NOT weights) through `compute_rollout_rejection_mask()`:\n\n- Computes safety-bounded ratios independently\n- Creates binary mask: tokens/sequences outside [lower_threshold, upper_threshold] → 0 (rejected)\n- Modified mask used for loss aggregation\n\n## Operation Modes\n\nThe framework provides **two operating modes** for computing π_old, which can be combined with different loss functions.\n\n### Operating Modes and Configuration\n\n| Configuration          | `bypass_mode` | `loss_type`            | Operating Mode | Loss Function | Description                                                       |\n| ---------------------- | ------------- | ---------------------- | -------------- | ------------- | ----------------------------------------------------------------- |\n| **Decoupled**          | `false`       | N/A                    | Decoupled      | PPO           | Computes `old_log_prob` separately via `actor.compute_log_prob()` |\n| **Bypass + PPO-clip**  | `true`        | `\"ppo_clip\"` (default) | Bypass         | PPO-clip      | PPO clipped objective (IS handled by ratio)                       |\n| **Bypass + REINFORCE** | `true`        | `\"reinforce\"`          | Bypass         | REINFORCE     | Policy gradient with explicit IS weights (no PPO clipping)        |\n\n### Operating Mode Details\n\n#### Decoupled Mode (Three Policies)\n\n**Policy setup:**\n\n- π_rollout: Behavior policy (data collection)\n- π_old: Proximal policy (computed via `actor.compute_log_prob()` at start of training epoch)\n- π_θ: Current policy (being updated)\n\n**Configuration:** `bypass_mode = false`\n\n**Properties:**\n\n- ✅ Achieves batch size invariance\n- ✅ Separately corrects Drift 1 (rollout→old) and Drift 2 (old→current)\n- ✅ Efficient stale data utilization\n- ❌ Extra forward pass needed (`actor.compute_log_prob()`)\n\n**Theory:** See [rollout_corr_math.md §3.1.1](rollout_corr_math.md#311-decoupled-mode-three-policies)\n\n#### Bypass Mode (Two Policies)\n\n**Policy setup:**\n\n- π_rollout: Behavior policy (data collection)\n- π_old = π_rollout: Proximal policy equals behavior policy\n- π_θ: Current policy (being updated)\n\n**Configuration:** `bypass_mode = true`\n\n**Properties:**\n\n- ✅ Skips `actor.compute_log_prob()` call (faster)\n- ✅ Handles off-policy correction via IS/RS (when using policy gradient with IS/RS)\n- ✅ Uses two policies instead of three (π_rollout = π_old)\n- ⚠️ Does not separate proximal policy from behavior policy (unlike decoupled mode)\n\n**Theory:** See [rollout_corr_math.md §3.1.2](rollout_corr_math.md#312-bypass-mode-two-policies)\n\n---\n\n### IS/RS Aggregation Levels (Orthogonal to Operating Mode)\n\nThe aggregation level can be chosen **independently** of the operating mode. Any aggregation level works in either decoupled or bypass mode.\n\n| `rollout_is`              | `rollout_rs`                                                       | Behavior                                                                          |\n| ------------------------- | ------------------------------------------------------------------ | --------------------------------------------------------------------------------- |\n| `null`                    | `null`                                                             | **Disabled**: No computation, no metrics, no rejection                            |\n| `null`                    | `\"token_k1\"`, `\"seq_sum_k1\"`, `\"seq_mean_k1\"`, `\"seq_max_k2\"`, etc | **Rejection only**: Compute metrics, NO weight correction, YES rejection sampling |\n| `\"token\"` or `\"sequence\"` | `null`                                                             | **IS weights only**: Weight correction enabled, NO rejection sampling             |\n| `\"token\"` or `\"sequence\"` | `\"token_k1\"`, `\"seq_sum_k1\"`, `\"seq_mean_k1\"`, `\"seq_max_k2\"`, etc | **Full correction**: Both weight correction and rejection sampling enabled        |\n\n### Key Insights\n\n- ✅ Any IS/RS aggregation level (token/sequence/geometric) can be used in **either** decoupled or bypass mode\n- ✅ You can use **rejection sampling alone** without IS weight correction (`rollout_is=null, rollout_rs=\"token_k1\"`)\n- ✅ You can use **IS weights alone** without outlier rejection (`rollout_is=\"token\", rollout_rs=null`)\n- ✅ You can use **both together** (`rollout_is=\"token\", rollout_rs=\"token_k1\"`)\n- ✅ You can **monitor metrics only** without any correction by setting both to `null` but still providing rollout_log_probs\n\n**Theory:** See [rollout_corr_math.md §3.3](rollout_corr_math.md#33-isrs-aggregation-levels) for details on aggregation levels.\n\n### Example Workflow\n\n**Recommended: Bypass Mode**\n\nThis workflow uses bypass mode for efficiency.\n\n1. **Start with metrics only** to understand the off-policy gap:\n\n   ```yaml\n    rollout_correction:\n      rollout_is: null\n      rollout_rs: null\n      bypass_mode: true # Bypass mode (recommended)\n      loss_type: ppo_clip # Default: PPO clipped objective\n   ```\n\n   Monitor `rollout_corr/kl`, `rollout_corr/log_ppl_abs_diff`, `rollout_corr/chi2_token` to assess off-policy gap.\n\n2. **Enable rejection sampling** if you see high outlier fractions:\n\n   ```yaml\n    rollout_correction:\n      rollout_is: null\n      rollout_rs: sequence # or \"geometric\" for higher sensitivity\n      rollout_rs_threshold: 2.0\n      bypass_mode: true # Bypass mode\n      loss_type: ppo_clip # or \"reinforce\" for explicit IS weights\n   ```\n\n   This excludes outliers from training without modifying gradients.\n\n3. **Enable full IS correction** (with REINFORCE loss) once comfortable with metrics:\n   ```yaml\n    rollout_correction:\n      rollout_is: sequence # Recommended: unbiased, suitable for most cases\n      rollout_is_threshold: 2.0\n      rollout_rs: sequence # or \"geometric\" for more aggressive filtering\n      rollout_rs_threshold: 2.0\n      bypass_mode: true # Bypass mode\n      loss_type: reinforce # REINFORCE with explicit IS weights\n   ```\n\n**Benefits of bypass mode:**\n\n- ✅ Skips expensive `actor.compute_log_prob()` forward pass (faster)\n- ✅ `loss_type` controls the loss function: \"ppo_clip\" (default) or \"reinforce\"\n- ✅ PPO-clip: IS handled by ratio (no explicit weights), RS mask applied\n- ✅ REINFORCE: Explicit IS weights computed on-the-fly (π_θ / π_rollout)\n- ✅ Both loss types work with all IS/RS combinations\n\n## Usage\n\n### Basic Setup\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token # Enable IS weights at token level\n    rollout_is_threshold: 2.0 # Threshold for IS weights\n    rollout_rs: null # No rejection sampling\n\nactor_rollout_ref:\n  rollout:\n    calculate_log_probs: true # Required!\n```\n\n### Additional Configurations for Bypass Mode\n\n- Set `actor_rollout_ref.actor.use_rollout_log_probs: true`\n- Set `actor_rollout_ref.actor.policy_loss.loss_mode: bypass_mode`\n- Set rollout correction config via `actor_rollout_ref.actor.policy_loss.rollout_correction`\n\n### Metrics\n\nAll metrics are prefixed with `rollout_corr/` in logs. For example, `rollout_is_mean` appears as `rollout_corr/rollout_is_mean`.\n\nThese metrics cover both:\n\n- **Diagnostic metrics**: KL divergence, perplexity differences (measuring off-policy gap)\n- **Correction statistics**: IS weights, rejection rates (measuring correction applied)\n\n#### **Core IS Weight Metrics**\n\n- **`rollout_is_mean`**: Mean importance sampling weight across all valid tokens\n\n  - Value close to 1.0 indicates minimal off-policy gap\n\n- **`rollout_is_std`**: Standard deviation of IS weights\n\n  - Higher values indicate greater variance in IS weights\n\n- **`rollout_is_min`**: Minimum IS weight observed\n\n  - Shows the most underweighted token/sequence\n  - For sequence/geometric: computed from unclamped log-space ratios (true minimum)\n  - For token: computed from safety-bounded weights\n\n- **`rollout_is_max`**: Maximum IS weight observed\n  - Shows the most overweighted token/sequence\n  - For sequence/geometric: computed from unclamped log-space ratios (true maximum before safety bound)\n  - For token: computed from safety-bounded weights (before threshold clamping)\n  - Compare with `rollout_is_threshold` to see truncation impact\n\n#### **Effective Sample Size**\n\n- **`rollout_is_eff_sample_size`**: Effective sample size after IS weighting\n  - **Formula**: `1 / mean(weights²)` where weights are normalized\n  - **Range**: 0.0 to 1.0 (as fraction of original batch)\n  - Lower values indicate weight concentration on fewer samples\n\n#### **Threshold Exceedance Metrics**\n\n- **`rollout_is_ratio_fraction_high`**: Fraction of weights exceeding upper threshold\n\n  - Shows how often truncation/masking occurs on high end\n  - For sequence/geometric: computed from unclamped log-space ratios (true exceedance)\n  - For token: computed from safety-bounded weights (before threshold clamping)\n\n- **`rollout_is_ratio_fraction_low`**: Fraction of weights below lower threshold (1/upper_threshold)\n  - Diagnostic metric showing how many weights are below the reciprocal threshold\n  - For sequence/geometric: computed from unclamped log-space ratios (true exceedance)\n  - For token: computed from safety-bounded weights (before truncation)\n\n#### **Sequence-Level Metrics** (for sequence aggregation)\n\n- **`rollout_is_seq_mean`**: Mean IS weight at sequence level\n\n  - Should match `rollout_is_mean` for sequence-level aggregation\n\n- **`rollout_is_seq_std`**: Standard deviation of sequence-level IS weights\n\n- **`rollout_is_seq_min`**: Minimum sequence-level IS weight\n\n- **`rollout_is_seq_max`**: Maximum sequence-level IS weight\n\n- **`rollout_is_seq_max_deviation`**: Maximum absolute deviation from 1.0 at sequence level\n\n  - Shows worst-case sequence off-policy gap\n\n- **`rollout_is_seq_fraction_high`**: Fraction of sequences exceeding upper threshold\n\n- **`rollout_is_seq_fraction_low`**: Fraction of sequences below lower threshold\n\n#### **Rejection Sampling Metrics** (when `rollout_rs` is enabled)\n\n- **`rollout_rs_masked_fraction`**: Fraction of tokens rejected via rejection sampling\n\n  - **Important**: Rejection sampling modifies `response_mask` (sets rejected tokens to 0)\n  - **Separate from IS weights**: IS weights are still truncated; rejection is an independent filtering step\n  - Only present when `rollout_rs` is enabled (token/sequence/geometric)\n\n- **`rollout_rs_seq_masked_fraction`**: Fraction of sequences with at least one rejected token\n  - Shows sequence-level impact of rejection sampling\n  - Token-level RS: sequence rejected if ANY token is outside [lower, upper]\n  - Sequence-level RS: entire sequence rejected or accepted based on sequence-level ratio\n  - Geometric RS: entire sequence rejected or accepted based on geometric mean\n\n#### **Off-Policy Diagnostic Metrics** (Training vs Rollout Policy)\n\n**Note on terminology:** These metrics use \"training\" to refer to the training reference policy and \"rollout\" to refer to π_rollout (the behavior policy used for data collection).\n\n- **Decoupled mode**: \"training\" = π_old (computed at start of training epoch)\n- **Bypass/Pure IS mode**: \"training\" = π_θ (current policy being trained)\n\nIn bypass/pure IS mode, metrics measure the drift between π_θ and π_rollout directly.\n\n- **`training_ppl`**: Perplexity of training reference policy (π_old in decoupled mode, π_θ in bypass/pure IS mode)\n\n  - **Formula**: `exp(-mean(log_probs))`\n  - Lower values indicate higher model confidence\n\n- **`rollout_ppl`**: Perplexity of rollout policy π_rollout (e.g., vLLM BF16)\n\n- **`ppl_ratio`**: Ratio of training PPL to rollout PPL\n\n  - **Formula**: `exp(mean(log(training_ppl / rollout_ppl)))`\n  - **Meaning**: > 1.0 means training is less confident than rollout\n\n- **`training_log_ppl`**: Log perplexity of training policy\n\n  - Useful for identifying trends (linear scale)\n\n- **`rollout_log_ppl`**: Log perplexity of rollout policy\n\n- **`log_ppl_diff`**: Mean difference in log perplexities\n\n  - **Formula**: `mean(log_ppl_rollout - log_ppl_training)`\n  - Sign indicates which policy is more confident\n\n- **`log_ppl_abs_diff`**: Mean absolute log perplexity difference\n\n  - Magnitude of off-policy gap regardless of direction\n\n- **`log_ppl_diff_max`**: Maximum log perplexity difference across sequences\n\n  - Identifies worst-case sequence\n\n- **`log_ppl_diff_min`**: Minimum log perplexity difference across sequences\n\n- **`kl`**: KL divergence KL(π_rollout || π_training)\n\n  - **Formula**: `mean(log_prob_rollout - log_prob_training)`\n  - **Note**: Can be negative (rollout is less confident)\n\n- **`k3_kl`**: K3 divergence (equals KL(π_rollout || π_training) in expectation)\n\n  - **Formula**: `mean(exp(log_ratio) - log_ratio - 1)`\n  - More stable than direct KL (non-negative per token)\n  - Always >= 0\n\n- **`chi2_token`**: Chi-squared divergence at token level\n\n  - **Formula**: `mean(ratio²) - 1` where ratio = π_training/π_rollout\n  - Measures second moment of IS weight distribution\n  - Always non-negative\n\n- **`chi2_seq`**: Chi-squared divergence at sequence level\n  - **Formula**: `mean((∏_t ratio_t)²) - 1`\n  - Sequence-level second moment of IS weights\n  - More sensitive than token-level chi-squared\n\n#### **Example: Accessing Metrics in Code**\n\n```python\n# Metrics are returned from compute_rollout_correction_and_rejection_mask\nfrom verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_rejection_mask\n\n# Returns 3 values (weights, modified_response_mask, metrics)\nweights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask(\n    old_log_prob=training_log_probs,      # from training policy\n    rollout_log_prob=rollout_log_probs,   # from rollout policy\n    response_mask=response_mask,\n    rollout_is=\"token\",  # Enable IS weights at token level\n    rollout_is_threshold=2.0,\n    rollout_rs=\"token_k1\",\n    rollout_rs_threshold=\"0.5_2.0\",\n)\n\n# Extract IS weights (processed, zeroed at padding)\nis_weights = weights_proto.batch[\"rollout_is_weights\"]\n\n# IS weights processing (with IS enabled at token level):\n# 1. Safety-bounded: exp(clamp(log_ratio, -20, 20)) per token\n# 2. Truncated: .clamp(max=2.0) to cap extreme weights\n# 3. Zeroed at padding positions\n# Note: Truncation is ALWAYS applied to IS weights (TIS: Truncated Importance Sampling)\n\n# modified_response_mask has rejection applied (since rollout_rs=\"token_k1\"):\n# 1. RS rejection: tokens outside [0.5, 2.0] masked to 0 via response_mask\n# Note: RS and IS are separate mechanisms - both can be enabled independently\n\n# All metrics have 'rollout_corr/' prefix\nprint(f\"Mean IS weight: {metrics['rollout_corr/rollout_is_mean']:.3f}\")\nprint(f\"Effective sample size: {metrics['rollout_corr/rollout_is_eff_sample_size']:.3f}\")\nprint(f\"RS masked fraction: {metrics['rollout_corr/rollout_rs_masked_fraction']:.3f}\")\nprint(f\"KL divergence: {metrics['rollout_corr/kl']:.3f}\")\n\n# Check IS weights for valid tokens (non-padding)\nvalid_weights = is_weights[response_mask.bool()]\nprint(f\"\\n✓ IS weights min (valid tokens): {valid_weights.min():.4f}\")\nprint(f\"✓ IS weights max (valid tokens): {valid_weights.max():.4f}\")\nprint(f\"✓ All valid IS weights > 0: {(valid_weights > 0).all()}\")\nprint(f\"✓ IS weights are capped at threshold: {(valid_weights <= 2.0).all()}\")\n\n# Check rejection via response_mask\nrejected_tokens = (response_mask == 1) & (modified_response_mask == 0)\nprint(f\"\\n✓ Rejected {rejected_tokens.sum()} tokens via response_mask\")\nprint(f\"✓ Rejection sampling modifies response_mask (separate from IS weight truncation)\")\nprint(f\"✓ IS weights are always truncated to [0, threshold] after safety bounding\")\n\n# Check for warning conditions\nif metrics['rollout_corr/rollout_is_mean'] < 0.5 or metrics['rollout_corr/rollout_is_mean'] > 2.0:\n    print(\"⚠️  Warning: Mean IS weight far from 1.0, significant off-policy gap detected\")\n\nif metrics['rollout_corr/rollout_is_eff_sample_size'] < 0.3:\n    print(\"⚠️  Warning: Low effective sample size, high weight concentration\")\n```\n\n#### **Example: Monitoring Metrics During Training**\n\n```python\n# In your training loop\nfor epoch in range(num_epochs):\n    for batch_idx, batch in enumerate(dataloader):\n        # ... rollout phase ...\n\n        # Compute IS weights and get metrics\n        rollout_corr_config = config.algorithm.get(\"rollout_correction\", None)\n        if rollout_corr_config is not None:\n            weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask(\n                old_log_prob=batch.old_log_prob,\n                rollout_log_prob=batch.rollout_log_prob,\n                response_mask=batch.response_mask,\n                rollout_is=rollout_corr_config.get(\"rollout_is\", None),\n                rollout_is_threshold=rollout_corr_config.get(\"rollout_is_threshold\", 2.0),\n                rollout_rs=rollout_corr_config.get(\"rollout_rs\", None),\n                rollout_rs_threshold=rollout_corr_config.get(\"rollout_rs_threshold\", None),\n            )\n\n        # Log to tensorboard/wandb\n        for metric_name, metric_value in metrics.items():\n            logger.log_scalar(metric_name, metric_value, step=global_step)\n\n        # IMPORTANT: Update batch response_mask with rejection applied\n        batch.response_mask = modified_response_mask\n\n        # Use IS weights in training (always safety-bounded, zeroed at padding)\n        is_weights = weights_proto.batch[\"rollout_is_weights\"]\n        # ... apply weights to policy gradient ...\n```\n\n#### **Example: Conditional Alerting Based on Metrics**\n\n```python\ndef check_rollout_correction_health(metrics, config):\n    \"\"\"Check if Rollout Correction metrics indicate healthy training.\"\"\"\n    warnings = []\n\n    # Check mean IS weight\n    mean_weight = metrics['rollout_corr/rollout_is_mean']\n    if mean_weight < 0.5 or mean_weight > 2.0:\n        warnings.append(f\"Mean IS weight {mean_weight:.3f} is far from 1.0\")\n\n    # Check effective sample size\n    ess = metrics['rollout_corr/rollout_is_eff_sample_size']\n    if ess < 0.3:\n        warnings.append(f\"Effective sample size {ess:.3f} is too low\")\n\n    # Check standard deviation\n    std = metrics['rollout_corr/rollout_is_std']\n    if std > 1.0:\n        warnings.append(f\"IS weight std {std:.3f} is too high\")\n\n    # Check KL divergence\n    kl = metrics['rollout_corr/kl']\n    if abs(kl) > 0.1:\n        warnings.append(f\"KL divergence {kl:.3f} indicates significant off-policy gap\")\n\n    # Check chi-squared divergence\n    if 'rollout_corr/chi2_token' in metrics:\n        chi2_token = metrics['rollout_corr/chi2_token']\n        if chi2_token > 1.0:\n            warnings.append(f\"Chi-squared divergence (token) {chi2_token:.3f} indicates severe distribution shift\")\n\n    if warnings:\n        print(\"⚠️  Rollout Correction Health Warnings:\")\n        for warning in warnings:\n            print(f\"  - {warning}\")\n        return False\n    else:\n        print(\"✅ Rollout Correction metrics look healthy\")\n        return True\n\n# Use in training\n_, _, metrics = compute_rollout_correction_and_rejection_mask(...)\nis_healthy = check_rollout_correction_health(metrics, config)\n\nif not is_healthy:\n    # Consider adjusting config or investigating issues\n    print(\"Consider:\")\n    print(\"  - Tightening rollout_is_threshold\")\n    print(\"  - Switching to geometric aggregation level\")\n    print(\"  - Checking if rollout and training policies are too different\")\n```\n\n### Running Examples\n\nStart with the basic token-level truncate configuration:\n\n```bash\nbash examples/rollout_correction/run_with_rollout_corr.sh\n```\n\nMonitor metrics for 1-2 epochs before adjusting parameters.\n\n## Configuration Examples\n\n### Example 1: IS Weights Only (Token Level)\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token\n    rollout_is_threshold: 2.0\n    rollout_rs: null # No rejection sampling\n```\n\n### Example 2: Rejection Sampling Only (No IS Weights)\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: null # No IS weights\n    rollout_rs: token_k1\n    rollout_rs_threshold: \"0.5_2.0\"\n```\n\n### Example 3: Both IS and RS (Token RS)\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token\n    rollout_is_threshold: 2.0\n    rollout_rs: token_k1\n    rollout_rs_threshold: \"0.5_2.0\"\n```\n\n### Example 5: Bypass Mode with PPO-clip (Default)\n\n```yaml\nalgorithm:\n  rollout_correction:\n    rollout_is: token\n    rollout_is_threshold: 2.0\n    rollout_rs: token_k1\n    rollout_rs_threshold: \"0.5_2.0\"\n    bypass_mode: true # Skip old_log_prob computation\n    loss_type: ppo_clip # PPO clipped objective (default)\n```\n\n**Skips expensive `actor.compute_log_prob()` forward pass. PPO ratio = π_θ/π_rollout handles IS.**\n\n### Example 6: Bypass Mode with REINFORCE\n\n```yaml\nrollout_correction:\n  rollout_is: sequence # Explicit IS correction in loss\n  rollout_is_threshold: 2.0\n  rollout_rs: null # Optional: can add rejection sampling\n  bypass_mode: true\n  loss_type: reinforce # REINFORCE with explicit IS weights\n```\n\n**No PPO clipping, pure policy gradient with IS correction**\n\n### Example 7: Bypass Mode with PPO-clip + Rejection Sampling\n\n```yaml\nrollout_correction:\n  rollout_is: sequence # Computed for metrics\n  rollout_is_threshold: 2.0\n  rollout_rs: seq_max_k2 # Sequence max χ²/2 guard\n  rollout_rs_threshold: 2.5\n  bypass_mode: true\n  loss_type: ppo_clip # PPO clipped objective (IS handled by ratio)\n```\n\n**PPO clipping with rejection sampling. IS handled by PPO ratio (no explicit IS weights).**\n\n## Troubleshooting\n\n### Issue: High spread in IS weights\n\n**Symptoms:** `rollout_is_std` > 1.0, `rollout_is_eff_sample_size` < 0.3\n\n**Solutions:**\n\n1. Switch from `sequence` to `geometric` level\n2. Tighten thresholds\n3. Verify rollout and training aren't too different\n\n### Issue: Mean IS weight far from 1.0\n\n**Symptoms:** `rollout_is_mean` < 0.5 or > 2.0\n\n**Solutions:**\n\n1. Verify `calculate_log_probs=True` is set\n2. Check rollout_log_probs are correctly passed\n3. Check for systematic distribution shift\n\n### Debugging: Visualizing Metrics\n\n**Example: Plot IS weight distribution**\n\n```python\nimport matplotlib.pyplot as plt\nimport numpy as np\n\ndef plot_is_metrics(metrics_history):\n    \"\"\"Plot rollout IS metrics over training steps.\"\"\"\n    fig, axes = plt.subplots(2, 3, figsize=(15, 10))\n\n    # Plot 1: Mean IS weight over time\n    axes[0, 0].plot(metrics_history['rollout_corr/rollout_is_mean'])\n    axes[0, 0].axhline(y=1.0, color='r', linestyle='--', label='Ideal')\n    axes[0, 0].set_title('Mean IS Weight')\n    axes[0, 0].set_xlabel('Step')\n    axes[0, 0].legend()\n\n    # Plot 2: Effective sample size\n    axes[0, 1].plot(metrics_history['rollout_corr/rollout_is_eff_sample_size'])\n    axes[0, 1].axhline(y=0.5, color='g', linestyle='--', label='Good')\n    axes[0, 1].axhline(y=0.3, color='r', linestyle='--', label='Warning')\n    axes[0, 1].set_title('Effective Sample Size')\n    axes[0, 1].set_xlabel('Step')\n    axes[0, 1].legend()\n\n    # Plot 3: KL divergence over time\n    axes[1, 0].plot(metrics_history['rollout_corr/kl'], label='KL')\n    axes[1, 0].plot(metrics_history['rollout_corr/k3_kl'], label='K3 KL')\n    axes[1, 0].axhline(y=0, color='g', linestyle='--', alpha=0.3)\n    axes[1, 0].set_title('KL Divergence')\n    axes[1, 0].set_xlabel('Step')\n    axes[1, 0].legend()\n\n    # Plot 4: PPL ratio over time\n    axes[1, 1].plot(metrics_history['rollout_corr/ppl_ratio'])\n    axes[1, 1].axhline(y=1.0, color='r', linestyle='--', label='Ideal')\n    axes[1, 1].set_title('PPL Ratio (Training/Rollout)')\n    axes[1, 1].set_xlabel('Step')\n    axes[1, 1].legend()\n\n    # Plot 5: Chi-squared divergence\n    if 'rollout_corr/chi2_token' in metrics_history:\n        axes[1, 2].plot(metrics_history['rollout_corr/chi2_token'], label='Token-level')\n        if 'rollout_corr/chi2_seq' in metrics_history:\n            axes[1, 2].plot(metrics_history['rollout_corr/chi2_seq'], label='Seq-level')\n        axes[1, 2].axhline(y=1.0, color='r', linestyle='--', label='Warning')\n        axes[1, 2].set_title('Chi-squared Divergence')\n        axes[1, 2].set_xlabel('Step')\n        axes[1, 2].legend()\n    else:\n        axes[1, 2].axis('off')\n\n    plt.tight_layout()\n    plt.savefig('rollout_is_metrics.png', dpi=150)\n    print(\"Saved plot to rollout_is_metrics.png\")\n```\n\n**Example: Metric collection during training**\n\n```python\n# Collect metrics over time\nmetrics_history = {\n    'rollout_corr/rollout_is_mean': [],\n    'rollout_corr/rollout_is_eff_sample_size': [],\n    'rollout_corr/kl': [],\n    'rollout_corr/k3_kl': [],\n    'rollout_corr/ppl_ratio': [],\n    'rollout_corr/chi2_token': [],\n    'rollout_corr/chi2_seq': [],\n}\n\n# In training loop\nfor step in range(num_steps):\n    # ... compute IS weights and rejection mask ...\n    _, _, metrics = compute_rollout_correction_and_rejection_mask(...)\n\n    # Store metrics\n    for key in metrics_history.keys():\n        if key in metrics:\n            metrics_history[key].append(metrics[key])\n\n    # Plot every 100 steps\n    if step % 100 == 0:\n        plot_is_metrics(metrics_history)\n```\n\n## Performance Impact\n\n- **Memory overhead**: ~1% of model memory\n- **Computational overhead**: 1-3% depending on level\n- **Training stability**: Significantly improved when off-policy gap exists\n\n## Testing\n\nRun the test suite to verify everything works:\n\n```bash\n# Basic unit tests\npython tests/trainer/ppo/test_rollout_corr.py\n\n# Integration tests (if pytest is available)\npytest tests/trainer/ppo/test_rollout_corr_integration.py -v\n```\n\nExpected output: All tests pass ✓\n\n## Additional Resources\n\n- **Implementation**: `verl/trainer/ppo/rollout_corr_helper.py`\n- **Examples**: `examples/rollout_correction/`\n- **DAPO Example**: `recipe/dapo/run_dapo_qwen2.5_32b_rollout_corr.sh`\n\n## Summary\n\nRollout Correction provides a unified framework for handling general off-policy problems in RL:\n\n- ✅ Corrects ANY distribution shift between data collection and training\n- ✅ Supports diverse scenarios: policy mismatch, staleness, replay buffers, off-policy algorithms\n- ✅ Numerical stability with safety bounds and rejection mechanisms\n- ✅ Comprehensive diagnostics: KL, perplexity, χ² divergence\n- ✅ Flexible methods from token-level to sequence-level aggregation\n- ✅ Memory-efficient implementation\n\n## References\n\n- **[Mathematical Formulations](rollout_corr_math.md)** - Detailed mathematical theory and derivations for all rollout correction methods\n- [Your Efficient RL Framework Secretly Brings You Off-Policy RL Training](https://fengyao.notion.site/off-policy-rl)\n"
  },
  {
    "path": "docs/algo/rollout_corr_math.md",
    "content": "# Mathematical Formulations of Rollout Correction Methods in `verl`\n\n**Author:** [Yingru Li](https://richardli.xyz)\n**Last updated:** 2025-11-04\n\n---\n\n> **📖 Documentation Structure**\n> - **This document** - Mathematical theory: formulations, derivations, and algorithmic foundations\n> - **[Rollout Correction Usage Guide](rollout_corr.md)** - Practical implementation: configurations, presets, troubleshooting\n>\n> Start here for theory and design rationale, refer to the usage guide for implementation.\n\n---\n\n### BibTeX Citation\n\n```bibtex\n@online{liu-li-2025-rl-collapse,\n  title = {When Speed Kills Stability: Demystifying {RL} Collapse from the Training-Inference Mismatch},\n  author = {Liu, Jiacai and Li, Yingru and Fu, Yuqian and Wang, Jiawei and Liu, Qian and Shen, Yu},\n  year = {2025},\n  month = sep,\n  url = {https://richardli.xyz/rl-collapse}\n}\n\n\n@article{li2025trust,\n  title={Trust Region Masking for Long-Horizon LLM Reinforcement Learning},\n  author={Li, Yingru and Liu, Jiacai and Xu, Jiawei and Tong, Yuxuan and Li, Ziniu and Liu, Qian and Wang, Baoxiang},\n  journal={arXiv preprint arXiv:2512.23075},\n  year={2025}\n}\n```\n\n### Blog Series\n\n- Main blog post: https://richardli.xyz/rl-collapse\n- [Part 1: Why Mismatch Breaks LLM-RL](https://richardli.xyz/rl-collapse-1) (analytical framework using TV distance for bias and χ²-divergence for variance)\n- [Part 2: The Gradient Estimator Trials](https://richardli.xyz/rl-collapse-2) (token-level vs sequence-level correction bias-variance tradeoff)\n- [Part 3: When Math Meets Reality—Toxic Tails and Length Traps](https://richardli.xyz/rl-collapse-3) (why rejection over clipping, and geometric-level RS)\n- Latest Paper: https://arxiv.org/abs/2512.23075\n\n## Abstract\n\nThis document provides the definitive mathematical formulations for rollout correction methods in `verl`, following the natural progression from **REINFORCE** to **PPO** to **Decoupled PPO**.\n\nRollout correction provides a unified framework to handle **general off-policy problems** in RL training - any scenario where the data collection distribution differs from the training distribution.\n\n**Applicable scenarios include:**\n- **Policy mismatch**: Different precision (FP8 vs FP16 vs BF16 vs FP32), different backends (vLLM vs SGLang vs FSDP vs Megatron)\n- **Temporal lag**: Model staleness, asynchronous rollout workers\n- **Replay buffers**: Training on historical trajectories from earlier policy versions\n- **Off-policy algorithms**: Behavioral cloning, DAPO, expert demonstrations\n- **Data filtering**: Reweighting, preference learning, curriculum learning\n\n---\n\n## Table of Contents\n\n1. [Theoretical Foundation: From REINFORCE to Decoupled PPO](#1-theoretical-foundation-from-reinforce-to-decoupled-ppo)\n2. [Implementation in verl: The Three-Policy Framework](#2-implementation-in-verl-the-three-policy-framework)\n3. [Algorithmic Components and Combinations](#3-algorithmic-components-and-combinations)\n4. [Off-Policy Diagnostic Metrics](#4-off-policy-diagnostic-metrics)\n5. [Summary and Decision Guide](#5-summary-and-decision-guide)\n6. [Implementation References](#6-implementation-references)\n\n---\n\n## 1. Theoretical Foundation: From REINFORCE to Decoupled PPO\n\nThis section establishes the theoretical progression that `verl` implements.\n\n### 1.1 REINFORCE: Policy Gradient Baseline\n\nThe REINFORCE algorithm ([Williams, 1992](https://doi.org/10.1007/BF00992696)) is the foundation of policy gradient methods.\n\n**Vanilla REINFORCE (On-Policy)**\n\nFor trajectories $\\tau = (s_0, a_0, s_1, a_1, \\ldots, s_T, a_T)$ sampled from the current policy $\\pi_\\theta$, the policy gradient is:\n\n$$\n\\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left[ \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere $A_t$ is the advantage function at timestep $t$.\n\n**Off-Policy REINFORCE**\n\nWhen trajectories are sampled from a different behavior policy $\\mu$, we apply importance sampling over the **joint trajectory distribution**:\n\n$$\n\\nabla_\\theta J(\\theta) = \\mathbb{E}_{\\tau \\sim \\mu} \\left[ \\frac{P_{\\pi_\\theta}(\\tau)}{P_\\mu(\\tau)} \\sum_{t=0}^T \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere the trajectory-level importance weight is:\n\n$$\n\\frac{P_{\\pi_\\theta}(\\tau)}{P_\\mu(\\tau)} = \\frac{p(s_0) \\prod_{t=0}^T \\pi_\\theta(a_t|s_t) p(s_{t+1}|s_t, a_t)}{p(s_0) \\prod_{t=0}^T \\mu(a_t|s_t) p(s_{t+1}|s_t, a_t)} = \\prod_{t=0}^T \\frac{\\pi_\\theta(a_t|s_t)}{\\mu(a_t|s_t)}\n$$\n\nThe transition dynamics $p(s_{t+1}|s_t, a_t)$ and initial state $p(s_0)$ cancel out, leaving only the product of per-step action probability ratios.\n\n**Key properties:**\n- **Off-policy capable**: Can learn from any behavior policy via importance sampling\n- **No trust region**: Policy updates not constrained\n\n**Implementation in verl:** The `bypass_pg_is` preset implements off-policy REINFORCE with truncated importance sampling.\n\n### 1.2 PPO: Adding Trust Region Control\n\nProximal Policy Optimization ([Schulman et al., 2017](https://arxiv.org/abs/1707.06347)) adds a clipped surrogate objective:\n\n$$\nL_{\\text{PPO}}(\\theta) = -\\mathbb{E}_{(s,a) \\sim \\mu} \\left[ \\min\\left( r_t(\\theta) A_t, \\text{clip}(r_t(\\theta), 1-\\epsilon, 1+\\epsilon) A_t \\right) \\right]\n$$\n\nwhere $r_t(\\theta) = \\frac{\\pi_\\theta(a_t|s_t)}{\\mu(a_t|s_t)}$ and $\\epsilon$ is the clip range (typically 0.2).\n\n**Key properties:**\n- **Two policies**: $\\mu$ (reference for clipping) and $\\pi_\\theta$ (being updated)\n- **Trust region via clipping**: Limits policy update magnitude via ratio $r_t(\\theta) = \\frac{\\pi_\\theta}{\\mu}$\n\n### 1.3 Decoupled PPO: Achieving Batch Size Invariance\n\nDecoupled PPO ([Hilton et al., 2021](https://arxiv.org/abs/2110.00641)) solves PPO's batch size sensitivity by **decoupling two roles**:\n1. **Proximal policy** $\\pi_{\\text{prox}}$: The anchor policy for PPO clipping (controls policy update size)\n2. **Behavior policy** $\\mu$: The policy that collected the data (for off-policy correction via importance sampling)\n\n**The problem**: Standard PPO controls policy update size via the ratio $\\frac{\\pi_\\theta}{\\pi_{\\text{old}}}$, where $\\pi_{\\text{old}}$ is assumed to be both the proximal policy *and* the behavior policy. This coupling makes the algorithm sensitive to batch size because aggregating data from multiple workers or using replay buffers changes the effective behavior policy.\n\n**The solution**: Decouple these two roles, leading to a **three-policy formulation**:\n\n$$\nL_{\\text{DecoupledPPO}}(\\theta) = -\\mathbb{E}_{(s,a) \\sim \\mu} \\left[ w_t \\cdot \\min\\left( r_t(\\theta) A_t, \\text{clip}(r_t(\\theta), 1-\\epsilon, 1+\\epsilon) A_t \\right) \\right]\n$$\n\nwhere:\n- $w_t = \\frac{\\pi_{\\text{prox}}(a_t|s_t)}{\\mu(a_t|s_t)}$: Importance sampling weight (corrects for behavior policy $\\mu$). Here $\\pi_{\\text{prox}}$ is frozen during training, so $w_t$ is constant (no stopgrad operator needed).\n- $r_t(\\theta) = \\frac{\\pi_\\theta(a_t|s_t)}{\\pi_{\\text{prox}}(a_t|s_t)}$: PPO ratio (controls policy update size against proximal policy $\\pi_{\\text{prox}}$)\n\n**Key properties**: By decoupling:\n- **Batch size invariance**: Policy update control (via $\\pi_{\\text{prox}}$) is independent of data aggregation\n- **Flexible behavior policy**: Any $\\mu$ can be used (different workers, replay buffers, or stale checkpoints)\n- **Stale data utilization**: Older trajectories can be corrected via importance sampling\n- **Clipping preserved**: Clipping against $\\pi_{\\text{prox}}$ limits update magnitude\n\n**This is the algorithm that `verl` implements via its three-policy framework.**\n\n---\n\n## 2. Implementation in verl: The Three-Policy Framework\n\nThe `verl` library implements decoupled PPO using three distinct policies, each serving a specific role.\n\n### 2.1 Policy Roles and Notation\n\n**$\\pi_{\\text{rollout}}$ (Behavior Policy $\\mu$)**\nThe policy used for data collection. This is the behavior distribution $\\mu$ from theory.\n\n- **When created**: During rollout/data collection phase\n- **Purpose**: Generate trajectories for training\n- **Common sources**:\n  - Policy mismatch: Same weights, different implementation (precision, backend)\n  - Temporal lag: Stale checkpoint from async workers\n  - Replay buffer: Historical data from earlier iterations\n  - Off-policy algorithms: Expert demonstrations, auxiliary policies (DAPO)\n  - Data filtering: Reweighted or filtered data\n- **Fixed**: Frozen during training on a batch\n\n**$\\pi_{\\text{old}}$ (Proximal Policy $\\pi_{\\text{prox}}$)**\nThe reference policy for PPO clipping. This is the \"proximal policy\" from decoupled PPO theory.\n\n- **When created**:\n  - **Decoupled mode**: Computed at start of training epoch via `actor.compute_log_prob()`\n  - **Bypass mode**: Set equal to $\\pi_{\\text{rollout}}$ (skips separate computation)\n- **Purpose**:\n  - Anchor point for PPO clipping (controls policy update size)\n  - When separate from $\\pi_{\\text{rollout}}$: Enables batch size invariance and efficient use of stale data\n- **Fixed**: Frozen during all PPO update epochs on the same batch\n\n**$\\pi_{\\theta}$ (Current Policy)**\nThe policy being actively optimized during training.\n\n- **Updated**: Every gradient step\n- **Purpose**: The policy we're improving\n\n### 2.2 Operating Modes\n\nThe three-policy framework can operate in two modes:\n\n**Decoupled Mode (Three Policies)**\n- Computes $\\pi_{\\text{old}}$ separately at the start of each training epoch\n- **Algorithm**: Full decoupled PPO with three policies (mathematically correct)\n- **Properties**: Achieves batch size invariance; separately corrects Drift 1 (rollout→old) and Drift 2 (old→current)\n\n**Bypass Mode (Two Policies)**\n- Sets $\\pi_{\\text{old}} = \\pi_{\\text{rollout}}$ (skips separate computation)\n- **Algorithm**: Uses $\\pi_{\\text{rollout}}$ as both behavior policy and proximal policy (mathematically correct)\n- **Key difference**: Proximal policy equals behavior policy, so no IS correction needed between them\n- **Properties**: Faster (skips `actor.compute_log_prob()` call); does not achieve batch size invariance\n\n### 2.3 Two Distribution Shifts\n\nThe three-policy framework handles two types of distribution drift:\n\n**Drift 1: $\\pi_{\\text{rollout}} \\to \\pi_{\\text{old}}$ (Off-Policy Gap)**\n\nThis is the distribution shift between the data collection policy and the training reference policy.\n\n- **Nature**: Ranges from negligible (same checkpoint, minor differences) to severe (replay buffers, expert data)\n- **Correction**: Importance sampling weight $w_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$\n- **Optional**: Can be ignored (bypass mode) when negligible\n\n**Drift 2: $\\pi_{\\text{old}} \\to \\pi_{\\theta}$ (Policy Update Drift)**\n\nThis is the drift from policy parameter updates during training.\n\n- **Nature**: Occurs as $\\pi_\\theta$ is updated via gradient descent\n- **Correction**: PPO clipping on ratio $r_t(\\theta) = \\frac{\\pi_\\theta(a_t|s_t)}{\\pi_{\\text{old}}(a_t|s_t)}$\n- **Universal**: Applies to both on-policy and off-policy training\n\n### 2.4 Notation Summary\n\n- $\\pi_{\\text{rollout}}$: Behavior policy (data collection)\n- $\\pi_{\\text{old}}$: Proximal policy (PPO anchor)\n- $\\pi_{\\theta}$: Current policy (being updated)\n- $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$: Per-token IS ratio (corrects Drift 1)\n- $r_t(\\theta) = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{old}}(a_t|s_t)}$: PPO ratio (corrects Drift 2)\n- $A_t$: Advantage at token $t$\n- $T$: Set of valid tokens in a sequence\n- $C_{\\text{IS}}$: Upper threshold for IS weights (e.g., 2.0)\n- $C_{\\text{RS-upper}}$: Upper threshold for RS mask (e.g., 2.0)\n- $C_{\\text{RS-lower}}$: Lower threshold for RS mask (typically $1/C_{\\text{RS-upper}}$)\n- $\\epsilon$: PPO clip range (typically 0.2)\n\n---\n\n## 3. Algorithmic Components and Combinations\n\nThe rollout correction framework in `verl` is built from **orthogonal components** that can be combined flexibly:\n\n1. **Operating Mode**: How $\\pi_{\\text{old}}$ is computed (Decoupled vs Bypass)\n2. **Loss Function**: PPO (with clipping) vs Pure IS (policy gradient only)\n3. **IS/RS Aggregation Level**: Token, Sequence, or Geometric\n\nThis section explains each component and their valid combinations.\n\n### 3.1 Operating Modes: Decoupled vs Bypass\n\nThe operating mode determines how the proximal policy $\\pi_{\\text{old}}$ is computed.\n\n#### 3.1.1 Decoupled Mode (Three Policies)\n\n**Configuration:** `bypass_mode = false`\n\n**Policy setup:**\n- $\\pi_{\\text{rollout}}$: Behavior policy (data collection)\n- $\\pi_{\\text{old}}$: Proximal policy (computed via `actor.compute_log_prob()` at start of training epoch)\n- $\\pi_{\\theta}$: Current policy (being updated)\n\n**IS ratio:** $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ (corrects Drift 1: rollout→old)\n\n**PPO ratio:** $r_t(\\theta) = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{old}}(a_t|s_t)}$ (corrects Drift 2: old→current)\n\n**Properties:**\n- ✅ Achieves batch size invariance\n- ✅ Separately corrects two distribution drifts\n- ✅ Efficient stale data utilization\n- ❌ Extra forward pass needed (`actor.compute_log_prob()`)\n\n#### 3.1.2 Bypass Mode (Two Policies)\n\n**Configuration:** `bypass_mode = true`\n\n**Policy setup:**\n- $\\pi_{\\text{rollout}}$: Behavior policy (data collection)\n- $\\pi_{\\text{old}} = \\pi_{\\text{rollout}}$: Proximal policy equals behavior policy\n- $\\pi_{\\theta}$: Current policy (being updated)\n\n**Ratios:**\n- **With PPO-clip loss** (`loss_type = \"ppo_clip\"`, default): PPO ratio $r_t(\\theta) = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ clips against rollout policy (IS handled by ratio)\n- **With REINFORCE loss** (`loss_type = \"reinforce\"`): IS ratio $\\rho_t = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ computed on-the-fly in loss function\n\n**Properties:**\n- ✅ Skips `actor.compute_log_prob()` call (faster)\n- ✅ Handles off-policy correction via IS/RS (when using policy gradient with IS/RS)\n- ✅ Uses two policies instead of three (π_rollout = π_old)\n- ⚠️ Does not separate proximal policy from behavior policy (unlike decoupled mode)\n\n---\n\n### 3.2 Loss Functions: PPO vs Policy Gradient\n\n#### 3.2.1 PPO Loss (with Clipping)\n\n**Configuration:** `loss_type = \"ppo_clip\"` (default in bypass mode)\n\n**Loss function:**\n\n$$\nL_{\\text{PPO}}(\\theta) = -\\mathbb{E}_t \\left[ w_t \\cdot \\min\\left( r_t(\\theta) A_t, \\text{clip}(r_t(\\theta), 1-\\epsilon, 1+\\epsilon) A_t \\right) \\right]\n$$\n\nwhere:\n- $w_t$: IS weight (depends on aggregation level, see Section 3.3). In decoupled mode, $w_t = \\frac{\\pi_{\\text{old}}}{\\pi_{\\text{rollout}}}$ where $\\pi_{\\text{old}}$ is frozen, so $w_t$ is constant (no stopgrad needed). In bypass mode with PPO loss, no separate IS weights are typically computed.\n- $r_t(\\theta) = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{old}}(a_t|s_t)}$: PPO ratio\n- $\\epsilon$: Clip range (typically 0.2)\n\n**Properties:**\n- Trust region control via clipping\n- Limits policy update magnitude\n- Standard in RL training\n\n#### 3.2.2 Policy Gradient Loss (with IS/RS Correction)\n\n**Configuration:** `loss_type = \"reinforce\"` (requires `bypass_mode = true`)\n\n**Loss function** (example with sequence-level IS):\n\n$$\nL_{\\text{PG}}(\\theta) = -\\mathbb{E}_{(s,a) \\sim \\pi_{\\text{rollout}}} \\left[ \\text{stopgrad}(w_{\\text{seq}}(\\theta)) \\cdot \\sum_{t \\in T} \\log \\pi_{\\theta}(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere:\n- $w_{\\text{seq}}(\\theta)$: Sample weight (IS or RS, see §3.3-3.4 for details)\n- For IS: $w_{\\text{seq}}(\\theta) = \\min\\left( \\prod_{t \\in T} \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}, C_{\\text{IS}} \\right)$\n- For RS: $w_{\\text{seq}}(\\theta) \\in \\{0, 1\\}$ (binary rejection mask)\n- **stopgrad operator**: The weight $w_{\\text{seq}}(\\theta)$ is computed using $\\pi_\\theta$ but treated as a **constant coefficient** when computing $\\nabla_\\theta L$. This is essential for importance sampling correctness (see theoretical justification below).\n\n**Effective gradient:**\n\n$$\n\\nabla_\\theta L_{\\text{PG}} = -\\mathbb{E}_{(s,a) \\sim \\pi_{\\text{rollout}}} \\left[ \\text{stopgrad}(w_{\\text{seq}}(\\theta)) \\cdot \\sum_{t \\in T} \\nabla_\\theta \\log \\pi_{\\theta}(a_t|s_t) \\cdot A_t \\right]\n$$\n\n**Theoretical Justification for stopgrad:**\n\nThe stopgrad operator is **mathematically required** by importance sampling theory, not an implementation detail. Here's why:\n\n**The fundamental principle**: Importance sampling is a technique to **change the measure** (reweight samples from one distribution to estimate expectations under another), not to optimize the reweighting function itself.\n\n**Formal derivation**:\n\n1. **Original objective**: We want to optimize $J(\\theta) = \\mathbb{E}_{\\tau \\sim \\pi_\\theta}[\\sum_t A_t]$.\n\n2. **Off-policy setting**: We only have samples from $\\pi_{\\text{rollout}}$, so we use importance sampling:\n   $$\n   J(\\theta) = \\mathbb{E}_{\\tau \\sim \\pi_{\\text{rollout}}} \\left[ \\underbrace{\\frac{P_{\\pi_\\theta}(\\tau)}{P_{\\pi_{\\text{rollout}}}(\\tau)}}_{w(\\tau;\\theta)} \\sum_t A_t \\right]\n   $$\n\n3. **Computing the policy gradient**: The correct gradient uses the **policy gradient theorem BEFORE importance sampling**:\n   $$\n   \\begin{aligned}\n   \\nabla_\\theta J(\\theta) &= \\nabla_\\theta \\mathbb{E}_{\\tau \\sim \\pi_\\theta}\\left[\\sum_t A_t\\right] \\\\\n   &= \\mathbb{E}_{\\tau \\sim \\pi_\\theta} \\left[\\sum_t A_t \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\right] \\quad \\text{(policy gradient theorem)} \\\\\n   &= \\mathbb{E}_{\\tau \\sim \\pi_{\\text{rollout}}} \\left[ w(\\tau;\\theta) \\sum_t A_t \\nabla_\\theta \\log \\pi_\\theta(a_t|s_t) \\right] \\quad \\text{(change of measure)}\n   \\end{aligned}\n   $$\n\n   In the final line, $w(\\tau;\\theta)$ appears as a **multiplicative coefficient** from the change of measure, not as something we differentiate.\n\n4. **What goes wrong without stopgrad**: If we naively compute $\\nabla_\\theta \\left[w(\\theta) \\log \\pi_\\theta \\right]$ in the loss, we get:\n   $$\n   \\nabla_\\theta \\left[w(\\theta) \\log \\pi_\\theta \\right] = \\underbrace{\\log \\pi_\\theta \\cdot \\nabla_\\theta w(\\theta)}_{\\text{WRONG: bias term}} + \\underbrace{w(\\theta) \\cdot \\nabla_\\theta \\log \\pi_\\theta}_{\\text{CORRECT: IS-weighted gradient}}\n   $$\n\n   The first term $\\log \\pi_\\theta \\cdot \\nabla_\\theta w(\\theta)$ is an artifact of the computational trick (using loss times log-prob), not part of the true policy gradient. It biases the gradient estimator and optimizes a different objective than $J(\\theta)$.\n\n5. **Implementation requirement**: In PyTorch, to compute only the second term, we must use:\n   ```python\n   loss = -advantages * log_prob * rollout_is_weights.detach()  # stopgrad on weights\n   ```\n   Without `.detach()`, autograd computes both terms, giving an incorrect gradient.\n\n**Intuition**: The IS weight $w(\\theta)$ tells us \"how much to trust this sample\" for estimating the gradient under $\\pi_\\theta$. We update $\\theta$ to maximize the reweighted objective, but we don't update $\\theta$ to maximize the weight itself—that would be circular reasoning (optimizing the correction factor instead of the actual objective).\n\n**Properties:**\n- **Algorithm**: Off-policy policy gradient with IS/RS correction\n- **Loss types** (`loss_type` config option in bypass mode):\n  - `\"ppo_clip\"` (default): PPO clipped objective\n    - $L = -\\mathbb{E}[\\min(r \\cdot A, \\text{clip}(r) \\cdot A)]$ where $r = \\pi_\\theta / \\pi_{\\text{rollout}}$\n    - Note: IS weights NOT applied (PPO ratio already handles it; would be double-counting)\n  - `\"reinforce\"`: Pure policy gradient with explicit IS weights, no PPO clipping\n    - $L = -\\mathbb{E}[w \\cdot \\log \\pi_\\theta(a|s) \\cdot A]$ where $w = \\pi_\\theta / \\pi_{\\text{rollout}}$\n- **Always uses bypass mode**: Direct $\\pi_\\theta$ to $\\pi_{\\text{rollout}}$ comparison\n- **Fast**: Single forward pass\n\n**Implementation:** `compute_policy_loss_bypass_mode()` and `compute_policy_loss_reinforce()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py)\n\n---\n\n### 3.3 IS/RS Aggregation Levels\n\nThe aggregation level determines how per-token probability ratios are combined into IS weights and/or rejection masks. This choice is **orthogonal to the operating mode** - you can use any aggregation level in either decoupled or bypass mode.\n\n#### 3.3.1 Token-Level Aggregation\n\n**IS weights:** $w_t = \\min(\\rho_t, C_{\\text{IS}})$ where $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ (decoupled) or $\\rho_t = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ (bypass/pure IS)\n\n**Configuration:**\n```python\nrollout_is = \"token\"  # IS weights\nrollout_rs = \"token_k1\"  # Optional: rejection sampling (ratio bounds)\n```\n\n**Properties:**\n- Independent truncation per token\n- Lower variance than sequence-level (product of ratios bounded individually)\n- **Bias-variance tradeoff**: Token-level correction has $O(T^2 \\Delta_{\\max})$ bias where $T$ is sequence length and $\\Delta_{\\max}$ is maximum per-token policy divergence. This bias becomes significant when the rollout policy deviates substantially from the training policy. Sequence-level correction is unbiased but has higher variance.\n- Typical threshold: 1.5 - 5.0\n- Optional batch normalization [§3.4](rollout_corr_math.md#34-batch-normalization): Normalizes over all token weights to ensure $\\mathbb{E}[\\tilde{w}_t] = 1$ (reduces variance)\n- **When to use**: Token-level works well when rollout policy stays within the trust region of training policy. When mismatch is significant, the bias becomes intolerable and sequence-level correction is preferred.\n\n**Loss function (REINFORCE + Token IS):**\n\n$$\nL_{\\text{REINFORCE+TIS}}(\\theta) = -\\mathbb{E}_t \\left[ \\text{stopgrad}(w_t) \\cdot \\log \\pi_\\theta(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere $w_t = \\min(\\rho_t, C_{\\text{IS}})$ are the truncated token-level IS weights. The stopgrad operator ensures that when computing $\\nabla_\\theta L$, the weights are treated as constants (see §3.2.2 for theoretical justification). This formulation can also be combined with PPO clipping by replacing the REINFORCE gradient with the clipped surrogate objective.\n\n**Implementation:**\n- IS weights: `compute_rollout_correction_weights()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L325-L402)\n- Loss: `compute_policy_loss()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py#L812-L884)\n\n#### 3.3.2 Sequence-Level Aggregation\n\n**IS weights:** $w_{\\text{seq}} = \\min\\left( \\prod_{t \\in T} \\rho_t, C_{\\text{IS}} \\right) = \\min\\left( \\exp\\left(\\sum_{t \\in T} \\log \\rho_t\\right), C_{\\text{IS}} \\right)$ (broadcast to all tokens)\n\n**Configuration:**\n```python\nrollout_is = \"sequence\"  # IS weights\nrollout_rs = \"seq_sum_k1\"  # Optional: rejection sampling\n```\n\n**Properties:**\n- Multiplicative aggregation across sequence\n- More sensitive to outliers than token-level\n- Typical threshold: 2.0 - 10.0\n- Optional batch normalization [§3.4](rollout_corr_math.md#34-batch-normalization): Normalizes over sequence means (one weight per sequence)\n\n**Terminology Note:**\n- **Seq-TIS (Sequence-Level Truncated IS)**: Clips the sequence ratio $\\rho(\\tau) \\to \\min(\\rho(\\tau), C)$. Maximizes information efficiency by extracting signal from all samples. Best for clean data with moderate mismatch.\n- **Seq-MIS (Sequence-Level Masked IS)**: Rejects (masks) sequences with $\\rho(\\tau) > C$ instead of clipping. Acts as a hard trust region filter. Best for severe mismatch or when the distribution tail is \"toxic\" (contains garbage/adversarial samples rather than signal).\n\n**Loss function (REINFORCE + Sequence IS):**\n\n$$\nL_{\\text{REINFORCE+SeqIS}}(\\theta) = -\\mathbb{E}_t \\left[ \\text{stopgrad}(w_{\\text{seq}}) \\cdot \\log \\pi_\\theta(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere $w_{\\text{seq}}$ is broadcast to all tokens in the sequence. The stopgrad operator ensures correct IS gradient computation (see §3.2.2). This formulation can also be combined with PPO clipping.\n\n#### 3.3.3 Geometric Mean Aggregation (Geo-RS)\n\n**Geometric mean ratio:** $\\rho_{\\text{geo}} = \\exp\\left( \\frac{1}{|T|} \\sum_{t \\in T} \\log \\rho_t \\right) = \\left(\\prod_{t \\in T} \\rho_t\\right)^{1/|T|}$ (broadcast to all tokens)\n\n**Configuration:**\n```python\nrollout_is = null  # No IS weights, pure rejection\nrollout_rs = \"seq_mean_k1\"  # Geometric mean rejection sampling (ratio bounds)\n```\n\n**Properties:**\n- Length-invariant (normalizes by sequence length)\n- Ideal ratio = 1.0 (policies match)\n- Typical bounds: `\"0.999_1.001\"` (~±0.1%)\n- **Used for rejection sampling only, not IS weighting**\n\n**The Length Trap Problem:**\n\nStandard IS estimators have a systematic **length bias** that penalizes long sequences. The importance ratio $\\rho(y)$ is multiplicative:\n\n$$\n\\rho(y) = \\prod_{t=1}^T \\frac{\\pi(y_t|y_{<t})}{\\mu(y_t|y_{<t})}\n$$\n\nAssume the new policy $\\pi$ differs slightly from $\\mu$, with average per-token ratio $\\approx 1.1$:\n- **Short sequence (10 tokens):** $\\rho \\approx 1.1^{10} \\approx 2.6$ → fits within threshold, **kept**\n- **Long sequence (100 tokens):** $\\rho \\approx 1.1^{100} \\approx 13,780$ → explodes past threshold, **rejected**\n\nThis creates **Context Collapse**: the model preferentially learns from short, shallow answers and rejects long chains of thought—even if per-step quality is identical. For reasoning models (CoT) and agents, this effectively penalizes \"thinking too long.\"\n\n**Geo-RS Solution:**\n\nGeometric-level rejection normalizes by sequence length, converting the extensive property (total probability product) to an intensive property (average per-token drift):\n\n$$\n\\rho_{\\text{geo}}(y) = \\rho(y)^{1/T}\n$$\n\nNow both sequences have the same \"trust score\":\n- **Short (10 tokens):** $(1.1^{10})^{1/10} = 1.1$\n- **Long (100 tokens):** $(1.1^{100})^{1/100} = 1.1$\n\n**Why tight thresholds?**\nFor 100 tokens with per-token log-ratio = 0.01 each:\n- Arithmetic product ratio: $e^{100 \\times 0.01} \\approx 2.7$\n- Geometric ratio: $e^{0.01} \\approx 1.010$\n\nA ratio bound of `\"0.999_1.001\"` rejects sequences whose average per-token log-deviation exceeds ≈0.1%.\n\n**Loss function (REINFORCE + Geometric RS):**\n\n$$\nL_{\\text{GeoRS}}(\\theta) = -\\mathbb{E}_{(s,a) \\mid \\text{seq} \\in \\mathcal{A}_{\\text{geo}}} \\left[ \\sum_{t \\in T} \\log \\pi_\\theta(a_t|s_t) \\cdot A_t \\right]\n$$\n\nwhere $\\mathcal{A}_{\\text{geo}} = \\{ \\text{seq} : C_{\\text{RS-lower}} \\leq \\rho_{\\text{geo}} \\leq C_{\\text{RS-upper}} \\}$ is the acceptance set (rejection mask). No IS weights are used, so no stopgrad needed. This formulation can also be combined with PPO clipping.\n\n**Combined Estimator (Geo-RS-Token-TIS):**\n\nFor best results, combine the **Geometric Filter** (length-invariant validity check) with **Token-level IS weights** (lower variance):\n\n$$\n\\hat{g}_{\\text{geo-rs-token-tis}}(y) = \\underbrace{\\mathbb{I}\\left( C_{\\text{low}} \\le \\rho(y)^{1/T} \\le C_{\\text{high}} \\right)}_{\\text{Geometric Filter}} \\cdot \\prod_t \\min(\\rho_t, C) \\cdot f(y)\n$$\n\nThis is implemented by combining `rollout_rs=\"seq_mean_k1\"` with `rollout_is=\"token\"`.\n\n#### 3.3.4 K2 Divergence Aggregation\n\n**Per-token statistic:**\n\n$$\nK2_t = \\frac{1}{2} \\left(\\log \\rho_t\\right)^2\n$$\n\nwhere $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ and the implementation clips $\\log \\rho_t$ to $[-20, 20]$ for numerical safety.\n\n**Sequence aggregations (share the same per-token $K2_t$):**\n- `seq_sum_k2`: $K2_{\\text{sum}} = \\sum_{t \\in T} K2_t$\n- `seq_mean_k2`: $K2_{\\text{mean}} = \\frac{1}{|T|} \\sum_{t \\in T} K2_t$\n- `seq_max_k2`: $K2_{\\text{max}} = \\max_{t \\in T} K2_t$\n\n**Configuration:**\n```python\nrollout_is = null            # Optional: pair with token IS weights for lower variance\nrollout_rs = \"token_k2\"      # or \"seq_sum_k2\", \"seq_mean_k2\", \"seq_max_k2\"\nrollout_rs_threshold = 2.0   # Positive upper bound only\n```\n\n**Properties:**\n- Symmetric quadratic penalty in $\\log \\rho_t$; equals zero when policies match.\n- Approximates $\\tfrac{1}{2}\\operatorname{Var}[\\log \\rho]$ for small policy drift, making it a smooth detector of mismatch.\n- Upper-threshold only: typical ranges are 1.5-3.0 for `token_k2`, 2.0-2.5 for `seq_mean_k2`, and 2.5-4.0 for `seq_sum_k2`.\n- `seq_max_k2` isolates single-token spikes even when the rest of the sequence is clean.\n- Can co-exist with token-level IS weights (`rollout_is=\"token\"`) to keep useful samples while clipping variance.\n\n**Combined Estimator (K2-RS-Token-TIS):**\n\nFor combined filtering and weighting, let $K2_{\\text{agg}}$ denote the selected aggregation (token, sum, mean, or max):\n\n$$\n\\hat{g}_{\\text{k2-rs-token-tis}}(y) = \\underbrace{\\mathbb{I}\\left( K2_{\\text{agg}}(y) \\le C_{\\text{k2}} \\right)}_{\\text{K2 Filter}} \\cdot \\prod_t \\min(\\rho_t, C) \\cdot f(y)\n$$\n\nThis is implemented via `rollout_rs=\"seq_mean_k2\"` (or another `k2` mode) together with `rollout_is=\"token\"`.\n\n#### 3.3.5 K3 Divergence Aggregation\n\n**K3 divergence at sequence level:**\n\n$$\nK3_{\\text{seq}} = \\frac{1}{|T|} \\sum_{t \\in T} \\left( \\rho_t - \\log \\rho_t - 1 \\right)\n$$\n\nwhere $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$ is the per-token ratio.\n\n**K3 equals the reverse KL:** In expectation, $K3 = \\text{KL}(\\pi_{\\text{rollout}} \\| \\pi_{\\text{old}})$. This follows from:\n- $\\mathbb{E}_{\\pi_\\text{rollout}}[\\rho] = 1$\n- $\\mathbb{E}_{\\pi_\\text{rollout}}[\\log \\rho] = -\\text{KL}(\\pi_{\\text{rollout}} \\| \\pi_{\\text{old}})$\n- Therefore: $K3 = 1 - (-\\text{KL}) - 1 = \\text{KL}(\\pi_{\\text{rollout}} \\| \\pi_{\\text{old}})$\n\n**Configuration:**\n```python\nrollout_is = null          # No IS weights, pure rejection\nrollout_rs = \"seq_mean_k3\" # K3 rejection sampling\n```\n\n**Properties:**\n- K3 divergence is always >= 0 per token (equals 0 when ρ = 1)\n- More stable than geometric ratio checks because each token term is non-negative\n- Only upper threshold applies (no lower threshold since K3 >= 0)\n- Typical threshold: 0.001 - 0.01\n\n**Why K3 over geometric ratio?**\n- Geometric ratio uses average log-ratio; small numerical bias can flip sign\n- K3 = E[ρ - log ρ - 1] is non-negative per token, offering a smoother detector\n- Both estimate the same quantity: KL(π_rollout || π_old)\n- For small divergences, K3 ≈ 0.5 × Var(log_ratio)\n\n**Combined Estimator (K3-RS-Token-TIS):**\n\nFor best results, combine K3 filter with token-level IS weights:\n\n$$\n\\hat{g}_{\\text{k3-rs-token-tis}}(y) = \\underbrace{\\mathbb{I}\\left( K3_{\\text{seq}} \\le C_{\\text{k3}} \\right)}_{\\text{K3 Filter}} \\cdot \\prod_t \\min(\\rho_t, C) \\cdot f(y)\n$$\n\nThis is implemented by combining `rollout_rs=\"seq_mean_k3\"` with `rollout_is=\"token\"`.\n\n\n---\n\n### 3.4 Batch Normalization\n\nAn optional variance reduction technique that normalizes IS weights to have mean 1.0 within each batch.\n\n**Configuration:**\n```python\nrollout_is_batch_normalize = True  # Default: False\n```\n\n**Normalization formula (aggregation-aware):**\n\nFor **token-level IS** (§3.3.1):\n\n$$\n\\tilde{w}_t = \\frac{w_t}{\\frac{1}{\\sum_{i,t} m_{i,t}} \\sum_{i,t} w_{i,t} \\cdot m_{i,t}}\n$$\n\nwhere $w_{i,t}$ are truncated token IS weights, $m_{i,t}$ is the response mask, and normalization is over **all tokens**.\n\nFor **sequence-level IS** (§3.3.2):\n\n$$\n\\tilde{w}_i = \\frac{w_i}{\\frac{1}{B}\\sum_{j=1}^B \\bar{w}_j}\n$$\n\nwhere $\\bar{w}_j = \\frac{1}{T_j}\\sum_{t=1}^{T_j} w_{j,t} \\cdot m_{j,t}$ is the per-sequence mean (all tokens in a sequence have the same weight), and normalization is over **sequences**.\n\n**Properties:**\n- Applied **after** truncation to preserve truncation semantics\n- Ensures $\\mathbb{E}[\\tilde{w}] = 1$ within each batch\n- **Aggregation-aware**: Token-level normalizes over tokens; sequence-level normalizes over sequences\n- Uses `masked_mean` to respect padding tokens\n- Reduces gradient magnitude variance by removing random batch-level scale fluctuations\n\n**Metrics:**\n- `rollout_is_batch_norm_factor`: The normalization factor applied (batch mean before normalization)\n\n**Implementation:** [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L401-L421)\n\n---\n\n### 3.5 Rejection Sampling (RS)\n\nRejection sampling can be added to **any combination** of operating mode and aggregation level. It modifies the `response_mask` to exclude outlier tokens/sequences.\n\n**Configuration examples:**\n```python\nrollout_rs = \"token_k1\"    # Token-level ratio bounds\nrollout_rs_threshold = \"0.6_1.6\"\n\nrollout_rs = \"seq_sum_k1\"  # Sequence sum of log ratios\nrollout_rs_threshold = \"0.5_2.0\"\n\nrollout_rs = \"seq_mean_k3\" # Sequence mean of K3 divergence\nrollout_rs_threshold = 0.01\n```\n\n**Acceptance set:**\n- **Token-level**: $\\mathcal{A}_{\\text{token}} = \\{ t : C_{\\text{RS-lower}} \\leq \\rho_t \\leq C_{\\text{RS-upper}} \\}$\n- **Sequence-level**: $\\mathcal{A}_{\\text{seq}} = \\{ \\text{seq} : C_{\\text{RS-lower}} \\leq \\prod_{t \\in T} \\rho_t \\leq C_{\\text{RS-upper}} \\}$\n- **Geometric**: $\\mathcal{A}_{\\text{geo}} = \\{ \\text{seq} : C_{\\text{RS-lower}} \\leq \\rho_{\\text{geo}} \\leq C_{\\text{RS-upper}} \\}$\n\n**Properties:**\n- Separate from IS weighting (can use RS without IS)\n- Reduces effective sample size\n- Filters extreme outliers\n\n**Implementation:** `compute_rollout_rejection_mask()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L80-L188)\n\n---\n\n### 3.6 Combination Matrix\n\n**Key insight:** Estimators (how IS/RS is computed) and operating modes (decoupled PPO vs bypass PG) are **orthogonal**. Any estimator can be combined with any operating mode.\n\n#### Estimator × Operating Mode\n\n| Estimator | Configuration | Compatible Modes |\n|-----------|---------------|------------------|\n| **Token-TIS** | `rollout_is=\"token\"` | Decoupled PPO, Bypass PG |\n| **Seq-TIS** | `rollout_is=\"sequence\"` | Decoupled PPO, Bypass PG |\n| **Seq-MIS** | `rollout_is=\"sequence\"` + `rollout_rs=\"seq_sum_k1\"` | Decoupled PPO, Bypass PG |\n| **Geo-RS** | `rollout_rs=\"seq_mean_k1\"` (geometric mean) | Decoupled PPO, Bypass PG |\n| **Geo-RS-Token-TIS** | `rollout_is=\"token\"` + `rollout_rs=\"seq_mean_k1\"` | Decoupled PPO, Bypass PG |\n| **K3-RS** | `rollout_rs=\"seq_mean_k3\"` | Decoupled PPO, Bypass PG |\n| **K3-RS-Token-TIS** | `rollout_is=\"token\"` + `rollout_rs=\"seq_mean_k3\"` | Decoupled PPO, Bypass PG |\n\n**Note:** In bypass mode, `loss_type` controls the loss function. Use \"ppo_clip\" (default) or \"reinforce\".\n\n#### Available Preset Methods\n\n| Preset Method | Estimator | Mode | Properties |\n|---------------|-----------|------|------------|\n| **Decoupled PPO Mode** (3 policies: π_rollout, π_old, π_θ) |\n| `decoupled_token_is()` | Token-TIS | Decoupled PPO | Per-token IS weights |\n| `decoupled_seq_is()` | Seq-TIS | Decoupled PPO | Sequence-level IS weights |\n| `decoupled_seq_is_rs()` | Seq-MIS | Decoupled PPO | Sequence IS + sequence RS |\n| `decoupled_geo_rs()` | Geo-RS | Decoupled PPO | Geometric RS |\n| `decoupled_geo_rs_token_tis()` | Geo-RS-Token-TIS | Decoupled PPO | Geometric filter + token IS |\n| **K3 KL Estimator** (more stable for small KL values) |\n| `decoupled_k3_rs()` | K3-RS | Decoupled PPO | K3 rejection, no IS weights |\n| `decoupled_k3_rs_token_tis()` | K3-RS-Token-TIS | Decoupled PPO | K3 filter + token clipped weight |\n| **Bypass Mode (PPO-clip)** (ratio handles IS, RS masks outliers) |\n| `bypass_ppo_clip()` | - | Bypass (PPO-clip) | PPO-clip only |\n| `bypass_ppo_clip_geo_rs()` | Geo-RS | Bypass (PPO-clip) | PPO-clip + Geo-RS (ratio) |\n| `bypass_ppo_clip_k3_rs()` | K3-RS | Bypass (PPO-clip) | PPO-clip + K3-RS |\n| **Bypass Mode (REINFORCE)** (explicit IS weights, no PPO clipping) |\n| `bypass_pg_is()` | Seq-TIS | Bypass (REINFORCE) | REINFORCE + Seq IS |\n| `bypass_pg_geo_rs()` | Geo-RS | Bypass (REINFORCE) | REINFORCE + Geo-RS (ratio) |\n| `bypass_pg_geo_rs_token_tis()` | Geo-RS-Token-TIS | Bypass (REINFORCE) | REINFORCE + Geo filter + token IS |\n| **Other** |\n| `disabled()` | - | - | Metrics only |\n\n**Note:** Bypass mode sets π_old = π_rollout and uses `loss_type` to select the loss function.\n\n#### Additional Supported Combinations (Manual Configuration)\n\nThese combinations are **fully supported** but require manual configuration:\n\n**1. Token IS + Token RS**\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=\"token\",\n    rollout_is_threshold=2.0,\n    rollout_rs=\"token_k1\",\n    rollout_rs_threshold=\"0.5_2.0\",\n)\n```\n**Properties:** Token-level IS weights + token-level RS mask.\n\n**2. Pure Token RS**\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=None,\n    rollout_rs=\"token_k1\",\n    rollout_rs_threshold=\"0.5_2.0\",\n)\n```\n**Properties:** Token-level RS mask only, no IS weights.\n\n**3. Pure Sequence RS**\n```python\nconfig = RolloutCorrectionConfig(\n    rollout_is=None,\n    rollout_rs=\"seq_sum_k1\",\n    rollout_rs_threshold=\"0.5_2.0\",\n)\n```\n**Properties:** Sequence-level RS mask only, no IS weights.\n\n**Key properties:**\n- Any IS aggregation level (token/sequence) can be used in either decoupled or bypass mode\n- Rejection sampling can be added to any combination\n- Geometric aggregation is typically used for RS only (not IS weighting)\n- Pure RS (`bypass_pg_rs`) uses bypass + geometric RS with `loss_type=\"reinforce\"` for REINFORCE (no IS weights)\n- All combinations in the table above are valid and supported by the implementation\n\n---\n\n### 3.7 Common Implementation Mistake\n\n#### Incorrect LLM-RL Implementation (PPO Without Rollout Correction)\n\n**Theory:** Naive LLM-RL implementation that incorrectly applies PPO by **ignoring the actual rollout policy** and assuming $\\pi_{\\text{old}} = \\pi_{\\text{rollout}}$.\n\n**Note:** This incorrect implementation pattern was identified in [Liu, Li, et al. (2025)](https://richardli.xyz/rl-collapse) as a key cause of training instability in LLM-RL systems, motivating the development of this rollout correction framework.\n\n**Loss Function:**\n\n$$\nL_{\\text{PPO}}(\\theta) = -\\mathbb{E}_t \\left[ \\min\\left( r_t(\\theta) A_t, \\text{clip}(r_t(\\theta), 1-\\epsilon, 1+\\epsilon) A_t \\right) \\right]\n$$\n\nwhere $r_t(\\theta) = \\frac{\\pi_{\\theta}(a_t|s_t)}{\\pi_{\\text{old}}(a_t|s_t)}$ (ignores $\\pi_{\\text{rollout}}$).\n\n**Why it's wrong:**\n- **Ignores $\\pi_{\\text{rollout}}$**: Uses $\\pi_{\\text{old}}$ as behavior policy instead of actual $\\pi_{\\text{rollout}}$\n- **Policy mismatch**: In LLM-RL, rollout typically uses different precision/backend/checkpoint than training, causing $\\pi_{\\text{rollout}} \\neq \\pi_{\\text{old}}$ even with same model weights\n- **Not PPO's fault**: PPO itself is correct; the issue is the incorrect assumption\n\n**Correct alternatives:**\n1. **Decoupled mode**: Three policies with IS correction from $\\pi_{\\text{rollout}}$ to $\\pi_{\\text{old}}$\n2. **Bypass mode**: Two policies using $\\pi_{\\text{rollout}}$ as both behavior policy and proximal policy\n3. **Bypass + Policy Gradient mode**: Two policies with IS/RS correction and no PPO clipping\n\n**Implementation:** `compute_policy_loss()` in [core_algos.py](../../verl/trainer/ppo/core_algos.py#L812-L884)\n\n---\n\n## 4. Off-Policy Diagnostic Metrics\n\nThese metrics quantify the severity of off-policy drift.\n\n**Note on notation:** Metrics use $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$. In bypass mode, $\\pi_{\\text{old}} = \\pi_{\\text{rollout}}$, so metrics measure rollout→current drift using $\\rho_t = \\frac{\\pi_{\\theta}}{\\pi_{\\text{rollout}}}$ instead.\n\n### 4.1 KL Divergence\n\n**Direct KL estimator:**\n\n$$\n\\text{KL}(\\pi_{\\text{rollout}} \\| \\pi_{\\text{old}}) = \\mathbb{E}_{t \\sim \\pi_{\\text{rollout}}} \\left[ \\log \\pi_{\\text{rollout}}(a_t|s_t) - \\log \\pi_{\\text{old}}(a_t|s_t) \\right]\n$$\n\n**K3 KL estimator** (alternative formulation):\n\n$$\n\\text{KL}_{\\text{K3}} = \\mathbb{E}_{t \\sim \\pi_{\\text{rollout}}} \\left[ \\rho_t - \\log \\rho_t - 1 \\right]\n$$\n\nwhere $\\rho_t = \\frac{\\pi_{\\text{old}}(a_t|s_t)}{\\pi_{\\text{rollout}}(a_t|s_t)}$.\n\n### 4.2 Perplexity\n\n**Old policy perplexity:**\n\n$$\n\\text{PPL}_{\\text{old}} = \\exp\\left( -\\frac{1}{|T|} \\sum_{t \\in T} \\log \\pi_{\\text{old}}(a_t|s_t) \\right)\n$$\n\n**Rollout policy perplexity:**\n\n$$\n\\text{PPL}_{\\text{rollout}} = \\exp\\left( -\\frac{1}{|T|} \\sum_{t \\in T} \\log \\pi_{\\text{rollout}}(a_t|s_t) \\right)\n$$\n\n**PPL ratio** (inverse of geometric mean IS weight):\n\n$$\n\\text{PPL}_{\\text{ratio}} = \\frac{\\text{PPL}_{\\text{old}}}{\\text{PPL}_{\\text{rollout}}} = \\exp\\left( -\\frac{1}{|T|} \\sum_{t \\in T} \\log \\rho_t \\right) = \\left(\\prod_{t \\in T} \\rho_t\\right)^{-1/|T|}\n$$\n\n**Interpretation:** Values > 1 mean $\\pi_{\\text{old}}$ assigns lower probability than $\\pi_{\\text{rollout}}$ to the observed actions (distribution shift).\n\n### 4.3 Chi-squared Divergence\n\nMeasures the second moment of the IS weight distribution.\n\n**Token-level:**\n\n$$\n\\chi^2_{\\text{token}} = \\mathbb{E}_{t \\sim \\pi_{\\text{rollout}}} \\left[ \\rho_t^2 \\right] - 1\n$$\n\n**Sequence-level:**\n\n$$\n\\chi^2_{\\text{seq}} = \\mathbb{E}_{\\text{seq} \\sim \\pi_{\\text{rollout}}} \\left[ \\left(\\prod_{t \\in T} \\rho_t\\right)^2 \\right] - 1\n$$\n\n**Interpretation:**\n- $\\chi^2 = 0$: Policies are identical\n- $\\chi^2 > 0$: Higher values indicate more severe off-policy distribution shift\n\n**Implementation:** `compute_offpolicy_metrics()` in [rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py#L670-L776)\n\n---\n\n## 5. Summary and Decision Guide\n\n### 5.1 Method Summary Table\n\n| Method | Theory | Policies | PPO Clip | IS Correction | Correctness | Speed |\n|--------|--------|----------|----------|---------------|-------------|-------|\n| **Bypass Mode** (π_old = π_rollout, `loss_type` selects algorithm) |\n| `loss_type=\"ppo_clip\"` (default) | PPO (ratio = π_θ/π_rollout) | 2 (rollout, θ) | ✅ | RS mask only (ratio handles IS) | ✅ Correct | **Fast** |\n| `loss_type=\"reinforce\"` | Off-policy REINFORCE | 2 (rollout, θ) | ❌ | ✅ (explicit IS weights) | ✅ Correct | **Fast** |\n| **Bypass Mode Presets (PPO-clip)** |\n| `bypass_ppo_clip` | PPO only | 2 (rollout, θ) | ✅ | - | ✅ Correct | **Fast** |\n| `bypass_ppo_clip_geo_rs` | PPO + Geo-RS | 2 (rollout, θ) | ✅ | Geo-RS mask (ratio) | ✅ Correct | **Fast** |\n| **Bypass Mode Presets (REINFORCE)** |\n| `bypass_pg_is` | REINFORCE + Seq-TIS | 2 (rollout, θ) | ❌ | ✅ Seq-TIS | ✅ Correct | **Fast** |\n| `bypass_pg_geo_rs` | REINFORCE + Geo-RS | 2 (rollout, θ) | ❌ | Geo-RS only (ratio) | ✅ Correct | **Fast** |\n| `bypass_pg_geo_rs_token_tis` | REINFORCE + Geo RS + Token IS | 2 (rollout, θ) | ❌ | ✅ Geo-RS-Token-TIS | ✅ Correct | **Fast** |\n| **Decoupled PPO Mode** (IS weights = π_old / π_rollout) |\n| `decoupled_token_is` | Decoupled PPO | 3 (rollout, old, θ) | ✅ | ✅ Token-TIS | ✅ Correct | Standard |\n| `decoupled_seq_is` | Decoupled PPO | 3 (rollout, old, θ) | ✅ | ✅ Seq-TIS | ✅ Correct | Standard |\n| `decoupled_seq_is_rs` | Decoupled PPO + RS | 3 (rollout, old, θ) | ✅ | ✅ Seq-MIS | ✅ Correct | Standard |\n| `decoupled_geo_rs` | Decoupled PPO + Geo-RS | 3 (rollout, old, θ) | ✅ | Geo-RS only (ratio) | ✅ Correct | Standard |\n| `decoupled_geo_rs_token_tis` | Decoupled PPO + Geo RS + Token IS | 3 (rollout, old, θ) | ✅ | ✅ Geo-RS-Token-TIS | ✅ Correct | Standard |\n| **Incorrect (for reference)** |\n| Naive LLM-RL | Incorrect PPO usage | 2 (old, θ) | ✅ | ❌ | ⚠️ Incorrect | Standard |\n\n**Notes:**\n- **Bypass mode** sets π_old = π_rollout and uses `loss_type` to select the loss function:\n  - `\"ppo_clip\"` (default): PPO clipped ratio (IS handled by ratio = π_θ/π_rollout, no explicit IS weights to avoid double-counting)\n  - `\"reinforce\"`: Explicit IS weights applied as $w \\cdot \\log \\pi \\cdot A$\n- Both loss types benefit from rejection sampling (RS) which masks out-of-distribution samples\n\n### 5.2 Estimator Hierarchy\n\nThese estimators define **how IS weights and rejection masks are computed**. They are orthogonal to the operating mode (decoupled PPO vs bypass policy gradient) and can be combined with either.\n\n| Estimator | Configuration | Mechanism | Best For |\n|-----------|---------------|-----------|----------|\n| **Token-TIS** | `rollout_is=\"token\"` | Clips per-token ratios | Lower variance IS with acceptable bias |\n| **Seq-TIS** | `rollout_is=\"sequence\"` | Clips sequence ratio $\\rho(\\tau) \\to \\min(\\rho(\\tau), C)$ | Clean data with moderate mismatch; unbiased |\n| **Seq-MIS** | `rollout_is=\"sequence\"` + `rollout_rs=\"seq_sum_k1\"` | Rejects sequences with $\\rho(\\tau) > C$ | Severe mismatch; filters \"toxic tail\" (garbage data) |\n| **Geo-RS** | `rollout_rs=\"seq_mean_k1\"` | Rejects on geometric mean ratio exp(E[log(r)]) | Length-invariant trust region |\n| **Geo-RS-Token-TIS** | `rollout_is=\"token\"` + `rollout_rs=\"seq_mean_k1\"` | Geometric filter + token IS weights | Ratio-based length normalization + lower variance IS |\n| **K3-RS** | `rollout_rs=\"seq_mean_k3\"` | Rejects on K3 KL divergence | Small KL values; smooth detector |\n| **K3-RS-Token-TIS** | `rollout_is=\"token\"` + `rollout_rs=\"seq_mean_k3\"` | K3 filter + token IS weights | Small KL + lower variance IS |\n\n**Note:** Each estimator can be used with either:\n- **Decoupled PPO** (`bypass_mode=false`): Three policies with PPO clipping\n- **Bypass Mode** (`bypass_mode=true`): Two policies with configurable loss type\n  - `loss_type=\"ppo_clip\"` (default): PPO clipped objective (IS via ratio, RS mask applied)\n  - `loss_type=\"reinforce\"`: REINFORCE with explicit IS weights\n\n### 5.3 Method Characteristics by Scenario\n\n**Choosing estimator by off-policy severity:**\n- **Negligible** (same checkpoint, minor differences): No IS correction needed; use bypass mode for efficiency\n- **Moderate** (async workers, slight staleness): Token-TIS provides per-token IS correction with lower variance\n- **Severe** (replay buffers, old data): Seq-TIS or Seq-MIS provides sequence-level IS correction; use Seq-MIS when high-weight samples are likely garbage\n\n**Choosing estimator by sequence length:**\n- **Short sequences** (standard chat): Seq-TIS is optimal\n- **Long sequences** (CoT, agents): K1-RS or K1-RS-Token-TIS to avoid Length Trap\n\n**Choosing operating mode:**\n- **Batch size invariance needed**: Use decoupled mode (`bypass_mode=false`)\n- **Computational efficiency needed**: Use bypass mode (`bypass_mode=true`) to skip `old_log_prob` computation\n- **No PPO clipping**: Use bypass mode with `loss_type=\"reinforce\"`\n\n### 5.4 Decoupled Mode vs Bypass Mode\n\n**Decoupled mode** (computes `old_log_prob` separately):\n- Implements full decoupled PPO with three policies (mathematically correct)\n- Separately measures and corrects Drift 1 (rollout→old) and Drift 2 (old→current)\n- Achieves batch size invariance and efficient stale data utilization\n- Enables accurate off-policy metrics monitoring\n\n**Bypass mode** (sets $\\pi_{\\text{old}} = \\pi_{\\text{rollout}}$):\n- Uses $\\pi_{\\text{rollout}}$ as both behavior policy and proximal policy (mathematically correct)\n- Computational efficiency: Skips separate `old_log_prob` computation\n- Does not achieve batch size invariance (proximal policy depends on data collection)\n\n---\n\n## 6. Implementation References\n\n- **[Rollout Correction Usage Guide](rollout_corr.md)** - Practical configuration and troubleshooting\n- **Config:** [verl/trainer/config/algorithm.py](../../verl/trainer/config/algorithm.py)\n- **IS/RS Helper:** [verl/trainer/ppo/rollout_corr_helper.py](../../verl/trainer/ppo/rollout_corr_helper.py)\n- **PPO Loss:** [verl/trainer/ppo/core_algos.py](../../verl/trainer/ppo/core_algos.py)\n- **Tests:** [tests/trainer/ppo/test_rollout_corr.py](../../tests/trainer/ppo/test_rollout_corr.py)\n\n---\n\n## References\n\n- **Williams, R. J. (1992).** \"Simple statistical gradient-following algorithms for connectionist reinforcement learning.\" *Machine Learning*, 8(3-4), 229-256. https://doi.org/10.1007/BF00992696\n- **Schulman, J., Wolski, F., Dhariwal, P., Radford, A., & Klimov, O. (2017).** \"Proximal policy optimization algorithms.\" *arXiv preprint arXiv:1707.06347.* https://arxiv.org/abs/1707.06347\n- **Hilton, J., Cobbe, K., & Schulman, J. (2021).** \"Batch size-invariance for policy optimization.\" *arXiv preprint arXiv:2110.00641.* https://arxiv.org/abs/2110.00641\n  - Introduced decoupled PPO: separating proximal policy (for controlling policy update size) from behavior policy (for off-policy correction) to achieve batch size invariance\n"
  },
  {
    "path": "docs/algo/spin.md",
    "content": "# Recipe: Self-Play Fine-Tuning (SPIN)\n\nLast updated: 05/31/2025.\n\n`verl` provides a recipe inspired by the paper **\"Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models\"** (SPIN). SPIN is a language model finetuning algorithm that enables iterative self-improvement through a self-play mechanism inspired by game theory.\n\n**Core Idea:** Models learn by playing against themselves, reducing reliance on external preference datasets or stronger teacher models:\n\n1.  **Synthetic Data Generation:** The current model generates responses, creating its own training data from previous iterations.\n2.  **Two-Player Game Setup:** A game involving two players acted by a single LLM.\n3.  **Iterative Training:** The model progressively improves by refining its policy, with each iteration's model becoming the opponent for the next iteration.\n\nPaper Authors: [Zixiang Chen](https://github.com/uclaml/SPIN)\\*, [Yihe Deng](https://github.com/uclaml/SPIN)\\*, [Huizhuo Yuan](https://scholar.google.com/citations?user=8foZzX4AAAAJ)\\*, [Kaixuan Ji](https://scholar.google.com/citations?user=FOoKDukAAAAJ), [Quanquan Gu](https://web.cs.ucla.edu/~qgu/)\n\n[[Webpage](https://uclaml.github.io/SPIN/)] [[Huggingface](https://huggingface.co/papers/2401.01335)] [[Paper](https://arxiv.org/abs/2401.01335)] [[Original Implementation](https://github.com/uclaml/SPIN)]\n\nverl Implementation Authors: [Chendong Wang](https://cdwang96.github.io/), [Chenyang Zhao](https://github.com/zhaochenyang20)\n\n---\n\n## Key Function (compute_online_dpo_loss) and Related works\nSPIN (Chen et al., 2024) proposes an iterative self-play mechanism to fine-tune language models. In each iteration, SPIN's training objective, when using a logistic loss function, is equivalent to Direct Preference Optimization (DPO) loss (Rafailov et al., 2023). \n\nThis `verl` recipe realizes SPIN's core concept by using DPO loss iteratively (Xu et al., 2023; Xiong et al., 2023; Snorkel AI, 2024). This means that in each iteration, we fine-tune the LLM using DPO loss for preference optimization. Notably, Xu et al. (2023) explored iterative preference optimization with pairwise cringe loss, while Xiong et al. (2023) discussed how to bridge theory and practice for RLHF under KL constraints using iterative training. The concept of iterative preference learning was also explored in online DPO (Guo et al., 2024), which focuses on direct alignment from online AI feedback. In online DPO, preference data is dynamically updated during training, allowing the model to learn from its own generated data.\n\nSpecifically, we developed the **`compute_online_dpo_loss`** function and built this SPIN recipe on top of it. By incorporating online preference generation, this approach enables continuously refining language models without relying on fixed external preference datasets.\n\n**Reference Papers:**\n* [Self-Play Fine-Tuning Converts Weak Language Models to Strong Language Models](https://arxiv.org/abs/2401.01335) (Chen et al., 2024) \n* [Direct Preference Optimization: Your Language Model is Secretly a Reward Model](https://arxiv.org/abs/2305.18290) (Rafailov et al., 2023) \n* [Somethings are more cringe than others: Preference optimization with the pairwise cringe loss](https://arxiv.org/abs/2312.16682) (Xu et al., 2023) \n* [Iterative preference learning from human feedback: Bridging theory and practice for rlhf under kl-constraint](https://arxiv.org/abs/2312.11456) (Xiong et al., 2023)\n* [Snorkel-Mistral-PairRM-DPO](https://huggingface.co/snorkelai/Snorkel-Mistral-PairRM-DPO) (Snorkel AI, 2024)\n* [Direct language model alignment from online ai feedback](https://arxiv.org/abs/2402.04792) (Guo et al., 2024)\n\n\n## Our Online DPO Implementation\n\nOur `compute_online_dpo_loss` function adapts `verl`'s existing PPO infrastructure (based on `verl` v0.3.0.post1) for this iterative online DPO. Key aspects of our implementation include:\n\n* **No Critic:** Unlike PPO, we omit the value function critic.\n* **Dynamic Reference Model:** An explicit reference policy (`ref_policy_wg`) is used for DPO loss. This reference model's weights can be periodically updated from the actor (`ref_update_freq`), providing a dynamic baseline.\n* **Online Preference Generation:** The `compute_onlineDPO_pref` function (in `core_algos.py`) dynamically creates chosen/rejected pairs based on a reward source (e.g., rule-based ranking for math problems).\n* **DPO Loss Integration:** We replace PPO's policy loss with our `compute_online_dpo_loss` (in `core_algos.py`) within the actor update (`dp_actor.py`), directly optimizing the policy using the generated preferences.\n* **Iterative Training Orchestration:** The `SpinTrainer` (in `spin_trainer.py`) manages the entire self-play loop: generation, preference labeling, optional reference model updates, and policy updates, enabling continuous self-improvement aligned with SPIN's principles.\n\n---\n## Algorithm\n\nThis recipe implements an Online algorithm adapted to the `verl` Reinforcement Learning framework, which provides an alternative to PPO for fine-tuning language models.\n\n**Online Loop:** Instead of maximizing a scalar reward signal in PPO, this approach directly optimizes the policy model to align with preference data generated *online* during training:\n\n1.  **Generation:** The current model generates multiple responses for each prompt in a batch.\n2.  **Preference Labeling:** A function evaluates these generated responses to determine which one is preferred (chosen) and which is dispreferred (rejected). This can be done using a reward function or implicit ranking based on specific rules. (In this recipe, we use rule-based ranking on the math problem).\n3.  **Update:** This preference tuple (`prompt`, `chosen_response`, `rejected_response`) is used to update the actor model using `compute_online_dpo_loss`, comparing against a reference model.\n\n**Connection with SPIN:**\nInstead of only using a fixed target data distribution, the online generation loop in step 2 will dynamically change the target data distribution by using a certain Preference Labeling method (rule-based ranking on the math problem by selecting the better one in this recipe). This explores the direction mentioned in SPIN's paper Section 7 about \"dynamically changing target data distribution\" to potentially elevate LLM performance beyond the fixed human-annotated data ceiling.\n\n---\n\n## Reproduce the Experiment (Example Setup)\n\nThe following steps outline how to set up the environment and run the SPIN recipe, based on the provided test log using GSM8K and Qwen2.5-3B-Instruct.\n\n1.  **Setup Environment (Example using Docker):**\n    ```bash\n    # Start a container with GPU access and shared memory\n    docker run -it --name spin_test --gpus all \\\n        --shm-size=32g \\\n        --ipc=host \\\n        -v /path/to/host/.cache:/root/.cache \\\n        -e HF_TOKEN=<YOUR_HUGGINGFACE_TOKEN> \\\n        lmsysorg/sglang:latest \\\n        /bin/bash\n\n    # Inside the container or on your host machine:\n    # Ensure /tmp is writable\n    mkdir -p /tmp\n    chmod 1777 /tmp\n\n    # Install Python 3.10 (if not present) and venv\n    sudo apt update\n    sudo apt install -y python3.10 python3.10-venv tmux\n    python3 -m ensurepip --upgrade\n\n    # Create and activate a virtual environment\n    python3 -m venv ~/.python/spin_env\n    source ~/.python/spin_env/bin/activate\n\n    # Install uv (fast package installer)\n    python3 -m pip install uv\n    ```\n\n2.  **Install verl and Dependencies:**\n    ```bash\n    # Clone the verl repository and checkout the spin branch\n    cd ~\n    git clone git@github.com:volcengine/verl.git && cd verl\n\n    # Install flash-attn (handle potential build issues)\n    python3 -m uv pip install wheel packaging\n    python3 -m uv pip install flash-attn --no-build-isolation --no-deps\n\n    # Install verl with sglang extras\n    python3 -m uv pip install -e \".[sglang]\"\n    ```\n    *Note: If `flash-attn` installation fails, try the manual steps again or consult its documentation.*\n\n3.  **Login & Download Data/Model:**\n    ```bash\n    # Login to Weights & Biases (optional, for logging)\n    export WANDB_API_KEY=<YOUR_WANDB_API_KEY>\n    # wandb login\n\n    # Download the GSM8K dataset\n    python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k # Adjusted path\n\n    # Download the base model (Example: Qwen2.5-3B-Instruct)\n    hf download Qwen/Qwen2.5-3B-Instruct --local-dir $HOME/models/Qwen2.5-3B-Instruct\n    ```\n\n4.  **Configure:**\n    * Modify the configuration file (e.g., `config/spin_trainer.yaml` or the one specified in the run script) with correct paths to your downloaded model, data, desired hyperparameters (`dpo_beta`, learning rate, etc.), and distributed training settings (nodes, GPUs per node).\n    * Pay attention to `actor_rollout_ref.model`, `data` paths, `reward_model` config (if using one), and `trainer.ref_update_freq`.\n\n5.  **Run Training:**\n    ```bash\n    # Set CUDA visible devices (adjust based on your hardware and config)\n    export CUDA_VISIBLE_DEVICES=0,1,2,3\n\n    # Launch the training script (e.g., test.sh or a custom script)\n    # Ensure test.sh points to the correct config and main script\n    bash recipe/spin/run_spin.sh\n    ```\n\n---\n\n## Configuration\n\n* The primary configuration is typically managed through a YAML file specified in the launch script (e.g., `config/spin_trainer.yaml`).\n* Key configuration sections:\n    * `data`: Paths to training/validation prompt files, batch sizes, sequence lengths.\n    * `actor_rollout_ref`: Paths to the base model (used for actor and initial reference), FSDP settings, optimization parameters (learning rate, scheduler).\n    * `reward_model`: Configuration for the reward model used for online preference labeling (path, batch size, etc.). Can be omitted if using a simpler reward function.\n    * `algorithm`: DPO-specific hyperparameters like `dpo_beta`, `dpo_loss_type`.\n    * `trainer`: Distributed training settings (nodes, GPUs per node), logging (WandB), checkpointing frequency, and `ref_update_freq` (set > 0 to enable periodic reference model updates from the actor).\n\n---\n\n## Key Files\n\n* `main_spin.py`: Main entry point using Hydra to load the config and launch the `SpinTrainer`.\n* `spin_trainer.py`: Defines the `SpinTrainer` class, orchestrating the Online DPO training loop.\n* `fsdp_workers.py`: Implements Ray workers (Actor, Reference) potentially using FSDP.\n* `dp_actor.py`: Contains the actor class, including the DPO policy update logic.\n* `core_algos.py`: Includes helper functions for `compute_online_dpo_loss` and `compute_onlineDPO_pref`.\n* `config/spin_trainer.yaml` (or similar): Main Hydra configuration file for the recipe.\n* `run_spin.sh` (or similar): Example bash script for launching a training run.\n* `README.md`: This file.\n\n---\n\n## Acknowledgement\n\nWe sincerely thank the contribution and guidance from the `verl` community and advisors, including (adapted from SPPO):\n\n* [Zixiang Chen](https://sites.google.com/view/zxchen)\n* [Yuhao Yang](https://github.com/yhyang201)\n* [Yifan Zhang](https://github.com/yifanzhang-pro)\n* [Yongan Xiang](https://github.com/BearBiscuit05)\n* [Junrong Lin](https://github.com/ocss884)\n* [Yuxuan Tong](https://github.com/tongyx361)\n* [Guangming Shen](https://github.com/PeterSH6)\n* [Biao He](https://www.linkedin.com/in/biao-he/)\n* [Qingquan Song](https://qingquansong.github.io/)\n* [Chenyang Zhao](https://zhaochenyang20.github.io/Chayenne/)\n* [Quanquan Gu](https://web.cs.ucla.edu/~qgu/)\n"
  },
  {
    "path": "docs/algo/sppo.md",
    "content": "# Recipe: Self-Play Preference Optimization (SPPO)\n\nLast updated: 05/28/2025.\n\nverl provides a community recipe implementation for the paper [Self-Play Preference Optimization for Language Model Alignment](https://arxiv.org/abs/2405.00675). SPPO can significantly enhance the performance of an LLM without strong external signals such as responses or preferences from GPT-4. It can outperform the model trained with iterative direct preference optimization (DPO), among other methods. SPPO is theoretically grounded, ensuring that the LLM can converge to the von Neumann winner (i.e., Nash equilibrium) under general, potentially intransitive preference, and empirically validated through extensive evaluations on multiple datasets.\n\nPaper Authors: [Yue Wu](https://yuewu.us/)\\*, [Zhiqing Sun](https://www.cs.cmu.edu/~zhiqings/)\\*, [Huizhuo Yuan](https://scholar.google.com/citations?user=8foZzX4AAAAJ)\\*, [Kaixuan Ji](https://scholar.google.com/citations?user=FOoKDukAAAAJ), [Yiming Yang](https://www.cs.cmu.edu/~yiming/), [Quanquan Gu](https://web.cs.ucla.edu/~qgu/)\n\nverl Implementation Authors: [Yuhao Yang](https://github.com/yhyang201), [Chenyang Zhao](https://github.com/zhaochenyang20)\n\n[[Webpage](https://uclaml.github.io/SPPO/)] [[Huggingface](https://huggingface.co/papers/2405.00675)] [[Paper](https://arxiv.org/abs/2405.00675)][[Original Implementation](https://github.com/uclaml/SPPO)]\n\n## Reproduce the Experiment\n\nWe evaluate the performance of SPPO on the MATH dataset. Starting from an initial score of 46.6 with Qwen2.5-7B-Instruct, we achieve a score of 65.6 after 20 epochs of training, placing our model approximately in the top 20 on the [MATH leaderboard](https://paperswithcode.com/sota/math-word-problem-solving-on-math). It's important to note that verl's internal evaluation metrics may not perfectly align with the official evaluation methodology for Qwen2.5-7B-Instruct. Therefore, for consistency and fair comparison, we report only the results based on verl's evaluation framework.\n\n```\ngit clone git@github.com:volcengine/verl.git\ncd verl\npython3 -m uv pip install -e \".[sglang]\"\n\nexport WANDB_API_KEY=<YOUR_WANDB_API_KEY>\n\npython3 examples/data_preprocess/math_dataset.py --local_dir ~/data/math\nhf download Qwen/Qwen2.5-7B-Instruct --local-dir $HOME/models/Qwen2.5-7B-Instruct\n\nexport CUDA_VISIBLE_DEVICES=0,1,2,3\nbash recipe/sppo/run_qwen2.5-7b_rm.sh\n```\n\nNote that the installation would occasionally fail to install flash-attn. If this happens, you can install it manually by running:\n\n```bash\npython3 -m uv pip install wheel\npython3 -m uv pip install packaging\npython3 -m uv pip install flash-attn --no-build-isolation --no-deps\n```\n\n## Acknowledgement\n\nWe sincerely thank the contribution and guidance from:\n\n- [Yue Wu](https://yuewu.us/)\n- [Chendong Wang](https://cdwang96.github.io/)\n- [Yifan Zhang](https://github.com/yifanzhang-pro)\n- [Yongan Xiang](https://github.com/BearBiscuit05)\n- [Junrong Lin](https://github.com/ocss884)\n- [Yuxuan Tong](https://github.com/tongyx361)\n- [Guangming Shen](https://github.com/PeterSH6)\n- [Biao He](https://www.linkedin.com/in/biao-he/)\n- [Qingquan Song](https://qingquansong.github.io/)\n- [Quanquan Gu](https://web.cs.ucla.edu/~qgu/)\n"
  },
  {
    "path": "docs/amd_tutorial/amd_build_dockerfile_page.rst",
    "content": "Getting started with AMD (ROCM Kernel)\n=====================================================\n\nLast updated: 07/06/2025.\n\nAuthor: `Yusheng Su <https://yushengsu-thu.github.io/>`_\n\nSetup\n-----\n\nIf you run on AMD GPUs (MI300) with ROCM platform, you cannot use the previous quickstart to run verl. You should follow the following steps to build a docker and set ``RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES`` or ``RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES`` when starting ray in verl's RLHF training.\n\n\ndocker/Dockerfile.rocm\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    FROM \"rlsys/rocm-6.3.4-patch:rocm6.3.4-numa-patch_ubuntu-22.04\"\n\n    SHELL [\"/bin/bash\", \"-ceuxo\", \"pipefail\"]\n\n    ENV MAX_JOBS=512\n\n    ENV PATH=\"/usr/local/python3.12/bin:$PATH\"\n    RUN ln -sf /usr/bin/python3.12 /usr/bin/python && \\\n        ln -sf /usr/bin/pip3.12 /usr/bin/pip\n\n    ############################################\n    RUN apt-get update\n    RUN apt-get install -y pkg-config liblzma-dev\n    ############################################\n\n    ###########################################\n    ##########Install TransformerEngine########\n    ###########################################\n    WORKDIR /workspace/\n    # transformer-engine install\n    # https://github.com/ROCm/TransformerEngine\n    RUN rm -rf TransformerEngine \n    RUN git clone --recursive https://github.com/ROCm/TransformerEngine.git\n    WORKDIR /workspace/TransformerEngine\n    git checkout 236178e5\n    # git checkout bb061ade\n    # git checkout 864405c\n    ENV NVTE_FRAMEWORK=pytorch \n    ENV NVTE_ROCM_ARCH=gfx942 \n    ENV NVTE_USE_HIPBLASLT=1\n    ENV NVTE_USE_ROCM=1  \n    # export CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr:${CMAKE_PREFIX_PATH:-}\"\n    ENV CMAKE_PREFIX_PATH=\"/opt/rocm:/opt/rocm/hip:/usr/local:/usr\"\n    RUN MAX_JOBS=$(MAX_JOBS) pip install . -vvv \n    WORKDIR /workspace/\n    ###########################################\n    ###########################################\n    ###########################################\n\n\n\n\n\n    ####################################################################################\n    ################Install vllm - sglang require vllm 0.6.7 dependency#################\n    ####################################################################################\n    #### Require vllm 0.6.7 - checkout 113274a0\n    WORKDIR /workspace/\n    RUN rm -rf vllm\n    RUN pip uninstall -y vllm\n    # Refer to here (down-grade vllm to 0.6.3): https://docs.vllm.ai/en/v0.6.3/getting_started/amd-installation.html\n    RUN git clone https://github.com/ROCm/vllm.git\n    # git clone https://github.com/vllm-project/vllm.git\n    WORKDIR /workspace/vllm\n    RUN git checkout 113274a0\n    ENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n    #ENV MAX_JOBS=512\n    ENV MAX_JOBS=${MAX_JOBS}\n    RUN pip install \"boto3>=1.26.0\"\n    RUN pip install setuptools_scm\n    # will add src into py. You can delete the repo\n    RUN python3 setup.py install\n    WORKDIR /workspace/\n    ####################################################################################\n    ####################################################################################\n    ####################################################################################\n\n\n\n    ###########################################\n    ############For hack docker################\n    ###########################################\n    RUN pip install setuptools==75.8.0\n    ###########################################\n    ###########################################\n    ###########################################\n\n\n\n    ###########################################\n    ############build sgalng###################\n    ###########################################\n    # Set environment variables\n    ENV BASE_DIR=/sgl-workspace\n    ENV BUILD_TYPE=all\n    ENV SGL_REPO=https://github.com/sgl-project/sglang\n    ENV SGL_BRANCH=v0.4.6.post5\n    ENV TRITON_REPO=https://github.com/ROCm/triton.git\n    ENV TRITON_COMMIT=improve_fa_decode_3.0.0\n    ENV AITER_REPO=https://github.com/ROCm/aiter.git\n    ENV AITER_COMMIT=v0.1.2\n    # v0.1.2 version - commit id: 9d11f47\n    # ENV AITER_COMMIT=9d11f47\n    ENV HIP_FORCE_DEV_KERNARG=1\n    ENV HSA_NO_SCRATCH_RECLAIM=1\n    ENV SGLANG_SET_CPU_AFFINITY=1\n    ENV SGLANG_ALLOW_OVERWRITE_LONGER_CONTEXT_LEN=1\n    ENV NCCL_MIN_NCHANNELS=112\n    ENV MOE_PADDING=1\n    ENV VLLM_FP8_PADDING=1\n    ENV VLLM_FP8_ACT_PADDING=1\n    ENV VLLM_FP8_WEIGHT_PADDING=1\n    ENV VLLM_FP8_REDUCE_CONV=1\n    ENV TORCHINDUCTOR_MAX_AUTOTUNE=1\n    ENV TORCHINDUCTOR_MAX_AUTOTUNE_POINTWISE=1\n    ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\n    ENV AMDGPU_TARGETS=gfx942\n    ENV ROCM_ARCH=gfx942\n    ENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n    # Switch to working directory\n    WORKDIR /sgl-workspace\n    # Clean and create directory\n    RUN rm -rf /sgl-workspace && mkdir -p /sgl-workspace\n\n    # Clone and build sglang\n    RUN git clone ${SGL_REPO} \\\n        && cd sglang \\\n        && git checkout ${SGL_BRANCH} || echo \"Using default branch\" \\\n        && cd sgl-kernel \\\n        && rm -f pyproject.toml \\\n        && mv pyproject_rocm.toml pyproject.toml \\\n        && python setup_rocm.py install \\\n        && cd .. \\\n        && if [ \"$BUILD_TYPE\" = \"srt\" ]; then \\\n            python -m pip --no-cache-dir install -e \"python[srt_hip]\"; \\\n        else \\\n            python -m pip --no-cache-dir install -e \"python[all_hip]\"; \\\n        fi \\\n        && cd /sgl-workspace \\\n        && cp -r /sgl-workspace/sglang /sglang \\\n        && python -m pip cache purge\n\n    # Install common Python packages\n    RUN pip install IPython orjson python-multipart torchao pybind11\n    # Rebuild Triton\n    RUN pip uninstall -y triton || true \\\n        && git clone ${TRITON_REPO} \\\n        && cd triton \\\n        && git checkout ${TRITON_COMMIT} \\\n        && cd python \\\n        && python3 setup.py install \\\n        && cd /sgl-workspace\n    # ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942 --amdgpu-lower-module-lds-strategy=1\"\n    # ENV HIPCC_COMPILE_FLAGS_APPEND=\"--offload-arch=gfx942\"\n\n    # Build aiter\n    #version: Commit 9d11f47\n        # && git checkout ${AITER_COMMIT} \\\n    RUN pip uninstall -y aiter || true\n    RUN git clone ${AITER_REPO} \\\n        && cd aiter \\\n        && git checkout ${AITER_COMMIT} \\\n        && git submodule sync \\\n        && git submodule update --init --recursive \\\n        && PREBUILD_KERNELS=1 GPU_ARCHS=gfx942 python3 setup.py install \\\n        && cd /sgl-workspace\n\n    # Copy MI300X config \n    RUN find /sgl-workspace/sglang/python/sglang/srt/layers/quantization/configs/ \\\n            /sgl-workspace/sglang/python/sglang/srt/layers/moe/fused_moe_triton/configs/ \\\n            -type f -name '*MI300X*' | \\\n            xargs -I {} sh -c 'vf_config=$(echo \"$1\" | sed \"s/MI300X/MI300X_VF/\"); cp \"$1\" \"$vf_config\"' -- {}\n\n    # Environment setup complete.\n    RUN echo \"Environment setup complete.\"\n\n    WORKDIR /workspace/\n    ###########################################\n    ###########################################\n    ###########################################\n\n\n\n\n\n\n    ###########################################\n    ###############vllm v0.8.5#################\n    ###########################################\n    WORKDIR /workspace/\n\n    ENV VLLM_TARGET_DEVICE=rocm \n    ENV ROCM_PATH=/opt/rocm \n    ENV SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev\n    # Find the repo path in: DockerFile/Dockerfile.rocm_yang\n    # RUN git clone https://github.com/RLFoundation/vllm-patch.git\n    RUN pip uninstall -y vllm || true\n    RUN rm -rf vllm-patch\n    RUN git clone https://github.com/RLFoundation/vllm-patch.git \\\n        && cd vllm-patch \\\n        && git checkout v0.8.5-sleep-numa \\\n        && rm -rf build/ dist/ *.egg-info \\\n        && ln -sf /opt/rocm/lib/libamdhip64.so /usr/lib/libamdhip64.so \\\n        && SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py install\n        # RUN SETUPTOOLS_SCM_PRETEND_VERSION=0.8.5.dev PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\" MAX_JOBS=${MAX_JOBS} python3 setup.py develop\n    WORKDIR /workspace/\n    ###########################################\n    ###########################################\n    ###########################################\n\n\n\n\n    #########################################\n    #### Install megatron-core###############\n    #########################################\n    RUN pip uninstall -y megatron-core && \\\n        git clone https://github.com/yushengsu-thu/Megatron-LM-amd_version.git && \\\n        cd Megatron-LM-amd_version && \\\n        pip install -vvv -e . && \\\n        cd /workspace/\n    #########################################\n    #########################################\n    #########################################\n\n\n\n\n    #######################################\n    ################apex###################\n    #######################################\n    WORKDIR /workspace/\n    RUN pip uninstall -y apex && \\\n        git clone git@github.com:ROCm/apex.git && \\\n        cd apex && \\\n        python setup.py install && \\\n        cd /workspace/ \n    #######################################\n    #######################################\n    #######################################\n\n\n    ################################################################################\n    ###########################Add torch_memory_saver###############################\n    ################################################################################\n    # Set environment variables\n    ENV HIPCC_COMPILE_FLAGS_APPEND=\"--amdgpu-target=gfx90a;gfx942 -D__HIP_PLATFORM_AMD__\"\n    ENV CFLAGS=\"-D__HIP_PLATFORM_AMD__\"\n    ENV CXXFLAGS=\"-D__HIP_PLATFORM_AMD__\"\n    RUN pip install \"git+https://github.com/YangWang92/torch_memory_saver_numa.git@numa\"\n    ################################################################################\n    ################################################################################\n    ################################################################################\n\n\n\n    ########################################\n    ######Install ray#######################\n    ########################################\n    # need to add this patch: https://github.com/ray-project/ray/pull/53531/files\n    RUN pip uninstall ray -y\n    RUN pip install \"ray[data,train,tune,serve]>=2.47.0\" \n    ########################################\n    ########################################\n    ########################################\n\n\n    ##########################################\n    #######Install other dependencies#########\n    ##########################################\n    RUN pip install \"tensordict==0.6.2\" --no-deps && \\\n        pip install accelerate \\\n        codetiming \\\n        datasets \\\n        dill \\\n        hydra-core \\\n        liger-kernel \\\n        numpy \\\n        pandas \\\n        peft \\\n        \"pyarrow>=15.0.0\" \\\n        pylatexenc \\\n        torchdata \\\n        wandb \\\n        orjson \\\n        pybind11\n        \n    WORKDIR /workspace/\n    RUN git clone https://github.com/volcengine/verl.git && \\\n        cd verl && \\\n        pip install -e . \n    ##########################################\n    ##########################################\n    ##########################################\n\n    WORKDIR /workspace/\n    CMD [\"/usr/bin/bash\"]\n\n\nBuild the image:\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    docker docker/build -t verl-rocm .\n\nRun the container\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nNote: You can pull the docker from this DockerHub: [RLSys Foundation](https://hub.docker.com/u/yushengsuthu)\nPull the image:\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    docker pull rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4\n\n    docker tag rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4 verl-rocm:latest\n\nRun the container\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n\nOptional: Running without root and with user permissions\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. code-block:: bash\n\n    docker run --rm -it \\\n      --device /dev/dri \\\n      --device /dev/kfd \\\n      -p 8265:8265 \\\n      --group-add video \\\n      --cap-add SYS_PTRACE \\\n      --security-opt seccomp=unconfined \\\n      --privileged \\\n      -v $HOME/.ssh:/root/.ssh \\\n      -v $HOME:$HOME \\\n      --shm-size 128G \\\n      -w $PWD \\\n      verl-rocm \\\n      /bin/bash\n\n(Optional): If you do not want to root mode and require assign yourself as the user\nPlease add ``-e HOST_UID=$(id -u)`` and ``-e HOST_GID=$(id -g)`` into the above docker launch script. \n\nExample\n-------\n\nDue to to special setting in AMD (ROCM) torch, \n1. If your ``ray>=2.45.0`` (default), you need to set ``RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES`` when starting ray in verl's RLHF training and add this [patch](https://github.com/ray-project/ray/pull/53531/files).\n2. If your ``ray<2.45.0``, you need to set ``RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES`` when starting ray in verl's RLHF training.\nInference ``$ENGINE`` can be ``vllm`` or ``sglang``. We choose ``vllm`` as default in the following examples.\n\n\n\nPPO\n~~~\n\n.. code-block:: bash\n\n    YOUR_PROJECT_NAME=r1-verl-ppo-upstream\n    YOUR_RUN_NAME=r1-training_ppo-upstream \n    # export HYDRA_FULL_ERROR=1\n\n    export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n    \n    # [ray] < 2.45.0\n    #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1\n\n    # [ray] >= 2.45.0\n    export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794\n\n    GPUS_PER_NODE=8\n    MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct\n    python3 examples/data_preprocess/gsm8k.py --local_save_dir data/gsm8k\n    python3 -c \"import transformers; transformers.pipeline('text-generation', model='$MODEL_PATH')\"\n    ENGINE=vllm #sglang\n\n    PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n     data.train_files=data/gsm8k/train.parquet \\\n     data.val_files=data/gsm8k/test.parquet \\\n     data.train_batch_size=256 \\\n     data.val_batch_size=1312 \\\n     data.max_prompt_length=512 \\\n     data.max_response_length=256 \\\n     actor_rollout_ref.model.path=$MODEL_PATH \\\n     actor_rollout_ref.actor.optim.lr=1e-6 \\\n     actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n     actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n     actor_rollout_ref.rollout.name=$ENGINE \\\n     actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n     actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n     critic.optim.lr=1e-5 \\\n     critic.model.path=$MODEL_PATH \\\n     critic.ppo_micro_batch_size_per_gpu=4 \\\n     algorithm.kl_ctrl.kl_coef=0.001 \\\n     trainer.logger=console \\\n     trainer.project_name=$YOUR_PROJECT_NAME \\\n     trainer.experiment_name=$YOUR_RUN_NAME \\\n     trainer.val_before_train=False \\\n     trainer.n_gpus_per_node=$GPUS_PER_NODE \\\n     trainer.nnodes=1 \\\n     trainer.save_freq=10 \\\n     trainer.test_freq=10 \\\n     trainer.total_epochs=15 #2>&1 | tee verl_demo.log\n\nGRPO\n~~~~\n\n.. code-block:: bash\n\n    YOUR_PROJECT_NAME=r1-verl-grpo-upstream\n    YOUR_RUN_NAME=r1-training_grpo-upstream\n    # export HYDRA_FULL_ERROR=1\n    # export FSDP_VERBOSE=1 \n\n    #export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n\n    # [ray] < 2.45.0\n    #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1\n\n    # [ray] >= 2.45.0\n    export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794\n\n    GPUS_PER_NODE=8\n    MODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct\n    # MODEL_PATH=Qwen/Qwen2-7B-Instruct\n    python3 examples/data_preprocess/gsm8k.py --local_save_dir data/gsm8k\n    python3 -c \"import transformers; transformers.pipeline('text-generation', model='$MODEL_PATH')\"\n    ENGINE=vllm #sglang\n    \n    python3 -m verl.trainer.main_ppo \\\n        algorithm.adv_estimator=grpo \\\n        data.train_files=data/gsm8k/train.parquet \\\n        data.val_files=data/gsm8k/test.parquet \\\n        data.train_batch_size=1024 \\\n        data.val_batch_size=1312 \\\n        data.max_prompt_length=512 \\\n        data.max_response_length=1024 \\\n        actor_rollout_ref.model.path=$MODEL_PATH \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n        actor_rollout_ref.actor.use_dynamic_bsz=True \\\n        actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n        actor_rollout_ref.actor.use_kl_loss=True \\\n        actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n        actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=Flase \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n        actor_rollout_ref.rollout.name=$ENGINE \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n        actor_rollout_ref.rollout.n=5 \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=False \\\n        algorithm.kl_ctrl.kl_coef=0.001 \\\n        trainer.critic_warmup=0 \\\n        trainer.logger=console \\\n        trainer.project_name=$YOUR_PROJECT_NAME \\\n        trainer.experiment_name=$YOUR_RUN_NAME \\\n        trainer.n_gpus_per_node=$GPUS_PER_NODE \\\n        trainer.val_before_train=False \\\n        trainer.nnodes=1 \\\n        trainer.save_freq=-1 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15\n\n\n\nMulti-node training: slurm with Docker/Podman container\n---------------------------------------------------------------------------------------\n\nIf you want to run multi-node training with slurm, you can use the following script. \n\n.. note::\n    1. You need to use ``podman`` or ``docker`` in the following script. We will release the apptainer script later.\n    2. If you want to use ``podman``, you just replace ``docker`` with ``podman`` in the following script.\n\nThe script includes the following steps:\n\n1. SLURM Configuration\n2. Environment Setup\n3. Docker/Podman Container Setup\n4. Ray Cluster Initialization\n5. Data Preprocessing\n6. Model Setup\n7. Training Launch\n\n\nslurm_script.sh\n~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    #!/bin/bash\n\n    #SBATCH --job-name=verl-ray-on-slurm\n    #SBATCH --nodes=2\n    #SBATCH --ntasks-per-node=2\n    #SBATCH --mem=200G\n    #SBATCH --time=30-00:00:00\n    #SBATCH --gpus-per-node=8\n    #SBATCH --cpus-per-task=28\n    #SBATCH --output=../verl_log/slurm-%j.out\n    #SBATCH --error=../verl_log/slurm-%j.err\n    #SBATCH --nodelist=gpu-[0,1]\n\n\n    # load necessary modules\n    ### Run this setup\n    # [Cluster]: Use docker\n    # docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n\n\n    ##########################################################################\n    ###The following setting should be set in different project and cluster###\n    ##########################################################################\n\n    ### Project\n    CONTAINER_NAME=\"multinode_verl_training\"\n    IMG=\"verl.rocm\"\n    DOCKERFILE=\"docker/Dockerfile.rocm\"\n    # echo $PWD\n    verl_workdir=\"${HOME}/projects/verl_upstream\"\n    export TRANSFORMERS_CACHE=\"${HOME}/.cache/huggingface\"\n    export HF_HOME=$TRANSFORMERS_CACHE\n\n    ### Cluster Network Setting\n    export NCCL_DEBUG=TRACE\n    export GPU_MAX_HW_QUEUES=2\n    export TORCH_NCCL_HIGH_PRIORITY=1\n    export NCCL_CHECKS_DISABLE=1\n    # export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 \n    export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9\n    export NCCL_IB_GID_INDEX=3\n    export NCCL_CROSS_NIC=0\n    export CUDA_DEVICE_MAX_CONNECTIONS=1\n    export NCCL_PROTO=Simple\n    export RCCL_MSCCL_ENABLE=0\n    export TOKENIZERS_PARALLELISM=false\n    export HSA_NO_SCRATCH_RECLAIM=1\n    ##########################################################################\n\n    ## Assign using GPUs\n    export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n\n    ### For rocm and training script\n    # [ray] < 2.45.0\n    #export RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1\n\n    # [ray] >= 2.45.0\n    export RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 # Patch with https://github.com/ray-project/ray/pull/52794\n\n\n    # Build and launch the Docker container\n    srun bash -c \"\n        # Exit on any error\n        set -e \n\n        # Clean up dangling images (images with <none> tag)\n        docker image prune -f\n\n        # Need to pull the docker first\n        docker pull rlsys/verl:verl-0.4.1_ubuntu-22.04_rocm6.3.4-numa-patch_vllm0.8.5_sglang0.4.6.post4\n        \n        if ! docker images --format \"{{.Repository}}:{{.Tag}}\" | grep -q \"${IMG}\"; then\n            echo \\\"Building ${IMG} image...\\\"\n            docker build -f \\\"${DOCKERFILE}\\\" -t \\\"${IMG}\\\" .\n        else\n            echo \\\"${IMG} image already exists, skipping build\\\"\n        fi\n\n        # Removing old container if exists\n        docker rm \\\"${CONTAINER_NAME}\\\" 2>/dev/null || true\n\n        # Checking network devices\n        ibdev2netdev\n\n        # Launch the docker\n        docker run --rm -d \\\n        -e HYDRA_FULL_ERROR=1 \\\n        -e RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES=1 \\\n        -e RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES=1 \\\n        -e NCCL_DEBUG=${NCCL_DEBUG} \\\n        -e GPU_MAX_HW_QUEUES=${GPU_MAX_HW_QUEUES} \\\n        -e TORCH_NCCL_HIGH_PRIORITY=${TORCH_NCCL_HIGH_PRIORITY} \\\n        -e NCCL_CHECKS_DISABLE=${NCCL_CHECKS_DISABLE} \\\n        -e NCCL_IB_HCA=${NCCL_IB_HCA} \\\n        -e NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX} \\\n        -e NCCL_CROSS_NIC=${NCCL_CROSS_NIC} \\\n        -e CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS} \\\n        -e NCCL_PROTO=${NCCL_PROTO} \\\n        -e RCCL_MSCCL_ENABLE=${RCCL_MSCCL_ENABLE} \\\n        -e TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM} \\\n        -e HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM} \\\n        -e TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} \\\n        -e HF_HOME=${HF_HOME} \\\n        --network host \\\n        --device /dev/dri \\\n        --device /dev/kfd \\\n        --device /dev/infiniband \\\n        --group-add video \\\n        --cap-add SYS_PTRACE \\\n        --security-opt seccomp=unconfined \\\n        --privileged \\\n        -v \\${HOME}:\\${HOME} \\\n        -v \\${HOME}/.ssh:/root/.ssh \\\n        -w \"${verl_workdir}\" \\\n        --shm-size 128G \\\n        --name \\\"${CONTAINER_NAME}\\\" \\\n        \\\"${IMG}\\\" \\\n        tail -f /dev/null\n\n        echo \\\"Container setup completed\\\"\n    \"\n        # (Optional): If you do not want to root mode and require assign yuorself as the user\n        # Please add `-e HOST_UID=$(id -u)` and `-e HOST_GID=$(id -g)` into the above docker launch script. \n\n\n\n\n\n    ### Ray launch the nodes before training\n\n    # Getting the node names\n    nodes_array=($(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | tr '\\n' ' '))\n\n    head_node=${nodes_array[0]}\n    head_node_ip=$(srun --nodes=1 --ntasks=1 -w \"$head_node\" hostname --ip-address)\n\n    # if we detect a space character in the head node IP, we'll\n    # convert it to an ipv4 address. This step is optional.\n    if [[ \"$head_node_ip\" == *\" \"* ]]; then\n        IFS=' ' read -ra ADDR <<<\"$head_node_ip\"\n    if [[ ${#ADDR[0]} -gt 16 ]]; then\n        head_node_ip=${ADDR[1]}\n    else\n        head_node_ip=${ADDR[0]}\n    fi\n        echo \"IPV6 address detected. We split the IPV4 address as $head_node_ip\"\n    fi\n\n    port=6379\n    ip_head=$head_node_ip:$port\n    export ip_head\n    echo \"IP Head: $ip_head\"\n\n    # make sure we set environment variables before Ray initialization\n\n    # Print out all env variables\n    printenv\n\n    echo \"Starting HEAD at $head_node\"\n    srun --nodes=1 --ntasks=1 -w \"$head_node\" \\\n        docker exec \"${CONTAINER_NAME}\" \\\n            ray start --head --node-ip-address=\"$head_node_ip\" --port=$port \\\n            --dashboard-port=8266 \\\n            --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n    # optional, though may be useful in certain versions of Ray < 1.0.\n    sleep 10\n\n    # number of nodes other than the head node\n    worker_num=$((SLURM_JOB_NUM_NODES - 1))\n\n    for ((i = 1; i <= worker_num; i++)); do\n        node_i=${nodes_array[$i]}\n        echo \"Debug: Starting worker on node_i = ${node_i}\"\n        if [ -z \"$node_i\" ]; then\n            echo \"Error: Empty node name for worker $i\"\n            continue\n        fi\n        echo \"Starting WORKER $i at $node_i\"\n        srun --nodes=1 --ntasks=1 -w \"$node_i\" \\\n            docker exec \"${CONTAINER_NAME}\" \\\n                ray start --address \"$ip_head\" --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n        sleep 5\n    done\n\n\n\n\n    # Ray initlization test (See whether any error in the above execution)\n    echo \"Testing Ray initialization in the slurm nodes...\"\n    docker exec \"${CONTAINER_NAME}\" python3 -c '\n    import ray\n    try:\n        ray.init(address=\"auto\")\n        print(\"\\n=== Ray Cluster Status ===\")\n        print(f\"Number of nodes: {len(ray.nodes())}\")\n        for node in ray.nodes():\n            print(\"Node: {}, Status: {}\".format(node[\"NodeManagerHostname\"], node[\"Alive\"]))\n            # print(f\"Node: {node}\")\n        ray.shutdown()\n        print(\"Ray initialization successful!\")\n    except Exception as e:\n        print(f\"Ray initialization failed: {str(e)}\")\n    '\n    echo \"=== Ray test completed ===\"\n    ######\n\n\n\n    # Run data preprocessing\n\n    echo \"Starting data preprocessing...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 \"examples/data_preprocess/gsm8k.py\" \"--local_save_dir\" \"../data/gsm8k\"\n\n    echo \"Starting data preprocessing...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 \"examples/data_preprocess/math_dataset.py\" \"--local_dir\" \"../data/math\"\n\n    train_files=\"../data/gsm8k/train.parquet\"\n    val_files=\"../data/gsm8k/test.parquet\"\n\n    # Download and test model\n    echo \"Loading model...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')\"\n    MODEL_PATH=\"Qwen/Qwen2-7B-Instruct\"\n\n    # Set model path after pipeline test\n    MODEL_PATH=\"Qwen/Qwen2.5-0.5B-Instruct\"\n\n    echo \"== Data and model loading Done ==\"\n\n    echo \"Start to train...\"\n\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')\"\n    MODEL_PATH=\"Qwen/Qwen2-7B-Instruct\"\n\n\n    PYTHONUNBUFFERED=1 srun --overlap --nodes=${SLURM_NNODES} --ntasks=1 -w \"$head_node\" \\\n        docker exec \"${CONTAINER_NAME}\" \\\n        python3 -m verl.trainer.main_ppo \\\n        data.train_files=$train_files \\\n        data.val_files=$val_files \\\n        data.train_batch_size=1024 \\\n        data.max_prompt_length=1024 \\\n        data.max_response_length=1024 \\\n        actor_rollout_ref.model.path=$MODEL_PATH \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n        actor_rollout_ref.rollout.name=vllm \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n        critic.optim.lr=1e-5 \\\n        critic.model.use_remove_padding=True \\\n        critic.model.path=$MODEL_PATH \\\n        critic.model.enable_gradient_checkpointing=False \\\n        critic.ppo_micro_batch_size_per_gpu=8 \\\n        critic.model.fsdp_config.param_offload=False \\\n        critic.model.fsdp_config.optimizer_offload=False \\\n        algorithm.kl_ctrl.kl_coef=0.0001 \\\n        trainer.critic_warmup=0 \\\n        trainer.logger='[\"console\",\"wandb\"]' \\\n        trainer.project_name='verl_example' \\\n        trainer.experiment_name='Qwen2.5-32B-Instruct_function_rm' \\\n        trainer.n_gpus_per_node=${SLURM_GPUS_PER_NODE} \\\n        trainer.val_before_train=False \\\n        trainer.nnodes=${SLURM_NNODES} \\\n        trainer.save_freq=-1 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15\n\n\nRun slurm_script.sh\n~~~~~~~~~~~~~~~~~~~~\nJust sbatch your slurm_script.sh\n\n.. code-block:: bash\n\n    sbatch slurm_script.sh\n\n"
  },
  {
    "path": "docs/amd_tutorial/amd_vllm_page.rst",
    "content": "verl performance tuning for AMD (ROCm Kernel)\n=====================================================\n\nLast updated: 11/13/2025.\n\nAuthor: `Yang Wang <https://github.com/YangWang92/>`_, `Songlin Jiang <https://github.com/HollowMan6/>`_\n\nUse vLLM Sleep Mode for AMD MI3xx series GPUs\n--------------------------------------------------------------\n\nBy default, verl requires vLLM to enable sleep mode, which allows vLLM to offload GPU memory to CPU memory after rollout. This feature has been merged into the main branch of vLLM for version later than 0.11.0.\n\nFor now, you can use the vLLM main branch and build it from the source code, or you can directly install vLLM from the pre-built ROCm wheels for vLLM version later than 0.11.0 when it's available.\n\n1. Clone the vLLM repository and build it with the following commands:\n\n.. code-block:: bash\n\n    git clone https://github.com/vllm-project/vllm.git\n    cd vllm\n    git reset --hard 4ca5cd5740c0cd7788cdfa8b7ec6a27335607a48 # You can also use a later commit as you wish\n    python -m pip install -r requirements/rocm.txt\n    VLLM_TARGET_DEVICE=rocm ROCM_PATH=/opt/rocm/ python3 setup.py develop\n\n2. Additionally, we recommend you to use the ROCm version later than or equal to ROCm 7.0.\n\nAfter the upgrade, you can verify whether sleep mode is working by trying out `these scripts <https://github.com/EmbeddedLLM/inference-experiment/tree/main/sleep_mode>`_.\n\nIf sleep mode is working, you should see the memory usage reduce after sleep.\n\nAfter applying the vLLM patch and completing the installation, you can enable sleep mode in verl to reduce memory overhead. This allows verl to offload unused GPU memory during rollout, significantly lowering the memory footprint during long-context training or multi-node reinforcement learning.\n\n\nEnable CUDA Graph and Bypass ROCm-related issues\n--------------------------------------------------------------\n\nDue to potential issues with CUDA graph capture in ROCm, we've found that vLLM's CUDA graph feature cannot be enabled on multiple nodes in verl on AMD platforms with vLLM V1 mode. This leads to significantly slower rollout performance.\n\nOur investigation shows that ROCm may trigger an unexpected crash when attempting to capture large batches with CUDA graph. One workaround is to set ``actor_rollout_ref.rollout.cudagraph_capture_sizes`` to values such as ``[1, 2, 4, 8, 16, 32, 64]`` (change depending on your GPU memory size).\n\nThen, you can choose to enable CUDA graph by setting ``actor_rollout_ref.rollout.enforce_eager`` to ``False`` in your verl configuration file.\n"
  },
  {
    "path": "docs/api/data.rst",
    "content": "Data interface\n=========================\n\nLast updated: 05/19/2025 (API docstrings are auto-generated).\n\nDataProto is the interface for data exchange.\n\nThe :class:`verl.DataProto` class contains two key members:\n\n- batch: a :class:`tensordict.TensorDict` object for the actual data\n- meta_info: a :class:`Dict` with additional meta information\n\nTensorDict\n~~~~~~~~~~~~\n\n:attr:`DataProto.batch` is built on top of :class:`tensordict`, a project in the PyTorch ecosystem.\nA TensorDict is a dict-like container for tensors. To instantiate a TensorDict, you must specify key-value pairs as well as the batch size.\n\n.. code-block:: python\n\n    >>> import torch\n    >>> from tensordict import TensorDict\n    >>> tensordict = TensorDict({\"zeros\": torch.zeros(2, 3, 4), \"ones\": torch.ones(2, 3, 5)}, batch_size=[2,])\n    >>> tensordict[\"twos\"] = 2 * torch.ones(2, 5, 6)\n    >>> zeros = tensordict[\"zeros\"]\n    >>> tensordict\n    TensorDict(\n    fields={\n        ones: Tensor(shape=torch.Size([2, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False),\n        twos: Tensor(shape=torch.Size([2, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),\n        zeros: Tensor(shape=torch.Size([2, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},\n    batch_size=torch.Size([2]),\n    device=None,\n    is_shared=False)\n\nOne can also index a tensordict along its batch_size. The contents of the TensorDict can be manipulated collectively as well.\n\n.. code-block:: python\n\n    >>> tensordict[..., :1]\n    TensorDict(\n    fields={\n        ones: Tensor(shape=torch.Size([1, 3, 5]), device=cpu, dtype=torch.float32, is_shared=False),\n        twos: Tensor(shape=torch.Size([1, 5, 6]), device=cpu, dtype=torch.float32, is_shared=False),\n        zeros: Tensor(shape=torch.Size([1, 3, 4]), device=cpu, dtype=torch.float32, is_shared=False)},\n    batch_size=torch.Size([1]),\n    device=None,\n    is_shared=False)\n    >>> tensordict = tensordict.to(\"cuda:0\")\n    >>> tensordict = tensordict.reshape(6)\n\nFor more about :class:`tensordict.TensorDict` usage, see the official tensordict_ documentation.\n\n.. _tensordict: https://pytorch.org/tensordict/stable/overview.html\n\n\nCore APIs\n~~~~~~~~~~~~~~~~~\n\n.. autoclass::  verl.DataProto\n   :members: to, select, union, make_iterator, concat\n"
  },
  {
    "path": "docs/api/single_controller.rst",
    "content": "Single Controller interface\n============================\n\nLast updated: 05/27/2025 (API docstrings are auto-generated).\n\nThe Single Controller provides a unified interface for managing distributed workers\nusing Ray or other backends and executing functions across them.\nIt simplifies the process of dispatching tasks and collecting results, particularly \nwhen dealing with data parallelism or model parallelism. \n\n\nCore APIs\n~~~~~~~~~~~~~~~~~\n\n.. autoclass:: verl.single_controller.Worker\n   :members: __init__, __new__, get_master_addr_port, get_cuda_visible_devices, world_size, rank\n\n.. autoclass:: verl.single_controller.WorkerGroup\n   :members: __init__,  world_size\n\n.. autoclass:: verl.single_controller.ClassWithInitArgs\n   :members: __init__, __call__\n\n.. autoclass:: verl.single_controller.ResourcePool\n   :members: __init__, world_size, local_world_size_list, local_rank_list\n\n.. autoclass:: verl.single_controller.ray.RayWorkerGroup\n   :members: __init__\n\n.. autofunction:: verl.single_controller.ray.create_colocated_worker_cls"
  },
  {
    "path": "docs/api/trainer.rst",
    "content": "Trainer Interface\n================================\n\nLast updated: 06/08/2025 (API docstrings are auto-generated).\n\nTrainers drive the training loop. Introducing new trainer classes in case of new training paradiam is encouraged.\n\n.. autosummary::\n   :nosignatures:\n\n   verl.trainer.ppo.ray_trainer.RayPPOTrainer\n\n\nCore APIs\n~~~~~~~~~~~~~~~~~\n\n.. autoclass:: verl.trainer.ppo.ray_trainer.RayPPOTrainer\n   :members: __init__, init_workers, fit\n\n.. automodule:: verl.utils.tokenizer\n   :members: hf_tokenizer\n\n.. automodule:: verl.trainer.ppo.core_algos\n   :members: agg_loss, kl_penalty, compute_policy_loss, kl_penalty\n\n.. automodule:: verl.trainer.ppo.reward\n   :members: load_reward_manager, compute_reward, compute_reward_async\n\n.. autoclass:: verl.workers.reward_manager.NaiveRewardManager\n\n.. autoclass:: verl.workers.reward_manager.DAPORewardManager\n"
  },
  {
    "path": "docs/api/utils.rst",
    "content": "Utilities\n============\n\nLast updated: 05/19/2025 (API docstrings are auto-generated).\n\nThis section documents the utility functions and classes in the VERL library.\n\nPython Functional Utilities\n------------------------------\n\n.. automodule:: verl.utils.py_functional\n   :members: append_to_dict\n\nFile System Utilities\n------------------------\n\n.. automodule:: verl.utils.fs\n   :members: copy_to_local\n\nTracking Utilities\n---------------------\n\n.. automodule:: verl.utils.tracking\n   :members: Tracking\n\nMetrics Utilities\n---------------------\n\n.. automodule::  verl.utils.metric\n   :members: reduce_metrics\n\nCheckpoint Management\n------------------------\n\n.. automodule:: verl.utils.checkpoint.checkpoint_manager\n   :members: find_latest_ckpt_path\n\n.. automodule:: verl.utils.checkpoint.fsdp_checkpoint_manager\n   :members: FSDPCheckpointManager\n\nDataset Utilities\n---------------------\n\n.. automodule:: verl.utils.dataset.rl_dataset\n   :members: RLHFDataset, collate_fn\n\nTorch Functional Utilities\n-----------------------------\n\n.. automodule:: verl.utils.torch_functional\n   :members: get_constant_schedule_with_warmup, masked_whiten, masked_mean, logprobs_from_logits\n\nSequence Length Balancing\n----------------------------\n\n.. automodule:: verl.utils.seqlen_balancing\n   :members: get_reverse_idx, rearrange_micro_batches\n\nUlysses Utilities\n--------------------\n\n.. automodule:: verl.utils.ulysses\n   :members: gather_outputs_and_unpad, ulysses_pad_and_slice_inputs\n\nFSDP Utilities\n------------------\n\n.. automodule:: verl.utils.fsdp_utils\n   :members: get_fsdp_wrap_policy, get_init_weight_context_manager, init_fn, load_fsdp_model_to_gpu, load_fsdp_optimizer, offload_fsdp_model_to_cpu, offload_fsdp_optimizer,\n\nDebug Utilities\n-------------------\n\n.. automodule:: verl.utils.profiler\n   :members: log_gpu_memory_usage, GPUMemoryLogger\n\n"
  },
  {
    "path": "docs/ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst",
    "content": "NPU-CI 添加指导\n===========\n\nLast updated: 02/02/2026.\n\n我们在 verl 上增加基于华为昇腾设备的CI用例添加指导。\n\nverl 仓库使用 GitHub Actions 作为 CI 平台，通过分层测试架构保障代码质量与系统稳定性。\nNPU 相关的工作流主要包括：\n\n* ``npu_unit_test.yml``：运行单元测试。\n* 以 ``_ascend.yml`` 结尾的文件：运行针对 Ascend NPU 的端到端测试或专项测试。\n\n添加新用例指南\n-----------------------------------\n\n1. 数据集与权重\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n流水机器上的权重与绝对路径：\n\n+---------------------------------------+-------------------------------------------------------------------+\n| 模型名称                              | 绝对路径                                                          |\n+=======================================+===================================================================+\n| Qwen3-30B-A3B-Instruct-2507           | ``${HOME}/.cache/models/Qwen/Qwen3-30B-A3B-Instruct-2507``        |\n+---------------------------------------+-------------------------------------------------------------------+\n| Qwen2.5-VL-3B-Instruct                | ``${HOME}/.cache/models/Qwen/Qwen2.5-VL-3B-Instruct``             |\n+---------------------------------------+-------------------------------------------------------------------+\n| Qwen2.5-0.5B                          | ``${HOME}/.cache/models/Qwen/Qwen2.5-0.5B``                       |\n+---------------------------------------+-------------------------------------------------------------------+\n| Qwen2.5-0.5B-Instruct                 | ``${HOME}/.cache/models/Qwen/Qwen2.5-0.5B-Instruct``              |\n+---------------------------------------+-------------------------------------------------------------------+\n| Qwen2.5-1.5B-Instruct                 | ``${HOME}/.cache/models/Qwen/Qwen2.5-1.5B-Instruct``              |\n+---------------------------------------+-------------------------------------------------------------------+\n| Skywork-Reward-V2-Llama-3.2-1B        | ``${HOME}/.cache/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B``  |\n+---------------------------------------+-------------------------------------------------------------------+\n\n流水机器上的数据集与绝对路径：\n\n+--------------+---------------------------------------------------+\n| 数据集名称   | 绝对路径                                          |\n+==============+===================================================+\n| gsm8k        | ``${HOME}/.cache/datasets/openai/gsm8k``          |\n+--------------+---------------------------------------------------+\n| geo3k        | ``${HOME}/.cache/datasets/hiyouga/geometry3k``    |\n+--------------+---------------------------------------------------+\n\n**Note**\n\n   {HOME}是root\n\n   gpu用例中权重在~/models/路径下，如需适配可以用软链接，``ln -s /root/.cache/models ~/models``\n\n   此处为原始数据集，按需进行数据处理，如下。\n   \n   ``python examples/data_preprocess/gsm8k_multiturn_sft.py --local_dataset_path ${HOME}/.cache/datasets/openai/gsm8k``\n\n\n2. 工作流 YAML 模板\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n如需新增一个工作流，可参考以下模板创建 ``.github/workflows/your_yml_ascend.yml`` 文件。\n\n主要修改部分包括：\n\n* 工作流名称（``name``）\n* 触发条件（``on``）\n* 运行环境（``runs-on``）\n* 容器镜像（``container.image``）\n* 具体执行步骤（``jobs.<job_id>.steps``）\n\n.. code-block:: yaml\n   :linenos:\n\n   name: your_yml_ascend  # 工作流唯一标识\n   # 触发条件配置\n   on:\n     push:\n       branches:\n         - main\n         - v0.*\n     pull_request:\n       branches:\n         - main\n       paths:\n         - \".github/workflows/your_yml_ascend.yml\"  # 必须包含此工作流文件路径\n         - \"path/to/affected_files\"               # 需监控的相关代码路径\n\n   # 并发控制策略\n   concurrency:\n     group: ${{ github.workflow }}-${{ github.ref }}\n     cancel-in-progress: ${{ github.ref != 'refs/heads/main' }}  # 仅非main分支取消进行中的任务\n\n   permissions:\n     contents: read  # 最小权限原则\n\n   jobs:\n     your_job_name:  # 任务唯一标识\n       if: github.repository_owner == 'verl-project'  # 仅在主仓库运行\n       runs-on: linux-aarch64-a2-4  # 硬件规格：a2实例，4卡NPU\n       timeout-minutes: 60          # 任务超时阈值（分钟）\n       container:\n         #运行镜像 该示例为vllm的镜像\n         image: swr.ap-southeast-1.myhuaweicloud.com/base_image/ascend-ci/verl/verl:verl-8.5.0-910b-ubuntu22.04-py3.11-latest\n         options: >-\n           --shm-size 16g  # 共享内存配置\n       env:\n         HF_ENDPOINT: \"https://hf-mirror.com\"\n         HF_HUB_ENABLE_HF_TRANSFER: \"0\"\n       steps:\n         - name: Check npu and CANN info\n           run: |\n             cat /usr/local/Ascend/ascend-toolkit/latest/\"$(uname -i)\"-linux/ascend_toolkit_install.info\n             npu-smi info\n         - name: Check initial pip list from image\n           run: pip list\n         - name: Checkout repository\n           uses: actions/checkout@v4\n           with:\n             fetch-depth: 0 \n             clean: true \n         - name: Install dependencies\n           run: |\n             pip install -r requirements-npu.txt\n             pip install -e .\n         - name: Verify environment\n           run: pip list\n         # 以下为具体测试步骤（根据需求定制）\n         - name: Preprocess dataset\n           run: python examples/data_preprocess/your_script.py --local_dataset_path ${HOME}/.cache/datasets/your_dataset\n         - name: Execute NPU test\n           run: |\n             ray stop --force \n             bash tests/special_npu/your_test_script.sh\n\n**Note**\n\n\n   ${HOME}/.cache/文件夹内一旦添加新内容，不会因CI跑完容器销毁而删除，请避免往该文件夹添加内容。\n\n\n3. 添加单元测试\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n步骤：\n\n(1) 在 ``tests/`` 目录下创建或修改单元测试文件（例如 ``test_xxx.py``）。\n(2) 若测试文件路径未被 ``npu_unit_test.yml`` 中的 ``--ignore-glob`` 规则排除，则会在以下命令中自动执行：\n\n   .. code-block:: yaml\n   \n      pytest -s -x --ignore-glob=\"xxx\" --ignore-glob=\"xxx\" tests/\n   \n(3) 若测试路径在 ``--ignore-glob`` 排除范围内，需在 ``npu_unit_test.yml`` 中新增一个 step 来显式运行该测试。\n(4) 如新增一批相关用例，建议单独创建专门的工作流文件以保持清晰。\n\n4. 添加端到端测试脚本\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n步骤：\n\n(1) 在 ``tests/special_npu/`` 目录下创建端到端测试脚本。\n(2) 在 ``.github/workflows/`` 目录中找到功能最接近的以 ``_ascend.yml`` 结尾的工作流文件，在其中添加一个 step 调用该脚本。\n(3) 若测试场景独立或较复杂，可考虑单独创建新的工作流文件。\n\n5. 测试策略建议\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n* **单元测试**：覆盖核心函数、类与方法，确保逻辑正确。\n* **集成/端到端测试**：覆盖典型训练、推理 pipeline，验证多模块协同与硬件适配。\n* **资源管理**：一个workflow里的多个job为并行运行，请合理设置超时时间，避免任务长时间挂起，请控制单个 job 的运行时间在 40min 以内。\n\n通过以上步骤，可系统化地为 verl 仓库添加 NPU 相关的自动化测试，确保代码变更在合并前经过充分验证。\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/ascend_performance_analysis_guide.md",
    "content": "# Ascend Performance Analysis Guide\n\nLast updated: 02/24/2026.\n\n## 背景介绍\n\n随着DeepSeek-R1的发布，大模型强化学习（RL）训练受到广泛关注。在昇腾NPU环境下，verl框架已积累了丰富的性能调优经验。本文系统总结了包括性能数据采集与分析在内的方法论，旨在帮助开发者更高效地运用MindStudio工具链，实现强化学习场景下的性能优化。\n\n### 强化学习计算流程概述\n\n1. **Rollout**：策略（actor）模型基于输入的prompt序列，推理生成回答（response序列）\n2. **ref logprob**：基于prompt和生成的response，reference模型计算ref logprob用于KL散度计算\n3. **logprob**：基于prompt和生成的response，actor模型计算logprob用于重要性采样\n4. **reward**：基于prompt和生成的response，奖励模型评估奖励值R_N。\n5. **update**：基于计算得到的R_N、ref logprob、logprob计算优化函数和策略梯度，对actor模型进行更新\n\n![rl_data_stream](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/rl_data_stream.png)\n\n## profilling工具使能\n\n### 使能方法\n\n使能和配置教程可参考：[verl/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst at main · verl-project/verl](https://github.com/verl-project/verl/raw/main/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst)\n\n## 性能分析方法论\n\n### 整体性能概览分析\n\n#### 1. 长耗时任务与资源空泡分析\n\n- **操作**：使用MindStudio Insight加载profiling数据，自动识别不同计算阶段，通过RL页签流水图定位长耗时任务与NPU资源空泡\n- **价值**：快速掌握不同阶段耗时占比\n- **效果展示**：\n\n![Bubble_analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Bubble_analysis.png)\n\n#### 2. 负载均衡分析\n\n- **操作**：通过MindStudio Insight直接查看MSTX打点数据，观察Rollout阶段不同DP Rank的负载均衡情况\n- **价值**：快速识别负载不均问题\n\n- **效果展示：**\n\n![Load_Balancing_Analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Load_Balancing_Analysis.gif)\n\n#### 3. 集群整体性能分析\n\n- **操作**：结合MSTT的rl_analysis功能，生成集群Timeline缩略图，观察各阶段整体耗时\n- **价值**：宏观掌握集群性能瓶颈\n- **操作指南**：[rl_analysis使用文档](https://gitcode.com/Ascend/mstt/raw/pre-research/profiler/msprof_analyze/docs/features/rl_analysis.md)\n- **效果展示**：\n\n![Cluster%20Performance%20Analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Cluster%20Performance%20Analysis.png)\n\n### 细粒度分析\n\n#### 性能分析\n\n- **操作**：可通过 MindStudio Insight Windows 或 Linux 版本加载 Profiling 数据\n\n- **价值**：MindStudio Insight 支持分析任务调度效率、算子执行性能、计算资源利用率、集合通信性能等。其 Timeline 视图具备任务拆解与 Overlap 分析功能（**为 MindStudio 独有核心特性，在 NV 及其他竞品中不具备，是 AI 调优的必备工具**），并支持鼠标交互式分析。\n\n- **效果展示**：\n\n![performance%20analysis](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/performance%20analysis.png)\n\n#### 内存分析\n\n##### **通过 Profiling 结合调用栈分析系统内存变化**\n\n- **操作**：采集数据时开启调用栈和内存视图功能。\n- **价值**：观察框架、CANN内存申请释放情况，可结合调用栈跟踪到前端python代码。\n- **效果展示**：结合调用栈进行内存变化分析。效果如下所示：\n\n![in-memory%20analytics](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/in-memory%20analytics.gif)\n\n##### **使用 msleaks 工具进行深层次内存分析**\n\n- **操作步骤**：参考 [msleaks 工具使用指南](https://www.hiascend.com/document/detail/zh/CANNCommunityEdition/83RC1alpha003/devaids/msleaks/atlas_msleaks_0001.html)。\n- **价值**：可以查看框架内存申请总量折线图/内存块图，并直接对应调用栈，可深层次分析框架内存使用情况。\n- **效果展示**：\n\n![msleaks](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/msleaks.gif)\n\n## 性能分析案例\n\n要做具体的性能分析，profiling要开启**level1**，否则算子的关键信息会缺失。\n\n### 1.host bound诊断\n\nhost bound是指CPU任务量综合大于NPU，导致NPU执行出现空泡的现象。可以通过看Host2Device的同步连线来判断，如果连线都是歪的，那证明这里的set信号早于wait信号，NPU一ready就执行了，那也是device bound：\n\n![host_bound_1](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/host_bound_1.png)\n\n如果确诊为host bound，那么我们可以打开CPU侧，找出各算子的下发耗时。注意找的时候需要找出所有CPU耗时的累加值，而不能找单层，因为首次调用的耗时是很长的。例如下图的GmmSwigluQuant，CPU上首次调用需要1ms，后续每次只需要200us。\n\n![host_bound_2](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/host_bound_2.png)\n\n此时有的算子在负重前行，有的算子拖了后腿，后者多于了前者。我们优先**找出来host耗时大于device的top算子，这些算子是拖后腿的**，可以交予算子团队重点分析。\n\n### 2.组网合理性分析\n\n有的时候，模型组网没有按照最高效的方式来，这一点在profiling中是非常易于识别的，下面会介绍一下分析思路并给出例子。\n\n通常来讲，LLM中的大的热点算子是Attention和FFN中的矩阵乘计算，二者加起来在prefill下可能达到计算耗时的70%+，decode下可能达到50%+。如果整体的耗时比例不符合预期，或者profiling中出现了一些新面孔，或者拼接类算子太多了，这都值得我们去分析一下模型组网，是不是使用算子的方式错了？尤其是拼接类算子，是值得我们逐一分析的。\n\n对于slice/split/concat这样的拼接类算子，还有transpose/cast这种转换算子，他们的存在往往是前后算子不直接配套造成的。如果前一个算子可以直接对输出做好尾处理，往往可以节省一个算子的启动开销和一次冗余读写。但这样的改变不一定符合算子的基本设计原则。\n\n举一个正例，对于某次Matmul的输出shape为[m, n0 + n1]，在这后面我们接了两个slice，输入均为这个[m, n0 + n1]的tensor，输出分别为[m, n0]和[m, n1]。第一个优化的思路是将两个slice改为一个split，这样耗时可以基本减半，[m, n0 + n1]的显存也可以尽早释放。进一步优化的思路是将矩阵乘的权重从[k, n0 + n1]分割为[k, n0]和[k, n1]，将原来的矩阵乘任务分成两个（前提是这两个的耗时加起来不比之前的劣化太多，分核策略不能出问题），从而彻底消除这个slice/split操作。\n\n![network_1](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/network_1.png)\n\n举一个反例，Rmsnorm(fp16)+Cast(fp16->fp32)+Matmul(fp32)，Rmsnorm虽然输入输出都是fp16，但考虑到累加运算的精度，内部是fp32做计算的。如果将Cast融到Rmsnorm内，本就内部使用fp32做计算的Rmsnorm就可以省去一个末尾fp32->fp16的cast，加上我们干掉的Cast，总共节省两个cast的同时避免了一次精度丢失。虽然这样看起来精度性能双收了，但fp16进，fp32出的Rmsnorm是反原则的（核心输入和输出需要是同数据类型），除非我们能在广大开源模型中频繁找到这样的结构，证明它的普适性，否则算子团队是不允许做这样的算子的。\n\n![network_2](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/network_2.png)\n\n### 3.算子性能初诊\n\n需要利用`\".\\ASCEND_PROFILER_OUTPUT\\operator_details.csv\"`来做分析，从而判断算子识否有性能问题。\n\nProfiling工具会统计这些流水线在不同核上的平均繁忙时间（xxx_time），与最慢核的完整kernel耗时（task_duration）做除法，得到流水线利用率（xxx_ratio）。这些流水线之间虽然互有依赖，且搬运类流水线会互抢带宽，但算子只要设计得当，是可以做到互相掩盖的。因此我们可以初步认为，**当算子的执行耗时大到一定程度上，算子应当在某一条流水线上形成bound**，即利用率要高到一定程度。经验上，在单算子耗时达到50μ时，就可以认为算子应当在bound流水线上，达成80%+的占用率了。\n\n以下图为例，第一行是一个FA算子，第二行是一个Matmul算子，FA在vec流水线上达到了88.1%的利用率，Matmul算子在mac流水线上达到了89.8%的利用率，他们的性能可以认为是合格的。\n\n![Operator%20performance](https://github.com/chengminhua/verl_data/raw/main/MindStudio_Insight_use/Operator%20performance.png)\n\n\n\n### 4.亲和shape调整\n\n对于一个模型而言，超参是我们控制不了的，但我们可以控制并发度、权重格式、切分策略等因素来迎合算子，使其发挥出最大的性能，这一节主要从算子搬运效率和负载均衡两个方面出发，讨论模型侧值得尝试的调整方向。\n\n#### 4.1 搬运效率亲和的shape\n\nmte2是一个自身效率严重受shape影响的流水线。要想让mte2保证最大搬运效率，我们需要保障如下两个条件至少满足其一：\n\n**（1）被搬运的矩阵使用nz作为format（最优）\n（2）被搬运的矩阵的尾轴512B对齐，且不为16KB的整数倍（近似最优）**\n\n对于权重矩阵来说，推理阶段尤其是decode，我们通常满足（1），训练阶段我们通常满足（2）。**如果我们做不到（1），我们就要迎合（2）**。典型的手段有：\n\n1，如果没达成B的矩阵的首轴是亲和的而尾轴不亲和，那么对它做transpose\n2，调整TP切分策略，避免出现不亲和的尾轴\n\n#### 4.2 负载均衡亲和的shape\n\n在算子shape不大时，受制于算子语义，我们有可能不能把所有核都利用起来，或者即使开满核，负载均衡却很差。这一小节主要是对decode阶段的小shape做分析。\n\n首先，我们明确出当前NPU卡是多少核的，如果不清楚，跑出来的profiling里都是20，40这样的数，就说明是20核，反之是24核。这里我的24核其实是代表了一个cube和两个vector组成的小组，我们可以认为是一个cube作为主核，带了两个vector作为从核。如果一个算子是纯vector算子，那么就不再有组的概念，40或48个vector核会作为主核直接独立去拿逻辑任务。\n\n对于LLM中的vector算子，它的一种常见分核策略有可能是分在最高维，也就是batch维，常见于对低维（也叫尾轴）有规约操作的norm类、动态量化类等算子；另一种是整体拍平，允许算子切分的非常细的算子，如elementwse算子。对于第一种，我们就可以在模型侧关注它的负载均衡问题。例如我们打48batch，而硬件却是个40个vector核，那这40个核会循环2次，第二次有多数的核会无事可做，这个batch数就可以认为是不友好的。如果将batch打到64或80，性能可以预见会是无损的。同样的情况下，如果是48核的卡，那我们可以认为这就是个非常友好的batch数。\n\n对于cube类算子，它常见的分核策略是以base快去切分M和N（K轴是累加轴，对它分核会引入确定性问题）。最常见的分块是baseM=128，baseN=256。在decode阶段，我们的耗时基本可以看做都是在搬权重，这是因为激活的M极小，M方向大概率只分了一块，那么右矩阵就只需要搬一次。所以我们在M≤128的范围内可以尽情提高M，对性能都基本是无损的，如果M大于128，可以认为(128, 256]是下一个性能分档。\n除了M外，N轴切分的任务也影响算子亲和性，以deepseekR1中的MLA预处理为例，它会使用同一个激活（shape为[batch_size, 7168]）与两个权重做矩阵乘(shape为[7168, 1536]和[7168, 576])。在batch_size打不大的情况下，即使baseN缩短为128，N轴都不能用满核数，所以此时这两个矩阵乘各自的耗时，会约等于将他们权重N轴拼起来乘(shape为[7168, 2112])的矩阵乘的耗时。如果仅考虑模型竞争力，我们更希望对这两个权重做合并，否则两个小的矩阵乘带宽利用率都会非常差。\n\n对于Attention算子，它常见的分核策略是q_seqlen、batch_size和kv_headnum。增量阶段q_seqlen会以MTP和GQA倍数做合并，但是通常也不会大过128，划分不出第二个任务，那么并行度基本就是batch_size * kv_headnum。\n\n总的来说，我们可以依据shape信息和算子类别，对算子是否有负载均衡问题作出识别，从而对我们切分策略选择，最高吞吐量的batch策略作出预判。\n\n\n\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/ascend_retool_best_pratice.rst",
    "content": "Ascend Retool Best Practice\n===================================\n\nLast updated: 03/01/2026.\n\n引言\n----------------------------------\n\nRetool论文参考([Retool](https://arxiv.org/pdf/2504.11536))\n集成代码解释器工具，通过多轮实时代码执行进行策略部署，并教会模型根据结果反馈学习何时以及如何调用工具。\n\n1. 环境构建\n2. 模型训练\n\n用例模型脚本以及其需要的硬件条件各自如下：\n\n===============    ============    ============    ===============\n模型                NPU型号         节点数量        训推后端\n===============    ============    ============    ===============\n``Qwen2.5-7B``     Atlas 900 A2         1          ``vllm + FSDP``\n===============    ============    ============    ===============\n\n环境构建\n-----------------------------------\n1.从自定义Conda环境进行构建\n\n============    ============================================================\nsoftware        version \n============    ============================================================\nPython          ``>= 3.10, <3.12``\nCANN            ``== 8.3.RC1``\ntorch           ``== 2.7.1``\ntorch_npu       ``== 2.7.1``\nverl            ``v0.6.1 commitId=d62da4950573d7a4b7ef2362337952e7ab59e78d``\nvllm            ``v0.11.0``\nvllm-ascend     ``v0.11.0-dev``\ntransformers    ``4.57.6``\n============    ============================================================\n\n模型训练与评估\n-----------------------------------\n1.模型数据准备\n^^^^^^^^^^^\n`Qwen2.5-7B`\n^^^^^^^^^^^\n**下载模型权重**\n\n--local-dir: 模型保存路径\n\n.. code-block:: bash\n\n  git clone https://huggingface.co/Qwen/Qwen2.5-7B-Instruct\n\n**下载训练数据集**\n\n.. code-block:: bash\n\n  git clone https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k\n\n**下载评估数据集**\n\n.. code-block:: bash\n\n  git clone https://huggingface.co/datasets/Maxwell-Jia/AIME_2024\n\n**下载预训练数据集**\n\n.. code-block:: bash\n\n  python3 recipe/retool/retool_sft_preprocess.py\n\n*注:自动下载ReTool-SFT，最后生成数据默认保存在~/ReTool-SFT/data目录下*\n\n**执行预训练脚本**\n\n.. code-block:: bash\n\n  bash recipe/retool/run_qwen2_7b_sft_npu.sh # 需适配脚本中路径\n\n**合并预训练权重生成checkpoint**\n\n.. code-block:: bash\n\n  python3 -m verl.model_merger merge --backend fsdp \\\n      --local_dir ${DATASETS}/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372 \\\n      --target_dir ${DATASETS}/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372/huggingface\n\n2.代码沙箱准备\n\n开源沙箱代码及部署参考\nhttps://github.com/bytedance/SandboxFusion\n\n**沙箱代码下载**\n\n.. code-block:: bash\n\n  git clone -b main https://github.com/bytedance/SandboxFusion.git\n\n**沙箱安装**\n\n.. code-block:: bash\n\n  cd SandboxFusion\n  conda create -n sandbox -y python=3.11\n  conda activate sandbox\n  pip install poetry\n  poetry lock\n  poetry install\n  mkdir -p docs/build\n  cd runtime/python\n  bash install-python-runtime.sh\n  cd ../../\n  make run-online\n\n3.训练\n\n示例配置文件如下，在recipe/retool目录下创建一个run_qwen2.5_7b_dapo_npu.sh\n根据开发者实际路径配置情况修改模型训练脚本中的以下参数\n\n.. code-block:: bash \n\n  set -x\n\n  export VLLM_USE_V1=1\n  export TORCHDYNAMO_DISABLE=1\n  export VLLM_ASCEND_ENABLE_NZ=0\n  export TASK_QUEUE_ENABLE=1\n  export VLLM_ENABLE_GRAPH_MODE=1\n  export HCCL_OP_EXPANSION_MODE=\"AIV\"\n  export VLLM_ASCEND_ENABLE_MLP_OPTIMIZE=1\n  export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\n  \n  # ================= data/model/tool =================\n  HDFS_ROOT=${HDFS_ROOT:-\"${PWD}\"}\n  DATA_ROOT=${DATA_ROOT:-\"${PWD}\"}\n  \n  dapo_math_17k=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k\n  aime_2024=$DATA_ROOT/dataset/Maxwell-Jia/AIME_2024\n  #aime_2025=$DATA_ROOT/dataset/yentinglin/aime_2025\n  model_path=$DATA_ROOT/dataset/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372/huggingface\n  \n  train_files=\"['$dapo_math_17k']\"\n  test_files=\"['$aime_2024']\"\n  \n  # tool\n  tool_config_path=recipe/retool/sandbox_fusion_tool_config.yaml\n  \n  # wandb\n  project_name=retool\n  experiment_name=qwen2.5-7b_dapo\n  default_local_dir=$DATA_ROOT/checkpoint/$experiment_name\n  \n  # 创建日志文件\n  export TIMESTAMP=$(date +%Y%m%d_%H%M%S)\n  LOG_DIR=\"$HDFS_ROOT/verl/logs/$project_name/$experiment_name\"\n  # 判断路径是否存在\n  if [ ! -d \"$LOG_DIR\" ]; then\n    # 路径不存在，创建路径\n    mkdir -p \"$LOG_DIR\"\n    echo \"Directory $LOG_DIR created.\"\n  else\n    echo \"Directory $LOG_DIR already exists.\"\n  fi\n  \n  LOG_FILE=\"${LOG_DIR}/${TIMESTAMP}.log\"\n  touch \"$LOG_FILE\"\n  echo \"Log file $LOG_FILE created.\"\n\n  # ================= algorithm =================\n  adv_estimator=grpo\n  \n  use_kl_in_reward=False\n  kl_coef=0.0\n  use_kl_loss=False\n  kl_loss_coef=0.0\n  \n  clip_ratio_low=0.2\n  clip_ratio_high=0.28\n  \n  max_turns=16\n  max_prompt_length=2048\n  max_response_length=20480\n  actor_lr=1e-6\n  \n  train_batch_size=32\n  ppo_mini_batch_size=16\n  \n  n_resp_per_prompt=16\n  n_resp_per_prompt_val=30\n  \n  # ================= performance =================\n  infer_tp=2 # vllm\n  train_sp=4 # train\n  offload=True\n  \n  actor_max_token_len_per_gpu=$(( (max_prompt_length + max_response_length) * 1 ))\n  log_prob_max_token_len_per_gpu=$(( actor_max_token_len_per_gpu * 4 ))\n\n  PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=$adv_estimator \\\n    algorithm.use_kl_in_reward=$use_kl_in_reward \\\n    algorithm.kl_ctrl.kl_coef=$kl_coef \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=$train_batch_size \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.custom_cls.path=recipe/retool/retool.py \\\n    data.custom_cls.name=CustomRLHFDataset \\\n    custom_reward_function.path=recipe/retool/retool.py \\\n    custom_reward_function.name=compute_score \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \\\n    actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.optim.lr=$actor_lr \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$train_sp \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=$offload \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=$offload \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$log_prob_max_token_len_per_gpu \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$actor_max_token_len_per_gpu \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.max_num_seqs=1024 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \\\n    actor_rollout_ref.rollout.multi_turn.enable=True \\\n    actor_rollout_ref.rollout.multi_turn.max_user_turns=$max_turns \\\n    actor_rollout_ref.rollout.multi_turn.max_assistant_turns=$max_turns \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=$tool_config_path \\\n    actor_rollout_ref.rollout.multi_turn.format=hermes \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.n=$n_resp_per_prompt \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    trainer.logger=['console'] \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$experiment_name \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.log_val_generations=20 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=100 \\\n    trainer.default_local_dir=$default_local_dir \\\n    trainer.test_freq=20 \\\n    trainer.device=npu \\\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.actor.entropy_checkpointing=True \\\n    actor_rollout_ref.ref.entropy_checkpointing=True \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    trainer.total_epochs=1 $@ > $LOG_FILE 2>&1 &\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/ascend_sglang_best_practices.rst",
    "content": "Ascend SGLang Best Practice\n===================================\n\nLast updated: 01/27/2026.\n\n.. _Qwen3-30B: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh\n.. _Qwen2.5-32B: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2-32b_sglang_fsdp_npu.sh\n引言\n----------------------------------\n\nSGLang 是当前主流的高性能开源推理引擎, 昇腾已经全面原生支持该推理引擎在verl中使用,\n仅需简单的构建流程，开发者即可完成环境构建，本文将提供两个经典用例来帮助开发者了解以下内容：\n\n1. 环境构建\n2. 模型训练与评估\n3. 性能采集\n\n两个用例模型脚本以及其需要的硬件条件各自如下：\n\n+----------------------+---------------------+----------+------------------------+\n| 模型                 | NPU型号             | 节点数量 | 训推后端               |\n+======================+=====================+==========+========================+\n| `Qwen3-30B`_         | Atlas 800T A3       | 1        | SGLang + Megatron      |\n+----------------------+---------------------+----------+------------------------+\n| `Qwen2.5-32B`_       | Atlas 900 A2        | 2        | SGLang + FSDP          |\n+----------------------+---------------------+----------+------------------------+\n\n环境构建\n-----------------------------------\n我们在quickstart中提供了两种构建环境的方法, 1.从镜像文件DockerFile进行构建 2.从自定义Conda环境进行构建\n\n在本实践中, 我们额外指定verl 的commit id 以避免引入其他问题\n\n.. code-block:: bash\n\n    cd verl\n    git checkout c98cb8cc\n模型训练与评估\n-----------------------------------\n1.模型数据准备\n^^^^^^^^^^^\n`Qwen3-30B`_\n^^^^^^^^^^^\n**下载模型权重**\n\nQwen3-30B: https://huggingface.co/Qwen/Qwen3-30B-A3B\n\n**下载数据集**\n\nDAPO-Math-17k: https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k\n\n**HuggingFace To Megatron权重转换(可选)**\n\n.. code-block:: bash\n\n  python scripts/converter_hf_to_mcore.py \\\n      --hf_model_path Qwen/Qwen3-30B-A3B \\\n      --output_path Qwen/Qwen3-30B-A3B-mcore \\\n      --use_cpu_initialization    # Only work for MoE models\n*注:verl当前已支持mbridge进行灵活的hf和mcore之间的权重转换,可以修改以下相关参数直接加载hf权重*\n\n.. code-block:: bash\n\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=False\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n\n`Qwen2.5-32B`_\n^^^^^^^^^^^\n**下载模型权重**\n\n--local-dir: 模型保存路径\n\n.. code-block:: bash\n\n  export HF_ENDPOINT=https://hf-mirror.com\n  hf download --resume-download Qwen/Qwen2.5-32B --local-dir /path/to/local_dir\n\n**下载及处理数据集**\n\n.. code-block:: bash\n\n    wget https://huggingface.co/datasets/agentica-org/DeepScaleR-Preview-Dataset/resolve/main/deepscaler.json\n    python recipe/r1_ascend/json_to_parquet.py --output_dir ./data/deepscaler --json_path path/to/deepscaler.json --train_data_ratio 0.9\n\n2.训练\n^^^^^^^^^^^\n根据开发者实际路径配置情况修改模型训练脚本中的以下参数\n\n.. code-block:: bash \n\n    # Model Weights Paths\n    MODEL_PATH=Qwen/Qwen3-30B-A3B\n    MCORE_MODEL_PATH=Qwen/Qwen3-30B-A3B-mcore\n    RAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n    CKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n    # File System Paths\n    TRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\n    TEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\n\n    #保存频率，-1默认不保存，如需评测请修改此参数\n    trainer.save_freq=-1\n\n对于单机任务 `Qwen3-30B`_ , 可以直接bash执行verl仓上示例脚本\n\n.. code-block:: bash \n\n  bash examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh\n对于多节点任务 `Qwen2.5-32B`_ ，我们推荐使用以下脚本进行大规模多节点训练拉起\n\n.. code-block:: bash\n\n  pkill -9 python\n  ray stop --force\n  rm -rf /tmp/ray\n  export RAY_DEDUP_LOGS=0\n  export HYDRA_FULL_ERROR=1\n  # TASK_QUEUE_ENABLE，下发优化，图模式设置为1，非图模式设置为2\n  export TASK_QUEUE_ENABLE=1\n  export HCCL_ASYNC_ERROR_HANDLING=0\n  export HCCL_EXEC_TIMEOUT=3600\n  export HCCL_CONNECT_TIMEOUT=3600\n  \n  export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\n  export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\n  export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\n  export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8\n  # 修改为当前需要跑的用例路径\n  DEFAULT_SH=\"./run_*.sh\"\n  echo \"Use $DEFAULT_SH\"\n  \n  ulimit -n 32768\n  mkdir logs\n  \n  NNODES=2\n  NPUS_PER_NODE=8\n  # 修改为对应主节点IP\n  MASTER_ADDR=\"IP FOR MASTER NODE\"\n  # 修改为当前节点的通信网卡\n  SOCKET_IFNAME=\"Your SOCKET IFNAME\"\n  export HCCL_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\n  export GLOO_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\n  # 获取当前IP\n  CURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\\.){3}[0-9]{1,3}' | awk '{print $NF}')\n  if [ \"$MASTER_ADDR\" = \"$CURRENT_IP\" ]; then\n    # 主节点启动\n    ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{\"NPU\": '$NPUS_PER_NODE'}'\n  \n    while true; do\n        ray_status_output=$(ray status)\n        npu_count=$(echo \"$ray_status_output\" | grep -oP '(?<=/)\\d+\\.\\d+(?=\\s*NPU)' | head -n 1)\n        npu_count_int=$(echo \"$npu_count\" | awk '{print int($1)}')\n        device_count=$((npu_count_int / $NPUS_PER_NODE))\n  \n        # 判断device_count 是否与 NNODES 相等\n        if [ \"$device_count\" -eq \"$NNODES\" ]; then\n            echo \"Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script.\"\n            ray status\n            bash $DEFAULT_SH\n            break\n        else\n            echo \"Waiting for Ray to allocate $NNODES devices. Current device count: $device_count\"\n            sleep 5\n        fi\n    done\n  else\n    # 子节点尝试往主节点注册 ray 直到成功\n    while true; do\n        # 尝试连接 ray 集群\n        ray start --address=\"$MASTER_ADDR:6766\" --resources='{\"NPU\": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP\n  \n        # 检查连接是否成功\n        ray status\n        if [ $? -eq 0 ]; then\n            echo \"Successfully connected to the Ray cluster!\"\n            break\n        else\n            echo \"Failed to connect to the Ray cluster. Retrying in 5 seconds...\"\n            sleep 5\n        fi\n    done\n  fi\n  \n  sleep 600\n\nDEFAULT_SH:修改为训练所用配置 sh 文件路径。在此案例中修改为 `Qwen2.5-32B`_ 路径。\n          \nNNODES 和 NPUS_PER_NODE:修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为2和8。\n          \nMASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。\n          \nSOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取：\n          \n.. code-block:: bash\n          \n  ifconfig |grep \"$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')\" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1\n\n3.模型评估\n^^^^^^^^^^^\n\n不同模型步骤一致,仅以Qwen3-30b为例列举\n\n我们通过 AISBenchmark 评估模型,该工具支持vllm/sglang多种推理后端的评估\n\n**安装方法**\n\n.. code-block:: bash\n\n  git clone https://gitee.com/aisbench/benchmark.git\n  cd benchmark\n  pip install -e .\n\n**下载评估数据集**\n\n.. code-block:: bash\n\n  cd path/to/benchmark/ais_bench/datasets\n  wget http://opencompass.oss-cn-shanghai.aliyuncs.com/datasets/data/math.zip\n  unzip math.zip\n  rm math.zip\n\n**修改AISBench配置代码使能sglang推理评测**\n\n打开 benchmark/ais_bench/benchmark/configs/models/vllm_api/vllm_api_stream_chat.py 文件，这是推理配置文件\n\n.. code-block:: bash\n\n    from ais_bench.benchmark.models import VLLMCustomAPIChatStream\n    from ais_bench.benchmark.utils.model_postprocessors import extract_non_reasoning_content\n    from ais_bench.benchmark.clients import OpenAIChatStreamClient, OpenAIChatStreamSglangClient\n\n    models = [\n        dict(\n            attr=\"service\",\n            type=VLLMCustomAPIChatStream,\n            abbr='sgl-api-stream-chat',\n            path=\"/path/to/Qwen3-30B\", # 修改为 Qwen3-30B 模型路径\n            model=\"qwen3-30b\",\n            request_rate = 0,\n            max_seq_len=2048,\n            retry = 2,\n            host_ip = \"localhost\", # 推理服务的IP\n            host_port = 8005, # 推理服务的端口\n            max_out_len = 8192,  # 最大输出tokens长度\n            batch_size=48, # 推理的最大并发数\n            trust_remote_code=False,\n            custom_client=dict(type=OpenAIChatStreamSglangClient), #使用sglang客户端\n            generation_kwargs = dict(\n                temperature = 0,\n                seed = 1234,\n            ),\n            pred_postprocessor=dict(type=extract_non_reasoning_content)\n        )\n    ]\n\n\n**启动sglang_server服务**\n\n.. code-block:: bash\n\n    python -m sglang.launch_server --model-path \"/path/to/Qwen3-30B\"  --tp-size 4 --dp-size 1 --port 8005 \n\n**启动sglang_client评测**\n\n.. code-block:: bash\n\n    ais_bench --models vllm_api_stream_chat --datasets math500_gen_0_shot_cot_chat_prompt\n\n**评测结果**\n\n经过训练,模型在Math-500上的评分显著上升\n\n+------+----------------------+---------+----------+------+----------------------+\n| iter | dataset              | version | metric   | mode | sgl-api-stream-chat  |\n+======+======================+=========+==========+======+======================+\n|   0  | math_prm800k_500     | c4b6f0  | accuracy | gen  | \t84.4             |\n+------+----------------------+---------+----------+------+----------------------+\n|  150 | math_prm800k_500     | c4b6f0  | accuracy | gen  |     91.7             |\n+------+----------------------+---------+----------+------+----------------------+\n\n性能采集\n-----------------------------------\n关于NPU profiling的详细文档请参考 `ascend_profiling_zh <https://github.com/volcengine/verl/blob/main/docs/ascend_tutorial/profiling/ascend_profiling_zh.rst>`_\n\n在 `Qwen3-30B`_ 的脚本中提供了基本的采集性能选项PROF_CONFIG，默认设置 global_profiler.steps=null 关闭采集， 开发者可根据实际需要进行参数修改\n\n采集完成后，开发者可以使用 `MindStudio Insight <https://www.hiascend.com/document/detail/zh/mindstudio/830/GUI_baseddevelopmenttool/msascendinsightug/Insight_userguide_0002.html>`_ 进行数据解析\n\n注: verl框架侧进行采集全量 Profiling 产生海量且重复的算子记录，可以根据文档修改代码仅采集关键阶段\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/dapo_multi_model_optimization_practice.md",
    "content": "# DAPO multi model optimization practice\n\n## DAPO 介绍\n\nLast updated: 03/04/2026.\n\nDAPO的论文可以参考：[DAPO](https://arxiv.org/pdf/2503.14476)，其中包含以下几个关键技术。\n\n* ​**Clip-Higher**​: 通过对重要性采样比的上限剪裁促进了系统的多样性并避免了熵坍缩（Entropy Collapse）。\n* ​**Dynamic Sampling**​: 提高了训练效率和稳定性。DAPO出了一种执行动态采样的策略，并过滤掉准确率等于1和0的提示组，从而保持批次间具有有效梯度的提示数量一致。\n* ​**Token-level Policy Gradient Loss**​: 在长链思维强化学习 (long-CoT RL) 场景中至关重要。\n* ​**Overlong Reward Shaping**​: 减少奖励噪声并稳定了训练。\n\n在verl中，可以进行如下设置，从而进行DAPO算法的运行。\n\n- **奖励模型的管理策略为 DAPO**\n  在dapo算法中，必须配置成dapo。\n\n```\nreward_model.reward_manager.name=dapo\n```\n\n- **Clip-Higher 更高裁剪**\n  `clip_ratio_low` 和 `clip_ratio_high` 用于指定 DAPO 目标函数中的 $\\varepsilon_{\\text {low }}$ 和 $\\varepsilon_{\\text {high }}$。\n\n```\nclip_ratio_low=0.2  # 裁剪比例下限，默认值为0.2\nclip_ratio_high=0.28 # 裁剪比例上限，默认值为0.28\n```\n\n- **动态采样的相关配置**\n  将 `filter_groups.enable` 设置为 `True` 会过滤掉输出 `metric` 完全相同的组，例如对于 `acc` 指标，过滤掉输出准确率全部为 1 或 0 的组。\n  训练器会使用 `gen_batch_size` 进行重复采样，直到生成足够数量的符合条件的组，或者达到 `max_num_gen_batches` 所指定的上限为止。\n\n```\ndata.gen_batch_size=${gen_prompt_bsz}\nalgorithm.filter_groups.enable=${enable_filter_groups} # 动态采样开关\nalgorithm.filter_groups.metric=${filter_groups_metric} # 使用准确率作为过滤标准\nalgorithm.filter_groups.max_num_gen_batches=${max_num_gen_batches} # 最大生成批次数量,最多重复生成数据的次数\n```\n\n- **Token-level Loss**\n  将 `loss_agg_mode` 设置为 `token-mean` 意味着计算一个批次中所有序列内所有 token 的（策略梯度）损失的平均值。\n\n```\nactor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode}\n# 注意：“token-mean”是默认行为。\n```\n\n- **奖励模型对超长回答的惩罚配置**\n  将 `overlong_buffer.enable` 设置为 `True` 将对输出长度过长但仍未超过硬上下文限制的输出进行惩罚。具体来说，当输出的长度超过 `max_response_length - overlong_buffer.len` 且超出 `0` 到 `overlong_buffer.len` 个 token 时，惩罚值会从 `0` 线性增加到 `overlong_buffer.penalty_factor`。\n\n```\nreward_model.overlong_buffer.enable=${enable_overlong_buffer} # 启用超长缓冲区惩罚,开启对超长输出的惩罚机制\nreward_model.overlong_buffer.len=${overlong_buffer_len}  # 缓冲区长度,定义缓冲区的toke,最大惩罚强度\nreward_model.overlong_buffer.penalty_factor=${overlong_penalty_factor}   #惩罚因子,最大惩罚强度\n```\n\n相关参数涉及的代码可以参考：[Recipe: Decoupled Clip and Dynamic Sampling Policy Optimization (DAPO)](https://github.com/verl-project/verl-recipe/blob/main/dapo/README.md)\n\n## 硬件要求\n\n当前支持Atlas 800T A3 与 Atlas 900 A3 SuperPoD。完成跑完本次最佳实践需要 2台Atlas 800T A3。关键软件版本可以参考：[Ascend Quickstart](https://github.com/volcengine/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_quick_start.rst)\n\n## 安装基础环境\n\n| software | version|\n| --- | --- |\n| Python| >= 3.10, <3.12 |\n| CANN | == 8.3.RC1 |\n| torch | == 2.7.1 |\n| torch_npu | == 2.7.1 |\n| verl | main分支 commitId=252d76908b903ad8fb6969eb3a5e5f873c95ea2b |\n| vllm | \tv0.11.0 |\n| vllm-ascend | v0.11.0-dev|\n| transformers | \t4.57.3|\n\n在本实践中, 我们通过指定 verl 的commit id 以避免引入其他问题\n```\ncd verl\ngit checkout 252d76908b903ad8fb6969eb3a5e5f873c95ea2b\n# 指定相应的recipe版本\ngit submodule update --init --recursive recipe\ncd recipe\ngit checkout main\n```\n\n## 模型训练\n\n### 数据集准备\n\nGeometry3k 数据集是由加利福尼亚大学洛杉矶分校与浙江大学联合研发的几何领域专用数据集，核心面向视觉问答（VQA）任务展开研究与模型训练。该数据集总计包含 3002 个样本，采用图像和文本两种模态数据形式构建，其中文本模态涵盖各类几何问题描述，图像则以可视化图表呈现问题中的几何图形信息，包括三角形、圆形、四边形等基础几何形状，以及不同图形间的位置、嵌套、相交等关联关系。可以从Hugging Face库下载对应的原始数据集：[Geometry3k ](https://huggingface.co/datasets/hiyouga/geometry3k)\n\n```python\n# 下载原始数据并预处理\npython ./examples/data_preprocess/geo3k.py --local_dir=./data/geo3k\n```\n\n### 权重下载\n\n从Hugging Face库下载对应的模型权重：[Qwen3-VL-30B-A3B-Instruct](https://huggingface.co/Qwen/Qwen3-VL-30B-A3B-Instruct/tree/main\n)\n\n### jemalloc安装\n\n为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理。\n\n#### Ubuntu 操作系统\n\n通过操作系统源安装jemalloc（注意： 要求ubuntu版本>=20.04）：\n\n```shell\nsudo apt install libjemalloc2\n```\n\n在启动任务前执行如下命令通过环境变量导入jemalloc，需先通过 **find /usr -name libjemalloc.so.2** 确认文件是否存在 ：\n\n```shell\n# arm64架构\nexport LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2\n# x86_64架构\nexport LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2\n```\n\n#### OpenEuler 操作系统\n\n执行如下命令重操作系统源安装jemalloc\n\n```shell\nyum install jemalloc\n```\n\n如果上述方法无法正常安装，可以通过源码编译安装 前往jemalloc官网下载最新稳定版本，官网地址:https://github.com/jemalloc/jemalloc/releases/\n\n```shell\ntar -xvf jemalloc-{version}.tar.bz2\ncd jemalloc-{version}\n./configure --prefix=/usr/local\nmake\nmake install\n```\n\n### 全局变量导入\n\n- 为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理，用于更好管理内存，避免长跑过程中内存 OOM。\n\n```\n# 根据实际安装路径设置 jemalloc 环境变量，例如安装路径为:/usr/local/lib/libjemalloc.so.2(可通过 find /usr -name libjemalloc.so.2 确认文件是否存在)\nexport LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\n```\n\n- 某些模型是通过 vllm ascend 进行优化的。但在某些情况下，优化后的模型可能并不适用。此时，将此值设置为 0 即可禁用优化后的模型。\n\n```\nexport USE_OPTIMIZED_MODEL=0\n```\n\n- 启用vLLM V1\n\n```\nexport VLLM_USE_V1=1\n```\n\n- 昇腾多卡通信的兜底配置，延长连接超时时间，避免集群环境下训练启动因连接慢而失败\n\n```\nexport HCCL_CONNECT_TIMEOUT=5400\n```\n\n- 控制 vLLM 在昇腾芯片上是否启用NZ优化\n\n```\nexport VLLM_ASCEND_ENABLE_NZ=0\n```\n\n### 训练\n```\n# Model Weights Paths\nMODEL_PATH=hf_weights/Qwen3-VL-30B-A3B-Instruct\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=$RAY_DATA_HOME/datasets/geo3k/train.parquet\nTEST_FILE=$RAY_DATA_HOME/datasets/geo3k/test.parquet\n\n# 保存频率，-1默认不保存，如需评测请修改此参数\ntrainer.save_freq=-1\n```\n\n对于单机任务 Qwen3-VL-30B , 修改脚本中参数`trainer.nnodes`为 1， `trainer.n_gpus_per_node` 为16，然后直接bash执行verl仓上示例脚本\n\n```\nbash recipe/dapo/run_dapo_qwen3_vl_30b_fsdp2_npu.sh\n```\n对于多节点任务 Qwen3-VL-30B ，我们推荐使用以下脚本进行大规模多节点训练拉起\n\n```\npkill -9 python\nray stop --force\nrm -rf /tmp/ray\nexport VLLM_USE_V1=1\nexport HCCL_CONNECT_TIMEOUT=5400\nexport VLLM_ASCEND_ENABLE_NZ=0\nexport LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\n# 修改为当前需要跑的用例路径\nDEFAULT_SH=\"./run_*.sh\"\necho \"Use $DEFAULT_SH\"\n\nulimit -n 32768\nmkdir logs\n\nNNODES=2\nNPUS_PER_NODE=8\n# 修改为对应主节点IP\nMASTER_ADDR=\"IP FOR MASTER NODE\"\n# 修改为当前节点的通信网卡\nSOCKET_IFNAME=\"Your SOCKET IFNAME\"\nexport HCCL_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\nexport GLOO_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\n# 获取当前IP\nCURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\\.){3}[0-9]{1,3}' | awk '{print $NF}')\nif [ \"$MASTER_ADDR\" = \"$CURRENT_IP\" ]; then\n  # 主节点启动\n  ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{\"NPU\": '$NPUS_PER_NODE'}'\n\n  while true; do\n      ray_status_output=$(ray status)\n      npu_count=$(echo \"$ray_status_output\" | grep -oP '(?<=/)\\d+\\.\\d+(?=\\s*NPU)' | head -n 1)\n      npu_count_int=$(echo \"$npu_count\" | awk '{print int($1)}')\n      device_count=$((npu_count_int / $NPUS_PER_NODE))\n\n      # 判断device_count 是否与 NNODES 相等\n      if [ \"$device_count\" -eq \"$NNODES\" ]; then\n          echo \"Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script.\"\n          ray status\n          bash $DEFAULT_SH\n          break\n      else\n          echo \"Waiting for Ray to allocate $NNODES devices. Current device count: $device_count\"\n          sleep 5\n      fi\n  done\nelse\n  # 子节点尝试往主节点注册 ray 直到成功\n  while true; do\n      # 尝试连接 ray 集群\n      ray start --address=\"$MASTER_ADDR:6766\" --resources='{\"NPU\": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP\n\n      # 检查连接是否成功\n      ray status\n      if [ $? -eq 0 ]; then\n          echo \"Successfully connected to the Ray cluster!\"\n          break\n      else\n          echo \"Failed to connect to the Ray cluster. Retrying in 5 seconds...\"\n          sleep 5\n      fi\n  done\nfi\n\nsleep 600\n```\nDEFAULT_SH: 修改为训练所用配置 sh 文件路径。在此案例中修改为 [Qwen3_VL_30B](https://github.com/verl-project/verl-recipe/blob/main/dapo/run%20dapo_qwen3_vl_30b_fsdp2_npu.sh) 路径。\n\nNNODES 和 NPUS_PER_NODE: 修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为2和8。\n\nMASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。\n\nSOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取：\n\n```\nifconfig |grep \"$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')\" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1\n```\n\n## 优化参考\n\n- **启动动态批次大小**\n  根据单 GPU 的最大 Token 总数（ppo_max_token_len_per_gpu）动态调整批次大小\n\n```\nactor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\nactor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\nactor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n```\n\n- **单个 GPU 能处理的最大 Token 总数**\n  当`use_dynamic_bsz=True`时，单 GPU 在一个微批次中能处理的最大 Token 数量\n\n```\nactor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}  \nactor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \nactor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n```\n\n- **单个 GPU 微批次大小**\n  当`use_dynamic_bsz=True`时，框架会以该值为​初始批次大小​，再根据`ppo_max_token_len_per_gpu`向上 / 向下调整\n\n```\nactor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \nactor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \nactor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2\n```\n\n- **启用 FSDP2 框架**\n  “将模型参数、梯度、优化器状态分片存储在不同 GPU 上”，避免单卡加载全量模型导致显存溢出。\n\n```\n# 启用 FSDP2 框架\nactor_rollout_ref.actor.strategy=fsdp2 \nactor_rollout_ref.ref.strategy=fsdp2 \ncritic.strategy=fsdp2\n\n# 仅用于 FSDP2：前向传播后重新分片以减少内存占用。\nactor_rollout_ref.actor.fsdp_config.reshard_after_forward=True\n# 仅用于 FSDP2：是否在模型前向传播后重新分片以节省内存。\nactor_rollout_ref.ref.fsdp_config.reshard_after_forward=True\n```\n\n- **启用专家并行配置**\n  指定有多少个 GPU用于并行计算不同的专家网络\n\n```\n# MoE 架构 Actor 模型的专家并行配置\nactor_rollout_ref.rollout.expert_parallel_size=8\n```\n\n\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/gspo_optimization_practice.md",
    "content": "# NPU Qwen3-32B GSPO Optimization Practice\r\n\r\nLast updated: 02/26/2026.\r\n\r\n本文章对应脚本地址：[qwen3_32b_gspo_npu](https://github.com/volcengine/verl/blob/main/examples/gspo_trainer/run_qwen3_32b_gspo_npu.sh)\r\n\r\n## 算法适配\r\n\r\nGSPO通过将优化颗粒度从**token级**提升到**sequence级**，规避了GRPO会遇到的**方差急剧增大**导致训练不稳定的情况，增加了训练的稳定性，同时该算法也在一定程度上提升了算法的收敛速度。\r\n\r\n想要成功在verl仓库中成功调用到GSPO算法，需要进行如下的必要配置\r\n\r\n~~~python\r\n# 核心算法配置  \r\nalgorithm.adv_estimator=grpo \\                    # 使用GRPO优势估计器    \r\nalgorithm.use_kl_in_reward=False \\                # 不在奖励中添加KL惩罚    \r\n# GSPO策略损失模式  \r\nactor_rollout_ref.actor.policy_loss.loss_mode=gspo \\ # 启用GSPO策略损失\r\n# 极小裁剪范围（GSPO特色）  \r\nactor_rollout_ref.actor.clip_ratio_low=0.0003 \\   # 裁剪下界，论文推荐值    \r\nactor_rollout_ref.actor.clip_ratio_high=0.0004 \\  # 裁剪上界，论文推荐值    \r\n# KL配置（GSPO不使用KL loss）  \r\nactor_rollout_ref.actor.use_kl_loss=False \\       # 禁用KL损失    \r\nactor_rollout_ref.actor.kl_loss_coef=0.0 \\        # KL损失系数设为0    \r\n# 序列级损失聚合模式（GSPO核心）  \r\nactor_rollout_ref.actor.loss_agg_mode=seq-mean-token-mean \\ # 序列级平均，GSPO论文推荐    \r\n# 批次配置  \r\nactor_rollout_ref.rollout.n=16 \\                  # 每个prompt生成16个响应（组采样）\r\n~~~\r\n\r\n一般选择入口函数为`verl.trainer.main_ppo`\r\n\r\n## 基础环境\r\n\r\n当前支持Atlas 800T A3 与 Atlas 900 A3 SuperPoD。完成跑完本次最佳实践需要 4台Atlas 800T A3。关键软件版本可以参考：[Ascend Quickstart](https://github.com/volcengine/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_quick_start.rst)\r\n\r\n### 安装基础环境\r\n\r\n| software     | version                                                    |\r\n| ------------ | ---------------------------------------------------------- |\r\n| Python       | >= 3.10, <3.12                                             |\r\n| CANN         | == 8.3.RC1                                                 |\r\n| torch        | == 2.7.1                                                   |\r\n| torch_npu    | == 2.7.1                                                   |\r\n| verl         | main分支 commitId=252d76908b903ad8fb6969eb3a5e5f873c95ea2b |\r\n| vllm         | v0.11.0                                                    |\r\n| vllm-ascend  | v0.11.0-dev                                                |\r\n| transformers | 4.57.3                                                     |\r\n\r\n在本实践中, 我们通过指定 verl 的commit id 以避免引入其他问题\r\n\r\n~~~bash\r\ncd verl\r\ngit checkout 252d76908b903ad8fb6969eb3a5e5f873c95ea2b\r\n# 指定相应的recipe版本\r\ngit submodule update --init --recursive recipe\r\n~~~\r\n\r\n### 权重获取\r\n\r\n从Hugging Face库下载对应的模型权重：[Qwen/Qwen3-32B · Hugging Face](https://huggingface.co/Qwen/Qwen3-32B)\r\n\r\n### 数据集准备\r\n\r\n~~~bash\r\n# 下载math-17k数据集\r\ngit clone https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k\r\n\r\n# 下载AIME_2024测试数据集\r\ngit clone https://huggingface.co/datasets/Maxwell-Jia/AIME_2024\r\n~~~\r\n\r\n### jemalloc安装\r\n\r\n为了确保 Ray 进程能够正常回收内存，需要安装并使能 jemalloc 库进行内存管理。\r\n\r\n#### Ubuntu 操作系统\r\n\r\n通过操作系统源安装jemalloc（注意： 要求ubuntu版本>=20.04）：\r\n\r\n```shell\r\nsudo apt install libjemalloc2\r\n```\r\n\r\n在启动任务前执行如下命令通过环境变量导入jemalloc，需先通过 **find /usr -name libjemalloc.so.2** 确认文件是否存在 ：\r\n\r\n```shell\r\n# arm64架构\r\nexport LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2\r\n# x86_64架构\r\nexport LD_PRELOAD=/usr/lib/x86_64-linux-gnu/libjemalloc.so.2\r\n```\r\n\r\n#### OpenEuler 操作系统\r\n\r\n执行如下命令重操作系统源安装jemalloc\r\n\r\n```shell\r\nyum install jemalloc\r\n```\r\n\r\n如果上述方法无法正常安装，可以通过源码编译安装 前往jemalloc官网下载最新稳定版本，官网地址:https://github.com/jemalloc/jemalloc/releases/\r\n\r\n```shell\r\ntar -xvf jemalloc-{version}.tar.bz2\r\ncd jemalloc-{version}\r\n./configure --prefix=/usr/local\r\nmake\r\nmake install\r\n```\r\n\r\n在启动任务前执行如下命令通过环境变量导入jemalloc：\r\n\r\n```shell\r\n#根据实际安装路径设置环境变量，例如安装路径为:/usr/local/lib/libjemalloc.so.2,可通过以下命令来设置环境变量(可通过 find /usr -name libjemalloc.so.2 确认文件是否存在)\r\nexport LD_PRELOAD=/usr/lib/aarch64-linux-gnu/libjemalloc.so.2\r\n```\r\n\r\n### 多机任务拉起\r\n\r\n针对本实践提供的多机任务，可用下面的脚本拉起\r\n\r\n~~~bash\r\npkill -9 python\r\nray stop --force\r\nrm -rf /tmp/ray\r\n\r\nexport RAY_DEDUP_LOGS=0\r\nexport HYDRA_FULL_ERROR=1\r\nexport TASK_QUEUE_ENABLE=1\r\nexport HCCL_EXEC_TIMEOUT=3600\r\nexport HCCL_CONNECT_TIMEOUT=3600\r\nexport HCCL_ASYNC_ERROR_HANDLING=0\r\nexport CPU_AFFINITY_CONF=1\r\nexport VLLM_USE_V1=1\r\nexport VLLM_ATTENTION_BACKEND=XFORMERS\r\nexport VLLM_ASCEND_ENABLE_FLASHCOMM=1\r\nexport VLLM_ASCEND_ENABLE_PREFETCH_MLP=1\r\nexport VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE=1\r\nexport LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\r\n\r\n# 修改为当前需要跑的用例路径\r\nDEFAULT_SH=\"./run_*.sh\"\r\necho \"Use $DEFAULT_SH\"\r\n\r\nulimit -n 32768\r\nmkdir logs\r\n\r\nNNODES=4\r\nNPUS_PER_NODE=16\r\n# 修改为对应主节点IP\r\nMASTER_ADDR=\"IP FOR MASTER NODE\"\r\n# 修改为当前节点的通信网卡\r\nSOCKET_IFNAME=\"Your SOCKET IFNAME\"\r\nexport HCCL_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\r\nexport GLOO_SOCKET_IFNAME=\"SOCKET IFNAME FOR CURRENT NODE\"\r\n# 获取当前IP\r\nCURRENT_IP=$(ifconfig $SOCKET_IFNAME | grep -Eo 'inet (addr:)?([0-9]{1,3}\\.){3}[0-9]{1,3}' | awk '{print $NF}')\r\nif [ \"$MASTER_ADDR\" = \"$CURRENT_IP\" ]; then\r\n  # 主节点启动\r\n  ray start --head --port 6766 --dashboard-host=$MASTER_ADDR --node-ip-address=$CURRENT_IP --dashboard-port=8260 --resources='{\"NPU\": '$NPUS_PER_NODE'}'\r\n\r\n  while true; do\r\n      ray_status_output=$(ray status)\r\n      npu_count=$(echo \"$ray_status_output\" | grep -oP '(?<=/)\\d+\\.\\d+(?=\\s*NPU)' | head -n 1)\r\n      npu_count_int=$(echo \"$npu_count\" | awk '{print int($1)}')\r\n      device_count=$((npu_count_int / $NPUS_PER_NODE))\r\n\r\n      # 判断device_count 是否与 NNODES 相等\r\n      if [ \"$device_count\" -eq \"$NNODES\" ]; then\r\n          echo \"Ray cluster is ready with $device_count devices (from $npu_count NPU resources), starting Python script.\"\r\n          ray status\r\n          bash $DEFAULT_SH\r\n          break\r\n      else\r\n          echo \"Waiting for Ray to allocate $NNODES devices. Current device count: $device_count\"\r\n          sleep 5\r\n      fi\r\n  done\r\nelse\r\n  # 子节点尝试往主节点注册 ray 直到成功\r\n  while true; do\r\n      # 尝试连接 ray 集群\r\n      ray start --address=\"$MASTER_ADDR:6766\" --resources='{\"NPU\": '$NPUS_PER_NODE'}' --node-ip-address=$CURRENT_IP\r\n\r\n      # 检查连接是否成功\r\n      ray status\r\n      if [ $? -eq 0 ]; then\r\n          echo \"Successfully connected to the Ray cluster!\"\r\n          break\r\n      else\r\n          echo \"Failed to connect to the Ray cluster. Retrying in 5 seconds...\"\r\n          sleep 5\r\n      fi\r\n  done\r\nfi\r\n\r\nsleep 600\r\n~~~\r\n\r\nDEFAULT_SH:修改为训练所用配置 sh 文件路径。在此案例中修改为 [Qwen2.5-32B](https://github.com/volcengine/verl/blob/main/examples/gspo_trainer/run_qwen3_32b_gspo_npu.sh) 路径。\r\n\r\nNNODES 和 NPUS_PER_NODE:修改为使用节点数量和每个节点 NPU 数量。在此案例中分别为4和16。\r\n\r\nMASTER_ADDR:修改为对应主节点 IP。即所有节点的 MASTER_ADDR 应该相同。\r\n\r\nSOCKET_IFNAME, HCCL_SOCKET_IFNAME, GLOO_SOCKET_IFNAME: 修改为对应通信网卡，通信网卡可以通过以下命令获取：\r\n\r\n```\r\nifconfig |grep \"$(hostname -I |awk '{print $1}'|awk -F '.' '{print $0}')\" -B 1|awk -F ':' '{print$1}' | head -1 | tail -1\r\n```\r\n\r\n## 性能调优\r\n\r\n优化从训练、推理、调度和其他四个方面入手。\r\n\r\n### 训练\r\n\r\n#### 动态bsz\r\n\r\n~~~bash\r\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\r\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\r\n~~~\r\n\r\n**这个优化点主要调整上面这两个参数，不过需要注意这两个参数调整的太大会导致OOM**\r\n\r\n**主要调整**`actor_ppo_max_token_len`,调大了会降低训练的耗时，调整`infer_ppo_max_token_len`没有明显的收益，可以不动\r\n\r\n**这两个参数的作用介绍如下：**\r\n\r\n**这两个参数用于控制动态批处理(dynamic batch size)模式下每个GPU处理的最大token数量**\r\n\r\n- **`actor_ppo_max_token_len`**: Actor模型在PPO更新(前向+反向传播)时每个GPU能处理的最大token数\r\n- **`infer_ppo_max_token_len`**: 推理阶段(Reference policy和Rollout)计算log概率时每个GPU能处理的最大token数\r\n\r\n### 推理\r\n\r\n#### ACLgraph+FULL_DECODE_ONLY\r\n\r\n推理算子下发方面的优化，平均能有`15%~20%`左右的性能收益。\r\n\r\n先看单开**ACLgraph**，如下：\r\n\r\n~~~bash\r\n# 开启ACLgraph+FULL_DECODE_ONLY（注意：当设置此参数为False时，TASK_QUEUE_ENABLE必须设置为1，不然会报错）\r\nactor_rollout_ref.rollout.enforce_eager=False\r\nactor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes='[8,16,32,64,128]' \\ \r\nactor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode='FULL_DECODE_ONLY' \\\r\n~~~\r\n\r\n`FULL_DECODE_ONLY`开启成功后有如下输出：\r\n\r\n![FULL_DECODE_ONLY result](https://github.com/wucong25/verl-data/blob/main/ascend_acl_graph.png)\r\n\r\n**`cudagraph_capture_sizes`参数设置指南**\r\n\r\ncudagraph_capture_sizes设置的值对应的是批大小，这里的批大小不是配置里的DP域对应的那个批次大小，这里是相较于vllm来说的批大小，单位为**token**\r\n\r\n默认生成的算法如下，可做参考\r\n\r\n![cudagraph_capture_sizes](https://github.com/wucong25/verl-data/blob/main/ascend_set_cudagraph_sizes.png)\r\n\r\n##### 推理后端切换\r\n\r\n使用方式：`export VLLM_ATTENTION_BACKEND=XFORMERS`\r\n\r\n![VLLM_ATTENTION_BACKEND](https://github.com/wucong25/verl-data/blob/main/ascend_vllm_attn_backend.png)\r\n\r\n注：需要注意某些后端在一些比较老的vllm-ascend版本内并不支持\r\n\r\n##### 使能vllm v1版本\r\n\r\n使用方式：`export VLLM_USE_V1=1`\r\n\r\n可以常开，一般都是正收益。\r\n\r\n### 调度\r\n\r\n#### AIV\r\n\r\n打开方式：设置`export HCCL_OP_EXPANSION_MODE=\"AIV\"`\r\n\r\nHCCL_OP_EXPANSION_MODE环境变量用于配置通信算法的编排展开位置，支持如下取值：\r\n\r\n- AI_CPU：代表通信算法的编排展开位置在Device侧的AI CPU计算单元。\r\n- AIV：代表通信算法的编排展开位置在Device侧的Vector Core计算单元。\r\n- HOST：代表通信算法的编排展开位置为Host侧CPU，Device侧根据硬件型号自动选择相应的调度器。\r\n- HOST_TS：代表通信算法的编排展开位置为Host侧CPU，Host向Device的Task Scheduler下发任务，Device的Task Scheduler进行任务调度执行。\r\n\r\n下面介绍两种展开机制\r\n\r\n##### HOST展开\r\n\r\n<img src=\"https://github.com/wucong25/verl-data/blob/main/ascend_task_queue1.png\" alt=\"image-20260113194257095\" style=\"zoom:50%;\" />\r\n\r\n- 软件栈工作在hostcpu，通信算法展开一个个task\r\n- 每个task调用runtime接口，下发到device的rtsqueue\r\n- STARS从rstqueue上顺序拿取task\r\n- 根据task类型分别调用掉SDMA和RDMA引擎。\r\n    **单算子瓶颈**：hostbound 每个task提交是2~5us，一个通信算子有几百个task，单算子场景不会在device上缓存，下发一个执行一个\r\n\r\n##### AICpu机制展开\r\n\r\n<img src=\"https://github.com/wucong25/verl-data/blob/main/ascend_task_queue3.png\" alt=\"image-20260113194333218\" style=\"zoom:50%;\" />\r\n\r\n- host侧不下发一个个task，把通信算子作为一个个kernel，放在通信算子kernel的队列上去。\r\n- STARS调度kernel队列流上的kernel，把kernel放到AiCPU上去执行。\r\n- AICPU调用函数（kernel），用一个线程执行kernel 函数，在函数内把通信task展开，把task放到rstqueue上，STARS调用。\r\n- 降低host和aicpu交互，由几百次降低为一次。\r\n- task的提交在AICPU上提交，做了提交的部分合并。\r\n\r\n#### TASK_QUEUE_ENABLE\r\n\r\n**使用方式：**`export TASK_QUEUE_ENABLE=2`\r\n\r\nTASK_QUEUE_ENABLE，下发优化，图模式设置为1（即开启图模式的时候这个要设置为1），非图模式设置为2\r\n\r\n示意图：\r\n\r\n![ascend task queue](https://github.com/wucong25/verl-data/blob/main/ascend_task_queue2.png)\r\n\r\n##### 绑核优化\r\n\r\n**使用方式：**`export CPU_AFFINITY_CONF=1`\r\n\r\n详细设置原理可看：https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0059.html\r\n\r\n### 其他\r\n\r\n以下内容汇总了若干全局环境变量的调优配置。由于这些参数在训练阶段与推理阶段往往都能带来正向收益，且目前尚缺乏足够精细的消融实验来严格区分它们各自对训练或推理的贡献占比，故统一归拢在此，供后续持续监控与进一步拆解分析。\r\n\r\n#### 使能jemalloc\r\n\r\n使用方式（注意需要先安装jemalloc库）：`export LD_PRELOAD=/usr/local/lib/libjemalloc.so.2`\r\n\r\n**安装使用教程：**[MindSpeed-RL/docs/install_guide.md · Ascend/MindSpeed-RL - AtomGit | GitCode](https://gitcode.com/Ascend/MindSpeed-RL/blob/master/docs/install_guide.md#高性能内存库-jemalloc-安装)\r\n\r\n#### 多流复用\r\n\r\n内存方面有优化\r\n\r\n使能方式：`export MULTI_STREAM_MEMORY_REUSE=1`\r\n\r\n原理介绍：https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0040.html\r\n\r\n#### VLLM_ASCEND_ENABLE_FLASHCOMM\r\n\r\n使用方式：`export VLLM_ASCEND_ENABLE_FLASHCOMM=1`\r\n\r\n启用昇腾 NPU 特有的FLASHCOMM高速通信优化技术\r\n\r\n地址：https://vllm-ascend.readthedocs.io/zh-cn/latest/user_guide/release_notes.html\r\n\r\n#### VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE\r\n\r\n使用方式：`export VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE=1`\r\n\r\n启用昇腾 NPU针对大模型推理的稠密计算优化\r\n\r\n地址：https://vllm-ascend.readthedocs.io/zh-cn/latest/user_guide/release_notes.html\r\n\r\n#### VLLM_ASCEND_ENABLE_PREFETCH_MLP\r\n\r\n使用方式：`export VLLM_ASCEND_ENABLE_PREFETCH_MLP=1`\r\n\r\n启用 MLP 层的权重预取机制\r\n\r\n<img src=\"https://github.com/wucong25/verl-data/blob/main/ascend_prefetch.png\" alt=\"image-20251124173132677\" style=\"zoom:50%;\" />\r\n\r\n### verl框架参数设置\r\n\r\n主要是内存方面的一些设置开关（注意，这个里面的优化都或多或少会导致吞吐量有一定程度的劣化）\r\n\r\n~~~bash\r\n# 梯度检查点 (Gradient Checkpointing)\r\n# 作用: 通过重新计算激活值来节省显存,以计算换内存。在前向传播时不保存中间激活值,反向传播时重新计算,可以显著降低显存占用,允许使用更大的batch size。\r\nactor_rollout_ref.model.enable_gradient_checkpointing=True\r\n\r\n# 参数卸载 (Parameter Offload)\r\n# 作用: 将模型参数卸载到CPU内存,训练时再加载回GPU。\r\nactor_rollout_ref.actor.fsdp_config.param_offload=${offload}  # True  \r\nactor_rollout_ref.ref.fsdp_config.param_offload=${offload}    # True\r\n\r\n# 优化器状态卸载 (Optimizer Offload)\r\n# 作用: 将优化器状态(如Adam的动量)卸载到CPU。优化器状态通常占用大量显存(对于Adam,每个参数需要额外8字节),卸载可以节省显存。\r\nactor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload}  # True\r\n\r\n# 释放推理引擎缓存 (Free Cache Engine)\r\n# 作用: 在训练阶段释放推理引擎的KV cache和权重。这是3D-HybridEngine的核心优化,允许在同一GPU上交替进行推理和训练,显著降低显存需求。\r\nactor_rollout_ref.rollout.free_cache_engine=True\r\n\r\n#  熵计算优化\r\n# entropy_checkpointing: 在训练时对熵计算启用重计算,降低显存峰值\r\n# entropy_from_logits_with_chunking: 分块处理logits张量(如2048 tokens一组),避免一次性加载整个[bsz*seq_len, vocab]张量\r\nactor_rollout_ref.actor.entropy_checkpointing=True  \r\nactor_rollout_ref.ref.entropy_checkpointing=True  \r\nactor_rollout_ref.actor.entropy_from_logits_with_chunking=True  \r\nactor_rollout_ref.ref.entropy_from_logits_with_chunking=True\r\n\r\n# 推理引擎显存配置\r\n# gpu_memory_utilization: 控制vLLM使用的GPU显存比例(0.90 = 90%)\r\n# enforce_eager=False: 启用CUDA graphs加速推理,但会占用额外显存\r\nactor_rollout_ref.rollout.gpu_memory_utilization=0.90  \r\nactor_rollout_ref.rollout.enforce_eager=False\r\n~~~\r\n\r\n## NPU调优参考文章\r\n\r\n环境变量相关：[环境变量列表-Ascend Extension for PyTorch6.0.0-昇腾社区](https://www.hiascend.com/document/detail/zh/Pytorch/600/apiref/Envvariables/Envir_001.html)\r\n\r\n社区性能调优教程：[性能调优流程-Ascend Extension for PyTorch6.0.0-昇腾社区](https://www.hiascend.com/document/detail/zh/Pytorch/600/ptmoddevg/trainingmigrguide/performance_tuning_0001.html)\r\n\r\n\r\n\r\n"
  },
  {
    "path": "docs/ascend_tutorial/examples/run_qwen3_32B_megatron_1k_256k_npu.md",
    "content": "# Long Sequence Qwen3-32B 1k-to-256k Example\n\nLast updated: 6/3/2026.\n\n本章对Qwen3-32B进行了长序列开发。Qwen3-32B的模型能力为最长推到40k\n\n## 全层实验\n\n对Qwen3-32B进行了长序列开发，脚本如下：\n\n```bash\nset -x\n\nexport USE_OPTIMIZED_MODEL=0\nexport VLLM_USE_V1=1\nexport VLLM_ASCEND_ENABLE_NZ=0\nexport VLLM_VERSION=\"0.13.0\"\nexport LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\nexport PYTORCH_NPU_ALLOC_CONF=\"max_split_size_mb:2048\"\n\nPROJECT_NAME=\"GRPO-Qwen3-32B\"\nEXPERIMENT_NAME=\"GRPO-Qwen3-32B-megatron-gsm8k\"\n\nSAVE_CHECKPOINT_DIR=$HOME/verl_checkpoints\nmath_train_path=$HOME/datasets/gsm8k/train.parquet\nmath_test_path=$HOME/datasets/gsm8k/test.parquet\ntrain_files=\"['$math_train_path']\"\ntest_files=\"['$math_test_path']\"\n\nuse_dynamic_bsz=False\nenable_chunked_prefill=True\ntp_size=8\nmax_prompt_length=1024\nmax_response_length=$((1024*256))\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / tp_size))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / tp_size))\ncp_size=4\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.shuffle=False \\\n    data.validation_shuffle=False \\\n    data.train_batch_size=64 \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.filter_overlong_prompts=False \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.path=$HOME/hf_weights/Qwen3-32B \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size=${CP} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=${enable_chunked_prefill} \\\n    actor_rollout_ref.rollout.enable_prefix_caching=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=False \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${SAVE_CHECKPOINT_DIR} \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=False \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${SAVE_CHECKPOINT_DIR} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=2 \\\n    trainer.save_freq=100 \\\n    trainer.test_freq=-1 \\\n    trainer.total_training_steps=100 \\\n    trainer.device=npu \\\n    trainer.project_name=${PROJECT_NAME} \\\n    trainer.experiment_name=${EXPERIMENT_NAME} \\\n    trainer.total_epochs=30\n```\n\n- 相关实验结果\n  \n![qwen3-32b-perfo](https://github.com/ChibiQuest/verl_data/blob/main/qwen3-32B-1k-256k/performance.png)\n\n## 减层实验\n\n在实际推理中，我们发现其最大在20k左右，因此对其进行减层实验，其response能到达到40k。\n\n在权重的`config.json`文件中，我们将`num_hidden_layers`从64减层到16\n\n```\n{\n  \"architectures\": [\n    \"Qwen3ForCausalLM\"\n  ],\n  \"attention_bias\": false,\n  \"attention_dropout\": 0.0,\n  \"bos_token_id\": 151643,\n  \"eos_token_id\": 151645,\n  \"head_dim\": 128,\n  \"hidden_act\": \"silu\",\n  \"hidden_size\": 5120,\n  \"initializer_range\": 0.02,\n  \"intermediate_size\": 25600,\n  \"max_position_embeddings\": 40960,\n  \"max_window_layers\": 64,\n  \"model_type\": \"qwen3\",\n  \"num_attention_heads\": 64,\n  \"num_hidden_layers\": 16,\n  \"num_key_value_heads\": 8,\n  \"rms_norm_eps\": 1e-06,\n  \"rope_scaling\": null,\n  \"rope_theta\": 1000000,\n  \"sliding_window\": null,\n  \"tie_word_embeddings\": false,\n  \"torch_dtype\": \"bfloat16\",\n  \"transformers_version\": \"4.51.0\",\n  \"use_cache\": true,\n  \"use_sliding_window\": false,\n  \"vocab_size\": 151936\n}\n\n```\n\n- 其实验结果如下：\n\n![qwen3-32b-function](https://github.com/ChibiQuest/verl_data/blob/main/qwen3-32B-1k-256k/function.png)"
  },
  {
    "path": "docs/ascend_tutorial/faq/faq.rst",
    "content": "Last updated: 03/16/2026.\r\n"
  },
  {
    "path": "docs/ascend_tutorial/features/ascend_backend_features.md",
    "content": "# Ascend Backend Features Guide\n==================================================================================\n\nLast updated: 03/03/2026.\n\n昇腾全面支持verl生态建设，本文将介绍NPU上对于verl的适配工作及后端特性支持供开发者进行参考\n\n---\n\n## 推理后端\n\n当前verl支持vllm/sglang这两种主流推理后端，均可在昇腾NPU上运行。\n\n### 1. vllm:\n\n昇腾通过vllm-ascend插件来支持vllm推理后端，该插件是 vLLM 社区支持 Ascend 后端的推荐方法。它遵循[[RFC]](https://github.com/vllm-project/vllm/issues/11162)，提供了一个可插拔接口，将 Ascend NPU 与 vLLM 解耦。\n\n##### 参数特性支持\n\n| vllm参数| verl对应通用参数 | 简介|\n| --- | --- | --- |\n| `model_path` | `actor_rollout_ref.model.path` |模型权重文件的路径|\n| `gpu_memory_utilization` | `actor_rollout_ref.rollout.gpu_memory_utilization` |用于控制每个阶段可使用的 GPU 内存量。它被指定为一个介于 0.0 和 1.0 之间的分数，其中：- 0.8 表示 GPU 总内存的 80%- 1.0 表示 GPU 总内存的 100%（不推荐，没有预留缓冲）|\n| `enforce_eager`| `actor_rollout_ref.rollout.enforce_eager` |禁用图模式，verl默认为False|\n| `enable_chunked_prefill`| `actor_rollout_ref.rollout.enable_chunked_prefill` | 分块预填充允许将大预填充分块成更小的块，并将它们与解码请求一起批处理。|\n| `free_cache_engine`| `actor_rollout_ref.rollout.free_cache_engine`  |在部署生成阶段之后卸载 KVCache，默认值为 True。|\n| `max_model_len` | `actor_rollout_ref.rollout.max_model_len` | 模型能够处理的最大序列长度。它限制了单个输入序列的最大长度 |\n| `tp_size`|  `actor_rollout_ref.rollout.tensor_model_parallel_size * data_parallel_size`|TP并行度|\n| `dp_size`| `actor_rollout_ref.rollout.data_parallel_size`|DP并行度|\n| `ep_size`| `actor_rollout_ref.rollout.expert_parallel_size`|EP并行度|\n| `node_rank`| `无，根据实际实例和卡数自动计算` |实例中的节点排序|\n| `load_format`|  `actor_rollout_ref.rollout.load_format` |要加载的模型权重格式|\n| `disable_log_stats`|  `actor_rollout_ref.rollout.disable_log_stats`|记录抢占请求的累积数量 |\n| `nnodes `|  `无，根据实际实例和卡数自动计算` | 每个实例包含的节点数量` |\n| `trust_remote_code`| `actor_rollout_ref.model.trust_remote_code`|是否允许在 Hub 上定义自定义模型，并将其写入自己的建模文件中|\n| `max_num_seqs` | `actor_rollout_ref.rollout.max_num_seqs` |正在运行的请求的最大数量|\n| `max_num_batched_tokens`| `actor_rollout_ref.rollout.max_num_batched_tokens` |在一次批处理（batch）中可以处理的最大总Token数|\n| `skip_tokenizer_init`| `actor_rollout_ref.rollout.skip_tokenizer_init` |跳过初始化分词器并将 input_ids 传递到推理请求中|\n| `enable_prefix_caching` | `actor_rollout_ref.rollout.enable_prefix_caching`|`用于启用自动前缀缓存` |\n| `quantization`| `actor_rollout_ref.rollout.quantization，默认为None`|`量化方法`|\n| `enforce_eager`|`actor_rollout_ref.rollout.enforce_eager`|标志用于强制使用PyTorch的eager执行模式，而非默认的图执行模式|\n\n### 2. sglang:\n\n对于sglang推理后端，昇腾通过直接向sglang社区进行持续建设与维护来支持相关功能。\n此外在verl中使用sglang还涉及以下组件, 我们在[quick start](https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/quick_start/ascend_sglang_quick_start.rst)中提供详细说明与一键安装脚本。\n\n| 组件| 描述|\n| --- | --- |\n| [sgl_kernel_npu](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/sgl_kernel_npu/README.md) | Ascend NPU  SGL 优化推理内核集合，包括注意力机制、归一化、激活函数、LoRA 适配器等。 |\n| [deepep](https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md) |  DeepEP的 Ascend 实现，为MoE模型提供高度优化的专家并行 (EP) 通信内核 |\n\n##### 参数特性支持\n\nverl中通过rollout config管理推理后端参数使能，包含通用参数和engine_kwargs自定义传参。\n以下列举在verl中常见设置的sglang特性参数，更多参数介绍请参考 [sglang社区NPU特性支持](https://docs.sglang.io/platforms/ascend_npu_support_features.html)\n\n| sglang参数| verl对应通用参数 | 简介|\n| --- | --- | --- |\n| model_path | actor_rollout_ref.model.path|模型权重文件的路径|\n| mem_fraction_static| actor_rollout_ref.rollout.gpu_memory_utilization |用于静态分配（模型权重和键值缓存内存池）的内存比例|\n| disable_cuda_graph| actor_rollout_ref.rollout.enforce_eager|禁用图模式，verl默认为False|\n| enable_memory_saver| 无，verl中默认设置为True | 允许使用 release_memory_occupation 和 resume_memory_occupation 来节省内存\n| base_gpu_id| 无，根据实际实例和卡数自动计算  |用于分配每个实例上计算卡资源时的的初始ID\n| gpu_id_step| 无，默认设置为1| 使用的连续计算卡ID 之间的差值\n| tp_size|  actor_rollout_ref.rollout.tensor_model_parallel_size * data_parallel_size|TP并行度|\n| dp_size| actor_rollout_ref.rollout.data_parallel_size|DP并行度|\n| ep_size| actor_rollout_ref.rollout.expert_parallel_size|EP并行度|\n| node_rank| 无，根据实际实例和卡数自动计算 |实例中的节点排序|\n| load_format|  actor_rollout_ref.rollout.load_format|要加载的模型权重格式|\n| dist_init_addr|  无，自动计算|用于初始化分布式后端的主机地址|\n| nnodes| 无，根据实际实例和卡数自动计算|每个实例包含的节点数量|\n| trust_remote_code| actor_rollout_ref.model.trust_remote_code|是否允许在 Hub 上定义自定义模型，并将其写入自己的建模文件中|\n| max_running_requests| actor_rollout_ref.rollout.max_num_seqs |正在运行的请求的最大数量|\n| log_level| 无，默认设置为error |日志记录器的日志级别|\n| skip_tokenizer_init| actor_rollout_ref.rollout.skip_tokenizer_init |跳过初始化分词器并将 input_ids 传递到推理请求中|\n| skip_server_warmup| 无，默认设置为True |跳过预热|\n| quantization| actor_rollout_ref.rollout.quantization，默认为None|量化方法|\n| attention_backend|actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend|attention内核,NPU应该设置为ascend|\n\n---\n\n## 训练后端\n\n### 1. FSDP\n\n昇腾通过torch_npu提供FSDP相关支持能力，当前pytorch api支持度参照[版本说明](https://www.hiascend.com/document/detail/zh/Pytorch/730/apiref/PyTorchNativeapi/docs/zh/native_apis/pytorch_2-7-1/torch-distributed-fsdp.md)。\n\n#### FSDP1\n##### 参数特性支持\n| verl参数 | 简介|\n| --- | --- |\n| `actor_rollout_ref.actor.fsdp_config.param_offload` |是否卸载模型权重到CPU，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.optimizer_offload` |是否卸载优化器状态到CPU，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` |控制前向计算后的参数行为，平衡内存与通信。默认值为True：前向后重新分片参数，反向时重新全收集|\n| `actor_rollout_ref.actor.fsdp_config.fsdp_size` | 每个FSDP分片组中的NPU数量；默认值-1表示自动。|\n| `actor_rollout_ref.actor.fsdp_config.forward_prefetch`  |在前向计算完成前预取下一次前向传播的 all-gather，仅用于FSDP1，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.use_orig_params` | FSDP是否会使用module的原始参数来初始化，仅用于FSDP1，默认值为False|\n| `actor_rollout_ref.actor.ulysses_sequence_parallel_size`|Ulysses序列并行大小|\n| `actor_rollout_ref.actor.entropy_from_logits_with_chunking`|通过分块计算熵以减少显存峰值，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.entropy_checkpointing`|在训练时对熵计算启用重计算,降低显存峰值，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.forward_only` |是否只进行前向计算，默认值为False|\n\n#### FSDP2\n##### 参数特性支持\n| verl参数 | 简介|\n| --- | --- |\n| `actor_rollout_ref.actor.fsdp_config.param_offload` |是否卸载模型权重到CPU，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.optimizer_offload` |是否卸载优化器状态到CPU，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.reshard_after_forward` |控制前向计算后的参数行为，平衡内存与通信。默认值为True：前向后重新分片参数，反向时重新全收集|\n| `actor_rollout_ref.actor.fsdp_config.fsdp_size` | 每个FSDP分片组中的NPU数量；默认值-1表示自动。|\n| `actor_rollout_ref.actor.ulysses_sequence_parallel_size`|Ulysses序列并行大小|\n| `actor_rollout_ref.actor.entropy_from_logits_with_chunking`|通过分块计算熵以减少显存峰值，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.entropy_checkpointing`|在训练时对熵计算启用重计算,降低显存峰值，默认值为False|\n| `actor_rollout_ref.actor.fsdp_config.forward_only` |是否只进行前向计算，默认值为False|\n\n\n\n### 2. Megatron\n\nMegatron 是 NVIDIA 推出的一个专注于模型并行的训练框架仓库。如果一个仓库（例如 Verl）的训练后端使用了 Megatron，同时又希望在 NPU 上运行该仓库，那么就需要额外安装 MindSpeed 来提供底层支持。下文将介绍 MindSpeed 是如何实现无感替换 Megatron 中的关键组件，从而使其能够适配 NPU 的。\n\nMindSpeed 底层的替换原理采用了 Monkey Patch 技术\n\n* MindSpeed Moneky Patch框架\n\n在verl里面通过`from mindspeed.megatron_adaptor import repatch  `触发patch，调用栈如下：\n\n~~~\nfrom mindspeed.megatron_adaptor import repatch\n├── 执行 megatron_adaptor.py 模块导入\n├── 导入 features_manager 模块\n├── 执行 mindspeed/features_manager/__init__.py  \n├── @AutoExecuteFunction 装饰器触发\n├── patch_features() 自动执行\n└── 进行`apply_features_pre_patches`和`apply_features_patches`操作\n~~~\n\n`Patch`类是整个patch系统的核心，实现了函数/类的动态替换\n\n~~~python\nclass Patch\n~~~\n\n`parse_path`方法实现了动态模块导入和创建\n\n~~~python\ndef parse_path(module_path, function_name, create_dummy)\n~~~\n\npatch系统支持多层装饰器叠加\n\n~~~\ndef apply_patch(self):  \n    final_patch_func = self.orig_func  \n    if self.patch_func is not None:  \n        final_patch_func = self.patch_func  \n\n    # 应用所有装饰器  \n    for wrapper in self.wrappers:  \n        final_patch_func = wrapper(final_patch_func)\n~~~\n\n* MindSpeedPatchesManager类\n\n`MindSpeedPatchesManager`作为全局单例管理所有patch\n\n~~~python\nclass MindSpeedPatchesManager:  \n    patches_info: Dict[str, Patch] = {}\n~~~\n\n* Feature集成模式\n\n各个Feature通过继承`MindSpeedFeature`基类集成patch系统\n\n~~~python\nclass MindSpeedFeature:\n    \"\"\"Base class for mindspeed features.\"\"\"\n\n    def __init__(self, feature_name: str, optimization_level: int = 2):\n        self.feature_name = feature_name.lower().strip().replace('-', '_')\n        self.optimization_level = optimization_level\n        self.default_patches = self.optimization_level == 0\n\n    def is_need_apply(self, args):\n        \"\"\"Check the feature is need to apply.\"\"\"\n        return (self.optimization_level <= args.optimization_level and getattr(args, self.feature_name, None)) \\\n            or self.default_patches\n\n    def register_args(self, parser: ArgumentParser):\n        \"\"\"Register cli arguments to enable the feature.\"\"\"\n        pass\n\n    def pre_validate_args(self, args: Namespace):\n        \"\"\"Validate the arguments of mindspeed before megatron args validation\n        and store some arguments of the mindspeed temporarily,\n        incase that megatron validate faile.\n        for example:\n            ```python\n            origin_context_parallel_size = args.context_parallel_size\n            args.context_parallel_size = 1\n            ```\n        \"\"\"\n        pass\n\n    def validate_args(self, args: Namespace):\n        \"\"\"Restore the arguments of the mindspeed.\n\n        for example:\n        ```python\n        args.context_parallel_size = origin_context_parallel_size\n        ```\n        \"\"\"\n        pass\n\n    def post_validate_args(self, args: Namespace):\n        \"\"\"validate mindspeed arguments after megatron arguments validation.\"\"\"\n        pass\n\n    def pre_register_patches(self, patch_manager: MindSpeedPatchesManager, args: Namespace):\n        \"\"\"Register all patch functions before import megatron\"\"\"\n        pass\n\n    def register_patches(self, patch_manager: MindSpeedPatchesManager, args: Namespace):\n        \"\"\"Register all patch functions the feature is related.\"\"\"\n        pass\n\n    def incompatible_check(self, global_args, check_args):\n        \"\"\"Register all incompatible functions the feature is related.\"\"\"\n        if getattr(global_args, self.feature_name, None) and getattr(global_args, check_args, None):\n            raise AssertionError('{} and {} are incompatible.'.format(self.feature_name, check_args))\n\n    def dependency_check(self, global_args, check_args):\n        \"\"\"Register all dependency functions the feature is related.\"\"\"\n        if getattr(global_args, self.feature_name, None) and not getattr(global_args, check_args, None):\n            raise AssertionError('{} requires {}.'.format(self.feature_name, check_args))\n\n    @staticmethod\n    def add_parser_argument_choices_value(parser, argument_name, new_choice):\n        \"\"\"Add a new choice value to the existing choices of a parser argument.\"\"\"\n        for action in parser._actions:\n            exist_arg = isinstance(action, argparse.Action) and argument_name in action.option_strings\n            if exist_arg and action.choices is not None and new_choice not in action.choices:\n                action.choices.append(new_choice)\n~~~\n\n##### 参数特性支持\n| verl参数 | 简介|\n| --- | --- |\n| `actor_rollout_ref.actor.megatron.optimizer_offload` |是否卸载模型优化器到CPU，默认值为False|\n| `actor_rollout_ref.actor.megatron.use_mbridge` |是否使用mbridge进行权重转换|\n| `actor_rollout_ref.actor.megatron.param_offload` |是否卸载模型权重到CPU，默认值为False|\n| `actor_rollout_ref.actor.megatron.tensor_model_parallel_size` | 张量并行大小；默认值为1。|\n| `actor_rollout_ref.actor.megatron.pipeline_model_parallel_size`  |流水并行大小，默认值为1|\n| `actor_rollout_ref.actor.megatron.expert_model_parallel_size` | 专家并行大小，默认值为1|\n| `actor_rollout_ref.actor.megatron.expert_tensor_parallel_size`|TP拓展EP大小，默认值为null|\n| `actor_rollout_ref.actor.context_parallel_size`|序列并行大小，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs`|张量在发送到下一个pp stage后,输出数据被释放，降低显存峰值，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm` |是否使用持久化 LayerNorm，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm` |是否使用持Group GEMM，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype` |用于路由和专家输出加权平均的数据类型。使用 fp32 或 fp64 可以提高稳定性，尤其是在专家数量较多时，默认值为fp32|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split` |如果设置为 True，在流水线并行的划分和放置策略中，loss 层会被视为一个标准的 Transformer 层来处理。默认为False。|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split` |如果设置为 True，在流水线并行的划分和放置策略中，输入embedding 层会被视为一个标准的 Transformer 层来处理。默认为False。|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity` |重新计算激活的粒度，可选项为'full', 'selective' and 'none'。其中full代表重新计算整个transformer layer，selective代表只计算transformer layer中的核心注意力部分。默认为'none'。|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method` |该参数需将recompute_granularity设置为'full'才生效，可选项为'uniform', 'block'。默认为None。|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers` |该参数需将recompute_granularity设置为'full'才生效，默认为None。若recompute_method设置为uniform，该参数含义为每个均匀划分的重新计算单元的transformer layers数量。例如你可以指定为--recompute_granularity full --recompute_method uniform --recompute_num_layers 4。recompute_num_layers越大，显存占用越小，计算成本越大。注意：当前进程中的模型层数需能被recompute_num_layers整除。默认为None。|\n| `actor_rollout_ref.actor.megatron.use_dist_checkpointing` |是否使用分布式权重，默认值为False|\n| `actor_rollout_ref.actor.megatron.dist_checkpointing_path` |分布式权重路径，默认值为null|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn` |是否使用fa，默认值为true|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb` |是否使用融合旋转位置编码，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu` |是否使用融合swiglu，默认值为False|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage` |第一个pipeline stage 的层数，默认值为none|\n| `actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage` |最后一个pipeline stage 的层数，默认值为none|\n"
  },
  {
    "path": "docs/ascend_tutorial/features/ascend_consistency.rst",
    "content": "推理一致性指导\n====================================\n\n在昇腾设备上对齐verl和vLLM两个框架下的推理结果。\n\nLast updated: 11/17/2025.\n\n这是一份在昇腾设备上对齐verl和vLLM两个框架下推理结果的教程。\n\n环境变量配置\n~~~~~~~~~~~~\n\n在多卡通信情况下：\n\n- HCCL通信下(默认场景):\n  \n  -  export CLOSE_MATMUL_K_SHIFT=1\n  -  export ATB_MATMUL_SHUFFLE_K_ENABLE=0\n  -  export HCCL_DETERMINISTIC=\"true\"\n  -  export VLLM_ENABLE_V1_MULTIPROCESSING=0\n\n- LCCL通信下(通过export HCCL_OP_EXPANSION_MODE=\"AIV\"使能）:\n\n  -  export CLOSE_MATMUL_K_SHIFT=1\n  -  export ATB_MATMUL_SHUFFLE_K_ENABLE=0\n  -  export LCCL_DETERMINISTIC=1\n  -  export ATB_LLM_LCOC_ENABLE=0\n  -  export VLLM_ENABLE_V1_MULTIPROCESSING=0\n\n在单卡无通信情况下：\n\n- HCCL和LCCL通信下:\n \n  -  export CLOSE_MATMUL_K_SHIFT=1\n  -  export ATB_MATMUL_SHUFFLE_K_ENABLE=0\n  -  export VLLM_ENABLE_V1_MULTIPROCESSING=0\n\nvLLM初始化参数\n~~~~~~~~~~~~\n\n需要对 SamplingParams 参数里单独设置seed, 保持vLLM和verl推理结果一致, 举例修改如下：\n\n.. code:: yaml\n\n      sampling_params = SamplingParams(n=1,\n                                       logprobs=0,  # can be set to 0 and let actor to recompute\n                                       max_tokens=config.response_length,\n                                       repetition_penalty=config.get(\"repetition_penalty\", 1.0),\n                                       seed=1234)\n\n"
  },
  {
    "path": "docs/ascend_tutorial/profiling/ascend_profiling_en.rst",
    "content": "Profiling Data Collection Guide\n==========================================================================================\n\nLast updated: 12/20/2025.\n\nThis is a tutorial for data collection using the GRPO or DAPO algorithm\nbased on FSDP or MindSpeed(Megatron) on Ascend devices.\n\nConfiguration\n-------------\n\nLeverage two levels of configuration to control data collection:\n\n- **Global profiler control**: Use parameters in ``verl/trainer/config/ppo_trainer.yaml`` (FSDP) or ``verl/trainer/config/ppo_megatron_trainer.yaml`` (MindSpeed) to control the collection mode and steps.\n- **Role profile control**: Use parameters in each role's ``profile`` field to control various parameters.\n\nGlobal collection control\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\nUse parameters in ppo_trainer.yaml to control the collection mode\nand steps.\n\n-  global_profiler: Control the ranks and mode of profiling\n\n   -  tool: The profiling tool to use, options are nsys, npu, torch,\n      torch_memory.\n   -  steps: This parameter can be set as a list that has\n      collection steps, such as [2, 4], which means it will collect steps 2\n      and 4. If set to null, no collection occurs.\n   -  save_path: The path to save the collected data. Default is\n      \"outputs/profile\".\n\n\nRole collection control\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nIn each role's ``profiler`` field, you can control the collection mode for that role.\n\n-  enable: Whether to enable profiling for this role.\n-  all_ranks: Whether to collect data from all ranks.\n-  ranks: A list of ranks to collect data from. If empty, no data is collected.\n-  tool_config: Configuration for the profiling tool used by this role.\n\nUse parameters in each role's ``profiler.tool_config.npu`` to control npu profiler behavior:\n\n-  level: Collection level—options are level_none, level0, level1, and\n   level2\n\n   -  level_none: Disables all level-based data collection (turns off profiler_level).\n   -  level0: Collect high-level application data, underlying NPU data, and operator execution details on NPU. After balancing data volume and analytical capability, Level 0 is recommended as the default configuration.\n   -  level1: Extends level0 by adding CANN-layer AscendCL data and AI Core performance metrics on NPU.\n   -  level2: Extends level1 by adding CANN-layer Runtime data and AI CPU metrics.\n\n-  contents: A list of options to control the collection content, such as\n   npu, cpu, memory, shapes, module, stack.\n   \n   -  npu: Whether to collect device-side performance data.\n   -  cpu: Whether to collect host-side performance data.\n   -  memory: Whether to enable memory analysis.\n   -  shapes: Whether to record tensor shapes.\n   -  module: Whether to record framework-layer Python call stack information. It is recommended to use 'module' instead of 'stack' for recording call stack information, as it costs less performance overhead.\n   -  stack: Whether to record operator call stack information.\n\n-  analysis: Enables automatic data parsing.\n-  discrete: Whether to enable discrete mode.\n\n\nExamples\n--------\n\nDisabling collection\n~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: null # disable profile\n\nEnd-to-End collection\n~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: [1, 2, 5]\n         save_path: ./outputs/profile\n      actor_rollout_ref:\n         actor:  # Set actor role profiler collection configuration parameters\n            profiler:\n               enable: True\n               all_ranks: True\n               tool_config:\n                  npu:\n                     discrete: False\n                     contents: [npu, cpu]  # Control collection list, default cpu, npu, can configure memory, shapes, module, etc.\n        # rollout & ref follow actor settings\n\n\nDiscrete Mode Collection\n~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: [1, 2, 5]\n         save_path: ./outputs/profile\n      actor_rollout_ref:\n         actor:\n            profiler:\n               enable: True  # Set to True to profile training\n               all_ranks: False\n               ranks: [0]  # Global Rank 0\n               tool_config:\n                  npu:\n                     discrete: True\n                     contents: [npu, cpu]\n         rollout:\n            profiler:\n               enable: True  # Set to True to profile inference\n               all_ranks: False\n               ranks: [0]  # In Agent Loop mode, this is the Replica Rank (e.g., 0-th instance)\n               tool_config:\n                  npu:\n                     discrete: True  # Must be enabled in Agent Loop mode\n         # ref follow actor settings\n\n**Agent Loop Mode Description**:\n\nWhen Rollout runs in `Agent Loop <../advance/agent_loop.rst>`_ mode, performance data for the Rollout phase **must be collected using discrete mode**. In this case, the Profiler is triggered by the inference engine backend.\n\n1. Rank Definition: ranks in the Rollout configuration refers to Replica Rank (inference instance index), not Global Rank.\n\n2. Inference Engine Support: Currently, vLLM and SGLang engines are supported without additional settings. Specific details are as follows:\n\n   - vLLM Engine: Automatically collects AsyncLLM scheduling stacks and inference process performance data. Does not support setting analysis (defaults to no analysis, requires offline analysis) and profiler_level (defaults to level1).\n   - SGLang Engine: Automatically collects inference process performance data. Does not support the memory option in contents. Does not support setting analysis (defaults to enabled) and profiler_level (defaults to level0).\n\n\nVisualization\n-------------\n\nCollected data is stored in the user-defined save_path and can be\nvisualized by using the `MindStudio Insight <https://www.hiascend.com/document/detail/zh/mindstudio/80RC1/GUI_baseddevelopmenttool/msascendinsightug/Insight_userguide_0002.html>`_ tool.\n\nAdditionally, in a Linux environment, the MindStudio Insight tool is provided in the form of a `JupyterLab Plugin <https://www.hiascend.com/document/detail/zh/mindstudio/82RC1/GUI_baseddevelopmenttool/msascendinsightug/Insight_userguide_0130.html>`_ ，offering a more intuitive and highly interactive user interface. The advantages of the JupyterLab plugin are as follows:\n\n- Seamless integration: Supports running the MindStudio Insight tool directly within the Jupyter environment, eliminating the need to switch platforms or copy data from the server, enabling data to be collected and used immediately.\n- Fast startup: Allows MindStudio Insight to be launched quickly via the JupyterLab command line or graphical interface.\n- Smooth operation: In a Linux environment, launching MindStudio Insight through JupyterLab effectively alleviates performance lag compared to the full-package communication mode, significantly improving the user experience.\n- Remote access: Supports remotely launching MindStudio Insight. Users can connect to the service via a local browser for direct visual analysis, reducing the difficulty of uploading and downloading data during large-model training or inference.\n\nIf the analysis parameter is set to False, offline parsing is required after data collection:\n\n.. code:: python\n\n    import torch_npu\n    # Set profiler_path to the parent directory of the \"localhost.localdomain_<PID>_<timestamp>_ascend_pt\" folder\n    torch_npu.profiler.profiler.analyse(profiler_path=profiler_path)\n\n\nAdvanced Guide: Fine-grained Collection\n---------------------------------------\n\nBackground and Challenges\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\nAlthough the configuration-based collection method mentioned above is convenient, it faces challenges in training scenarios with **long sequences (Long Context)** or **large global batch sizes (Large Global Batch Size)**. Within a complete training step (Step), model computation exhibits high-frequency and repetitive characteristics:\n\n1. **Rollout phase**: Sequence generation (Generate Sequence) is an autoregressive process involving thousands of forward computations of the Decoder model.\n2. **Training phase**: To control peak memory usage, verl typically adopts a Micro-Batch strategy, dividing large data streams into multiple micro-batches for computation.\n\n   - **compute_log_prob (Actor/Ref)**: Involves multiple rounds of pure forward propagation.\n   - **update_policy (Actor/Critic)**: Involves multiple rounds of forward and backward propagation.\n\nThis characteristic leads to massive and repetitive operator records from full profiling. As shown in the image below:\n\n.. image:: https://raw.githubusercontent.com/mengchengTang/verl-data/master/verl_ascend_profiler.png\n\nEven with ``discrete`` mode enabled, performance data files for a single stage can still reach several TB, leading to **parsing failures** or **visualization tool lag**.\n\nSolution: Critical Path Sampling\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nTo solve the above problems, we can adopt a **critical path sampling** strategy: Based on the API interface provided by `torch_npu.profiler <https://www.hiascend.com/document/detail/zh/canncommercial/80RC2/devaids/auxiliarydevtool/atlasprofiling_16_0038.html>`_, directly modify Python source code to collect only representative data segments (such as specific Decode Steps or the first Micro-Batch).\n\n    **Important Notes**\n\n    1. This chapter involves direct source code modification. It is recommended to back up files before modification and restore them after debugging.\n    2. When using code instrumentation for collection, be sure to **disable global collection** (``global_profiler: steps: null``) in ``ppo_trainer.yaml`` or ``ppo_megatron_trainer.yaml`` to avoid Profiler conflicts.\n\n1. Fine-grained Collection in Rollout Phase\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nFor vLLM or SGLang inference engines, we can control the ``schedule`` parameter to collect model forward propagation performance data for specific tokens.\n\n**vLLM Engine**\n\n- **Reference Version**: vLLM v0.11.0, vLLM-Ascend v0.11.0rc1\n- **Modified File**: ``vllm-ascend/vllm_ascend/worker/worker_v1.py``\n\n.. code-block:: diff\n\n      class NPUWorker(WorkerBase):\n\n          def __init__(self, *args, **kwargs):\n              # ... existing code ...\n\n  +           # Initialize profiler\n  +           import torch_npu\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(\n  +               profiler_level=torch_npu.profiler.ProfilerLevel.Level1,\n  +               export_type=torch_npu.profiler.ExportType.Db,  # You can choose torch_npu.profiler.ExportType.Text format\n  +           )\n  +           self.profiler_npu = torch_npu.profiler.profile(\n  +               activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU],\n  +               with_modules=False,  # Collect call stack\n  +               profile_memory=False,  # Collect memory\n  +               experimental_config=experimental_config,\n  +               # Skip first step, warmup one step, collect 3 steps, repeat 1 time. If you want to collect decode steps 30~70, set schedule=torch_npu.profiler.schedule(wait=29, warmup=1, active=30, repeat=1)\n  +               schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/vllm_profile\", analyse_flag=True)  # Data save path and whether to parse online\n  +           )\n  +           self.profiler_npu.start()\n\n              # ... existing code ...\n\n          def execute_model(self, scheduler_output=None, intermediate_tensors=None, **kwargs):\n              # ... existing code ...\n              output = self.model_runner.execute_model(scheduler_output,\n                                                  intermediate_tensors)\n\n  +           self.profiler_npu.step()  # Drive schedule to collect partial decode steps\n\n              # ... existing code ...\n\n**SGLang Engine**\n\n- **Reference Version**: SGLang master branch\n- **Modified File**: ``sglang/python/sglang/srt/model_executor/model_runner.py``\n\n.. code-block:: diff\n\n      # ... existing imports ...\n  +   import torch_npu\n\n      class ModelRunner:\n\n          def __init__(self, *args, **kwargs):\n              # ... existing init code ...\n\n  +           # Initialize profiler (same configuration as above, omitted)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.profiler_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # Skip first step, warmup one step, collect 3 steps, repeat 1 time.\n  +               schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/sglang_profile\", analyse_flag=True)\n  +           )\n  +           self.profiler_npu.start()\n\n          def forward(self, forward_batch, **kwargs):\n              # ... existing code ...\n\n  +           self.profiler_npu.step()  # Drive schedule to collect partial decode steps\n              return output\n\n2. Fine-grained Collection in compute_log_prob (Actor & Ref) Phase\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThis phase computes probability distributions for new and old policies.\n\n**FSDP Backend**\n\nThe FSDP backend allows fine-grained control at the Micro-Batch level.\n\n- **Modified File**: ``verl/workers/actor/dp_actor.py``\n\n.. code-block:: diff\n\n      # ... import dependencies ...\n  +   import torch_npu\n\n      class DataParallelPPOActor(BasePPOActor):\n\n          def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n\n  +           role = \"Ref\" if self.actor_optimizer is None else \"Actor\"\n  +           # Prepare profiler (same configuration as above, omitted)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # wait=0, warmup=0, active=1: directly collect first micro-batch\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f\"./outputs/{role}_compute_log_prob\", analyse_flag=True)\n  +           )\n\n\n  +           # This function is shared by ref and actor, set role flag to distinguish. If you want to collect actor_compute_log_prob, set if role==\"Actor\":\n  +           if role==\"Ref\":\n  +               self.prof_npu.start()\n\n              for micro_batch in micro_batches:\n\n                  # ... original computation logic ...\n                  with torch.no_grad():\n                      entropy, log_probs = self._forward_micro_batch(...)\n\n  +                   # Drive schedule to collect micro batch\n  +                   if role==\"Ref\":\n  +                       self.prof_npu.step()\n\n                  # ...\n\n\n**Megatron Backend**\n\nThe Micro-Batch scheduling in the Megatron backend is managed internally by the framework and does not currently support fine-grained collection at the Micro-Batch level through simple code instrumentation. It is recommended to use global configuration for collection.\n\n3. Fine-grained Collection in update_policy (Actor & Critic) Phase\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThe Update phase includes forward and backward propagation.\n\n**FSDP Backend**\n\nThe FSDP backend supports collection at both Mini-Batch and Micro-Batch granularities.\n\n- **Modified File**: ``verl/workers/actor/dp_actor.py``\n\n.. code-block:: diff\n\n      # ... import dependencies ...\n  +   import torch_npu\n\n      class DataParallelPPOActor(BasePPOActor):\n\n          def update_policy(self, data: DataProto):\n\n  +           # Prepare profiler (same configuration as above, omitted)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # Only collect first Mini Batch (including all Micro-Batch computations and one optimizer update)\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/fsdp_actor_update_profile\", analyse_flag=True)\n  +           )\n  +           self.prof_npu.start()\n\n              # ... PPO Epochs loop ...\n              for _ in range(self.config.ppo_epochs):\n                  # ... Mini Batch loop ...\n                  for batch_idx, mini_batch in enumerate(mini_batches):\n                      # ... mini_batches split ...\n\n                      for i, micro_batch in enumerate(micro_batches):\n                          # ... Original Forward & Backward logic ...\n                          # ... loss.backward() ...\n                          pass\n\n                      grad_norm = self._optimizer_step()\n\n  +                   # Drive schedule to collect mini batch, if you want micro batch collection, move self.prof_npu.step() inside the micro_batch loop\n  +                   self.prof_npu.step()\n\n\n**Megatron Backend**\n\nThe Megatron backend supports collection at the Mini-Batch granularity.\n\n- **Modified File**: ``verl/workers/actor/megatron_actor.py``\n\n.. code-block:: diff\n\n      class MegatronPPOActor(BasePPOActor):\n\n          def update_policy(self, dataloader: Iterable[DataProto]) -> dict:\n              # ...\n  +           # Prepare profiler (same configuration as above, omitted)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # Only collect computation of first Mini Batch (including all Micro-Batches) and one optimizer update\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/megatron_actor_update_profile\", analyse_flag=True)\n  +           )\n  +           self.prof_npu.start()\n\n              for data in dataloader:\n                  # ... internally calls self.forward_backward_batch for computation ...\n                  # ... metric_micro_batch = self.forward_backward_batch(...)\n\n                  # ... self.actor_optimizer.step() ...\n\n  +               # Drive schedule to collect mini batch\n  +               self.prof_npu.step()"
  },
  {
    "path": "docs/ascend_tutorial/profiling/ascend_profiling_zh.rst",
    "content": "Profiling采集指导\n==================================================================================\n\nLast updated: 12/20/2025.\n\n这是一份在昇腾设备上基于FSDP或MindSpeed(Megatron)后端，使用GRPO或DAPO算法进行数据采集的教程。\n\n配置\n----\n\n使用两级profile设置来控制数据采集\n\n- 全局采集控制：使用verl/trainer/config/ppo_trainer.yaml(FSDP)，或verl/trainer/config/ppo_megatron_trainer.yaml(MindSpeed)中的配置项控制采集的模式和步数。\n- 角色profile控制：通过每个角色中的配置项控制等参数。\n\n全局采集控制\n~~~~~~~~~~~~\n\n通过 ppo_trainer.yaml 中的参数控制采集步数和模式：\n\n-  global_profiler: 控制采集的rank和模式\n\n   -  tool: 使用的采集工具，选项有 nsys、npu、torch、torch_memory。\n   -  steps: 此参数可以设置为包含采集步数的列表，例如 [2, 4]，表示将采集第2步和第4步。如果设置为 null，则不进行采集。\n   -  save_path: 保存采集数据的路径。默认值为 \"outputs/profile\"。\n\n角色profiler控制\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n在每个角色的 ``profiler`` 字段中，您可以控制该角色的采集模式。\n\n-  enable: 是否为此角色启用性能分析。\n-  all_ranks: 是否从所有rank收集数据。\n-  ranks: 要收集数据的rank列表。如果为空，则不收集数据。\n-  tool_config: 此角色使用的性能分析工具的配置。\n\n通过每个角色的 ``profiler.tool_config.npu`` 中的参数控制具体采集行为：\n\n-  level: 采集级别—选项有 level_none、level0、level1 和 level2\n\n   -  level_none: 禁用所有基于级别的数据采集（关闭 profiler_level）。\n   -  level0: 采集高级应用数据、底层NPU数据和NPU上的算子执行详情。在权衡数据量和分析能力后，level0是推荐的默认配置。\n   -  level1: 在level0基础上增加CANN层AscendCL数据和NPU上的AI Core性能指标。\n   -  level2: 在level1基础上增加CANN层Runtime数据和AI CPU指标。\n\n-  contents: 控制采集内容的选项列表，例如\n   npu、cpu、memory、shapes、module、stack。\n   \n   -  npu: 是否采集设备端性能数据。\n   -  cpu: 是否采集主机端性能数据。\n   -  memory: 是否启用内存分析。\n   -  shapes: 是否记录张量形状。\n   -  module: 是否记录框架层Python调用栈信息。相较于stack，更推荐使用module记录调用栈信息，因其产生的性能膨胀更低。\n   -  stack: 是否记录算子调用栈信息。\n\n-  analysis: 启用自动数据解析。\n-  discrete: 使用离散模式。\n\n示例\n----\n\n禁用采集\n~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: null # disable profile\n\n端到端采集\n~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: [1, 2, 5]\n         save_path: ./outputs/profile\n      actor_rollout_ref:\n         actor:  # 设置 actor role 的 profiler 采集配置参数\n            profiler:\n               enable: True\n               all_ranks: True\n               tool_config:\n                  npu:\n                     discrete: False\n                     contents: [npu, cpu]  # 控制采集列表，默认cpu、npu，可配置memory、shapes、module等\n\n        # rollout & ref follow actor settings\n\n\n离散模式采集\n~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n      global_profiler:\n         steps: [1, 2, 5]\n         save_path: ./outputs/profile\n      actor_rollout_ref:\n         actor:\n            profiler:\n               enable: True  # 设置为 True 以采集训练阶段\n               all_ranks: False\n               ranks: [0]  # 全局 Rank 0\n               tool_config:\n                  npu:\n                     discrete: True\n                     contents: [npu, cpu]\n         rollout:\n            profiler:\n               enable: True  # 设置为 True 以采集推理阶段\n               all_ranks: False\n               ranks: [0]  # 在 Agent Loop 模式下，此处指推理实例的 Replica Rank (例如第 0 个实例)\n               tool_config:\n                  npu:\n                     discrete: True  # Agent Loop 模式下必须开启离散模式\n         # ref follow actor settings\n\n**Agent Loop 模式说明**：\n\n在 `Agent Loop <../advance/agent_loop.rst>`_ 模式下，Rollout 阶段的性能数据 **必须使用离散模式** 采集，此时 Profiler 由推理引擎后端触发。\n\n1. Rank 定义：Rollout 配置中的 ranks 指代 Replica Rank（推理实例索引），而非全局 Rank。\n\n2. 推理引擎支持：当前支持vLLM和SGLang引擎，无需额外设置。具体说明如下：\n\n   - vLLM 引擎：自动采集 AsyncLLM 调度栈及推理进程性能数据。不支持设置 analysis（默认不解析，需离线解析）和 profiler_level（默认 level1）。\n   - SGLang 引擎：自动采集推理进程性能数据。不支持 contents 中的 memory 配置项。不支持设置 analysis（默认解析）和 profiler_level（默认 level0）。\n\n可视化\n------\n\n采集后的数据存放在用户设置的save_path下，可通过 `MindStudio Insight <https://www.hiascend.com/document/detail/zh/mindstudio/80RC1/GUI_baseddevelopmenttool/msascendinsightug/Insight_userguide_0002.html>`_ 工具进行可视化。\n\n另外在Linux环境下，MindStudio Insight工具提供了 `JupyterLab插件 <https://www.hiascend.com/document/detail/zh/mindstudio/82RC1/GUI_baseddevelopmenttool/msascendinsightug/Insight_userguide_0130.html>`_ 形态，提供更直观和交互式强的操作界面。JupyterLab插件优势如下：\n\n- 无缝集成：支持在Jupyter环境中直接运行MindStudio Insight工具，无需切换平台，无需拷贝服务器上的数据，实现数据即采即用。\n- 快速启动：通过JupyterLab的命令行或图形界面，可快速启动MindStudio Insight工具。\n- 运行流畅：在Linux环境下，通过JupyterLab环境启动MindStudio Insight，相较于整包通信，有效解决了运行卡顿问题，操作体验显著提升。\n- 远程访问：支持远程启动MindStudio Insight，可通过本地浏览器远程连接服务直接进行可视化分析，缓解了大模型训练或推理数据上传和下载的困难。\n\n如果analysis参数设置为False，采集之后需要进行离线解析：\n\n.. code:: python\n\n    import torch_npu\n    # profiler_path请设置为\"localhost.localdomain_<PID>_<timestamp>_ascend_pt\"目录的上一级目录\n    torch_npu.profiler.profiler.analyse(profiler_path=profiler_path)\n\n\n进阶指南：精细化采集\n--------------------\n\n背景与挑战\n~~~~~~~~~~\n\n上述基于配置文件的采集方式虽然便捷，但在 **长序列 (Long Context)** 或 **大全局批量 (Large Global Batch Size)** 的训练场景中面临挑战。\n在一个完整的训练步 (Step) 内，模型计算呈现出高频次、重复性的特征：\n\n1. Rollout 阶段：序列生成 (Generate Sequence) 是一个自回归过程，涉及成千上万次 Decoder 模型的前向计算。\n2. Training 阶段：为了控制显存峰值，verl 通常采用 Micro-Batch 策略，将庞大的数据流切分为多个微批次进行计算。\n   \n  - compute_log_prob (Actor/Ref)：涉及多轮纯前向传播。\n  - update_policy (Actor/Critic)：涉及多轮前向与反向传播。\n\n这种特性会导致全量 Profiling 产生海量且重复的算子记录。如下图所示：\n\n.. image:: https://raw.githubusercontent.com/mengchengTang/verl-data/master/verl_ascend_profiler.png\n\n即使使用了 ``discrete`` 模式，单个阶段的性能数据文件仍可能达到数 TB，导致 **解析失败** 或 **可视化工具卡顿** 。\n\n解决方案：关键路径采样\n~~~~~~~~~~~~~~~~~~~~~~\n\n为了解决上述问题，我们可以采用 **关键路径采样** 策略：基于 `torch_npu.profiler <https://www.hiascend.com/document/detail/zh/canncommercial/80RC2/devaids/auxiliarydevtool/atlasprofiling_16_0038.html>`_ 提供的API接口，直接修改 Python 源码，仅采集具有代表性的数据片段（如特定 Decode Step 或首个 Micro-Batch）。\n\n    **重要提示**\n\n    1. 本章节涉及直接修改源码。建议修改前备份文件，调试完成后恢复。\n    2. 使用代码插桩采集时，请务必在 ``ppo_trainer.yaml`` 或 ``ppo_megatron_trainer.yaml`` 中**禁用全局采集** (``global_profiler: steps: null``)，以避免 Profiler 冲突。\n\n1. Rollout 阶段精细化采集\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\n对于 vLLM 或 SGLang 推理引擎，我们可以通过控制 ``schedule`` 参数来控制采集模型在特定token的前向传播性能数据。\n\n**vLLM 引擎**\n\n- **参考版本**：vLLM v0.11.0, vLLM-Ascend v0.11.0rc1\n- **修改文件**：``vllm-ascend/vllm_ascend/worker/worker_v1.py``\n\n.. code-block:: diff\n\n      class NPUWorker(WorkerBase):\n  \n          def __init__(self, *args, **kwargs):\n              # ... existing code ...\n  \n  +           # Initialize profiler\n  +           import torch_npu\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(\n  +               profiler_level=torch_npu.profiler.ProfilerLevel.Level1,\n  +               export_type=torch_npu.profiler.ExportType.Db,  # 可选择torch_npu.profiler.ExportType.Text格式\n  +           )\n  +           self.profiler_npu = torch_npu.profiler.profile(\n  +               activities=[torch_npu.profiler.ProfilerActivity.CPU, torch_npu.profiler.ProfilerActivity.NPU],\n  +               with_modules=False,  # 采集调用栈\n  +               profile_memory=False,  # 采集内存\n  +               experimental_config=experimental_config,\n  +               # 跳过第一步，warmup一步，采集3步，重复1次。如果想采集第30~70个decode step，可以设置为schedule=torch_npu.profiler.schedule(wait=29, warmup=1, active=30, repeat=1)\n  +               schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/vllm_profile\", analyse_flag=True)  # 采集数据保存路径，是否在线解析\n  +           )\n  +           self.profiler_npu.start()\n  \n              # ... existing code ...\n  \n          def execute_model(self, scheduler_output=None, intermediate_tensors=None, **kwargs):\n              # ... existing code ...\n              output = self.model_runner.execute_model(scheduler_output,\n                                                  intermediate_tensors)\n              \n  +           self.profiler_npu.step()  # 驱动 schedule，对部分decode step进行采集\n              \n              # ... existing code ...\n\n**SGLang 引擎**\n\n- **参考版本**：SGLang master 分支\n- **修改文件**：``sglang/python/sglang/srt/model_executor/model_runner.py``\n\n.. code-block:: diff\n\n      # ... existing imports ...\n  +   import torch_npu\n  \n      class ModelRunner:\n  \n          def __init__(self, *args, **kwargs):\n              # ... existing init code ...\n              \n  +           # Initialize profiler (配置同上，略)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.profiler_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # 跳过第一步，warmup一步，采集3步，重复1次。\n  +               schedule=torch_npu.profiler.schedule(wait=1, warmup=1, active=3, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/sglang_profile\", analyse_flag=True)\n  +           )\n  +           self.profiler_npu.start()\n  \n          def forward(self, forward_batch, **kwargs):\n              # ... existing code ...\n\n  +           self.profiler_npu.step()  # 驱动 schedule，对部分decode step进行采集\n              return output\n\n2. compute_log_prob (Actor & Ref) 阶段精细化采集\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n该阶段计算新旧策略的概率分布。\n\n**FSDP 后端**\n\nFSDP 后端允许在 Micro-Batch 级别进行精细控制。\n\n- **修改文件**：``verl/workers/actor/dp_actor.py``\n\n.. code-block:: diff\n\n      # ... 引入依赖 ...\n  +   import torch_npu\n  \n      class DataParallelPPOActor(BasePPOActor):\n  \n          def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n\n  +           role = \"Ref\" if self.actor_optimizer is None else \"Actor\"\n  +           # 准备 profiler (配置同上，略)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # wait=0, warmup=0, active=1: 直接采集第一个 micro-batch\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(f\"./outputs/{role}_compute_log_prob\", analyse_flag=True)\n  +           )\n\n\n  +           # 此函数ref和actor共用，设置role标志位来区分。如果想采集actor_compute_log_prob，可设置if role==\"Actor\":\n  +           if role==\"Ref\":\n  +               self.prof_npu.start()\n  \n              for micro_batch in micro_batches:\n  \n                  # ... 原始计算逻辑 ...\n                  with torch.no_grad():\n                      entropy, log_probs = self._forward_micro_batch(...)\n                      \n  +                   # 驱动 schedule，对micro batch进行采集\n  +                   if role==\"Ref\":\n  +                       self.prof_npu.step()\n                  \n                  # ...\n\n\n**Megatron 后端**\n\nMegatron 后端的 Micro-Batch 调度由框架内部管理，暂不支持通过简单的代码插桩进行 Micro-Batch 级别的精细化采集。建议使用全局配置进行采集。\n\n3. update_policy (Actor & Critic) 阶段精细化采集\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nUpdate 阶段包含前向和反向传播。\n\n**FSDP 后端**\n\nFSDP 后端支持设置对 Mini-Batch 和 Micro-Batch 的粒度进行采集。\n\n- **修改文件**：``verl/workers/actor/dp_actor.py``\n\n.. code-block:: diff\n\n      # ... 引入依赖 ...\n  +   import torch_npu\n  \n      class DataParallelPPOActor(BasePPOActor):\n  \n          def update_policy(self, data: DataProto):\n              \n  +           # 准备 profiler (配置同上，略)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # 仅采集第一个 Mini Batch（包含所有 Micro-Batch 的计算和一次优化器更新）\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/fsdp_actor_update_profile\", analyse_flag=True)\n  +           )\n  +           self.prof_npu.start()\n              \n              # ... PPO Epochs 循环 ...\n              for _ in range(self.config.ppo_epochs):\n                  # ... Mini Batch 循环 ...\n                  for batch_idx, mini_batch in enumerate(mini_batches):\n                      # ... mini_batches 切分 ...\n  \n                      for i, micro_batch in enumerate(micro_batches):\n                          # ... 原始 Forward & Backward 逻辑 ...\n                          # ... loss.backward() ...\n                          pass\n      \n                      grad_norm = self._optimizer_step()\n                      \n  +                   # 驱动 schedule，对mini batch进行采集，如果想对micro batch进行，则将self.prof_npu.step()移动到micro_batch的循环内\n  +                   self.prof_npu.step()\n  \n\n**Megatron 后端**\n\nMegatron 后端支持以 Mini-Batch 的粒度进行采集。\n\n- **修改文件**：``verl/workers/actor/megatron_actor.py``\n\n.. code-block:: diff\n\n      class MegatronPPOActor(BasePPOActor):\n          \n          def update_policy(self, dataloader: Iterable[DataProto]) -> dict:\n              # ...\n  +           # 准备 profiler (配置同上，略)\n  +           experimental_config = torch_npu.profiler._ExperimentalConfig(...)\n  +           self.prof_npu = torch_npu.profiler.profile(\n  +               # ...\n  +               # 仅采集第一个 Mini Batch 的计算（含所有 Micro-Batch）和一次优化器更新\n  +               schedule=torch_npu.profiler.schedule(wait=0, warmup=0, active=1, repeat=1),\n  +               on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(\"./outputs/megatron_actor_update_profile\", analyse_flag=True)\n  +           )\n  +           self.prof_npu.start()\n              \n              for data in dataloader:\n                  # ... 内部会调用 self.forward_backward_batch 进行计算 ...\n                  # ... metric_micro_batch = self.forward_backward_batch(...)\n                  \n                  # ... self.actor_optimizer.step() ...\n                  \n  +               # 驱动 schedule，对mini batch进行采集\n  +               self.prof_npu.step()\n"
  },
  {
    "path": "docs/ascend_tutorial/quick_start/ascend_quick_start.rst",
    "content": "Ascend Quickstart\n===================================\n\nLast updated: 03/03/2026.\n\n\n\n关键更新\n----------------------------------\n\n2025/12/11：verl 存量场景目前支持自动识别 NPU 设备类型， GPU 脚本在昇腾上运行，原则上不再需要显式设置 trainer.device=npu 参数，新增特性通过设置 trainer.device 仍可优先使用，逐步适配自动识别能力。\n\n    [说明] 自动识别 NPU 设备类型的前提，是运行程序所在环境包含 torch_npu 软件包。如不包含该软件包，仍需显式指定 trainer.device=npu 参数。\n\n硬件支持\n-----------------------------------\n\nAtlas 200T A2 Box16\n\nAtlas 900 A2 PODc\n\nAtlas 800T A3\n\n\n安装流程\n-----------------------------------\n\n\nDockerFile镜像构建 & 获取 & 使用 \n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n如需要通过 DockerFile 构建镜像，或希望使用基于 verl 构建的镜像，请参考 `文档 <https://github.com/volcengine/verl/tree/main/docs/ascend_tutorial/quick_start/dockerfile_build_guidance.rst>`_ \n如果想直接获取镜像，请前往`quay.io/ascend/verl <https://quay.io/repository/ascend/verl?tab=tags&tag=latest>`_ 进行获取，镜像中已包含基础环境和依赖软件包。\n\n安装基础环境\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. 基础环境涉及以下软件包，请参考 `文档 <https://gitcode.com/Ascend/pytorch>`_ 安装。\n\n    +---------------+----------------------+\n    | software      | version              |\n    +---------------+----------------------+\n    | Python        | >= 3.10, <3.12       |\n    +---------------+----------------------+\n    | CANN          | == 8.5.0             |\n    +---------------+----------------------+\n    | torch         | == 2.8.0             |\n    +---------------+----------------------+\n    | torch_npu     | == 2.8.0             |\n    +---------------+----------------------+\n\n2. （可选）在 x86 平台安装时，pip 需要配置额外的源，指令如下：\n\n    .. code-block:: bash\n\n        pip config set global.extra-index-url \"https://download.pytorch.org/whl/cpu/\"\n\n\n安装其他软件包\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n基础环境准备完毕后，需要通过指令安装以下软件包：\n\n    +---------------+----------------------+\n    | torchvision   | == 0.22.1            |\n    +---------------+----------------------+\n    | triton-ascend | == 3.2.0             |\n    +---------------+----------------------+\n    | transformers  | == 4.57.6            |\n    +---------------+----------------------+\n    \n    tips: verl is not support transformers 5.0.0 or higher\n    安装指令：\n    \n    .. code-block:: bash\n    \n        # 安装torchvision，版本需要和torch匹配\n        pip install torchvision==0.22.1\n    \n        # 清理环境上可能存在的历史triton/triton-ascend软件包残留\n        pip uninstall -y triton triton-ascend\n    \n        # 安装triton-ascend，不需要单独安装triton\n        pip install triton-ascend==3.2.0\n\n\n安装 vllm & vllm-ascend\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. 需确保CANN ascend-toolkit 和 nnal 环境变量被激活，对于CANN默认安装路径 /usr/local/Ascend 而言，激活指令如下：\n\n    .. code-block::\n\n        source /usr/local/Ascend/ascend-toolkit/set_env.sh\n        source /usr/local/Ascend/nnal/atb/set_env.sh\n\n2. vllm 源码安装指令：\n\n    .. code-block:: bash\n\n        git clone --depth 1 --branch v0.13.0 https://github.com/vllm-project/vllm.git\n        cd vllm && pip install -r requirements/build.txt\n        VLLM_TARGET_DEVICE=empty pip install -v -e. && cd ..\n\n3. vllm-ascend 源码安装指令：\n\n    .. code-block:: bash\n\n        git clone -b releases/v0.13.0 https://github.com/vllm-project/vllm-ascend.git\n        cd vllm-ascend && pip install -r requirements.txt    \n        export COMPILE_CUSTOM_KERNELS=1 && pip install -v -e . && cd ..\n\n\n安装 MindSpeed\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nMindSpeed 源码安装指令：\n\n    .. code-block:: bash\n    \n        # 下载 MindSpeed，切换到指定commit-id，并下载 Megatron-LM\n        git clone https://gitcode.com/Ascend/MindSpeed.git\n        cd MindSpeed && git checkout 2.3.0_core_r0.12.1 && cd ..\n        git clone --depth 1 --branch core_v0.12.1 https://github.com/NVIDIA/Megatron-LM.git\n    \n        # 安装 MindSpeed & Megatron\n        pip install -e MindSpeed\n        pip install -e Megatron-LM\n    \n        # 安装 mbridge\n        pip install mbridge\n\nMindSpeed 对应 Megatron-LM 后端使用场景，使用方式如下：\n\n    1. 使能 verl worker 模型 ``strategy`` 配置为 ``megatron`` ，例如 ``actor_rollout_ref.actor.strategy=megatron``。\n    \n    2. MindSpeed 自定义入参可通过 ``override_transformer_config`` 参数传入，例如对 actor 模型开启 FA 特性可使用 ``+actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True``。\n    \n    3. 更多特性信息可参考 `MindSpeed & verl 文档 <https://gitcode.com/Ascend/MindSpeed/blob/master/docs/user-guide/verl.md>`_ 。\n\n\n安装verl\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n.. code-block:: bash\n\n    git clone --recursive https://github.com/volcengine/verl.git\n    cd verl && pip install -r requirements-npu.txt && pip install -v -e . && cd ..\n\n    # （可选）提示：为了更佳的使用体验，最好将recipe子模块更新至最新commit\n    cd recipe && git checkout main && cd ..\n\n昇腾暂不支持生态库说明\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nverl 中昇腾暂不支持生态库如下：\n\n    +---------------+----------------+\n    | software      | description    |\n    +---------------+----------------+\n    | flash_attn    | not supported  |\n    +---------------+----------------+\n    | liger-kernel  | not supported  |\n    +---------------+----------------+\n    \n    1. 不支持通过 flash_attn 使能 flash attention 加速，支持通过 transformers 使用。\n    2. 不支持 liger-kernel 使能。\n\n\n快速开始\n-----------------------------------\n正式使用前，建议您通过对Qwen2.5-0.5B GRPO的训练尝试以检验环境准备和安装的正确性。\n\n1.下载数据集并将数据集预处理为parquet格式，以便包含计算RL奖励所需的必要字段\n\n    .. code-block:: bash\n    \n        python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k\n\n2.执行训练\n\n    .. code-block:: bash\n    \n        set -x\n    \n        export VLLM_ATTENTION_BACKEND=XFORMERS\n    \n        python3 -m verl.trainer.main_ppo \\\n            algorithm.adv_estimator=grpo \\\n            data.train_files=$HOME/data/gsm8k/train.parquet \\\n            data.val_files=$HOME/data/gsm8k/test.parquet \\\n            data.train_batch_size=128 \\\n            data.max_prompt_length=512 \\\n            data.max_response_length=128 \\\n            data.filter_overlong_prompts=True \\\n            data.truncation='error' \\\n            actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n            actor_rollout_ref.actor.optim.lr=5e-7 \\\n            actor_rollout_ref.model.use_remove_padding=False \\\n            actor_rollout_ref.actor.entropy_coeff=0.001 \\\n            actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n            actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=20 \\\n            actor_rollout_ref.actor.use_kl_loss=True \\\n            actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n            actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n            actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n            actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n            actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n            actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n            actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n            actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n            actor_rollout_ref.rollout.name=vllm \\\n            actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n            actor_rollout_ref.rollout.n=5 \\\n            actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n            actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n            algorithm.kl_ctrl.kl_coef=0.001 \\\n            trainer.critic_warmup=0 \\\n            trainer.logger=console \\\n            trainer.project_name='verl_grpo_example_gsm8k' \\\n            trainer.experiment_name='qwen2_7b_function_rm' \\\n            trainer.n_gpus_per_node=8 \\\n            trainer.nnodes=1 \\\n            trainer.save_freq=-1 \\\n            trainer.test_freq=5 \\\n            trainer.total_epochs=1 $@\n\n\n\n算法支持现状\n-----------------------------------\n\n**表1** RL类算法\n\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    | algorithm             |         model           | download link                                                    |   actor.strategy  |   rollout.name    |   shell location                                                                                                                             |     hardware             |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen2.5-7B-instruct     |`7B <https://huggingface.co/Qwen/Qwen2.5-7B-Instruct>`_           |        FSDP       |    vllm-ascend    |`qwen2_5_7b_grpo_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_7b_grpo_npu.sh>`_                        |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen2.5-32B-instruct    |`32B <https://huggingface.co/Qwen/Qwen2.5-32B-Instruct>`_         |        FSDP       |    vllm-ascend    |`qwen2_5_32b_grpo_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_32b_grpo_npu.sh>`_                      |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen2.5-VL-3B-instruct  |`3B <https://huggingface.co/Qwen/Qwen2.5-VL-3B-Instruct>`_        |        FSDP       |    vllm-ascend    |`qwen2_5_vl_3b_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_3b_npu.sh>`_                            |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen2.5-VL-7B-instruct  |`7B <https://huggingface.co/Qwen/Qwen2.5-VL-7B-Instruct>`_        |        FSDP       |    vllm-ascend    |`qwen2_5_vl_7b_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_7b_npu.sh>`_                            |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen2.5-VL-32B-instruct |`32B <https://huggingface.co/Qwen/Qwen2.5-VL-32B-Instruct>`_      |        FSDP       |    vllm-ascend    |`qwen2_5_vl_32b_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen2_5_vl_32b_npu.sh>`_                          |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen3-4B                |`4B <https://huggingface.co/Qwen/Qwen3-4B>`_                      |        FSDP       |    vllm-ascend    |`qwen3-4B_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3_4b_grpo_vllm_1k_npu.sh>`_                         |    Atlas 800T A3         |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen3-8B                |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP       |    vllm-ascend    |`qwen3_8b_vllm_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3-8b_npu.sh>`_                                 |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen3-8B                |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP       |    sglang         |`qwen3_8b_sglang_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_grpo_sglang_32k_spmd_npu.sh>`_          |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | Qwen3-32B               |`32B <https://huggingface.co/Qwen/Qwen3-32B>`_                    |        FSDP       |    vllm-ascend    |`qwen3-32B_npu <https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3-32b_npu.sh>`_                                    |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   GRPO                | DeepSeekv3-671B         |`671B <https://huggingface.co/deepseek-ai/DeepSeek-V3>`_          |        Megatron   |    vllm-ascend    |`deepseek_v3_megatron_npu <https://github.com/verl-project/verl-recipe/blob/main//r1_ascend/run_deepseekv3_671b_grpo_megatron_npu.sh>`_       |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen2.5-7B-instruct     |`7B <https://huggingface.co/Qwen/Qwen2.5-7B-Instruct>`_           |        FSDP       |    vllm-ascend    |`qwen2.5_7b_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen2.5_7b_npu.sh>`_                                    |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen2.5-32B             |`32B <https://huggingface.co/Qwen/Qwen2.5-32B>`_                  |        FSDP       |    vllm-ascend    |`qwen2.5_32b_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen2.5_32b_npu.sh>`_                                  |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen3-8B-base           |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP       |    vllm-ascend    |`qwen3_8b_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen3_8b_base_npu.sh>`_                                   |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen3-14B-base          |`14B <https://huggingface.co/Qwen/Qwen3-14B>`_                    |        FSDP       |    vllm-ascend    |`qwen3_14b_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen3_14b_base_npu.sh>`_                                 |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen3-30B-A3B-base      |`30B <https://huggingface.co/Qwen/Qwen3-30B-A3B>`_                |        FSDP       |    vllm-ascend    |`qwen3_30b_fsdp_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen3_moe_30b_base_fsdp_npu.sh>`_                   |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   DAPO                | Qwen3-30B-A3B-base      |`30B <https://huggingface.co/Qwen/Qwen3-30B-A3B>`_                |        Megatron   |    vllm-ascend    |`qwen3_30b_megatron_npu <https://github.com/verl-project/verl-recipe/blob/main//dapo/run_dapo_qwen3_moe_30b_megatron_npu.sh>`_                |    Atlas 200T A2 Box16   |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   PPO                 | Qwen3-8B                |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP       |    vllm-ascend    |`qwen3_8b_ppo_npu <https://github.com/volcengine/verl/blob/main/examples/ppo_trainer/run_qwen3-8b_npu.sh>`_                                   |    Atlas 900 A2 PODc     |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n    |   One_Step_Off_Policy | Qwen3-8B                |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP2      |    vllm-ascend    |`qwen3_8b_fsdp2_npu <https://github.com/verl-project/verl-recipe/blob/main//one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh>`_ |    Atlas 800T A3         |\n    +-----------------------+-------------------------+------------------------------------------------------------------+-------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+--------------------------+\n\n**表2** SFT类算法\n\n    +-----------+-------------------------+------------------------------------------------------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+----------------------+\n    | algorithm |         model           |  download link                                                   |   actor.strategy  |    shell location                                                                                                                            |     hardware         |\n    +-----------+-------------------------+------------------------------------------------------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+----------------------+\n    |  SFT-PEFT | Qwen3-8B                |`8B <https://huggingface.co/Qwen/Qwen3-8B>`_                      |        FSDP       |`sft_peft_sp2_npu <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/run_qwen3_8b_sft_peft_sp2_npu.sh>`_                        |   Atlas 900 A2 PODc  |\n    +-----------+-------------------------+-------------------------+----------------------------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+----------------------+\n    | ReTool-SFT| Qwen2-7B-instruct       |`7B <https://huggingface.co/Qwen/Qwen2-7B-Instruct>`_             |        FSDP       |`qwen2_7b_sft_npu <https://github.com/verl-project/verl-recipe/blob/main/retool/run_qwen2_7b_sft_npu.sh>`_                                    |   Atlas 900 A2 PODc  |\n    +-----------+-------------------------+-------------------------+----------------------------------------+-------------------+----------------------------------------------------------------------------------------------------------------------------------------------+----------------------+\n\n\n声明\n-----------------------------------\nverl中提供的ascend支持代码、Dockerfile、镜像皆为参考样例，如在生产环境中使用请通过官方正式途径沟通，谢谢。\n"
  },
  {
    "path": "docs/ascend_tutorial/quick_start/ascend_sglang_quick_start.rst",
    "content": "Ascend Quickstart with SGLang Backend\n===================================\n\nLast updated: 01/27/2026.\n\n我们在 verl 上增加对华为昇腾设备的支持。\n\n硬件支持\n-----------------------------------\n\nAtlas 200T A2 Box16\n\nAtlas 900 A2 PODc\n\nAtlas 800T A3\n\n\n安装\n-----------------------------------\n关键支持版本\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n+-----------+-----------------+\n| software  | version         |\n+===========+=================+\n| Python    | == 3.11         |\n+-----------+-----------------+\n| HDK       | >= 25.3.RC1     |\n+-----------+-----------------+\n| CANN      | >= 8.3.RC1      |\n+-----------+-----------------+\n| torch     | >= 2.7.1        |\n+-----------+-----------------+\n| torch_npu | >= 2.7.1.post2  |\n+-----------+-----------------+\n| sglang    | v0.5.8          |\n+-----------+-----------------+\n\n从 Docker 镜像进行安装\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n我们提供了DockerFile进行构建,详见 `dockerfile_build_guidance <https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/quick_start/dockerfile_build_guidance.rst>`_ ，请根据设备自行选择对应构建文件\n\n从自定义环境安装\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n**1. 安装HDK&CANN依赖并激活**\n\n异构计算架构CANN(Compute Architecture for Neural Networks)是昇腾针对AI场景推出的异构计算架构, 为了使训练和推理引擎能够利用更好、更快的硬件支持, 我们需要安装以下 `先决条件 <https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/softwareinst/instg/instg_quick.html?Mode=PmIns&InstallType=netconda&OS=openEuler&Software=cannToolKit>`_\n\n+-----------+-------------+\n| HDK       | >= 25.3.RC1 |\n+-----------+-------------+\n| CANN      | >= 8.3.RC1  |\n+-----------+-------------+\n安装完成后请激活环境\n\n.. code-block:: bash\n\n    source /usr/local/Ascend/ascend-toolkit/set_env.sh\n    source /usr/local/Ascend/nnal/atb/set_env.sh\n\n**2. 创建conda环境**\n\n.. code-block:: bash\n    \n    # create conda env\n    conda create -n verl-sglang python==3.11\n    conda activate verl-sglang\n\n**3. 然后，执行我们在 verl 中提供的脚本** `install_sglang_mcore_npu.sh <https://github.com/verl-project/verl/blob/main/scripts/install_sglang_mcore_npu.sh>`_\n\n如果在此步骤中遇到错误，请检查脚本并手动按照脚本中的步骤操作。\n\n.. code-block:: bash\n\n    git clone https://github.com/volcengine/verl.git  \n    # Make sure you have activated verl conda env\n    # NPU_DEVICE=A3 or A2 depends on your device\n    # USE_MEGATRON=1 if you need to install megatron backend\n    NPU_DEVICE=A3 USE_MEGATRON=1 bash verl/scripts/install_sglang_mcore_npu.sh\n\n**4. 安装verl**\n\n.. code-block:: bash\n\n    cd verl\n    pip install --no-deps -e .\n    pip install -r requirements-npu.txt \n\n\n快速开始\n-----------------------------------\n\n**1.当前NPU sglang脚本一览**\n\n.. _Qwen3-30B: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh\n.. _Qwen2.5-32B: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen2-32b_sglang_fsdp_npu.sh\n.. _Qwen3-8B-1k: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_grpo_sglang_1k_spmd_npu.sh\n.. _Qwen3-8B-32k: https://github.com/verl-project/verl/blob/main/examples/grpo_trainer/run_qwen3_8b_grpo_sglang_32k_spmd_npu.sh\n\n   +-----------------+----------------+----------+-------------------+\n   | 模型            | 推荐NPU型号    | 节点数量 | 训推后端          |\n   +=================+================+==========+===================+\n   | `Qwen3-30B`_    | Atlas 800T A3  | 1        | SGLang + Megatron |\n   +-----------------+----------------+----------+-------------------+\n   | `Qwen2.5-32B`_  | Atlas 900 A2   | 2        | SGLang + FSDP     |\n   +-----------------+----------------+----------+-------------------+\n   | `Qwen3-8B-1k`_  | Atlas A3/A2    | 1        | SGLang + FSDP     |\n   +-----------------+----------------+----------+-------------------+\n   | `Qwen3-8B-32k`_ | Atlas A3/A2    | 1        | SGLang + FSDP     |\n   +-----------------+----------------+----------+-------------------+\n\n**2.最佳实践**\n\n我们提供基于verl+sglang `Qwen3-30B`_ 以及 `Qwen2.5-32B`_ 的 `最佳实践 <https://github.com/verl-project/verl/blob/main/docs/ascend_tutorial/examples/ascend_sglang_best_practices.rst>`_ 作为参考\n\n**3.环境变量与参数**\n\n当前NPU上支持sglang后端必须添加以下环境变量\n\n.. code-block:: bash\n\n    #支持NPU单卡多进程 https://www.hiascend.com/document/detail/zh/canncommercial/850/commlib/hcclug/hcclug_000091.html\n    export HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\n    export HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\n    #规避ray在device侧调用无法根据is_npu_available接口识别设备可用性\n    export RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\n    #根据当前设备和需要卡数定义\n    export ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15\n    #使能推理EP时需要\n    export SGLANG_DEEPEP_BF16_DISPATCH=1\n\n\n\n当前verl已解析推理常见参数, 详见 `async_sglang_server.py <https://github.com/verl-project/verl/blob/main/verl/workers/rollout/sglang_rollout/async_sglang_server.py>`_  中 ServerArgs初始化传参,其他 `sglang参数 <https://github.com/sgl-project/sglang/blob/main/docs/advanced_features/server_arguments.md>`_ 均可通过engine_kwargs 进行参数传递\n\nvllm后端推理脚本转换为sglang, 需要添加修改以下参数\n\n.. code-block:: bash\n\n    #必须\n    actor_rollout_ref.rollout.name=sglang\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=\"ascend\"\n    #可选\n    #使能推理EP，详细使用方法见 https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README_CN.md\n    ++actor_rollout_ref.rollout.engine_kwargs.sglang.deepep_mode=\"auto\" \n    ++actor_rollout_ref.rollout.engine_kwargs.sglang.moe_a2a_backend=\"deepep\"\n    #Moe模型多DP时必须设置为True\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False\n    #chunked_prefill默认关闭\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1\n\n\n\n"
  },
  {
    "path": "docs/ascend_tutorial/quick_start/dockerfile_build_guidance.rst",
    "content": "Ascend Dockerfile Build Guidance\n===================================\n\nLast updated: 03/03/2025.\n\n\n镜像获取 & 公开镜像地址\n--------------------\n\n昇腾在 `quay.io/ascend/verl <https://quay.io/repository/ascend/verl?tab=tags&tag=latest>`_ 中托管每日构建的 A2/A3 镜像，基于上述 Dockerfile 构建。\n\n每日构建镜像名格式：verl-{CANN版本}-{NPU设备类型}-{操作系统版本}-{python版本}-latest\n\nverl release版本镜像名格式：verl-{CANN版本}-{NPU设备类型}-{操作系统版本}-{python版本}-{verl release版本号}\n\n\n\n镜像硬件支持\n-----------------------------------\n\nAtlas 200T A2 Box16\n\nAtlas 900 A2 PODc\n\nAtlas 800T A3\n\n\n镜像内各组件版本信息清单\n----------------\n\n================= ============\n组件        版本\n================= ============\n基础镜像            Ubuntu 22.04\nPython             3.11\nCANN               8.5.0\ntorch              2.8.0\ntorch_npu          2.8.0\ntorchvision        0.22.1\nvLLM               0.13.0\nvLLM-ascend        0.13.0\nMegatron-LM        v0.12.1\nMindSpeed          2.3.0_core_r0.12.1\ntriton-ascend      3.2.0\nmbridge            latest version\nSGLang             v0.5.8\nsgl-kernel-npu     (46b73de)\n================= ============\n\n\nDockerfile构建镜像脚本清单\n---------------------------\n\n============== ============== ============== ==============================================================\n设备类型         基础镜像版本     推理后端        参考文件\n============== ============== ============== ==============================================================\nA2              8.2.RC1        vLLM            `Dockerfile.ascend_8.2.rc1_a2 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.2.rc1_a2>`_\nA2              8.3.RC1        vLLM            `Dockerfile.ascend_8.3.rc1_a2 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.3.rc1_a2>`_\nA2              8.5.0          vLLM            `Dockerfile.ascend_8.5.0_a2 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.5.0_a2>`_\nA2              8.3.RC1        SGLang          `Dockerfile.ascend.sglang_8.3.rc1_a2 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a2>`_\nA3              8.2.RC1        vLLM            `Dockerfile.ascend_8.2.rc1_a3 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.2.rc1_a3>`_\nA3              8.3.RC1        vLLM            `Dockerfile.ascend_8.3.rc1_a3 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.3.rc1_a3>`_\nA3              8.5.0          vLLM            `Dockerfile.ascend_8.5.0_a3 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend_8.5.0_a3>`_\nA3              8.3.RC1         SGLang          `Dockerfile.ascend.sglang_8.3.rc1_a3 <https://github.com/volcengine/verl/blob/main/docker/ascend/Dockerfile.ascend.sglang_8.3.rc1_a3>`_\n============== ============== ============== ==============================================================\n\n\n镜像构建命令示例\n--------------------\n\n.. code:: bash\n\n   # Navigate to the directory containing the Dockerfile \n   cd {verl-root-path}/docker/ascend\n\n   # Build the image\n   # vLLM\n   docker build -f Dockerfile.ascend_8.3.rc1_a2 -t verl-ascend:8.3.rc1-a2 .\n   # SGLang\n   docker build -f Dockerfile.ascend.sglang_8.3.rc1_a2 -t verl-ascend-sglang:8.3.rc1-a2 .\n\n\n声明\n--------------------\nverl中提供的ascend相关Dockerfile、镜像皆为参考样例，可用于尝鲜体验，如在生产环境中使用请通过官方正式途径沟通，谢谢。"
  },
  {
    "path": "docs/blog/v0.7.md",
    "content": "# verl 0.7 release blog\n\n**Author:** verl team\n\nLast updated: 01/03/2026.\n\n## Overview\nverl adopts a Hybrid-Controller architecture (also known as HybridFlow). Sharing design principles with asynchronous sharded dataflow systems like Google Pathways, verl models Reinforcement Learning (RL) algorithms, such as PPO, GRPO, DAPO, and others, as a multi-stage, multi-model and parallelizable dataflow graph.\n\nTo balance flexibility with performance, verl unifies two distinct programming models:\n\n**High-Level Single-Controller (MPMD)**: At the orchestration level, a single process `RLTrainer` manages the global computation graph. It handles macro-tasks such as scheduling rollout generation, triggering reward scoring, and dispatching distributed training jobs.\n\n**Internal Multi-Controller (SPMD)**: Internally, the Model Engine operates in standard distributed training mode. Workers execute identical programs, via trainer backends like FSDP, Megatron, or VeOmni, or rollout executors (not rollout server) like vLLM/SGLang/TensorRT-LLM, to perform heavy distributed computation, synchronizing via collective communication.\n\n<div align=\"center\">\n <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/hybridflow.png?raw=true\" width=\"800\" alt=\"hybridflow.png\">\n</div>\n\nThis hybrid approach offers significant advantages:\n\n**Flexible Orchestration**: The single-controller design allows verl to dynamically manage complex constraints within the computation graph, including flexible data dependencies, diverse resource allocation and model placement, and fine-grained asynchronous staleness control.\n\n**Abstraction of Complexity**: We encapsulate complex parallel strategies—such as 5D parallelism (DP, TP, CP, PP, and EP)—strictly within the Model Engine. This allows users to focus entirely on RL algorithm implementation without getting bogged down by the details of distributed training.\n\nFurthermore, leveraging Ray placement groups, verl provides `ResourcePool` and `WorkerGroup` abstractions. These enable flexible GPU sharing among the various roles in the RL process—such as actor, critic, reward, and rollout—allowing components to share resources efficiently while remaining isolated.\n\nAs illustrated in the diagram below, the overall architecture of verl is divided into two layers:\n\n- **verl-core**: provides four components required for the RL pipeline: model engine, rollout engine, checkpoint engine, and transfer queue. Each component exposes abstract interfaces, making them both extensible and pluggable.\n- **verl-trainer**: builds upon these components, construct various RL pipelines—such as on-policy, one-step-off-policy, and fully asynchronous—tailored to meet the demands of diverse scenarios.\n\n<div align=\"center\">\n <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/verl-arch.png?raw=true\" width=\"400\" alt=\"verl-arch.png\">\n</div>\n\n\n## verl-core\n### Model Engine\n\nThe Model Engine serves as verl's core training engine, defining a set of abstract interfaces that support pluggable backends. It operates in SPMD mode:\n- SFT: Workers are launched via torchrun.\n- RL: Workers are executed via the WorkerGroup API, invoked by the single-controller.\n\nThe abstract interfaces include methods like `initialize`, `forward`, `optimizer_step`, and `load`/`offload`. Integrating a new training engine simply requires inheriting and implementing these interfaces. Crucially, because all backends adhere to this unified abstraction, adding a new Model Engine requires absolutely no code modification on the caller side. The RLTrainer remains completely agnostic to the backend's specific parallel strategy when calling these interfaces, while the WorkerGroup automatically handles data dispatch and collection based on the underlying parallelism.\n\nCurrently, the Model Engine supports the following backends (more backend maybe supported in future, e.g torchtitan):\n|Backend|Parallelism|Performance|Support Model|New Model Support Time\n|-----|-----|----|----|----|\n|FSDP| FSDP+SP|Dense medium/MoE low| all transformer models|Day 0\n|MCore| DP+TP+PP+EP+CP|High| see [Megatron-Bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) support model list|few weeks or month\n|VeOmni| FSDP+SP+EP|Medium| see [VeOmni](https://github.com/ByteDance-Seed/VeOmni) support model list|~1 week\n\n```python\nclass BaseEngine:\n    def initialize(self):\n        \"\"\"Instantiate or load the model, optimizer, and learning rate scheduler.\"\"\"\n        raise NotImplementedError\n\n    def optimizer_zero_grad(self):\n        \"\"\"Zero the gradients of the optimizer.\"\"\"\n        raise NotImplementedError\n\n    def optimizer_step(self):\n        \"\"\"Perform an optimization step using the optimizer.\"\"\"\n        raise NotImplementedError\n\n    def lr_scheduler_step(self):\n        \"\"\"Advance the learning rate scheduler by one step.\"\"\"\n        raise NotImplementedError\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:\n        \"\"\"Perform a forward pass and optionally a backward pass on a batch of data.\"\"\"\n        raise NotImplementedError\n\n    def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None, None], Optional[dict]]:\n        \"\"\"Get a generator that yields per-tensor parameters and optional peft config.\"\"\"\n        raise NotImplementedError\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"Move model parameters, optimizer states, or both to the specified device.\"\"\"\n        raise NotImplementedError\n```\n\n\n### Rollout Engine\nAs LLM reinforcement learning evolves from single-turn, static tasks to multi-turn, dynamic, and interactive agentic tasks, the legacy SPMD rollout mode previously used by verl has become insufficient. Consequently, in verl v0.7, we have removed the SPMD rollout mode and switched to rollout server mode by default.\n\n<div align=\"center\">\n    <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/rollout_engine.png?raw=true\" width=\"300\" alt=\"rollout_engine.png\">\n</div>\n\nIn the server mode, the LLM server operates as online serving rather than the traditional offline batch inference. Clients send per-sample requests to the server, enabling the engine to utilize dynamic batching. This significantly enhances throughput efficiency for multi-turn conversation. Furthermore, the server-based approach eliminates the need for intrusive modifications to the LLM inference engine, allowing for the seamless integration of modern inference backends such as vLLM, SGLang, and TensorRT-LLM.\n\nOn the client side, verl introduces an extensible **AgentLoop** abstraction designed to define custom agentic task loops. This abstraction manages the cycle of requesting responses from the LLM server and interacting with external environments to obtain feedback. We provide two default implementations:\n- **SingleTurnAgentLoop**: Designed for standard single-turn tasks.\n- **ToolAgentLoop**: Designed for classic ReAct architectures involving multi-turn tool invocation.\n\nUsers can implement custom AgentLoop logic tailored to their specific needs, such as [SWEAgentLoop](https://github.com/volcengine/verl/pull/4080) or GUIAgentLoop.\n\n```python\nclass AgentLoopBase(ABC):\n    @abstractmethod\n    async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:\n        \"\"\"Run agent loop to interact with LLM server and environment.\n\n        Args:\n            sampling_params (Dict[str, Any]): LLM sampling params.\n            **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`.\n\n        Returns:\n            AgentLoopOutput: Agent loop output.\n        \"\"\"\n        raise NotImplementedError\n```\n\n### TransferQueue\nAs mentioned, verl uses a global single-controller RLTrainer to orchestrate the computation graph. A major limitation in the current implementation is that the RLTrainer handles both control and data flow, creating a bottleneck when dispatching data between components. This issue is amplified by the massive data volumes in multimodal training (images, video, audio) and complex algorithms like router replay, which requires transmitting large tensors per sample. Our earlier attempt to solve this using the Ray object store yielded poor performance due to the lack of tensor optimization and fine-grained column access.\n\n<div align=\"center\">\n    <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/transfer_queue.png?raw=true\" width=\"500\" alt=\"transfer_queue.png\">\n</div>\n\nIn v0.7, we experimentally introduced **TransferQueue** to decouple control flow from data flow. The RLTrainer now only dispatch instructions and metadata, while TransferQueue handles data transmission via reference passing. TransferQueue is specifically optimized for PyTorch tensors (supporting zero-copy and RDMA) and allows for backend extensions like ZeroMQ, NIXL, and Ray RDT. We plan to make this the default transmission method in v0.8.\n\n```python\n# In PPOTrainer\ndef fit(self):\n    batch = next(dataloader)\n    gen_batch: BatchMeta = self.rollout_manager.generate_sequences(batch)\n    output: BatchMeta = self.actor_rollout_wg.compute_log_prob(gen_batch)\n    gen_batch = gen_batch.union(output)\n    output = self.actor_rollout_wg.update_actor(gen_batch)\n\n# In Worker\ndef compute_log_prob(self, batch: BatchMeta) -> BatchMeta:\n    data = tq.get(batch)\n    output = self.actor.infer_batch(data=data)\n    return tq.put(output)\n```\n\n### Checkpoint Engine\n\nWith the increase in LLM context lengths and the evolution of agentic tasks, the \"long-tail\" problem in rollout has become prominent, limiting the overall efficiency of RL training.\n\nTo mitigate this, a viable strategy is moving from on-policy synchronous training to off-policy asynchronous training, e.g [Laminar](https://arxiv.org/abs/2510.12633), [Areal](https://arxiv.org/abs/2505.24298), [StreamRL](https://arxiv.org/abs/2504.15930), [LlamaRL](https://arxiv.org/pdf/2505.24034), [PipelineRL](https://arxiv.org/abs/2509.19128). This involves separating the rollout and model engines onto different nodes (a disaggregated architecture, as opposed to colocated), with data transmitted via queues. This separation alleviates the rollout long-tail issue and enables rollout elastic scaling, fault tolerance, and heterogeneous hardware. However, it introduces a new challenge: efficient cross-node parameter synchronization.\n\n<div align=\"center\">\n    <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/checkpoint_engine.png?raw=true\" width=\"500\" alt=\"checkpoint_engine.png\">\n</div>\n\nTo address this, we introduce the Checkpoint Engine: a unified abstraction layer designed to synchronize weights between various training and inference backends.\n- It provides three unified APIs to implement the streaming transmission of parameters.\n- Users can extend the Transport Layer implementation based on their specific infrastructure requirements (device, network, local cache, etc.).\n\nCurrently, we provide two transport backends: NCCL (for broadcast collective communication) and NIXL (for P2P point-to-point communication).\n\n```python\nclass CheckpointEngine(ABC):\n    @abstractmethod\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    async def receive_weights(self) -> Generator[tuple[str, torch.Tensor], None, None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        raise NotImplementedError\n```\n\n## verl-trainer\nBuilding upon the four core components provided by verl-core, verl-trainer constructs several RL training pipelines tailored to specific scenarios. These pipelines are designed to address training efficiency challenges across varying scales and requirements:\n\n**On-policy (Synchronous)**\n  - Main Features: Executes rollout and training serially, typically sharing GPU resources (Colocate). It strictly adheres to standard on-policy algorithm definitions, where training must wait for all samples to be generated.\n  - Scenarios: Best for baseline implementations, scenarios where strict algorithmic correctness is prioritized over training throughput.\n\n**One-step-off-policy (Async)**\n  - Main Features: Parallelizes generation and training by overlapping the current training step with the next batch's generation. It employs resource isolation and uses parameters from the previous step for rollout to minimize GPU idle time.\n  - Scenarios: Ideal for scenarios requiring moderate efficiency gains (20%–40%) while maintaining training stability very close to strict on-policy methods.\n\n**Fully async (Decoupled & Streaming)**\n  - Main Features: Completely decouples the Trainer and Rollouter onto separate nodes. It utilizes streaming data transfer, staleness control, and partial rollout mechanisms to maximize throughput and mitigate long-tail generation latency.\n  - Scenarios: Essential for large-scale training (e.g., 128+ GPUs) or complex reasoning tasks (e.g., long chain-of-thought) where generation latency significantly bottlenecks performance.\n\n<div align=\"center\">\n    <img src=\"https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/fully_async.png?raw=true\" width=\"1000\" alt=\"fully_async.png\">\n</div>\n\n## roadmap\n### v0.7 release\n\n**Model Engine**\n- Integrate Megatron-Bridge and support LoRA/PEFT, see blog post: [How We Build Trillion Parameter Reasoning RL with 10% GPUs](https://macaron.im/mindlab/research/building-trillion-parameter-reasoning-rl-with-10-gpus)\n- Support experimental fp8 training for megatron backend\n- Support new model for megatron backend: GPT-OSS, Qwen3-Next\n- Comprehensive support for new mode engine, FSDP and Megatron engine are production ready.\n  - Dispatch tensordict with nested tensor instead of padded DataProto\n  - Add TrainingWorker that resembles Tinker-like API\n  - Add VLM support for model engine, SFT and RL trainer\n  - Add model engine based critic model\n  - Implement ActorRolloutRefWorker by TrainingWorker, support different backend in one worker\n- New VeOmni engine added, still in alpha status.\n\n**Rollout Engine**\n- Remove SPMD rollout mode\n- Support blockwise fp8 rollout for vllm and sglang; support online quant for vllm with torchao\n- Experimental router replay support for vllm\n- Optimize multi-modal data fetch and preprocess, support video input\n- Upgrade to vllm==0.12.0; sglang==0.5.6\n\n**Reward**\n- Support hybrid reward scenarios, including generative, discriminative, rule-based rewards, and their combinations.\n- Refactor reward models into server mode, supporting both colocated and standalone deployments.\n- Introduce new reward managers to handle more complex scenarios, limited mode for request rate control and remote mode for CPU-intensive tasks.\n\n**Algorithm**\n- Add [CISPO](https://arxiv.org/pdf/2506.13585): Clipped IS-weight Policy Optimization\n- Add [SAPO](https://arxiv.org/abs/2511.20347): Soft Adaptive Policy Optimization\n\n**Recipe**\n- [NEW] VLA: add experimental support for VLA model\n- [NEW] [rhymerl](https://arxiv.org/abs/2508.18588): History Rhymes: Accelerating LLM Reinforcement Learning with RhymeRL\n- TransferQueue: support multiple data partition and optimize tensor zero-copy serialization\n- One-step-off-policy/Fully async: optimize weight synchronization by checkpoint engine with bucket and pipeline support.\n\n### v0.8\n\n**Model Engine**\n- Deprecate DataProto by Tensordict for zero padding transmission\n- Switch default to new model engine, mark legacy engine (fsdp_workers.py, megatron_workers.py) as deprecated\n- Feature parity between new and legacy model engine: LoRA/PEFT, etc\n- Polish VeOmni engine to production ready status\n- Support MTP RL training\n- Optimize GPU memory for long context: fine-grained activation recompuation/offload\n- New model support: DeepSeek V3.2, etc\n\n**Rollout Engine**\n- New rollout engine TensorRT-LLM\n- Separate vllm worker from trainer process, update weights by cuda ipc\n\n**TransferQueue**\n- Merge TransferQueue recipe into main\n- Optimize e2e image/video vlm training pipeline by TransferQueue\n- Optimize router replay transmission by TransferQueue\n\n**Checkpoint Engine**\n- Add checkpoint engine abstract interface\n- Add NCCL and NIXL transport backend\n- Add more transport backend\n\n### v0.9\n\n**Trainer**\n- Merge Full async into main: refactor with verl-core component\n\n**Model Engine**\n- Remove legacy model engine (fsdp_workers.py, megatron_workers.py)\n- Support omni-model RL training: Qwen3-Omni, BAGEL, etc\n\n**Rollout Engine**\n- New rollout engine vllm-omni\n\n**More agentic training recipe**\n- SWEAgent\n- GUIAgent\n"
  },
  {
    "path": "docs/conf.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\n# Configuration file for the Sphinx documentation builder.\r\n#\r\n# This file only contains a selection of the most common options. For a full\r\n# list see the documentation:\r\n# https://www.sphinx-doc.org/en/master/usage/configuration.html\r\n\r\n# -- Path setup --------------------------------------------------------------\r\n\r\n# If extensions (or modules to document with autodoc) are in another directory,\r\n# add these directories to sys.path here. If the directory is relative to the\r\n# documentation root, use os.path.abspath to make it absolute, like shown here.\r\n#\r\n# import os\r\n# import sys\r\n# sys.path.insert(0, os.path.abspath('.'))\r\n\r\n\r\n# -- Project information -----------------------------------------------------\r\n\r\nproject = \"verl\"\r\ncopyright = \"2024 ByteDance Seed Foundation MLSys Team\"\r\nauthor = \"Guangming Sheng, Chi Zhang, Yanghua Peng, Haibin Lin\"\r\n\r\n\r\n# -- General configuration ---------------------------------------------------\r\n# The master toctree document.\r\nmaster_doc = \"index\"\r\n\r\n# Add any Sphinx extension module names here, as strings. They can be\r\n# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom\r\n# ones.\r\nextensions = [\r\n    \"myst_parser\",\r\n    \"sphinx.ext.autodoc\",\r\n    \"sphinx.ext.autosummary\",\r\n    \"sphinx.ext.autosectionlabel\",\r\n    \"sphinx.ext.napoleon\",\r\n    \"sphinx.ext.viewcode\",\r\n]\r\n\r\n# MyST-Parser settings\r\nmyst_enable_extensions = [\r\n    \"dollarmath\",  # Enables $...$ and $$...$$ syntax\r\n    \"amsmath\",  # Enables amsmath environments\r\n]\r\n\r\n# Use Google style docstrings instead of NumPy docstrings.\r\nnapoleon_google_docstring = True\r\nnapoleon_numpy_docstring = False\r\n\r\n# The suffix(es) of source filenames.\r\n# You can specify multiple suffix as a list of string:\r\nsource_suffix = {\r\n    \".rst\": \"restructuredtext\",\r\n    \".md\": \"markdown\",\r\n}\r\n\r\n# Add any paths that contain templates here, relative to this directory.\r\ntemplates_path = [\"_templates\"]\r\n\r\n# The language for content autogenerated by Sphinx. Refer to documentation\r\n# for a list of supported languages.\r\n#\r\n# This is also used if you do content translation via gettext catalogs.\r\n# Usually you set \"language\" from the command line for these cases.\r\nlanguage = \"en\"\r\n\r\n# List of patterns, relative to source directory, that match files and\r\n# directories to ignore when looking for source files.\r\n# This pattern also affects html_static_path and html_extra_path.\r\nexclude_patterns = [\"_build\", \"Thumbs.db\", \".DS_Store\"]\r\n\r\n\r\n# -- Options for HTML output -------------------------------------------------\r\n\r\n# The theme to use for HTML and HTML Help pages.  See the documentation for\r\n# a list of builtin themes.\r\n#\r\nhtml_theme = \"sphinx_rtd_theme\"\r\n\r\n# Add any paths that contain custom static files (such as style sheets) here,\r\n# relative to this directory. They are copied after the builtin static files,\r\n# so a file named \"default.css\" will overwrite the builtin \"default.css\".\r\nhtml_static_path = [\"_static\"]\r\n\r\n# Add the JavaScript file\r\nhtml_js_files = [\r\n    \"js/runllm-widget.js\",\r\n    \"js/resizable-sidebar.js\",\r\n]\r\n\r\n# Add custom CSS file for full-width layout\r\nhtml_css_files = [\r\n    \"custom.css\",\r\n]\r\n\r\nexclude_patterns += [\"README.md\", \"README_vllm0.7.md\"]\r\n\r\nsuppress_warnings = [\"ref.duplicate\", \"ref.myst\"]\r\n"
  },
  {
    "path": "docs/data/transfer_queue.md",
    "content": "# TransferQueue Data System\n\nLast updated: 01/07/2026.\n\nThis doc introduce [TransferQueue](https://gitcode.com/Ascend/TransferQueue), an asynchronous streaming data management system for efficient post-training.\n\n🔥 **Now TransferQueue is formally open-sourced at [GitCode](https://gitcode.com/Ascend/TransferQueue). We will soon provide a [Github Mirror Repo](https://github.com/Ascend/TransferQueue) for community contributions. <span style=\"color: #FF0000;\">You are welcome to submit contributions or propose new ideas on either platform!**</span>\n\n\n> At the mean time, the early development history remains accessible at: https://github.com/TransferQueue/TransferQueue.\n\n<h2 id=\"overview\"> Overview</h2>\n\nTransferQueue is a high-performance data storage and transfer module with panoramic data visibility and streaming scheduling capabilities, optimized for efficient dataflow in post-training workflows.\n\n<p align=\"center\">\n  <img src=\"https://github.com/TransferQueue/community_doc/blob/main/docs/tq_arch.png?raw=true\" width=\"70%\">\n</p>\n\nTransferQueue offers **fine-grained, sample-level** data management and **load-balancing** (on the way) capabilities, serving as a data gateway that decouples explicit data dependencies across computational tasks. This enables a divide-and-conquer approach, significantly simplifies the algorithm controller design.\n\n<p align=\"center\">\n  <img src=\"https://github.com/TransferQueue/community_doc/blob/main/docs/main_func.png?raw=true\" width=\"70%\">\n</p>\n\n<h2 id=\"updates\"> Updates</h2>\n\n - **Dec 30, 2025**:  **TransferQueue x verl** integration is tested with the DAPO algorithm at scale **(64 nodes, 1024 cards)**. It significantly optimizes host memory utilization and accelerates data transfers. Stay tuned for more details!\n - **Dec 20, 2025**: 🔥 The official [tutorial](https://github.com/TransferQueue/TransferQueue/tree/main/tutorial) is released! Feel free to check it out.\n - **Nov 10, 2025**: We disentangle the data retrieval logic from TransferQueueController [PR#101](https://github.com/TransferQueue/TransferQueue/pull/101). Now you can implement your own `Sampler` to control how to consume the data.\n - **Nov 5, 2025**: We provide a `KVStorageManager` that simplifies the integration with KV-based storage backends [PR#96](https://github.com/TransferQueue/TransferQueue/pull/96). The first available KV-based backend is [Yuanrong](https://gitee.com/openeuler/yuanrong-datasystem).\n - **Nov 4, 2025**: Data partition capability is available in [PR#98](https://github.com/TransferQueue/TransferQueue/pull/98). Now you can define logical data partitions to manage your train/val/test datasets.\n - **Oct 25, 2025**: We make storage backends pluggable in [PR#66](https://github.com/TransferQueue/TransferQueue/pull/66). You can try to integrate your own storage backend with TransferQueue now!\n - **Oct 21, 2025**: Official integration into verl is ready [verl/pulls/3649](https://github.com/volcengine/verl/pull/3649). Following PRs will optimize the single controller architecture by fully decoupling data & control flows.\n - **July 22, 2025**: We present a series of Chinese blogs on <a href=\"https://zhuanlan.zhihu.com/p/1930244241625449814\">Zhihu 1</a>, <a href=\"https://zhuanlan.zhihu.com/p/1933259599953232589\">2</a>.\n - **July 21, 2025**: We started an RFC on verl community [verl/RFC#2662](https://github.com/volcengine/verl/discussions/2662).\n - **July 2, 2025**: We publish the paper [AsyncFlow](https://arxiv.org/abs/2507.01663).\n\n<h2 id=\"components\"> Components</h2>\n\n### Control Plane: Panoramic Data Management\n\nIn the control plane, `TransferQueueController` tracks the **production status** and **consumption status** of each training sample as metadata. When all the required data fields are ready (i.e., written to the `TransferQueueStorageManager`), we know that this data sample can be consumed by downstream tasks.\n\nFor consumption status, we record the consumption records for each computational task (e.g., `generate_sequences`, `compute_log_prob`, etc.). Therefore, even when different computation tasks require the same data field, they can consume the data independently without interfering with each other.\n\n<p align=\"center\">\n  <img src=\"https://github.com/TransferQueue/community_doc/blob/main/docs/control_plane.png?raw=true\" width=\"70%\">\n</p>\n\nTo make the data retrieval process more customizable, we provide a `Sampler` class that allows users to define their own data retrieval and consumption logic. Refer to the [Customize](#customize) section for details.\n\n> In the future, we plan to support **load-balancing** and **dynamic batching** capabilities in the control plane. Additionally, we will support data management for disaggregated frameworks where each rank manages the data retrieval by itself, rather than coordinated by a single controller.\n\n### Data Plane: Distributed Data Storage\n\nIn the data plane, we provide a pluggable design that enables TransferQueue to integrate with different storage backends according to user requirements.\n\nSpecifically, we provide a `TransferQueueStorageManager` abstraction class that defines the core APIs as follows:\n\n- `async def put_data(self, data: TensorDict, metadata: BatchMeta) -> None`\n- `async def get_data(self, metadata: BatchMeta) -> TensorDict`\n- `async def clear_data(self, metadata: BatchMeta) -> None`\n\nThis class encapsulates the core interaction logic within the TransferQueue system. You only need to write a simple subclass to integrate your own storage backend. Refer to the [Customize](#customize) section for details.\n\nCurrently, we support the following storage backends:\n\n- SimpleStorageUnit: A basic CPU memory storage with minimal data format constraints and easy usability.\n- [Yuanrong](https://gitcode.com/openeuler/yuanrong-datasystem) (beta, [#PR107](https://github.com/TransferQueue/TransferQueue/pull/107), [#PR96](https://github.com/TransferQueue/TransferQueue/pull/96)): An Ascend native data system that provides hierarchical storage interfaces including HBM/DRAM/SSD.\n- [Mooncake Store](https://github.com/kvcache-ai/Mooncake) (alpha, [#PR162](https://github.com/TransferQueue/TransferQueue/pull/162)): A high-performance, KV-based hierarchical storage that supports RDMA transport between GPU and DRAM.\n- [Ray Direct Transport](https://docs.ray.io/en/master/ray-core/direct-transport.html) (alpha, [#PR167](https://github.com/TransferQueue/TransferQueue/pull/167)): Ray's new feature that allows Ray to store and pass objects directly between Ray actors.\n\nAmong them, `SimpleStorageUnit` serves as our default storage backend, coordinated by the `AsyncSimpleStorageManager` class. Each storage unit can be deployed on a separate node, allowing for distributed data management.\n\n`SimpleStorageUnit` employs a 2D data structure as follows:\n\n- Each row corresponds to a training sample, assigned a unique index within the corresponding global batch.\n- Each column represents the input/output data fields for computational tasks.\n\nThis data structure design is motivated by the computational characteristics of the post-training process, where each training sample is generated in a relayed manner across task pipelines. It provides an accurate addressing capability, which allows fine-grained, concurrent data read/write operations in a streaming manner.\n\n<p align=\"center\">\n  <img src=\"https://github.com/TransferQueue/community_doc/blob/main/docs/data_plane.png?raw=true\" width=\"70%\">\n</p>\n\n### User Interface: Asynchronous & Synchronous Client\n\nThe interaction workflow of TransferQueue system is as follows:\n\n1. A process sends a read request to the `TransferQueueController`.\n2. `TransferQueueController` scans the production and consumption metadata for each sample (row), and dynamically assembles a micro-batch metadata according to the load-balancing policy. This mechanism enables sample-level data scheduling.\n3. The process retrieves the actual data from distributed storage units using the metadata provided by the controller.\n\nTo simplify the usage of TransferQueue, we have encapsulated this process into `AsyncTransferQueueClient` and `TransferQueueClient`. These clients provide both asynchronous and synchronous interfaces for data transfer, allowing users to easily integrate TransferQueue into their framework.\n\n> In the future, we will provide a `StreamingDataLoader` interface for disaggregated frameworks as discussed in [issue#85](https://github.com/TransferQueue/TransferQueue/issues/85) and [verl/RFC#2662](https://github.com/volcengine/verl/discussions/2662). Leveraging this abstraction, each rank can automatically get its own data like `DataLoader` in PyTorch. The TransferQueue system will handle the underlying data scheduling and transfer logic caused by different parallelism strategies, significantly simplifying the design of disaggregated frameworks.\n\n<h2 id=\"show-cases\">🔥 Showcases</h2>\n\n### General Usage\n\nThe primary interaction points are `AsyncTransferQueueClient` and `TransferQueueClient`, serving as the communication interface with the TransferQueue system.\n\nCore interfaces:\n\n- `(async_)get_meta(data_fields: list[str], batch_size:int, partition_id: str, mode: str, task_name:str, sampling_config: Optional[dict[str, Any]]) -> BatchMeta`\n- `(async_)get_data(metadata: BatchMeta) -> TensorDict`\n- `(async_)put(data: TensorDict, metadata: Optional[BatchMeta], partition_id: Optional[str])`\n- `(async_)clear_partition(partition_id: str)` and `(async_)clear_samples(metadata: BatchMeta)`\n\n<span style=\"color: #FF0000;\">**Refer to our [tutorial](https://github.com/TransferQueue/TransferQueue/tree/main/tutorial) for detailed examples.**</span>\n\n\n### verl Example\n\nThe primary motivation for integrating TransferQueue to verl now is to **alleviate the data transfer bottleneck of the single controller `RayPPOTrainer`**. Currently, all `DataProto` objects must be routed through `RayPPOTrainer`, resulting in a single point bottleneck of the whole post-training system. \n\n![verl_dataflow_DataProto](https://github.com/TransferQueue/community_doc/blob/main/docs/verl_workflow.jpeg?raw=true)\n\n\nLeveraging TransferQueue, we separate experience data transfer from metadata dispatch by\n\n- Replacing `DataProto` with `BatchMeta` (metadata) and `TensorDict` (actual data) structures\n- Preserving verl's original Dispatch/Collect logic via BatchMeta (maintaining single-controller debuggability)\n- Accelerating data transfer by TransferQueue's distributed storage units\n\n![verl_dataflow_TransferQueue](https://github.com/TransferQueue/community_doc/blob/main/docs/verl_workflow_with_tq.jpeg?raw=true)\n\n\nYou may refer to the [recipe](https://github.com/TransferQueue/TransferQueue/tree/dev/recipe/simple_use_case), where we mimic the verl usage in both async & sync scenarios. Official integration to verl is also available now at [verl/pulls/3649](https://github.com/volcengine/verl/pull/3649) (with subsequent PRs to further optimize the integration).\n\n\n### Use Python package\n```bash\npip install TransferQueue\n```\n\n### Build wheel package from source code\n\nFollow these steps to build and install:\n1. Clone the source code from the GitHub repository\n   ```bash\n   git clone https://github.com/TransferQueue/TransferQueue/\n   cd TransferQueue\n   ```\n\n2. Install dependencies\n   ```bash\n   pip install -r requirements.txt\n   ```\n\n3. Build and install\n   ```bash\n   python -m build --wheel\n   pip install dist/*.whl\n   ```\n\n<h2 id=\"performance\">📊 Performance</h2>\n\n<p align=\"center\">\n  <img src=\"https://github.com/TransferQueue/community_doc/blob/main/docs/performance_0.1.1.dev2.png?raw=true\" width=\"100%\">\n</p>\n\n> Note: The above benchmark for TransferQueue is based on our naive `SimpleStorageUnit` backend. By introducing high-performance storage backends and optimizing serialization/deserialization, we expect to achieve even better performance. Warmly welcome contributions from the community!\n\nFor detailed performance benchmarks, please refer to [this blog](https://www.yuque.com/haomingzi-lfse7/hlx5g0/tml8ke0zkgn6roey?singleDoc#).\n\nWe also provide a [stress test report](https://www.yuque.com/haomingzi-lfse7/hlx5g0/ydbwgo5k2umaag78?singleDoc#) that demonstrates **768 concurrent clients writing 1.4 TB of data** into TransferQueue across 4 nodes. The system remains stable without any crashes or data loss, achieving 80% bandwidth.\n\n<h2 id=\"customize\"> 🛠️ Customize TransferQueue</h2>\n\n### Define your own data retrieval logic\nWe provide a `BaseSampler` abstraction class, which defines the following interface:\n\n```python3\n@abstractmethod\ndef sample(\n    self,\n    ready_indexes: list[int],\n    batch_size: int,\n    *args: Any,\n    **kwargs: Any,\n) -> tuple[list[int], list[int]]:\n    \"\"\"Sample a batch of indices from the ready indices.\n\n    Args:\n        ready_indexes: List of global indices for which all required fields of the\n        corresponding samples have been produced, and the samples are not labeled as\n        consumed in the corresponding task.\n        batch_size: Number of samples to select\n        *args: Additional positional arguments for specific sampler implementations\n        **kwargs: Additional keyword arguments for specific sampler implementations\n\n    Returns:\n        List of sampled global indices of length batch_size\n        List of global indices of length batch_size that should be labeled as consumed\n        (will never be retrieved in the future)\n\n    Raises:\n        ValueError: If batch_size is invalid or ready_indexes is insufficient\n    \"\"\"\n    raise NotImplementedError(\"Subclasses must implement sample\")\n```\n\nIn this design, we separate data retrieval and data consumption through the two return values, which enables us to easily control sample replacement. We have implemented two reference designs: `SequentialSampler` and `GRPOGroupNSampler`.\n\nThe `Sampler` class or instance should be passed to the `TransferQueueController` during initialization. During each `get_meta` call, you can provide dynamic sampling parameters to the `Sampler`.\n\n```python3\nfrom transfer_queue import TransferQueueController, TransferQueueClient, GRPOGroupNSampler, process_zmq_server_info\n\n# Option 1: Pass the sampler class to the TransferQueueController\ncontroller = TransferQueueController.remote(GRPOGroupNSampler)\n\n# Option 2: Pass the sampler instance to the TransferQueueController (if you need custom configuration)\nyour_own_sampler = YourOwnSampler(config)\ncontroller = TransferQueueController.remote(your_own_sampler)\n\n# Use the sampler\nbatch_meta = client.get_meta(\n    data_fields=[\"input_ids\", \"attention_mask\"],\n    batch_size=8,\n    partition_id=\"train_0\",\n    task_name=\"generate_sequences\",\n    sampling_config={\"n_samples_per_prompt\": 4}  # Put the required sampling parameters here\n)\n```\n\n<span style=\"color: #FF0000;\">**Refer to [tutorial/04_custom_sampler.py](https://github.com/TransferQueue/TransferQueue/blob/main/tutorial/04_custom_sampler.py) for more details.**</span>\n\n\n### How to integrate a new storage backend\n\nThe data plane is organized as follows:\n```text\n  transfer_queue/\n  ├── storage/\n  │   ├── __init__.py\n  │   │── simple_backend.py             # Default distributed storage backend (SimpleStorageUnit) by TQ \n  │   ├── managers/                     # Managers are upper level interfaces that encapsulate the interaction logic with TQ system.\n  │   │   ├── __init__.py\n  │   │   ├──base.py                    # TransferQueueStorageManager, KVStorageManager\n  │   │   ├──simple_backend_manager.py  # AsyncSimpleStorageManager\n  │   │   ├──yuanrong_manager.py        # YuanrongStorageManager\n  │   │   ├──mooncake_manager.py        # MooncakeStorageManager\n  │   │   └──factory.py                 # TransferQueueStorageManagerFactory\n  │   └── clients/                      # Clients are lower level interfaces that directly manipulate the target storage backend.\n  │   │   ├── __init__.py\n  │   │   ├── base.py                   # TransferQueueStorageKVClient\n  │   │   ├── yuanrong_client.py        # YuanrongStorageClient\n  │   │   ├── mooncake_client.py        # MooncakeStorageClient\n  │   │   ├── ray_storage_client.py     # RayStorageClient\n  │   │   └── factory.py                # TransferQueueStorageClientFactory\n```\n\nTo integrate TransferQueue with a custom storage backend, start by implementing a subclass that inherits from `TransferQueueStorageManager`. This subclass acts as an adapter between the TransferQueue system and the target storage backend. For KV-based storage backends, you can simply inherit from `KVStorageManager`, which can serve as the general manager for all KV-based backends.\n\nDistributed storage backends often come with their own native clients serving as the interface of the storage system. In such cases, a low-level adapter for this client can be written, following the examples provided in the `storage/clients` directory.\n\nFactory classes are provided for both `StorageManager` and `StorageClient` to facilitate easy integration. Adding necessary descriptions of required parameters in the factory class helps enhance the overall user experience.\n\n<h2 id=\"contribution\"> ✏️ Contribution Guide</h2>\n\n<span style=\"color: #FF0000;\">**Contributions are warmly welcome!**</span>\n\nNew ideas, feature suggestions, and user experience feedback are all encouraged—feel free to submit issues or PRs. We will respond as soon as possible.\n\nWe recommend using pre-commit for better code format.\n\n```bash\n# install pre-commit\npip install pre-commit\n\n# run the following command in your repo folder, then fix the check before committing your code\npre-commit install && pre-commit run --all-files --show-diff-on-failure --color=always\n```\n\n\n<h2 id=\"citation\"> Citation</h2>\nPlease kindly cite our paper if you find this repo is useful:\n\n```bibtex\n@article{han2025asyncflow,\n  title={AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training},\n  author={Han, Zhenyu and You, Ansheng and Wang, Haibo and Luo, Kui and Yang, Guang and Shi, Wenqi and Chen, Menglong and Zhang, Sicheng and Lan, Zeshun and Deng, Chunshi and others},\n  journal={arXiv preprint arXiv:2507.01663},\n  year={2025}\n}\n```"
  },
  {
    "path": "docs/examples/config.rst",
    "content": ".. _config-explain-page:\n\nConfig Explanation\n===================\n\nLast updated: 06/18/2025.\n\nppo_trainer.yaml for RL FSDP Backend\n-------------------------------------\n\nData\n~~~~\n\n.. code:: yaml\n\n   data:\n     tokenizer: null\n     train_files: ~/data/rlhf/gsm8k/train.parquet\n     val_files: ~/data/rlhf/gsm8k/test.parquet\n     train_max_samples: -1  # set to -1 to use full dataset\n     val_max_samples: -1  # set to -1 to use full dataset\n     prompt_key: prompt\n     max_prompt_length: 512\n     max_response_length: 512\n     train_batch_size: 1024\n     return_raw_input_ids: False  # This should be set to true when the tokenizer between policy and rm differs\n     return_raw_chat: False\n     return_full_prompt: False\n     shuffle: True\n     seed: 42\n     filter_overlong_prompts: False\n     filter_overlong_prompts_workers: 1\n     truncation: error\n     image_key: images\n     trust_remote_code: True\n     custom_cls:\n        path: null\n        name: null\n\n- ``data.train_files``: Training set parquet. Can be a list or a single\n  file. The program will read all files into memory, so it can't be too\n  large (< 100GB). The path can be either local path or HDFS path. For\n  HDFS path, we provide utils to download it to DRAM and convert the\n  HDFS path to local path.\n- ``data.val_files``: Validation parquet. Can be a list or a single\n  file.\n- ``data.train_max_samples``: Maximum number of samples to use from the\n  training dataset. Set to -1 to use the full dataset.\n- ``data.val_max_samples``: Maximum number of samples to use from the\n  validation dataset. Set to -1 to use the full dataset.\n- ``data.prompt_key``: The field in the dataset where the prompt is\n  located. Default is 'prompt'.\n- ``data.max_prompt_length``: Maximum prompt length. All prompts will be\n  left-padded to this length. An error will be reported if the length is\n  too long\n- ``data.max_response_length``: Maximum response length. Rollout in RL\n  algorithms (e.g. PPO) generates up to this length\n- ``data.train_batch_size``: Batch size sampled for one training\n  iteration of different RL algorithms.\n- ``data.return_raw_input_ids``: Whether to return the original\n  input_ids without adding chat template. This is mainly used to\n  accommodate situations where the reward model's chat template differs\n  from the policy. It needs to be decoded first, then apply the RM's\n  chat template. If using a model-based RM, and the policy and RM\n  chat_templates are different, this flag needs to be set\n- ``data.return_raw_chat``: Whether to return the original chat (prompt)\n  without applying chat template.\n- ``data.return_full_prompt``: Whether to return the full prompt with chat template\n- ``data.shuffle``: Whether to shuffle the data in the dataloader.\n- ``data.seed``: An integer seed to use when shuffling the data. If not set or set to\n  `null`, the data shuffling will not be seeded, resulting in a different data order on each run.\n- ``data.filter_overlong_prompts``: Default don't filter.\n- ``data.filter_overlong_prompts_workers``: For large-scale dataset, filtering\n  overlong prompts could be timeconsuming. You cat set the ``filter_overlong_prompts_workers``\n  to use multiprocessing for speed up. Default to 1.\n- ``data.truncation``: Truncate the input_ids or prompt length if they\n  exceed max_prompt_length. Default is 'error', not allow exceed the\n  max_prompt_length. The users should increase the max_prompt_length if\n  throwing the error. You can also set ``left``, ``right`` and ``middle``. \n  When ``middle`` is selected, the logic splits the allowed max length roughly in half \n  and keeps the head and tail of the sequence, effectively discarding the middle section.\n- ``data.image_key``: The field in the multi-modal dataset where the image is\n  located. Default is 'images'.\n- ``data.trust_remote_code``: If the remote tokenizer has python file, we can use this field to allow \n  using remote tokenizer. For example: moonshotai/Moonlight-16B-A3B-Instruct\n\nCustomized Dataset\n~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nCustomized dataset extension is implemented for the SFT trainer and can be extended to other trainers with similar changes.\n\n.. code:: yaml\n\n   custom_cls:\n     path: null\n     name: null\n\n- ``data.custom_cls.path``: The path to the file containing your customized dataset class. If not specified, pre-implemented dataset will be used.\n- ``data.custom_cls.name``: The name of the dataset class within the specified file.\n\nActor/Rollout/Reference Policy\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n\n   actor_rollout_ref:\n    hybrid_engine: True\n    model:\n      path: ~/models/deepseek-llm-7b-chat\n      external_lib: null\n      override_config:\n        attn_implementation: flash_attention_2  # or eager, sdpa - attention implementation override\n        model_config: {}\n        moe_config:  # Megatron only, can adjust moe configuration\n          freeze_moe_router: False  # Megatron only, can freeze moe router (no grad)\n      enable_gradient_checkpointing: False\n      enable_activation_offload: False\n      trust_remote_code: False\n      use_remove_padding: False\n    actor:\n      strategy: fsdp  # This is for backward-compatibility\n      ppo_mini_batch_size: 256\n      ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu\n      ppo_micro_batch_size_per_gpu: 8\n      use_dynamic_bsz: False\n      ppo_max_token_len_per_gpu: 16384 # n * ${data.max_prompt_length} + ${data.max_response_length}\n      grad_clip: 1.0\n      clip_ratio: 0.2\n      entropy_coeff: 0.0\n      use_kl_loss: False # True for GRPO\n      # Rollout Correction (corrects distribution mismatch between rollout and training)\n      rollout_correction:\n        rollout_is: token # IS weights\n        rollout_is_threshold: 2.0 # Upper threshold for IS weights\n        rollout_rs: null # Rejection sampling\n        rollout_rs_threshold: null # RS upper threshold\n      use_torch_compile: True # False to disable torch compile\n      kl_loss_coef: 0.001 # for grpo\n      kl_loss_type: low_var_kl # for grpo\n      ppo_epochs: 1\n      data_loader_seed: null\n      shuffle: False\n      ulysses_sequence_parallel_size: 1 # sp size\n      optim:\n        lr: 1e-6\n        lr_warmup_steps: -1 # Prioritized. Negative values mean delegating to lr_warmup_steps_ratio.\n        lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n        min_lr_ratio: 0.0   # only used with cosine lr scheduler, default to 0.0\n        num_cycles: 0.5     # only used with cosine lr scheduler, default to 0.5\n        lr_scheduler_type: constant  # select from constant/cosine\n        total_training_steps: -1  # must be override by program\n      fsdp_config:\n        wrap_policy:\n          # transformer_layer_cls_to_wrap: None\n          min_num_params: 0\n        param_offload: False\n        optimizer_offload: False\n        fsdp_size: -1\n      checkpoint:\n        # What to include in saved checkpoints\n        # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n        save_contents: ['model', 'optimizer', 'extra']\n        # For more flexibility, you can specify the contents to load from the checkpoint.\n        load_contents: ${actor_rollout_ref.actor.checkpoint.save_contents}\n    ref:\n      fsdp_config:\n        param_offload: False\n        wrap_policy:\n          # transformer_layer_cls_to_wrap: None\n          min_num_params: 0\n      log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n      log_prob_micro_batch_size_per_gpu: 16\n      log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n      log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n      ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size\n    rollout:\n      name: vllm\n      temperature: 1.0\n      top_k: -1 # 0 for hf rollout, -1 for vllm rollout\n      top_p: 1\n      prompt_length: ${data.max_prompt_length}  # not use for opensource\n      response_length: ${data.max_response_length}\n      # for vllm rollout\n      dtype: bfloat16 # should align with FSDP\n      gpu_memory_utilization: 0.5\n      ignore_eos: False\n      enforce_eager: True\n      free_cache_engine: True\n      load_format: dummy_dtensor\n      tensor_model_parallel_size: 2\n      max_num_batched_tokens: 8192\n      max_num_seqs: 1024\n      log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n      log_prob_micro_batch_size_per_gpu: 16\n      log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n      log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n      # for hf rollout\n      do_sample: True\n      engine_kwargs: # inference engine parameters, please refer vllm/sglang official doc for detail\n        vllm: {}\n        sglang: {}\n\n      n: 1 # for each prompt, sample n responses (i.e. num sample times). set it to values > 1 for grpo, rloo\n      calculate_log_probs: False # set to True for computing log probs via rollouts\n      val_kwargs:\n        # sampling parameters for validation\n        top_k: -1 # 0 for hf rollout, -1 for vllm rollout\n        top_p: 1.0\n        temperature: 0\n        n: 1\n        do_sample: False # default eager for validation\n\n      agent:\n        custom_async_server: # Use custom async server implementation for rollout\n          path: null\n          name: null\n\n**Common config for actor, rollout and reference model**\n\n- ``actor_rollout_ref.hybrid_engine``: Whether it's a hybrid engine,\n  currently only supports hybrid engine\n- ``actor_rollout_ref.model.path``: Huggingface model path. This can be\n  either local path or HDFS path. For HDFS path, we provide utils to\n  download it to DRAM and convert the HDFS path to local path.\n- ``actor_rollout_ref.model.external_libs``: Additional Python packages\n  that need to be imported. Used to register models or tokenizers into\n  the Huggingface system.\n- ``actor_rollout_ref.model.override_config``: Used to override some of\n  the model's original configurations. Common overrides include:\n  \n  - ``attn_implementation``: Override the attention implementation. Default is ``flash_attention_2``.\n    Supported values: ``flash_attention_2``, ``eager``, ``sdpa``. Use ``eager`` for debugging or\n    compatibility issues. See :ref:`attention-implementation-override` for detailed usage.\n\n- ``actor_rollout_ref.model.enable_gradient_checkpointing``: FSDP only, decide\n  Whether to enable gradient checkpointing for the actor,\n  Megatron uses recompute options in ``override_transformer_config`` to set this\n- ``actor_rollout_ref.model.enable_activation_offload``: Whether to enable\n  activation offloading for the actor\n- ``actor_rollout_ref.model.trust_remote_code``: Whether to enable loading\n  a remote code model\n- ``actor_rollout_ref.model.use_fused_kernels``: Whether to use fused\n  kernels in the model. If set to True, the following parameters will be\n  used.\n\n  - ``actor_rollout_ref.model.fused_kernel_options.impl_backend``: The\n    implementation backend for fused kernels. Options: \"triton\" or\n    \"torch\". Default is \"torch\".\n    While in megatron, we only support \"triton\" as the\n    implementation backend, so there is no need for this option.\n\n- ``actor_rollout_ref.model.use_remove_padding``: Whether to use remove\n  padding in the model. If set to True, the model will remove padding\n  tokens in the input_ids and response_ids. This helps a lot in improving model running efficiency.\n\n- ``actor_rollout_ref.model.tiled_mlp``: TiledMLP configuration for memory-efficient\n  MLP computation. Reduces peak memory by processing MLP forward/backward in tiles.\n  Only compatible with FSDP2 (requires ``actor_rollout_ref.actor.strategy=fsdp2``).\n\n  - ``actor_rollout_ref.model.tiled_mlp.enabled``: Whether to enable TiledMLP.\n    Default is False.\n  - ``actor_rollout_ref.model.tiled_mlp.num_shards``: Number of shards to split\n    the input. Higher values reduce peak memory but may slightly impact performance.\n    Default is 4.\n\n**Actor model**\n\n- ``actor_rollout_ref.actor.strategy``: fsdp or megatron. In this\n  example, we use fsdp backend.\n\n- ``actor_rollout_ref.actor.ppo_mini_batch_size``: One sample is split\n  into multiple sub-batches with batch_size=ppo_mini_batch_size for PPO\n  updates. The ppo_mini_batch_size is a global num across all workers/gpus\n\n- ``actor_rollout_ref.actor.ppo_micro_batch_size``: [Will be deprecated, use ppo_micro_batch_size_per_gpu] \n  Similar to gradient accumulation, the micro_batch_size_per_gpu for one forward pass,\n  trading speed for GPU memory. The value represent the global view.\n\n- ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``: Similar to gradient\n  accumulation, the micro_batch_size_per_gpu for one forward pass, trading speed\n  for GPU memory. The value represent the local num per gpu.\n\n- ``actor_rollout_ref.actor.grad_clip``: Gradient clipping for actor\n  updates\n- ``actor_rollout_ref.actor.use_kl_loss``: to use kl loss in actor. When used, we are not applying KL in the reward function.\n\n- ``actor_rollout_ref.actor.clip_ratio``: PPO clip ratio\n\n- ``actor_rollout_ref.actor.use_torch_compile``: Whether to use torch compile in actor\n\n- ``actor_rollout_ref.actor.entropy_coeff``: The weight of entropy when\n  calculating PPO loss. The default value is changed to 0.0 since v0.3.x\n\n- ``actor_rollout_ref.actor.ppo_epochs``: Number of epochs for PPO\n  updates on one set of sampled data\n\n- ``actor_rollout_ref.actor.data_loader_seed``: From torch 2.6.0 Megatron backend can get wrong seed generated by pytorch \n  between cp ranks and cause misalignment between data on these ranks, so we shall manually set the seed to avoid hanging\n  issue. if ``actor_rollout_ref.actor.shuffle`` is not null, this must be set.\n\n- ``actor_rollout_ref.actor.shuffle``: Whether to shuffle data when\n  there are multiple epochs\n\n- ``actor_rollout_ref.actor.optim``: Actor's optimizer parameters\n\n- ``actor_rollout_ref.actor.fsdp_config``: FSDP config for actor\n  training\n\n  - ``wrap_policy``: FSDP wrap policy. By default, it uses Huggingface's\n    wrap policy, i.e., wrapping by DecoderLayer\n\n    - No need to set transformer_layer_cls_to_wrap, so we comment it.\n\n  - ``*_offload``: Whether to enable parameter, gradient and optimizer\n    offload\n\n    - Trading speed for GPU memory.\n\n- ``actor_rollout_ref.actor.use_kl_loss``: Whether to enable kl loss. Default is False.\n\n- ``actor_rollout_ref.actor.kl_loss_coef``: The coefficient of kl loss. Default is 0.001. \n\n- ``actor_rollout_ref.actor.kl_loss_type``: Support ``kl`` (``k1``), ``abs``, ``mse`` (``k2``), ``low_var_kl`` (``k3``) and ``full``. Appending ``+`` in the end (e.g., ``k1+`` and ``k3+``) would use straight-through to employ ``k2`` for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/volcengine/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty()` in `core_algos.py <https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/core_algos.py>`_ . See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\n- ``actor_rollout_ref.actor.checkpoint``: The configurations of checkpoint function in actor\n\n  - ``save_contents``: The contents to save in the checkpoint. By default, we save model, optimizer and extra information in the checkpoint.\n    The extra information includes Rng states currently, FSDP supported lr_scheduler, and Megatron opt_param_scheduler will coming soon.\n    We do not store hf_model in checkpoint by default, but we provide a tool in ``scripts/model_merge.py`` to convert checkpoint format to hf format.\n\n  - ``load_contents``: The contents to load in the checkpoint, you can specify different checkpoint loading contents. By default, it is the same with ``save_checkpoint``.\n\n**Reference Model**\n\nReference model will be enabled when ``actor.use_kl_loss`` or/and ``algorithm.use_kl_in_reward`` is/are True.\n\n- ``actor_rollout_ref.ref``: FSDP config same as actor. **For models\n  larger than 7B, it's recommended to turn on offload for ref by\n  default**\n\n- ``actor_rollout_ref.ref.log_prob_micro_batch_size``: [Will be deprecate, use log_prob_micro_batch_size_per_gpu]\n  The batch size for one forward pass in the computation of ``ref_log_prob``. The value represent the global num.\n\n- ``actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu``: The batch size\n  for one forward pass in the computation of ``ref_log_prob``. The value represent the local num per gpu.\n\n**Rollout Model**\n\n- ``actor_rollout_ref.rollout.name``: hf/vllm/sglang.\n\n- Rollout (Auto-regressive) parameters. The key should be equal to the\n  property name in vLLM's ``SamplingParams``.\n\n  - ``temperature``, ``top_k``, ``top_p`` and others: Sampling\n    parameters in ``SamplingParams``.\n\n- ``actor_rollout_ref.rollout.dtype``: Rollout model parameters type. This should be align with\n  the actor model parameter type in FSDP/Megatron backend.\n\n- ``actor_rollout_ref.rollout.gpu_memory_utilization``:\n\n  - For vLLM v0.7.0 and later: The fraction of **total** GPU memory to be used for the vLLM instance.\n  - For SGLang: Corresponding to ``mem_fraction_static``, the fraction of the free GPU memory used for **static** memory like model weights and KV cache. \n\n- ``actor_rollout_ref.rollout.tensor_model_parallel_size``: TP size for rollout. Only effective\n  for vllm.\n\n- ``actor_rollout_ref.rollout.log_prob_micro_batch_size``: [Will be deprecate, use log_prob_micro_batch_size_per_gpu]\n  The batch size for one forward pass in the computation of ``log_prob``. The value represent the global num.\n\n- ``actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu``: Micro batch size per gpu (The batch size for\n  one forward pass) for recalculating ``log_prob``. The value represent the local num per gpu.\n\n- ``actor_rollout_ref.rollout.do_sample``: Whether to sample during training rollout. If set to False, the rollout model\n  will perform greedy sampling.\n\n- ``actor_rollout_ref.rollout.val_kwargs```: Sampling parameters used specifically during validation.\n\n  - ``top_k``: Top-k sampling parameter. Default to -1 for vLLM rollout or 0 for HF rollout.\n  - ``top_p``: Top-p sampling parameter. Default is 1.0 (disabled).\n  - ``temperature``: Sampling temperature. Default is 0 (deterministic greedy).\n  - ``n``: Number of responses to generate during validation. Default is 1.\n  - ``do_sample``: Whether to use sampling during validation. Default is False for\n    deterministic outputs. When set to True, the rollout will use the ``actor_rollout_ref.rollout.val_kwargs`` parameters\n    (top_k, top_p, temperature) to control the sampling behavior.\n\n- ``actor_rollout_ref.rollout.engine_kwargs.vllm``: extra vllm engine args, please refer vllm official doc for detail\n\n- ``actor_rollout_ref.rollout.engine_kwargs.sglang``: extra sglang engine args, please refer sglang official doc for detail\n\n- ``actor_rollout_ref.rollout.ignore_eos``: Whether to ignore the EOS\n  token and continue generating tokens after the EOS token is generated.\n\n- ``actor_rollout_ref.rollout.free_cache_engine``: Offload the KVCache\n  after rollout generation stage. Default is True. When set to True,\n  for vllm v0.5.4 and v0.6.3, we need to disable the usage of CUDAGraph\n  (set ``enforce_eager`` to True.)\n\n- ``actor_rollout_ref.rollout.enforce_eager``: Whether to use CUDAGraph\n  in vLLM generation. Default set to True to disable CUDAGraph.\n\n- ``actor_rollout_ref.rollout.load_format``: Which weight loader to use\n  to load the actor model weights to the rollout model.\n\n  - ``auto``: Use Megatron weight loader.\n  - ``megatron``: Use Megatron weight loader. Deployed with Megatron\n    backend. The input model ``state_dict()`` is already partitioned\n    along TP dimension and already gathered along PP dimension. This\n    weight loader requires that the Rollout model and Actor model's\n    parameters shape and name should be identical.\n  - ``dtensor``: Default solution when using Huggingface weight loader.\n    Deployed with FSDP backend and the state_dict_type is\n    ``StateDictType.SHARDED_STATE_DICT``. Recommend to use this weight\n    loader\n  - ``hf``: Use Huggingface weight loader. Deployed with FSDP backend\n    and the state_dict_type is ``StateDictType.FULL_STATE_DICT``. This\n    solution doesn't need to rewrite the weight loader for each model\n    implemented in vLLM but it results in larger peak memory usage.\n  - ``dummy_hf``, ``dummy_megatron``, ``dummy_dtensor``: Random\n    initialization.\n\n.. note:: **NOTED**: In this config field, users only need to select from ``dummy_megatron``, ``dummy_dtensor``, ``dummy_hf`` for rollout initialization and our hybrid engine will select the corresponding weight loader (i.e., ``megatron``, ``dtensor``, ``hf``) during actor/rollout weight synchronization.\n\n\nMegatron Optimizer and Optimizer Parameter Scheduler\n____________________________________________________\n\n.. code:: yaml\n\n    optim:\n      optimizer: adam\n      lr: 1e-6\n      clip_grad: 1.0\n      total_training_steps: -1  # must be override by program\n      lr_warmup_init: 0.0  # initial learning rate for warmup, default to 0.0\n      lr_warmup_steps: -1 # Prioritized. Negative values mean delegating to lr_warmup_steps_ratio.\n      lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n      lr_decay_steps: null\n      lr_decay_style: constant # select from constant/linear/cosine/inverse_square_root\n      min_lr: 0.0 # minimum learning rate, default to 0.0\n      weight_decay: 0.01\n      weight_decay_incr_style: constant # select from constant/linear/cosine\n      lr_wsd_decay_style: exponential # select from constant/exponential/cosine\n      lr_wsd_decay_steps: null\n      use_checkpoint_opt_param_scheduler: False # use checkpoint optimizer parameter scheduler\n\n\nNotice that there are some differences in APIs between Megatron optimizer and FSDP optimizer.\n\n- Megatron optimizer scheduler names the period after lr_warmup as lr_decay_steps, so the ``lr_scheduler_type`` actually means the style of lr decay after warmup.\n- Megatron optimizer also support weight decay decay mechanism\n- ``use_checkpoint_opt_param_scheduler`` determines whether to use the checkpoint optimizer parameter scheduler. If set to True, the optimizer parameter scheduler will be saved in the checkpoint and loaded from the checkpoint during resuming training.\n\nFor learning rate decay, original Megatron pretrain default option of ``lr_decay_style`` is ``linear``,\nmeaning that the learning rate will be linearly decayed from the initial learning rate to ``min_lr`` within the\n``lr_decay_steps``. However, in verl, to align with FSDP's default behavior, we set the default\n``lr_decay_style`` to ``constant``, meaning that the learning rate will be kept constant after the warmup stage.\n\n\nCritic Model\n~~~~~~~~~~~~\n\nMost parameters for Critic are similar to Actor Model.\n\nReward Model\n~~~~~~~~~~~~\n\n.. code:: yaml\n\n   reward_model:\n     enable: False\n     model:\n       input_tokenizer: ${actor_rollout_ref.model.path}  # set this to null if the chat template is identical\n       path: ~/models/Anomy-RM-v0.1\n       external_lib: ${actor_rollout_ref.model.external_lib}\n       trust_remote_code: False\n       fsdp_config:\n         min_num_params: 0\n         param_offload: False\n     micro_batch_size_per_gpu: 16\n     max_length: null\n     reward_manager: naive\n\n- ``reward_model.enable``: Whether to enable reward model. If False, we\n  compute the reward only with the user-defined reward functions. In\n  GSM8K and Math examples, we disable reward model. For RLHF alignment\n  example using full_hh_rlhf, we utilize reward model to assess the\n  responses. If False, the following parameters are not effective.\n- ``reward_model.model``\n\n  - ``input_tokenizer``: Input tokenizer. If the reward model's chat\n    template is inconsistent with the policy, we need to first decode to\n    plaintext, then apply the rm's chat_template. Then score with RM. If\n    chat_templates are consistent, it can be set to null.\n  - ``path``: RM's HDFS path or local path. Note that RM only supports\n    AutoModelForSequenceClassification. Other model types need to define\n    their own RewardModelWorker and pass it from the code.\n  - ``trust_remote_code``: Whether to enable loading a remote code model,\n    default to False.\n- ``reward_model.reward_manager``:  Reward Manager. This defines the mechanism\n  of computing rule-based reward and handling different reward sources. Default\n  is ``naive``. If all verification functions are multiprocessing-safe, the reward\n  manager can be set to ``prime`` for parallel verification.\n\nCustomized Reward Function\n~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n  \n   custom_reward_function:\n     path: null\n     name: compute_score\n\n- ``custom_reward_function.path``: The path to the file containing your customized reward function. If not specified, pre-implemented reward functions will be used.\n- ``custom_reward_function.name`` (Optional) : The name of the reward function within the specified file. Default is 'compute_score'.\n\nAlgorithm\n~~~~~~~~~\n\n.. code:: yaml\n\n   algorithm:\n     gamma: 1.0\n     lam: 1.0\n     adv_estimator: gae\n     use_kl_in_reward: False\n     kl_penalty: kl  # how to estimate kl divergence\n     kl_ctrl:\n       type: fixed\n       kl_coef: 0.005\n       horizon: 10000\n       target_kl: 0.1\n     # Rollout Correction\n     rollout_correction:\n       rollout_is: null  # IS weights\n       rollout_is_threshold: 2.0  # Upper threshold for IS weights\n       rollout_rs: null  # Rejection sampling\n       rollout_rs_threshold: null  # RS upper threshold\n\n- ``gamma``: discount factor\n- ``lam``: Trade-off between bias and variance in the GAE estimator\n- ``adv_estimator``: Support ``gae``, ``grpo``, ``reinforce_plus_plus``, ``reinforce_plus_plus_baseline``, ``rloo``, ``rloo_vectorized``, ``grpo_vectorized``\n- ``use_kl_in_reward``: Whether to enable in-reward kl penalty. Default is False.\n- ``kl_penalty``: Support ``kl``, ``abs``, ``mse``, ``low_var_kl`` and ``full``. How to\n  calculate the kl divergence between actor and reference policy. For\n  specific options, refer to `kl_penalty()` in `core_algos.py <https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/core_algos.py>`_ .\n- ``kl_ctrl``: Config for in-reward kl_penalty controller\n\n  - ``kl_coef``: The (initial) coefficient of in-reward kl_penalty. Default is 0.001.\n  - ``type``: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController.\n  - ``horizon`` and ``target_kl``: See source code of AdaptiveKLController for details.\n\n- ``rollout_correction``: Rollout Correction configuration (nested dict). Set to ``null`` to disable.\n  When enabled, contains:\n\n  - ``rollout_is``: IS weights aggregation level, ``null`` to disable IS weights.\n  - ``rollout_is_threshold``: Upper threshold for IS weights (e.g., 2.0).\n  - ``rollout_rs``: Rejection sampling mode, ``null`` to disable RS.\n  - ``rollout_rs_threshold``: RS upper threshold.\n\n  Note: Rollout Correction requires setting ``actor_rollout_ref.rollout.calculate_log_probs=True``.\n\nTrainer\n~~~~~~~\n\n.. code:: yaml\n\n   trainer:\n     total_epochs: 30\n     project_name: verl_examples\n     experiment_name: gsm8k\n     logger: ['console', 'wandb']\n     log_val_generations: 0\n     nnodes: 1\n     n_gpus_per_node: 8\n     save_freq: -1\n     val_before_train: True\n     test_freq: 2\n     critic_warmup: 0\n     default_hdfs_dir: null # hdfs checkpoint path\n     default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name} # local checkpoint path\n     resume_mode: auto # or disable or resume_path if resume_from_path is set\n     resume_from_path: null\n     remove_previous_ckpt_in_save: False\n     del_local_ckpt_after_load: False\n     ray_wait_register_center_timeout: 300\n\n- ``trainer.total_epochs``: Number of epochs in training.\n- ``trainer.project_name``: For wandb, swanlab, mlflow\n- ``trainer.experiment_name``: For wandb, swanlab, mlflow\n- ``trainer.logger``: Support console and wandb, swanlab, mlflow, tensorboard, trackio\n- ``trainer.log_val_generations``: The number of logged generation during validation (default ``0``)\n- ``trainer.nnodes``: Number of nodes used in the training.\n- ``trainer.n_gpus_per_node``: Number of GPUs per node.\n- ``trainer.save_freq``: The frequency (by iteration) to save checkpoint\n  of the actor and critic model.\n- ``trainer.val_before_train``: Whether to run validation before training.\n- ``trainer.test_freq``: The validation frequency (by iteration).\n- ``trainer.critic_warmup``: The number of iteration to train the critic\n  model before actual policy learning.\n- ``trainer.resume_mode``: The mode of resuming training. Support\n  ``disable``, ``auto`` and ``resume_path``. If set to ``auto`` as default, the\n  program will automatically resume from the latest checkpoint in the\n  ``default_local_dir``. If set to ``resume_path``, the program will resume\n  from the path specified in ``resume_from_path``.\n- ``trainer.resume_from_path``: The path to resume training from. Only\n  effective when ``resume_mode`` is set to ``resume_path``.\n- ``trainer.remove_previous_ckpt_in_save``: Whether to remove previous\n  checkpoints in the save directory. Default is False.\n- ``trainer.del_local_ckpt_after_load``: Whether to delete local\n  checkpoints after loading them. Default is False.\n- ``trainer.ray_wait_register_center_timeout``: The timeout for waiting\n  for the ray register center to be ready. Default is 300 seconds.\n\n\nThis figure illustrates how the configurations affect the training.\n\nhttps://excalidraw.com/#json=pfhkRmiLm1jnnRli9VFhb,Ut4E8peALlgAUpr7E5pPCA\n\n.. image:: https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d\n\n\nevaluation.yaml\n---------------\n\nData\n~~~~\n\n.. code:: yaml\n\n   data:\n     path: /tmp/math_Qwen2-7B-Instruct.parquet\n     prompt_key: prompt\n     response_key: responses\n     data_source_key: data_source\n     reward_model_key: reward_model\n\n- ``data.path``: Path to the dataset file (Parquet format).\n- ``data.prompt_key``: The field in the dataset where the prompt is located. Default is 'prompt'.\n- ``data.response_key``: The key holds the generated responses. This should be a list of strings representing the responses. Default is 'responses'.\n- ``data.data_source_key``: This is used to separate metric calculations for different data sources, ensuring that metrics are calculated independently for each source.\n- ``data.reward_model_key``: The key holds the reference answers. These reference answers typically serve as the ground truth or test cases for the task.\n\nCustomized Reward Function\n~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code:: yaml\n  \n   custom_reward_function:\n     path: null\n     name: compute_score\n\n- ``custom_reward_function.path``: The path to the file containing your customized reward function. If not specified, pre-implemented reward functions will be used.\n- ``custom_reward_function.name`` (Optional) : The name of the reward function within the specified file. Default is 'compute_score'.\n\nsft_trainer.yaml for SFT FSDP Backend\n--------------------------------------\n\n\nOptim\n~~~~~~~\n\n.. code:: yaml\n\n   optim:\n     optimizer: AdamW\n     optimizer_impl: torch.optim\n     lr: 1e-5\n     weight_decay: 0.01\n     lr_warmup_steps_ratio: 0.1\n     clip_grad: 1.0\n     lr_scheduler: cosine\n     override_optimizer_config: null\n\n- ``optimizer``: Optimizer class name (e.g., ``\"AdamW\"``, ``\"AdamW8bit\"``, ``\"_AdamW\"``). The class name as it appears in the module.\n- ``optimizer_impl``: Module path to import optimizer from (e.g., ``\"torch.optim\"``, ``\"torchao.optim\"``, ``\"bitsandbytes.optim\"``).\n- ``optim.lr``: Learning rate for the optimizer.\n- ``optim.weight_decay``: Weight decay for the optimizer.\n- ``optim.lr_warmup_steps_ratio``: Ratio of warmup steps to total training steps.\n- ``optim.clip_grad``: Gradient clipping value.\n- ``optim.lr_scheduler``: Learning rate scheduler type. Options:\n\n  - ``cosine``: Cosine learning rate scheduler with warmup (default).\n  - ``wsd``: Warmup-Stable-Decay scheduler that provides a stable learning rate phase between warmup and decay phases.\n\n- ``override_optimizer_config``: Dictionary of additional optimizer-specific keyword arguments. For example, to use ``torchao.optim``'s ``_AdamW`` with BF16 stochastic rounding: ``{\"bf16_stochastic_round\": true}``\n\nModel\n~~~~~~~~~~~~\n\nMost parameters for Model are similar to Reward Model.\n\n.. code:: yaml\n\n   model:\n     partial_pretrain: ~/models/gemma-1.1-7b-it\n     fsdp_config:\n       model_dtype: fp32\n       wrap_policy:\n         min_num_params: 0\n       cpu_offload: False\n       offload_params: False\n     external_lib: null\n     enable_gradient_checkpointing: False\n     trust_remote_code: False\n     lora_rank: 0\n     lora_alpha: 16\n     target_modules: all-linear\n     use_liger: False\n\n- ``partial_pretrain``: HDFS path or local path for the pretrained model.\n- ``fsdp_config``\n\n  - ``model_dtype``: Model parameters type, default to ``fp32``.\n    Support: ``bf16``, ``fp16``, ``fp32``.\n  - ``cpu_offload``: Whether to enable CPU offloading for FSDP. If True,\n    the offload_params will be used as argument.\n  - ``offload_params``: Whether to offload parameters to CPU\n    when not involved in computation. If True, then this offloads gradients\n    to CPU as well, meaning that the optimizer step runs on CPU.\n\n- ``lora_rank``: The rank of the LoRA model, default to 0. If ``lora_rank``>0,\n  we will train LoRA modules instead of tuning the full model.\n- ``lora_alpha``: The alpha parameter for LoRA scaling, default to 16.\n- ``target_modules``: The names of the modules to apply the adapter to,\n  default to ``all-linear``. See `peft docs <https://huggingface.co/docs/peft/v0.15.0/en/package_reference/lora#peft.LoraConfig.target_modules>`_ for detail.\n\n- ``use_liger``: Whether to enable Liger kernel, default to False. If True,\n  we apply Liger kernel to the model (depends on `liger-kernel`).\n"
  },
  {
    "path": "docs/examples/gsm8k_example.rst",
    "content": "GSM8K Example\n=============\n\nLast updated: 03/25/2025.\n\nIntroduction\n------------\n\nIn this example, we train an LLM to tackle the GSM8k task.\n\nPaper: https://arxiv.org/pdf/2110.14168\n\nDataset: https://huggingface.co/datasets/openai/gsm8k\n\nNote that the original paper mainly focuses on training a verifier (a\nreward model) to solve math problems via Best-of-N sampling. In this\nexample, we train an RLHF agent using a rule-based reward model.\n\nDataset Introduction\n--------------------\n\nGSM8k is a math problem dataset. The prompt is an elementary school\nproblem. The LLM model is required to answer the math problem.\n\nThe training set contains 7473 samples and the test set contains 1319\nsamples.\n\n**An example**\n\nPrompt\n\n   Katy makes coffee using teaspoons of sugar and cups of water in the\n   ratio of 7:13. If she used a total of 120 teaspoons of sugar and cups\n   of water, calculate the number of teaspoonfuls of sugar she used.\n\nSolution\n\n   The total ratio representing the ingredients she used to make the\n   coffee is 7+13 = <<7+13=20>>20 Since the fraction representing the\n   number of teaspoons she used is 7/20, she used 7/20\\ *120 =\n   <<7/20*\\ 120=42>>42 #### 42\n\nStep 1: Prepare dataset\n-----------------------\n\n.. code:: bash\n\n   cd examples/data_preprocess\n   python3 gsm8k.py --local_save_dir ~/data/gsm8k\n\nStep 2: Download Model\n----------------------\n\nThere're three ways to prepare the model checkpoints for post-training:\n\n- Download the required models from huggingface or modelscope\n\n.. code:: bash\n\n   hf download deepseek-ai/deepseek-math-7b-instruct --local-dir ~/models/deepseek-math-7b-instruct --local-dir-use-symlinks False\n   # or\n   modelscope download --model deepseek-ai/deepseek-math-7b-instruct --local_dir ~/models/deepseek-math-7b-instruct\n\n- Already store your store model in the local directory or HDFS path.\n- Also, you can directly use the model name in huggingface (e.g.,\n  deepseek-ai/deepseek-math-7b-instruct) in\n  ``actor_rollout_ref.model.path`` and ``critic.model.path`` field in\n  the run script. You can also download models from modelscope by setting environmental variable ``VERL_USE_MODELSCOPE=True``.\n  See examples/ppo_trainer/run_deepseek7b_llm_modelscope.sh for example.\n\nNoted that users should prepare checkpoints for actor, critic and reward\nmodel.\n\n[Optional] Step 3: SFT your Model\n---------------------------------\n\nWe provide a SFT Trainer using PyTorch FSDP in\n`sft_trainer.py <https://github.com/volcengine/verl/blob/main/verl/trainer/sft_trainer.py>`_. \nUsers can customize their own SFT\nscript using our FSDP SFT Trainer.\n\nWe also provide various training scripts for SFT on GSM8K dataset in `gsm8k sft directory <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/>`_.\n\n.. code:: shell\n\n   set -x\n\n   torchrun -m verl.trainer.sft_trainer \\\n       data.train_files=$HOME/data/gsm8k/train.parquet \\\n       data.val_files=$HOME/data/gsm8k/test.parquet \\\n       data.messages_key=messages \\\n       data.micro_batch_size_per_gpu=8 \\\n       model.path=deepseek-ai/deepseek-coder-6.7b-instruct \\\n       trainer.project_name=gsm8k-sft \\\n       trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \\\n       trainer.total_epochs=4 \\\n       trainer.logger='[\"console\",\"wandb\"]'\n\n\nIf you use AMD GPUs (ROCm kernel), you need to add the following environment variables into the run script:\n\n    .. code-block:: bash\n\n        export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n        export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n        export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n\n\nStep 4: Perform PPO training with your model on GSM8K Dataset\n-------------------------------------------------------------\n\n- Prepare your own run.sh script. Here's an example for GSM8k dataset\n  and deepseek-llm-7b-chat model.\n- Users could replace the ``data.train_files`` ,\\ ``data.val_files``,\n  ``actor_rollout_ref.model.path`` and ``critic.model.path`` based on\n  their environment.\n- See :doc:`config` for detailed explanation of each config field.\n\n**Reward Model/Function**\n\nWe use a rule-based reward model. We force the model to produce a final\nanswer following 4 “#” as shown in the solution. We extract the final\nanswer from both the solution and model's output using regular\nexpression matching. We compare them and assign a reward of 1 to correct\nanswer, 0.1 to incorrect answer and 0 to no answer.\n\n**Training Script**\n\nThe training script example for FSDP and Megatron-LM backend are stored in examples/ppo_trainer directory.\n\n.. code:: bash\n\n   cd ../ppo_trainer\n   bash run_deepseek7b_llm.sh\n\nThe script of run_deepseek7b_llm.sh\n\n.. code:: bash\n\n   set -x\n\n   python3 -m verl.trainer.main_ppo \\\n      data.train_files=$HOME/data/gsm8k/train.parquet \\\n      data.val_files=$HOME/data/gsm8k/test.parquet \\\n      data.train_batch_size=1024 \\\n      data.max_prompt_length=512 \\\n      data.max_response_length=512 \\\n      actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n      actor_rollout_ref.actor.optim.lr=1e-6 \\\n      actor_rollout_ref.model.use_remove_padding=True \\\n      actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n      actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n      actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n      actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n      actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n      actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n      actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n      actor_rollout_ref.rollout.name=vllm \\\n      actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n      actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n      actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n      critic.optim.lr=1e-5 \\\n      critic.model.use_remove_padding=True \\\n      critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n      critic.model.enable_gradient_checkpointing=True \\\n      critic.ppo_micro_batch_size_per_gpu=32 \\\n      critic.model.fsdp_config.param_offload=False \\\n      critic.model.fsdp_config.optimizer_offload=False \\\n      algorithm.kl_ctrl.kl_coef=0.001 \\\n      trainer.critic_warmup=0 \\\n      trainer.logger='[\"console\",\"wandb\"]' \\\n      trainer.project_name='verl_example_gsm8k' \\\n      trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n      trainer.n_gpus_per_node=8 \\\n      trainer.nnodes=1 \\\n      trainer.save_freq=-1 \\\n      trainer.test_freq=1 \\\n      trainer.total_epochs=15 $@\n\n\nIf you use AMD GPUs (ROCm kernel), you need to add the following environment variables into the run script:\n\n    .. code-block:: bash\n\n        export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n        export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n        export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n\nIf you encounter any issues in using AMD GPUs running VeRL, feel free to contact me - `Yusheng Su <https://yushengsu-thu.github.io/>`_."
  },
  {
    "path": "docs/examples/multi_modal_example.rst",
    "content": "Multi-Modal Example Architecture\n=================================\n\nLast updated: 04/28/2025.\n\nIntroduction\n------------\n\nNow, verl has supported multi-modal training. You can use fsdp and \nvllm/sglang to start a multi-modal RL task. Megatron supports is also \non the way.\n\nFollow the steps below to quickly start a multi-modal RL task.\n\nStep 1: Prepare dataset\n-----------------------\n\n.. code:: python\n\n    # it will be saved in the $HOME/data/geo3k folder\n    python examples/data_preprocess/geo3k.py\n\nStep 2: Download Model\n----------------------\n\n.. code:: bash\n\n    # download the model from huggingface\n    python3 -c \"import transformers; transformers.pipeline(model='Qwen/Qwen2.5-VL-7B-Instruct')\"\n\nStep 3: Perform GRPO training with multi-modal model on Geo3K Dataset\n---------------------------------------------------------------------\n\n.. code:: bash\n\n    # run the task\n    bash examples/grpo_trainer/run_qwen2_5_vl-7b.sh\n\n\n\n\n\n\n\n\n"
  },
  {
    "path": "docs/examples/ppo_code_architecture.rst",
    "content": "PPO Example Architecture\n========================\n\nLast updated: 02/17/2025.\n\nLet's start with the Proximal Policy Optimization algorithm, which is\nmost widely used algorithm in LLM post-training.\n\nThe main entry point of the PPO algorithm example is:\n`main_ppo.py <https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py>`_.\nIn this tutorial, we will go through the code architecture in `main_ppo.py <https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py>`_.\n\nDefine the data\n---------------\n\nUsers need to preprocess and store the dataset in parquet files.\nAnd we implement `RLHFDataset` to load and tokenize the parquet files.\n\nFor ``RLHFDataset`` (Default), at least 1 fields are required:\n\n- ``prompt``: Contains the string prompt\n\nWe already provide some examples of processing the datasets to parquet\nfiles in `data_preprocess directory <https://github.com/volcengine/verl/blob/main/examples/data_preprocess>`_. Currently, we support\npreprocess of GSM8k, MATH, Hellasage, Full_hh_rlhf datasets. See :doc:`../preparation/prepare_data` for\nmore information.\n\nDefine the reward functions for different datasets\n--------------------------------------------------\n\nIn this main entry point, the users only need to define their own reward\nfunction based on the datasets (or applications) utilized in PPO\ntraining.\n\nFor example, we already provide reward functions for `GSM8k <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score/gsm8k.py>`_ \nand `MATH <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score/math.py>`_\ndatasets in the ``_select_rm_score_fn``. In the ``RewardManager``, we\nwill compute the reward score based on the data_source to select\ncorresponding reward functions. For some RLHF datasets (e.g.,\nfull_hh_rlhf), the reward model is utilized to assess the responses\nwithout any reward functions. In this case, the ``RewardManager`` will\nreturn the ``rm_score`` computed by the reward model directly.\n\nSee `reward functions <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score>`_ for detailed implementation.\n\nDefine worker classes\n---------------------\n\n.. code:: python\n\n   if config.actor_rollout_ref.actor.strategy in {\"fsdp\", \"fsdp2\"}: # for FSDP backend\n       assert config.critic.strategy in {\"fsdp\", \"fsdp2\"}\n       from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker\n       from verl.single_controller.ray import RayWorkerGroup\n       ray_worker_group_cls = RayWorkerGroup\n\n   elif config.actor_rollout_ref.actor.strategy == 'megatron': # for Megatron backend\n       assert config.actor_rollout_ref.actor.strategy == config.critic.strategy\n       from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker\n       from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup\n       ray_worker_group_cls = NVMegatronRayWorkerGroup # Ray worker class for Megatron-LM\n\n   else:\n       raise NotImplementedError\n\n   from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role\n\n   role_worker_mapping = {\n       Role.ActorRollout: ActorRolloutRefWorker,\n       Role.Critic: CriticWorker,\n       Role.RefPolicy: ActorRolloutRefWorker\n   }\n\n   global_pool_id = 'global_pool'\n   resource_pool_spec = {\n       global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,\n   }\n   mapping = {\n       Role.ActorRollout: global_pool_id,\n       Role.Critic: global_pool_id,\n       Role.RefPolicy: global_pool_id,\n   }\n\nStep 1: Construct the mapping between roles and workers\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nA role represents a group of workers in the same process. We have\npre-defined several roles in `ray_trainer.py <https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/ray_trainer.py#L38>`_.\n\n.. code:: python\n\n   class Role(Enum):\n       \"\"\"\n       To create more roles dynamically, you can subclass Role and add new members\n       \"\"\"\n       Actor = 0  # This worker only has Actor\n       Rollout = 1 # This worker only has Rollout\n       ActorRollout = 2 # This worker has both actor and rollout, it's a HybridEngine\n       Critic = 3 # This worker only has critic\n       RefPolicy = 4 # This worker only has reference policy\n       RewardModel = 5 # This worker only has reward model\n       ActorRolloutRef = 6 # This worker contains actor, rollout and reference policy simultaneously \n\nStep 2: Define the worker class corresponding to this role\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n- We have pre-implemented the ``ActorRolloutRefWorker``. Through\n  different configs, it can be a standalone actor, a standalone rollout,\n  an ActorRollout HybridEngine, or an ActorRolloutRef HybridEngine\n- We also pre-implemented workers for ``Actor``, ``Rollout``,\n  ``Critic``, ``Reward Model`` and ``Reference model`` on two different\n  backend: PyTorch FSDP\n  and Megatron-LM.\n  See `FSDP Workers <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py>`_ \n  and `Megatron-LM Workers <https://github.com/volcengine/verl/blob/main/verl/workers/megatron_workers.py>`_\n  for more information.\n\nStep 3: Define resource pool id and resource pool spec\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n- Resource pool is a division of global GPU resources,\n  ``resource_pool_spec`` is a dict, mapping from id to # of GPUs\n\n  - In the above example, we defined a global resource pool:\n    global_pool_id, and then put all roles on this one resource pool\n    with all the GPUs in this post-training task. This refers to\n    *co-locate* placement where all the models share the same set of\n    GPUs.\n\n- See resource pool and placement for advance usage.\n\nDefining reward model/function\n------------------------------\n\n.. code:: python\n\n   # we should adopt a multi-source reward function here\n   # - for rule-based rm, we directly call a reward score\n   # - for model-based rm, we call a model\n   # - for code related prompt, we send to a sandbox if there are test cases\n   # - finally, we combine all the rewards together\n   # - The reward type depends on the tag of the data\n   if config.reward_model.enable:\n       from verl.workers.fsdp_workers import RewardModelWorker\n       role_worker_mapping[Role.RewardModel] = RewardModelWorker\n       mapping[Role.RewardModel] = global_pool_id\n    \n   reward_fn = RewardManager(tokenizer=tokenizer, num_examine=0)\n\n   # Note that we always use function-based RM for validation\n   val_reward_fn = RewardManager(tokenizer=tokenizer, num_examine=1)\n\n   resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n\nSince not all tasks use model-based RM, users need to define here\nwhether it's a model-based RM or a function-based RM\n\n- If it's a model-based RM, directly add the ``RewardModel`` role in the\n  resource mapping and add it to the resource pool mapping.\n\n  - Note that the pre-defined ``RewardModelWorker`` only supports models\n    with the structure of huggingface\n    ``AutoModelForSequenceClassification``. If it's not this model, you\n    need to define your own RewardModelWorker in `FSDP Workers <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py>`_ \n    and `Megatron-LM Workers <https://github.com/volcengine/verl/blob/main/verl/workers/megatron_workers.py>`_.\n\n- If it's a function-based RM, the users are required to classified the\n  reward function for each datasets.\n\n.. code:: python\n\n   def _select_rm_score_fn(data_source):\n       if data_source == 'openai/gsm8k':\n           return gsm8k.compute_score\n       elif data_source == 'lighteval/MATH':\n           return math.compute_score\n       else:\n           raise NotImplementedError\n\nSee reward functions implemented in `directory <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score/>`_ \nfor more information.\n\nDefine, init and run the PPO Trainer\n------------------------------------\n\n.. code:: python\n\n   trainer = RayPPOTrainer(config=config,\n                           tokenizer=tokenizer,\n                           role_worker_mapping=role_worker_mapping,\n                           resource_pool_manager=resource_pool_manager,\n                           ray_worker_group_cls=ray_worker_group_cls,\n                           reward_fn=reward_fn,\n                           val_reward_fn=val_reward_fn)\n   trainer.init_workers()\n   trainer.fit()\n\n- We first initialize the ``RayPPOTrainer`` with user config, tokenizer\n  and all the above worker mapping, resource pool, worker group and\n  reward functions\n- We first call the ``trainer.init_workers()`` to initialize the models\n  on the allocated GPUs (in the resource pool)\n- The actual PPO training will be executed in ``trainer.fit()``\n\nverl can be easily extended to other RL algorithms by reusing the Ray\nmodel workers, resource pool and reward functions. See :doc:`extension<../advance/dpo_extension>` for\nmore information.\n\nDetails of the ``RayPPOTrainer`` is discussed in :doc:`Ray Trainer<../workers/ray_trainer>`.\n"
  },
  {
    "path": "docs/examples/sandbox_fusion_example.rst",
    "content": "Sandbox Fusion Example\n============================\n\nLast updated: 06/27/2025.\n\nIntroduction\n------------\n\nSandbox Fusion is a remote code sandbox service that provides a secure environment for running and evaluating code generated by Large Language Models (LLMs). This example demonstrates how to train an LLM and use Sandbox Fusion to verify generated code, enhancing both security and performance.\n\nBy leveraging a remote code sandbox service with greater CPU resources for concurrent code verification, you can reduce the reward stage time by 10-30%, depending on the quality of the generated code.\n\nStep 1: Prepare the Dataset\n---------------------------\n\nWe use the Eurus-2-RL-Data dataset for training. This dataset combines math and code questions, making it suitable for LLM training tasks. You can download it from HuggingFace: `Eurus-2-RL-Data Dataset <https://huggingface.co/datasets/PRIME-RL/Eurus-2-RL-Data>`_.\n\nStep 2: Set Up the Sandbox Fusion Service\n-----------------------------------------\n\nSandbox Fusion is a remote code sandbox service designed to securely run and evaluate LLM-generated code. To use it:\n\n1. **Access Full Documentation**: For detailed setup instructions, refer to the `Sandbox Fusion Documentation <https://bytedance.github.io/SandboxFusion/>`_.\n2. **Deploy the Service**: Choose one of the following deployment methods:\n\n   - **Local Deployment**: Follow the guide `here <https://bytedance.github.io/SandboxFusion/docs/docs/get-started#local-deployment>`_.\n   - **FaaS Instance (Volcengine)**: Create an instance using the `Volcengine Documentation <https://www.volcengine.com/docs/6662/1539235>`_.\n\nAfter deployment, you will receive an API endpoint in the format: ``https://<ip-address-or-domain-name>/run_code``.\n\nStep 3: Configure the Training Script\n-------------------------------------\n\nTo integrate Sandbox Fusion into your training script, configure the following parameters:\n\n**Key Settings for Sandbox Fusion**\n\n- ``reward_model.sandbox_fusion.url='<API-endpoint>'``: Enable Sandbox Fusion by specifying the API endpoint (must end with ``/run_code``).\n- ``reward_model.sandbox_fusion.max_concurrent=256``: Set the maximum number of concurrent API requests to the Sandbox Fusion service.\n- ``reward_model.sandbox_fusion.memory_limit_mb=1024``: Set the memory limit (in MB) for each sandbox instance. Defaults to 1024MB if not specified.\n\n**Additional Optimization**\n\nTo further reduce code verification time, enable parallel processing with:  \n\n- ``reward_model.reward_manager=prime``: The Prime reward manager verifies code across multiple subprocesses concurrently.\n\n**Example Script**\n\nFor a practical implementation, refer to the example script:  \n\n``examples/ppo_trainer/run_deepseek7b_llm_sandbox_fusion.sh``\n\nOnce you’ve set your API endpoint in the script, you can start the training job."
  },
  {
    "path": "docs/examples/skypilot_examples.rst",
    "content": "SkyPilot Examples\n=================\n\nLast updated: 09/04/2025.\n\nThis guide provides examples of running VERL reinforcement learning training on Kubernetes clusters or cloud platforms with GPU nodes using `SkyPilot <https://github.com/skypilot-org/skypilot>`_.\n\nInstallation and Configuration\n-------------------------------\n\nStep 1: Install SkyPilot\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\nChoose the installation based on your target platform:\n\n.. code-block:: bash\n\n   # For Kubernetes only\n   pip install \"skypilot[kubernetes]\"\n   \n   # For AWS\n   pip install \"skypilot[aws]\"\n   \n   # For Google Cloud Platform\n   pip install \"skypilot[gcp]\"\n   \n   # For Azure\n   pip install \"skypilot[azure]\"\n   \n   # For multiple platforms\n   pip install \"skypilot[kubernetes,aws,gcp,azure]\"\n\nStep 2: Configure Your Platform\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nSee https://docs.skypilot.co/en/latest/getting-started/installation.html\n\nStep 3: Set Up Environment Variables\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nExport necessary API keys for experiment tracking:\n\n.. code-block:: bash\n\n   # For Weights & Biases tracking\n   export WANDB_API_KEY=\"your-wandb-api-key\"\n   \n   # For HuggingFace gated models (if needed)\n   export HF_TOKEN=\"your-huggingface-token\"\n\nExamples\n--------\n\nAll example configurations are available in the `examples/skypilot/ <https://github.com/volcengine/verl/tree/main/examples/skypilot>`_ directory on GitHub. See the `README <https://github.com/volcengine/verl/blob/main/examples/skypilot/README.md>`_ for additional details.\n\nPPO Training\n~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky launch -c verl-ppo verl-ppo.yaml --secret WANDB_API_KEY -y\n\nRuns PPO training on GSM8K dataset using Qwen2.5-0.5B-Instruct model across 2 nodes with H100 GPUs. Based on examples in ``examples/ppo_trainer/``.\n\n`View verl-ppo.yaml on GitHub <https://github.com/volcengine/verl/blob/main/examples/skypilot/verl-ppo.yaml>`_\n\nGRPO Training\n~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky launch -c verl-grpo verl-grpo.yaml --secret WANDB_API_KEY -y\n\nRuns GRPO (Group Relative Policy Optimization) training on MATH dataset using Qwen2.5-7B-Instruct model. Memory-optimized configuration for 2 nodes. Based on examples in ``examples/grpo_trainer/``.\n\n`View verl-grpo.yaml on GitHub <https://github.com/volcengine/verl/blob/main/examples/skypilot/verl-grpo.yaml>`_\n\nMulti-turn Tool Usage Training\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky launch -c verl-multiturn verl-multiturn-tools.yaml \\\n     --secret WANDB_API_KEY --secret HF_TOKEN -y\n\nSingle-node training with 8xH100 GPUs for multi-turn tool usage with Qwen2.5-3B-Instruct. Includes tool and interaction configurations for GSM8K. Based on examples in ``examples/sglang_multiturn/`` but uses vLLM instead of sglang.\n\n`View verl-multiturn-tools.yaml on GitHub <https://github.com/volcengine/verl/blob/main/examples/skypilot/verl-multiturn-tools.yaml>`_\n\nConfiguration\n-------------\n\nThe example YAML files are pre-configured with:\n\n- **Infrastructure**: Kubernetes clusters (``infra: k8s``) - can be changed to ``infra: aws`` or ``infra: gcp``, etc.\n- **Docker Image**: VERL's official Docker image with CUDA 12.6 support\n- **Setup**: Automatically clones and installs VERL from source\n- **Datasets**: Downloads required datasets during setup phase\n- **Ray Cluster**: Configures distributed training across nodes\n- **Logging**: Supports Weights & Biases via ``--secret WANDB_API_KEY``\n- **Models**: Supports gated HuggingFace models via ``--secret HF_TOKEN``\n\nLaunch Command Options\n----------------------\n\n- ``-c <name>``: Cluster name for managing the job\n- ``--secret KEY``: Pass secrets for API keys (can be used multiple times)\n- ``-y``: Skip confirmation prompt\n\nMonitoring Your Jobs\n--------------------\n\nCheck Cluster Status\n~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky status\n\nView Logs\n~~~~~~~~~\n\n.. code-block:: bash\n\n   sky logs verl-ppo  # View logs for the PPO job\n\nSSH into Head Node\n~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   ssh verl-ppo\n\nAccess Ray Dashboard\n~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky status --endpoint 8265 verl-ppo  # Get dashboard URL\n\nStop a Cluster\n~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   sky down verl-ppo\n"
  },
  {
    "path": "docs/faq/faq.rst",
    "content": "Frequently Asked Questions\n====================================\n\nLast updated: 09/24/2025.\n\nRay related\n------------\n\nHow to add breakpoint for debugging with distributed Ray?\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nPlease checkout the official debugging guide from Ray: https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html\n\n\n\"Unable to register worker with raylet\"\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nThe cause of this issue is due to some system setting, e.g., SLURM added some constraints on how the CPUs are shared on a node. \nWhile `ray.init()` tries to launch as many worker processes as the number of CPU cores of the machine,\nsome constraints of SLURM restricts the `core-workers` seeing the `raylet` process, leading to the problem.\n\nTo fix this issue, you can set the config term ``ray_init.num_cpus`` to a number allowed by your system.\n\nDistributed training\n------------------------\n\nHow to run multi-node post-training with Ray?\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nYou can start a ray cluster and submit a ray job, following the official guide from Ray: https://docs.ray.io/en/latest/ray-core/starting-ray.html\n\nThen in the configuration, set the ``trainer.nnode`` config to the number of machines for your job.\n\nHow to use verl on a Slurm-managed cluster?\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nRay provides users with `this <https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html>`_ official\ntutorial to start a Ray cluster on top of Slurm. We have verified the :doc:`GSM8K example<../examples/gsm8k_example>`\non a Slurm cluster under a multi-node setting with the following steps.\n\n1. [Optional] If your cluster support `Apptainer or Singularity <https://apptainer.org/docs/user/main/>`_ and you wish\nto use it, convert verl's Docker image to an Apptainer image. Alternatively, set up the environment with the package\nmanager available on your cluster or use other container runtimes (e.g. through `Slurm's OCI support <https://slurm.schedmd.com/containers.html>`_) available to you.\n\n.. code:: bash\n\n    apptainer pull /your/dest/dir/vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3.sif docker://verlai/verl:vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3\n\n2. Follow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints.\n\n3. Modify `examples/slurm/ray_on_slurm.slurm <https://github.com/volcengine/verl/blob/main/examples/slurm/ray_on_slurm.slurm>`_ with your cluster's own information.\n\n4. Submit the job script to the Slurm cluster with `sbatch`.\n\nPlease note that Slurm cluster setup may vary. If you encounter any issues, please refer to Ray's\n`Slurm user guide <https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html>`_ for common caveats.\n\nIf you changed Slurm resource specifications, please make sure to update the environment variables in the job script if necessary.\n\n\nInstall related\n------------------------\n\nNotImplementedError: TensorDict does not support membership checks with the `in` keyword. \n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nDetail error information: \n\n.. code:: bash\n\n    NotImplementedError: TensorDict does not support membership checks with the `in` keyword. If you want to check if a particular key is in your TensorDict, please use `key in tensordict.keys()` instead.\n\nCause of the problem: There is no suitable version of tensordict package for the linux-arm64 platform. The confirmation method is as follows:\n\n.. code:: bash\n\n    pip install tensordict==0.6.2\n\nOutput example:\n\n.. code:: bash\n\n    ERROR: Could not find a version that satisfies the requirement tensordict==0.6.2 (from versions: 0.0.1a0, 0.0.1b0, 0.0.1rc0, 0.0.2a0, 0.0.2b0, 0.0.3, 0.1.0, 0.1.1, 0.1.2, 0.8.0, 0.8.1, 0.8.2, 0.8.3)\n    ERROR: No matching distribution found for tensordict==0.6.2\n\nSolution 1st:\n  Install tensordict from source code:\n\n.. code:: bash\n\n    pip uninstall tensordict\n    git clone https://github.com/pytorch/tensordict.git\n    cd tensordict/\n    git checkout v0.6.2\n    python setup.py develop\n    pip install -v -e .\n\nSolution 2nd:\n  Temperally modify the error takeplace codes: tensordict_var -> tensordict_var.keys()\n\n\nIllegal memory access\n---------------------------------\n\nIf you encounter the error message like ``CUDA error: an illegal memory access was encountered`` during rollout, please check the vLLM documentation for troubleshooting steps specific to your vLLM version.\n\nCheckpoints\n------------------------\n\nIf you want to convert the model checkpoint into huggingface safetensor format, please refer to ``verl/model_merger``.\n\n\nTriton ``compile_module_from_src`` error\n------------------------------------------------\n\nIf you encounter triton compilation error similar to the stacktrace below, please set the ``use_torch_compile`` flag according to\nhttps://verl.readthedocs.io/en/latest/examples/config.html to disable just-in-time compilation for fused kernels.\n\n.. code:: bash\n\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/jit.py\", line 345, in <lambda>\n    return lambda *args, **kwargs: self.run(grid=grid, warmup=False, *args, **kwargs)\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/autotuner.py\", line 338, in run\n    return self.fn.run(*args, **kwargs)\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/jit.py\", line 607, in run\n    device = driver.active.get_current_device()\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py\", line 23, in __getattr__\n    self._initialize_obj()\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py\", line 20, in _initialize_obj\n    self._obj = self._init_fn()\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/driver.py\", line 9, in _create_driver\n    return actives[0]()\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py\", line 371, in __init__\n    self.utils = CudaUtils()  # TODO: make static\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py\", line 80, in __init__\n    mod = compile_module_from_src(Path(os.path.join(dirname, \"driver.c\")).read_text(), \"cuda_utils\")\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/backends/nvidia/driver.py\", line 57, in compile_module_from_src\n    so = _build(name, src_path, tmpdir, library_dirs(), include_dir, libraries)\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/site-packages/triton/runtime/build.py\", line 48, in _build\n    ret = subprocess.check_call(cc_cmd)\n  File \"/data/lbh/conda_envs/verl/lib/python3.10/subprocess.py\", line 369, in check_call\n    raise CalledProcessError(retcode, cmd)\n\nWhat is the meaning of train batch size, mini batch size, and micro batch size?\n------------------------------------------------------------------------------------------\n\nThis figure illustrates the relationship between different batch size configurations.\n\nhttps://excalidraw.com/#json=pfhkRmiLm1jnnRli9VFhb,Ut4E8peALlgAUpr7E5pPCA\n\n.. image:: https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d\n\nHow to generate ray timeline to analyse performance of a training job?\n------------------------------------------------------------------------------------------\n\nTo generate the ray timeline file, you can set the config term ``ray_init.timeline_json_file`` to a json file path.\nFor example:\n\n.. code:: bash\n\n    ray_init.timeline_json_file=/tmp/ray_timeline.json\n  \nThe file will be generated in the specified path at the end of a training job.\nYou can use tools like chrome://tracing or the Perfetto UI and view the ray timeline file.\n\nThis figure shows the ray timeline file generated by from a training job on 1 node with 4 GPUs\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray_timeline.png?raw=true\n\nHow to set proxy only for wandb?\n------------------------------------------------------------------------------------------\n\nIf you need a proxy to access wandb, you can add below config in your training job script.\nComparing to using global https_proxy env variable, this approach won't mess up other http requests, such as ChatCompletionScheduler.\n\n.. code:: bash\n\n  +trainer.wandb_proxy=http://<your proxy and port>\n\nMissmatch between inference and training sequence (high actor/grad_norm)\n------------------------------------------------------------------------------------------\n\nIf you encounter the issue of actor/grad_norm metric continuously increasing during training, it might be caused by a significant precision mismatching between the inference engine and training. You can use the following parameter to confirm this:\n\n.. code:: bash\n\n    actor_rollout_ref.rollout.calculate_log_probs=True\n\nThis parameter will add metrics like training/rollout_probs_diff_mean , which can be used to verify if there is a precision difference between inference and training.\n\nUnder normal circumstances, the value of training/rollout_probs_diff_mean should be below 0.005. If you observe this value to be higher than 0.01, it indicates a precision issue from the inference engine.\nThe precision issue is known to occur under the following conditions:\n\n1. Using non-Hopper architecture GPUs, such as A100, L20, B200, etc.\n\n2. Using vLLM `with issue 22103 <https://github.com/vllm-project/vllm/issues/22103>`_ as the inference engine.\n\n3. The input and output texts are long, for example, in multi-turn scenarios using reasioning models like Qwen3 for RL training.\n\nIf all three conditions above are met and you observe that rollout_probs_diff_mean is too high, it is recommended to add the following parameter to resolve the precision issue:\n\n.. code:: bash\n\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_cascade_attn=True\n\nThe root cause of this issue is a bug in the flash attention used by vLLM. Although it has been fixed, the fix has not yet been released in the latest version of vLLM (v0.10.2).\nFor a more detailed explanation of this issue, please refer to `Fix LSE output error in FA2 kv-split <https://github.com/vllm-project/flash-attention/pull/87>`_.\n\nUntil vLLM releases a new version with this fix, it is recommended to use the configuration above to disable cascade attention as a workaround.\n"
  },
  {
    "path": "docs/hybrid_flow.rst",
    "content": "=========================================================\nHybridFlow Programming Guide\n=========================================================\n\nLast updated: 06/02/2025.\n\n.. _vermouth: https://github.com/vermouth1992\n\nAuthor: `Chi Zhang <https://github.com/vermouth1992>`_\n\nverl is an open source implementation of the paper `HybridFlow <https://arxiv.org/abs/2409.19256v2>`_ [1]_. In this section, we will introduce the basic concepts of HybridFlow, the motivation and how to program with verl APIs.\n\nMotivation and Design\n------------------------\nWe use dataflow to represent RL systems. [4]_.\n\nDataFlow\n~~~~~~~~~~~~~~~~~~~~\n\nDataflow is an abstraction of computations. Neural Network training is a typical dataflow. It can be represented by computational graph. \n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/dataflow.jpeg?raw=true\n   :alt: The dataflow graph from CS231n 2024 lecture 4\n\nThis figure [2]_ represents the computation graph of a polynomial function followed by a sigmoid function. In the data flow of neural network computation, each node represents an operator, and each edge represents the direction of forward/backward propagation. The computation graph determines the architecture of the neural network.\n\nRL as a dataflow problem\n++++++++++++++++++++++++++++++++++++++++++++++\n\nReinforcement learning (RL) training can also be represented as a dataflow. Below is the dataflow graph that represents the PPO algorithm used in RLHF [3]_:\n\n.. image:: https://picx.zhimg.com/70/v2-cb8ab5ee946a105aab6a563e92682ffa_1440w.avis?source=172ae18b&biz_tag=Post\n  :alt: PPO dataflow graph, credit to Zhihu 低级炼丹师\n\nHowever, the dataflow of RL has fundamental differences compared with dataflow of neural network training as follows:\n\n+--------------------------+--------------------------------------------------+---------------------+\n| Workload                 | Node                                             | Edge                |\n+--------------------------+--------------------------------------------------+---------------------+\n| Neural Network Training  | Operator (+/-/matmul/softmax)                    | Tensor movement     |\n+--------------------------+--------------------------------------------------+---------------------+\n| Reinforcement Learning   | High-level operators (rollout/model forward)     | Data Movement       |\n+--------------------------+--------------------------------------------------+---------------------+\n\nIn the case of tabular reinforcement learning, each operator is a simple scalar math operation (e.g., bellman update). In deep reinforcement learning(DRL), each operator is a high-level neural network computation such as model inference/update. This makes RL a two-level dataflow problem:\n\n- Control flow: defines how the high-level operators are executed (e.g., In PPO, we first perform rollout. Then, we perform advantage computation. Finally, we perform training). It expresses the **core logics of RL algorithms**.\n- Computation flow: defines the dataflow of **neural network computation** (e.g., model forward/backward/optimizer).\n\n\nDesign Choices\n~~~~~~~~~~~~~~~~~~~~\nThe model size used in DRL before the LLM era is typically small. Thus, the high-level neural network computation can be done in a single process. This enables embedding the computation flow inside the control flow as a single process.\n\nHowever, in the LLM era, the computation flow (e.g., training neural network) becomes a multi-process program. This naturally leads to two design choices:\n\n1. Convert the control flow into a multi-process program as well. Then colocate with computation flow (unified multi-controller)\n\n- Advantages:\n\n  - Achieves the **optimal performance** under fixed computation flow and control flow as the communication overhead in both training and data transfer is minimized.\n\n- Disadvantages:\n\n  - The computation and/or control flow is **hard to reuse** from software perspective as computation code is coupled with specific controller code. For example, the training loop of PPO is generic. Say we have an PPO training flow implemented with a specific computation flow such as FSDP. Neither the control flow or computation flow can be reused if we want to switch the computation flow from FSDP to Megatron, due to the coupling of control and computation flows.\n  - Requires more efforts from the user under flexible and dynamic control flows, due to the multi-process nature of the program.\n\n2. Separate the flows: single process for the control flow and multi-process for computation flow\n\n- Advantages:\n\n  - The computation flow defined elsewhere can be **easily reused** after the decoupling.\n  - The controller runs on a single process. Implementing a new RL algorithm with a **different control flow is simple and easy**.\n\n- Disadvantages:\n\n  - Additional **data communication overhead** each time the controller process and computatation processes interact. The data has to be sent back and forth.\n\nIn verl, the latter strategy with separate control flow and computation flow is adopted. verl is designed to decouple the control flow of RL algorithms, and the implementation of computation engines.\n\nOverall Execution Diagram\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nBelow is a simplified diagram denoting the execution of a reinforcement learning job. In the diagram, the controller runs on a single process, while the generator/actor workers, critic workers run on multiple processes, placed with specific resource groups. For rollout, the controller passes the data to the generator to perform sample generation. When the rollout is done, the data is passed back to controller for the next step of the algorithm. Similar execution is done for other workers. With the hybrid controller design, the data flow and computation is decoupled to provide both efficiency in computation and flexibility in defining algorithm training loops.\n\n.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/driver_worker.png?raw=true\n   :alt: The execution diagram\n\nCodebase walkthrough (PPO)\n------------------------------------------------\n\nEntry function\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\nCode: https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py\n\nIn this file, we define a remote function `main_task` that serves as the controller (driver) process as shown in the above figure. We also define a ``RewardManager``, where users can customize their reward function based on the data source in the dataset. Note that `RewardManager` should return the final token-level reward that is optimized by RL algorithms. Note that users can combine model-based rewards and rule-based rewards.\nThe ``main_task`` constructs a RayPPOTrainer instance and launch the fit. Note that ``main_task`` **runs as a single process**.\n\nWe highly recommend that the ``main_task`` is NOT scheduled on the head of the ray cluster because ``main_task`` will consume a lot of memory but the head usually contains very few resources.\n\nRay trainer\n~~~~~~~~~~~~~~~~~~~~\nCode: https://github.com/volcengine/verl/blob/main/verl/trainer/ppo/ray_trainer.py\n\nThe RayPPOTrainer manages \n\n- Worker and WorkerGroup construction\n- Runs the main loop of PPO algorithm\n\nNote that, the fit function of RayPPOTrainer **runs as a single process**.\n\nWorker and WorkerGroup construction\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nEach workerGroup manages a list of workers that runs remotely. Note that the worker group runs in the process of its constructor.\nEach worker inside the WorkerGroup runs on a GPU. The worker group serves as a proxy for the controller process to interact with a list of workers, in order to perform certain computations. **In order to do so, we have to bind the methods of the worker into the method of the WorkerGroup and define the data dispatch and data collection**. This is done via simple decoration that will be introduced in the Worker definition section.\n\nFor example, in PPO, we define 3 worker groups:\n\n- ActorRolloutRef: manages actor, rollout and reference policy. ActorRolloutRefWorker can be instantiated as a single actor, a single rollout, a single reference policy, a combined actor/rollout or a combined actor/rollout/ref. This design is aimed for the maximum code reuse in various scenarios. The reason for colocating actor and rollout is for fast weight transfer using nccl. The reason for coloating actor and reference is to implement an efficient lora PPO as the reference policy is simply the base model of PPO in lora. The colocation is done via ``verl.single_controller.ray.base.create_colocated_worker_cls``, where it creates a single ray remote class exposing all class methods from these roles.\n- Critic: manages the critic model\n- Reward: manages the reward model\n\nThe worker group will be constructed on the resource pool it designates. The resource pool is a set of GPUs in the ray cluster.\n\nWorker definition\n~~~~~~~~~~~~~~~~~~~~\n\n.. _ActorRolloutRefWorker: https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py\n\nWe take `ActorRolloutRefWorker <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py>`_ for an example.\nThe APIs it should expose to the controller process are:\n\n- init_model: build the underlying model\n- generate_sequences: given prompts, generate responses\n- compute_log_prob: compute the log-probability of a generated sequence using actor\n- compute_ref_log_prob: compute the log-probability of a generated sequence using reference policy\n- save_checkpoint: save the checkpoint\n\nNote that these methods are defined in the worker that can only be invoked via remote calls. For example, if the controller process wants to initialize the model, it has to call\n\n.. code-block:: python\n\n   for worker in actor_rollout_ref_wg:\n       worker.init_model.remote()\n\nIf the controller process wants to generate sequences, it has to call\n\n.. code-block:: python\n\n   data = xxx\n   # split the data into dp chunks\n   data_dp_lst = data.split(dp_size)\n   output_dp_lst = []\n   for i, worker in enumerate(actor_rollout_ref_wg):\n       output_future = worker.generate_sequences.remote(data_dp_lst[i])\n       output_dp_lst.append(output_future)\n   output = torch.cat(ray.get(output_dp_lst), dim=0)\n\nWe observe that controller process calling worker group methods in general can be divided into 3 parts:\n\n- Split the data into data parallel sizes\n- Dispatch the corresponding data into each worker\n- Collect and concatenate the data when the computation finishes\n\nIn verl, we design a syntax sugar to encapsulate the 3 processes into a single call from the controller process.\n\n.. code-block:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def generate_sequences(data):\n       ...\n\n   # on the driver\n   output = actor_rollout_ref_wg.generate_sequences(data)\n\nWe decorate the method of the worker with a ``register`` that explicitly defines how the input data should be split and dispatched to each worker, and how the output data should be collected and concatenated by the controller. For example, ``Dispatch.DP_COMPUTE_PROTO`` splits the input data into dp chunks, dispatch each data to each worker, collect the output and concatenate the results. Note that this function requires the input and output to be a DataProto defined here (https://github.com/volcengine/verl/blob/main/verl/protocol.py).\n\n\nPPO main loop\n~~~~~~~~~~~~~~~~~~~~\nWith the aforementioned APIs, we can implement the main loop of PPO as if it is a single process program\n\n.. code-block:: python\n\n   for prompt in dataloader:\n       output = actor_rollout_ref_wg.generate_sequences(prompt)\n       old_log_prob = actor_rollout_ref_wg.compute_log_prob(output)\n       ref_log_prob = actor_rollout_ref_wg.compute_ref_log_prob(output)\n       values = critic_wg.compute_values(output)\n       rewards = reward_wg.compute_scores(output)\n       # compute_advantages is running directly on the control process\n       advantages = compute_advantages(values, rewards)\n       output = output.union(old_log_prob)\n       output = output.union(ref_log_prob)\n       output = output.union(values)\n       output = output.union(rewards)\n       output = output.union(advantages)\n       # update actor\n       actor_rollout_ref_wg.update_actor(output)\n       critic.update_critic(output)\n\nTakeaways\n~~~~~~~~~~~~~~~~~~~~\n- This programming paradigm enables users to use different computation backend without modification of the control process.\n- This programming paradigm enables flexible placement (by changing the mapping of WorkerGroup and ResourcePool) without modification of the control process.\n\nRepository organization\n------------------------------------------------\n\nImportant code files in the repository are organized as below:\n\n.. code-block:: bash\n\n   verl # the verl package\n     trainer\n       main_ppo.py  # the entrypoint for RL training\n       ppo\n         ray_trainer.py  # the training loop for RL algorithms such as PPO\n       sft_trainer.py  # the SFT trainer with FSDP backend\n     config\n       generation.yaml  # configuration template for rollout\n       ppo_trainer.yaml  # configuration template for the RL trainer\n     workers\n       protocol.py  # the interface of DataProto\n       fsdp_workers.py   # the FSDP worker interfaces: ActorRolloutRefWorker, CriticWorker, RewardModelWorker\n       megatron_workers.py  # the Megatron worker interfaces: ActorRolloutRefWorker, CriticWorker, RewardModelWorker\n       actor\n         dp_actor.py  #  data parallel actor with FSDP backend\n         megatron_actor.py  # nD parallel actor with Megatron backend\n       critic\n         dp_critic.py  # data parallel critic with FSDP backend\n         megatron_critic.py  # nD parallel critic with FSDP backend\n       reward_model\n         megatron\n           reward_model.py  # reward model with Megatron backend\n       rollout\n         vllm\n           vllm_rollout.py  # rollout with vllm backend\n         hf_rollout.py  # rollout with huggingface TGI backend\n       sharding_manager\n         fsdp_ulysses.py  # data and model resharding when using FSDP + ulysses\n         fsdp_vllm.py  # data and model resharding when using FSDP + ulysses + vllm\n         megatron_vllm.py  # data and model resharding when using Megatron + vllm\n     utils\n       dataset  # datasets for SFT/RM/RL\n       reward_score  # function based reward\n         gsm8k.py  # reward function for gsm8k dataset\n         math.py  # reward function for math dataset\n       seqlen_balancing.py  # the sequence balance optimization\n     models\n       llama  # Megatron implementation for llama, deepseek, mistral, etc\n       transformers  # ulysses integration with transformer models such as llama, qwen, etc\n       weight_loader_registery.py  # registry of weight loaders for loading hf ckpt into Megatron\n     third_party\n       vllm  # adaptor for vllm's usage in RL\n         vllm_spmd  # vllm >= v0.7 adaptor\n   examples  # example scripts\n   tests  # integration and unit tests\n   .github  # the configuration of continuous integration tests\n\n\n.. [1] HybridFlow: A Flexible and Efficient RLHF Framework: https://arxiv.org/abs/2409.19256v2\n.. [2] Data flow graph credit to CS231n 2024 lecture 4: https://cs231n.stanford.edu/slides/2024/lecture_4.pdf\n.. [3] PPO dataflow graph credit to 低级炼丹师 from Zhihu​: https://zhuanlan.zhihu.com/p/635757674\n.. [4] RLFlow\n"
  },
  {
    "path": "docs/index.rst",
    "content": "Welcome to verl's documentation!\n================================================\n\nverl is a flexible, efficient and production-ready RL training framework designed for large language models (LLMs) post-training. It is an open source implementation of the `HybridFlow <https://arxiv.org/pdf/2409.19256>`_ paper.\n\nverl is flexible and easy to use with:\n\n- **Easy extension of diverse RL algorithms**: The hybrid programming model combines the strengths of single-controller and multi-controller paradigms to enable flexible representation and efficient execution of complex Post-Training dataflows. Allowing users to build RL dataflows in a few lines of code.\n\n- **Seamless integration of existing LLM infra with modular APIs**: Decouples computation and data dependencies, enabling seamless integration with existing LLM frameworks, such as PyTorch FSDP, Megatron-LM, vLLM and SGLang. Moreover, users can easily extend to other LLM training and inference frameworks.\n\n- **Flexible device mapping and parallelism**: Supports various placement of models onto different sets of GPUs for efficient resource utilization and scalability across different cluster sizes.\n\n- Ready integration with popular HuggingFace models\n\n\nverl is fast with:\n\n- **State-of-the-art throughput**: By seamlessly integrating existing SOTA LLM training and inference frameworks, verl achieves high generation and training throughput.\n\n- **Efficient actor model resharding with 3D-HybridEngine**: Eliminates memory redundancy and significantly reduces communication overhead during transitions between training and generation phases.\n\n--------------------------------------------\n\n.. _Contents:\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Quickstart\n\n   start/install\n   start/quickstart\n   start/multinode\n   start/ray_debug_tutorial\n   start/more_resources\n   start/agentic_rl\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Programming guide\n\n   hybrid_flow\n   single_controller\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Data Preparation\n\n   preparation/prepare_data\n   preparation/reward_function\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Configurations\n\n   examples/config\n\n.. toctree::\n   :maxdepth: 1\n   :caption: PPO Example\n\n   examples/ppo_code_architecture\n   examples/gsm8k_example\n   examples/multi_modal_example\n   examples/skypilot_examples\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Algorithms\n\n   algo/ppo.md\n   algo/grpo.md\n   algo/collabllm.md\n   algo/dapo.md\n   algo/spin.md\n   algo/sppo.md\n   algo/entropy.md\n   algo/opo.md\n   algo/baseline.md\n   algo/gpg.md\n   algo/rollout_corr.md\n   algo/rollout_corr_math.md\n   algo/otb.md\n   algo/dppo.md\n\n.. toctree::\n   :maxdepth: 1\n   :caption: PPO Trainer and Workers\n\n   workers/ray_trainer\n   workers/fsdp_workers\n   workers/megatron_workers\n   workers/automodel_workers\n   workers/sglang_worker\n   workers/trtllm_worker\n   workers/model_engine\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Performance Tuning Guide\n\n   perf/dpsk.md\n   perf/best_practices\n   perf/perf_tuning\n   perf/perf_tuning_on_ascend.rst\n   README_vllm0.8.md\n   perf/device_tuning\n   perf/verl_profiler_system.md\n   perf/nsight_profiling.md\n   perf/torch_profiling.md\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Adding new models\n\n   advance/fsdp_extension\n   advance/megatron_extension\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Advanced Features\n\n   advance/checkpoint\n   advance/rope\n   advance/attention_implementation\n   advance/ppo_lora.rst\n   sglang_multiturn/multiturn.rst\n   sglang_multiturn/interaction_system.rst\n   advance/placement\n   advance/dpo_extension\n   examples/sandbox_fusion_example\n   advance/rollout_trace.rst\n   advance/rollout_skip.rst\n   advance/one_step_off\n   advance/agent_loop\n   advance/reward_loop\n   advance/fully_async\n   data/transfer_queue.md\n   advance/grafana_prometheus.md\n   advance/fp8.md\n   advance/async-on-policy-distill\n   advance/mtp.md\n\n.. toctree::\n   :maxdepth: 2\n   :caption: Hardware Support\n\n   amd_tutorial/amd_build_dockerfile_page.rst\n   amd_tutorial/amd_vllm_page.rst\n   ascend_tutorial/contribution_guide/ascend_ci_guide_zh.rst\n   ascend_tutorial/quick_start/ascend_quick_start.rst\n   ascend_tutorial/quick_start/dockerfile_build_guidance.rst\n   ascend_tutorial/quick_start/ascend_sglang_quick_start.rst\n   ascend_tutorial/features/ascend_consistency.rst\n   ascend_tutorial/features/ascend_backend_features.md\n   ascend_tutorial/profiling/ascend_profiling_zh.rst\n   ascend_tutorial/profiling/ascend_profiling_en.rst\n   ascend_tutorial/examples/gspo_optimization_practice.md\n   ascend_tutorial/examples/ascend_performance_analysis_guide.md\n   ascend_tutorial/examples/dapo_multi_model_optimization_practice.md\n   ascend_tutorial/examples/ascend_sglang_best_practices.rst\n   ascend_tutorial/examples/ascend_retool_best_pratice.rst\n   ascend_tutorial/examples/run_qwen3_32B_megatron_1k_256k_npu.md\n   ascend_tutorial/faq/faq.rst\n\n.. toctree::\n   :maxdepth: 1\n   :caption: API References\n\n   api/data\n   api/single_controller.rst\n   api/trainer.rst\n   api/utils.rst\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Blog\n\n   blog/v0.7.md\n\n.. toctree::\n   :maxdepth: 2\n   :caption: FAQ\n\n   faq/faq\n\n.. toctree::\n   :maxdepth: 1\n   :caption: Development Notes\n\n   sglang_multiturn/sandbox_fusion.rst\n\nContribution\n-------------\n\nverl is free software; you can redistribute it and/or modify it under the terms\nof the Apache License 2.0. We welcome contributions.\nJoin us on `GitHub <https://github.com/volcengine/verl>`_, `Slack <https://join.slack.com/t/verlgroup/shared_invite/zt-2w5p9o4c3-yy0x2Q56s_VlGLsJ93A6vA>`_ and `Wechat <https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/WeChat.JPG>`_ for discussions.\n\nContributions from the community are welcome! Please check out our `project roadmap <https://github.com/volcengine/verl/issues/710>`_ and `good first issues <https://github.com/volcengine/verl/issues?q=is%3Aissue%20state%3Aopen%20label%3A%22good%20first%20issue%22>`_ to see where you can contribute.\n\nCode Linting and Formatting\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nWe use pre-commit to help improve code quality. To initialize pre-commit, run:\n\n.. code-block:: bash\n\n   pip install pre-commit\n   pre-commit install\n\nTo resolve CI errors locally, you can also manually run pre-commit by:\n\n.. code-block:: bash\n\n   pre-commit run\n\nAdding CI tests\n^^^^^^^^^^^^^^^^^^^^^^^^\n\nIf possible, please add CI test(s) for your new feature:\n\n1. Find the most relevant workflow yml file, which usually corresponds to a ``hydra`` default config (e.g. ``ppo_trainer``, ``ppo_megatron_trainer``, ``sft_trainer``, etc).\n2. Add related path patterns to the ``paths`` section if not already included.\n3. Minimize the workload of the test script(s) (see existing scripts for examples).\n\nWe are HIRING! Send us an `email <mailto:haibin.lin@bytedance.com>`_ if you are interested in internship/FTE opportunities in MLSys/LLM reasoning/multimodal alignment.\n"
  },
  {
    "path": "docs/perf/best_practices.rst",
    "content": "Verl LLM Best Practices (DAPO + Qwen3-235B)\n===========================================\n\nLast updated: 11/03/2025.\n\nPurpose\n-------\n\nThis guide uses DAPO training on Qwen3-235B as a concrete example. We unpack every parameter that appears in the optimization objective, map it to Verl configuration entries, and share field-tested recommendations so you can derive sensible settings for your own workloads.\n\n.. note::\n\n   1. The guide only covers the subset of parameters required to reproduce the DAPO experiments discussed here. For the full list, refer to the ``config`` components in the Verl source tree: https://github.com/volcengine/verl/tree/main/verl/trainer/config\n   2. PPO and GRPO introduce KL-constrained policies. We therefore include that setup in the explanations below. You can treat all configurations mentioned here as a DAPO pipeline augmented with a KL penalty.\n\nOptimization Objectives\n-----------------------\n\nDAPO objective\n~~~~~~~~~~~~~~\n\n.. math::\n\n   \\begin{aligned}\n   \\mathcal{J}_{\\mathrm{DAPO}}(\\theta)= & \\mathbb{E}_{(q, a) \\sim \\mathcal{D},\\left\\{o_i\\right\\}_{i=1}^G \\sim \\pi_{\\theta_{\\text {old }}}(\\cdot \\mid q)} \\\n    {\\left[\\frac{1}{\\sum_{i=1}^G\\left|o_i\\right|} \\sum_{i=1}^G \\sum_{t=1}^{\\left|o_i\\right|} \\min \\left(r_{i, t}(\\theta) \\hat{A}_{i, t}, \\operatorname{clip}\\left(r_{i, t}(\\theta), 1-\\varepsilon_{\\text {low }}, 1+\\varepsilon_{\\text {high }}\\right) \\hat{A}_{i, t}\\right)\\right] } \\\\\n   \\end{aligned}\n\n.. math::\n   \\text { s.t. } \\quad 0<\\mid\\left\\{o_i \\mid \\text { is_equivalent }\\left(a, o_i\\right)\\right\\} \\mid<G,\n\n.. math::\n\n   \\text {where} \\quad r_{i, t}(\\theta)=\\frac{\\pi_\\theta\\left(o_{i, t} \\mid q, o_{i,<t}\\right)}{\\pi_{\\theta_{\\text {old }}}\\left(o_{i, t} \\mid q, o_{i,<t}\\right)}, \\quad \\hat{A}_{i, t}=\\frac{R_i-\\operatorname{mean}\\left(\\left\\{R_i\\right\\}_{i=1}^G\\right)}{\\operatorname{std}\\left(\\left\\{R_i\\right\\}_{i=1}^G\\right)}\n\nGRPO objective\n~~~~~~~~~~~~~~\n\n.. math::\n\n   \\begin{aligned}\n   \\mathcal{J}_{G R P O}(\\theta) & =\\mathbb{E}_{q \\sim P(Q),\\left\\{o_i\\right\\}_{i=1}^G \\sim \\pi_{\\theta_{\\text {old }}}(O \\mid q)} \\\n   \\frac{1}{G} \\sum_{i=1}^G \\frac{1}{\\left|o_i\\right|} \\sum_{t=1}^{\\left|o_i\\right|}\\left\\{\\min \\left[\\frac{\\pi_\\theta\\left(o_{i, t} \\mid q, o_{i,<t}\\right)}{\\pi_{\\theta_{\\text {old }}}\\left(o_{i, t} \\mid q, o_{i,<t}\\right)} \\hat{A}_{i, t}, \\operatorname{clip}\\left(\\frac{\\pi_\\theta\\left(o_{i, t} \\mid q, o_{i,<t}\\right)}{\\pi_{\\theta_{\\text {old }}}\\left(o_{i, t} \\mid q, o_{i,<t}\\right)}, 1-\\varepsilon, 1+\\varepsilon\\right) \\hat{A}_{i, t}\\right]-\\beta \\mathbb{D}_{K L}\\left[\\pi_\\theta \\| \\pi_{r e f}\\right]\\right\\},\n   \\end{aligned}\n\nNotation Overview\n-----------------\n\n:math:`(q,a)\\sim D`\n  - :math:`D` denotes the training dataset. For each sample, :math:`q` is the input prompt (for math tasks, the question) and :math:`a` is the target output—typically the final answer without intermediate reasoning steps.\n\n:math:`G`\n  - Group size. For every prompt we sample :math:`G` independent responses.\n\n:math:`\\theta`\n  - Actor model parameters.\n\n:math:`\\pi`\n  - Sampling strategy that bundles the rollout backend (vLLM, sglang, etc.) and all generation hyperparameters. Because LLMs generate tokens autoregressively, rollout dominates runtime, so backend-specific tuning is critical.\n\n:math:`\\pi_\\theta`\n  - Actor policy obtained by instantiating :math:`\\pi` with parameters :math:`\\theta`.\n\n:math:`\\hat{A}_{i,t}`\n  - Advantage of the :math:`i`-th sample within the group at timestep :math:`t`.\n\n:math:`R_i`\n  - Reward assigned to the :math:`i`-th sample in the group.\n\n:math:`\\mathbb{D}_{KL}`\n  - KL divergence between two policies.\n\n:math:`\\beta`\n  - Coefficient that weights the KL term.\n\n:math:`\\pi_{old}`\n  - Frozen “old” policy, updated after every :math:`\\texttt{train_batch_size}` samples.\n\n:math:`\\pi_{ref}`\n  - Reference policy used to compute the KL divergence.\n\n:math:`o_i, |o_i|`\n  - :math:`o_i` is the generated output sequence for the :math:`i`-th prompt; :math:`|o_i|` is its token length.\n\n:math:`\\pi_\\theta(o_{i,t} \\mid q_i, o_{i,<t})`\n  - Probability of emitting token :math:`o_{i,t}` at timestep :math:`t` given prompt :math:`q_i` and the previously generated prefix under parameters :math:`\\theta`. In practice, the rollout engine first generates full responses, then concatenates prompts and outputs for each model; with attention masks we can compute all token probabilities in one pass.\n\n:math:`\\varepsilon_{low}` and :math:`\\varepsilon_{high}`\n  - Lower and upper clipping bounds for importance sampling. DAPO adopts a clip-higher strategy, so the upper bound is different from the lower bound to prevent overly large policy updates.\n\nParameter Reference\n-------------------\n\n:math:`(q,a)\\sim D`\n  - ``data.train_files`` / ``data.val_files``:\n    Training and validation datasets. They must be stored as ``.parquet``. Use the conversion scripts under ``examples/data_preprocess`` and make sure your ``data_source`` implements the matching reward function. You can also reuse the HuggingFace dataset ``BytedTsinghua-SIA/DAPO-Math-17k``.\n  - ``data.prompt_key``:\n    Column name for prompts. Keep the default ``prompt`` unless you have a clearer schema.\n  - ``data.max_prompt_length``:\n    Upper bound on prompt length. Set it to cover the longest prompt in the corpus; when long-tail samples make it too large, lower the value and combine with ``data.truncation``.\n  - ``data.truncation``:\n    Policy for over-length inputs (truncate-left/right or raise). ``left`` works for most runs. If training logs show large ``clip_ratio`` and poor metrics, increase ``data.max_prompt_length`` or clean the data. Set to ``error`` when strict validation is required.\n\n:math:`G`\n  - ``actor_rollout_ref.rollout.n``:\n    Number of generations per prompt. Typical values: GRPO 64, DAPO 16.\n\n:math:`\\theta`\n  - ``actor_rollout_ref.model.path``:\n    Path to the actor checkpoint in HuggingFace-compatible format.\n  - ``actor_rollout_ref.actor.megatron.use_mbridge``:\n    Enable mbridge format conversion when the model was trained with Megatron. Use the latest mbridge release: https://github.com/ISEEKYAN/mbridge.\n    Now it must be True.\n  - ``actor_rollout_ref.actor.megatron.vanilla_mbridge``:\n    If set to True, use mbridge, else use Megatron-Bridge https://github.com/NVIDIA-NeMo/Megatron-Bridge.\n    Now it is True by default. and it will defaultly be set to False in the future(v0.8).\n\n:math:`\\pi`\n  - ``actor_rollout_ref.rollout.name``:\n    Rollout backend. Verl currently supports ``vllm`` and ``sglang``—benchmark and tune according to your infrastructure.\n  - ``actor_rollout_ref.rollout.response_length`` / ``data.max_response_length``:\n    Maximum generated tokens (rollout setting takes precedence). Larger values improve quality but consume more memory and latency. Monitor ``clip_ratio``; values above 0.1 often mean you are truncating too much.\n  - ``actor_rollout_ref.rollout.gpu_memory_utilization``:\n    Target GPU memory usage during rollout. Push it as high as possible without triggering OOM; with parameter/gradient/optimizer offload enabled, 0.8–0.9 is common.\n  - ``actor_rollout_ref.rollout.tensor_model_parallel_size``:\n    Tensor parallel degree for the inference engine. Ensure ``(memory_per_gpu * gpu_memory_utilization * TP) > 2 * model_parameters`` (bf16/fp16). Increase TP gradually to expand KV cache capacity while watching communication cost—especially once TP > 8.\n  - ``actor_rollout_ref.rollout.temperature`` / ``top_p`` / ``top_k``:\n    Sampling knobs for rollout. Keep enough randomness; ``temperature=1.0``, ``top_p=1.0``, ``top_k=-1`` are good defaults.\n  - ``actor_rollout_ref.rollout.val_kwargs.temperature`` / ``top_p`` / ``top_k`` / ``do_sample`` / ``n``:\n    Sampling options for validation. Set ``temperature > 0`` to prevent repetitive thinking chains. For small test sets (e.g., AIME24) raise ``n`` (64 is a common choice) to reduce variance. A practical starting point is ``temperature=1.0``, ``top_p=0.7``, ``top_k=-1``, ``do_sample=True``, ``n=1`` and then increase ``n`` as needed.\n  - ``+actor_rollout_ref.rollout.engine_kwargs.vllm.*`` / ``+actor_rollout_ref.rollout.engine_kwargs.sglang.*``:\n    Extra backend options injected via the ``+`` syntax. Consult backend docs for exact semantics. Some switches (for example ``pipeline_parallel_size``) may not be supported yet; when TP=32, ``enable_expert_parallel=True`` can even slow down DeepSeek-V3 rollout, so benchmark carefully.\n\n:math:`\\pi_\\theta`\n  - ``data.train_batch_size``:\n    Total batch size per training iteration. Each rollout produces ``train_batch_size * n`` samples. Larger values reduce the number of rollouts but increase off-policy drift.\n  - ``actor_rollout_ref.actor.ppo_mini_batch_size``:\n    Mini-batch size per optimization step. Tune it the same way you would for standard deep learning workloads.\n  - ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``:\n    Samples processed per forward pass on one GPU group (a Megatron group contains TP * PP * CP GPUs). Keep it ≤ ``ppo_mini_batch_size`` and as large as memory allows.\n  - ``actor_rollout_ref.actor.use_dynamic_bsz``:\n    Enable dynamic batch sizing to adapt to sequence length and improve throughput.\n  - ``actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu``:\n    Maximum tokens per GPU when computing log probabilities under dynamic batching. Set it to at least a multiple of ``max_prompt_length + max_response_length`` to prevent truncation.\n  - Megatron parallelism parameters (``pipeline_model_parallel_size`` / ``tensor_model_parallel_size`` / ``expert_model_parallel_size`` / ``expert_tensor_parallel_size`` / ``context_parallel_size``):\n    Balance PP/TP/EP/ETP/CP to match memory and network constraints. In bf16/fp16, each parameter consumes roughly ``2 / TP`` bytes; if you keep FP32 master weights or skip optimizer offload, reserve another 4–8 bytes for Adam. Activations scale with ``micro_batch_size × sequence_length × hidden_size`` and can be mitigated with gradient checkpointing, dynamic batches, or offload. Prefer increasing TP first, add PP when necessary, extend sequence capacity with CP, align EP/ETP with TP for MoE models, and keep DP minimal on constrained clusters while combining with offload. Always align the setup with hardware topology and communication cost.\n  - ``actor_rollout_ref.model.use_fused_kernels``:\n    Enable Verl’s fused kernels for supported models to squeeze out additional performance.\n\n:math:`\\hat{A}_{i,t}`\n  - ``algorithm.adv_estimator``:\n    Advantage estimator. Set to ``grpo`` for DAPO/GRPO.\n\n:math:`R_i`\n  - ``reward_model.reward_manager``:\n    Reward aggregation strategy. Use ``dapo`` for DAPO and ``naive`` for GRPO.\n\n:math:`D_{KL}`\n  - ``algorithm.use_kl_in_reward``:\n    Whether to add a KL term to the reward. ``True`` for PPO, ``False`` for GRPO and DAPO.\n  - ``actor_rollout_ref.actor.use_kl_loss``:\n    Whether to include a KL loss term. ``False`` for PPO, ``True`` for GRPO, ``False`` for DAPO.\n\n:math:`\\beta`\n  - ``actor_rollout_ref.actor.kl_loss_coef``:\n    Weight of the KL loss. Start around 0.001. Larger values curb reward hacking but reduce exploration.\n  - ``algorithm.kl_ctrl.kl_coef``:\n    KL coefficient applied within the reward. Adjust to match your tolerance for divergence.\n\n:math:`\\pi_{old}`\n  - ``actor_rollout_ref.rollout.log_prob_use_dynamic_bsz``:\n    Enable dynamic batching when the old policy computes log-probabilities. Recommended.\n\n:math:`\\pi_{ref}`\n  - ``actor_rollout_ref.ref.log_prob_use_dynamic_bsz``:\n    Enable dynamic batching for the reference policy. Recommended.\n  - Reference Megatron parallelism:\n    Keep ``pipeline_model_parallel_size``, ``tensor_model_parallel_size``, ``expert_model_parallel_size``, ``expert_tensor_parallel_size``, and ``context_parallel_size`` in sync with the actor.\n  - ``actor_rollout_ref.ref.megatron.param_offload``:\n    Offload reference parameters to CPU when the actor does so. Even without gradients or optimizer states, parity helps with capacity planning.\n\n:math:`o_i` / :math:`|o_i|`\n  - ``actor_rollout_ref.actor.loss_agg_mode``:\n    Loss aggregation mode. Token-level ``token-mean`` matches the recommendations from Dr.GRPO and DAPO; use ``seq-mean-token-mean`` to reproduce the original GRPO behavior.\n\n:math:`\\pi_\\theta(o_{i,t} \\mid q_i,o_{i,<t})`\n  - ``actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu`` / ``actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu``:\n    Batch size while computing token probabilities. Rollout engines generate outputs and then concatenate inputs for each model, so balance memory against throughput.\n\n:math:`\\epsilon_{low}` / :math:`\\epsilon_{high}`\n  - ``actor_rollout_ref.actor.clip_ratio_low`` / ``actor_rollout_ref.actor.clip_ratio_high``:\n    Importance sampling clipping bounds. For DAPO, use ``clip_ratio_low=0.2`` and ``clip_ratio_high=0.28``.\n\nvLLM inference optimizations\n  - ``actor_rollout_ref.rollout.enable_chunked_prefill``:\n    Enables chunked prefill to boost GPU utilization (vLLM only). Tune together with ``max_num_batched_tokens``.\n  - ``actor_rollout_ref.rollout.max_num_batched_tokens``:\n    Maximum tokens per batch. A practical rule of thumb is ``max(8192, max_prompt_length + max_response_length, max_model_len)``; see the vLLM docs for details.\n  - ``actor_rollout_ref.rollout.enforce_eager``:\n    Disables CUDA graphs. By default vLLM leverages CUDA graphs for speed at the cost of extra memory (not limited by ``gpu_memory_utilization``); set this to ``True`` when memory is tight.\n  - ``actor_rollout_ref.rollout.cudagraph_capture_sizes``:\n    Explicit capture batch sizes for CUDA graphs. Default is ``null``; on memory-constrained systems try ``[1, 2, 4, 8, 16, 32]``.\n\nOptimizer settings\n  - ``actor_rollout_ref.actor.optim.lr``:\n    Learning rate. Start around ``1e-5`` or ``1e-6``.\n  - ``actor_rollout_ref.actor.optim.lr_warmup_steps``:\n    Number of warmup steps (e.g., 10).\n  - ``actor_rollout_ref.actor.optim.weight_decay``:\n    Weight decay coefficient, typically 0.1.\n  - ``actor_rollout_ref.actor.optim.clip_grad``:\n    Gradient clipping threshold, commonly 1.\n  - ``+actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction``:\n    Portion of optimizer updates executed on CPU. Large models such as DeepSeek benefit from enabling it with value 1.\n  - ``+actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d`` / ``+...use_precision_aware_optimizer`` / ``+...optimizer_cpu_offload``:\n    Companion switches for hybrid optimizers. Turn them on alongside CPU offload.\n\nMegatron-related parameters\n  - ``actor_rollout_ref.actor.megatron.param_offload`` / ``optimizer_offload`` / ``grad_offload``:\n    Offload parameters, optimizer states, and gradients to CPU when GPU memory is insufficient.\n  - ``+actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method`` / ``recompute_granularity`` / ``recompute_num_layers``:\n    Gradient checkpointing controls. Enable (e.g., ``uniform``, ``full``, ``1``) to trade computation for memory.\n  - ``+actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype`` / ``moe_shared_expert_overlap`` / ``moe_permute_fusion`` / ``moe_enable_deepep`` / ``moe_token_dispatcher_type``:\n    Recommended MoE knobs (sample values: ``fp32``, ``False``, ``True``, ``True``, ``flex``) for stable performance.\n  - ``+actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion``:\n    Enables fused gradient accumulation for additional speedup.\n  - ``+actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split`` / ``account_for_loss_in_pipeline_split`` / ``num_layers_in_last_pipeline_stage``:\n    Pipeline-parallel adjustments when layer counts do not divide evenly. Treat embedding and loss as standalone stages and set ``num_layers_in_last_pipeline_stage`` (0 or ``${LAST_LAYER}``) when you need manual control.\n\nTrainer\n  - ``trainer.logger``:\n    Logging backends. Use ``['console', 'wandb']`` or, on Volcano Engine ML Platform, ``['console', 'vemlp_wandb']``.\n  - ``trainer.project_name`` / ``trainer.experiment_name``:\n    Hierarchical naming for projects and experiments so you can locate runs quickly.\n  - ``trainer.n_gpus_per_node`` / ``trainer.nnodes``:\n    Number of GPUs per node and total node count. Match your cluster allocation.\n  - ``trainer.test_freq`` / ``trainer.save_freq`` / ``trainer.total_epochs``:\n    Evaluation interval, checkpoint interval, and total epochs—configure for your SLA.\n  - ``trainer.log_val_generations``:\n    Number of validation samples stored in logs. Start with 10 and adjust as needed.\n  - ``trainer.val_before_train``:\n    Run validation before training begins when you require a baseline checkpoint.\n"
  },
  {
    "path": "docs/perf/device_tuning.rst",
    "content": "Hardware Resource Needed for RL\n===============================\n\nLast updated: 06/25/2025.\n\nSince RL requires more resources compared to regular training, \ndetermining how much resources are needed to successfully run it before training \nis a relatively difficult task. To provide more people with reference points for \nresource selection when dealing with different models and tasks, this section is \nmainly dedicated to introducing the environmental requirements based on experiments \nwe have conducted.\n\nHowever, due to limited staff and equipment resources, we also hope for more \ncontributions from the open-source community. When submitting a PR, it is necessary \nto provide a script to be added to the example/tuning scripts.\n\nWe need two types of scripts: one is the configuration that can run with the **minimum \nresources(min)**, and the other is the configuration that runs with **recommended resources(recommended)**. For the former, \nit can be understood as a script that can run after applying all memory optimization techniques \n(e.g., offload, gradient checkpointing). For the latter, it can be understood as a script that \ncan run while avoiding operations that incur additional time overhead as much as possible (targetting best throughput).\n\nWhen defining script names, please follow this format: \n``[model]_[task]_[gpunums]_[device]_[train]_[infer].sh``. This will effectively improve \nthe script's recognizability. You can place the script under the ``examples/tuning/`` directory.\n\nIf you happen to have a configuration that has already been tested, we welcome you to submit \na PR and include a screenshot from Wandb or other verifiable evidence.\n\n----------------------------------------\n\n0.5B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2.5-0.5B\n      - GRPO-LoRA\n      - 1*H100\n      - 116\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-0.5b_grpo-lora_1_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/0.5b/qwen2-0.5b_grpo-lora_1_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n1.5B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2.5-1.5B\n      - GRPO-LoRA\n      - 1*H100\n      - 128\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-1.5b_grpo-lora_1_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/1.5b/qwen2-1.5b_grpo-lora_1_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n3B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2.5-3B\n      - GRPO-LoRA\n      - 1*H100\n      - 62\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-3b_grpo-lora_1_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/3b/qwen2-3b_grpo-lora_1_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n7B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2-7B\n      - GRPO\n      - 2*H800\n      - \\\n      - fsdp\n      - vllm0.8.2\n      - `qwen2-7b_grpo_2_h800_fsdp_vllm <https://github.com/volcengine/verl/blob/main/examples/tuning/7b/qwen2-7b_grpo_2_h800_fsdp_vllm.sh>`_\n      - `Xiangyongan <xiangyongan@bytedance.com>`_\n    * - MIN\n      - Qwen2.5-7B\n      - GRPO-LoRA\n      - 1*H100\n      - 16\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-7b_grpo-lora_1_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/7b/qwen2-7b_grpo-lora_1_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n14B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2-14B\n      - GRPO\n      - 4*H800\n      - \\\n      - fsdp\n      - vllm0.8.2\n      - `qwen2-14b_grpo_4_h800_fsdp_vllm <https://github.com/volcengine/verl/blob/main/examples/tuning/14b/qwen2-14b_grpo_4_h800_fsdp_vllm.sh>`_\n      - `Xiangyongan <xiangyongan@bytedance.com>`_\n    * - MIN\n      - Qwen2.5-14B\n      - GRPO-LoRA\n      - 2*H100\n      - 116\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-14b_grpo-lora_2_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/14b/qwen2-14b_grpo-lora_2_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n32B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n    \n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2-32B\n      - GRPO\n      - 8*H20\n      - \\\n      - megatron\n      - vllm0.8.2\n      - `qwen2-32b_grpo_8_h20_megatron_vllm <https://github.com/volcengine/verl/tree/main/examples/tuning/32b/qwen2_32B_grpo_8_h20_megatron_vllm.sh>`_\n      - `Xiangyongan <xiangyongan@bytedance.com>`_\n    * - MIN\n      - Qwen2.5-32B\n      - GRPO-LoRA\n      - 4*H100\n      - 180\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-32b_grpo-lora_4_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/32b/qwen2-32b_grpo-lora_4_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n70B\n~~~\n\n.. list-table::\n    :widths: auto\n    :header-rows: 1\n\n    * - Tag\n      - Model\n      - Task\n      - Resource\n      - MaxBatch\n      - Train\n      - Infer\n      - Link\n      - Contributor\n    * - MIN\n      - Qwen2-70B\n      - GRPO\n      - 32*H20\n      - \\\n      - fsdp\n      - vllm0.8.2\n      - `qwen2-70b_grpo_32_h20_fsdp_vllm <https://github.com/volcengine/verl/blob/main/examples/tuning/70b/qwen2-70b_grpo_32_h20_fsdp_vllm.sh>`_\n      - `Xiangyongan <xiangyongan@bytedance.com>`_\n    * - MIN\n      - Qwen2-70B\n      - GRPO\n      - 32*H800\n      - \\\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-70b_grpo_32_h800_fsdp_vllm <https://github.com/volcengine/verl/blob/main/examples/tuning/70b/qwen2-70b_grpo_32_h800_fsdp_vllm.sh>`_\n      - `Xiangyongan <xiangyongan@bytedance.com>`_\n    * - MIN\n      - Qwen2.5-72B\n      - GRPO-LoRA\n      - 8*H100\n      - 176\n      - fsdp\n      - vllm0.8.3\n      - `qwen2-72b_grpo-lora_8_h100_fsdp_vllm.sh <https://github.com/volcengine/verl/blob/main/examples/tuning/70b/qwen2-72b_grpo-lora_8_h100_fsdp_vllm.sh>`_\n      - `SimonHuang <thelongestusernameofall@gmail.com>`_\n\n405B\n~~~~\n\n.. table::\n   :widths: auto\n\n   ====== ====== ====== ======== ======== ====== ====== ======\n   tag    model  task   resource MaxBatch train  infer  link\n   ====== ====== ====== ======== ======== ====== ====== ======\n   \\      \\      \\        \\        \\      \\      \\\n   ====== ====== ====== ======== ======== ====== ====== ======\n\n671B\n~~~~\n\n.. table::\n   :widths: auto\n\n   ====== ====== ====== ======== ======== ====== ====== ======\n   tag    model  task   resource MaxBatch train  infer  link\n   ====== ====== ====== ======== ======== ====== ====== ======\n   \\      \\      \\        \\        \\      \\      \\\n   ====== ====== ====== ======== ======== ====== ====== ======\n"
  },
  {
    "path": "docs/perf/dpsk.md",
    "content": "# Training DeepSeek 671b\n\nLast updated: 08/20/2025.\n\nverl integrates Megatron to support large MoE models such as `Qwen3-235B-A22B` and `deepseek-ai/DeepSeek-V3`. This is an ongoing community effort.\n\nIn the journey the community added the following features and optimizations that enable verl with larger models:\n- per tensor weight resharding between rollout and training\n- context parallelism and expert parallelism enabled via megatron\n- dynamic batch size (sequence balance) for megatron\n- reduced ray-related serialization overhead\n- optimizer offloading, recomputation, and efficient kernels\n- various debugging metrics and utils\n- hybrid optimizer\n\nand the megatron backend now has a wider list of models supported:\n- DeepSeek-V3\n- Moonlight\n- Qwen3\n- Qwen2.5-VL (to be merged soon)\n- Qwen2\n- Mixtral\n\n## Getting Started\n\n### preparation\nThe recommended image with pre-built Megatron dependency is `verlai/verl:app-verl0.4-vllm0.8.5-mcore0.13.0-preview`, which is built using the Dockerfile at [docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview](https://github.com/volcengine/verl/blob/main/docker/verl0.4-cu124-torch2.6-fa2.7.4/Dockerfile.app.vllm.mcore0.13.preview).\n\nThe image is build in Hopper GPUs with DeepEP. It does not support None-Hopper GPUs, such as A100. You may need to reinstall DeepEP to work with A100.\n\nWith `OFFLOAD_FRACTION=1`, the system's minimum requirements are lowered. It can run on as few as 96 H20 (96GB) GPUs for DeepSeek-V3, and on as few as 32 H20 (96GB) GPUs for Qwen3-235B-A22B. However, this configuration will use 1.6TB CPU memory per node. If you run out of CPU memory or require faster training speed, you can add more nodes.\n\n### DeepSeek 671b\n\nFor DeepSeek-V3 671b, please refer to [examples/grpo_trainer/run_deepseek671b_math_megatron_96gb.sh](https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_deepseek671b_math_megatron_96gb.sh).\n\nMTP and quantilization is disabled during RL training.\n\nTo train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware.\n| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | LAST_LAYER |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n| 96 | 12 | 8 | 12 | 8 | 1. | False | 6 |\n| 128 | 16 | 8 | 16 | 8 | 0.5 | True | 1 |\n| 256 | 32 | 8 | 16 | 8 | 0. | True | 1 |\n| 512 | 64 | 1 | 16 | 32 | 0 | True | 1 |\n\n### Qwen3 235b\n\nFor Qwen3-235b, please refer to [examples/grpo_trainer/run_qwen3-235b_megatron_96gb.sh](https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3-235b_megatron_96gb.sh).\n\nTo train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware.\n| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | LAST_LAYER |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n| 32 | 4 | 4 | 8 | 4 | 1. | False | 6 |\n| 64 | 8 | 4 | 8 | 4 | 0.5 | True | 6 |\n| 128 | 16 | 4 | 8 | 4 | 0 | True | 6 |\n| 256 | 32 | 4 | 8 | 4 | 0 | True | 6 |\n\n### Benchmark\nHere are some benchmark results for DeepSeek / Qwen3-235B. All configurations match the recommended settings based on the number of GPUs.\n\n| model | num gpus | mean response length | rollout time(s) | GPU memory(GB) | CPU memory(GB) | MFU | step time(s) |\n| -- | -- | -- | -- | -- | -- | -- | -- |\n| DeepSeek 671b | 96 | 1960 | 1050 | 66 | 1500 | 0.19 | 1700 |\n\n### Qwen3-30B-A3B MOE\n\nFor Qwen3-30b, please refer to [examples/grpo_trainer/run_qwen3moe-30b_megatron_96gb.sh](https://github.com/volcengine/verl/blob/main/examples/grpo_trainer/run_qwen3moe-30b_megatron_96gb.sh).\n\nTo train your project, configure the following environment variables based on the number of available GPUs. These are recommended settings and can be adjusted based on your specific hardware.\n| num gpus | NNODES | TP | PP | EP | OFFLOAD_FRACTION | OFFLOAD_OPTIM | MFU |\n| -- | -- | -- | -- | -- | -- | -- | -- | \n| 8 | 1 | 1 | 1 | 8 | 1. | True | 0.4 |\n| 16 | 2 | 1 | 1 | 8 | 1. | True | 0.37 |\n| 32 | 4 | 1 | 1 | 8 | 1. | True | 0.31 |\n\n\n## Upcoming Optimizations\n\nThe community continue to optimize large MoE models further, ongoing efforts include:\n- further optimizing memory consumption, and provide recommended/tuned configurations with various machine types\n- optimizing long context RL training performance\n- performance improvement with SGLang x Megatron\n\nWe invite the community to try and improve verl together. Get connected with us on [slack](https://join.slack.com/t/verlgroup/shared_invite/zt-2w5p9o4c3-yy0x2Q56s_VlGLsJ93A6vA)/[wechat](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/WeChat.JPG)/[Github issues](https://github.com/volcengine/verl/issues/708)!\n\n## Acknowledgement\n@vermouth1992 @ISEEKYAN @ETOgaosion @yzlnew @ShareLer @BearBiscuit05 @ccclyu @ann-qin-lu @SwordFaith @zzong2006 @zhaochenyang20 @ocss884 @eric-haibin-lin @chenhaiq @techkang\n"
  },
  {
    "path": "docs/perf/nsight_profiling.md",
    "content": "# NVIDIA Nsight Systems profiling in verl\n\nLast updated: 06/20/2025.\n\nThis guide explains how to use NVIDIA Nsight Systems for profiling verl training runs.\n\n## Configuration\n\nProfiling in verl can be configured through several parameters in the trainer configuration file (ppo_trainer.yaml or other files like dapo_trainer.yaml):\n\n### Prerequisites\n\nNsight Systems version is important, please reference `docker/Dockerfile.vllm.sglang.megatron` for the version we used.\n\n### Global profiling control\n\nverl has one single controller process and multiple worker processes. Both controller and worker processes can be profiled. Since the controller process can be executed in any nodes in the cluster, there is a message printed in the logging to indicate the controller process node hostname and process id.\n\nIn `global_profiler`, three new config entries control the profiler behaviors:\n\n* **`global_profiler.steps`**. List of step numbers at which profiling should be performed. For example: [1, 2, 5] will profile steps 1, 2, and 5. And ``null`` means no profiling.\n\n* **`global_profiler.profile_continuous_steps`**. If true, and the following `global_profiler.discrete==False`, then the continuous steps in `global_profiler.steps` will be combined into one database. For example the above step 1 and 2 are in one database, and 5 in another. If false, every step occupies at least one database. The reason for this config is to observe the program behaviors between steps.\n\nNsys options in controller nodes and worker nodes are configured in `global_profiler.global_tool_config.nsys`:\n\n* **`global_profiler.global_tool_config.nsys.controller_nsight_options`**. This config group is for the single controller. All fields in this config group will be just sent to Nsight Systems when Ray starts the controller process. `ppo_trainer.yaml` provides a workable example. Users can reference [Nsight Systems manual](https://docs.nvidia.com/nsight-systems/UserGuide/index.html) and [Ray user guide](https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html) for more details.\n* **`global_profiler.global_tool_config.nsys.worker_nsight_options`**. This config group is for the worker processes. Similarly all fields in this config group will be just sent to Nsight Systems when Ray starts the controller process. Capture range is used to control the profiler when to start and stop. So `capture-range: \"cudaProfilerApi\"` is fixed and does not change it. Users can change `capture-range-end` with some accurate calculation or just leave it `null`.\n\n### Worker process profiling\n\nVerl manages mulitiple RL roles, _Actor_, _Ref_, _Rollout_, _Critic_, _Reward_, which are implemented in different Worker classes. And these workers can be combined into one Ray Actor, running in a process group. Each RL role has its own profiling config group, `profiler`, which consists of three fields:\n\n* **`all_ranks` and `ranks`**. When `all_ranks` is set `True` then all ranks will be profiled; when set `False`, `ranks` will be profiled. By default, verl profiles the whole training process in a series ` worker_process_<PID>.<RID>.nsys-rep` files for each process rank. PID is the process ID; RID is the capture range ID.\n* **`discrete`**. When set `False`, all the roles actions in one training step will be dumped in one database. When set `True`, the actions annotated by `DistProfiler.annotate` will be dumped into a discrete database. In this case, each role's action occupies one `<RID>`.\n* **Verl collocate mode**. Verl can combine two Worker sub classes to one Worker Actor. In this case, the user should take care that the combined Workers have consistent `discrete`. The Nsight Systems profiler uses a `torch.cuda.profiler.start()` and `stop()` pair to dump a `<step>` database anyway.\n\n### where to find the profiling data\n\nBy default the `*.nsys-rep` files are saved in the directory `/tmp/ray/session_latest/logs/nsight/` at each node. According to the Ray manual, this default directory is not changeable. [&#34;however, Ray preserves the `--output` option of the default config&#34;](https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html).\n\nSome users may think it is not convenient, but it is understandable that Ray may start hundreds of processes and it would be a big network file system pressure if we save the files in one central place.\n\n## Usage Example\n\nTo enable profiling for specific components and steps, modify your ppo_trainer.yaml like this:\n\n### Disable profiler\n\n```yaml\n    profiler:\n        steps: null # disable profile\n```\n\n### Enable profiler and one database for one training step\n\n```yaml\n    global_profiler:\n        steps: [1, 2, 5]\n        discrete: False\n    actor_rollout_ref:\n        actor:\n            profiler:\n                enable: True\n                all_ranks: True\n        # rollout & ref follow actor settings\n    critic:\n            profiler:\n                enable: True\n                all_ranks: True\n    reward_model:\n            profiler:\n                enable: True\n                all_ranks: True\n```\n\n### Enable profiler and multiple databases for one training step\n\n```yaml\n    profiler:\n        steps: [1, 2, 5]\n        discrete: True\n```\n\n## Profiling Output\n\nWhen profiling is enabled, verl will generate Nsight Systems profiles for the specified components and steps. The profiles will include:\n\n- CUDA kernel execution\n- Memory operations\n- CPU-GPU synchronization\n- NVTX markers for key operations\n\nNsight Systems supports multi-report view, to open multiple databases together. In this mode, different processes and steps can be aligned in one time line for better analysis.\n"
  },
  {
    "path": "docs/perf/perf_tuning.rst",
    "content": "Performance Tuning Guide\n==============================\n\nLast updated: 07/17/2025.\n\nAuthor: `Guangming Sheng <https://github.com/PeterSH6>`_, `Jiali Zheng <https://github.com/CurryRice233>`_\n\nIn this section, we will discuss how to tune the performance of all the stages in verl, including:\n\n1. Rollout generation throughput.\n\n2. Enable ``use_remove_padding=True`` for sequence packing (i.e., data packing and remove padding).\n\n3. Batch size tuning for forward and backward computation\n\n4. Enable ``use_dynamic_bsz=True`` for higher throughput.\n\n5. Utilize Ulysses Sequence Parallel for Long Context Training\n\n6. LigerKernel for SFT performance optimization\n\n7. Forward prefetch in FSDP training backend\n\n8. Memory optimization for entropy calculation from logits\n\nRollout Generation Tuning\n--------------------------\n\nverl currently supports two rollout backends: vLLM and TGI (with SGLang support coming soon). \n\nBelow are key factors for tuning vLLM-based rollout. Before tuning, we recommend setting ``actor_rollout_ref.rollout.disable_log_stats=False`` so that rollout statistics are logged.\n\n- Increase ``gpu_memory_utilization``.\n\n  - For vLLM v0.7.0 and later, the vLLM instance will only use gpu_memory_utilization of the **total** memory.\n  - For SGLang, it's the fraction of the free GPU memory used for **static** memory like model weights and KV cache. However, the remaining (1-gpu_memory_utilization) will also be used during inference.\n\n  However, if model parameters and optimizer states are not offloaded, using too high a fraction can lead to OOM. \n  A value between 0.5 and 0.7 often strikes a good balance between high throughput and avoiding OOM.\n\n  Note: since the definition of ``gpu_memory_utilization`` varies across inference engines, a value that works well for one engine may cause OOM for another.\n\n- Adjust ``max_num_seqs`` or ``max_num_batched_tokens``.\n  If the GPU cache utilization is relatively low in the log, increase ``max_num_seqs`` or ``max_num_batched_tokens`` \n  can enlarge the effective batch size in the decoding stage, allowing more concurrent requests per batch. \n  We recommend setting ``max_num_batched_tokens > 2048`` for higher throughput.\n\n- Use a smaller ``tensor_parallel_size``. \n  When GPU resources allow, a smaller tensor parallel size spawns more vLLM replicas. \n  Data parallelism (DP) can yield higher throughput than tensor parallelism (TP), but also increases KVCache consumption. \n  Carefully balance the trade-off between more replicas and higher memory usage.\n  Our experiment in Sec. 8.4 of `HybridFlow paper <https://arxiv.org/pdf/2409.19256v2>`_ evaluate this trade-off.\n\n- Balance performance and memory using ``cudagraph_capture_sizes``.\n  If ``cudagraph_capture_sizes`` is set, vLLM will try to capture the model execution graph for different batch sizes.\n  Since cudagraph memory can not be offloaded to cpu, The memory stay in gpu when update actor is running. \n  Using smaller batch sizes can avoid OOM but slightly reduce throughput.\n  Must to set ``enforce_eager=False`` to use ``cudagraph_capture_sizes``.\n\nMore tuning details such as dealing with Preemption and Chunked-prefill\ncan be found in `vLLM official tuning guide <https://docs.vllm.ai/en/latest/performance/optimization.html>`_ \n\nFor optimal performance, we recommend using vLLM v0.8.3 or later. See https://github.com/volcengine/verl/blob/main/docs/README_vllm0.8.md for details.\n\nEnable remove padding (sequence packing)\n-----------------------------------------\n\nCurrently, for llama, mistral, gemma1 and qwen based models, users can enable `use_remove_padding=True` to utilize the \nsequence packing implementation provided by transformers library.\n\nFor other models, transformers library may also support it but we haven't tested it yet.\nUsers can add the desired model config to the  `test_transformer.py <https://github.com/volcengine/verl/blob/main/tests/models/test_transformer.py#L24>`_ file.\nAnd test its functionality by running the following command:\n\n.. code-block:: bash\n\n  pytest -s tests/models/test_transformer.py\n\nIf the test passes, you can add your desired model into the model `registry.py <https://github.com/volcengine/verl/blob/main/verl/models/registry.py#L24>`_ file.\nThen, you can enjoy the performance boost of sequence packing\nand welcome to PR your tested model to verl!\n\n\nBatch Size Tuning\n-----------------\n\nTo achieve higher throughput in experience preparation (i.e., model fwd) and model update (i.e., actor/critic fwd/bwd), \nusers may need to tune the ``*micro_batch_size_per_gpu`` for different computation.\n\nIn verl, the core principle for setting batch sizes is:\n\n- **Algorithmic metrics** (train batch size, PPO mini-batch size) are *global* (from a single-controller perspective), \n  normalized in each worker. See the `normalization code <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py#L120-L122>`_.\n\n- **Performance-related parameters** (micro batch size, max token length for dynamic batch size) are *local* parameters that define the per-GPU data allocations. \n  See the `normalization code <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py#L127>`_.\n\n.. note:: In your training script, please use ``*micro_batch_size_per_gpu`` instead of ``*micro_batch_size``. \n  So that you don't need to consider the normalization of the ``micro_batch_size`` and ``micro_batch_size`` will be deprecated.\n\nBatch Size Tuning tips\n\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\n\nTherefore, users may need to tune the ``*micro_batch_size_per_gpu`` to accelerate training. Here're some tips:\n\n1. **Enable gradient checkpointing**: \n   Set ``actor_rollout_ref.model.enable_gradient_checkpointing=True`` and ``critic.model.enable_gradient_checkpointing=True``. \n   This often allows for larger micro-batch sizes and will be beneficial for large mini-batch training.\n\n2. Increase the ``*micro_batch_size_per_gpu`` as much as possible till equals to normalized ``mini_batch_size``.\n\n3. **Use larger forward-only parameters**: \n   Forward only parameter, such as ``actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu``, \n   ``actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu``, ``critic.forward_micro_batch_size_per_gpu`` could be larger (e.g., 2x) than training related micro batch sizes,\n   such as ``actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu``, ``critic.ppo_micro_batch_size_per_gpu``.\n\n4. **Allow larger micro-batch sizes for Critic and Reward models**:\n   micro batch size of Critic and Reward model could be larger than Actor model. This is because the actor model has much larger vocab size in the final layer.\n\n5. **Enable activation offloading**:\n   Set ``actor_rollout_ref.model.enable_activation_offload=True`` and ``critic.model.enable_activation_offload=True``.\n   This often works together with gradient checkpointing to get larger micro-batch sizes and it's only available in FSDP backend now.\n\nTuning for Dynamic Batch Size\n-----------------------------\n\nDynamic batch size is a technique that allows the model to process similar number of tokens in a single forward pass (with different actual batch sizes).\nThis can significantly improve the training efficiency and reduce the memory usage.\n\nTo utilize this technique, users can set ``use_dynamic_bsz=True`` in actor, ref, critic and reward models.\nWith ``use_dynamic_bsz=True``, users don't need to tune ``*micro_batch_size_per_gpu``. \nInstead, users should tune the following parameters:\n\n- ``actor_rollout_ref.actor.ppo_max_token_len_per_gpu``, ``critic.ppo_max_token_len_per_gpu``: \n  The maximum number of tokens to be processed in fwd and bwd of ``update_policy`` and ``update_critic``.\n\n- ``actor_rollout_ref.ref.log_prob_max_token_len_per_gpu`` and ``actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu``: \n  The maximum number of tokens to be processed in a the fwd computation of ``compute_log_prob`` and ``compute_ref_log_prob``.\n\n- ``critic.forward_micro_batch_size_per_gpu``, ``reward_model.forward_micro_batch_size_per_gpu``: \n  The maximum number of tokens to be processed in a the fwd computation of ``compute_values``, ``compute_rm_score``.\n\nDynamic Batch Size Tuning tips\n\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\n\nHere're some tips to tune the above parameters:\n\n1. **Increase** ``actor_rollout_ref.actor.ppo_max_token_len_per_gpu``  \n   Make it at least 2 x (max_prompt_length + max_response_length). We set it to 3x in `run_qwen2-7b_rm_seq_balance.sh <https://github.com/volcengine/verl/blob/main/examples/ppo_trainer/run_qwen2-7b_rm_seq_balance.sh#L25>`_.\n   Try to increase it to get higher throughput.\n\n2. **Forward-only parameters can be larger**: \n   Similar to the non-dynamic-batch scenario, forward-only token limits can exceed those used in forward/backward operations.\n \n3. **Use larger limits for Critic and Reward models**:\n   Critic and Reward parameters can be set at least 2× the Actor’s limits. For instance, we set them to 4× here:  \n   `run_qwen2-7b_rm_seq_balance.sh <https://github.com/volcengine/verl/blob/main/examples/ppo_trainer/run_qwen2-7b_rm_seq_balance.sh#L40>`_\n   \n.. :math:`\\text{critic.ppo_max_token_len_per_gpu}  = 2 \\times  \\text{actor.ppo_max_token_len_per_gpu})`.\n\nUlysses Sequence Parallel for Long Context Training\n----------------------------------------------------\n\nTo utilize this technique, users can set ``ulysses_sequence_parallel_size>1`` in actor, ref, critic and reward models.\n\nWe support different model utilize different ulysses_sequence_parallel_size sizes.\n\nTo train long sequence (>32k), users may need to decrease the ``*micro_batch_size_per_gpu`` and ``*max_token_len_per_gpu`` to avoid OOM.\n\nLigerKernel for SFT\n----------------------\n\nLigerKernel is a high-performance kernel for Supervised Fine-Tuning (SFT) that can improve training efficiency. To enable LigerKernel in your SFT training:\n\n1. Install liger-kernel via ``pip3 install liger-kernel``. In your SFT configuration file (e.g., ``verl/trainer/config/sft_trainer.yaml``), set the ``use_liger`` parameter:\n\n   .. code-block:: yaml\n\n      model:\n        use_liger: True  # Enable LigerKernel for SFT\n\n2. The default value is ``False``. Enable it only when you want to use LigerKernel's optimizations.\n\n3. LigerKernel is particularly useful for improving training performance in SFT scenarios.\n\nForward prefetch in FSDP training backend\n----------------------\n\nDuring the training phase, users can enable forward prefetching in FSDP by setting ``fsdp_config.forward_prefetch=True``. For example, ``actor_rollout_ref.actor.fsdp_config.forward_prefetch=True``. This configuration prefetches the next forward-pass all-gather operation before completing the current forward computation, overlapping communication with computation and improving efficiency. For further details, refer to the `FSDP forward_prefetch <https://docs.pytorch.org/docs/stable/fsdp.html#module-torch.distributed.fsdp>`_ documentation.\n\n.. note::\n    Backward prefetch is unsupported because the ``BACKWARD_POST`` policy may prefetch incorrectly in nested-module cases. For details, see the `FSDP documentation <https://github.com/pytorch/torchtitan/blob/main/docs/fsdp.md?plain=1#L70>`_\n\nMigrating to FSDP2\n----------------------\n\nFSDP2 offers notable improvements over FSDP1. According to `PyTorch TorchTitan benchmarks <https://arxiv.org/abs/2410.06511v1>`_:\n\n- 7% lower GPU memory usage on average\n- 1.5% throughput improvement with BF16 training\n- Better composability with DTensor and per-parameter sharding\n\n**Enabling FSDP2 in VERL:**\n\n   .. code-block:: python\n\n    # Enable FSDP2 in actor configuration\n    actor_rollout_ref.actor.strategy=\"fsdp2\"\n\n.. note:: \n   FSDP2 requires PyTorch 2.1+ and is recommended for models with transformer architecture.\n\nMemory optimization for entropy calculation from logits\n----------------------\n\nThe ``logits`` tensor (typically of shape ``[bsz*seq_len, voc]``) can consume significant memory. When using ``compute_entropy_from_logits``, memory usage reaches approximately ``[bsz*seq_len, voc] × (4 bytes (float32) + 2 bytes (autocast for softmax+logsumexp) + 1 byte (softmax output))``.\n\nTo reduce this memory peak, enable chunked computation by setting:\n``actor_rollout_ref.ref.entropy_from_logits_with_chunking = True``\nThis processes the tensor in chunks of shape ``[chunk_size, voc]`` (e.g., 2048) rather than the full sequence length, exclusively during the model's forward pass.\n\nAdditionally, during training, standard gradient checkpointing (``enable_gradient_checkpointing=True``) does not apply to entropy calculations. To reduce memory peaks in this context, set:\n``actor_rollout_ref.actor.entropy_checkpointing = True``\nThis enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training.\n"
  },
  {
    "path": "docs/perf/perf_tuning_on_ascend.rst",
    "content": "Performance Tuning Guide on Ascend\n====================================\n\nLast updated:  01/29/2026.\n\nAuthor:  `Xiaobo Hu <https: //github.com/tardis-key>`_, `Haozhe Li <https: //github.com/ZLiao097>`_\n\n`Perf Tuning <https: //github.com/verl-project/verl/blob/main/docs/perf/perf_tuning.rst>`_ 中介绍的性能调优方法在昇腾设备中同样适用。本文重点介绍了昇腾特有的一些调优手段，包括融合算子优化、特定硬件配置和昇腾亲和特性等。\n\n融合算子\n--------------------------\n\n常用融合算子列表\n**********************************\n\n融合算子的优化原理为，通过数学意义上的等价替换，将多个算子融为一个算子的计算，减少冗余计算，同时减少下发次数，从而提高性能。几个典型的NPU融合算子列举如下，目前均已在 npu_patch.py 中对 Qwen2、Qwen3 系列模型完成替换。\n\n当前verl中使用的全量融合算子请查阅 `npu_patch.py <https: //github.com/verl-project/verl/blob/main/verl/models/transformers/npu_patch.py>`_ \n\nMatrix Computation-Communication operator fusion (MC2) \n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\nMC2 是 CANN 中一系列计算通信融合算子的统称，这些算子将原本串行的通信和计算操作融合在一起，通过内部的切分和流水线并行执行来优化性能。\n\n在 vllm-ascend 中，可以通过指定环境变量：\n\n.. code-block:: sh\n\n    export VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE=1\n\n在前向计算的 ``RowParallelLinear`` 中使能 ``torch_npu.npu_mm_all_reduce_base`` ，将分离的 ``matmul`` 和 ``allreduce`` 合并为一个融合算子。\n\n`RotaryMul&RotaryMulGrad <https: //www.hiascend.com/document/detail/zh/Pytorch/730/ptmoddevg/trainingmigrguide/performance_tuning_0030.html>`_\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\ntorch_npu 接口:  ``torch_npu.npu_rotary_mul(x, r1, r2)``\n\n参数说明: \n\n- x: q，k，shape要求输入为4维，一般为 ``[B, N, S, D]`` 或 ``[B, S, N, D]`` 或 ``[S, B, N, D]`` 。\n\n- r1: cos值 ，shape要求输入为4维，一般为 ``[1, 1, S, D]`` 或 ``[1, S, 1, D]`` 或 ``[S, 1, 1, D]`` 。\n\n- r2: sin 值，shape要求输入为4维，一般为 ``[1, 1, S, D]`` 或 ``[1, S, 1, D]`` 或 ``[S, 1, 1, D]`` 。\n\n`RmsNorm&RmsNormGrad <https: //www.hiascend.com/document/detail/zh/Pytorch/730/ptmoddevg/trainingmigrguide/performance_tuning_0031.html>`_\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\ntorch_npu 接口:  ``torch_npu.npu_rms_norm(self, gamma, epsilon=1e-06) -> (Tensor, Tensor)`` \n参数说明: \n\n- self: Tensor 类型，shape 支持 1-8 维。\n\n- gamma: Tensor 类型，通常为weight，shape 要求与 self 的后几维保持一致。\n\n- epsilon: Float 数据类型，用于防止除 0 错误。\n\n输出说明: \n\n- 第 1 个输出为 Tensor，计算公式的最终输出y。\n\n- 第 2 个输出为 Tensor， rms_norm 的中间结果 rstd ，用于反向计算。\n\n`Swiglu <https: //www.hiascend.com/document/detail/zh/Pytorch/730/ptmoddevg/trainingmigrguide/performance_tuning_0035.html>`_\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\ntorch_npu 接口:  ``torch_npu.npu_swiglu(Tensor self, int dim=-1) -> (Tensor)`` \n\n参数说明: \n\n- self: Tensor 类型，shape支持 1-8 维。\n\n- dim: Int 类型，默认为 -1。\n\n输出说明: \n\n- 输出为 Tensor，计算公式的最终输出 y。\n\n`GroupMatMul <https: //www.hiascend.com/document/detail/zh/Pytorch/730/apiref/torchnpuCustomsapi/docs/context/torch_npu-npu_grouped_matmul.md>`_\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n函数原型：\n\n.. code:: python\n\n    npu_grouped_matmul(\n        x, \n        weight, \n        *, \n        bias=None, \n        scale=None, \n        offset=None, \n        antiquant_scale=None, \n        antiquant_offset=None, \n        per_token_scale=None, \n        group_list=None, \n        activation_input=None, \n        activation_quant_scale=None, \n        activation_quant_offset=None, \n        split_item=0, group_type=None, \n        group_list_type=0, \n        act_type=0, \n        output_dtype=None, \n        tuning_config=None\n    ) -> List[Tensor]\n\n详细使用方法见标题文档链接\n\nFSDP后端融合算子使用方法\n**********************************\n\n在 ``verl/models/transformers/npu_patch.py`` 目录中，已经把可用的融合算子通过 patch 的形式进行替换，无需进行其他操作即可默认进行使用\n\nMegatron后端融合算子使用方法\n**********************************\n\nMegatron 的融合算子集成在 MindSpeed 中，需要添加特定参数开启: \n\n1. **Flash Attention（必须开启）**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True\n       ++actor_rollout_ref.ref.megatron.override_transformer_config.use_flash_attn=True\n\n2. **RotaryMul**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True\n       +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rotary_pos_emb=True\n\n3. **RMSNorm**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_rmsnorm=True\n\n4. **GroupMatMul**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True\n\n5. **Swiglu**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.use_fused_swiglu=True\n\n6. **Permute/Unpermute**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.fused_permute_unpermute=True\n\n7. **MC2**\n   ::\n\n       +actor_rollout_ref.actor.megatron.override_transformer_config.use_ascend_mc2=True\n\n昇腾通用配置\n--------------------------\n\n`算子下发 <https: //www.hiascend.com/document/detail/zh/Pytorch/730/comref/Envvariables/docs/zh/environment_variable_reference/TASK_QUEUE_ENABLE.md>`_\n************************************************************************************************************************************************************************************************************\n\n通过 ``TASK_QUEUE_ENABLE`` 可配置 task_queue 算子下发队列优化等级，默认为 Level 1 优化。该配置可以减少host下发时间，可用于缓解由下发导致的整体free过大问题。\n\n.. image :: https://github.com/verl-project/verl-data/blob/main/images/ascend/perf_tuning_task_queue.png\n    :width: 500px\n\nLevel 0 : 不开启下发流水优化。\n\nLevel 1 : \\ 将算子下发任务分为两段，一部分任务（主要是 aclnn 算子的调用）放在新增的二级流水上，一、二级流水通过算子队列传递任务，相互并行，通过部分掩盖减少整体的下发耗时，提升端到端性能。\n\nLevel 2 : \\ 基于 Level 1 的优化进一步平衡了一、二级流水的任务负载，主要是将 workspace 相关任务迁移至二级流水，掩盖效果更好，性能收益更大。该配置仅在二进制场景生效，建议配置值为 Level 2 优化。\n\n`通讯算法编排展开 <https: //www.hiascend.com/document/detail/zh/canncommercial/850/maintenref/envvar/envref_07_0096.html>`_\n************************************************************************************************************************************************************************************************************\n使用环境变量 ``HCCL_OP_EXPANSION_MODE=AIV`` 用于配置通信算法的编排展开位置，支持如下取值: \n\n- **AI_CPU:** 代表通信算法的编排展开位置在 Device 侧的 AI CPU，Device 侧根据硬件型号自动选择相应的调度器。\n\n- **AIV:** 代表通信算法的编排展开位置在 Device 侧的 Vector Core，执行也在 Vector Core。\n\n- **HOST:** 代表通信算法的编排展开位置为 Host 侧 CPU，Device 侧根据硬件型号自动选择相应的调度器。\n\n- **HOST_TS:** 代表通信算法的编排展开位置为 Host 侧 CPU，Host 向 Device 的 Task Scheduler 下发任务，Device 的 Task Scheduler 进行任务调度执行。\n\n推理阶段调优\n--------------------------\n\nChunked Prefill in V1\n***************************\n\nVLLM 当前版本已默认启用 VLLM V1，使用以下配置启用 Chunked Prefill：\n\n.. code-block:: sh\n\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n\n原理参考 `VLLM 官方文档 <https://docs.vllm.ai/en/v0.4.2/models/performance.html>`_。\n\nGraph Mode\n***************************\n\n与 CUDA 类似，NPU 通过以下配置启用 **ACL Graph**：\n\n.. code-block:: sh\n\n    actor_rollout_ref.rollout.enforce_eager=False\n\n文档：`ACL Graph <https://docs.vllm.ai/projects/ascend/en/latest/developer_guide/feature_guide/ACL_Graph.html>`_\n\n.. note::\n    ACL Graph 与 ``taskqueue Level 2`` 原理冲突，**二者无法同时开启**。\n\n训练阶段调优\n--------------------------\n\nFSDP\n**********************************\n\n.. csv-table::\n   :header: \"FSDP\", \"说明\"\n   :widths: 30, 60\n\n   \"/\",\"仅切分优化器(Zero-1)\"\n   SHARD_GRAD_OP,切分梯度和优化器(Zero-2)\n   \"HYBRID_SHARD\",\"切分权重、梯度和优化器(Zero-3)\"\n   \"2D device_mesh+HYBRID_SHARD\",\"又称HSDP（FSDP+DDP）例如device_mesh=[2,8], 每8个rank为一个FSDP组，组内进行FSDP切分，共有两个组，两个组间进行DDP，通过allreduce同步梯度。\"\n   \"2D device_mesh+HYBRID_SHARD_ZERO2\",\"HSDP的Zero2版本\"\n   NO_SHARD,DDP\n\nFSDP 不支持 Zero-1， VeRL中会根据卡数和 ``actor_rollout_ref.actor.fsdp_config.fsdp_size``  来决定 device mesh 的取值，默认使用 Zero-3 进行切分；如果模型较小（建议小于 7B 时），可以通过控制参数 ``actor_rollout_ref.actor.fsdp_config.reshard_after_forward`` 为 ``True`` 在 FSDP/FSDP2 上使用 Zero-2 来优化性能.\n\nMegatron\n**********************************\n\n在模型较大时，使用 Megatron 作为训练后端可以更灵活的进行性能调优。\n\n当 DP 并行显存无法容纳模型时，优先开启 TP 来切分模型权重，如果模型仍然过大，再开启 PP 来进一步切分；如果序列过长导致激活太大，则可以开启 CP 和 SP 来进行优化；在 MoE 模型中则可以额外开启 EP 来控制对专家的切分，如果专家过小，为了避免将权重切的果味细碎，则可以开启 ETP 来避免 MoE 部分的 TP 切分，而将多个完整的专家分布到 DP 和 TP 上。\n\nTP、PP、EP、ETP和 Megatron 使用方式一样，CP 和 SP 在 NPU 上开启方式: \n\n- SP: ``Sequence Parallel`` 在 Tensor Parallel 的基础上进一步提高计算效率，是一种通过将输入数据的序列维度进行切分的并行计算方式。在 NPU 上通过 MindSpeed 来调用SP:\n  ::\n\n      actor_rollout_ref.actor.megatron.override_transformer_config.sequence_parallel=True\n\n- CP: ``Context Parallel`` 是一种在多个 GPU/NPU 上并行处理神经网络激活值的方法，他通过在序列维度上对输入张量进行划分来实现。在 NPU 上通过 MindSpeed 来调用 CP （两个参数必须同时添加）:\n  ::\n\n      actor_rollout_ref.actor.megatron.context_parallel_size\n      actor_rollout_ref.actor.megatron.override_transformer_config.context_parallel_size\n\nMegatron-distributed optimizer\n**********************************\n\n在面对较大尺寸模型时，通常需要将优化器分片到一个 DP 域内的每张卡上来节省显存。Megatron 后端下在 NPU 上开启分布式优化器:\n\n::\n\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_distributed_optimizer=True\n"
  },
  {
    "path": "docs/perf/torch_profiling.md",
    "content": "# PyTorch Profiling in verl\n\nLast updated: 01/13/2026.\n\nThis guide explains how to use the native [PyTorch Profiler](https://pytorch.org/tutorials/recipes/recipes/profiler_recipe.html) for profiling verl training runs.\n\n## Configuration\n\nProfiling in verl can be configured through parameters in the trainer configuration file (e.g., `ppo_trainer.yaml`).\n\n### Global Profiling Control\n\nIn `global_profiler`, you can control when and how profiling occurs globally:\n\n* **`global_profiler.steps`**: List of step numbers to profile. E.g., `[1, 2, 5]` profiles steps 1, 2, and 5. Set to `null` to disable.\n* **`global_profiler.save_path`**: Directory to save the profiling results. Default is `outputs/profile`.\n\n### Role Profiling Control\n\nEach RL role (Actor, Critic, etc.) has its own `profiler` configuration:\n\n* **`enable`**: Whether to enable profiling for this role.\n* **`all_ranks`**: If `True`, profiles all ranks.\n* **`ranks`**: List of specific ranks to profile if `all_ranks` is `False`.\n* **`tool_config.torch`**: Configuration specific to the PyTorch Profiler.\n\n#### PyTorch Profiler Options (`tool_config.torch`)\n\nYou can customize the PyTorch Profiler behavior using the following fields under `tool_config.torch`:\n\n* **`contents`**: List of contents to profile.\n    *   **`cpu`**: Profile CPU activities.\n    *   **`cuda`**: Profile CUDA activities.\n    *   **`memory`**: Track tensor memory allocation/free.\n    *   **`shapes`**: Record shapes of operator inputs.\n    *   **`stack`**: Record source code file and line number.\n* **`schedule`**: (Advanced) configuration for `wait`, `warmup`, `active`, `repeat` cycles.\n\n## Examples\n\n### 1. End-to-End Collection\n\nCollects performance data for all steps in a single trace file.\n\n```yaml\nglobal_profiler:\n  steps: [1, 2, 5]\n  save_path: ./outputs/profile\n\nactor_rollout_ref:\n  actor:\n    profiler:\n      enable: True\n      all_ranks: True\n      tool_config:\n        torch:\n          discrete: False\n          contents: [cpu, cuda]\n  # rollout & ref follow actor settings\n```\n\n### 2. Discrete Mode Collection\n\nDiscrete mode saves separate trace files for each step. This is useful for detailed analysis and is **mandatory** when using Agent Loop.\n\n**Configuration Example**\n\nThis configuration supports profiling both Training (Actor) and Inference (Rollout). You can enable/disable them independently.\n\n```yaml\nactor_rollout_ref:\n  actor:\n    profiler:\n      enable: True # Set to True to profile training\n      all_ranks: False\n      ranks: [0] # Global Rank 0\n      tool_config:\n        torch:\n          discrete: True\n          contents: [cpu, cuda]\n  rollout:\n    profiler:\n      enable: True # Set to True to profile inference\n      all_ranks: False\n      ranks: [0] # In Agent Loop, this is the Replica Rank (e.g. 0-th instance)\n      tool_config:\n        torch:\n          discrete: True # REQUIRED \n  # ref follow actor settings\n```\n\n**Agent Loop Mode Description**\n\nWhen Rollout runs in [Agent Loop](../advance/agent_loop.rst) mode, performance data for the Rollout phase **must be collected using discrete mode**. In this case, the Profiler is triggered by the inference engine backend.\n\n1. Rank Definition: ranks in the Rollout configuration refers to Replica Rank (inference instance index), not Global Rank.\n\n2. Inference Engine Support: Currently, vLLM and SGLang engines are supported without additional settings. Specific details are as follows:\n\n   *   **vLLM Engine**: Automatically collects AsyncLLM scheduling stacks and inference process performance data.\n   *   **SGLang Engine**: Automatically collects inference process performance data. Does not support the memory option in contents.\n\n## Visualization\n\nCollected trace files (usually `.json` or `.json.gz`) are stored in the configured `save_path`.\n\nYou can visualize them using:\n\n1.  **Chrome Tracing**: Open `chrome://tracing` in a Chrome browser and load the JSON file.\n2.  **Perfetto**: Open [ui.perfetto.dev](https://ui.perfetto.dev/) and load the file (recommended for large traces).\n3.  **TensorBoard**: If using the TensorBoard plugin for PyTorch Profiler.\n"
  },
  {
    "path": "docs/perf/verl_profiler_system.md",
    "content": "# verl Profiler System\n\nLast updated: 08/18/2025.\n\n## Architecture\n\nThe architecture of verl profiler system is like below:\n\n![verl-profiler-arch](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/2bc7ed0ba2f37f21707bfac3b241eca4b86d1bc6/docs/verl_profiler_arch.png)\n\nThere is a global profiler and tool configuration to set some common config in single controller level, deciding\n\n- `tool`: which tool to use\n- `steps`: which steps to profile\n- `save_path`: results saving path\n\nWhen some tool need to profile behavior of each role, configurations in role-level is needed:\n\n- `tool`: which tool to use\n- `enable`: whether enable profiling on this role\n- rank info: `all_ranks` and `rank` to decide which rank to profile or log output\n\nFor tool config in role-level, there are some detailed behavior needed to control, like the `discrete` mode in nsys profiler.\n\nEvery role has a profiler config, and by default, rollout/ref/reward models follow the Actor's behavior.\n\n## To Add a new profiling tool\n\nNew added profiling tool shall reuse the current APIs as much as possible.\n\n1. The logic of **whether to use the tool**: `tool == [new tool]`.\n2. Add the global and local tool config to `ppo_trainer.yaml`/`ppo_megatron_trainer.yaml` and each `[role].yaml`, under `global_tool_config.[new tool]` and `tool_config.[new tool]`\n3. The tool config should be implemented in `verl/utils/profiler/config.py`, inherit the `BaseConfig` class.\n4. Implement profiling tool initialization logic using configurations in `global_profiler.global_tool_config.[new tool]` and the results saving logics (can also save in role-level profile)\n5. For role function-level profiling, please follow the nsys profiler way in `nvtx_profiler.py`, implement a profiler class inherit `DistProfiler` and import new profiler in `verl/utils/profiler/__init__.py`\n6. Add unit test and examples for others to use in convinience."
  },
  {
    "path": "docs/preparation/prepare_data.rst",
    "content": "Prepare Data for Post-Training\n========================================\n\nLast updated: 02/09/2025.\n\nBefore starting the post-training job, we need to prepare the data for\nthe policy training. The data should be stored in the parquet format.\n\nWe provide several data preprocess scripts for different datasets,\nincluding GSM8K, MATH, HelloSwag, Full_hh_rlhf. To prepare other datasets, we need\nto follow the following steps: The data preprocess script can be divided\ninto two parts:\n\n1. The first part is the common part, which loads the dataset from\n   huggingface's ``datasets`` package. Then preprocess the datasets with\n   the ``make_map_fn`` and then store in the parquet format.\n\n.. code:: python\n\n   import re\n   import os\n   import datasets\n\n   from verl.utils.hdfs_io import copy, makedirs\n   import argparse\n\n   # To extract the solution for each prompts in the dataset\n   # def extract_solution(solution_str): \n   # ...\n\n\n   if __name__ == '__main__':\n       parser = argparse.ArgumentParser()\n       parser.add_argument('--local_dir', default='/opt/tiger/gsm8k')\n       parser.add_argument('--hdfs_dir', default=None)\n\n       args = parser.parse_args()\n\n       num_few_shot = 5\n       data_source = 'openai/gsm8k'\n\n       dataset = datasets.load_dataset(data_source, 'main')\n\n       train_dataset = dataset['train']\n       test_dataset = dataset['test']\n\n           # Construct a `def make_map_fn(split)` for the corresponding datasets.\n       # ...\n           \n       train_dataset = train_dataset.map(function=make_map_fn('train'), with_indices=True)\n       test_dataset = test_dataset.map(function=make_map_fn('test'), with_indices=True)\n\n       local_dir = args.local_dir\n       hdfs_dir = args.hdfs_dir\n\n       train_dataset.to_parquet(os.path.join(local_dir, 'train.parquet'))\n       test_dataset.to_parquet(os.path.join(local_dir, 'test.parquet'))\n\n       makedirs(hdfs_dir)\n\n       copy(src=local_dir, dst=hdfs_dir)\n\n2. The users are required to implement the ``make_map_fn()`` function\n   (as well as the ``extract_solution``) on their own to support\n   different datasets or tasks.\n\nWe already implemented the data preprocess of GSM8k, MATH, Hellaswag and Full_hh_rlhf\ndatasets. And we take the GSM8k dataset as an example:\n\n**GSM8K**\n\nIn the ``make_map_fn``, each data field should consist of the following\n5 fields:\n\n1. ``data_source``: The name of the dataset. To index the corresponding\n   reward function in the ``RewardModel``\n2. ``prompt``: This field should be constructed in the format of\n   huggingface chat_template. The tokenizer in ``RLHFDataset`` will\n   apply chat template and tokenize the prompt.\n3. ``ability``: Define the task category.\n4. ``reward_model``: Currently, we only utilize the ``ground_truth``\n   field during evaluation. The ``ground_truth`` is computed by the\n   ``extract_solution`` function. **NOTED** that the implementation of\n   the corresponding reward function should align with this extracted\n   ``ground_truth``.\n5. ``extra_info``: Record some information of the current prompt. Not\n   use for now.\n\n.. code:: python\n\n   def extract_solution(solution_str):\n       solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str) # extract the solution after ####\n       assert solution is not None\n       final_solution = solution.group(0)\n       final_solution = final_solution.split('#### ')[1].replace(',', '')\n       return final_solution\n\n   instruction_following = \"Let's think step by step and output the final answer after \\\"####\\\".\"\n\n   # add a row to each data item that represents a unique id\n   def make_map_fn(split):\n\n       def process_fn(example, idx):\n           question = example.pop('question')\n\n           question = question + ' ' + instruction_following\n\n           answer = example.pop('answer')\n           solution = extract_solution(answer)\n           data = {\n               \"data_source\": data_source,\n               \"prompt\": [{\n                   \"role\": \"user\",\n                   \"content\": question\n               }],\n               \"ability\": \"math\",\n               \"reward_model\": {\n                   \"style\": \"rule\",\n                   \"ground_truth\": solution\n               },\n               \"extra_info\": {\n                   'split': split,\n                   'index': idx\n               }\n           }\n           return data\n\n       return process_fn\n"
  },
  {
    "path": "docs/preparation/reward_function.rst",
    "content": "Implement Reward Function for Dataset\n======================================\n\nLast updated: 06/02/2025.\n\nFor each dataset, we need to implement a reward function or utilize a reward model to compute the rewards for the generated responses.\nWe already pre-implemented some reward functions in `reward_score directory <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score>`_.\nYou can also use customized reward functions.\n\nCurrently, we support reward functions for GSM8k and MATH datasets. For RLHF datasets (e.g.,\nfull_hh_rlhf) and Code Generation (e.g., APPS), we utilize reward model\nand SandBox (will opensource soon) for evaluation respectively.\n\nRewardManager\n-------------\n\nIn the entrypoint of the PPO Post-Training script `main_ppo.py <https://github.com/volcengine/verl/blob/main/verl/trainer/main_ppo.py#L33>`_,\nwe implement a ``RewardManager`` that utilize pre-implemented reward functions to compute the scores for each response.\n\nIn the ``RewardManager``, we implemented a ``__call__`` function to\ncompute the score for each response. \nAll the reward functions are executed by ``compute_score_fn``.\nThe input is a ``DataProto``, which includes:\n\n- ``input_ids``, ``attention_mask``: ``input_ids`` and ``attention_mask`` after applying\n  chat_template, including prompt and response\n- ``responses``: response tokens\n- ``ground_truth``: The ground truth string of the current prompt.\n  Stored in ``non_tensor_batch`` in the ``DataProto``, which should be\n  preprocessed in the parquet files.\n- ``data_source``: The dataset name of the current prompt. Stored in\n  ``non_tensor_batch`` in the ``DataProto``, which should be\n  preprocessed in the parquet files.\n\nAfter detokenize the responses, the responses string and the ground\ntruth string will be input to the ``compute_score_fn`` to compute the\nscore for each response.\n\nReward Functions\n----------------\n\nPre-implemented\n~~~~~~~~~~~~~~~\n\nWe already pre-implemented some reward functions in `reward_score directory <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score>`_.\n\n- In the `GSM8k example <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score/gsm8k.py>`_, we\n  force the response to output the final answer after four ####, then\n  use string matching to compare with the ground truth. If completely\n  correct, score 1 point; if the format is correct, score 0.1 points; if\n  the format is incorrect, score 0 points.\n- In the `MATH example <https://github.com/volcengine/verl/blob/main/verl/utils/reward_score/math.py>`_, we follow\n  the implementation in `lm-evaluation-harness repository <https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py>`_.\n\nCustomized\n~~~~~~~~~~\n\nYou can implement customized reward functions in a separate file and specify them using ``custom_reward_function.path`` and ``custom_reward_function.name``. For the set of them, please refer to :ref:`config-explain-page`.\n\nThe parameters of your reward function should be ``data_source``, ``solution_str``, ``ground_truth``, and ``extra_info``.\nFor example:\n\n.. code:: python\n\n  def my_reward_fn(data_source, solution_str, ground_truth, extra_info=None):\n    return len(solution_str)/100\n\nIf you are testing only a single customized reward function, you can simply name it 'compute_score' and leave ``custom_reward_function.name`` unset.\n\nTo run multiple tests with different customized reward functions, you can modify both ``custom_reward_function.path`` and ``custom_reward_function.name`` for each trial. \nFor instance, you might create a single `my_reward.py` file and implement multiple reward functions within it. This way, for different trials, you only need to adjust ``custom_reward_function.name``, making it more convenient to conduct multiple tests within scripts.\n"
  },
  {
    "path": "docs/requirements-docs.txt",
    "content": "# markdown support\r\nrecommonmark\r\nmyst_parser\r\n# markdown table support\r\nsphinx-markdown-tables\r\n\r\n# theme default rtd\r\n\r\n# crate-docs-theme\r\nsphinx-rtd-theme\r\n\r\n# pin tokenizers version to avoid env_logger version req\r\ntokenizers==0.21\r\n"
  },
  {
    "path": "docs/sglang_multiturn/interaction_system.rst",
    "content": "Interaction System for Multi-turn RL Training\n=============================================\n\nLast updated: 06/25/2025.\n\nOverview\n--------\n\nThe verl interaction system enables dynamic, multi-turn conversational feedback during reinforcement learning training. This system allows models to engage in iterative problem-solving scenarios where interaction agents can provide corrective feedback, guidance, or evaluation based on the model's responses.\n\n**New in Multi-Interaction Support**: The system now supports multiple named interactions within a single training session, enabling sophisticated training scenarios where different samples can use different interaction strategies. This allows for curriculum learning, domain-specific feedback, and flexible agent switching at the sample level.\n\nKey features:\n\n- **Async-based Architecture**: Non-blocking interaction processing for distributed training\n- **Instance Management**: Stateful session handling with unique instance IDs for concurrent interactions\n- **SGLang Integration**: Seamless integration with SGLang rollout system for multi-turn conversations\n- **Configuration-driven**: Dynamic agent loading via YAML configuration files\n- **Multi-Interaction Support**: Registry system enabling multiple named interactions per rollout\n- **Sample-Level Selection**: Each sample can specify which interaction to use via configuration\n- **Reward Integration**: Turn-level scoring mechanism integrated with verl's reward system\n\nArchitecture\n------------\n\nThe interaction system follows a plugin-based architecture with clear separation of concerns:\n\n.. code-block::\n\n    Interaction Registry System\n         ↓\n    BaseInteraction (Abstract Interface)\n         ↓\n    Multiple Named Interactions (e.g., Gsm8kInteraction, CustomInteraction)\n         ↓\n    SGLang Rollout Integration (interaction_map)\n         ↓\n    Sample-Level Interaction Selection\n         ↓\n    Async Request Lifecycle Management\n\nCore Components\n~~~~~~~~~~~~~~~\n\n**Interaction Registry System**\n\nThe interaction registry system allows loading and managing multiple named interactions:\n\n.. code-block:: python\n\n    from verl.interactions.utils.interaction_registry import initialize_interactions_from_config\n    \n    # Load multiple interactions from config\n    interaction_map = initialize_interactions_from_config(\"config.yaml\")\n    \n    # Access specific interaction by name\n    gsm8k_interaction = interaction_map[\"gsm8k\"]\n    custom_interaction = interaction_map[\"custom_solver\"]\n\n**BaseInteraction Interface**\n\nAll interaction agents must implement the ``BaseInteraction`` abstract class:\n\n.. code-block:: python\n\n    from verl.interactions.base import BaseInteraction\n    from typing import Dict, Any, List, Tuple, Optional\n\n    class BaseInteraction:\n        def __init__(self, config: Dict[str, Any]):\n            self.config = config\n            self.name: str = config.get(\"name\", \"interaction_agent\")\n        \n        async def start_interaction(self, instance_id: Optional[str] = None, **kwargs) -> str:\n            \"\"\"Initialize interaction session, return instance_id\"\"\"\n            \n        async def generate_response(self, instance_id: str, messages: List[Dict[str, Any]], **kwargs) -> Tuple[bool, str, float, Dict[str, Any]]:\n            \"\"\"Generate response, return (should_terminate, response, score, metadata)\"\"\"\n            \n        async def calculate_score(self, instance_id: str, **kwargs) -> float:\n            \"\"\"Calculate turn-level score for RL training\"\"\"\n            \n        async def finalize_interaction(self, instance_id: str, **kwargs) -> None:\n            \"\"\"Clean up resources\"\"\"\n\n**Request Lifecycle**\n\nThe interaction system integrates with SGLang's async rollout via state management:\n\n1. ``PENDING`` → Initialize interaction via ``start_interaction()``\n2. ``GENERATING`` → Model generates response\n3. ``INTERACTING`` → Process response via ``generate_response()``\n4. ``GENERATING`` → Continue if not terminated, otherwise ``COMPLETED``\n\nConfiguration\n-------------\n\n**Basic Setup**\n\nEnable interaction in your rollout configuration:\n\n.. code-block:: yaml\n\n    actor_rollout_ref:\n        rollout:\n            multi_turn:\n                enable: true\n                interaction_config_path: \"path/to/interaction_config.yaml\"\n                max_user_turns: 10\n                max_assistant_turns: 10\n\n**Interaction Configuration File**\n\nCreate an interaction configuration file (e.g., ``interaction_config.yaml``):\n\n**Single Interaction (Legacy Format)**\n\n.. code-block:: yaml\n\n    interaction:\n      - name: \"gsm8k\"\n        class_name: \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\"\n        config: {}\n\n**Multiple Interactions (New Format)**\n\n.. code-block:: yaml\n\n    interaction:\n      - name: \"gsm8k\"\n        class_name: \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\"\n        config: {}\n      - name: \"custom_solver\"\n        class_name: \"custom.interactions.CustomInteraction\"\n        config: \n          solver_type: \"advanced\"\n          timeout: 30\n      - name: \"code_verifier\"\n        class_name: \"verl.interactions.base.BaseInteraction\"\n        config: \n          verification_mode: \"strict\"\n\n**Automatic Name Generation**\n\nIf no ``name`` field is provided, the system will automatically generate one from the class name:\n\n.. code-block:: yaml\n\n    interaction:\n      - class_name: \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\"\n        config: {}\n        # Automatically generates name: \"gsm8k\"\n\nThe system will dynamically load all specified interaction classes and make them available by name.\n\nImplementation Example: GSM8K\n-----------------------------\n\nThe GSM8K interaction demonstrates a complete implementation for math problem-solving scenarios:\n\n.. code-block:: python\n\n    from verl.interactions.base import BaseInteraction\n    from verl.utils.reward_score import gsm8k\n    from uuid import uuid4\n\n    class Gsm8kInteraction(BaseInteraction):\n        def __init__(self, config: dict):\n            super().__init__(config)\n            self._instance_dict = {}\n\n        async def start_interaction(self, instance_id=None, ground_truth=None, **kwargs):\n            if instance_id is None:\n                instance_id = str(uuid4())\n            self._instance_dict[instance_id] = {\n                \"response\": \"\",\n                \"ground_truth\": ground_truth,\n                \"reward\": 0.0,\n            }\n            return instance_id\n\n        async def generate_response(self, instance_id, messages, **kwargs):\n            # Extract last assistant message content\n            content = \"\"\n            for item in reversed(messages):\n                if item.get(\"role\") == \"assistant\":\n                    content = item.get(\"content\", \"\")\n                    break\n\n            # Ensure GSM8K format (#### prefix)\n            self._instance_dict[instance_id][\"response\"] = content\n\n            reward = await self.calculate_score(instance_id)\n            if reward == 1.0:\n                return True, \"Your response is correct!\", 1.0, {}\n            else:\n                return False, \"Your response is incorrect! You need to reflect on your answer and try again.\", 0.0, {}\n\n        async def calculate_score(self, instance_id, **kwargs):\n            return gsm8k.compute_score(\n                self._instance_dict[instance_id][\"response\"],\n                self._instance_dict[instance_id][\"ground_truth\"],\n                method=\"strict\", format_score=0.0, score=1.0,\n            )\n\n        async def finalize_interaction(self, instance_id, **kwargs):\n            del self._instance_dict[instance_id]\n\nTraining Integration\n--------------------\n\n**Training Script Configuration**\n\nInclude interaction configuration in your training command:\n\n.. code-block:: bash\n\n    python3 -m verl.trainer.main_ppo \\\\\n        --config-path=\"$CONFIG_PATH\" \\\\\n        --config-name='gsm8k_multiturn_grpo_w_interaction' \\\\\n        algorithm.adv_estimator=grpo \\\\\n        data.train_batch_size=512 \\\\\n        data.return_raw_chat=True \\\\\n        actor_rollout_ref.rollout.name=sglang \\\\\n        actor_rollout_ref.rollout.multi_turn.interaction_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml\" \\\\\n        trainer.total_epochs=15\n\n**Data Requirements**\n\nEnsure your dataset includes interaction parameters with the ``name`` field for interaction selection:\n\n.. code-block:: python\n\n    # Dataset should include interaction_kwargs in non_tensor_batch\n    interaction_kwargs = [\n        {\"name\": \"gsm8k\", \"query\": \"What is 2+2?\", \"ground_truth\": \"4\"},\n        {\"name\": \"custom_solver\", \"query\": \"Solve: x^2 + 5x + 6 = 0\", \"ground_truth\": \"x = -2, -3\"},\n        {\"name\": \"gsm8k\", \"query\": \"What is 3+3?\", \"ground_truth\": \"6\"},\n    ]\n\n**Sample-Level Interaction Selection**\n\nEach sample can specify which interaction to use via the ``name`` field. This enables flexible training scenarios where different samples use different interaction strategies:\n\n.. code-block:: python\n\n    # Example: Math problems use GSM8K interaction, code problems use code verifier\n    data_samples = [\n        {\n            \"prompt\": \"What is 15% of 200?\",\n            \"interaction_kwargs\": {\n                \"name\": \"gsm8k\",\n                \"query\": \"What is 15% of 200?\", \n                \"ground_truth\": \"30\"\n            }\n        },\n        {\n            \"prompt\": \"Write a function to check if a number is prime\",\n            \"interaction_kwargs\": {\n                \"name\": \"code_verifier\",\n                \"code_type\": \"python\",\n                \"expected_behavior\": \"return True for prime numbers\"\n            }\n        }\n    ]\n\n**Backward Compatibility**\n\nIf no ``name`` field is provided in ``interaction_kwargs``, the system defaults to ``\"gsm8k\"`` for backward compatibility.\n\nBest Practices\n--------------\n\n**Resource Management**\n\n- Always implement proper cleanup in ``finalize_interaction()``\n- Use unique instance IDs to avoid conflicts in concurrent training\n- Handle edge cases like empty messages or malformed content\n\n**Performance Optimization**\n\n- Keep interaction logic lightweight to avoid blocking training\n- Use async/await properly to maintain non-blocking behavior\n- Consider caching expensive computations within interaction instances\n\n**Testing**\n\nComprehensive testing is essential for interaction systems:\n\n.. code-block:: python\n\n    import pytest\n    from unittest.mock import patch\n\n    @pytest.mark.asyncio\n    async def test_interaction_workflow():\n        interaction = YourInteraction({})\n        \n        # Test complete workflow\n        instance_id = await interaction.start_interaction(ground_truth=\"expected_answer\")\n        \n\n        messages = [{\"role\": \"user\", \"content\": \"user_content\"}, {\"role\": \"assistant\", \"content\": \"assistant_content\"}]\n        should_terminate, response, reward, metadata = await interaction.generate_response(instance_id, messages)\n        \n        assert should_terminate in [True, False]\n        assert isinstance(reward, float)\n        \n        await interaction.finalize_interaction(instance_id)\n\nAdvanced Usage\n--------------\n\n**Multi-Interaction Training Strategies**\n\nYou can design sophisticated training scenarios using multiple interactions:\n\n.. code-block:: python\n\n    # Example: Progressive difficulty with different interaction agents\n    class MathTrainingPipeline:\n        def create_interaction_config(self):\n            return {\n                \"interaction\": [\n                    {\n                        \"name\": \"basic_math\",\n                        \"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\",\n                        \"config\": {\"difficulty\": \"easy\"}\n                    },\n                    {\n                        \"name\": \"advanced_math\", \n                        \"class_name\": \"custom.interactions.AdvancedMathInteraction\",\n                        \"config\": {\"difficulty\": \"hard\", \"allow_hints\": True}\n                    },\n                    {\n                        \"name\": \"competition_math\",\n                        \"class_name\": \"custom.interactions.CompetitionMathInteraction\", \n                        \"config\": {\"time_limit\": 300, \"show_steps\": False}\n                    }\n                ]\n            }\n    \n        def create_curriculum_data(self, epoch):\n            if epoch < 5:\n                return [{\"name\": \"basic_math\", ...} for _ in samples]\n            elif epoch < 10:\n                return [{\"name\": \"advanced_math\", ...} for _ in samples]\n            else:\n                return [{\"name\": \"competition_math\", ...} for _ in samples]\n\n**Custom Scoring Functions**\n\nYou can integrate custom reward functions:\n\n.. code-block:: python\n\n    async def calculate_score(self, instance_id, **kwargs):\n        response = self._instance_dict[instance_id][\"response\"]\n        ground_truth = self._instance_dict[instance_id][\"ground_truth\"]\n        \n        # Custom evaluation logic\n        if custom_evaluation_function(response, ground_truth):\n            return 1.0\n        else:\n            return 0.0\n\n**Multi-step Interactions**\n\nFor complex scenarios requiring multiple feedback rounds:\n\n.. code-block:: python\n\n    async def generate_response(self, instance_id, messages, **kwargs):\n        instance = self._instance_dict[instance_id]\n        instance[\"attempts\"] += 1\n        \n        # Evaluate current response\n        reward = await self.calculate_score(instance_id)\n        \n        if reward > 0.8:\n            return True, \"Excellent work!\", reward, {}\n        elif instance[\"attempts\"] < 3:\n            return False, \"Good attempt, but try to improve...\", reward, {}\n        else:\n            return True, \"Maximum attempts reached.\", reward, {}\n\nTroubleshooting\n---------------\n\n**Common Issues**\n\n1. **Instance ID Conflicts**: Ensure unique instance IDs across concurrent sessions\n2. **Memory Leaks**: Always call ``finalize_interaction()`` to clean up resources\n3. **Blocking Operations**: Keep interaction logic async and non-blocking\n4. **Configuration Errors**: Verify interaction config path and class name are correct\n5. **Interaction Name Conflicts**: Ensure all interactions have unique names in the configuration\n6. **Missing Interaction**: Verify the ``name`` field in ``interaction_kwargs`` matches available interactions\n7. **Backward Compatibility**: When migrating from single to multi-interaction, add ``name`` fields to existing data\n\n**Debugging**\n\nEnable debug logging to trace interaction flow:\n\n.. code-block:: bash\n\n    export VERL_LOGGING_LEVEL=DEBUG\n\n**Performance Monitoring**\n\nMonitor interaction performance impact on training throughput and adjust accordingly.\n\nRelated Documentation\n--------------------\n\n- :doc:`multiturn`: Basic multi-turn rollout configuration\n- :doc:`sandbox_fusion`: Tool integration with SGLang\n- :doc:`search_tool_example`: Search tool implementation example"
  },
  {
    "path": "docs/sglang_multiturn/multiturn.rst",
    "content": "Multi-turn Rollout Support\n==========================\n\nLast updated: 06/27/2025.\n\nBasic Configuration\n~~~~~~~~~~~~~~~~~~~\n\nTo enable multi-turn rollout, make sure to configure the following fields in your rollout configuration:\n\n.. code-block:: yaml\n\n    actor_rollout_ref: \n        rollout: \n            multi_turn: True\n            name: \"sglang\"\n\nThese configuration activates the sglang engine for multi-turn interaction during rollout.\n\nCustom Tool Configuration\n~~~~~~~~~~~~~~~~~~~~~~~~~\n\nFor custom environment interaction tools, you can implement your own tools based on ``verl.tools.base_tool.BaseTool``. Then, specify your tool configurations in a YAML file:\n\n.. code-block:: yaml\n\n    tools:\n      - class_name: \"\"\n        config: \n            type: native\n        tool_schema:\n\nYou may refer to GSM8KTool_example_configuration_, which is one example of the tool configurations. Its implementation can be found in gsm8k_tool.py_.\n\nFinally, set the ``tools_config_file`` in your rollout config:\n\n.. code-block:: yaml\n\n    actor_rollout_ref:\n        rollout:\n            tool_kwargs:\n                tools_config_file: <path_to_tool_yaml_file>\n\nThis allows integration of customized tool behaviors during actor rollout steps.\n\nIf you want rollout with simulated interaction, you can set the ``interaction_config_file`` in your rollout config:\n\n.. code-block:: yaml\n\n    interaction:\n      - class_name: \"\"\n        config: {}\n\n.. code-block:: yaml\n\n    actor_rollout_ref:\n        rollout:\n            interaction_config_file: <path_to_interaction_yaml_file>\n\nIf your tool creates multi-modal inputs, you should return a list of multi-modal inputs in your tool.execute() implementation.\n\nImage and video should be processed before returning. For example, if you are using Qwen2.5-VL, you can use the following code to get the representations:\n\n.. code-block:: python\n\n    async def create(self, ...) -> tuple[str, ToolResponse]:\n        ...\n        from verl.utils.dataset.vision_utils import process_image, process_video\n\n        img1 = process_image(img1)\n        video1 = process_video(video1)\n\n        # due to the (image | video) key is (\"image\" | \"video\") instead of (\"images\" | \"videos\") in vllm, we need to use (\"image\" | \"video\") to specify list of images/videos\n        # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205\n        return instance_id, ToolResponse(image=[img1, ...], video=[video1, ...], text=\"...\")\n\n    async def execute(self, ...) -> Tuple[str | Dict[str, Any], float, dict]:\n        ...\n        from verl.utils.dataset.vision_utils import process_image, process_video\n\n        img1 = process_image(img1)\n        video1 = process_video(video1)\n\n        # due to the (image | video) key is (\"image\" | \"video\") instead of (\"images\" | \"videos\") in vllm, we need to use (\"image\" | \"video\") to specify list of images/videos\n        # link: https://github.com/vllm-project/vllm/blob/3c545c0c3b98ee642373a308197d750d0e449403/vllm/multimodal/parse.py#L205\n        return ToolResponse(image=[img1, ...], video=[video1, ...], text=\"...\"), 0, {}\n\nremeber to set ``return_multi_modal_inputs: False`` in your dataset config in order to process the multi-modal inputs in the rollout correctly.\nRefer to the `Handling Multi-Modal Inputs in Datasets`_ section for more details.\n\nMCP Tool Configuration\n~~~~~~~~~~~~~~~~~~~~~~\n\nFor MCP interaction tools, you can flexibly configure them using a YAML file. The typical setup is as follows:\n\n.. code-block:: yaml\n\n    tools:\n      - class_name: \"\"\n        config:\n            type: mcp\n        mcp:\n            mcp_servers_config_path: ./mcp_server.json\n            tool_selected_list: {}\n\nThe ``tool_selected_list`` field is optional and specifies which tools to use from the servers. If you want to enable all available tools, simply omit this attribute. Besides, ``mcp_servers_config_path`` points to a JSON file containing the MCP server configurations. For example:\n\n.. code-block:: json\n\n      {\n          \"mcpServers\": {\n              \"SSE Server\": {\n                  \"url\": \"your_server_url\",\n                  \"auth_token\": \"your_server_api_token\"\n              },\n              \"STDIO Server\": {\n                  \"command\": \"npx\",\n                  \"args\": [\"-y\", \"server-mcp@0.2.1\"],\n                  \"env\": {\n                    \"SERVER_API_KEY\": \"your_server_api_token\"\n                  }\n              }\n          }\n      }\n\nSince the content formats returned by the MCP server may vary, users can inherit from ``MCPBaseTool`` and override the ``_parse_tool_result`` method to implement custom parsing logic.\n\n.. code-block:: python\n\n   class MCPYourTool(MCPBaseTool):\n       def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n           super().__init__(config, tool_schema)\n\n       def _parse_tool_result(self, content: list) -> Tuple[str, dict]:\n           ...\n\nOverall, you may refer to mcp_search_tool.py_ and mcp_tool_config.yaml_ for custom implementation and configuration.\n\nMulti-turn Tokenization\n~~~~~~~~~~~~~~~~~~~~~~~\n\nTokenizing multi-turn rollouts poses a challenge: after applying the chat template and tokenizing the full message list, it's hard to identify which tokens belong to assistant messages. Since the token list is flat, it lacks direct alignment with the message roles.\n\nTo address this, we adopt a **delta-based tokenization** strategy. Each time the LLM generates a new message, we:\n\n1. Apply the chat template to all prior messages (`messages[:i]`).\n2. Apply the chat template again including the latest message (`messages[:i+1]`).\n3. Tokenize only the *delta* between these two serialized message strings.\n\nThis ensures that only tokens generated by the assistant are included in the loss mask.\n\n.. code-block:: python\n\n   # When using tokenizer\n   # Exclude the assistant prompt (e.g., \"<|im_start|>assistant\") from the loss by setting add_generation_prompt=True\n   prev = tokenizer.apply_chat_template(messages[:i], add_generation_prompt=True, tokenize=False)\n   curr = tokenizer.apply_chat_template(messages[:i+1], add_generation_prompt=False, tokenize=False)\n   token_ids += tokenizer.encode(curr[len(prev):], add_special_tokens=False)\n   loss_mask += [1] * len(token_ids)  # Mask only the new assistant tokens\n\n.. code-block:: python\n\n   # When using processor\n   # Exclude the assistant prompt (e.g., \"<|im_start|>assistant\") from the loss by setting add_generation_prompt=True\n   prev = processor.apply_chat_template(messages[:i], add_generation_prompt=True, tokenize=False)\n   prev_model_inputs = processor(text=prev, images=images, videos=videos, return_tensors=\"pt\")[0].tolist()\n   curr = processor.apply_chat_template(messages[:i+1], add_generation_prompt=False, tokenize=False)\n   curr_model_inputs = processor(text=curr, images=images, videos=videos, return_tensors=\"pt\")[0].tolist()\n   token_ids += curr_model_inputs[\"input_ids\"][len(prev_model_inputs[\"input_ids\"]):]\n   loss_mask += [1] * len(token_ids)  # Mask only the new assistant tokens\n\nWhile we've validated this produces consistent results with full message tokenization, future models' chat template could break compatibility. To guard against silent inconsistencies, we compare the delta-based tokenization with full-tokenization results by default at the end of each rollout.\n\nIf you see the following warning, you can check the mismatched substring in the log:\n\n.. code-block::\n\n    Inconsistent training and inference tokenization detected. This may lead to unexpected behavior during training. Please review your chat template to determine if this is intentional. For more information, refer to the multiturn README.md.\n\nThe tokenization sanity check mode can be configured using the ``actor_rollout_ref.rollout.multi_turn.tokenization_sanity_check_mode`` parameter, which accepts the following values:\n\n- ``strict`` (default): Performs strict comparison between delta-based and full tokenization results, raising warnings for any differences.\n\n- ``ignore_strippable``: Ignores differences in whitespace characters (``\\n``, ``\\t``, ``\\r``, spaces) while still checking for meaningful text mismatches. This is useful when debugging chat template issues where whitespace variations are expected and acceptable.\n\n- ``disable``: Completely disables the tokenization sanity check. Only use this if you have thoroughly validated that tokenization discrepancies are expected and won't impact training.\n\nExample configuration:\n\n.. code-block:: yaml\n\n    actor_rollout_ref:\n        rollout:\n            multi_turn:\n                tokenization_sanity_check_mode: \"ignore_strippable\"  # Choose from: \"disable\", \"ignore_strippable\", \"strict\"\n\nHandling Multi-Modal Inputs in Datasets\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nIf your dataset includes multi-modal inputs (such as images or videos), you can control whether these are pre-processed and included in each sample by setting the return_multi_modal_inputs flag in your dataset config (used by RLHFDataset).\n\n- ``return_multi_modal_inputs: True`` (default): The dataset will pre-process and include a multi_modal_inputs dictionary for each sample. This dict contains the model-ready representations (e.g., image tensors, video tensors, etc.) as produced by your processor. This is useful for single-turn or SFT-style training, where the model expects all modalities to be present in the batch.\n\n- ``return_multi_modal_inputs: False``: The dataset will not include the multi_modal_inputs field. This is recommended for multi-turn RL or tool-augmented rollouts, where the model may generate new multi-modal inputs dynamically during rollout, and you want to avoid conflicts or redundant data in the batch.\n\n\nSpecial Cases\n^^^^^^^^^^^^^\n\nSome models (e.g., Qwen/QwQ-32B and Qwen3 series) remove internal reasoning content during chat template rendering. As a result, the message content can vary across turns, making the delta-based tokenization inaccurate.\n\nFor example, for the following conversation:\n\n.. code-block:: python\n\n    messages = [\n        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n        {\"role\": \"user\", \"content\": \"What is 2 + 2?\"},\n        {\"role\": \"assistant\", \"content\": \"<think>user asked about a simple math question.</think> 2 + 2 = 4.\"},\n        {\"role\": \"user\", \"content\": \"Explain why.\"},\n        {\"role\": \"assistant\", \"content\": \"<think>user wants to know the reasoning behind the answer. Search for a good explanation</think>\",\n         \"tool_calls\": [{\"id\": \"tool1\", \"type\": \"search\", \"arguments\": {\"query\": \"Why is 2 + 2 = 4?\"}}]},\n        {\"role\": \"tool\", \"content\": \"The sum of two and two is four because it is a basic arithmetic operation.\"},\n        {\"role\": \"assistant\", \"content\": \"<think>The tool provided a good explanation.</think>The sum of two and two is four because it is a basic arithmetic operation.\"}\n    ]\n\n1. Qwen/QwQ-32B will remove all reasoning content except the last assistant message after applying the chat template.\n\n.. code-block:: text\n\n    <|im_start|>system\n    You are a helpful assistant.<|im_end|>\n    <|im_start|>user\n    What is 2 + 2?<|im_end|>\n    <|im_start|>assistant\n     2 + 2 = 4.<|im_end|>\n    <|im_start|>user\n    Explain why.<|im_end|>\n    <|im_start|>assistant\n    <tool_call>\n    {\"name\": \"\", \"arguments\": {\"query\": \"Why is 2 + 2 = 4?\"}}\n    </tool_call><|im_end|>\n    <|im_start|>user\n    <tool_response>\n    The sum of two and two is four because it is a basic arithmetic operation.\n    </tool_response><|im_end|>\n    <|im_start|>assistant\n    <think>The tool provided a good explanation.</think> The sum of two and two is four because it is a basic arithmetic operation.<|im_end|>\n\n2. Qwen3 series will remove all reasoning content before the last user message.\n\n.. code-block:: text\n\n    <|im_start|>system\n    You are a helpful assistant.<|im_end|>\n    <|im_start|>user\n    What is 2 + 2?<|im_end|>\n    <|im_start|>assistant\n     2 + 2 = 4.<|im_end|>\n    <|im_start|>user\n    Explain why.<|im_end|>\n    <|im_start|>assistant\n    <think>\n    user wants to know the reasoning behind the answer. Search for a good explanation\n    </think>\n\n    <tool_call>\n    {\"name\": \"\", \"arguments\": {\"query\": \"Why is 2 + 2 = 4?\"}}\n    </tool_call><|im_end|>\n    <|im_start|>user\n    <tool_response>\n    The sum of two and two is four because it is a basic arithmetic operation.\n    </tool_response><|im_end|>\n    <|im_start|>assistant\n    <think>\n    The tool provided a good explanation.\n    </think>\n\n    The sum of two and two is four because it is a basic arithmetic operation.<|im_end|>\n\nTo handle this, we fall back to a **fixed base conversation** containing only a single system and user message. Since this base doesn't include assistant messages or reasoning content, it remains consistent across turns.\n\n.. code-block:: python\n\n    BASE_CHAT_HISTORY = [\n        {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n        {\"role\": \"user\", \"content\": \"I am a user.\"}\n    ]\n    prev = tokenizer.apply_chat_template(BASE_CHAT_HISTORY, add_generation_prompt=True, tokenize=False)\n    curr = tokenizer.apply_chat_template([*BASE_CHAT_HISTORY, messages[i]], add_generation_prompt=False, tokenize=False)\n    token_ids += tokenizer.encode(curr[len(prev):], add_special_tokens=False)\n    loss_mask += [1] * len(token_ids)\n\nThis method works well for Qwen3 series. However, Qwen/QwQ-32B currently has a bug in its chat template. A fix_ has been proposed but not yet adopted. Until then, use the following command to download the fixed model revision:\n\n.. code-block:: bash\n\n    pip install huggingface_hub\n    hf download Qwen/QwQ-32B --revision refs/pr/81\n\n.. _fix: https://huggingface.co/Qwen/QwQ-32B/discussions/81\n\nDiscrepancy Between Training and Inference Templates\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nAlthough the above approach fixes the delta mismatch issue, the removal of reasoning content in the inference-time chat template introduces a new discrepancy: training uses the full reasoning content, while inference does not.\n\nThis mismatch can affect model performance in unpredictable ways. To avoid it, we default to using the full response (including reasoning) for both training and rollout.\n\nHowever, this approach comes with trade-offs:\n\n1. Long reasoning contents can easily exceed the model's context window, especially in multi-turn rollout.\n2. There's a mismatch between rollout and production environment now—models will not have reasoning content from past turns if you use the default chat template in production.\n\nWe are still evaluating the impact of these issues. If you experience context length problems or prefer rollouts that match production (i.e., exclude reasoning), you can enable:\n\n``actor_rollout_ref.rollout.multi_turn.use_inference_chat_template = True``\n\nGSM8K Multi-turn Training Performance  \n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nSee the training performance of multi-turn rollout on the GSM8K task HERE_.\n\n.. _HERE: https://wandb.ai/zhaochenyang20/gsm8k_async_rl/runs/1ro1r7om?nw=nwuserzhaochenyang20\n\n.. _GSM8KTool_example_configuration: https://github.com/volcengine/verl/blob/main/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\n\n.. _gsm8k_tool.py: https://github.com/volcengine/verl/blob/main/verl/tools/gsm8k_tool.py\n\n.. _mcp_search_tool.py: https://github.com/volcengine/verl/blob/main/verl/tools/mcp_search_tool.py\n\n.. _mcp_tool_config.yaml: https://github.com/volcengine/verl/blob/main/examples/sglang_multiturn/config/tool_config/mcp_tool_config.yaml\n\nInteraction System\n~~~~~~~~~~~~~~~~~~\n\nFor dynamic conversational feedback during RL training, see:\n\n.. toctree::\n   :maxdepth: 1\n\n   interaction_system\n\nSearch Tool Integration\n~~~~~~~~~~~~~~~~~~~~~~~\n\n.. toctree::\n   :maxdepth: 1\n\n   search_tool_example\n\nCode Walkthrough\n~~~~~~~~~~~~~~~~~~~~~~~\nIf you want to learn more in depth about the code execution flow, please read https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/tree/main/rlhf/verl/multi-turn/code-walk-through\n"
  },
  {
    "path": "docs/sglang_multiturn/sandbox_fusion.rst",
    "content": "===============================\nSandbox Fusion Tool Integration\n===============================\n\nLast updated: 06/10/2025.\n\nMotivations\n===========\n\n- As users of verl, we want to allow the model to call certain tools during Actor rollout, incorporating the results into the training process.\n- A colleague from ByteDance proposed a paper aimed at enhancing model capability through code execution tools.\n- We aim to support tool-calling capabilities of inference engines using `sandbox-fusion` as the code execution system, providing the community with a reimplementation of `retools`.\n\nReward Compute with Sandbox Fusion + FaaS Integration\n=====================================================\n\n- In current datasets and tasks, similar work already exists (e.g., Prime), which uses local processes as runners to execute model-generated code for reward computation.\n- On this basis, #1429 has advanced the design by integrating FaaS as the runner for reward computation.\n\nGoals\n=====\n\n- Adapt to the `sglang` tool-calling protocol and define tools for sandbox fusion.\n- Integrate with the `async-rollout` process, ensuring sandbox fusion tools follow asyncIO conventions.\n- Design and implement a basic rate limiter to prevent issues such as 429 errors.\n\nNon-Goals\n=========\n\n- Training effectiveness is out of scope.\n- Observability metrics are not considered.\n- Distributed failover and component fault tolerance are not addressed.\n\nDesign Details\n==============\n\nTool Schema Definition\n----------------------\n\n- Currently, only code execution is considered, requiring a `code` field in the JSON from the model.\n- Only Python code is supported for now, so no `language` parameter is defined.\n\n.. code-block:: python\n\n   OpenAIFunctionToolSchema(\n       type=\"function\",\n       function=OpenAIFunctionSchema(\n           name=\"code_interpreter\",\n           description=\"A tool for executing code.\",\n           parameters=OpenAIFunctionParametersSchema(\n               type=\"object\",\n               properties={\n                   \"code\": OpenAIFunctionPropertySchema(\n                       type=\"string\",\n                       description=\"The code to execute.\",\n                       enum=None,\n                   )\n               },\n               required=[\"code\"],\n           ),\n           strict=False,\n       )\n   )\n\nConfiguration Parameters\n--------------------------\n\n+----------------------------+--------------------------------------------------------------+\n| Parameter Name             | Description                                                  |\n+============================+==============================================================+\n| `num_workers`              | Number of worker threads/processes per DP to request runner. |\n+----------------------------+--------------------------------------------------------------+\n| `rate_limit`               | Global limit of concurrent code executions. Default: 10      |\n+----------------------------+--------------------------------------------------------------+\n| `default_timeout`          | Timeout (in seconds) for each code execution. Default: 30    |\n+----------------------------+--------------------------------------------------------------+\n| `default_language`         | Default programming language. Default: \"python\"              |\n+----------------------------+--------------------------------------------------------------+\n| `enable_global_rate_limit` | Whether to enable global rate limiting. Default: True        |\n+----------------------------+--------------------------------------------------------------+\n| `sandbox_fusion_url`       | URL for the veFaas sandbox execution service                 |\n+----------------------------+--------------------------------------------------------------+\n\nRate Limiting Design\n-----------------------\n\nObjective:\n\n- Limit the number of inflight requests using a token bucket model.\n\n- Ensure ordered submission to code runners to avoid starvation due to backoff.\n\nDesign Highlights:\n\n- Use Ray Global Actor as a singleton distributed counter at cluster level.\n  \n- Semaphore used for counting, with `acquire` and `release` in separate thread pools to preserve order.\n  \n- Use Ray’s cloud-pickle to serialize functions for decoupled `ExecutionWorker`.\n\n.. code-block:: python\n\n   @ray.remote(concurrency_groups={\"acquire\": 1,\"release\": 10})\n   class TokenBucketWorker:\n       def __init__(self, rate_limit: int):\n           self.rate_limit = rate_limit\n           self.current_count = 0\n           self._semaphore = threading.Semaphore(rate_limit)\n\n       @ray.method(concurrency_group=\"acquire\")\n       def acquire(self):\n           self._semaphore.acquire()\n           self.current_count += 1\n\n       @ray.method(concurrency_group=\"release\")\n       def release(self):\n           self._semaphore.release()\n           self.current_count -= 1\n\n       def get_current_count(self):\n           return self.current_count\n\n   class ExecutionWorker:\n       def __init__(self, enable_global_rate_limit=True, rate_limit=10):\n           self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None\n\n       def _init_rate_limit(self, rate_limit):\n           return TokenBucketWorker.options(name=\"rate-limiter\", get_if_exists=True).remote(rate_limit)\n\n       def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:\n           with ExitStack() as stack:\n               stack.callback(self.rate_limit_worker.release.remote)\n               ray.get(self.rate_limit_worker.acquire.remote())\n               try:\n                   return fn(*fn_args, **fn_kwargs)\n               except Exception as e:\n                   logger.warning(f\"Error when executing code: {e}\")\n\n   def init_execution_pool(num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode=PoolMode.ThreadMode):\n       if mode == PoolMode.ThreadMode:\n           return ray.remote(ExecutionWorker).options(max_concurrency=num_workers).remote(\n               enable_global_rate_limit=enable_global_rate_limit,\n               rate_limit=rate_limit\n           )\n       else:\n           raise NotImplementedError(\"Process mode is not implemented yet\")\n\nTool Implementation\n-------------------\n\n- Use `instance_id` to identify requests across multiple dialogue rounds.\n  \n- Use `execution_pool` to implement async invocation.\n  \n- Cleanup state after rollout completion.\n\n.. code-block:: python\n\n   class SandboxFusionTool(BaseTool):\n       def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n           ...\n           self.execution_pool = init_execution_pool(...)\n           ...\n\n       async def create(self, instance_id: Optional[str] = None, ...):\n           ...\n\n        async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> Tuple[str, float, dict]:\n            code = parameters.get(\"code\", \"\")\n            timeout = parameters.get(\"timeout\", self.default_timeout)\n            language = parameters.get(\"language\", self.default_language)\n            if not isinstance(code, str):\n                code = str(code)\n\n            result = await self.execution_pool.execute.remote(self.execute_code,instance_id,code,timeout,language)\n            self._instance_dict[instance_id][\"reward\"].append(result.strip())\n\n            return result, result, {}\n\n        def execute_code(self,instance_id,code,timeout=30,language=\"python\"):\n            result_status, metadata  = _process_single_case(0, None, None,self.sandbox_fusion_url, code, timeout, language)\n            # we should always expect this since we don't have correct answer\n            if metadata[\"run_status\"] == \"Finished\":\n                actual_output = metadata[\"stdout\"] if metadata[\"stdout\"] is not None else \"\"\n                return actual_output\n            else:\n                return \"no stdout here\"\n\n       async def calc_reward(self, instance_id: str, ...):\n           ...\n\n       async def release(self, instance_id: str, ...):\n           ...\n\nTest Plan\n=========\n\nUnit Tests\n----------\n\n- **test_tools_registration**: Test tool registration and initialization.\n- **test_rollout_req_creation**: Validate that `AsyncRolloutReq` is built correctly.\n- **test_over_size_case**: Ensure rollout terminates early when exceeding `max_seq_len`.\n- **test_tool_call_basic_case**: Mock `sglang` output, validate tool call and result.\n- **test_tool_call_batch_case**: Test batch processing of tool calls.\n- **test_basic_multi_process_init**: Validate Ray global actor behaves as singleton.\n- **TestSingleNodeRateLimiterCase**: Verify rate limiter works in single-node mode.\n- **test_rotten_execution**: Ensure rate limiter recovers from function errors.\n- **TestMultiNodeRateLimiterCase**: Verify behavior in multi-node environments.\n\ne2e Tests\n----------\nwe provide e2e test scripts in `tests/special_e2e` folder, named `tests/special_e2e/run_gsm8k_fsdp_sgl_multiturn_sf_tool.sh`\n\nby setting 'trainer.rollout_data_dir' you can dump the rollout data to local disk. here is an sample taken from the rollout data:\n\n.. code-block:: python\n\n   {\n     \"input\": \"\n     \n     system\\nYou are a math expert. You are given a question and you need to solve it step by step. Reasoning step by step before any tool call. You should use the `calc_gsm8k_reward` tool after step by step solving the question, before generate final answer at least once and refine your answer if necessary. Put your final answer in the format of `#### <answer>`.\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\\n{\\\"type\\\": \\\"function\\\", \\\"function\\\": {\\\"name\\\": \\\"code_interpreter\\\", \\\"description\\\": \\\"A tool for executing code.\\\", \\\"parameters\\\": {\\\"type\\\": \\\"object\\\", \\\"properties\\\": {\\\"code\\\": {\\\"type\\\": \\\"string\\\", \\\"description\\\": \\\"The code to execute.\\\", \\\"enum\\\": null}}, \\\"required\\\": [\\\"code\\\"]}, \\\"strict\\\": false}}\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call>\\n\n     \n     user\\nJohn gets a bonus that's the same percentage every year.  Last year he made $100,000 and got a $10,000 bonus.  This year he makes $200,000.  How much will John make this year when adding both his total pay and bonus together? Let's think step by step and output the final answer after `####`.\\n\n     \n     assistant\\n\",\n     \"output\": \"assistant\\n<think>\\nOkay, let's see. John gets a bonus that's the same percentage every year. Last year, he made $100,000 and got a $10,000 bonus. This year, he makes $200,000. The question is asking how much he'll make this year when adding both his total pay and bonus together.\\n\\nFirst, I need to figure out the percentage bonus he gets. Since the bonus is the same percentage each year, if I can find that percentage based on last year's details, I can apply it to this year's earnings to find this year's bonus. Then, I'll add that bonus to his total pay this year.\\n\\nLast year, he earned $100,000 and got a $10,000 bonus. To find the percentage, I can divide the bonus by his total earnings. So, 10,000 divided by 100,000. Let me calculate that.\\n\\n10,000 / 100,000 equals 0.1, which is 10% as a percentage. So, his bonus is 10% of his pay each year.\\n\\nThis year, he makes $200,000. So, the bonus this year would be 10% of 200,000. Let me compute that using the code interpreter to ensure accuracy.\\n\n     \n     <tool_call>\\n{\\\"name\\\": \\\"code_interpreter\\\", \\\"arguments\\\": {\\\"code\\\": \\\"total_pay_this_year = 200000\\\\nbonus_percentage = 10 / 100\\\\nbonus_this_year = total_pay_this_year * bonus_percentage\\\\ntotal_income_this_year = total_pay_this_year + bonus_this_year\\\\nprint(total_income_this_year)\\\", \\\"executes\\\": \\\"True\\\"}}\\n</tool_call>\\n\n     \n     tool\\n220000.0\\n\\n\n     \n     assistant\\n<think>\\nSo the code shows that John's total income this year would be $220,000. This is consistent with the 10% bonus on last year's pay of $100,000, which was $10,000. Therefore, this year's bonus of $20,000 added to his total pay of $200,000 results in a total income of $220,000.\\n\\nI need to verify that there are no other factors at play. The problem states it's a constant percentage bonus each year, so there's no mention of changing percentages or other variables. Therefore, applying a straightforward percentage of 10% on this year's partial salary amount would result in $20,000. Adding that to his total pay gives $220,000. The calculation is correct and aligns with the question's parameters.\\n</think>\\n\\n#### 220000.0\",\n     \"score\": 0,\n     \"step\": 1\n   }\n\nhere is the readable format version:\n\n.. code-block:: python\n\n   [system]\n   \n   You are a math expert. You are given a question and you need to solve it step by step. Reasoning step by step before any tool call. You should use the `calc_gsm8k_reward` tool after step by step solving the question, before generate final answer at least once and refine your answer if necessary. Put your final answer in the format of `#### <answer>`.\n   \n   # Tools\n   \n   You may call one or more functions to assist with the user query.\n   \n   You are provided with function signatures within <tools></tools> XML tags:\n   <tools>\n   {\"type\": \"function\", \"function\": {\"name\": \"code_interpreter\", \"description\": \"A tool for executing code.\", \"parameters\": {\"type\": \"object\", \"properties\": {\"code\": {\"type\": \"string\", \"description\": \"The code to execute.\", \"enum\": null}}, \"required\": [\"code\"]}, \"strict\": false}}\n   </tools>\n   \n   For each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\n   <tool_call>\n   {\"name\": <function-name>, \"arguments\": <args-json-object>}\n   </tool_call>\n   \n   [user]\n   \n   John gets a bonus that's the same percentage every year.  Last year he made $100,000 and got a $10,000 bonus.  This year he makes $200,000.  How much will John make this year when adding both his total pay and bonus together? Let's think step by step and output the final answer after `####`.\n   \n   [assistant]\n   \n   <think>\n   Okay, let's see. John gets a bonus that's the same percentage every year. Last year, he made $100,000 and got a $10,000 bonus. This year, he makes $200,000. The question is asking how much he'll make this year when adding both his total pay and bonus together.\n   \n   First, I need to figure out the percentage bonus he gets. Since the bonus is the same percentage each year, if I can find that percentage based on last year's details, I can apply it to this year's earnings to find this year's bonus. Then, I'll add that bonus to his total pay this year.\n   \n   Last year, he earned $100,000 and got a $10,000 bonus. To find the percentage, I can divide the bonus by his total earnings. So, 10,000 divided by 100,000. Let me calculate that.\n   \n   10,000 / 100,000 equals 0.1, which is 10% as a percentage. So, his bonus is 10% of his pay each year.\n   \n   This year, he makes $200,000. So, the bonus this year would be 10% of 200,000. Let me compute that using the code interpreter to ensure accuracy.\n   <tool_call>\n   {\"name\": \"code_interpreter\", \"arguments\": {\"code\": \"total_pay_this_year = 200000\\nbonus_percentage = 10 / 100\\nbonus_this_year = total_pay_this_year * bonus_percentage\\ntotal_income_this_year = total_pay_this_year + bonus_this_year\\nprint(total_income_this_year)\", \"executes\": \"True\"}}\n   </tool_call>\n   \n   [tool]\n   \n   220000.0\n   \n   [assistant]\n   \n   <think>\n   So the code shows that John's total income this year would be $220,000. This is consistent with the 10% bonus on last year's pay of $100,000, which was $10,000. Therefore, this year's bonus of $20,000 added to his total pay of $200,000 results in a total income of $220,000.\n   \n   I need to verify that there are no other factors at play. The problem states it's a constant percentage bonus each year, so there's no mention of changing percentages or other variables. Therefore, applying a straightforward percentage of 10% on this year's partial salary amount would result in $20,000. Adding that to his total pay gives $220,000. The calculation is correct and aligns with the question's parameters.\n   </think>\n   \n   #### 220000.0\n\n\nYou can also use the `RolloutViewer` TUI tool to view the dumped rollout data:\n\n\n.. code-block:: bash\n\n    python scripts/rollout_viewer.py ${trainer.rollout_data_dir}\n\n\n.. image:: https://github.com/user-attachments/assets/e34e5157-2880-4a21-afb2-73885d0dfb11\n   :alt: RolloutViewer screenshot"
  },
  {
    "path": "docs/sglang_multiturn/search_tool_example.rst",
    "content": "=======================\r\nSearch Tool Integration\r\n=======================\r\n\r\nLast updated: 05/30/2025.\r\n\r\nIntroduction\r\n------------\r\n- We have added a search tool calling function to Multi-Turn RL, enabling the model to initiate retrieval requests during Actor rollout and directly use retrieval results for training. **We support using a local dense retriever as the retrieval tool, as well as integrating with your own local retrieval engine.**\r\n\r\n\r\n\r\nQuick Reproduction\r\n------------------\r\n\r\nCreate a New Docker Container\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\n.. code:: bash\r\n\r\n   docker run \\\r\n       -it \\\r\n       --shm-size 32g \\\r\n       --gpus all \\\r\n       -v {Huggingface-Cache-Path}:/root/.cache \\\r\n       --ipc=host \\\r\n       --network=host \\\r\n       --privileged \\\r\n       --name sglang_{your-name} \\\r\n       lmsysorg/sglang:dev \\\r\n       /bin/zsh\r\n\r\nIf you need to restart after exiting the container:\r\n\r\n.. code:: bash\r\n\r\n   docker start -i sglang_{your-name}\r\n\r\nUpdate Python and Configure the Virtual Environment using uv\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\n.. code:: bash\r\n\r\n   apt update\r\n   apt install -y python3.10 python3.10-venv\r\n\r\n   # Create a virtual environment\r\n   python3 -m venv ~/.python/verl-multiturn-rollout\r\n\r\n   # Activate the virtual environment\r\n   source ~/.python/verl-multiturn-rollout/bin/activate\r\n\r\n   # Install uv\r\n   python3 -m pip install uv\r\n\r\nInstall verl Upstream\r\n~~~~~~~~~~~~~~~~~~~~~\r\n\r\n.. code:: bash\r\n\r\n   cd ~\r\n   git clone https://github.com/volcengine/verl.git\r\n   cd verl\r\n\r\n   # Install verl\r\n   python3 -m uv pip install .\r\n   python3 -m uv pip install -r ./requirements_sglang.txt\r\n\r\n   # Manually install flash-attn\r\n   python3 -m uv pip install wheel\r\n   python3 -m uv pip install packaging\r\n   python3 -m uv pip install flash-attn --no-build-isolation --no-deps\r\n\r\nSet Up a Local Retrieval Engine\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\nIf you are using your own local retrieval service, you can skip this\r\nstep. We chose the local dense retriever provided in the search-R1\r\nexample; detailed instructions are in the `searchR1\r\ndocs <https://raw.githubusercontent.com/PeterGriffinJin/Search-R1/refs/heads/main/docs/retriever.md>`__.\r\nIn brief:\r\n\r\n-  The GPU version offers higher accuracy and speed; each GPU uses about\r\n   5–7 GB of memory.\r\n-  The CPU version can be used for simple testing but has lower\r\n   retrieval precision, which will degrade training performance. See the\r\n   `retriever\r\n   documentation <https://github.com/PeterGriffinJin/Search-R1/blob/main/docs/retriever.md>`__\r\n   in search-R1 for details.\r\n-  Recommend using Conda to install faiss-gpu=1.8.0; venv may cause errors.\r\n\r\n**Note**: To start both the training process and the local retrieval\r\nservice, we launch two separate Python environments. The training uses\r\nuv in the verl-multiturn-rollout environment, while the retriever uses\r\nconda to install ``faiss-gpu``.\r\n\r\n.. code:: bash\r\n\r\n   # Download the Miniconda installer script\r\n   wget https://repo.anaconda.com/miniconda/Miniconda3-latest-Linux-x86_64.sh -O ~/miniconda.sh\r\n\r\n   # Install to $HOME/miniconda3 in batch mode\r\n   bash ~/miniconda.sh -b -p $HOME/miniconda3\r\n\r\n   # Activate conda (only in the current shell)\r\n   eval \"$($HOME/miniconda3/bin/conda shell.bash hook)\"\r\n\r\n   # (Optional) Add conda to your default shell startup\r\n   conda init\r\n\r\n   # Reload shell config\r\n   source ~/.bashrc\r\n\r\n   # Create and activate the retriever environment with Python 3.10\r\n   conda create -n retriever python=3.10 -y\r\n   conda activate retriever\r\n\r\n   # Install PyTorch (with GPU support) and related libraries\r\n   conda install pytorch==2.4.0 torchvision==0.19.0 torchaudio==2.4.0 pytorch-cuda=12.1 -c pytorch -c nvidia -y\r\n\r\n   # Install other Python packages\r\n   pip install transformers datasets pyserini huggingface_hub\r\n\r\n   # Install the GPU version of faiss\r\n   conda install faiss-gpu=1.8.0 -c pytorch -c nvidia -y\r\n\r\n   # Install the API service framework\r\n   pip install uvicorn fastapi\r\n\r\nDownload the Indexing and Corpus\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\nThe local retrieval files are large—prepare sufficient disk space.\r\nDownloading is about 60–70 GB, and uncompressed takes about 132 GB:\r\n\r\n.. code:: bash\r\n\r\n   conda activate retriever\r\n\r\n   save_path=/the/path/to/save\r\n   python examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py --save_path $save_path\r\n   cat $save_path/part_* > $save_path/e5_Flat.index\r\n   gzip -d $save_path/wiki-18.jsonl.gz\r\n\r\nStart the Local flat e5 Retrieval Server\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\n1. The first startup will download models and load the index.\r\n2. Apart from the download, startup takes about 1–2 minutes.\r\n3. After startup, each GPU uses about 5–7 GB of memory, leaving the rest\r\n   for multi-turn RL training.\r\n\r\n.. code:: bash\r\n\r\n   conda activate retriever\r\n\r\n   index_file=$save_path/e5_Flat.index\r\n   corpus_file=$save_path/wiki-18.jsonl\r\n   retriever_name=e5\r\n   retriever_path=intfloat/e5-base-v2\r\n\r\n   python examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py \\\r\n     --index_path $index_file \\\r\n     --corpus_path $corpus_file \\\r\n     --topk 3 \\\r\n     --retriever_name $retriever_name \\\r\n     --retriever_model $retriever_path \\\r\n     --faiss_gpu\r\n\r\nSet Up WANDB_API_KEY\r\n~~~~~~~~~~~~~~~~~~~~\r\n\r\n.. code:: bash\r\n\r\n   export WANDB_API_KEY={YOUR_WANDB_API_KEY}\r\n\r\n   # Define a timestamp function\r\n   function now() {\r\n       date '+%Y-%m-%d-%H-%M'\r\n   }\r\n\r\n**Preprocess the Dataset**\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\n   **Note:** The following data processing and training commands must be\r\n   run in the verl-multiturn-rollout environment.\r\n\r\n.. code:: bash\r\n\r\n   python3 examples/data_preprocess/preprocess_search_r1_dataset.py\r\n\r\nTesting on 8 x H20\r\n~~~~~~~~~~~~~~~~~~\r\n\r\n.. code:: bash\r\n\r\n   # Ensure the now() function is defined\r\n   # Create a logs directory\r\n   mkdir -p logs\r\n\r\n   # Set GPUs and run with a suitable log path\r\n   export CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\r\n\r\n   nohup bash examples/sglang_multiturn/search_r1_like/run_qwen2.5-3b_instruct_search_multiturn.sh \\\r\n     trainer.experiment_name=qwen2.5-3b-it_rm-searchR1-like-sgl-multiturn-$(now) \\\r\n     > logs/searchR1-like$(now).log 2>&1 &\r\n\r\nCustom Search Configuration\r\n---------------------------\r\n\r\nTo enable multi-turn reasoning, set the following fields in your config:\r\n\r\n.. code:: yaml\r\n\r\n   actor_rollout_ref:\r\n     rollout:\r\n       name: \"sglang\"\r\n       multi_turn:\r\n         enable: True\r\n\r\nYou must specify ``retrieval_service_url`` in ``examples/sglang_multiturn/config/tool_config/search_tool_config.yaml``, and properly configure concurrency. For more details on concurrency, refer to the Sandbox Fusion example:\r\n\r\n.. code:: yaml\r\n\r\n   tools:\r\n     - class_name: verl.tools.search_tool.SearchTool\r\n       config:\r\n         retrieval_service_url: http://127.0.0.1:8000/retrieve\r\n         num_workers: 120\r\n         rate_limit: 120\r\n         timeout: 30\r\n\r\nThe retriever input/output formats are as follows. If your service\r\nparameters match, only modify ``retrieval_service_url``. You can also\r\ncustomize in ``search_r1_like_utils.py``.\r\n\r\n.. code:: python\r\n\r\n   Input format:\r\n   {\r\n     \"queries\": [\"What is Python?\", \"Tell me about neural networks.\"],\r\n     \"topk\": 3,\r\n     \"return_scores\": true\r\n   }\r\n\r\n   Output format (when return_scores=True, similarity scores are returned):\r\n   {\r\n       \"result\": [\r\n           [   # Results for each query\r\n               {\r\n                   \"document\": doc, \"score\": score\r\n               },\r\n               # ... more documents\r\n           ],\r\n           # ... results for other queries\r\n       ]\r\n   }\r\n\r\nNotes\r\n-----\r\n\r\n1. The total training time is about 27 hours; meanwhile, the validation\r\n   dataset is very large (51 k), and each validation takes about 6000 s.\r\n   (Therefore, ``val_before_train=False`` by default)\r\n"
  },
  {
    "path": "docs/single_controller.rst",
    "content": "The Design of ``verl.single_controller``\n==============================================\n\nLast updated: 05/21/2025.\n\n**Author:**\\  `Wang Zhang <https://github.com/zw0610>`__\n\nPreface\n-------\n\nWe prepared this document for developers of ``verl``, particularly those\ninterested in understanding or contributing to the\n``verl.single_controller`` module. It is not intended for end users, but\nfor contributors seeking to understand the architectural rationale and\ninternal mechanics.\n\n--------------\n\nOrigin\n------\n\nThe ``single_controller`` module originated from a request I received —\nto adapt a toy single-process RLHF script into a distributed system with\nminimal changes, while maintaining ease of debugging.\n\nCommon practice — such as using PyTorch’s Distributed Data Parallel\n(DDP) — typically involves wrapping ``nn.Module`` and launching multiple\nprocesses that execute the same function under different ranks. However,\nthis approach presents two main limitations in the context of\ndistributed RLHF: - Difficulty representing multiple DAGs as required by\nPPO; - Difficulty inspecting intermediate tensors during training.\n\nTo maintain debuggability, we opted for a different approach — breaking\nthe training loop into well-defined stages like ``generate_sequences``,\n``compute_advantages``, and so on.\n\nWe selected `Ray <https://www.ray.io/>`__ as the initial backend for\n``verl`` due to its ability to expose Python class methods as RPC\nendpoints. However, Ray’s default model only supports **one method call,\none RPC**, while training LLMs typically requires coordination across\nmultiple processes.\n\nTo hide this multi-Ray actors invocation for a single method from users,\nwe introduced the following components:\n\n-  ``WorkerGroup`` – manages a group of remote workers and provides\n   a unified interface for multi-process distributed computation;\n-  ``ResourcePool`` – binds computational resources to worker\n   processes;\n-  ``ClassWithArgs`` – enables delayed remote instantiation with\n   specified initialization arguments.\n\n--------------\n\nA Running Example: ``generate_sequences``\n-----------------------------------------\n\nTo illustrate the design, we walk through how the ``generate_sequences``\nmethod in the ``ActorRolloutRefWorker`` class is registered and invoked\nacross distributed workers.\n\n--------------\n\nStep 1: Register with a Decorator\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThe first step is to define the ``generate_sequences`` and decorate it\nwith ``@register`` as it will be called in driver script.\n\n**Source:**\n`fsdp_workers.py <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/workers/fsdp_workers.py#L528>`__\n\n.. code:: python\n\n   class ActorRolloutRefWorker(Worker):\n       ...\n       @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n       def generate_sequences(self, prompts: DataProto):\n           prompts = prompts.to(torch.cuda.current_device())\n           ...\n\nThe ``@register`` decorator adds metadata to the ``generate_sequences``\nmethod. Currently, it doesn’t alter functionality, but attaches\nattributes via a magic key (``MAGIC_ATTR``):\n\n**Source:**\n`decorator.py <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/base/decorator.py#L411>`__\n\n.. code:: python\n\n   def register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True):\n       ...\n       def decorator(func):\n           @wraps(func)\n           def inner(*args, **kwargs):\n               if materialize_futures:\n                   args, kwargs = _materialize_futures(*args, **kwargs)\n               return func(*args, **kwargs)\n\n           attrs = {\"dispatch_mode\": dispatch_mode, \"execute_mode\": execute_mode, \"blocking\": blocking}\n           setattr(inner, MAGIC_ATTR, attrs)\n           return inner\n\n       return decorator\n\nAs the code shows, values of ``dispatch_mode``, ``execute_mode`` and\n``blocking`` is attached the ``generate_sequences`` method.\n\n--------------\n\nStep 2: Binding During Initialization\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\nThese attached attributes are extracted and utilized when\n``ActorRolloutRefWorker``, wrapped in a ``RayClassWithArgs``, is passed\ninto a ``RayWorkerGroup``.\n\n**Source:**\n`main_generation.py <https://github.com/volcengine/verl/blob/4ae9a0fdab229f75f080e9478807783ed4c97154/verl/trainer/main_generation.py#L82>`__\n\n.. code:: python\n\n   ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(ActorRolloutRefWorker), config=config, role=\"rollout\")\n   resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes)\n   wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init)\n\nDuring the\n`initialization <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/ray/base.py#L184>`__\nof ``RayWorkerGroup``, two key steps occur:\n\n1. Worker instances (Ray actors) are created:\n   `RayWorkerGroup._init_with_resource_pool <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/ray/base.py#L211>`__\n2. Methods decorated with ``@register`` are bound to ``RayWorkerGroup``:\n   `RayWorkerGroup._bind_worker_method <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/ray/base.py#L214>`__\n\n.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/worker_group_init.png?raw=true\n   :alt: initialization_and_binding_of_worker_group\n\n   initialization_and_binding_of_worker_group\n\nThe binding procedure is the heart of ``verl.single_controller``.\n\n**Key function:**\n`WorkerGroup._bind_worker_method <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/base/worker_group.py#L143>`__\n\n.. code:: python\n\n   def _bind_worker_method(self, user_defined_cls, func_generator):\n       ...\n       for method_name in dir(user_defined_cls):\n           try:\n               method = getattr(user_defined_cls, method_name)\n               assert callable(method)\n           except Exception:\n               continue  # Skip properties\n           <<<to be continue 1>>>\n\nWhen a method has the ``MAGIC_ATTR``, the attributes set by\n``@register`` are extracted:\n\n.. code:: python\n\n           <<<continue 1>>>\n           if hasattr(method, MAGIC_ATTR):\n               attribute = getattr(method, MAGIC_ATTR)\n               dispatch_mode = attribute[\"dispatch_mode\"]\n               execute_mode = attribute[\"execute_mode\"]\n               blocking = attribute[\"blocking\"]\n\n               <<<to be continue 2>>>\n\nAs show in the flow chart above, these attributes are fed into\n``func_generator``. However, ``func_generator`` takes ``method_name``,\n``dispatch_fn``, ``collect_fn``, ``execute_fn``, ``blocking``. We need\nto find the corresponding ``dispatch_fn`` and ``collect_fn`` associated\nwith the ``dispatch_mode`` (``DP_COMPUTE_PROTO``) from\n`DISPATCH_MODE_FN_REGISTRY <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/base/decorator.py#L387>`__:\n\n.. code:: python3\n\n   DISPATCH_MODE_FN_REGISTRY = {\n       Dispatch.ONE_TO_ALL: {\n           \"dispatch_fn\": dispatch_one_to_all,\n           \"collect_fn\": collect_all_to_all,\n       },\n       ...\n       Dispatch.DP_COMPUTE_PROTO: {\n           \"dispatch_fn\": dispatch_dp_compute_data_proto,\n           \"collect_fn\": collect_dp_compute_data_proto,\n       },\n       ...\n   }\n\nSimilarly, the ``execute_fn`` is selected by ``execute_mode`` and\nextracted by:\n\n.. code:: python\n\n               <<<continue 2>>>\n               # get execute_fn_name\n               execute_mode = get_predefined_execute_fn(execute_mode=execute_mode)\n               wg_execute_fn_name = execute_mode[\"execute_fn_name\"]\n\n               # get execute_fn from string\n               try:\n                   execute_fn = getattr(self, wg_execute_fn_name)\n                   assert callable(execute_fn), \"execute_fn must be callable\"\n               except Exception:\n                   print(f\"execute_fn {wg_execute_fn_name} is invalid\")\n                   raise\n               <<<to be continue 3>>>\n\nIn this ``generate_sequences`` cases: -\n``dispatch_mode = Dispatch.DP_COMPUTE_PROTO`` -\n``dispatch_fn = dispatch_dp_compute_data_proto`` -\n``collect_fn = collect_dp_compute_data_proto`` -\n``execute_fn = RayWorkerGroup.execute_all``\n\nONE_TO_ALL v.s. DP_COMPUTE_PROTO\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n``dispatch_mode`` is associated with a ``dispatch_fn`` and a\n``collect_fn``. As the name implies, ``dispatch_fn`` processes the input\narguments in ``WorkerGroup`` and generate a batch (list) of input\narguments, each of which will be fed into a worker attached to the\n``WorkerGroup``.\n\n``dispatch_fn`` of ``ONE_TO_ALL`` is\n`dispatch_one_to_all <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/base/decorator.py#L119>`__,\nwhich just duplicates all the input arguments into N replicas, where N\nequals the number of Workers attached to the ``worker_group``:\n\n.. code:: python\n\n   def dispatch_one_to_all(worker_group, *args, **kwargs):\n       args = tuple([arg] * worker_group.world_size for arg in args)\n       kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}\n       return args, kwargs\n\n``dispatch_fn`` of ``DP_COMPUTE_PROTO`` is\n`dispatch_dp_compute_data_proto <https://github.com/volcengine/verl/blob/c59ab2f4788f9a910836a9f2f53dcdb62dfa314e/verl/single_controller/base/decorator.py#L350>`__,\nwhich uses ``DataProto.chunk`` to split a large ``DataProto`` into N\nsmaller ``DataProto``, where N equals the world_size (number of the\nworkers) of the ``worker_group``:\n\n.. code:: python\n\n   def dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):\n       from verl.single_controller.base.worker_group import WorkerGroup\n\n       assert isinstance(worker_group, WorkerGroup)\n       # Note: enable auto padding for dp compute DatapProto\n       splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding(\n           worker_group.world_size,\n           *args,\n           **kwargs,\n       )\n       return splitted_args, splitted_kwargs\n\nThe ``collect_fn`` follows the same pattern and process a batch (list)\nof returned value from all workers of a ``WorkerGroup`` and merge it\ninto a list as ``collect_all_to_all`` does or a large ``DataProto`` as\n``collect_dp_compute_data_proto`` does.\n\nFinally, a new method is dynamically generated using ``func_generator``\nand added to the ``WorkerGroup`` instance:\n\n.. code:: python\n\n               <<<continue 3>>>\n               # bind a new method to the RayWorkerGroup\n               func = func_generator(\n                   self,\n                   method_name,\n                   dispatch_fn=dispatch_fn,\n                   collect_fn=collect_fn,\n                   execute_fn=execute_fn,\n                   blocking=blocking,\n               )\n\n               try:\n                   setattr(self, method_name, func)\n                   method_names.append(method_name)\n               except Exception as e:\n                   raise ValueError(f\"Fail to set method_name {method_name}\") from e\n\nThis makes the method invocable via the ``WorkerGroup`` interface.\n\n--------------\n\nStep 3: Call Chain\n~~~~~~~~~~~~~~~~~~\n\nAll the machinery above ensures that distributed calls feel identical to\nsingle-process ones. In the original single-process script, the code\nlooks like:\n\n.. code:: python\n\n   rollout = Rollout()\n   rollout.generate_sequences(batch)\n\nWith ``verl``, the multiprocess program becomes:\n\n.. code:: python\n\n   rollout = RayWorkerGroup(resource_pool=[4], RayClassWithArgs(Rollout))\n   rollout.generate_sequences(batch)\n\n.. figure:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/call_generate_sequences.png?raw=true\n   :alt: call_chain_of_generate_sequences\n\n   call_chain_of_generate_sequences\n\nBehind this simple call: - ``dispatch_fn`` splits input across workers -\n``execute_fn`` performs the actual remote invocation - ``collect_fn``\ngathers the results\n\nAll of this is abstracted away, enabling developers to write distributed\ncode with minimal changes to their existing logic.\n\n--------------\n\nBeyond RL Post-Training: Generalizing ``verl.single_controller``\n----------------------------------------------------------------\n\nThe ``verl.single_controller`` module generalizes well beyond\nreinforcement learning. It provides a clean abstraction to batch-process\nremote method calls, with automatic input/output handling.\n\nBy minimizing the gap between single-process and multi-process scripts,\n``verl.single_controller`` opens the door to distributed computing in\nbroader domains — not limited to RL post-training.\n\nWe hope this design inspires more examples and extensions from the\ncommunity.\n"
  },
  {
    "path": "docs/start/agentic_rl.rst",
    "content": "Agentic RL Training\n===================\n\nLast updated: 07/15/2025.\n\nOverview\n----------\nThe goal of Agentic RL is to improve the performance of backend models from reinforcement learning to the Agent. During the training process, a series of features are developed:\n\n1. Server-based asynchronous rollout\n2. Multi-turn conversations and tool calls\n3. LangGraph-based Agent\n\n\nThis document explains the system principles and usage involved to help users implement Agentic RL.\n\n\nServer-based Asynchronous Rollout\n---------------------------------\n\nSince Agents need to interact with the environment through various tool calls, in order to avoid GPU idling while waiting for tool call return results, an asyncio based co-routing mechanism is utilized to execute each rollout requests asynchronously, thereby improving training performance. To support asynchronous rollout, the inference engine (server) and the agent (client) are architecturally separated, implementing a server-based system with the following objectives:\n\n1. Enabling load balancing mechanisms to balance loads across multiple GPUs and reduce the impact of long-tail requests on performance. For this purpose, scheduling capabilities in stream mode (recipe\\stream_mode) are implemented as a recipe.\n2. Preventing agent specific features such as tracing from affecting the inference engine.\n\nSystem Architecture\n~~~~~~~~~~~~~~~~~~~\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/agent_loop.png?raw=true\n\nFor more detail on internal design, please refer to :doc:`Agent Loop<../advance/agent_loop>`.\n\nSystem Components\n~~~~~~~~~~~~~~~~~\n\n+--------------------------+----------------------------------------------------------------------------+\n| Component                | Role                                                                       |\n+==========================+============================================================================+\n| AgentLoop                | Client, implements Agent functions                                         |\n+--------------------------+----------------------------------------------------------------------------+\n| AsyncLLMServerManager    | Inference gateway, provides generate interface for AgentLoop               |\n+--------------------------+----------------------------------------------------------------------------+\n| AsyncServer              | Server, each instance is connected to one DP group of the inference engine |\n+--------------------------+----------------------------------------------------------------------------+\n\n**\"generate\" Interface**\n\nThe \"generate\" function based on ray actor is used between the Client and Server instead of the standard chat completion API. This is because the conversion between tokens and text can be irreversible. For example, the token converted from \"<think>\" will be different from that generated by the LLM. During the training phase, it is necessary to strictly use the tokens generated by LLM inference to avoid inaccurate in computing advantage, which may affect model performance. Having the Server provide a token-based API helps the Client maintain the relationship between the text generated by tool calls and the tokens returned by the LLM, so as to output correct tokens for training.\n\n\n**Inference Engine Adaptation**\nAsyncServer uniformly provides a generate function to the upper layer, with separate implementations for SGLang and vLLM to hide underlying differences:\n\n1. The SGLang AsyncServer uses the async_generate interface of the SGLang engine, which is located on the first GPU of each TP group. Therefore, AsyncServer needs to remotely call async_generate through ray actor.\n2. The vLLM AsyncServer uses the generate interface of the vLLM engine, which can communicate with the GPUs in the TP group through ZMQ and can be directly called in AsyncServer.\n\n\nUsage Example\n~~~~~~~~~~~~~\n\nFollow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints.\n\nThere are two options required to use agent loop:\n\n- `data.return_raw_chat=True`\n- `actor_rollout_ref.rollout.mode=async`\n\nThis example uses the sglang inference engine by default, and you can also modify rollout_name to use vllm.\n\n.. code-block:: bash\n\n    bash examples/grpo_trainer/run_qwen2-7b_seq_balance.sh\n\n\nMulti-turn Conversations and Tool Calls\n---------------------------------------\n\nFollow :doc:`Multi-turn Rollout Support<../sglang_multiturn/multiturn>` to prepare tool and configuration files.\n\nThe Tool Agent Loop has an additional requirement: adding an \"agent_name\" field to the dataset. During rollout, it will choose to use tool_agent_loop or single_turn_agent (default) based on this field.\n\nUsage Example\n~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    # install mlflow to view toolcall and llm trace\n    pip install mlflow\n\n    # This will download and preprocess the GSM8K dataset into ~/data/gsm8k/ and add the \"agent_name\" field.\n    python examples/data_preprocess/gsm8k_tool_agent_loop.py\n\n    # Start training with tool calls and enabled mlflow based trace helping to debug the rollout details\n    bash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_tool_agent_mlflow.sh\n\n    # When training is done, start a mlflow server to view trace\n    mlflow ui -h 0.0.0.0 -p 5000 --backend-store-uri sqlite:////tmp/mlruns.db\n\n    # then you can open http://<your ip address>:5000 from browser to view trace\n\n\nNote: During training, because the model may sometimes fail to generate correct toolcall tags, an error message \"Failed to decode tool call\" will be output to the console, which does not indicate an abnormality in training.\n\n\nFollow :doc:`Rollout trace<../advance/rollout_trace>` to known more about trace feature.\n\n\n\nAgent Framework\n---------------\n\nSystem Architecture\n~~~~~~~~~~~~~~~~~~~\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/langgraph_agent.png?raw=true\n\nSystem Components\n~~~~~~~~~~~~~~~~~\n\n+--------------------------+-----------------------------------------------------------------------------------------------+\n| Component                | Role                                                                                          |\n+==========================+===============================================================================================+\n| ChatModel                | LLM object of LangChain, used to adapt to the “generate” api provided by AsyncLLMServerManager|\n+--------------------------+-----------------------------------------------------------------------------------------------+\n| ReactAgentLoop           | Agent adaptation layer, which by default supports a naive LangGraph Agentic.                  |\n|                          | New classes can be derived to support user-defined Agents, and the run function needs to be   |\n|                          | implemented to complete Agent calls.                                                          |\n+--------------------------+-----------------------------------------------------------------------------------------------+\n| AsyncServer              | Server, each instance is connected to one DP group of the inference engine.                   |\n+--------------------------+-----------------------------------------------------------------------------------------------+\n\n\nFollow doc \"recipe/langgraph_agent/example/README.md\" for more details."
  },
  {
    "path": "docs/start/install.rst",
    "content": "Installation\n============\n\nRequirements\n------------\n\n- **Python**: Version >= 3.10\n- **CUDA**: Version >= 12.8\n\nverl supports various backends. Currently, the following configurations are available:\n\n- **FSDP** and **Megatron-LM** (optional) for training.\n- **SGLang**, **vLLM** and **TGI** for rollout generation.\n\nChoices of Backend Engines\n----------------------------\n\n1. Training:\n\nWe recommend using **FSDP** backend to investigate, research and prototype different models, datasets and RL algorithms. The guide for using FSDP backend can be found in :doc:`FSDP Workers<../workers/fsdp_workers>`.\n\nFor users who pursue better scalability, we recommend using **Megatron-LM** backend. Currently, we support `Megatron-LM v0.13.1 <https://github.com/NVIDIA/Megatron-LM/tree/core_v0.13.1>`_. The guide for using Megatron-LM backend can be found in :doc:`Megatron-LM Workers<../workers/megatron_workers>`.\n\n\n2. Inference:\n\nFor inference, vllm 0.8.3 and later versions have been tested for stability. We recommend turning on env var `VLLM_USE_V1=1` for optimal performance.\n\nFor SGLang, refer to the :doc:`SGLang Backend<../workers/sglang_worker>` for detailed installation and usage instructions. SGLang rollout is under extensive development and offers many advanced features and optimizations. We encourage users to report any issues or provide feedback via the `SGLang Issue Tracker <https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/issues/106>`_.\n\nFor huggingface TGI integration, it is usually used for debugging and single GPU exploration.\n\nInstall from docker image\n-------------------------\n\nStart from v0.6.0, we use vllm and sglang release image as our base image.\n\nBase Image\n::::::::::\n\n- vLLM: https://hub.docker.com/r/vllm/vllm-openai\n- SGLang: https://hub.docker.com/r/lmsysorg/sglang\n\nApplication Image\n:::::::::::::::::\n\nUpon base image, the following packages are added:\n\n- flash_attn\n- Megatron-LM\n- Apex\n- TransformerEngine\n- DeepEP\n\nLatest docker file:\n\n- `Dockerfile.stable.vllm <https://github.com/volcengine/verl/blob/main/docker/Dockerfile.stable.vllm>`_\n- `Dockerfile.stable.sglang <https://github.com/volcengine/verl/blob/main/docker/Dockerfile.stable.sglang>`_\n\nAll pre-built images are available in dockerhub: `verlai/verl <https://hub.docker.com/r/verlai/verl>`_. For example, ``verlai/verl:sgl055.latest``, ``verlai/verl:vllm011.latest``.\n\nYou can find the latest images used for development and ci in our github workflows:\n\n- `.github/workflows/vllm.yml <https://github.com/volcengine/verl/blob/main/.github/workflows/vllm.yml>`_\n- `.github/workflows/sgl.yml <https://github.com/volcengine/verl/blob/main/.github/workflows/sgl.yml>`_\n\n\nInstallation from Docker\n::::::::::::::::::::::::\n\nAfter pulling the desired Docker image and installing desired inference and training frameworks, you can run it with the following steps:\n\n1. Launch the desired Docker image and attach into it:\n\n.. code:: bash\n\n    docker create --runtime=nvidia --gpus all --net=host --shm-size=\"10g\" --cap-add=SYS_ADMIN -v .:/workspace/verl --name verl <image:tag> sleep infinity\n    docker start verl\n    docker exec -it verl bash\n\n\n2.\tIf you use the images provided, you only need to install verl itself without dependencies:\n\n.. code:: bash\n\n    # install the nightly version (recommended)\n    git clone https://github.com/volcengine/verl && cd verl\n    pip3 install --no-deps -e .\n\n[Optional] If you hope to switch between different frameworks, you can install verl with the following command:\n\n.. code:: bash\n\n    # install the nightly version (recommended)\n    git clone https://github.com/volcengine/verl && cd verl\n    pip3 install -e \".[vllm]\"\n    pip3 install -e \".[sglang]\"\n\n\nInstall from custom environment\n---------------------------------------------\n\nWe recommend to use docker images for convenience. However, if your environment is not compatible with the docker image, you can also install verl in a python environment.\n\n.. note::\n\n    - Dockerfile provides more details than this installation instructions. You can find examples in each Dockerfile, for example `verl0.6-cu128-torch2.8.0-fa2.7.4 Dockerfile.base <https://github.com/volcengine/verl/blob/v0.6.0/docker/verl0.6-cu128-torch2.8.0-fa2.7.4/Dockerfile.base>`_ .\n\n\nPre-requisites\n::::::::::::::\n\nFor training and inference engines to utilize better and faster hardware support, CUDA/cuDNN and other dependencies are required,\nand some of the dependencies are easy to be overridden when installing other packages,\nso we put them in the :ref:`Post-installation` step.\n\n.. note::\n\n    - The installation steps below are recommended configurations for the latest version of verl.\n\n    If you are trying to customize your own environment, please ignore the strict constraints.\n\nWe need to install the following pre-requisites:\n\n- **CUDA**: Version >= 12.8\n- **cuDNN**: Version >= 9.10.0\n- **Apex**\n\nCUDA above 12.8 is recommended to use as the docker image,\nplease refer to `NVIDIA's official website <https://developer.nvidia.com/cuda-toolkit-archive>`_ for other version of CUDA.\n\n.. code:: bash\n\n    # change directory to anywher you like, in verl source code directory is not recommended\n    wget https://developer.download.nvidia.com/compute/cuda/12.8.1/local_installers/cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb\n    dpkg -i cuda-repo-ubuntu2204-12-8-local_12.8.1-570.124.06-1_amd64.deb\n    cp /var/cuda-repo-ubuntu2204-12-8-local/cuda-*-keyring.gpg /usr/share/keyrings/\n    apt-get update\n    apt-get -y install cuda-toolkit-12-8\n    update-alternatives --set cuda /usr/local/cuda-12-8\n\n\ncuDNN can be installed via the following command,\nplease refer to `NVIDIA's official website <https://developer.nvidia.com/rdp/cudnn-archive>`_ for other version of cuDNN.\n\n.. code:: bash\n\n    # change directory to anywher you like, in verl source code directory is not recommended\n    wget https://developer.download.nvidia.com/compute/cudnn/9.10.2/local_installers/cudnn-local-repo-ubuntu2204-9.10.2_1.0-1_amd64.deb\n    dpkg -i cudnn-local-repo-ubuntu2204-9.10.2_1.0-1_amd64.deb\n    cp /var/cudnn-local-repo-ubuntu2204-9.10.2/cudnn-*-keyring.gpg /usr/share/keyrings/\n    apt-get update\n    apt-get -y install cudnn-cuda-12\n\nInstall dependencies\n::::::::::::::::::::\n\n.. note::\n\n    We recommend to use a fresh new conda environment to install verl and its dependencies.\n\n    **Notice that the inference frameworks often strictly limit your pytorch version and will directly override your installed pytorch if not paying enough attention.**\n\n    As a countermeasure, it is recommended to install inference frameworks first with the pytorch they needed. For vLLM, if you hope to use your existing pytorch,\n    please follow their official instructions\n    `Use an existing PyTorch installation <https://docs.vllm.ai/en/latest/getting_started/installation/gpu.html#build-wheel-from-source>`_ .\n\n\n1. First of all, to manage environment, we recommend using conda:\n\n.. code:: bash\n\n   conda create -n verl python==3.12\n   conda activate verl\n\n\n2. Then, execute the ``install.sh`` script that we provided in verl:\n\n.. code:: bash\n\n    # Make sure you have activated verl conda env\n    # If you need to run with megatron\n    bash scripts/install_vllm_sglang_mcore.sh\n    # Or if you simply need to run with FSDP\n    USE_MEGATRON=0 bash scripts/install_vllm_sglang_mcore.sh\n\n\nIf you encounter errors in this step, please check the script and manually follow the steps in the script.\n\n[Optional] NVIDIA Apex is recommended for Megatron-LM training, but it's not needed if you only use FSDP backend.\nYou can install it via the following command, but notice that this steps can take a very long time.\nIt is recommended to set the ``MAX_JOBS`` environment variable to accelerate the installation process,\nbut do not set it too large, otherwise the memory will be overloaded and your machines may hang.\n\n.. code:: bash\n\n    # change directory to anywher you like, in verl source code directory is not recommended\n    git clone https://github.com/NVIDIA/apex.git && \\\n    cd apex && \\\n    MAX_JOB=32 pip install -v --disable-pip-version-check --no-cache-dir --no-build-isolation --config-settings \"--build-option=--cpp_ext\" --config-settings \"--build-option=--cuda_ext\" ./\n\nInstall verl\n::::::::::::\n\nFor installing the latest version of verl, the best way is to clone and\ninstall it from source. Then you can modify our code to customize your\nown post-training jobs.\n\n.. code:: bash\n\n   git clone https://github.com/volcengine/verl.git\n   cd verl\n   pip install --no-deps -e .\n\n\nPost-installation\n:::::::::::::::::\n\nPlease make sure that the installed packages are not overridden during the installation of other packages.\n\nThe packages worth checking are:\n\n- **torch** and torch series\n- **vLLM**\n- **SGLang**\n- **pyarrow**\n- **tensordict**\n- **nvidia-cudnn-cu12**: For Magetron backend\n\nIf you encounter issues about package versions during running verl, please update the outdated ones.\n\n\nInstall with AMD GPUs - ROCM kernel support\n------------------------------------------------------------------\n\nWhen you run on AMD GPUs (MI300) with ROCM platform, you cannot use the previous quickstart to run verl. You should follow the following steps to build a docker and run it.\nIf you encounter any issues in using AMD GPUs running verl, feel free to contact me - `Yusheng Su <https://yushengsu-thu.github.io/>`_.\n\nFind the docker for AMD ROCm: `docker/Dockerfile.rocm <https://github.com/volcengine/verl/blob/main/docker/Dockerfile.rocm>`_\n::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::::\n\n.. code-block:: bash\n\n    #  Build the docker in the repo dir:\n    # docker build -f docker/Dockerfile.rocm -t verl-rocm:03.04.2015 .\n    # docker images # you can find your built docker\n    FROM rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n\n    # Set working directory\n    # WORKDIR $PWD/app\n\n    # Set environment variables\n    ENV PYTORCH_ROCM_ARCH=\"gfx90a;gfx942\"\n\n    # Install vllm\n    RUN pip uninstall -y vllm && \\\n        rm -rf vllm && \\\n        git clone -b v0.6.3 https://github.com/vllm-project/vllm.git && \\\n        cd vllm && \\\n        MAX_JOBS=$(nproc) python3 setup.py install && \\\n        cd .. && \\\n        rm -rf vllm\n\n    # Copy the entire project directory\n    COPY . .\n\n    # Install dependencies\n    RUN pip install \"tensordict<0.6\" --no-deps && \\\n        pip install accelerate \\\n        codetiming \\\n        datasets \\\n        dill \\\n        hydra-core \\\n        liger-kernel \\\n        numpy \\\n        pandas \\\n        datasets \\\n        peft \\\n        \"pyarrow>=15.0.0\" \\\n        pylatexenc \\\n        \"ray[data,train,tune,serve]\" \\\n        torchdata \\\n        transformers \\\n        wandb \\\n        orjson \\\n        pybind11 && \\\n        pip install -e . --no-deps\n\nBuild the image\n::::::::::::::::::::::::\n\n.. code-block:: bash\n\n    docker build -t verl-rocm .\n\nLaunch the container\n::::::::::::::::::::::::::::\n\n.. code-block:: bash\n\n    docker run --rm -it \\\n      --device /dev/dri \\\n      --device /dev/kfd \\\n      -p 8265:8265 \\\n      --group-add video \\\n      --cap-add SYS_PTRACE \\\n      --security-opt seccomp=unconfined \\\n      --privileged \\\n      -v $HOME/.ssh:/root/.ssh \\\n      -v $HOME:$HOME \\\n      --shm-size 128G \\\n      -w $PWD \\\n      verl-rocm \\\n      /bin/bash\n\nIf you do not want to root mode and require assign yourself as the user,\nPlease add ``-e HOST_UID=$(id -u)`` and ``-e HOST_GID=$(id -g)`` into the above docker launch script.\n\nverl with AMD GPUs currently supports FSDP as the training engine, vLLM and SGLang as the inference engine. We will support Megatron in the future.\n"
  },
  {
    "path": "docs/start/more_resources.rst",
    "content": "More Resources\n==============\n\nLast updated: 06/30/2025.\n\n- Introduction to verl (`Slides <https://tongyx361.github.io/blogs/posts/verl-intro>`_)\n- verl Code Walkthrough (`Slides <https://tongyx361.github.io/blogs/posts/verl-tutorial>`_, `Talk in Chinese <https://hcqnc.xetlk.com/sl/3vACOK>`_) \n"
  },
  {
    "path": "docs/start/multinode.rst",
    "content": "Multinode Training\n==================\n\nLast updated: 06/10/2025.\n\n.. _wuxibin89: https://github.com/wuxibin89\n\nAuthor: `Xibin Wu <https://github.com/wuxibin89>`_, `Yusheng Su <https://yushengsu-thu.github.io/>`_.\n\nOption 1: Launch Manually\n------------------------------\n\nSet up multinode ray cluster\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n1. Start head node with ``ray start --head --dashboard-host=0.0.0.0``, there're 2 address you should care about:\n\n- GCS address: ``ray start --address=<address>``, where worker node should connect to.\n- Dashboard address: ``<address>:8265``, where you should submit job to the cluster.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/head.png?raw=true\n\n2. Start worker node with ``ray start --address=<address>`` you get above.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/worker.png?raw=true\n\n3. Now you should see the cluster have 2 nodes with ``ray status``.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/status.png?raw=true\n\n4. Additionally, you can access dashboard in the browser with the address you get above. \n\n*Firewall rules maybe need configure to access the dashboard, if there's any trouble, please contact your network administrator.*\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/overview.png?raw=true\n\nSubmit job to ray cluster\n~~~~~~~~~~~~~~~~~~~~~~~~~\n1. Submit ray job to cluster with the dashboard address you get above.\n\n.. code-block:: bash\n\n    ray job submit --address=\"http://127.0.0.1:8265\" \\\n        --runtime-env=verl/trainer/runtime_env.yaml \\\n        --no-wait \\\n        -- \\\n        python3 -m verl.trainer.main_ppo \\\n        trainer.n_gpus_per_node=8 \\\n        trainer.nnodes=2 \\\n        ...\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/submit.png?raw=true\n\n2. Then you can check the job status with the following commands:\n\n- ray job list: list all jobs submitted to the cluster.\n- ray job logs <Submission ID>: query the logs of the job.\n- ray job status <Submission ID>: query the status of the job.\n- ray job stop <Submission ID>: request the job to be stopped.\n- ray job list | grep submission_id | grep JobStatus | grep RUNNING | grep -oP 'raysubmit_[^'\\''\"]+' | head -n 1: get the latest job submission ID of the running job.\n- ray job logs <Submission ID> --follow: added ``--follow`` parameter to ray job logs command to enable continuous log streaming.\n\n3. You can also access driver/task/actor logs in ``/tmp/ray/session_latest/logs/``, driver log is ``job-driver-raysubmit_<Submission ID>.log``.\n\n4. We strongly recommend you to view job detail from dashboard in multinode training, because it provide more structure way to view the job information.\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/job.png?raw=true\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/job_detail.png?raw=true\n\nOption 2: Launch via SkyPilot on Kubernetes or clouds\n------------------------------------------------------\n\n.. note::\n   Ready-to-use SkyPilot example configurations are available in the `examples/skypilot/ <https://github.com/volcengine/verl/tree/main/examples/skypilot>`_ directory:\n   \n   - ``verl-ppo.yaml`` - PPO training with GSM8K dataset\n   - ``verl-grpo.yaml`` - GRPO training with MATH dataset  \n   - ``verl-multiturn-tools.yaml`` - Multi-turn tool usage training\n   \n   See the `SkyPilot examples README <https://github.com/volcengine/verl/tree/main/examples/skypilot>`_ for detailed usage instructions.\n\nStep 1: Setup SkyPilot\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~\nSkyPilot can support different clouds, here we use GCP as example. `install skypilot <https://docs.skypilot.co/en/latest/getting-started/installation.html>`_\n\n.. code-block:: bash\n\n    conda create -y -n sky python=3.10\n    conda activate sky\n    pip install \"skypilot[gcp]\"\n\n    conda install -c conda-forge google-cloud-sdk\n    gcloud init\n\n    # Run this if you don't have a credential file.\n    # This will generate ~/.config/gcloud/application_default_credentials.json.\n    gcloud auth application-default login\n\n    # Check if the GCP credential is correctly setup.\n    sky check gcp\n\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/setup_skypilot.png?raw=true\n\nStep 2: Prepare dataset\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n   git clone https://github.com/volcengine/verl.git\n   cd examples/data_preprocess\n   python3 gsm8k.py --local_save_dir ~/data/gsm8k\n\n\nStep 3: Submit a job with SkyPilot\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n1. Create a SkyPilot YAML ``verl-cluster.yml`` with the following content:\n\n.. parsed-literal:: workdir: .  will sync all the data in the current dir to the remote cluster.\n\n.. code-block:: yaml\n\n   resources:\n     accelerators: L4:1 # every node has 1 L4 GPU\n     image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4\n     memory: 64+        # every node has 64 GB memory\n     ports: 8265        # expose port for ray dashboard\n\n   num_nodes: 2         # cluster size\n\n   # --------------- Work Directory Synchronization (workdir) ---------------\n   # Defines the local working directory to be synchronized to the remote cluster.\n   # Here, '.' means synchronizing the directory where the sky submit command is currently run.\n   workdir: .\n\n   # --------------- (secrets) ---------------\n   secrets:\n     ## your wandb api key ##\n     WANDB_API_KEY: null\n\n   # --------------- File Mounts/Data Upload (file_mounts) ---------------\n   # If your dataset (gsm8k folder) is local, it needs to be uploaded to the remote cluster.\n   file_mounts:\n     # Remote path (relative to remote user's home directory): Local path\n     # /remote/dir1/file: /local/dir1/file\n     data/gsm8k: ~/data/gsm8k\n\n   # --------------- Environment Setup (setup) ---------------\n   # Commands run on each node of the remote cluster to set up the environment (e.g., install dependencies). These are run directly inside Docker.\n   setup: |\n     rm -rf verl\n     git clone https://github.com/volcengine/verl.git\n     cd verl\n     pip3 install -v -e .[vllm]\n\n   # --------------- Run Command (run) ---------------\n   # The actual task commands to be executed on the remote cluster.\n   # This script will first start the Ray cluster (different ray start commands are executed on Head and Worker nodes).\n   # Then, your training script will only be run on the Head node (SKYPILOT_NODE_RANK == 0).\n   run: |\n     # Get the Head node's IP and total number of nodes (environment variables injected by SkyPilot).\n     head_ip=`echo \"$SKYPILOT_NODE_IPS\" | head -n1`\n     num_nodes=`echo \"$SKYPILOT_NODE_IPS\" | wc -l` # Here num_nodes should be equal to 2.\n\n     # login wandb\n     python3 -c \"import wandb; wandb.login(relogin=True, key='$WANDB_API_KEY')\"\n\n     # Start Ray based on node role (Head=0, Worker>0).\n     # This logic is a standard Ray cluster startup script.\n     if [ \"$SKYPILOT_NODE_RANK\" == \"0\" ]; then\n       # Head node starts Ray Head.\n       echo \"Starting Ray head node...\"\n       # Check if a Ray Head is already running to avoid duplicate starts.\n       ps aux | grep ray | grep 6379 &> /dev/null ||  ray start --head --disable-usage-stats \\\n             --port=6379 \\\n             --dashboard-host=0.0.0.0 \\\n             --dashboard-port=8265\n\n       # Wait for all worker nodes to join the cluster.\n       while [ $(ray nodes | grep NODE_ID | wc -l) -lt $num_nodes ]; do\n         echo \"Waiting for all nodes to join... ($(ray nodes | grep NODE_ID | wc -l)/$num_nodes)\"\n         sleep 5\n       done\n\n       # Head node executes the training script.\n       echo \"Executing training script on head node...\"\n\n       python3 -m verl.trainer.main_ppo \\\n        data.train_files=data/gsm8k/train.parquet \\\n        data.val_files=data/gsm8k/test.parquet \\\n        data.train_batch_size=256 \\\n        data.max_prompt_length=512 \\\n        data.max_response_length=256 \\\n        actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n        actor_rollout_ref.rollout.name=vllm \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n        critic.optim.lr=1e-5 \\\n        critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n        critic.ppo_micro_batch_size_per_gpu=4 \\\n        algorithm.kl_ctrl.kl_coef=0.001 \\\n        trainer.logger=['console','wandb'] \\\n        trainer.val_before_train=False \\\n        trainer.default_hdfs_dir=null \\\n        trainer.n_gpus_per_node=1 \\\n        trainer.nnodes=2 \\\n        trainer.save_freq=20 \\\n        trainer.test_freq=20 \\\n        trainer.total_epochs=2 \\\n        trainer.project_name=verl_examples \\\n        trainer.experiment_name=experiment_name_gsm8k\n\n     else\n       # Wait for Ray Head to start.\n       sleep 10 # Increase waiting time to ensure Head finishes starting.\n       # Worker node starts Ray Worker.\n       echo \"Starting Ray worker node...\"\n\n       # Check if a Ray Worker is already running to avoid duplicate starts.\n       ps aux | grep ray | grep $head_ip:6379 &> /dev/null || ray start --address $head_ip:6379 --disable-usage-stats\n\n       # Add sleep to after `ray start` to give ray enough time to daemonize\n       sleep 5 # Ensure Worker successfully connects to Head.\n     fi\n\n     # No commands are added to the Worker node here; the Worker's main task is to start Ray and wait for the Head node to assign tasks.\n     echo \"Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK.\"\n\n\n.. code-block:: bash\n\n    export WANDB_API_KEY=<your-wandb-api-key>\n    sky launch -c verl --secret WANDB_API_KEY verl-cluster.yml\n\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/running_job.png?raw=true\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/running_job_1.png?raw=true\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/finished.png?raw=true\n\n**Check the cluster on GCP**\n\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/gcp_instances.png?raw=true\n\n**Check Ray Dashboard**\n\nWe can see the cluster on the RAY Dashboard with the GCP head node:\n\n```console\n$ sky status --endpoint 8265 verl\n1.2.3.4:8265\n```\n\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_overview.png?raw=true\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_jobs.png?raw=true\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/ray_dashboard_cluster.png?raw=true\n\n\n**Check the checkpoint of model**\n\n.. code-block:: bash\n\n    # login the head node\n    ssh verl\n    # The global step will vary. Find the correct path from the training logs.\n    cd ~/sky_workdir/checkpoints/verl_examples/gsm8k/\n    # Then list contents to find the checkpoint, e.g.:\n    ls -R .\n\n.. image:: https://github.com/yottalabsai/open-source/blob/main/static/verl/saved_model.png?raw=true\n\n\nOption 3: Launch via Slurm\n------------------------------\n\nRay provides users with `this <https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html>`_ official\ntutorial to start a Ray cluster on top of Slurm. We have verified the :doc:`GSM8K example<../examples/gsm8k_example>`\non a Slurm cluster under a multi-node setting with the following steps.\n\n1. [Optional] If your cluster support `Apptainer or Singularity <https://apptainer.org/docs/user/main/>`_ and you wish\nto use it, convert verl's Docker image to an Apptainer image. Alternatively, set up the environment with the package\nmanager available on your cluster or use other container runtimes (e.g. through `Slurm's OCI support <https://slurm.schedmd.com/containers.html>`_) available to you.\n\n.. code:: bash\n\n    apptainer pull /your/dest/dir/vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3.sif docker://verlai/verl:vemlp-th2.4.0-cu124-vllm0.6.3-ray2.10-te1.7-v0.0.3\n\n2. Follow :doc:`GSM8K example<../examples/gsm8k_example>` to prepare the dataset and model checkpoints.\n\n3. Modify `examples/slurm/ray_on_slurm.slurm <https://github.com/volcengine/verl/blob/main/examples/slurm/ray_on_slurm.slurm>`_ with your cluster's own information.\n\n4. Submit the job script to the Slurm cluster with `sbatch`.\n\nPlease note that Slurm cluster setup may vary. If you encounter any issues, please refer to Ray's\n`Slurm user guide <https://docs.ray.io/en/latest/cluster/vms/user-guides/community/slurm.html>`_ for common caveats.\n\nIf you changed Slurm resource specifications, please make sure to update the environment variables in the job script if necessary.\n\n\nOption 4: Launch via dstack\n------------------------------\n\n`dstackai/dstack <https://github.com/dstackai/dstack>`_ is an open-source container orchestrator that simplifies distributed training across cloud providers and on-premises environments\nwithout the need to use K8S or Slurm.\n\nPrerequisite\n~~~~~~~~~~~~\nOnce dstack is `installed <https://dstack.ai/docs/installation>`_, initialize the directory as a repo with ``dstack init``. \n\n.. code-block:: bash\n\n    mkdir myproject && cd myproject\n    dstack init\n\n**Create a fleet**\n\nBefore submitting distributed training jobs, create a `dstack` `fleet <https://dstack.ai/docs/concepts/fleets>`_.\n\nRun a Ray cluster task\n~~~~~~~~~~~~~~~~~~~~~~\n\nOnce the fleet is created, define a Ray cluster task, e.g. in ``ray-cluster.dstack.yml``:\n\n.. code-block:: yaml\n\n    type: task\n    name: ray-verl-cluster\n\n    nodes: 2\n\n    env:\n        - WANDB_API_KEY\n        - PYTHONUNBUFFERED=1\n        - CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n    \n    image: verlai/verl:app-verl0.6-transformers4.56.1-sglang0.5.2-mcore0.13.0-te2.2\n    commands:\n        - git clone https://github.com/volcengine/verl\n        - cd verl\n        - pip install --no-deps -e .\n        - pip install hf_transfer hf_xet\n        - |\n        if [ $DSTACK_NODE_RANK = 0 ]; then\n            python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k\n            python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-7B-Instruct')\" \n            ray start --head --port=6379;\n        else\n            ray start --address=$DSTACK_MASTER_NODE_IP:6379\n        fi\n\n    # Expose Ray dashboard port\n    ports:\n        - 8265\n\n    resources:\n        gpu: 80GB:8\n        shm_size: 128GB\n\n    # Save checkpoints on the instance\n    volumes:\n        - /checkpoints:/checkpoints\n\nNow, if you run this task via `dstack apply`, it will automatically forward the Ray's dashboard port to `localhost:8265`.\n\n.. code-block:: bash\n\n    dstack apply -f ray-cluster.dstack.yml\n\nAs long as the `dstack apply` is attached, you can use `localhost:8265` to submit Ray jobs for execution\n\nSubmit Ray jobs\n~~~~~~~~~~~~~~~\n\nBefore you can submit Ray jobs, ensure to install `ray` locally:\n   \n.. code-block:: shell\n\n    pip install ray\n\nNow you can submit the training job to the Ray cluster which is available at ``localhost:8265``:\n   \n.. code-block:: shell\n\n    $ RAY_ADDRESS=http://localhost:8265\n    $ ray job submit \\\n        -- python3 -m verl.trainer.main_ppo \\\n        data.train_files=/root/data/gsm8k/train.parquet \\\n        data.val_files=/root/data/gsm8k/test.parquet \\\n        data.train_batch_size=256 \\\n        data.max_prompt_length=512 \\\n        data.max_response_length=256 \\\n        actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n        critic.optim.lr=1e-5 \\\n        critic.model.path=Qwen/Qwen2.5-7B-Instruct \\\n        critic.ppo_micro_batch_size_per_gpu=4 \\\n        algorithm.kl_ctrl.kl_coef=0.001 \\\n        trainer.project_name=ppo_training \\\n        trainer.experiment_name=qwen-2.5-7B \\\n        trainer.val_before_train=False \\\n        trainer.n_gpus_per_node=8 \\\n        trainer.nnodes=2 \\\n        trainer.default_local_dir=/checkpoints \\\n        trainer.save_freq=10 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15 2>&1 | tee verl_demo.log \\\n        trainer.resume_mode=disable\n\n\nFor more details on how `dstack` works, check out its `documentation <https://dstack.ai/docs>`_.\n\nHow to debug?\n---------------------\n\n\nRay Distributed Debugger VSCode Extension (Recommended)\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\n\n1. Starting with Ray 2.39, Anyscale has introduced the `Ray Distributed Debugger <https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html>`_ VSCode extension. Follow the extension’s installation instructions, then add your cluster using the dashboard URL you obtained earlier.\n\n   .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/debugger.png?raw=true\n      :alt: Ray Distributed Debugger VSCode extension screenshot\n\n2. Prerequisites.\n\n   Ensure the following are installed (see the extension README for more detail):\n\n   - Visual Studio Code  \n   - `ray[default]` >= 2.9.1  \n   - `debugpy` >= 1.8.0  \n\n   .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/c7098b755ff689859837773a916c857.png?raw=true\n      :alt: VSCode with Ray prerequisites\n\n3. Environment Variables.\n\n   To enable post‑mortem debugging, set:\n\n   .. code-block:: bash\n\n      export RAY_DEBUG_POST_MORTEM=1\n\n   .. admonition:: Note\n      :class: important\n\n      Be sure to remove any legacy flags before starting Ray:\n\n      - `RAY_DEBUG=legacy`  \n      - `--ray-debugger-external`\n\n4. Configuring BreakpointsSet up breakpoint() in your code, and submit job to cluster. Then the extension will show the breakpoint information.\n\n\n   1. Insert `breakpoint()` calls into your remote functions.  \n   2. Submit your job to the cluster.  \n\n   The extension will detect active breakpoints and display them in VSCode.\n\n   .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/4ddad74395c79a1402331c0ce73316f.png?raw=true\n      :alt: Detected breakpoint in VSCode\n\n   **Note:** Breakpoints are only supported inside functions decorated with `@ray.remote`.\n\n5. Launching the Debugger.\n\n   Run your job directly from the command line (do not use a `launch.json`):\n\n   .. code-block:: bash\n\n      python job.py\n\n6. Attaching to a Breakpoint.\n\n Once the process hits the first `breakpoint()`, click the Ray Distributed Debugger icon in the VSCode sidebar to attach the debugger.\n\n   .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/4ddad74395c79a1402331c0ce73316f.png?raw=true\n      :alt: Attaching VSCode debugger to Ray process\n\n7. Debugging With Multiple breakpoint().\n\n   For each subsequent task, first disconnect the current debugger session, then click the extension icon again to attach to the next breakpoint.\n\n   .. image:: https://github.com/aoshen524/verl/blob/main/docs/start/6e83c910a62c82fecb89c6619e001cd.png?raw=true\n      :alt: Disconnecting and reconnecting the debugger\n\nLegacy Ray Debugger\n~~~~~~~~~~~~~~~~~~~\n1. Ray has a builtin legacy `debugger <https://docs.ray.io/en/latest/ray-observability/user-guides/debug-apps/ray-debugging.html>`_ that allows you to debug your distributed applications. To enable debugger, start ray cluster with ``RAY_DEBUG=legacy`` and ``--ray-debugger-external``.\n\n.. code-block:: bash\n\n    # start head node\n    RAY_DEBUG=legacy ray start --head --dashboard-host=0.0.0.0 --ray-debugger-external\n    # start worker node\n    RAY_DEBUG=legacy ray start --address='10.124.46.192:6379' --ray-debugger-external\n\n2. Set up breakpoint in your code, and submit job to cluster. Then run ``ray debug`` to wait breakpoint:\n\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/legacy.png?raw=true\n\n\nMulti-node training on AMD clusters\n---------------------------------------------------------------------------------------\n\nIf you want to run multi-node training with slurm with Docker/Podman container on AMD Cluster, you can use the following script. \n\nIf you encounter any issues in using AMD GPUs running verl, please contact `Yusheng Su <https://yushengsu-thu.github.io/>`_.\n\n.. note::\n    1. You need to use ``podman`` or ``docker`` in the following script. We will release the apptainer script later.\n    2. If you want to use ``podman``, you just replace ``docker`` with ``podman`` in the following script.\n\nThe script includes the following steps:\n\n1. SLURM Configuration\n2. Environment Setup\n3. Docker/Podman Container Setup\n4. Ray Cluster Initialization\n5. Data Preprocessing\n6. Model Setup\n7. Training Launch\n\n\nslurm_script.sh\n~~~~~~~~~~~~~~~~~~~~\n\n.. code-block:: bash\n\n    #!/bin/bash\n\n    #SBATCH --job-name=verl-ray-on-slurm\n    #SBATCH --nodes=2\n    #SBATCH --ntasks-per-node=2\n    #SBATCH --mem=200G\n    #SBATCH --time=30-00:00:00\n    #SBATCH --gpus-per-node=8\n    #SBATCH --cpus-per-task=28\n    #SBATCH --output=../verl_log/slurm-%j.out\n    #SBATCH --error=../verl_log/slurm-%j.err\n    #SBATCH --nodelist=gpu-[0,1]\n\n\n    # load necessary modules\n    ### Run this setup\n    # [Cluster]: Use docker\n    # docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n\n\n    ##########################################################################\n    ###The following setting should be set in different project and cluster###\n    ##########################################################################\n\n    ### Project\n    CONTAINER_NAME=\"multinode_verl_training\"\n    IMG=\"verl.rocm\"\n    DOCKERFILE=\"docker/Dockerfile.rocm\"\n    # echo $PWD\n    verl_workdir=\"${HOME}/projects/verl_upstream\"\n    export TRANSFORMERS_CACHE=\"${HOME}/.cache/huggingface\"\n    export HF_HOME=$TRANSFORMERS_CACHE\n\n    ### Cluster Network Setting\n    export NCCL_DEBUG=TRACE\n    export GPU_MAX_HW_QUEUES=2\n    export TORCH_NCCL_HIGH_PRIORITY=1\n    export NCCL_CHECKS_DISABLE=1\n    # export NCCL_IB_HCA=rdma0,rdma1,rdma2,rdma3,rdma4,rdma5,rdma6,rdma7 \n    export NCCL_IB_HCA=mlx5_0,mlx5_1,mlx5_2,mlx5_3,mlx5_4,mlx5_5,mlx5_8,mlx5_9\n    export NCCL_IB_GID_INDEX=3\n    export NCCL_CROSS_NIC=0\n    export CUDA_DEVICE_MAX_CONNECTIONS=1\n    export NCCL_PROTO=Simple\n    export RCCL_MSCCL_ENABLE=0\n    export TOKENIZERS_PARALLELISM=false\n    export HSA_NO_SCRATCH_RECLAIM=1\n    ##########################################################################\n\n    ### For rocm and training script\n    export HIP_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n    export ROCR_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n    export CUDA_VISIBLE_DEVICES=$HIP_VISIBLE_DEVICES\n\n\n    # Build and launch the Docker container\n    srun bash -c \"\n        # Exit on any error\n        set -e \n\n        # Clean up dangling images (images with <none> tag)\n        docker image prune -f\n\n        # Need to pull the docker first\n        docker pull docker.io/rocm/vllm:rocm6.2_mi300_ubuntu20.04_py3.9_vllm_0.6.4\n        \n        if ! docker images --format \"{{.Repository}}:{{.Tag}}\" | grep -q \"${IMG}\"; then\n            echo \\\"Building ${IMG} image...\\\"\n            docker build -f \\\"${DOCKERFILE}\\\" -t \\\"${IMG}\\\" .\n        else\n            echo \\\"${IMG} image already exists, skipping build\\\"\n        fi\n\n        # Removing old container if exists\n        docker rm \\\"${CONTAINER_NAME}\\\" 2>/dev/null || true\n\n        # Checking network devices\n        ibdev2netdev\n\n        # Launch the docker\n        docker run --rm -d \\\n        -e HYDRA_FULL_ERROR=1 \\\n        -e HIP_VISIBLE_DEVICES=${HIP_VISIBLE_DEVICES} \\\n        -e ROCR_VISIBLE_DEVICES=${ROCR_VISIBLE_DEVICES} \\\n        -e CUDA_VISIBLE_DEVICES=${CUDA_VISIBLE_DEVICES} \\\n        -e NCCL_DEBUG=${NCCL_DEBUG} \\\n        -e GPU_MAX_HW_QUEUES=${GPU_MAX_HW_QUEUES} \\\n        -e TORCH_NCCL_HIGH_PRIORITY=${TORCH_NCCL_HIGH_PRIORITY} \\\n        -e NCCL_CHECKS_DISABLE=${NCCL_CHECKS_DISABLE} \\\n        -e NCCL_IB_HCA=${NCCL_IB_HCA} \\\n        -e NCCL_IB_GID_INDEX=${NCCL_IB_GID_INDEX} \\\n        -e NCCL_CROSS_NIC=${NCCL_CROSS_NIC} \\\n        -e CUDA_DEVICE_MAX_CONNECTIONS=${CUDA_DEVICE_MAX_CONNECTIONS} \\\n        -e NCCL_PROTO=${NCCL_PROTO} \\\n        -e RCCL_MSCCL_ENABLE=${RCCL_MSCCL_ENABLE} \\\n        -e TOKENIZERS_PARALLELISM=${TOKENIZERS_PARALLELISM} \\\n        -e HSA_NO_SCRATCH_RECLAIM=${HSA_NO_SCRATCH_RECLAIM} \\\n        -e TRANSFORMERS_CACHE=${TRANSFORMERS_CACHE} \\\n        -e HF_HOME=${HF_HOME} \\\n        --network host \\\n        --device /dev/dri \\\n        --device /dev/kfd \\\n        --device /dev/infiniband \\\n        --group-add video \\\n        --cap-add SYS_PTRACE \\\n        --security-opt seccomp=unconfined \\\n        --privileged \\\n        -v \\${HOME}:\\${HOME} \\\n        -v \\${HOME}/.ssh:/root/.ssh \\\n        -w \"${verl_workdir}\" \\\n        --shm-size 128G \\\n        --name \\\"${CONTAINER_NAME}\\\" \\\n        \\\"${IMG}\\\" \\\n        tail -f /dev/null\n\n        echo \\\"Container setup completed\\\"\n    \"\n        # (Optional): If you do not want to root mode and require assign yuorself as the user\n        # Please add `-e HOST_UID=$(id -u)` and `-e HOST_GID=$(id -g)` into the above docker launch script. \n\n\n\n\n\n    ### Ray launch the nodes before training\n\n    # Getting the node names\n    nodes_array=($(scontrol show hostnames \"$SLURM_JOB_NODELIST\" | tr '\\n' ' '))\n\n    head_node=${nodes_array[0]}\n    head_node_ip=$(srun --nodes=1 --ntasks=1 -w \"$head_node\" hostname --ip-address)\n\n    # if we detect a space character in the head node IP, we'll\n    # convert it to an ipv4 address. This step is optional.\n    if [[ \"$head_node_ip\" == *\" \"* ]]; then\n        IFS=' ' read -ra ADDR <<<\"$head_node_ip\"\n    if [[ ${#ADDR[0]} -gt 16 ]]; then\n        head_node_ip=${ADDR[1]}\n    else\n        head_node_ip=${ADDR[0]}\n    fi\n        echo \"IPV6 address detected. We split the IPV4 address as $head_node_ip\"\n    fi\n\n    port=6379\n    ip_head=$head_node_ip:$port\n    export ip_head\n    echo \"IP Head: $ip_head\"\n\n    # make sure we set environment variables before Ray initialization\n\n    # Print out all env variables\n    printenv\n\n    echo \"Starting HEAD at $head_node\"\n    srun --nodes=1 --ntasks=1 -w \"$head_node\" \\\n        docker exec \"${CONTAINER_NAME}\" \\\n            ray start --head --node-ip-address=\"$head_node_ip\" --port=$port \\\n            --dashboard-port=8266 \\\n            --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n    # optional, though may be useful in certain versions of Ray < 1.0.\n    sleep 10\n\n    # number of nodes other than the head node\n    worker_num=$((SLURM_JOB_NUM_NODES - 1))\n\n    for ((i = 1; i <= worker_num; i++)); do\n        node_i=${nodes_array[$i]}\n        echo \"Debug: Starting worker on node_i = ${node_i}\"\n        if [ -z \"$node_i\" ]; then\n            echo \"Error: Empty node name for worker $i\"\n            continue\n        fi\n        echo \"Starting WORKER $i at $node_i\"\n        srun --nodes=1 --ntasks=1 -w \"$node_i\" \\\n            docker exec \"${CONTAINER_NAME}\" \\\n                ray start --address \"$ip_head\" --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n        sleep 5\n    done\n\n\n\n\n    # Ray initlization test (See whether any error in the above execution)\n    echo \"Testing Ray initialization in the slurm nodes...\"\n    docker exec \"${CONTAINER_NAME}\" python3 -c '\n    import ray\n    try:\n        ray.init(address=\"auto\")\n        print(\"\\n=== Ray Cluster Status ===\")\n        print(f\"Number of nodes: {len(ray.nodes())}\")\n        for node in ray.nodes():\n            print(\"Node: {}, Status: {}\".format(node[\"NodeManagerHostname\"], node[\"Alive\"]))\n            # print(f\"Node: {node}\")\n        ray.shutdown()\n        print(\"Ray initialization successful!\")\n    except Exception as e:\n        print(f\"Ray initialization failed: {str(e)}\")\n    '\n    echo \"=== Ray test completed ===\"\n    ######\n\n\n\n    # Run data preprocessing\n\n    echo \"Starting data preprocessing...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 \"examples/data_preprocess/gsm8k.py\" \"--local_save_dir\" \"../data/gsm8k\"\n\n    echo \"Starting data preprocessing...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 \"examples/data_preprocess/math_dataset.py\" \"--local_dir\" \"../data/math\"\n\n    train_files=\"../data/gsm8k/train.parquet\"\n    val_files=\"../data/gsm8k/test.parquet\"\n\n    # Download and test model\n    echo \"Loading model...\"\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')\"\n    MODEL_PATH=\"Qwen/Qwen2-7B-Instruct\"\n\n    # Set model path after pipeline test\n    MODEL_PATH=\"Qwen/Qwen2.5-0.5B-Instruct\"\n\n    echo \"== Data and model loading Done ==\"\n\n    echo \"Start to train...\"\n\n    docker exec \"${CONTAINER_NAME}\" \\\n        python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2-7B-Instruct')\"\n    MODEL_PATH=\"Qwen/Qwen2-7B-Instruct\"\n\n\n    PYTHONUNBUFFERED=1 srun --overlap --nodes=${SLURM_NNODES} --ntasks=1 -w \"$head_node\" \\\n        docker exec \"${CONTAINER_NAME}\" \\\n        python3 -m verl.trainer.main_ppo \\\n        data.train_files=$train_files \\\n        data.val_files=$val_files \\\n        data.train_batch_size=1024 \\\n        data.max_prompt_length=1024 \\\n        data.max_response_length=1024 \\\n        actor_rollout_ref.model.path=$MODEL_PATH \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n        actor_rollout_ref.rollout.name=vllm \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n        critic.optim.lr=1e-5 \\\n        critic.model.use_remove_padding=True \\\n        critic.model.path=$MODEL_PATH \\\n        critic.model.enable_gradient_checkpointing=False \\\n        critic.ppo_micro_batch_size_per_gpu=8 \\\n        critic.model.fsdp_config.param_offload=False \\\n        critic.model.fsdp_config.optimizer_offload=False \\\n        algorithm.kl_ctrl.kl_coef=0.0001 \\\n        trainer.critic_warmup=0 \\\n        trainer.logger='[\"console\",\"wandb\"]' \\\n        trainer.project_name='verl_example' \\\n        trainer.experiment_name='Qwen2.5-32B-Instruct_function_rm' \\\n        trainer.n_gpus_per_node=${SLURM_GPUS_PER_NODE} \\\n        trainer.val_before_train=False \\\n        trainer.nnodes=${SLURM_NNODES} \\\n        trainer.save_freq=-1 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15\n\n\nRun multi-node training with above slurm_script.sh\n~~~~~~~~~~~~~~~~~~~~\nJust sbatch your slurm_script.sh\n\n.. code-block:: bash\n\n    sbatch slurm_script.sh\n\n"
  },
  {
    "path": "docs/start/quickstart.rst",
    "content": ".. _quickstart:\n\n=========================================================\nQuickstart: PPO training on GSM8K dataset\n=========================================================\n\nPost-train a LLM using GSM8K dataset.\n\nIntroduction\n------------\n\n.. _hf_dataset_gsm8k: https://huggingface.co/datasets/openai/gsm8k\n\nIn this example, we train an LLM to tackle the `GSM8k <hf_dataset_gsm8k>`_ task with function-based rewards. [1]_\n\nPrerequisite:\n\n- the latest version of ``verl`` and its dependencies installed following the installation guide. Using the docker image is recommended.\n\n- a GPU with at least 24 GB HBM\n\n\nDataset Introduction\n--------------------\n\nGSM8k is a math problem dataset. The prompt is an elementary school\nproblem. The LLM model is asked to solve the math problem. Below is an example:\n\nPrompt\n\n   Katy makes coffee using teaspoons of sugar and cups of water in the\n   ratio of 7:13. If she used a total of 120 teaspoons of sugar and cups\n   of water, calculate the number of teaspoonfuls of sugar she used.\n\nSolution\n\n   The total ratio representing the ingredients she used to make the\n   coffee is 7+13 = <<7+13=20>>20 Since the fraction representing the\n   number of teaspoons she used is 7/20, she used 7/20\\ *120 =\n   <<7/20*\\ 120=42>>42 #### 42\n\nStep 1: Prepare the dataset\n----------------------------\n\nWe preprocess the dataset in parquet format so that (1) it contains necessary fields for computing RL rewards and (2) is faster to read.\n\n.. code-block:: bash\n\n   python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k\n\nStep 2: Download a model for post-training\n-------------------------------------------\n\nIn this example, we start with the ``Qwen2.5-0.5B-Instruct`` model.\n\nIf you want to perform SFT before RL, refer to the :doc:`Complete GSM8K Example<../examples/gsm8k_example>`, the `sft directory <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k>`_ and `SFT Trainer <https://github.com/volcengine/verl/blob/main/verl/trainer/sft_trainer.py>`_ for further details.\n\n.. code-block:: bash\n\n   python3 -c \"import transformers; transformers.pipeline('text-generation', model='Qwen/Qwen2.5-0.5B-Instruct')\"\n\nStep 3: Perform PPO training with the instruct model\n----------------------------------------------------------------------\n\n**Reward Model/Function**\n\nWe use a pre-defined rule-based reward model. We force the model to produce a final\nanswer following 4 “#” as shown in the solution. We extract the final\nanswer from both the solution and model's output using regular\nexpression matching. We assign a reward of 1 to correct\nanswer, 0.0 to incorrect answer and 0 to no answer. \n\nFor more details, please refer to `verl/utils/reward_score/gsm8k.py <https://github.com/volcengine/verl/blob/v0.4.1/verl/utils/reward_score/gsm8k.py>`_.\n\n**Training Script**\n\nNow let's run PPO training with the dataset and model above. [2]_\n\n\nSet the ``data.train_files`` ,\\ ``data.val_files``, ``actor_rollout_ref.model.path`` and ``critic.model.path`` based on your dataset and model names or paths.\nYou may set ``VERL_USE_MODELSCOPE=True`` to download models from `modelscope <https://www.modelscope.cn>`_ instead of `huggingface <https://huggingface.co>`_.\n\n.. code-block:: bash\n\n   PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.logger=console \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=10 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 2>&1 | tee verl_demo.log\n\nYou are expected to see the following logs, indicating training in progress. The key metric ``val/test_score/openai/gsm8k`` is computed every ``trainer.test_freq`` steps:\n\n.. code-block:: bash\n\n    step:0 - timing/gen:21.470 - timing/ref:4.360 - timing/values:5.800 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty_coeff:0.001 - timing/adv:0.109 - timing/update_critic:15.664 - critic/vf_loss:14.947 - critic/vf_clipfrac:0.000 - critic/vpred_mean:-2.056 - critic/grad_norm:1023.278 - critic/lr(1e-4):0.100 - timing/update_actor:20.314 - actor/entropy_loss:0.433 - actor/pg_loss:-0.005 - actor/pg_clipfrac:0.000 - actor/ppo_kl:0.000 - actor/grad_norm:1.992 - actor/lr(1e-4):0.010 - critic/score/mean:0.004 - critic/score/max:1.000 - critic/score/min:0.000 - critic/rewards/mean:0.004 - critic/rewards/max:1.000 - critic/rewards/min:0.000 - critic/advantages/mean:-0.000 - critic/advantages/max:2.360 - critic/advantages/min:-2.280 - critic/returns/mean:0.003 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.045 - critic/values/max:9.500 - critic/values/min:-14.000 - response_length/mean:239.133 - response_length/max:256.000 - response_length/min:77.000 - prompt_length/mean:104.883 - prompt_length/max:175.000 - prompt_length/min:68.000\n    step:1 - timing/gen:23.020 - timing/ref:4.322 - timing/values:5.953 - actor/reward_kl_penalty:0.000 - actor/reward_kl_penalty:0.001 - timing/adv:0.118 - timing/update_critic:15.646 - critic/vf_loss:18.472 - critic/vf_clipfrac:0.384 - critic/vpred_mean:1.038 - critic/grad_norm:942.924 - critic/lr(1e-4):0.100 - timing/update_actor:20.526 - actor/entropy_loss:0.440 - actor/pg_loss:0.000 - actor/pg_clipfrac:0.002 - actor/ppo_kl:0.000 - actor/grad_norm:2.060 - actor/lr(1e-4):0.010 - critic/score/mean:0.000 - critic/score/max:0.000 - critic/score/min:0.000 - critic/rewards/mean:0.000 - critic/rewards/max:0.000 - critic/rewards/min:0.000 - critic/advantages/mean:0.000 - critic/advantages/max:2.702 - critic/advantages/min:-2.616 - critic/returns/mean:0.000 - critic/returns/max:0.000 - critic/returns/min:0.000 - critic/values/mean:-2.280 - critic/values/max:11.000 - critic/values/min:-16.000 - response_length/mean:232.242 - response_length/max:256.000 - response_length/min:91.000 - prompt_length/mean:102.398 - prompt_length/max:185.000 - prompt_length/min:70.000\n\nCheckout ``Algorithm Baselines`` page for full training and validation logs for reference.\n\nThe checkpoint is saved at the following dir by default: ``checkpoints/${trainer.project_name}/${trainer.experiment_name}``. You can merge the saved checkpoints to huggingface model using ``verl.model_merger`` module, for example:\n\n.. code-block:: bash\n\n    python3 -m verl.model_merger merge \\\n        --backend fsdp \\\n        --local_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor \\\n        --target_dir checkpoints/${trainer.project_name}/${trainer.experiment_name}/global_step_1/actor/huggingface\n\nFor more details about checkpoint and model merging, please refer to :ref:`checkpoint-page`.\n\nTo enable ``wandb`` for experiment tracking, set the following configs:\n\n.. code-block:: bash\n\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=$YOUR_PROJECT_NAME \\\n    trainer.experiment_name=$YOUR_RUN_NAME \\\n\nIf you encounter out of memory issues with HBM less than 32GB, enable the following configs would help:\n\n.. code-block:: bash\n\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    critic.ppo_micro_batch_size_per_gpu=1 \\\n\nFor the full set of configs, please refer to :ref:`config-explain-page` for detailed explanation and performance tuning.\n\n\n.. [1] The original paper (https://arxiv.org/pdf/2110.14168) mainly focuses on training a verifier (a reward model) to solve math problems via Best-of-N sampling. In this example, we train an RL agent using a rule-based reward model.\n.. [2] More training script examples for FSDP and Megatron-LM backend are stored in `examples/ppo_trainer <https://github.com/volcengine/verl/tree/main/examples/ppo_trainer>`_ directory.\n"
  },
  {
    "path": "docs/start/ray_debug_tutorial.rst",
    "content": "Ray Debug Tutorial\r\n==================\r\n\r\nLast updated: 04/23/2025\r\n\r\n\r\n.. _wuxibin89: https://github.com/wuxibin89\r\n\r\nAuthor: `Ao Shen <https://aoshen524.github.io/>`_.\r\n\r\nHow to debug?\r\n---------------------\r\n\r\n\r\nRay Distributed Debugger VSCode Extension (Recommended)\r\n~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~\r\n\r\n1. Starting with Ray 2.39, Anyscale has introduced the `Ray Distributed Debugger <https://docs.ray.io/en/latest/ray-observability/ray-distributed-debugger.html>`_ VSCode extension. Follow the extension’s installation instructions, then add your cluster using the dashboard URL you obtained earlier.\r\n\r\n   .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/debugger.png?raw=true\r\n      :alt: Ray Distributed Debugger VSCode extension screenshot\r\n\r\n2. Prerequisites.\r\n\r\n   Ensure the following are installed (see the extension README for more detail):\r\n\r\n   - Visual Studio Code  \r\n   - `ray[default]` >= 2.9.1  \r\n   - `debugpy` >= 1.8.0  \r\n\r\n   .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/readme.png?raw=true\r\n      :alt: VSCode with Ray prerequisites\r\n\r\n3. Environment Variables.\r\n\r\n   To enable post‑mortem debugging, set:\r\n\r\n   .. code-block:: bash\r\n\r\n      export RAY_DEBUG_POST_MORTEM=1\r\n\r\n   .. admonition:: Note\r\n      :class: important\r\n\r\n      Be sure to remove any legacy flags before starting Ray:\r\n\r\n      - `RAY_DEBUG=legacy`  \r\n      - `--ray-debugger-external`\r\n\r\n4. Configuring BreakpointsSet up breakpoint() in your code, and submit job to cluster. Then the extension will show the breakpoint information.\r\n\r\n\r\n   1. Insert `breakpoint()` calls into your remote functions.  \r\n   2. Submit your job to the cluster.  \r\n\r\n   The extension will detect active breakpoints and display them in VSCode.\r\n\r\n   **Note:** Breakpoints are only supported inside functions decorated with `@ray.remote`.\r\n\r\n5. Launching the Debugger.\r\n\r\n   Run your job directly from the command line (do not use a `launch.json`):\r\n\r\n   .. code-block:: bash\r\n\r\n      python job.py\r\n\r\n6. Attaching to a Breakpoint.\r\n\r\n Once the process hits the first `breakpoint()`, click the Ray Distributed Debugger icon in the VSCode sidebar to attach the debugger.\r\n\r\n   .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/launch.png?raw=true\r\n      :alt: Attaching VSCode debugger to Ray process\r\n\r\n7. Debugging With Multiple breakpoint().\r\n\r\n   For each subsequent task, first disconnect the current debugger session, then click the extension icon again to attach to the next breakpoint.\r\n\r\n   .. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/disconnect.png?raw=true\r\n      :alt: Disconnecting and reconnecting the debugger\r\n\r\nLegacy Ray Debugger\r\n~~~~~~~~~~~~~~~~~~~\r\n1. Ray has a builtin legacy `debugger <https://docs.ray.io/en/latest/ray-observability/user-guides/debug-apps/ray-debugging.html>`_ that allows you to debug your distributed applications. To enable debugger, start ray cluster with ``RAY_DEBUG=legacy`` and ``--ray-debugger-external``.\r\n\r\n.. code-block:: bash\r\n\r\n    # start head node\r\n    RAY_DEBUG=legacy ray start --head --dashboard-host=0.0.0.0 --ray-debugger-external\r\n    # start worker node\r\n    RAY_DEBUG=legacy ray start --address='10.124.46.192:6379' --ray-debugger-external\r\n\r\n2. Set up breakpoint in your code, and submit job to cluster. Then run ``ray debug`` to wait breakpoint:\r\n\r\n.. image:: https://github.com/eric-haibin-lin/verl-community/blob/main/docs/ray/legacy.png?raw=true\r\n\r\n"
  },
  {
    "path": "docs/workers/automodel_workers.rst",
    "content": "Automodel Backend\n=================\n\nLast updated: 03/07/2026.\n\nWe support the Automodel (nemo_automodel) backend by implementing the\n``AutomodelEngine`` and ``AutomodelEngineWithLMHead`` engine classes.\nThe Automodel backend delegates model building, parallelization, optimizer\nsharding, LR scheduling, gradient clipping, and checkpointing to\nnemo_automodel's infrastructure while using verl's training loop,\ndata pipeline, and loss function.\n\n**Requirements**\n\n- Automodel r0.3.0\n- transformers v5.0.0\n\n**Pros**\n\n- Supports FSDP2 and TP distributed strategies out of\n  the box.\n\n- Native support for Mixture-of-Experts (MoE) models with Expert\n  Parallelism (EP) via DeepEP.\n\n- TransformerEngine (TE) integration for optimized attention, linear\n  layers, and RMSNorm.\n\n- Readily supports any HuggingFace model without checkpoint conversion.\n\n**Cons**\n\n- Pipeline parallelism is not yet supported.\n\n\nSFT Examples\n------------\n\nWe provide example SFT training scripts using the Automodel backend in\n`examples/sft/gsm8k/ <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/>`_.\n\nBasic: Qwen2.5-0.5B with FSDP2\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nA minimal example using ``Qwen/Qwen2.5-0.5B-Instruct`` with FSDP2 and\nno parallelism:\n\n.. code:: shell\n\n   bash examples/sft/gsm8k/run_qwen_05_automodel.sh 4 /tmp/automodel_sft_test\n\nSee `run_qwen_05_automodel.sh <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/run_qwen_05_automodel.sh>`_.\n\nAdvanced: Qwen3-30B MoE with Expert Parallelism\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\nA larger-scale example using ``Qwen/Qwen3-30B-A3B-Base`` (MoE model)\nwith Expert Parallelism (EP=8), DeepEP, TransformerEngine backend, and\ntorch_mm experts backend:\n\n.. code:: shell\n\n   bash examples/sft/gsm8k/run_qwen3_30b_automodel.sh 8 /tmp/automodel_sft_30b\n\nSee `run_qwen3_30b_automodel.sh <https://github.com/volcengine/verl/blob/main/examples/sft/gsm8k/run_qwen3_30b_automodel.sh>`_.\n"
  },
  {
    "path": "docs/workers/fsdp_workers.rst",
    "content": "PyTorch FSDP Backend\n======================\n\nLast updated: 12/01/2025.\n\nWe support PyTorch FSDP Backend by implementing various workers for\nactor, critic, reference, rollout and reward models.\n\n**Pros**\n\n- Readily support various models.\n\n  - Users only need to implement the corresponding\n    ``dtensor_weight_loader`` for weight synchronization between FSDP\n    and vLLM. While for ``hf_weight_loader``, users can directly apply\n    any models supported both in HF and vLLM without any code change.\n\n- Easy to organize the forward and backward computation for each model.\n\n**Cons**\n\n- Poor scalability when it comes to large-scale models (e.g. Llama 70B\n  and 405B)\n- The resharding overhead between actor and rollout could be larger than\n  Megatron-LM backend.\n\nDue to the simplicity, we recommend using FSDP backend for algorithm\nresearch and prototyping.\n\nFSDP Workers\n--------------\n\nActorRolloutRefWorker\n^^^^^^^^^^^^^^^^^^^^^\n\nActor/Rollout HybridEngine\n''''''''''''''''''''''''''\n\n1. HybridEngine, Actor and Rollout initialization API.\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n   def init_model(self):\n\n``ONE_TO_ALL``: when calling the ``init_model`` function from the driver\nprocess, each worker (on a GPU) will execute the following model\ninitialization process.\n\nThe initialization details of HybridEngine, Actor and Rollout are\nhighlighted below:\n\n1. ``DataParallelPPOActor`` implements the simple PPO computation logics\n   when the model is built with FSDP, including compute log prob, model\n   update.\n2. ``vLLMRollout`` support generation with vLLM. We modify the vLLM\n   Engine and make it executed under SPMD to fit into our\n   ``WorkerGroup`` design.\n\nSee `source code <https://github.com/volcengine/verl/blob/main/verl/workers/fsdp_workers.py>`_. for more information.\n\n1. Generate sequence and recompute log prob\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def generate_sequences(self, prompts: DataProto):\n\n- ``Dispatch.DP_COMPUTE_PROTO``: The data will be dispatched and\n  collected along the DP dimension\n\n- In this function, the rollout model will perform auto-regressive\n  generation and the actor model will recompute the old log prob for the\n  generated response.\n\n3. Update actor model\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def update_actor(self, data: DataProto):\n\n- Update the actor model weight using PPO & entropy loss.\n\nReferenceModel\n''''''''''''''\n\n1. Reference model initialization\n\nThe reference model is initialized using the same function as the actor\nmodel without initializing the HybridEngine and Optimizer. Then the\nactor model is also wrapped by the ``DataParallelPPOActor``.\n\n2. Compute reference log prob\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def compute_ref_log_prob(self, data: DataProto):\n\n- In this function, the reference model will call the compute log prob\n  function in ``DataParallelPPOActor`` to compute the reference log\n  prob.\n\nCriticWorker and RewardWorker\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. Model initialization\n\nQuite similar to reference model. The CriticWorker will perform\nadditional initialization for the Optimizer.\n\n2. Compute Values for CriticWorker\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def compute_values(self, data: DataProto):\n\n3. Update Critic\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def update_critic(self, data: DataProto):\n\n4. Compute Reward\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n   def compute_rm_score(self, data: DataProto):\n\n\nHybridShard\n------------\n\nWe didn't support FSDP `HybridShard`. To support this, we may need to\nconstruct a 2D device mesh and test the corresponding\n``dtensor_weight_loader`` and ``hf_weight_loader`` for each model.\n"
  },
  {
    "path": "docs/workers/megatron_workers.rst",
    "content": "Megatron-LM Backend\n===================\n\nLast updated: 12/01/2025.\n\nWe support Megatron Backend by implementing various workers for actor,\ncritic, reference, rollout and reward models. We also implement the\n``3DHybridEngine`` using Megatron-LM and vLLM/SGLang in\n`megatron_vllm.py <https://github.com/volcengine/verl/blob/main/verl/workers/sharding_manager/megatron_vllm.py>`_\nand `megatron_sglang.py <https://github.com/volcengine/verl/blob/main/verl/workers/sharding_manager/megatron_sglang.py>`_.\n\n**Pros**\n\n- Support 5D parallelism (TP, EP, CP, DP, PP) and sequence parallelism\n  for best scalablility and throughput.\n- 3D HybridEngine can significantly reduce peak memory usage and reduce\n  weight synchronize overhead between actor and rollout.\n\n**Cons**\n\n- Huggingface Models and Megatron checkpoints need tools for conversion.\n\n\nDevelopment Progress\n--------------------\n\n\nNote that [Deprecated] means that the feature is not supported in the latest\nversion of verl.\n[To-Optimize] means that the feature is implemented but not optimized yet.\n[WIP] means that the feature is working in progress.\n[In-Release] means that the feature is ready and in review process,\ncoming at any time.\n\n\n+---------------+-----------------------------------------------------------+\n| [Deprecated]  | Megatron 3D Parallelism with custom models                |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron 0.11.0 ``GPTModel`` support                      |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron GRPO support                                     |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron with vLLM 0.8.2, with per-tensor weights loading |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron with Context Parallel                            |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Qwen2MoE model support                                    |\n+---------------+-----------------------------------------------------------+\n| [To-Optimize] | Megatron dist Checkpoint                                  |\n+---------------+-----------------------------------------------------------+\n| [To-Optimize] | Huggingface and Megatron Checkpoint Converter             |\n+---------------+-----------------------------------------------------------+\n| [To-Optimize] | Efficient fused linear, entropy and cross entropy         |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron offload(param, grad, optimizer)                  |\n+---------------+-----------------------------------------------------------+\n| [Done]        | Megatron Profiler                                         |\n+---------------+-----------------------------------------------------------+\n| [In-Release]  | Megatron 0.12.0, TE 2.2 with vLLM 0.8.3 and Fused Attn    |\n+---------------+-----------------------------------------------------------+\n| [WIP]         | Moonlight/DeepSeek-V3 model support                       |\n+---------------+-----------------------------------------------------------+\n| [WIP]         | Expert Parallel support                                   |\n+---------------+-----------------------------------------------------------+\n| [WIP]         | Megatron support dynamic batch size                       |\n+---------------+-----------------------------------------------------------+\n| [To-Do]       | Performance tuning                                        |\n+---------------+-----------------------------------------------------------+\n| [MileStone]   | Runnable with DeepSeek-V3 671B post-training              |\n+---------------+-----------------------------------------------------------+\n\n\n\nUtils of Megatron Workers\n-------------------------\n\nMegatronWorker\n^^^^^^^^^^^^^^\n\n``MegatronWorker`` is the base class of different megatron worker\nclasses. In this class, ``get_megatron_global_info`` and\n``get_megatron_rank_info`` function to retrieve the 3D parallel world\nsize and rank of each ``Worker`` running on specific GPU. These information\nwill be used in transfer protocol for Megatron Backend.\n\nThe following ``Worker`` class for different models will be utilized to\nconstruct the ``WorkerGroup`` .\n\nWe implement various of APIs for each ``Worker`` class decorated by the\n``@register(dispatch_mode=)`` . These APIs can be called by the ray\ndriver process. The data can be correctly collect and dispatch following\nthe ``dispatch_mode`` on each function. The supported dispatch_model\n(i.e., transfer protocols) can be found in `decorator.py <https://github.com/volcengine/verl/blob/main/verl/single_controller/base/decorator.py>`_.\n\nActorRolloutRefWorker\n^^^^^^^^^^^^^^^^^^^^^\n\nThis class is implemented for Actor/Rollout HybridEngine or for the\nreference model to initialize their model and perform computation.\n\nActor/Rollout HybridEngine\n''''''''''''''''''''''''''\n\n1. HybridEngine, Actor and Rollout initialization API.\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n   def init_model(self):\n\n``ONE_TO_ALL``: when calling the ``init_model`` function from the driver\nprocess, each worker (on a GPU) will execute the following model\ninitialization process.\n\nThe initialization details of HybridEngine, Actor and Rollout are\nhighlighted below:\n\n1. ``MegatronPPOActor`` implements the simple PPO computation logics\n   when the model is built with Megatron, including compute log prob,\n   model update.\n2. ``vLLMRollout`` support generation with vLLM. We modify the vLLM\n   Engine and make it executed under SPMD to fit into our\n   ``WorkerGroup`` design.\n\nSee `source code <https://github.com/volcengine/verl/blob/main/verl/workers/megatron_workers.py#L63>`_ for more information.\n\n.. code:: python\n\n   # build actor model\n   self.actor = MegatronPPOActor(config=self.config.actor,\n                                 model_config=self.actor_model_config,\n                                 megatron_config=megatron_config,\n                                 actor_module=self.actor_module,\n                                 actor_optimizer=self.actor_optimizer,\n                                 actor_optimizer_config=self.actor_optim_config)\n\n   # build rollout\n   # rollout initialization\n   rollout = vLLMRollout(actor_module=params,\n                        config=self.config.rollout,\n                        tokenizer=self.tokenizer,\n                        model_hf_config=self.actor_model_config,\n                        train_tp=mpu.get_tensor_model_parallel_world_size())\n   ...\n\n1. Generate sequence and recompute log prob\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_PP_AS_DP_PROTO)\n   def generate_sequences(self, prompts: DataProto):\n\n- ``Dispatch.MEGATRON_PP_AS_DP_PROTO``: The PP dimension of the actor\n  model will be regarded as DP dimension. Then the driver process will\n  dispatch and collect the data according to this reorganization. This\n  is because, in HybridEngine, the actor weight, which usually applied\n  larger 3D parallel sizes, will be gathered along the PP dimension and\n  TP dimension. Therefore, the corresponding data should be dispatched\n  and collected through the 3D parallel group of the rollout model,\n  rather than the actor model. However, the world_size and rank\n  information can only be retrieved from ``get_megatron_global_info`` and\n  ``get_megatron_rank_info``, which records the 3D information for the\n  actor model. Moreover, the data resharding inside TP dimension will be\n  processed within the HybridEngine.\n\n- In this function, the rollout model will perform auto-regressive\n  generation and the actor model will recompute the old log prob for the\n  generated response.\n\n3. Update actor model\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)\n   def update_actor(self, data: DataProto):\n\n- ``Dispatch.MEGATRON_COMPUTE_PROTO``: User passes the data partitioned\n  by DP dimension. The data is dispatched to all tp/pp ranks within the\n  same dp group, and ultimately only collects output data from tp=0 and\n  the last pp.\n- Update the actor model weight using PPO & entropy loss.\n\n\n..note:: \n\n   Currently, training Tensor Parallel Size can be different from inference\n   Tensor Parallel Size.\n\n\nReferenceModel\n''''''''''''''\n\n1. Reference model initialization\n\nThe reference model is initialized using the same function as the actor\nmodel without initializing the HybridEngine and Optimizer. Then the\nactor model is also wrapped by the ``MegatronPPOActor``.\n\n2. Compute reference log prob\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)\n   def compute_ref_log_prob(self, data: DataProto):\n\n- In this function, the reference model will call the compute log prob\n  function in ``MegatronPPOActor`` to compute the reference log prob.\n\nCriticWorker and RewardWorker\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. Model initialization\n\nQuite similar to reference model. The CriticWorker will perform\nadditional initialization for the Optimizer.\n\n2. Compute Values for CriticWorker\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)\n   def compute_values(self, data: DataProto):\n\n3. Update Critic\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)\n   def update_critic(self, data: DataProto):\n\n4. Compute Reward\n\n.. code:: python\n\n   @register(dispatch_mode=Dispatch.MEGATRON_COMPUTE_PROTO)\n   def compute_rm_score(self, data: DataProto):\n\n\nUtils of Train Optimization\n---------------------------\n\nOffload\n^^^^^^^\nWhen resources are tight, the offload method can lower GPU memory \nusage, helping training and inference frameworks work well under verl. \nIt moves parameters, gradients, and optimizers to CPU memory and only \nloads them back to the GPU when needed.\n\nIf you want to use the offload, you can add the following parameters \nfor the actor and ref separately. \n\n.. code:: python\n\n   # For the actor\n   actor_rollout_ref.actor.megatron.param_offload=True \\\n   actor_rollout_ref.actor.megatron.grad_offload=True \\\n   actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n   # For the ref w/o grad and optimizer\n   actor_rollout_ref.ref.megatron.param_offload=True \\\n\n\nFor the critic, you can include these parameters.\n\n.. code:: python\n\n   # For the critic\n   critic.megatron.param_offload=True \\\n   critic.megatron.grad_offload=True \\\n   critic.megatron.optimizer_offload=True \\\n\n\nRelated MCore Document\n----------------------\n\nThere is also a detailed document of using MCore to train different\nkinds of models, please refer to `MCore Document <https://github.com/volcengine/verl/blob/main/verl/models/mcore/readme.md>`_.\n"
  },
  {
    "path": "docs/workers/model_engine.rst",
    "content": "Model Engine\n============\n\n.. _vermouth: https://github.com/vermouth1992\n\nAuthor: `Chi Zhang <https://github.com/vermouth1992>`_\n\nLast updated: 09/25/2025.\n\nCurrent Support Matrix\n----------------------\n\n+----------+-----------+--------------+-------------+--------------------------+\n| Backends | Model     | Scalability  | Model       | Pain points              |\n|          | Supported |              | Definition  |                          |\n|          |           |              |             |                          |\n+==========+===========+==============+=============+==========================+\n| FSDP     | Day 1     | - Dense is OK| Huggingface | Monkey patch can be      |\n| +        | support   |              | + monkey    | easily impacted by       |\n| ulysses  | HF model  | - MoE is bad | patch       | transformers version     |\n+----------+-----------+--------------+-------------+--------------------------+\n| MCore    | Limited   | Best         | GPTModel    | Supporting new models is |\n|          |           |              | (One model  | difficult                |\n|          |           |              | for all)    |                          |\n+----------+-----------+--------------+-------------+--------------------------+\n\n-  We monkey patch attention function to support ulysses\n-  We monkey patch VLM models to support FSDP with mixed data with and\n   without images\n\nClass Hierarchy\n---------------\n\nNote that all the workers and trainers run in **SPMD** mode. SFT/DPO/RM\ntrainer is directly invoked by ``torchrun``. The Actor/Critic worker can\nalso be invoked by a RayWorkerGroup and provides APIs to a single\ncontroller.\n\n-  Base Engine level: implement model init, optimizer init, lr scheduler\n   init, sharding, checkpoint manager.\n-  Full Engine level: subclass base engine and implement\n   ``forward_step``.\n-  Worker/SPMD trainer level: **engine agnostic**, implement training\n   logics using abstract engine APIs\n\nRL trainer utilizes workers to construct HybridFlow program. This is out\nof the scope of model engine.\n\nExisting Model Types\n--------------------\n\n========== ====================== ======================\nModel type Language model         Value model\n========== ====================== ======================\nInput      text/image/video/audio text/image/video/audio\nOutput     logits for next token  logits as value\n========== ====================== ======================\n\nCurrently, we have two model types: language model and value model. We\nexpect to expand the category to include Qwen-Omni family (output both\ntext and audio) and VLA models.\n\nData Format\n-----------\n\nCurrently, verl adopts left-right padding data format in RL trainer.\nThis creates massive padding when the discrepancy between response\nlength is large. We will start to implement no-padding format throughout\nthe whole system.\n\n.. image:: https://github.com/vermouth1992/verl-data/blob/master/images/data_format.png?raw=true\n   :alt: Data Format\n\nHere is the migration plan:\n- Implement no-padding format in engine\n- Add a transformation layer in Actor/Critic worker.\n- Replace Actor/Critic Worker in RL trainer\n- Implement no-padding throughput system\n\nCheckpoint System\n-----------------\n\n.. image:: https://github.com/vermouth1992/verl-data/blob/master/images/verl-ckpt.png?raw=true\n   :alt: Model Engine Checkpoint System\n\nThe engine constructs the model using huggingface config, then load\nweights from huggingface checkpoint. If the engine directly uses\nhuggingface model definition, it can use function provided by\n``transformers``. Otherwise, each engine has to write their own\ncheckpoint load logic (e.g.,\n`mbridge <https://github.com/ISEEKYAN/mbridge>`__). During model\ntraining, each engine has to implement save_checkpoint and\nload_checkpoint that save/load intermediate sharded checkpoint including\nmodel, optimizer and lr scheduler states. Each engine has to implement a\ncheckpoint merge script, that merges the intermediate sharded checkpoint\nback to huggingface format.\n\nAPI\n---\n\nA tentative model engine API can be found:\nhttps://github.com/volcengine/verl/blob/main/verl/workers/engine/base.py#L24\n\nExtension\n---------\n\nAdd a new backend\n~~~~~~~~~~~~~~~~~\n\n-  Start a new folder under ``verl/workers/engine``. Then, implement\n   ``transformer_impl.py``. If you want to implement a non-transformer\n   model, please contact us in advance.\n-  Add the engine config to the GSM8k SFT trainer script:\n   https://github.com/volcengine/verl/blob/main/tests/special_e2e/sft/run_sft_engine_gsm8k.sh\n-  Invoke the tests with your backend:\n   https://github.com/volcengine/verl/blob/main/tests/special_e2e/sft/test_sft_engine_all.sh.\n   This test script will run various backends and various\n   configurations, and compare the loss and grad norm of the first step\n   to make sure they are close.\n\nAdd a new model type\n~~~~~~~~~~~~~~~~~~~~\n\n-  This is mainly reserved for models whose the output is not just text\n   (e.g., Qwen3-Omni). Please discuss with us before you proceed.\n"
  },
  {
    "path": "docs/workers/ray_trainer.rst",
    "content": "PPO Ray Trainer\n===============\n\nLast updated: 02/12/2025.\n\nWe implement the RayPPOTrainer, which is a trainer runs on the driver\nprocess on a single CPU/GPU node (default is CPU).\n\nThe PPORayTrainer include 3 core functions for data preparation,\nWorkerGroup initialization and PPO training loop.\n\nData Preparation\n----------------\n\nThe ``PPORayTrainer``, as a single process, is responsible for loading a\ncomplete batch of samples (prompts) from the dataset and then dispatch\nto different worker_groups running on different GPUs.\n\nTo generalize the data loading, we implement the ``RLHFDataset`` class\nto load the preprocessed parquet files, apply chat templates to the\nprompts, add padding, truncate prompts that exceed max prompt length and\nthen tokenize.\n\n.. code:: python\n\n   self.train_dataset = RLHFDataset(data_files=self.config.data.train_files,\n                                       tokenizer=self.tokenizer,\n                                       config=self.config.data)\n\nThen, the dataloader will iterate the dataset under PPO mini batch size.\n\nWorkerGroup Initialization\n--------------------------\n\nWe first introduce a basic implementation of initializing the\n``WorkerGroup`` of the actor model on a given set of GPUs.\n\n.. code:: python\n\n   # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool\n   # For FSDP backend, we recommend using max_colocate_count=1 that merge all WorkerGroups into one.\n   # For Megatron backend, we recommend using max_colocate_count>1 that can utilize different WorkerGroup for differnt models\n   resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes,\n                                   use_gpu=True,\n                                   max_colocate_count=1)\n   # define actor rollout cls to be init on remote\n   actor_rollout_cls = RayClassWithInitArgs(cls=ActorRolloutWorker)\n   # define actor_rollout worker group\n   actor_rollout_worker_group = MegatronRayWorkerGroup(resource_pool=resource_pool,\n                                                       ray_cls_with_init=actor_rollout_cls,\n                                                       default_megatron_kwargs=config.actor_rollout.megatron)\n\nDifferent WorkerGroups, like ``actor_rollout_worker_group`` ,\n``critic_worker_group`` and ``ref_worker_group`` lies on a separate\nprocess in the above implementation.\n\nThe driver process can then call the distributed compute function within\nthe ``actor_rollout_worker_group`` and other roles to construct the RL\ntraining loop.\n\nFor models colocated in the same set of GPUs, we further provide a\nfine-grain optimization, which merge the ``worker_group`` of different roles\nin the same process. This optimization can save the redundant\nCUDA/distributed context in different processes.\n\n.. code:: python\n\n   # initialize WorkerGroup\n   # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,\n   # you should not use `create_colocated_worker_cls`. Instead, directly pass different resource pool to different worker groups.\n   # See TODO(url) for more information.\n   all_wg = {}\n   for resource_pool, class_dict in self.resource_pool_to_cls.items():\n       worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n       wg_dict = self.ray_worker_group_cls(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)\n       spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n       all_wg.update(spawn_wg)\n\n   if self.use_critic:\n       self.critic_wg = all_wg['critic']\n       self.critic_wg.init_model()\n\n   if self.use_reference_policy:\n       self.ref_policy_wg = all_wg['ref']\n       self.ref_policy_wg.init_model()\n\n   if self.use_rm:\n       self.rm_wg = all_wg['rm']\n       self.rm_wg.init_model()\n\n   # we should create rollout at the end so that vllm can have a better estimation of kv cache memory\n   self.actor_rollout_wg = all_wg['actor_rollout']\n   self.actor_rollout_wg.init_model()\n\n.. note:: For megatron backend, if we merge the ``worker_groups`` into the same processes, all the roles will utilize the same 3D parallel size. To optimize this, we may need to maintain several 3D process groups for each role in the same distributed context. If you want to use different 3D parallel size for different roles, please follow the similar architecture of the first code block to initialize each role's ``worker_group``\n\n\nPPO Training Loop\n-----------------\n\nWe implement the PPO training loop by calling the functions in\nworker_group of each role. The input and output data of each function is\na ``DataProto`` object implemented in `protocol.py <https://github.com/volcengine/verl/blob/main/verl/protocol.py>`_. In the training\nloop, trainer will dispatch/collect the data to/from different GPUs\nfollowing the transfer protocols wrapped in the workers' functions. The\ncomputation of PPO micro batches is processed in ``update_actor`` and\n``update_critic`` functions.\n\nTo extend to other RLHF algorithms, such as DPO, GRPO, please refer to\n:doc:`../advance/dpo_extension`.\n\n.. code:: python\n\n   def fit(self):\n       \"\"\"\n       The training loop of PPO.\n       The driver process only need to call the compute functions of the worker group through RPC to construct the PPO dataflow.\n       The light-weight advantage computation is done on the driver process.\n       \"\"\"\n       from verl.utils.tracking import Tracking\n       from omegaconf import OmegaConf\n\n       logger = Tracking(project_name=self.config.trainer.project_name,\n                           experiment_name=self.config.trainer.experiment_name,\n                           default_backend=self.config.trainer.logger,\n                           config=OmegaConf.to_container(self.config, resolve=True))\n\n       global_steps = 0\n\n       # perform validation before training\n       # currently, we only support validation using the reward_function.\n       if self.val_reward_fn is not None:\n           val_metrics = self._validate()\n           pprint(f'Initial validation metrics: {val_metrics}')\n\n       for epoch in range(self.config.trainer.total_epochs):\n           for batch_dict in self.train_dataloader:\n               metrics = {}\n\n               batch: DataProto = DataProto.from_single_dict(batch_dict)\n               # batch = batch.to('cuda')\n\n               # pop those keys for generation\n               gen_batch = batch.pop(batch_keys=['input_ids', 'attention_mask', 'position_ids'])\n\n               # generate a batch\n               with Timer(name='gen', logger=None) as timer:\n                   gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)\n               metrics['timing/gen'] = timer.last\n\n               batch = batch.union(gen_batch_output)\n\n               if self.use_reference_policy:\n                   # compute reference log_prob\n                   with Timer(name='ref', logger=None) as timer:\n                       ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)\n                       batch = batch.union(ref_log_prob)\n                   metrics['timing/ref'] = timer.last\n\n               # compute values\n               with Timer(name='values', logger=None) as timer:\n                   values = self.critic_wg.compute_values(batch)\n                   batch = batch.union(values)\n               metrics['timing/values'] = timer.last\n\n               with Timer(name='adv', logger=None) as timer:\n                   # compute scores. Support both model and function-based.\n                   # We first compute the scores using reward model. Then, we call reward_fn to combine\n                   # the results from reward model and rule-based results.\n                   if self.use_rm:\n                       # we first compute reward model score\n                       reward_tensor = self.rm_wg.compute_rm_score(batch)\n                       batch = batch.union(reward_tensor)\n\n                   # we combine with rule-based rm\n                   reward_tensor = self.reward_fn(batch)\n                   batch.batch['token_level_scores'] = reward_tensor\n\n                   # compute rewards. apply_kl_penalty if available\n                   batch, kl_metrics = apply_kl_penalty(batch,\n                                                           kl_ctrl=self.kl_ctrl_in_reward,\n                                                           kl_penalty=self.config.algorithm.kl_penalty)\n                   metrics.update(kl_metrics)\n\n                   # compute advantages, executed on the driver process\n                   batch = compute_advantage(batch,\n                                               self.config.algorithm.gamma,\n                                               self.config.algorithm.lam,\n                                               adv_estimator=self.config.algorithm.adv_estimator)\n               metrics['timing/adv'] = timer.last\n\n               # update critic\n               if self.use_critic:\n                   with Timer(name='update_critic', logger=None) as timer:\n                       critic_output = self.critic_wg.update_critic(batch)\n                   metrics['timing/update_critic'] = timer.last\n                   critic_output_metrics = reduce_metrics(critic_output.meta_info['metrics'])\n                   metrics.update(critic_output_metrics)\n\n               # implement critic warmup\n               if self.config.trainer.critic_warmup <= global_steps:\n                   # update actor\n                   with Timer(name='update_actor', logger=None) as timer:\n                       actor_output = self.actor_rollout_wg.update_actor(batch)\n                   metrics['timing/update_actor'] = timer.last\n                   actor_output_metrics = reduce_metrics(actor_output.meta_info['metrics'])\n                   metrics.update(actor_output_metrics)\n\n               # validate\n               if self.val_reward_fn is not None and (global_steps + 1) % self.config.trainer.test_freq == 0:\n                   with Timer(name='testing', logger=None) as timer:\n                       val_metrics: dict = self._validate()\n                       val_metrics = {f'val/{key}': val for key, val in val_metrics.items()}\n                   metrics['timing/testing'] = timer.last\n                   metrics.update(val_metrics)\n\n               # collect metrics\n               data_metrics = compute_data_metrics(batch=batch)\n               metrics.update(data_metrics)\n\n               # TODO: make a canonical logger that supports various backend\n               logger.log(data=metrics, step=global_steps)\n\n               if self.config.trainer.save_freq > 0 and (global_steps + 1) % self.config.trainer.save_freq == 0:\n                   actor_local_path = os.path.join(self.config.trainer.default_local_dir, 'actor',\n                                                   f'global_step_{global_steps}')\n                   actor_remote_path = os.path.join(self.config.trainer.default_hdfs_dir, 'actor')\n                   self.actor_rollout_wg.save_checkpoint(actor_local_path, actor_remote_path)\n\n                   if self.use_critic:\n                       critic_local_path = os.path.join(self.config.trainer.default_local_dir, 'critic',\n                                                           f'global_step_{global_steps}')\n                       critic_remote_path = os.path.join(self.config.trainer.default_hdfs_dir, 'critic')\n                       self.critic_wg.save_checkpoint(critic_local_path, critic_remote_path)\n\n               global_steps += 1\n\n       # perform validation after training\n       if self.val_reward_fn is not None:\n           val_metrics = self._validate()\n           pprint(f'Final validation metrics: {val_metrics}')\n"
  },
  {
    "path": "docs/workers/sglang_worker.rst",
    "content": "SGLang Backend\n==============\n\nLast updated: 05/31/2025.\n\n**Authored By SGLang RL Team and listed alphabetically by last name**\n\n`Jingyi Chen <https://github.com/fzyzcjy>`_, `Yitong Guan <https://github.com/minleminzui>`_, `Zhuobin Huang <https://zobinhuang.github.io/sec_about/>`_, `Jiajun Li <https://github.com/guapisolo>`_, `Ji Li <https://github.com/GeLee-Q>`_, `Shenggui Li <https://franklee.xyz/about>`_, `Junrong Lin <https://github.com/ocss884>`_, `Xiang Long <https://github.com/SwordFaith>`_, `Rui Lu <https://scholar.google.com/citations?user=-MGuqDcAAAAJ>`_, `Jin Pan <https://jhinpan.github.io/>`_, `Shuai Shi <https://github.com/shuaills>`_, `Yushen Su <https://yushengsu-thu.github.io/>`_, `Xinyuan Tong <https://github.com/JustinTong0323>`_, `Chendong Wang <https://github.com/cedricbeta>`_, `Hanchen Zhang <https://scholar.google.com/citations?user=pGcJcagAAAAJ>`_, `Haoran Wang <https://ubecc.github.io/about/>`_, `Yongan Xiang <https://github.com/BearBiscuit05>`_, `Chengxing Xie <https://yitianlian.github.io/>`_, `Yuhao Yang <https://github.com/yhyang201>`_, `Jinwei Yao <https://kivi-yao.github.io/>`_, `Qiaolin Yu <https://github.com/Qiaolin-Yu>`_, `Yuzhen Zhou <https://github.com/zyzshishui>`_, `Chenyang Zhao <https://github.com/zhaochenyang20>`_\n\n\n\nIntroduction\n------------\n`SGLang <https://github.com/sgl-project/sglang>`_ is an open-source state-of-the-art inference service engine, fully adopted by xAI to support all inference needs of Grok during research and serving processes.\n\nCurrently, verl fully supports using SGLang as the inference engine during the rollout phase. As a rollout engine, SGLang provides the same feature coverage as vLLM., including memory saving and multi-node rollout features. After installing verl and SGLang, simply add ``actor_rollout_ref.rollout.name=sglang`` at startup script to seamlessly switch between the two inference frameworks.\n\nIn addition, the SGLang team is actively working on supporting features such as Multi-Turn Agentic RL, VLM RLHF, Server-Based RLHF, and Partial Rollout. You can track the related development progress in the `Tracking Roadmap <https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/issues/74>`_.\n\nInstallation\n------------\nPlease always follow the following command to install SGLang with verl. \n\n.. code-block:: bash\n    \n    pip install --upgrade pip\n    # Currently 0.4.8, subject to updates at any time, please refer to the latest version specified in `setup.py`\n    pip install -e \".[sglang]\"\n\nYou can check the following dependencies are in your environment:\n\n.. note::\n\n    - **PyTorch**: 2.6.0+cu124\n    - **CUDA**: 12.4\n    - **flashinfer-python**: 0.2.5+cu124torch2.6\n    - **SGLang**: 0.4.6.post5\n    - **sgl-kernel**: 0.1.4\n\nUsing SGLang as the Inference Backend for PPO Training on a Single Machine\n-------------------------------------------------------------------------\nWe use Qwen/Qwen2-7B-Instruct on the gsm8k dataset for a simple test.\n\n1. Run the following command to prepare the gsm8k dataset:\n\n.. code-block:: bash\n\n    python3 examples/data_preprocess/gsm8k.py\n\n2. Run the following script to conduct a PPO experiment on a single machine with 4 GPUs:\n\n.. code-block:: bash\n\n    export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK=True\n    PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n        data.train_files=$HOME/data/gsm8k/train.parquet \\\n        data.val_files=$HOME/data/gsm8k/test.parquet \\\n        data.train_batch_size=4096 \\\n        data.max_prompt_length=4096 \\\n        data.max_response_length=4096 \\\n        actor_rollout_ref.rollout.name=sglang \\\n        actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n        critic.optim.lr=1e-5 \\\n        critic.model.path=Qwen/Qwen2-7B-Instruct \\\n        critic.ppo_micro_batch_size_per_gpu=4 \\\n        critic.model.fsdp_config.param_offload=True \\\n        critic.model.fsdp_config.optimizer_offload=True \\\n        algorithm.kl_ctrl.kl_coef=0.001 \\\n        trainer.logger=console \\\n        trainer.val_before_train=False \\\n        trainer.n_gpus_per_node=4 \\\n        trainer.nnodes=1 \\\n        trainer.save_freq=-1 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15 2>&1 | tee verl_demo.log\n\nWhy export SGL_DISABLE_TP_MEMORY_INBALANCE_CHECK?\n^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^\n\n1. ``verl`` initializes a ``SGLangRollout`` module during rollout, which is used to evaluate/generate samples.\n\n2. ``SGLangRollout`` will initialize ``Engine``, and further initialize a ``torch.distributed.DeviceMesh``, used to support Tensor Parallel (TP).\n\n3. ``DeviceMesh.init()`` internally checks the free GPU memory of all participating devices. If the difference is too large (more than ~10%), it directly reports an error to avoid initialization failures or deadlocks.\n\nWhy might there be inconsistent GPU memory?\n\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\"\n\n**1. Ray Distributed Actor loads the model at different times**\n\n``verl`` uses Ray-based multi-process, multi-GPU concurrent training. Each ``WorkerDict`` may be called at different times:\n\n.. code-block:: python\n\n    self.rollout = SGLangRollout(...)\n\nDifferent workers initialize the model at different times → different memory usage.\n\n**2. Delayed initialization causes memory bias**\n\nSome workers start model loading/inference (e.g., ``generate_sequences()``, ``compute_log_prob()``) earlier than others.  \nEarly workers already use up GPU memory → late workers still have empty memory → memory difference appears.\n\n**3. SGLang's TP init uses \"all-device broadcast\", but there's no uniform release timing**\n\nAlthough ``SGLangRollout`` may only involve subset of GPUs, its ``Engine`` initialization calls ``torch.distributed.init_process_group()`` and broadcasts weights, so:\n\n- Non-rollout GPUs also join the communication.\n- Later on, ``DeviceMesh`` init will fail due to \"inconsistent memory\".\n\n**4. Different FSDP/TP loading behaviors also lead to mismatch**\n\nIf using:\n\n.. code-block:: bash\n\n    actor.fsdp_config.param_offload=True  \n    ref.fsdp_config.param_offload=True\n\nThen some workers keep params on CPU while others already sharded to GPU → leads to asymmetric memory layout.\n\nUsing SGLang as the Inference Backend for PPO Training Across Multiple Machines\n------------------------------------------------------------------------------\nSGLang also supports running verl's RAY-based cross-machine inference in IPv4 and IPv6 scenarios. In the script below, we use TP=16 for cross-machine inference. Suppose we have two interconnected machines: node0 with IP 10.94.16.4 and node1 with IP 10.94.16.5.\n\n1. Start Ray on node0:\n\n.. code-block:: bash\n\n    ray start --head --dashboard-host=0.0.0.0\n\nYou will see the following prompt:\n\n.. code-block:: bash\n\n    Usage stats collection is enabled. To disable this, add `--disable-usage-stats` to the command that starts the cluster, or run the following command: `ray disable-usage-stats` before starting the cluster. See https://docs.ray.io/en/master/cluster/usage-stats.html for more details.\n\n    Local node IP: 10.94.16.4\n\n    --------------------\n    Ray runtime started.\n    --------------------\n\n    Next steps\n    To add another node to this Ray cluster, run\n        ray start --address='10.94.16.4:6379'\n\n2. Have node1 join the Ray cluster:\n\nRun the following command on node1:\n\n.. code-block:: bash\n\n    ray start --address='10.94.16.4:6379'\n\nRun the following command to confirm that the Ray cluster now has two nodes:\n\n.. code-block:: bash\n\n    ray status\n\nYou can see that the cluster has two nodes with 16 GPUs:\n\n.. code-block:: bash\n\n    ======== Autoscaler status: 2025-04-09 09:25:37.694016 ========\n    Node status\n    ---------------------------------------------------------------\n    Active:\n     1 node_ef382ffd687d8f6b060c1b68e63ada7341b936fe5b1901dd04de1027\n     1 node_1eb4d7d07e793114c23a89d1a41f1f76acf6ef5b35af844a4ee8e4ba\n    Pending:\n     (no pending nodes)\n    Recent failures:\n     (no failures)\n\n    Resources\n    ---------------------------------------------------------------\n    Usage:\n     0.0/360.0 CPU\n     0.0/16.0 GPU\n     0B/3.39TiB memory\n     0B/372.53GiB object_store_memory\n\n3. Run the following script to train meta-llama/Llama-3.1-8B-Instruct with TP=16 across 2 machines using 16 GPUs:\n\n.. code-block:: bash\n\n    DATA_DIR=$HOME/data/gsm8k\n\n    python3 -m verl.trainer.main_ppo \\\n        actor_rollout_ref.rollout.name=sglang \\\n        data.train_files=$DATA_DIR/train.parquet \\\n        data.val_files=$DATA_DIR/test.parquet \\\n        data.train_batch_size=4096 \\\n        data.max_prompt_length=4096 \\\n        data.max_response_length=4096 \\\n        actor_rollout_ref.model.path=meta-llama/Llama-3.1-8B-Instruct \\\n        actor_rollout_ref.actor.optim.lr=1e-6 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=16 \\\n        actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n        actor_rollout_ref.rollout.free_cache_engine=True \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size=16 \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n        critic.optim.lr=1e-5 \\\n        critic.model.use_remove_padding=True \\\n        critic.model.path=meta-llama/Llama-3.1-8B-Instruct \\\n        critic.model.enable_gradient_checkpointing=True \\\n        critic.ppo_micro_batch_size=16 \\\n        critic.model.fsdp_config.param_offload=True \\\n        critic.model.fsdp_config.optimizer_offload=True \\\n        algorithm.kl_ctrl.kl_coef=0.001 \\\n        trainer.critic_warmup=0 \\\n        trainer.logger=console \\\n        trainer.val_before_train=True \\\n        trainer.n_gpus_per_node=8 \\\n        trainer.nnodes=2 \\\n        trainer.save_freq=-1 \\\n        trainer.test_freq=10 \\\n        trainer.total_epochs=15 2>&1 | tee verl_demo.log\n"
  },
  {
    "path": "docs/workers/trtllm_worker.rst",
    "content": "TensorRT-LLM Backend\n====================\n\nLast updated: 12/31/2025.\n\n**Authored By TensorRT-LLM Team**\n\nIntroduction\n------------\n`TensorRT-LLM <https://github.com/NVIDIA/TensorRT-LLM>`_ is a high-performance LLM inference engine with state-of-the-art optimizations for NVIDIA GPUs.\nThe verl integration of TensorRT-LLM is based on TensorRT-LLM's `Ray orchestrator <https://github.com/NVIDIA/TensorRT-LLM/tree/main/examples/ray_orchestrator>`_. This integration is in its early stage, with more features and optimizations to come.\n\nThe TensorRT-LLM rollout engine primarily targets the colocated mode. Instead of relying purely on standard colocated mode, we adopted a mixed design combining aspects of the hybrid engine and colocated mode.\n\nInstallation\n------------\nWe provide ``docker/Dockerfile.stable.trtllm`` for building a docker image with TensorRT-LLM pre-installed. The verl integration is supported from ``nvcr.io/nvidia/tensorrt-llm/release:1.2.0rc6``, and you can choose other TensorRT-LLM versions via ``TRTLLM_BASE_IMAGE`` from the `NGC Catalog <https://catalog.ngc.nvidia.com/orgs/nvidia/teams/tensorrt-llm/containers/release>`_.\n\nAlternatively, refer to the `TensorRT-LLM installation guide <https://nvidia.github.io/TensorRT-LLM/installation/index.html>`_ for compatible environments if you want to build your own.\n\nInstall verl with TensorRT-LLM:\n\n.. code-block:: bash\n\n    pip install --upgrade pip\n    pip install -e \".[trtllm]\"\n\n.. note::\n\n    Using the TensorRT-LLM rollout requires setting the following environment variables before launching the Ray cluster. These have been included in all the example scripts:\n\n    .. code-block:: bash\n\n        # Clean all SLURM/MPI/PMIx env to avoid PMIx mismatch error.\n        for v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n            unset \"$v\"\n        done\n\nUsing TensorRT-LLM as the Rollout Engine for GRPO\n-------------------------------------------------\n\nWe provide the following GRPO recipe scripts for you to test the performance and accuracy curve of TensorRT-LLM as the rollout engine:\n\n.. code-block:: bash\n\n    ## For FSDP training engine\n    bash examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh\n    ## For Megatron-Core training engine\n    bash examples/grpo_trainer/run_qwen2-7b_math_megatron_trtllm.sh\n\nUsing TensorRT-LLM as the Rollout Engine for DAPO\n-------------------------------------------------\n\nWe provide a DAPO recipe script ``recipe/dapo/test_dapo_7b_math_trtllm.sh``.\n\n.. code-block:: bash\n\n    ## For FSDP training engine\n    bash recipe/dapo/test_dapo_7b_math_trtllm.sh\n    ## For Megatron-Core training engine\n    TRAIN_ENGINE=megatron bash recipe/dapo/test_dapo_7b_math_trtllm.sh\n\n"
  },
  {
    "path": "examples/cispo_trainer/run_cispo_qwen2_5_0_5b_gsm8k.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\n\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=cispo \\\n    actor_rollout_ref.actor.clip_ratio_low=10 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.2 \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.model.torch_dtype=bfloat16 \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_cispo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_0_5b_cispo' \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=5 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=3 $@\n"
  },
  {
    "path": "examples/data_preprocess/aime2024_multiturn_w_tool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the DAPO-Math-17k dataset to multiturn format\n\"\"\"\n\nimport argparse\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/retool_aime2024\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_path = \"BytedTsinghua-SIA/AIME-2024\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"default\")\n    else:\n        dataset = datasets.load_dataset(data_path, \"default\")\n\n    train_dataset = dataset[\"train\"]\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            orig_extra_info = example.pop(\"extra_info\")\n            extra_info = orig_extra_info.copy()\n            extra_info[\"need_tools_kwargs\"] = True\n            extra_info[\"tools_kwargs\"] = {\n                \"code_interpreter\": {\n                    \"create_kwargs\": {\n                        \"ground_truth\": example[\"reward_model\"][\"ground_truth\"],\n                    },\n                },\n            }\n            example[\"extra_info\"] = extra_info\n            return example\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/dapo_multiturn_w_tool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the DAPO-Math-17k dataset to multiturn format\n\"\"\"\n\nimport argparse\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/retool_dapo\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_path = \"BytedTsinghua-SIA/DAPO-Math-17k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"default\")\n    else:\n        dataset = datasets.load_dataset(data_path, \"default\")\n\n    train_dataset = dataset[\"train\"]\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            orig_extra_info = example.pop(\"extra_info\")\n            extra_info = orig_extra_info.copy()\n            extra_info[\"need_tools_kwargs\"] = True\n            extra_info[\"tools_kwargs\"] = {\n                \"code_interpreter\": {\n                    \"create_kwargs\": {\n                        \"ground_truth\": example[\"reward_model\"][\"ground_truth\"],\n                    },\n                },\n            }\n            example[\"extra_info\"] = extra_info\n            return example\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/full_hh_rlhf.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\n- Preprocess data and split the training set into 75% for training RM and 25% for validting RM.\n- All the training data is used to train SFT and RL.\n- Both chosen and rejected is used to train SFT\n\"\"\"\n\nimport argparse\nimport os\n\nimport pandas as pd\nfrom datasets import load_dataset\nfrom tqdm.auto import tqdm\n\nfrom verl.utils.fs import copy, makedirs\n\n\ndef generate_sft_dataset(target_hdfs_path_dir, local_dir=\"~/data/full_hh_rlh/sft\", local_dataset_path=None):\n    if local_dataset_path is not None:\n        dataset = load_dataset(local_dataset_path)\n    else:\n        dataset = load_dataset(\"Dahoas/full-hh-rlhf\")\n    output = {\"prompt\": [], \"response\": []}\n    for data in tqdm(dataset[\"train\"]):\n        # add chosen\n        output[\"prompt\"].append(data[\"prompt\"])\n        output[\"response\"].append(data[\"chosen\"])\n\n        # add rejection\n        output[\"prompt\"].append(data[\"prompt\"])\n        output[\"response\"].append(data[\"rejected\"])\n\n    df = pd.DataFrame(output)\n\n    local_dir = os.path.expanduser(local_dir)\n    os.makedirs(local_dir, exist_ok=True)\n\n    local_path = os.path.join(local_dir, \"train.parquet\")\n\n    df.to_parquet(path=local_path)\n\n    if target_hdfs_path_dir is not None:\n        hdfs_dir = target_hdfs_path_dir + \"/\" + \"train.parquet\"\n        makedirs(hdfs_dir)\n\n        copy(local_path, hdfs_dir)\n\n\ndef generate_rm_dataset(target_hdfs_path_dir, local_dir=\"~/data/full_hh_rlh/rm\", local_dataset_path=None):\n    if local_dataset_path is not None:\n        train_dataset = load_dataset(local_dataset_path, split=\"train[:75%]\")\n        test_dataset = load_dataset(local_dataset_path, split=\"train[-25%:]\")\n    else:\n        train_dataset = load_dataset(\"Dahoas/full-hh-rlhf\", split=\"train[:75%]\")\n        test_dataset = load_dataset(\"Dahoas/full-hh-rlhf\", split=\"train[-25%:]\")\n\n    local_dir = os.path.expanduser(local_dir)\n    os.makedirs(local_dir, exist_ok=True)\n\n    for dataset, name in zip([train_dataset, test_dataset], [\"train\", \"test\"], strict=True):\n        output = {\"prompt\": [], \"chosen\": [], \"rejected\": []}\n        for data in tqdm(dataset):\n            # add chosen\n            output[\"prompt\"].append(data[\"prompt\"])\n            output[\"chosen\"].append(data[\"chosen\"])\n            output[\"rejected\"].append(data[\"rejected\"])\n\n        df = pd.DataFrame(output)\n\n        local_path = os.path.join(local_dir, name + \".parquet\")\n\n        df.to_parquet(path=local_path)\n\n        if target_hdfs_path_dir is not None:\n            hdfs_dir = target_hdfs_path_dir + \"/\" + name + \".parquet\"\n            makedirs(hdfs_dir)\n\n            copy(local_path, hdfs_dir)\n\n\ndef generate_rl_dataset(target_hdfs_path_dir, local_dir=\"~/data/full_hh_rlhf/rl\", local_dataset_path=None):\n    if local_dataset_path is not None:\n        dataset = load_dataset(local_dataset_path)\n    else:\n        dataset = load_dataset(\"Dahoas/full-hh-rlhf\")\n    train_dataset = dataset[\"train\"]\n\n    data_source = \"Dahoas/full-hh-rlhf\"\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            prompt = example.pop(\"prompt\")\n            response = example.pop(\"response\")\n\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [{\"role\": \"user\", \"content\": prompt}],\n                \"ability\": \"alignment\",\n                \"reward_model\": {\n                    \"style\": \"model\",\n                    \"ground_truth\": response,  # should not be used\n                },\n                \"extra_info\": {\"split\": split, \"index\": idx},\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    local_dir = os.path.expanduser(local_dir)\n    local_path = os.path.join(local_dir, \"train.parquet\")\n    train_dataset.to_parquet(local_path)\n\n    if target_hdfs_path_dir is not None:\n        hdfs_dir = target_hdfs_path_dir + \"/\" + \"train.parquet\"\n        makedirs(hdfs_dir)\n\n        copy(local_path, hdfs_dir)\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--split\", type=str, choices=[\"sft\", \"rm\", \"rl\"], required=True)\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", type=str, required=False, default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\",\n        type=str,\n        default=\"~/data/full_hh_rlhf\",\n        help=\"The save directory for the preprocessed dataset.\",\n    )\n\n    args = parser.parse_args()\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    if args.split == \"sft\":\n        generate_sft_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path)\n    elif args.split == \"rm\":\n        generate_rm_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path)\n    elif args.split == \"rl\":\n        generate_rl_dataset(args.hdfs_dir, os.path.join(local_save_dir, args.split), args.local_dataset_path)\n    else:\n        raise NotImplementedError\n"
  },
  {
    "path": "examples/data_preprocess/geo3k.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the Geometry3k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None)\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/geo3k\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"hiyouga/geometry3k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(\n            local_dataset_path,\n        )\n    else:\n        dataset = datasets.load_dataset(\n            data_source,\n        )\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = (\n        r\"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. \"\n        r\"The reasoning process MUST BE enclosed within <think> </think> tags. \"\n        r\"The final answer MUST BE put in \\boxed{}.\"\n    )\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            problem = example.pop(\"problem\")\n            prompt = problem + \" \" + instruction_following\n            answer = example.pop(\"answer\")\n            images = example.pop(\"images\")\n\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [\n                    {\n                        \"role\": \"user\",\n                        \"content\": prompt,\n                    }\n                ],\n                \"images\": images,\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": answer},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer,\n                    \"question\": problem,\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True, num_proc=8)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True, num_proc=8)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/geo3k_multiturn_w_tool.py",
    "content": "# Copyright 2023-2025 SGLang Team\n# Copyright Amazon.com, Inc. or its affiliates.\n# Copyright 2025 Reallm Labs Ltd. or its affiliates\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nPreprocess the Geometry3k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\",\n        default=\"~/data/geo3k_multiturn_w_tool\",\n        help=\"The save directory for the preprocessed dataset.\",\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"hiyouga/geometry3k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path)\n    else:\n        dataset = datasets.load_dataset(data_source)\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = (\n        r\"You FIRST think about the reasoning process as an internal monologue and then provide the final answer. \"\n        r\"The reasoning process MUST BE enclosed within <think> </think> tags. \"\n        r\"The final answer MUST BE put in \\boxed{}.\"\n    )\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            problem = example.pop(\"problem\")\n            prompt = problem + \" \" + instruction_following\n            answer = example.pop(\"answer\")\n            images = example.pop(\"images\")\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [\n                    {\n                        \"role\": \"system\",\n                        \"content\": (\n                            \"You are a math expert. You are given a question and you need to solve it step by step. \"\n                            \"Reasoning step by step before any tool call. \"\n                            \"You should use the `calc_geo3k_reward` tool after step by step solving the question, \"\n                            \"before generate final answer at least once and refine your answer if necessary. \"\n                        ),\n                    },\n                    {\n                        \"role\": \"user\",\n                        \"content\": prompt,\n                    },\n                ],\n                \"images\": images,\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": answer},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer,\n                    \"question\": problem,\n                    \"need_tools_kwargs\": True,\n                    \"tools_kwargs\": {\n                        \"calc_geo3k_reward\": {\n                            \"create_kwargs\": {\"ground_truth\": answer},\n                            # \"execute_kwargs\": {},\n                            # \"calc_reward_kwargs\": {},\n                            # \"release_kwargs\": {},\n                        },\n                    },\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True, num_proc=8)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True, num_proc=8)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/gsm8k.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the GSM8k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef extract_solution(solution_str):\n    solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n    assert solution is not None\n    final_solution = solution.group(0)\n    final_solution = final_solution.split(\"#### \")[1].replace(\",\", \"\")\n    return final_solution\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/gsm8k\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"openai/gsm8k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"main\")\n    else:\n        dataset = datasets.load_dataset(data_source, \"main\")\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = 'Let\\'s think step by step and output the final answer after \"####\".'\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question_raw = example.pop(\"question\")\n\n            question = question_raw + \" \" + instruction_following\n\n            answer_raw = example.pop(\"answer\")\n            solution = extract_solution(answer_raw)\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [\n                    {\n                        \"role\": \"user\",\n                        \"content\": question,\n                    }\n                ],\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": solution},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer_raw,\n                    \"question\": question_raw,\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/gsm8k_multiturn_sft.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the GSM8k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef extract_solution(solution_str):\n    solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n    assert solution is not None\n    final_solution = solution.group(0)\n    final_solution = final_solution.split(\"#### \")[1].replace(\",\", \"\")\n    return final_solution\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/gsm8k_sft\", help=\"The save directory for the preprocessed dataset.\"\n    )\n    parser.add_argument(\"--hdfs_dir\", default=None)\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"openai/gsm8k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"main\")\n    else:\n        dataset = datasets.load_dataset(data_source, \"main\")\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = 'Let\\'s think step by step and output the final answer after \"####\".'\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question_raw = example.pop(\"question\")\n\n            question = question_raw + \" \" + instruction_following\n\n            answer_raw = example.pop(\"answer\")\n            data = {\n                \"messages\": [\n                    {\n                        \"role\": \"user\",\n                        \"content\": question,\n                    },\n                    {\n                        \"role\": \"assistant\",\n                        \"content\": answer_raw,\n                    },\n                ],\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    local_save_dir = os.path.expanduser(local_save_dir)\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/gsm8k_multiturn_w_interaction.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the GSM8k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef extract_solution(solution_str):\n    solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n    assert solution is not None\n    final_solution = solution.group(0)\n    final_solution = final_solution.split(\"#### \")[1].replace(\",\", \"\")\n    return final_solution\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/gsm8k\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"openai/gsm8k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"main\")\n    else:\n        dataset = datasets.load_dataset(data_source, \"main\")\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = \"Let's think step by step and output the final answer after `####`.\"\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question_raw = example.pop(\"question\")\n\n            question = question_raw + \" \" + instruction_following\n\n            answer_raw = example.pop(\"answer\")\n            solution = extract_solution(answer_raw)\n            data = {\n                \"data_source\": data_source,\n                \"agent_name\": \"tool_agent\",\n                \"prompt\": [\n                    {\n                        \"role\": \"system\",\n                        \"content\": (\n                            \"You are a math expert. You are given a question and you need to solve it step by step. \"\n                            \"You should rethinking carefully if user point out your answer is wrong. \"\n                            \"Put your final answer in the format of `#### <answer>`.\"\n                        ),\n                    },\n                    {\n                        \"role\": \"user\",\n                        \"content\": question,\n                    },\n                ],\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": solution},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer_raw,\n                    \"question\": question_raw,\n                    \"interaction_kwargs\": {\n                        \"name\": \"gsm8k\",\n                        \"query\": question,\n                        \"ground_truth\": solution,\n                    },\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/gsm8k_multiturn_w_tool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the GSM8k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef extract_solution(solution_str):\n    solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n    assert solution is not None\n    final_solution = solution.group(0)\n    final_solution = final_solution.split(\"#### \")[1].replace(\",\", \"\")\n    return final_solution\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/gsm8k\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"openai/gsm8k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"main\")\n    else:\n        dataset = datasets.load_dataset(data_source, \"main\")\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = \"Let's think step by step and output the final answer after `####`.\"\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question_raw = example.pop(\"question\")\n\n            question = question_raw + \" \" + instruction_following\n\n            answer_raw = example.pop(\"answer\")\n            solution = extract_solution(answer_raw)\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [\n                    {\n                        \"role\": \"system\",\n                        \"content\": (\n                            \"You are a math expert. You are given a question and you need to solve it step by step. \"\n                            \"Reasoning step by step before any tool call. \"\n                            \"You should use the `calc_gsm8k_reward` tool after step by step solving the question, \"\n                            \"before generate final answer at least once and refine your answer if necessary. \"\n                            \"Put your final answer in the format of `#### <answer>`.\"\n                        ),\n                    },\n                    {\n                        \"role\": \"user\",\n                        \"content\": question,\n                    },\n                ],\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": solution},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer_raw,\n                    \"question\": question_raw,\n                    \"need_tools_kwargs\": True,\n                    \"tools_kwargs\": {\n                        \"calc_gsm8k_reward\": {\n                            \"create_kwargs\": {\"ground_truth\": solution},\n                            # \"execute_kwargs\": {},\n                            # \"calc_reward_kwargs\": {},\n                            # \"release_kwargs\": {},\n                        },\n                    },\n                    \"interaction_kwargs\": {\n                        \"query\": question,\n                        \"ground_truth\": solution,\n                    },\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/gsm8k_tool_agent_loop.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the GSM8k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef extract_solution(solution_str):\n    solution = re.search(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n    assert solution is not None\n    final_solution = solution.group(0)\n    final_solution = final_solution.split(\"#### \")[1].replace(\",\", \"\")\n    return final_solution\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/gsm8k\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"openai/gsm8k\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path, \"main\")\n    else:\n        dataset = datasets.load_dataset(data_source, \"main\")\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = \"Let's think step by step and output the final answer after `####`.\"\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question_raw = example.pop(\"question\")\n\n            question = question_raw + \" \" + instruction_following\n\n            answer_raw = example.pop(\"answer\")\n            solution = extract_solution(answer_raw)\n            data = {\n                \"data_source\": data_source,\n                \"agent_name\": \"tool_agent\",\n                \"prompt\": [\n                    {\n                        \"role\": \"system\",\n                        \"content\": (\n                            \"You are a math expert. You are given a question and you need to solve it step by step. \"\n                            \"Reasoning step by step before any tool call. \"\n                            \"You should use the `calc_gsm8k_reward` tool after step by step solving the question, \"\n                            \"before generate final answer at least once and refine your answer if necessary. \"\n                            \"Put your final answer in the format of `#### <answer>`.\"\n                        ),\n                    },\n                    {\n                        \"role\": \"user\",\n                        \"content\": question,\n                    },\n                ],\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": solution},\n                \"extra_info\": {\n                    \"split\": split,\n                    \"index\": idx,\n                    \"answer\": answer_raw,\n                    \"question\": question_raw,\n                    \"need_tools_kwargs\": True,\n                    \"tools_kwargs\": {\n                        \"calc_gsm8k_reward\": {\n                            \"create_kwargs\": {\"ground_truth\": solution},\n                            # \"execute_kwargs\": {},\n                            # \"calc_reward_kwargs\": {},\n                            # \"release_kwargs\": {},\n                        },\n                    },\n                    \"interaction_kwargs\": {\n                        \"query\": question,\n                        \"ground_truth\": solution,\n                    },\n                },\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/hellaswag.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess Hellaswag dataset.\n\n\"\"\"\n\nimport argparse\nimport os\nimport re\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef preprocess(text):\n    text = text.strip()\n    # NOTE: Brackets are artifacts of the WikiHow dataset portion of HellaSwag.\n    text = text.replace(\" [title]\", \". \")\n    text = re.sub(\"\\\\[.*?\\\\]\", \"\", text)\n    text = text.replace(\"  \", \" \")\n    return text\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None, help=\"The save directory for the preprocessed dataset.\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/hellaswag\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"Rowan/hellaswag\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(local_dataset_path)\n    else:\n        dataset = datasets.load_dataset(data_source, trust_remote_code=True)\n\n    train_dataset = dataset[\"train\"]\n    val_dataset = dataset[\"validation\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction = \"Please complete the following sentence.\\n\"\n\n    def make_map_fn(split):\n        def process_fn(doc, idx):\n            ctx = doc[\"ctx_a\"] + \" \" + doc[\"ctx_b\"].capitalize()\n            query = preprocess(doc[\"activity_label\"] + \": \" + ctx)\n            choices = [preprocess(ending) for ending in doc[\"endings\"]]\n            gold = int(doc[\"label\"])\n\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [{\"role\": \"user\", \"content\": query}],\n                \"ability\": \"nlp\",\n                \"reward_model\": {\n                    \"style\": \"model\",\n                    \"eval\": \"multiple_choice\",  # using loglikelihood\n                    \"ground_truth\": gold,\n                    \"choices\": choices,\n                },\n                \"extra_info\": {\"split\": split, \"index\": idx},\n            }\n            return data\n\n        return process_fn\n\n    # filter data that doesn't have a label\n    train_dataset = train_dataset.filter(lambda x: len(x[\"label\"]) > 0)\n    val_dataset = val_dataset.filter(lambda x: len(x[\"label\"]) > 0)\n    test_dataset = test_dataset.filter(lambda x: len(x[\"label\"]) > 0)\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    val_dataset = val_dataset.map(function=make_map_fn(\"validation\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    val_dataset.to_parquet(os.path.join(local_save_dir, \"validation.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/math_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the MATH-lighteval dataset to parquet format\n\"\"\"\n\nimport argparse\nimport json\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\nfrom verl.utils.reward_score.math_reward import last_boxed_only_string, remove_boxed\n\n\ndef extract_solution(solution_str):\n    return remove_boxed(last_boxed_only_string(solution_str))\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None)\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/math\", help=\"The save directory for the preprocessed dataset.\"\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    # 'lighteval/MATH' is no longer available on huggingface.\n    # Use mirror repo: DigitalLearningGmbH/MATH-lighteval\n    data_source = \"DigitalLearningGmbH/MATH-lighteval\"\n    print(f\"Loading the {data_source} dataset from huggingface...\", flush=True)\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(\n            local_dataset_path,\n        )\n    else:\n        dataset = datasets.load_dataset(\n            data_source,\n        )\n\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    instruction_following = \"Let's think step by step and output the final answer within \\\\boxed{}.\"\n\n    # add a row to each data item that represents a unique id\n    def make_map_fn(split):\n        def process_fn(example, idx):\n            question = example.pop(\"problem\")\n\n            question = question + \" \" + instruction_following\n\n            answer = example.pop(\"solution\")\n            solution = extract_solution(answer)\n            data = {\n                \"data_source\": data_source,\n                \"prompt\": [{\"role\": \"user\", \"content\": question}],\n                \"ability\": \"math\",\n                \"reward_model\": {\"style\": \"rule\", \"ground_truth\": solution},\n                \"extra_info\": {\"split\": split, \"index\": idx},\n            }\n            return data\n\n        return process_fn\n\n    train_dataset = train_dataset.map(function=make_map_fn(\"train\"), with_indices=True)\n    test_dataset = test_dataset.map(function=make_map_fn(\"test\"), with_indices=True)\n\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    local_dir = os.path.expanduser(local_save_dir)\n    hdfs_dir = args.hdfs_dir\n\n    train_dataset.to_parquet(os.path.join(local_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_dir, \"test.parquet\"))\n    # Save one example as JSON for reference\n    example = train_dataset[0]\n    with open(os.path.join(local_dir, \"train_example.json\"), \"w\") as f:\n        json.dump(example, f, indent=2)\n    example = test_dataset[0]\n    with open(os.path.join(local_dir, \"test_example.json\"), \"w\") as f:\n        json.dump(example, f, indent=2)\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n\n        copy(src=local_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/multiturn.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nCreate a simple multi-turn dataset for testing\n\"\"\"\n\nimport argparse\nimport os\n\nimport pandas as pd\n\n\ndef main():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=\"~/data/multiturn\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    args = parser.parse_args()\n\n    # Create example conversations\n    conversations = []\n\n    # Conversation 1\n    conversations.append(\n        {\n            \"messages\": [\n                {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n                {\"role\": \"user\", \"content\": \"What is the capital of France?\"},\n                {\"role\": \"assistant\", \"content\": \"The capital of France is Paris.\"},\n                {\"role\": \"user\", \"content\": \"And what about Germany?\"},\n                {\"role\": \"assistant\", \"content\": \"The capital of Germany is Berlin.\"},\n            ]\n        }\n    )\n\n    # Conversation 2\n    conversations.append(\n        {\n            \"messages\": [\n                {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n                {\"role\": \"user\", \"content\": \"Can you explain quantum computing?\"},\n                {\n                    \"role\": \"assistant\",\n                    \"content\": \"Quantum computing is a type of computing that uses quantum-mechanical phenomena, \"\n                    \"such as superposition and entanglement, to perform operations on data.\",\n                },\n                {\"role\": \"user\", \"content\": \"How is it different from classical computing?\"},\n                {\n                    \"role\": \"assistant\",\n                    \"content\": \"Classical computing uses bits that are either 0 or 1, while quantum computing uses \"\n                    \"quantum bits or qubits that can exist in multiple states simultaneously due to superposition.\",\n                },\n            ]\n        }\n    )\n\n    # Conversation 3\n    conversations.append(\n        {\n            \"messages\": [\n                {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n                {\"role\": \"user\", \"content\": \"Write a simple Python function to calculate factorial.\"},\n                {\n                    \"role\": \"assistant\",\n                    \"content\": (\n                        \"```python\\ndef factorial(n):\\n    if n == 0 or n == 1:\\n        return 1\\n    else:\\n        \"\n                        \"return n * factorial(n-1)\\n```\\n\\nThis is a recursive function to calculate the \"\n                        \"factorial of a number.\"\n                    ),\n                },\n                {\"role\": \"user\", \"content\": \"Can you make it iterative instead?\"},\n                {\n                    \"role\": \"assistant\",\n                    \"content\": (\n                        \"```python\\ndef factorial(n):\\n    result = 1\\n    for i in range(1, n+1):\\n        \"\n                        \"result *= i\\n    return result\\n```\\n\\nThis is an iterative version of the factorial function.\"\n                    ),\n                },\n            ]\n        }\n    )\n\n    # Create train and test datasets\n    train_data = conversations[:2]  # First 2 conversations for training\n    test_data = conversations[2:]  # Last conversation for testing\n\n    # Create output directory\n    local_dir = os.path.expanduser(args.local_dir)\n    os.makedirs(local_dir, exist_ok=True)\n\n    # Save to parquet files\n    train_df = pd.DataFrame(train_data)\n    test_df = pd.DataFrame(test_data)\n\n    train_df.to_parquet(os.path.join(local_dir, \"train.parquet\"))\n    test_df.to_parquet(os.path.join(local_dir, \"test.parquet\"))\n\n    # Handle HDFS if specified\n    if args.hdfs_dir is not None:\n        try:\n            from verl.utils.hdfs_io import copy, makedirs\n\n            makedirs(args.hdfs_dir)\n            copy(src=local_dir, dst=args.hdfs_dir)\n        except ImportError:\n            print(\"Warning: HDFS support not available. Skipping HDFS copy.\")\n\n    # Print statistics\n    print(f\"Train dataset size: {len(train_df)}\")\n    print(f\"Test dataset size: {len(test_df)}\")\n    print(f\"Data saved to {local_dir}\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "examples/data_preprocess/pokemon.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\"\"\"\nPreprocess the llamafactory/pokemon-gpt4o-captions dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\n\nimport datasets\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=None)\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--local_dataset_path\", default=None, help=\"The local path to the raw dataset, if it exists.\")\n    parser.add_argument(\n        \"--local_save_dir\",\n        default=\"~/data/pokemon-gpt4o-captions\",\n        help=\"The save directory for the preprocessed dataset.\",\n    )\n\n    args = parser.parse_args()\n    local_dataset_path = args.local_dataset_path\n\n    data_source = \"llamafactory/pokemon-gpt4o-captions\"\n\n    if local_dataset_path is not None:\n        dataset = datasets.load_dataset(\n            local_dataset_path,\n        )\n    else:\n        dataset = datasets.load_dataset(\n            data_source,\n        )\n\n    def map_fn(row: dict):\n        messages = []\n        conversation = row.pop(\"conversations\")\n        for conv in conversation:\n            if conv[\"from\"] == \"gpt\":\n                role = \"assistant\"\n            elif conv[\"from\"] == \"human\":\n                role = \"user\"\n            else:\n                raise ValueError(f\"Unknown role: {conv['from']}\")\n            messages.append(\n                {\n                    \"role\": role,\n                    \"content\": conv[\"value\"],\n                }\n            )\n\n        row[\"messages\"] = messages\n        return row\n\n    dataset = dataset[\"train\"].map(map_fn, num_proc=16)\n    dataset = dataset.train_test_split(test_size=0.1)\n    train_dataset = dataset[\"train\"]\n    test_dataset = dataset[\"test\"]\n\n    hdfs_dir = args.hdfs_dir\n    local_save_dir = args.local_dir\n    if local_save_dir is not None:\n        print(\"Warning: Argument 'local_dir' is deprecated. Please use 'local_save_dir' instead.\")\n    else:\n        local_save_dir = args.local_save_dir\n\n    train_dataset.to_parquet(os.path.join(local_save_dir, \"train.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_save_dir, \"test.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n        copy(src=local_save_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/data_preprocess/preprocess_search_r1_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\r\n# Copyright 2023-2024 SGLang Team\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\nimport argparse\r\nimport logging\r\nimport os\r\nimport tempfile\r\n\r\nimport pandas as pd\r\nfrom huggingface_hub import hf_hub_download\r\nfrom huggingface_hub.utils import EntryNotFoundError\r\n\r\nfrom verl.utils.hdfs_io import copy, makedirs\r\n\r\n# Setup logging\r\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(levelname)s - %(message)s\")\r\nlogger = logging.getLogger(__name__)\r\n\r\n# Configuration constants\r\nDEFAULT_SYSTEM_CONTENT = \"You are a helpful and harmless assistant.\"\r\nDEFAULT_USER_CONTENT_PREFIX = (\r\n    \"Answer the given question. You must conduct reasoning inside <think> and </think> \"\r\n    \"first every time you get new information. After reasoning, if you find you lack \"\r\n    \"some knowledge, you can call a search engine by <tool_call> query </tool_call> \"\r\n    \"and it will return the top searched results between <tool_response> and \"\r\n    \"</tool_response>. You can search as many times as your want. If you find no \"\r\n    \"further external knowledge needed, you can directly provide the answer inside \"\r\n    \"<answer> and </answer>, without detailed illustrations. For example, \"\r\n    \"<answer> Beijing </answer>. Question: \"\r\n)\r\n\r\n\r\ndef process_single_row(row, current_split_name, row_index):\r\n    \"\"\"\r\n    Process a single row of data for SearchR1-like format.\r\n\r\n    Args:\r\n        row: DataFrame row containing the original data\r\n        current_split_name: Name of the current split (train/test)\r\n        row_index: Index of the row in the DataFrame\r\n\r\n    Returns:\r\n        pd.Series: Processed row data in the required format\r\n    \"\"\"\r\n    question = row.get(\"question\", \"\")\r\n\r\n    # Build prompt structure\r\n    user_content = user_content_prefix.rstrip(\"\\n\") + question\r\n    prompt = [{\"role\": \"system\", \"content\": system_content}, {\"role\": \"user\", \"content\": user_content}]\r\n\r\n    # Extract ground truth from reward_model or fallback to golden_answers\r\n    reward_model_data = row.get(\"reward_model\")\r\n    if isinstance(reward_model_data, dict) and \"ground_truth\" in reward_model_data:\r\n        ground_truth = reward_model_data.get(\"ground_truth\")\r\n    else:\r\n        ground_truth = row.get(\"golden_answers\", [])\r\n\r\n    # Process data source\r\n    data_source_tagged = \"searchR1_\" + str(row.get(\"data_source\", \"\"))\r\n\r\n    # Build tools kwargs structure\r\n    tools_kwargs = {\r\n        \"search\": {\r\n            \"create_kwargs\": {\"ground_truth\": ground_truth, \"question\": question, \"data_source\": data_source_tagged}\r\n        }\r\n    }\r\n\r\n    # Build complete extra_info structure\r\n    extra_info = {\r\n        \"index\": row_index,\r\n        \"need_tools_kwargs\": True,\r\n        \"question\": question,\r\n        \"split\": current_split_name,\r\n        \"tools_kwargs\": tools_kwargs,\r\n    }\r\n\r\n    return pd.Series(\r\n        {\r\n            \"data_source\": data_source_tagged,\r\n            \"prompt\": prompt,\r\n            \"ability\": row.get(\"ability\"),\r\n            \"reward_model\": reward_model_data,\r\n            \"extra_info\": extra_info,\r\n            \"metadata\": row.get(\"metadata\"),\r\n        }\r\n    )\r\n\r\n\r\ndef main():\r\n    local_save_dir = os.path.expanduser(args.local_dir)\r\n    os.makedirs(local_save_dir, exist_ok=True)\r\n\r\n    processed_files = []\r\n\r\n    # Download and process files using temporary directory\r\n    with tempfile.TemporaryDirectory() as tmp_download_dir:\r\n        for split in [\"train\", \"test\"]:\r\n            parquet_filename = f\"{split}.parquet\"\r\n            logger.info(f\"Processing {split} split...\")\r\n\r\n            try:\r\n                # Download Parquet file from HuggingFace\r\n                logger.info(f\"Downloading {parquet_filename} from {args.hf_repo_id}\")\r\n                local_parquet_filepath = hf_hub_download(\r\n                    repo_id=args.hf_repo_id,\r\n                    filename=parquet_filename,\r\n                    repo_type=\"dataset\",\r\n                    local_dir=tmp_download_dir,\r\n                    local_dir_use_symlinks=False,\r\n                )\r\n\r\n                # Load and process Parquet file\r\n                df_raw = pd.read_parquet(local_parquet_filepath)\r\n                logger.info(f\"Loaded {len(df_raw)} rows from {parquet_filename}\")\r\n\r\n                def apply_process_row(row, split_name=split):\r\n                    return process_single_row(row, current_split_name=split_name, row_index=row.name)\r\n\r\n                df_processed = df_raw.apply(apply_process_row, axis=1)\r\n\r\n                # Save processed DataFrame\r\n                output_file_path = os.path.join(local_save_dir, f\"{split}.parquet\")\r\n                df_processed.to_parquet(output_file_path, index=False)\r\n                logger.info(f\"Saved {len(df_processed)} processed rows to {output_file_path}\")\r\n                processed_files.append(output_file_path)\r\n\r\n            except EntryNotFoundError:\r\n                logger.warning(f\"{parquet_filename} not found in repository {args.hf_repo_id}\")\r\n            except Exception as e:\r\n                logger.error(f\"Error processing {split} split: {e}\")\r\n\r\n    if not processed_files:\r\n        logger.warning(\"No data was processed or saved\")\r\n        return\r\n\r\n    logger.info(f\"Successfully processed {len(processed_files)} files to {local_save_dir}\")\r\n\r\n    # Copy to HDFS if specified\r\n    if args.hdfs_dir:\r\n        try:\r\n            makedirs(args.hdfs_dir)\r\n            copy(src=local_save_dir, dst=args.hdfs_dir)\r\n            logger.info(f\"Successfully copied files to HDFS: {args.hdfs_dir}\")\r\n        except Exception as e:\r\n            logger.error(f\"Error copying files to HDFS: {e}\")\r\n\r\n\r\nif __name__ == \"__main__\":\r\n    parser = argparse.ArgumentParser(description=\"Download Search-R1 from HuggingFace, process, and save to Parquet.\")\r\n    parser.add_argument(\r\n        \"--hf_repo_id\", default=\"PeterJinGo/nq_hotpotqa_train\", help=\"HuggingFace dataset repository ID.\"\r\n    )\r\n    parser.add_argument(\r\n        \"--local_dir\",\r\n        default=\"~/data/searchR1_processed_direct\",\r\n        help=\"Local directory to save the processed Parquet files.\",\r\n    )\r\n    parser.add_argument(\"--hdfs_dir\", default=None, help=\"Optional HDFS directory to copy the Parquet files to.\")\r\n\r\n    args = parser.parse_args()\r\n\r\n    # System and user content configuration\r\n    system_content = DEFAULT_SYSTEM_CONTENT\r\n    user_content_prefix = DEFAULT_USER_CONTENT_PREFIX\r\n\r\n    main()\r\n"
  },
  {
    "path": "examples/dppo_trainer/dppo.md",
    "content": "# Divergence Proximal Policy Optimization (DPPO)\n\n\n<div align=\"center\">\n\n## Rethinking the Trust Region in LLM Reinforcement Learning\n\n[![Paper](https://img.shields.io/badge/paper-A42C25?style=for-the-badge&logo=arxiv&logoColor=white )](https://arxiv.org/pdf/2602.04879)\n[![Github](https://img.shields.io/badge/Stable_RL-000000?style=for-the-badge&logo=github&logoColor=000&logoColor=white)](https://github.com/sail-sg/Stable-RL)\n[![Twitter](https://img.shields.io/badge/Twitter-%23000000.svg?style=for-the-badge&logo=twitter&logoColor=white)](https://x.com/QPHutu/status/2019435642539897303)\n\n</div>\n\n\n## ✨Getting started\n\n1. Prepare the datasets by running [prepare_dapo_data.sh](https://github.com/verl-project/verl-recipe/blob/3490a22a0a3adeb7e4787fe70b1060b642efbae4/dapo/prepare_dapo_data.sh):\n\n```bash\nbash prepare_dapo_data.sh # This downloads the datasets to ${HOME}/verl/data by default\n```\n\n2. Prepare the model:\n\n```bash\nhf download Qwen/Qwen3-30B-A3B-Base --local-dir ${HOME}/verl/models/Qwen3-30B-A3B-Base\n```\n\n3. Run the script:\n\n```bash\n# run DPPO-Binary-KL\nLOSS_MODE=dppo_kl bash examples/dppo_trainer/run_qwen30b_dppo.sh\n\n# run DPPO-Binary-TV\nLOSS_MODE=dppo_tv bash examples/dppo_trainer/run_qwen30b_dppo.sh\n\n# run GRPO baseline\nLOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.2 bash examples/dppo_trainer/run_qwen30b_dppo.sh\n# or GRPO with clip higher\nLOSS_MODE=vanilla CLIP_LOW=0.2 CLIP_HIGH=0.28 bash examples/dppo_trainer/run_qwen30b_dppo.sh\n```\n\n## 📖Introduction\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/ppo_vs_dppo.jpg?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nComparison of **PPO** and the proposed **DPPO** (the Binary-TV variant). **(Left)** The surrogate objective and corresponding masks for PPO and DPPO. PPO (and variants like GRPO) employs a heuristic mask based on the probability ratio. In contrast, DPPO utilizes a more principled mask based on a direct approximation of policy divergence (e.g., Total Variation), ensuring updates stay within a theoretically grounded trust region. **(Right)** Experimental results on the AIME24 using Qwen3-30B-A3B-Base. DPPO significantly outperforms GRPO baselines, achieving superior training stability and final performance even without rollout routing replay (R3).\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/sanity_test.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nDPPO variants achieve stable training while controlling the training-inference mismatch at a low level. In contrast, methods without a trust region (PG-IS, CISPO) or with a misspecified one (MiniRL) suffer from growing mismatch and eventual collapse.\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/moe_prob_ratio_tv.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nThe plots show numerical differences between a training and an inference engine for Qwen3-30B-A3B-Base with identical parameters. **(Left)** The probability ratio (used in PPO) is highly volatile for low-probability tokens. **(Right)** In contrast, the TV divergence is more stable. This highlights a key flaw of PPO's clipping mechanism: it **over-penalizes low-probability tokens**, which can slow down learning; and **under-penalizes high-probability tokens**, which can permit large, destabilizing updates.\n\n\n<div align=\"left\">\n  <img src=\"https://github.com/sail-sg/Stable-RL/blob/main/figures/clipped_tokens.png?raw=true\" alt=\"issue\" style=\"width: 96%; height: auto;\">\n</div>\n\nThe most frequently clipped tokens (by GRPO) are important to the reasoning task! \nThey are dominated by:\n- numbers, like 1, 4\n- mathematical symbols, like +, -, =\n- reasoning and structural Words: Wait, Thus, Next\n\n## Top-K divergence approximation\n\nWe only implement the DPPO-Binary-TV/DPPO-Binary-KL here due to their simplicity.\n\nFor the TopK divergence approximation, please refer to the [the original repo](https://github.com/sail-sg/Stable-RL) for a complete implementation.\n\n## Citation\nIf you find our works useful for your research, please consider citing:\n\n```bibtex\n@article{qi2026dppo,\n  title={Rethinking the Trust Region in LLM Reinforcement Learning},\n  author={Qi, Penghui and Zhou, Xiangxin and Liu, Zichen and Pang, Tianyu and Du, Chao and Lin, Min and Lee, Wee Sun},\n  journal={arXiv preprint arXiv:2602.04879},\n  year={2026}\n}\n```\n\n## 🌻Acknowledgement\nWe implement our reinforcement learning algorithm extending from [verl](https://github.com/volcengine/verl). We utilize [vLLM](https://github.com/vllm-project/vllm) and [sglang](https://github.com/sgl-project/sglang) for inference. Our models are trained primarily on [Qwen3 family](https://huggingface.co/collections/Qwen/qwen3). Our training data is built from [DAPO-MATH](https://huggingface.co/datasets/BytedTsinghua-SIA/DAPO-Math-17k). Thanks for their great contributions!\n"
  },
  {
    "path": "examples/dppo_trainer/run_qwen30b_dppo.sh",
    "content": "# run Qwen3-30B-A3B-Base on dapo-math-17k dataset\nset -x\n\n# ================================ DPPO Specific Parameters ===========================\n\n# Why from GRPO to DPPO?\n\"\"\"\nThe ratio clipping mechanism in GRPO/PPO is structurally ill-suited due to the large, \nlong-tailed vocabularies inherent to LLMs. It over-penalizes low-probability tokens \nand under-penalizes high-probability ones, leading to training inefficiency and \ninstability. For example, increasing a rare token’s probability from 1e−5 to 1e−3 \ngenerates a massive ratio of 100 that triggers clipping, even though the actual \ndivergence is negligible. Conversely, small ratio changes on high-probability tokens \ncan make catastrophic shifts in probability mass (e.g., a drop from 0.99 to 0.8), yet \nit often remains unpenalized by the clipping mechanism.\n\nDPPO addresses this issue by using a divergence-based clipping mechanism, achieving \nsuperior training stability and final performance compared to existing methods.\n\nDPPO paper: https://arxiv.org/pdf/2602.04879\n\"\"\"\n\nLOSS_MODE=${LOSS_MODE:-\"dppo_tv\"}\n\nif [[ $LOSS_MODE == \"dppo_kl\" ]]; then\n    # The KL divergence threshold for DPPO.\n    clip_ratio=0.05\n    clip_ratio_low=${CLIP_LOW:-0.05}\n    clip_ratio_high=${CLIP_HIGH:-0.05}\nelif [[ $LOSS_MODE == \"dppo_tv\" ]]; then\n    # The TV divergence threshold for DPPO.\n    clip_ratio=0.15\n    clip_ratio_low=${CLIP_LOW:-0.15}\n    clip_ratio_high=${CLIP_HIGH:-0.15}\nelif [[ $LOSS_MODE == \"vanilla\" ]]; then\n    # GRPO baseline\n    clip_ratio=0.2\n    clip_ratio_low=${CLIP_LOW:-0.2}\n    clip_ratio_high=${CLIP_HIGH:-0.28}\nelse\n    echo \"Invalid loss mode: $LOSS_MODE\"\n    exit 1\nfi\n\n# Disable dual-clip PPO and TIS for a fair comparison between GRPO and DPPO.\nclip_ratio_c=10000.0\n\n# ===================================== Algorithm =====================================\nadv_estimator=grpo\n\n# We recommand directly clipping the ratio/divergence with respect to the original \n# rollout policy (implemented by bypass_mode=True), instead of the recomputed one. \n# This can not only save the computation cost, but also improve the training stability \n# for both GRPO and DPPO by controlling the training-inference mismatch at a low level.\n# See Section 5.2 in https://arxiv.org/pdf/2602.04879 for more details.\nbypass_mode=True\n\n# We recommand using Dr.GRPO to remove the length and difficulty bias in original GRPO.\n# See Section 3.1 in https://arxiv.org/pdf/2503.20783 for more details.\nnorm_adv_by_std_in_grpo=False               # remove the difficulty bias\nloss_agg_mode=\"seq-mean-token-sum-norm\"     # remove the length bias\n\n# reference policy\nuse_kl_in_reward=False\nkl_coef=0.001\nuse_kl_loss=False\nkl_loss_coef=0.001\n\nactor_lr=1e-6\ncritic_lr=2e-6\ngae_gamma=1.0\ngae_lam=0.95\ncritic_warmup=0\n\n\n# ================================== Data/Model/Config =================================\n\n# Node Info\nNNODES=${NNODES:-2}\n\n# wandb\nbackend=megatron # fsdp, fsdp2, megatron\nproject_name=Qwen3-30B-A3B-Base-dapo-math-17k\nexperiment_name=\"${backend}-${NNODES}nodes-${LOSS_MODE}-low${clip_ratio_low}-high${clip_ratio_high}\"\n\n# Paths\nDATA_ROOT=${DATA_ROOT:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${DATA_ROOT}/ckpts/${project_name}/${experiment_name}\"}\nMODEL_PATH=${MODEL_PATH:-\"${DATA_ROOT}/models/Qwen3-30B-A3B-Base\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${DATA_ROOT}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${DATA_ROOT}/data/aime-2024.parquet\"}\n\n\nactor_model_path=$MODEL_PATH\ncritic_model_path=$MODEL_PATH\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=False\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\ntrain_batch_size=256\nppo_mini_batch_size=32\nn_resp_per_prompt=16\nn_resp_per_prompt_val=1\n\n# ===================================== Training ======================================\nactor_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 1))\ncritic_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 1))\n\n# FSDP parallelism config\nUSP_SIZE=4\nACTOR_FSDP_CONFIG=\"\n    actor_rollout_ref.actor.fsdp_config.strategy=$backend \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$USP_SIZE\"\n\n# Megatron parallelism config\nTP_SIZE=2\nCP_SIZE=1\nPP_SIZE=1\nVPP_SIZE=null\nEP_SIZE=8\nETP_SIZE=1\nACTOR_MEGATRON_CONFIG=\"\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP_SIZE \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP_SIZE \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP_SIZE \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$VPP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP_SIZE \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True\"\n\n# Actor model config\nACTOR_CONFIG=\"\n    actor_rollout_ref.actor.optim.lr=$actor_lr \\\n    actor_rollout_ref.model.path=$actor_model_path \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.clip_ratio=$clip_ratio \\\n    actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \\\n    actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \\\n    actor_rollout_ref.actor.clip_ratio_c=$clip_ratio_c \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.calculate_entropy=True \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${LOSS_MODE} \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu\"\n\n# Critic model config\nCIRITC_CONFIG=\"\n    critic.optim.lr=$critic_lr \\\n    critic.model.path=$critic_model_path \\\n    critic.model.use_remove_padding=True \\\n    critic.ppo_max_token_len_per_gpu=$critic_max_token_len_per_gpu \\\n    critic.ulysses_sequence_parallel_size=$USP_SIZE\"\n\nCRITIC_FSDP_CONFIG=\"${ACTOR_FSDP_CONFIG//actor_rollout_ref.actor/critic.model}\"\nCRITIC_MEGATRON_CONFIG=\"${ACTOR_MEGATRON_CONFIG//actor_rollout_ref.actor/critic}\"\n\nif [[ $backend == \"megatron\" ]]; then\n    CONFIG_NAME=ppo_megatron_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_MEGATRON_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_MEGATRON_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nelse # fsdp, fsdp2\n    CONFIG_NAME=ppo_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_FSDP_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_FSDP_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nfi\n\n# ===================================== Inference =====================================\nrollout_name=vllm\nif [ \"$rollout_name\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\ninfer_tp=4\ninfer_dp=1\ninfer_ep=1\ngpu_memory_utilization=0.7\n\nROLLOUT_CONFIG=\"\n    actor_rollout_ref.rollout.name=$rollout_name \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \\\n    actor_rollout_ref.rollout.data_parallel_size=$infer_dp \\\n    actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \\\n    actor_rollout_ref.rollout.n=$n_resp_per_prompt \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val\"\n\n# ===================================== Reward =====================================\nREWARD_CONFIG=\"\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length}\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=./config \\\n    --config-name=$CONFIG_NAME \\\n    algorithm.adv_estimator=$adv_estimator \\\n    algorithm.use_kl_in_reward=$use_kl_in_reward \\\n    algorithm.kl_ctrl.kl_coef=$kl_coef \\\n    algorithm.gamma=$gae_gamma \\\n    algorithm.lam=$gae_lam \\\n    algorithm.rollout_correction.bypass_mode=$bypass_mode \\\n    algorithm.norm_adv_by_std_in_grpo=$norm_adv_by_std_in_grpo \\\n    data.train_files=\"$TRAIN_FILE\" \\\n    data.val_files=\"$TEST_FILE\" \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=$train_batch_size \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=False \\\n    data.filter_overlong_prompts_workers=64 \\\n    data.truncation='error' \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=$critic_warmup \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$experiment_name \\\n    trainer.default_local_dir=$CKPTS_DIR \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=$NNODES \\\n    trainer.val_before_train=False \\\n    trainer.log_val_generations=100 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=500 \\\n    $ACTOR_CONFIG \\\n    $CIRITC_CONFIG \\\n    $ROLLOUT_CONFIG \\\n    $REWARD_CONFIG\n"
  },
  {
    "path": "examples/fapo_trainer/README.md",
    "content": "<p align=\"center\">\n<h1 align=\"center\">FAPO: Flawed-Aware Policy Optimization for Efficient and Reliable Reasoning</h1>\n\n<p align=\"center\">\n    <a href=\"https://fapo-rl.github.io/\"><img alt=\"Project Page\" src=\"https://img.shields.io/badge/📒-Project Page-blue\"></a>\n    <a href=\"https://verl.readthedocs.io/en/latest/advance/reward_loop.html\"><img alt=\"Infra Design\" src=\"https://img.shields.io/badge/🏗️-Infra Design-teal\">\n    <a href=\"https://huggingface.co/collections/dyyyyyyyy/fapo\"><img alt=\"Resources\" src=\"https://img.shields.io/badge/🤗 HuggingFace-Data & Models-green\"></a>\n    <a href=\"https://arxiv.org/abs/2510.22543\"><img alt=\"Paper\" src=\"https://img.shields.io/badge/📄-Arxiv Paper-orange\"></a>\n    <a href=\"https://github.com/yyDing1/FAPO\"><img alt=\"Code\" src=\"https://img.shields.io/badge/💻-Code-blueviolet\"></a>\n</p>\n\nThis example include a runnable and fully reproducible example that demonstrates how to use reward model to optimize a policy.\n\n## Quick start\n\nFirst construct the training and evaluation datasets by:\n\n```bash\npython examples/fapo_trainer/prepare_data.py --local_dir ${RAY_DATA_HOME}/data/\n```\n\nOr you can directly use the data available [here](https://huggingface.co/datasets/dyyyyyyyy/FAPO-Reasoning-Dataset).\n\nTo integrate the GRM into the final training, we provide two options:\n\n1. **Colocate Mode:** The reward model is colocated with the trainer and runs synchronously.\n2. **Standalone Mode:** A separate resource pool is allocated to deploy the GenRM, which runs asynchronously.\n\nThe following list the most-relevant parameters in the config file:\n\n```yaml\nreward:\n  reward_model: \n    model_path: \"/path/to/your/reward_model\" # your reward model path\n    # whether to enable resource pool for the reward model\n    # True -> Standalone Mode, False -> Colocate Mode\n    enable_resource_pool: True\n    # the number of nodes to deploy the reward model\n    # only effective when enable_resource_pool is True\n    nnodes: 1\n    # the number of GPUs to deploy the reward model on each node\n    # only effective when enable_resource_pool is True\n    n_gpus_per_node: 8\n    # inference engine configs, similar to those in rollout configs\n    rollout:\n      # set to True in colocate mode, False in standalone mode\n      free_cache_engine: True\n      # ... (ommitted)\n\n  # customized reward function, where user should implement the invocation logic\n  # of the specified reward model (both generative and discriminative)\n  custom_reward_function:\n    path: null\n    name: compute_score\n  \n```\n\n![](https://github.com/yyDing1/verl-materials/blob/main/reward_loop.svg)\n\n## Choice 1: Colocate Reward Model\n\n```bash\ncd verl # Repo root\nexport RAY_ADDRESS=\"...\" # The Ray cluster address to connect to\nexport WORKING_DIR=\"${PWD}\" # The local directory to package to the Ray cluster\nexport NNODES=xxx\n\nbash examples/fapo_trainer/run_qwen_7b_rm_colocate.sh  # 7b fapo model\nbash examples/fapo_trainer/run_qwen_32b_rm_colocate.sh  # 32b fapo model\n```\n\n## Choice 2: Standalone Reward Model\n\n```bash\ncd verl # Repo root\nexport RAY_ADDRESS=\"...\" # The Ray cluster address to connect to\nexport WORKING_DIR=\"${PWD}\" # The local directory to package to the Ray cluster\nexport NNODES=xxx  # for actor/rollout/trainer\nexport RM_NODES=xxx  # for standalone reward model\n\nbash examples/fapo_trainer/run_qwen_7b_rm_standalone.sh  # 7b fapo model\nbash examples/fapo_trainer/run_qwen_32b_rm_standalone.sh  # 32b fapo model\n```\n\n## Results\n\nCompared with baseline (no reward model), FAPO significantly improves the reasoning ability of the model.\n\n![fapo-result](https://fapo-rl.github.io/_astro/intro_main.DKe72RHX_1Us2HB.webp)\n\n## Use discriminative reward model\n\nIf you would like to use discriminative reward models, the usage is essentially similar to GenRM. You only need to replace the \"/v1/chat/completions\" endpoint in the custom reward function with the reward model's endpoint.\n\nWe provide a standard way to compute the DisRM reward score, with the implementation in `RewardLoopWorker::compute_score_disrm`.\n\nYou can enable this computation method by not specifying a custom reward function.\n\n## More complex reward model scenarios\n\nBoth GenRM and DisRM can obtain reward scores via HTTP requests in the custom reward function.\nThis allows users to flexibly combine rule-based rewards with reward models to construct more sophisticated reward logic.\n\nFor more detailed usage instructions and infrastructure design, please refer to the [Reward Loop](https://verl.readthedocs.io/en/latest/advance/reward_loop.html) document.\n\n## Citation\n\nIf you find our works useful for your research, please consider citing:\n\n```bibtex\n@article{ding2025fapo,\n  title={FAPO: Flawed-Aware Policy Optimization for Efficient and Reliable Reasoning},\n  author={Ding, Yuyang and Zhang, Chi and Li, Juntao and Lin, Haibin and Zhang, Min},\n  journal={arXiv preprint arXiv:2510.22543},\n  year={2025}\n}\n```"
  },
  {
    "path": "examples/fapo_trainer/prepare_data.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nfrom functools import partial\n\nfrom datasets import concatenate_datasets, load_dataset\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\n\ndef example_map_fn(example, idx, process_fn, data_source, ability, split):\n    question, prompt, ground_truth = process_fn(example)\n    data = {\n        \"data_source\": data_source,\n        \"prompt\": [{\"role\": \"user\", \"content\": prompt}],\n        \"ability\": ability,\n        \"reward_model\": {\"style\": \"rule\", \"ground_truth\": ground_truth},\n        \"extra_info\": {\"split\": split, \"index\": idx, \"question\": question},\n    }\n    return data\n\n\ndef build_aime2024_dataset():\n    def process_aime2024(example):\n        question, ground_truth = example[\"Problem\"], str(example[\"Answer\"])\n        prompt = question.strip() + \"\\n\\n\" + \"Please reason step by step, and put your final answer within \\\\boxed{}.\"\n        return question, prompt, ground_truth\n\n    data_source = \"Maxwell-Jia/AIME_2024\"\n    print(f\"Loading the {data_source} dataset from huggingface...\", flush=True)\n    dataset = load_dataset(data_source, split=\"train\")\n    map_fn = partial(example_map_fn, process_fn=process_aime2024, data_source=\"aime24\", ability=\"Math\", split=\"test\")\n    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names)\n    return dataset\n\n\ndef build_aime2025_dataset():\n    def process_aime2025(example):\n        question, ground_truth = example[\"problem\"], str(example[\"solution\"])\n        prompt = question.strip() + \"\\n\\n\" + \"Please reason step by step, and put your final answer within \\\\boxed{}.\"\n        return question, prompt, ground_truth\n\n    data_source = \"yentinglin/aime_2025\"\n    print(f\"Loading the {data_source} dataset from huggingface...\", flush=True)\n    dataset = load_dataset(data_source, split=\"train\")\n    map_fn = partial(example_map_fn, process_fn=process_aime2025, data_source=\"aime25\", ability=\"Math\", split=\"test\")\n    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names)\n    return dataset\n\n\ndef build_gpqa_diamond_dataset():\n    import random\n\n    GPQA_QUERY_TEMPLATE = (\n        \"{Question}\\n\"\n        \"A. {A}\\nB. {B}\\nC. {C}\\nD. {D}\\n\\n\"\n        \"Please reason step by step, and put your final answer (only the choice letter) within \\\\boxed{{}}.\"\n    )\n\n    def process_gpqa_diamond(example):\n        choices = [\n            example[\"Incorrect Answer 1\"].strip(),\n            example[\"Incorrect Answer 2\"].strip(),\n            example[\"Incorrect Answer 3\"].strip(),\n        ]\n        random.shuffle(choices)\n        gold_index = random.randint(0, 3)\n        choices.insert(gold_index, example[\"Correct Answer\"].strip())\n        question = example[\"Question\"]\n        query_prompt = GPQA_QUERY_TEMPLATE.format(\n            A=choices[0],\n            B=choices[1],\n            C=choices[2],\n            D=choices[3],\n            Question=question,\n        )\n        gold_choice = \"ABCD\"[gold_index]\n        return question, query_prompt, gold_choice\n\n    data_source = \"Idavidrein/gpqa\"\n    print(f\"Loading the {data_source} dataset from huggingface...\", flush=True)\n\n    dataset = load_dataset(data_source, \"gpqa_diamond\", split=\"train\")\n    map_fn = partial(\n        example_map_fn, process_fn=process_gpqa_diamond, data_source=\"gpqa-diamond\", ability=\"General\", split=\"test\"\n    )\n    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names)\n    return dataset\n\n\ndef build_dapo_train_dataset():\n    def process_dapo(example):\n        question, ground_truth = example[\"prompt\"], example[\"solution\"]\n        prompt = question.strip() + \"\\n\\n\" + \"Please reason step by step, and put your final answer within \\\\boxed{}.\"\n        return question, prompt, ground_truth\n\n    data_source = \"open-r1/DAPO-Math-17k-Processed\"\n    print(f\"Loading the {data_source} dataset from huggingface...\", flush=True)\n    dataset = load_dataset(data_source, \"all\", split=\"train\")\n    map_fn = partial(example_map_fn, process_fn=process_dapo, data_source=\"math-dapo\", ability=\"Math\", split=\"train\")\n    dataset = dataset.map(map_fn, with_indices=True, remove_columns=dataset.column_names)\n    return dataset\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--local_dir\", default=\"~/data/genrm\")\n    parser.add_argument(\"--hdfs_dir\", default=None)\n    parser.add_argument(\"--tasks\", default=\"all\")\n\n    args = parser.parse_args()\n\n    train_dataset = build_dapo_train_dataset()\n    train_dataset = concatenate_datasets([train_dataset for _ in range(20)])\n\n    test_datasets = []\n    # AIME 2024\n    aime24_dataset = build_aime2024_dataset()\n    test_datasets.extend([aime24_dataset for _ in range(32)])\n    # AIME 2025\n    aime25_dataset = build_aime2025_dataset()\n    test_datasets.extend([aime25_dataset for _ in range(32)])\n    # GPQA Diamond\n    gpqa_dataset = build_gpqa_diamond_dataset()\n    test_datasets.extend([gpqa_dataset for _ in range(4)])\n    test_dataset = concatenate_datasets(test_datasets)\n\n    local_dir = args.local_dir\n    hdfs_dir = args.hdfs_dir\n\n    train_dataset.to_parquet(os.path.join(local_dir, \"fapo-train-boxed.parquet\"))\n    test_dataset.to_parquet(os.path.join(local_dir, \"fapo-test-full-boxed.parquet\"))\n\n    if hdfs_dir is not None:\n        makedirs(hdfs_dir)\n\n        copy(src=local_dir, dst=hdfs_dir)\n"
  },
  {
    "path": "examples/fapo_trainer/reward_fn.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nimport os\n\nimport aiohttp\nfrom transformers import PreTrainedTokenizer\n\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.reward_score.math_dapo import last_boxed_only_string, normalize_final_answer, remove_boxed\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef verify(\n    solution_str: str,\n    gt: str,\n) -> tuple[bool, str]:\n    solution_str = solution_str[-300:]\n    boxed_answer = last_boxed_only_string(solution_str)\n    if boxed_answer is not None:\n        extracted_answer = remove_boxed(boxed_answer)\n    else:\n        extracted_answer = \"[INVALID]\"\n\n    pred = normalize_final_answer(extracted_answer)\n    gt = normalize_final_answer(gt)\n    return (pred == gt), pred\n\n\nasync def compute_score_baseline(\n    solution_str: str,\n    ground_truth: str,\n    **kwargs,\n):\n    loop = get_event_loop()\n    \"\"\"Compute the reward score for Baseline.\"\"\"\n    correct, pred = await loop.run_in_executor(None, lambda: verify(solution_str, ground_truth))\n    reward_score = 1.0 if correct else -1.0\n    return {\"score\": reward_score, \"acc\": correct, \"pred\": pred}\n\n\n# FAPO Hyper-parameters\nFAPO_GENRM_TEMPLATE = (\n    \"The following is a math problem with its ground truth answer, along with an AI solution (split into steps):\\n\\n\"\n    \"[Math Problem]\\n\\n\"\n    \"{problem}\\n\\n\"\n    \"[Ground Truth]\\n\\n\"\n    \"{ground_truth}\\n\\n\"\n    \"[AI Solution]\\n\\n\"\n    \"{solution}\\n\\n\"\n    \"Your task is to review and critique the solution step by step. \"\n    \"Once you identify an error in a step, return the index of the step where the earliest error occurs. \"\n    \"Otherwise, return the index of -1 (which typically denotes 'not found').\\n\\n\"\n    \"Please reason step by step, put your final answer (i.e., the index) in \\\\boxed{{}}.\"\n)\nMAX_TOKENS = 16384\nFLAWED_REWARD_PENALTY = 1.0\n\n\n# async def generate_aiohttp(router_address: str, prompt_ids: list[int], sampling_params: dict):\n#     payload = {\n#         \"input_ids\": prompt_ids,\n#         \"sampling_params\": sampling_params,\n#     }\n#     url = f\"http://{router_address}/generate\"\n#     try:\n#         session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=None))\n#         async with session.post(url, json=payload) as resp:\n#             output = await resp.text()\n#             try:\n#                 output = json.loads(output)\n#                 return output\n#             except Exception:\n#                 logger.error(f\"Failed to parse JSON response: {output}\")\n#                 return {}\n#     finally:\n#         await session.close()\n\n\nasync def post_request(router_address: str, payload: dict, endpoint: str, max_retries: int = 5):\n    url = f\"http://{router_address}/{endpoint}\"\n    last_exception = None\n    for attempt in range(max_retries):\n        try:\n            # It's safer to have a timeout instead of None, which can hang indefinitely.\n            timeout = aiohttp.ClientTimeout(total=None)\n            async with aiohttp.ClientSession(timeout=timeout) as session:\n                async with session.post(url, json=payload) as resp:\n                    resp.raise_for_status()\n                    return await resp.json()\n        except aiohttp.ClientResponseError as e:\n            # Do not retry on 4xx client errors, but retry on 5xx server errors.\n            if 400 <= e.status < 500:\n                logger.error(f\"Request to {url} failed with client error HTTP {e.status}: {e}. Not retrying.\")\n                raise\n            last_exception = e\n            logger.warning(\n                f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with HTTP {e.status}: {e}. Retrying...\"\n            )\n        except (asyncio.TimeoutError, aiohttp.ClientConnectorError) as e:\n            last_exception = e\n            logger.warning(f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed: {e}. Retrying...\")\n        except Exception as e:\n            last_exception = e\n            logger.warning(\n                f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with unexpected error: {e}. Retrying...\"\n            )\n\n        if attempt < max_retries - 1:\n            # Using exponential backoff is generally better than a fixed sleep.\n            backoff_seconds = 2**attempt\n            await asyncio.sleep(min(backoff_seconds, 30))\n\n    logger.error(f\"Max retries ({max_retries}) reached for request to {url}.\")\n    if last_exception:\n        raise last_exception\n\n\nasync def compute_score_fapo(\n    data_source: str,\n    solution_str: str,\n    ground_truth: str,\n    extra_info: dict,\n    reward_router_address: str,\n    reward_model_tokenizer: PreTrainedTokenizer,\n):\n    \"\"\"Compute the reward score for FAPO.\"\"\"\n    loop = get_event_loop()\n\n    question, split = extra_info[\"question\"], extra_info[\"split\"]\n    correct, pred = await loop.run_in_executor(None, lambda: verify(solution_str, ground_truth))\n    reward_score = 1.0 if correct else -1.0\n    is_flawed_positive = False\n\n    # for test set or incorrect solution, directly return the reward score\n    if split == \"test\" or not correct:\n        return {\"score\": reward_score, \"acc\": correct, \"pred\": pred, \"is_flawed_positive\": is_flawed_positive}\n\n    grm_prompt = FAPO_GENRM_TEMPLATE.format(\n        problem=question,\n        ground_truth=ground_truth,\n        solution=solution_str,\n    )\n    messages = [{\"role\": \"user\", \"content\": grm_prompt}]\n    grm_outputs = await post_request(\n        router_address=reward_router_address,\n        payload={\n            \"messages\": messages,\n            \"max_tokens\": MAX_TOKENS,\n        },\n        endpoint=\"v1/chat/completions\",\n    )\n    grm_response = grm_outputs[\"choices\"][0][\"message\"][\"content\"]\n    try:\n        err_location = remove_boxed(last_boxed_only_string(grm_response))\n        is_flawed_positive = int(err_location) != -1\n    except Exception:\n        is_flawed_positive = False\n\n    if is_flawed_positive:\n        reward_score -= FLAWED_REWARD_PENALTY\n\n    return {\"score\": reward_score, \"acc\": correct, \"pred\": pred, \"is_flawed_positive\": is_flawed_positive}\n"
  },
  {
    "path": "examples/fapo_trainer/run_qwen_7b_rm_colocate.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='FAPO-Reproduce'\nexp_name='FAPO-7B'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=8\ntrain_prompt_mini_bsz=32\n\n# Ray\nRAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-4}\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nGRM_PATH=${GRM_PATH:-\"${RAY_DATA_HOME}/models/FAPO-GenRM-4B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/fapo-train-boxed.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/fapo-test-full-boxed.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=1\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=1\nfsdp_size=8\n\nray job submit --no-wait --runtime-env=\"${RUNTIME_ENV}\" \\\n    --address \"${RAY_ADDRESS}\" \\\n    --working-dir \"${WORKING_DIR}\" \\\n    -- python3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.return_raw_chat=True \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.enable_resource_pool=False \\\n    reward.reward_model.model_path=${GRM_PATH} \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.90 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.free_cache_engine=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=True \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    reward.custom_reward_function.path=examples/fapo_trainer/reward_fn.py \\\n    reward.custom_reward_function.name=compute_score_fapo \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=200 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/fapo_trainer/run_qwen_7b_rm_standalone.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='FAPO-Reproduce'\nexp_name='FAPO-7B'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=8\ntrain_prompt_mini_bsz=32\n\n# Ray\nRAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nRM_NODES=${RM_NODES:-2}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nGRM_PATH=${GRM_PATH:-\"${RAY_DATA_HOME}/models/FAPO-GenRM-4B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/fapo-train-boxed.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/fapo-test-full-boxed.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=1\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=1\nfsdp_size=8\n\nray job submit --no-wait --runtime-env=\"${RUNTIME_ENV}\" \\\n    --address \"${RAY_ADDRESS}\" \\\n    --working-dir \"${WORKING_DIR}\" \\\n    -- python3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.return_raw_chat=True \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.enable_resource_pool=True \\\n    reward.reward_model.n_gpus_per_node=8 \\\n    reward.reward_model.nnodes=${RM_NODES} \\\n    reward.reward_model.model_path=${GRM_PATH} \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.90 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.free_cache_engine=False \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=True \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    reward.custom_reward_function.path=examples/fapo_trainer/reward_fn.py \\\n    reward.custom_reward_function.name=compute_score_fapo \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=200 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/gdpo_trainer/run_qwen1_5b_gdpo.sh",
    "content": "export HCCL_ASYNC_ERROR_HANDLING=0\n\nexport DATA_DIR=\"./dataset/rlla_4k\"\nexport BASE_MODEL=\"/path/to/your/Qwen2.5-1.5B-Instruct\"\nexport EXPERIMENT_NAME=\"qwen2.5-1.5B-GDPO\"\nexport CKPT_DIR=\"./results/gdpo\"\n\n# Env variables for computing score in rlla.py\nexport REFINEDREWARD=0\nexport COARSEREWARD=0\nexport CORRECTMAX1=0\nexport MAX1STEP30MAX3=0\nexport SCHEDULEREWARD=0\nexport SCHEDULELENGTH=0\n\nPROJECT_DIR=\"$(pwd)\"\n\ntrainer_n_gpus_per_node=8\ntrainer_nnodes=1\n\npython3 -u -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gdpo \\\n    +algorithm.gdpo_reward_keys='[\"accuracy_reward\", \"format_reward\"]' \\\n    data.train_files=$DATA_DIR/train.parquet \\\n    data.val_files=$DATA_DIR/test.parquet \\\n    data.train_batch_size=32 \\\n    data.val_batch_size=16 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    actor_rollout_ref.model.path=$BASE_MODEL \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=4 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.prompt_length=2048 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    reward.custom_reward_function.path=\"$PROJECT_DIR/verl/utils/reward_score/rlla.py\" \\\n    reward.custom_reward_function.name=compute_score \\\n    reward.reward_manager.name=gdpo \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=\"GDPO-qwen2.5\" \\\n    trainer.n_gpus_per_node=$trainer_n_gpus_per_node \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.nnodes=$trainer_nnodes \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=10 \\\n    trainer.default_local_dir=$CKPT_DIR \\\n    trainer.total_epochs=15 \\\n    trainer.val_before_train=False 2>&1\n"
  },
  {
    "path": "examples/generation/run_deepseek7b_mutli_node.sh",
    "content": "set -x\n\ndata_path=$HOME/data/rlhf/gsm8k/test.parquet\nsave_path=$HOME/data/rlhf/math/deepseek_v2_lite_gen_test.parquet\nmodel_path=deepseek-ai/deepseek-llm-7b-chat\n\npython3 -m verl.trainer.main_generation \\\n    trainer.nnodes=2 \\\n    trainer.n_gpus_per_node=8 \\\n    data.path=$data_path \\\n    data.prompt_key=prompt \\\n    data.n_samples=1 \\\n    data.output_path=$save_path \\\n    model.path=$model_path\\\n    +model.trust_remote_code=True \\\n    rollout.temperature=1.0 \\\n    rollout.top_k=50 \\\n    rollout.top_p=0.7 \\\n    rollout.prompt_length=2048 \\\n    rollout.response_length=1024 \\\n    rollout.tensor_model_parallel_size=16 \\\n    rollout.gpu_memory_utilization=0.8\n"
  },
  {
    "path": "examples/generation/run_deepseek_v2_lite_math.sh",
    "content": "set -x\n\ndata_path=$HOME/data/gsm8k/test.parquet\nsave_path=$HOME/data/gsm8k/deepseek_v2_lite_gen_test.parquet\nmodel_path=deepseek-ai/deepseek-llm-7b-chat\n\npython3 -m verl.trainer.main_generation \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=8 \\\n    data.path=$data_path \\\n    data.prompt_key=prompt \\\n    data.n_samples=1 \\\n    data.output_path=$save_path \\\n    model.path=$model_path \\\n    +model.trust_remote_code=True \\\n    rollout.temperature=1.0 \\\n    rollout.top_k=50 \\\n    rollout.top_p=0.7 \\\n    rollout.prompt_length=2048 \\\n    rollout.response_length=1024 \\\n    rollout.tensor_model_parallel_size=2 \\\n    rollout.gpu_memory_utilization=0.8\n"
  },
  {
    "path": "examples/gmpo_trainer/README.md",
    "content": "<div align=center>\n  \n# Geometric-Mean Policy Optimization\n</div>\n\nThis is the official implementaion of paper [***Geometric-Mean Policy Optimization***](https://arxiv.org/abs/2507.20673).\n\n<div align=center>\n<img width=\"3092\" height=\"864\" alt=\"image\" src=\"https://github.com/user-attachments/assets/20b04c4e-7ee8-4775-9af8-33c0158336e2\" />\n</div>\n\n## 1. Contents\n- Geometric-Mean Policy Optimization\n  - [1. Contents](#1-contents)\n  - [2. Introduction](#2-introduction)\n  - [3. Code Usage](#3-code-usage)\n  - [4. Contacts](#4-contacts)\n  - [5. Citation](#5-citation)\n\n## 2. Introduction\n\nGroup Relative Policy Optimization (GRPO) has significantly enhanced the reasoning capability of large language models by optimizing the arithmetic mean of token-level rewards. Unfortunately, GRPO is observed to suffer from unstable policy updates when facing tokens with outlier importance-weighted rewards, which manifest as extreme importance sampling ratios during training. In this study, we propose Geometric-Mean Policy Optimization (GMPO), with the aim to improve the stability of GRPO through suppressing token reward outliers. Instead of optimizing the arithmetic mean, GMPO maximizes the geometric mean of token-level rewards, which is inherently less sensitive to outliers and maintains a more stable range of importance sampling ratio. GMPO is plug-and-play—simply replacing GRPO's arithmetic mean with the geometric mean of token-level rewards, as the latter is inherently less sensitive to outliers. GMPO is theoretically plausible—analysis reveals that both GMPO and GRPO are weighted forms of the policy gradient while the former enjoys more stable weights, which consequently benefits policy optimization and performance. Experiments on multiple mathematical reasoning benchmarks show that GMPO-7B improves the average Pass@1 of GRPO by up to 4.1%, outperforming many state-of-the-art approaches.\n\n## 3. Code Usage\n\nThe key configurations are:\n```\nclip_ratio_low=0.4\nclip_ratio_high=0.4\nloss_mode=geo_mean\n```\nWe observed that using a large clip ratio during Mixture-of-Experts (MoE) model training often leads to optimization instability. When training MoE models, consider lowering the clip ratio to achieve more stable convergence.\nTo get started quickly, run:\n```\nbash examples/gmpo_trainer/run_qwen2_5-7b_math.sh\n```\n\nGMPO can be combined with other methods such as DAPO (experimental - not fully tested):\n```\nbash examples/gmpo_trainer/test_dapo_7b_math.sh \nbash examples/gmpo_trainer/test_dapo_qwen3_30b_math.sh\n```\n\n## 4. Contacts\nIf you have any question about our work or this repository, please don't hesitate to contact us by emails or open an issue under this project.\n- [zhaoyuzhong20@mails.ucas.ac.cn](zhaoyuzhong20@mails.ucas.ac.cn)\n- [liuyue171@mails.ucas.ac.cn](liuyue171@mails.ucas.ac.cn)\n- [lecu@microsoft.com](lecu@microsoft.com)\n- [wanfang@ucas.ac.cn](wanfang@ucas.ac.cn)\n\n## 5. Citation\n```\n@article{zhao2025geometric,\n  title={Geometric-mean policy optimization},\n  author={Zhao, Yuzhong and Liu, Yue and Liu, Junpeng and Chen, Jingye and Wu, Xun and Hao, Yaru and Lv, Tengchao and Huang, Shaohan and Cui, Lei and Ye, Qixiang and others},\n  journal={arXiv preprint arXiv:2507.20673},\n  year={2025}\n}\n```\n"
  },
  {
    "path": "examples/gmpo_trainer/run_qwen2_5-7b_math.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\nuse_kl_loss=False\nloss_mode=geo_mean\nclip_ratio=0.4\nsave_contents=\"['model', 'optimizer', 'extra']\"\n\nexport WANDB_MODE=offline\nsave_contents=\"['hf_model']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-Math-7B \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.checkpoint.save_contents=${save_contents} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_gmpo_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_5_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/gmpo_trainer/test_dapo_7b_math.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.4\nclip_ratio_high=0.4\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=4\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=4\nfsdp_size=32\n\nloss_mode=geo_mean\n\n# export WANDB_MODE=offline\nsave_contents=\"['model', 'optimizer', 'extra']\"\n# save_contents=\"['hf_model']\"\n\n# reference run wandb: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.actor.checkpoint.save_contents=\"${save_contents}\" \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=10 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=200 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/gmpo_trainer/test_dapo_qwen3_30b_math.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen3-30B-A3B-Base-MATH-0527a1'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.4\nclip_ratio_high=0.4\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\nloss_mode=geo_mean\n\n# export WANDB_MODE=offline\nsave_contents=\"['model', 'optimizer', 'extra']\"\n# save_contents=\"['hf_model']\"\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=4\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=4\nfsdp_size=32\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.actor.checkpoint.save_contents=\"${save_contents}\" \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=10 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=300 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/gpg_trainer/gpg.md",
    "content": "# GPG: Group Policy Gradient\n\nGroup Policy Gradient (GPG) is a minimalist reinforcement learning (RL) method that enhances the reasoning ability of large language models without relying on supervised fine-tuning or complex tricks. GPG revisits traditional policy gradients and directly optimizes the RL objective—no surrogate losses, no KL penalties, no critic, and no reference model. Compared to GRPO, GPG is simpler, more efficient, and achieves better results on many tasks. For more details, please refer to the original paper [GPG: A Simple and Strong Reinforcement Learning Baseline for Model Reasoning\n](https://arxiv.org/abs/2504.02546).\n\n## Key Components\n- Use a corrected advantage function to improve policy gradient accuracy and training efficiency.\n- By eliminating the critic and reference models, avoiding KL divergence constraints, significantly simplifies the training process compared to Group Relative Policy Optimization (GRPO)\n\n## Configuration\nTo configure GPG within the framework, use the following YAML settings.\n\n```yaml\nalgorithm:\n  adv_estimator: gpg \nactor_rollout_ref:\n  actor:\n    policy_loss:\n      loss_mode: \"gpg\"\n```\n\n## Advanced Extensions\nGPG is a simple and strong baseline for model reasoning. Although it avoids using KL loss in its original form, you can still use KL loss to further improve the performance.\n\n```yaml\nalgorithm:\n  adv_estimator: gpg\nactor_rollout_ref:\n  actor:\n    use_kl_loss: True # enable kl regularization\n    kl_loss_coef: 0.01\n    policy_loss:\n      loss_mode: \"gpg\"\n```"
  },
  {
    "path": "examples/gpg_trainer/run_qwen2-7b_math.sh",
    "content": "set -x\n\n# If you are using vllm<=0.6.3, you might need to set the following environment variable to avoid bugs:\n# export VLLM_ATTENTION_BACKEND=XFORMERS\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gpg \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=gpg \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_gpg_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/gpg_trainer/run_qwen2-7b_math_megatron.sh",
    "content": "set -x\n\n# If you are using vllm<=0.6.3, you might need to set the following environment variable to avoid bugs:\n# export VLLM_ATTENTION_BACKEND=XFORMERS\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=gpg \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=gpg \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_gpg_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/README.md",
    "content": "# Group Relative Policy Optimization (GRPO)\n\nIn reinforcement learning, classic algorithms like PPO rely on a \"critic\" model to estimate the value of actions, guiding the learning process. However, training this critic model can be resource-intensive. \n\nGRPO simplifies this process by eliminating the need for a separate critic model. Instead, it operates as follows:\n- Group Sampling: For a given problem, the model generates multiple possible solutions, forming a \"group\" of outputs.\n- Reward Assignment: Each solution is evaluated and assigned a reward based on its correctness or quality.\n- Baseline Calculation: The average reward of the group serves as a baseline. \n- Policy Update: The model updates its parameters by comparing each solution's reward to the group baseline, reinforcing better-than-average solutions and discouraging worse-than-average ones.\n\nThis approach reduces computational overhead by avoiding the training of a separate value estimation model, making the learning process more efficient. For more details, refer to the original paper [DeepSeekMath: Pushing the Limits of Mathematical Reasoning in Open Language Models](https://arxiv.org/pdf/2402.03300)\n\n## Key Components\n\n- No Value Function (Critic-less): unlike PPO, GRPO does not train a separate value network (critic)\n- Group Sampling (Grouped Rollouts): instead of evaluating one rollout per input, GRPO generates multiple completions (responses) from the current policy for each prompt. This set of completions is referred to as a group.\n- Relative Rewards: within each group, completions are scored (e.g., based on correctness), and rewards are normalized relative to the group.\n\n## Configuration\n\nNote that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior.\n\nDespite that many configurations start with the `ppo_` prefix, they work across different RL algorithms in verl, as the GRPO training loop is similar to that of PPO (without critic).\n\n![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d)\n\n- `actor_rollout.ref.rollout.n`: For each prompt, sample n times. Default to 1. For GRPO, please set it to a value larger than 1 for group sampling.\n\n- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n`\n\n- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers.\n\n- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for GRPO updates on one set of sampled trajectories for actor\n\n- `actor_rollout_ref.actor.clip_ratio`: The GRPO clip range. Default to 0.2\n\n- `algorithm.adv_estimator`: Default is gae. Please set it to grpo instead\n\n- `actor_rollout_ref.actor.loss_agg_mode`: Default is \"token-mean\". Options include \"token-mean\", \"seq-mean-token-sum\", \"seq-mean-token-mean\". The original GRPO paper takes the sample-level loss (seq-mean-token-mean), which may be unstable in long-CoT scenarios. All GRPO example scripts provided in verl uses the default configuration \"token-mean\" for loss aggregation instead.\n\nInstead of adding KL penalty in the reward, GRPO regularizes by directly adding the KL divergence between the trained policy and the reference policy to the loss:\n\n- `actor_rollout_ref.actor.use_kl_loss`: To use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False. Please set it to True for GRPO.\n\n- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001.\n\n- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending \"+\" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/volcengine/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\n## Advanced Extensions\n\n### DrGRPO\n\nThe work [Understanding R1-Zero-Like Training: A Critical Perspective](https://arxiv.org/pdf/2503.20783) claims there's optimization bias in GRPO, that leads to artificially longer responses, especially for incorrect outputs. This inefficiency stems from the way GRPO calculates advantages using group-based reward normalization, which can inadvertently favor longer, less accurate responses. Instead, DrGRPO aggregates token-level losses by normalizing with a global constant to eliminate length bias.\n\nConfigure the following to enable DrGRPO, with all other parameters the same as GRPO's:\n\n- `actor_rollout_ref.actor.loss_agg_mode`: \"seq-mean-token-sum-norm\", which turns off seq-dim averaging\n- `actor_rollout_ref.actor.loss_scale_factor`: (Optional) Set to a constant integer (e.g., max response length) to ensure consistent normalization throughout training. If not set, uses the current batch's response length.\n- `actor_rollout_ref.actor.use_kl_loss`: Please set it to False for DrGRPO\n- `algorithm.norm_adv_by_std_in_grpo`: False, which turns off standard deviation norm\n\n## Reference Example\n\nQwen2.5 GRPO training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/qwen2-7b-fsdp2.log)\n\n```bash\nbash examples/grpo_trainer/run_qwen3-8b.sh\n```\n\nFor more reference performance, please see https://verl.readthedocs.io/en/latest/algo/baseline.html\n"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek671b_math_megatron_80gb.sh",
    "content": "set -x\n\n# # 0. download HF checkpoint\n# # remove the `quantization_config` in the `config.json`\n# # set `num_nextn_predict_layers=0` to disable MTP, which is not currently supported\n# hf download deepseek-ai/DeepSeek-V3-0324\n\n# no offline dist checkpoint needed, now with mbridge>=0.13.0, we can directly init model from huggingface downloaded fp8 weights\n# tested on docker://verlai/verl:app-verl0.5-transformers4.55.4-vllm0.10.0-mcore0.13.0-te2.2\nLLM=\"<path_to_dsv3_config>\"\n\n\n# 2. run the script\ngsm8k_train_path=/root/data/gsm8k/train.parquet\ngsm8k_test_path=/root/data/gsm8k/test.parquet\ntrain_files=$gsm8k_train_path\ntest_files=$gsm8k_test_path\n\nALL_OFFLOAD=${ALL_OFFLOAD:-True}\nCOMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD}\n\nACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nREF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_PARAM_OFFLOAD=${CRITIC_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_GRAD_OFFLOAD=${CRITIC_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nCRITIC_OPTIMIZER_OFFLOAD=${CRITIC_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nRM_PARAM_OFFLOAD=${RM_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\n\n# 256 H100(80GB)\nNODES=32\nPP=16\nTP=1\nEP=16\nETP=1\nINFER_TP=32\n# consider TP/ETP, and enable recompute if short of memory\n\n# full recompute\n\nn_resp_per_prompt=4\nmax_prompt_length=2048\nmax_response_length=4096\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\n# RAY_ADDRESS='auto' ray job submit --working-dir . --\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=grpo \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$LLM \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.temperature=1.0 \\\n    actor_rollout_ref.rollout.top_p=1.0 \\\n    actor_rollout_ref.rollout.top_k=-1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$INFER_TP \\\n    trainer.logger='[\"console\",\"tensorboard\"]' \\\n    trainer.project_name='verl_megatron_gsm8k_examples' \\\n    trainer.experiment_name='dsv3-32nodes' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=$NODES \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend='fused' \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=4 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} \\\n    actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    trainer.default_local_dir=$CKPT_DIR \\\n    trainer.val_before_train=False \\\n    trainer.total_epochs=100 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek671b_math_megatron_96gb.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n## !!!!!!!important!!!!!!\n# 1. set the following environment variables on all your nodes\n# env_vars:\n#   CUDA_DEVICE_MAX_CONNECTIONS: \"1\"\n#   NCCL_NVLS_ENABLE: \"0\"\n#   VLLM_USE_V1: 1\n# 2. install mbridge=0.1.13 on all your node with the following command: \n# pip3 install git+https://github.com/ISEEKYAN/mbridge\n# 3. remove the `quantization_config` in the DeepSeek-V3's `config.json` and \n# set `num_nextn_predict_layers=0` to disable MTP, which is not currently supported\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\n[ -f \"${SCRIPT_DIR}/env.sh\" ] && source \"${SCRIPT_DIR}/env.sh\"\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1204 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=96\nn_resp_per_prompt=8\ntrain_prompt_mini_bsz=32\n\n\n# minimum nodes for DeepSeek-V3: 12 nodes\nNNODES=${NNODES:-12}\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n\nMODEL_PATH=$RAY_DATA_HOME/models/DeepSeek-V3-config-verl\n\nTRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 10 / 10))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1))\noffload=True\noptim_offload=${OFFLOAD_OPTIM:-True}\ngen_tp=32\ntrain_tp=${TP:-8}\ntrain_pp=${PP:-12}\n\nEP=${EP:-8}\nETP=1\nCP=1\noptimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}\nLAST_LAYER=${LAST_LAYER:-6}\n\n\nproject_name='verl-deepseek-v3'\nexp_name=\"671B-${NNODES}-pp${train_pp}-tp${train_tp}-ep${EP}-actor-length${actor_ppo_max_token_len}\"\nCKPTS_DIR=$RAY_DATA_HOME/ckpt/${project_name}/${exp_name}\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${optim_offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.nccl_timeout=1200 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${CP} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=False \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_shared_expert_overlap=False \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split=False \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split=False \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=${LAST_LAYER} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=100 \\\n    trainer.total_epochs=10 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek7b_llm.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=80 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=160 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=160 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek7b_llm_math.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm_math' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek7b_llm_math_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='deepseek_llm_7b_math_megatron' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_deepseek7b_llm_seq_balance.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm_seq_packing' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_glm41v_9b.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=zai-org/GLM-4.1V-9B-Thinking \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='glm41v_9b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_gptoss_20b.sh",
    "content": "#!/bin/bash\n\ncat > get_model.py << EOF\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, Mxfp4Config\n\nmodel_id = \"openai/gpt-oss-20b\"\noutput_dir = \"$HOME/models/gpt-oss-20b-bf16\"\n\nquantization_config = Mxfp4Config(dequantize=True)\nmodel_kwargs = dict(\n    attn_implementation=\"eager\",\n    torch_dtype=torch.bfloat16,\n    quantization_config=quantization_config,\n    use_cache=False,\n    device_map=\"auto\",\n)\n\nmodel = AutoModelForCausalLM.from_pretrained(model_id, **model_kwargs)\n\n# Patch config with custom attribute before saving\nmodel.config.attn_implementation = \"eager\"\n\nmodel.save_pretrained(output_dir)\ntokenizer = AutoTokenizer.from_pretrained(model_id)\ntokenizer.save_pretrained(output_dir)\nEOF\n\npython get_model.py\n# or you can use lmsys/gpt-oss-20b-bf16\n# recommend to use same value for train_batch_size and ppo_mini_batch_size\n# to avoid MOE training instability\n# use large value for max_response_length if you want to use reasoning effort high.\n\n\nmodel_dir=$HOME/models/gpt-oss-20b-bf16\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$gsm8k_train_path\" \\\n    data.val_files=\"$gsm8k_test_path\" \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=8192 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    +data.apply_chat_template_kwargs.reasoning_effort=medium \\\n    actor_rollout_ref.model.path=${model_dir} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    +actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=triton \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='oai_oss_20b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=50 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_minicpmo2_6.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=128 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=False \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    data.trust_remote_code=True \\\n    data.custom_cls.path=recipe/minicpmo/rl_dataset.py \\\n    data.custom_cls.name=RLHFDataset \\\n    actor_rollout_ref.model.path=openbmb/MiniCPM-o-2_6 \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.use_orig_params=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='minicpmo2_6_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_mistral13b_skyworkrm_hhrlhf.sh",
    "content": "train_files=data/full_hh_rlhf/rl/train.parquet\ntest_files=data/full_hh_rlhf/rl/train.parquet # no use\n\nmax_prompt_length=4096\nmax_response_length=2048\n\ngen_tp=4\nn_per_prompt=5\nadv_estimator=\"grpo\"\n\nproject_name=verl_full_hh_rlhf_examples\nexp_name=\"grpo_mistral13B-skyworkLlama8b-hhrlhf\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=$adv_estimator \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=512 \\\n    data.prompt_key=\"prompt\" \\\n    data.return_raw_chat=True \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=mistralai/Mistral-Nemo-Instruct-2407 \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.n=$n_per_prompt \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=Skywork/Skywork-Reward-Llama-3.1-8B \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=8192 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.val_before_train=False \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$exp_name \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=10 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=5 $@"
  },
  {
    "path": "examples/grpo_trainer/run_moonlight16b_math_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\nHF_MODEL_PATH=moonshotai/Moonlight-16B-A3B\nDIST_CKPT_PATH=${DIST_CKPT_PATH}\n\ntrain_path=$HOME/data/gsm8k/train.parquet\ntest_path=$HOME/data/gsm8k/test.parquet\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=192 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.trust_remote_code=True \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=3 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=4 \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=1 \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=3 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=4 \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=1 \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='moonlight_megatron_ep' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=3 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_nemotron_nano_v3_megatron.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n################################################### environment ###################################################\n### # 1. use docker image `verlai/verl:vllm015.dev`` and install correct dependencies:\n# pip install nvidia-modelopt\n# MAX_JOBS=32 pip install git+https://github.com/Dao-AILab/causal-conv1d.git --no-build-isolation --no-cache-dir\n# MAX_JOBS=32 pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation --no-cache-dir\n# pip install --no-deps git+https://github.com/NVIDIA-NeMo/Megatron-Bridge \n# pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_dev_r0.16.0\n# unset ROCR_VISIBLE_DEVICES\n# unset PYTORCH_CUDA_ALLOC_CONF\n\n################################################### quick config ###################################################\n\nrollout_mode=\"async\"\nreturn_raw_chat=\"False\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\ndtype=\"bfloat16\"\n\nproject_name='DAPO'\nexp_name='nano_30b'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 2))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 1))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=32\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\n# Ray\nRAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-1}\n# Paths\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/NVIDIA-Nemotron-3-Nano-30B-A3B-BF16\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 10 / 10))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1))\noffload=True\ngen_tp=8\ntrain_tp=8\ntrain_pp=1\nEP=8\nETP=1\n\n################################################### start of config ###################################################\n\nFP8=(\n    # train\n    # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8=\"e4m3\" # e4m3 or hybrid\n    # +actor_rollout_ref.actor.megatron.override_transformer_config.fp8_recipe=\"blockwise\"\n    # +actor_rollout_ref.actor.optim.override_optimizer_config.fp8_recipe=\"blockwise\"\n    # rollout\n    actor_rollout_ref.actor.megatron.dtype=${dtype}\n    actor_rollout_ref.rollout.dtype=${dtype}\n    # +actor_rollout_ref.rollout.quantization=\"fp8\"\n)\n\nDATA=(\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    data.prompt_key=prompt\n    data.return_raw_chat=$return_raw_chat\n    data.truncation='left'\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.train_batch_size=${train_prompt_bsz}\n)\n\nREWARD_MODEL=(\n    +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer}\n    +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len}\n    +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor}\n    +reward_model.reward_kwargs.overlong_buffer_cfg.log=False\n    +reward_model.reward_kwargs.max_resp_len=${max_response_length}\n    reward_model.reward_manager=dapo\n)\n\nPERF_OPT=(\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True\n    actor_rollout_ref.model.use_fused_kernels=False\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n)\n\nACTOR=(\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low}\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high}\n    actor_rollout_ref.actor.clip_ratio_c=10.0\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10\n    actor_rollout_ref.actor.optim.weight_decay=0.1\n    actor_rollout_ref.actor.optim.clip_grad=1.0\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.megatron.param_offload=${offload}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload}\n    actor_rollout_ref.actor.megatron.grad_offload=${offload}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode}\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=False\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.name=${rollout_name}\n    actor_rollout_ref.rollout.mode=${rollout_mode}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.70\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length))\n    actor_rollout_ref.rollout.temperature=${temperature}\n    actor_rollout_ref.rollout.top_p=${top_p}\n    actor_rollout_ref.rollout.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature}\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p}\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.calculate_log_probs=True\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n)\n\nTRAINER=(\n    trainer.logger=['console','wandb']\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.n_gpus_per_node=8\n    trainer.nnodes=\"${NNODES}\"\n    trainer.val_before_train=False\n    trainer.test_freq=10\n    trainer.save_freq=-1\n    trainer.total_epochs=10\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n    trainer.resume_mode=auto\n    trainer.log_val_generations=10\n)\n\nFORWARD_ONLY_SETS=(\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.model.trust_remote_code=True\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n################################################### start script ###################################################\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${REWARD_MODEL[@]}\" \\\n    \"${FP8[@]}\" \\\n    \"${PERF_OPT[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    \"${FORWARD_ONLY_SETS[@]}\" \\"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-32b_sglang_fsdp_npu.sh",
    "content": "#!/bin/bash\nset -xeuo pipefail\nmkdir -p logs\n\n# Project Configuration\nproject_name='GRPO-Qwen2.5-32B-BASE-SGLang'\nexp_name='GRPO-Qwen2.5-32B-BASE-FSDP-SGLang'\n\n# Necessary env\nexport HCCL_CONNECT_TIMEOUT=1500\nexport HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\nexport HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\n\nexport RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\n# If the number of nodes is 16, ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15\nexport ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n\nexport DISABLE_L2_CACHE=1\nexport TASK_QUEUE_ENABLE=1\n\n# Node Info\nNNODES=${NNODES:-2}\nNPUS_PER_NODE=${NPUS_PER_NODE:-8}\n\n# Model Weights Paths\nMODEL_PATH=Qwen/Qwen2.5-32B\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/datasets/deepscaler/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/datasets/deepscaler/test.parquet\"}\n\n# Data Configuration\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\n\n# Training Batch Configuration\ntrain_prompt_bsz=32\ntrain_prompt_mini_bsz=32\nn_resp_per_prompt=8\n\n# Algorithm Configuration\nadv_estimator=grpo\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\n# Performance and Memory Management Configuration\nall_offload=True\nuse_dynamic_bsz=False\n\n# SGLang Configuration\ngen_tp=4\ngen_sp=1\ngen_dp=1\ngen_ep=1\ngpu_memory_utilization=0.5\n\n# Data Configuration\nDATA_CONFIG=(\n    # File Paths\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    # Data Structure\n    data.prompt_key=prompt\n    # Batch and Length Configuration\n    data.train_batch_size=${train_prompt_bsz}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    # Preprocessing\n    data.filter_overlong_prompts=False\n    data.truncation='left'\n)\n\n# Model Configuration\nMODEL_CONFIG=(\n    # Model Path\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    # Model Processing\n    actor_rollout_ref.model.use_remove_padding=True\n    actor_rollout_ref.model.enable_gradient_checkpointing=True\n)\n\n# Reinforcement Learning Algorithm Configuration\nALGORITHM_CONFIG=(\n    # Advantage Estimation\n    algorithm.adv_estimator=${adv_estimator}\n    # KL Divergence Control\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n)\n\n# Actor Model Configuration\nACTOR_CONFIG=(\n    # Core Runtime Settings\n    actor_rollout_ref.actor.use_torch_compile=False\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    # Loss Function Configuration\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.entropy_coeff=0\n    # PPO Training Parameters\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    # Optimizer Settings\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.fsdp_config.param_offload=${all_offload}\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${all_offload}\n    )\n\n# Reference Model Configuration\nREF_CONFIG=(\n    # Core Runtime Settings\n    actor_rollout_ref.ref.use_torch_compile=False\n    # Log Probability Inference\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    # Memory Optimization\n    actor_rollout_ref.ref.fsdp_config.param_offload=${all_offload}\n)\n\n# Rollout Configuration\nROLLOUT_CONFIG=(\n    # Rollout Engine\n    actor_rollout_ref.rollout.name=sglang\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=\"ascend\"\n    # Generation Parameters\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.top_p=1.0\n    actor_rollout_ref.rollout.top_k=-1\n    actor_rollout_ref.rollout.temperature=1.0\n    # Log Probability Inference\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    # Memory Management\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp}\n    actor_rollout_ref.rollout.expert_parallel_size=${gen_ep}\n    actor_rollout_ref.rollout.enable_chunked_prefill=False\n    actor_rollout_ref.rollout.multi_stage_wake_up=True\n    # Validation Generation\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.top_p=1.0\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0\n    actor_rollout_ref.nccl_timeout=1800\n)\n\n# Trainer Configuration\nTRAINER_CONFIG=(\n    trainer.logger='[\"console\"]'\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.nnodes=\"${NNODES}\"\n    trainer.n_gpus_per_node=\"${NPUS_PER_NODE}\"\n    trainer.total_epochs=5\n    trainer.val_before_train=False\n    trainer.test_freq=-1\n    trainer.save_freq=100\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n    trainer.critic_warmup=0\n)\n\n# Main GRPO Training Command\n# Add the reward function processing for the DeepScaler dataset here\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_trainer.yaml' \\\n    reward.custom_reward_function.path=recipe/r1_ascend/deepscaler.py \\\n    reward.custom_reward_function.name=compute_score \\\n    \"${DATA_CONFIG[@]}\" \\\n    \"${MODEL_CONFIG[@]}\" \\\n    \"${ACTOR_CONFIG[@]}\" \\\n    \"${REF_CONFIG[@]}\" \\\n    \"${ROLLOUT_CONFIG[@]}\" \\\n    \"${ALGORITHM_CONFIG[@]}\" \\\n    \"${TRAINER_CONFIG[@]}\" \\\n    \"$@\" | tee logs/run_qwen2_5-32b_grpo_fsdp_sglang_npu.log"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b.sh",
    "content": "set -x\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_math.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_math_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\nrollout_mode=\"async\"\nexport VLLM_USE_V1=1\nreturn_raw_chat=\"True\"\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\nUSE_FUSED_KERNELS=True\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.return_raw_chat=$return_raw_chat \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=$rollout_mode \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_math_megatron_lora.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n# Need to install Megatron-Bridge\n# NOTE: Make sure you use Megatron-Bridge later than 0.2.0 \n# (Recommend https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/83a7c1134c562d8c6decd10a1f0a6e6a7a8a3a44 or later)\n# for proper MoE LoRA support.\n\n# For Megatron communication/computation overlapping\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\n############################ Quick Config ############################\n\nrollout_name=\"vllm\" # sglang or vllm\nproject_name='verl_grpo_example_gsm8k_math'\nexp_name='qwen2_7b_megatron_lora'\n\nadv_estimator=grpo\n\nmax_prompt_length=1024\nmax_response_length=1024\ntrain_prompt_bsz=128\n\n############################ Paths ############################\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\n############################ Parameter Groups ############################\n\nDATA=(\n    data.train_files=\"$train_files\"\n    data.val_files=\"$test_files\"\n    data.max_prompt_length=$max_prompt_length\n    data.max_response_length=$max_response_length\n    data.train_batch_size=$train_prompt_bsz\n    data.filter_overlong_prompts=True\n    data.truncation='error'\n    data.shuffle=False\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct\n    actor_rollout_ref.model.lora.rank=256\n    actor_rollout_ref.model.lora.alpha=512\n    actor_rollout_ref.model.lora.lora_A_init_method=kaiming\n    # # Optional: Use canonical LoRA\n    # actor_rollout_ref.model.lora.type=\"canonical_lora\"\n    # actor_rollout_ref.model.lora.target_modules='[\"linear_q\",\"linear_k\",\"linear_v\",\"linear_proj\",\"linear_fc1_up\",\"linear_fc1_gate\",\"linear_fc2\"]'\n\n    # # Optional: Add dropout to LoRA layers\n    # actor_rollout_ref.model.lora.dropout=0.05\n    # actor_rollout_ref.model.lora.dropout_position=pre\n)\n\nACTOR=(\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.ppo_mini_batch_size=16\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2\n    actor_rollout_ref.actor.use_dynamic_bsz=True\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=False\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=4\n    actor_rollout_ref.actor.use_kl_loss=True\n    actor_rollout_ref.actor.kl_loss_coef=0.001\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.entropy_coeff=0\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2\n    actor_rollout_ref.rollout.name=$rollout_name\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6\n    actor_rollout_ref.rollout.n=4\n)\n\nREF=(\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=1\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=4\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=$adv_estimator\n    algorithm.use_kl_in_reward=False\n)\n\nTRAINER=(\n    trainer.logger='[\"console\",\"wandb\"]'\n    trainer.project_name=$project_name\n    trainer.experiment_name=$exp_name\n    trainer.n_gpus_per_node=8\n    trainer.nnodes=1\n    trainer.save_freq=20\n    trainer.test_freq=5\n    trainer.total_epochs=15\n    trainer.val_before_train=False\n)\n\n############################ Launch ############################\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${REF[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    \"$@\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_math_megatron_trtllm.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\n# Clean all slurm / MPI / PMIx env to avoid pmix mismatch error\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n\nexport RAY_DEDUP_LOGS=0\n\n# -----\n# Config\n# -----\nTP=${1:-4}\nACTOR_TP=${ACTOR_TP:-4}\nPROJECT_NAME=${PROJECT_NAME:-\"verl_grpo_example_gsm8k_math\"}\nEXP_NAME=megatron-trtllm-qwen2-7b-tp${TP}-8gpus\n\nif [ $TP -eq 4 ]; then\n    MAX_BATCH_SIZE=1024\nelse\n    MAX_BATCH_SIZE=384\nfi\n\n# -----\n# Data\n# -----\nDATADIR=${DATADIR:-$PWD/data}\n\nGSM8K_TRAIN_PATH=${DATADIR}/gsm8k/train.parquet\nGSM8K_TEST_PATH=${DATADIR}/gsm8k/test.parquet\nMATH_TRAIN_PATH=${DATADIR}/math/train.parquet\nMATH_TEST_PATH=${DATADIR}/math/test.parquet\n\nTRAIN_FILES=\"['$GSM8K_TRAIN_PATH', '$MATH_TRAIN_PATH']\"\nTEST_FILES=\"['$GSM8K_TEST_PATH', '$MATH_TEST_PATH']\"\n\nUSE_FUSED_KERNELS=True\n\n# -----\n# Launch\n# -----\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$TRAIN_FILES\" \\\n    data.val_files=\"$TEST_FILES\" \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.model.use_fused_kernels=$USE_FUSED_KERNELS \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${TP} \\\n    actor_rollout_ref.rollout.name=trtllm \\\n    actor_rollout_ref.rollout.mode=\"async\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=${MAX_BATCH_SIZE} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=32768 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${PROJECT_NAME}\" \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.resume_mode=disable \\\n    trainer.total_epochs=15 \\\n    \"${@:2}\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh",
    "content": "set -x\n\n# Clean all slurm / MPI / PMIx env to avoid pmix mismatch error\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n\nexport RAY_DEDUP_LOGS=0\n\n# -----\n# Config\n# -----\nTP=${1:-4}\nPROJECT_NAME=${PROJECT_NAME:-\"verl_grpo_example_gsm8k_math\"}\nEXP_NAME=trtllm-qwen2-7b-tp${TP}-8gpus${EXP_NAME_SUFFIX:+\"-\"}${EXP_NAME_SUFFIX}\n\nif [ $TP -eq 4 ]; then\n    MAX_BATCH_SIZE=1024\nelse\n    MAX_BATCH_SIZE=384\nfi\n\n# -----\n# Data\n# -----\nDATADIR=${DATADIR:-$PWD/data}\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen2-7B-Instruct\"}\n\nGSM8K_TRAIN_PATH=${DATADIR}/gsm8k/train.parquet\nGSM8K_TEST_PATH=${DATADIR}/gsm8k/test.parquet\nMATH_TRAIN_PATH=${DATADIR}/math/train.parquet\nMATH_TEST_PATH=${DATADIR}/math/test.parquet\n\nTRAIN_FILES=\"['$GSM8K_TRAIN_PATH', '$MATH_TRAIN_PATH']\"\nTEST_FILES=\"['$GSM8K_TEST_PATH', '$MATH_TEST_PATH']\"\n\n# -----\n# Launch\n# -----\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.rollout_correction.rollout_is_threshold=2.0 \\\n    data.train_files=\"$TRAIN_FILES\" \\\n    data.val_files=\"$TEST_FILES\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=1024 \\\n    data.return_raw_chat=True \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.hybrid_engine=True \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${TP} \\\n    actor_rollout_ref.rollout.name=trtllm \\\n    actor_rollout_ref.rollout.mode=\"async\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=${MAX_BATCH_SIZE} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=32768 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${PROJECT_NAME}\" \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.resume_mode=disable \\\n    trainer.total_epochs=15 \\\n    \"${@:2}\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_seq_balance.sh",
    "content": "set -x\n\n\n# For async rollout mode, dataset should return raw chat.\nrollout_mode=\"async\"\nrollout_name=\"sglang\" # sglang or vllm\nreturn_raw_chat=\"True\"\nif [ \"$rollout_name\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.return_raw_chat=$return_raw_chat \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$rollout_name \\\n    actor_rollout_ref.rollout.mode=$rollout_mode \\\n    actor_rollout_ref.rollout.multi_turn.format=hermes \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm_kl1e-3' \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_seq_balance_math_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\noffload=True\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=12000 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2-7b_sgl_megatron.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5-32b_grpo_megatron_vllm_npu.sh",
    "content": "#!/bin/bash\nset -xeuo pipefail\nmkdir -p logs\n\n# Project Configuration\nproject_name='GRPO-Qwen2.5-32B-BASE-MATH'\nexp_name='GRPO-Qwen2.5-32B-BASE-Megatron-vLLM'\n\n# Node Info\nNNODES=${NNODES:-1}\nNPUS_PER_NODE=${NPUS_PER_NODE:-16}\n\n# Model Weights Paths\nMODEL_PATH=Qwen/Qwen2.5-32B\nMCORE_MODEL_PATH=Qwen/Qwen2.5-32B-dist\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=$RAY_DATA_HOME/dataset/gsm8k/train.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/gsm8k/test.parquet\n\n# Data Configuration\nmax_prompt_length=$((1024 * 1))\nmax_response_length=$((1024 * 1))\n\n# Training Batch Configuration\ntrain_prompt_bsz=128\ntrain_prompt_mini_bsz=32\nn_resp_per_prompt=16\n\n# Algorithm Configuration\nadv_estimator=grpo\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\n# Performance and Memory Management Configuration\nall_offload=True\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 8))\noptimizer_offload_fraction=1\n\n# Megatron Configuration\ntrain_tp=4\ntrain_ep=1\ntrain_etp=1\ntrain_pp=4\ntrain_cp=1\n\n# vLLM Configuration\ngen_tp=2\ngen_dp=1\ngen_ep=1\ngpu_memory_utilization=0.8\nmax_model_len=$((max_prompt_length + max_response_length))\nmax_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1))\n\n# Data Configuration\nDATA_CONFIG=(\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    data.prompt_key=prompt\n    data.train_batch_size=${train_prompt_bsz}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.filter_overlong_prompts=False\n    data.truncation='left'\n)\n\n# Model Configuration\nMODEL_CONFIG=(\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.model.use_remove_padding=True\n)\n\n# Algorithm Configuration\nALGORITHM_CONFIG=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n\n# Actor Model Configuration\nACTOR_CONFIG=(\n    actor_rollout_ref.actor.use_torch_compile=False\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.ppo_epochs=1\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.optim.lr=1e-6\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction}\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${train_etp}\n    actor_rollout_ref.actor.megatron.param_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.grad_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=False\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True\n)\n\n# Reference Model Configuration\nREF_CONFIG=(\n    actor_rollout_ref.ref.use_torch_compile=False\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${train_etp}\n    actor_rollout_ref.ref.megatron.param_offload=${all_offload}\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=False\n)\n\n# Rollout Configuration\nROLLOUT_CONFIG=(\n    actor_rollout_ref.rollout.name=vllm\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.top_p=1.0\n    actor_rollout_ref.rollout.top_k=-1\n    actor_rollout_ref.rollout.temperature=1.0\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization}\n    actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_batched_tokens}\n    actor_rollout_ref.rollout.max_model_len=${max_model_len}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp}\n    actor_rollout_ref.rollout.expert_parallel_size=${gen_ep}\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.enable_prefix_caching=True\n    actor_rollout_ref.rollout.enforce_eager=True\n    actor_rollout_ref.rollout.free_cache_engine=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.top_p=1.0\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0\n)\n\n# Trainer Configuration\nTRAINER_CONFIG=(\n    trainer.logger='[\"console\",\"tensorboard\"]'\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.nnodes=\"${NNODES}\"\n    trainer.n_gpus_per_node=\"${NPUS_PER_NODE}\"\n    trainer.device='npu'\n    trainer.total_epochs=15\n    trainer.val_before_train=False\n    trainer.test_freq=-1\n    trainer.save_freq=-1\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n)\n\n# Main GRPO Training Command\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA_CONFIG[@]}\" \\\n    \"${MODEL_CONFIG[@]}\" \\\n    \"${ACTOR_CONFIG[@]}\" \\\n    \"${REF_CONFIG[@]}\" \\\n    \"${ROLLOUT_CONFIG[@]}\" \\\n    \"${ALGORITHM_CONFIG[@]}\" \\\n    \"${TRAINER_CONFIG[@]}\" \\\n    \"$@\" | tee logs/run_qwen2_5-32b_grpo_megatron_vllm_npu.log\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5-3b_gsm8k_grpo_lora.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    trainer.val_before_train=False \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.model.lora_rank=64 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2.5_3b_grpo_lora' \\\n    trainer.n_gpus_per_node=2 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n\n    # actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    # data.train_batch_size=1024 \\\n    # trainer.n_gpus_per_node=8 \\\n    # actor_rollout_ref.model.use_shm=True \\\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5-3b_gsm8k_grpo_lora_from_adapter.sh",
    "content": "set -x\n\nlora_adapter_path=${lora_adapter_path:-/path/saved/lora_adapter}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.model.use_shm=True \\\n    actor_rollout_ref.model.lora_adapter_path=${lora_adapter_path} \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2.5_3b_grpo_lora' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5-7b_math_megatron_diff_tp.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name='qwen2_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_32b_grpo_npu.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-32B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6\\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_32b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=2 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_7b_grpo_discrete_prof_npu.sh",
    "content": "set -x\n\n# profiling configuration\nPROFILE_STEPS=\"[2,4]\"\nPROFILE_RANKS_ALL=False\nDISCRETE=True\nPROFILE_RANKS=\"[1,2]\"\n\n# profiling NPU options\nSAVE_PATH=\"$HOME/profile_data\"\nLEVEL=\"level0\"\nCONTENTS=['npu','cpu']\nANALYSIS=True\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.optim.lr=5e-8 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=2 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.profiler.enable=True \\\n    actor_rollout_ref.ref.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=5 \\\n    global_profiler.tool=npu \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.save_path=$SAVE_PATH\n    $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_7b_grpo_e2e_prof_npu.sh",
    "content": "set -x\n\n# profiling configuration\nPROFILE_STEPS=\"[2,4]\"\nPROFILE_RANKS_ALL=True\nDISCRETE=False\n\n# profiling NPU options\nSAVE_PATH=\"$HOME/profile_data\"\nLEVEL=\"level0\"\nCONTENTS=['npu','cpu']\nANALYSIS=True\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=5e-8 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=2 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.profiler.enable=True \\\n    actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=5 \\\n    global_profiler.tool=npu \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.save_path=$SAVE_PATH\n    $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_7b_grpo_npu.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=5e-8 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_7b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=5 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b-megatron.sh",
    "content": "set -x\nENGINE=${1:-vllm}\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\nHF_MODEL_PATH=Qwen/Qwen2.5-VL-7B-Instruct\nDIST_CKPT_PATH=${DIST_CKPT_PATH}\n\n# convert HF model to megatron format offlinely\n# python scripts/converter_hf_to_mcore.py --hf_model_path $HF_MODEL_PATH --output_path $DIST_CKPT_PATH\n\n\n# megatron tuning guide:\n# 1. recommend to offload all states by setting ALL_OFFLOAD=True\n# 2. enable dynamic batch size by setting actor_rollout_ref.actor.use_dynamic_bsz=True ref.log_prob_use_dynamic_bsz=True rollout.log_prob_use_dynamic_bsz=True\n# 3. set ppo_max_token_len_per_gpu and log_prob_max_token_len_per_gpu as large as possible for better MFU (limited by GPU memory). assure ppo_max_token_len_per_gpu > max_prompt_length+max_response_length, if sequence length is too long, you can increase the TP/PP size\n# 4. if memory is very limited, enable full recompute, but the mfu will be 30% lower\n#        full recompute settings:\n#        +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n#        +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n#        +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n\nALL_OFFLOAD=${ALL_OFFLOAD:-True}\nCOMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD}\n\nACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nREF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\n\n\ntrain_path=$HOME/data/geo3k/train.parquet\ntest_path=$HOME/data/geo3k/test.parquet\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=5120 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=20480 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=20480 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=1 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} \\\n    actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b-sglang.sh",
    "content": "set -x\n\n# python examples/data_preprocess/geo3k.py --local_dir ~/data/geo3k\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \\\n    actor_rollout_ref.rollout.multi_stage_wake_up=True \\\n    global_profiler.tool=torch_memory \\\n    global_profiler.save_path=./mem_snapshots \\\n    global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries=100000 \\\n    global_profiler.global_tool_config.torch_memory.stack_depth=32 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.mode=async \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b-trtllm.sh",
    "content": "set -x\n\n# Clean all slurm / MPI / PMIx env to avoid pmix mismatch error\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n\nexport RAY_DEDUP_LOGS=0\n\n# -----\n# Config\n# -----\nTP=${1:-4}\nPROJECT_NAME=${PROJECT_NAME:-\"verl_grpo_example_gsm8k_math\"}\nEXP_NAME=trtllm-qwen2.5-vl-7b-tp${TP}-8gpus${EXP_NAME_SUFFIX:+\"-\"}${EXP_NAME_SUFFIX}\n\nif [ $TP -eq 4 ]; then\n    MAX_BATCH_SIZE=1024\nelse\n    MAX_BATCH_SIZE=384\nfi\n\n# -----\n# Data\n# -----\nDATADIR=${DATADIR:-$PWD/data}\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen2.5-VL-7B-Instruct\"}\n\nGEO3K_TRAIN_PATH=${DATADIR}/geo3k/train.parquet\nGEO3K_TEST_PATH=${DATADIR}/geo3k/test.parquet\nTRAIN_FILES=\"['$GEO3K_TRAIN_PATH']\"\nTEST_FILES=\"['$GEO3K_TEST_PATH']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.rollout_correction.rollout_is_threshold=2.0 \\\n    data.train_files=\"$TRAIN_FILES\" \\\n    data.val_files=\"$TEST_FILES\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.return_raw_chat=True \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.trust_remote_code=True \\\n    actor_rollout_ref.hybrid_engine=True \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    +actor_rollout_ref.model.override_config.attn_implementation=eager \\\n    +actor_rollout_ref.ref.model.override_config.attn_implementation=eager \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${TP} \\\n    actor_rollout_ref.rollout.name=trtllm \\\n    actor_rollout_ref.rollout.mode=\"async\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=${MAX_BATCH_SIZE} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=16384 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.ref.strategy=fsdp2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\"]' \\\n    trainer.project_name=\"${PROJECT_NAME}\" \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=10 \\\n    trainer.test_freq=5 \\\n    trainer.resume_mode=disable \\\n    trainer.total_epochs=10 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b_freeze_vision.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.freeze_vision_tower=True \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b_lora.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n# If you are using vllm<=0.6.3, you might need to set the following environment variable to avoid bugs:\n# export VLLM_ATTENTION_BACKEND=XFORMERS\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.model.lora_rank=64 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.model.exclude_modules='.*visual.*' \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl-7b_seq_balance.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=6144 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=6144 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl_32b_npu.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-32B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_32b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=2 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl_3b_npu.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_3b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl_3b_trtllm.sh",
    "content": "set -x\n\n# Clean all slurm / MPI / PMIx env to avoid pmix mismatch error\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n\nexport RAY_DEDUP_LOGS=0\n\n# -----\n# Config\n# -----\nTP=${1:-4}\nPROJECT_NAME=${PROJECT_NAME:-\"verl_grpo_example_gsm8k_math\"}\nEXP_NAME=trtllm-qwen2.5-vl-3b-tp${TP}-8gpus${EXP_NAME_SUFFIX:+\"-\"}${EXP_NAME_SUFFIX}\n\nif [ $TP -eq 4 ]; then\n    MAX_BATCH_SIZE=1024\nelse\n    MAX_BATCH_SIZE=384\nfi\n\n# -----\n# Data\n# -----\nDATADIR=${DATADIR:-$PWD/data}\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen2.5-VL-3B-Instruct\"}\n\nGSM8K_TRAIN_PATH=${DATADIR}/gsm8k/train.parquet\nGSM8K_TEST_PATH=${DATADIR}/gsm8k/test.parquet\nMATH_TRAIN_PATH=${DATADIR}/math/train.parquet\nMATH_TEST_PATH=${DATADIR}/math/test.parquet\n\nTRAIN_FILES=\"['$GSM8K_TRAIN_PATH', '$MATH_TRAIN_PATH']\"\nTEST_FILES=\"['$GSM8K_TEST_PATH', '$MATH_TEST_PATH']\"\n\n# -----\n# Launch\n# -----\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.rollout_correction.rollout_is_threshold=2.0 \\\n    data.train_files=\"$TRAIN_FILES\" \\\n    data.val_files=\"$TEST_FILES\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=1024 \\\n    data.return_raw_chat=True \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.hybrid_engine=True \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${TP} \\\n    actor_rollout_ref.rollout.name=trtllm \\\n    actor_rollout_ref.rollout.mode=\"async\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=${MAX_BATCH_SIZE} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=32768 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${PROJECT_NAME}\" \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.resume_mode=disable \\\n    trainer.total_epochs=15 \\\n    \"${@:2}\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen2_5_vl_7b_npu.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-235b_megatron_96gb.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n## !!!!!!!important!!!!!!\n## set the following environment variables on all your nodes\n# env_vars:\n#   CUDA_DEVICE_MAX_CONNECTIONS: \"1\"\n#   NCCL_NVLS_ENABLE: \"0\"\n#   VLLM_USE_V1: 1\n# install mbridge=0.1.13 on all your node with the following command: \n# pip3 install git+https://github.com/ISEEKYAN/mbridge\n\nSCRIPT_DIR=\"$(cd \"$(dirname \"${BASH_SOURCE[0]}\")\" && pwd)\"\n[ -f \"${SCRIPT_DIR}/env.sh\" ] && source \"${SCRIPT_DIR}/env.sh\"\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1204 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 1))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=${TRAIN_BS:-32}\nn_resp_per_prompt=8\ntrain_prompt_mini_bsz=16\n\n# minimum nodes need for qwen3-235B-A22B\nNNODES=${NNODES:-4}\n# Paths\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n\nMODEL_PATH=$RAY_DATA_HOME/models/Qwen3-235B-A22B\n\nTRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 10 / 10))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1))\noffload=True\nOPTIM_OFFLOAD=${OPTIM_OFFLOAD:-True}\ngen_tp=8\ntrain_tp=${TP:-4}\ntrain_pp=${PP:-8}\n\nEP=${EP:-4}\nETP=1\nCP=1\noptimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}\nlast_layer=${LAST_LAYER:-10}\n\nproject_name='verl-qwen3'\nexp_name=\"235B-${NNODES}-pp${train_pp}-tp${train_tp}-ep${EP}-actor-length${actor_ppo_max_token_len}\"\nCKPTS_DIR=$RAY_DATA_HOME/ckpt/${project_name}/${exp_name}\n\n# TODO: support cuda graph for rollout by setting the following config\n    # actor_rollout_ref.rollout.cudagraph_capture_sizes=[1,2,4,8,16,32]\n    # actor_rollout_ref.rollout.enforce_eager=False\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${OPTIM_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.nccl_timeout=1200 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${CP} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=\"flex\" \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=100 \\\n    trainer.total_epochs=10 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-30b_dapo_megatron_fp8_trtllm.sh",
    "content": "set -x\n\n# Clean all slurm / MPI / PMIx env to avoid pmix mismatch error\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n\nexport RAY_DEDUP_LOGS=0\n\n# ----------------------\n# Config for GB200 node\n# ----------------------\nTP=${INFER_TP:-4}\nACTOR_TP=${ACTOR_TP:-4}\nACTOR_PP=${ACTOR_PP:-2}\nACTOR_VPP=${ACTOR_VPP:-2}\nACTOR_EP=${ACTOR_EP:-2}\nACTOR_CP=${ACTOR_CP:-1}\nREF_TP=${REF_TP:-4}\nREF_PP=${REF_PP:-2}\nREF_VPP=${REF_VPP:-2}\nREF_EP=${REF_EP:-2}\nREF_CP=${REF_CP:-1}\nGEN_MOE_TP=${GEN_MOE_TP:-2}\nGEN_MOE_EP=${GEN_MOE_EP:-2}\nPROJECT_NAME=${PROJECT_NAME:-\"Qwen3-30B-A3B-DAPO-GB200\"}\nNNODES=${NNODES:-4}\nGPUS_PER_NODE=${GPUS_PER_NODE:-4}\n# MOE backend for TRTLLM when using FP8 quantization:\n#   - Blackwell: use DEEPGEMM\n#   - Hopper: use CUTLASS\nTRTLLM_MOE_BACKEND=${TRTLLM_MOE_BACKEND:-\"DEEPGEMM\"}\nEXP_NAME=qwen3-30b-dapo-megatron-fp8-trtllm-n${NNODES}-tp${TP}-moe-tp${GEN_MOE_TP}-moe-ep${GEN_MOE_EP}${EXP_NAME_SUFFIX:+\"-\"}${EXP_NAME_SUFFIX}\n\nif [ $TP -eq 4 ] || [ $TP -eq 2 ]; then\n    MAX_NUM_SEQS=1024\nelse\n    MAX_NUM_SEQS=384\nfi\n\n# -----\n# Data\n# -----\nDATA_DIR=${DATA_DIR:-\"$PWD\"}\n\nDAPO_MATH_TRAIN=${DAPO_MATH_TRAIN:-\"${DATA_DIR}/data/DAPO-Math-17k/data/dapo-math-17k.parquet\"}\nAIME_VAL=${AIME_VAL:-\"${DATA_DIR}/data/AIME-2024/data/aime-2024.parquet\"}\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen3-30B-A3B-Base\"}\n\n# When PP=1, Megatron interleaved schedule is invalid; pass null so PP=1 works (e.g. 2-node)\n[ \"${ACTOR_PP}\" -gt 1 ] && ACTOR_VPP_OVERRIDE=${ACTOR_VPP} || ACTOR_VPP_OVERRIDE=null\n[ \"${REF_PP}\" -gt 1 ] && REF_VPP_OVERRIDE=${REF_VPP} || REF_VPP_OVERRIDE=null\n\n# -----\n# Launch\n# -----\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.use_kl_in_reward=False \\\n    algorithm.kl_ctrl.kl_coef=0.0 \\\n    reward_model.reward_manager=dapo \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.enable=True \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.len=4096 \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=1.0 \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward_model.reward_kwargs.max_resp_len=8192 \\\n    data.train_files=\"${DAPO_MATH_TRAIN}\" \\\n    data.val_files=\"${AIME_VAL}\" \\\n    data.prompt_key=prompt \\\n    data.return_raw_chat=True \\\n    data.truncation=left \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=8192 \\\n    data.train_batch_size=512 \\\n    data.filter_overlong_prompts=False \\\n    actor_rollout_ref.hybrid_engine=True \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.optim.lr=1e-5 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=30720 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.0 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.loss_agg_mode=token-mean \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP_OVERRIDE} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=40960 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP_OVERRIDE} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \\\n    actor_rollout_ref.rollout.name=trtllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=40960 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${TP} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.max_num_seqs=${MAX_NUM_SEQS} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=10240 \\\n    actor_rollout_ref.rollout.max_model_len=10240 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_timeout_iters=32 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.batch_wait_max_tokens_ratio=0.5 \\\n    +actor_rollout_ref.rollout.engine_kwargs.trtllm.moe_config.backend=${TRTLLM_MOE_BACKEND} \\\n    +actor_rollout_ref.rollout.moe_tensor_parallel_size=${GEN_MOE_TP} \\\n    actor_rollout_ref.rollout.expert_parallel_size=${GEN_MOE_EP} \\\n    +actor_rollout_ref.rollout.quantization=fp8 \\\n    actor_rollout_ref.rollout.prompt_length=2048 \\\n    actor_rollout_ref.rollout.response_length=8192 \\\n    actor_rollout_ref.rollout.temperature=1.0 \\\n    actor_rollout_ref.rollout.top_p=1 \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${PROJECT_NAME}\" \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=${GPUS_PER_NODE} \\\n    trainer.nnodes=${NNODES} \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.resume_mode=auto \\\n    trainer.total_epochs=1000 \\\n    trainer.val_before_train=False \\\n    trainer.log_val_generations=10 \\\n    \"${@}\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-32b_npu.sh",
    "content": "set -x\n\nproject_name='GRPO-Qwen3'\nexp_name='GRPO-Qwen3-32b-npu'\ngen_tp=4\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-32B\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/test.parquet\"}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=4 \\\n    +actor_rollout_ref.actor.fsdp_config.mixed_precision.param_dtype=bf16 \\\n    +actor_rollout_ref.actor.fsdp_config.mixed_precision.reduce_dtype=bf16 \\\n    +actor_rollout_ref.actor.fsdp_config.mixed_precision.buffer_dtype=fp32 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=32768 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=4 \\\n    trainer.resume_from_path=checkpoints/ \\\n    trainer.save_freq=500 \\\n    trainer.test_freq=50 \\\n    trainer.total_epochs=50 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-4b_gsm8k_grpo_lora_merge.sh",
    "content": "set -x\n\n# initial \"val-core/openai/gsm8k/acc/mean@1\":0.378316906747536\n# after training: \"val-core/openai/gsm8k/acc/mean@1\":0.9264594389689158\n\nTIMESTAMP=$(date +%Y%m%d.%H%M%S)\nproject_name=verl_grpo_example_gsm8k\nexperiment_name=qwen3_4b_grpo-lora-merged-${TIMESTAMP}\ntrain_dir=outputs/$project_name/$experiment_name/\nmkdir -p $train_dir\nexport TENSORBOARD_DIR=$train_dir/tensorboard_log/\nexport VERL_FILE_LOGGER_PATH=$train_dir/metrics.jsonl\n\nmax_token_len_per_gpu=24576\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    trainer.val_before_train=True \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=128 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-4B \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    +actor_rollout_ref.model.lora.merge=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=64 \\\n    actor_rollout_ref.actor.optim.lr=1.0e-05 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${max_token_len_per_gpu} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bf16 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${max_token_len_per_gpu} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${max_token_len_per_gpu} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.strategy=fsdp2 \\\n    actor_rollout_ref.ref.fsdp_config.model_dtype=bf16 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"tensorboard\",\"file\"]' \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$experiment_name \\\n    trainer.default_local_dir=$train_dir \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 \\\n    2>&1 | tee $train_dir/train_log.txt\n\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-8b.sh",
    "content": "# Tested successfully on the hiyouga/verl:ngc-th2.6.0-cu126-vllm0.8.4-flashinfer0.2.2-cxx11abi0 image.\n# It outperforms the Qwen2 7B base model by two percentage points on the test set of GSM8K.\n\nset -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-8B \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen3_8b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3-8b_npu.sh",
    "content": "set -x\n\nproject_name='GRPO-Qwen3'\nexp_name='GRPO-Qwen3-8B-npu'\ngen_tp=2\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-8B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.resume_mode=auto \\\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=True \\\n    ++actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \\\n    ++actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \\\n    ++actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=5 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_235b_megatron_npu.sh",
    "content": "set -xeuo pipefail\n\nproject_name='GRPO'\nexp_name='GRPO-qwen3-235b-megatron'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nenable_filter_groups=False\nmax_num_gen_batches=32\nfilter_groups_metric=acc\nmax_prompt_length=$((1024 * 8))\nmax_response_length=$((1024 * 4))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=128 # must be > n_gpus. need to fix\ngen_prompt_bsz=$((train_prompt_bsz * 1))\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=128 # mini_bsz * n >= micro_bsz * pp * dp\n\nNNODES=8\n\nMODEL_PATH=${WORK_DIR}/Qwen3-235B-A22B\nMCORE_MODEL_PATH=${WORK_DIR}/Qwen3-235B-A22B-Mcore\n\nCKPTS_DIR=\".ckpt\"\n\nTRAIN_FILE=${WORK_DIR}/gsm8k/train.parquet\nTEST_FILE=${WORK_DIR}/gsm8k/test.parquet\n\nval_top_p=0.7\nUSE_MBRIDGE=False\nUSE_CKPT=True\noffload=True\n\n\ngen_tp=8\ngen_dp=8\nrollout_max_num_seqs=64\nmax_num_batched_tokens=$((1024))\ntrain_tp=4\ntrain_ep=4\ntrain_pp=8\n\nactor_ppo_max_token_len=$((max_prompt_length + max_response_length))\ninfer_ppo_max_token_len=$((max_prompt_length + max_response_length))\nactor_ppo_max_token_len=$((actor_ppo_max_token_len))\ninfer_ppo_max_token_len=$((infer_ppo_max_token_len))\n\npython3 -m verl.trainer.main_ppo --config-path=config  --config-name='ppo_megatron_trainer' \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.filter_overlong_prompts=False \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep} \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_CKPT \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enable_prefix_caching=True \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_first_pipeline_stage=11 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=11 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp} \\\n    actor_rollout_ref.rollout.expert_parallel_size=64 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    +actor_rollout_ref.rollout.enable_expert_parallel=True \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.75 \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_batched_tokens} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=$USE_CKPT \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=-1 \\\n    trainer.save_freq=100 \\\n    trainer.total_epochs=1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    actor_rollout_ref.nccl_timeout=7200 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True \\\n    +actor_rollout_ref.ref.megatron.override_transformer_config.use_flash_attn=True \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    trainer.device=npu \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes=\"[8, 16, 32, 64, 128]\" \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_DECODE_ONLY\" 2>&1 | tee \"logs/verl_qwen3_235b_sy$(date +%Y%m%d_%H%M).log\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_4b_grpo_vllm_1k_npu.sh",
    "content": "set -xeuo pipefail\nsource /usr/local/Ascend/ascend-toolkit/set_env.sh\nsource /usr/local/Ascend/nnal/atb/set_env.sh\n\n# 使用v1引擎\nexport VLLM_USE_V1=1\n# 指定vllm 版本\nexport VLLM_VERSION=0.9.1\n\n# 开启二级流水\nexport TASK_QUEUE_ENABLE=2\n# 开启细绑核\nexport CPU_AFFINITY_CONF=1\n# 使用jemalloc优化内存访问（依赖安装jemalloc）\nexport LD_PRELOAD=\"/usr/lib/aarch64-linux-gnu/libjemalloc.so.2${LD_PRELOAD:+:$LD_PRELOAD}\"\n\n# A3 机器单机8卡\ntrainer_n_gpus_per_node=16\ntrainer_nnodes=1\ntrainer_project_name='verl_grpo_example_gsm8k'\ntrainer_experiment_name=\"qwen3_4b_grpo_8npu}\"\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-4B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${trainer_project_name}/${trainer_experiment_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/test.parquet\"}\n\nexport TENSORBOARD_DIR=\"${RAY_DATA_HOME}/tensorboard_dir/${trainer_project_name}/${trainer_experiment_name}\"\nmkdir -p \"${RAY_DATA_HOME}/logs/${trainer_project_name}\"\nLOG_PATH=\"${RAY_DATA_HOME}/logs/${trainer_project_name}/${trainer_experiment_name}.log\"\n\nuse_dynamic_bsz=True\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=${TRAIN_FILE} \\\n    data.val_files=${TEST_FILE} \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=5e-7 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=3000 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.use_torch_compile=True \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.project_name=${trainer_project_name} \\\n    trainer.experiment_name=${trainer_experiment_name} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.n_gpus_per_node=$trainer_n_gpus_per_node \\\n    trainer.nnodes=$trainer_nnodes \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 \\\n    trainer.val_before_train=False 2>&1 | tee ${LOG_PATH}"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_5-35b-megatron.sh",
    "content": "#!/usr/bin/env bash\n# Qwen3.5-35B-A3B MoE GRPO RL with Megatron (single node, 8 GPUs, geo3k dataset)\n#\n# notes on vllm:\n#     by 20260225, the latest vllm nightly does not support qwen3.5 rollout, to use this script, you need to \n#         1. wait until vllm supports qwen3.5 officially, and build a verl docker with that version of vllm\n#         2. self build a verl docker image with vllm from source code with qwen3.5 support (main branch 20260225 is OK)\n#     I succeeded in running this script with the main branch of vllm on 20260225, yet there are still some minor issues\n#     the vllm qwen3.5 during initialization, need to be fixed. Also, the cuda_graph is somehow not working, need to be \n#     fixed, either by verl team with supoorts to vllm0.16, or by vllm team.\n# Requirements:\n#   - 8 GPUs (80GB each, e.g. 1x8 H100/H200)\n#   - Additional packages on top of the base image:\n#       pip install --upgrade transformers\n#       pip install flash-linear-attention\n#       pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n#   - Megatron-LM==0.16.0\n#\n# Qwen3.5 architecture notes:\n#   Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does\n#   NOT support packed sequences (THD format) in Megatron-LM. Therefore:\n#     - model.use_remove_padding=False           (deprecated option, will be removed in the future forces bshd compute format)\n#     - actor.megatron.use_remove_padding=False  (forces bshd compute format)\n#     - actor.use_dynamic_bsz=False              (required for bshd mode)\n#\n#   Once Megatron-LM adds THD support for Qwen3.5 GDN, use_remove_padding\n#   can be set to True for better performance.\n#\n# Tested parallelism config (8 GPUs / 1 node):\n#   TP=2 PP=1 CP=1 EP=8 ETP=1 GEN_TP=8\n#\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\nexport VLLM_USE_V1=1\nexport VLLM_ALLREDUCE_USE_SYMM_MEM=0\n\nset -xeuo pipefail\n\n########################### Quick Config ###########################\n\nTP=${TP:-2}\nPP=${PP:-1}\nCP=${CP:-1}\nEP=${EP:-8}\nETP=${ETP:-1}\nGEN_TP=${GEN_TP:-8}\n\nALL_OFFLOAD=${ALL_OFFLOAD:-True}\n\nrollout_name=\"vllm\"\nproject_name='verl_grpo_qwen3_5_35b_geo3k'\nexp_name='qwen3_5_35b_megatron'\nadv_estimator=grpo\n\nHF_MODEL_PATH=${HF_MODEL_PATH:-\"Qwen3.5-35B-A3B\"}\ntrain_path=${train_path:-$HOME/data/geo3k/train.parquet}\ntest_path=${test_path:-$HOME/data/geo3k/test.parquet}\n\n########################### Parameter Arrays ###########################\n\nDATA=(\n    data.train_files=${train_path}\n    data.val_files=${test_path}\n    data.train_batch_size=32\n    data.max_prompt_length=1024\n    data.max_response_length=2048\n    data.truncation='error'\n    data.filter_overlong_prompts=True\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=${HF_MODEL_PATH}\n    actor_rollout_ref.model.trust_remote_code=True\n    actor_rollout_ref.model.use_remove_padding=False\n)\n\nACTOR=(\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.ppo_mini_batch_size=32\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=4096\n    actor_rollout_ref.actor.use_dynamic_bsz=False\n    actor_rollout_ref.actor.use_kl_loss=True\n    actor_rollout_ref.actor.kl_loss_coef=0.01\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=True\n    actor_rollout_ref.actor.megatron.use_remove_padding=False\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.dtype=bfloat16\n    ++actor_rollout_ref.actor.megatron.override_transformer_config.attention_backend=auto\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_aux_loss_coeff=0.01\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_z_loss_coeff=0.001\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.name=${rollout_name}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${GEN_TP}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6\n    actor_rollout_ref.rollout.n=5\n    actor_rollout_ref.rollout.mode=async\n    actor_rollout_ref.rollout.dtype=bfloat16\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=False\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096\n)\n\nREF=(\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=False\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=4096\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD}\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=False\n)\n\nTRAINER=(\n    trainer.critic_warmup=0\n    trainer.logger='[\"console\",\"wandb\"]'\n    trainer.project_name=${project_name}\n    trainer.experiment_name=${exp_name}\n    trainer.n_gpus_per_node=8\n    trainer.nnodes=1\n    trainer.save_freq=20\n    trainer.val_before_train=False\n    trainer.test_freq=5\n    trainer.total_epochs=15\n)\n\n########################### Launch ###########################\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${REF[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    \"$@\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_8b_grpo_sglang_1k_spmd_npu.sh",
    "content": "set -x\nexport HCCL_CONNECT_TIMEOUT=1500\nexport HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\nexport HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\nexport RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\nexport ASCEND_RT_VISIBLE_DEVICES=0,1,2,3\n# WORKSPACE_HOME and DATA_HOME support custom path configuration.\nWORKSPACE_HOME=$pwd\nDATA_HOME=$pwd\n\nsp_size=4\nnum_npu=4\ntp_size=4\ntrain_prompt_bsz=16\ntrain_prompt_mini_bsz=16\n\nmax_prompt_length=512\nmax_response_length=1024\n\nCKPTS_DIR=$WORKSPACE_HOME/logs/ckpt/qwen3_8b\nmodel_path=$DATA_HOME/models/Qwen3-8B\ntrain_data=$DATA_HOME/datasets/processed_gsm8k/train.parquet\nvalid_data=$DATA_HOME/datasets/processed_gsm8k/test.parquet\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$train_data \\\n    data.val_files=$valid_data \\\n    data.train_batch_size=$train_prompt_bsz \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$train_prompt_mini_bsz \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$tp_size \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \\\n    actor_rollout_ref.rollout.n=5 \\\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=\"ascend\" \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.nccl_timeout=1800 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.val_before_train=False \\\n    trainer.project_name='verl_grpo_example_512_1024_gsm8k' \\\n    trainer.experiment_name='qwen3_8b_function_rm' \\\n    trainer.n_gpus_per_node=$num_npu \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=1000 \\\n    trainer.test_freq=10000 \\\n    trainer.total_epochs=5 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_8b_grpo_sglang_32k_spmd_npu.sh",
    "content": "set -x\nexport HCCL_CONNECT_TIMEOUT=1500\nexport HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\nexport HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\nexport RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\nexport ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\n# WORKSPACE_HOME and DATA_HOME support custom path configuration.\nWORKSPACE_HOME=$pwd\nDATA_HOME=$pwd\n\nsp_size=4\nnum_gpu=8\ntp_size=4\ntrain_prompt_bsz=16\ntrain_prompt_mini_bsz=16\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 32))\n\nCKPTS_DIR=$WORKSPACE_HOME/logs/ckpt/qwen3_8b\nmodel_path=$DATA_HOME/models/Qwen3-8B\ntrain_data=$DATA_HOME/datasets/dapo/dapo-math-17k.parquet\nvalid_data=$DATA_HOME/datasets/dapo/aime-2024.parquet\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$train_data \\\n    data.val_files=$valid_data \\\n    data.train_batch_size=$train_prompt_bsz \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=False \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$train_prompt_mini_bsz \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$tp_size \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \\\n    actor_rollout_ref.rollout.n=5 \\\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=\"ascend\" \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.nccl_timeout=3600 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.val_before_train=False \\\n    trainer.project_name='verl_grpo_example_2k_32k' \\\n    trainer.experiment_name='qwen3_8b_function_rm' \\\n    trainer.n_gpus_per_node=$num_gpu \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=1000 \\\n    trainer.test_freq=10000 \\\n    trainer.total_epochs=5 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_vl-235b-megatron.sh",
    "content": "set -x\nENGINE=${1:-vllm}\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\nexport VLLM_ALLREDUCE_USE_SYMM_MEM=0 # for vllm0.11.0 with TP\n\n\nHF_MODEL_PATH=${HF_MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-VL-235B-A22B-Instruct\"}\n\nGEN_TP=${GEN_TP:-16}\nCP=${CP:-2}\nTP=${TP:-4}\nPP=${PP:-8}\nEP=${EP:-8}\nETP=${ETP:-1}\n\ntrain_path=$HOME/data/geo3k/train.parquet\ntest_path=$HOME/data/geo3k/test.parquet\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$GEN_TP \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_loss_in_pipeline_split=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.account_for_embedding_in_pipeline_split=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen3_vl_235b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=8 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_vl-30b-megatron.sh",
    "content": "set -x\nENGINE=${1:-vllm}\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\nexport VLLM_ALLREDUCE_USE_SYMM_MEM=0 # for vllm0.11.0 with TP\n\n\nHF_MODEL_PATH=${HF_MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-VL-30B-A3B-Instruct\"}\n\nGEN_TP=${GEN_TP:-4}\nCP=${CP:-2}\nTP=${TP:-2}\nPP=${PP:-1}\nEP=${EP:-8}\nETP=${ETP:-1}\n\ntrain_path=$HOME/data/geo3k/train.parquet\ntest_path=$HOME/data/geo3k/test.parquet\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$GEN_TP \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    # Use aux_loss and z_loss to mitigate expert load imbalance when training MoE models\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_aux_loss_coeff=0.01 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_z_loss_coeff=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen3_vl_30b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_vl-8b-megatron.sh",
    "content": "set -x\nENGINE=${1:-vllm}\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\n# dependency: vllm>=0.11.0, megatron-lm>=0.13, mbridge with qwen3vl_cp branch\n# environment option1: use a stable container later than docker://verlai/verl:vllm011.dev6 \n    # and install mbridge in it by following the instruction in the container\n            # pip remove mbridge if you have installed it\n            # pip install git+https://github.com/ISEEKYAN/mbridge.git@qwen3vl_cp # for correct mbridge\n# environment option2: use container docker://verlai/verl:vllm011.dev_qwenvl_cp\n \n\nexport VLLM_ALLREDUCE_USE_SYMM_MEM=0 # for vllm0.11.0 with TP\n\n\nHF_MODEL_PATH=${HF_MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-VL-8B-Instruct\"}\n\nGEN_TP=${GEN_TP:-4}\nCP=${CP:-2}\nTP=${TP:-2}\nPP=${PP:-2}\n\ntrain_path=$HOME/data/geo3k/train.parquet\ntest_path=$HOME/data/geo3k/test.parquet\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$GEN_TP \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=4096 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen3_vl_8b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_vl-8b_npu.sh",
    "content": "set -x\n\nproject_name='GRPO-Qwen3_vl'\nexp_name='GRPO-Qwen3_vl-8B-npu'\ngen_tp=1\nsp_size=1\nENGINE=${1:-vllm}\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-VL-8B-Instruct\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/geo3k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/geo3k/test.parquet\"}\n\n# Rollout Correction parameters (sequence-level TIS + geometric RS)\nrollout_is=sequence\nrollout_is_threshold=2.0\nrollout_is_batch_normalize=true\nrollout_rs=token_k1\nrollout_rs_threshold=0.6_1.6\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=20000 \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=32 \\\n    actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True \\\n    actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True \\\n    actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True \\\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.model.enable_activation_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$sp_size \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=$sp_size \\\n    actor_rollout_ref.rollout.n=5 \\\n    data.shuffle=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=20 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=4 \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.resume_mode=auto \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \\\n    algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} \\\n    algorithm.rollout_correction.rollout_rs=${rollout_rs} \\\n    algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=5 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3_vl_30b_vllm_fsdp_npu.sh",
    "content": "set -x\n\nproject_name='GRPO-Qwen3_vl'\nexp_name='GRPO-Qwen3_vl-30B-npu'\ngen_tp=8\nsp_size=2\nENGINE=${1:-vllm}\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-VL-30B-A3B-Instruct\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/geo3k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/geo3k/test.parquet\"}\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n\n# Rollout Correction parameters\nrollout_is=sequence\nrollout_is_threshold=2.0\nrollout_is_batch_normalize=true\nrollout_rs=token_k1\nrollout_rs_threshold=0.6_1.6\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=32 \\\n    actor_rollout_ref.actor.fsdp_config.reshard_after_forward=True \\\n    actor_rollout_ref.ref.fsdp_config.reshard_after_forward=True \\\n    actor_rollout_ref.actor.fsdp_config.entropy_checkpointing=True \\\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$sp_size \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=$sp_size \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=20000 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    algorithm.use_kl_in_reward=False \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \\\n    algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} \\\n    algorithm.rollout_correction.rollout_rs=${rollout_rs} \\\n    algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\", \"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=2 \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.resume_mode=auto \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=5 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 \\"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3moe-30b_grpo_megatron_vllm_npu.sh",
    "content": "#!/bin/bash\nset -xeuo pipefail\nmkdir -p logs\n\n# Project Configuration\nproject_name='GRPO-Qwen3-30b-A3B-BASE-MATH'\nexp_name='GRPO-Qwen3-30B-A3B-BASE-Megatron-vLLM'\n\n# Node Info\nNNODES=${NNODES:-1}\nNPUS_PER_NODE=${NPUS_PER_NODE:-16}\n\n# Model Weights Paths\nMODEL_PATH=Qwen/Qwen3-30B-A3B-Base\nMCORE_MODEL_PATH=Qwen/Qwen3-30B-A3B-Base-dist\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=$RAY_DATA_HOME/dataset/gsm8k/train.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/gsm8k/test.parquet\n\n# Data Configuration\nmax_prompt_length=$((1024 * 1))\nmax_response_length=$((1024 * 1))\n\n# Training Batch Configuration\ntrain_prompt_bsz=128\ntrain_prompt_mini_bsz=32\nn_resp_per_prompt=16\n\n# Algorithm Configuration\nadv_estimator=grpo\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\n# Performance and Memory Management Configuration\nall_offload=True\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 4))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 8))\noptimizer_offload_fraction=1\n\n# Megatron Configuration\ntrain_tp=2\ntrain_ep=8\ntrain_etp=1\ntrain_pp=2\ntrain_cp=1\n\n# vLLM Configuration\ngen_tp=2\ngen_dp=1\ngen_ep=1\ngpu_memory_utilization=0.8\nmax_model_len=$((max_prompt_length + max_response_length))\nmax_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1))\n\n# Data Configuration\nDATA_CONFIG=(\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    data.prompt_key=prompt\n    data.train_batch_size=${train_prompt_bsz}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.filter_overlong_prompts=False\n    data.truncation='left'\n)\n\n# Model Configuration\nMODEL_CONFIG=(\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.model.use_remove_padding=True\n)\n\n# Algorithm Configuration\nALGORITHM_CONFIG=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n\n# Actor Model Configuration\nACTOR_CONFIG=(\n    actor_rollout_ref.actor.use_torch_compile=False\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.ppo_epochs=1\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.optim.lr=1e-6\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction}\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${train_etp}\n    actor_rollout_ref.actor.megatron.param_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.grad_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=False\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n)\n\n# Reference Model Configuration\nREF_CONFIG=(\n    actor_rollout_ref.ref.use_torch_compile=False\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${train_etp}\n    actor_rollout_ref.ref.megatron.param_offload=${all_offload}\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=False\n)\n\n# Rollout Configuration\nROLLOUT_CONFIG=(\n    actor_rollout_ref.rollout.name=vllm\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.top_p=1.0\n    actor_rollout_ref.rollout.top_k=-1\n    actor_rollout_ref.rollout.temperature=1.0\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization}\n    actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_batched_tokens}\n    actor_rollout_ref.rollout.max_model_len=${max_model_len}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp}\n    actor_rollout_ref.rollout.expert_parallel_size=${gen_ep}\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.enable_prefix_caching=True\n    actor_rollout_ref.rollout.enforce_eager=True\n    actor_rollout_ref.rollout.free_cache_engine=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.top_p=1.0\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0\n)\n\n# Trainer Configuration\nTRAINER_CONFIG=(\n    trainer.logger='[\"console\",\"tensorboard\"]'\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.nnodes=\"${NNODES}\"\n    trainer.n_gpus_per_node=\"${NPUS_PER_NODE}\"\n    trainer.device='npu'\n    trainer.total_epochs=15\n    trainer.val_before_train=False\n    trainer.test_freq=-1\n    trainer.save_freq=-1\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n)\n\n# Main GRPO Training Command\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA_CONFIG[@]}\" \\\n    \"${MODEL_CONFIG[@]}\" \\\n    \"${ACTOR_CONFIG[@]}\" \\\n    \"${REF_CONFIG[@]}\" \\\n    \"${ROLLOUT_CONFIG[@]}\" \\\n    \"${ALGORITHM_CONFIG[@]}\" \\\n    \"${TRAINER_CONFIG[@]}\" \\\n    \"$@\" | tee logs/run_qwen3moe-30b_grpo_megatron_vllm_npu.log\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3moe-30b_megatron_96gb.sh",
    "content": "set -x\n\n# tested in NNODES=1~4 * 96G H20 GPU\nNNODES=${NNODES:-1}\nNGPUS_PER_NODES=${NGPUS_PER_NODES:-8}\n\nproject_name='DAPO-Qwen3-30b-MATH'\nexp_name='DAPO-Qwen3-30b-MATH-megatron'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=128\ntrain_ppo_micro_batch_size_per_gpu=2\ninfer_ppo_micro_batch_size_per_gpu=2\n# Paths\nMODEL_PATH=Qwen/Qwen3-30B-A3B-Base\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nTRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\nTEST_FILE=\"['$aime24_test_path']\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\noffload=True\n\noptimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}\n\nCOMMON_PP=${COMMON_PP:-1}\nCOMMON_VPP=${COMMON_VPP:-null}\nCOMMON_CP=${COMMON_CP:-1}\nCOMMON_TP=${COMMON_TP:-1}\nCOMMON_EP=${COMMON_EP:-8}\nCOMMON_ETP=${COMMON_ETP:-1}\n\nTRAIN_TP=${TRAIN_TP:-$COMMON_TP}\nINFER_TP=${INFER_TP:-4}\n\nACTOR_PP=${ACTOR_PP:-$COMMON_PP}\nACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}\nACTOR_CP=${ACTOR_CP:-$COMMON_CP}\nACTOR_TP=${ACTOR_TP:-$TRAIN_TP}\nACTOR_EP=${ACTOR_EP:-$COMMON_EP}\nACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}\nROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}\nREF_PP=${REF_PP:-$COMMON_PP}\nREF_VPP=${REF_VPP:-$COMMON_VPP}\nREF_CP=${REF_CP:-$COMMON_CP}\nREF_TP=${REF_TP:-$TRAIN_TP}\nREF_EP=${REF_EP:-$COMMON_EP}\nREF_ETP=${REF_ETP:-$COMMON_ETP}\nCRITIC_PP=${CRITIC_PP:-$COMMON_PP}\nCRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}\nCRITIC_CP=${CRITIC_CP:-$COMMON_CP}\nCRITIC_TP=${CRITIC_TP:-$TRAIN_TP}\nCRITIC_EP=${CRITIC_EP:-$COMMON_EP}\nCRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}\nRM_PP=${RM_PP:-$COMMON_PP}\nRM_VPP=${RM_VPP:-$COMMON_VPP}\nRM_CP=${RM_CP:-$COMMON_CP}\nRM_TP=${RM_TP:-$TRAIN_TP}\nRM_EP=${RM_EP:-$COMMON_EP}\nRM_ETP=${RM_ETP:-$COMMON_ETP}\n\n# install mbridge\n# pip3 install git+https://github.com/ISEEKYAN/mbridge\nUSE_MBRIDGE=True\nUSE_DIST_CKPT=False\n\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    +actor_rollout_ref.model.override_config.model_config.max_position_embeddings=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.model.use_fused_kernels=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.lr_decay_style='constant' \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=\"flex\" \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODES}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=100 \\\n    trainer.total_epochs=10 \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3moe-30b_megatron_lora.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n# Need to install Megatron-Bridge\n# NOTE: Make sure you use Megatron-Bridge later than 0.2.0 \n# (Recommend https://github.com/NVIDIA-NeMo/Megatron-Bridge/commit/83a7c1134c562d8c6decd10a1f0a6e6a7a8a3a44 or later)\n# for proper MoE LoRA support.\n\n# For Megatron communication/computation overlapping\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\n########################### Quick Config ###########################\n\nTP=${TP:-2}\nPP=${PP:-2}\nCP=${CP:-2}\nEP=${EP:-4}\nETP=${ETP:-1}\n\nALL_OFFLOAD=${ALL_OFFLOAD:-True}\n\n\nrollout_name=\"vllm\"\nproject_name='verl_grpo_example_gsm8k_math'\nexp_name='qwen3_30b_a3b_megatron_lora'\nadv_estimator=grpo\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\n\n########################### Parameter Arrays ###########################\n\nDATA=(\n    data.train_files=${gsm8k_train_path}\n    data.val_files=${gsm8k_test_path}\n    data.train_batch_size=128\n    data.max_prompt_length=1024\n    data.max_response_length=1024\n    data.truncation='error'\n    data.filter_overlong_prompts=True\n    data.shuffle=False\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=Qwen/Qwen3-30B-A3B-Instruct-2507\n    actor_rollout_ref.model.use_fused_kernels=True\n    actor_rollout_ref.model.lora.rank=32\n    actor_rollout_ref.model.lora.alpha=64\n    actor_rollout_ref.model.lora.lora_A_init_method=kaiming\n    # # Optional: Use canonical LoRA\n    # actor_rollout_ref.model.lora.type=\"canonical_lora\"\n    # actor_rollout_ref.model.lora.target_modules='[\"linear_q\",\"linear_k\",\"linear_v\",\"linear_proj\",\"linear_fc1_up\",\"linear_fc1_gate\",\"linear_fc2\"]'\n\n    # # Optional: Add dropout to LoRA layers\n    # actor_rollout_ref.model.lora.dropout=0.05\n    # actor_rollout_ref.model.lora.dropout_position=pre\n)\n\nACTOR=(\n    actor_rollout_ref.actor.optim.lr=3e-6\n    actor_rollout_ref.actor.ppo_mini_batch_size=16\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=False\n    actor_rollout_ref.actor.use_dynamic_bsz=True\n    actor_rollout_ref.actor.use_kl_loss=True\n    actor_rollout_ref.actor.kl_loss_coef=0.001\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD}\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True\n    actor_rollout_ref.rollout.name=${rollout_name}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.25\n    actor_rollout_ref.rollout.enforce_eager=True\n    actor_rollout_ref.rollout.free_cache_engine=True\n    actor_rollout_ref.rollout.n=4\n)\n\nREF=(\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD}\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=${adv_estimator}\n)\n\nTRAINER=(\n    trainer.critic_warmup=0\n    trainer.logger='[\"console\",\"wandb\"]'\n    trainer.project_name=${project_name}\n    trainer.experiment_name=${exp_name}\n    trainer.n_gpus_per_node=8\n    trainer.nnodes=1\n    trainer.save_freq=20\n    trainer.test_freq=5\n    trainer.total_epochs=15\n)\n\n########################### Launch ###########################\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${REF[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    \"$@\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3moe-30b_megatron_lora_fp16.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\npwd=`pwd`\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\nTP=${TP:-2}\nPP=${PP:-2}\nCP=${CP:-2}\nEP=${EP:-4}\nETP=${ETP:-1}\n\nALL_OFFLOAD=${ALL_OFFLOAD:-True}\n\noptimizer_offload_fraction=1.\n\ndtype=\"float16\" # [\"bfloat16\", \"float16\"]\nrollout_name=\"vllm\"\nproject_name='verl_grpo_example_gsm8k_math_fp16'\nexp_name='qwen3_30b_a3b_megatron_lora'\nadv_estimator=grpo\n\n# Paths\nMODEL_PATH=$HOME/Qwen/Qwen3-30B-A3B-Instruct-2507\nCKPTS_DIR=${pwd}/ckpt/${exp_name}\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\n\n########################### Parameter Arrays ###########################\n\nDATA=(\n    data.train_files=${gsm8k_train_path}\n    data.val_files=${gsm8k_test_path}\n    data.train_batch_size=128\n    data.max_prompt_length=1024\n    data.max_response_length=1024\n    data.truncation='error'\n    data.filter_overlong_prompts=True\n    data.shuffle=False\n    data.return_raw_chat=$return_raw_chat\n    data.filter_overlong_prompts_workers=128\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=${MODEL_PATH}\n    actor_rollout_ref.model.lora.rank=16\n    actor_rollout_ref.model.lora.alpha=32\n    actor_rollout_ref.model.lora.dtype=${dtype}\n    actor_rollout_ref.model.use_fused_kernels=True\n)\n\nACTOR=(\n    actor_rollout_ref.actor.optim.lr=3e-6\n    actor_rollout_ref.actor.ppo_mini_batch_size=16\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2\n    actor_rollout_ref.actor.megatron.use_mbridge=True\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=False\n    actor_rollout_ref.actor.use_dynamic_bsz=True\n    actor_rollout_ref.actor.use_kl_loss=True\n    actor_rollout_ref.actor.kl_loss_coef=0.001\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.actor.megatron.param_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.grad_offload=${ALL_OFFLOAD}\n    actor_rollout_ref.actor.megatron.dtype=${dtype}\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True\n    +actor_rollout_ref.actor.megatron.override_ddp_config.grad_reduce_in_fp32=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction}\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=${ALL_OFFLOAD}\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True\n    actor_rollout_ref.rollout.name=${rollout_name}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5\n    actor_rollout_ref.rollout.enforce_eager=True\n    actor_rollout_ref.rollout.free_cache_engine=True\n    actor_rollout_ref.rollout.n=4\n    actor_rollout_ref.rollout.dtype=${dtype}\n)\n\nREF=(\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True\n    actor_rollout_ref.ref.megatron.dtype=${dtype}\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${TP}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${PP}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${EP}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${CP}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${ETP}\n    actor_rollout_ref.ref.megatron.param_offload=${ALL_OFFLOAD}\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=${adv_estimator}\n)\n\nTRAINER=(\n    trainer.critic_warmup=0\n    trainer.logger='[\"console\",\"wandb\"]'\n    trainer.project_name=${project_name}\n    trainer.experiment_name=${exp_name}\n    trainer.n_gpus_per_node=8\n    trainer.nnodes=1\n    trainer.save_freq=20\n    trainer.test_freq=5\n    trainer.total_epochs=15\n    trainer.val_before_train=False\n    trainer.max_actor_ckpt_to_keep=1\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n    trainer.log_val_generations=10\n)\n\n########################### Launch ###########################\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${REF[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    2>&1 | tee ${pwd}/log/${exp_name}_$(date +'%Y%m%d_%H%M%S').log"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3moe-30b_sglang_megatron_npu.sh",
    "content": "#!/bin/bash\nset -xeuo pipefail\n# Project Configuration\nproject_name='DAPO-Qwen3-30b-A3B-BASE-MATH'\nexp_name='DAPO-Qwen3-30B-A3B-BASE-Megatron-SGLang'\n\n# Necessary env\nexport HCCL_CONNECT_TIMEOUT=1500\nexport HCCL_HOST_SOCKET_PORT_RANGE=60000-60050\nexport HCCL_NPU_SOCKET_PORT_RANGE=61000-61050\n\nexport RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES=1\nexport ASCEND_RT_VISIBLE_DEVICES=0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,15\n\nexport DISABLE_L2_CACHE=1\nexport TASK_QUEUE_ENABLE=1\n\n# Node Info\nNNODES=${NNODES:-1}\nNPUS_PER_NODE=${NPUS_PER_NODE:-16}\n\n# Model Weights Paths\nMODEL_PATH=Qwen/Qwen3-30B-A3B\nMCORE_MODEL_PATH=Qwen/Qwen3-30B-A3B-mcore\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\n# Data Length Configuration\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\n\n# Training Batch Configuration\ntrain_prompt_bsz=16\ntrain_prompt_mini_bsz=16\nn_resp_per_prompt=8\n\n# Algorithm Configuration\nadv_estimator=grpo\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\n# Performance and Memory Management Configuration\nall_offload=True\nuse_dynamic_bsz=False\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\n\n# Megatron Parallelism Configuration\ntrain_tp=4\ntrain_ep=4\ntrain_etp=4\ntrain_pp=1\ntrain_cp=1\n\n# SGLang Generation Configuration\ngen_tp=4\ngen_dp=1\ngen_ep=1\ngpu_memory_utilization=0.5\nmax_model_len=$((max_prompt_length + max_response_length))\nmax_num_batched_tokens=$(((max_prompt_length + max_response_length) * 1))\n\n# Data Configuration\nDATA_CONFIG=(\n    # File Paths\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    # Data Structure\n    data.prompt_key=prompt\n    # Batch and Length Configuration\n    data.train_batch_size=${train_prompt_bsz}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    # Preprocessing\n    data.filter_overlong_prompts=False\n    data.truncation='left'\n)\n\n# Model Configuration\nMODEL_CONFIG=(\n    # Model Path\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    # Model Processing\n    actor_rollout_ref.model.use_remove_padding=True\n)\n\n# Reinforcement Learning Algorithm Configuration\nALGORITHM_CONFIG=(\n    # Advantage Estimation\n    algorithm.adv_estimator=${adv_estimator}\n    # KL Divergence Control\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n\nACTOR_CONFIG=(\n    # Core Runtime Settings\n    actor_rollout_ref.actor.use_torch_compile=False\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    # Loss Function Configuration\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.entropy_coeff=0\n    # PPO Training Parameters\n    actor_rollout_ref.actor.ppo_epochs=1\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    # Optimizer Settings\n    actor_rollout_ref.actor.optim.lr=1e-6\n    # Megatron Parallelism Strategy\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${train_etp}\n    # Memory Optimization\n    actor_rollout_ref.actor.megatron.param_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.optimizer_offload=${all_offload}\n    actor_rollout_ref.actor.megatron.grad_offload=${all_offload}\n    # Model Weights Management\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True\n    actor_rollout_ref.actor.megatron.use_mbridge=False\n    # Transformer Architecture Optimizations\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1\n)\n\nREF_CONFIG=(\n    # Core Runtime Settings\n    actor_rollout_ref.ref.use_torch_compile=False\n    # Log Probability Inference\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    # Megatron Parallelism Strategy\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp}\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp}\n    actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp}\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${train_ep}\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${train_etp}\n    # Memory Optimization\n    actor_rollout_ref.ref.megatron.param_offload=${all_offload}\n    # Model Weights Management\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${MCORE_MODEL_PATH}\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True\n    actor_rollout_ref.ref.megatron.use_mbridge=False\n)\n\nROLLOUT_CONFIG=(\n    # Rollout Engine\n    actor_rollout_ref.rollout.name=sglang\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.attention_backend=\"ascend\"\n    # Generation Parameters\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.top_p=1.0\n    actor_rollout_ref.rollout.top_k=-1\n    actor_rollout_ref.rollout.temperature=1.0\n    # Log Probability Inference\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    # Memory Management\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization}\n    # Parallelism Strategy\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp}\n    actor_rollout_ref.rollout.expert_parallel_size=${gen_ep}\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.enable_dp_attention=False\n    # Performance Optimization\n    +actor_rollout_ref.rollout.engine_kwargs.sglang.chunked_prefill_size=-1\n    actor_rollout_ref.rollout.enforce_eager=False\n    # Validation Generation\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.top_p=1.0\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0\n)\n\nTRAINER_CONFIG=(\n    # Logger Configuration\n    trainer.logger='[\"console\"]'\n    # Project Settings\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    # Hardware Configuration\n    trainer.nnodes=\"${NNODES}\"\n    trainer.n_gpus_per_node=\"${NPUS_PER_NODE}\"\n    trainer.device='npu'\n    # Training Schedule\n    trainer.total_epochs=15\n    trainer.val_before_train=False\n    trainer.test_freq=-1\n    trainer.save_freq=-1\n    # Checkpoint Directory\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n)\n\n# profiling configuration\nPROF_CONFIG=(\n    global_profiler.tool=npu \n    global_profiler.steps=null \n    global_profiler.save_path=/profpath \n    actor_rollout_ref.actor.profiler.enable=True \n    actor_rollout_ref.actor.profiler.ranks=\"[0]\" \n    actor_rollout_ref.actor.profiler.all_ranks=False \n    actor_rollout_ref.actor.profiler.tool_config.npu.discrete=True \n    actor_rollout_ref.actor.profiler.tool_config.npu.contents=['npu','cpu'] \n    actor_rollout_ref.actor.profiler.tool_config.npu.level=level0 \n    actor_rollout_ref.actor.profiler.tool_config.npu.analysis=True \n    actor_rollout_ref.rollout.profiler.enable=True \n    actor_rollout_ref.rollout.profiler.ranks=\"[0]\" \n    actor_rollout_ref.rollout.profiler.all_ranks=False \n)\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    \"${DATA_CONFIG[@]}\" \\\n    \"${MODEL_CONFIG[@]}\" \\\n    \"${ACTOR_CONFIG[@]}\" \\\n    \"${REF_CONFIG[@]}\" \\\n    \"${ROLLOUT_CONFIG[@]}\" \\\n    \"${ALGORITHM_CONFIG[@]}\" \\\n    \"${TRAINER_CONFIG[@]}\" \\\n    \"${PROF_CONFIG[@]}\" \\\n    \"$@\"\n"
  },
  {
    "path": "examples/grpo_trainer/run_qwen3next_80b_fsdp_npu.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name=\"verl_grpo_qwen3-next-80b\"\nexperiment_name=\"Qwen3_Next_80B_Instruct\"\n\n# Paths\nWORK_DIR=${WORK_DIR:-\"${HOME}/verl\"}\nMODEL_PATH=${WORK_DIR}/Qwen3-Next-80B-A3B-Instruct\nTRAIN_FILE=${WORK_DIR}/datasets/dapo-math-17k/dapo-math-17k.parquet\nTEST_FILE=${WORK_DIR}/datasets/aime/aime-2024.parquet\n\n# algorithm\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# batch\ntrain_batch_size=16\nrollout_n=16\nppo_mini_batch_size=8\n\n# length\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 20))\n\n# algorithm\nlearning_rate=1e-6\nwarmup_steps=0\n# enable_filter_groups=True\n\n# performance\nsp_size=8\ngen_tp=4\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\noffload=True\n\nDATA=(\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    data.train_batch_size=${train_batch_size}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.truncation='error'\n)\n\nACTOR=(\n    actor_rollout_ref.actor.strategy=fsdp2\n    actor_rollout_ref.nccl_timeout=14400\n\n    # fsdp\n    actor_rollout_ref.actor.fsdp_config.use_orig_params=True\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload}\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload}\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=False\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1\n    +actor_rollout_ref.actor.fsdp_config.mixed_precision.reduce_dtype=bf16\n\n    # optimizer\n    actor_rollout_ref.actor.optim.lr=${learning_rate}\n    actor_rollout_ref.actor.optim.lr_warmup_steps=${warmup_steps}\n\n    # ppo config\n    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size}\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size}\n\n    # entropy\n    actor_rollout_ref.actor.entropy_checkpointing=True\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True\n\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low}\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high}\n    actor_rollout_ref.actor.clip_ratio_c=10.0\n\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.actor.use_torch_compile=False\n)\n\nROLLOUT=(\n    actor_rollout_ref.rollout.name=vllm\n    actor_rollout_ref.rollout.n=${rollout_n}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8\n    actor_rollout_ref.rollout.load_format=auto\n    actor_rollout_ref.rollout.enforce_eager=True\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length))\n    actor_rollout_ref.rollout.calculate_log_probs=True\n\n    actor_rollout_ref.rollout.temperature=${temperature}\n    actor_rollout_ref.rollout.top_p=${top_p}\n    actor_rollout_ref.rollout.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature}\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p}\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n)\n\nREF=(\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size}\n    actor_rollout_ref.ref.use_torch_compile=False\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload}\n    actor_rollout_ref.ref.fsdp_config.optimizer_offload=${offload}\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=False\n\n    actor_rollout_ref.ref.entropy_checkpointing=True\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True\n\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n)\n\nTRAINER=(\n    trainer.logger='[\"console\"]'\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${experiment_name}\"\n    trainer.n_gpus_per_node=16\n    trainer.nnodes=4\n    trainer.val_before_train=False\n    trainer.save_freq=5\n    trainer.test_freq=-1\n    trainer.total_epochs=1\n    trainer.device=npu\n)\n\nMODEL=(\n    actor_rollout_ref.model.path=${MODEL_PATH}\n    actor_rollout_ref.model.use_remove_padding=True\n    actor_rollout_ref.model.enable_activation_offload=${offload}\n)\n\nALGORITHM=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n\n# =========================================================\necho \"Starting Training with:\"\necho \"Project: ${project_name}, Exp: ${experiment_name}\"\necho \"Rollout N: ${rollout_n}, Batch Size: ${train_batch_size}, LR: ${learning_rate}\"\n\n\npython3 -m verl.trainer.main_ppo \\\n    \"${DATA[@]}\" \\\n    \"${ACTOR[@]}\" \\\n    \"${ROLLOUT[@]}\" \\\n    \"${REF[@]}\" \\\n    \"${TRAINER[@]}\" \\\n    \"${ALGORITHM[@]}\" \\\n    \"${MODEL[@]}\" \\\n"
  },
  {
    "path": "examples/grpo_trainer/run_seed_oss_36b.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=64 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=ByteDance-Seed/Seed-OSS-36B-Base \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.ref.strategy=fsdp2 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\"]' \\\n    trainer.project_name='verl_grpo_seed_oss_36b' \\\n    trainer.experiment_name='seed_oss_36b' \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@"
  },
  {
    "path": "examples/gspo_trainer/run_qwen30b_gspo.sh",
    "content": "# run Qwen3-30B GSPO with new model engine\nset -x\n\nHDFS_ROOT=${HDFS_ROOT:-$PWD}\nDATA_ROOT=${DATA_ROOT:-$PWD}\n\n# wandb\nbackend=megatron # fsdp, fsdp2, megatron\nproject_name=wuxibin_gspo\nexperiment_name=qwen3-30B-base-grpo-$backend\ndefault_local_dir=$DATA_ROOT/checkpoint/$project_name/$experiment_name\n\n# ===================================== Algorithm =====================================\nadv_estimator=grpo\nloss_mode=gspo\n\n# reference policy\nuse_kl_in_reward=False\nkl_coef=0.001\nuse_kl_loss=False\nkl_loss_coef=0.001\n\nclip_ratio_low=3e-4\nclip_ratio_high=4e-4\n\nactor_lr=1e-6\ncritic_lr=2e-6\ngae_gamma=1.0\ngae_lam=0.95\ncritic_warmup=0\n\n# ===================================== Data/Model =====================================\ntrain_files=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k/data/dapo-math-17k.parquet\ntest_files=$DATA_ROOT/dataset/aime-2024.parquet\n\nactor_model_path=$HDFS_ROOT/model/Qwen3-30B-A3B-Base\ncritic_model_path=$actor_model_path\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\ntrain_batch_size=256\nppo_mini_batch_size=32\nn_resp_per_prompt=16\nn_resp_per_prompt_val=1\n\n# ===================================== Training =====================================\nactor_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 3))\ncritic_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 4))\n\n# FSDP parallelism config\nUSP_SIZE=4\nACTOR_FSDP_CONFIG=\"\n    actor_rollout_ref.actor.fsdp_config.strategy=$backend \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$USP_SIZE\"\n\n# Megatron parallelism config\nTP_SIZE=2\nCP_SIZE=1\nPP_SIZE=1\nVPP_SIZE=null\nEP_SIZE=8\nETP_SIZE=1\nACTOR_MEGATRON_CONFIG=\"\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP_SIZE \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP_SIZE \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP_SIZE \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$VPP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP_SIZE \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True\"\n\n# Actor model config\nACTOR_CONFIG=\"\n    actor_rollout_ref.actor.optim.lr=$actor_lr \\\n    actor_rollout_ref.model.path=$actor_model_path \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \\\n    actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode}\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu\"\n\n# Critic model config\nCIRITC_CONFIG=\"\n    critic.optim.lr=$critic_lr \\\n    critic.model.path=$critic_model_path \\\n    critic.model.use_remove_padding=True \\\n    critic.ppo_max_token_len_per_gpu=$critic_max_token_len_per_gpu \\\n    critic.ulysses_sequence_parallel_size=$USP_SIZE\"\n\nCRITIC_FSDP_CONFIG=\"${ACTOR_FSDP_CONFIG//actor_rollout_ref.actor/critic.model}\"\nCRITIC_MEGATRON_CONFIG=\"${ACTOR_MEGATRON_CONFIG//actor_rollout_ref.actor/critic}\"\n\nif [[ $backend == \"megatron\" ]]; then\n    CONFIG_NAME=ppo_megatron_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_MEGATRON_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_MEGATRON_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nelse # fsdp, fsdp2\n    CONFIG_NAME=ppo_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_FSDP_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_FSDP_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nfi\n\n# ===================================== Inference =====================================\nrollout_name=vllm\nif [ \"$rollout_name\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\ninfer_tp=4\ninfer_dp=1\ninfer_ep=1\ngpu_memory_utilization=0.8\n\nROLLOUT_CONFIG=\"\n    actor_rollout_ref.rollout.name=$rollout_name \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \\\n    actor_rollout_ref.rollout.data_parallel_size=$infer_dp \\\n    actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \\\n    actor_rollout_ref.rollout.n=$n_resp_per_prompt \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val\"\n\n# ===================================== Reward =====================================\nREWARD_CONFIG=\"\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length}\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=./config \\\n    --config-name=$CONFIG_NAME \\\n    algorithm.adv_estimator=$adv_estimator \\\n    algorithm.use_kl_in_reward=$use_kl_in_reward \\\n    algorithm.kl_ctrl.kl_coef=$kl_coef \\\n    algorithm.gamma=$gae_gamma \\\n    algorithm.lam=$gae_lam \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=$train_batch_size \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.filter_overlong_prompts_workers=64 \\\n    data.truncation='error' \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=$critic_warmup \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$experiment_name \\\n    trainer.default_local_dir=$default_local_dir \\\n    trainer.n_gpus_per_node=$ARNOLD_WORKER_GPU \\\n    trainer.nnodes=$ARNOLD_WORKER_NUM \\\n    trainer.val_before_train=False \\\n    trainer.log_val_generations=100 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=500 \\\n    $ACTOR_CONFIG \\\n    $CIRITC_CONFIG \\\n    $ROLLOUT_CONFIG \\\n    $REWARD_CONFIG\n"
  },
  {
    "path": "examples/gspo_trainer/run_qwen3_32b_gspo_npu.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\nmkdir -p logs\nulimit -n 32768\n\n## Basic Environment Settings\nexport RAY_DEDUP_LOGS=0\nexport HYDRA_FULL_ERROR=1\nexport TASK_QUEUE_ENABLE=1\nexport HCCL_EXEC_TIMEOUT=3600\nexport HCCL_CONNECT_TIMEOUT=3600\nexport HCCL_ASYNC_ERROR_HANDLING=0\nexport CPU_AFFINITY_CONF=1\nexport VLLM_USE_V1=1\nexport VLLM_ATTENTION_BACKEND=XFORMERS\nexport VLLM_ASCEND_ENABLE_FLASHCOMM=1\nexport VLLM_ASCEND_ENABLE_PREFETCH_MLP=1\nexport VLLM_ASCEND_ENABLE_DENSE_OPTIMIZE=1\nexport LD_PRELOAD=/usr/local/lib/libjemalloc.so.2\n\n# Project Configuration\nproject_name='GSPO-Qwen3-32B-BASE-MATH'\nexp_name='GSPO-Qwen3-32B-BASE-Megatron-vLLM'\n\n# Node Info\nNNODES=${NNODES:-4}\nNPUS_PER_NODE=${NPUS_PER_NODE:-16}\n\n# Model Weights Paths\nMODEL_PATH=Qwen/Qwen3-32B\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n\n# File System Paths\nTRAIN_FILE=$RAY_DATA_HOME/dataset/dapo-math-17k.parquet\nTEST_FILE=$RAY_DATA_HOME/dataset/aime-2024.parquet\n\n# Ray Configuration\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n\n# Data Length Configuration\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\n\n# Training Batch Configuration\ntrain_prompt_bsz=256\ngen_prompt_bsz=$((train_prompt_bsz * 1))\ntrain_prompt_mini_bsz=64\nn_resp_per_prompt=16\n\n# GSPO Loss Configuration\nadv_estimator=grpo\nloss_mode=gspo\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\nclip_ratio_low=0.0003\nclip_ratio_high=0.0004\nloss_agg_mode=\"seq-mean-token-mean\"\n\n# FSDP Parallelism Configuration\nactor_strategy=fsdp2\nref_strategy=fsdp2\nsp_size=4\nfsdp_size=-1\n\n# Performance and Memory Management Configuration\noffload=True\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\n\n# vLLM Configuration\ngen_tp=4\ngpu_memory_utilization=0.7\nmax_model_len=$((max_prompt_length + max_response_length))\nmax_num_batched_tokens=$((max_prompt_length + max_response_length))\n\n\n# Data Configuration\nDATA_CONFIG=(\n    data.train_files=\"${TRAIN_FILE}\"\n    data.val_files=\"${TEST_FILE}\"\n    data.prompt_key=prompt\n    data.train_batch_size=${train_prompt_bsz}\n    +data.gen_batch_size=${gen_prompt_bsz}\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.truncation='left'\n)\n\n# Model Configuration\nMODEL_CONFIG=(\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.model.use_remove_padding=True\n    actor_rollout_ref.model.enable_gradient_checkpointing=True\n)\n\n# Algorithm Configuration\nALGORITHM_CONFIG=(\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n)\n\n# Actor Model Configuration\nACTOR_CONFIG=(\n    actor_rollout_ref.actor.use_torch_compile=False\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.actor.strategy=${actor_strategy}\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode}\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low}\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high}\n    actor_rollout_ref.actor.clip_ratio_c=10.0\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode}\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.grad_clip=1.0\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len}\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10\n    actor_rollout_ref.actor.optim.weight_decay=0.1\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size}\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size}\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload}\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload}\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=True\n    actor_rollout_ref.actor.entropy_checkpointing=True\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True\n)\n\n# Reference Model Configuration\nREF_CONFIG=(\n    actor_rollout_ref.ref.use_torch_compile=False\n    actor_rollout_ref.ref.strategy=${ref_strategy}\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size}\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload}\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=True\n    actor_rollout_ref.ref.entropy_checkpointing=True\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True\n)\n\n# Rollout Configuration\nROLLOUT_CONFIG=(\n    actor_rollout_ref.rollout.name=vllm\n    actor_rollout_ref.rollout.calculate_log_probs=True\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.top_p=1.0\n    actor_rollout_ref.rollout.top_k=-1\n    actor_rollout_ref.rollout.temperature=1.0\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz}\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len}\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization}\n    actor_rollout_ref.rollout.max_num_batched_tokens=${max_num_batched_tokens}\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp}\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.enforce_eager=False\n    actor_rollout_ref.rollout.free_cache_engine=True\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_capture_sizes=\"[8, 16, 32, 64, 128, 192, 256]\"\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_DECODE_ONLY\"\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.7\n    actor_rollout_ref.rollout.val_kwargs.top_k=-1\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096\n)\n\n# Trainer Configuration\nTRAINER_CONFIG=(\n    trainer.logger='[\"console\"]'\n    trainer.project_name=\"${project_name}\"\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.nnodes=\"${NNODES}\"\n    trainer.n_gpus_per_node=\"${NPUS_PER_NODE}\"\n    trainer.device='npu'\n    trainer.total_epochs=10\n    trainer.val_before_train=False\n    trainer.test_freq=-1\n    trainer.save_freq=100\n    trainer.default_local_dir=\"${CKPTS_DIR}\"\n    trainer.resume_mode=auto\n    trainer.balance_batch=True\n)\n\n# Main GSPO Training Command\npython3 -m verl.trainer.main_ppo \\\n    \"${DATA_CONFIG[@]}\" \\\n    \"${MODEL_CONFIG[@]}\" \\\n    \"${ACTOR_CONFIG[@]}\" \\\n    \"${REF_CONFIG[@]}\" \\\n    \"${ROLLOUT_CONFIG[@]}\" \\\n    \"${ALGORITHM_CONFIG[@]}\" \\\n    \"${TRAINER_CONFIG[@]}\" \\\n    \"$@\" | tee logs/run_qwen3_32b_gspo_megatron_vllm_npu.log"
  },
  {
    "path": "examples/gspo_trainer/test_gspo_3b_math.sh",
    "content": "#!/usr/bin/env bash\n#SBATCH --job-name=rl-gspo-3B\n#SBATCH --partition=main\n#SBATCH --nodes=1                # Number of nodes\n#SBATCH --ntasks-per-node=1      # One task per node\n#SBATCH --cpus-per-task=128      # cpu-cores per task\n#SBATCH --gres=gpu:8\n#SBATCH --mem=0\n#SBATCH --exclusive\n#SBATCH --time=500:00:00\n#SBATCH --output=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.out\n#SBATCH --error=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.err\n\nset -xeuo pipefail\n\n# activate the venv\necho \"Activating verl environment...\"\neval \"$(conda shell.bash hook)\"\nconda deactivate\nconda activate verl\n\n# can make training faster, depends on your infrastructure\nexport NCCL_IBEXT_DISABLE=1\nexport NCCL_NVLS_ENABLE=1\nexport NCCL_IB_HCA=mlx5\nexport UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1\n\n# Set how many GPUs we actually have on this node.\nexport GPUS_PER_NODE=8\n\nNNODES=${SLURM_JOB_NUM_NODES}\nexport NNODES\n\nexport VLLM_ATTENTION_BACKEND=FLASH_ATTN\nexport RAY_LOGGING_LEVEL=DEBUG\nexport HYDRA_FULL_ERROR=1\nexport WANDB_API_KEY=... # your wandb API key\n\necho \"Using $NNODES nodes for training...\"\n\n# ------------------------------------- Setup xp params ---------------------------------------\nproject_name='RL-GSPO'\n\nadv_estimator=grpo\nloss_mode=gspo\nloss_agg_mode=\"seq-mean-token-mean\"\nMODEL_PATH=Qwen/Qwen2.5-3B-Instruct\noffload=false # it's a small model, offloading will just slow-down training\nrollout_engine=vllm\nrollout_mode=async\nreturn_raw_chat=\"True\"\nif [ \"$rollout_engine\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\ngpu_memory_utilization=0.8\nreward_manager=dapo\nadv_estimator=grpo\nshuffle_dataset=true\nfirst_time_dataset_prep=true # prepare dataset\n\ntest_freq=10\nsave_freq=10\ntotal_epochs=10\ntotal_training_steps=500\nval_before_train=false\n\nuse_kl_in_reward=false\nkl_coef=0.0\nuse_kl_loss=false\nkl_loss_coef=0.0\n\nclip_ratio_low=0.0003 # as recommended by the paper, see Sec. 5.1\nclip_ratio_high=0.0004 # as recommended by the paper, see Sec. 5.1\ntrain_batch_size=512\nppo_mini_batch_size=128 # maintain 4 mini-batches as recommended by the paper, see Sec. 5.1\nppo_micro_batch_size_per_gpu=8 # setup depending on your GPU memory\nn_resp_per_prompt=16\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\n# dapo reward manager params\nenable_overlong_buffer=false # true\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Paths and namings\nSFT_MODEL=$(basename $MODEL_PATH)\nexp_name=\"${loss_mode}-epslow-${clip_ratio_low}-epshigh-${clip_ratio_high}-${SFT_MODEL}-RL\"\nCKPTS_DIR=/rl/checkpoints/experimental/4b/${loss_mode}/${exp_name}\n\n# Sampling params at rollouts\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=1\nuse_dynamic_bsz=true\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=true\ngen_tp=1\nentropy_checkpointing=true # This enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training.\n\n# ------------------------------------- train/val data preparation ---------------------------------------\nif [ \"$first_time_dataset_prep\" = true ]; then\n    echo \"Preprocessing GSM8K dataset...\"\n    python examples/data_preprocess/gsm8k.py --local_save_dir /data/gsm8k/\nfi\n\ngsm8k_train_path=/data/gsm8k/train.parquet\ngsm8k_test_path=/data/gsm8k/test.parquet\n\n# set the paths\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    data.train_files=\"${train_files}\" \\\n    data.val_files=\"${test_files}\" \\\n    data.shuffle=$shuffle_dataset \\\n    data.prompt_key=prompt \\\n    data.truncation='error' \\\n    data.filter_overlong_prompts=true \\\n    data.return_raw_chat=${return_raw_chat} \\\n    data.train_batch_size=${train_batch_size} \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.model.use_remove_padding=true \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.name=${rollout_engine} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=true \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.05 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=true \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=true \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.entropy_checkpointing=${entropy_checkpointing} \\\n    reward.reward_manager.name=${reward_manager} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=false \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${GPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=${val_before_train} \\\n    trainer.test_freq=${test_freq} \\\n    trainer.save_freq=${save_freq} \\\n    trainer.total_epochs=${total_epochs} \\\n    trainer.total_training_steps=${total_training_steps} \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=2 \\\n    $@\n"
  },
  {
    "path": "examples/gspo_trainer/test_gspo_3b_math_slurm.sh",
    "content": "#!/usr/bin/env bash\n#SBATCH --job-name=rl-gspo-3B\n#SBATCH --partition=main\n#SBATCH --nodes=1                # Number of nodes\n#SBATCH --ntasks-per-node=1      # One task per node\n#SBATCH --cpus-per-task=128      # cpu-cores per task\n#SBATCH --gres=gpu:8\n#SBATCH --mem=0\n#SBATCH --exclusive\n#SBATCH --time=500:00:00\n#SBATCH --output=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.out\n#SBATCH --error=/rl/logs/Qwen2.5-3B/gspo/math/vllm_%x_%j.err\n\nset -xeuo pipefail\n\n# activate the venv\necho \"Activating verl environment...\"\neval \"$(conda shell.bash hook)\"\nconda deactivate\nconda activate verl\n\n# can make training faster, depends on your infrastructure\nexport NCCL_IBEXT_DISABLE=1\nexport NCCL_NVLS_ENABLE=1\nexport NCCL_IB_HCA=mlx5\nexport UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1\n\n# Set how many GPUs we actually have on this node.\nexport GPUS_PER_NODE=8\n\nNNODES=${SLURM_JOB_NUM_NODES}\nexport NNODES\n\nexport VLLM_ATTENTION_BACKEND=FLASH_ATTN\nexport RAY_memory_monitor_refresh_ms=0\nexport RAY_LOGGING_LEVEL=DEBUG\nexport HYDRA_FULL_ERROR=1\nexport WANDB_API_KEY=... # your wandb API key\n\n# Let Ray know how many nodes to expect\nexport RAY_NUM_NODES=$NNODES\n\necho \"Using $NNODES nodes for training...\"\n\n# ------------------------------------- Setup xp params ---------------------------------------\nproject_name='RL-GSPO'\n\nadv_estimator=grpo\nloss_mode=gspo\nloss_agg_mode=\"seq-mean-token-mean\"\nMODEL_PATH=Qwen/Qwen2.5-3B-Instruct\noffload=false # it's a small model, offloading will just slow-down training\nrollout_engine=vllm\nrollout_mode=async\nreturn_raw_chat=\"True\"\nif [ \"$rollout_engine\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\ngpu_memory_utilization=0.8\nreward_manager=dapo\nadv_estimator=grpo\nshuffle_dataset=true\nfirst_time_dataset_prep=true # prepare dataset\n\ntest_freq=10\nsave_freq=10\ntotal_epochs=10\ntotal_training_steps=500\nval_before_train=false\n\nuse_kl_in_reward=false\nkl_coef=0.0\nuse_kl_loss=false\nkl_loss_coef=0.0\n\nclip_ratio_low=0.0003 # as recommended by the paper, see Sec. 5.1\nclip_ratio_high=0.0004 # as recommended by the paper, see Sec. 5.1\ntrain_batch_size=512\nppo_mini_batch_size=128 # maintain 4 mini-batches as recommended by the paper, see Sec. 5.1\nppo_micro_batch_size_per_gpu=8 # setup depending on your GPU memory\nn_resp_per_prompt=16\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\n# dapo reward manager params\nenable_overlong_buffer=false # true\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Paths and namings\nSFT_MODEL=$(basename $MODEL_PATH)\nexp_name=\"${loss_mode}-epslow-${clip_ratio_low}-epshigh-${clip_ratio_high}-${SFT_MODEL}-RL\"\nCKPTS_DIR=/rl/checkpoints/experimental/4b/${loss_mode}/${exp_name}\n\n# Sampling params at rollouts\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nsp_size=1\nuse_dynamic_bsz=true\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=true\ngen_tp=1\nentropy_checkpointing=true # This enables entropy recomputation specifically for the entropy calculation, lowering memory usage during training.\n\n# ------------------------------------- train/val data preparation ---------------------------------------\nif [ \"$first_time_dataset_prep\" = true ]; then\n    echo \"Preprocessing GSM8K dataset...\"\n    python examples/data_preprocess/gsm8k.py --local_save_dir /data/gsm8k/\nfi\n\ngsm8k_train_path=/data/gsm8k/train.parquet\ngsm8k_test_path=/data/gsm8k/test.parquet\n\n# set the paths\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    data.train_files=\"${train_files}\" \\\n    data.val_files=\"${test_files}\" \\\n    data.shuffle=$shuffle_dataset \\\n    data.prompt_key=prompt \\\n    data.truncation='error' \\\n    data.filter_overlong_prompts=true \\\n    data.return_raw_chat=${return_raw_chat} \\\n    data.train_batch_size=${train_batch_size} \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.model.use_remove_padding=true \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.name=${rollout_engine} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=true \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.05 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=${gpu_memory_utilization} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=true \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=true \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.entropy_checkpointing=${entropy_checkpointing} \\\n    reward.reward_manager.name=${reward_manager} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=false \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${GPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=${val_before_train} \\\n    trainer.test_freq=${test_freq} \\\n    trainer.save_freq=${save_freq} \\\n    trainer.total_epochs=${total_epochs} \\\n    trainer.total_training_steps=${total_training_steps} \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=2 \\\n    $@\n"
  },
  {
    "path": "examples/gspo_trainer/test_gspo_qwen30b_a3b_ep.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nexport NCCL_DEBUG=WARN\n# export VERL_LOGGING_LEVEL=DEBUG\n\nproject_name='DAPO'\nexp_name='GSPO-Qwen3-30B-A3B-Base-MATH'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=3e-4\nclip_ratio_high=4e-4\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\nloss_mode=gspo\n\ntrain_prompt_bsz=256\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\n# RAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# MODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base\"}\n# CKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\n# TRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\n# TEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nMODEL_PATH=$HDFS_ROOT/model/Qwen3-30B-A3B-Base\nCKPTS_DIR=$DATA_ROOT/checkpoint/${project_name}/${exp_name}\nTRAIN_FILE=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k/data/dapo-math-17k.parquet\naime24_test_path=$DATA_ROOT/dataset/aime-2024.parquet\n\nTEST_FILE=\"['$aime24_test_path']\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 1))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\n\n# gen\nrollout_name=vllm # vllm or sglang\nif [ \"$rollout_name\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\ngen_tp=1\ngen_dp=4\ngen_ep=4\n\n# train\ntrain_tp=4\ntrain_pp=1\nEP=4\nETP=1\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.return_raw_chat=True \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.data_parallel_size=${gen_dp} \\\n    actor_rollout_ref.rollout.expert_parallel_size=${gen_ep} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}-tp${gen_tp}-ep${gen_ep}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=30 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=300 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "examples/mtp_trainer/runtime_env.yaml",
    "content": "working_dir: ./\n\nexcludes:\n  - \".git/\"\n\nenv_vars:\n  VLLM_USE_V1: \"1\"\n  HYDRA_FULL_ERROR: \"1\"\n  NCCL_NVLS_ENABLE: \"0\"\n  NCCL_SOCKET_IFNAME: \"eth0\"\n  TMPDIR: \"/tmp\"\n  CUDA_HOME: \"/usr/local/cuda\"\n  CUDA_TMPDIR: \"/tmp\"\n  CUDA_CACHE_PATH: \"/tmp/cuda_cache\"\n  # For distributed training, the path must be set on a distributed file system (DFS) to ensure visibility across all nodes.\n  HF_HOME: \"/tmp/hf_home_mimo\"\n  PYTHONPATH: \"/tmp/hf_home_mimo/modules/\"\n"
  },
  {
    "path": "examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron.sh",
    "content": "#!/usr/bin/env bash\n\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-mimo-7b-rl-megatron'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=128\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/examples/mtp_trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-16}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/MiMo-7B-RL\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=4\ntrain_tp=2\ntrain_pp=2\ntrain_cp=2\n\ncommon_params=(\nactor_rollout_ref.model.mtp.enable=True\nactor_rollout_ref.model.mtp.enable_train=True\nactor_rollout_ref.model.mtp.mtp_loss_scaling_factor=0.1\nactor_rollout_ref.model.mtp.detach_encoder=True\n)\n\npython -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger='[\"console\",\"tensorboard\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=False \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    actor_rollout_ref.rollout.disable_log_stats=False \\\n    actor_rollout_ref.rollout.prometheus.enable=True \\\n    actor_rollout_ref.rollout.prometheus.port=44398 \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    data.trust_remote_code=True \\\n    trainer.total_training_steps=400 \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    \"${common_params[@]}\"\n"
  },
  {
    "path": "examples/mtp_trainer/test_dapo_mimo_7b_with_mtp_math_megatron_4_4.sh",
    "content": "#!/usr/bin/env bash\n\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-mimo-7b-rl-megatron'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/examples/mtp_trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-16}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/MiMo-7B-RL\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=False\ngen_tp=2\ntrain_tp=2\ntrain_pp=1\ntrain_cp=1\n\ntrain_prompt_bsz=128\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\nmtp_params=(\n  actor_rollout_ref.actor.megatron.use_mbridge=True\n  actor_rollout_ref.model.mtp.enable=True\n  actor_rollout_ref.model.mtp.enable_train=True\n  actor_rollout_ref.model.mtp.mtp_loss_scaling_factor=0.1\n  actor_rollout_ref.model.mtp.detach_encoder=True\n  actor_rollout_ref.model.mtp.enable_rollout=True\n  )\n\nfully_async=(\n  data.train_batch_size=0\n  data.gen_batch_size=1\n  trainer.test_freq=10\n  actor_rollout_ref.hybrid_engine=False\n  actor_rollout_ref.rollout.calculate_log_probs=True\n  actor_rollout_ref.actor.optim.lr_decay_steps=51200\n  rollout.total_rollout_steps=$(((512*100)))\n  trainer.nnodes=1\n  trainer.n_gpus_per_node=4\n  rollout.nnodes=1\n  rollout.n_gpus_per_node=4\n  async_training.staleness_threshold=0.5\n  async_training.trigger_parameter_sync_step=4\n  async_training.require_batches=1\n  async_training.partial_rollout=True\n)\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    --config-path=config \\\n    --config-name='fully_async_ppo_megatron_trainer.yaml'\\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${train_cp} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${train_cp} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    reward_model.reward_manager=dapo \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward_model.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward_model.reward_kwargs.max_resp_len=${max_response_length} \\\n    actor_rollout_ref.rollout.disable_log_stats=False \\\n    actor_rollout_ref.rollout.prometheus.enable=True \\\n    actor_rollout_ref.rollout.prometheus.port=44398 \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    data.trust_remote_code=True \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.total_epochs=10 \\\n    \"${mtp_params[@]}\" \\\n    \"${fully_async[@]}\""
  },
  {
    "path": "examples/otb_trainer/run_qwen2_5-7b.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=optimal_token_baseline \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=128 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.calculate_sum_pi_squared=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.75 \\\n    actor_rollout_ref.rollout.n=8 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5-7b-otb' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/README.md",
    "content": "# Proximal Policy Optimization (PPO)\n\nProximal Policy Optimization (PPO) is a family of policy gradient methods for reinforcement learning, proposed by OpenAI in 2017. PPO strikes a balance between simplicity, stability, and performance, making it one of the most widely used algorithms in modern RL applications, including large-scale language model fine-tuning.\n\nTraditional policy gradient methods like REINFORCE or Vanilla Policy Gradient suffer from:\n\n- High variance and sample inefficiency.\n- Instability due to large policy updates.\n\nPPO addresses this problem using a clipped surrogate objective that avoids overly large updates without requiring second-order derivatives.\n\nFor more technical details regarding PPO, we suggest reading the introduction in the [OpenAI spinning up tutorial](https://spinningup.openai.com/en/latest/algorithms/ppo.html), and the paper [Proximal Policy Optimization Algorithms](https://arxiv.org/abs/1707.06347).\n\n## Key Components\n\n- Actor-Critic Architecture: PPO requires both an actor model (policy) and a critic model (value function). This differs from other algorithms like GRPO and RLOO that don't require a critic model.\n\n- Generalized Advantage Estimation (GAE): PPO uses GAE for computing advantage values, which helps reduce variance in policy gradient estimates while maintaining low bias.\n\n- Clipped Surrogate Objective: The core of PPO is implemented through the clipped surrogate objective function that limits policy updates.\n\n## Configuration\n\nNote that all configs containing `micro_batch_size` are used to configure the maximum sample or token count per forward or backward pass to avoid GPU OOMs, whose value should not change algorithmic/convergence behavior.\n\nMost critic configs are similar to those of actors. Note that the critic model is omitted from the figure below.\n\n![image](https://github.com/user-attachments/assets/16aebad1-0da6-4eb3-806d-54a74e712c2d)\n\n- `data.train_batch_size`: The global batch size of prompts used to generate a set of sampled trajectories/rollouts. The number of responses/trajectories is `data.train_batch_size * actor_rollout.ref.rollout.n`\n\n- `actor_rollout_ref.actor.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO actor updates. The ppo_mini_batch_size is a global size across all workers\n\n- `critic.ppo_mini_batch_size`: The set of sampled trajectories is split into multiple mini-batches with batch_size=ppo_mini_batch_size for PPO critic updates. The ppo_mini_batch_size is a global size across all workers\n\n- `actor_rollout_ref.actor.clip_ratio`: The PPO clip range. Default to 0.2\n\n- `actor_rollout_ref.actor.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for actor\n\n- `critic.ppo_epochs`: Number of epochs for PPO updates on one set of sampled trajectories for critic. Defaults to `actor_rollout_ref.actor.ppo_epochs`\n\n- `algorithm.gamma`: discount factor\n\n- `algorithm.lam`: The lambda term that trades off between bias and variance in the GAE estimator\n\n- `algorithm.adv_estimator`: Support gae, grpo, reinforce_plus_plus, reinforce_plus_plus_baseline, rloo, rloo_vectorized\n\n## Advanced Extensions\n\n### KL Divergence Control\n\nOptions to prevent the policy from diverging too far from a reference policy. Two mechanisms are available: KL reward penalty and KL loss. For more technical details, see [Training language models to follow instructions with human feedback](https://arxiv.org/abs/2203.02155)\n\nOptions to use KL loss for KL divergence control: \n\n- `actor_rollout_ref.actor.use_kl_loss`: to use kl loss in the actor. When used, we are not applying KL in the reward function. Default is False\n\n- `actor_rollout_ref.actor.kl_loss_coef`: The coefficient of kl loss. Default is 0.001.\n\n- `actor_rollout_ref.actor.kl_loss_type`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. Appending \"+\" in the end (e.g., 'k1+' and 'k3+') would apply straight through to employ k2 for unbiased gradient estimation, regardless of the kl value estimation (see https://github.com/volcengine/verl/pull/2953#issuecomment-3162113848 for more details). How to calculate the kl divergence between actor and reference policy. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\nOptions to use KL penalty in the reward:\n\n- `algorithm.use_kl_in_reward`: Whether to enable in-reward kl penalty. Default is False.\n\n- `algorithm.kl_penalty`: Support kl(k1), abs, mse(k2), low_var_kl(k3) and full. This defines the way to calculate the kl divergence between actor and reference policy. For specific options, refer to `kl_penalty` in core_algos.py. See this blog post for detailed analysis: http://joschu.net/blog/kl-approx.html\n\n- `algorithm.kl_ctrl.kl_coef`: The (initial) coefficient of in-reward kl_penalty. Default is 0.001.\n- `algorithm.kl_ctrl.type`: 'fixed' for FixedKLController and 'adaptive' for AdaptiveKLController.\n- `algorithm.kl_ctrl.horizon`: See source code of AdaptiveKLController for details.\n- `algorithm.kl_ctrl.target_kl`: See source code of AdaptiveKLController for details.\n\n### Dual-clip PPO\n\nThe Dual-Clip PPO introduces a approach by applying a lower bound to the policy ratio when the advantage is less than zero, when multiplied by a large raito, does not exceed a specified lower bound.\n\n![image](https://github.com/user-attachments/assets/fc232181-d8b0-4307-8dd2-4dc0a4c1c139)\n\n- `actor_rollout_ref.actor.clip_ratio_c`: lower bound of the value for Dual-clip PPO, defaults to 3.0\n\n## Reference Example\n\nQwen2.5 training log and commands: [link](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log)\n\n```bash\nbash run_gemma.sh\n  trainer.n_gpus_per_node=1 \\\n  actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n  trainer.logger=console \\\n  critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  data.train_batch_size=256 \\\n  actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n  actor_rollout_ref.actor.ppo_micro_batch_size=2 \\\n  critic.ppo_micro_batch_size=2\n```\n\nReference performance with verl v0.2:\n\n| Model                          | Method          | Score | Link                                                                                           |\n|-------------------------------|------------------|-------|------------------------------------------------------------------------------------------------|\n| Qwen/Qwen2.5-0.5B-Instruct     | pretrained model | 36.4  | [Qwen Blog](https://qwenlm.github.io/blog/qwen2.5-llm/)                                        |\n| Qwen/Qwen2.5-0.5B-Instruct     | PPO              | 56.7  | [PPO Command and Logs](https://github.com/eric-haibin-lin/verl-data/blob/experiments/gsm8k/Qwen2.5-0.5B-bsz256_2-prompt1024-resp512-0.567.log) |\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek7b_llm.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=1 \\\n    trainer.use_legacy_worker_impl=auto \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek7b_llm_modelscope.sh",
    "content": "set -x\n\nVERL_USE_MODELSCOPE=True \\\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=1 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek7b_llm_pfppo.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    algorithm.use_pf_ppo=True \\\n    algorithm.pf_ppo.reweight_method=pow \\  # [\"pow\", \"max_min\", \"max_random\"]\n    algorithm.pf_ppo.weight_pow=2.0 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.rollout.n=5 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=1 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek7b_llm_sandbox_fusion.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    reward.sandbox_fusion.url='https://xxxxxxxxx.apigateway-cn-beijing.volceapi.com/run_code' \\\n    reward.sandbox_fusion.max_concurrent=128 \\\n    reward.reward_manager.name=prime \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/Eurus-2-RL-Data/train.parquet \\\n    data.val_files=$HOME/data/Eurus-2-RL-Data/validation.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_sandbox_fusion' \\\n    trainer.experiment_name='deepseek_llm_7b_function_sandbox_fusion' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=1 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek7b_llm_sp2.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=64 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    critic.optim.lr=1e-5 \\\n    critic.ulysses_sequence_parallel_size=2 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=64 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm_sp2' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek_full_hh_rlhf.sh",
    "content": "set -x\n\ntrain_files=$HOME/data/full_hh_rlhf/rl/train.parquet\ntest_files=$HOME/data/full_hh_rlhf/rl/train.parquet # no use\n\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=128 \\\n    data.max_response_length=128 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=deepseek-ai/deepseek-llm-7b-chat \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=4 \\\n    reward.reward_model.rollout.prompt_length=256 \\\n    reward.reward_model.rollout.response_length=128 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_megatron_full_hh_rlhf_examples' \\\n    trainer.experiment_name='deepseek_llm_7b_model_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=100 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek_math_gsm8k_megatron.sh",
    "content": "set -x\n\n# Example runnable on H20 * 8\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_ppo_gsm8k_math_examples' \\\n    trainer.experiment_name='deepseek_llm_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=100 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_deepseek_math_gsm8k_megatron_nsys.sh",
    "content": "set -x\n\n# Example runnable on H20 * 8\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=${train_files:-\"$gsm8k_train_path\"}\ntest_files=${test_files:-\"$gsm8k_test_path\"}\n\n# Nsight profiling configuration\nPROFILE_STEPS=\"[1]\" # or [] or null\nPROFILE_RANKS_ALL=False # or True\nPROFILE_RANKS=[0,4]\nDISCRETE=True  # or True\n\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    critic.profiler.enable=True \\\n    critic.profiler.ranks=$PROFILE_RANKS \\\n    critic.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_ppo_gsm8k_math_examples' \\\n    trainer.experiment_name='deepseek_llm_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=100 \\\n    trainer.total_training_steps=1 \\\n    global_profiler.tool=nsys \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.global_tool_config.nsys.discrete=$DISCRETE $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_gemma.sh",
    "content": "set -x\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=google/gemma-2-2b-it \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=False \\\n    critic.model.path=google/gemma-2-2b-it \\\n    critic.model.enable_gradient_checkpointing=False \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example' \\\n    trainer.experiment_name='gemma2b_function_rm' \\\n    trainer.n_gpus_per_node=2 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_moonlight16b_a3b_gsm8k_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\n\n# 0. download the model\nhf download moonshotai/Moonlight-16B-A3B-Instruct\n\n# 1. convert the model to mcore format\n# change the HF_MODEL_PATH and DIST_CKPT_PATH to your own path\nHF_MODEL_PATH=/data/models/moonshotai/Moonlight-16B-A3B-Instruct\nDIST_CKPT_PATH=/data/mcore_ckpt/Moonlight-16B-A3B-Instruct\npython scripts/converter_hf_to_mcore.py --hf_model_path $HF_MODEL_PATH --output_path $DIST_CKPT_PATH\n\n\n# 2. run the script\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\ntrain_files=$gsm8k_train_path\ntest_files=$gsm8k_test_path\n\nALL_OFFLOAD=${ALL_OFFLOAD:-False}\nCOMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD}\n\nACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nREF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_PARAM_OFFLOAD=${CRITIC_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_GRAD_OFFLOAD=${CRITIC_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nCRITIC_OPTIMIZER_OFFLOAD=${CRITIC_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nRM_PARAM_OFFLOAD=${RM_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\n\n\nNODES=4\nPP=2\nTP=8\nEP=8\nETP=1\nVLLM_TP=4\n\n# RAY_ADDRESS='auto' ray job submit --working-dir . -- \npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.trust_remote_code=True \\\n    actor_rollout_ref.model.path=$LLM \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=$LLM \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_megatron_gsm8k_examples' \\\n    trainer.experiment_name='moonlight_16b_a3b_instruct_1node' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=$NODES \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    critic.model.trust_remote_code=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.num_layers_in_last_pipeline_stage=13 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$VLLM_TP \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \\\n    critic.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \\\n    critic.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    critic.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    critic.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} \\\n    actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} \\\n    critic.megatron.param_offload=${CRITIC_PARAM_OFFLOAD} \\\n    critic.megatron.optimizer_offload=${CRITIC_OPTIMIZER_OFFLOAD} \\\n    critic.megatron.grad_offload=${CRITIC_GRAD_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \\\n    critic.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    critic.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    trainer.val_before_train=False \\\n    trainer.total_epochs=100 $@\n    "
  },
  {
    "path": "examples/ppo_trainer/run_qwen1.5_moe_a2.7b-gsm8k_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\n# 0. download the model\n#hf download Qwen/Qwen1.5-MoE-A2.7B-Chat\n\n# 1. convert the model to mcore format\n# change the HF_MODEL_PATH and DIST_CKPT_PATH to your own path\nHF_MODEL_PATH=/data/models/Qwen/Qwen1.5-MoE-A2.7B-Chat\nDIST_CKPT_PATH=/data/mcore_ckpt/Qwen1.5-MoE-A2.7B-Chat\npython scripts/converter_hf_to_mcore.py --hf_model_path $HF_MODEL_PATH --output_path $DIST_CKPT_PATH\n\n# 2. run the script\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\ntrain_files=$gsm8k_train_path\ntest_files=$gsm8k_test_path\n\nNODES=4\nPP=2\nTP=4\nCP=1\nVLLM_TP=4\n\n# RAY_ADDRESS='auto' ray job submit --working-dir . -- \npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=$CP \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$VLLM_TP \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=$HF_MODEL_PATH \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    critic.megatron.tensor_model_parallel_size=$TP \\\n    critic.megatron.pipeline_model_parallel_size=$PP \\\n    critic.megatron.context_parallel_size=$CP \\\n    critic.megatron.use_dist_checkpointing=True \\\n    critic.megatron.dist_checkpointing_path=$DIST_CKPT_PATH \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_megatron_gsm8k_examples' \\\n    trainer.experiment_name='qwen1.5_moe_nochat' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=$NODES \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=100 $@\n    "
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_math_gsm8k_megatron.sh",
    "content": "set -x\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo --config-path=./config --config-name='ppo_megatron_trainer'\\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_ppo_gsm8k_math_examples' \\\n    trainer.experiment_name='qwen2_7b_megatron' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=100 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_rm.sh",
    "content": "# Discliamer: the model used in the script is only for academic purpose.\nset -x\n\n# Data preparation scripts are available in ``examples/data_preprocess``.\n# Example usage:\n#\n#   python3 examples/data_preprocess/math_dataset.py --local_dir ~/data/math\n#   python3 examples/data_preprocess/gsm8k.py --local_save_dir ~/data/gsm8k\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\n\n# prepare model ckpt\nhf download Qwen/Qwen2-7B-Instruct --local-dir $HOME/models/Qwen2-7B-Instruct &\nhf download sfairXC/FsfairX-LLaMA3-RM-v0.1 --local-dir $HOME/models/FsfairX-LLaMA3-RM-v0.1 &\nwait\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=\"$HOME/models/Qwen2-7B-Instruct\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.optim.lr_warmup_steps_ratio=0.05 \\\n    critic.model.path=\"$HOME/models/Qwen2-7B-Instruct\" \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=\"$HOME/models/FsfairX-LLaMA3-RM-v0.1\" \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=2048 \\\n    reward.reward_model.rollout.response_length=1024 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example' \\\n    trainer.val_before_train=False \\\n    trainer.experiment_name='Qwen2-7B-Instruct_hybrid_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_rm_reward_loop_colocate.sh",
    "content": "# download datasets and models\n# python3 examples/data_preprocess/gsm8k.py\n# python3 examples/data_preprocess/math_dataset.py\n# hf download Skywork/Skywork-Reward-V2-Llama-3.2-3B --local-dir $HOME/models/Skywork-Reward-V2-Llama-3.2-3B\n# hf download Qwen/Qwen2.5-3B-Instruct --local-dir $HOME/models/Qwen2.5-3B-Instruct\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=\"$HOME/models/Qwen2.5-3B-Instruct\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.optim.lr_warmup_steps_ratio=0.05 \\\n    critic.model.path=\"$HOME/models/Qwen2.5-3B-Instruct\" \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=\"$HOME/models/Skywork-Reward-V2-Llama-3.2-3B\" \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.prompt_length=4096 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_test_qwen25_rm' \\\n    trainer.val_before_train=False \\\n    trainer.experiment_name='reward_loop_colocate_reward_model' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_rm_seq_balance.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=4096 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=4096 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24000 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.use_dynamic_bsz=True \\\n    critic.ppo_max_token_len_per_gpu=98304 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=sfairXC/FsfairX-LLaMA3-RM-v0.1\\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=8192 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='qwen2-7b_hybrid_rm_bsz8k_p4k_r4k_seq_packing' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_rm_seq_balance_fused_kernels.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\nFUSED_KERNEL_BACKEND=triton # or 'torch' for torch backend\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=4096 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=4096 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.model.fused_kernel_options.impl_backend=$FUSED_KERNEL_BACKEND \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24000 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.use_dynamic_bsz=True \\\n    critic.ppo_max_token_len_per_gpu=98304 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=sfairXC/FsfairX-LLaMA3-RM-v0.1 \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=8192 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='qwen2-7b_hybrid_rm_bsz8k_p4k_r4k_seq_packing_fused_kernel' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_rm_seq_balance_nsys.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=${train_files:-\"$gsm8k_train_path\"}\ntest_files=${test_files:-\"$gsm8k_test_path\"}\n\nPROFILE_STEPS=\"[1,2,5]\" # or [] or null\nPROFILE_RANKS_ALL=False # or True\nPROFILE_RANKS=[0,4]\nDISCRETE=True  # or True\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=4096 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=4096 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=12000 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24000 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=2 \\\n    critic.use_dynamic_bsz=True \\\n    critic.ppo_max_token_len_per_gpu=98304 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    critic.profiler.enable=True \\\n    critic.profiler.ranks=$PROFILE_RANKS \\\n    critic.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=sfairXC/FsfairX-LLaMA3-RM-v0.1\\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=8192 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='qwen2-7b_hybrid_rm_bsz8k_p4k_r4k_seq_packing' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=15 \\\n    trainer.total_training_steps=6 \\\n    global_profiler.profile_continuous_steps=True \\\n    global_profiler.tool=nsys \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.global_tool_config.nsys.discrete=$DISCRETE $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_seq_balance.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\n# For async rollout mode, dataset should return raw chat.\nrollout_mode=\"async\"\nreturn_raw_chat=\"True\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.return_raw_chat=$return_raw_chat \\\n    data.train_batch_size=4096 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=4096 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=$rollout_mode \\\n    actor_rollout_ref.rollout.multi_turn.format=hermes \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24000 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_max_token_len_per_gpu=98304 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='qwen2-7b_function_rm_bsz8k_p4k_r4k_seq_packing' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2-7b_sglang_seq_balance.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=4096 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=4096 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=24000 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2-7B-Instruct \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_max_token_len_per_gpu=98304 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='qwen2-7b_function_rm_bsz8k_p4k_r4k_seq_packing' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.val_before_train=False \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2.5-32b.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-32B-Instruct \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen2.5-32B-Instruct \\\n    critic.model.enable_gradient_checkpointing=False \\\n    critic.ppo_micro_batch_size_per_gpu=8 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example' \\\n    trainer.experiment_name='Qwen2.5-32B-Instruct_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=4 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen2.5-3b_rm_reward_loop_colocate.sh",
    "content": "# download datasets and models\n# python3 examples/data_preprocess/gsm8k.py\n# python3 examples/data_preprocess/math_dataset.py\n# hf download Skywork/Skywork-Reward-V2-Llama-3.2-3B --local-dir $HOME/models/Skywork-Reward-V2-Llama-3.2-3B\n# hf download Qwen/Qwen2.5-3B-Instruct --local-dir $HOME/models/Qwen2.5-3B-Instruct\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=\"$HOME/models/Qwen2.5-3B-Instruct\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.optim.lr_warmup_steps_ratio=0.05 \\\n    critic.model.path=\"$HOME/models/Qwen2.5-3B-Instruct\" \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=32 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=\"$HOME/models/Skywork-Reward-V2-Llama-3.2-3B\" \\\n    reward.reward_model.rollout.name=vllm \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=4096 \\\n    reward.reward_model.rollout.response_length=4096 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_test_qwen25_rm' \\\n    trainer.val_before_train=True \\\n    trainer.experiment_name='reward_loop_colocate_reward_model' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ppo_trainer/run_qwen3-8b_npu.sh",
    "content": "set -x\n\nexport VLLM_USE_V1=1\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/dapo-math-17k.parquet \\\n    data.val_files=$HOME/data/dapo-math-17k.parquet \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=2000 \\\n    data.max_response_length=12000 \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-8B \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=14000 \\\n    actor_rollout_ref.rollout.max_num_seqs=64 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=Qwen/Qwen3-8B \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=1 \\\n    critic.ulysses_sequence_parallel_size=2 \\\n    critic.model.fsdp_config.param_offload=True \\\n    critic.model.fsdp_config.optimizer_offload=True \\\n    critic.use_dynamic_bsz=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_example_dapo_math_17k' \\\n    trainer.experiment_name='qwen3_8b_fsdp' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=-1 \\\n    trainer.val_before_train=False \\\n    trainer.max_actor_ckpt_to_keep=1 \\\n    trainer.max_critic_ckpt_to_keep=1 \\\n    trainer.total_training_steps=100 $@"
  },
  {
    "path": "examples/prefix_grouper/README.md",
    "content": "# PrefixGrouper Examples\n\nThis directory contains examples for using **PrefixGrouper**, an optimization technique that groups samples by shared prompts to reduce redundant computations in GRPO.\n\n## Introduction\n\n> Official Repository: [https://github.com/johncaged/PrefixGrouper](https://github.com/johncaged/PrefixGrouper)\n\n``PrefixGrouper`` is a plug-and-play efficient GRPO training tool that requires minimal modifications to existing codebases to achieve reduced computation, lower device memory consumption, and accelerated training.\n\nIn current mainstream GRPO training pipelines, policy model training primarily involves copying prefixes (typically questions, multimodal inputs, etc.) `G` times. Consequently, when training data prefixes are sufficiently long (e.g., long-context reasoning, image/long-video inference), redundant computation during training becomes non-negligible.\n\n**PrefixGrouper** decomposes the original redundant self-attention operation into prefix self-attention + suffix concat-attention.\n\n<h3 align=\"center\">\n    <img src=\"https://raw.githubusercontent.com/johncaged/PrefixGrouper/main/assets/images/method.jpg\">\n</h3>\n\n## Installation\n\n```bash\npip install prefix_grouper\n```\n\n## Limitations\n\n- Currently only supports FSDP worker (Megatron worker is not supported yet).\n- Incompatible with `use_dynamic_bsz=True`.\n- Incompatible with `use_remove_padding=True` (Flash Attention V2 variable length).\n- Incompatible with `use_fused_kernels=True`.\n- Incompatible with Ulysses sequence parallelism (`use_ulysses_sp=True`) and ring-attention.\n\nNote: `balance_batch=True` is now supported with group-level balancing, which keeps samples with the same uid together on the same rank. However, this requires `batch_size % (world_size * rollout.n) == 0`. For example, with `world_size=8` and `rollout.n=4`, you need `batch_size` to be a multiple of 32.\n\n## How to Use\n\n### 1. Enable PrefixGrouper in Config\n\nSimply set `use_prefix_grouper=True` in your training config:\n\n```yaml\nactor_rollout_ref:\n  actor:\n    use_prefix_grouper: True\n  model:\n    use_remove_padding: False \n```\n\nOptionally enable balance_batch for better load distribution:\n```yaml\ntrainer:\n  balance_batch: True  # Now supported with group-level balancing\n```\n\n### 2. Run Training\n\nUse the provided script `run_qwen3_prefix_grouper.sh` as an example:\n\n```bash\nbash examples/prefix_grouper/run_qwen3_prefix_grouper.sh\n```\n\n## How It Works\n\nWhen `use_prefix_grouper=True`, verl automatically patches the attention functions in `transformers.modeling_utils.ALL_ATTENTION_FUNCTIONS` to support the `prefix_grouper` parameter. No model code modifications are needed.\n\nThe patch wraps each attention function to:\n1. Extract `prefix_grouper` from kwargs\n2. If `prefix_grouper` is None, call original attention\n3. If `prefix_grouper` is provided, use PrefixGrouper's optimized attention computation\n\n## Performance\n\n**Benchmark Results** (Qwen3-4B, 4×H800, `rollout.n=4`):\n\n| Context Length | Metric | PG | No PG | Speedup |\n|----------------|--------|-----|-------|---------|\n| **4K** | `old_log_prob` | 1.31s | 1.70s | **1.30x** |\n| | `update_actor` | 4.80s | 6.07s | **1.26x** |\n| | `step` | 17.08s | 19.40s | **1.14x** |\n| **8K** | `old_log_prob` | 1.69s | 2.63s | **1.56x** |\n| | `update_actor` | 5.98s | 10.18s | **1.70x** |\n| | `step` | 19.48s | 24.71s | **1.27x** |\n\nAs context length increases, the speedup becomes more pronounced.\n"
  },
  {
    "path": "examples/prefix_grouper/run_qwen3_prefix_grouper.sh",
    "content": "set -x\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-8B \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.use_prefix_grouper=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen3_function_rm_pg' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.balance_batch=True \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/ray/tutorial.ipynb",
    "content": "{\n \"cells\": [\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"0ddc582b\",\n   \"metadata\": {},\n   \"source\": [\n    \"# VeRL Ray API Tutorial\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"71fe3b94\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Chapter 1: Ray Basics\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 144,\n   \"id\": \"1347d381\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"import os\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 145,\n   \"id\": \"e75b9d44\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"import warnings\\n\",\n    \"\\n\",\n    \"import ray\\n\",\n    \"import torch\\n\",\n    \"\\n\",\n    \"warnings.filterwarnings(\\\"ignore\\\")\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 146,\n   \"id\": \"2e90ae00\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stderr\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"2024-11-01 17:27:19,132\\tINFO worker.py:1752 -- Started a local Ray instance.\\n\"\n     ]\n    },\n    {\n     \"data\": {\n      \"application/vnd.jupyter.widget-view+json\": {\n       \"model_id\": \"9cc9d2ccbdfb48918c8fd6cd13a0807a\",\n       \"version_major\": 2,\n       \"version_minor\": 0\n      },\n      \"text/html\": [\n       \"<div class=\\\"lm-Widget p-Widget lm-Panel p-Panel jp-Cell-outputWrapper\\\">\\n\",\n       \"    <div style=\\\"margin-left: 50px;display: flex;flex-direction: row;align-items: center\\\">\\n\",\n       \"        <div class=\\\"jp-RenderedHTMLCommon\\\" style=\\\"display: flex; flex-direction: row;\\\">\\n\",\n       \"  <svg viewBox=\\\"0 0 567 224\\\" fill=\\\"none\\\" xmlns=\\\"http://www.w3.org/2000/svg\\\" style=\\\"height: 3em;\\\">\\n\",\n       \"    <g clip-path=\\\"url(#clip0_4338_178347)\\\">\\n\",\n       \"        <path d=\\\"M341.29 165.561H355.29L330.13 129.051C345.63 123.991 354.21 112.051 354.21 94.2307C354.21 71.3707 338.72 58.1807 311.88 58.1807H271V165.561H283.27V131.661H311.8C314.25 131.661 316.71 131.501 319.01 131.351L341.25 165.561H341.29ZM283.29 119.851V70.0007H311.82C331.3 70.0007 342.34 78.2907 342.34 94.5507C342.34 111.271 331.34 119.861 311.82 119.861L283.29 119.851ZM451.4 138.411L463.4 165.561H476.74L428.74 58.1807H416L367.83 165.561H380.83L392.83 138.411H451.4ZM446.19 126.601H398L422 72.1407L446.24 126.601H446.19ZM526.11 128.741L566.91 58.1807H554.35L519.99 114.181L485.17 58.1807H472.44L514.01 129.181V165.541H526.13V128.741H526.11Z\\\" fill=\\\"var(--jp-ui-font-color0)\\\"/>\\n\",\n       \"        <path d=\\\"M82.35 104.44C84.0187 97.8827 87.8248 92.0678 93.1671 87.9146C98.5094 83.7614 105.083 81.5067 111.85 81.5067C118.617 81.5067 125.191 83.7614 130.533 87.9146C135.875 92.0678 139.681 97.8827 141.35 104.44H163.75C164.476 101.562 165.622 98.8057 167.15 96.2605L127.45 56.5605C121.071 60.3522 113.526 61.6823 106.235 60.3005C98.9443 58.9187 92.4094 54.9203 87.8602 49.0574C83.3109 43.1946 81.0609 35.8714 81.5332 28.4656C82.0056 21.0599 85.1679 14.0819 90.4252 8.8446C95.6824 3.60726 102.672 0.471508 110.08 0.0272655C117.487 -0.416977 124.802 1.86091 130.647 6.4324C136.493 11.0039 140.467 17.5539 141.821 24.8501C143.175 32.1463 141.816 39.6859 138 46.0505L177.69 85.7505C182.31 82.9877 187.58 81.4995 192.962 81.4375C198.345 81.3755 203.648 82.742 208.33 85.3976C213.012 88.0532 216.907 91.9029 219.616 96.5544C222.326 101.206 223.753 106.492 223.753 111.875C223.753 117.258 222.326 122.545 219.616 127.197C216.907 131.848 213.012 135.698 208.33 138.353C203.648 141.009 198.345 142.375 192.962 142.313C187.58 142.251 182.31 140.763 177.69 138L138 177.7C141.808 184.071 143.155 191.614 141.79 198.91C140.424 206.205 136.44 212.75 130.585 217.313C124.731 221.875 117.412 224.141 110.004 223.683C102.596 223.226 95.6103 220.077 90.3621 214.828C85.1139 209.58 81.9647 202.595 81.5072 195.187C81.0497 187.779 83.3154 180.459 87.878 174.605C92.4405 168.751 98.9853 164.766 106.281 163.401C113.576 162.035 121.119 163.383 127.49 167.19L167.19 127.49C165.664 124.941 164.518 122.182 163.79 119.3H141.39C139.721 125.858 135.915 131.673 130.573 135.826C125.231 139.98 118.657 142.234 111.89 142.234C105.123 142.234 98.5494 139.98 93.2071 135.826C87.8648 131.673 84.0587 125.858 82.39 119.3H60C58.1878 126.495 53.8086 132.78 47.6863 136.971C41.5641 141.163 34.1211 142.972 26.7579 142.059C19.3947 141.146 12.6191 137.574 7.70605 132.014C2.79302 126.454 0.0813599 119.29 0.0813599 111.87C0.0813599 104.451 2.79302 97.2871 7.70605 91.7272C12.6191 86.1673 19.3947 82.5947 26.7579 81.6817C34.1211 80.7686 41.5641 82.5781 47.6863 86.7696C53.8086 90.9611 58.1878 97.2456 60 104.44H82.35ZM100.86 204.32C103.407 206.868 106.759 208.453 110.345 208.806C113.93 209.159 117.527 208.258 120.522 206.256C123.517 204.254 125.725 201.276 126.771 197.828C127.816 194.38 127.633 190.677 126.253 187.349C124.874 184.021 122.383 181.274 119.205 179.577C116.027 177.88 112.359 177.337 108.826 178.042C105.293 178.746 102.113 180.654 99.8291 183.44C97.5451 186.226 96.2979 189.718 96.3 193.32C96.2985 195.364 96.7006 197.388 97.4831 199.275C98.2656 201.163 99.4132 202.877 100.86 204.32ZM204.32 122.88C206.868 120.333 208.453 116.981 208.806 113.396C209.159 109.811 208.258 106.214 206.256 103.219C204.254 100.223 201.275 98.0151 197.827 96.97C194.38 95.9249 190.676 96.1077 187.348 97.4873C184.02 98.8669 181.274 101.358 179.577 104.536C177.879 107.714 177.337 111.382 178.041 114.915C178.746 118.448 180.653 121.627 183.439 123.911C186.226 126.195 189.717 127.443 193.32 127.44C195.364 127.443 197.388 127.042 199.275 126.259C201.163 125.476 202.878 124.328 204.32 122.88ZM122.88 19.4205C120.333 16.8729 116.981 15.2876 113.395 14.9347C109.81 14.5817 106.213 15.483 103.218 17.4849C100.223 19.4868 98.0146 22.4654 96.9696 25.9131C95.9245 29.3608 96.1073 33.0642 97.4869 36.3922C98.8665 39.7202 101.358 42.4668 104.535 44.1639C107.713 45.861 111.381 46.4036 114.914 45.6992C118.447 44.9949 121.627 43.0871 123.911 40.301C126.195 37.515 127.442 34.0231 127.44 30.4205C127.44 28.3772 127.038 26.3539 126.255 24.4664C125.473 22.5788 124.326 20.8642 122.88 19.4205ZM19.42 100.86C16.8725 103.408 15.2872 106.76 14.9342 110.345C14.5813 113.93 15.4826 117.527 17.4844 120.522C19.4863 123.518 22.4649 125.726 25.9127 126.771C29.3604 127.816 33.0638 127.633 36.3918 126.254C39.7198 124.874 42.4664 122.383 44.1635 119.205C45.8606 116.027 46.4032 112.359 45.6988 108.826C44.9944 105.293 43.0866 102.114 40.3006 99.8296C37.5145 97.5455 34.0227 96.2983 30.42 96.3005C26.2938 96.3018 22.337 97.9421 19.42 100.86ZM100.86 100.86C98.3125 103.408 96.7272 106.76 96.3742 110.345C96.0213 113.93 96.9226 117.527 98.9244 120.522C100.926 123.518 103.905 125.726 107.353 126.771C110.8 127.816 114.504 127.633 117.832 126.254C121.16 124.874 123.906 122.383 125.604 119.205C127.301 116.027 127.843 112.359 127.139 108.826C126.434 105.293 124.527 102.114 121.741 99.8296C118.955 97.5455 115.463 96.2983 111.86 96.3005C109.817 96.299 107.793 96.701 105.905 97.4835C104.018 98.2661 102.303 99.4136 100.86 100.86Z\\\" fill=\\\"#00AEEF\\\"/>\\n\",\n       \"    </g>\\n\",\n       \"    <defs>\\n\",\n       \"        <clipPath id=\\\"clip0_4338_178347\\\">\\n\",\n       \"            <rect width=\\\"566.93\\\" height=\\\"223.75\\\" fill=\\\"white\\\"/>\\n\",\n       \"        </clipPath>\\n\",\n       \"    </defs>\\n\",\n       \"  </svg>\\n\",\n       \"</div>\\n\",\n       \"\\n\",\n       \"        <table class=\\\"jp-RenderedHTMLCommon\\\" style=\\\"border-collapse: collapse;color: var(--jp-ui-font-color1);font-size: var(--jp-ui-font-size1);\\\">\\n\",\n       \"    <tr>\\n\",\n       \"        <td style=\\\"text-align: left\\\"><b>Python version:</b></td>\\n\",\n       \"        <td style=\\\"text-align: left\\\"><b>3.9.2</b></td>\\n\",\n       \"    </tr>\\n\",\n       \"    <tr>\\n\",\n       \"        <td style=\\\"text-align: left\\\"><b>Ray version:</b></td>\\n\",\n       \"        <td style=\\\"text-align: left\\\"><b>2.10.0</b></td>\\n\",\n       \"    </tr>\\n\",\n       \"    \\n\",\n       \"</table>\\n\",\n       \"\\n\",\n       \"    </div>\\n\",\n       \"</div>\\n\"\n      ],\n      \"text/plain\": [\n       \"RayContext(dashboard_url='', python_version='3.9.2', ray_version='2.10.0', ray_commit='09abba26b5bf2707639bb637c208d062a47b46f6')\"\n      ]\n     },\n     \"execution_count\": 146,\n     \"metadata\": {},\n     \"output_type\": \"execute_result\"\n    },\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"\\u001b[36m(GPUAccumulator pid=224400)\\u001b[0m rank 0, value: tensor([1.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulator pid=225234)\\u001b[0m rank 2, value: tensor([3.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulator pid=225607)\\u001b[0m rank 0, value: tensor([2.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulator pid=226423)\\u001b[0m rank 1, value: tensor([3.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulator pid=226857)\\u001b[0m rank 3, value: tensor([6.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulatorDecorator pid=227475)\\u001b[0m 10\\n\",\n      \"\\u001b[36m(GPUAccumulatorDecorator pid=227475)\\u001b[0m rank 0, value: tensor([10.], device='cuda:0')\\n\",\n      \"\\u001b[36m(GPUAccumulatorDecorator pid=227655)\\u001b[0m rank 1, value: tensor([11.], device='cuda:0')\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# Build a local ray cluster. The head node and worker node are on this machine\\n\",\n    \"ray.init()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"a127e4e4\",\n   \"metadata\": {},\n   \"source\": [\n    \"Implement an Accumulator class.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 147,\n   \"id\": \"20e7b9a3\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"@ray.remote\\n\",\n    \"class Accumulator:\\n\",\n    \"    def __init__(self):\\n\",\n    \"        self.value = 0\\n\",\n    \"\\n\",\n    \"    def add(self, x):\\n\",\n    \"        self.value += x\\n\",\n    \"\\n\",\n    \"    def get_value(self):\\n\",\n    \"        return self.value\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 148,\n   \"id\": \"3b80098c\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Instantiate an accumulator. Accumulator can be viewed as a process, acting as an RPC service.\\n\",\n    \"accumulator = Accumulator.remote()\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 149,\n   \"id\": \"b14b1009\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"0\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"value_ref = accumulator.get_value.remote()  # Check the current value. Note that this function returns immediately and does not actually wait for the remote execution to complete.\\n\",\n    \"# Get the value\\n\",\n    \"value = ray.get(value_ref)\\n\",\n    \"print(value)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 150,\n   \"id\": \"513a84b3\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"10\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# Accumulate, then check the result.\\n\",\n    \"accumulator.add.remote(10)  # Similarly, the 'add' here will return immediately.\\n\",\n    \"new_value = ray.get(accumulator.get_value.remote())\\n\",\n    \"print(new_value)\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"3c332fe0\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Chapter 2: Resource Pool and RayWorkerGroup\\n\",\n    \"In the previous example, it was a simple single-process worker. \\n\",\n    \"In this example, we implement a worker with a GPU and form a RayWorkerGroup. Within this RayWorkerGroup, we implement a simple operation of an accumulator.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 151,\n   \"id\": \"04229afb\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"from verl.single_controller.base import Worker\\n\",\n    \"from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, merge_resource_pool\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 152,\n   \"id\": \"0d0dbd58\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"resource_pool = RayResourcePool([4], use_gpu=True)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 153,\n   \"id\": \"68f6838a\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"@ray.remote\\n\",\n    \"class GPUAccumulator(Worker):\\n\",\n    \"    def __init__(self) -> None:\\n\",\n    \"        super().__init__()\\n\",\n    \"        # The initial value of each rank is the same as the rank\\n\",\n    \"        self.value = torch.zeros(size=(1,), device=\\\"cuda\\\") + self.rank\\n\",\n    \"\\n\",\n    \"    def add(self, x):\\n\",\n    \"        self.value += x\\n\",\n    \"        print(f\\\"rank {self.rank}, value: {self.value}\\\")\\n\",\n    \"        return self.value.cpu()\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 154,\n   \"id\": \"23aad8fe\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"[tensor([1.]), tensor([2.]), tensor([3.]), tensor([4.])]\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# Each worker's initial value is its rank, and then each rank's value is incremented by 1, so the values obtained on each rank are [1, 2, 3, 4]\\n\",\n    \"class_with_args = RayClassWithInitArgs(cls=GPUAccumulator)\\n\",\n    \"worker_group = RayWorkerGroup(resource_pool, class_with_args)\\n\",\n    \"print(worker_group.execute_all_sync(\\\"add\\\", x=[1, 1, 1, 1]))\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"e6705284\",\n   \"metadata\": {},\n   \"source\": [\n    \"The principle of parameter passing: The input parameter is a list of length world_size, where each element in the list is dispatched respectively to each worker in the RayWorkerGroup. \\n\",\n    \"The return parameter is also a list, corresponding to the return value of each worker.\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"d25c2412\",\n   \"metadata\": {},\n   \"source\": [\n    \"### GPU Resource Sharing\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"f74f6d24\",\n   \"metadata\": {},\n   \"source\": [\n    \"RayWorkerGroups mapped to the same resource pool share the GPU. In this example, we implement three resource pools: the first occupies 4 GPUs, the second also occupies 4 GPUs, and the last occupies all 8 GPUs. Among them, the first resource pool reuses the resource pool mentioned above.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 155,\n   \"id\": \"49f9c06f\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Create a new resource pool and then merge the newly created resource pool with the previous one.\\n\",\n    \"resource_pool_1 = RayResourcePool([4], use_gpu=True, name_prefix=\\\"a\\\")\\n\",\n    \"resource_pool_merge = merge_resource_pool(resource_pool, resource_pool_1)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 156,\n   \"id\": \"05c2e305\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Establish a RayWorkerGroup on the newly created resource pool.\\n\",\n    \"worker_group_1 = RayWorkerGroup(resource_pool_1, class_with_args)\\n\",\n    \"worker_group_merge = RayWorkerGroup(resource_pool_merge, class_with_args)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 157,\n   \"id\": \"6b9b13f4\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"[tensor([2.]), tensor([3.]), tensor([4.]), tensor([5.])]\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# Run 'add' on the second set of 4 GPUs; the result should be [2, 3, 4, 5].\\n\",\n    \"output_1 = worker_group_1.execute_all_sync(\\\"add\\\", x=[2, 2, 2, 2])\\n\",\n    \"print(output_1)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 158,\n   \"id\": \"d856d030\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"[tensor([3.]), tensor([4.]), tensor([5.]), tensor([6.]), tensor([7.]), tensor([8.]), tensor([9.]), tensor([10.])]\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# Run 'add' on the merged set of 8 GPUs; the result should be [3, 4, 5, 6, 7, 8, 9, 10].\\n\",\n    \"output_merge = worker_group_merge.execute_all_sync(\\\"add\\\", x=[3, 3, 3, 3, 3, 3, 3, 3])\\n\",\n    \"print(output_merge)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 159,\n   \"id\": \"33a4628c\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"4 4 8\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"print(worker_group.world_size, worker_group_1.world_size, worker_group_merge.world_size)\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"3df19d13\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Chapter 3: Data Dispatch, Execution and Collection\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"acb22d9d\",\n   \"metadata\": {},\n   \"source\": [\n    \"In the above example, we used the `execute_all_sync` function in the RayWorkerGroup to dispatch data from the driver to each worker. This is very inconvenient for coding. \\n\",\n    \"In this chapter, we use the form of function decorators to allow RayWorkerGroup to directly call functions written in the Worker, and to greatly simplify parameter passing.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 160,\n   \"id\": \"35237432\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"from verl.single_controller.base.decorator import Dispatch, Execute, register\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 161,\n   \"id\": \"88b8ba3b\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"@ray.remote\\n\",\n    \"class GPUAccumulatorDecorator(Worker):\\n\",\n    \"    def __init__(self) -> None:\\n\",\n    \"        super().__init__()\\n\",\n    \"        # The initial value of each rank is the same as the rank\\n\",\n    \"        self.value = torch.zeros(size=(1,), device=\\\"cuda\\\") + self.rank\\n\",\n    \"\\n\",\n    \"    # map from a single input to all the worker\\n\",\n    \"    @register(Dispatch.ONE_TO_ALL)\\n\",\n    \"    def add(self, x):\\n\",\n    \"        print(x)\\n\",\n    \"        self.value = self.value + x\\n\",\n    \"        print(f\\\"rank {self.rank}, value: {self.value}\\\")\\n\",\n    \"        return self.value.cpu()\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 162,\n   \"id\": \"eddaa043\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"class_with_args = RayClassWithInitArgs(cls=GPUAccumulatorDecorator)\\n\",\n    \"gpu_accumulator_decorator = RayWorkerGroup(resource_pool_merge, class_with_args)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 163,\n   \"id\": \"10087c91\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"[tensor([10.]), tensor([11.]), tensor([12.]), tensor([13.]), tensor([14.]), tensor([15.]), tensor([16.]), tensor([17.])]\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"# As we can see, 10 is automatically dispatched to each Worker in this RayWorkerGroup.\\n\",\n    \"print(gpu_accumulator_decorator.add(x=10))\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"540ee6ad\",\n   \"metadata\": {},\n   \"source\": [\n    \"### Custom Dispatch, Collection\\n\",\n    \"Users can customize `dispatch` and `collection` function. You only need to write the `dispatch_fn` and `collect_fn` functions yourself. We also support executing RPC only on rank_zero, with specific examples provided below.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 164,\n   \"id\": \"8e041270\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"from verl.single_controller.base.decorator import Dispatch, collect_all_to_all, register\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 165,\n   \"id\": \"43b5be31\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"def two_to_all_dispatch_fn(worker_group, *args, **kwargs):\\n\",\n    \"    \\\"\\\"\\\"\\n\",\n    \"    Assume the input is a list of 2. Duplicate the input interleaved and pass to each worker.\\n\",\n    \"    \\\"\\\"\\\"\\n\",\n    \"    for arg in args:\\n\",\n    \"        assert len(arg) == 2\\n\",\n    \"        for i in range(worker_group.world_size - 2):\\n\",\n    \"            arg.append(arg[i % 2])\\n\",\n    \"    for k, v in kwargs.items():\\n\",\n    \"        assert len(v) == 2\\n\",\n    \"        for i in range(worker_group.world_size - 2):\\n\",\n    \"            v.append(v[i % 2])\\n\",\n    \"    return args, kwargs\\n\",\n    \"\\n\",\n    \"\\n\",\n    \"@ray.remote\\n\",\n    \"class TestActor(Worker):\\n\",\n    \"    # TODO: pass *args and **kwargs is bug prone and not very convincing\\n\",\n    \"    def __init__(self, x) -> None:\\n\",\n    \"        super().__init__()\\n\",\n    \"        self._x = x\\n\",\n    \"\\n\",\n    \"    def foo(self, y):\\n\",\n    \"        return self._x + y\\n\",\n    \"\\n\",\n    \"    @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)\\n\",\n    \"    def foo_rank_zero(self, x, y):\\n\",\n    \"        return self._x + y + x\\n\",\n    \"\\n\",\n    \"    @register(dispatch_mode={\\\"dispatch_fn\\\": two_to_all_dispatch_fn, \\\"collect_fn\\\": collect_all_to_all})\\n\",\n    \"    def foo_custom(self, x, y):\\n\",\n    \"        return self._x + y + x\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 166,\n   \"id\": \"83ec6609\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"class_with_args = RayClassWithInitArgs(cls=TestActor, x=2)\\n\",\n    \"worker_group = RayWorkerGroup(resource_pool, class_with_args)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 167,\n   \"id\": \"62c58d8a\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"output_ref = worker_group.foo_custom(x=[1, 2], y=[5, 6])\\n\",\n    \"assert output_ref == [8, 10, 8, 10]\\n\",\n    \"\\n\",\n    \"output_ref = worker_group.foo_rank_zero(x=1, y=2)\\n\",\n    \"assert output_ref == 5\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 168,\n   \"id\": \"14689353\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"8\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"print(gpu_accumulator_decorator.world_size)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 169,\n   \"id\": \"2c80bbf4\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Shutdown ray cluster\\n\",\n    \"ray.shutdown()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"a5c8151c\",\n   \"metadata\": {},\n   \"source\": [\n    \"## Chapter 4: NVMegatronRayWorkerGroup\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"cd5680e9\",\n   \"metadata\": {},\n   \"source\": [\n    \"Due to the Ray issue, we can only support max_colocate_count=1 in RayResourcePool for now. \\n\",\n    \"This means that each GPU can only have one process.\\n\",\n    \"We can support max_colocate > 1 when applying this pull request: https://github.com/ray-project/ray/pull/44385\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"92724419\",\n   \"metadata\": {},\n   \"source\": [\n    \"Therefore, we need to restart the ray and initialize a new resource_pool to demonstrate the **NVMegatronRayWorkerGroup**\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": null,\n   \"id\": \"9b038538\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Build a local ray cluster. The head node and worker node are on this machine\\n\",\n    \"ray.init()\"\n   ]\n  },\n  {\n   \"cell_type\": \"markdown\",\n   \"id\": \"ebfd8798\",\n   \"metadata\": {},\n   \"source\": [\n    \"Finally, we implement a `NVMegatronRayWorkerGroup`, within which we create a Megatron and then run a tensor parallel (tp) split Llama mlp layer. Here, we use a complex dispatch mode, `Megatron_COMPUTE`. This dispatch mode assumes that user passes the data partitioned by DP dimension. The data is dispatched to all tp/pp ranks within the same dp group, and ultimately only collects output data from tp=0 and the last pp. In this way, for users that only write code on the driver, the Megatron behind the RPC becomes transparent.\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 171,\n   \"id\": \"5a032154\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"/opt/tiger/Megatron-LM\\n\",\n      \"/opt/tiger/Megatron-LM/megatron/__init__.py\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"import sys\\n\",\n    \"\\n\",\n    \"current_pythonpath = os.environ.get(\\\"PYTHONPATH\\\", \\\"\\\")\\n\",\n    \"\\n\",\n    \"new_path = \\\"/opt/tiger/Megatron-LM\\\"\\n\",\n    \"\\n\",\n    \"new_pythonpath = f\\\"{new_path}:{current_pythonpath}\\\" if current_pythonpath else new_path\\n\",\n    \"\\n\",\n    \"os.environ[\\\"PYTHONPATH\\\"] = new_pythonpath\\n\",\n    \"\\n\",\n    \"print(new_path)\\n\",\n    \"sys.path.append(new_path)\\n\",\n    \"\\n\",\n    \"import megatron\\n\",\n    \"\\n\",\n    \"print(megatron.__file__)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 172,\n   \"id\": \"8c84cd5a\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"from megatron.core import parallel_state as mpu\\n\",\n    \"from omegaconf import OmegaConf\\n\",\n    \"\\n\",\n    \"from verl.single_controller.base.decorator import Dispatch, Execute, register\\n\",\n    \"from verl.single_controller.base.megatron.worker import MegatronWorker\\n\",\n    \"from verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\\n\",\n    \"from verl.single_controller.ray.megatron import NVMegatronRayWorkerGroup\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 173,\n   \"id\": \"1b1debcc\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"resource_pool = RayResourcePool([4], use_gpu=True, max_colocate_count=1)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 174,\n   \"id\": \"bccbe081\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"@ray.remote\\n\",\n    \"class MLPLayerWorker(MegatronWorker):\\n\",\n    \"    def __init__(self):\\n\",\n    \"        super().__init__()\\n\",\n    \"        rank = int(os.environ[\\\"LOCAL_RANK\\\"])\\n\",\n    \"        torch.distributed.init_process_group(backend=\\\"nccl\\\")\\n\",\n    \"        torch.cuda.set_device(rank)\\n\",\n    \"\\n\",\n    \"        mpu.initialize_model_parallel(\\n\",\n    \"            tensor_model_parallel_size=4,\\n\",\n    \"            pipeline_model_parallel_size=1,\\n\",\n    \"            virtual_pipeline_model_parallel_size=None,\\n\",\n    \"            pipeline_model_parallel_split_rank=None,\\n\",\n    \"            use_sharp=False,\\n\",\n    \"            context_parallel_size=1,\\n\",\n    \"            expert_model_parallel_size=1,\\n\",\n    \"            nccl_communicator_config_path=None,\\n\",\n    \"        )\\n\",\n    \"        from megatron.core import tensor_parallel\\n\",\n    \"\\n\",\n    \"        tensor_parallel.model_parallel_cuda_manual_seed(10)\\n\",\n    \"\\n\",\n    \"    @register(Dispatch.ONE_TO_ALL)\\n\",\n    \"    def init_model(self, config):\\n\",\n    \"        from omegaconf import OmegaConf\\n\",\n    \"\\n\",\n    \"        from verl.models.llama.megatron.layers import ParallelLlamaMLP\\n\",\n    \"        from verl.utils.megatron_utils import init_model_parallel_config\\n\",\n    \"\\n\",\n    \"        megatron_config = OmegaConf.create(\\n\",\n    \"            {\\n\",\n    \"                \\\"sequence_parallel\\\": False,\\n\",\n    \"                \\\"param_dtype\\\": \\\"fp32\\\",\\n\",\n    \"                \\\"tensor_model_parallel_size\\\": mpu.get_tensor_model_parallel_world_size(),\\n\",\n    \"                \\\"pipeline_model_parallel_rank\\\": mpu.get_pipeline_model_parallel_rank(),\\n\",\n    \"                \\\"pipeline_model_parallel_size\\\": mpu.get_pipeline_model_parallel_world_size(),\\n\",\n    \"                \\\"virtual_pipeline_model_parallel_rank\\\": mpu.get_virtual_pipeline_model_parallel_rank(),\\n\",\n    \"                \\\"virtual_pipeline_model_parallel_size\\\": mpu.get_virtual_pipeline_model_parallel_world_size(),\\n\",\n    \"            }\\n\",\n    \"        )\\n\",\n    \"\\n\",\n    \"        megatron_config = init_model_parallel_config(megatron_config)\\n\",\n    \"        self.parallel_layer = ParallelLlamaMLP(config=config, megatron_config=megatron_config)\\n\",\n    \"\\n\",\n    \"    @register(Dispatch.ONE_TO_ALL)\\n\",\n    \"    def get_weights(self):\\n\",\n    \"        output = {}\\n\",\n    \"        for key, val in self.parallel_layer.named_parameters():\\n\",\n    \"            output[key] = val\\n\",\n    \"        return output\\n\",\n    \"\\n\",\n    \"    @register(Dispatch.MEGATRON_COMPUTE)\\n\",\n    \"    def run_layer(self, x):\\n\",\n    \"        x = x.to(\\\"cuda\\\")\\n\",\n    \"        y = self.parallel_layer(x)\\n\",\n    \"        return y\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 175,\n   \"id\": \"a655271d\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"layer_cls = RayClassWithInitArgs(cls=MLPLayerWorker)\\n\",\n    \"layer_worker_group = NVMegatronRayWorkerGroup(\\n\",\n    \"    resource_pool=resource_pool,\\n\",\n    \"    ray_cls_with_init=layer_cls,\\n\",\n    \")\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 176,\n   \"id\": \"f105ebee\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"4 4 1 1\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"print(layer_worker_group.world_size, layer_worker_group.tp_size, layer_worker_group.pp_size, layer_worker_group.dp_size)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 177,\n   \"id\": \"38655091\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"ffn_hidden_size = 11008\\n\",\n    \"batch_size = 16\\n\",\n    \"seq_len = 2048\\n\",\n    \"hidden_size = 4096\\n\",\n    \"\\n\",\n    \"config = OmegaConf.create(\\n\",\n    \"    {\\n\",\n    \"        \\\"hidden_size\\\": hidden_size,\\n\",\n    \"        \\\"intermediate_size\\\": ffn_hidden_size,\\n\",\n    \"        \\\"hidden_act\\\": \\\"silu\\\",\\n\",\n    \"        \\\"pretraining_tp\\\": 1,\\n\",\n    \"        \\\"tp\\\": layer_worker_group.tp_size,\\n\",\n    \"    }\\n\",\n    \")\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 178,\n   \"id\": \"a026efca\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"x = torch.rand(size=(seq_len, batch_size, hidden_size), dtype=torch.float32)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 179,\n   \"id\": \"f5fcaf13\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"data\": {\n      \"text/plain\": [\n       \"[None, None, None, None]\"\n      ]\n     },\n     \"execution_count\": 179,\n     \"metadata\": {},\n     \"output_type\": \"execute_result\"\n    }\n   ],\n   \"source\": [\n    \"layer_worker_group.init_model(config)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 180,\n   \"id\": \"3f5cc9b4\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [\n    {\n     \"name\": \"stdout\",\n     \"output_type\": \"stream\",\n     \"text\": [\n      \"torch.Size([2048, 16, 4096])\\n\"\n     ]\n    }\n   ],\n   \"source\": [\n    \"output = layer_worker_group.run_layer(\\n\",\n    \"    [x]\\n\",\n    \")  # This must be a list of size 1, ensuring that the input equals the data parallel (dp).\\n\",\n    \"print(output[0].shape)\"\n   ]\n  },\n  {\n   \"cell_type\": \"code\",\n   \"execution_count\": 181,\n   \"id\": \"49792210\",\n   \"metadata\": {\n    \"tags\": []\n   },\n   \"outputs\": [],\n   \"source\": [\n    \"# Shutdown ray cluster\\n\",\n    \"ray.shutdown()\"\n   ]\n  }\n ],\n \"metadata\": {\n  \"kernelspec\": {\n   \"display_name\": \"Python 3 (ipykernel)\",\n   \"language\": \"python\",\n   \"name\": \"python3\"\n  },\n  \"language_info\": {\n   \"codemirror_mode\": {\n    \"name\": \"ipython\",\n    \"version\": 3\n   },\n   \"file_extension\": \".py\",\n   \"mimetype\": \"text/x-python\",\n   \"name\": \"python\",\n   \"nbconvert_exporter\": \"python\",\n   \"pygments_lexer\": \"ipython3\",\n   \"version\": \"3.9.2\"\n  }\n },\n \"nbformat\": 4,\n \"nbformat_minor\": 5\n}\n"
  },
  {
    "path": "examples/reinforce_plus_plus_trainer/run_qwen2-7b_math_rf.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=reinforce_plus_plus \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=1024 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=mse \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/reinforce_plus_plus_trainer/run_qwen2-7b_math_rf_baseline.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/gsm8k/test.parquet\nmath_train_path=$HOME/data/math/train.parquet\nmath_test_path=$HOME/data/math/test.parquet\n\ntrain_files=\"['$gsm8k_train_path', '$math_train_path']\"\ntest_files=\"['$gsm8k_test_path', '$math_test_path']\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=reinforce_plus_plus_baseline \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=1024 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=mse \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/remax_trainer/run_qwen2.5-3b_seq_balance.sh",
    "content": "set -x\n\nexport HF_DATASETS_OFFLINE=1\nexport TRANSFORMERS_OFFLINE=1\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=remax \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=30000 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=True \\\n    algorithm.kl_penalty=kl \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_remax_example_gsm8k' \\\n    trainer.experiment_name='qwen2.5_3b_function_rm_kl1e-3' \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=5 $@\n"
  },
  {
    "path": "examples/remax_trainer/run_qwen2.5-7b_seq_balance.sh",
    "content": "set -x\n\nexport HF_DATASETS_OFFLINE=1\nexport TRANSFORMERS_OFFLINE=1\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=remax \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=4 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=True \\\n    algorithm.kl_penalty=kl \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_remax_example_gsm8k' \\\n    trainer.experiment_name='qwen2.5_7b_function_rm_kl1e-3' \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=10 $@\n"
  },
  {
    "path": "examples/rloo_trainer/run_qwen2-7b.sh",
    "content": "set -x\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=rloo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=Qwen/Qwen2-7B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=80 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=160 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=160 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=True \\\n    algorithm.kl_penalty=kl \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_rloo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/rollout_correction/run_with_rollout_corr.sh",
    "content": "#!/usr/bin/env bash\n# Example: RLOO (REINFORCE Leave-One-Out) with Rollout Correction\n# This demonstrates self-normalized sequence-level IS with pure policy gradient\n#\n# References:\n#   - Rollout Correction Docs: https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr.md\n#   - Rollout Correction Math: https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr_math.md\n\nset -xeuo pipefail\n\n# ==============================================================================\n# Rollout Correction Configuration (RLOO)\n# ==============================================================================\n\n# Importance Sampling (IS) weights configuration\nrollout_is=\"sequence\"                     # Self-normalized sequence-level IS\nrollout_is_threshold=2.0                  # Upper threshold for IS weights\nrollout_is_batch_normalize=\"true\"        # Self-normalization (mean=1.0)\n\n# Rejection Sampling (RS) configuration\nrollout_rs=\"null\"                         # No rejection sampling for basic RLOO\nrollout_rs_threshold=\"null\"               # RS threshold spec (string or float)\n\n# Bypass mode with REINFORCE loss (no PPO clipping)\nbypass_mode=\"true\"     # Skip old_log_prob computation\nloss_type=\"reinforce\"  # REINFORCE with explicit IS weights (alternative: \"ppo_clip\")\n\n# ==============================================================================\n# Model and Data Configuration\n# ==============================================================================\n\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen2.5-7B\"}\nTRAIN_FILE=${TRAIN_FILE:-\"data/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"data/test.parquet\"}\n\nmax_prompt_length=2048\nmax_response_length=4096\n\n# ==============================================================================\n# Training Configuration\n# ==============================================================================\n\ntrain_batch_size=128\nppo_mini_batch_size=32\nppo_epochs=1\nlearning_rate=5e-7\n\n# ==============================================================================\n# Algorithm Configuration (RLOO)\n# ==============================================================================\n\nadv_estimator=rloo                        # RLOO advantage estimator\ngamma=1.0\n\n# ==============================================================================\n# Launch Training\n# ==============================================================================\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_batch_size} \\\n    data.truncation='left' \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.gamma=${gamma} \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \\\n    algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} \\\n    algorithm.rollout_correction.rollout_rs=${rollout_rs} \\\n    algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} \\\n    algorithm.rollout_correction.bypass_mode=${bypass_mode} \\\n    algorithm.rollout_correction.loss_type=${loss_type} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=${learning_rate} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.ppo_epochs=${ppo_epochs} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"rollout_corr_rloo_example\" \\\n    trainer.experiment_name=\"rloo_seq_is_pure\" \\\n    trainer.total_epochs=10\n\necho \"Training completed!\"\necho \"\"\necho \"RLOO Configuration:\"\necho \"  - Algorithm: RLOO (REINFORCE Leave-One-Out)\"\necho \"  - Advantage estimator: ${adv_estimator}\"\necho \"  - IS mode: ${rollout_is} (self-normalized: ${rollout_is_batch_normalize})\"\necho \"  - IS threshold: ${rollout_is_threshold}\"\necho \"  - Bypass mode: ${bypass_mode}, loss_type: ${loss_type}\"\necho \"\"\necho \"Monitor these key metrics in wandb:\"\necho \"  - rollout_corr/rollout_is_mean (should be ~1.0 before batch norm)\"\necho \"  - rollout_corr/rollout_is_batch_norm_factor (normalization factor applied)\"\necho \"  - rollout_corr/rollout_is_eff_sample_size (should be >0.5)\"\n"
  },
  {
    "path": "examples/rollout_correction/run_with_rollout_corr_multi_rs.sh",
    "content": "#!/usr/bin/env bash\n# Example: PPO-clip with Rollout Correction using multiple RS criteria\n# Demonstrates chaining token-level and sequence-level rejection sampling\n# (token_k1 + seq_max_k2) alongside optional IS metrics.\n#\n# References:\n#   - Rollout Correction Docs: https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr.md\n#   - Rollout Correction Math: https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr_math.md\n\nset -xeuo pipefail\n\n# ==============================================================================\n# Rollout Correction Configuration (PPO-clip + multi RS)\n# ==============================================================================\n\n# Importance Sampling (IS) weights configuration\nrollout_is=\"token\"                       # Token-level IS for metrics/analysis\nrollout_is_threshold=2.0                 # Upper threshold for IS weights\nrollout_is_batch_normalize=\"false\"       # Keep raw truncated weights\n\n# Rejection Sampling (RS) configuration (multi-criteria)\n# - token_k1 keeps per-token ratios inside [lower, upper]\n# - seq_max_k2 rejects sequences with extreme chi-square spikes\nrollout_rs=\"token_k1,seq_max_k2\"\nrollout_rs_threshold=\"0.6_1.6,2.5\"\n\n# Bypass PPO mode (reuse rollout_log_prob)\nbypass_mode=\"true\"\nloss_type=\"ppo_clip\"\n\n# ==============================================================================\n# Model and Data Configuration\n# ==============================================================================\n\nMODEL_PATH=${MODEL_PATH:-\"Qwen/Qwen2.5-7B\"}\nTRAIN_FILE=${TRAIN_FILE:-\"data/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"data/test.parquet\"}\n\nmax_prompt_length=2048\nmax_response_length=4096\n\n# ==============================================================================\n# Training Configuration\n# ==============================================================================\n\ntrain_batch_size=128\nppo_mini_batch_size=32\nppo_epochs=1\nlearning_rate=3e-6\n\n# ==============================================================================\n# Algorithm Configuration\n# ==============================================================================\n\nadv_estimator=grpo\ngamma=1.0\n\n# ==============================================================================\n# Launch Training\n# ==============================================================================\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_batch_size} \\\n    data.truncation='left' \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.gamma=${gamma} \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \\\n    algorithm.rollout_correction.rollout_is_batch_normalize=${rollout_is_batch_normalize} \\\n    algorithm.rollout_correction.rollout_rs=\\'${rollout_rs}\\' \\\n    algorithm.rollout_correction.rollout_rs_threshold=\\'${rollout_rs_threshold}\\' \\\n    algorithm.rollout_correction.bypass_mode=${bypass_mode} \\\n    algorithm.rollout_correction.loss_type=${loss_type} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=${learning_rate} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${ppo_mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.ppo_epochs=${ppo_epochs} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=\"rollout_corr_multi_rs_example\" \\\n    trainer.experiment_name=\"ppo_clip_multi_rs\" \\\n    trainer.total_epochs=5\n\necho \"Training completed!\"\necho \"\"\necho \"Multi-RS Configuration:\"\necho \"  - rollout_is: ${rollout_is} (threshold=${rollout_is_threshold}, batch_norm=${rollout_is_batch_normalize})\"\necho \"  - rollout_rs: ${rollout_rs}\"\necho \"  - rollout_rs_threshold: ${rollout_rs_threshold}\"\necho \"  - bypass_mode: ${bypass_mode}, loss_type: ${loss_type}\"\necho \"\"\necho \"Track these metrics in wandb:\"\necho \"  - rollout_corr/rollout_rs_token_k1_mean\"\necho \"  - rollout_corr/rollout_rs_seq_max_k2_mean\"\necho \"  - rollout_corr/rollout_rs_masked_fraction\"\n"
  },
  {
    "path": "examples/router_replay/README.md",
    "content": "# Router Replay\n\nRouter Replay is an advanced routing replay functionality within the Verl framework designed for Mixture of Experts (MoE) models. It enables deterministic training by recording and replaying routing decisions, ensuring consistent model behavior across training runs.\n\n\n## Key Features\n\n### Multiple Operating Modes\n- **`disabled`**: Router replay functionality is completely disabled\n- **`R2`**: Standard router replay mode for recording and replaying routing decisions\n- **`R3`**: Rollout-specific router replay mode optimized for reinforcement learning workflows\n\n### Core Capabilities\n- **Seamless Integration**: Works with reinforcement learning pipelines including PPO\n- **Distributed Training Support**: Compatible with multi-GPU and multi-node training environments\n- **Flexible Configuration**: Easy to configure via YAML files or command-line parameters\n\n## Configuration\n\n### RouterReplayConfig Parameters\n\n```yaml\nrouter_replay:\n  mode: \"disabled\"  # Available options: disabled, R2, R3\n  record_file: null  # Path for recording routing decisions\n  replay_file: null   # Path for replaying recorded decisions\n```\n\n## Quick Start Guide\n\n### Enabling R2 Mode\n\n#### Configuration File Method\nAdd the following to your training configuration:\n\n```yaml\nactor:\n  router_replay:\n    mode: \"R2\"\n```\n\n#### Command Line Method\nEnable R2 mode via command-line parameters:\n\n```bash\nactor_rollout_ref.actor.router_replay.mode=\"R2\"\n```\n\n### Enabling R3 Mode\n\n#### Configuration File Method\nConfigure both actor and rollout settings:\n\n```yaml\n# Actor configuration\nrouter_replay:\n  mode: \"R3\"\n\n# Rollout configuration  \nenable_rollout_routing_replay: True\n```\n\n#### Command Line Method\nEnable R3 mode via command-line parameters:\n\n```bash\nactor_rollout_ref.actor.router_replay.mode=\"R3\"\nactor_rollout_ref.rollout.enable_rollout_routing_replay=True\n```\n\nR3 mode requires the rollout backend to support returning router selection results. Currently, this functionality is being tested based on the vllm implementation at https://github.com/vllm-project/vllm/pull/28284 as well as bug fix at https://github.com/vllm-project/vllm/pull/33013 and SGLang implementation at https://github.com/sgl-project/sglang/commit/bed301a5acaa9577c9aa706468bdf242f6a43051.\n"
  },
  {
    "path": "examples/router_replay/run_qwen30_a3b_megatron_sglang.sh",
    "content": "\nset -x\n\nNODES=6\n\n# R2: enable routing replay\n# R3: enable rollout routing replay\n# If enabling R3, please set actor_rollout_ref.rollout.enable_rollout_routing_replay=True\n# R3 example is based on SGLang related commit https://github.com/sgl-project/sglang/commit/bed301a5acaa9577c9aa706468bdf242f6a43051\n\nROUTING_REPLAY_MODE=\"R3\" \n\nif [ \"$ROUTING_REPLAY_MODE\" = \"R3\" ]; then\n    ENABLE_ROLLOUT_ROUTING_REPLAY=True\nelse\n    ENABLE_ROLLOUT_ROUTING_REPLAY=False\nfi\n\nDIST_CKPT_PATH=\"\"\nHF_MODEL_PATH=\"\"\nTRAIN_DATA_PATH=\"\"\nTEST_DATA_PATH=\"\"\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\nPP=6\nVPP=None\nTP=1\nEP=8\nETP=1\nSGLANG_INFER_TP=4\noffload=True\ngpu_memory_utilization=0.65\nbs=3\nmicro_bs=3\nuse_dynamic_bsz=False\nmax_prompt_length=512\nmax_response_length=512\nppo_mini_batch_size=3\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\n\nUSE_LEGACY_WORKER_IMPL=\"enable\" # disable, enable\nif [ \"$USE_LEGACY_WORKER_IMPL\" = \"disable\" ]; then\n    ROUTING_REPLAY_MODE_ARG=\"actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTING_REPLAY_MODE}\"\n    remove_padding=True\nelse\n    ROUTING_REPLAY_MODE_ARG=\"actor_rollout_ref.actor.router_replay.mode=${ROUTING_REPLAY_MODE}\"\n    remove_padding=False\nfi\nexper_name=Node${NODES}_bs${bs}_${PP}${TP}${EP}${ETP}_${SGLANG_INFER_TP}_minbs${ppo_mini_batch_size}_micro_bs${micro_bs}\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$TRAIN_DATA_PATH \\\n    data.val_files=$TEST_DATA_PATH \\\n    data.train_batch_size=$bs \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.model.use_remove_padding=${remove_padding} \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    ${ROUTING_REPLAY_MODE_ARG} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=False \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$SGLANG_INFER_TP \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.enable_rollout_routing_replay=True \\\n    actor_rollout_ref.rollout.skip_tokenizer_init=True \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=['console'] \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name=\"$exper_name\" \\\n    trainer.nnodes=$NODES \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_training_steps=50000 \\\n    trainer.balance_batch=False \\\n    trainer.use_legacy_worker_impl=${USE_LEGACY_WORKER_IMPL} \\\n    trainer.val_before_train=False 2>&1\n"
  },
  {
    "path": "examples/router_replay/run_qwen30_a3b_megatron_vllm.sh",
    "content": "\nset -x\n\nNODES=1\n\n# R2: enable routing replay\n# R3: enable rollout routing replay\n# If enabling R3, please set actor_rollout_ref.rollout.enable_rollout_routing_replay=True\n# R3 example is based on vllm related pr:\n#   - https://github.com/vllm-project/vllm/pull/28284\n#   - https://github.com/vllm-project/vllm/pull/33013\n\nROUTING_REPLAY_MODE=\"R3\" \n\nif [ \"$ROUTING_REPLAY_MODE\" = \"R3\" ]; then\n    ENABLE_ROLLOUT_ROUTING_REPLAY=True\nelse\n    ENABLE_ROLLOUT_ROUTING_REPLAY=False\nfi\n\nDIST_CKPT_PATH=\"\"\nHF_MODEL_PATH=\"\"\nTRAIN_DATA_PATH=\"\"\nTEST_DATA_PATH=\"\"\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\nPP=1\nVPP=None\nTP=2\nEP=8\nETP=1\nVLLM_INFER_TP=2\noffload=True\ngpu_memory_utilization=0.65\nbs=8\nmicro_bs=3\nuse_dynamic_bsz=True\nmax_prompt_length=1024\nmax_response_length=1024\nppo_mini_batch_size=8\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\n\nUSE_LEGACY_WORKER_IMPL=\"disable\" # disable, enable\nif [ \"$USE_LEGACY_WORKER_IMPL\" = \"disable\" ]; then\n    ROUTING_REPLAY_MODE_ARG=\"actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTING_REPLAY_MODE}\"\n    remove_padding=True\nelse\n    ROUTING_REPLAY_MODE_ARG=\"actor_rollout_ref.actor.router_replay.mode=${ROUTING_REPLAY_MODE}\"\n    remove_padding=False\nfi\nexper_name=Node${NODES}_bs${bs}_${PP}${TP}${EP}${ETP}_${VLLM_INFER_TP}_minbs${ppo_mini_batch_size}_micro_bs${micro_bs}\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$TRAIN_DATA_PATH \\\n    data.val_files=$TEST_DATA_PATH \\\n    data.train_batch_size=$bs \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.model.use_remove_padding=$remove_padding \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    ${ROUTING_REPLAY_MODE_ARG} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=flex \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$VLLM_INFER_TP \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=$micro_bs \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$PP \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$TP \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$ETP \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=['console'] \\\n    trainer.project_name='verl_grpo_example_gsm8k_math' \\\n    trainer.experiment_name=\"$exper_name\" \\\n    trainer.nnodes=$NODES \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.total_training_steps=50000 \\\n    trainer.balance_batch=False \\\n    trainer.use_legacy_worker_impl=${USE_LEGACY_WORKER_IMPL} \\\n    trainer.val_before_train=False 2>&1\n"
  },
  {
    "path": "examples/sapo_trainer/run_qwen30b_sapo.sh",
    "content": "#!/bin/bash\n#SBATCH --job-name=sapo-30B\n#SBATCH --partition=main\n#SBATCH --nodes=1                # Number of nodes\n#SBATCH --ntasks-per-node=1      # One task per node\n#SBATCH --cpus-per-task=128      # cpu-cores per task (>1 if multi-threaded tasks)\n#SBATCH --gres=gpu:8\n#SBATCH --gpus-per-node=8\n#SBATCH --mem=0\n#SBATCH --exclusive\n#SBATCH --time=500:00:00\n#SBATCH --output=logs/sapo/30B/frugal_math/%x_%j.out\n#SBATCH --error=logs/sapo/30B/frugal_math/%x_%j.err\n\n# This script runs the training of RL on multi-nodes. It does resume automatically from latest checkpoint if the run crashes.\n# Example run with Qwen3-30B SAPO with new model engine\nset -x\n\nexport WANDB_API_KEY=YOUR_WANDB_API_KEY_HERE\nENV_NAME=verl_0_6_1\n\n# Ensure Python can import the top-level verl package even when the script is relocated by Slurm\nif [[ -n \"$SLURM_SUBMIT_DIR\" && -d \"$SLURM_SUBMIT_DIR\" ]]; then\n    cd \"$SLURM_SUBMIT_DIR\"\n    SCRIPT_SOURCE_DIR=\"$SLURM_SUBMIT_DIR\"\nelse\n    SCRIPT_SOURCE_DIR=$(cd -- \"$(dirname \"${BASH_SOURCE[0]}\")\" >/dev/null 2>&1 && pwd)\nfi\nREPO_ROOT=$(cd -- \"$SCRIPT_SOURCE_DIR/../..\" >/dev/null 2>&1 && pwd)\nVERL_REPO_ROOT=\"$REPO_ROOT\"\n\nadd_repo_to_pythonpath() {\n    if [[ -z \"$PYTHONPATH\" ]]; then\n        export PYTHONPATH=\"$VERL_REPO_ROOT\"\n    else\n        case \":$PYTHONPATH:\" in\n            *\":$VERL_REPO_ROOT:\"*) ;;\n            *) export PYTHONPATH=\"$VERL_REPO_ROOT:$PYTHONPATH\" ;;\n        esac\n    fi\n}\n\nadd_repo_to_pythonpath\n\n# can make training faster depending on clusters\nexport NCCL_IBEXT_DISABLE=1\nexport NCCL_NVLS_ENABLE=1\nexport NCCL_IB_HCA=mlx5\nexport UCX_NET_DEVICES=mlx5_0:1,mlx5_1:1,mlx5_2:1,mlx5_3:1,mlx5_4:1,mlx5_5:1,mlx5_6:1,mlx5_7:1\n\n# Determine how many nodes were allocated. \nNNODES=${SLURM_JOB_NUM_NODES}\nexport NNODES\n\n# Determine how many GPUs we actually have on the master node.\n# Carefull! Assumes all nodes have same number of GPUs! \n# SLURM sets SLURM_GPUS_PER_NODE only when #SBATCH --gpus-per-node is used, not with --gres.\n# uncomment below line to manually set number of gpus per node if not using --gpus-per-node\n# export SLURM_GPUS_PER_NODE=8\n# SLURM_GPUS_PER_NODE=${SLURM_GPUS_PER_NODE:-$(nvidia-smi -L | wc -l)} # 8\n# export SLURM_GPUS_PER_NODE\necho \"SLURM_GPUS_PER_NODE: $SLURM_GPUS_PER_NODE\"\n\n# Set DATA_ROOT to current working directory if not set\nDATA_ROOT=${DATA_ROOT:-$PWD}\necho \"DATA_ROOT: $DATA_ROOT\"\n\n# wandb logging\nbackend=fsdp # fsdp, fsdp2, megatron\nproject_name=RL4LLM\n# experiment_name=qwen3-30B-base-sapo-$backend\nexperiment_name=qwen3-30B-base-vanilla-$backend\ndefault_local_dir=$DATA_ROOT/checkpoint/$project_name/$experiment_name\n\n# ===================================== Algorithm =====================================\nadv_estimator=grpo\nloss_mode=sapo # explicitly specify sapo! default is vanilla and is not compatible with SAPO. It uses clipping instead of smoothing.\n\n# reference policy\nuse_kl_in_reward=False\nkl_coef=0.001\nuse_kl_loss=False\nkl_loss_coef=0.001\n\n# Positive and negative tau for smoothing function in SAPO (https://arxiv.org/pdf/2511.20347)\n# default values used in the paper with Qwen3-30B-A3B-Base\n# clipping is not used in SAPO!\ntau_pos=1.0\ntau_neg=1.05\n\nactor_lr=1e-6\ncritic_lr=2e-6\ngae_gamma=1.0\ngae_lam=0.95\ncritic_warmup=0\n\n# ===================================== Data/Model =====================================\n\nfirst_time_dataset_prep=true\nHF_DATA_PATH=\"BytedTsinghua-SIA/DAPO-Math-17k\"\nSTAGE=\"stage-1\"\n\nif [ \"$first_time_dataset_prep\" = true ]; then\n    echo \"Preparing training dataset...\"\n    python $VERL_REPO_ROOT/examples/data_preprocess/dapo_multiturn_w_tool.py \\\n        --local_save_dir $DATA_ROOT/dataset/dapo/ \n    echo \"Training dataset prepared.\"\n\n    echo \"Preparing testing dataset...\"\n    python $VERL_REPO_ROOT/examples/data_preprocess/aime2024_multiturn_w_tool.py \\\n        --local_save_dir $DATA_ROOT/dataset/test/aime_24/\n    echo \"Testing dataset prepared.\"\n\n    echo \"Dataset preparation completed.\"\nfi\n\ntrain_files=$DATA_ROOT/dataset/dapo/train.parquet\ntest_files=$DATA_ROOT/dataset/test/aime_24/train.parquet\n\nactor_model_path=Qwen/Qwen3-30B-A3B-Base\ncritic_model_path=$actor_model_path\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\ntrain_batch_size=256\nppo_mini_batch_size=32\nn_resp_per_prompt=16\nn_resp_per_prompt_val=1\n\n# ===================================== Training =====================================\nactor_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 3))\ncritic_max_token_len_per_gpu=$(((max_prompt_length + max_response_length) * 4))\n\nenable_gradient_checkpointing=True\nparam_offload=False\noptimizer_offload=False\n\n\nVAL_BEFORE_TRAIN=False\nSAVE_FREQ=-1 # we do not save!\nTEST_FREQ=10\nTOTAL_EPOCHS=10\nTOTAL_TRAINING_STEPS=2000\n\n# FSDP parallelism config\nUSP_SIZE=4\nACTOR_FSDP_CONFIG=\"\n    actor_rollout_ref.actor.fsdp_config.strategy=$backend \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=$param_offload \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=$optimizer_offload \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$USP_SIZE\"\n\n# Megatron parallelism config\nTP_SIZE=1\nCP_SIZE=1\nPP_SIZE=1\nVPP_SIZE=null\nEP_SIZE=8\nETP_SIZE=1\nACTOR_MEGATRON_CONFIG=\"\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$TP_SIZE \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$CP_SIZE \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$PP_SIZE \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$VPP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$EP_SIZE \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ETP_SIZE \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_granularity=full \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.recompute_num_layers=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True\"\n\n# Actor model config\nACTOR_CONFIG=\"\n    actor_rollout_ref.actor.optim.lr=$actor_lr \\\n    actor_rollout_ref.model.path=$actor_model_path \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=${enable_gradient_checkpointing} \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.tau_pos=$tau_pos \\\n    actor_rollout_ref.actor.tau_neg=$tau_neg \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode}\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu\"\n\n# Critic model config\nCIRITC_CONFIG=\"\n    critic.optim.lr=$critic_lr \\\n    critic.model.path=$critic_model_path \\\n    critic.model.use_remove_padding=True \\\n    critic.ppo_max_token_len_per_gpu=$critic_max_token_len_per_gpu \\\n    critic.ulysses_sequence_parallel_size=$USP_SIZE\"\n\nCRITIC_FSDP_CONFIG=\"${ACTOR_FSDP_CONFIG//actor_rollout_ref.actor/critic.model}\"\nCRITIC_MEGATRON_CONFIG=\"${ACTOR_MEGATRON_CONFIG//actor_rollout_ref.actor/critic}\"\n\nif [[ $backend == \"megatron\" ]]; then\n    CONFIG_NAME=ppo_megatron_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_MEGATRON_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_MEGATRON_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nelse # fsdp, fsdp2\n    CONFIG_NAME=ppo_trainer\n    ACTOR_CONFIG=\"$ACTOR_CONFIG $ACTOR_FSDP_CONFIG\"\n    if [[ $adv_estimator == \"gae\" ]]; then\n        CIRITC_CONFIG=\"$CIRITC_CONFIG $CRITIC_FSDP_CONFIG\"\n    else\n        CIRITC_CONFIG=\"\"\n    fi\nfi\n\n# ===================================== Inference =====================================\nrollout_engine=vllm\ninfer_tp=4\ninfer_dp=1\ninfer_ep=1\ngpu_memory_utilization=0.8\n\nROLLOUT_CONFIG=\"\n    actor_rollout_ref.rollout.name=$rollout_engine \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \\\n    actor_rollout_ref.rollout.data_parallel_size=$infer_dp \\\n    actor_rollout_ref.rollout.expert_parallel_size=$infer_ep \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=$gpu_memory_utilization \\\n    actor_rollout_ref.rollout.n=$n_resp_per_prompt \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.7 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val\"\n\n# ===================================== Reward =====================================\nREWARD_CONFIG=\"\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length}\"\n\n\n# ============================= Prepare RAY on Slurm ===============================\n\n# we should activate it before we start ray to avoid errors\necho \"Activating $ENV_NAME environment...\"\neval \"$(conda shell.bash hook)\"\nconda deactivate\nconda activate \"$ENV_NAME\"\nadd_repo_to_pythonpath\n\nexport VLLM_ATTENTION_BACKEND=FLASH_ATTN\nexport RAY_memory_monitor_refresh_ms=0\nexport RAY_LOGGING_LEVEL=DEBUG\nexport HYDRA_FULL_ERROR=1\n\n# Let Ray know how many nodes to expect\nexport RAY_NUM_NODES=$NNODES\n\n# Get head node and its IP\nnodes=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\")\nnodes_array=($nodes)\nhead_node=${nodes_array[0]}\nhead_node_ip=$(srun --nodes=1 --ntasks=1 -w \"$head_node\" hostname --ip-address)\n\n# Convert to IPv4 if needed\nif [[ \"$head_node_ip\" == *\" \"* ]]; then\n  IFS=' ' read -ra ADDR <<<\"$head_node_ip\"\n  if [[ ${#ADDR[0]} -gt 16 ]]; then\n    head_node_ip=${ADDR[1]}\n  else\n    head_node_ip=${ADDR[0]}\n  fi\n  echo \"IPV6 address detected. Using IPV4: $head_node_ip\"\nfi\n\nport=6379\nip_head=$head_node_ip:$port\nexport MASTER_ADDR=$head_node_ip\nexport MASTER_PORT=$port\nexport ip_head\n\necho \"Starting Ray HEAD at $head_node ($ip_head)\"\nuntil nvidia-smi > /dev/null 2>&1; do\n  echo \"Waiting for GPU visibility...\"\n  sleep 2\ndone\nsrun --nodes=1 --ntasks=1 -w \"$head_node\" \\\n  ray start --head --node-ip-address=\"$head_node_ip\" --port=$port \\\n  --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n\nsleep 10\n\nworker_num=$((SLURM_JOB_NUM_NODES - 1))\nfor ((i = 1; i <= worker_num; i++)); do\n  node_i=${nodes_array[$i]}\n  echo \"Starting WORKER $i at $node_i\"\n  until nvidia-smi > /dev/null 2>&1; do\n    echo \"Waiting for GPU visibility...\"\n    sleep 2\n  done\n  srun --nodes=1 --ntasks=1 -w \"$node_i\" \\\n    ray start --address \"$ip_head\" --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n  sleep 5\ndone\n\n# Final launch barrier\nsleep 10\n\n# ================================= Launch Training ================================\n\necho \"Using $SLURM_NNODES nodes for training...\"\n\necho \"==== Confirming Ray sees all GPUs ====\"\npython -c \"import ray; ray.init(address='auto'); print(ray.cluster_resources())\"\necho \"==== Done checking resources ====\"\n\n# we should activate it before we start ray to avoid errors\necho \"Activating $ENV_NAME environment...\"\neval \"$(conda shell.bash hook)\"\nconda deactivate\nconda activate \"$ENV_NAME\"\nadd_repo_to_pythonpath\n\nsrun --overlap --nodes=${NNODES} --ntasks=1 -w \"$head_node\"\\\n    python -m verl.trainer.main_ppo \\\n        --config-path=./config \\\n        --config-name=$CONFIG_NAME \\\n        algorithm.adv_estimator=$adv_estimator \\\n        algorithm.use_kl_in_reward=$use_kl_in_reward \\\n        algorithm.kl_ctrl.kl_coef=$kl_coef \\\n        algorithm.gamma=$gae_gamma \\\n        algorithm.lam=$gae_lam \\\n        data.train_files=\"$train_files\" \\\n        data.val_files=\"$test_files\" \\\n        data.return_raw_chat=True \\\n        data.train_batch_size=$train_batch_size \\\n        data.max_prompt_length=$max_prompt_length \\\n        data.max_response_length=$max_response_length \\\n        data.filter_overlong_prompts=True \\\n        data.filter_overlong_prompts_workers=64 \\\n        data.truncation='error' \\\n        trainer.use_legacy_worker_impl=disable \\\n        trainer.critic_warmup=$critic_warmup \\\n        trainer.logger=['console','wandb'] \\\n        trainer.project_name=$project_name \\\n        trainer.experiment_name=$experiment_name \\\n        trainer.default_local_dir=$default_local_dir \\\n        trainer.n_gpus_per_node=$SLURM_GPUS_PER_NODE \\\n        trainer.nnodes=$NNODES \\\n        trainer.val_before_train=$VAL_BEFORE_TRAIN \\\n        trainer.log_val_generations=100 \\\n        trainer.save_freq=$SAVE_FREQ \\\n        trainer.test_freq=$TEST_FREQ \\\n        trainer.total_epochs=$TOTAL_EPOCHS \\\n        trainer.total_training_steps=$TOTAL_TRAINING_STEPS \\\n        $ACTOR_CONFIG \\\n        $CIRITC_CONFIG \\\n        $ROLLOUT_CONFIG \\\n        $REWARD_CONFIG\n"
  },
  {
    "path": "examples/sapo_trainer/run_qwen3_8b_sapo_npu.sh",
    "content": "set -euxo pipefail\n\nulimit -n 32768\n\n## Basic Environment Settings\nexport RAY_DEDUP_LOGS=0\nexport HYDRA_FULL_ERROR=1\nexport TASK_QUEUE_ENABLE=1\nexport HCCL_EXEC_TIMEOUT=3600\nexport HCCL_CONNECT_TIMEOUT=3600\nexport HCCL_ASYNC_ERROR_HANDLING=0\nexport CPU_AFFINITY_CONF=1\nexport VLLM_USE_V1=1\n\nproject_name='SAPO-Qwen3'\nexp_name='SAPO-Qwen3-8B-npu'\ngen_tp=2\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/Qwen3-8B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/dataset/dapo_processed/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/dataset/aime-24_processed/train.parquet\"}\n\n# reference policy\nuse_kl_in_reward=False\nkl_coef=0.001\nuse_kl_loss=False\nkl_loss_coef=0.001\n\n# ------Algorithm settings-------\n# Positive and negative tau for smoothing function in SAPO (https://arxiv.org/pdf/2511.20347)\n# default values used in the paper with Qwen3-30B-A3B-Base\n# clipping is not used in SAPO!\n\nloss_mode=sapo # explicitly specify sapo! default is vanilla and is not compatible with SAPO. It uses clipping instead of smoothing.\n\ntau_pos=1.0\ntau_neg=1.05\n\ngae_gamma=1.0\ngae_lam=0.95\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.use_kl_in_reward=$use_kl_in_reward \\\n    algorithm.kl_ctrl.kl_coef=$kl_coef \\\n    algorithm.gamma=$gae_gamma \\\n    algorithm.lam=$gae_lam \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.filter_overlong_prompts_workers=64 \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=10 \\\n    actor_rollout_ref.actor.tau_pos=$tau_pos \\\n    actor_rollout_ref.actor.tau_neg=$tau_neg \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=${loss_mode} \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\"]' \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.resume_mode=auto \\\n    actor_rollout_ref.actor.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.ref.fsdp_config.forward_prefetch=True \\\n    actor_rollout_ref.actor.entropy_from_logits_with_chunking=True \\\n    actor_rollout_ref.ref.entropy_from_logits_with_chunking=True \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15\n"
  },
  {
    "path": "examples/sft/gsm8k/run_deepseek_6b7.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_deepseek_6b7.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    model.path=deepseek-ai/deepseek-coder-6.7b-instruct \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-deepseek-coder-6.7b-instruct \\\n    trainer.total_epochs=4 \\\n    trainer.logger='[\"console\",\"wandb\"]' $@"
  },
  {
    "path": "examples/sft/gsm8k/run_gemma_2b.sh",
    "content": "# Tested with 2 & 4 GPUs\n\nset -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_gemma_2b.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    model.path=google/gemma-2b-it \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-gemma-2b-it \\\n    trainer.total_epochs=2 \\\n    trainer.logger='[\"console\",\"wandb\"]' $@"
  },
  {
    "path": "examples/sft/gsm8k/run_gemma_7b.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_gemma_7b.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    model.path=google/gemma-1.1-7b-it \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-gemma-1.1-7b-it \\\n    trainer.total_epochs=4 \\\n    trainer.logger='[\"console\",\"wandb\"]' $@\n"
  },
  {
    "path": "examples/sft/gsm8k/run_mimo_megatron_mtp.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nNUM_GPUS=${NUM_GPUS:-8}\nSP_SIZE=${SP_SIZE:-1}\nTP_SIZE=${TP_SIZE:-1}\nPP_SIZE=${PP_SIZE:-1}\nVPP_SIZE=${VPP_SIZE:-null}\nCP_SIZE=${CP_SIZE:-1}\nPAD_MODE=${PAD_MODE:-no_padding}\nUSE_REMOVE_PADDING=${USE_REMOVE_PADDING:-False}\nLR=\"1e-5\"\nMINLR=\"1e-6\"\n\nexport VERL_SFT_LOGGING_LEVEL=INFO\n\nbackend=${BACKEND:-megatron}\n\nTENSORBOARD_DIR=~/tensorboard\n\nMASTER_ADDR=${MASTER_ADDR:-localhost}\nMASTER_PORT=${MASTER_PORT:-29500}\nNNODES=${NNODES:-1}\nRANK=${RANK:-0}\n\nENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n\n# Note the default MultiturnSFT Dataset requires all the sys/user/assistant in 'data.message_key'\nDATASET_DIR=${DATASET_DIR:-~/dataset/rl/gsm8k}\nTRAIN_FILES=${DATASET_DIR}/train.parquet\nVAL_FILES=${DATASET_DIR}/eval.parquet\n\nproject_name=verl_sft_test\n\nRESUME_MODE=disable\n\nMODEL_PATH=\"XiaomiMiMo/MiMo-7B-RL\"\nckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-${backend}}\n\n# currently relies on these two commits that is not on master\nPYPATH=$HOME/pythonpath\nmkdir -p $PYPATH && cd $PYPATH\n[ -d Megatron-LM ] || git clone https://github.com/NVIDIA/Megatron-LM -b dev && (cd Megatron-LM; git checkout 23e092f41ec8bc659020e401ddac9576c1cfed7e)\n[ -d mbridge ] || git clone https://github.com/ArronHZG/mbridge -b feature/verl_mtp && (cd mbridge; git checkout 6bf2d45a15dc4fb52d2f0c38ff546bee33447d10)\ncd -\nexport PYTHONPATH=$PYTHONPATH:$PYPATH/mbridge:$PYPATH/Megatron-LM\n\n\nMEGATRON_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    optim=${backend} \\\n    optim.lr=${LR} \\\n    optim.min_lr=${MINLR} \\\n    optim.lr_warmup_steps=10 \\\n    optim.weight_decay=0.1 \\\n    optim.betas='[0.9,0.95]' \\\n    optim.clip_grad=1.0 \\\n    optim.lr_warmup_init=0 \\\n    optim.lr_decay_style=cosine \\\n    engine.override_transformer_config.recompute_method=uniform \\\n    engine.override_transformer_config.recompute_granularity=full \\\n    engine.override_transformer_config.recompute_num_layers=1 \\\n    engine.use_dist_checkpointing=False \\\n    engine.tensor_model_parallel_size=${TP_SIZE} \\\n    engine.pipeline_model_parallel_size=${PP_SIZE} \\\n    engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \\\n    engine.context_parallel_size=${CP_SIZE} \\\n    engine.use_mbridge=True \\\n    \"\n\nENGINE_CONFIG=\"$MEGATRON_ENGINE_CONFIG\"\necho \"Using megatron engine\"\nexp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-lr-${MINLR}-${LR}\n\nmkdir -p \"${ckpts_home}\"\n\n$COMMAND \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${TRAIN_FILES}\" \\\n    data.train_batch_size=64 \\\n    data.micro_batch_size_per_gpu=2 \\\n    data.pad_mode=${PAD_MODE} \\\n    data.truncation=error \\\n    data.max_length=1024 \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=2048 \\\n    data.messages_key=prompt \\\n    data.num_workers=0 \\\n    model.path=$MODEL_PATH \\\n    model.use_remove_padding=${USE_REMOVE_PADDING} \\\n    model.trust_remote_code=True \\\n    model.mtp.enable=True \\\n    ${ENGINE_CONFIG} \\\n    trainer.test_freq=after_each_epoch \\\n    trainer.save_freq=-1 \\\n    trainer.logger=\"['console']\" \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_epochs=1 \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.resume_mode=${RESUME_MODE}\n    "
  },
  {
    "path": "examples/sft/gsm8k/run_nemotron_nano_v3.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n################################################### environment ###################################################\n### # 1. use docker image `verlai/verl:vllm015.dev`` and install correct dependencies:\n# pip install nvidia-modelopt\n# MAX_JOBS=32 pip install git+https://github.com/Dao-AILab/causal-conv1d.git --no-build-isolation --no-cache-dir\n# MAX_JOBS=32 pip install git+https://github.com/state-spaces/mamba.git --no-build-isolation --no-cache-dir\n# pip install --no-deps git+https://github.com/NVIDIA-NeMo/Megatron-Bridge \n# pip install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_dev_r0.16.0\n# unset ROCR_VISIBLE_DEVICES\n# unset PYTORCH_CUDA_ALLOC_CONF\n\nENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n\nTRAIN_FILES=$HOME/data/gsm8k/train.parquet\nVAL_FILES=$HOME/data/gsm8k/eval.parquet\n\nbackend=${BACKEND:-megatron}\n\nproject_name=verl_sft_gsm8k\n\nRESUME_MODE=auto\nMODEL_NAME=${MODEL_NAME:-NVIDIA-Nemotron-3-Nano-30B-A3B-BF16}\n\nTP_SIZE=${TP_SIZE:-8}\nPP_SIZE=${PP_SIZE:-1}\nVPP_SIZE=${VPP_SIZE:-null}\nCP_SIZE=${CP_SIZE:-1}\nEP_SIZE=${EP_SIZE:-8}\nETP_SIZE=${ETP_SIZE:-1}\n\nPAD_MODE=${PAD_MODE:-no_padding}\n\nUSE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True}\n\nDTYPE=${DTYPE:-\"bfloat16\"}\n\n\nMEGATRON_ENGINE_CONFIG=(\n    engine=${backend}\n    optim=${backend}\n    optim.lr=2e-5\n    optim.lr_warmup_steps=5\n    optim.weight_decay=0.1\n    optim.betas=\"[0.9,0.95]\"\n    optim.clip_grad=1.0\n    optim.lr_warmup_init=0\n    optim.lr_decay_style=cosine\n    optim.min_lr=2e-6\n    +optim.override_optimizer_config.optimizer_offload_fraction=1\n    +optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True\n    +optim.override_optimizer_config.use_precision_aware_optimizer=True\n    +optim.override_optimizer_config.optimizer_cpu_offload=True\n    engine.tensor_model_parallel_size=${TP_SIZE}\n    engine.pipeline_model_parallel_size=${PP_SIZE}\n    engine.virtual_pipeline_model_parallel_size=${VPP_SIZE}\n    engine.context_parallel_size=${CP_SIZE}\n    engine.use_mbridge=True\n    engine.dtype=${DTYPE}\n    engine.vanilla_mbridge=False\n    engine.expert_model_parallel_size=${EP_SIZE}\n    engine.expert_tensor_parallel_size=${ETP_SIZE}\n    engine.override_transformer_config.attention_backend=auto\n    +engine.override_transformer_config.recompute_method=uniform\n    +engine.override_transformer_config.recompute_granularity=full\n    +engine.override_transformer_config.recompute_num_layers=1\n)\n\nENGINE_CONFIG=\"${MEGATRON_ENGINE_CONFIG[@]}\"\necho \"Using megatron engine\"\nexp_name=${MODEL_NAME}-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-megatron-20260210\n\ntorchrun --nnodes=1 --nproc_per_node=8 ${ENTRYPOINT} \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=96 \\\n    data.max_length=2048 \\\n    data.pad_mode=${PAD_MODE} \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=2048 \\\n    data.messages_key=messages \\\n    data.ignore_input_ids_mismatch=True \\\n    model.path=$MODEL_PATH \\\n    model.use_remove_padding=${USE_REMOVE_PADDING} \\\n    model.trust_remote_code=True \\\n    ${ENGINE_CONFIG} \\\n    trainer.test_freq=-1 \\\n    trainer.save_freq=500 \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_epochs=1 \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.resume_mode=${RESUME_MODE} \\\n    trainer.max_ckpt_to_keep=10 \\\n    checkpoint.save_contents=[model,optimizer,extra]"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen3_30b_automodel.sh",
    "content": "# Requires: Automodel, transformers>=5.3.0, torchao\n# MoE also requires: grouped_gemm (github.com/fanshiqing/grouped_gemm v1.1.4)\n\nset -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen3_30b_automodel.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/hellaswag_sft/hellaswag_sft.parquet \\\n    data.val_files=$HOME/data/hellaswag_sft/hellaswag_sft.parquet \\\n    data.train_batch_size=512 \\\n    data.max_length=2048 \\\n    data.truncation=left \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=8192 \\\n    data.messages_key=messages \\\n    data.ignore_input_ids_mismatch=True \\\n    data.train_max_samples=-1 \\\n    data.val_max_samples=1024 \\\n    model=hf_model \\\n    model.path=Qwen/Qwen3-30B-A3B-Base \\\n    model.trust_remote_code=True \\\n    model.use_remove_padding=True \\\n    engine=automodel \\\n    engine.distributed_strategy=fsdp2 \\\n    engine.tp_size=1 \\\n    engine.pp_size=1 \\\n    engine.cp_size=1 \\\n    engine.ep_size=8 \\\n    engine.backend_config.dispatcher=deepep \\\n    engine.backend_config.attn=te \\\n    engine.backend_config.linear=te \\\n    engine.backend_config.rms_norm=torch_fp32 \\\n    engine.backend_config.enable_fsdp_optimizations=True \\\n    engine.backend_config.experts=torch_mm \\\n    engine.activation_checkpointing=True \\\n    engine.model_dtype=bf16 \\\n    engine.attn_implementation=te \\\n    engine.use_torch_compile=False \\\n    optim=automodel \\\n    optim.optimizer=FusedAdam \\\n    optim.optimizer_impl=transformer_engine.pytorch.optimizers.fused_adam \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.1 \\\n    optim.weight_decay=0 \\\n    optim.betas='[0.9,0.95]' \\\n    optim.clip_grad=1.0 \\\n    optim.init_lr_ratio=0.1 \\\n    optim.min_lr_ratio=0.01 \\\n    optim.lr_scheduler_type=cosine \\\n    optim.master_weights=true \\\n    optim.store_param_remainders=true \\\n    optim.exp_avg_dtype=bf16 \\\n    optim.exp_avg_sq_dtype=bf16 \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=hellaswag-sft \\\n    trainer.experiment_name=hellaswag-sft-qwen3-30b-automodel \\\n    trainer.total_epochs=2 \\\n    trainer.total_training_steps=100 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=10 \\\n    trainer.logger=console \\\n    trainer.seed=1111 \\\n    trainer.nnodes=1 \\\n    trainer.resume_mode=disable $@\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen3_5_megatron.sh",
    "content": "#!/usr/bin/env bash\n# Qwen3.5-397B-A17B SFT with Megatron backend + mbridge\n#\n# Requirements:\n#   - 128+ GPUs (80GB each, e.g. 16x8 H100/H200)\n#   - Docker: verlai/verl:vllm015 (or equivalent)\n#   - Additional packages on top of the base image:\n#       pip install --upgrade transformers\n#       pip install flash-linear-attention\n#       pip install -U git+https://github.com/ISEEKYAN/mbridge.git\n#   - Megatron-LM==0.16.0\n#\n# Qwen3.5 architecture notes:\n#   Qwen3.5 uses Gated Delta Net (GDN) linear attention which currently does\n#   NOT support packed sequences (THD format) in Megatron-LM. Therefore:\n#     - engine.use_remove_padding=False  (forces bshd compute format)\n#     - data.use_dynamic_bsz=False       (required for bshd mode)\n#\n#   Once https://github.com/NVIDIA/Megatron-LM/pull/2644 is merged, THD\n#   format will be supported and engine.use_remove_padding can be set to True\n#   for better performance.\n#\n# Tested parallelism config (128 GPUs / 16 nodes):\n#   TP=2 PP=4 EP=32 CP=1\n\nset -xeuo pipefail\n\n# ============================================================\n# Distributed\n# ============================================================\nNUM_GPUS=${NUM_GPUS:-8}\nMASTER_ADDR=${MASTER_ADDR:-localhost}\nMASTER_PORT=${MASTER_PORT:-29500}\nNNODES=${NNODES:-16}\nNODE_RANK=${NODE_RANK:-0}\n\n# ============================================================\n# Data\n# ============================================================\nDATASET_DIR=${DATASET_DIR:-~/dataset}\nTRAIN_FILES=${TRAIN_FILES:-${DATASET_DIR}/train.parquet}\n\n# ============================================================\n# Model\n# ============================================================\nMODEL_PATH=${MODEL_PATH:-Qwen/Qwen3.5-397B-A17B}\n\n# ============================================================\n# Parallelism\n# ============================================================\nTP_SIZE=${TP_SIZE:-2}\nPP_SIZE=${PP_SIZE:-4}\nVPP_SIZE=${VPP_SIZE:-null}\nCP_SIZE=${CP_SIZE:-1}\nEP_SIZE=${EP_SIZE:-32}\nETP_SIZE=${ETP_SIZE:-1}\n\n# ============================================================\n# Training\n# ============================================================\nTRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-128}\nMICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-2}\nMAX_LENGTH=${MAX_LENGTH:-2048}\nLR=${LR:-2e-5}\nMIN_LR=${MIN_LR:-2e-6}\nDTYPE=${DTYPE:-bfloat16}\n\nBACKEND=megatron\nRESUME_MODE=${RESUME_MODE:-disable}\n\nproject_name=verl_sft_qwen3_5\nexp_name=qwen3_5-${BACKEND}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-ep${EP_SIZE}\nckpts_home=${ckpts_home:-~/verl/checkpoints/${project_name}/${exp_name}}\nmkdir -p \"${ckpts_home}\"\n\n# ============================================================\n# Engine config\n# ============================================================\n# Key Qwen3.5 settings:\n#   engine.use_remove_padding=False   - GDN requires bshd format (no THD)\n#   engine.vanilla_mbridge=True       - use mbridge (not megatron-bridge)\nENGINE_CONFIG=\"\\\n    engine=${BACKEND} \\\n    optim=${BACKEND} \\\n    optim.lr=${LR} \\\n    optim.min_lr=${MIN_LR} \\\n    optim.lr_warmup_steps=10 \\\n    optim.weight_decay=0.1 \\\n    optim.betas='[0.9,0.95]' \\\n    optim.clip_grad=1.0 \\\n    optim.lr_warmup_init=0 \\\n    optim.lr_decay_style=cosine \\\n    +optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    engine.tensor_model_parallel_size=${TP_SIZE} \\\n    engine.pipeline_model_parallel_size=${PP_SIZE} \\\n    engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \\\n    engine.context_parallel_size=${CP_SIZE} \\\n    engine.expert_model_parallel_size=${EP_SIZE} \\\n    engine.expert_tensor_parallel_size=${ETP_SIZE} \\\n    engine.use_mbridge=True \\\n    engine.vanilla_mbridge=True \\\n    engine.dtype=${DTYPE} \\\n    engine.use_remove_padding=False \\\n    engine.override_transformer_config.attention_backend=auto \\\n    +engine.override_transformer_config.recompute_method=uniform \\\n    +engine.override_transformer_config.recompute_granularity=full \\\n    +engine.override_transformer_config.recompute_num_layers=1\"\n\n# ============================================================\n# Launch\n# ============================================================\ntorchrun \\\n    --nproc_per_node=${NUM_GPUS} \\\n    --nnodes=${NNODES} \\\n    --node_rank=${NODE_RANK} \\\n    --master_addr=${MASTER_ADDR} \\\n    --master_port=${MASTER_PORT} \\\n    -m verl.trainer.sft_trainer \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.train_batch_size=${TRAIN_BATCH_SIZE} \\\n    data.micro_batch_size_per_gpu=${MICRO_BATCH_SIZE} \\\n    data.max_length=${MAX_LENGTH} \\\n    data.pad_mode=no_padding \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=False \\\n    data.max_token_len_per_gpu=${MAX_LENGTH} \\\n    data.messages_key=messages \\\n    model.path=${MODEL_PATH} \\\n    model.use_remove_padding=False \\\n    model.trust_remote_code=True \\\n    ${ENGINE_CONFIG} \\\n    trainer.test_freq=-1 \\\n    trainer.save_freq=500 \\\n    trainer.logger=\"['console']\" \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_epochs=1 \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.resume_mode=${RESUME_MODE}\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen3_8b_sft_peft_sp2_npu.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen3_8b_sft_peft_sp2_npu.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.micro_batch_size_per_gpu=64 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    model.path=Qwen/Qwen3-8B \\\n    model.use_remove_padding=true \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen3-8b-instruct \\\n    trainer.logger=console \\\n    trainer.total_epochs=2 $@ \\\n    model.lora_rank=32 \\\n    model.lora_alpha=16 \\\n    model.target_modules=all-linear\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen_05_automodel.sh",
    "content": "# Requires: Automodel, transformers>=5.3.0, torchao\n# MoE also requires: grouped_gemm (github.com/fanshiqing/grouped_gemm v1.1.4)\n\nset -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen_05_automodel.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k_sft/train.parquet \\\n    data.val_files=$HOME/data/gsm8k_sft/test.parquet \\\n    data.train_batch_size=128 \\\n    data.pad_mode=no_padding \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=2048 \\\n    data.messages_key=messages \\\n    data.ignore_input_ids_mismatch=True \\\n    model=hf_model \\\n    model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    model.use_remove_padding=True \\\n    engine=automodel \\\n    engine.distributed_strategy=fsdp2 \\\n    engine.tp_size=1 \\\n    engine.pp_size=1 \\\n    engine.cp_size=1 \\\n    engine.ep_size=1 \\\n    engine.use_torch_compile=False \\\n    optim=automodel \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas='[0.9,0.95]' \\\n    optim.clip_grad=1.0 \\\n    optim.init_lr_ratio=0 \\\n    optim.min_lr_ratio=0.1 \\\n    optim.lr_scheduler_type=cosine \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-automodel \\\n    trainer.total_epochs=2 \\\n    trainer.test_freq=-1 \\\n    trainer.save_freq=-1 \\\n    trainer.logger=console \\\n    trainer.seed=1111 \\\n    trainer.resume_mode=disable $@\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen_05_peft.sh",
    "content": "# Tested with 2 & 4 GPUs\n\nset -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen_05_peft.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-instruct \\\n    trainer.logger=console \\\n    trainer.total_epochs=1 $@ \\\n    model.lora_rank=32\\\n    model.lora_alpha=16 \\\n    model.target_modules=all-linear\n\n    # Or you can do this:\n    # model.target_modules=[q_proj,v_proj] \\\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen_05_sp2.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen_05_sp2.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    model.use_remove_padding=true \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-instruct-sp2 \\\n    trainer.logger=console \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "examples/sft/gsm8k/run_qwen_05_sp2_liger.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen_05_sp2.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    model.use_liger=True \\\n    model.use_remove_padding=true \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-instruct-sp2-liger \\\n    trainer.logger=console $@\n"
  },
  {
    "path": "examples/sft/gsm8k/run_seed_oss_36b_sft.sh",
    "content": "set -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_seed_oss_36b_sft.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.micro_batch_size_per_gpu=4 \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    model.path=ByteDance-Seed/Seed-OSS-36B-Base \\\n    model.use_remove_padding=true \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-seed-oss-36b \\\n    trainer.logger=console \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "examples/sft/multiturn/run_qwen_05_sp2.sh",
    "content": "#!/bin/bash\nset -x\n\nif [ \"$#\" -lt 2 ]; then\n    echo \"Usage: run_qwen_05_sp2.sh <nproc_per_node> <save_path> [other_configs...]\"\n    exit 1\nfi\n\nnproc_per_node=$1\nsave_path=$2\n\n# Shift the arguments so $@ refers to the rest\nshift 2\n\ntorchrun --nnodes=1 --nproc_per_node=$nproc_per_node \\\n     -m verl.trainer.sft_trainer \\\n    data.train_files=$HOME/data/multiturn/train.parquet \\\n    data.val_files=$HOME/data/multiturn/test.parquet \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=4 \\\n    model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    model.use_remove_padding=true \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    trainer.default_local_dir=$save_path \\\n    trainer.project_name=multiturn-sft \\\n    trainer.experiment_name=multiturn-sft-qwen-2.5-0.5b-instruct-sp2 \\\n    trainer.logger=console \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "examples/sft/vlm/run_qwen3_vl_2b.sh",
    "content": "#!/usr/bin/env bash\n# python examples/data_preprocess/pokemon.py\nset -xeuo pipefail\n\nHDFS_ROOT=${HDFS_ROOT:-$PWD}\nDATA_ROOT=${DATA_ROOT:-$PWD}\n\nENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n\nTRAIN_FILES=${HOME}/data/pokemon-gpt4o-captions/train.parquet\n\nbackend=${BACKEND:-fsdp}\n\nproject_name=verl_sft_test\n\nRESUME_MODE=auto\nMODEL_ID=${HDFS_ROOT}/model/Qwen3-VL-2B-Instruct\n# MODEL_ID=${HDFS_ROOT}/model/Qwen3-VL-30B-A3B-Instruct\n\nSP_SIZE=${SP_SIZE:-2}\nFSDP_SIZE=${FSDP_SIZE:--1}\nFSDP_STRATEGY=${FSDP_STRATEGY:-\"fsdp2\"}\n\nTP_SIZE=${TP_SIZE:-2}\nPP_SIZE=${PP_SIZE:-2}\nVPP_SIZE=${VPP_SIZE:-null}\nCP_SIZE=${CP_SIZE:-1}\n\nPAD_MODE=${PAD_MODE:-no_padding}\n\nUSE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True}\n\nFSDP_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    optim=${backend} \\\n    optim.lr=2e-5 \\\n    optim.lr_warmup_steps_ratio=0.01 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.min_lr_ratio=0.1 \\\n    optim.warmup_style=cosine \\\n    engine.ulysses_sequence_parallel_size=${SP_SIZE} \\\n    engine.strategy=${FSDP_STRATEGY} \\\n    engine.fsdp_size=${FSDP_SIZE}\"\n\n\nMEGATRON_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    optim=${backend} \\\n    optim.lr=2e-5 \\\n    optim.lr_warmup_steps_ratio=0.01 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.lr_warmup_init=0 \\\n    optim.lr_decay_style=cosine \\\n    optim.min_lr=2e-6 \\\n    engine.tensor_model_parallel_size=${TP_SIZE} \\\n    engine.pipeline_model_parallel_size=${PP_SIZE} \\\n    engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \\\n    engine.context_parallel_size=${CP_SIZE} \\\n    engine.use_mbridge=True \\\n    engine.vanilla_mbridge=True\"\n\nif [ \"$backend\" = \"fsdp\" ]; then\n    ENGINE_CONFIG=\"$FSDP_ENGINE_CONFIG\"\n    echo \"Using fsdp engine\"\n    exp_name=pokemon-qwen3-2b-${backend}-${FSDP_STRATEGY}-sp${SP_SIZE}-fsdp-1202a1\nelse\n    ENGINE_CONFIG=\"$MEGATRON_ENGINE_CONFIG\"\n    echo \"Using megatron engine\"\n    exp_name=pokemon-qwen3-2b-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-megatron-1202a1\nfi\n\nCKPT_HOME=${CKPT_HOME:-$HOME/open_verl/sft/${project_name}/${exp_name}}\nmkdir -p \"${CKPT_HOME}\"\n\ntorchrun --standalone --nnodes=1 --nproc-per-node=${NUM_TRAINERS:-8} \\\n    ${ENTRYPOINT} \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.train_batch_size=96 \\\n    data.max_length=2048 \\\n    data.pad_mode=${PAD_MODE} \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=65536 \\\n    model.path=$MODEL_ID \\\n    model.use_remove_padding=${USE_REMOVE_PADDING} \\\n    ${ENGINE_CONFIG} \\\n    trainer.test_freq=-1 \\\n    trainer.save_freq=4000 \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_epochs=10 \\\n    trainer.default_local_dir=\"${CKPT_HOME}\" \\\n    trainer.resume_mode=${RESUME_MODE} \\\n    trainer.max_ckpt_to_keep=5 \\\n    checkpoint.save_contents=[model,optimizer,extra]"
  },
  {
    "path": "examples/sglang_multiturn/README.md",
    "content": "# Multi-Turn Rollout Example (GSM8K)\n\nThis example demonstrates how to perform **multi-turn rollout** using SGLang with a tool-calling capable model (e.g., Qwen2.5-3B) on the GSM8K dataset.\n\n## Usage\n\n### Step 1: Download GSM8K Dataset\n\n```bash\ncd examples/data_preprocess\npython3 gsm8k_multiturn_w_tool.py\n```\n\nThis will download and preprocess the GSM8K dataset into ~/data/gsm8k/.\n\n### Step 2: Run Multi-Turn Rollout\n\nIf you have 8 GPUs\nUse the standard 8-GPU script:\n\n```bash\ncd your_verl_root_dir\nbash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh\n```\n\nIf you have only 4 GPUs\nUse the fallback 4-GPU script:\n\n```bash\ncd your_verl_root_dir\nbash examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn_4xgpu.sh \n```\n\n## Notes\n\n- The rollout supports multi-turn conversations with tool-calling capabilities.\n- Current tools are used for GSM8K answer evaluation.\n- Future versions may extend to search and code interpreter tools.\n"
  },
  {
    "path": "examples/sglang_multiturn/config/geo3k_multiturn_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 2048\n  max_response_length: 2048\n  train_batch_size: 256\n  return_raw_chat: True\n  return_multi_modal_inputs: False\n\nactor_rollout_ref:\n  hybrid_engine: True\n  model:\n    custom_chat_template: \"{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{%- if tools %}{{- '<|im_start|>system\\\\n' }}{%- if messages[0]['role'] == 'system' %}{{- messages[0]['content'] }}{%- else %}{{- 'You are a helpful assistant.' }}{%- endif %}{{- \\\"\\\\n\\\\n# Tools\\\\n\\\\nYou may call one or more functions to assist with the user query.\\\\n\\\\nYou are provided with function signatures within <tools></tools> XML tags:\\\\n<tools>\\\" }}{%- for tool in tools %}{{- \\\"\\\\n\\\" }}{{- tool | tojson }}{%- endfor %}{{- \\\"\\\\n</tools>\\\\n\\\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\\\n<tool_call>\\\\n{\\\\\\\"name\\\\\\\": <function-name>, \\\\\\\"arguments\\\\\\\": <args-json-object>}\\\\n</tool_call><|im_end|>\\\\n\\\" }}{% for message in messages %}{% if message['role'] != 'system' or loop.first == false %}{%- if (message.role == \\\"user\\\") or (message.role == \\\"system\\\" and not loop.first) or (message.role == \\\"assistant\\\" and not message.tool_calls) %}<|im_start|>{{ message['role'] }}\\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\\n{% endif %}{%- elif message.role == \\\"assistant\\\" %}{{- '<|im_start|>' + message.role }}{%- if message.content %}{{- '\\\\n' + message.content }}{%- endif %}{%- for tool_call in message.tool_calls %}{%- if tool_call.function is defined %}{%- set tool_call = tool_call.function %}{%- endif %}{{- '\\\\n<tool_call>\\\\n{\\\"name\\\": \\\"' }}{{- tool_call.name }}{{- '\\\", \\\"arguments\\\": ' }}{{- tool_call.arguments | tojson }}{{- '}\\\\n</tool_call>' }}{%- endfor %}{{- '<|im_end|>\\\\n' }}{%- elif message.role == \\\"tool\\\" %}{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \\\"tool\\\") %}{{- '<|im_start|>user' }}{%- endif %}{{- '\\\\n<tool_response>\\\\n' }}{% if message['content'] is string %}{{ message.content }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif content['type'] == 'text' or 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}{{- '\\\\n</tool_response>' }}{%- if loop.last or (messages[loop.index0 + 1].role != \\\"tool\\\") %}{{- '<|im_end|>\\\\n' }}{%- endif %}{%- endif %}{% endif %}{% endfor %}{%- else %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n{% endif %}{%- if (message.role == \\\"user\\\") or (message.role == \\\"system\\\" and not loop.first) or (message.role == \\\"assistant\\\" and not message.tool_calls) %}<|im_start|>{{ message['role'] }}\\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\\n{% endif %}{%- elif message.role == \\\"assistant\\\" %}{{- '<|im_start|>' + message.role }}{%- if message.content %}{{- '\\\\n' + message.content }}{%- endif %}{%- for tool_call in message.tool_calls %}{%- if tool_call.function is defined %}{%- set tool_call = tool_call.function %}{%- endif %}{{- '\\\\n<tool_call>\\\\n{\\\"name\\\": \\\"' }}{{- tool_call.name }}{{- '\\\", \\\"arguments\\\": ' }}{{- tool_call.arguments | tojson }}{{- '}\\\\n</tool_call>' }}{%- endfor %}{{- '<|im_end|>\\\\n' }}{%- elif message.role == \\\"tool\\\" %}{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \\\"tool\\\") %}{{- '<|im_start|>user' }}{%- endif %}{{- '\\\\n<tool_response>\\\\n' }}{% if message['content'] is string %}{{ message.content }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif content['type'] == 'text' or 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}{{- '\\\\n</tool_response>' }}{%- if loop.last or (messages[loop.index0 + 1].role != \\\"tool\\\") %}{{- '<|im_end|>\\\\n' }}{%- endif %}{%- endif %}{% endfor %}{%- endif %}{% if add_generation_prompt %}<|im_start|>assistant\\n{% endif %}\"\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n      # tool_config_path: \"./config/tool_config/gsm8k_tool_config.yaml\"\n"
  },
  {
    "path": "examples/sglang_multiturn/config/geo3k_multiturn_megatron_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_megatron_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 2048\n  max_response_length: 2048\n  train_batch_size: 256\n  return_raw_chat: True\n  return_multi_modal_inputs: False\n\nactor_rollout_ref:\n  hybrid_engine: True\n  model:\n    custom_chat_template: \"{% set image_count = namespace(value=0) %}{% set video_count = namespace(value=0) %}{%- if tools %}{{- '<|im_start|>system\\\\n' }}{%- if messages[0]['role'] == 'system' %}{{- messages[0]['content'] }}{%- else %}{{- 'You are a helpful assistant.' }}{%- endif %}{{- \\\"\\\\n\\\\n# Tools\\\\n\\\\nYou may call one or more functions to assist with the user query.\\\\n\\\\nYou are provided with function signatures within <tools></tools> XML tags:\\\\n<tools>\\\" }}{%- for tool in tools %}{{- \\\"\\\\n\\\" }}{{- tool | tojson }}{%- endfor %}{{- \\\"\\\\n</tools>\\\\n\\\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\\\n<tool_call>\\\\n{\\\\\\\"name\\\\\\\": <function-name>, \\\\\\\"arguments\\\\\\\": <args-json-object>}\\\\n</tool_call><|im_end|>\\\\n\\\" }}{% for message in messages %}{% if message['role'] != 'system' or loop.first == false %}{%- if (message.role == \\\"user\\\") or (message.role == \\\"system\\\" and not loop.first) or (message.role == \\\"assistant\\\" and not message.tool_calls) %}<|im_start|>{{ message['role'] }}\\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\\n{% endif %}{%- elif message.role == \\\"assistant\\\" %}{{- '<|im_start|>' + message.role }}{%- if message.content %}{{- '\\\\n' + message.content }}{%- endif %}{%- for tool_call in message.tool_calls %}{%- if tool_call.function is defined %}{%- set tool_call = tool_call.function %}{%- endif %}{{- '\\\\n<tool_call>\\\\n{\\\"name\\\": \\\"' }}{{- tool_call.name }}{{- '\\\", \\\"arguments\\\": ' }}{{- tool_call.arguments | tojson }}{{- '}\\\\n</tool_call>' }}{%- endfor %}{{- '<|im_end|>\\\\n' }}{%- elif message.role == \\\"tool\\\" %}{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \\\"tool\\\") %}{{- '<|im_start|>user' }}{%- endif %}{{- '\\\\n<tool_response>\\\\n' }}{% if message['content'] is string %}{{ message.content }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif content['type'] == 'text' or 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}{{- '\\\\n</tool_response>' }}{%- if loop.last or (messages[loop.index0 + 1].role != \\\"tool\\\") %}{{- '<|im_end|>\\\\n' }}{%- endif %}{%- endif %}{% endif %}{% endfor %}{%- else %}{% for message in messages %}{% if loop.first and message['role'] != 'system' %}<|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n{% endif %}{%- if (message.role == \\\"user\\\") or (message.role == \\\"system\\\" and not loop.first) or (message.role == \\\"assistant\\\" and not message.tool_calls) %}<|im_start|>{{ message['role'] }}\\n{% if message['content'] is string %}{{ message['content'] }}<|im_end|>\\n{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}<|im_end|>\\n{% endif %}{%- elif message.role == \\\"assistant\\\" %}{{- '<|im_start|>' + message.role }}{%- if message.content %}{{- '\\\\n' + message.content }}{%- endif %}{%- for tool_call in message.tool_calls %}{%- if tool_call.function is defined %}{%- set tool_call = tool_call.function %}{%- endif %}{{- '\\\\n<tool_call>\\\\n{\\\"name\\\": \\\"' }}{{- tool_call.name }}{{- '\\\", \\\"arguments\\\": ' }}{{- tool_call.arguments | tojson }}{{- '}\\\\n</tool_call>' }}{%- endfor %}{{- '<|im_end|>\\\\n' }}{%- elif message.role == \\\"tool\\\" %}{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \\\"tool\\\") %}{{- '<|im_start|>user' }}{%- endif %}{{- '\\\\n<tool_response>\\\\n' }}{% if message['content'] is string %}{{ message.content }}{% else %}{% for content in message['content'] %}{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}{% set image_count.value = image_count.value + 1 %}{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>{% elif content['type'] == 'video' or 'video' in content %}{% set video_count.value = video_count.value + 1 %}{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>{% elif content['type'] == 'text' or 'text' in content %}{{ content['text'] }}{% endif %}{% endfor %}{% endif %}{{- '\\\\n</tool_response>' }}{%- if loop.last or (messages[loop.index0 + 1].role != \\\"tool\\\") %}{{- '<|im_end|>\\\\n' }}{%- endif %}{%- endif %}{% endfor %}{%- endif %}{% if add_generation_prompt %}<|im_start|>assistant\\n{% endif %}\"\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n      # tool_config_path: \"./config/tool_config/gsm8k_tool_config.yaml\"\n"
  },
  {
    "path": "examples/sglang_multiturn/config/gsm8k_multiturn_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n"
  },
  {
    "path": "examples/sglang_multiturn/config/gsm8k_multiturn_grpo_server.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n    sglang_rollout_mode: server\n    server:\n      timeout: 60\n      max_attempts: 3\n      retry_delay: 2\n      max_connections: 1000\n      max_start_wait_time: 300.0"
  },
  {
    "path": "examples/sglang_multiturn/config/gsm8k_multiturn_grpo_w_interaction.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_user_turns: 5\n"
  },
  {
    "path": "examples/sglang_multiturn/config/gsm8k_multiturn_megatron_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_megatron_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n  \n"
  },
  {
    "path": "examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml",
    "content": "interaction:\n  - name: \"gsm8k\"\n    class_name: \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\"\n    config: {}"
  },
  {
    "path": "examples/sglang_multiturn/config/retool_multiturn_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 5\n      tool_config_path: \"./config/tool_config/sandbox_fusion_tool_config.yaml\"\n"
  },
  {
    "path": "examples/sglang_multiturn/config/search_multiturn_grpo.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n  shuffle: False\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 2\n      format: qwen\n"
  },
  {
    "path": "examples/sglang_multiturn/config/search_multiturn_grpo_one_step_off.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  max_prompt_length: 1024\n  max_response_length: 1024\n  train_batch_size: 256\n  return_raw_chat: True\n  shuffle: False\n\nactor_rollout_ref:\n  hybrid_engine: True\n  rollout:\n    name: sglang\n    multi_turn:\n      enable: True\n      max_assistant_turns: 2\n      format: qwen\n"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/geo3k_tool_config.yaml",
    "content": "tools:\n  - class_name: \"verl.tools.geo3k_tool.Geo3kTool\"\n    config: \n      type: native\n    tool_schema:\n      type: \"function\"\n      function:\n        name: \"calc_geo3k_reward\"\n        description: \"A tool for calculating the reward of geo3k. (1.0 if parsed answer is correct, 0.0 if parsed answer is incorrect or not correctly parsed)\"\n        parameters:\n          type: \"object\"\n          properties:\n            answer:\n              type: \"string\"\n              description: \"The model's answer to the geo3k problem, must be a digits\"\n          required: [\"answer\"]"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml",
    "content": "tools:\n  - class_name: \"verl.tools.gsm8k_tool.Gsm8kTool\"\n    config: \n      type: native\n    tool_schema:\n      type: \"function\"\n      function:\n        name: \"calc_gsm8k_reward\"\n        description: \"A tool for calculating the reward of gsm8k. (1.0 if parsed answer is correct, 0.0 if parsed answer is incorrect or not correctly parsed)\"\n        parameters:\n          type: \"object\"\n          properties:\n            answer:\n              type: \"string\"\n              description: \"The model's answer to the GSM8K math problem, must be a digits\"\n          required: [\"answer\"]\n"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/mcp_server.json",
    "content": "{\n    \"mcpServers\": {\n        \"Tavily Expert\": {\n            \"url\": \"your_tavily_expert_url\",\n            \"auth_token\": \"your_tavily_api_token\"\n        }\n    }\n}"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/mcp_tool_config.yaml",
    "content": "tools:\n  - class_name: verl.tools.mcp_search_tool.MCPSearchTool\n    config:\n      rate_limit: 120\n      timeout: 120\n      type: mcp\n    mcp:\n      mcp_servers_config_path: ./mcp_server.json\n      # optional\n      tool_selected_list: \n        - tavily_search_tool"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/sandbox_fusion_tool_config.yaml",
    "content": "tools:\n  - class_name: \"verl.tools.sandbox_fusion_tools.SandboxFusionTool\"\n    config: \n      sandbox_fusion_url: \"https://xxx.apigateway-cn-beijing.volceapi.com/run_code\"\n      num_workers: 10\n      enable_global_rate_limit: true\n      rate_limit: 10\n      default_timeout: 30\n      default_language: \"python\"\n      memory_limit_mb: 1024\n      type: native\n\n    tool_schema:\n      type: \"function\"\n      function:\n        name: \"code_interpreter\"\n        description: \"A tool for executing code.\"\n        parameters:\n          type: \"object\"\n          properties:\n            code:\n              type: \"string\"\n              description: \"The code to execute.\"\n          required: [\"code\"]"
  },
  {
    "path": "examples/sglang_multiturn/config/tool_config/search_tool_config.yaml",
    "content": "tools:\n  - class_name: verl.tools.search_tool.SearchTool\n    config:\n      retrieval_service_url: http://127.0.0.1:8000/retrieve\n      num_workers: 120\n      rate_limit: 120\n      timeout: 30\n      type: native\n    tool_schema:\n      type: function\n      function:\n        name: search\n        description: Searches the web for relevant information based on the given query.\n        parameters:\n          type: object\n          properties:\n            query_list:\n              type: array\n              item:\n                type: string\n              description: A list of fully-formed semantic queries. The tool will return search results for each query.\n          required: \n            - query_list"
  },
  {
    "path": "examples/sglang_multiturn/geo3k/run_qwen2.5-3b_geo3k_multiturn.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='geo3k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='geo3k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-geo3k-sgl-multi-w-tool-verify-n16' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/geo3k_multiturn_w_tool/train.parquet \\\n    data.val_files=$HOME/data/geo3k_multiturn_w_tool/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/geo3k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/geo3k/run_qwen2.5-3b_geo3k_multiturn_4xgpu.sh",
    "content": "# run on 4xH100\n# make sure your current working directory is the root of the project\n\nset -x\nexport HYDRA_FULL_ERROR=1\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='geo3k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='geo3k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-geo3k-async-sgl-multi-w-tool-verify-n16-4cards' \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.total_epochs=15 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \\\n    critic.ppo_max_token_len_per_gpu=8192 \\\n    critic.forward_max_token_len_per_gpu=8192 \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/geo3k_tool_config.yaml\" \\\n    $@"
  },
  {
    "path": "examples/sglang_multiturn/geo3k/run_qwen2.5-3b_megatron_geo3k_multiturn.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n# this is a verification training script, the parallel setting should be tuned to your model\n\nset -x\n\nexport PYTHONUNBUFFERED=1\nexport RAY_DEDUP_LOGS=0\nexport RUST_BACKTRACE=1\nexport HYDRA_FULL_ERROR=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='geo3k_multiturn_megatron_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-VL-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.megatron.seed=42 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='geo3k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-geo3k-sgl-multi-w-tool-n8-mcore-v2505201745_seed42' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/geo3k_multiturn_w_tool/train.parquet \\\n    data.val_files=$HOME/data/geo3k_multiturn_w_tool/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/geo3k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/gsm8k_toolcall_shaping/gsm8k_toolcall_shaping.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom __future__ import annotations\n\nfrom typing import Any, Optional\n\nfrom verl.utils.reward_score.gsm8k import compute_score as gsm8k_compute_score\n\n\ndef toolcall_shaping_reward(\n    data_source: Optional[str],\n    solution_str: str,\n    ground_truth: str,\n    extra_info: Optional[dict[str, Any]] = None,\n    *,\n    method: str = \"strict\",\n    format_score: float = 0.1,\n    score: float = 1.0,\n    shaping_reward: float = 0.1,\n    trigger_substring: str = \"<tool_call>\",\n    **kwargs,\n) -> float:\n    \"\"\"\n    GSM8K reward + tool-call shaping reward (trajectory-level).\n    \"\"\"\n    base = gsm8k_compute_score(solution_str, ground_truth, method, format_score, score)\n\n    bonus = shaping_reward if (trigger_substring and trigger_substring in solution_str) else 0.0\n    return float(base + bonus)\n\n\n# Optional: keep a default name for convenience in verl config (default is compute_score) [web:59][web:65]\ndef compute_score(\n    data_source: Optional[str],\n    solution_str: str,\n    ground_truth: str,\n    extra_info: Optional[dict[str, Any]] = None,\n    **kwargs,\n) -> float:\n    return toolcall_shaping_reward(\n        data_source=data_source,\n        solution_str=solution_str,\n        ground_truth=ground_truth,\n        extra_info=extra_info,\n        **kwargs,\n    )\n"
  },
  {
    "path": "examples/sglang_multiturn/gsm8k_toolcall_shaping/run_gsm8k_grpo_toolcall_shaping.sh",
    "content": "# make sure your current working directory is the root of the project\n\n\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.sampler.class_name=\"RandomCurriculumSampler\" \\\n    data.sampler.class_path=\"pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu\" \\\n    data.dataloader_num_workers=0 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.train_batch_size=256 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen0.5b_gsm8k_toolcall_shaping' \\\n    reward.custom_reward_function.path=\"$PROJECT_DIR/examples/sglang_multiturn/gsm8k_toolcall_shaping/gsm8k_toolcall_shaping.py\" \\\n    reward.custom_reward_function.name=compute_score \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen0.5b_gsm8k_multiturn_curriculum.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.sampler.class_name=\"RandomCurriculumSampler\" \\\n    data.sampler.class_path=\"pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu\" \\\n    data.dataloader_num_workers=0 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.train_batch_size=256 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen3-4b_function_rm-gsm8k-sgl-multi-w-tool-verify-n16' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-0.5b_gsm8k_multiturn_w_interaction.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\nTRAIN_BATCH_SIZE=${TRAIN_BATCH_SIZE:-512}\nMICRO_BATCH_SIZE=${MICRO_BATCH_SIZE:-8}\nOFFLOAD=${OFFLOAD:-False}\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo_w_interaction' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=$TRAIN_BATCH_SIZE \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=$((1024 * 3)) \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.enable_activation_offload=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$TRAIN_BATCH_SIZE \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$MICRO_BATCH_SIZE \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=$OFFLOAD \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=$OFFLOAD \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=$MICRO_BATCH_SIZE \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=$MICRO_BATCH_SIZE \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=$OFFLOAD \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen2.5-0.5b_function_rm-gsm8k-sgl-multi-w-interaction-n8' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/gsm8k_verl_sgl_multi_turn_w_interaction/train.parquet \\\n    data.val_files=$HOME/data/gsm8k_verl_sgl_multi_turn_w_interaction/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.interaction_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\nfunction now() {\n    date '+%d-%H-%M'\n}\n\nEXPERIMENT_NAME=\"qwen2.5-3b_baseline_$(now)\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    global_profiler.tool=torch_memory \\\n    global_profiler.save_path=./mem_snapshots \\\n    global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries=100000 \\\n    global_profiler.global_tool_config.torch_memory.stack_depth=32 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \\\n    actor_rollout_ref.rollout.multi_stage_wake_up=True \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.over_sample_rate=0.1 \\\n    actor_rollout_ref.rollout.mode=async \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='multi-turn-grpo-qwen2.5-3b-sglang' \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.val_before_train=True \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn_4xgpu.sh",
    "content": "# run on 4xH100\n# make sure your current working directory is the root of the project\n\nset -x\nexport HYDRA_FULL_ERROR=1\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-gsm8k-async-sgl-multi-w-tool-verify-n16-4cards' \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.total_epochs=15 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \\\n    critic.ppo_max_token_len_per_gpu=8192 \\\n    critic.forward_max_token_len_per_gpu=8192 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    actor_rollout_ref.rollout.multi_turn.interaction_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml\" \\\n    actor_rollout_ref.rollout.multi_turn.max_user_turns=1 \\\n    $@"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn_4xgpu_server.sh",
    "content": "# run on 4xH100\n# make sure your current working directory is the root of the project\n\nset -x\nexport HYDRA_FULL_ERROR=1\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo_server' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\", \"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl_server' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-gsm8k-async-sgl-multi-w-tool-verify-n16-4cards' \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.total_epochs=15 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \\\n    critic.ppo_max_token_len_per_gpu=8192 \\\n    critic.forward_max_token_len_per_gpu=8192 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    actor_rollout_ref.rollout.multi_turn.interaction_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml\" \\\n    actor_rollout_ref.rollout.multi_turn.max_user_turns=1 \\\n    $@"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn_server.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\nfunction now() {\n    date '+%d-%H-%M'\n}\n\nEXPERIMENT_NAME=\"qwen2.5-3b_baseline_$(now)\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo_server' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \\\n    actor_rollout_ref.rollout.multi_stage_wake_up=True \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.over_sample_rate=0 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='multi-turn-grpo-qwen2.5-3b-sglang' \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.val_before_train=True \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_multiturn_vllm_fsdp.sh",
    "content": "# run on Ascend 910\n# make sure your current working directory is the root of the project\n\nset -x\nulimit -n 65535\n\n#set vllm v1 env\nexport VLLM_USE_V1=1\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\nTRAIN_BATCH_SIZE=32\nMICRO_BATCH_SIZE=8\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    actor_rollout_ref.rollout.name=vllm \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=${TRAIN_BATCH_SIZE} \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=\"Qwen/Qwen2.5-3B-Instruct\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${TRAIN_BATCH_SIZE} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${MICRO_BATCH_SIZE} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${MICRO_BATCH_SIZE} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9\\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${MICRO_BATCH_SIZE} \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-gsm8k-sgl-multi-w-tool-verify-n16' \\\n    trainer.n_gpus_per_node=16 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.logger='[\"console\"]' \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    trainer.total_epochs=15 \\\n    actor_rollout_ref.rollout.trace.token2text=False \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.multi_turn.enable=true \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    actor_rollout_ref.rollout.free_cache_engine=True"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_gsm8k_tool_agent_mlflow.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.rollout.trace.backend=mlflow \\\n    actor_rollout_ref.rollout.trace.token2text=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"mlflow\"]' \\\n    trainer.project_name='gsm8k_tool-agent' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-gsm8k-sgl-tool-agent-verify-n16' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.total_training_steps=2 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen2.5-3b_megatron_gsm8k_multiturn.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n# this is a verification training script, the parallel setting should be tuned to your model\n\nset -x\n\nexport PYTHONUNBUFFERED=1\nexport RAY_DEDUP_LOGS=0\nexport RUST_BACKTRACE=1\nexport HYDRA_FULL_ERROR=1\nexport CUDA_DEVICE_MAX_CONNECTIONS=1\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_megatron_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=/user/longxiang1/models/Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.megatron.seed=42 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b_function_rm-gsm8k-sgl-multi-w-tool-n8-mcore-v2505201745_seed42' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=/user/longxiang1/data/gsm8k_verl_sgl_multi_turn_preprocessed_v2/train.parquet \\\n    data.val_files=/user/longxiang1/data/gsm8k_verl_sgl_multi_turn_preprocessed_v2/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen3-4b_gsm8k_multiturn.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-4B \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=16 \\\n    actor_rollout_ref.rollout.over_sample_rate=0.1 \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name='qwen3-4b_function_rm-gsm8k-sgl-multi-w-tool-verify-n16' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.total_epochs=15 $@\n\n"
  },
  {
    "path": "examples/sglang_multiturn/run_qwen3_4b_dapo_multiturn.sh",
    "content": "set -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npip install --upgrade \"huggingface-hub>=0.34.0\"\nhf download \\\n    BytedTsinghua-SIA/DAPO-Math-17k \\\n    --repo-type dataset \\\n    --local-dir $HOME/data/BytedTsinghua-SIA/DAPO-Math-17k\n\n\nhf download \\\n    Maxwell-Jia/AIME_2024 \\\n    --repo-type dataset \\\n    --local-dir $HOME/data/Maxwell-Jia/AIME_2024\n\n\n# Note:\n# 1. \n# a sandbox fusion server is needed to run the code interpreter tool.\n# docker run -it -p 8080:8080 volcengine/sandbox-fusion:server-20250609\n\n# 2. \n# The model located at font-info/qwen3-4b-sft-SGLang-RL (https://huggingface.co/font-info/qwen3-4b-sft-SGLang-RL)\n# is a fine-tuned version provided by the SGLang RL team. Without supervised fine-tuning (SFT)\n# on the Retool dataset, Dapo training will not converge.\n\n# If you still wish to perform SFT from scratch, follow the steps below:\n\n# Step 1: Download the SFT dataset\n#hf download JoeYing/ReTool-SFT --repo-type dataset --local-dir ./ReTool-SFT\n\n# Step 2: Preprocess the data for SFT\n#python3 recipe/retool/retool_sft_preprocess.py\n\n# Step 3: Run SFT training\n#bash recipe/retool/run_qwen2-32b_sft.sh\n\n# having trouble setup? see https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/blob/main/rlhf/verl/multi-turn/release_log/latest_sglang.md for more details.\n\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.use_kl_in_reward=False \\\n    algorithm.kl_ctrl.kl_coef=0.0 \\\n    data.train_files=$HOME/data/BytedTsinghua-SIA/DAPO-Math-17k \\\n    data.val_files=$HOME/data/Maxwell-Jia/AIME_2024 \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=16384 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.custom_cls.path=$PROJECT_DIR/recipe/retool/retool.py \\\n    data.custom_cls.name=CustomRLHFDataset \\\n    reward.custom_reward_function.path=$PROJECT_DIR/recipe/retool/retool.py \\\n    reward.custom_reward_function.name=compute_score \\\n    actor_rollout_ref.model.path=font-info/qwen3-4b-sft-SGLang-RL \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.0 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=32768 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.85 \\\n    actor_rollout_ref.rollout.multi_stage_wake_up=True \\\n    actor_rollout_ref.rollout.multi_turn.enable=True \\\n    actor_rollout_ref.rollout.multi_turn.max_user_turns=16 \\\n    actor_rollout_ref.rollout.multi_turn.max_assistant_turns=16 \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=$PROJECT_DIR/recipe/retool/sandbox_fusion_tool_config.yaml \\\n    actor_rollout_ref.rollout.multi_turn.format=hermes \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=30 \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=sglang-dapo-multiturn \\\n    trainer.experiment_name=qwen3_4b_sft_dapo_multiturn \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.log_val_generations=20 \\\n    trainer.val_before_train=True \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=20 \\\n    trainer.total_epochs=15 \\\n    $@\n"
  },
  {
    "path": "examples/sglang_multiturn/search_r1_like/local_dense_retriever/download.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 Search-R1 Contributors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/scripts/download.py\n\n\nimport argparse\n\nfrom huggingface_hub import hf_hub_download\n\nparser = argparse.ArgumentParser(description=\"Download files from a Hugging Face dataset repository.\")\nparser.add_argument(\"--repo_id\", type=str, default=\"PeterJinGo/wiki-18-e5-index\", help=\"Hugging Face repository ID\")\nparser.add_argument(\"--save_path\", type=str, required=True, help=\"Local directory to save files\")\n\nargs = parser.parse_args()\n\nrepo_id = \"PeterJinGo/wiki-18-e5-index\"\nfor file in [\"part_aa\", \"part_ab\"]:\n    hf_hub_download(\n        repo_id=repo_id,\n        filename=file,  # e.g., \"e5_Flat.index\"\n        repo_type=\"dataset\",\n        local_dir=args.save_path,\n    )\n\nrepo_id = \"PeterJinGo/wiki-18-corpus\"\nhf_hub_download(\n    repo_id=repo_id,\n    filename=\"wiki-18.jsonl.gz\",\n    repo_type=\"dataset\",\n    local_dir=args.save_path,\n)\n"
  },
  {
    "path": "examples/sglang_multiturn/search_r1_like/local_dense_retriever/retrieval_server.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 Search-R1 Contributors\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/search_r1/search/retrieval_server.py\n\nimport argparse\nimport json\nimport warnings\nfrom typing import Optional\n\nimport datasets\nimport faiss\nimport numpy as np\nimport torch\nimport uvicorn\nfrom fastapi import FastAPI\nfrom pydantic import BaseModel\nfrom tqdm import tqdm\nfrom transformers import AutoModel, AutoTokenizer\n\n\ndef load_corpus(corpus_path: str):\n    corpus = datasets.load_dataset(\"json\", data_files=corpus_path, split=\"train\", num_proc=4)\n    return corpus\n\n\ndef load_docs(corpus, doc_idxs):\n    results = [corpus[int(idx)] for idx in doc_idxs]\n    return results\n\n\ndef load_model(model_path: str, use_fp16: bool = False):\n    model = AutoModel.from_pretrained(model_path, trust_remote_code=True)\n    model.eval()\n    model.cuda()\n    if use_fp16:\n        model = model.half()\n    tokenizer = AutoTokenizer.from_pretrained(model_path, use_fast=True, trust_remote_code=True)\n    return model, tokenizer\n\n\ndef pooling(pooler_output, last_hidden_state, attention_mask=None, pooling_method=\"mean\"):\n    if pooling_method == \"mean\":\n        last_hidden = last_hidden_state.masked_fill(~attention_mask[..., None].bool(), 0.0)\n        return last_hidden.sum(dim=1) / attention_mask.sum(dim=1)[..., None]\n    elif pooling_method == \"cls\":\n        return last_hidden_state[:, 0]\n    elif pooling_method == \"pooler\":\n        return pooler_output\n    else:\n        raise NotImplementedError(\"Pooling method not implemented!\")\n\n\nclass Encoder:\n    def __init__(self, model_name, model_path, pooling_method, max_length, use_fp16):\n        self.model_name = model_name\n        self.model_path = model_path\n        self.pooling_method = pooling_method\n        self.max_length = max_length\n        self.use_fp16 = use_fp16\n\n        self.model, self.tokenizer = load_model(model_path=model_path, use_fp16=use_fp16)\n        self.model.eval()\n\n    @torch.no_grad()\n    def encode(self, query_list: list[str], is_query=True) -> np.ndarray:\n        # processing query for different encoders\n        if isinstance(query_list, str):\n            query_list = [query_list]\n\n        if \"e5\" in self.model_name.lower():\n            if is_query:\n                query_list = [f\"query: {query}\" for query in query_list]\n            else:\n                query_list = [f\"passage: {query}\" for query in query_list]\n\n        if \"bge\" in self.model_name.lower():\n            if is_query:\n                query_list = [\n                    f\"Represent this sentence for searching relevant passages: {query}\" for query in query_list\n                ]\n\n        inputs = self.tokenizer(\n            query_list, max_length=self.max_length, padding=True, truncation=True, return_tensors=\"pt\"\n        )\n        inputs = {k: v.cuda() for k, v in inputs.items()}\n\n        if \"T5\" in type(self.model).__name__:\n            # T5-based retrieval model\n            decoder_input_ids = torch.zeros((inputs[\"input_ids\"].shape[0], 1), dtype=torch.long).to(\n                inputs[\"input_ids\"].device\n            )\n            output = self.model(**inputs, decoder_input_ids=decoder_input_ids, return_dict=True)\n            query_emb = output.last_hidden_state[:, 0, :]\n        else:\n            output = self.model(**inputs, return_dict=True)\n            query_emb = pooling(\n                output.pooler_output, output.last_hidden_state, inputs[\"attention_mask\"], self.pooling_method\n            )\n            if \"dpr\" not in self.model_name.lower():\n                query_emb = torch.nn.functional.normalize(query_emb, dim=-1)\n\n        query_emb = query_emb.detach().cpu().numpy()\n        query_emb = query_emb.astype(np.float32, order=\"C\")\n\n        del inputs, output\n        torch.cuda.empty_cache()\n\n        return query_emb\n\n\nclass BaseRetriever:\n    def __init__(self, config):\n        self.config = config\n        self.retrieval_method = config.retrieval_method\n        self.topk = config.retrieval_topk\n\n        self.index_path = config.index_path\n        self.corpus_path = config.corpus_path\n\n    def _search(self, query: str, num: int, return_score: bool):\n        raise NotImplementedError\n\n    def _batch_search(self, query_list: list[str], num: int, return_score: bool):\n        raise NotImplementedError\n\n    def search(self, query: str, num: int = None, return_score: bool = False):\n        return self._search(query, num, return_score)\n\n    def batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):\n        return self._batch_search(query_list, num, return_score)\n\n\nclass BM25Retriever(BaseRetriever):\n    def __init__(self, config):\n        super().__init__(config)\n        from pyserini.search.lucene import LuceneSearcher\n\n        self.searcher = LuceneSearcher(self.index_path)\n        self.contain_doc = self._check_contain_doc()\n        if not self.contain_doc:\n            self.corpus = load_corpus(self.corpus_path)\n        self.max_process_num = 8\n\n    def _check_contain_doc(self):\n        return self.searcher.doc(0).raw() is not None\n\n    def _search(self, query: str, num: int = None, return_score: bool = False):\n        if num is None:\n            num = self.topk\n        hits = self.searcher.search(query, num)\n        if len(hits) < 1:\n            if return_score:\n                return [], []\n            else:\n                return []\n        scores = [hit.score for hit in hits]\n        if len(hits) < num:\n            warnings.warn(\"Not enough documents retrieved!\", stacklevel=2)\n        else:\n            hits = hits[:num]\n\n        if self.contain_doc:\n            all_contents = [json.loads(self.searcher.doc(hit.docid).raw())[\"contents\"] for hit in hits]\n            results = [\n                {\n                    \"title\": content.split(\"\\n\")[0].strip('\"'),\n                    \"text\": \"\\n\".join(content.split(\"\\n\")[1:]),\n                    \"contents\": content,\n                }\n                for content in all_contents\n            ]\n        else:\n            results = load_docs(self.corpus, [hit.docid for hit in hits])\n\n        if return_score:\n            return results, scores\n        else:\n            return results\n\n    def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):\n        results = []\n        scores = []\n        for query in query_list:\n            item_result, item_score = self._search(query, num, True)\n            results.append(item_result)\n            scores.append(item_score)\n        if return_score:\n            return results, scores\n        else:\n            return results\n\n\nclass DenseRetriever(BaseRetriever):\n    def __init__(self, config):\n        super().__init__(config)\n        self.index = faiss.read_index(self.index_path)\n        if config.faiss_gpu:\n            co = faiss.GpuMultipleClonerOptions()\n            co.useFloat16 = True\n            co.shard = True\n            self.index = faiss.index_cpu_to_all_gpus(self.index, co=co)\n\n        self.corpus = load_corpus(self.corpus_path)\n        self.encoder = Encoder(\n            model_name=self.retrieval_method,\n            model_path=config.retrieval_model_path,\n            pooling_method=config.retrieval_pooling_method,\n            max_length=config.retrieval_query_max_length,\n            use_fp16=config.retrieval_use_fp16,\n        )\n        self.topk = config.retrieval_topk\n        self.batch_size = config.retrieval_batch_size\n\n    def _search(self, query: str, num: int = None, return_score: bool = False):\n        if num is None:\n            num = self.topk\n        query_emb = self.encoder.encode(query)\n        scores, idxs = self.index.search(query_emb, k=num)\n        idxs = idxs[0]\n        scores = scores[0]\n        results = load_docs(self.corpus, idxs)\n        if return_score:\n            return results, scores.tolist()\n        else:\n            return results\n\n    def _batch_search(self, query_list: list[str], num: int = None, return_score: bool = False):\n        if isinstance(query_list, str):\n            query_list = [query_list]\n        if num is None:\n            num = self.topk\n\n        results = []\n        scores = []\n        for start_idx in tqdm(range(0, len(query_list), self.batch_size), desc=\"Retrieval process: \"):\n            query_batch = query_list[start_idx : start_idx + self.batch_size]\n            batch_emb = self.encoder.encode(query_batch)\n            batch_scores, batch_idxs = self.index.search(batch_emb, k=num)\n            batch_scores = batch_scores.tolist()\n            batch_idxs = batch_idxs.tolist()\n\n            # load_docs is not vectorized, but is a python list approach\n            flat_idxs = sum(batch_idxs, [])\n            batch_results = load_docs(self.corpus, flat_idxs)\n            # chunk them back\n            batch_results = [batch_results[i * num : (i + 1) * num] for i in range(len(batch_idxs))]\n\n            results.extend(batch_results)\n            scores.extend(batch_scores)\n\n            del batch_emb, batch_scores, batch_idxs, query_batch, flat_idxs, batch_results\n            torch.cuda.empty_cache()\n\n        if return_score:\n            return results, scores\n        else:\n            return results\n\n\ndef get_retriever(config):\n    if config.retrieval_method == \"bm25\":\n        return BM25Retriever(config)\n    else:\n        return DenseRetriever(config)\n\n\n#####################################\n# FastAPI server below\n#####################################\n\n\nclass Config:\n    \"\"\"\n    Minimal config class (simulating your argparse)\n    Replace this with your real arguments or load them dynamically.\n    \"\"\"\n\n    def __init__(\n        self,\n        retrieval_method: str = \"bm25\",\n        retrieval_topk: int = 10,\n        index_path: str = \"./index/bm25\",\n        corpus_path: str = \"./data/corpus.jsonl\",\n        dataset_path: str = \"./data\",\n        data_split: str = \"train\",\n        faiss_gpu: bool = True,\n        retrieval_model_path: str = \"./model\",\n        retrieval_pooling_method: str = \"mean\",\n        retrieval_query_max_length: int = 256,\n        retrieval_use_fp16: bool = False,\n        retrieval_batch_size: int = 128,\n    ):\n        self.retrieval_method = retrieval_method\n        self.retrieval_topk = retrieval_topk\n        self.index_path = index_path\n        self.corpus_path = corpus_path\n        self.dataset_path = dataset_path\n        self.data_split = data_split\n        self.faiss_gpu = faiss_gpu\n        self.retrieval_model_path = retrieval_model_path\n        self.retrieval_pooling_method = retrieval_pooling_method\n        self.retrieval_query_max_length = retrieval_query_max_length\n        self.retrieval_use_fp16 = retrieval_use_fp16\n        self.retrieval_batch_size = retrieval_batch_size\n\n\nclass QueryRequest(BaseModel):\n    queries: list[str]\n    topk: Optional[int] = None\n    return_scores: bool = False\n\n\napp = FastAPI()\n\n\n@app.post(\"/retrieve\")\ndef retrieve_endpoint(request: QueryRequest):\n    \"\"\"\n    Endpoint that accepts queries and performs retrieval.\n\n    Input format:\n    {\n      \"queries\": [\"What is Python?\", \"Tell me about neural networks.\"],\n      \"topk\": 3,\n      \"return_scores\": true\n    }\n\n    Output format (when return_scores=True，similarity scores are returned):\n    {\n        \"result\": [\n            [   # Results for each query\n                {\n                    {\"document\": doc, \"score\": score}\n                },\n                # ... more documents\n            ],\n            # ... results for other queries\n        ]\n    }\n    \"\"\"\n    if not request.topk:\n        request.topk = config.retrieval_topk  # fallback to default\n\n    # Perform batch retrieval\n    results, scores = retriever.batch_search(\n        query_list=request.queries, num=request.topk, return_score=request.return_scores\n    )\n\n    # Format response\n    resp = []\n    for i, single_result in enumerate(results):\n        if request.return_scores:\n            # If scores are returned, combine them with results\n            combined = []\n            for doc, score in zip(single_result, scores[i], strict=True):\n                combined.append({\"document\": doc, \"score\": score})\n            resp.append(combined)\n        else:\n            resp.append(single_result)\n    return {\"result\": resp}\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser(description=\"Launch the local faiss retriever.\")\n    parser.add_argument(\n        \"--index_path\", type=str, default=\"/home/peterjin/mnt/index/wiki-18/e5_Flat.index\", help=\"Corpus indexing file.\"\n    )\n    parser.add_argument(\n        \"--corpus_path\",\n        type=str,\n        default=\"/home/peterjin/mnt/data/retrieval-corpus/wiki-18.jsonl\",\n        help=\"Local corpus file.\",\n    )\n    parser.add_argument(\"--topk\", type=int, default=3, help=\"Number of retrieved passages for one query.\")\n    parser.add_argument(\"--retriever_name\", type=str, default=\"e5\", help=\"Name of the retriever model.\")\n    parser.add_argument(\n        \"--retriever_model\", type=str, default=\"intfloat/e5-base-v2\", help=\"Path of the retriever model.\"\n    )\n    parser.add_argument(\"--faiss_gpu\", action=\"store_true\", help=\"Use GPU for computation\")\n\n    args = parser.parse_args()\n\n    # 1) Build a config (could also parse from arguments).\n    #    In real usage, you'd parse your CLI arguments or environment variables.\n    config = Config(\n        retrieval_method=args.retriever_name,  # or \"dense\"\n        index_path=args.index_path,\n        corpus_path=args.corpus_path,\n        retrieval_topk=args.topk,\n        faiss_gpu=args.faiss_gpu,\n        retrieval_model_path=args.retriever_model,\n        retrieval_pooling_method=\"mean\",\n        retrieval_query_max_length=256,\n        retrieval_use_fp16=True,\n        retrieval_batch_size=512,\n    )\n\n    # 2) Instantiate a global retriever so it is loaded once and reused.\n    retriever = get_retriever(config)\n\n    # 3) Launch the server. By default, it listens on http://127.0.0.1:8000\n    uvicorn.run(app, host=\"0.0.0.0\", port=8000)\n"
  },
  {
    "path": "examples/sglang_multiturn/search_r1_like/run_qwen2.5-3b_instruct_search_multiturn.sh",
    "content": "# run on 8xH20\n# make sure your current working directory is the root of the project\n\nset -x\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\n\nTRAIN_DATA=\"$HOME/data/searchR1_processed_direct/train.parquet\"\nVAL_DATA=\"$HOME/data/searchR1_processed_direct/test.parquet\"\n\nTOOL_CONFIG=\"$CONFIG_PATH/tool_config/search_tool_config.yaml\"\n\n\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='search_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=512 \\\n    data.val_batch_size=256 \\\n    data.max_prompt_length=4096 \\\n    data.max_response_length=3000 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.285 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.max_model_len=15000 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.multi_turn.max_assistant_turns=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.val_before_train=False \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='search_r1_like_async_rl' \\\n    trainer.experiment_name='qwen2.5-3b-instruct_function_rm-search-async-sgl-multi-w-searchtool-verify-n16' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=100 \\\n    trainer.test_freq=50 \\\n    data.train_files=\"$TRAIN_DATA\" \\\n    data.val_files=\"$VAL_DATA\"  \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$TOOL_CONFIG\" \\\n    trainer.total_epochs=1 $@\n\n"
  },
  {
    "path": "examples/skypilot/README.md",
    "content": "# verl with SkyPilot\n\nRun verl reinforcement learning training jobs on Kubernetes clusters or cloud platforms with GPU nodes using [SkyPilot](https://github.com/skypilot-org/skypilot).\n\n## Installation and Configuration\n\n### Step 1: Install SkyPilot\n\nChoose the installation based on your target platform:\n\n```bash\n# For Kubernetes only\npip install \"skypilot[kubernetes]\"\n\n# For AWS\npip install \"skypilot[aws]\"\n\n# For Google Cloud Platform\npip install \"skypilot[gcp]\"\n\n# For Azure\npip install \"skypilot[azure]\"\n\n# For multiple platforms\npip install \"skypilot[kubernetes,aws,gcp,azure]\"\n```\n\n### Step 2: Configure Your Platform\n\nSee https://docs.skypilot.co/en/latest/getting-started/installation.html\n\n### Step 3: Set Up Environment Variables\n\nExport necessary API keys for experiment tracking:\n\n```bash\n# For Weights & Biases tracking\nexport WANDB_API_KEY=\"your-wandb-api-key\"\n\n# For HuggingFace gated models (if needed)\nexport HF_TOKEN=\"your-huggingface-token\"\n```\n\n## Examples\n\n### PPO Training\n```bash\nsky launch -c verl-ppo verl-ppo.yaml --secret WANDB_API_KEY -y\n```\nRuns PPO training on GSM8K dataset using Qwen2.5-0.5B-Instruct model across 2 nodes with H100 GPUs. Based on examples in [`../ppo_trainer/`](../ppo_trainer/).\n\n### GRPO Training  \n```bash\nsky launch -c verl-grpo verl-grpo.yaml --secret WANDB_API_KEY -y\n```\nRuns GRPO (Group Relative Policy Optimization) training on MATH dataset using Qwen2.5-7B-Instruct model. Memory-optimized configuration for 2 nodes. Based on examples in [`../grpo_trainer/`](../grpo_trainer/).\n\n### Multi-turn Tool Usage Training\n```bash\nsky launch -c verl-multiturn verl-multiturn-tools.yaml --secret WANDB_API_KEY --secret HF_TOKEN -y\n```\nSingle-node training with 8xH100 GPUs for multi-turn tool usage with Qwen2.5-3B-Instruct. Includes tool and interaction configurations for GSM8K. Based on examples in [`../sglang_multiturn/`](../sglang_multiturn/) but uses vLLM instead of sglang.\n\n## Configuration\n\nThe example YAML files are pre-configured with:\n\n- **Infrastructure**: Kubernetes clusters (`infra: k8s`) - can be changed to `infra: aws` or `infra: gcp`, etc.\n- **Docker Image**: verl's official Docker image with CUDA 12.6 support\n- **Setup**: Automatically clones and installs verl from source\n- **Datasets**: Downloads required datasets during setup phase\n- **Ray Cluster**: Configures distributed training across nodes\n- **Logging**: Supports Weights & Biases via `--secret WANDB_API_KEY`\n- **Models**: Supports gated HuggingFace models via `--secret HF_TOKEN`\n\n## Launch Command Options\n\n- `-c <name>`: Cluster name for managing the job\n- `--secret KEY`: Pass secrets for API keys (can be used multiple times)\n- `-y`: Skip confirmation prompt\n\n## Monitoring Your Jobs\n\n### Check cluster status\n```bash\nsky status\n```\n\n### View logs\n```bash\nsky logs verl-ppo  # View logs for the PPO job\n```\n\n### SSH into head node\n```bash\nssh verl-ppo\n```\n\n### Access Ray dashboard\n```bash\nsky status --endpoint 8265 verl-ppo  # Get dashboard URL\n```\n\n### Stop a cluster\n```bash\nsky down verl-ppo\n```\n"
  },
  {
    "path": "examples/skypilot/verl-grpo.yaml",
    "content": "resources:\n  infra: k8s\n  accelerators: H100:1 \n  memory: 128+\n  image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4\n  ports: 8265\n\nnum_nodes: 2\n\nsecrets:\n  WANDB_API_KEY: \n\nsetup: |  \n  rm -rf verl\n  git clone https://github.com/volcengine/verl.git\n  cd verl\n  pip3 install -v -e .[vllm]\n  pip3 install flashinfer-python\n  echo \"Downloading Math dataset...\"\n  mkdir -p ~/data/math\n  python3 \"$(pwd)/examples/data_preprocess/math_dataset.py\" --local_dir ~/data/math\n  echo \"Math dataset download completed\"\n\nrun: |\n  HEAD_IP=$(echo \"$SKYPILOT_NODE_IPS\" | head -n1)\n  NUM_NODES=$SKYPILOT_NUM_NODES\n  NUM_GPUS_PER_NODE=$SKYPILOT_NUM_GPUS_PER_NODE\n  \n  if [ \"$SKYPILOT_NODE_RANK\" == \"0\" ]; then\n    echo \"Starting Ray head node...\"\n    ps aux | grep ray | grep 6379 &> /dev/null ||  ray start --head --disable-usage-stats \\\n          --port=6379 \\\n          --dashboard-host=0.0.0.0 \\\n          --dashboard-port=8265\n\n    # Wait for all worker nodes to join\n    retry_count=0\n    max_retries=30\n    while [ $retry_count -lt $max_retries ]; do\n      connected_nodes=$(ray status 2>/dev/null | grep -c \"node_\" || echo \"0\")\n      echo \"Connected nodes: $connected_nodes/$NUM_NODES (attempt $((retry_count+1))/$max_retries)\"\n      \n      if [ \"$connected_nodes\" -ge \"$NUM_NODES\" ]; then\n        echo \"All nodes connected to Ray cluster\"\n        break\n      fi\n      \n      retry_count=$((retry_count+1))\n      sleep 10\n    done\n\n    python3 -m verl.trainer.main_ppo \\\n     algorithm.adv_estimator=grpo \\\n     data.train_files=$HOME/data/math/train.parquet \\\n     data.val_files=$HOME/data/math/test.parquet \\\n     data.train_batch_size=32 \\\n     data.max_prompt_length=256 \\\n     data.max_response_length=256 \\\n     data.filter_overlong_prompts=True \\\n     data.truncation='error' \\\n     actor_rollout_ref.model.path=Qwen/Qwen2.5-7B-Instruct \\\n     actor_rollout_ref.actor.optim.lr=1e-6 \\\n     actor_rollout_ref.model.use_remove_padding=True \\\n     actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n     actor_rollout_ref.actor.ppo_epochs=1 \\\n     actor_rollout_ref.actor.use_kl_loss=False \\\n     actor_rollout_ref.actor.entropy_coeff=0 \\\n     actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n     actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n     actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n     actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n     actor_rollout_ref.rollout.name=vllm \\\n     actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n     actor_rollout_ref.rollout.n=1 \\\n     actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n     actor_rollout_ref.rollout.max_num_batched_tokens=2048 \\\n     actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n     actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n     algorithm.use_kl_in_reward=False \\\n     trainer.critic_warmup=0 \\\n     trainer.logger=[console,wandb] \\\n     trainer.project_name=verl_math_grpo_demo \\\n     trainer.experiment_name=qwen25_7b_grpo \\\n     trainer.n_gpus_per_node=$NUM_GPUS_PER_NODE \\\n     trainer.nnodes=$NUM_NODES \\\n     trainer.save_freq=-1 \\\n     trainer.test_freq=-1 \\\n     trainer.total_epochs=1\n\n  else\n    sleep 15\n    echo \"Starting Ray worker node...\"\n    ps aux | grep ray | grep $HEAD_IP:6379 &> /dev/null || ray start --address $HEAD_IP:6379 --disable-usage-stats\n    sleep 10\n  fi\n\n  echo \"Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK.\""
  },
  {
    "path": "examples/skypilot/verl-multiturn-tools.yaml",
    "content": "resources:\n  infra: k8s\n  accelerators: H100:8\n  memory: 128+\n  image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4\n  ports: 8265\n\nnum_nodes: 1\n\nsecrets:\n  WANDB_API_KEY: \n  HF_TOKEN: # in case you're using gated models from the HF hub\n\nsetup: |\n  rm -rf verl\n  git clone https://github.com/volcengine/verl.git\n  cd verl\n  pip3 install -v -e .[vllm]\n  pip3 install flashinfer-python\n  pip install \"transformers<4.54.0\" # https://github.com/vllm-project/vllm-ascend/issues/2046\n  # Download GSM8K dataset for multiturn tool training\n  echo \"Downloading GSM8K dataset...\"\n  mkdir -p ~/data/gsm8k\n  python3 \"$(pwd)/examples/data_preprocess/gsm8k.py\" --local_dir ~/data/gsm8k\n  echo \"GSM8K dataset download completed\"\n\nrun: |\n  NUM_GPUS_PER_NODE=$SKYPILOT_NUM_GPUS_PER_NODE\n  PROJECT_DIR=\"$(pwd)/verl\"\n  CONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n  \n  # Single node setup - no worker coordination needed\n  echo \"Starting Ray head node...\"\n  ps aux | grep ray | grep 6379 &> /dev/null || ray start --head --disable-usage-stats \\\n        --port=6379 \\\n        --dashboard-host=0.0.0.0 \\\n        --dashboard-port=8265\n\n  cd verl\n\n  python3 -m verl.trainer.main_ppo \\\n     --config-path=\"$CONFIG_PATH\" \\\n     --config-name='gsm8k_multiturn_grpo' \\\n     algorithm.adv_estimator=grpo \\\n     data.train_batch_size=512 \\\n     data.max_prompt_length=1024 \\\n     data.max_response_length=1024 \\\n     data.filter_overlong_prompts=True \\\n     data.truncation='error' \\\n     data.return_raw_chat=True \\\n     data.train_files=$HOME/data/gsm8k/train.parquet \\\n     data.val_files=$HOME/data/gsm8k/test.parquet \\\n     actor_rollout_ref.model.path=Qwen/Qwen2.5-3B-Instruct \\\n     actor_rollout_ref.actor.optim.lr=1e-6 \\\n     actor_rollout_ref.model.use_remove_padding=True \\\n     actor_rollout_ref.actor.ppo_mini_batch_size=512 \\\n     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n     actor_rollout_ref.actor.use_kl_loss=True \\\n     actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n     actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n     actor_rollout_ref.actor.entropy_coeff=0 \\\n     actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n     actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n     actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=64 \\\n     actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n     actor_rollout_ref.rollout.name=vllm \\\n     actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n     actor_rollout_ref.rollout.n=16 \\\n     actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=64 \\\n     actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n     algorithm.use_kl_in_reward=False \\\n     trainer.critic_warmup=0 \\\n     trainer.logger=[console,wandb] \\\n     trainer.project_name=verl_multiturn_tools \\\n     trainer.experiment_name=qwen25_7b_gsm8k_multiturn_tools \\\n     trainer.n_gpus_per_node=$NUM_GPUS_PER_NODE \\\n     trainer.nnodes=1 \\\n     trainer.save_freq=10 \\\n     trainer.test_freq=5 \\\n     trainer.total_epochs=10 \\\n     actor_rollout_ref.actor.ppo_max_token_len_per_gpu=8192 \\\n     actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=8192 \\\n     actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=8192 \\\n     critic.ppo_max_token_len_per_gpu=8192 \\\n     critic.forward_max_token_len_per_gpu=8192 \\\n     actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n     actor_rollout_ref.rollout.multi_turn.interaction_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/interaction_config/gsm8k_interaction_config.yaml\" \\\n     actor_rollout_ref.rollout.multi_turn.max_user_turns=1\n\n  echo \"Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK.\""
  },
  {
    "path": "examples/skypilot/verl-ppo.yaml",
    "content": "resources:\n  infra: k8s\n  accelerators: H100:1 \n  memory: 128+\n  image_id: docker:verlai/verl:base-verl0.5-cu126-cudnn9.8-torch2.7.0-fa2.7.4\n  ports: 8265\n\nnum_nodes: 2\n\nsecrets:\n  WANDB_API_KEY: \n\nsetup: |  \n  rm -rf verl\n  git clone https://github.com/volcengine/verl.git\n  cd verl\n  pip3 install -v -e .[vllm]\n  pip3 install flashinfer-python\n  # Download GSM8K dataset - alternative approach\n  echo \"Downloading GSM8K dataset...\"\n  mkdir -p ~/data/gsm8k\n  # Check if the script exists and use absolute path\n  if [ -f \"$(pwd)/examples/data_preprocess/gsm8k.py\" ]; then\n    python3 \"$(pwd)/examples/data_preprocess/gsm8k.py\" --local_dir ~/data/gsm8k\n  else\n    echo \"Warning: gsm8k.py script not found, skipping dataset download\"\n    # You might want to download the dataset manually or use a different approach\n  fi\n  echo \"GSM8K dataset download completed\"\n\nrun: |\n  # Get the Head node's IP and total number of nodes\n  HEAD_IP=$(echo \"$SKYPILOT_NODE_IPS\" | head -n1)\n  NUM_NODES=$SKYPILOT_NUM_NODES\n  \n  # login wandb\n  # python3 -c \"import wandb; wandb.login(relogin=True, key='$WANDB_API_KEY')\"\n\n  if [ \"$SKYPILOT_NODE_RANK\" == \"0\" ]; then\n    # Head node starts Ray Head\n    echo \"Starting Ray head node...\"\n    ps aux | grep ray | grep 6379 &> /dev/null ||  ray start --head --disable-usage-stats \\\n          --port=6379 \\\n          --dashboard-host=0.0.0.0 \\\n          --dashboard-port=8265\n\n    # Wait for all worker nodes to join the cluster with better checking\n    echo \"Waiting for all nodes to join Ray cluster...\"\n    retry_count=0\n    max_retries=30\n    while [ $retry_count -lt $max_retries ]; do\n      connected_nodes=$(ray status 2>/dev/null | grep -c \"node_\" || echo \"0\")\n      echo \"Connected nodes: $connected_nodes/$NUM_NODES (attempt $((retry_count+1))/$max_retries)\"\n      \n      if [ \"$connected_nodes\" -ge \"$NUM_NODES\" ]; then\n        echo \"All nodes connected to Ray cluster\"\n        break\n      fi\n      \n      retry_count=$((retry_count+1))\n      sleep 10\n    done\n\n    if [ $retry_count -eq $max_retries ]; then\n      echo \"WARNING: Not all nodes connected to Ray cluster after $max_retries attempts\"\n      echo \"Current Ray status:\"\n      ray status\n    fi\n\n    python3 -m verl.trainer.main_ppo \\\n     data.train_files=$HOME/data/gsm8k/train.parquet \\\n     data.val_files=$HOME/data/gsm8k/test.parquet \\\n     data.train_batch_size=256 \\\n     data.max_prompt_length=512 \\\n     data.max_response_length=256 \\\n     actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n     actor_rollout_ref.actor.optim.lr=1e-6 \\\n     actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n     actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n     actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n     actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n     actor_rollout_ref.rollout.name=vllm \\\n     actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n     actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n     critic.optim.lr=1e-5 \\\n     critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n     critic.ppo_micro_batch_size_per_gpu=4 \\\n     algorithm.kl_ctrl.kl_coef=0.001 \\\n     trainer.logger=[console,wandb] \\\n     trainer.val_before_train=False \\\n     trainer.default_hdfs_dir=null \\\n     trainer.n_gpus_per_node=1 \\\n     trainer.nnodes=2 \\\n     trainer.save_freq=20 \\\n     trainer.test_freq=20 \\\n     trainer.total_epochs=2 \\\n     trainer.project_name=verl_examples \\\n     trainer.experiment_name=experiment_name_gsm8k\n\n  else\n    # Wait for Ray Head to start\n    sleep 15\n    # Worker node starts Ray Worker\n    echo \"Starting Ray worker node...\"\n    ps aux | grep ray | grep $HEAD_IP:6379 &> /dev/null || ray start --address $HEAD_IP:6379 --disable-usage-stats\n    sleep 10\n  fi\n\n  echo \"Node setup and Ray start script finished for rank $SKYPILOT_NODE_RANK.\""
  },
  {
    "path": "examples/slurm/ray_on_slurm.slurm",
    "content": "#!/bin/bash\n#SBATCH --job-name=verl-ray-on-slurm\n#SBATCH --nodes=2\n#SBATCH --ntasks-per-node=1\n#SBATCH --mem=200G\n#SBATCH --partition=your-partition\n#SBATCH --time=01:00:00\n#SBATCH --account=your-account\n#SBATCH --gpus-per-node=4\n#SBATCH --cpus-per-task=64\n#SBATCH --output=slurm-%j.out\n#SBATCH --error=slurm-%j.err\n\n# load necessary modules\n\n# replace these information with your own\nverl_workdir=/path/to/verl\ntrain_files=/path/to/gsm8k/train.parquet\nval_files=/path/to/gsm8k/test.parquet\napptainer_image_path=/path/to/verl-ngc.sif\n# replace these information with your own\n\n# Getting the node names\nnodes=$(scontrol show hostnames \"$SLURM_JOB_NODELIST\")\nnodes_array=(\"$nodes\")\n\nhead_node=${nodes_array[0]}\nhead_node_ip=$(srun --nodes=1 --ntasks=1 -w \"$head_node\" hostname --ip-address)\n\n# if we detect a space character in the head node IP, we'll\n# convert it to an ipv4 address. This step is optional.\nif [[ \"$head_node_ip\" == *\" \"* ]]; then\nIFS=' ' read -ra ADDR <<<\"$head_node_ip\"\nif [[ ${#ADDR[0]} -gt 16 ]]; then\n  head_node_ip=${ADDR[1]}\nelse\n  head_node_ip=${ADDR[0]}\nfi\necho \"IPV6 address detected. We split the IPV4 address as $head_node_ip\"\nfi\n\nport=6379\nip_head=$head_node_ip:$port\nexport ip_head\necho \"IP Head: $ip_head\"\n\n# make sure we set environment variables before Ray initialization\n\nprintenv\n\necho \"Starting HEAD at $head_node\"\nsrun --nodes=1 --ntasks=1 -w \"$head_node\" \\\n    apptainer run --nv --bind $verl_workdir $apptainer_image_path \\\n        ray start --head --node-ip-address=\"$head_node_ip\" --port=$port \\\n        --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n# optional, though may be useful in certain versions of Ray < 1.0.\nsleep 10\n\n# number of nodes other than the head node\nworker_num=$((SLURM_JOB_NUM_NODES - 1))\n\nfor ((i = 1; i <= worker_num; i++)); do\n    node_i=${nodes_array[$i]}\n    echo \"Starting WORKER $i at $node_i\"\n    srun --nodes=1 --ntasks=1 -w \"$node_i\" \\\n        apptainer run --nv --bind $verl_workdir $apptainer_image_path \\\n            ray start --address \"$ip_head\" --num-cpus \"${SLURM_CPUS_PER_TASK}\" --num-gpus \"${SLURM_GPUS_PER_NODE}\" --block &\n    sleep 5\ndone\n\nPYTHONUNBUFFERED=1 srun --overlap --nodes=1 --ntasks=1 -w \"$head_node\" \\\n    apptainer run --nv --bind $verl_workdir $apptainer_image_path \\\n    python3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$train_files \\\n    data.val_files=$val_files \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=256 \\\n    actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n    critic.ppo_micro_batch_size_per_gpu=4 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.logger=console \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=\"${SLURM_GPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${SLURM_NNODES}\" \\\n    trainer.save_freq=10 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=15 2>&1 | tee verl_demo_slurm.log\n"
  },
  {
    "path": "examples/split_placement/README.md",
    "content": "# Split Placement Example\nHere we introduce how to run the naive implementation of the split placement of PPO algorithm.\nWe will release the complete version of flexible placement in the near future.\n\n For quickstart, you can only follow Step 2 to modify the code and then follow Step 4 to execute the split placement example.\n\n### Step 1: Placing the models to different GPUs\nSpecify the placement and resource allocation. In the example, we place the actor and reference in the first half of the GPUs while map the critic and reward model (if any) to the second half of the GPUs.\n```python\nactor_rollout_ref_pool_id = 'actor_rollout_ref_pool'\ncritic_pool_id = 'critic_pool'\nif config.trainer.nnodes // 2 == 0 and config.trainer.n_gpus_per_node // 2 > 0:\n    resource_pool_spec = {\n        actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes,\n        critic_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes,\n    }\nelse:\n    resource_pool_spec = {\n        actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2),\n        critic_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2),\n    }\nprint(f'resource_pool_spec: {resource_pool_spec}')\nmapping = {\n    Role.ActorRollout: actor_rollout_ref_pool_id,\n    Role.Critic: critic_pool_id,\n    Role.RefPolicy: actor_rollout_ref_pool_id,\n}\nmapping[Role.RewardModel] = critic_pool_id\n```\n\n### Step 2: Make the models executed asynchronously\nBased on the model placement, we need to make the models executed asynchronously.\n\nTo do so, you need to turn off the `blocking` flag (i.e., `blocking=False`) in our decorator of some model operations.\nFor example, we hope the actor update and critic update can be executed in parallel, then we need to make the following modification in `fsdp_workers.py`\n\n```\n@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False)\ndef update_actor(self, data: DataProto):\n    ...\n\n@register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False)\ndef update_critic(self, data: DataProto):\n    ...\n```\n\nWe can also parallelize the computation of `ref_log_prob` and `values` and `rewards` in the split placement. For simplicity of the tutorial, we don't do this in this example.\n\n### Step 3: Execute these operation in parallel in the single controller process\nTo implement the parallel execution of the actor and critic update, the only thing we need to modify in the `ray_trainer.py` is to `get` the concurrent  `futures` on the single controller process.\n\n```python\ncritic_output = critic_output.get()\nactor_output = actor_output.get()\n```\n\n### Step 4: Run the split placement example\n\n```\nbash run_deepseek7b_llm.sh\n```\n"
  },
  {
    "path": "examples/split_placement/config/ppo_trainer_split.yaml",
    "content": "# the ppo trainer split config will override default ppo_trainer.yaml\n\nhydra:\n  searchpath:\n    - file://../../verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ndata:\n  tokenizer: null\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1  # set to -1 to use full dataset\n  val_max_samples: -1  # set to -1 to use full dataset\n  prompt_key: prompt\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null # DEPRECATED: Validation datasets are sent to inference engines as a whole batch, which will schedule the memory themselves\n  return_raw_input_ids: False  # This should be set to true when the tokenizer between policy and rm differs\n  return_raw_chat: False\n  return_full_prompt: False\n  shuffle: True\n  seed: 42\n\nactor_rollout_ref:\n  hybrid_engine: True\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    external_lib: null\n    override_config: { }\n    enable_gradient_checkpointing: True\n    use_remove_padding: False\n  actor:\n    strategy: fsdp  # This is for backward-compatibility\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: False\n    ppo_max_token_len_per_gpu: 16384 # n * ${data.max_prompt_length} + ${data.max_response_length}\n    grad_clip: 1.0\n    clip_ratio: 0.2\n    entropy_coeff: 0.0\n    use_kl_loss: False # True for GRPO\n    kl_loss_coef: 0.001 # for grpo\n    kl_loss_type: low_var_kl # for grpo\n    ppo_epochs: 1\n    shuffle: False\n    ulysses_sequence_parallel_size: 1 # sp size\n    optim:\n      lr: 1e-6\n      lr_warmup_steps: -1 # Prioritized. Negative values mean delegating to lr_warmup_steps_ratio.\n      lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n      min_lr_ratio: null   # only useful for warmup with cosine\n      lr_scheduler_type: constant  # select from constant/cosine\n      total_training_steps: -1  # must be override by program\n    fsdp_config:\n      wrap_policy:\n        # transformer_layer_cls_to_wrap: None\n        min_num_params: 0\n      param_offload: False\n      optimizer_offload: False\n      fsdp_size: -1\n  ref:\n    fsdp_config:\n      param_offload: False\n      wrap_policy:\n        # transformer_layer_cls_to_wrap: None\n        min_num_params: 0\n    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n    ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size} # sp size\n  rollout:\n    name: vllm\n    temperature: 1.0\n    top_k: -1 # 0 for hf rollout, -1 for vllm rollout\n    top_p: 1\n    prompt_length: ${data.max_prompt_length}  # not use for opensource\n    response_length: ${data.max_response_length}\n    # for vllm rollout\n    dtype: bfloat16 # should align with FSDP\n    gpu_memory_utilization: 0.5\n    ignore_eos: False\n    enforce_eager: True\n    free_cache_engine: True\n    load_format: dummy_dtensor\n    tensor_model_parallel_size: 2\n    max_num_batched_tokens: 8192\n    max_num_seqs: 1024\n    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n    disable_log_stats: True\n    enable_chunked_prefill: True # could get higher throughput\n    # for hf rollout\n    do_sample: True\n    # number of responses (i.e. num sample times)\n    n: 1 # > 1 for grpo\n\ncritic:\n  strategy: fsdp\n  optim:\n    lr: 1e-5\n    lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n    min_lr_ratio: null   # only useful for warmup with cosine\n    lr_scheduler_type: constant  # select from constant/cosine\n    total_training_steps: -1  # must be override by program\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${actor_rollout_ref.model.path}\n    override_config: { }\n    external_lib: ${actor_rollout_ref.model.external_lib}\n    enable_gradient_checkpointing: True\n    use_remove_padding: False\n    fsdp_config:\n      param_offload: False\n      optimizer_offload: False\n      wrap_policy:\n        # transformer_layer_cls_to_wrap: None\n        min_num_params: 0\n      fsdp_size: -1\n  ppo_mini_batch_size: ${actor_rollout_ref.actor.ppo_mini_batch_size}\n  ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu\n  ppo_micro_batch_size_per_gpu: null\n  forward_micro_batch_size: ${critic.ppo_micro_batch_size}\n  forward_micro_batch_size_per_gpu: ${critic.ppo_micro_batch_size_per_gpu}\n  use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n  ppo_max_token_len_per_gpu: 32768 # (${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}) * 2\n  forward_max_token_len_per_gpu: ${critic.ppo_max_token_len_per_gpu}\n  ulysses_sequence_parallel_size: 1 # sp size\n  ppo_epochs: ${actor_rollout_ref.actor.ppo_epochs}\n  shuffle: ${actor_rollout_ref.actor.shuffle}\n  grad_clip: 1.0\n  cliprange_value: 0.5\n\nreward_model:\n  enable: False\n  strategy: fsdp\n  model:\n    input_tokenizer: ${actor_rollout_ref.model.path}  # set this to null if the chat template is identical\n    path: ~/models/FsfairX-LLaMA3-RM-v0.1\n    external_lib: ${actor_rollout_ref.model.external_lib}\n    use_remove_padding: False\n    fsdp_config:\n      min_num_params: 0\n      param_offload: False\n      fsdp_size: -1\n  micro_batch_size: null # will be deprecated, use micro_batch_size_per_gpu\n  micro_batch_size_per_gpu: null # set a number\n  max_length: null\n  ulysses_sequence_parallel_size: 1 # sp size\n  use_dynamic_bsz: ${critic.use_dynamic_bsz}\n  forward_max_token_len_per_gpu: ${critic.forward_max_token_len_per_gpu}\n  reward_manager: naive\n\nalgorithm:\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  use_kl_in_reward: False\n  kl_penalty: kl  # how to estimate kl divergence\n  kl_ctrl:\n    type: fixed\n    kl_coef: 0.001\n\ntrainer:\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger: [ 'console', 'wandb' ]\n  log_val_generations: 0\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  # auto: find the last ckpt to resume. If can't find, start from scratch\n  resume_mode: auto # or disable or resume_path if resume_from_path is set\n  resume_from_path: null\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n\nray_kwargs:\n  ray_init:\n    num_cpus: null # `None` means using all CPUs, which might cause hang if limited in systems like SLURM. Please set to a number allowed then.\n"
  },
  {
    "path": "examples/split_placement/main_ppo_split.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nNote that we don't combine the main with ray_trainer as ray_trainer is used by other main.\n\"\"\"\n\nimport hydra\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\nfrom split_monkey_patch import fit\n\nfrom verl import DataProto\nfrom verl.trainer.ppo.ray_trainer import RayPPOTrainer\nfrom verl.trainer.ppo.utils import need_reference_policy\nfrom verl.utils.reward_score import gsm8k, math_reward\n\n\ndef _select_rm_score_fn(data_source):\n    if data_source == \"openai/gsm8k\":\n        return gsm8k.compute_score\n    elif data_source == \"lighteval/MATH\":\n        return math_reward.compute_score\n    else:\n        raise NotImplementedError\n\n\nclass RewardManager:\n    def __init__(self, tokenizer, num_examine) -> None:\n        self.tokenizer = tokenizer\n        self.num_examine = num_examine  # the number of batches of decoded responses to print to the console\n\n    def __call__(self, data: DataProto, return_dict: bool = False):\n        \"\"\"We will expand this function gradually based on the available datasets\"\"\"\n\n        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn\n        if \"rm_scores\" in data.batch.keys():\n            return data.batch[\"rm_scores\"]\n\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n\n        already_print_data_sources = {}\n\n        for i in range(len(data)):\n            data_item = data[i]  # DataProtoItem\n\n            prompt_ids = data_item.batch[\"prompts\"]\n\n            prompt_length = prompt_ids.shape[-1]\n\n            valid_prompt_length = data_item.batch[\"attention_mask\"][:prompt_length].sum()\n            valid_prompt_ids = prompt_ids[-valid_prompt_length:]\n\n            response_ids = data_item.batch[\"responses\"]\n            valid_response_length = data_item.batch[\"attention_mask\"][prompt_length:].sum()\n            valid_response_ids = response_ids[:valid_response_length]\n\n            # decode\n            sequences = torch.cat((valid_prompt_ids, valid_response_ids))\n            sequences_str = self.tokenizer.decode(sequences)\n\n            ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n\n            # select rm_score\n            data_source = data_item.non_tensor_batch[\"data_source\"]\n            compute_score_fn = _select_rm_score_fn(data_source)\n\n            score = compute_score_fn(solution_str=sequences_str, ground_truth=ground_truth)\n            reward_tensor[i, valid_response_length - 1] = score\n\n            if data_source not in already_print_data_sources:\n                already_print_data_sources[data_source] = 0\n\n            if already_print_data_sources[data_source] < self.num_examine:\n                already_print_data_sources[data_source] += 1\n                print(sequences_str)\n\n        if return_dict:\n            return {\"reward_tensor\": reward_tensor}\n        else:\n            return reward_tensor\n\n\n@hydra.main(config_path=\"config\", config_name=\"ppo_trainer_split\", version_base=None)\ndef main(config):\n    if not ray.is_initialized():\n        # this is for local ray cluster\n        default_runtime_env = {\"env_vars\": {\"TOKENIZERS_PARALLELISM\": \"true\", \"NCCL_DEBUG\": \"WARN\"}}\n        ray_init_kwargs = config.ray_kwargs.get(\"ray_init\", {})\n        runtime_env_kwargs = ray_init_kwargs.get(\"runtime_env\", {})\n        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)\n        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, \"runtime_env\": runtime_env})\n        print(f\"ray init kwargs: {ray_init_kwargs}\")\n        ray.init(**OmegaConf.to_container(ray_init_kwargs))\n\n    ray.get(main_task.remote(config))\n\n\n@ray.remote\ndef main_task(config):\n    # print initial config\n    from pprint import pprint\n\n    from omegaconf import OmegaConf\n\n    from verl.utils.fs import copy_to_local\n\n    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values\n    OmegaConf.resolve(config)\n\n    # download the checkpoint from hdfs\n    local_path = copy_to_local(config.actor_rollout_ref.model.path)\n\n    # instantiate tokenizer\n    from verl.utils import hf_tokenizer\n\n    tokenizer = hf_tokenizer(local_path)\n\n    # define worker classes\n    if config.actor_rollout_ref.actor.strategy in {\"fsdp\", \"fsdp2\"}:\n        assert config.critic.strategy in {\"fsdp\", \"fsdp2\"}\n        from verl.single_controller.ray import RayWorkerGroup\n        from verl.workers.fsdp_workers import ActorRolloutRefWorker, CriticWorker\n\n        ray_worker_group_cls = RayWorkerGroup\n\n    elif config.actor_rollout_ref.actor.strategy == \"megatron\":\n        assert config.actor_rollout_ref.actor.strategy == config.critic.strategy\n        from verl.single_controller.ray import RayWorkerGroup\n        from verl.workers.megatron_workers import ActorRolloutRefWorker, CriticWorker\n\n        ray_worker_group_cls = RayWorkerGroup\n\n    else:\n        raise NotImplementedError\n\n    from verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role\n\n    role_worker_mapping = {\n        Role.ActorRollout: ray.remote(ActorRolloutRefWorker),\n        Role.Critic: ray.remote(CriticWorker),\n    }\n\n    # NOTE: initialze two resource pool\n    actor_rollout_ref_pool_id = \"actor_rollout_ref_pool\"\n    critic_pool_id = \"critic_pool\"\n    if config.trainer.nnodes // 2 == 0 and config.trainer.n_gpus_per_node // 2 > 0:\n        resource_pool_spec = {\n            actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes,\n            critic_pool_id: [config.trainer.n_gpus_per_node // 2] * config.trainer.nnodes,\n        }\n    else:\n        resource_pool_spec = {\n            actor_rollout_ref_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2),\n            critic_pool_id: [config.trainer.n_gpus_per_node] * (config.trainer.nnodes // 2),\n        }\n    print(f\"resource_pool_spec: {resource_pool_spec}\")\n    mapping = {\n        Role.ActorRollout: actor_rollout_ref_pool_id,\n        Role.Critic: critic_pool_id,\n    }\n\n    # use reference model\n    if need_reference_policy(config):\n        role_worker_mapping[Role.RefPolicy] = ray.remote(ActorRolloutRefWorker)\n        mapping[Role.RefPolicy] = actor_rollout_ref_pool_id\n\n    reward_fn = RewardManager(tokenizer=tokenizer, num_examine=0)\n\n    # Note that we always use function-based RM for validation\n    val_reward_fn = RewardManager(tokenizer=tokenizer, num_examine=1)\n\n    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n\n    RayPPOTrainer.fit = fit\n    trainer = RayPPOTrainer(\n        config=config,\n        tokenizer=tokenizer,\n        role_worker_mapping=role_worker_mapping,\n        resource_pool_manager=resource_pool_manager,\n        ray_worker_group_cls=ray_worker_group_cls,\n        reward_fn=reward_fn,\n        val_reward_fn=val_reward_fn,\n    )\n    trainer.init_workers()\n    trainer.fit()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "examples/split_placement/run_deepseek7b_llm.sh",
    "content": "set -x\n\npython3 main_ppo_split.py \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=deepseek-ai/deepseek-llm-7b-chat \\\n    critic.model.enable_gradient_checkpointing=False \\\n    critic.ppo_micro_batch_size_per_gpu=8 \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_example_gsm8k' \\\n    trainer.experiment_name='deepseek_llm_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/split_placement/split_monkey_patch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nAn naive implementation of split placment example\n\"\"\"\n\nimport uuid\nfrom copy import deepcopy\nfrom pprint import pprint\n\nimport numpy as np\nimport torch\n\nfrom verl import DataProto\nfrom verl.trainer.ppo.ray_trainer import (\n    AdvantageEstimator,\n    apply_kl_penalty,\n    compute_advantage,\n    compute_data_metrics,\n    compute_timing_metrics,\n    marked_timer,\n)\nfrom verl.trainer.ppo.reward import compute_reward\nfrom verl.utils.metric import reduce_metrics\n\n\ndef fit(self):\n    \"\"\"\n    The training loop of PPO.\n    The driver process only need to call the compute functions of the worker group through RPC\n    to construct the PPO dataflow.\n    The light-weight advantage computation is done on the driver process.\n    \"\"\"\n    from omegaconf import OmegaConf\n\n    from verl.utils.tracking import Tracking\n\n    logger = Tracking(\n        project_name=self.config.trainer.project_name,\n        experiment_name=self.config.trainer.experiment_name,\n        default_backend=self.config.trainer.logger,\n        config=OmegaConf.to_container(self.config, resolve=True),\n    )\n\n    self.global_steps = 0\n\n    # load checkpoint before doing anything\n    self._load_checkpoint()\n\n    # perform validation before training\n    # currently, we only support validation using the reward_function.\n    if self.val_reward_fn is not None and self.config.trainer.get(\"val_before_train\", True):\n        val_metrics = self._validate()\n        pprint(f\"Initial validation metrics: {val_metrics}\")\n        logger.log(data=val_metrics, step=self.global_steps)\n        if self.config.trainer.get(\"val_only\", False):\n            return\n\n    # we start from step 1\n    self.global_steps += 1\n    last_val_metrics = None\n\n    for epoch in range(self.config.trainer.total_epochs):\n        for batch_dict in self.train_dataloader:\n            metrics = {}\n            timing_raw = {}\n\n            batch: DataProto = DataProto.from_single_dict(batch_dict)\n\n            # pop those keys for generation\n            gen_batch = batch.pop(batch_keys=[\"input_ids\", \"attention_mask\", \"position_ids\"])\n            is_last_step = self.global_steps >= self.total_training_steps\n\n            with marked_timer(\"step\", timing_raw):\n                # generate a batch\n                with marked_timer(\"gen\", timing_raw):\n                    gen_batch_output = self.actor_rollout_wg.generate_sequences(gen_batch)\n                    timing_raw.update(gen_batch_output.meta_info[\"timing\"])\n                    gen_batch_output.meta_info.pop(\"timing\", None)\n\n                if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:\n                    with marked_timer(\"gen_max\", timing_raw):\n                        gen_baseline_batch = deepcopy(gen_batch)\n                        gen_baseline_batch.meta_info[\"do_sample\"] = False\n                        gen_baseline_output = self.actor_rollout_wg.generate_sequences(gen_baseline_batch)\n\n                        batch = batch.union(gen_baseline_output)\n                        # compute reward model score on batch\n                        rm_scores = None\n                        if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                            rm_scores = self.rm_wg.compute_rm_score(batch)\n                            batch = batch.union(rm_scores)\n                        reward_baseline_tensor, _ = compute_reward(batch, self.reward_fn)\n                        reward_baseline_tensor = reward_baseline_tensor.sum(dim=-1)\n\n                        keys_to_pop = set(gen_baseline_output.batch.keys())\n                        if rm_scores is not None:\n                            keys_to_pop.update(rm_scores.batch.keys())\n                        batch.pop(batch_keys=list(keys_to_pop))\n\n                        batch.batch[\"reward_baselines\"] = reward_baseline_tensor\n\n                        del rm_scores, gen_baseline_batch, gen_baseline_output\n\n                batch.non_tensor_batch[\"uid\"] = np.array(\n                    [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object\n                )\n                # repeat to align with repeated responses in rollout\n                batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n                batch = batch.union(gen_batch_output)\n\n                # Balance the number of valid tokens across DP ranks.\n                # NOTE: This usually changes the order of data in the `batch`,\n                # which won't affect the advantage calculation (since it's based on uid),\n                # but might affect the loss calculation (due to the change of mini-batching).\n                # TODO: Decouple the DP balancing and mini-batching.\n                self._balance_batch(batch, metrics=metrics)\n\n                # compute global_valid tokens\n                batch.meta_info[\"global_token_num\"] = torch.sum(batch.batch[\"attention_mask\"], dim=-1).tolist()\n\n                # recompute old_log_probs\n                with marked_timer(\"old_log_prob\", timing_raw):\n                    old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)\n                    batch = batch.union(old_log_prob)\n\n                if self.use_reference_policy:\n                    # compute reference log_prob\n                    with marked_timer(\"ref\", timing_raw):\n                        ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)\n                        batch = batch.union(ref_log_prob)\n\n                # compute values\n                if self.use_critic:\n                    with marked_timer(\"values\", timing_raw):\n                        values = self.critic_wg.compute_values(batch)\n                        batch = batch.union(values)\n\n                with marked_timer(\"adv\", timing_raw):\n                    # compute scores. Support both model and function-based.\n                    # We first compute the scores using reward model. Then, we call reward_fn to combine\n                    # the results from reward model and rule-based results.\n                    if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                        # we first compute reward model score\n                        reward_tensor = self.rm_wg.compute_rm_score(batch)\n                        batch = batch.union(reward_tensor)\n\n                    # we combine with rule-based rm\n                    reward_tensor, _ = compute_reward(batch, self.reward_fn)\n                    batch.batch[\"token_level_scores\"] = reward_tensor\n\n                    # compute rewards. apply_kl_penalty if available\n                    if self.config.algorithm.use_kl_in_reward:\n                        batch, kl_metrics = apply_kl_penalty(\n                            batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty\n                        )\n                        metrics.update(kl_metrics)\n                    else:\n                        batch.batch[\"token_level_rewards\"] = batch.batch[\"token_level_scores\"]\n\n                    # compute advantages, executed on the driver process\n                    norm_adv_by_std_in_grpo = self.config.algorithm.get(\"norm_adv_by_std_in_grpo\", True)\n                    batch = compute_advantage(\n                        batch,\n                        adv_estimator=self.config.algorithm.adv_estimator,\n                        gamma=self.config.algorithm.gamma,\n                        lam=self.config.algorithm.lam,\n                        num_repeat=self.config.actor_rollout_ref.rollout.n,\n                        norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n                        config=self.config.algorithm,\n                    )\n\n                # implement critic warmup\n                if self.config.trainer.critic_warmup <= self.global_steps:\n                    # update actor\n                    with marked_timer(\"update_actor_call\", timing_raw):\n                        actor_output = self.actor_rollout_wg.update_actor(batch)\n                else:\n                    actor_output = None\n\n                # update critic\n                if self.use_critic:\n                    with marked_timer(\"update_critic_call\", timing_raw):\n                        critic_output = self.critic_wg.update_critic(batch)\n\n                    # NOTE: make sure you set blocking=False in update_actor and update_crtic in the worker class\n                    with marked_timer(\"update_actor_critic\", timing_raw):\n                        critic_output = critic_output.get()\n                        critic_output_metrics = reduce_metrics(critic_output.meta_info[\"metrics\"])\n                        metrics.update(critic_output_metrics)\n\n                if actor_output is not None:\n                    actor_output = actor_output.get()\n                    actor_output_metrics = reduce_metrics(actor_output.meta_info[\"metrics\"])\n                    metrics.update(actor_output_metrics)\n\n            # validate\n            if (\n                self.val_reward_fn is not None\n                and self.config.trainer.test_freq > 0\n                and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)\n            ):\n                with marked_timer(\"testing\", timing_raw):\n                    val_metrics: dict = self._validate()\n                    if is_last_step:\n                        last_val_metrics = val_metrics\n                metrics.update(val_metrics)\n\n            if self.config.trainer.save_freq > 0 and (\n                is_last_step or self.global_steps % self.config.trainer.save_freq == 0\n            ):\n                with marked_timer(\"save_checkpoint\", timing_raw):\n                    self._save_checkpoint()\n\n            # collect metrics\n            metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))\n            metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))\n\n            # TODO: make a canonical logger that supports various backend\n            logger.log(data=metrics, step=self.global_steps)\n\n            if self.global_steps >= self.total_training_steps:\n                pprint(f\"Final validation metrics: {last_val_metrics}\")\n                return\n\n            self.global_steps += 1\n"
  },
  {
    "path": "examples/tuning/0.5b/qwen2-0.5b_grpo-lora_1_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=4\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-0.5b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=0.5b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-0.5B-Instruct\n\nset -x\nnproc_per_gpu=1\nnnodes=1\nngpu_per_node=1\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    trainer.val_before_train=False \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.1 \\\n    actor_rollout_ref.rollout.n=1 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/1.5b/qwen2-1.5b_grpo-lora_1_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-1.5b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=1.5b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-1.5B-Instruct\n\nset -x\nnproc_per_gpu=128\nnnodes=1\nngpu_per_node=1\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.1 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/14b/qwen2-14b_grpo-lora_2_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-14b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=14b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-14B-Instruct\n\nset -x\nnproc_per_gpu=58 # 32√ → 64× → 48√ → 56√ → 60× → 58√ → 59×\nnnodes=1\nngpu_per_node=2\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.25 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=2 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/14b/qwen2_14b_grpo_4_h800_fsdp_vllm.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/rlhf/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/rlhf/math/test.parquet\nmodel_path=Qwen/Qwen2.5-Coder-14B-Instruct\n\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\nPYTHONPATH=/opt/tiger/open_verl python3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_14b_function_rm' \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@\n"
  },
  {
    "path": "examples/tuning/32b/qwen2-32b_grpo-lora_4_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-32b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=32b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-32B-Instruct\n\nset -x\nnproc_per_gpu=45 # 32√ → 64× → 48× → 40√ → 44√ → 46× → 45×\nnnodes=1\nngpu_per_node=4\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.3 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=4 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/32b/qwen2_32B_grpo_8_h20_megatron_vllm.sh",
    "content": "set -x\n\n# we need this to avoid fragmentation of GPU memory\nexport PYTORCH_CUDA_ALLOC_CONF=max_split_size_mb:256\n\ngsm8k_train_path=$HOME/data/rlhf/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/rlhf/math/test.parquet\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\nmodel_path=Qwen/Qwen2.5-32B\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=512 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=6144 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='megatron_vllm_qwen2_32b' \\\n    trainer.experiment_name='qwen2_32b_grpo_8_h20' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/tuning/3b/qwen2-3b_grpo-lora_1_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-3b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=3b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-3B-Instruct\n\nset -x\nnproc_per_gpu=62\nnnodes=1\nngpu_per_node=1\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.1 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/70b/qwen2-70b_grpo_32_h20_fsdp_vllm.sh",
    "content": "set -x\n\ngsm8k_train_path=$HOME/data/rlhf/gsm8k/train.parquet\ngsm8k_val_path=$HOME/data/rlhf/math/test.parquet\nmodel_path=Qwen/Qwen2-72B-Instruct\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$data_path \\\n    data.val_files=$gsm8k_val_path \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=16 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='Qwen2_72B_Instruct' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=4 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@"
  },
  {
    "path": "examples/tuning/70b/qwen2-70b_grpo_32_h800_fsdp_vllm.sh",
    "content": "set -x\n\n#### important: vllm version must be >= 0.8.3\n\ngsm8k_train_path=$HOME/data/rlhf/gsm8k/train.parquet\ngsm8k_val_path=$HOME/data/rlhf/math/test.parquet\nmodel_path=Qwen/Qwen2-72B-Instruct\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$gsm8k_train_path \\\n    data.val_files=$gsm8k_val_path \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=24000 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=16 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='Qwen2_72B_Instruct' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=4 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@"
  },
  {
    "path": "examples/tuning/70b/qwen2-72b_grpo-lora_8_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-72b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=72b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-72B-Instruct\n\nset -x\nnproc_per_gpu=22 # 16√ → 32× → 24× → 20√ → 22√ → 23×\nnnodes=1\nngpu_per_node=8\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/7b/qwen2-7b_grpo-lora_1_h100_fsdp_vllm.sh",
    "content": "# -*- coding: utf-8 -*-\nexport CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7\nNOW=$(date +%Y%m%d)\nexport WANDB_DIR=gsm8k-grpo-lora-qwen2.5-7b-${NOW}\nexport WANDB_PROJECT=${WANDB_DIR}\nexport WANDB_EXP=7b-${NOW}\nMODEL_PATH=Qwen/Qwen2.5-7B-Instruct\n\nset -x\nnproc_per_gpu=16 # 64√ → 128× → 96√ → 112× → 104× → 100√ → 102× → 101×\nnnodes=1\nngpu_per_node=1\ntotal_procs=$(( nproc_per_gpu * nnodes * ngpu_per_node ))\nmini_batch_size=$(( total_procs ))\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=data/gsm8k/train.parquet \\\n    data.val_files=data/gsm8k/test.parquet \\\n    data.train_batch_size=${total_procs} \\\n    data.val_batch_size=${total_procs} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=$MODEL_PATH  \\\n    actor_rollout_ref.model.use_shm=True  \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.model.lora_rank=32 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.model.target_modules=all-linear \\\n    actor_rollout_ref.actor.optim.lr=3e-5 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=-1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.2 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.max_num_seqs=512 \\\n    actor_rollout_ref.rollout.max_model_len=1536 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=1536 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=${mini_batch_size} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n    actor_rollout_ref.actor.entropy_coeff=0.001 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name=${WANDB_PROJECT} \\\n    trainer.experiment_name=${WANDB_EXP} \\\n    trainer.n_gpus_per_node=1 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=20 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@ 2>&1 | tee ${WANDB_PROJECT}.log\n"
  },
  {
    "path": "examples/tuning/7b/qwen2-7b_grpo_2_h800_fsdp_vllm.sh",
    "content": "set -x\n\n\ngsm8k_train_path=$HOME/data/rlhf/gsm8k/train.parquet\ngsm8k_test_path=$HOME/data/rlhf/math/test.parquet\nmodel_path=Qwen/Qwen2-7B-Instruct\n\ntrain_files=\"['$gsm8k_train_path']\"\ntest_files=\"['$gsm8k_test_path']\"\n\nPYTHONPATH=/opt/tiger/open_verl python3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=1024 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=2 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=15 $@\n"
  },
  {
    "path": "examples/tutorial/agent_loop_get_started/agent_loop_tutorial.ipynb",
    "content": "{\n    \"cells\": [\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"# Train ReAct agent with code sandbox\\n\",\n                \"\\n\",\n                \"In this tutorial, we will demonstrate how to train a [ReAct](https://arxiv.org/abs/2210.03629) agent to solve math problem with code sandbox.\\n\",\n                \"\\n\",\n                \"The agent works as follows:\\n\",\n                \"1. Given a math problem, the agent first query LLM to generate response and tool calls, which are python code to be executed in sandbox.\\n\",\n                \"2. If there is a tool call, the agent execute the python code in code sandbox.\\n\",\n                \"3. After code execution, the agent get the result from sandbox and append to chat history.\\n\",\n                \"4. The agent query LLM again until no tool call or max context length reached.\\n\",\n                \"\\n\",\n                \"\\n\",\n                \"<figure>\\n\",\n                \"  <img src=\\\"https://langchain-ai.github.io/langgraph/agents/assets/agent.png\\\" alt=\\\"ReAct\\\" width=\\\"400\\\">\\n\",\n                \"  <figcaption style=\\\"font-style: italic; color: #666;\\\">\\n\",\n                \"    source: <a href=\\\"https://langchain-ai.github.io/langgraph/agents/overview/\\\" target=\\\"_blank\\\">LangGraph</a>\\n\",\n                \"  </figcaption>\\n\",\n                \"</figure>\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"## 1. Prerequisite\\n\",\n                \"\\n\",\n                \"To run the examples in this notebook, you need to install the verl package first.\\n\",\n                \"```bash\\n\",\n                \"git clone https://github.com/volcengine/verl\\n\",\n                \"cd verl\\n\",\n                \"pip install -e .\\n\",\n                \"```\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 1,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stderr\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"2025-10-16 23:20:11,956\\tINFO worker.py:2004 -- Started a local Ray instance. View the dashboard at \\u001b[1m\\u001b[32mhttp://127.0.0.1:8265 \\u001b[39m\\u001b[22m\\n\",\n                        \"/usr/local/lib/python3.12/dist-packages/ray/_private/worker.py:2052: FutureWarning: Tip: In future versions of Ray, Ray will no longer override accelerator visible devices env var if num_gpus=0 or num_gpus=None (default). To enable this behavior and turn off this error message, set RAY_ACCEL_ENV_VAR_OVERRIDE_ON_ZERO=0\\n\",\n                        \"  warnings.warn(\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"import asyncio\\n\",\n                \"import sys\\n\",\n                \"import tempfile\\n\",\n                \"import os\\n\",\n                \"import socket\\n\",\n                \"import json\\n\",\n                \"\\n\",\n                \"import requests\\n\",\n                \"import ray\\n\",\n                \"import fastapi\\n\",\n                \"import uvicorn\\n\",\n                \"from starlette.requests import Request\\n\",\n                \"from starlette.responses import JSONResponse\\n\",\n                \"from pprint import pprint\\n\",\n                \"\\n\",\n                \"import verl\\n\",\n                \"\\n\",\n                \"ray.init()\\n\",\n                \"verl_config_dir = os.path.join(os.path.dirname(verl.__file__), \\\"trainer/config\\\")\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"For demo purpose, we will use Qwen/Qwen3-1.7B as the LLM. First, let's download required model and dataset used in this tutorial.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": null,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"import pyarrow.parquet as pq\\n\",\n                \"from huggingface_hub import snapshot_download\\n\",\n                \"\\n\",\n                \"snapshot_download(\\n\",\n                \"    repo_id=\\\"verl-team/lighteval-MATH-preprocessed\\\",\\n\",\n                \"    repo_type=\\\"dataset\\\",\\n\",\n                \"    local_dir=os.path.expanduser(\\\"~/verl-team/lighteval-MATH-preprocessed\\\"),\\n\",\n                \")\\n\",\n                \"snapshot_download(\\n\",\n                \"    repo_id=\\\"Qwen/Qwen3-1.7B\\\",\\n\",\n                \"    repo_type=\\\"model\\\",\\n\",\n                \"    local_dir=os.path.expanduser(\\\"~/Qwen/Qwen3-1.7B\\\"),\\n\",\n                \")\\n\",\n                \"\\n\",\n                \"model_path = os.path.expanduser(\\\"~/Qwen/Qwen3-1.7B\\\")\\n\",\n                \"train_file = os.path.expanduser(\\\"~/verl-team/lighteval-MATH-preprocessed/train.parquet\\\")\\n\",\n                \"test_file = os.path.expanduser(\\\"~/verl-team/lighteval-MATH-preprocessed/test.parquet\\\")\\n\",\n                \"\\n\",\n                \"test = pq.read_table(test_file)\\n\",\n                \"test_file = os.path.expanduser(\\\"~/verl-team/lighteval-MATH-preprocessed/test_100.parquet\\\")\\n\",\n                \"pq.write_table(test[:100], test_file)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"verl support both vllm and sglang rollout server for high performance inference. This tutorial has been tested on both vllm and sglang, you can choose either of them to run the tutorial.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 3,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"rollout_name = \\\"???\\\"  # vllm or sglang\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"## 2. Basic tool call\\n\",\n                \"For beginning, let's see how we can do basic tool call in verl with example from [Transformer tool use](https://huggingface.co/docs/transformers/main/chat_extras#tool-use). To use tool in verl, we need to define a tool class that inherits from `BaseTool`, and implement the following methods:\\n\",\n                \"- `get_openai_tool_schema`: return the schema of the tool in `OpenAIFunctionToolSchema` format.\\n\",\n                \"- `execute`: execute the tool with the given parameters, and return the result in `ToolResponse` format.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 4,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"{\\n\",\n                        \"  \\\"type\\\": \\\"function\\\",\\n\",\n                        \"  \\\"function\\\": {\\n\",\n                        \"    \\\"name\\\": \\\"get_current_temperature\\\",\\n\",\n                        \"    \\\"description\\\": \\\"Get current temperature at a location.\\\",\\n\",\n                        \"    \\\"parameters\\\": {\\n\",\n                        \"      \\\"type\\\": \\\"object\\\",\\n\",\n                        \"      \\\"properties\\\": {\\n\",\n                        \"        \\\"location\\\": {\\n\",\n                        \"          \\\"type\\\": \\\"string\\\",\\n\",\n                        \"          \\\"description\\\": \\\"The location to get the temperature for, in the format \\\\\\\"City, State, Country\\\\\\\".\\\"\\n\",\n                        \"        },\\n\",\n                        \"        \\\"unit\\\": {\\n\",\n                        \"          \\\"type\\\": \\\"string\\\",\\n\",\n                        \"          \\\"description\\\": \\\"The unit to return the temperature in. Defaults to \\\\\\\"celsius\\\\\\\".\\\",\\n\",\n                        \"          \\\"enum\\\": [\\n\",\n                        \"            \\\"celsius\\\",\\n\",\n                        \"            \\\"fahrenheit\\\"\\n\",\n                        \"          ]\\n\",\n                        \"        }\\n\",\n                        \"      },\\n\",\n                        \"      \\\"required\\\": [\\n\",\n                        \"        \\\"location\\\"\\n\",\n                        \"      ]\\n\",\n                        \"    }\\n\",\n                        \"  }\\n\",\n                        \"}\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"from transformers.utils import get_json_schema\\n\",\n                \"from verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema, ToolResponse\\n\",\n                \"\\n\",\n                \"\\n\",\n                \"class WeatherTool(BaseTool):\\n\",\n                \"    def get_current_temperature(self, location: str, unit: str = \\\"celsius\\\"):\\n\",\n                \"        \\\"\\\"\\\"Get current temperature at a location.\\n\",\n                \"\\n\",\n                \"        Args:\\n\",\n                \"            location: The location to get the temperature for, in the format \\\"City, State, Country\\\".\\n\",\n                \"            unit: The unit to return the temperature in. Defaults to \\\"celsius\\\". (choices: [\\\"celsius\\\", \\\"fahrenheit\\\"])\\n\",\n                \"\\n\",\n                \"        Returns:\\n\",\n                \"            the temperature, the location, and the unit in a dict\\n\",\n                \"        \\\"\\\"\\\"\\n\",\n                \"        return {\\n\",\n                \"            \\\"temperature\\\": 26.1,\\n\",\n                \"            \\\"location\\\": location,\\n\",\n                \"            \\\"unit\\\": unit,\\n\",\n                \"        }\\n\",\n                \"\\n\",\n                \"    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\\n\",\n                \"        schema = get_json_schema(self.get_current_temperature)\\n\",\n                \"        return OpenAIFunctionToolSchema(**schema)\\n\",\n                \"\\n\",\n                \"    async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[ToolResponse, float, dict]:\\n\",\n                \"        try:\\n\",\n                \"            result = self.get_current_temperature(**parameters)\\n\",\n                \"            return ToolResponse(text=json.dumps(result)), 0, {}\\n\",\n                \"        except Exception as e:\\n\",\n                \"            return ToolResponse(text=str(e)), 0, {}\\n\",\n                \"\\n\",\n                \"\\n\",\n                \"weather_tool = WeatherTool(config={}, tool_schema=None)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"Next, let's launch a standalone rollout server without hybrid engine (which is more heavy to start) to test the basic tool call.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": null,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"from hydra import compose, initialize_config_dir\\n\",\n                \"from verl.workers.rollout.replica import get_rollout_replica_class\\n\",\n                \"\\n\",\n                \"with initialize_config_dir(config_dir=verl_config_dir):\\n\",\n                \"    config = compose(\\n\",\n                \"        config_name=\\\"ppo_trainer\\\",\\n\",\n                \"        overrides=[\\n\",\n                \"            \\\"actor_rollout_ref.rollout.name=\\\" + rollout_name,\\n\",\n                \"            \\\"actor_rollout_ref.rollout.mode=async\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.tensor_model_parallel_size=1\\\",\\n\",\n                \"            \\\"actor_rollout_ref.model.path=\\\" + model_path,\\n\",\n                \"            \\\"actor_rollout_ref.rollout.response_length=4096\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.skip_tokenizer_init=False\\\",\\n\",\n                \"            \\\"+actor_rollout_ref.rollout.engine_kwargs.vllm.enable_auto_tool_choice=True\\\",\\n\",\n                \"            \\\"+actor_rollout_ref.rollout.engine_kwargs.vllm.tool_call_parser=hermes\\\",\\n\",\n                \"            \\\"+actor_rollout_ref.rollout.engine_kwargs.sglang.tool_call_parser=qwen25\\\",\\n\",\n                \"        ],\\n\",\n                \"    )\\n\",\n                \"\\n\",\n                \"rollout_server_class = get_rollout_replica_class(config.actor_rollout_ref.rollout.name)\\n\",\n                \"rollout_server = rollout_server_class(\\n\",\n                \"    replica_rank=0,\\n\",\n                \"    config=config.actor_rollout_ref.rollout,\\n\",\n                \"    model_config=config.actor_rollout_ref.model,\\n\",\n                \")\\n\",\n                \"\\n\",\n                \"await rollout_server.init_standalone()\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"Then, we can query LLM with openai client. Note that we need to pass the tool schema to server to guide LLM generating tool calls. We can see that the LLM correctly generates a tool call to get the temperature in Paris.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 6,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"[{'content': \\\"Hey, what's the temperature in Paris right now?\\\", 'role': 'user'},\\n\",\n                        \" {'role': 'assistant',\\n\",\n                        \"  'tool_calls': [{'function': {'arguments': '{\\\"location\\\": \\\"Paris, France\\\"}',\\n\",\n                        \"                               'name': 'get_current_temperature'},\\n\",\n                        \"                  'id': 'call_b10bdde504a0411690e96b55',\\n\",\n                        \"                  'index': -1,\\n\",\n                        \"                  'type': 'function'}]}]\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"from openai import AsyncOpenAI\\n\",\n                \"\\n\",\n                \"client = AsyncOpenAI(\\n\",\n                \"    api_key=\\\"dummy\\\",\\n\",\n                \"    base_url=f\\\"http://{rollout_server._server_address}/v1\\\",\\n\",\n                \")\\n\",\n                \"\\n\",\n                \"messages = [{\\\"role\\\": \\\"user\\\", \\\"content\\\": \\\"Hey, what's the temperature in Paris right now?\\\"}]\\n\",\n                \"completion = await client.chat.completions.create(\\n\",\n                \"    model=config.actor_rollout_ref.model.path,\\n\",\n                \"    messages=messages,\\n\",\n                \"    tools=[weather_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\\n\",\n                \"    extra_body={\\n\",\n                \"        \\\"chat_template_kwargs\\\": {\\\"enable_thinking\\\": False},\\n\",\n                \"    },\\n\",\n                \")\\n\",\n                \"\\n\",\n                \"message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\\n\",\n                \"messages.append(message)\\n\",\n                \"pprint(messages)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"We can execute the tool call with arguments generated by LLM and get the temperature in Paris.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 7,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"text='{\\\"temperature\\\": 26.1, \\\"location\\\": \\\"Paris, France\\\", \\\"unit\\\": \\\"celsius\\\"}' image=None video=None\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"args = json.loads(message[\\\"tool_calls\\\"][0][\\\"function\\\"][\\\"arguments\\\"])\\n\",\n                \"tool_response, _, _ = await weather_tool.execute(\\\"\\\", args)\\n\",\n                \"print(tool_response)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"Then, we can add the tool response to chat history and query LLM again. With the tool response, LLM can generate a final response to the user.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 8,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"[{'content': \\\"Hey, what's the temperature in Paris right now?\\\", 'role': 'user'},\\n\",\n                        \" {'role': 'assistant',\\n\",\n                        \"  'tool_calls': [{'function': {'arguments': '{\\\"location\\\": \\\"Paris, France\\\"}',\\n\",\n                        \"                               'name': 'get_current_temperature'},\\n\",\n                        \"                  'id': 'call_b10bdde504a0411690e96b55',\\n\",\n                        \"                  'index': -1,\\n\",\n                        \"                  'type': 'function'}]},\\n\",\n                        \" {'content': '{\\\"temperature\\\": 26.1, \\\"location\\\": \\\"Paris, France\\\", \\\"unit\\\": '\\n\",\n                        \"             '\\\"celsius\\\"}',\\n\",\n                        \"  'role': 'tool'},\\n\",\n                        \" {'content': 'The current temperature in Paris is 26.1°C.',\\n\",\n                        \"  'role': 'assistant'}]\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"messages.append({\\\"role\\\": \\\"tool\\\", \\\"content\\\": tool_response.text})\\n\",\n                \"completion = await client.chat.completions.create(\\n\",\n                \"    model=config.actor_rollout_ref.model.path,\\n\",\n                \"    messages=messages,\\n\",\n                \"    tools=[weather_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\\n\",\n                \"    extra_body={\\n\",\n                \"        \\\"chat_template_kwargs\\\": {\\\"enable_thinking\\\": False},\\n\",\n                \"    },\\n\",\n                \")\\n\",\n                \"\\n\",\n                \"message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\\n\",\n                \"messages.append(message)\\n\",\n                \"pprint(messages)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"## 2. Advanced tool call with code sandbox\\n\",\n                \"\\n\",\n                \"Now, let's see a more realistic example of tool call with code sandbox, which is widely used in real-world applications.\\n\",\n                \"\\n\",\n                \"### 2.1 Implement a naive code sandbox\\n\",\n                \"\\n\",\n                \"To execute python code snippet generated by LLM, we need a code sandbox environment. In this tutorial, we will implement a very naive code sandbox, which is\\n\",\n                \"a FastAPI http server with `/run_code` endpoint. The server works as follows:\\n\",\n                \"1. Receive a http request, write the python code snippet to a temp file.\\n\",\n                \"2. Spawn a subprocess to execute the code, and get stdout and stderr of the subprocess.\\n\",\n                \"3. Return the stdout and stderr of the subprocess as http response.\\n\",\n                \"\\n\",\n                \"> 🚨 **WARNING:** This naive code sandbox is for demonstration purpose only, do not use it in production. Please use docker/kata container for stronger isolation and security restriction.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 9,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"@ray.remote(num_cpus=1)\\n\",\n                \"class Sandbox:\\n\",\n                \"    \\\"\\\"\\\"Sandbox to execute python code.\\\"\\\"\\\"\\n\",\n                \"\\n\",\n                \"    def __init__(self):\\n\",\n                \"        self.address = ray._private.services.get_node_ip_address()\\n\",\n                \"        self.port = self._get_free_port()\\n\",\n                \"        asyncio.create_task(self._start_fastapi_server())\\n\",\n                \"\\n\",\n                \"    async def code_execution(self, request: Request):\\n\",\n                \"        request_json = await request.json()\\n\",\n                \"        code = request_json[\\\"code\\\"]\\n\",\n                \"        # print(f\\\"execute code:\\\\n{code}\\\")\\n\",\n                \"\\n\",\n                \"        _, temp_file = tempfile.mkstemp(suffix=\\\".py\\\", prefix=\\\"temp_code\\\", dir=None, text=True)\\n\",\n                \"        with open(temp_file, \\\"w\\\") as f:\\n\",\n                \"            f.write(code)\\n\",\n                \"\\n\",\n                \"        try:\\n\",\n                \"            process = await asyncio.create_subprocess_exec(\\n\",\n                \"                sys.executable, temp_file, stdout=asyncio.subprocess.PIPE, stderr=asyncio.subprocess.PIPE\\n\",\n                \"            )\\n\",\n                \"\\n\",\n                \"            stdout, stderr = await process.communicate()\\n\",\n                \"\\n\",\n                \"            response = {\\n\",\n                \"                \\\"status\\\": \\\"Success\\\" if process.returncode == 0 else \\\"Failed\\\",\\n\",\n                \"                \\\"run_result\\\": {\\n\",\n                \"                    \\\"status\\\": \\\"Finished\\\",\\n\",\n                \"                    \\\"stdout\\\": stdout.decode(),\\n\",\n                \"                    \\\"stderr\\\": stderr.decode(),\\n\",\n                \"                    \\\"return_code\\\": process.returncode,\\n\",\n                \"                },\\n\",\n                \"            }\\n\",\n                \"            return JSONResponse(content=response)\\n\",\n                \"        finally:\\n\",\n                \"            try:\\n\",\n                \"                os.unlink(temp_file)\\n\",\n                \"            except Exception:\\n\",\n                \"                pass\\n\",\n                \"\\n\",\n                \"    def _get_free_port(self):\\n\",\n                \"        with socket.socket() as sock:\\n\",\n                \"            sock.bind((\\\"\\\", 0))\\n\",\n                \"            return sock.getsockname()[1]\\n\",\n                \"\\n\",\n                \"    async def _start_fastapi_server(self):\\n\",\n                \"        app = fastapi.FastAPI()\\n\",\n                \"        app.router.add_api_route(\\\"/run_code\\\", self.code_execution, methods=[\\\"POST\\\"])\\n\",\n                \"\\n\",\n                \"        config = uvicorn.Config(app, host=[\\\"::\\\", \\\"0.0.0.0\\\"], port=self.port, log_level=\\\"warning\\\")\\n\",\n                \"        server = uvicorn.Server(config)\\n\",\n                \"        await server.serve()\\n\",\n                \"\\n\",\n                \"    async def get_server_address(self) -> str:\\n\",\n                \"        \\\"\\\"\\\"Get FastAPI server address.\\\"\\\"\\\"\\n\",\n                \"        return f\\\"{self.address}:{self.port}\\\"\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 10,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"sandbox = Sandbox.remote()\\n\",\n                \"sandbox_address = ray.get(sandbox.get_server_address.remote())\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"### 2.2 Define sandbox tool\\n\",\n                \"\\n\",\n                \"As shown in the previous section, we also defined a tool for the code sandbox. In the `execute` method, we send the code snippet to code sandbox by http request and get the output.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 11,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"{\\n\",\n                        \"  \\\"type\\\": \\\"function\\\",\\n\",\n                        \"  \\\"function\\\": {\\n\",\n                        \"    \\\"name\\\": \\\"code_interpreter\\\",\\n\",\n                        \"    \\\"description\\\": \\\"Execute the code in the sandbox.\\\",\\n\",\n                        \"    \\\"parameters\\\": {\\n\",\n                        \"      \\\"type\\\": \\\"object\\\",\\n\",\n                        \"      \\\"properties\\\": {\\n\",\n                        \"        \\\"code\\\": {\\n\",\n                        \"          \\\"type\\\": \\\"string\\\",\\n\",\n                        \"          \\\"description\\\": \\\"The code to be executed.\\\"\\n\",\n                        \"        }\\n\",\n                        \"      },\\n\",\n                        \"      \\\"required\\\": [\\n\",\n                        \"        \\\"code\\\"\\n\",\n                        \"      ]\\n\",\n                        \"    }\\n\",\n                        \"  }\\n\",\n                        \"}\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"import re\\n\",\n                \"import aiohttp\\n\",\n                \"\\n\",\n                \"\\n\",\n                \"class SandboxTool(BaseTool):\\n\",\n                \"    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\\n\",\n                \"        super().__init__(config, tool_schema)\\n\",\n                \"        # Different model may use different code pattern, e.g. python, py, etc.\\n\",\n                \"        self.code_pattern = re.compile(r\\\"```py(.*?)```\\\", re.DOTALL)\\n\",\n                \"\\n\",\n                \"    async def code_interpreter(self, code: str) -> str:\\n\",\n                \"        \\\"\\\"\\\"Execute the code in the sandbox.\\n\",\n                \"\\n\",\n                \"        Args:\\n\",\n                \"            code: The code to be executed.\\n\",\n                \"\\n\",\n                \"        Returns:\\n\",\n                \"            str: The output of the code execution.\\n\",\n                \"        \\\"\\\"\\\"\\n\",\n                \"        async with aiohttp.ClientSession() as session:\\n\",\n                \"            async with session.post(\\n\",\n                \"                self.config.get(\\\"sandbox_fusion_url\\\"),\\n\",\n                \"                json={\\\"code\\\": code},\\n\",\n                \"            ) as resp:\\n\",\n                \"                resp.raise_for_status()\\n\",\n                \"                result = await resp.json()\\n\",\n                \"                stdout, stderr = result[\\\"run_result\\\"][\\\"stdout\\\"], result[\\\"run_result\\\"][\\\"stderr\\\"]\\n\",\n                \"                return stdout + stderr\\n\",\n                \"\\n\",\n                \"    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\\n\",\n                \"        schema = get_json_schema(self.code_interpreter)\\n\",\n                \"        return OpenAIFunctionToolSchema(**schema)\\n\",\n                \"\\n\",\n                \"    async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[str, float, dict]:\\n\",\n                \"        code = parameters[\\\"code\\\"]\\n\",\n                \"        matches = self.code_pattern.findall(code)\\n\",\n                \"        if matches:\\n\",\n                \"            code = matches[0].strip()\\n\",\n                \"\\n\",\n                \"        # NOTE: Some script may not explicitly print result, we need to add a print statement to the end of the script.\\n\",\n                \"        # More better way is to SFT the model to make it print result by default, we skip SFT stage in this tutorial.\\n\",\n                \"        lines = code.split(\\\"\\\\n\\\")\\n\",\n                \"        for i, line in reversed(list(enumerate(lines))):\\n\",\n                \"            if line == \\\"\\\":\\n\",\n                \"                continue\\n\",\n                \"            if not lines[i].startswith(\\\"print\\\"):\\n\",\n                \"                lines[i] = f\\\"print({line})\\\"\\n\",\n                \"            break\\n\",\n                \"        code = \\\"\\\\n\\\".join(lines)\\n\",\n                \"\\n\",\n                \"        result = await self.code_interpreter(code)\\n\",\n                \"        return ToolResponse(text=result), 0.0, {}\\n\",\n                \"\\n\",\n                \"\\n\",\n                \"sandbox_tool = SandboxTool(config={\\\"sandbox_fusion_url\\\": f\\\"http://{sandbox_address}/run_code\\\"}, tool_schema=None)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"First, let's try to execute a valid code and check the response with stdout.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 12,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"(ToolResponse(text='sqrt(3)\\\\n', image=None, video=None), 0.0, {})\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"code = \\\"\\\"\\\"```py\\n\",\n                \"import sympy\\n\",\n                \"\\n\",\n                \"print(sympy.sqrt(3))\\n\",\n                \"```\\\"\\\"\\\"\\n\",\n                \"\\n\",\n                \"print(await sandbox_tool.execute(instance_id=\\\"\\\", parameters={\\\"code\\\": code}))\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"Then, let's try to execute an invalid code and check the response with stderr. The error message is important to inform LLM to fix code in next generation.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 13,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"(ToolResponse(text='Traceback (most recent call last):\\\\n  File \\\"/tmp/temp_code3e2f638_.py\\\", line 2, in <module>\\\\n    print(sympy.sqrt(3))\\\\n          ^^^^^\\\\nNameError: name \\\\'sympy\\\\' is not defined\\\\n', image=None, video=None), 0.0, {})\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"code_invalid = \\\"\\\"\\\"\\n\",\n                \"print(sympy.sqrt(3))\\n\",\n                \"\\\"\\\"\\\"\\n\",\n                \"\\n\",\n                \"print(await sandbox_tool.execute(instance_id=\\\"\\\", parameters={\\\"code\\\": code_invalid}))\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"### 2.3 Test sandbox tool\\n\",\n                \"\\n\",\n                \"Now, we can test sandbox tool with real math problem. In this tutorial, we will use the [DigitalLearningGmbH/MATH-lighteval](https://huggingface.co/datasets/DigitalLearningGmbH/MATH-lighteval) dataset, which consists of problems from mathematics competitions, including the AMC 10, AMC 12, AIME, and more.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 14,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"data\": {\n                        \"application/vnd.jupyter.widget-view+json\": {\n                            \"model_id\": \"ebd09c8816b140a59a879e5a5e218950\",\n                            \"version_major\": 2,\n                            \"version_minor\": 0\n                        },\n                        \"text/plain\": [\n                            \"Generating train split: 0 examples [00:00, ? examples/s]\"\n                        ]\n                    },\n                    \"metadata\": {},\n                    \"output_type\": \"display_data\"\n                }\n            ],\n            \"source\": [\n                \"from datasets import load_dataset\\n\",\n                \"\\n\",\n                \"dataset = load_dataset(\\\"parquet\\\", data_files=test_file)[\\\"train\\\"]\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"For debug purpose, we can implement ReAct agent as a simple loop. For RL training, there are more subtle issue and corner case to deal with, we provide a built-in ReAct agent loop which will be discussed in next section.\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 15,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stdout\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"No tool calls, finish_reason: stop\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"messages = dataset[\\\"prompt\\\"][0]\\n\",\n                \"\\n\",\n                \"while True:\\n\",\n                \"    # 1. Chat with the model\\n\",\n                \"    completion = await client.chat.completions.create(\\n\",\n                \"        model=config.actor_rollout_ref.model.path,\\n\",\n                \"        messages=messages,\\n\",\n                \"        tools=[sandbox_tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True)],\\n\",\n                \"        extra_body={\\n\",\n                \"            \\\"chat_template_kwargs\\\": {\\\"enable_thinking\\\": False},\\n\",\n                \"        },\\n\",\n                \"    )\\n\",\n                \"\\n\",\n                \"    message = completion.choices[0].message.model_dump(exclude_unset=True, exclude_none=True)\\n\",\n                \"    messages.append(message)\\n\",\n                \"\\n\",\n                \"    # 2. Call tools\\n\",\n                \"    finish_reason = completion.choices[0].finish_reason\\n\",\n                \"    if finish_reason != \\\"tool_calls\\\":\\n\",\n                \"        print(f\\\"No tool calls, finish_reason: {finish_reason}\\\")\\n\",\n                \"        break\\n\",\n                \"\\n\",\n                \"    try:\\n\",\n                \"        tool_calls = completion.choices[0].message.tool_calls[0]\\n\",\n                \"        args = json.loads(tool_calls.function.arguments)\\n\",\n                \"        result, _, _ = await sandbox_tool.execute(\\\"\\\", args)\\n\",\n                \"    except Exception as e:\\n\",\n                \"        print(f\\\"Error: {e}\\\")\\n\",\n                \"\\n\",\n                \"    # 3. Add tool response to messages\\n\",\n                \"    messages.append(\\n\",\n                \"        {\\n\",\n                \"            \\\"role\\\": \\\"tool\\\",\\n\",\n                \"            \\\"content\\\": result.text,\\n\",\n                \"        }\\n\",\n                \"    )\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 16,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"data\": {\n                        \"text/plain\": [\n                            \"[{'content': \\\"How many vertical asymptotes does the graph of $y=\\\\\\\\frac{2}{x^2+x-6}$ have? Let's think step by step and output the final answer within \\\\\\\\boxed{}.\\\",\\n\",\n                            \"  'role': 'user'},\\n\",\n                            \" {'content': \\\"To determine the number of vertical asymptotes for the function $ y = \\\\\\\\frac{2}{x^2 + x - 6} $, we need to find the values of $ x $ where the denominator equals zero, as these points are where the function is undefined and potentially where it has vertical asymptotes.\\\\n\\\\nThe denominator is $ x^2 + x - 6 $. To find the vertical asymptotes, we need to solve the equation:\\\\n\\\\n$$ x^2 + x - 6 = 0 $$\\\\n\\\\nThis is a quadratic equation, and we can solve it using the quadratic formula:\\\\n\\\\n$$ x = \\\\\\\\frac{-b \\\\\\\\pm \\\\\\\\sqrt{b^2 - 4ac}}{2a} $$\\\\n\\\\nwhere $ a = 1 $, $ b = 1 $, and $ c = -6 $. Let's solve this equation to find the values of $ x $ where the denominator is zero, which will give us the vertical asymptotes.\\\",\\n\",\n                            \"  'role': 'assistant',\\n\",\n                            \"  'tool_calls': [{'id': 'call_4d873672ff8445159e4e5e45',\\n\",\n                            \"    'function': {'arguments': '{\\\"code\\\": \\\"from sympy import symbols, solve\\\\\\\\nx = symbols(\\\\'x\\\\')\\\\\\\\nroots = solve(x**2 + x - 6, x)\\\\\\\\nroots\\\"}',\\n\",\n                            \"     'name': 'code_interpreter'},\\n\",\n                            \"    'type': 'function',\\n\",\n                            \"    'index': -1}]},\\n\",\n                            \" {'role': 'tool', 'content': '[-3, 2]\\\\n'},\\n\",\n                            \" {'content': 'The roots of the equation $ x^2 + x - 6 = 0 $ are $ x = -3 $ and $ x = 2 $. These are the values of $ x $ where the denominator is zero, which means the function $ y = \\\\\\\\frac{2}{x^2 + x - 6} $ is undefined at these points. \\\\n\\\\nSince the denominator is zero at these values, the function has vertical asymptotes at $ x = -3 $ and $ x = 2 $. Therefore, the graph of the function has two vertical asymptotes.\\\\n\\\\nThe final answer is $\\\\\\\\boxed{2}$.',\\n\",\n                            \"  'role': 'assistant'}]\"\n                        ]\n                    },\n                    \"execution_count\": 16,\n                    \"metadata\": {},\n                    \"output_type\": \"execute_result\"\n                }\n            ],\n            \"source\": [\n                \"messages\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"We can see that the ReAct agent properly query LLM, execute sandbox tool call, finally generate the answer.\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"## 3. End-to-end training with tool agent loop\\n\",\n                \"\\n\",\n                \"After tool has been implemented and tested, we can do end-to-end RL training to tune the model to properly use the tool. To simplify agentic RL training, verl provide [Agent Loop](https://verl.readthedocs.io/en/latest/advance/agent_loop.html) abstraction, which allow user to define custom agent loop:\\n\",\n                \"- Search agent\\n\",\n                \"- Math agent\\n\",\n                \"- SWE agent\\n\",\n                \"- GUI agent\\n\",\n                \"- ...\\n\",\n                \"\\n\",\n                \"For ease of use, verl provide two pre-defined agent loop:\\n\",\n                \"- SingleTurnAgentLoop: single-turn conversation without tool calling\\n\",\n                \"- ToolAgentLoop: multi-turn conversation with tool calling, interaction\\n\",\n                \"\\n\",\n                \"To use ToolAgentLoop, user only need to provide tools configuration in json/yaml file. In the configuration file, user should specify following fields for each tool:\\n\",\n                \"- class_name: fully qualified class name of the tool used to dynamically load the custom tool class\\n\",\n                \"- config: key-word arguments used to initialize the tool instance\\n\",\n                \"\\n\",\n                \"Let's dump our sandbox tool configuration to a json file:\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 17,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stderr\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"2025-10-16 23:07:16,868\\tINFO worker.py:2004 -- Started a local Ray instance. View the dashboard at \\u001b[1m\\u001b[32mhttp://127.0.0.1:8265 \\u001b[39m\\u001b[22m\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"ray.shutdown()\\n\",\n                \"\\n\",\n                \"sandbox = Sandbox.remote()\\n\",\n                \"sandbox_address = ray.get(sandbox.get_server_address.remote())\\n\",\n                \"\\n\",\n                \"tool_config = {\\n\",\n                \"    \\\"tools\\\": [\\n\",\n                \"        {\\n\",\n                \"            \\\"class_name\\\": \\\"sandbox.SandboxTool\\\",\\n\",\n                \"            \\\"config\\\": {\\n\",\n                \"                \\\"type\\\": \\\"native\\\",\\n\",\n                \"                \\\"sandbox_fusion_url\\\": f\\\"http://{sandbox_address}/run_code\\\",\\n\",\n                \"            },\\n\",\n                \"        },\\n\",\n                \"    ],\\n\",\n                \"}\\n\",\n                \"\\n\",\n                \"tool_config_path = \\\"tool_config.json\\\"\\n\",\n                \"with open(tool_config_path, \\\"w\\\") as f:\\n\",\n                \"    json.dump(tool_config, f)\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": 18,\n            \"metadata\": {},\n            \"outputs\": [\n                {\n                    \"name\": \"stderr\",\n                    \"output_type\": \"stream\",\n                    \"text\": [\n                        \"/tmp/ipykernel_174199/3963810189.py:3: UserWarning: \\n\",\n                        \"The version_base parameter is not specified.\\n\",\n                        \"Please specify a compatability version level, or None.\\n\",\n                        \"Will assume defaults for version 1.1\\n\",\n                        \"  with initialize_config_dir(config_dir=verl_config_dir):\\n\"\n                    ]\n                }\n            ],\n            \"source\": [\n                \"from hydra import compose, initialize_config_dir\\n\",\n                \"\\n\",\n                \"with initialize_config_dir(config_dir=verl_config_dir):\\n\",\n                \"    config = compose(\\n\",\n                \"        config_name=\\\"ppo_trainer\\\",\\n\",\n                \"        overrides=[\\n\",\n                \"            \\\"algorithm.adv_estimator=grpo\\\",\\n\",\n                \"            \\\"data.train_files=\\\" + train_file,\\n\",\n                \"            \\\"data.val_files=\\\" + test_file,\\n\",\n                \"            \\\"data.return_raw_chat=True\\\",\\n\",\n                \"            \\\"data.train_batch_size=32\\\",\\n\",\n                \"            \\\"data.max_prompt_length=1024\\\",\\n\",\n                \"            \\\"data.max_response_length=1024\\\",\\n\",\n                \"            \\\"+data.apply_chat_template_kwargs.enable_thinking=False\\\",\\n\",\n                \"            # actor related\\n\",\n                \"            \\\"actor_rollout_ref.model.path=\\\" + model_path,\\n\",\n                \"            \\\"actor_rollout_ref.actor.ppo_mini_batch_size=8\\\",\\n\",\n                \"            \\\"actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8\\\",\\n\",\n                \"            \\\"actor_rollout_ref.actor.fsdp_config.param_offload=True\\\",\\n\",\n                \"            \\\"actor_rollout_ref.actor.fsdp_config.optimizer_offload=True\\\",\\n\",\n                \"            # rollout related\\n\",\n                \"            \\\"actor_rollout_ref.rollout.name=\\\" + rollout_name,\\n\",\n                \"            \\\"actor_rollout_ref.rollout.mode=async\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.tensor_model_parallel_size=1\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.n=8\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.multi_turn.tool_config_path=\\\" + tool_config_path,\\n\",\n                \"            \\\"actor_rollout_ref.rollout.agent.default_agent_loop=tool_agent\\\",\\n\",\n                \"            \\\"actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8\\\",\\n\",\n                \"            # trainer related\\n\",\n                \"            \\\"trainer.val_before_train=True\\\",\\n\",\n                \"            \\\"trainer.log_val_generations=10\\\",\\n\",\n                \"            \\\"trainer.n_gpus_per_node=8\\\",\\n\",\n                \"            \\\"trainer.test_freq=-1\\\",\\n\",\n                \"            \\\"trainer.total_training_steps=5\\\",\\n\",\n                \"            \\\"trainer.logger=['console','tensorboard', 'wandb']\\\",\\n\",\n                \"            \\\"trainer.project_name=verl\\\",\\n\",\n                \"            \\\"trainer.experiment_name=\\\" + os.path.basename(model_path),\\n\",\n                \"        ],\\n\",\n                \"    )\"\n            ]\n        },\n        {\n            \"cell_type\": \"code\",\n            \"execution_count\": null,\n            \"metadata\": {},\n            \"outputs\": [],\n            \"source\": [\n                \"from verl.trainer.main_ppo import main\\n\",\n                \"\\n\",\n                \"main(config)\"\n            ]\n        },\n        {\n            \"cell_type\": \"markdown\",\n            \"metadata\": {},\n            \"source\": [\n                \"For demo purpose, we only train 5 steps, you can verify the training process by checking wandb metrics:\\n\",\n                \"- num_turns: min/max/mean chat conversation turns in each step.\\n\",\n                \"- critic rewards: min/max/mean critic rewards in each step.\\n\",\n                \"\\n\",\n                \"For more realistic agentic RL training, please refer to our recipe:\\n\",\n                \"- [retool](https://github.com/volcengine/verl-recipe/tree/main/retool): implementation of paper [ReTool: Reinforcement Learning for Strategic Tool Use in LLMs](https://arxiv.org/abs/2504.11536)\\n\",\n                \"- [collabllm](https://github.com/volcengine/verl-recipe/tree/main/collabllm): implementation of paper [CollabLLM: From Passive Responders to Active Collaborators](https://arxiv.org/pdf/2502.00640)\\n\",\n                \"- [deepeyes](https://github.com/volcengine/verl-recipe/tree/main/deepeyes): implementation of paper [DeepEyes: Incentivizing \\\"Thinking with Images\\\" via Reinforcement Learning](https://arxiv.org/abs/2505.14362)\"\n            ]\n        }\n    ],\n    \"metadata\": {\n        \"fileId\": \"398ea641-8a51-4a0b-b64e-6b7cd6b72164\",\n        \"filePath\": \"/opt/tiger/open_verl/examples/agent_loop_tutorial.ipynb\",\n        \"kernelspec\": {\n            \"display_name\": \"Python 3\",\n            \"language\": \"python\",\n            \"name\": \"python3\"\n        },\n        \"language_info\": {\n            \"codemirror_mode\": {\n                \"name\": \"ipython\",\n                \"version\": 3\n            },\n            \"file_extension\": \".py\",\n            \"mimetype\": \"text/x-python\",\n            \"name\": \"python\",\n            \"nbconvert_exporter\": \"python\",\n            \"pygments_lexer\": \"ipython3\",\n            \"version\": \"3.12.3\"\n        }\n    },\n    \"nbformat\": 4,\n    \"nbformat_minor\": 2\n}"
  },
  {
    "path": "examples/tutorial/agent_loop_get_started/sandbox.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport re\n\nimport aiohttp\nfrom transformers.utils import get_json_schema\n\nfrom verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema, ToolResponse\n\n\nclass SandboxTool(BaseTool):\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        super().__init__(config, tool_schema)\n        # Different model may use different code pattern, e.g. python, py, etc.\n        self.code_pattern = re.compile(r\"```py(.*?)```\", re.DOTALL)\n\n    async def code_interpreter(self, code: str) -> str:\n        \"\"\"Execute the code in the sandbox.\n\n        Args:\n            code: The code to be executed.\n\n        Returns:\n            str: The output of the code execution.\n        \"\"\"\n        async with aiohttp.ClientSession() as session:\n            async with session.post(\n                self.config.get(\"sandbox_fusion_url\"),\n                json={\"code\": code},\n            ) as resp:\n                resp.raise_for_status()\n                result = await resp.json()\n                stdout, stderr = result[\"run_result\"][\"stdout\"], result[\"run_result\"][\"stderr\"]\n                return stdout + stderr\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        schema = get_json_schema(self.code_interpreter)\n        return OpenAIFunctionToolSchema(**schema)\n\n    async def execute(self, instance_id: str, parameters: dict, **kwargs) -> tuple[str, float, dict]:\n        code = parameters[\"code\"]\n        matches = self.code_pattern.findall(code)\n        if matches:\n            code = matches[0].strip()\n\n        # NOTE: Some script may not explicitly print result, we need to add a print statement to the end of the script.\n        # More better way is to SFT the model to make it print result by default, we skip SFT stage in this tutorial.\n        lines = code.split(\"\\n\")\n        for i, line in reversed(list(enumerate(lines))):\n            if line == \"\":\n                continue\n            if not lines[i].startswith(\"print\"):\n                lines[i] = f\"print({line})\"\n            break\n        code = \"\\n\".join(lines)\n\n        result = await self.code_interpreter(code)\n        return ToolResponse(text=result), 0.0, {}\n"
  },
  {
    "path": "pyproject.toml",
    "content": "# -------------------------------\n# build-system\n# -------------------------------\n[build-system]\nrequires = [\n    \"setuptools>=61.0\",\n    \"wheel\"\n]\nbuild-backend = \"setuptools.build_meta\"\n\n# -------------------------------\n# project (PEP 621 metadata)\n# -------------------------------\n[project]\nname = \"verl\"\n# We'll mark the version as \"dynamic\" because it's read from the file \"verl/version/version\" \n# (PEP 621 calls this \"dynamic version\"). \n# The actual version is specified in the [tool.setuptools.dynamic] section below.\ndynamic = [\"version\", \"dependencies\", \"optional-dependencies\", \"authors\", \"urls\"]\n\ndescription = \"verl: Volcano Engine Reinforcement Learning for LLM\"\nlicense = {text = \"Apache-2.0\"}  # Changed from file to text format\nreadme = {file = \"README.md\", content-type = \"text/markdown\"}\nrequires-python = \">=3.10\"\n\n# -------------------------------\n# tool.ruff - Linting configuration\n# -------------------------------\n[tool.ruff]\n# Note: While the formatter will attempt to format lines such that they remain within the line-length,\n# it isn't a hard upper bound, and formatted lines may exceed the line-length.\nline-length = 120\nexclude = [\"scripts/legacy_model_merger.py\"]\n\n[tool.ruff.lint]\nisort = {known-first-party = [\"verl\"]}\n# c.f. https://github.com/vllm-project/vllm/blob/ce8d6b75fc0586045df75ee1568a5b5f9957251b/pyproject.toml\nselect = [\n    # pycodestyle\n    \"E\",\n    # Pyflakes\n    \"F\",\n    # pyupgrade\n    \"UP\",\n    # flake8-bugbear\n    \"B\",\n    # isort\n    \"I\",\n    \"G\",\n]\nignore = [\n    # star imports\n    \"F405\", \"F403\",\n    # lambda expression assignment\n    \"E731\",\n    # Loop control variable not used within loop body\n    \"B007\",\n    # f-string format\n    \"UP032\",\n    # `.log()` statement uses f-string\n    \"G004\",\n    # X | None for type annotations\n    \"UP045\",\n    # deprecated import\n    \"UP035\",\n]\n\n# -------------------------------\n# tool.mypy - typechecking config\n# -------------------------------\n[tool.mypy]\npretty            = true\nignore_missing_imports = true\nexplicit_package_bases = true\nfollow_imports = \"skip\"\n\n# Blanket silence\nignore_errors = true\n\n[[tool.mypy.overrides]]\nmodule = [\n\"verl.trainer.config.algorithm\",\n\"verl.trainer.ppo.core_algos\",\n\"verl.trainer.ppo.reward\",\n\"verl.workers.reward_manager\",\n\"verl.workers.reward_manager.*\",\n]\nignore_errors = false\n\n# -------------------------------\n# tool.setuptools - Additional config\n# -------------------------------\n[tool.setuptools]\n# True means `setuptools` will attempt to include all relevant files in package_data automatically.\n# This corresponds to `include_package_data=True` in setup.py.\ninclude-package-data = true\n\n# We read the version from a file in 'verl/version/version'\n[tool.setuptools.dynamic]\nversion = {file = \"verl/version/version\"}\n\n# If you need to mimic `package_dir={'': '.'}`:\n[tool.setuptools.package-dir]\n\"\" = \".\"\n\n# If you need to include specific non-Python data (like YAML files or version file):\n# This is the rough equivalent of package_data={'': ['version/*'], 'verl': ['trainer/config/*.yaml']}\n[tool.setuptools.package-data]\nverl = [\n  \"version/*\",\n  \"trainer/config/*.yaml\",\n  \"trainer/config/*/*.yaml\",\n  \"experimental/*/config/*.yaml\",\n]\n"
  },
  {
    "path": "requirements-cuda.txt",
    "content": "flash-attn"
  },
  {
    "path": "requirements-npu.txt",
    "content": "# requirements.txt records the full set of dependencies for development\naccelerate\ncodetiming\ndatasets\ndill\nhydra-core\nnumpy<2.0.0\npandas\npeft>=0.15.2\npyarrow>=15.0.0\npybind11\npylatexenc\ntensordict>=0.8.0,<=0.10.0,!=0.9.0\nray[default]\nwandb\nmathruler\ntorchdata\neinops\nqwen_vl_utils\nhf_transfer\ntriton-ascend==3.2.0\n"
  },
  {
    "path": "requirements-test.txt",
    "content": "pytest\npre-commit\npy-spy\npytest-asyncio\npytest-rerunfailures\n"
  },
  {
    "path": "requirements.txt",
    "content": "# requirements.txt records the full set of dependencies for development\naccelerate\ncodetiming\ndatasets\ndill\nhydra-core\nliger-kernel\nnumpy<2.0.0\npandas\npeft\npyarrow>=19.0.0\npybind11\npylatexenc\npre-commit\nray[default]\ntensordict>=0.8.0,<=0.10.0,!=0.9.0\ntorchdata\ntransformers\n# vllm==0.8.4\nwandb\npackaging>=20.0\nuvicorn\nfastapi\nlatex2sympy2_extended\nmath_verify\ntensorboard\n"
  },
  {
    "path": "requirements_sglang.txt",
    "content": "# requirements.txt records the full set of dependencies for development\naccelerate\ncodetiming\ndatasets\ndill\nflash-attn\nhydra-core\nnumpy<2.0.0\npandas\npeft\npyarrow>=19.0.0\npybind11\npylatexenc\nray[default]>=2.10\ntensordict>=0.8.0,<=0.10.0,!=0.9.0\ntorchdata\ntorchvision\ntransformers\nwandb\nsglang[all]==0.5.2\nhuggingface_hub\n"
  },
  {
    "path": "scripts/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "scripts/converter_hf_to_mcore.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport os\nimport warnings\nfrom contextlib import contextmanager\nfrom importlib.metadata import version\nfrom typing import Any, Callable, ContextManager, Optional\n\nimport numpy as np\nimport torch\nimport torch.distributed as dist\n\ntry:\n    # NPU patch\n    import mindspeed.megatron_adaptor  # noqa: F401\n    from mindspeed.megatron_adaptor import repatch\nexcept ImportError:\n    repatch = None\n    pass\n\nfrom accelerate import init_empty_weights\nfrom megatron.core import dist_checkpointing\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.dist_checkpointing.mapping import ShardedTensor\nfrom megatron.core.dist_checkpointing.serialization import StrictHandling\nfrom megatron.core.models.gpt.gpt_model import ModelType\nfrom megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed\nfrom packaging.version import Version\nfrom transformers import AutoConfig\n\nfrom verl.model_merger.megatron_model_merger import get_dynamic_pipeline_shards\nfrom verl.models.mcore import hf_to_mcore_config\nfrom verl.utils.device import get_device_name, get_torch_device\nfrom verl.utils.megatron_utils import get_model\n\n\ndef _init_args():\n    \"\"\"\n    Examples:\n\n    1. single rank conversion for any model:\n        > python converter_hf_to_mcore.py --hf_model_path %{hf_model} --output_path ${output_path}\n    2. distributed conversion for DeepseekV3 671B:\n        > torchrun --nproc_per_node 1 --nnodes 4 --node_rank ${RANK} converter_hf_to_mcore.py \\\n          --hf_model_path %{hf_model} --output_path ${output_path}\n    \"\"\"\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--hf_model_path\", type=str, required=True, help=\"The path for the huggingface model\")\n    parser.add_argument(\"--output_path\", type=str, required=True, help=\"The path for the output mcore model\")\n    parser.add_argument(\"--pp_size\", type=int, default=1, help=\"pipeline model parallel size\")\n    parser.add_argument(\"--ep_size\", type=int, default=1, help=\"expert model parallel size\")\n    parser.add_argument(\"--use_cpu_initialization\", action=\"store_true\", help=\"Whether to use cpu initialization\")\n    parser.add_argument(\"--test\", action=\"store_true\", help=\"Whether to test the conversion\")\n    parser.add_argument(\"--trust_remote_code\", action=\"store_true\", help=\"Whether to trust remote code\")\n    args = parser.parse_args()\n    return args\n\n\ndef test_conversion(megatron_model_provider, tfconfig, output_path, model):\n    ########### test ###########\n    # load model\n    model_test = get_model(\n        model_provider_func=megatron_model_provider,\n        model_type=ModelType.encoder_or_decoder,\n        wrap_with_ddp=True,\n        transformer_config=tfconfig,\n    )\n    ref_state_dict = model_test[0].module.sharded_state_dict()\n    dist_checkpointing.load(ref_state_dict, output_path, strict=StrictHandling.ASSUME_OK_UNEXPECTED)\n\n    dut_state_dict = model[0].module.state_dict()\n    for name in dut_state_dict.keys():\n        if dut_state_dict[name] is None:\n            print(f\"[Warning] {name} is none in dut_state_dict\")\n            continue\n        dut_data = dut_state_dict[name].data\n        if name in ref_state_dict:\n            ref_data = ref_state_dict[name]\n            if isinstance(ref_data, ShardedTensor):\n                ref_data = ref_data.data.view(ref_data.local_shape)\n            else:\n                ref_data = ref_data.data\n            assert dut_data.shape == ref_data.shape, f\"{name=} {dut_data.shape=} {ref_data.shape=}\"\n            assert (dut_data == ref_data).all(), f\"{name} is not equal\"\n            print(f\"{name} is equal\")\n        else:\n            print(f\"[Warning] {name} is not in ref_state_dict\")\n    for name in ref_state_dict.keys():\n        if ref_state_dict[name] is None:\n            print(f\"[Warning] {name} is none in ref_state_dict\")\n            continue\n        ref_data = ref_state_dict[name]\n        if isinstance(ref_data, ShardedTensor):\n            ref_data = ref_data.data.view(ref_data.local_shape)\n        else:\n            ref_data = ref_data.data\n        if name in dut_state_dict:\n            dut_data = dut_state_dict[name].data\n            assert dut_data.shape == ref_data.shape, f\"{name=} {dut_data.shape=} {ref_data.shape=}\"\n            assert (dut_data == ref_data).all(), f\"{name} is not equal\"\n            print(f\"{name} is equal\")\n        else:\n            print(f\"[Warning] {name} is not in dut_state_dict\")\n    print(\"Conversion test passed!\")\n\n\n@torch.inference_mode()\ndef convert_checkpoint_from_transformers_to_megatron(\n    hf_model, model, hf_config, layer_start_end: Optional[tuple[int, int]] = None\n):\n    if layer_start_end is None:\n        layer_start_end = (0, len(model.decoder.layers))\n    layer_start, layer_end = layer_start_end\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    ep_rank = mpu.get_expert_model_parallel_rank()\n    ep_size = mpu.get_expert_model_parallel_world_size()\n    numel = 0\n\n    num_attention_heads = hf_config.num_attention_heads\n    num_key_value_heads = hf_config.num_key_value_heads\n    hidden_dim = hf_config.hidden_size\n    head_dim = getattr(hf_config, \"head_dim\", hidden_dim // num_attention_heads)\n    if num_attention_heads != num_key_value_heads:\n        print(\"[WARNING] Converting GQA model\")\n    has_qkv_bias = getattr(hf_config, \"qkv_bias\", False) or getattr(hf_config, \"attention_bias\", False)\n    has_share_expert = getattr(hf_config, \"shared_expert_intermediate_size\", None)\n    if pp_rank == 0:\n        numel += safe_copy(hf_model.model.embed_tokens.weight, model.embedding.word_embeddings.weight)\n\n    assert len(model.decoder.layers) == (layer_end - layer_start), (\n        f\"Expected {len(model.decoder.layers)} layers, but got {layer_end - layer_start}\"\n    )\n    for layer_idx, (layer, hf_layer) in enumerate(\n        zip(model.decoder.layers, hf_model.model.layers[layer_start:layer_end], strict=True)\n    ):\n        global_layer_idx = layer_idx + layer_start\n        numel_cur = numel\n        numel += safe_copy(hf_layer.input_layernorm.weight, layer.self_attention.linear_qkv.layer_norm_weight)\n\n        q = hf_layer.self_attn.q_proj.weight.view(\n            [num_key_value_heads, head_dim * num_attention_heads // num_key_value_heads, -1]\n        )\n        k = hf_layer.self_attn.k_proj.weight.view([num_key_value_heads, head_dim, -1])\n        v = hf_layer.self_attn.v_proj.weight.view([num_key_value_heads, head_dim, -1])\n        qkv = torch.cat([q, k, v], dim=1).view(-1, hidden_dim).contiguous()\n        numel += safe_copy(qkv, layer.self_attention.linear_qkv.weight)\n\n        if has_qkv_bias:\n            q_bias = hf_layer.self_attn.q_proj.bias.view([num_key_value_heads, -1])\n            k_bias = hf_layer.self_attn.k_proj.bias.view([num_key_value_heads, -1])\n            v_bias = hf_layer.self_attn.v_proj.bias.view([num_key_value_heads, -1])\n            qkv_bias = torch.cat([q_bias, k_bias, v_bias], dim=1).view(-1).contiguous()\n            numel += safe_copy(qkv_bias, layer.self_attention.linear_qkv.bias)\n\n        if hasattr(hf_layer.self_attn, \"q_norm\"):\n            numel += safe_copy(hf_layer.self_attn.q_norm.weight.data, layer.self_attention.q_layernorm.weight)\n            numel += safe_copy(hf_layer.self_attn.k_norm.weight.data, layer.self_attention.k_layernorm.weight)\n\n        numel += safe_copy(hf_layer.self_attn.o_proj.weight, layer.self_attention.linear_proj.weight)\n        numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight)\n\n        numel += safe_copy(hf_layer.mlp.gate.weight, layer.mlp.router.weight)\n\n        for idx, hf_expert in enumerate(hf_layer.mlp.experts):\n            num_experts = len(hf_layer.mlp.experts)\n            num_local_experts = num_experts // ep_size\n            expert_idx_start = ep_rank * num_local_experts\n            expert_idx_end = (ep_rank + 1) * num_local_experts\n            if idx < expert_idx_start or idx >= expert_idx_end:\n                continue\n            local_expert_idx = idx - expert_idx_start\n\n            fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight])\n            numel += safe_copy(fc1_weight, layer.mlp.experts.linear_fc1._parameters[f\"weight{local_expert_idx}\"])\n            numel += safe_copy(\n                hf_expert.down_proj.weight, layer.mlp.experts.linear_fc2._parameters[f\"weight{local_expert_idx}\"]\n            )\n\n        if has_share_expert:\n            numel += safe_copy(hf_layer.mlp.shared_expert_gate.weight, layer.mlp.shared_experts.gate_weight)\n            shared_fc1_weight = torch.cat(\n                [hf_layer.mlp.shared_expert.gate_proj.weight, hf_layer.mlp.shared_expert.up_proj.weight]\n            )\n            numel += safe_copy(shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight)\n            numel += safe_copy(hf_layer.mlp.shared_expert.down_proj.weight, layer.mlp.shared_experts.linear_fc2.weight)\n        print(f\"{pp_rank=} {global_layer_idx=} {layer_idx=} {numel=} numel this layer={numel - numel_cur}\")\n\n    if pp_rank == pp_size - 1:\n        numel += safe_copy(hf_model.model.norm.weight, model.decoder.final_layernorm.weight)\n        numel += safe_copy(hf_model.lm_head.weight, model.output_layer.weight)\n    return numel\n\n\ndef safe_copy(\n    src_tensor: torch.Tensor,\n    dst_tensor: torch.Tensor,\n    skip_dtype_assert: bool = False,\n):\n    if not skip_dtype_assert:\n        if src_tensor.dtype != dst_tensor.dtype:\n            raise ValueError(f\"Get source dtype {src_tensor.dtype}, but target dtype {dst_tensor.dtype}\")\n    assert src_tensor.shape == dst_tensor.shape\n    dst_tensor.data.copy_(src_tensor.data)\n    return src_tensor.numel()\n\n\n@torch.inference_mode()\ndef convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hfmodel, mgmodel, hf_config):\n    mgmodel = mgmodel.bfloat16()\n    hfmodel = hfmodel.bfloat16()\n    num_attention_heads = hf_config.num_attention_heads\n    num_query_groups = hf_config.num_key_value_heads\n    hidden_size = hf_config.hidden_size\n    head_dim = hidden_size // num_attention_heads\n\n    # 1. vision model\n    if Version(version(\"transformers\")) < Version(\"4.52.0\"):\n        print(\"Using transformers < 4.52 API to load vision model\")\n        hfvision = hfmodel.visual\n    else:\n        hfvision = hfmodel.model.visual\n    mgvision = mgmodel.vision_model\n    vision_hidden_size = mgvision.config.hidden_size\n    vision_num_query_groups = mgvision.config.num_query_groups\n    vision_head_dim = vision_hidden_size // mgvision.config.num_attention_heads\n    copied_numel = 0\n    safe_copy(hfvision.rotary_pos_emb.inv_freq, mgvision.rotary_pos_emb.inv_freq)\n    copied_numel += safe_copy(hfvision.patch_embed.proj.weight, mgvision.patch_embed.proj.weight)\n    for hfblock, mgblock in zip(hfvision.blocks, mgvision.decoder.layers, strict=True):\n        # norm1 --> linear_qkv.norm\n        copied_numel += safe_copy(hfblock.norm1.weight, mgblock.self_attention.linear_qkv.layer_norm_weight)\n        # norm2 --> mlp.linear_fc1.norm\n        copied_numel += safe_copy(hfblock.norm2.weight, mgblock.mlp.linear_fc1.layer_norm_weight)\n        # qkv --> self_attention.linear_qkv\n        converted_weight = (\n            hfblock.attn.qkv.weight.view(3, vision_num_query_groups, -1, vision_head_dim, vision_hidden_size)\n            .transpose(0, 1)\n            .flatten(1, 2)\n            .reshape(-1, vision_hidden_size)\n            .contiguous()\n        )\n        copied_numel += safe_copy(converted_weight, mgblock.self_attention.linear_qkv.weight)\n        converted_bias = (\n            hfblock.attn.qkv.bias.view(3, vision_num_query_groups, -1)\n            .transpose(0, 1)\n            .flatten(1, 2)\n            .view(-1)\n            .contiguous()\n        )\n        copied_numel += safe_copy(converted_bias, mgblock.self_attention.linear_qkv.bias)\n        # proj --> self_attention.linear_proj\n        copied_numel += safe_copy(hfblock.attn.proj.weight, mgblock.self_attention.linear_proj.weight)\n        copied_numel += safe_copy(hfblock.attn.proj.bias, mgblock.self_attention.linear_proj.bias)\n        # mlp --> mlp: gate\n        fc1_weight = torch.cat([hfblock.mlp.gate_proj.weight, hfblock.mlp.up_proj.weight])\n        fc1_bias = torch.cat([hfblock.mlp.gate_proj.bias, hfblock.mlp.up_proj.bias])\n        copied_numel += safe_copy(fc1_weight, mgblock.mlp.linear_fc1.weight)\n        copied_numel += safe_copy(fc1_bias, mgblock.mlp.linear_fc1.bias)\n        copied_numel += safe_copy(hfblock.mlp.down_proj.weight, mgblock.mlp.linear_fc2.weight)\n        copied_numel += safe_copy(hfblock.mlp.down_proj.bias, mgblock.mlp.linear_fc2.bias)\n\n    # 2. vision projector\n    hfprojector = hfvision.merger\n    mgprojector = mgvision.projection\n    copied_numel += safe_copy(hfprojector.ln_q.weight, mgvision.decoder.final_layernorm.weight)\n\n    copied_numel += safe_copy(hfprojector.mlp[0].weight, mgprojector.encoder.linear_fc1.weight)\n    copied_numel += safe_copy(hfprojector.mlp[0].bias, mgprojector.encoder.linear_fc1.bias)\n    copied_numel += safe_copy(hfprojector.mlp[2].weight, mgprojector.encoder.linear_fc2.weight)\n    copied_numel += safe_copy(hfprojector.mlp[2].bias, mgprojector.encoder.linear_fc2.bias)\n    n_params = sum([t.numel() for t in hfvision.state_dict().values()])\n    assert n_params == copied_numel, f\"n_params={n_params} != copied_numel={copied_numel}\"\n    # 3. llm [just Qwen2]\n    if Version(version(\"transformers\")) < Version(\"4.52.0\"):\n        print(\"Using transformers < 4.52 API to load llm\")\n        hfllm = hfmodel.model\n    else:\n        hfllm = hfmodel.model.language_model\n    mgllm = mgmodel.language_model\n    copied_numel = 0\n    copied_numel += safe_copy(hfllm.embed_tokens.weight, mgllm.embedding.word_embeddings.weight)\n    layermaps = zip(mgllm.decoder.layers, hfllm.layers, strict=True)\n    for mglayer, hflayer in layermaps:\n        copied_numel += safe_copy(hflayer.input_layernorm.weight, mglayer.self_attention.linear_qkv.layer_norm_weight)\n\n        q_proj_weight = hflayer.self_attn.q_proj.weight.view(num_query_groups, -1, head_dim, hidden_size)\n        k_proj_weight = hflayer.self_attn.k_proj.weight.view(num_query_groups, -1, head_dim, hidden_size)\n        v_proj_weight = hflayer.self_attn.v_proj.weight.view(num_query_groups, -1, head_dim, hidden_size)\n        qkv_proj = torch.cat([q_proj_weight, k_proj_weight, v_proj_weight], dim=1).view(-1, hidden_size).contiguous()\n        copied_numel += safe_copy(qkv_proj, mglayer.self_attention.linear_qkv.weight)\n\n        q_proj_bias = hflayer.self_attn.q_proj.bias.view(num_query_groups, -1)\n        k_proj_bias = hflayer.self_attn.k_proj.bias.view(num_query_groups, -1)\n        v_proj_bias = hflayer.self_attn.v_proj.bias.view(num_query_groups, -1)\n        qkv_bias = torch.cat([q_proj_bias, k_proj_bias, v_proj_bias], dim=1).view(-1).contiguous()\n        copied_numel += safe_copy(qkv_bias, mglayer.self_attention.linear_qkv.bias)\n        copied_numel += safe_copy(hflayer.self_attn.o_proj.weight, mglayer.self_attention.linear_proj.weight)\n\n        fc1_weight = torch.cat([hflayer.mlp.gate_proj.weight, hflayer.mlp.up_proj.weight])\n        copied_numel += safe_copy(fc1_weight, mglayer.mlp.linear_fc1.weight)\n\n        copied_numel += safe_copy(hflayer.mlp.down_proj.weight, mglayer.mlp.linear_fc2.weight)\n        copied_numel += safe_copy(hflayer.post_attention_layernorm.weight, mglayer.mlp.linear_fc1.layer_norm_weight)\n\n    copied_numel += safe_copy(hfllm.norm.weight, mgllm.decoder.final_layernorm.weight)\n    if not hf_config.tie_word_embeddings:\n        safe_copy(hfmodel.lm_head.weight, mgllm.output_layer.weight)\n\n    n_params = sum([t.numel() for t in hfllm.state_dict().values()])\n\n    assert n_params == copied_numel, f\"n_params={n_params} != copied_numel={copied_numel}\"\n\n\n@torch.inference_mode()\ndef convert_checkpoint_from_transformers_to_megatron_dpskv3(\n    hf_model,\n    model,\n    hf_config,\n    tfconfig,\n    layer_start_end: Optional[tuple[int, int]] = None,\n):\n    warnings.warn(\"MTP model is not supported yet\", stacklevel=2)\n    if layer_start_end is None:\n        layer_start_end = (0, len(model.decoder.layers))\n    layer_start, layer_end = layer_start_end\n    numel: int = 0\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    ep_rank = mpu.get_expert_model_parallel_rank()\n    ep_size = mpu.get_expert_model_parallel_world_size()\n\n    if pp_rank == 0:\n        numel += safe_copy(hf_model.model.embed_tokens.weight, model.embedding.word_embeddings.weight)\n\n    assert len(model.decoder.layers) == (layer_end - layer_start), (\n        f\"Expected {len(model.decoder.layers)} layers, but got {layer_end - layer_start}\"\n    )\n    for layer_idx, (layer, hf_layer) in enumerate(\n        zip(model.decoder.layers, hf_model.model.layers[layer_start:layer_end], strict=True)\n    ):\n        global_layer_idx = layer_idx + layer_start\n        numel_cur: int = numel\n        numel += safe_copy(hf_layer.input_layernorm.weight, layer.input_layernorm.weight)\n\n        if hf_config.q_lora_rank is None:\n            numel += safe_copy(hf_layer.self_attn.q_proj.weight, layer.self_attention.linear_q_proj.weight)\n        else:\n            numel += safe_copy(hf_layer.self_attn.q_a_proj.weight, layer.self_attention.linear_q_down_proj.weight)\n            numel += safe_copy(hf_layer.self_attn.q_b_proj.weight, layer.self_attention.linear_q_up_proj.weight)\n            numel += safe_copy(\n                hf_layer.self_attn.q_a_layernorm.weight, layer.self_attention.linear_q_up_proj.layer_norm_weight\n            )\n\n        numel += safe_copy(\n            hf_layer.self_attn.kv_a_proj_with_mqa.weight, layer.self_attention.linear_kv_down_proj.weight\n        )\n        numel += safe_copy(hf_layer.self_attn.kv_b_proj.weight, layer.self_attention.linear_kv_up_proj.weight)\n        numel += safe_copy(\n            hf_layer.self_attn.kv_a_layernorm.weight, layer.self_attention.linear_kv_up_proj.layer_norm_weight\n        )\n        numel += safe_copy(hf_layer.self_attn.o_proj.weight, layer.self_attention.linear_proj.weight)\n\n        if not hasattr(layer.mlp, \"router\"):\n            numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.mlp.linear_fc1.layer_norm_weight)\n            numel += safe_copy(\n                torch.cat([hf_layer.mlp.gate_proj.weight, hf_layer.mlp.up_proj.weight]), layer.mlp.linear_fc1.weight\n            )\n            numel += safe_copy(hf_layer.mlp.down_proj.weight, layer.mlp.linear_fc2.weight)\n        else:\n            numel += safe_copy(hf_layer.mlp.gate.weight, layer.mlp.router.weight)\n            # NOTE: the e_score_correction_bias in mcore model will be initialized with bfloat16 and \\\n            # recover to fp32 in the first forward. There is always a diff in the bias between two models (~0.3%)\n            numel += safe_copy(\n                hf_layer.mlp.gate.e_score_correction_bias, layer.mlp.router.expert_bias, skip_dtype_assert=True\n            )\n            if tfconfig.moe_grouped_gemm:\n                for i, hf_expert in enumerate(hf_layer.mlp.experts):\n                    num_experts = len(hf_layer.mlp.experts)\n                    num_local_experts = num_experts // ep_size\n                    expert_idx_start = ep_rank * num_local_experts\n                    expert_idx_end = (ep_rank + 1) * num_local_experts\n                    if i < expert_idx_start or i >= expert_idx_end:\n                        continue\n                    local_expert_idx = i - expert_idx_start\n\n                    fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight])\n                    linear_fc1_weighti = getattr(layer.mlp.experts.linear_fc1, \"weight\" + str(local_expert_idx))\n                    numel += safe_copy(fc1_weight, linear_fc1_weighti)\n                    linear_fc2_weighti = getattr(layer.mlp.experts.linear_fc2, \"weight\" + str(local_expert_idx))\n                    numel_w2 = safe_copy(hf_expert.down_proj.weight, linear_fc2_weighti)\n                    numel += numel_w2\n            else:\n                for i, hf_expert in enumerate(hf_layer.mlp.experts):\n                    expert = layer.mlp.experts.local_experts[i]\n                    fc1_weight = torch.cat([hf_expert.gate_proj.weight, hf_expert.up_proj.weight])\n                    numel += safe_copy(fc1_weight, expert.linear_fc1.weight)\n                    numel += safe_copy(hf_expert.down_proj.weight, expert.linear_fc2.weight)\n            numel += safe_copy(hf_layer.post_attention_layernorm.weight, layer.pre_mlp_layernorm.weight)\n            shared_fc1_weight = torch.cat(\n                [hf_layer.mlp.shared_experts.gate_proj.weight, hf_layer.mlp.shared_experts.up_proj.weight]\n            )\n            numel += safe_copy(shared_fc1_weight, layer.mlp.shared_experts.linear_fc1.weight)\n            numel += safe_copy(hf_layer.mlp.shared_experts.down_proj.weight, layer.mlp.shared_experts.linear_fc2.weight)\n        print(f\"{pp_rank=} {global_layer_idx=} {layer_idx=} {numel=} numel this layer={numel - numel_cur}\")\n        numel_hf_one_layer = sum([i.numel() for i in hf_layer.state_dict().values()])\n        if hasattr(layer.mlp, \"router\"):\n            numel_hf_one_layer -= numel_w2 * 3 * len(hf_layer.mlp.experts) // ep_size * (ep_size - 1)\n        assert numel - numel_cur == numel_hf_one_layer, \"numel mismatch\"\n\n    if pp_rank == pp_size - 1:\n        numel += safe_copy(hf_model.model.norm.weight, model.decoder.final_layernorm.weight)\n        if not hf_config.tie_word_embeddings:\n            numel += safe_copy(hf_model.lm_head.weight, model.output_layer.weight)\n    print(f\"{pp_rank=} {numel=}\")\n    return numel\n\n\n@contextmanager\ndef noop_context() -> Any:\n    yield\n\n\ndef support_distributed_convert(hf_config: AutoConfig) -> bool:\n    for arch in [\"DeepseekV3ForCausalLM\", \"Qwen3MoeForCausalLM\", \"Qwen2MoeForCausalLM\"]:\n        if arch in hf_config.architectures:\n            return True\n    return False\n\n\ndef convert_hf_to_mcore(\n    hf_model_path, output_path, pp_size=1, ep_size=1, use_cpu_initialization=False, test=False, trust_remote_code=False\n):\n    os.makedirs(output_path, exist_ok=True)\n    if len(os.listdir(output_path)) > 0 and not test:\n        print(f\"Output path {output_path} is not empty, skipping conversion\")\n        return\n\n    # init torch distributed and mpu\n    if \"WORLD_SIZE\" not in os.environ:\n        os.environ[\"RANK\"] = \"0\"\n        os.environ[\"WORLD_SIZE\"] = \"1\"\n        os.environ[\"MASTER_ADDR\"] = \"localhost\"\n        os.environ[\"MASTER_PORT\"] = \"12355\"\n\n    torch.distributed.init_process_group(\"nccl\")\n\n    local_rank = os.getenv(\"LOCAL_RANK\", 0)\n    world_size = dist.get_world_size()\n    get_torch_device().set_device(f\"{get_device_name()}:{local_rank}\")\n    if ep_size * pp_size != world_size:\n        pp_size = world_size\n        print(f\"pp_size is set to {pp_size}\")\n\n    mpu.initialize_model_parallel(\n        tensor_model_parallel_size=1,\n        pipeline_model_parallel_size=pp_size,\n        virtual_pipeline_model_parallel_size=None,\n        context_parallel_size=1,\n        expert_model_parallel_size=ep_size,\n    )\n    model_parallel_cuda_manual_seed(0)\n\n    # init hf config\n    hf_config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code)\n    print(hf_config, flush=True)\n\n    if repatch:\n        if hf_config.architectures[0] == \"DeepseekV3ForCausalLM\":\n            config_repatch = dict(multi_head_latent_attention=True)\n            repatch(config_repatch)\n\n    if world_size > 1 and not support_distributed_convert(hf_config):\n        raise NotImplementedError(f\"distributed conversion is not supported for {hf_config.architectures} yet.\")\n\n    pipeline_shards = get_dynamic_pipeline_shards(hf_config.num_hidden_layers, pp_size)\n    print(f\"Pipeline shards: {pipeline_shards}\", flush=True)\n\n    tfconfig = hf_to_mcore_config(\n        hf_config,\n        torch.bfloat16,\n        num_layers_in_first_pipeline_stage=pipeline_shards[0] if len(pipeline_shards) > 1 else None,\n        num_layers_in_last_pipeline_stage=pipeline_shards[-1] if len(pipeline_shards) > 2 else None,\n    )\n    tfconfig.use_cpu_initialization = use_cpu_initialization\n    tie_word_embeddings = getattr(hf_config, \"tie_word_embeddings\", False)\n\n    # init megatron model\n    def megatron_model_provider(pre_process, post_process):\n        from verl.models.mcore import init_mcore_model\n\n        parallel_model = init_mcore_model(\n            tfconfig,\n            hf_config,\n            pre_process,\n            post_process,\n            share_embeddings_and_output_weights=tie_word_embeddings,\n            value=False,\n        )\n        return parallel_model\n\n    context: Callable[..., ContextManager] = init_empty_weights if use_cpu_initialization else noop_context\n    with context():\n        model = get_model(\n            model_provider_func=megatron_model_provider,\n            model_type=ModelType.encoder_or_decoder,\n            wrap_with_ddp=False,\n            transformer_config=tfconfig,\n        )\n\n    if use_cpu_initialization:\n        # convert meta device to empty tensor so it can use `copy_` function\n        model[0].module = model[0].module.to_empty(device=\"cpu\")\n\n    with warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\")\n    from transformers import AutoModelForCausalLM, AutoModelForImageTextToText\n\n    # init hf model\n    if \"Qwen2_5_VLForConditionalGeneration\" in hf_config.architectures:\n        hf_model = AutoModelForImageTextToText.from_pretrained(\n            hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code\n        )\n    else:\n        hf_model = AutoModelForCausalLM.from_pretrained(\n            hf_model_path, torch_dtype=torch.bfloat16, trust_remote_code=trust_remote_code\n        )\n    hf_state_dict = hf_model.state_dict()\n\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n\n    # distributed convert\n    if world_size > 1 and support_distributed_convert(hf_config):\n        pipeline_cumsum = np.cumsum(pipeline_shards)\n        layer_start = 0 if pp_rank == 0 else pipeline_cumsum[pp_rank - 1]\n        layer_end = pipeline_cumsum[pp_rank]\n        if \"DeepseekV3ForCausalLM\" in hf_config.architectures:\n            numel_partial: int = convert_checkpoint_from_transformers_to_megatron_dpskv3(\n                hf_model, model[0].module, hf_config, tfconfig=tfconfig, layer_start_end=(layer_start, layer_end)\n            )\n        elif \"Qwen3MoeForCausalLM\" in hf_config.architectures or \"Qwen2MoeForCausalLM\" in hf_config.architectures:\n            numel_partial: int = convert_checkpoint_from_transformers_to_megatron(\n                hf_model, model[0].module, hf_config, layer_start_end=(layer_start, layer_end)\n            )\n        else:\n            raise NotImplementedError(f\"Distributed conversion is not supported for {hf_config.architectures} yet.\")\n\n        numel_tensor = torch.tensor([numel_partial]).to(get_device_name())\n        dist.all_reduce(numel_tensor, op=dist.ReduceOp.SUM)\n        numel = int(numel_tensor.cpu().item())\n        print(f\"total numel={numel} vs {hf_model.num_parameters()=}\")\n        if numel != hf_model.num_parameters():\n            warnings.warn(f\"numel mismatch: {numel=} != {hf_model.num_parameters()=}\", stacklevel=1)\n\n    # load hf state dict to megatron model\n    elif \"Qwen2MoeForCausalLM\" in hf_config.architectures:\n        convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config)\n    elif \"Qwen2_5_VLForConditionalGeneration\" in hf_config.architectures:\n        convert_checkpoint_from_transformers_to_megatron_qwen2_5_vl(hf_model, model[0].module, hf_config)\n    elif \"DeepseekV3ForCausalLM\" in hf_config.architectures:\n        convert_checkpoint_from_transformers_to_megatron_dpskv3(hf_model, model[0].module, hf_config, tfconfig=tfconfig)\n    elif \"Qwen3MoeForCausalLM\" in hf_config.architectures:\n        convert_checkpoint_from_transformers_to_megatron(hf_model, model[0].module, hf_config)\n    else:\n        assert not use_cpu_initialization, \"use_cpu_initialization is only supported for MoE model\"\n        from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel\n\n        load_state_dict_to_megatron_gptmodel(\n            state_dict=hf_state_dict,\n            wrapped_models=model,\n            config=hf_config,\n            params_dtype=torch.bfloat16,\n            is_value_model=False,\n        )\n\n    megatron_state_dict = model[0].module.sharded_state_dict()\n    del hf_state_dict, hf_model\n\n    # save megatron model\n    if len(os.listdir(output_path)) == 0:\n        dist_checkpointing.save(megatron_state_dict, output_path, sharded_strategy=None, async_sharded_save=False)\n    if test:\n        test_conversion(megatron_model_provider, tfconfig, output_path, model)\n\n\nif __name__ == \"__main__\":\n    args = _init_args()\n    convert_hf_to_mcore(\n        args.hf_model_path,\n        args.output_path,\n        args.pp_size,\n        args.ep_size,\n        args.use_cpu_initialization,\n        args.test,\n        args.trust_remote_code,\n    )\n"
  },
  {
    "path": "scripts/diagnose.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Diagnose script for checking OS/hardware/python/pip/verl/network.\nThe output of this script can be a very good hint to issue/problem.\n\"\"\"\n\nimport os\nimport platform\nimport socket\nimport subprocess\nimport sys\nimport time\n\nimport psutil\n\ntry:\n    from urllib.parse import urlparse\n    from urllib.request import urlopen\nexcept ImportError:\n    from urllib2 import urlopen\n    from urlparse import urlparse\nimport argparse\nimport importlib.metadata\n\nimport torch\n\nURLS = {\n    \"PYPI\": \"https://pypi.python.org/pypi/pip\",\n}\n\nREGIONAL_URLS = {\n    \"cn\": {\n        \"PYPI(douban)\": \"https://pypi.douban.com/\",\n        \"Conda(tsinghua)\": \"https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/free/\",\n    }\n}\n\n\ndef test_connection(name, url, timeout=10):\n    \"\"\"Simple connection test\"\"\"\n    urlinfo = urlparse(url)\n    start = time.time()\n    try:\n        socket.gethostbyname(urlinfo.netloc)\n    except Exception as e:\n        print(\"Error resolving DNS for {}: {}, {}\".format(name, url, e))\n        return\n    dns_elapsed = time.time() - start\n    start = time.time()\n    try:\n        _ = urlopen(url, timeout=timeout)\n    except Exception as e:\n        print(\"Error open {}: {}, {}, DNS finished in {} sec.\".format(name, url, e, dns_elapsed))\n        return\n    load_elapsed = time.time() - start\n    print(\"Timing for {}: {}, DNS: {:.4f} sec, LOAD: {:.4f} sec.\".format(name, url, dns_elapsed, load_elapsed))\n\n\ndef check_python():\n    print(\"----------Python Info----------\")\n    print(\"Version      :\", platform.python_version())\n    print(\"Compiler     :\", platform.python_compiler())\n    print(\"Build        :\", platform.python_build())\n    print(\"Arch         :\", platform.architecture())\n\n\ndef check_pip():\n    print(\"------------Pip Info-----------\")\n    try:\n        import pip\n\n        print(\"Version      :\", pip.__version__)\n        print(\"Directory    :\", os.path.dirname(pip.__file__))\n    except ImportError:\n        print(\"No corresponding pip install for current python.\")\n\n\ndef _get_current_git_commit():\n    try:\n        result = subprocess.run([\"git\", \"rev-parse\", \"HEAD\"], capture_output=True, text=True, check=True)\n        return result.stdout.strip()\n    except subprocess.CalledProcessError as e:\n        print(f\"Error running git command: {e.stderr.strip()}\")\n        return None\n    except FileNotFoundError:\n        print(\"Did not find command: git\")\n        return None\n\n\ndef check_verl():\n    print(\"----------verl Info-----------\")\n    try:\n        sys.path.insert(0, os.getcwd())\n        import verl\n\n        print(\"Version      :\", verl.__version__)\n        verl_dir = os.path.dirname(verl.__file__)\n        print(\"Directory    :\", verl_dir)\n        try:\n            commit_hash = _get_current_git_commit()\n            print(\"Commit Hash  :\", commit_hash)\n        except AttributeError:\n            print(\"Commit hash not found. \")\n    except ImportError as e:\n        print(f\"No verl installed: {e}\")\n    except Exception as e:\n        import traceback\n\n        if not isinstance(e, IOError):\n            print(\"An error occurred trying to import verl.\")\n            print(\"This is very likely due to missing or incompatible library files.\")\n        print(traceback.format_exc())\n\n\ndef check_os():\n    print(\"----------Platform Info----------\")\n    print(\"Platform     :\", platform.platform())\n    print(\"system       :\", platform.system())\n    print(\"node         :\", platform.node())\n    print(\"release      :\", platform.release())\n    print(\"version      :\", platform.version())\n\n\ndef check_hardware():\n    print(\"----------Hardware Info----------\")\n    print(\"machine      :\", platform.machine())\n    print(\"processor    :\", platform.processor())\n    if sys.platform.startswith(\"darwin\"):\n        pipe = subprocess.Popen((\"sysctl\", \"-a\"), stdout=subprocess.PIPE)\n        output = pipe.communicate()[0]\n        for line in output.split(b\"\\n\"):\n            if b\"brand_string\" in line or b\"features\" in line:\n                print(line.strip())\n    elif sys.platform.startswith(\"linux\"):\n        subprocess.call([\"lscpu\"])\n    elif sys.platform.startswith(\"win32\"):\n        subprocess.call([\"wmic\", \"cpu\", \"get\", \"name\"])\n\n\ndef check_network(args):\n    print(\"----------Network Test----------\")\n    if args.timeout > 0:\n        print(\"Setting timeout: {}\".format(args.timeout))\n        socket.setdefaulttimeout(10)\n    for region in args.region.strip().split(\",\"):\n        r = region.strip().lower()\n        if not r:\n            continue\n        if r in REGIONAL_URLS:\n            URLS.update(REGIONAL_URLS[r])\n        else:\n            import warnings\n\n            warnings.warn(\"Region {} do not need specific test, please refer to global sites.\".format(r), stacklevel=2)\n    for name, url in URLS.items():\n        test_connection(name, url, args.timeout)\n\n\ndef check_environment():\n    print(\"----------Environment----------\")\n    for k, v in os.environ.items():\n        if k.startswith(\"VERL_\") or k.startswith(\"OMP_\") or k.startswith(\"KMP_\") or k == \"CC\" or k == \"CXX\":\n            print('{}=\"{}\"'.format(k, v))\n\n\ndef check_pip_package_versions():\n    packages = [\"vllm\", \"sglang\", \"ray\", \"torch\"]\n    for package in packages:\n        try:\n            version = importlib.metadata.version(package)\n            print(f\"{package}\\t     : {version}\")\n        except importlib.metadata.PackageNotFoundError:\n            print(f\"{package}\\t     : not found.\")\n\n\ndef check_cuda_versions():\n    if torch.cuda.is_available():\n        try:\n            cuda_runtime_version = torch.version.cuda\n            print(f\"CUDA Runtime : {cuda_runtime_version}\")\n            import subprocess\n\n            nvcc_output = subprocess.check_output([\"nvcc\", \"--version\"]).decode(\"utf-8\")\n            cuda_compiler_version = next((line for line in nvcc_output.splitlines() if \"release\" in line), None)\n            if cuda_compiler_version:\n                print(f\"CUDA Compiler : {cuda_compiler_version.strip()}\")\n            else:\n                print(\"Could not determine CUDA compiler version.\")\n        except FileNotFoundError as e:\n            print(f\"CUDA compiler : Not found: {e}\")\n        except Exception as e:\n            print(f\"An error occurred while checking CUDA versions: {e}\")\n    else:\n        print(\"CUDA is not available.\")\n\n\ndef _get_cpu_memory():\n    \"\"\"\n    Get the total CPU memory capacity in GB.\n    \"\"\"\n    memory = psutil.virtual_memory()\n    return memory.total / (1024**3)\n\n\ndef _get_gpu_info():\n    \"\"\"\n    Get GPU type, GPU memory, and GPU count using nvidia-smi command.\n    \"\"\"\n    try:\n        result = subprocess.run(\n            [\"nvidia-smi\", \"--query-gpu=gpu_name,memory.total\", \"--format=csv,noheader,nounits\"],\n            capture_output=True,\n            text=True,\n            check=True,\n        )\n        gpu_lines = result.stdout.strip().split(\"\\n\")\n        gpu_count = len(gpu_lines)\n        gpu_info = []\n        for line in gpu_lines:\n            gpu_name, gpu_memory = line.split(\", \")\n            gpu_info.append(\n                {\n                    \"type\": gpu_name,\n                    \"memory\": float(gpu_memory) / 1024,  # Convert to GB\n                }\n            )\n        return gpu_count, gpu_info\n    except (subprocess.CalledProcessError, FileNotFoundError):\n        print(\"Failed to execute nvidia-smi command.\")\n        return 0, []\n\n\ndef _get_system_info():\n    \"\"\"\n    Get CPU memory capacity, GPU type, GPU memory, and GPU count.\n    \"\"\"\n    cpu_memory = _get_cpu_memory()\n    gpu_count, gpu_info = _get_gpu_info()\n    return {\"cpu_memory\": cpu_memory, \"gpu_count\": gpu_count, \"gpu_info\": gpu_info}\n\n\ndef check_system_info():\n    print(\"----------System Info----------\")\n    system_info = _get_system_info()\n    print(f\"CPU Memory\\t: {system_info['cpu_memory']:.2f} GB\")\n    print(f\"GPU Count\\t: {system_info['gpu_count']}\")\n    for i, gpu in enumerate(system_info[\"gpu_info\"]):\n        print(f\"GPU {i + 1}\\tType    : {gpu['type']}\")\n        print(f\"GPU {i + 1}\\tMemory  : {gpu['memory']:.2f} GB\")\n\n\ndef parse_args():\n    \"\"\"Parse arguments.\"\"\"\n    parser = argparse.ArgumentParser(\n        formatter_class=argparse.ArgumentDefaultsHelpFormatter,\n        description=\"Diagnose script for checking the current system.\",\n    )\n    choices = [\"python\", \"pip\", \"verl\", \"system\", \"os\", \"environment\"]\n    for choice in choices:\n        parser.add_argument(\"--\" + choice, default=1, type=int, help=\"Diagnose {}.\".format(choice))\n    parser.add_argument(\"--network\", default=0, type=int, help=\"Diagnose network.\")\n    parser.add_argument(\"--hardware\", default=0, type=int, help=\"Diagnose hardware.\")\n    parser.add_argument(\n        \"--region\",\n        default=\"\",\n        type=str,\n        help=\"Additional sites in which region(s) to test. \\\n                        Specify 'cn' for example to test mirror sites in China.\",\n    )\n    parser.add_argument(\"--timeout\", default=10, type=int, help=\"Connection test timeout threshold, 0 to disable.\")\n    args = parser.parse_args()\n    return args\n\n\nif __name__ == \"__main__\":\n    args = parse_args()\n    if args.python:\n        check_python()\n\n    if args.pip:\n        check_pip()\n        check_pip_package_versions()\n\n    if args.verl:\n        check_verl()\n\n    if args.os:\n        check_os()\n\n    if args.hardware:\n        check_hardware()\n\n    if args.network:\n        check_network(args)\n\n    if args.environment:\n        check_environment()\n        check_cuda_versions()\n\n    if args.system:\n        check_system_info()\n"
  },
  {
    "path": "scripts/generate_trainer_config.sh",
    "content": "#!/usr/bin/env bash\nset -euox pipefail\n\n\n# Define config specifications: \"config_name:output_file:config_arg\"\nCONFIG_SPECS=(\n    \"ppo_trainer:_generated_ppo_trainer.yaml:\"\n    \"ppo_megatron_trainer:_generated_ppo_megatron_trainer.yaml:--config-name=ppo_megatron_trainer.yaml\"\n    \"ppo_trainer:_generated_ppo_veomni_trainer.yaml:model_engine=veomni\"\n    \"ppo_trainer:_generated_ppo_torchtitan_trainer.yaml:model_engine=torchtitan\"\n)\n\ngenerate_config() {\n    local config_name=\"$1\"\n    local output_file=\"$2\"\n    local config_arg=\"$3\"\n\n    local target_cfg=\"verl/trainer/config/${output_file}\"\n    local tmp_header=$(mktemp)\n    local tmp_cfg=$(mktemp)\n\n    echo \"# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'\" > \"$tmp_header\"\n    echo \"# in which it invokes 'python3 scripts/print_cfg.py --cfg job ${config_arg}' to flatten the 'verl/trainer/config/${config_name}.yaml' config fields into a single file.\" >> \"$tmp_header\"\n    echo \"# Do not modify this file directly.\" >> \"$tmp_header\"\n    echo \"# The file is usually only for reference and never used.\" >> \"$tmp_header\"\n    echo \"\" >> \"$tmp_header\"\n\n    python3 scripts/print_cfg.py --cfg job ${config_arg} > \"$tmp_cfg\"\n\n    cat \"$tmp_header\" > \"$target_cfg\"\n    sed -n '/^actor_rollout_ref/,$p' \"$tmp_cfg\" >> \"$target_cfg\"\n\n    rm \"$tmp_cfg\" \"$tmp_header\"\n\n    echo \"Generated: $target_cfg\"\n}\n\nfor spec in \"${CONFIG_SPECS[@]}\"; do\n    IFS=':' read -r config_name output_file config_arg <<< \"$spec\"\n    generate_config \"$config_name\" \"$output_file\" \"$config_arg\"\ndone\n\nfor spec in \"${CONFIG_SPECS[@]}\"; do\n    IFS=':' read -r config_name output_file config_arg <<< \"$spec\"\n    target_cfg=\"verl/trainer/config/${output_file}\"\n    if ! git diff --exit-code -- \"$target_cfg\" >/dev/null; then\n        echo \"✖ $target_cfg is out of date. Please regenerate via 'scripts/generate_trainer_config.sh' and commit the changes.\"\n        exit 1\n    fi\ndone\n\necho \"All good\"\nexit 0\n"
  },
  {
    "path": "scripts/init_random_model.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script override a model with custom config and random weights, mainly for create small models for \ndebugging purposes.\n\nUsage:\n    python scripts/init_random_model.py \\\n        --hf_model_path <path_to_hf_model> \\\n        --new_config_path <path_to_new_config.json> \\\n        --output_path <path_to_output_model>\n\n\"\"\"\n\nimport argparse\nimport json\nimport os\nimport warnings\nfrom typing import Any\n\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer, PretrainedConfig\n\n\ndef _init_args():\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--hf_model_path\", type=str, required=True, help=\"The path for the huggingface model\")\n    parser.add_argument(\"--new_config_path\", type=str, required=True, help=\"The path for the new config file\")\n    parser.add_argument(\"--output_path\", type=str, required=True, help=\"The path for the output random model\")\n    parser.add_argument(\n        \"--trust_remote_code\",\n        action=\"store_true\",\n        help=\"Whether to trust remote code when loading HF model. Disabled by default for security.\",\n    )\n    args = parser.parse_args()\n    return args\n\n\ndef check_output_path(output_path: str):\n    if os.path.exists(output_path):\n        warnings.warn(f\"Output path '{output_path}' already exists. Will do nothing.\", stacklevel=2)\n        exit()\n    else:\n        os.makedirs(output_path, exist_ok=True)\n        print(f\"Output path '{output_path}' created.\")\n\n\ndef check_configs(original_config: dict[str, Any], new_config: dict[str, Any]) -> bool:\n    \"\"\"\n    Check if the original config and new config are compatible.\n    This is a placeholder function; actual implementation may vary based on requirements.\n    \"\"\"\n    # Example check: ensure 'model_type' is the same\n    if new_config.get(\"model_type\", None) is not None and original_config.get(\"model_type\") != new_config.get(\n        \"model_type\"\n    ):\n        raise RuntimeError(\"Model types do not match.\")\n    for key in new_config:\n        if key not in original_config:\n            warnings.warn(\n                f\"Key '{key}' in new config does not exist in original config, may not take effect.\", stacklevel=2\n            )\n\n\ndef init_random_model(hf_model_path, new_config_path, output_path, trust_remote_code: bool = False):\n    config = AutoConfig.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code)\n    tokenizer = AutoTokenizer.from_pretrained(hf_model_path, trust_remote_code=trust_remote_code)\n    config_dict = PretrainedConfig.get_config_dict(hf_model_path)[0]\n    print(config_dict)\n    with open(new_config_path) as f:\n        new_config_dict = json.load(f)\n    check_configs(config_dict, new_config_dict)\n    config_dict.update(new_config_dict)\n    new_confg = config.from_dict(config_dict)\n    print(f\"new_config: {new_confg}\")\n    if trust_remote_code:\n        model = AutoModelForCausalLM.from_pretrained(\n            hf_model_path, config=new_confg, trust_remote_code=trust_remote_code, torch_dtype=new_confg.torch_dtype\n        )\n    else:\n        model = AutoModelForCausalLM.from_config(new_confg, torch_dtype=new_confg.torch_dtype)\n    model.save_pretrained(output_path)\n    tokenizer.save_pretrained(output_path)\n    new_confg.save_pretrained(output_path)\n    print(f\"Random model initialized and saved to {output_path}\")\n\n\nif __name__ == \"__main__\":\n    args = _init_args()\n    check_output_path(args.output_path)\n    init_random_model(\n        hf_model_path=args.hf_model_path,\n        new_config_path=args.new_config_path,\n        output_path=args.output_path,\n        trust_remote_code=args.trust_remote_code,\n    )\n"
  },
  {
    "path": "scripts/install_sglang_mcore_npu.sh",
    "content": "#!/bin/bash\nset -e\nNPU_DEVICE=${NPU_DEVICE:=A3}\nUSE_MEGATRON=${USE_MEGATRON:-1}\n\nexport MAX_JOBS=32\n\necho \"1. install SGLang from source\"\ngit clone -b v0.5.8 https://github.com/sgl-project/sglang.git\ncd sglang\nmv python/pyproject_other.toml python/pyproject.toml\npip install -e python[srt_npu]\ncd ..\n\necho \"2. install torch & torch_npu & triton_ascend & other basic packages\"\npip install torch==2.7.1 torch_npu==2.7.1.post2 torchvision==0.22.1\npip install pybind11 click==8.2.1 mbridge \"numpy<2.0.0\" cachetools\n\n\necho \"3. install sgl-kernel-npu form source, detailed readme in https://github.com/sgl-project/sgl-kernel-npu/blob/main/python/deep_ep/README.md\"\ngit clone https://github.com/sgl-project/sgl-kernel-npu.git\ncd sgl-kernel-npu\ngit checkout 46b73de\nsed -i '101s/^/# /' build.sh\nif [ \"$NPU_DEVICE\" = \"A3\" ]; then\n    bash build.sh\nfi\nif [ \"$NPU_DEVICE\" = \"A2\" ]; then\n    bash build.sh -a deepep2\nfi\npip install output/torch_memory_saver*.whl\npip install output/sgl_kernel_npu*.whl\npip install output/deep_ep*.whl\ncd \"$(pip show deep-ep | grep -E '^Location:' | awk '{print $2}')\" && ln -s deep_ep/deep_ep_cpp*.so && cd -\npython -c \"import deep_ep; print(deep_ep.__path__)\"\ncd ..\n# install sgl-kernel-npu from release whl\n# if [ \"$NPU_DEVICE\" = \"A3\" ]; then\n#     wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.01.21/sgl-kernel-npu_2026.01.21_8.5.0_a3.zip\n# fi\n# if [ \"$NPU_DEVICE\" = \"A2\" ]; then\n#     wget https://github.com/sgl-project/sgl-kernel-npu/releases/download/2026.01.21/sgl-kernel-npu_2026.01.21_8.5.0_910b.zip\n# fi\n# unzip sgl-kernel-npu*.zip\n# pip install output/torch_memory_saver*.whl\n# pip install output/sgl_kernel_npu*.whl\n# pip install output/deep_ep*.whl\n\nif [ $USE_MEGATRON -eq 1 ]; then\n    echo \"4. install Megatron and MindSpeed\"\n    git clone -b 2.3.0_core_r0.12.1 https://gitcode.com/Ascend/MindSpeed.git \n    pip install -e MindSpeed \n    pip install git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.12.1 \nfi\n\necho \"5. May need to uninstall timm & triton\"\npip uninstall -y timm triton\necho \"Successfully installed all packages\"\n"
  },
  {
    "path": "scripts/install_vllm_sglang_mcore.sh",
    "content": "#!/bin/bash\n\nUSE_MEGATRON=${USE_MEGATRON:-1}\nUSE_SGLANG=${USE_SGLANG:-1}\n\nexport MAX_JOBS=32\n\necho \"1. install inference frameworks and pytorch they need\"\nif [ $USE_SGLANG -eq 1 ]; then\n    pip install \"sglang[all]==0.5.2\" --no-cache-dir && pip install torch-memory-saver --no-cache-dir\nfi\npip install --no-cache-dir \"vllm==0.11.0\"\n\necho \"2. install basic packages\"\npip install \"transformers[hf_xet]>=4.51.0\" accelerate datasets peft hf-transfer \\\n    \"numpy<2.0.0\" \"pyarrow>=15.0.0\" pandas \"tensordict>=0.8.0,<=0.10.0,!=0.9.0\" torchdata \\\n    ray[default] codetiming hydra-core pylatexenc qwen-vl-utils wandb dill pybind11 liger-kernel mathruler \\\n    pytest py-spy pre-commit ruff tensorboard \n\necho \"pyext is lack of maintainace and cannot work with python 3.12.\"\necho \"if you need it for prime code rewarding, please install using patched fork:\"\necho \"pip install git+https://github.com/ShaohonChen/PyExt.git@py311support\"\n\npip install \"nvidia-ml-py>=12.560.30\" \"fastapi[standard]>=0.115.0\" \"optree>=0.13.0\" \"pydantic>=2.9\" \"grpcio>=1.62.1\"\n\n\necho \"3. install FlashAttention and FlashInfer\"\n# Install flash-attn-2.8.1 (cxx11abi=False)\nwget -nv https://github.com/Dao-AILab/flash-attention/releases/download/v2.8.1/flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl && \\\n    pip install --no-cache-dir flash_attn-2.8.1+cu12torch2.8cxx11abiFALSE-cp312-cp312-linux_x86_64.whl\n\npip install --no-cache-dir flashinfer-python==0.3.1\n\n\nif [ $USE_MEGATRON -eq 1 ]; then\n    echo \"4. install TransformerEngine and Megatron\"\n    echo \"Notice that TransformerEngine installation can take very long time, please be patient\"\n    pip install \"onnxscript==0.3.1\"\n    NVTE_FRAMEWORK=pytorch pip3 install --no-deps git+https://github.com/NVIDIA/TransformerEngine.git@v2.6\n    pip3 install --no-deps git+https://github.com/NVIDIA/Megatron-LM.git@core_v0.13.1\nfi\n\n\necho \"5. May need to fix opencv\"\npip install opencv-python\npip install opencv-fixer && \\\n    python -c \"from opencv_fixer import AutoFix; AutoFix()\"\n\n\nif [ $USE_MEGATRON -eq 1 ]; then\n    echo \"6. Install cudnn python package (avoid being overridden)\"\n    pip install nvidia-cudnn-cu12==9.10.2.21\nfi\n\necho \"Successfully installed all packages\"\n"
  },
  {
    "path": "scripts/legacy_model_merger.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis script is used to merge huggingface model and test verl checkpoints from FSDP and Megatron backends.\n\nTo merge FSDP checkpoints:\n```sh\npython scripts/legacy_model_merger.py merge \\\n    --backend fsdp \\\n    --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \\\n    --target_dir /path/to/merged_hf_model\n```\n\nTo merge Megatron checkpoints:\n```sh\npython scripts/legacy_model_merger.py merge \\\n    --backend megatron \\\n    --tie-word-embedding \\\n    --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \\\n    --target_dir /path/to/merged_hf_model\n```\n\nFor more details, please refer to documentation:\nhttps://verl.readthedocs.io/en/latest/advance/checkpoint.html#convert-fsdp-and-megatron-checkpoints-to-huggingface-format-model\n\"\"\"\n\nimport argparse\nimport os\nimport re\nimport warnings\nfrom abc import ABC, abstractmethod\nfrom concurrent.futures import ThreadPoolExecutor\nfrom dataclasses import dataclass, field\nfrom pathlib import Path\nfrom typing import Optional, Union\n\nimport numpy as np\nimport torch\nfrom accelerate import init_empty_weights\nfrom safetensors.torch import load_file\nfrom torch.distributed._tensor import Placement, Shard\nfrom transformers import (\n    AutoConfig,\n    AutoModelForCausalLM,\n    AutoModelForTokenClassification,\n    GenerationConfig,\n    PretrainedConfig,\n)\n\ntry:\n    # for torch 2.5+\n    from torch.distributed.tensor import DTensor\nexcept ImportError:\n    from torch.distributed._tensor import DTensor\n\nfrom tqdm import tqdm\n\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\n\nAutoModelForVision2Seq = get_auto_model_for_vision2seq()\n\n\n@dataclass\nclass ModelMergerConfig:\n    operation: str  # 'merge' or 'test'\n    backend: str\n    local_dir: str\n    hf_model_config_path: str\n    target_dir: Optional[str] = \"tmp\"\n    hf_upload_path: Optional[str] = None\n    private: bool = False\n    test_hf_dir: Optional[str] = None\n    tie_word_embedding: bool = False\n    is_value_model: bool = False\n    hf_model_path: Optional[str] = None\n    hf_upload: bool = field(init=False)\n\n    def __post_init__(self):\n        self.hf_upload = self.operation == \"merge\" and bool(self.hf_upload_path)\n        if self.operation == \"test\":\n            self.target_dir = None\n            self.hf_upload_path = None\n            self.private = False\n\n\nclass BaseModelMerger(ABC):\n    def __init__(self, config: ModelMergerConfig):\n        self.config = config\n        self.hf_model_config_path = config.hf_model_config_path\n\n        if config.hf_model_path:\n            print(\n                \"Warning: --hf_model_path is deprecated and will be removed in a future version. Currently verl will save huggingface model configuration files into checkpoint directories. Therefore, there is no need to provide --hf_model_path. \"\n            )\n            self.hf_model_config_path = config.hf_model_path\n\n        # Auto-detect huggingface subdirectory if it exists\n        huggingface_subdir = os.path.join(self.hf_model_config_path, \"huggingface\")\n        if os.path.isdir(huggingface_subdir):\n            self.hf_model_config_path = huggingface_subdir\n\n        self.model_config = AutoConfig.from_pretrained(self.hf_model_config_path)\n\n    def get_transformers_auto_model_class(self):\n        # Handle case where architectures might be None or empty\n        if self.model_config.architectures is None or len(self.model_config.architectures) == 0:\n            # Try to infer from model_type if architectures is missing\n            model_type = getattr(self.model_config, 'model_type', '').lower()\n            if 'vision' in model_type or 'vl' in model_type:\n                return AutoModelForVision2Seq\n            elif 'causal' in model_type or 'gpt' in model_type or 'llama' in model_type or 'qwen' in model_type:\n                return AutoModelForCausalLM\n            else:\n                raise NotImplementedError(\n                    f\"Cannot determine model class: architectures is None and model_type '{model_type}' is not recognized\"\n                )\n        \n        architecture = self.model_config.architectures[0]\n        if \"ForTokenClassification\" in architecture:\n            return AutoModelForTokenClassification\n        elif \"ForCausalLM\" in architecture:\n            return AutoModelForCausalLM\n        elif \"ForConditionalGeneration\" in architecture:\n            return AutoModelForVision2Seq\n\n        raise NotImplementedError(f\"Unknown architecture {self.model_config.architectures}\")\n\n    def patch_model_generation_config(self, model):\n        \"\"\"\n        The generation_config created from model config may be different to the pretrained model,\n        this may lead to error when generating: https://github.com/volcengine/verl/issues/1246\n\n        This function patch the generation_config created from model config to the pretrained model.\n        \"\"\"\n        if model.can_generate():\n            try:\n                model.generation_config = GenerationConfig.from_pretrained(self.hf_model_config_path)\n            except OSError:\n                print(\n                    f\"Warning: Generation config file not found in {self.hf_model_config_path}, using a generation config created from the model config.\"\n                )\n        return model\n\n    def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]):\n        \"\"\"\n        Save lora adapter to safetensors.\n\n        Returns:\n            lora_path: str, the path to the lora adapter. None if no lora adapter found.\n\n        Note:\n            This function change the 'state_dict' in place.\n        \"\"\"\n        lora_params_names = [name for name in state_dict.keys() if \"lora_\" in name]\n\n        if len(lora_params_names) == 0:\n            return None\n\n        import json\n        from typing import OrderedDict\n\n        import peft\n        from safetensors.torch import save_file\n\n        lora_params = OrderedDict()\n        target_modules = set()\n        lora_key = None\n\n        for name in lora_params_names:\n            lora_key = name.replace(\".default.weight\", \".weight\")\n            target_modules.add(lora_key.split(\".\")[-3])\n            lora_params[lora_key] = state_dict.pop(name)\n\n        lora_rank = min(lora_params[lora_key].shape[0], lora_params[lora_key].shape[1])\n        peft_dict = {\n            \"r\": lora_rank,\n            \"lora_alpha\": 0,  # lora_alpha is not set. An error should be raised to inform the user to set it manually.\n            \"target_modules\": list(target_modules),\n        }\n        peft_config = peft.LoraConfig(**peft_dict).to_dict()\n        peft_config[\"task_type\"] = peft_config[\"task_type\"].value if peft_config[\"task_type\"] else None\n        peft_config[\"peft_type\"] = peft_config[\"peft_type\"].value if peft_config[\"peft_type\"] else None\n        peft_config[\"target_modules\"] = list(peft_config[\"target_modules\"])\n\n        lora_path = os.path.join(self.config.target_dir, \"lora_adapter\")\n        os.makedirs(lora_path, exist_ok=True)\n        with open(os.path.join(lora_path, \"adapter_config.json\"), \"w\", encoding=\"utf-8\") as f:\n            json.dump(peft_config, f, ensure_ascii=False, indent=4)\n        save_file(lora_params, os.path.join(lora_path, \"adapter_model.safetensors\"))\n\n        for name in list(state_dict.keys()):\n            key = (\n                name.replace(\"base_model.model.\", \"\")\n                .replace(\".base_layer.weight\", \".weight\")\n                .replace(\".base_layer.bias\", \".bias\")\n            )\n            state_dict[key] = state_dict.pop(name)\n\n        return lora_path\n\n    def save_hf_model_and_tokenizer(self, state_dict: dict[str, torch.Tensor]):\n        auto_model_class = self.get_transformers_auto_model_class()\n        with init_empty_weights():\n            model = auto_model_class.from_config(self.model_config, torch_dtype=torch.bfloat16)\n        model.to_empty(device=\"cpu\")\n        model = self.patch_model_generation_config(model)\n\n        lora_path = self.save_lora_adapter(state_dict)\n        if lora_path:\n            print(f\"Saving lora adapter to {lora_path}\")\n\n        print(f\"Saving model to {self.config.target_dir}\")\n        model.save_pretrained(self.config.target_dir, state_dict=state_dict)\n        del state_dict\n        del model\n\n        processor = hf_processor(self.hf_model_config_path)\n        try:\n            tokenizer = hf_tokenizer(self.hf_model_config_path)\n        except Exception as e:\n            warnings.warn(f\"Failed to create tokenizer: {e}. This may affect tokenizer saving\", stacklevel=1)\n            tokenizer = None\n        if processor is not None:\n            print(f\"Saving processor to {self.config.target_dir}\")\n            processor.save_pretrained(self.config.target_dir)\n        if tokenizer is not None:\n            print(f\"Saving tokenizer to {self.config.target_dir}\")\n            tokenizer.save_pretrained(self.config.target_dir)\n\n    def upload_to_huggingface(self):\n        from huggingface_hub import HfApi\n\n        api = HfApi()\n        api.create_repo(repo_id=self.config.hf_upload_path, private=self.config.private, exist_ok=True)\n        api.upload_folder(folder_path=self.config.target_dir, repo_id=self.config.hf_upload_path, repo_type=\"model\")\n\n    @abstractmethod\n    def merge_and_save(self):\n        raise NotImplementedError(\"Subclasses should implement this method\")\n\n\nclass FSDPModelMerger(BaseModelMerger):\n    def _get_world_size(self) -> int:\n        \"\"\"Extracts the FSDP world_size from checkpoint filenames (e.g., 'model_world_size_8_rank_0.pt').\"\"\"\n        for filename in os.listdir(self.config.local_dir):\n            match = re.match(r\"model_world_size_(\\d+)_rank_0\\.pt\", filename)\n            if match:\n                return int(match.group(1))\n        raise FileNotFoundError(\n            f\"Could not determine world size. No file matching 'model_world_size_(\\\\d+)_rank_0.pt' found in {self.config.local_dir}\"\n        )\n\n    def _load_rank_zero_state_dict(self, world_size: int) -> dict:\n        return torch.load(\n            Path(self.config.local_dir) / f\"model_world_size_{world_size}_rank_0.pt\",\n            map_location=\"cpu\",\n            weights_only=False,\n        )\n\n    def _extract_device_mesh_info(self, state_dict: dict, world_size: int) -> tuple[np.ndarray, tuple[str, ...]]:\n        \"\"\"\n        Retrieves sharding information (device_mesh, mesh_dim_names) from a DTensor in the state_dict.\n        If no DTensor is found, infers a simple FSDP mesh based on world_size.\n        \"\"\"\n        pivot_key = sorted(list(state_dict.keys()))[0]\n        weight = state_dict[pivot_key]\n\n        if isinstance(weight, DTensor):\n            # get sharding info\n            device_mesh = weight.device_mesh\n            mesh = device_mesh.mesh\n            mesh_dim_names = device_mesh.mesh_dim_names\n        else:\n            # for non-DTensor\n            mesh = np.array([world_size], dtype=np.int64)\n            mesh_dim_names = (\"fsdp\",)\n\n        return mesh, mesh_dim_names\n\n    def _calculate_shard_configuration(\n        self, mesh: np.ndarray, mesh_dim_names: tuple[str, ...]\n    ) -> tuple[int, tuple[int, ...]]:\n        \"\"\"Calculates the total number of shards and the shape of the device mesh.\"\"\"\n        assert mesh_dim_names in ((\"fsdp\",), (\"ddp\", \"fsdp\")), f\"Unsupported mesh_dim_names {mesh_dim_names}\"\n\n        if \"tp\" in mesh_dim_names:\n            # TODO: \"tp\" is not supported yet due to the above assert\n            total_shards = mesh.shape[-1] * mesh.shape[-2]\n            mesh_shape = (mesh.shape[-2], mesh.shape[-1])\n        else:\n            total_shards = mesh.shape[-1]\n            mesh_shape = (mesh.shape[-1],)\n\n        return total_shards, mesh_shape\n\n    def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) -> torch.Tensor:\n        \"\"\"Merges a list of tensors based on their DTensor placement\"\"\"\n        if placement.is_replicate():\n            return tensors[0]\n        elif placement.is_partial():\n            raise NotImplementedError(\"Partial placement is not supported yet\")\n        elif placement.is_shard():\n            return torch.cat(tensors, dim=placement.dim).contiguous()\n\n        raise NotImplementedError(f\"Unsupported placement: {placement}\")\n\n    def _load_and_merge_state_dicts(\n        self, world_size: int, total_shards: int, mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...]\n    ) -> dict[str, torch.Tensor]:\n        model_state_dict_lst = [None] * total_shards\n\n        def process_one_shard(rank: int, model_state_dict_lst: list):\n            model_path = Path(self.config.local_dir) / f\"model_world_size_{world_size}_rank_{rank}.pt\"\n            state_dict = torch.load(model_path, map_location=\"cpu\", weights_only=False)\n            model_state_dict_lst[rank] = state_dict\n            return state_dict\n\n        with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor:\n            futures = [executor.submit(process_one_shard, rank, model_state_dict_lst) for rank in range(total_shards)]\n            for future in tqdm(futures, desc=f\"Loading {total_shards} FSDP shards\", total=total_shards):\n                future.result()\n\n        # Merge state dicts from all shards\n        state_dict = {}\n        param_placements: dict[str, list] = {}\n\n        for key in set(model_state_dict_lst[0].keys()):\n            state_dict[key] = []\n            for model_state_shard in model_state_dict_lst:\n                # add tensor shard in order of rank to state_dict[key]\n                tensor = model_state_shard.pop(key)\n                if isinstance(tensor, DTensor):\n                    state_dict[key].append(tensor._local_tensor.bfloat16())\n\n                    placements = tuple(tensor.placements)\n                    # replicated placement at dp dimension can be discarded\n                    if mesh_dim_names[0] in (\"dp\", \"ddp\"):\n                        placements = placements[1:]\n\n                    if key not in param_placements:\n                        param_placements[key] = placements\n                    else:\n                        assert param_placements[key] == placements\n                else:\n                    state_dict[key].append(tensor.bfloat16())\n\n        del model_state_dict_lst\n\n        # Merge tensors\n        for key in sorted(state_dict):\n            if not isinstance(state_dict[key], list):\n                print(f\"No need to merge key {key}\")\n                continue\n            if key in param_placements:\n                # merge shards\n                placements: tuple[Shard] = param_placements[key]\n                if len(mesh_shape) == 1:\n                    # 1-D list, FSDP without TP\n                    assert len(placements) == 1\n                    shards = state_dict[key]\n                    state_dict[key] = self._merge_by_placement(shards, placements[0])\n                else:\n                    # 2-D list, FSDP + TP\n                    raise NotImplementedError(\"FSDP + TP is not supported yet\")\n            else:\n                state_dict[key] = torch.cat(state_dict[key], dim=0)\n\n        return state_dict\n\n    def merge_and_save(self):\n        world_size = self._get_world_size()\n        rank_zero_state_dict = self._load_rank_zero_state_dict(world_size)\n\n        mesh, mesh_dim_names = self._extract_device_mesh_info(rank_zero_state_dict, world_size)\n        print(f\"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}\")\n\n        total_shards, mesh_shape = self._calculate_shard_configuration(mesh, mesh_dim_names)\n        print(f\"Processing model shards with {total_shards} {mesh_shape} in total\")\n\n        merged_state_dict = self._load_and_merge_state_dicts(world_size, total_shards, mesh_shape, mesh_dim_names)\n\n        if self.config.operation == \"test\":\n            if not self.config.test_hf_dir:\n                raise ValueError(\"test_hf_dir must be provided for test operation\")\n            self._test_state_dict(merged_state_dict)\n        elif self.config.operation == \"merge\":\n            self.save_hf_model_and_tokenizer(merged_state_dict)\n            if self.config.hf_upload:\n                self.upload_to_huggingface()\n        else:\n            raise ValueError(f\"Unknown operation: {self.config.operation}\")\n\n    def _test_state_dict(self, state_dict: dict[str, torch.Tensor]):\n        auto_model_class = self.get_transformers_auto_model_class()\n\n        hf_model = auto_model_class.from_pretrained(self.config.test_hf_dir, torch_dtype=torch.bfloat16)\n        hf_state_dict = hf_model.state_dict()\n        del hf_model\n\n        hf_model_keys = set(hf_state_dict.keys())\n        collected_keys = set(state_dict.keys())\n\n        missing_keys = hf_model_keys - collected_keys\n        assert len(missing_keys) == 0, f\"Missing keys in collected state dict: {list(sorted(missing_keys))}\"\n\n        extra_keys = collected_keys - hf_model_keys\n        assert len(extra_keys) == 0, f\"Extra keys in collected state dict: {list(sorted(extra_keys))}\"\n\n        for key in hf_model_keys:\n            hf_shape = hf_state_dict[key].shape\n            collected_shape = state_dict[key].shape\n            assert hf_shape == collected_shape, (\n                f\"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}\"\n            )\n\n            hf_dtype = hf_state_dict[key].dtype\n            collected_dtype = state_dict[key].dtype\n            assert hf_dtype == collected_dtype, (\n                f\"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}\"\n            )\n\n            torch.testing.assert_close(hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6)\n\n        print(\"FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager.\")\n\n\nclass MegatronModelMerger(BaseModelMerger):\n    def __init__(self, config: ModelMergerConfig):\n        from verl.utils.megatron_utils import get_hf_config_and_tokenizer_checkpoint_path\n\n        config.hf_model_config_path = get_hf_config_and_tokenizer_checkpoint_path(config.local_dir)\n        super().__init__(config)\n\n        self.params_mapping = {\n            # megatron core gpt model name, huggingface model name\n            # NOTICE: It's a little bit tricky, when 2 keys have the same prefix, we need to make sure the longer key within the containing relationship is processed first.\n            \"embedding.word_embeddings\": \"model.embed_tokens\",\n            # attn\n            \"self_attention.linear_qkv.layer_norm_weight\": \"input_layernorm.weight\",\n            \"self_attention.linear_qkv.layer_norm_bias\": \"input_layernorm.bias\",\n            \"self_attention.linear_qkv\": \"self_attn.qkv_proj\",\n            \"self_attention.q_layernorm\": \"self_attn.q_norm\",\n            \"self_attention.k_layernorm\": \"self_attn.k_norm\",\n            \"self_attention.linear_proj\": \"self_attn.o_proj\",\n            # mla\n            \"self_attention.linear_q_proj\": \"self_attn.q_proj\",\n            \"self_attention.linear_q_down_proj\": \"self_attn.q_a_proj\",\n            \"self_attention.linear_q_up_proj.layer_norm_weight\": \"self_attn.q_a_layernorm.weight\",\n            \"self_attention.linear_q_up_proj\": \"self_attn.q_b_proj\",\n            \"self_attention.linear_kv_down_proj\": \"self_attn.kv_a_proj_with_mqa\",\n            \"self_attention.linear_kv_up_proj.layer_norm_weight\": \"self_attn.kv_a_layernorm.weight\",\n            \"self_attention.linear_kv_up_proj\": \"self_attn.kv_b_proj\",\n            # mlp\n            \"pre_mlp_layernorm\": \"post_attention_layernorm\",\n            \"mlp.linear_fc1.layer_norm_weight\": \"post_attention_layernorm.weight\",\n            \"mlp.linear_fc1.layer_norm_bias\": \"post_attention_layernorm.bias\",\n            \"mlp.linear_fc1\": \"mlp.gate_up_proj\",\n            \"mlp.linear_fc2\": \"mlp.down_proj\",\n            # moe\n            \"mlp.router.expert_bias\": \"mlp.gate.e_score_correction_bias\",\n            \"mlp.router\": \"mlp.gate\",\n            \"mlp.shared_experts.linear_fc1\": \"mlp.shared_experts.gate_up_proj\",\n            \"mlp.shared_experts.linear_fc2\": \"mlp.shared_experts.down_proj\",\n            \"linear_fc1\": \"gate_up_proj\",\n            \"linear_fc2\": \"down_proj\",\n            # output\n            \"final_layernorm\": \"norm\",\n            \"output_layer\": \"lm_head\",\n        }\n\n    def _get_tp_pp_rank_from_sharded_dir(self, sharded_dir: str) -> tuple[int, int]:\n        tp_rank = pp_rank = None\n        rank_list = sharded_dir.split(\"_\")[2:]\n        if re.match(r\"mp_rank_(\\d\\d)_(\\d\\d\\d)\", sharded_dir):\n            tp_rank = int(rank_list[0])\n            pp_rank = int(rank_list[1])\n        elif re.match(r\"mp_rank_(\\d\\d)\", sharded_dir):\n            tp_rank = int(rank_list[0])\n            pp_rank = 0\n\n        assert tp_rank is not None and pp_rank is not None, f\"Invalid sharded dir {sharded_dir}\"\n\n        return tp_rank, pp_rank\n\n    def _check_megatron_checkpoint_path(self, model_path: str) -> tuple[list[str], int, int]:\n        \"\"\"\n        Validates the Megatron checkpoint structure (presence of 'model.pt' in sharded directories).\n        Determines TP and PP sizes from directory names.\n        \"\"\"\n        tp_size = 0\n        pp_size = 0\n        sharded_dirs = sorted(os.listdir(model_path))\n        for sharded_dir in sharded_dirs:\n            assert \"model.pt\" in os.listdir(Path(model_path) / sharded_dir), f\"model.pt not found in {sharded_dir}\"\n            tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir)\n            tp_size = max(tp_size, tp_rank + 1)\n            pp_size = max(pp_size, pp_rank + 1)\n        return sharded_dirs, tp_size, pp_size\n\n    def _merge_across_tp(\n        self,\n        key: str,\n        tp_data: list[torch.Tensor],\n        config: PretrainedConfig,\n        tp_size: int,\n        is_value_model: bool = False,\n    ) -> Union[torch.Tensor, list[torch.Tensor]]:\n        if \"linear_fc1.weight\" in key:\n            # if the tensor is gate and proj\n            gate_lst = []\n            up_lst = []\n            for infer_param in tp_data:\n                gate, up = infer_param.chunk(2)\n                gate_lst.append(gate)\n                up_lst.append(up)\n            gate = torch.cat(gate_lst, dim=0)\n            up = torch.cat(up_lst, dim=0)\n            return [gate, up]\n        elif \"self_attention.linear_qkv.\" in key and \"layer_norm\" not in key:\n            # if the tensor is qkv, for each param on tp, split into q, k, v\n            # concat q, k, v separately.\n            q_lst = []\n            k_lst = []\n            v_lst = []\n            assert config.num_attention_heads % config.num_key_value_heads == 0\n            num_q_per_kv = config.num_attention_heads // config.num_key_value_heads\n            assert tp_data[0].shape[0] % (num_q_per_kv + 2) == 0\n            kv_size_per_tp = tp_data[0].shape[0] // (num_q_per_kv + 2)\n            split_size = [kv_size_per_tp * num_q_per_kv, kv_size_per_tp, kv_size_per_tp]\n\n            for infer_param in tp_data:\n                num_query_groups_per_partition = config.num_key_value_heads // tp_size\n                for chunk in infer_param.chunk(num_query_groups_per_partition):\n                    split_size = [\n                        kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition,\n                        kv_size_per_tp // num_query_groups_per_partition,\n                        kv_size_per_tp // num_query_groups_per_partition,\n                    ]\n                    q, k, v = chunk.split(split_size)\n                    q_lst.append(q)\n                    k_lst.append(k)\n                    v_lst.append(v)\n\n            q = torch.cat(q_lst, dim=0)\n            k = torch.cat(k_lst, dim=0)\n            v = torch.cat(v_lst, dim=0)\n            return [q, k, v]\n        elif \"layer_norm\" in key or \"layernorm\" in key or \"router\" in key or (\"output_layer\" in key and is_value_model):\n            return tp_data[0]\n        else:\n            dim = 0\n            if \"linear_fc2.weight\" in key or \"self_attention.linear_proj\" in key:\n                dim = 1\n            return torch.cat(tp_data, dim=dim)\n\n    def _load_state_dicts(\n        self, model_ckpt_path: str, sharded_dirs: list[str], tp_size: int, pp_size: int\n    ) -> list[list[dict]]:\n        model_state_dict_lst = [[None for _ in range(tp_size)] for _ in range(pp_size)]\n\n        def _process_one_megatron_shard(sharded_dir: str):\n            model_file_path = Path(model_ckpt_path) / sharded_dir / \"model.pt\"\n            state_dict = torch.load(model_file_path, map_location=\"cpu\", weights_only=False)\n            tp_rank, pp_rank = self._get_tp_pp_rank_from_sharded_dir(sharded_dir)\n            model_state_dict_lst[pp_rank][tp_rank] = state_dict\n\n        with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor:\n            futures = [executor.submit(_process_one_megatron_shard, sharded_dir) for sharded_dir in sharded_dirs]\n            for future in tqdm(futures, desc=f\"Loading {len(sharded_dirs)} Megatron shards\", total=len(sharded_dirs)):\n                future.result()\n\n        return model_state_dict_lst\n\n    def _check_megatron_state_key(self, key: str) -> bool:\n        \"\"\"\n        Checks if the key is a valid Megatron state key.\n\n        Now the model merger only supports keys that start with \"decoder/embedding/output_layer\" in TransformerLayer.\n        Shall not use key starts with \"model.\"\n        \"\"\"\n        if key.startswith(\"model.\"):\n            raise ValueError(\n                f\"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder/embedding/output_layer' in TransformerLayer.\"\n            )\n\n        skip_checking_keys = [\"embedding.word_embeddings\", \"output_layer\"]\n        for skip_key in skip_checking_keys:\n            if skip_key in key:\n                print(f\"skip checking key {key}\")\n                return\n\n        # Exclude extra state keys\n        if not key.startswith(\"decoder\"):\n            raise ValueError(\n                f\"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer.\"\n            )\n\n    def _merge_state_dicts(\n        self, model_state_dict_lst: list[list[dict]], tp_size: int, pp_size: int\n    ) -> dict[str, torch.Tensor]:\n        state_dict = {}\n        vpp_size = len(model_state_dict_lst[0][0])\n        layers_cum = 0\n\n        for vpp_rank in range(vpp_size):\n            for pp_rank in range(pp_size):\n                layers_handled = 0\n                keys = model_state_dict_lst[pp_rank][0][vpp_rank].keys()\n                for key in keys:\n                    if \"extra_state\" in key:\n                        continue\n                    if self.config.tie_word_embedding and (\"output_layer\" in key):\n                        print(\"skip lm_head and reward_head loading because of tie_word_embeddings\")\n                        continue\n\n                    self._check_megatron_state_key(key)\n                    hf_name = self._replace_name(key, self.params_mapping)\n                    assert hf_name is not None, f\"Failed to convert layer name [{key}] from megatron to huggingface.\"\n                    if \"model.layers.\" in hf_name:\n                        local_layer_no = int(hf_name.split(\".\")[2])\n                        layers_handled = max(local_layer_no, layers_handled)\n                        global_layer_no = local_layer_no + layers_cum\n                        new_key_list = hf_name.split(\".\")\n                        new_key_list[2] = str(global_layer_no)\n                        hf_name = \".\".join(new_key_list)\n                    else:\n                        warnings.warn(f\"hf_name {hf_name} will not be fixed with layer number\", stacklevel=2)\n\n                    tp_data = [model_state_dict_lst[pp_rank][tp_rank][vpp_rank][key] for tp_rank in range(tp_size)]\n                    merged = self._merge_across_tp(key, tp_data, self.model_config, tp_size, self.config.is_value_model)\n\n                    if not isinstance(merged, list):\n                        state_dict[hf_name] = merged\n                    elif len(merged) == 3:\n                        # split qkv\n                        for n, d in zip([\"q\", \"k\", \"v\"], merged):\n                            state_dict[hf_name.replace(\"qkv\", n)] = d\n                    elif len(merged) == 2:\n                        # split gate up\n                        state_dict[hf_name.replace(\"gate_up\", \"gate\")] = merged[0]\n                        state_dict[hf_name.replace(\"gate_up\", \"up\")] = merged[1]\n                    print(\n                        f\"converted {key} to {hf_name} with shape {merged.shape if isinstance(merged, torch.Tensor) else [t.shape for t in merged]}\"\n                    )\n\n                layers_cum += layers_handled + 1  # zero based\n\n        return state_dict\n\n    def merge_and_save(self):\n        from verl.utils.megatron_utils import get_model_checkpoint_path\n\n        model_ckpt_path = get_model_checkpoint_path(self.config.local_dir)\n        sharded_dirs, tp_size, pp_size = self._check_megatron_checkpoint_path(model_ckpt_path)\n        print(f\"sharded_dirs: {sharded_dirs}, tp_size: {tp_size}, pp_size: {pp_size}, mp_size: {len(sharded_dirs)}\")\n\n        model_state_dict_lst = self._load_state_dicts(model_ckpt_path, sharded_dirs, tp_size, pp_size)\n        merged_state_dict = self._merge_state_dicts(model_state_dict_lst, tp_size, pp_size)\n        del model_state_dict_lst\n\n        if self.config.operation == \"test\":\n            if not self.config.test_hf_dir:\n                raise ValueError(\"test_hf_dir must be provided for test operation\")\n            self._test_state_dict(merged_state_dict)\n        elif self.config.operation == \"merge\":\n            self.save_hf_model_and_tokenizer(merged_state_dict)\n            if self.config.hf_upload:\n                self.upload_to_huggingface()\n        else:\n            raise ValueError(f\"Unknown operation: {self.config.operation}\")\n\n    def _test_state_dict(self, state_dict: dict[str, torch.Tensor]):\n        \"\"\"\n        Compares the merged Megatron state_dict against a reference safetensors model.\n        Applies necessary name mappings from Megatron to Hugging Face conventions using _replace_name.\n        \"\"\"\n        ref_state_dict = load_file(Path(self.config.test_hf_dir) / \"model.safetensors\")\n\n        for name, loaded_weight in state_dict.items():\n            # name = self._replace_name(original_name, self.params_mapping)\n            if not name or name.endswith(\".bias\") and name not in ref_state_dict:\n                continue\n            if \"rotary_emb.inv_freq\" in name:\n                continue\n            if self.config.tie_word_embedding and \"lm_head.weight\" in name:\n                continue\n            if name not in ref_state_dict:\n                raise RuntimeError(f\"key: {name} not exist in state_dict\")\n            param = ref_state_dict[name]\n            assert loaded_weight.dtype == param.dtype\n            torch.testing.assert_close(loaded_weight, param, atol=1e-2, rtol=5e-2)\n\n    def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str:\n        for m_name, v_name in name_mapping.items():\n            if m_name not in megatron_name:\n                continue\n\n            megatron_name = megatron_name.replace(\"decoder\", \"model\")\n            param_name = megatron_name.replace(m_name, v_name)\n            return param_name\n\n        return None  # Return None if no mapping found\n\n\ndef main():\n    parser = argparse.ArgumentParser(description=\"verl model merger\")\n    subparsers = parser.add_subparsers(dest=\"operation\", required=True, help=\"Specify 'merge' or 'test' operation.\")\n\n    base_op_parser = argparse.ArgumentParser(add_help=False)\n    base_op_parser.add_argument(\n        \"--backend\", type=str, required=True, choices=[\"fsdp\", \"megatron\"], help=\"The backend of the model\"\n    )\n    base_op_parser.add_argument(\"--local_dir\", type=str, required=True, help=\"Path to the saved model checkpoints\")\n    base_op_parser.add_argument(\n        \"--hf_model_path\",\n        type=str,\n        default=None,\n        help=\"(Deprecated) Path to the original Hugging Face model for config.\",\n    )\n    base_op_parser.add_argument(\n        \"--tie-word-embedding\",\n        action=\"store_true\",\n        help=\"Whether to tie word embedding weights (currently only Megatron supported)\",\n    )\n    base_op_parser.add_argument(\n        \"--is-value-model\",\n        action=\"store_true\",\n        help=\"Whether the model is a value model (currently only Megatron supported)\",\n    )\n\n    merge_parser = subparsers.add_parser(\"merge\", parents=[base_op_parser], help=\"Merge model checkpoints and save.\")\n    merge_parser.add_argument(\n        \"--target_dir\", default=\"tmp\", type=str, help=\"Directory to save the merged huggingface model\"\n    )\n    merge_parser.add_argument(\n        \"--hf_upload_path\", default=None, type=str, help=\"Hugging Face repository ID to upload the model\"\n    )\n    merge_parser.add_argument(\n        \"--private\", action=\"store_true\", help=\"Whether to upload the model to a private Hugging Face repository\"\n    )\n\n    test_parser = subparsers.add_parser(\n        \"test\", parents=[base_op_parser], help=\"Test merged model against a reference Hugging Face model\"\n    )\n    test_parser.add_argument(\n        \"--test_hf_dir\", type=str, required=True, help=\"Path to the reference Hugging Face model directory for testing\"\n    )\n\n    args = parser.parse_args()\n\n    common_config_args = {\n        \"operation\": args.operation,\n        \"backend\": args.backend,\n        \"tie_word_embedding\": args.tie_word_embedding,\n        \"is_value_model\": args.is_value_model,\n        \"local_dir\": args.local_dir,\n        \"hf_model_path\": args.hf_model_path,\n        \"hf_model_config_path\": args.local_dir,\n    }\n\n    if args.operation == \"merge\":\n        config = ModelMergerConfig(\n            **common_config_args,\n            target_dir=args.target_dir,\n            hf_upload_path=args.hf_upload_path,\n            private=args.private,\n            test_hf_dir=None,\n        )\n        os.makedirs(config.target_dir, exist_ok=True)\n    elif args.operation == \"test\":\n        config = ModelMergerConfig(\n            **common_config_args,\n            test_hf_dir=args.test_hf_dir,\n            # the following args are not used by test operation\n            target_dir=None,\n            hf_upload_path=None,\n            private=False,\n        )\n    else:\n        raise NotImplementedError(f\"Unknown operation: {args.operation}\")\n\n    if config.backend == \"fsdp\":\n        merger = FSDPModelMerger(config)\n    elif config.backend == \"megatron\":\n        merger = MegatronModelMerger(config)\n    else:\n        raise NotImplementedError(f\"Unknown backend: {config.backend}\")\n\n    merger.merge_and_save()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "scripts/megatron_merge_lora.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom pprint import pprint\n\nimport hydra\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.megatron_utils import get_hf_model_checkpoint_path, load_megatron_model_to_gpu\nfrom verl.workers.megatron_workers import ActorRolloutRefWorker\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n\n\nclass CustomSaveWorker(ActorRolloutRefWorker):\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_merged_weights(self, hf_ckpt_path):\n        import os\n\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor_module)\n\n        torch.distributed.barrier()\n\n        print(f\"[Rank {os.environ.get('RANK', '?')}] Saving weights to {hf_ckpt_path}...\")\n\n        if self.vanilla_bridge:\n            self.bridge.save_weights(\n                self.actor_module, hf_ckpt_path, distributed_filesystem=True, memory_efficient=True\n            )\n        else:\n            self.bridge.save_hf_weights(self.actor_module, hf_ckpt_path)\n\n        return True\n\n\n@hydra.main(config_path=\"../verl/trainer/config\", config_name=\"ppo_megatron_trainer\", version_base=None)\ndef main(config):\n    assert config.actor_rollout_ref.model.lora.adapter_path is not None, \"adapter_path must be specified\"\n\n    if (\n        config.actor_rollout_ref.actor.optim.lr_decay_steps is None\n        or config.actor_rollout_ref.actor.optim.lr_decay_steps < 1\n    ):\n        # set to bypass OptimizerParamScheduler checks\n        config.actor_rollout_ref.actor.optim.lr_decay_steps = 100000\n\n    run_merge(config)\n\n\ndef run_merge(config) -> None:\n    if not ray.is_initialized():\n        # this is for local ray cluster\n        default_runtime_env = {\"env_vars\": {\"TOKENIZERS_PARALLELISM\": \"true\", \"NCCL_DEBUG\": \"WARN\"}}\n        ray_init_kwargs = config.ray_kwargs.get(\"ray_init\", {})\n        runtime_env_kwargs = ray_init_kwargs.get(\"runtime_env\", {})\n        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)\n        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, \"runtime_env\": runtime_env})\n        print(f\"ray init kwargs: {ray_init_kwargs}\")\n        ray.init(**OmegaConf.to_container(ray_init_kwargs))\n\n    ray.get(main_task.remote(config))\n\n\n@ray.remote(num_cpus=1)\ndef main_task(config):\n    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values\n    OmegaConf.resolve(config)\n\n    ray_cls_with_init = RayClassWithInitArgs(\n        cls=ray.remote(CustomSaveWorker), config=config.actor_rollout_ref, role=\"actor\"\n    )\n    resource_pool = RayResourcePool(process_on_nodes=[config.trainer.n_gpus_per_node] * config.trainer.nnodes)\n\n    worker = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=ray_cls_with_init,\n        device_name=config.trainer.device,\n    )\n    worker.init_model()\n\n    adapter_path = config.actor_rollout_ref.model.lora.adapter_path\n    hf_ckpt_path = get_hf_model_checkpoint_path(os.path.dirname(adapter_path))\n    worker.save_merged_weights(hf_ckpt_path)\n\n\nif __name__ == \"__main__\":\n    \"\"\"\n    Use the same config as your training script, besides **specifying the adapter_path**.\n\n    For example, your training script starts with: \n        `python3 -m verl.trainer.main_ppo --config-name=ppo_megatron_trainer ...`\n    Now replace it with \n        `python3 ./scripts/megatron_merge_lora.py --config-name=ppo_megatron_trainer ...`\n    \"\"\"\n    main()\n"
  },
  {
    "path": "scripts/print_cfg.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\ntry:\n    import hydra\nexcept ImportError as e:\n    raise ImportError(\"Please install hydra-core via 'pip install hydra-core' and retry.\") from e\n\n\n@hydra.main(config_path=\"../verl/trainer/config\", config_name=\"ppo_trainer\", version_base=None)\ndef main(config):\n    \"\"\"Main entry point for PPO training with Hydra configuration management.\n\n    Args:\n        config_dict: Hydra configuration dictionary containing training parameters.\n    \"\"\"\n    print(config)\n    from verl.utils.config import omega_conf_to_dataclass\n\n    profiler_config = omega_conf_to_dataclass(config.critic.profiler)\n    print(profiler_config)\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "scripts/rollout_viewer.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport re\nimport traceback\nfrom pathlib import Path\nfrom typing import Annotated, Optional\n\nimport aiofiles\n\ntry:\n    import ujson as json\nexcept ImportError:\n    import json\nimport typer\nfrom rich.highlighter import ReprHighlighter\nfrom rich.markdown import Markdown\nfrom rich.table import Table\nfrom rich.text import Text\nfrom textual import on\nfrom textual.app import App, ComposeResult\nfrom textual.containers import Horizontal, Vertical, VerticalScroll\nfrom textual.widgets import Input, ProgressBar, Select, SelectionList, Static\n\nINDEX_KEY = \"__IDX\"\nFILE_SUFFIX = \".jsonl\"\n\n\ndef check_textual_version():\n    # check if textual version is equal to 0.52.1\n    import textual\n    from packaging.version import Version\n\n    if Version(textual.__version__) != Version(\"0.52.1\"):\n        raise ImportError(f\"Textual version {textual.__version__} is not supported, please pip install textual==0.52.1\")\n\n\ncheck_textual_version()\n\n\nasync def load_path(p: Path, data: dict, mask_strs: str, idx: int, pbar):\n    samples = []\n    async with aiofiles.open(p, encoding=\"utf-8\") as f:\n        async for line in f:\n            d = json.loads(line)\n            for k in d:\n                if isinstance(d[k], str):\n                    if mask_strs:\n                        d[k] = re.sub(rf\"{mask_strs}\", \"*\", d[k])\n                else:\n                    d[k] = json.dumps(d[k], ensure_ascii=False, indent=4)\n\n            d[INDEX_KEY] = len(samples)\n            samples.append(d)\n        data[idx] = {\"samples\": samples}\n\n    print(f\"path {p} loaded\")\n    pbar.advance(1)\n\n\nasync def load_dir(path: Path, data: dict[int, dict], pbar, mask_strs: str = \"\"):\n    paths = list(path.glob(f\"*{FILE_SUFFIX}\"))\n    paths = sorted(paths, key=lambda x: int(x.stem))\n\n    tasks = [load_path(p, data, mask_strs, i, pbar) for i, p in enumerate(paths)]\n\n    await asyncio.gather(*tasks)\n\n\nclass Highlighter(ReprHighlighter):\n    highlights = ReprHighlighter.highlights + [\n        r\"(?P<tag_name>[][\\<\\>{}()\\|（）【】\\[\\]=`])\",\n        r\"\\<\\|(?P<tag_name>[\\w\\W]*?)\\|\\>\",\n    ]\n\n\ndef center_word_with_equals_exactly(word: str, total_length: int, char: str = \"=\") -> str:\n    if len(word) > total_length:\n        return word\n\n    padding = total_length - len(word)\n    left_pad = (padding) // 2\n    right_pad = (padding + 1) // 2\n    return char * left_pad + \" \" + word + \" \" + char * right_pad\n\n\ndef highlight_keyword(content: str, keyword: Optional[str]):\n    if not keyword:\n        return Text(content)\n    text = Text()\n    parts = content.split(keyword)\n    for i, part in enumerate(parts):\n        text.append(part, style=None)\n        if i < len(parts) - 1:\n            # text.append(keyword, style=Style(color=\"#d154d1\", bgcolor=\"yellow\", bold=True))\n            text.append(keyword, style=\"on #8f51b5\")\n    return text\n\n\nhelp_doc = \"\"\"\n⌨️   keybinds：\n\n- `f/esc`: find/cancel\n- `tab/←/→`: change focus\n- `j/k`: page down/up\n- `g/G`: scroll home/end\n- `n/N`: next sample/step\n- `p/P`: previous sample/step\n- `s`: switch display mode\n  - plain text\n  - rich table\n\n\"\"\"\n\n\nclass JsonLineViewer(App):\n    BINDINGS = [\n        (\"left\", \"focus_previous\", \"Focus Previous\"),\n        (\"right\", \"focus_next\", \"Focus Next\"),\n        (\"s\", \"swith_render\", \"switch render\"),\n        # control\n        (\"n\", \"next_sample\", \"Next Sample\"),\n        (\"N\", \"next_step\", \"Next Step\"),\n        (\"p\", \"previous_sample\", \"Previous Sample\"),\n        (\"P\", \"previous_step\", \"Previous Step\"),\n        # search\n        (\"f\", \"toggle_search\", \"find\"),\n        (\"enter\", \"next_search\", \"find next\"),\n        (\"escape\", \"cancel_search\", \"cancel find\"),\n        # scroll\n        (\"j\", \"page_down\", \"page down\"),\n        (\"k\", \"page_up\", \"page up\"),\n        (\"g\", \"page_home\", \"page home\"),\n        (\"G\", \"page_end\", \"page end\"),\n    ]\n\n    CSS = \"\"\"\n\n    Select:focus > SelectCurrent {\n        border: tall #8f51b5;\n    }\n    Select.-expanded > SelectCurrent {\n        border: tall #8f51b5;\n    }\n    #select-container {\n        width: 15%;\n        height: 100%;\n        align: center top;\n    }\n    #search-container {\n        height: 10%;\n        align: center top;\n    }\n    #search-box {\n        width: 50%;\n    }\n    #reqid-box {\n        width: 50%;\n    }\n    \"\"\"\n\n    def __init__(self, step_num: int, data: dict[int, dict], pbar):\n        super().__init__()\n        self.step_num = step_num\n\n        self.data = data\n        self.render_table = False\n        self.selected_step_index = 0\n        self.selected_sample_index = 0\n        self.pbar = pbar\n\n        self.matches = []\n        self.current_match_index = 0\n\n        self.highlighter = Highlighter()\n\n        first_samples = data[list(data.keys())[0]][\"samples\"]\n        # Prepare the initial field filter list (all keys from the first sample)\n        self.filter_fields = [(f, f, True) for f in first_samples[0].keys()]\n\n        # Internal set used for fast membership checks when we add new fields on the fly.\n        # We keep it here so that when new columns appear in later steps (e.g. `request_id`),\n        # they can be added to the UI automatically without restarting the viewer.\n        self._field_set: set[str] = set(first_samples[0].keys())\n        self.sample_num = len(first_samples)\n\n    def compose(self) -> ComposeResult:\n        with Horizontal(id=\"search-container\"):\n            yield Input(placeholder=\"find something...\", id=\"search-box\")\n            yield Input(placeholder=\"request id...\", id=\"reqid-box\")\n            with Vertical(id=\"search-container2\"):\n                yield self.pbar\n                yield Static(\"\", id=\"search-status\")\n\n        with Horizontal():\n            with Vertical(id=\"select-container\"):\n                yield Static(\"\\n\")\n                yield Static(\n                    renderable=Markdown(\n                        help_doc,\n                    ),\n                    markup=False,\n                )\n                yield Static(\"\\n\")\n                yield Select(\n                    id=\"step-select\",\n                    value=0,\n                    prompt=\"select step\",\n                    options=[(\"step: 1\", 0)],\n                    allow_blank=False,\n                )\n                yield Select(\n                    id=\"sample-select\",\n                    value=0,\n                    prompt=\"select sample\",\n                    options=[(\"sample: 1\", 0)],\n                    allow_blank=False,\n                )\n                yield Select(\n                    id=\"sample-sort\",\n                    value=0,\n                    prompt=\"排序\",\n                    options=[\n                        (\"sort\", 0),\n                        (\"score asc\", 1),\n                        (\"score desc\", 2),\n                    ],\n                    allow_blank=False,\n                )\n\n                yield SelectionList[int]((\"Select ALL\", 1, True), id=\"fields-select-all\")\n                with VerticalScroll(id=\"scroll-view2\"):\n                    yield SelectionList[str](*self.filter_fields, id=\"fields-select\")\n            with VerticalScroll(id=\"scroll-view\"):\n                yield Static(id=\"content\", markup=False)\n\n    async def on_mount(self) -> None:\n        self.step_select = self.query_one(\"#step-select\", Select)\n        self.sample_select = self.query_one(\"#sample-select\", Select)\n        self.sample_sort = self.query_one(\"#sample-sort\", Select)\n        self.content_display = self.query_one(\"#content\", Static)\n        self.search_box = self.query_one(\"#search-box\", Input)\n        self.reqid_box = self.query_one(\"#reqid-box\", Input)\n        self.scroll_view = self.query_one(\"#scroll-view\", VerticalScroll)\n        self.search_status = self.query_one(\"#search-status\", Static)\n        self.fields_select = self.query_one(\"#fields-select\", SelectionList)\n        self.fields_select.border_title = \"field filter\"\n\n        if self.data:\n            self.step_select.set_options([(f\"step: {i + 1}\", i) for i in range(self.step_num)])\n            self.sample_select.set_options([(f\"sample: {i + 1}\", i) for i in range(self.sample_num)])\n            self.step_select.focus()\n            await self.update_content()\n\n    def update_result_options(self, offset: int = 0, sort_desc: Optional[bool] = None):\n        options = []\n        if isinstance(self.selected_step_index, int) and self.selected_step_index < len(self.data):\n            if self.sample_num is None or sort_desc is not None:\n                samples = self.data[self.selected_step_index].get(\"samples\", [])\n                if not samples:\n                    self.selected_sample_index = offset\n                    return\n                if sort_desc is not None:\n                    samples = sorted(\n                        samples,\n                        key=lambda x: x.get(\"score\", x.get(\"score_1\", 0)),\n                        reverse=sort_desc,\n                    )\n\n                options = [(f\"sample: {r[INDEX_KEY] + 1}\", r[INDEX_KEY]) for r in samples]\n                self.sample_select.set_options(options)\n                self.sample_num = len(samples)\n\n            if sort_desc is not None and options:\n                self.selected_sample_index = options[0][1]\n            else:\n                self.selected_sample_index = offset\n\n    async def update_content(self, search_keyword: Optional[str] = None):\n        content = \"\"\n        try:\n            samples = self.data[self.selected_step_index].get(\"samples\", [])\n            content_dict_full = samples[self.selected_sample_index]\n\n            # Dynamically track any NEW keys that appear and add them to the field filter.\n            self._update_fields_select(content_dict_full.keys())\n\n            # Apply field selection filter (only show selected fields)\n            content_dict = {k: v for k, v in content_dict_full.items() if k in self.fields_select.selected}\n            if self.render_table:\n                content = Table(\"key\", \"value\", show_lines=True)\n                for k in content_dict:\n                    v = content_dict[k]\n                    v = f\"{v}\"\n                    content.add_row(\n                        k,\n                        self.highlighter(highlight_keyword(v, search_keyword)),\n                    )\n            else:\n                text = Text()\n                for k in content_dict:\n                    v = content_dict[k]\n                    s = center_word_with_equals_exactly(k, 64) + f\"\\n{v}\\n\"\n                    text.append(highlight_keyword(s, search_keyword))\n                content = self.highlighter(text)\n        except KeyError:\n            content = f\"Loading data asynchronously, progress: {len(self.data)}/{self.step_num} step\"\n\n        except Exception:\n            content = self.highlighter(traceback.format_exc())\n\n        self.content_display.update(content)\n\n    # ---------------------------------------------------------------------\n    # Request-ID jump logic\n    # ---------------------------------------------------------------------\n\n    @on(Input.Submitted, \"#reqid-box\")\n    async def on_reqid_submitted(self, event: Input.Submitted) -> None:\n        \"\"\"Jump to the sample that has a matching `request_id`.\"\"\"\n\n        req_id_raw = event.value.strip()\n        # Remove hyphens so search is tolerant to different id formats\n        req_id = req_id_raw.replace(\"-\", \"\")\n        if not req_id:\n            return\n\n        found = False\n        for step_idx, step_data in self.data.items():\n            for sample in step_data.get(\"samples\", []):\n                sample_id = str(sample.get(\"request_id\", \"\"))\n                if sample_id.replace(\"-\", \"\") == req_id:\n                    # Update selected indices\n                    self.selected_step_index = step_idx\n                    self.step_select.value = step_idx\n\n                    # Ensure sample list is updated and select sample\n                    self.update_result_options(offset=sample[INDEX_KEY])\n                    self.selected_sample_index = sample[INDEX_KEY]\n                    self.sample_select.value = sample[INDEX_KEY]\n\n                    await self._clear_search()\n                    await self.update_content()\n\n                    found = True\n                    break\n            if found:\n                break\n\n        if not found:\n            self.search_status.update(Text(f\"request_id '{req_id_raw}' not found\", style=\"bold red\"))\n        else:\n            # Keep the typed id in the input box so users see what was searched.\n            pass\n\n    # ---------------------------------------------------------------------\n    # Helper: add new fields to SelectionList on-the-fly\n    # ---------------------------------------------------------------------\n\n    def _update_fields_select(self, keys):\n        \"\"\"Add any unseen *keys* to the field-selection widget so they can be toggled.\n\n        The viewer is often launched with only the first step loaded. Later steps may\n        introduce new columns (e.g. `request_id`). This helper ensures those fields\n        become visible without requiring a restart.\n        \"\"\"\n        # Ensure we have the widget (only after on_mount)\n        if not hasattr(self, \"fields_select\"):\n            return\n\n        for k in keys:\n            if k not in self._field_set:\n                self._field_set.add(k)\n                try:\n                    # By default, new fields are selected so they appear immediately.\n                    self.fields_select.add_option(k, k, selected=True)\n                except Exception:\n                    # Fallback for older textual versions where signature is different.\n                    self.fields_select.add_option((k, k, True))\n\n    @on(Select.Changed, \"#step-select\")\n    async def step_changed(self, event):\n        self.selected_step_index = event.value\n        self.update_result_options()\n        await self.update_content()\n\n    @on(Select.Changed, \"#sample-select\")\n    async def sample_changed(self, event):\n        self.selected_sample_index = event.value\n        await self._clear_search()\n        await self.update_content()\n\n    @on(Select.Changed, \"#sample-sort\")\n    async def sort_changed(self, event):\n        v = event.value\n        self.update_result_options(sort_desc=None if v == 0 else False if v == 1 else True)\n        await self.update_content()\n\n    @on(SelectionList.SelectedChanged, \"#fields-select\")\n    async def fields_changed(self, event):\n        await self.update_content()\n\n    @on(SelectionList.SelectedChanged, \"#fields-select-all\")\n    async def fields_all_changed(self, event):\n        s = self.query_one(\"#fields-select-all\", SelectionList)\n        if s.selected:\n            self.fields_select.select_all()\n        else:\n            self.fields_select.deselect_all()\n\n    def action_focus_previous(self):\n        self.screen.focus_previous()\n\n    def action_focus_next(self):\n        self.screen.focus_next()\n\n    async def action_next_step(self) -> None:\n        self.selected_step_index += 1\n        if self.selected_step_index >= self.step_num:\n            self.selected_step_index = 0\n        self.step_select.value = self.selected_step_index\n        self.update_result_options()\n        await self.update_content()\n\n    async def action_next_sample(self) -> None:\n        self.selected_sample_index += 1\n        if not self.sample_num or self.selected_sample_index >= self.sample_num:\n            self.selected_sample_index = 0\n        self.sample_select.value = self.selected_sample_index\n        await self._clear_search()\n        await self.update_content()\n\n    async def action_previous_step(self) -> None:\n        self.selected_step_index -= 1\n        if self.selected_step_index < 0:\n            self.selected_step_index = self.step_num - 1\n        self.step_select.value = self.selected_step_index\n        self.update_result_options()\n        await self.update_content()\n\n    async def action_previous_sample(self) -> None:\n        self.selected_sample_index -= 1\n        if self.selected_sample_index < 0:\n            self.selected_sample_index = self.sample_num - 1\n        self.sample_select.value = self.selected_sample_index\n        await self._clear_search()\n        await self.update_content()\n\n    async def action_swith_render(self):\n        self.render_table = not self.render_table\n        await self.update_content()\n\n    def action_toggle_search(self) -> None:\n        self.search_box.focus()\n\n    async def action_cancel_search(self) -> None:\n        self.search_box.value = \"\"\n        await self._clear_search()\n        await self.update_content()\n\n    async def _clear_search(self):\n        self.matches = []\n        self.search_status.update(\"\")\n        self.current_match_index = 0\n\n    @on(Input.Submitted, \"#search-box\")\n    async def on_search_submitted(self, event: Input.Submitted) -> None:\n        self.matches = []\n        self.current_match_index = 0\n        if event.value:\n            await self.update_content(event.value)\n            renderable = self.content_display.render()\n            if isinstance(renderable, Table):\n                return\n\n            assert isinstance(renderable, Text)\n            console = self.content_display._console\n            lines = renderable.wrap(console, self.scroll_view.container_size.width)\n            line_idx_recorded = set()\n            for line_idx, line in enumerate(lines):\n                if line_idx in line_idx_recorded:\n                    continue\n                if event.value in line:\n                    self.matches.append(\n                        {\n                            \"line\": line_idx,\n                            \"word\": event.value,\n                        }\n                    )\n                    line_idx_recorded.add(line_idx)\n            self.scroll_view.focus()\n            await self.action_next_search()\n\n    async def action_next_search(self) -> None:\n        if not self.matches or self.current_match_index >= len(self.matches):\n            return\n\n        target_line = self.matches[self.current_match_index][\"line\"]\n        self.scroll_view.scroll_to(x=0, y=target_line * 1, animate=False)\n        self.current_match_index = (self.current_match_index + 1) % len(self.matches)\n        self.search_status.update(\n            Text(\n                f\"Find ：{self.current_match_index + 1}/{len(self.matches)}\",\n                style=\"bold on #8f51b5\",\n            )\n        )\n\n    def action_page_up(self):\n        self.scroll_view.scroll_page_up(animate=False)\n\n    def action_page_down(self):\n        self.scroll_view.scroll_page_down(animate=False)\n\n    def action_page_home(self):\n        self.scroll_view.scroll_home(animate=False)\n\n    def action_page_end(self):\n        self.scroll_view.scroll_end(animate=False)\n\n\nasync def _run(path: Path, mask_str: str):\n    assert path.exists(), f\"{path} not exist\"\n\n    paths = list(path.glob(f\"*{FILE_SUFFIX}\"))\n    paths = sorted(paths, key=lambda x: int(x.stem))\n\n    if not paths:\n        raise ValueError(f\"no available reward dump files under f{path}\")\n\n    print(f\"get jsonl file nums: {len(paths)}\")\n\n    pbar = ProgressBar(total=len(paths), name=\"data load progress\")\n    data = {}\n    await load_path(paths[0], data, mask_str, 0, pbar)\n    app = JsonLineViewer(step_num=len(paths), data=data, pbar=pbar)\n    await asyncio.gather(load_dir(path, data, pbar, mask_str), app.run_async())\n\n\napp = typer.Typer()\n\n\n@app.command(help=\"launch TUI APP\")\ndef run(\n    rollout_data_dir: Path,\n    mask_str: Annotated[str, typer.Option(help=\"string that will be masked to *\")] = r\"<\\|image_pad\\|>|<\\|imgpad\\|>\",\n):\n    loop = asyncio.get_event_loop()\n    loop.run_until_complete(_run(rollout_data_dir, mask_str))\n\n\nif __name__ == \"__main__\":\n    app()\n"
  },
  {
    "path": "scripts/veomni/moe_merge.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nMerge individual MoE expert weights into stacked tensors for efficient loading.\n\nThis script takes a HuggingFace checkpoint with individual expert weights\n(e.g., model.layers.{i}.mlp.experts.{j}.gate_proj.weight) and merges them\ninto stacked tensors (e.g., model.layers.{i}.mlp.experts.gate_proj) for\nfaster loading and better memory efficiency in VeOmni.\n\nThe merging process:\n1. Loads individual expert weights from the HF checkpoint\n2. Stacks them into single tensors for each projection type\n3. Handles all three projection types: gate_proj, up_proj, down_proj\n4. Supports both Qwen3-MoE (num_experts) and DeepSeek (n_routed_experts) formats\n5. Handles models with initial dense layers (first_k_dense_replace)\n\nUsage: python moe_merge.py --raw_hf_path <input_checkpoint> --merge_hf_path <output_dir>\n\"\"\"\n\nimport os\nfrom argparse import ArgumentParser\nfrom dataclasses import dataclass\nfrom glob import glob\nfrom typing import Generator\n\nimport torch\nfrom safetensors.torch import safe_open\nfrom tqdm import tqdm\nfrom transformers import AutoConfig\nfrom veomni.models import build_tokenizer, save_model_weights\n\n\n@dataclass\nclass StateDictIterator:\n    filepath: str\n\n    def __iter__(self) -> Generator[tuple[str, \"torch.Tensor\"], None, None]:\n        if self.filepath.endswith(\".safetensors\"):\n            with safe_open(self.filepath, framework=\"pt\", device=\"cpu\") as f:\n                for key in f.keys():\n                    yield key, f.get_tensor(key)\n\n        else:\n            state_dict = torch.load(self.filepath, map_location=\"cpu\", weights_only=True, mmap=True)\n            for key in state_dict.keys():\n                yield key, state_dict[key]\n\n\ndef main(raw_hf_path, merge_hf_path):\n    torch.set_default_dtype(torch.bfloat16)\n    os.makedirs(merge_hf_path, exist_ok=True)\n\n    config = AutoConfig.from_pretrained(raw_hf_path)\n    tokenizer = build_tokenizer(raw_hf_path)\n\n    safetensor_files = list(glob(os.path.join(raw_hf_path, \"*.safetensors\")))\n    safetensor_files.sort()\n    state_dict_iterators = [StateDictIterator(shard_file) for shard_file in safetensor_files]\n    new_state_dict = {}\n    for state_dict_iterator in tqdm(state_dict_iterators, desc=\"Loading checkpoint shards\"):\n        for name, tensor in state_dict_iterator:\n            new_state_dict[name] = tensor.cpu()\n\n    print(new_state_dict.keys())\n\n    if hasattr(config, \"num_experts\"):\n        # qwen3moe\n        num_experts = config.num_experts\n    elif hasattr(config, \"n_routed_experts\"):\n        # deepseek\n        num_experts = config.n_routed_experts\n    else:\n        raise RuntimeError(\"could not find how many experts to assign\")\n    num_hidden_layers = config.num_hidden_layers\n\n    if hasattr(config, \"first_k_dense_replace\"):\n        # deepseek first k dense layer\n        moe_layer_start_idx = config.first_k_dense_replace\n    else:\n        # moe layer only in the model\n        moe_layer_start_idx = 0\n\n    for i in range(moe_layer_start_idx, num_hidden_layers):\n        gate_proj = []\n        for j in range(num_experts):\n            gate_proj.append(new_state_dict.pop(f\"model.layers.{i}.mlp.experts.{j}.gate_proj.weight\"))\n\n        new_state_dict[f\"model.layers.{i}.mlp.experts.gate_proj\"] = torch.stack(gate_proj)\n        up_proj = []\n        for j in range(num_experts):\n            up_proj.append(new_state_dict.pop(f\"model.layers.{i}.mlp.experts.{j}.up_proj.weight\"))\n\n        new_state_dict[f\"model.layers.{i}.mlp.experts.up_proj\"] = torch.stack(up_proj)\n        down_proj = []\n        for j in range(num_experts):\n            down_proj.append(new_state_dict.pop(f\"model.layers.{i}.mlp.experts.{j}.down_proj.weight\"))\n\n        new_state_dict[f\"model.layers.{i}.mlp.experts.down_proj\"] = torch.stack(down_proj)\n\n    model_assets = [config, tokenizer]\n    save_model_weights(merge_hf_path, new_state_dict, model_assets=model_assets)\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\"--raw_hf_path\", type=str, required=True)\n    parser.add_argument(\"--merge_hf_path\", type=str, required=True)\n    args = parser.parse_args()\n    main(args.raw_hf_path, args.merge_hf_path)\n"
  },
  {
    "path": "scripts/veomni/moe_split.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nReverse process of moe_merge.py - splits merged MoE expert weights back to individual experts.\n\nThis script takes a HF checkpoint that has been processed by moe_merge.py (where expert weights\nare stacked into single tensors) and splits them back to the original format with individual\nexpert weights.\n\nThe process reverses the merging by:\n1. Loading stacked tensors like model.layers.{i}.mlp.experts.gate_proj\n2. Unstacking them back to individual experts model.layers.{i}.mlp.experts.{j}.gate_proj.weight\n3. Handling all three projection types: gate_proj, up_proj, down_proj\n\nUsage: python moe_split.py --merge_hf_path <merged_checkpoint> --split_hf_path <output_dir>\n\"\"\"\n\nimport os\nfrom argparse import ArgumentParser\nfrom dataclasses import dataclass\nfrom glob import glob\nfrom typing import Generator\n\nimport torch\nfrom safetensors.torch import safe_open\nfrom tqdm import tqdm\nfrom transformers import AutoConfig\nfrom veomni.models import build_tokenizer, save_model_weights\n\n\n@dataclass\nclass StateDictIterator:\n    filepath: str\n\n    def __iter__(self) -> Generator[tuple[str, \"torch.Tensor\"], None, None]:\n        if self.filepath.endswith(\".safetensors\"):\n            with safe_open(self.filepath, framework=\"pt\", device=\"cpu\") as f:\n                for key in f.keys():\n                    yield key, f.get_tensor(key)\n\n        else:\n            state_dict = torch.load(self.filepath, map_location=\"cpu\", weights_only=True, mmap=True)\n            for key in state_dict.keys():\n                yield key, state_dict[key]\n\n\ndef main(merge_hf_path, split_hf_path):\n    torch.set_default_dtype(torch.bfloat16)\n    os.makedirs(split_hf_path, exist_ok=True)\n\n    config = AutoConfig.from_pretrained(merge_hf_path)\n    tokenizer = build_tokenizer(merge_hf_path)\n\n    safetensor_files = list(glob(os.path.join(merge_hf_path, \"*.safetensors\")))\n    safetensor_files.sort()\n    state_dict_iterators = [StateDictIterator(shard_file) for shard_file in safetensor_files]\n    new_state_dict = {}\n    for state_dict_iterator in tqdm(state_dict_iterators, desc=\"Loading checkpoint shards\"):\n        for name, tensor in state_dict_iterator:\n            new_state_dict[name] = tensor.cpu()\n\n    num_experts = config.num_experts\n    num_hidden_layers = config.num_hidden_layers\n    for i in range(num_hidden_layers):\n        print(f\"Converting layer {i}\")\n        for proj_name in [\"gate_proj\", \"up_proj\", \"down_proj\"]:\n            stacked_key = f\"model.layers.{i}.mlp.experts.{proj_name}\"\n            if stacked_key in new_state_dict:\n                stacked_tensor = new_state_dict.pop(stacked_key)\n                for j in range(num_experts):\n                    expert_key = f\"model.layers.{i}.mlp.experts.{j}.{proj_name}.weight\"\n                    new_state_dict[expert_key] = stacked_tensor[j]\n\n    model_assets = [config, tokenizer]\n\n    print(\"Saving to safetensors\")\n    save_model_weights(split_hf_path, new_state_dict, model_assets=model_assets)\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\"--merge_hf_path\", type=str, required=True)\n    parser.add_argument(\"--split_hf_path\", type=str, required=True)\n    args = parser.parse_args()\n    main(args.merge_hf_path, args.split_hf_path)\n"
  },
  {
    "path": "setup.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# setup.py is the fallback installation script when pyproject.toml does not work\nimport os\nfrom pathlib import Path\n\nfrom setuptools import find_packages, setup\n\nversion_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))\n\nwith open(os.path.join(version_folder, \"verl/version/version\")) as f:\n    __version__ = f.read().strip()\n\ninstall_requires = [\n    \"accelerate\",\n    \"codetiming\",\n    \"datasets\",\n    \"dill\",\n    \"hydra-core\",\n    \"numpy<2.0.0\",\n    \"pandas\",\n    \"peft\",\n    \"pyarrow>=19.0.0\",\n    \"pybind11\",\n    \"pylatexenc\",\n    \"ray[default]>=2.41.0\",\n    \"torchdata\",\n    \"tensordict>=0.8.0,<=0.10.0,!=0.9.0\",\n    \"transformers\",\n    \"wandb\",\n    \"packaging>=20.0\",\n    \"tensorboard\",\n]\n\nTEST_REQUIRES = [\"pytest\", \"pre-commit\", \"py-spy\", \"pytest-asyncio\", \"pytest-rerunfailures\"]\nPRIME_REQUIRES = [\"pyext\"]\nGEO_REQUIRES = [\"mathruler\", \"torchvision\", \"qwen_vl_utils\"]\nGPU_REQUIRES = [\"liger-kernel\", \"flash-attn\"]\nMATH_REQUIRES = [\"math-verify\"]  # Add math-verify as an optional dependency\nVLLM_REQUIRES = [\"tensordict>=0.8.0,<=0.10.0,!=0.9.0\", \"vllm>=0.8.5,<=0.12.0\"]\nTRTLLM_REQUIRES = [\"tensorrt-llm>=1.2.0rc6\"]\nSGLANG_REQUIRES = [\n    \"tensordict>=0.8.0,<=0.10.0,!=0.9.0\",\n    \"sglang[srt,openai]==0.5.6\",\n    \"torch==2.9.1\",\n]\nTRL_REQUIRES = [\"trl<=0.9.6\"]\nMCORE_REQUIRES = [\"mbridge\"]\n\nextras_require = {\n    \"test\": TEST_REQUIRES,\n    \"prime\": PRIME_REQUIRES,\n    \"geo\": GEO_REQUIRES,\n    \"gpu\": GPU_REQUIRES,\n    \"math\": MATH_REQUIRES,\n    \"vllm\": VLLM_REQUIRES,\n    \"sglang\": SGLANG_REQUIRES,\n    \"trl\": TRL_REQUIRES,\n    \"mcore\": MCORE_REQUIRES,\n    \"trtllm\": TRTLLM_REQUIRES,\n}\n\n\nthis_directory = Path(__file__).parent\nlong_description = (this_directory / \"README.md\").read_text()\n\nsetup(\n    name=\"verl\",\n    version=__version__,\n    package_dir={\"\": \".\"},\n    packages=find_packages(where=\".\"),\n    url=\"https://github.com/volcengine/verl\",\n    license=\"Apache 2.0\",\n    author=\"Bytedance - Seed - MLSys\",\n    author_email=\"zhangchi.usc1992@bytedance.com, gmsheng@connect.hku.hk\",\n    description=\"verl: Volcano Engine Reinforcement Learning for LLM\",\n    install_requires=install_requires,\n    extras_require=extras_require,\n    package_data={\n        \"\": [\"version/*\"],\n        \"verl\": [\n            \"trainer/config/*.yaml\",\n            \"trainer/config/*/*.yaml\",\n            \"experimental/*/config/*.yaml\",\n        ],\n    },\n    include_package_data=True,\n    long_description=long_description,\n    long_description_content_type=\"text/markdown\",\n)\n"
  },
  {
    "path": "tests/README.md",
    "content": "# Tests layout\n\nEach folder under tests/ corresponds to a test category for a sub-namespace in verl. For instance:\n- `tests/trainer` for testing functionality related to `verl/trainer`\n- `tests/models` for testing functionality related to `verl/models`\n- ...\n\nThere are a few folders with `special_` prefix, created for special purposes:\n- `special_distributed`: unit tests that must run with multiple GPUs\n- `special_e2e`: end-to-end tests with training/generation scripts\n- `special_npu`: tests for NPUs\n- `special_sanity`: a suite of quick sanity tests\n- `special_standalone`: a set of test that are designed to run in dedicated environments\n\nAccelerators for tests \n- By default tests are run with GPU available, except for the ones under `special_npu`, and any test script whose name ends with `on_cpu.py`.\n- For test scripts with `on_cpu.py` name suffix would be tested on CPU resources in linux environment.\n\n# Workflow layout\n\nAll CI tests are configured by yaml files in `.github/workflows/`. Here's an overview of all test configs:\n1. A list of always triggered CPU sanity tests: `check-pr-title.yml`, `secrets_scan.yml`, `check-pr-title,yml`, `pre-commit.yml`, `doc.yml`\n2. Some heavy multi-GPU unit tests, such as `model.yml`, `vllm.yml`, `sgl.yml`\n3. End-to-end tests: `e2e_*.yml`\n4. Unit tests\n  - `cpu_unit_tests.yml`, run pytest on all scripts with file name pattern `tests/**/test_*_on_cpu.py`\n  - `gpu_unit_tests.yml`, run pytest on all scripts with file without the `on_cpu.py` suffix.\n  - Since cpu/gpu unit tests by default runs all tests under `tests`, please make sure tests are manually excluded in them when\n    - new workflow yaml is added to `.github/workflows`\n    - new tests are added to workflow mentioned in 2."
  },
  {
    "path": "tests/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/checkpoint_engine/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/checkpoint_engine/test_correctness_on_gpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport pytest\nimport ray\n\nfrom tests.checkpoint_engine.test_utils import create_rollout_worker_group, create_trainer_worker_group\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.single_controller.ray.base import (\n    RayResourcePool,\n    split_resource_pool,\n)\nfrom verl.utils.device import get_device_name\nfrom verl.utils.ray_utils import auto_await\nfrom verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig\n\n\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"rebuild_group\", [False, True])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(2, 6)])\n@auto_await\nasync def test_nccl_checkpoint_engine(\n    rebuild_group,\n    num_trainer,\n    num_rollout,\n    num_nodes=1,\n    num_gpus_per_node=8,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-8B-Base\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"UCX_TLS\": \"rc,tcp,cuda\",\n                \"UCX_MAX_RNDV_RAILS\": \"4\",\n                \"UCX_LOG_LEVEL\": \"INFO\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(\n        backend=\"nccl\", engine_kwargs={\"nccl\": {\"rebuild_group\": rebuild_group}}\n    )\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\n@pytest.mark.skip(reason=\"temporary skip since our ci environment is not ready\")\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"device\", [\"cuda\", \"cpu\"])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(2, 6)])\n@auto_await\nasync def test_nixl_checkpoint_engine(\n    num_trainer,\n    num_rollout,\n    device,\n    num_nodes=1,\n    num_gpus_per_node=8,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-8B-Base\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                # TODO: it's pretty hard to set these environment variables right, please consult\n                # with your network admin. Maybe auto adjust UCX_* according to NCCL_IB_*?\n                \"UCX_TLS\": \"rc,ud,cuda\",\n                # \"UCX_IB_GID_INDEX\": \"3\", # NCCL_IB_GID_INDEX\n                # \"UCX_IB_DEVICES\": \"mlx5_1:1,mlx5_2:1,mlx5_3:1\", # NCCL_IB_HCA\n                \"UCX_RC_TIMEOUT\": \"30s\",  # NCCL_IB_TIMEOUT\n                \"UCX_RC_RETRY_COUNT\": \"7\",  # NCCL_IB_RETRY_COUNT\n                \"UCX_KEEPALIVE_INTERVAL\": \"1s\",\n                \"UCX_KEEPALIVE_NUM_EPS\": \"10\",\n                \"UCX_MAX_RNDV_RAILS\": \"4\",\n                \"UCX_IB_ROCE_REACHABILITY_MODE\": \"all\",\n                \"UCX_LOG_LEVEL\": \"INFO\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(backend=\"nixl\", engine_kwargs={\"nixl\": {\"device\": device}})\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\n@pytest.mark.skip(reason=\"temporary skip since our ci environment is not ready\")\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"rebuild_group\", [False])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(2, 6)])\n@auto_await\nasync def test_kimi_checkpoint_engine(\n    rebuild_group,\n    num_trainer,\n    num_rollout,\n    num_nodes=1,\n    num_gpus_per_node=8,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-8B-Base\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"NCCL_IB_HCA\": \"mlx5\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(\n        backend=\"kimi_ckpt_engine\", engine_kwargs={\"kimi_ckpt_engine\": {\"rebuild_group\": rebuild_group}}\n    )\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    resource_pool.get_placement_groups(device_name=get_device_name())\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_nccl_checkpoint_engine(\n        rebuild_group=False,\n        num_trainer=2,\n        num_rollout=30,\n        num_nodes=4,\n        num_gpus_per_node=8,\n        check_allclose=False,\n        model_path=os.environ[\"HDFS_ROOT\"] + \"/model/Qwen3-30B-A3B-Base\",\n    )\n"
  },
  {
    "path": "tests/checkpoint_engine/test_correctness_on_npu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport pytest\nimport ray\n\nfrom tests.checkpoint_engine.test_utils import create_rollout_worker_group, create_trainer_worker_group\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.single_controller.ray.base import (\n    RayResourcePool,\n    split_resource_pool,\n)\nfrom verl.utils.device import get_device_name\nfrom verl.utils.ray_utils import auto_await\nfrom verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig\n\n\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"rebuild_group\", [False])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(2, 6)])\n@auto_await\nasync def test_hccl_checkpoint_engine(\n    rebuild_group,\n    num_trainer,\n    num_rollout,\n    num_nodes=1,\n    num_gpus_per_node=8,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-8B-Base\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"HCCL_CONNECT_TIMEOUT\": \"1500\",\n                \"HCCL_HOST_SOCKET_PORT_RANGE\": \"60000-60050\",\n                \"HCCL_NPU_SOCKET_PORT_RANGE\": \"61000-61050\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(\n        backend=\"nccl\", engine_kwargs={\"nccl\": {\"rebuild_group\": rebuild_group}}\n    )\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    resource_pool.get_placement_groups(device_name=get_device_name())\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\n@pytest.mark.skip(reason=\"temporary skip since our ci environment is not ready\")\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"rebuild_group\", [False])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(4, 28)])\nasync def test_kimi_checkpoint_engine(\n    rebuild_group,\n    num_trainer,\n    num_rollout,\n    num_nodes=2,\n    num_gpus_per_node=16,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-32B\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"HCCL_CONNECT_TIMEOUT\": \"1500\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(\n        backend=\"kimi_ckpt_engine\", engine_kwargs={\"kimi_ckpt_engine\": {\"rebuild_group\": rebuild_group}}\n    )\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    resource_pool.get_placement_groups(device_name=get_device_name())\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\n@pytest.mark.skip(reason=\"temporary skip since our ci environment is not ready\")\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"device\", [\"npu\"])\n@pytest.mark.parametrize(\"num_trainer, num_rollout\", [(2, 6)])\nasync def test_mooncake_checkpoint_engine(\n    rebuild_group,\n    num_trainer,\n    num_rollout,\n    device,\n    num_nodes=1,\n    num_gpus_per_node=8,\n    check_allclose=True,\n    model_path=\"~/models/Qwen/Qwen3-8B-Base\",\n):\n    model_path = os.path.expanduser(model_path)\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"ASCEND_USE_SHORT_CONNECTION\": \"1\",\n                \"VERL_LOGGING_LEVEL\": \"DEBUG\",\n            }\n        }\n    )\n\n    # initialize config\n    checkpoint_engine_config = CheckpointEngineConfig(\n        backend=\"mooncake\", engine_kwargs={\"mooncake\": {\"device\": device, \"rebuild_group\": rebuild_group}}\n    )\n    model_config = HFModelConfig(path=model_path, use_remove_padding=True)\n    rollout_config = RolloutConfig(name=\"vllm\", checkpoint_engine=checkpoint_engine_config)\n\n    # create trainer and rollout worker group\n    resource_pool = RayResourcePool(process_on_nodes=[num_gpus_per_node] * num_nodes, max_colocate_count=3)\n    resource_pool.get_placement_groups(device_name=get_device_name())\n    trainer_pool, rollout_pool = split_resource_pool(resource_pool, [num_trainer, num_rollout])\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n    rollout, replicas = await create_rollout_worker_group(rollout_pool, model_config, rollout_config, check_allclose)\n\n    # create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(config=checkpoint_engine_config, trainer=trainer, replicas=replicas)\n    for _ in range(3):\n        await checkpoint_manager.update_weights()\n        rollout.check_weights()\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_hccl_checkpoint_engine(\n        rebuild_group=False,\n        num_trainer=2,\n        num_rollout=6,\n        num_nodes=1,\n        num_gpus_per_node=8,\n        check_allclose=False,\n        model_path=os.environ[\"HDFS_ROOT\"] + \"/model/Qwen3-30B-A3B-Base\",\n    )\n"
  },
  {
    "path": "tests/checkpoint_engine/test_special_server_adapter.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport os\n\nimport pytest\nimport ray\nfrom omegaconf import DictConfig\nfrom transformers import PreTrainedTokenizer\n\nfrom tests.checkpoint_engine.test_utils import create_trainer_worker_group\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.agent_loop.agent_loop import AgentLoopManager, AsyncLLMServerManager\nfrom verl.experimental.fully_async_policy.agent_loop.agent_loop import FullyAsyncLLMServerManager\nfrom verl.single_controller.ray import (\n    RayResourcePool,\n)\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.workers.config import CheckpointEngineConfig, HFModelConfig\n\n\n@pytest.fixture\ndef init_config() -> DictConfig:\n    from hydra import compose, initialize_config_dir\n\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(\n            config_name=\"ppo_trainer\",\n            overrides=[\n                \"+async_training.partial_rollout=True\",\n            ],\n        )\n\n    config.actor_rollout_ref.model.path = os.path.expanduser(\"~/models/Qwen/Qwen3-VL-2B-Instruct\")\n    config.actor_rollout_ref.rollout.name = os.environ[\"ROLLOUT_NAME\"]\n    config.actor_rollout_ref.rollout.max_num_seqs = 256\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.checkpoint_engine.backend = \"nccl\"\n    config.actor_rollout_ref.rollout.nnodes = 1\n    config.trainer.n_gpus_per_node = 4\n    config.trainer.nnodes = 1\n\n    return config\n\n\nasync def _run_update_weights_with_global_steps_none(\n    server_manager: AsyncLLMServerManager,\n    checkpoint_manager: CheckpointEngineManager,\n    tokenizer: PreTrainedTokenizer,\n):\n    await checkpoint_manager.update_weights(global_steps=None)\n    prompt = [{\"role\": \"user\", \"content\": \"How to make a sandwich?\"}]\n    prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True)\n    output = await server_manager.generate(\n        request_id=\"test_0\",\n        prompt_ids=prompt_ids,\n        sampling_params={\n            \"temperature\": 1.0,\n            \"logprobs\": True,\n        },\n    )\n    assert output.stop_reason not in (\"aborted\", \"abort\"), (\n        f\"output.stop_reason is {output.stop_reason}, expected not abort\"\n    )\n    assert output.extra_fields[\"global_steps\"] is None, (\n        f\"output.extra_fields['global_steps'] is {output.extra_fields['global_steps']}, expected None\"\n    )\n    print(\"========== [update_weights with global_steps=None] ==========\")\n    print(\"[RESPONSE]\", tokenizer.decode(output.token_ids, skip_special_tokens=True))\n\n\nasync def _run_server_manager_without_resume(\n    initial_steps: int,\n    train_steps: int,\n    server_manager: AsyncLLMServerManager,\n    checkpoint_manager: CheckpointEngineManager,\n    prompts: list[list[dict]],\n    tokenizer: PreTrainedTokenizer,\n):\n    for global_steps in range(initial_steps, initial_steps + train_steps):\n        tasks = []\n        for i, prompt in enumerate(prompts):\n            prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True)\n            tasks.append(\n                asyncio.create_task(\n                    server_manager.generate(\n                        request_id=f\"test_{global_steps}_{i}\",\n                        prompt_ids=prompt_ids,\n                        sampling_params={\n                            \"temperature\": 1.0,\n                            \"logprobs\": True,\n                        },\n                    )\n                )\n            )\n\n        # wait a while and update weights to interrupt the generation\n        await asyncio.sleep(2)\n        await checkpoint_manager.update_weights(global_steps=global_steps)\n\n        outputs = await asyncio.gather(*tasks)\n        expected_steps = global_steps - 1\n        for output in outputs:\n            global_steps = output.extra_fields[\"global_steps\"]\n            assert output.stop_reason in (\"aborted\", \"abort\"), (\n                f\"output.stop_reason is {output.stop_reason}, expected in abort\"\n            )\n            assert global_steps == expected_steps, f\"output.global_steps is {global_steps}, expected {expected_steps}\"\n        print(f\"========== [{initial_steps=}, {train_steps=}] ==========\")\n        print(\"[RESPONSE]\", tokenizer.decode(outputs[0].token_ids, skip_special_tokens=True))\n\n\nasync def _run_server_manager_with_resume(\n    initial_steps: int,\n    train_steps: int,\n    server_manager: FullyAsyncLLMServerManager,\n    checkpoint_manager: CheckpointEngineManager,\n    prompts: list[list[dict]],\n    tokenizer: PreTrainedTokenizer,\n):\n    # 1. rollout generate responses\n    tasks = []\n    for i, prompt in enumerate(prompts):\n        prompt_ids = tokenizer.apply_chat_template(prompt, add_generation_prompt=True, tokenize=True)\n        tasks.append(\n            asyncio.create_task(\n                server_manager.generate(\n                    request_id=f\"test_{initial_steps}_{i}\",\n                    prompt_ids=prompt_ids,\n                    sampling_params={\n                        \"temperature\": 1.0,\n                        \"logprobs\": True,\n                    },\n                )\n            )\n        )\n\n    # 2. trainer update weights to rollout multiple times\n    for global_steps in range(initial_steps, initial_steps + train_steps):\n        # wait a while and update weights to interrupt the generation\n        await asyncio.sleep(2)\n        await checkpoint_manager.update_weights(global_steps=global_steps)\n\n    # 3. wait for rollout generate responses finished\n    outputs = await asyncio.gather(*tasks)\n    expected_min_steps = initial_steps - 1\n    for output in outputs:\n        min_global_steps = output.extra_fields[\"min_global_steps\"]\n        max_global_steps = output.extra_fields[\"max_global_steps\"]\n        assert min_global_steps == expected_min_steps, (\n            f\"output.min_global_steps is {min_global_steps}, expected {expected_min_steps}\"\n        )\n        assert max_global_steps > expected_min_steps, (\n            f\"output.max_global_steps is {max_global_steps}, expected > {expected_min_steps}\"\n        )\n        assert output.stop_reason not in (\"aborted\", \"abort\"), (\n            f\"output.stop_reason is {output.stop_reason}, expected not abort\"\n        )\n    print(f\"========== [{initial_steps=}, {train_steps=}] ==========\")\n    print(\"[RESPONSE]\", tokenizer.decode(outputs[0].token_ids, skip_special_tokens=True))\n\n\n@pytest.mark.asyncio\nasync def test_server_adapter(init_config):\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n                \"VLLM_DISABLE_COMPILE_CACHE\": \"1\",\n            }\n        }\n    )\n\n    # 1. create trainer worker group\n    model_config: HFModelConfig = omega_conf_to_dataclass(init_config.actor_rollout_ref.model)\n    checkpoint_engine_config: CheckpointEngineConfig = omega_conf_to_dataclass(\n        init_config.actor_rollout_ref.rollout.checkpoint_engine\n    )\n    trainer_pool = RayResourcePool(process_on_nodes=[init_config.trainer.n_gpus_per_node], max_colocate_count=3)\n    trainer = create_trainer_worker_group(trainer_pool, model_config, checkpoint_engine_config)\n    trainer.reset()\n\n    # 2. create standalone rollout with AgentLoopManager\n    agent_loop_manager = await AgentLoopManager.create(config=init_config)\n    servers = list(\n        zip(\n            agent_loop_manager.server_addresses,\n            [server._server_handle for server in agent_loop_manager.rollout_replicas],\n            strict=True,\n        )\n    )\n    load_balancer_handle = agent_loop_manager.global_load_balancer\n\n    # 3. create checkpoint engine manager\n    checkpoint_manager = CheckpointEngineManager(\n        config=checkpoint_engine_config, trainer=trainer, replicas=agent_loop_manager.rollout_replicas\n    )\n\n    n = 4\n    prompts = [\n        [{\"role\": \"user\", \"content\": \"Please write an article about the history of China, at least 1000 words.\"}],\n        [{\"role\": \"user\", \"content\": \"Please write an article about the history of America, at least 1000 words.\"}],\n        [{\"role\": \"user\", \"content\": \"Please write an article about the geography of China, at least 1000 words.\"}],\n        [{\"role\": \"user\", \"content\": \"Please write an article about the geography of America, at least 1000 words.\"}],\n    ] * n\n\n    server_manager = AsyncLLMServerManager(\n        config=init_config, servers=servers, load_balancer_handle=load_balancer_handle\n    )\n\n    # 4. test update_weights with global_steps=None\n    await _run_update_weights_with_global_steps_none(\n        server_manager=server_manager,\n        checkpoint_manager=checkpoint_manager,\n        tokenizer=model_config.tokenizer,\n    )\n\n    # 5. test AsyncLLMServerManager without partial rollout resume\n    await checkpoint_manager.update_weights(global_steps=0)\n    await _run_server_manager_without_resume(\n        initial_steps=1,\n        train_steps=3,\n        server_manager=server_manager,\n        checkpoint_manager=checkpoint_manager,\n        prompts=prompts,\n        tokenizer=model_config.tokenizer,\n    )\n\n    # 6. test FullyAsyncLLMServerManager with partial rollout resume\n    server_manager = FullyAsyncLLMServerManager(\n        config=init_config, servers=servers, load_balancer_handle=load_balancer_handle\n    )\n    await _run_server_manager_with_resume(\n        initial_steps=4,\n        train_steps=3,\n        server_manager=server_manager,\n        checkpoint_manager=checkpoint_manager,\n        prompts=prompts,\n        tokenizer=model_config.tokenizer,\n    )\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/checkpoint_engine/test_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nfrom typing import Generator\n\nimport ray\nimport torch\nfrom transformers import AutoModelForCausalLM\n\nfrom verl.checkpoint_engine import CheckpointEngineRegistry, CheckpointEngineWorker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.device import get_device_name\nfrom verl.utils.fs import copy_to_local\nfrom verl.workers.config import CheckpointEngineConfig, FSDPEngineConfig, HFModelConfig, RolloutConfig\nfrom verl.workers.engine_workers import TrainingWorker, TrainingWorkerConfig\nfrom verl.workers.rollout import BaseRollout, RolloutReplica\n\n\nclass TrainingWorkerTest(TrainingWorker):\n    def __init__(self, config: TrainingWorkerConfig, checkpoint_engine_config: CheckpointEngineConfig) -> None:\n        super().__init__(config)\n\n        backend = checkpoint_engine_config.backend\n        bucket_size = checkpoint_engine_config.update_weights_bucket_megabytes << 20\n        engine_kwargs = checkpoint_engine_config.engine_kwargs.get(backend, {})\n        if torch.distributed.get_rank() == 0:\n            engine_kwargs[\"is_master\"] = True\n        self.checkpoint_engine = CheckpointEngineRegistry.new(backend, bucket_size=bucket_size, **engine_kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    async def update_weights(self, global_steps: int = None):\n        per_tensor_param, _ = self.engine.get_per_tensor_param()\n        await self.checkpoint_engine.send_weights(per_tensor_param)\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False)\n    def execute_checkpoint_engine(self, method: str, *args, **kwargs):\n        return getattr(self.checkpoint_engine, method)(*args, **kwargs)\n\n\nclass MockServerAdapter(BaseRollout):\n    def __init__(self, config: RolloutConfig, model_config: HFModelConfig, check_allclose: bool = True):\n        super().__init__(config, model_config, device_mesh=None)\n        self.check_allclose = check_allclose\n        self.model = None\n        self.received_weights: dict[str, torch.Tensor] = {}\n\n    async def resume(self, tags: list[str]):\n        raise NotImplementedError()\n\n    async def release(self):\n        raise NotImplementedError()\n\n    async def update_weights(\n        self,\n        weights: Generator[tuple[str, torch.Tensor], None, None],\n        **kwargs,\n    ):\n        async for name, weight in weights:\n            weight = weight.clone()\n            if self.check_allclose:\n                self.received_weights[name] = weight.clone()\n\n    def check_weights(self):\n        if not self.check_allclose:\n            return\n\n        if self.model is None:\n            local_path = copy_to_local(self.model_config.path)\n            self.model = AutoModelForCausalLM.from_pretrained(local_path, torch_dtype=torch.bfloat16, device_map=\"cpu\")\n\n        for name, weight in self.model.state_dict().items():\n            assert name in self.received_weights, f\"weight {name} not received\"\n            received = self.received_weights[name]\n            assert torch.allclose(weight.to(received.device), received), f\"weight {name} not equal\"\n        self.received_weights.clear()\n\n\nclass MockReplica(RolloutReplica):\n    async def init_hybrid(self, worker_group: RayWorkerGroup):\n        \"\"\"Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process.\n\n        Args:\n            worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized.\n        \"\"\"\n        self.workers = worker_group.workers[\n            self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1)\n        ]\n\n    def get_ray_class_with_init_args(self) -> RayClassWithInitArgs:\n        \"\"\"Get rollout worker actor class for colocated and standalone mode.\"\"\"\n        raise NotImplementedError\n\n    async def launch_servers(self):\n        \"\"\"Launch http server in each node.\"\"\"\n        raise NotImplementedError\n\n\nclass CheckpointEngineWorkerTest(CheckpointEngineWorker):\n    def __init__(\n        self, rollout_config: RolloutConfig, model_config: HFModelConfig, check_allclose: bool = True, *args, **kwargs\n    ) -> None:\n        server_adapter = MockServerAdapter(rollout_config, model_config, check_allclose)\n        super().__init__(rollout_config, model_config, server_adapter, *args, **kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def check_weights(self):\n        self.server_adapter.check_weights()\n\n\ndef create_trainer_worker_group(\n    resource_pool: RayResourcePool, model_config: HFModelConfig, checkpoint_engine_config: CheckpointEngineConfig\n) -> RayWorkerGroup:\n    engine_config = FSDPEngineConfig(forward_only=True, fsdp_size=resource_pool.world_size, strategy=\"fsdp\")\n    trainer_config = TrainingWorkerConfig(\n        model_type=\"language_model\",\n        model_config=model_config,\n        engine_config=engine_config,\n    )\n\n    ray_cls_with_init = RayClassWithInitArgs(\n        cls=ray.remote(TrainingWorkerTest),\n        config=trainer_config,\n        checkpoint_engine_config=checkpoint_engine_config,\n    )\n    ray_cls_with_init.update_options(\n        {\n            \"runtime_env\": {\n                \"env_vars\": {\n                    \"PYTORCH_CUDA_ALLOC_CONF\": \"expandable_segments:True\",\n                }\n            }\n        }\n    )\n    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name())\n    return wg\n\n\nasync def create_rollout_worker_group(\n    resource_pool: RayResourcePool,\n    model_config: HFModelConfig,\n    rollout_config: RolloutConfig,\n    check_allclose: bool = True,\n) -> tuple[RayWorkerGroup, list[MockReplica]]:\n    # create rollout worker group\n    ray_cls_with_init = RayClassWithInitArgs(\n        cls=ray.remote(CheckpointEngineWorkerTest),\n        model_config=model_config,\n        rollout_config=rollout_config,\n        check_allclose=check_allclose,\n    )\n    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name())\n\n    # create rollout replicas\n    rollout_world_size = (\n        rollout_config.tensor_model_parallel_size\n        * rollout_config.data_parallel_size\n        * rollout_config.pipeline_model_parallel_size\n    )\n    num_replicas = wg.world_size // rollout_world_size\n    replicas = []\n    for replica_rank in range(num_replicas):\n        replica = MockReplica(\n            replica_rank=replica_rank,\n            config=rollout_config,\n            model_config=model_config,\n        )\n        replicas.append(replica)\n    await asyncio.gather(*[replica.init_hybrid(wg) for replica in replicas])\n\n    return wg, replicas\n"
  },
  {
    "path": "tests/experimental/agent_loop/agent_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ray\nfrom omegaconf import DictConfig\n\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.agent_loop import AgentLoopManager\nfrom verl.experimental.reward_loop import RewardLoopManager\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\nfrom verl.single_controller.ray.base import create_colocated_worker_cls\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager, Role\nfrom verl.utils import omega_conf_to_dataclass\nfrom verl.workers.fsdp_workers import AsyncActorRolloutRefWorker\n\n\ndef init_agent_loop_manager(config: DictConfig) -> AgentLoopManager | RayWorkerGroup:\n    # =========================== 1. Create hybrid ActorRollout workers ===========================\n    actor_rollout_cls = AsyncActorRolloutRefWorker\n    role_worker_mapping = {\n        Role.ActorRollout: ray.remote(actor_rollout_cls),\n    }\n\n    global_pool_id = \"global_pool\"\n    resource_pool_spec = {\n        global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,\n    }\n    mapping = {\n        Role.ActorRollout: global_pool_id,\n    }\n    if config.reward.reward_model.enable_resource_pool:\n        mapping[Role.RewardModel] = \"reward_pool\"\n        if config.reward.reward_model.n_gpus_per_node <= 0:\n            raise ValueError(\"config.reward.reward_model.n_gpus_per_node must be greater than 0\")\n        if config.reward.reward_model.nnodes <= 0:\n            raise ValueError(\"config.reward.reward_model.nnodes must be greater than 0\")\n\n        reward_pool = [config.reward.reward_model.n_gpus_per_node] * config.reward.reward_model.nnodes\n        resource_pool_spec[\"reward_pool\"] = reward_pool\n    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n    resource_pool_manager.create_resource_pool()\n    resource_pool_to_cls = {pool: {} for pool in resource_pool_manager.resource_pool_dict.values()}\n\n    # create actor and rollout\n    resource_pool = resource_pool_manager.get_resource_pool(Role.ActorRollout)\n    actor_rollout_cls = RayClassWithInitArgs(\n        cls=role_worker_mapping[Role.ActorRollout], config=config.actor_rollout_ref, role=\"actor_rollout\"\n    )\n    resource_pool_to_cls[resource_pool][\"actor_rollout\"] = actor_rollout_cls\n\n    all_wg = {}\n    for resource_pool, class_dict in resource_pool_to_cls.items():\n        worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n        wg_dict = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=worker_dict_cls)\n        spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n        all_wg.update(spawn_wg)\n    actor_rollout_wg = all_wg[\"actor_rollout\"]\n    actor_rollout_wg.init_model()\n\n    if config.actor_rollout_ref.rollout.mode == \"sync\":\n        raise ValueError(\"Agent loop tests require async rollout mode. Please set rollout.mode=async.\")\n\n    # =========================== 2. Create AgentLoopManager ===========================\n    rm_resource_pool = (\n        resource_pool_manager.get_resource_pool(Role.RewardModel) if config.reward.reward_model.enable else None\n    )\n    reward_loop_manager = RewardLoopManager(\n        config=config,\n        rm_resource_pool=rm_resource_pool,\n    )\n    agent_loop_manager = AgentLoopManager.create(\n        config=config,\n        worker_group=actor_rollout_wg,\n        reward_loop_worker_handles=reward_loop_manager.reward_loop_workers,\n    )\n    checkpoint_manager = CheckpointEngineManager(\n        config=omega_conf_to_dataclass(config.actor_rollout_ref.rollout.checkpoint_engine),\n        trainer=actor_rollout_wg,\n        replicas=agent_loop_manager.rollout_replicas,\n    )\n    checkpoint_manager.sleep_replicas()\n    checkpoint_manager.update_weights()\n\n    return agent_loop_manager\n"
  },
  {
    "path": "tests/experimental/agent_loop/qwen_vl_tool_chat_template.jinja2",
    "content": "{% set image_count = namespace(value=0) %}\n{% set video_count = namespace(value=0) %}\n{%- if tools %}\n{{- '<|im_start|>system\\n' }}\n{%- if messages[0]['role'] == 'system' %}\n{%- if messages[0]['content'] is string %}\n{{- messages[0]['content'] }}\n{%- else %}\n{{- messages[0]['content'][0]['text'] }}\n{%- endif %}\n{%- else %}\n{{- 'You are a helpful assistant.' }}\n{%- endif %}\n{{- \"\\n\\n# Tools\\n\\nYou may call one or more functions to assist with the user query.\\n\\nYou are provided with function signatures within <tools></tools> XML tags:\\n<tools>\" }}\n{%- for tool in tools %}\n{{- \"\\n\" }}\n{{- tool | tojson }}\n{%- endfor %}\n{{- \"\\n</tools>\\n\\nFor each function call, return a json object with function name and arguments within <tool_call></tool_call> XML tags:\\n<tool_call>\\n{\\\"name\\\": <function-name>, \\\"arguments\\\": <args-json-object>}\\n</tool_call><|im_end|>\\n\" }}\n{% for message in messages %}\n{% if message['role'] != 'system' or loop.first == false %}\n{%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}\n{{ message['content'] }}<|im_end|>\n{% else %}\n{% for content in message['content'] %}\n{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}\n{% set image_count.value = image_count.value + 1 %}\n{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>\n{% elif content['type'] == 'video' or 'video' in content %}\n{% set video_count.value = video_count.value + 1 %}\n{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>\n{% elif 'text' in content %}\n{{ content['text'] }}\n{% endif %}\n{% endfor %}<|im_end|>\n{% endif %}\n{%- elif message.role == \"assistant\" %}\n{{- '<|im_start|>' + message.role }}\n{%- if message.content %}\n{{- '\\n' + message.content }}\n{%- endif %}\n{%- for tool_call in message.tool_calls %}\n{%- if tool_call.function is defined %}\n{%- set tool_call = tool_call.function %}\n{%- endif %}\n{{- '\\n<tool_call>\\n{\"name\": \"' }}\n{{- tool_call.name }}\n{{- '\", \"arguments\": ' }}\n{{- tool_call.arguments | tojson }}\n{{- '}\\n</tool_call>' }}\n{%- endfor %}\n{{- '<|im_end|>\\n' }}\n{%- elif message.role == \"tool\" %}\n{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n{{- '<|im_start|>user' }}\n{%- endif %}\n{{- '\\n<tool_response>\\n' }}\n{% if message['content'] is string %}\n{{ message.content }}\n{% else %}\n{% for content in message['content'] %}\n{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}\n{% set image_count.value = image_count.value + 1 %}\n{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>\n{% elif content['type'] == 'video' or 'video' in content %}\n{% set video_count.value = video_count.value + 1 %}\n{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>\n{% elif content['type'] == 'text' or 'text' in content %}\n{{ content['text'] }}\n{% endif %}\n{% endfor %}\n{% endif %}\n{{- '\\n</tool_response>' }}\n{%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n{{- '<|im_end|>\\n' }}\n{%- endif %}\n{%- endif %}\n{% endif %}\n{% endfor %}\n{%- else %}\n{% for message in messages %}\n{% if loop.first and message['role'] != 'system' %}\n<|im_start|>system\nYou are a helpful assistant.<|im_end|>\n{% endif %}\n{%- if (message.role == \"user\") or (message.role == \"system\" and not loop.first) or (message.role == \"assistant\" and not message.tool_calls) %}\n<|im_start|>{{ message['role'] }}\n{% if message['content'] is string %}\n{{ message['content'] }}<|im_end|>\n{% else %}\n{% for content in message['content'] %}\n{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}\n{% set image_count.value = image_count.value + 1 %}\n{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>\n{% elif content['type'] == 'video' or 'video' in content %}\n{% set video_count.value = video_count.value + 1 %}\n{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>\n{% elif 'text' in content %}\n{{ content['text'] }}\n{% endif %}\n{% endfor %}<|im_end|>\n{% endif %}\n{%- elif message.role == \"assistant\" %}\n{{- '<|im_start|>' + message.role }}\n{%- if message.content %}\n{{- '\\n' + message.content }}\n{%- endif %}\n{%- for tool_call in message.tool_calls %}\n{%- if tool_call.function is defined %}\n{%- set tool_call = tool_call.function %}\n{%- endif %}\n{{- '\\n<tool_call>\\n{\"name\": \"' }}\n{{- tool_call.name }}\n{{- '\", \"arguments\": ' }}\n{{- tool_call.arguments | tojson }}\n{{- '}\\n</tool_call>' }}\n{%- endfor %}\n{{- '<|im_end|>\\n' }}\n{%- elif message.role == \"tool\" %}\n{%- if (loop.index0 == 0) or (messages[loop.index0 - 1].role != \"tool\") %}\n{{- '<|im_start|>user' }}\n{%- endif %}\n{{- '\\n<tool_response>\\n' }}\n{% if message['content'] is string %}\n{{ message.content }}\n{% else %}\n{% for content in message['content'] %}\n{% if content['type'] == 'image' or 'image' in content or 'image_url' in content %}\n{% set image_count.value = image_count.value + 1 %}\n{% if add_vision_id %}Picture {{ image_count.value }}: {% endif %}<|vision_start|><|image_pad|><|vision_end|>\n{% elif content['type'] == 'video' or 'video' in content %}\n{% set video_count.value = video_count.value + 1 %}\n{% if add_vision_id %}Video {{ video_count.value }}: {% endif %}<|vision_start|><|video_pad|><|vision_end|>\n{% elif content['type'] == 'text' or 'text' in content %}\n{{ content['text'] }}\n{% endif %}\n{% endfor %}\n{% endif %}\n{{- '\\n</tool_response>' }}\n{%- if loop.last or (messages[loop.index0 + 1].role != \"tool\") %}\n{{- '<|im_end|>\\n' }}\n{%- endif %}\n{%- endif %}\n{% endfor %}\n{%- endif %}\n{% if add_generation_prompt %}\n<|im_start|>assistant\n{% endif %}"
  },
  {
    "path": "tests/experimental/agent_loop/test_agent_loop_extra_fields_schema_on_cpu.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport warnings\nfrom typing import Any, Optional\n\nimport numpy as np\nimport pytest\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom verl.experimental.agent_loop.agent_loop import (\n    AgentLoopMetrics,\n    AgentLoopOutput,\n    AgentLoopWorker,\n    DictConfigWrap,\n    _InternalAgentLoopOutput,\n)\nfrom verl.experimental.agent_loop.single_turn_agent_loop import SingleTurnAgentLoop\nfrom verl.utils.dataset.rl_dataset import RLHFDataset\nfrom verl.workers.rollout.replica import TokenOutput\n\n\nclass _FakeServerManager:\n    async def generate(\n        self,\n        request_id: str,\n        *,\n        prompt_ids: list[int],\n        sampling_params: dict[str, Any],\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> TokenOutput:\n        del request_id, sampling_params, image_data, video_data\n        # Return a short, deterministic \"generation\" for testing.\n        return TokenOutput(token_ids=prompt_ids[-1:] + [11, 12, 13], log_probs=[0.0, 0.0, 0.0, 0.0])\n\n    async def generate_for_partial(\n        self,\n        request_id: str,\n        *,\n        prompt_ids: list[int],\n        sampling_params: dict[str, Any],\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> tuple[list[int], list[float], bool]:\n        del request_id, sampling_params, image_data, video_data\n        # Return a short partial generation and \"not cancelled\".\n        response_ids = prompt_ids[-1:] + [21, 22]\n        response_logprobs = [0.0] * len(response_ids)\n        return response_ids, response_logprobs, False\n\n\nclass _FakeTokenizer:\n    padding_side = \"right\"\n\n    def apply_chat_template(\n        self,\n        messages: list[dict[str, Any]],\n        *,\n        tools: Optional[list[dict]] = None,\n        add_generation_prompt: bool = True,\n        tokenize: bool = True,\n        **kwargs,\n    ) -> list[int]:\n        del messages, tools, add_generation_prompt, tokenize, kwargs\n        # Minimal tokenization: return a small prompt.\n        return [101, 102]\n\n    def pad(\n        self,\n        encoded_inputs: dict[str, list[int]],\n        *,\n        padding: str,\n        max_length: int,\n        return_tensors: str,\n        return_attention_mask: bool,\n    ) -> dict[str, torch.Tensor]:\n        del padding, return_tensors\n        input_ids = encoded_inputs[\"input_ids\"]\n        if len(input_ids) > max_length:\n            if self.padding_side == \"left\":\n                input_ids = input_ids[-max_length:]\n            else:\n                input_ids = input_ids[:max_length]\n\n        pad_len = max_length - len(input_ids)\n        if self.padding_side == \"left\":\n            padded_ids = [0] * pad_len + input_ids\n            attention_mask = [0] * pad_len + [1] * len(input_ids)\n        else:\n            padded_ids = input_ids + [0] * pad_len\n            attention_mask = [1] * len(input_ids) + [0] * pad_len\n\n        output = {\"input_ids\": torch.tensor([padded_ids], dtype=torch.long)}\n        if return_attention_mask:\n            output[\"attention_mask\"] = torch.tensor([attention_mask], dtype=torch.long)\n        return output\n\n    def decode(self, ids: list[int] | torch.Tensor, skip_special_tokens: bool = True) -> str:\n        del ids, skip_special_tokens\n        return \"<decoded>\"\n\n\ndef _pad_1d(ids: list[int], *, length: int, pad_id: int = 0) -> list[int]:\n    if len(ids) > length:\n        return ids[:length]\n    return ids + [pad_id] * (length - len(ids))\n\n\ndef _to_internal(\n    *,\n    output_prompt_ids: list[int],\n    output_response_ids: list[int],\n    output_response_mask: list[int],\n    metrics: AgentLoopMetrics,\n    extra_fields: dict[str, Any],\n    num_turns: int,\n    prompt_len: int,\n    response_len: int,\n) -> _InternalAgentLoopOutput:\n    prompt_ids = _pad_1d(output_prompt_ids, length=prompt_len, pad_id=0)\n    response_ids = _pad_1d(output_response_ids, length=response_len, pad_id=0)\n    response_mask = _pad_1d(output_response_mask, length=response_len, pad_id=0)\n\n    seq_len = prompt_len + response_len\n    attention_mask = _pad_1d([1] * len(output_prompt_ids), length=prompt_len, pad_id=0) + _pad_1d(\n        [1] * len(output_response_ids),\n        length=response_len,\n        pad_id=0,\n    )\n    input_ids = prompt_ids + response_ids\n    position_ids = list(range(seq_len))\n\n    def t(x: list[int]) -> torch.Tensor:\n        return torch.tensor([x], dtype=torch.long)\n\n    return _InternalAgentLoopOutput(\n        prompt_ids=t(prompt_ids),\n        response_ids=t(response_ids),\n        response_mask=t(response_mask),\n        attention_mask=t(attention_mask),\n        input_ids=t(input_ids),\n        position_ids=t(position_ids),\n        response_logprobs=None,\n        routed_experts=None,\n        multi_modal_inputs=None,\n        multi_modal_data=None,\n        reward_score=None,\n        num_turns=num_turns,\n        metrics=metrics,\n        extra_fields=extra_fields,\n    )\n\n\n@pytest.mark.asyncio\nasync def test_agent_loop_extra_fields_schema_stable_for_training_concat_on_cpu():\n    # Minimal config surface used by the agent loops.\n    config = OmegaConf.create(\n        {\n            \"actor_rollout_ref\": {\n                \"rollout\": {\"prompt_length\": 16, \"response_length\": 16, \"multi_turn\": {\"tool_config_path\": None}},\n                \"model\": {},\n            },\n            \"data\": {\n                \"tool_config_path\": None,\n                \"apply_chat_template_kwargs\": {},\n            },\n        }\n    )\n\n    server_manager = _FakeServerManager()\n    tokenizer = _FakeTokenizer()\n    processor = None\n\n    trainer_config = DictConfigWrap(config)\n    data_config = DictConfigWrap(config.data)\n\n    single_turn = SingleTurnAgentLoop(\n        trainer_config=trainer_config,\n        server_manager=server_manager,\n        tokenizer=tokenizer,\n        processor=processor,\n        dataset_cls=RLHFDataset,\n        data_config=data_config,\n    )\n\n    raw_prompt = [{\"role\": \"user\", \"content\": \"hi\"}]\n    sampling_params: dict[str, Any] = {}\n\n    out = await single_turn.run(sampling_params=sampling_params, raw_prompt=raw_prompt)\n\n    # Agent loop outputs should always contain these fields with consistent types.\n    assert out.extra_fields[\"turn_scores\"] == []\n    assert out.extra_fields[\"tool_rewards\"] == []\n\n    internal_a = _to_internal(\n        output_prompt_ids=out.prompt_ids,\n        output_response_ids=out.response_ids,\n        output_response_mask=out.response_mask,\n        metrics=out.metrics,\n        extra_fields=out.extra_fields,\n        num_turns=out.num_turns,\n        prompt_len=len(out.prompt_ids),\n        response_len=len(out.response_ids),\n    )\n\n    # Mimic two \"worker chunks\" and concatenate as in training.\n    dummy_worker = type(\"_DummyWorker\", (), {\"reward_loop_worker_handles\": None})()\n    merged = AgentLoopWorker._postprocess(\n        dummy_worker,\n        inputs=[internal_a],\n        input_non_tensor_batch={\n            \"index\": np.array([0], dtype=object),\n            \"agent_name\": np.array([\"single_turn_agent\"], dtype=object),\n        },\n    )\n\n    # Stable schema: present regardless of which loop produced a sample.\n    stable_keys = (\n        \"turn_scores\",\n        \"tool_rewards\",\n        \"min_global_steps\",\n        \"max_global_steps\",\n        \"extras\",\n    )\n    for key in stable_keys:\n        assert key in merged.non_tensor_batch, f\"missing key in merged batch: {key}\"\n        assert merged.non_tensor_batch[key].shape == (1,), (\n            f\"invalid shape for {key}: {merged.non_tensor_batch[key].shape}\"\n        )\n\n    # And the list-typed fields are actually lists (not missing / scalar).\n    assert merged.non_tensor_batch[\"turn_scores\"][0] == []\n    assert merged.non_tensor_batch[\"tool_rewards\"][0] == []\n\n\n@pytest.mark.asyncio\nasync def test_agent_loop_postprocess_accepts_read_only_routed_experts_on_cpu():\n    class _DummyWorker:\n        _compute_multi_modal_inputs = AgentLoopWorker._compute_multi_modal_inputs\n        _compute_position_ids = AgentLoopWorker._compute_position_ids\n        _compute_score = AgentLoopWorker._compute_score\n\n        def __init__(self):\n            self.tokenizer = _FakeTokenizer()\n            self.rollout_config = OmegaConf.create({\"prompt_length\": 4, \"response_length\": 4})\n            self.processor = None\n            self.reward_loop_worker_handles = None\n\n    routed_experts = np.arange(8, dtype=np.int64).reshape(4, 2, 1)\n    routed_experts.setflags(write=False)\n    assert not routed_experts.flags.writeable\n\n    output = AgentLoopOutput(\n        prompt_ids=[101, 102],\n        response_ids=[11, 12],\n        response_mask=[1, 1],\n        routed_experts=routed_experts,\n        metrics=AgentLoopMetrics(),\n        extra_fields={},\n    )\n\n    with warnings.catch_warnings():\n        warnings.filterwarnings(\n            \"error\",\n            message=\"The given NumPy array is not writable.*\",\n            category=UserWarning,\n        )\n        internal = await AgentLoopWorker._agent_loop_postprocess(\n            _DummyWorker(),\n            output,\n            raw_prompt=[{\"role\": \"user\", \"content\": \"hi\"}],\n        )\n\n    expected = torch.tensor(routed_experts.copy()).unsqueeze(0)\n    assert internal.routed_experts is not None\n    assert internal.routed_experts.shape == (1, 8, 2, 1)\n    torch.testing.assert_close(internal.routed_experts[:, 2:6], expected)\n    assert torch.count_nonzero(internal.routed_experts[:, :2]) == 0\n    assert torch.count_nonzero(internal.routed_experts[:, 6:]) == 0\n"
  },
  {
    "path": "tests/experimental/agent_loop/test_basic_agent_loop.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport os\nfrom typing import Any\n\nimport numpy as np\nimport pytest\nimport ray\nfrom omegaconf import DictConfig\nfrom transformers.utils import get_json_schema\n\nfrom tests.experimental.agent_loop.agent_utils import init_agent_loop_manager\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.agent_loop.agent_loop import GlobalRequestLoadBalancer, get_trajectory_info\nfrom verl.protocol import DataProto\nfrom verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema\nfrom verl.tools.schemas import ToolResponse\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.workers.config import CheckpointEngineConfig\n\n\n@pytest.fixture\ndef init_config() -> DictConfig:\n    from hydra import compose, initialize_config_dir\n\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(\n            config_name=\"ppo_trainer\",\n            overrides=[\n                \"actor_rollout_ref.actor.use_dynamic_bsz=true\",\n                # test sleep/wake_up with fsdp offload\n                \"actor_rollout_ref.actor.fsdp_config.param_offload=True\",\n                \"actor_rollout_ref.actor.fsdp_config.optimizer_offload=True\",\n                \"reward.reward_manager.name=dapo\",\n                \"+reward.reward_kwargs.overlong_buffer_cfg.enable=False\",\n                \"+reward.reward_kwargs.overlong_buffer_cfg.len=3072\",\n                \"+reward.reward_kwargs.max_resp_len=4096\",\n            ],\n        )\n\n    model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n    config.actor_rollout_ref.model.path = model_path\n    config.actor_rollout_ref.rollout.name = os.environ[\"ROLLOUT_NAME\"]\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.enforce_eager = True\n    config.actor_rollout_ref.rollout.prompt_length = 4096\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.n = 4\n    config.actor_rollout_ref.rollout.agent.num_workers = 2\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = True\n\n    return config\n\n\ndef test_single_turn(init_config):\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n\n    agent_loop_manager = init_agent_loop_manager(init_config)\n\n    raw_prompts = [\n        [\n            {\n                \"role\": \"user\",\n                \"content\": \"Let's play a role playing game. Your name is Alice, your favorite color is blue.\",\n            }\n        ],\n        [{\"role\": \"user\", \"content\": \"Let's play a role playing game. Your name is Bob, your favorite color is red.\"}],\n    ]\n    batch = DataProto(\n        non_tensor_batch={\n            \"raw_prompt\": np.array(raw_prompts),\n            \"agent_name\": np.array([\"single_turn_agent\"] * len(raw_prompts)),\n            \"data_source\": np.array([\"openai/gsm8k\"] * len(raw_prompts)),\n            \"reward_model\": np.array([{\"style\": \"rule\", \"ground_truth\": \"1.0\"}] * len(raw_prompts)),\n        },\n    )\n    n = init_config.actor_rollout_ref.rollout.n\n    batch = batch.repeat(n)\n    result = agent_loop_manager.generate_sequences(prompts=batch)\n    assert len(result) == len(raw_prompts) * n\n\n    # check result\n    seq_len = result.batch[\"prompts\"].size(1) + result.batch[\"responses\"].size(1)\n    assert result.batch[\"input_ids\"].size(1) == seq_len\n    assert result.batch[\"attention_mask\"].size(1) == seq_len\n    assert result.batch[\"position_ids\"].size(1) == seq_len\n\n    if init_config.actor_rollout_ref.rollout.calculate_log_probs:\n        assert result.batch[\"rollout_log_probs\"].size(1) == result.batch[\"responses\"].size(1)\n\n    # check compute score\n    assert result.batch[\"rm_scores\"].shape == result.batch[\"responses\"].shape\n    reward_tensor = result.batch[\"rm_scores\"]\n    reward_extra_keys = result.meta_info.get(\"reward_extra_keys\", [])\n    reward_extra_info = {key: result.non_tensor_batch[key] for key in reward_extra_keys}\n    assert reward_tensor.shape == result.batch[\"responses\"].shape\n    assert \"acc\" in reward_extra_info, f\"reward_extra_info {reward_extra_info} should contain 'acc'\"\n    assert reward_extra_info[\"acc\"].shape == (len(result),), f\"invalid acc: {reward_extra_info['acc']}\"\n\n    # check turns\n    num_turns = result.non_tensor_batch[\"__num_turns__\"]\n    assert np.all(num_turns == 2)\n\n    print(\"Test passed!\")\n    ray.shutdown()\n\n\nclass WeatherTool(BaseTool):\n    def get_current_temperature(self, location: str, unit: str = \"celsius\"):\n        \"\"\"Get current temperature at a location.\n\n        Args:\n            location: The location to get the temperature for, in the format \"City, State, Country\".\n            unit: The unit to return the temperature in. Defaults to \"celsius\". (choices: [\"celsius\", \"fahrenheit\"])\n\n        Returns:\n            the temperature, the location, and the unit in a dict\n        \"\"\"\n        print(f\"[DEBUG] get_current_temperature: {location}, {unit}\")\n        return {\n            \"temperature\": 26.1,\n            \"location\": location,\n            \"unit\": unit,\n        }\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        schema = get_json_schema(self.get_current_temperature)\n        return OpenAIFunctionToolSchema(**schema)\n\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        try:\n            result = self.get_current_temperature(**parameters)\n            return ToolResponse(text=json.dumps(result)), 0, {}\n        except Exception as e:\n            return ToolResponse(text=str(e)), 0, {}\n\n\nclass WeatherToolWithData(BaseTool):\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        schema = get_json_schema(self.get_temperature_date)\n        return OpenAIFunctionToolSchema(**schema)\n\n    def get_temperature_date(self, location: str, date: str, unit: str = \"celsius\"):\n        \"\"\"Get temperature at a location and date.\n\n        Args:\n            location: The location to get the temperature for, in the format \"City, State, Country\".\n            date: The date to get the temperature for, in the format \"Year-Month-Day\".\n            unit: The unit to return the temperature in. Defaults to \"celsius\". (choices: [\"celsius\", \"fahrenheit\"])\n\n        Returns:\n            the temperature, the location, the date and the unit in a dict\n        \"\"\"\n        print(f\"[DEBUG] get_temperature_date: {location}, {date}, {unit}\")\n        return {\n            \"temperature\": 25.9,\n            \"location\": location,\n            \"date\": date,\n            \"unit\": unit,\n        }\n\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        try:\n            result = self.get_temperature_date(**parameters)\n            return ToolResponse(text=json.dumps(result)), 0, {}\n        except Exception as e:\n            return ToolResponse(text=str(e)), 0, {}\n\n\ndef test_tool_agent(init_config):\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        },\n        ignore_reinit_error=True,\n    )\n\n    # =========================== 1. Init rollout manager ===========================\n    tool_config = {\n        \"tools\": [\n            {\n                \"class_name\": \"tests.experimental.agent_loop.test_basic_agent_loop.WeatherTool\",\n                \"config\": {\"type\": \"native\"},\n            },\n            {\n                \"class_name\": \"tests.experimental.agent_loop.test_basic_agent_loop.WeatherToolWithData\",\n                \"config\": {\"type\": \"native\"},\n            },\n        ]\n    }\n    tool_config_path = \"/tmp/tool_config.json\"\n    with open(tool_config_path, \"w\") as f:\n        json.dump(tool_config, f)\n\n    n = 2\n    init_config.actor_rollout_ref.rollout.n = n\n    init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path\n    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 2\n    init_config.actor_rollout_ref.rollout.calculate_log_probs = True\n    agent_loop_manager = init_agent_loop_manager(init_config)\n\n    # =========================== 2. Generate sequences  ===========================\n    raw_prompts = [\n        [\n            {\"role\": \"user\", \"content\": \"How are you?\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What's the temperature in Los Angeles now?\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What's the temperature in New York now?\"},\n        ],\n        [\n            {\n                \"role\": \"system\",\n                \"content\": \"You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\\n\\n\"\n                \"Current Date: 2024-09-30\",\n            },\n            {\"role\": \"user\", \"content\": \"What's the temperature in San Francisco now? How about tomorrow?\"},\n        ],\n    ]\n    batch = DataProto(\n        non_tensor_batch={\n            \"raw_prompt\": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),\n            \"agent_name\": np.array([\"tool_agent\"] * len(raw_prompts)),\n            \"data_source\": np.array([\"openai/gsm8k\"] * len(raw_prompts)),\n            \"reward_model\": np.array([{\"style\": \"rule\", \"ground_truth\": \"1.0\"}] * len(raw_prompts)),\n        },\n    )\n    batch = batch.repeat(n)\n    result = agent_loop_manager.generate_sequences(prompts=batch)\n    assert len(result) == len(raw_prompts) * n\n\n    # Check turns\n    num_turns = result.non_tensor_batch[\"__num_turns__\"]\n    print(f\"num_turns: {num_turns}\")\n    for i in range(len(num_turns)):\n        if i // n == 0:\n            # [user, assistant]\n            assert num_turns[i] == 2\n        else:\n            # [user, assistant, tool, assistant]\n            assert num_turns[i] == 4\n\n    # Check response_mask\n    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)\n    responses = result.batch[\"responses\"]\n    response_mask = result.batch[\"response_mask\"]\n    attention_mask = result.batch[\"attention_mask\"]\n    assert result.batch[\"rm_scores\"].size(1) == responses.size(1)\n    assert responses.size() == response_mask.size(), f\"{responses.size()} != {response_mask.size()}\"\n    assert result.batch[\"rollout_log_probs\"].size(1) == result.batch[\"responses\"].size(1)\n\n    response_length = response_mask.size(1)\n    for i in range(len(responses)):\n        # response with tool response\n        valid_tokens = responses[i][attention_mask[i][-response_length:].bool()]\n        response_with_obs = tokenizer.decode(valid_tokens)\n\n        # response without tool response\n        valid_tokens = responses[i][response_mask[i].bool()]\n        response_without_obs = tokenizer.decode(valid_tokens)\n\n        assert \"<tool_response>\" not in response_without_obs, (\n            f\"found <tool_response> in response: {response_without_obs}\"\n        )\n        assert \"</tool_response>\" not in response_without_obs, (\n            f\"found </tool_response> in response: {response_without_obs}\"\n        )\n        print(\"=========================\")\n        print(response_with_obs)\n        print(\"---\")\n        print(response_without_obs)\n\n    print(\"Test passed!\")\n    ray.shutdown()\n\n\ndef test_tool_agent_with_interaction(init_config):\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n\n    # =========================== 1. Init rollout manager ===========================\n    tool_config = {\n        \"tools\": [\n            {\n                \"class_name\": \"tests.experimental.agent_loop.test_basic_agent_loop.WeatherTool\",\n                \"config\": {\"type\": \"native\"},\n            },\n            {\n                \"class_name\": \"tests.experimental.agent_loop.test_basic_agent_loop.WeatherToolWithData\",\n                \"config\": {\"type\": \"native\"},\n            },\n        ]\n    }\n    tool_config_path = \"/tmp/tool_config.json\"\n    with open(tool_config_path, \"w\") as f:\n        json.dump(tool_config, f)\n\n    interaction_config = {\n        \"interaction\": [\n            {\"name\": \"weather\", \"class_name\": \"verl.interactions.weather_interaction.WeatherInteraction\", \"config\": {}}\n        ]\n    }\n    interaction_config_path = \"/tmp/interaction_config.json\"\n    with open(interaction_config_path, \"w\") as f:\n        json.dump(interaction_config, f)\n\n    n = 2\n    init_config.actor_rollout_ref.rollout.n = n\n    init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path\n    init_config.actor_rollout_ref.rollout.multi_turn.interaction_config_path = interaction_config_path\n    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 2\n    agent_loop_manager = init_agent_loop_manager(init_config)\n    checkpoint_engine_config = omega_conf_to_dataclass(\n        init_config.actor_rollout_ref.rollout.checkpoint_engine, CheckpointEngineConfig\n    )\n    checkpoint_manager = CheckpointEngineManager(\n        config=checkpoint_engine_config,\n        trainer=agent_loop_manager.worker_group,\n        replicas=agent_loop_manager.rollout_replicas,\n    )\n    checkpoint_manager.sleep_replicas()\n    checkpoint_manager.update_weights()\n\n    # =========================== 2. Generate sequences  ===========================\n    raw_prompts = [\n        [\n            {\"role\": \"user\", \"content\": \"How are you?\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What's the temperature in Los Angeles now?\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What's the temperature in New York now?\"},\n        ],\n        [\n            {\n                \"role\": \"system\",\n                \"content\": \"You are Qwen, created by Alibaba Cloud. You are a helpful assistant.\\n\\n\"\n                \"Current Date: 2024-09-30\",\n            },\n            {\"role\": \"user\", \"content\": \"What's the temperature in San Francisco now? How about tomorrow?\"},\n        ],\n    ]\n    batch = DataProto(\n        non_tensor_batch={\n            \"raw_prompt\": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),\n            \"agent_name\": np.array([\"tool_agent\"] * len(raw_prompts)),\n            \"data_source\": np.array([\"openai/gsm8k\"] * len(raw_prompts)),\n            \"reward_model\": np.array([{\"style\": \"rule\", \"ground_truth\": \"1.0\"}] * len(raw_prompts)),\n            \"extra_info\": np.array(\n                [\n                    {\"interaction_kwargs\": {\"name\": \"weather\"}},\n                    {\"interaction_kwargs\": {\"name\": \"weather\"}},\n                    {\"interaction_kwargs\": {\"name\": \"weather\"}},\n                    {\"interaction_kwargs\": {\"name\": \"weather\"}},\n                ]\n            ),\n        },\n    )\n    batch = batch.repeat(n)\n    result = agent_loop_manager.generate_sequences(prompts=batch)\n    assert len(result) == len(raw_prompts) * n\n\n    # Check turns\n    num_turns = result.non_tensor_batch[\"__num_turns__\"]\n    print(f\"num_turns: {num_turns}\")\n    for i in range(len(num_turns)):\n        if i // n == 0:\n            # [user, assistant, user]\n            assert num_turns[i] == 3\n        else:\n            # [user, assistant, tool, assistant, user]\n            assert num_turns[i] == 5\n\n    # Check response_mask\n    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)\n    responses = result.batch[\"responses\"]\n    response_mask = result.batch[\"response_mask\"]\n    attention_mask = result.batch[\"attention_mask\"]\n    assert responses.size() == response_mask.size(), f\"{responses.size()} != {response_mask.size()}\"\n    response_length = response_mask.size(1)\n\n    for i in range(len(responses)):\n        # response with tool response\n        valid_tokens = responses[i][attention_mask[i][-response_length:].bool()]\n        response_with_obs = tokenizer.decode(valid_tokens)\n\n        # response without tool response\n        valid_tokens = responses[i][response_mask[i].bool()]\n        response_without_obs = tokenizer.decode(valid_tokens)\n\n        assert \"\\udb82\\udc89\" not in response_without_obs, f\"found \\udb82\\udc89 in response: {response_without_obs}\"\n        assert \"\\udb82\\udc8a\" not in response_without_obs, f\"found \\udb82\\udc8a in response: {response_without_obs}\"\n        print(\"=========================\")\n        print(response_with_obs)\n        print(\"---\")\n        print(response_without_obs)\n\n    print(\"Test passed!\")\n    ray.shutdown()\n\n\n@pytest.mark.asyncio\nasync def test_get_trajectory_info():\n    \"\"\"Tests the get_trajectory_info method.\"\"\"\n    # Initialize the class to set up class-level attributes\n    step = 10\n    index = [1, 1, 3, 3]\n    expected_info = [\n        {\"step\": step, \"sample_index\": 1, \"rollout_n\": 0, \"validate\": False},\n        {\"step\": step, \"sample_index\": 1, \"rollout_n\": 1, \"validate\": False},\n        {\"step\": step, \"sample_index\": 3, \"rollout_n\": 0, \"validate\": False},\n        {\"step\": step, \"sample_index\": 3, \"rollout_n\": 1, \"validate\": False},\n    ]\n\n    trajectory_info = await get_trajectory_info(step, index, validate=False)\n\n    assert trajectory_info == expected_info\n\n\n# ──────────────────────────────────────────────────────────────────────\n# GlobalRequestLoadBalancer unit tests (lightweight, no GPU required)\n# ──────────────────────────────────────────────────────────────────────\n\n\n@pytest.fixture(scope=\"module\")\ndef ray_for_lb():\n    ray.init(ignore_reinit_error=True)\n    yield\n    ray.shutdown()\n\n\nclass TestLoadBalancerRouting:\n    \"\"\"Least-loaded selection.\"\"\"\n\n    def test_distributes_across_servers(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\", \"s2\"])\n        servers = [ray.get(lb.acquire_server.remote(request_id=f\"r{i}\")) for i in range(3)]\n        assert sorted(servers) == [\"s0\", \"s1\", \"s2\"]\n\n    def test_new_requests_route_to_least_loaded(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\", \"s2\"])\n        # Load s0 with 3 inflight requests\n        ray.get(lb.acquire_server.remote(request_id=\"a\"))  # -> s0\n        ray.get(lb.acquire_server.remote(request_id=\"a\"))  # sticky -> s0\n        ray.get(lb.acquire_server.remote(request_id=\"a\"))  # sticky -> s0\n        # Load s1 with 1 inflight request\n        ray.get(lb.acquire_server.remote(request_id=\"b\"))  # -> s1\n        # s2 has 0 inflight, so next new request must go to s2\n        s_new = ray.get(lb.acquire_server.remote(request_id=\"d\"))\n        assert s_new == \"s2\"\n\n    def test_release_rebalances(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\"])\n        s0 = ray.get(lb.acquire_server.remote(request_id=\"r0\"))\n        s1 = ray.get(lb.acquire_server.remote(request_id=\"r1\"))\n        assert s0 != s1\n        ray.get(lb.release_server.remote(server_id=s0))\n        ray.get(lb.release_server.remote(server_id=s1))\n        s2 = ray.get(lb.acquire_server.remote(request_id=\"r2\"))\n        s3 = ray.get(lb.acquire_server.remote(request_id=\"r3\"))\n        assert s2 != s3\n\n    def test_release_invalid_server_raises(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\"])\n        with pytest.raises(ray.exceptions.RayTaskError, match=\"Invalid server_id\") as excinfo:\n            ray.get(lb.release_server.remote(server_id=\"nonexistent\"))\n        assert \"Invalid server_id\" in str(excinfo.value)\n\n    def test_release_without_inflight_raises(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\"])\n        with pytest.raises(ray.exceptions.RayTaskError, match=\"no inflight\") as excinfo:\n            ray.get(lb.release_server.remote(server_id=\"s1\"))\n        assert \"no inflight\" in str(excinfo.value)\n\n\nclass TestLoadBalancerStickySession:\n    \"\"\"Request-level sticky session.\"\"\"\n\n    def test_same_request_id_same_server(self, ray_for_lb):\n        lb = GlobalRequestLoadBalancer.remote(server_actor_ids=[\"s0\", \"s1\", \"s2\", \"s3\"])\n        s0 = ray.get(lb.acquire_server.remote(request_id=\"conv-abc\"))\n        ray.get(lb.release_server.remote(server_id=s0))\n        s1 = ray.get(lb.acquire_server.remote(request_id=\"conv-abc\"))\n        assert s0 == s1\n"
  },
  {
    "path": "tests/experimental/agent_loop/test_gpt_oss_tool_parser.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport pytest\nfrom transformers import AutoTokenizer\n\nfrom verl.experimental.agent_loop.tool_parser import GptOssToolParser\n\n\n@pytest.mark.asyncio\n@pytest.mark.skip(reason=\"local test only\")\nasync def test_gpt_oss_tool_parser():\n    example_text = \"\"\"\n<|start|>assistant<|channel|>commentary to=functions.get_current_weather \\\n<|constrain|>json<|message|>{\"location\": \"Tokyo\"}<|call|>\n<|start|>functions.get_current_weather to=assistant<|channel|>commentary<|message|>\\\n{ \"temperature\": 20, \"sunny\": true }<|end|>\"\"\"\n    tokenizer = AutoTokenizer.from_pretrained(\"openai/gpt-oss-20b\")\n    response_ids = tokenizer.encode(example_text)\n    tool_parser = GptOssToolParser(tokenizer)\n    _, function_calls = await tool_parser.extract_tool_calls(response_ids)\n    assert len(function_calls) == 1\n    assert function_calls[0].name == \"get_current_weather\"\n    assert function_calls[0].arguments == '{\"location\": \"Tokyo\"}'\n"
  },
  {
    "path": "tests/experimental/agent_loop/test_multi_modal.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport os\nfrom typing import Any\n\nimport numpy as np\nimport pytest\nimport ray\nfrom omegaconf import DictConfig\nfrom PIL import Image\nfrom transformers.utils import get_json_schema\n\nfrom tests.experimental.agent_loop.agent_utils import init_agent_loop_manager\nfrom verl.protocol import DataProto\nfrom verl.tools.base_tool import BaseTool, OpenAIFunctionToolSchema\nfrom verl.tools.schemas import ToolResponse\nfrom verl.utils import hf_tokenizer\n\n\ndef parse_multi_modal_type(messages: list[dict]) -> str:\n    message = messages[-1]\n    if isinstance(message[\"content\"], str):\n        return \"text\"\n\n    for content in message[\"content\"]:\n        if content[\"type\"] == \"image\":\n            return \"image\"\n        elif content[\"type\"] == \"video\":\n            return \"video\"\n\n    return \"text\"\n\n\n@pytest.fixture\ndef init_config() -> DictConfig:\n    from hydra import compose, initialize_config_dir\n\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(\n            config_name=\"ppo_trainer\",\n            overrides=[\n                \"actor_rollout_ref.actor.use_dynamic_bsz=true\",\n                # test sleep/wake_up with fsdp offload\n                \"actor_rollout_ref.actor.fsdp_config.param_offload=True\",\n                \"actor_rollout_ref.actor.fsdp_config.optimizer_offload=True\",\n            ],\n        )\n\n    model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-VL-3B-Instruct\")\n    config.actor_rollout_ref.model.path = model_path\n    config.actor_rollout_ref.rollout.name = os.environ[\"ROLLOUT_NAME\"]\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.enforce_eager = True\n    config.actor_rollout_ref.rollout.prompt_length = 10240\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.n = 4\n    config.actor_rollout_ref.rollout.agent.num_workers = 2\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = True\n\n    return config\n\n\nclass ImageGeneratorTool(BaseTool):\n    def generate_image(self, description: str, size: str = \"256x256\"):\n        \"\"\"Generate a simple image based on description.\n\n        Args:\n            description: The description of the image to generate.\n            size: The size of the image. Defaults to \"256x256\". (choices: [\"256x256\", \"512x512\"])\n\n        Returns:\n            A generated image\n        \"\"\"\n        print(f\"[DEBUG] generate_image: {description}, {size}\")\n        # Create a simple colored image for testing\n        width, height = map(int, size.split(\"x\"))\n\n        # Create different colors based on description\n        if \"red\" in description.lower():\n            color = (255, 0, 0)\n        elif \"blue\" in description.lower():\n            color = (0, 0, 255)\n        elif \"green\" in description.lower():\n            color = (0, 255, 0)\n        else:\n            color = (128, 128, 128)  # gray\n\n        # Create image\n        image = Image.new(\"RGB\", (width, height), color)\n\n        # Add some pattern to make it more interesting\n        for i in range(0, width, 50):\n            for j in range(0, height, 50):\n                # Add white squares in a grid pattern\n                for x in range(i, min(i + 20, width)):\n                    for y in range(j, min(j + 20, height)):\n                        image.putpixel((x, y), (255, 255, 255))\n\n        return image\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        schema = get_json_schema(self.generate_image)\n        return OpenAIFunctionToolSchema(**schema)\n\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        try:\n            image = self.generate_image(**parameters)\n            # Return the PIL Image directly - the framework should handle the conversion\n            return ToolResponse(image=[image]), 0, {}\n        except Exception as e:\n            return ToolResponse(text=str(e)), 0, {}\n\n\n@pytest.mark.flaky(reruns=3)\ndef test_multimodal_tool_agent(init_config):\n    \"\"\"Test agent loop with multimodal tool that returns images using Qwen VL model.\"\"\"\n    ray.shutdown()\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        },\n        ignore_reinit_error=True,\n    )\n\n    # Add custom chat template to enable tool calling support (same as recipe/deepeyes)\n    template_path = os.path.join(os.path.dirname(__file__), \"qwen_vl_tool_chat_template.jinja2\")\n    with open(template_path, encoding=\"utf-8\") as f:\n        custom_chat_template = f.read()\n\n    init_config.actor_rollout_ref.model.custom_chat_template = custom_chat_template\n\n    # =========================== 1. Init rollout manager with image tool ===========================\n    tool_config = {\n        \"tools\": [\n            {\n                \"class_name\": \"tests.experimental.agent_loop.test_multi_modal.ImageGeneratorTool\",\n                \"config\": {\"type\": \"native\"},\n            },\n        ]\n    }\n    tool_config_path = \"/tmp/multimodal_tool_config.json\"\n    with open(tool_config_path, \"w\") as f:\n        json.dump(tool_config, f)\n\n    n = 2\n    init_config.actor_rollout_ref.rollout.n = n\n    init_config.actor_rollout_ref.rollout.multi_turn.tool_config_path = tool_config_path\n    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1\n    init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1\n    agent_loop_manager = init_agent_loop_manager(init_config)\n\n    # =========================== 2. Generate sequences with multimodal prompts ===========================\n    raw_prompts = [\n        [\n            {\"role\": \"user\", \"content\": \"How are you?\"},\n        ],\n        [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\n                        \"type\": \"video\",\n                        \"video\": os.path.expanduser(\"~/models/hf_data/test-videos/space_woaudio.mp4\"),\n                        \"min_pixels\": 4 * 32 * 32,\n                        \"max_pixels\": 256 * 32 * 32,\n                        \"total_pixels\": 4096 * 32 * 32,\n                    },\n                    {\n                        \"type\": \"text\",\n                        \"text\": \"Describe this video. Then you must call the \"\n                        \"image generator tool to generate a green image for me.\",\n                    },\n                ],\n            },\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"Please generate a red image for me.\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"Can you create a blue picture with size 512x512?\"},\n        ],\n        [\n            {\n                \"role\": \"system\",\n                \"content\": (\n                    \"You are Qwen VL, created by Alibaba Cloud. You are a helpful \"\n                    \"assistant that can generate and analyze images.\"\n                ),\n            },\n            {\"role\": \"user\", \"content\": \"Generate a green landscape image and describe what you see in it.\"},\n        ],\n    ]\n\n    batch = DataProto(\n        non_tensor_batch={\n            \"raw_prompt\": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),\n            \"agent_name\": np.array([\"tool_agent\"] * len(raw_prompts)),\n            \"data_source\": np.array([\"openai/gsm8k\"] * len(raw_prompts)),\n            \"reward_model\": np.array([{\"style\": \"rule\", \"ground_truth\": \"1.0\"}] * len(raw_prompts)),\n        },\n    )\n    batch = batch.repeat(n)\n    result = agent_loop_manager.generate_sequences(prompts=batch)\n    assert len(result) == len(raw_prompts) * n\n\n    # Check turns\n    num_turns = result.non_tensor_batch[\"__num_turns__\"]\n    multi_modal_inputs = result.non_tensor_batch[\"multi_modal_inputs\"]\n    print(f\"num_turns: {num_turns}\")\n    for i in range(len(num_turns)):\n        multi_modal_type = parse_multi_modal_type(raw_prompts[i // n])\n        if multi_modal_type == \"video\":\n            assert \"pixel_values_videos\" in multi_modal_inputs[i], f\"Sample {i} should have pixel_values_videos\"\n            assert \"video_grid_thw\" in multi_modal_inputs[i], f\"Sample {i} should have video_grid_thw\"\n\n        if i // n == 0:\n            # First prompt: \"How are you?\" - should have 2 turns [user, assistant]\n            assert num_turns[i] == 2, f\"Expected 2 turns but got {num_turns[i]} for sample {i}\"\n        elif i // n == 1:\n            # TODO: prompt with video not generate tool call as expected\n            assert num_turns[i] == 2 or num_turns[i] == 4, (\n                f\"Expected 2 or 4 turns but got {num_turns[i]} for sample {i}\"\n            )\n        else:\n            # Tool-calling prompts should have 4 turns [user, assistant, tool, assistant]\n            assert num_turns[i] == 4, f\"Expected 4 turns but got {num_turns[i]} for sample {i}\"\n            assert \"pixel_values\" in multi_modal_inputs[i], f\"Sample {i} should have pixel_values\"\n            assert \"image_grid_thw\" in multi_modal_inputs[i], f\"Sample {i} should have image_grid_thw\"\n\n    # Check that images were properly returned in the tool responses\n    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)\n    responses = result.batch[\"responses\"]\n    response_mask = result.batch[\"response_mask\"]\n    attention_mask = result.batch[\"attention_mask\"]\n    assert responses.size() == response_mask.size(), f\"{responses.size()} != {response_mask.size()}\"\n    response_length = response_mask.size(1)\n\n    image_found_count = 0\n    for i in range(len(responses)):\n        # response with tool response (including images)\n        valid_tokens = responses[i][attention_mask[i][-response_length:].bool()]\n        response_with_obs = tokenizer.decode(valid_tokens)\n\n        # response without tool response\n        valid_tokens = responses[i][response_mask[i].bool()]\n        response_without_obs = tokenizer.decode(valid_tokens)\n\n        # Check that tool responses were properly masked out from training\n        assert \"<tool_response>\" not in response_without_obs, (\n            f\"found <tool_response> in response: {response_without_obs}\"\n        )\n        assert \"</tool_response>\" not in response_without_obs, (\n            f\"found </tool_response> in response: {response_without_obs}\"\n        )\n\n        # Check that images were included in the full response\n        if \"<image>\" in response_with_obs or \"image\" in response_with_obs.lower():\n            image_found_count += 1\n\n        print(\"=========================\")\n        print(\"Response with tool observations:\")\n        print(response_with_obs)\n        print(\"---\")\n        print(\"Response without tool observations:\")\n        print(response_without_obs)\n\n    # Verify that tool-calling responses contained image-related content\n    print(f\"Found {image_found_count} responses with image content out of {len(responses)}\")\n    # We should have at least some image content from the tool-calling prompts\n    # Note: First prompt might not use tools, so we don't expect 100% image content\n    expected_tool_calls = sum(1 for i in range(len(num_turns)) if num_turns[i] == 4)\n    assert image_found_count >= 0, (\n        f\"No image-related content found, but expected at least some from {expected_tool_calls} tool calls\"\n    )\n\n    print(\"Multimodal tool test passed!\")\n    ray.shutdown()\n\n\ndef test_multimodal_single_turn_agent(init_config):\n    \"\"\"Test single turn agent loop with multimodal inputs using Qwen VL model.\"\"\"\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        },\n        ignore_reinit_error=True,\n    )\n\n    # =========================== 1. Init rollout manager ===========================\n    n = 2\n    init_config.actor_rollout_ref.rollout.n = n\n    init_config.actor_rollout_ref.rollout.multi_turn.max_parallel_calls = 1\n    init_config.actor_rollout_ref.rollout.multi_turn.max_user_turns = 1\n    agent_loop_manager = init_agent_loop_manager(init_config)\n\n    # =========================== 2. Generate sequences with multimodal prompts ===========================\n    # Create a simple test image\n    test_image = Image.new(\"RGB\", (256, 256), (100, 150, 200))\n    test_image2 = Image.new(\"RGB\", (512, 512), (100, 150, 200))\n\n    raw_prompts = [\n        # text\n        [\n            {\"role\": \"user\", \"content\": \"Hello, how are you?\"},\n        ],\n        # image\n        [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"image\", \"image\": test_image},\n                    {\"type\": \"text\", \"text\": \"What color is this image?\"},\n                ],\n            },\n        ],\n        # system + image\n        [\n            {\n                \"role\": \"system\",\n                \"content\": \"You are Qwen VL, created by Alibaba Cloud. You are a helpful assistant.\",\n            },\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"image\", \"image\": test_image2},\n                    {\"type\": \"text\", \"text\": \"Describe this image in detail.\"},\n                ],\n            },\n        ],\n        # video\n        [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\n                        \"type\": \"video\",\n                        \"video\": os.path.expanduser(\"~/models/hf_data/test-videos/space_woaudio.mp4\"),\n                        \"min_pixels\": 4 * 32 * 32,\n                        \"max_pixels\": 256 * 32 * 32,\n                        \"total_pixels\": 4096 * 32 * 32,\n                    },\n                    {\"type\": \"text\", \"text\": \"Describe this video.\"},\n                ],\n            },\n        ],\n    ]\n\n    batch = DataProto(\n        non_tensor_batch={\n            \"raw_prompt\": np.array([np.array(prompt) for prompt in raw_prompts], dtype=object),\n            \"agent_name\": np.array([\"single_turn_agent\"] * len(raw_prompts)),\n            \"data_source\": np.array([\"openai/gsm8k\"] * len(raw_prompts)),\n            \"reward_model\": np.array([{\"style\": \"rule\", \"ground_truth\": \"1.0\"}] * len(raw_prompts)),\n        },\n    )\n\n    batch = batch.repeat(n)\n    result = agent_loop_manager.generate_sequences(prompts=batch)\n    assert len(result) == len(raw_prompts) * n\n\n    # Check turns - all should be single turn (2: user + assistant)\n    num_turns = result.non_tensor_batch[\"__num_turns__\"]\n    print(f\"num_turns: {num_turns}\")\n    for i in range(len(num_turns)):\n        assert num_turns[i] == 2, f\"Expected 2 turns but got {num_turns[i]} for sample {i}\"\n\n    # Verify responses\n    tokenizer = hf_tokenizer(init_config.actor_rollout_ref.model.path)\n    prompts = result.batch[\"prompts\"]\n    responses = result.batch[\"responses\"]\n    response_mask = result.batch[\"response_mask\"]\n    input_ids = result.batch[\"input_ids\"]\n    position_ids = result.batch[\"position_ids\"]\n    multi_modal_inputs = result.non_tensor_batch[\"multi_modal_inputs\"]\n    assert responses.size() == response_mask.size(), f\"{responses.size()} != {response_mask.size()}\"\n    assert position_ids.size() == (input_ids.size(0), 4, input_ids.size(1))  # (batch_size, 4, seq_len)\n\n    # Check for image pads in prompts\n    image_pad_count = 0\n    for i in range(len(prompts)):\n        prompt_ids = prompts[i][prompts[i] != tokenizer.pad_token_id].tolist()\n        prompt_text = tokenizer.decode(prompt_ids)\n\n        # Check if this sample should have image pads (samples with index 1 and 2 in each repeat have images)\n        sample_idx = i // n\n        has_image_pad = \"<|image_pad|>\" in prompt_text or \"<|vision_start|>\" in prompt_text\n\n        print(\"=========================\")\n        print(f\"Sample {i} (original prompt index: {sample_idx}):\")\n        print(f\"Prompt length: {len(prompt_ids)} tokens\")\n        print(f\"Has image_pad: {has_image_pad}\")\n\n        # Check multi-modal type\n        multi_modal_type = parse_multi_modal_type(raw_prompts[sample_idx])\n\n        if multi_modal_type == \"text\":\n            assert len(multi_modal_inputs[i]) == 0, f\"Sample {i} should not have multi-modal inputs\"\n        elif multi_modal_type == \"image\":\n            assert \"pixel_values\" in multi_modal_inputs[i], f\"Sample {i} should have pixel_values\"\n            assert \"image_grid_thw\" in multi_modal_inputs[i], f\"Sample {i} should have image_grid_thw\"\n        else:\n            assert \"pixel_values_videos\" in multi_modal_inputs[i], f\"Sample {i} should have pixel_values_videos\"\n            assert \"video_grid_thw\" in multi_modal_inputs[i], f\"Sample {i} should have video_grid_thw\"\n\n        # Show first 200 chars of prompt\n        print(f\"Prompt text (first 200 chars): {prompt_text[:200]}...\")\n\n    for i in range(len(responses)):\n        valid_tokens = responses[i][response_mask[i].bool()]\n        response_text = tokenizer.decode(valid_tokens)\n        print(f\"Sample {i} response: {response_text[:100]}...\")\n\n    # Verify that we found image pads in multimodal samples\n    expected_multimodal_samples = 2 * n  # 2 prompts with images, repeated n times\n    print(f\"\\nFound {image_pad_count} samples with image_pad out of {expected_multimodal_samples} expected\")\n\n    print(\"Single turn multimodal test passed!\")\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/agent_loop/test_standalone_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport os\n\nimport pytest\nimport ray\nfrom omegaconf import DictConfig\nfrom openai import AsyncOpenAI, OpenAI\n\nfrom tests.experimental.agent_loop.agent_utils import init_agent_loop_manager\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.utils import omega_conf_to_dataclass\nfrom verl.workers.rollout.replica import get_rollout_replica_class\n\n\n@pytest.fixture\ndef init_config() -> DictConfig:\n    from hydra import compose, initialize_config_dir\n\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    config.trainer.n_gpus_per_node = 4\n    config.trainer.nnodes = 2\n    config.actor_rollout_ref.actor.use_dynamic_bsz = True\n    config.actor_rollout_ref.model.path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n    config.actor_rollout_ref.rollout.name = os.environ[\"ROLLOUT_NAME\"]\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = False\n\n    return config\n\n\n@pytest.mark.asyncio\n@pytest.mark.parametrize(\"tp_size\", [2, 4])\nasync def test_standalone_rollout(init_config, tp_size):\n    \"\"\"Test standalone rollout single node and multi nodes.\"\"\"\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n                \"NCCL_P2P_DISABLE\": \"1\",  # disable p2p in L20\n            }\n        }\n    )\n\n    init_config.actor_rollout_ref.rollout.tensor_model_parallel_size = tp_size\n    num_replicas = (init_config.trainer.n_gpus_per_node * init_config.trainer.nnodes) // tp_size\n    rollout_config = init_config.actor_rollout_ref.rollout\n    model_config = init_config.actor_rollout_ref.model\n\n    # create standalone rollout server\n    rollout_server_class = get_rollout_replica_class(init_config.actor_rollout_ref.rollout.name)\n    rollout_servers = [\n        rollout_server_class(\n            replica_rank=replica_rank,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=init_config.trainer.n_gpus_per_node,\n        )\n        for replica_rank in range(num_replicas)\n    ]\n    await asyncio.gather(*[server.init_standalone() for server in rollout_servers])\n\n    server_handles = [server._server_handle for server in rollout_servers]\n    server_addresses = [server._server_address for server in rollout_servers]\n    assert len(server_handles) == num_replicas\n    assert len(server_addresses) == num_replicas\n\n    os.environ.pop(\"HTTPS_PROXY\", None)\n    os.environ.pop(\"HTTP_PROXY\", None)\n    os.environ.pop(\"NO_PROXY\", None)\n\n    client = AsyncOpenAI(\n        api_key=\"123-abc\",\n        base_url=f\"http://{server_addresses[0]}/v1\",\n    )\n\n    completion = await client.chat.completions.create(\n        model=init_config.actor_rollout_ref.model.path,\n        messages=[{\"role\": \"user\", \"content\": \"What can you do?\"}],\n    )\n    print(completion.choices[0].message.content)\n\n    ray.shutdown()\n\n\n@pytest.mark.skip(reason=\"local test only\")\ndef test_hybrid_rollout_with_ep(init_config):\n    \"\"\"Test hybrid rollout with expert parallelism, DP=2, TP=4, EP=8.\"\"\"\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n\n    model_path = os.path.expanduser(\"~/models/Qwen/Qwen3-30B-A3B-Instruct-2507\")\n    init_config.actor_rollout_ref.model.path = model_path\n\n    # parallelism config\n    init_config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n    init_config.actor_rollout_ref.rollout.data_parallel_size = 4\n    init_config.actor_rollout_ref.rollout.expert_parallel_size = 8\n\n    # 1. init hybrid worker: FSDP+rollout\n    # - build FSDP model and optimizer\n    # - offload FSDP model and optimizer, build rollout\n    # - sleep rollout and load FSDP model and optimizer\n    agent_loop_manager = init_agent_loop_manager(init_config)\n    checkpoint_manager = CheckpointEngineManager(\n        config=omega_conf_to_dataclass(init_config.actor_rollout_ref.rollout.checkpoint_engine),\n        trainer=agent_loop_manager.worker_group,\n        replicas=agent_loop_manager.rollout_replicas,\n    )\n    checkpoint_manager.sleep_replicas()\n    checkpoint_manager.update_weights()\n\n    # 3. test async openai call\n    server_address = agent_loop_manager.server_addresses[0]\n    client = OpenAI(\n        api_key=\"123-abc\",\n        base_url=f\"http://{server_address}/v1\",\n    )\n\n    smapling_params = {\n        \"temperature\": 1.0,\n        \"top_p\": 1.0,\n        \"max_tokens\": 512,\n    }\n\n    response = client.chat.completions.create(\n        model=model_path,\n        messages=[{\"role\": \"user\", \"content\": \"What can you do?\"}],\n        **smapling_params,\n    )\n\n    completion = response.choices[0].message.content\n    print(f\"response: {completion}\")\n\n    print(\"Test passed!\")\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/reward_loop/reward_fn.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\n\nimport aiohttp\nfrom openai.types.chat import ChatCompletion\nfrom transformers import PreTrainedTokenizer\n\nGRM_PROMPT_TEMPLATE = \"\"\"\nYou are given a problem and a proposed solution.\n\nProblem:\n{problem}\n\nSolution:\n{solution}\n\nPlease evaluate how well the solution addresses the problem. \nGive a score from 1 to 10, where:\n- 1 means the solution is completely irrelevant or incorrect.\n- 5 means the solution is partially correct but incomplete or not well reasoned.\n- 10 means the solution is fully correct, well-reasoned, and directly solves the problem.\n\nOnly output the score as a single number (integer).\n\"\"\".strip()\n\n\nasync def chat_complete(router_address: str, chat_complete_request: dict):\n    url = f\"http://{router_address}/v1/chat/completions\"\n    try:\n        timeout = aiohttp.ClientTimeout(total=None)\n        session = aiohttp.ClientSession(timeout=timeout)\n        async with session.post(url, json=chat_complete_request) as resp:\n            output = await resp.text()\n            output = json.loads(output)\n            return ChatCompletion(**output)\n    except Exception as e:\n        raise e\n    finally:\n        await session.close()\n\n\nasync def compute_score_gsm8k(\n    data_source: str,\n    solution_str: str,\n    ground_truth: str,\n    extra_info: dict,\n    reward_router_address: str,\n    reward_model_tokenizer: PreTrainedTokenizer,\n):\n    \"\"\"Compute the reward score.\"\"\"\n\n    grm_prompt = GRM_PROMPT_TEMPLATE.format(problem=extra_info[\"question\"], solution=solution_str)\n    messages = [{\"role\": \"user\", \"content\": grm_prompt}]\n    sampling_params = {\"temperature\": 0.7, \"top_p\": 0.8, \"max_tokens\": 4096}\n    model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n    chat_complete_request = {\n        \"messages\": messages,\n        \"model\": model_name,\n        **sampling_params,\n    }\n    result = await chat_complete(\n        router_address=reward_router_address,\n        chat_complete_request=chat_complete_request,\n    )\n    grm_response = result.choices[0].message.content\n    try:\n        score = int(grm_response.split(\"\\n\\n\")[-1].strip())\n    except Exception:\n        score = 0\n    return {\"score\": score, \"acc\": score == 10, \"genrm_response\": grm_response}\n\n\ndef compute_score_math_verify(\n    data_source: str,\n    solution_str: str,\n    ground_truth: str,\n    extra_info: dict,\n    **kwargs,\n):\n    \"\"\"Compute the reward score.\"\"\"\n    from verl.utils.reward_score.math_verify import compute_score\n\n    return compute_score(\n        model_output=solution_str,\n        ground_truth=ground_truth,\n    )\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_agent_reward_loop_colocate.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport ray\nfrom hydra import compose, initialize_config_dir\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom transformers import AutoTokenizer\n\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.agent_loop import AgentLoopManager\nfrom verl.experimental.reward_loop import RewardLoopManager\nfrom verl.protocol import DataProto\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\nfrom verl.trainer.main_ppo import create_rl_sampler\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.utils import omega_conf_to_dataclass\nfrom verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn\nfrom verl.utils.device import get_device_name\nfrom verl.workers.fsdp_workers import ActorRolloutRefWorker, AsyncActorRolloutRefWorker\n\n\ndef test_agent_reward_loop_standalone():\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    rollout_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    reward_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n\n    # actor_rollout_ref config\n    config.data.return_raw_chat = True\n    config.data.max_prompt_length = 1024\n    config.data.max_response_length = 4096\n    config.actor_rollout_ref.model.path = rollout_model_path\n    config.actor_rollout_ref.actor.use_dynamic_bsz = True\n    config.actor_rollout_ref.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n    config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.8\n    config.actor_rollout_ref.rollout.enforce_eager = True\n    config.actor_rollout_ref.rollout.prompt_length = 1024\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = True\n    config.trainer.n_gpus_per_node = 8\n    config.trainer.nnodes = 1\n\n    config.reward.reward_manager.name = \"dapo\"\n    config.reward.reward_model.enable = True\n    config.reward.reward_model.enable_resource_pool = False\n    config.reward.reward_model.n_gpus_per_node = 8\n    config.reward.reward_model.model_path = reward_model_path\n    config.reward.reward_model.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.reward.reward_model.rollout.gpu_memory_utilization = 0.8\n    config.reward.reward_model.rollout.tensor_model_parallel_size = 2\n    config.reward.reward_model.rollout.skip_tokenizer_init = False\n    config.reward.reward_model.rollout.prompt_length = 5120\n    config.reward.reward_model.rollout.response_length = 4096\n    config.reward.custom_reward_function.path = \"tests/experimental/reward_loop/reward_fn.py\"\n    config.reward.custom_reward_function.name = \"compute_score_gsm8k\"\n\n    # 1. init reward model manager\n    actor_rollout_cls = (\n        AsyncActorRolloutRefWorker if config.actor_rollout_ref.rollout.mode == \"async\" else ActorRolloutRefWorker\n    )\n    global_pool_id = \"global_pool\"\n    resource_pool_spec = {\n        global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,\n    }\n    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=None)\n    resource_pool_manager.create_resource_pool()\n    resource_pool = resource_pool_manager.resource_pool_dict[global_pool_id]\n    actor_rollout_cls = RayClassWithInitArgs(\n        cls=ray.remote(actor_rollout_cls), config=config.actor_rollout_ref, role=\"actor_rollout\"\n    )\n    actor_rollout_wg = RayWorkerGroup(\n        resource_pool=resource_pool, ray_cls_with_init=actor_rollout_cls, device_name=get_device_name()\n    )\n    actor_rollout_wg.init_model()\n\n    agent_loop_manager = AgentLoopManager.create(\n        config=config,\n        worker_group=actor_rollout_wg,\n    )\n    # sleep rollout replicas\n    checkpoint_manager = CheckpointEngineManager(\n        config=omega_conf_to_dataclass(config.actor_rollout_ref.rollout.checkpoint_engine),\n        trainer=actor_rollout_wg,\n        replicas=agent_loop_manager.rollout_replicas,\n    )\n    checkpoint_manager.sleep_replicas()\n    reward_loop_manager = RewardLoopManager(config, rm_resource_pool=resource_pool)\n\n    # 2. init test data\n    local_folder = os.path.expanduser(\"~/data/gsm8k/\")\n\n    data_files = [os.path.join(local_folder, \"train.parquet\")]\n    tokenizer = AutoTokenizer.from_pretrained(rollout_model_path)\n\n    dataset = RLHFDataset(\n        data_files=data_files,\n        tokenizer=tokenizer,\n        config=config.data,\n        processor=None,\n    )\n\n    batch_size = 64\n    sampler = create_rl_sampler(config.data, dataset)\n    dataloader = StatefulDataLoader(\n        dataset=dataset,\n        batch_size=batch_size,\n        num_workers=config.data.dataloader_num_workers,\n        drop_last=True,\n        collate_fn=collate_fn,\n        sampler=sampler,\n    )\n\n    # 3. generate responses\n    batch_dict = next(iter(dataloader))\n    batch = DataProto.from_single_dict(batch_dict)\n\n    def _get_gen_batch(batch: DataProto) -> DataProto:\n        reward_keys = set({\"data_source\", \"reward_model\", \"extra_info\", \"uid\"}) & batch.non_tensor_batch.keys()\n\n        # pop those keys for generation\n        batch_keys_to_pop = []\n        non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_keys\n        gen_batch = batch.pop(\n            batch_keys=batch_keys_to_pop,\n            non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),\n        )\n\n        # For agent loop, we need reward model keys to compute score.\n        gen_batch.non_tensor_batch.update(batch.non_tensor_batch)\n\n        return gen_batch\n\n    # wake up rollout replicas via update_weight\n    checkpoint_manager.update_weights()\n    gen_batch = _get_gen_batch(batch)\n    gen_batch = agent_loop_manager.generate_sequences(gen_batch)\n    checkpoint_manager.sleep_replicas()\n\n    batch = batch.union(gen_batch)\n    rm_outputs = reward_loop_manager.compute_rm_score(batch)\n\n    for output in rm_outputs[:5]:\n        print(output.non_tensor_batch)\n\n    print(\"done\")\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_agent_reward_loop_standalone.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport ray\nfrom hydra import compose, initialize_config_dir\nfrom torchdata.stateful_dataloader import StatefulDataLoader\n\nfrom verl.experimental.agent_loop import AgentLoopManager\nfrom verl.experimental.reward_loop import RewardLoopManager\nfrom verl.protocol import DataProto\nfrom verl.trainer.main_ppo import create_rl_sampler\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn\n\n\ndef test_agent_reward_loop_standalone():\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    rollout_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    reward_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n\n    # actor_rollout_ref config\n    config.data.return_raw_chat = True\n    config.data.max_prompt_length = 1024\n    config.data.max_response_length = 4096\n    config.actor_rollout_ref.model.path = rollout_model_path\n    config.actor_rollout_ref.actor.use_dynamic_bsz = True\n    config.actor_rollout_ref.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n    config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.9\n    config.actor_rollout_ref.rollout.enforce_eager = True\n    config.actor_rollout_ref.rollout.prompt_length = 1024\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = True\n    config.actor_rollout_ref.rollout.nnodes = 1\n    config.trainer.n_gpus_per_node = 4\n    config.trainer.nnodes = 1\n\n    config.reward.reward_manager.name = \"dapo\"\n    config.reward.reward_model.enable = True\n    config.reward.reward_model.enable_resource_pool = True\n    config.reward.reward_model.n_gpus_per_node = 4\n    config.reward.reward_model.nnodes = 1\n    config.reward.reward_model.model_path = reward_model_path\n    config.reward.reward_model.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.reward.reward_model.rollout.gpu_memory_utilization = 0.9\n    config.reward.reward_model.rollout.tensor_model_parallel_size = 2\n    config.reward.reward_model.rollout.skip_tokenizer_init = False\n    config.reward.reward_model.rollout.prompt_length = 5120\n    config.reward.reward_model.rollout.response_length = 4096\n    config.reward.custom_reward_function.path = \"tests/experimental/reward_loop/reward_fn.py\"\n    config.reward.custom_reward_function.name = \"compute_score_gsm8k\"\n\n    # 1. init reward model manager\n    reward_loop_manager = RewardLoopManager(config)\n    agent_loop_manager = AgentLoopManager.create(\n        config=config,\n        reward_loop_worker_handles=reward_loop_manager.reward_loop_workers,\n    )\n\n    # 2. init test data\n    local_folder = os.path.expanduser(\"~/data/gsm8k/\")\n    data_files = [os.path.join(local_folder, \"train.parquet\")]\n    tokenizer = hf_tokenizer(rollout_model_path)\n\n    dataset = RLHFDataset(\n        data_files=data_files,\n        tokenizer=tokenizer,\n        config=config.data,\n        processor=None,\n    )\n\n    batch_size = 64\n    sampler = create_rl_sampler(config.data, dataset)\n    dataloader = StatefulDataLoader(\n        dataset=dataset,\n        batch_size=batch_size,\n        num_workers=config.data.dataloader_num_workers,\n        drop_last=True,\n        collate_fn=collate_fn,\n        sampler=sampler,\n    )\n\n    # 3. generate responses\n    batch_dict = next(iter(dataloader))\n    batch = DataProto.from_single_dict(batch_dict)\n\n    # standalone reward model should wake up for agent_reward_loop\n    gen_batch = agent_loop_manager.generate_sequences(prompts=batch)\n\n    rm_scores = gen_batch.batch[\"rm_scores\"]\n    sample_scores = rm_scores.sum(dim=1)\n    print(sample_scores)\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_async_token_bucket_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport time\n\nimport pytest\n\nfrom verl.experimental.reward_loop.reward_manager.limited import AsyncTokenBucket\n\n\nclass TestAsyncTokenBucket:\n    \"\"\"Unit tests for AsyncTokenBucket rate limiter.\"\"\"\n\n    @pytest.mark.asyncio\n    async def test_basic_acquire(self):\n        \"\"\"Test basic token acquisition.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        # Should be able to acquire tokens immediately when bucket is full\n        start = time.time()\n        await bucket.acquire(5.0)\n        elapsed = time.time() - start\n\n        assert elapsed < 0.1, \"Initial acquire should be immediate\"\n        assert bucket.tokens == pytest.approx(5.0, abs=0.1)\n\n    @pytest.mark.asyncio\n    async def test_refill_mechanism(self):\n        \"\"\"Test that tokens refill over time.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        # Consume all tokens\n        await bucket.acquire(10.0)\n        assert bucket.tokens == pytest.approx(0.0, abs=0.1)\n\n        # Wait for refill (should get ~5 tokens in 0.5 seconds at 10 tokens/sec)\n        await asyncio.sleep(0.5)\n\n        # Try to acquire 4 tokens (should succeed without waiting)\n        start = time.time()\n        await bucket.acquire(4.0)\n        elapsed = time.time() - start\n\n        assert elapsed < 0.1, \"Acquire should be quick after refill\"\n\n    @pytest.mark.asyncio\n    async def test_waiting_for_tokens(self):\n        \"\"\"Test that acquire waits when insufficient tokens available.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        # Consume all tokens\n        await bucket.acquire(10.0)\n\n        # Try to acquire more tokens (should wait ~0.5 seconds for 5 tokens)\n        start = time.time()\n        await bucket.acquire(5.0)\n        elapsed = time.time() - start\n\n        # Should wait approximately 0.5 seconds (5 tokens / 10 tokens per second)\n        assert 0.4 < elapsed < 0.7, f\"Expected ~0.5s wait, got {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_max_tokens_cap(self):\n        \"\"\"Test that tokens don't exceed max_tokens capacity.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=5.0)\n\n        # Wait for potential overflow\n        await asyncio.sleep(1.0)\n\n        # Tokens should be capped at max_tokens\n        await bucket.acquire(1.0)\n\n        # After 1 second at 10 tokens/sec, should have max_tokens (5.0)\n        # After acquiring 1, should have 4.0 remaining\n        assert bucket.tokens <= 5.0, \"Tokens should not exceed max_tokens\"\n\n    @pytest.mark.asyncio\n    async def test_fractional_tokens(self):\n        \"\"\"Test acquiring fractional tokens.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0)\n\n        # Acquire fractional amounts\n        await bucket.acquire(0.5)\n        await bucket.acquire(1.5)\n        await bucket.acquire(2.3)\n\n        assert bucket.tokens == pytest.approx(100.0 - 0.5 - 1.5 - 2.3, abs=0.1)\n\n    @pytest.mark.asyncio\n    async def test_concurrent_acquires(self):\n        \"\"\"Test multiple concurrent acquire operations.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        async def acquire_task(num_tokens: float, task_id: int):\n            await bucket.acquire(num_tokens)\n            return task_id\n\n        # Launch 5 concurrent tasks, each acquiring 3 tokens (15 total)\n        # Bucket only has 10, so some will need to wait\n        start = time.time()\n        tasks = [acquire_task(3.0, i) for i in range(5)]\n        results = await asyncio.gather(*tasks)\n        elapsed = time.time() - start\n\n        # Should take at least 0.5 seconds to refill 5 tokens\n        # (15 needed - 10 available) / 10 tokens per second = 0.5 seconds\n        assert elapsed >= 0.4, f\"Expected >=0.4s for concurrent acquires, got {elapsed:.3f}s\"\n        assert len(results) == 5, \"All tasks should complete\"\n\n    @pytest.mark.asyncio\n    async def test_high_rate_limit(self):\n        \"\"\"Test with high rate limit (simulating high-throughput scenarios).\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=1000.0, max_tokens=1000.0)\n\n        # Rapidly acquire tokens\n        start = time.time()\n        for _ in range(100):\n            await bucket.acquire(10.0)  # 1000 tokens total\n        elapsed = time.time() - start\n\n        # Should complete in approximately 1 second\n        assert elapsed < 1.5, f\"High rate limit test took too long: {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_zero_initial_state(self):\n        \"\"\"Test that bucket starts with full tokens.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        assert bucket.tokens == 10.0, \"Bucket should start full\"\n        assert bucket.last_update is None, \"last_update should be None initially\"\n\n        # After first acquire, last_update should be set\n        await bucket.acquire(1.0)\n        assert bucket.last_update is not None, \"last_update should be set after acquire\"\n\n    @pytest.mark.asyncio\n    async def test_rate_limit_accuracy(self):\n        \"\"\"Test rate limit accuracy over time.\"\"\"\n        rate = 50.0  # 50 tokens per second\n        bucket = AsyncTokenBucket(rate_limit=rate, max_tokens=rate)\n\n        # Consume all tokens and measure refill time for 25 tokens\n        await bucket.acquire(50.0)\n\n        start = time.time()\n        await bucket.acquire(25.0)\n        elapsed = time.time() - start\n\n        expected_time = 25.0 / rate  # 0.5 seconds\n        # Allow 20% margin for timing inaccuracy\n        assert abs(elapsed - expected_time) < expected_time * 0.2, f\"Expected ~{expected_time:.3f}s, got {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_sequential_acquires(self):\n        \"\"\"Test sequential acquire operations.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=20.0, max_tokens=20.0)\n\n        # Sequential acquires without waiting\n        await bucket.acquire(5.0)\n        await bucket.acquire(5.0)\n        await bucket.acquire(5.0)\n        await bucket.acquire(5.0)\n\n        # Bucket should be empty\n        assert bucket.tokens == pytest.approx(0.0, abs=0.1)\n\n        # Next acquire should wait\n        start = time.time()\n        await bucket.acquire(10.0)\n        elapsed = time.time() - start\n\n        assert elapsed >= 0.4, \"Should wait for token refill\"\n\n    @pytest.mark.asyncio\n    async def test_default_max_tokens(self):\n        \"\"\"Test that max_tokens defaults to rate_limit.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=15.0)\n\n        assert bucket.max_tokens == 15.0, \"max_tokens should default to rate_limit\"\n        assert bucket.tokens == 15.0, \"Initial tokens should equal max_tokens\"\n\n    @pytest.mark.asyncio\n    async def test_single_token_acquire(self):\n        \"\"\"Test default acquire of 1 token.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        await bucket.acquire()  # Default num_tokens=1.0\n\n        assert bucket.tokens == pytest.approx(9.0, abs=0.1)\n\n    @pytest.mark.asyncio\n    async def test_large_token_acquire(self):\n        \"\"\"Test acquiring more tokens than bucket capacity.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        # Try to acquire 50 tokens (5x capacity)\n        start = time.time()\n        await bucket.acquire(50.0)\n        elapsed = time.time() - start\n\n        # Should wait for: (50 - 10) / 10 = 4 seconds\n        assert 3.5 < elapsed < 5.0, f\"Expected ~4s wait for large acquire, got {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_thread_safety_with_lock(self):\n        \"\"\"Test that lock prevents race conditions.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0)\n        results = []\n\n        async def acquire_and_record():\n            await bucket.acquire(10.0)\n            results.append(1)\n\n        # Launch many concurrent tasks\n        tasks = [acquire_and_record() for _ in range(10)]\n        await asyncio.gather(*tasks)\n\n        # All tasks should complete\n        assert len(results) == 10, \"All tasks should complete successfully\"\n\n        # Bucket should have consumed exactly 100 tokens\n        assert bucket.tokens == pytest.approx(0.0, abs=0.5)\n\n    @pytest.mark.asyncio\n    async def test_multiple_wait_cycles(self):\n        \"\"\"Test multiple wait cycles in the acquire loop.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=10.0, max_tokens=10.0)\n\n        # Consume all tokens\n        await bucket.acquire(10.0)\n\n        # Acquire tokens that require multiple refill cycles\n        start = time.time()\n        await bucket.acquire(15.0)\n        elapsed = time.time() - start\n\n        # Should wait for 15 tokens / 10 tokens per second = 1.5 seconds\n        assert 1.3 < elapsed < 1.8, f\"Expected ~1.5s for multiple refill cycles, got {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_rapid_small_acquires(self):\n        \"\"\"Test many rapid small acquisitions.\"\"\"\n        bucket = AsyncTokenBucket(rate_limit=100.0, max_tokens=100.0)\n\n        start = time.time()\n        for _ in range(50):\n            await bucket.acquire(2.0)  # 100 tokens total\n        elapsed = time.time() - start\n\n        # Should complete quickly since we're within capacity\n        assert elapsed < 0.5, f\"Rapid small acquires took too long: {elapsed:.3f}s\"\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\"])\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_math_verify.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport ray\nfrom hydra import compose, initialize_config_dir\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom transformers import AutoTokenizer\n\nfrom tests.experimental.agent_loop.agent_utils import init_agent_loop_manager\nfrom verl.protocol import DataProto\nfrom verl.trainer.main_ppo import create_rl_sampler\nfrom verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn\n\n\ndef test_agent_reward_loop_standalone():\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    rollout_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n\n    # actor_rollout_ref config\n    config.data.return_raw_chat = True\n    config.data.max_prompt_length = 1024\n    config.data.max_response_length = 4096\n    config.actor_rollout_ref.model.path = rollout_model_path\n    config.actor_rollout_ref.actor.use_dynamic_bsz = True\n    config.actor_rollout_ref.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n    config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.9\n    config.actor_rollout_ref.rollout.enforce_eager = True\n    config.actor_rollout_ref.rollout.prompt_length = 2048\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.skip_tokenizer_init = True\n    config.trainer.n_gpus_per_node = 8\n    config.trainer.nnodes = 1\n\n    config.reward.reward_manager.name = \"remote\"\n    config.reward.num_workers = 2\n    config.reward.custom_reward_function.path = \"tests/experimental/reward_loop/reward_fn.py\"\n    config.reward.custom_reward_function.name = \"compute_score_math_verify\"\n\n    # 1. init reward model manager\n    agent_loop_manager = init_agent_loop_manager(config)\n\n    # 2. init test data\n    local_folder = os.path.expanduser(\"~/data/math/\")\n    data_files = [os.path.join(local_folder, \"train.parquet\")]\n    tokenizer = AutoTokenizer.from_pretrained(rollout_model_path)\n\n    dataset = RLHFDataset(\n        data_files=data_files,\n        tokenizer=tokenizer,\n        config=config.data,\n        processor=None,\n    )\n\n    batch_size = 64\n    sampler = create_rl_sampler(config.data, dataset)\n    dataloader = StatefulDataLoader(\n        dataset=dataset,\n        batch_size=batch_size,\n        num_workers=config.data.dataloader_num_workers,\n        drop_last=True,\n        collate_fn=collate_fn,\n        sampler=sampler,\n    )\n\n    # 3. generate responses\n    batch_dict = next(iter(dataloader))\n    batch = DataProto.from_single_dict(batch_dict)\n    gen_batch = agent_loop_manager.generate_sequences(prompts=batch)\n\n    rm_scores = gen_batch.batch[\"rm_scores\"]\n    accuracy = rm_scores.sum(dim=-1).mean()\n    print(accuracy)\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_rate_limited_reward_manager_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport os.path\nimport time\n\nimport pytest\nimport torch\nfrom omegaconf import DictConfig\nfrom transformers import AutoTokenizer\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager.limited import RateLimitedRewardManager\n\n\n# Mock API reward functions for testing\nclass MockAPICounter:\n    \"\"\"Shared counter to track API calls across tests.\"\"\"\n\n    def __init__(self):\n        self.call_count = 0\n        self.call_times = []\n        self.lock = asyncio.Lock()\n\n    async def record_call(self):\n        async with self.lock:\n            self.call_count += 1\n            self.call_times.append(time.time())\n\n    def reset(self):\n        self.call_count = 0\n        self.call_times.clear()\n\n    def get_rate_per_second(self, window_start: float = None):\n        \"\"\"Calculate API call rate over a time window.\"\"\"\n        if window_start is None:\n            if not self.call_times:\n                return 0.0\n            window_start = self.call_times[0]\n\n        if not self.call_times:\n            return 0.0\n\n        window_end = self.call_times[-1]\n        duration = window_end - window_start\n\n        if duration <= 0:\n            return 0.0\n\n        calls_in_window = sum(1 for t in self.call_times if t >= window_start)\n        return calls_in_window / duration\n\n\n# Global counter instance\napi_counter = MockAPICounter()\n\n\ndef mock_sync_reward_function(\n    data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n) -> float:\n    \"\"\"Synchronous mock reward function that simulates API call.\"\"\"\n    # Simulate API processing time\n    time.sleep(0.01)\n\n    # Simple scoring logic\n    score = 1.0 if solution_str.strip() == ground_truth.strip() else 0.0\n    return score\n\n\nasync def mock_async_reward_function(\n    data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n) -> float:\n    \"\"\"Asynchronous mock reward function that simulates API call.\"\"\"\n    # Record API call for rate tracking\n    await api_counter.record_call()\n\n    # Simulate async API call (e.g., HTTP request)\n    await asyncio.sleep(0.01)\n\n    # Simple scoring logic\n    score = 1.0 if solution_str.strip() == ground_truth.strip() else 0.0\n    return score\n\n\nasync def mock_slow_api_function(\n    data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n) -> float:\n    \"\"\"Slow mock API function for timeout testing.\"\"\"\n    await asyncio.sleep(2.0)  # Simulate slow API\n    return 0.5\n\n\nasync def mock_failing_api_function(\n    data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n) -> float:\n    \"\"\"Mock API function that raises an exception.\"\"\"\n    await api_counter.record_call()\n    raise ValueError(\"Simulated API error\")\n\n\nasync def mock_dict_result_function(\n    data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n) -> dict:\n    \"\"\"Mock API function that returns dict result.\"\"\"\n    await api_counter.record_call()\n    await asyncio.sleep(0.01)\n\n    correct = solution_str.strip() == ground_truth.strip()\n    return {\"score\": 1.0 if correct else 0.0, \"correct\": correct, \"reasoning\": \"Mock reasoning\"}\n\n\ndef create_test_data_proto(tokenizer, response_text: str, ground_truth: str, data_source: str = \"test\"):\n    \"\"\"Helper to create DataProto for testing.\"\"\"\n    response_ids = tokenizer.encode(response_text, add_special_tokens=False)\n    response_tensor = torch.tensor([response_ids], dtype=torch.long)\n    attention_mask = torch.ones_like(response_tensor)\n\n    data = DataProto.from_dict(\n        {\n            \"responses\": response_tensor,\n            \"attention_mask\": attention_mask,\n        }\n    )\n\n    # Wrap non-tensor values in lists to match batch dimension\n    data.non_tensor_batch = {\"data_source\": [data_source], \"reward_model\": [{\"ground_truth\": ground_truth}]}\n\n    return data\n\n\nclass TestRateLimitedRewardManager:\n    \"\"\"Integration tests for RateLimitedRewardManager with mock API functions.\"\"\"\n\n    @pytest.fixture(autouse=True)\n    def setup_and_teardown(self):\n        \"\"\"Reset global state before each test.\"\"\"\n        api_counter.reset()\n        # Reset class state\n        RateLimitedRewardManager._class_initialized = False\n        RateLimitedRewardManager._semaphore = None\n        RateLimitedRewardManager._rpm_limiter = None\n        RateLimitedRewardManager._tpm_limiter = None\n        yield\n        # Cleanup\n        api_counter.reset()\n\n    @pytest.fixture\n    def tokenizer(self):\n        \"\"\"Load a simple tokenizer for testing.\"\"\"\n        return AutoTokenizer.from_pretrained(os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\"))\n\n    @pytest.mark.asyncio\n    async def test_basic_reward_computation(self, tokenizer):\n        \"\"\"Test basic reward computation without rate limiting.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        # Create test data\n        data = create_test_data_proto(tokenizer, \"correct answer\", \"correct answer\")\n\n        # Compute reward\n        result = await manager.run_single(data)\n\n        assert \"reward_score\" in result\n        assert result[\"reward_score\"] == 1.0\n        assert api_counter.call_count == 1\n\n    @pytest.mark.asyncio\n    async def test_rpm_rate_limiting(self, tokenizer):\n        \"\"\"Test request per minute (RPM) rate limiting.\"\"\"\n        # Set RPM limit to 60 (1 request per second)\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 10,\n                    \"max_rpm\": 60,  # 1 request per second\n                    \"timeout\": 10.0,\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        # Create test data\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Make 3 requests - should be rate limited\n        start_time = time.time()\n\n        results = []\n        for _ in range(3):\n            result = await manager.run_single(data)\n            results.append(result)\n\n        elapsed = time.time() - start_time\n\n        # Should take at least ~2 seconds for 3 requests at 1 req/sec\n        assert elapsed >= 1.8, f\"RPM limiting failed: {elapsed:.3f}s for 3 requests\"\n        assert all(r[\"reward_score\"] == 1.0 for r in results)\n        assert api_counter.call_count == 3\n\n    @pytest.mark.asyncio\n    async def test_tpm_rate_limiting(self, tokenizer):\n        \"\"\"Test tokens per minute (TPM) rate limiting.\"\"\"\n        # Set TPM limit to 6000 (100 tokens per second)\n        # With 2000 tokens per request, that's 0.05 req/sec or 20 seconds per request\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 10,\n                    \"max_tpm\": 6000,  # 100 tokens per second\n                    \"estimated_tokens_per_request\": 2000,  # Each request = 2000 tokens\n                    \"timeout\": 30.0,\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Make 2 requests\n        start_time = time.time()\n\n        result1 = await manager.run_single(data)\n        result2 = await manager.run_single(data)\n\n        elapsed = time.time() - start_time\n\n        # First request: consumes 2000 tokens (immediate)\n        # Second request: needs 2000 tokens, waits for refill\n        # Wait time: 2000 tokens / 100 tokens per second = 20 seconds\n        assert elapsed >= 18.0, f\"TPM limiting failed: {elapsed:.3f}s for 2 requests\"\n        assert result1[\"reward_score\"] == 1.0\n        assert result2[\"reward_score\"] == 1.0\n\n    @pytest.mark.asyncio\n    async def test_concurrency_limiting(self, tokenizer):\n        \"\"\"Test concurrent request limiting.\"\"\"\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 2,  # Only 2 concurrent requests\n                    \"timeout\": 10.0,\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Launch 5 concurrent requests\n        start_time = time.time()\n\n        tasks = [manager.run_single(data) for _ in range(5)]\n        results = await asyncio.gather(*tasks)\n\n        elapsed = time.time() - start_time\n\n        # All should succeed\n        assert len(results) == 5\n        assert all(r[\"reward_score\"] == 1.0 for r in results)\n\n        # With concurrency=2 and 0.01s per request, should take at least 0.03s\n        # (3 batches: 2+2+1)\n        assert elapsed >= 0.02, f\"Concurrency limiting may not be working: {elapsed:.3f}s\"\n\n    @pytest.mark.asyncio\n    async def test_timeout_handling(self, tokenizer):\n        \"\"\"Test timeout handling for slow API.\"\"\"\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 10,\n                    \"timeout\": 0.5,  # 500ms timeout\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_slow_api_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Should timeout and return 0.0\n        result = await manager.run_single(data)\n\n        assert result[\"reward_score\"] == 0.0\n        assert result[\"reward_extra_info\"].get(\"timeout\") is True\n        assert result[\"reward_extra_info\"].get(\"acc\") == 0.0\n\n    @pytest.mark.asyncio\n    async def test_error_handling(self, tokenizer):\n        \"\"\"Test error handling for failing API.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_failing_api_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Should catch exception and return 0.0\n        result = await manager.run_single(data)\n\n        assert result[\"reward_score\"] == 0.0\n        assert \"error\" in result[\"reward_extra_info\"]\n        assert \"Simulated API error\" in result[\"reward_extra_info\"][\"error\"]\n        assert result[\"reward_extra_info\"].get(\"acc\") == 0.0\n        assert api_counter.call_count == 1\n\n    @pytest.mark.asyncio\n    async def test_dict_result_format(self, tokenizer):\n        \"\"\"Test handling of dict return format from reward function.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_dict_result_function)\n\n        data = create_test_data_proto(tokenizer, \"correct\", \"correct\")\n\n        result = await manager.run_single(data)\n\n        assert result[\"reward_score\"] == 1.0\n        assert result[\"reward_extra_info\"][\"score\"] == 1.0\n        assert result[\"reward_extra_info\"][\"correct\"] is True\n        assert result[\"reward_extra_info\"][\"reasoning\"] == \"Mock reasoning\"\n\n    @pytest.mark.asyncio\n    async def test_sync_reward_function(self, tokenizer):\n        \"\"\"Test that synchronous reward functions work correctly.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_sync_reward_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        result = await manager.run_single(data)\n\n        assert result[\"reward_score\"] == 1.0\n        assert manager.is_async_reward_score is False\n\n    @pytest.mark.asyncio\n    async def test_combined_rate_limits(self, tokenizer):\n        \"\"\"Test all three rate limiting layers together.\"\"\"\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 2,\n                    \"max_rpm\": 120,  # 2 requests per second\n                    \"max_tpm\": 12000,  # 200 tokens per second\n                    \"estimated_tokens_per_request\": 100,  # 0.5 seconds per request\n                    \"timeout\": 10.0,\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Make 6 requests to exceed burst capacity (RPM bucket starts with 2 tokens)\n        start_time = time.time()\n\n        tasks = [manager.run_single(data) for _ in range(6)]\n        results = await asyncio.gather(*tasks)\n\n        elapsed = time.time() - start_time\n\n        # Bucket starts with 2 RPM tokens and 200 TPM tokens\n        # First 2 requests: use burst capacity (2 RPM tokens, 200 TPM tokens)\n        # Next 4 requests: need 4 RPM tokens (wait 2 seconds) and 400 TPM tokens (wait 2 seconds)\n        # Limiting factor: RPM at 2 seconds\n        assert elapsed >= 1.8, f\"Combined rate limiting: {elapsed:.3f}s\"\n        assert all(r[\"reward_score\"] == 1.0 for r in results)\n        assert api_counter.call_count == 6\n\n    @pytest.mark.asyncio\n    async def test_correct_vs_incorrect_answers(self, tokenizer):\n        \"\"\"Test scoring of correct vs incorrect answers.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        # Test correct answer\n        data_correct = create_test_data_proto(tokenizer, \"right answer\", \"right answer\")\n        result_correct = await manager.run_single(data_correct)\n\n        # Test incorrect answer\n        data_incorrect = create_test_data_proto(tokenizer, \"wrong answer\", \"right answer\")\n        result_incorrect = await manager.run_single(data_incorrect)\n\n        assert result_correct[\"reward_score\"] == 1.0\n        assert result_incorrect[\"reward_score\"] == 0.0\n\n    @pytest.mark.asyncio\n    async def test_high_throughput(self, tokenizer):\n        \"\"\"Test high throughput with many concurrent requests.\"\"\"\n        config = DictConfig(\n            {\n                \"reward\": {\n                    \"max_concurrent\": 20,\n                    \"max_rpm\": 6000,  # 100 requests per second\n                    \"timeout\": 10.0,\n                }\n            }\n        )\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(config=config, tokenizer=tokenizer, compute_score=mock_async_reward_function)\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n\n        # Launch 200 concurrent requests (more than burst capacity of 100)\n        start_time = time.time()\n\n        tasks = [manager.run_single(data) for _ in range(200)]\n        results = await asyncio.gather(*tasks)\n\n        elapsed = time.time() - start_time\n\n        assert len(results) == 200\n        assert all(r[\"reward_score\"] == 1.0 for r in results)\n\n        # Bucket starts with 100 tokens (burst capacity)\n        # First 100 requests: use burst capacity instantly\n        # Next 100 requests: need to wait for refill at 100 tokens/sec = 1 second minimum\n        # Total time should be at least 1 second\n        assert elapsed >= 0.9, f\"Should take at least 0.9s for rate limiting, took {elapsed:.3f}s\"\n\n        # Calculate actual rate over the time window\n        actual_rate = api_counter.call_count / elapsed\n\n        # Average rate should not significantly exceed 100 req/sec\n        # Allow some burst overhead due to initial capacity\n        assert actual_rate <= 200, f\"Rate limiting failed: {actual_rate:.1f} req/sec (max 200)\"\n\n    @pytest.mark.asyncio\n    async def test_class_initialization_once(self, tokenizer):\n        \"\"\"Test that class initialization only happens once.\"\"\"\n        config = DictConfig({\"reward\": {\"max_concurrent\": 5, \"timeout\": 10.0}})\n\n        # Initialize multiple times\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        first_semaphore = RateLimitedRewardManager._semaphore\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        second_semaphore = RateLimitedRewardManager._semaphore\n\n        # Should be the same object\n        assert first_semaphore is second_semaphore\n\n    @pytest.mark.asyncio\n    async def test_extra_info_handling(self, tokenizer):\n        \"\"\"Test that extra_info is properly passed to reward function.\"\"\"\n        received_extra_info = {}\n\n        async def mock_reward_with_extra_info(\n            data_source: str, solution_str: str, ground_truth: str, extra_info: dict, **kwargs\n        ):\n            received_extra_info.update(extra_info)\n            return 1.0\n\n        config = DictConfig({\"reward\": {\"max_concurrent\": 10, \"timeout\": 10.0}})\n\n        RateLimitedRewardManager.init_class(config, tokenizer)\n        manager = RateLimitedRewardManager(\n            config=config, tokenizer=tokenizer, compute_score=mock_reward_with_extra_info\n        )\n\n        data = create_test_data_proto(tokenizer, \"answer\", \"answer\")\n        data.non_tensor_batch[\"extra_info\"] = [{\"custom_field\": \"test_value\"}]\n\n        await manager.run_single(data)\n\n        assert \"custom_field\" in received_extra_info\n        assert received_extra_info[\"custom_field\"] == \"test_value\"\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\", \"-s\"])\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_reward_model_disrm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport ray\nimport torch\nfrom hydra import compose, initialize_config_dir\n\nfrom verl.experimental.reward_loop import RewardLoopManager\nfrom verl.protocol import DataProto\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.model import compute_position_id_with_mask\nfrom verl.utils.tokenizer import normalize_token_ids\n\n\ndef create_data_samples(tokenizer) -> DataProto:\n    convs = [\n        [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the range of the numeric output of a sigmoid node in a neural network?\",\n            },\n            {\"role\": \"assistant\", \"content\": \"Between -1 and 1.\"},\n        ],\n        [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the range of the numeric output of a sigmoid node in a neural network?\",\n            },\n            {\"role\": \"assistant\", \"content\": \"Between 0 and 1.\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What is the capital of Australia?\"},\n            {\n                \"role\": \"assistant\",\n                \"content\": \"Canberra is the capital city of Australia.\",\n            },\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What is the capital of Australia?\"},\n            {\n                \"role\": \"assistant\",\n                \"content\": \"Sydney is the capital of Australia.\",\n            },\n        ],\n    ]\n    raw_prompt = [conv[:1] for conv in convs]\n    data_source = [\"gsm8k\"] * len(convs)\n    reward_info = [{\"ground_truth\": \"Not Used\"}] * len(convs)\n    extra_info = [{\"question\": conv[0][\"content\"]} for conv in convs]\n\n    prompt_length, response_length = 1024, 4096\n    pad_token_id = tokenizer.pad_token_id\n    prompts, responses, input_ids, attention_masks = [], [], [], []\n    for conv in convs:\n        prompt_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv[:1], tokenize=True))\n        response_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv, tokenize=True))[len(prompt_tokens) :]\n\n        padded_prompt = [pad_token_id] * (prompt_length - len(prompt_tokens)) + prompt_tokens\n        padded_response = response_tokens + [pad_token_id] * (response_length - len(response_tokens))\n        attention_mask = (\n            [0] * (prompt_length - len(prompt_tokens))\n            + [1] * len(prompt_tokens)\n            + [1] * len(response_tokens)\n            + [0] * (response_length - len(response_tokens))\n        )\n        prompts.append(torch.tensor(padded_prompt))\n        responses.append(torch.tensor(padded_response))\n        input_ids.append(torch.tensor(padded_prompt + padded_response))\n        attention_masks.append(torch.tensor(attention_mask))\n\n    prompts = torch.stack(prompts)\n    responses = torch.stack(responses)\n    input_ids = torch.stack(input_ids)\n    attention_masks = torch.stack(attention_masks)\n    position_ids = compute_position_id_with_mask(attention_masks)\n\n    data = DataProto.from_dict(\n        tensors={\n            \"prompts\": prompts,\n            \"responses\": responses,\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_masks,\n            \"position_ids\": position_ids,\n        },\n        non_tensors={\n            \"data_source\": data_source,\n            \"reward_model\": reward_info,\n            \"raw_prompt\": raw_prompt,\n            \"extra_info\": extra_info,\n        },\n    )\n    return data, convs\n\n\ndef test_reward_model_manager():\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    rollout_model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n    reward_model_name = os.path.expanduser(\"~/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B\")\n\n    config.actor_rollout_ref.model.path = rollout_model_name\n    config.reward.num_workers = 1\n    config.reward.reward_manager.name = \"dapo\"\n    config.reward.reward_model.enable = True\n    config.reward.reward_model.enable_resource_pool = True\n    config.reward.reward_model.n_gpus_per_node = 8\n    config.reward.reward_model.nnodes = 1\n    config.reward.reward_model.model_path = reward_model_name\n    config.reward.reward_model.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.reward.reward_model.rollout.gpu_memory_utilization = 0.9\n    config.reward.reward_model.rollout.tensor_model_parallel_size = 2\n    config.reward.reward_model.rollout.skip_tokenizer_init = False\n    config.reward.reward_model.rollout.prompt_length = 2048\n    config.reward.reward_model.rollout.response_length = 4096\n\n    # 1. init reward model manager\n    reward_loop_manager = RewardLoopManager(config)\n\n    # 2. init test data\n    rollout_tokenizer = hf_tokenizer(rollout_model_name)\n    data, convs = create_data_samples(rollout_tokenizer)\n\n    # 3. generate responses\n    outputs = reward_loop_manager.compute_rm_score(data)\n\n    for idx, (conv, output) in enumerate(zip(convs, outputs, strict=True)):\n        print(f\"Problem {idx}:\\n{conv[0]['content']}\\n\")\n        print(f\"AI Solution {idx}:\\n{conv[1]['content']}\\n\")\n        print(f\"DisRM Score {idx}:\\n{output.batch['rm_scores'].sum(dim=-1).item()}\\n\")\n        print(\"=\" * 50 + \"\\n\")\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/reward_loop/test_reward_model_genrm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport ray\nimport torch\nfrom hydra import compose, initialize_config_dir\n\nfrom verl.experimental.reward_loop import RewardLoopManager\nfrom verl.protocol import DataProto\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.model import compute_position_id_with_mask\nfrom verl.utils.tokenizer import normalize_token_ids\n\n\ndef create_data_samples(tokenizer) -> DataProto:\n    convs = [\n        [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the range of the numeric output of a sigmoid node in a neural network?\",\n            },\n            {\"role\": \"assistant\", \"content\": \"Between -1 and 1.\"},\n        ],\n        [\n            {\n                \"role\": \"user\",\n                \"content\": \"What is the range of the numeric output of a sigmoid node in a neural network?\",\n            },\n            {\"role\": \"assistant\", \"content\": \"Between 0 and 1.\"},\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What is the capital of Australia?\"},\n            {\n                \"role\": \"assistant\",\n                \"content\": \"Canberra is the capital city of Australia.\",\n            },\n        ],\n        [\n            {\"role\": \"user\", \"content\": \"What is the capital of Australia?\"},\n            {\n                \"role\": \"assistant\",\n                \"content\": \"Sydney is the capital of Australia.\",\n            },\n        ],\n    ]\n    raw_prompt = [conv[:1] for conv in convs]\n    data_source = [\"gsm8k\"] * len(convs)\n    reward_info = [{\"ground_truth\": \"Not Used\"}] * len(convs)\n    extra_info = [{\"question\": conv[0][\"content\"]} for conv in convs]\n\n    prompt_length, response_length = 1024, 4096\n    pad_token_id = tokenizer.pad_token_id\n    prompts, responses, input_ids, attention_masks = [], [], [], []\n    for conv in convs:\n        prompt_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv[:1], tokenize=True))\n        response_tokens = normalize_token_ids(tokenizer.apply_chat_template(conv, tokenize=True))[len(prompt_tokens) :]\n\n        padded_prompt = [pad_token_id] * (prompt_length - len(prompt_tokens)) + prompt_tokens\n        padded_response = response_tokens + [pad_token_id] * (response_length - len(response_tokens))\n        attention_mask = (\n            [0] * (prompt_length - len(prompt_tokens))\n            + [1] * len(prompt_tokens)\n            + [1] * len(response_tokens)\n            + [0] * (response_length - len(response_tokens))\n        )\n        prompts.append(torch.tensor(padded_prompt))\n        responses.append(torch.tensor(padded_response))\n        input_ids.append(torch.tensor(padded_prompt + padded_response))\n        attention_masks.append(torch.tensor(attention_mask))\n\n    prompts = torch.stack(prompts)\n    responses = torch.stack(responses)\n    input_ids = torch.stack(input_ids)\n    attention_masks = torch.stack(attention_masks)\n    position_ids = compute_position_id_with_mask(attention_masks)\n\n    data = DataProto.from_dict(\n        tensors={\n            \"prompts\": prompts,\n            \"responses\": responses,\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_masks,\n            \"position_ids\": position_ids,\n        },\n        non_tensors={\n            \"data_source\": data_source,\n            \"reward_model\": reward_info,\n            \"raw_prompt\": raw_prompt,\n            \"extra_info\": extra_info,\n        },\n    )\n    return data, convs\n\n\ndef test_reward_model_manager():\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        }\n    )\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(config_name=\"ppo_trainer\")\n\n    rollout_model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    reward_model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n\n    config.actor_rollout_ref.model.path = rollout_model_name\n    config.reward.custom_reward_function.path = \"tests/experimental/reward_loop/reward_fn.py\"\n    config.reward.custom_reward_function.name = \"compute_score_gsm8k\"\n    config.reward.num_workers = 1\n    config.reward.reward_manager.name = \"dapo\"\n    config.reward.reward_model.enable = True\n    config.reward.reward_model.enable_resource_pool = True\n    config.reward.reward_model.n_gpus_per_node = 8\n    config.reward.reward_model.nnodes = 1\n    config.reward.reward_model.model_path = reward_model_name\n    config.reward.reward_model.rollout.name = os.getenv(\"ROLLOUT_NAME\", \"vllm\")\n    config.reward.reward_model.rollout.gpu_memory_utilization = 0.9\n    config.reward.reward_model.rollout.tensor_model_parallel_size = 2\n    config.reward.reward_model.rollout.skip_tokenizer_init = False\n    config.reward.reward_model.rollout.prompt_length = 2048\n    config.reward.reward_model.rollout.response_length = 4096\n\n    # 1. init reward model manager\n    reward_loop_manager = RewardLoopManager(config)\n\n    # 2. init test data\n    rollout_tokenizer = hf_tokenizer(rollout_model_name)\n    data, convs = create_data_samples(rollout_tokenizer)\n\n    # 3. generate responses\n    outputs = reward_loop_manager.compute_rm_score(data)\n\n    for idx, (conv, output) in enumerate(zip(convs, outputs, strict=True)):\n        print(f\"Problem {idx}:\\n{conv[0]['content']}\\n\")\n        print(f\"AI Solution {idx}:\\n{conv[1]['content']}\\n\")\n        print(f\"GRM Response {idx}:\\n{output.non_tensor_batch['genrm_response']}\\n\")\n        print(\"=\" * 50 + \"\\n\")\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/experimental/vla/test_sim_envs.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport unittest\n\nimport numpy as np\nimport pytest\nfrom omegaconf import OmegaConf\n\n\n# @pytest.mark.parametrize(\"simulator_type\", [\"libero\", \"isaac\"])\n@pytest.mark.parametrize(\"simulator_type\", [\"isaac\"])\ndef test_sim_env_creation_and_step(simulator_type):\n    num_envs = 8\n    actions = np.array(\n        [\n            [5.59112417e-01, 8.06460073e-02, 1.36817226e-02, -4.64279854e-04, -1.72158767e-02, -6.57548380e-04, -1],\n            [2.12711899e-03, -3.13366604e-01, 3.41386353e-04, -4.64279854e-04, -8.76528812e-03, -6.57548380e-04, -1],\n            [7.38182960e-02, -4.64548351e-02, -6.63602950e-02, -4.64279854e-04, -2.32520114e-02, -6.57548380e-04, -1],\n            [7.38182960e-02, -1.60845593e-01, 3.41386353e-04, -4.64279854e-04, 1.05503430e-02, -6.57548380e-04, -1],\n            [7.38182960e-02, -3.95982152e-01, -7.97006313e-02, -5.10713711e-03, 3.22804279e-02, -6.57548380e-04, -1],\n            [2.41859427e-02, -3.64206941e-01, -6.63602950e-02, -4.64279854e-04, 1.05503430e-02, -6.57548380e-04, -1],\n            [4.62447664e-02, -5.16727952e-01, -7.97006313e-02, -4.64279854e-04, 1.05503430e-02, 8.73740975e-03, -1],\n            [4.62447664e-02, -5.73923331e-01, 3.41386353e-04, -4.64279854e-04, 6.92866212e-03, -6.57548380e-04, -1],\n        ]\n    )\n    cfg = OmegaConf.create(\n        {\n            \"max_episode_steps\": 512,\n            \"only_eval\": False,\n            \"reward_coef\": 1.0,\n            \"init_params\": {\n                \"camera_names\": [\"agentview\"],\n            },\n            \"video_cfg\": {\n                \"save_video\": True,\n                \"video_base_dir\": \"/tmp/test_sim_env_creation_and_step\",\n            },\n            \"task_suite_name\": \"libero_10\",\n            \"num_envs\": num_envs,\n            \"num_group\": 1,\n            \"group_size\": num_envs,\n            \"seed\": 0,\n        },\n    )\n\n    sim_env = None\n    if simulator_type == \"isaac\":\n        from verl.experimental.vla.envs.isaac_env.isaac_env import IsaacEnv\n\n        sim_env = IsaacEnv(cfg, rank=0, world_size=1)\n    elif simulator_type == \"libero\":\n        from verl.experimental.vla.envs.libero_env.libero_env import LiberoEnv\n\n        sim_env = LiberoEnv(cfg, rank=0, world_size=1)\n    else:\n        raise ValueError(f\"simulator_type {simulator_type} is not supported\")\n\n    video_count = 0\n    for i in [0]:\n        # The first call to step with actions=None will reset the environment\n        step = 0\n        sim_env.reset_envs_to_state_ids([0] * num_envs, [i] * num_envs)\n        for action in actions:\n            obs_venv, reward_venv, terminated_venv, truncated_venv, info_venv = sim_env.step(\n                np.array([action] * num_envs)\n            )\n\n            assert isinstance(obs_venv, dict)\n            assert reward_venv.shape == (num_envs,)\n            assert terminated_venv.shape == (num_envs,)\n            assert truncated_venv.shape == (num_envs,)\n            assert isinstance(info_venv, dict)\n\n            if terminated_venv.any() or truncated_venv.any():\n                break\n            step += 1\n\n        sim_env.flush_video(video_sub_dir=f\"task_{i}\")\n        assert os.path.exists(os.path.join(cfg.video_cfg.video_base_dir, f\"rank_0/task_{i}/{video_count}.mp4\"))\n        os.remove(os.path.join(cfg.video_cfg.video_base_dir, f\"rank_0/task_{i}/{video_count}.mp4\"))\n        video_count += 1\n\n    print(\"test passed\")\n    sim_env.close()\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/interactions/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/interactions/test_gsm8k_interaction.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom unittest.mock import patch\n\nimport pytest\n\nfrom verl.interactions.gsm8k_interaction import Gsm8kInteraction\n\n\nclass TestGsm8kInteraction:\n    \"\"\"Test cases for Gsm8kInteraction class.\"\"\"\n\n    def setup_method(self):\n        \"\"\"Set up test environment before each test method.\"\"\"\n        self.config = {\"name\": \"gsm8k\"}\n        self.interaction = Gsm8kInteraction(self.config)\n\n    def test_init(self):\n        \"\"\"Test Gsm8kInteraction initialization.\"\"\"\n        assert self.interaction._instance_dict == {}\n        assert self.interaction.config == self.config\n        assert self.interaction.name == \"gsm8k\"\n\n    @pytest.mark.asyncio\n    async def test_start_interaction_with_instance_id(self):\n        \"\"\"Test start_interaction with provided instance_id.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        result_id = await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        assert result_id == instance_id\n        assert instance_id in self.interaction._instance_dict\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"\"\n        assert self.interaction._instance_dict[instance_id][\"ground_truth\"] == ground_truth\n        assert self.interaction._instance_dict[instance_id][\"reward\"] == 0.0\n\n    @pytest.mark.asyncio\n    async def test_start_interaction_without_instance_id(self):\n        \"\"\"Test start_interaction without provided instance_id (auto-generated).\"\"\"\n        ground_truth = \"42\"\n\n        result_id = await self.interaction.start_interaction(ground_truth=ground_truth)\n\n        assert result_id is not None\n        assert len(result_id) == 36  # UUID4 length\n        assert result_id in self.interaction._instance_dict\n        assert self.interaction._instance_dict[result_id][\"ground_truth\"] == ground_truth\n\n    @pytest.mark.asyncio\n    async def test_start_interaction_without_ground_truth(self):\n        \"\"\"Test start_interaction without ground_truth parameter.\"\"\"\n        instance_id = \"test_instance\"\n\n        result_id = await self.interaction.start_interaction(instance_id=instance_id)\n\n        assert result_id == instance_id\n        assert self.interaction._instance_dict[instance_id][\"ground_truth\"] is None\n\n    @pytest.mark.asyncio\n    async def test_generate_response_correct_answer_with_prefix(self):\n        \"\"\"Test generate_response with correct answer already having #### prefix.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [{\"role\": \"assistant\", \"content\": \"#### 42\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is True\n        assert response == \"Your response is correct!\"\n        assert reward == 1.0\n        assert metadata == {}\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"#### 42\"\n\n    @pytest.mark.asyncio\n    async def test_generate_response_correct_answer_without_prefix(self):\n        \"\"\"Test generate_response with correct answer missing #### prefix.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [{\"role\": \"assistant\", \"content\": \"42\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is True\n        assert response == \"Your response is correct!\"\n        assert reward == 1.0\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"42\"\n\n    @pytest.mark.asyncio\n    async def test_generate_response_incorrect_answer(self):\n        \"\"\"Test generate_response with incorrect answer.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [{\"role\": \"assistant\", \"content\": \"24\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is False\n        assert response == \"Your response is incorrect! You need to reflect on your answer and try again.\"\n        assert reward == 0.0\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"24\"\n\n    @pytest.mark.asyncio\n    async def test_generate_response_multiple_messages(self):\n        \"\"\"Test generate_response with multiple messages (should use last assistant message).\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [\n            {\"role\": \"user\", \"content\": \"What is 2+2?\"},\n            {\"role\": \"assistant\", \"content\": \"### 4\"},\n            {\"role\": \"user\", \"content\": \"What is 40+2?\"},\n            {\"role\": \"assistant\", \"content\": \"#### 42\"},\n        ]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is True\n        assert response == \"Your response is correct!\"\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"#### 42\"\n\n    @pytest.mark.asyncio\n    async def test_generate_response_no_assistant_message(self):\n        \"\"\"Test generate_response with no assistant messages.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [{\"role\": \"user\", \"content\": \"Hello!\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is False\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"\"\n\n    @pytest.mark.asyncio\n    async def test_calculate_score_direct_call(self):\n        \"\"\"Test calculate_score method directly.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        # Set a response\n        self.interaction._instance_dict[instance_id][\"response\"] = \"#### 42\"\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0) as mock_compute:\n            score = await self.interaction.calculate_score(instance_id)\n\n            assert score == 1.0\n            mock_compute.assert_called_once_with(\"#### 42\", \"42\", method=\"strict\", format_score=0.0, score=1.0)\n\n    @pytest.mark.asyncio\n    async def test_calculate_score_with_kwargs(self):\n        \"\"\"Test calculate_score method with additional kwargs.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        # Set a response\n        self.interaction._instance_dict[instance_id][\"response\"] = \"#### 24\"\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0) as mock_compute:\n            score = await self.interaction.calculate_score(instance_id, extra_param=\"test\")\n\n            assert score == 0.0\n            mock_compute.assert_called_once_with(\"#### 24\", \"42\", method=\"strict\", format_score=0.0, score=1.0)\n\n    @pytest.mark.asyncio\n    async def test_finalize_interaction(self):\n        \"\"\"Test finalize_interaction method.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        assert instance_id in self.interaction._instance_dict\n\n        await self.interaction.finalize_interaction(instance_id)\n\n        assert instance_id not in self.interaction._instance_dict\n\n    @pytest.mark.asyncio\n    async def test_finalize_interaction_with_kwargs(self):\n        \"\"\"Test finalize_interaction method with additional kwargs.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        assert instance_id in self.interaction._instance_dict\n\n        await self.interaction.finalize_interaction(instance_id, extra_param=\"test\")\n\n        assert instance_id not in self.interaction._instance_dict\n\n    @pytest.mark.asyncio\n    async def test_finalize_nonexistent_interaction(self):\n        \"\"\"Test finalize_interaction with non-existent instance_id.\"\"\"\n        instance_id = \"nonexistent_instance\"\n\n        # This should raise KeyError\n        with pytest.raises(KeyError):\n            await self.interaction.finalize_interaction(instance_id)\n\n    @pytest.mark.asyncio\n    async def test_full_interaction_workflow_correct(self):\n        \"\"\"Test complete interaction workflow with correct answer.\"\"\"\n        ground_truth = \"42\"\n\n        # Start interaction\n        instance_id = await self.interaction.start_interaction(ground_truth=ground_truth)\n\n        # Generate response with correct answer\n        messages = [{\"role\": \"assistant\", \"content\": \"42\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is True\n        assert reward == 1.0\n\n        # Finalize interaction\n        await self.interaction.finalize_interaction(instance_id)\n        assert instance_id not in self.interaction._instance_dict\n\n    @pytest.mark.asyncio\n    async def test_full_interaction_workflow_incorrect(self):\n        \"\"\"Test complete interaction workflow with incorrect answer.\"\"\"\n        ground_truth = \"42\"\n\n        # Start interaction\n        instance_id = await self.interaction.start_interaction(ground_truth=ground_truth)\n\n        # Generate response with incorrect answer\n        messages = [{\"role\": \"assistant\", \"content\": \"24\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is False\n        assert reward == 0.0\n\n        # Continue with another attempt\n        messages.append({\"role\": \"user\", \"content\": response})\n        messages.append({\"role\": \"assistant\", \"content\": \"42\"})\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=1.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is True\n        assert reward == 1.0\n\n        # Finalize interaction\n        await self.interaction.finalize_interaction(instance_id)\n        assert instance_id not in self.interaction._instance_dict\n\n    @pytest.mark.asyncio\n    async def test_multiple_concurrent_interactions(self):\n        \"\"\"Test multiple concurrent interaction instances.\"\"\"\n        ground_truth_1 = \"42\"\n        ground_truth_2 = \"24\"\n\n        # Start multiple interactions\n        instance_id_1 = await self.interaction.start_interaction(ground_truth=ground_truth_1)\n        instance_id_2 = await self.interaction.start_interaction(ground_truth=ground_truth_2)\n\n        assert len(self.interaction._instance_dict) == 2\n        assert instance_id_1 in self.interaction._instance_dict\n        assert instance_id_2 in self.interaction._instance_dict\n\n        # Test responses for both instances\n        messages_1 = [{\"role\": \"assistant\", \"content\": \"42\"}]\n        messages_2 = [{\"role\": \"assistant\", \"content\": \"24\"}]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", side_effect=[1.0, 1.0]):\n            should_terminate_1, _, reward_1, _ = await self.interaction.generate_response(instance_id_1, messages_1)\n            should_terminate_2, _, reward_2, _ = await self.interaction.generate_response(instance_id_2, messages_2)\n\n        assert should_terminate_1 is True\n        assert should_terminate_2 is True\n        assert reward_1 == 1.0\n        assert reward_2 == 1.0\n\n        # Finalize both interactions\n        await self.interaction.finalize_interaction(instance_id_1)\n        await self.interaction.finalize_interaction(instance_id_2)\n\n        assert len(self.interaction._instance_dict) == 0\n\n    @pytest.mark.asyncio\n    async def test_edge_case_empty_messages(self):\n        \"\"\"Test edge case with empty messages list.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = []\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is False\n        assert reward == 0.0\n        assert self.interaction._instance_dict[instance_id][\"response\"] == \"\"\n\n    @pytest.mark.asyncio\n    async def test_edge_case_message_without_content(self):\n        \"\"\"Test edge case with message without content field.\"\"\"\n        instance_id = \"test_instance\"\n        ground_truth = \"42\"\n\n        # Setup instance\n        await self.interaction.start_interaction(instance_id=instance_id, ground_truth=ground_truth)\n\n        messages = [\n            {\"role\": \"assistant\"}  # Missing content field\n        ]\n\n        with patch(\"verl.utils.reward_score.gsm8k.compute_score\", return_value=0.0):\n            should_terminate, response, reward, metadata = await self.interaction.generate_response(\n                instance_id, messages\n            )\n\n        assert should_terminate is False\n        assert reward == 0.0\n        assert self.interaction._instance_dict[instance_id][\"response\"] is None\n\n    def test_inheritance_from_base_interaction(self):\n        \"\"\"Test that Gsm8kInteraction properly inherits from BaseInteraction.\"\"\"\n        from verl.interactions.base import BaseInteraction\n\n        assert isinstance(self.interaction, BaseInteraction)\n\n        # Test that all required methods are implemented\n        assert hasattr(self.interaction, \"start_interaction\")\n        assert hasattr(self.interaction, \"generate_response\")\n        assert hasattr(self.interaction, \"calculate_score\")\n        assert hasattr(self.interaction, \"finalize_interaction\")\n\n        # Test that methods are callable\n        assert callable(self.interaction.start_interaction)\n        assert callable(self.interaction.generate_response)\n        assert callable(self.interaction.calculate_score)\n        assert callable(self.interaction.finalize_interaction)\n\n    def test_name_attribute_initialization(self):\n        \"\"\"Test name attribute initialization with different configs.\"\"\"\n        # Test with explicit name in config\n        config_with_name = {\"name\": \"custom_gsm8k\"}\n        interaction_with_name = Gsm8kInteraction(config_with_name)\n        assert interaction_with_name.name == \"custom_gsm8k\"\n\n        # Test with default name when not provided in config\n        config_without_name = {}\n        interaction_without_name = Gsm8kInteraction(config_without_name)\n        assert interaction_without_name.name == \"interaction_agent\"  # Default from BaseInteraction\n\n        # Test that name is accessible as attribute\n        assert hasattr(self.interaction, \"name\")\n        assert self.interaction.name == \"gsm8k\"\n"
  },
  {
    "path": "tests/interactions/test_interaction_registry.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport tempfile\n\nimport pytest\nfrom omegaconf import OmegaConf\n\nfrom verl.interactions.base import BaseInteraction\nfrom verl.interactions.gsm8k_interaction import Gsm8kInteraction\nfrom verl.interactions.utils.interaction_registry import (\n    get_interaction_class,\n    initialize_interactions_from_config,\n)\n\n\nclass TestInteractionRegistry:\n    def test_get_interaction_class(self):\n        \"\"\"Test getting interaction class by name.\"\"\"\n        # Test getting base interaction class\n        base_cls = get_interaction_class(\"verl.interactions.base.BaseInteraction\")\n        assert base_cls == BaseInteraction\n\n        # Test getting gsm8k interaction class\n        gsm8k_cls = get_interaction_class(\"verl.interactions.gsm8k_interaction.Gsm8kInteraction\")\n        assert gsm8k_cls == Gsm8kInteraction\n\n    def test_initialize_single_interaction_from_config(self):\n        \"\"\"Test initializing single interaction from config.\"\"\"\n        # Create temporary config file\n        config_content = {\n            \"interaction\": [\n                {\n                    \"name\": \"test_gsm8k\",\n                    \"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\",\n                    \"config\": {},\n                }\n            ]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            interaction_map = initialize_interactions_from_config(temp_config_path)\n\n            # Check that interaction was created\n            assert len(interaction_map) == 1\n            assert \"test_gsm8k\" in interaction_map\n            assert isinstance(interaction_map[\"test_gsm8k\"], Gsm8kInteraction)\n            assert interaction_map[\"test_gsm8k\"].name == \"test_gsm8k\"\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_initialize_multiple_interactions_from_config(self):\n        \"\"\"Test initializing multiple interactions from config.\"\"\"\n        config_content = {\n            \"interaction\": [\n                {\n                    \"name\": \"gsm8k_solver\",\n                    \"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\",\n                    \"config\": {},\n                },\n                {\n                    \"name\": \"base_agent\",\n                    \"class_name\": \"verl.interactions.base.BaseInteraction\",\n                    \"config\": {\"custom_param\": \"test_value\"},\n                },\n            ]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            interaction_map = initialize_interactions_from_config(temp_config_path)\n\n            # Check that both interactions were created\n            assert len(interaction_map) == 2\n            assert \"gsm8k_solver\" in interaction_map\n            assert \"base_agent\" in interaction_map\n\n            # Check types\n            assert isinstance(interaction_map[\"gsm8k_solver\"], Gsm8kInteraction)\n            assert isinstance(interaction_map[\"base_agent\"], BaseInteraction)\n\n            # Check names were injected\n            assert interaction_map[\"gsm8k_solver\"].name == \"gsm8k_solver\"\n            assert interaction_map[\"base_agent\"].name == \"base_agent\"\n\n            # Check custom config was passed\n            assert interaction_map[\"base_agent\"].config.get(\"custom_param\") == \"test_value\"\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_initialize_interaction_without_explicit_name(self):\n        \"\"\"Test that interaction name is derived from class name when not specified.\"\"\"\n        config_content = {\n            \"interaction\": [{\"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\", \"config\": {}}]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            interaction_map = initialize_interactions_from_config(temp_config_path)\n\n            # Check that interaction name was derived from class name\n            assert len(interaction_map) == 1\n            assert \"gsm8k\" in interaction_map  # Should be \"gsm8k\" after removing \"interaction\" suffix\n            assert isinstance(interaction_map[\"gsm8k\"], Gsm8kInteraction)\n            assert interaction_map[\"gsm8k\"].name == \"gsm8k\"\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_initialize_empty_config(self):\n        \"\"\"Test initializing from empty config.\"\"\"\n        config_content = {\"interaction\": []}\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            interaction_map = initialize_interactions_from_config(temp_config_path)\n            assert len(interaction_map) == 0\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_invalid_class_name(self):\n        \"\"\"Test handling of invalid class name.\"\"\"\n        config_content = {\n            \"interaction\": [{\"name\": \"invalid\", \"class_name\": \"invalid.module.InvalidClass\", \"config\": {}}]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            with pytest.raises(ModuleNotFoundError):\n                initialize_interactions_from_config(temp_config_path)\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_duplicate_interaction_names(self):\n        \"\"\"Test handling of duplicate interaction names.\"\"\"\n        config_content = {\n            \"interaction\": [\n                {\"name\": \"duplicate\", \"class_name\": \"verl.interactions.base.BaseInteraction\", \"config\": {}},\n                {\n                    \"name\": \"duplicate\",\n                    \"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\",\n                    \"config\": {},\n                },\n            ]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            with pytest.raises(ValueError, match=\"Duplicate interaction name 'duplicate' found\"):\n                initialize_interactions_from_config(temp_config_path)\n        finally:\n            os.unlink(temp_config_path)\n\n    def test_auto_name_generation_edge_cases(self):\n        \"\"\"Test automatic name generation for various class name patterns.\"\"\"\n        config_content = {\n            \"interaction\": [\n                {\"class_name\": \"verl.interactions.base.BaseInteraction\", \"config\": {}},\n                {\"class_name\": \"verl.interactions.gsm8k_interaction.Gsm8kInteraction\", \"config\": {}},\n            ]\n        }\n\n        with tempfile.NamedTemporaryFile(mode=\"w\", suffix=\".yaml\", delete=False) as f:\n            OmegaConf.save(config_content, f.name)\n            temp_config_path = f.name\n\n        try:\n            interaction_map = initialize_interactions_from_config(temp_config_path)\n\n            # Check that names were generated correctly\n            assert len(interaction_map) == 2\n            assert \"base\" in interaction_map  # BaseInteraction -> base\n            assert \"gsm8k\" in interaction_map  # Gsm8kInteraction -> gsm8k\n        finally:\n            os.unlink(temp_config_path)\n"
  },
  {
    "path": "tests/kill_github_tests.sh",
    "content": "#!/bin/bash\n\nif [ \"$#\" -ne 1 ]; then\n    echo \"Usage: $0 YOUR_GITHUB_TOKEN\"\n    echo \"Please provide exactly one input argument for your github token.\"\n    exit 1\nfi\n\n# Set your GitHub repository details\nOWNER=\"volcengine\"\nREPO=\"verl\"\nTOKEN=$1\n\n# API URL for workflow runs\nAPI_URL=\"https://api.github.com/repos/$OWNER/$REPO/actions/runs?status=queued\"\n\n# Check required commands\ncommand -v jq >/dev/null 2>&1 || { echo \"jq is required but not installed. Aborting.\"; exit 1; }\n\n# Get queued workflow runs\nresponse=$(curl -s -H \"Authorization: token $TOKEN\" -H \"Accept: application/vnd.github.v3+json\" \"$API_URL\")\n\n# Run this for debugging\n# echo $response\n\n# Extract run IDs\nqueued_run_ids=$(echo \"$response\" | jq -r '.workflow_runs[] | .id')\n\nif [ -z \"$queued_run_ids\" ]; then\n    echo \"No queued workflow runs found.\"\n    exit 0\nfi\n\n# Cancel each queued run\nfor run_id in $queued_run_ids; do\n    echo \"Cancelling run $run_id\"\n    cancel_url=\"https://api.github.com/repos/$OWNER/$REPO/actions/runs/$run_id/cancel\"\n    curl -s -X POST -H \"Authorization: token $TOKEN\" -H \"Accept: application/vnd.github.v3+json\" \"$cancel_url\"\ndone\n\necho \"Cancelled all queued workflow runs.\"\n"
  },
  {
    "path": "tests/models/test_engine.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\n\nfrom functools import partial\n\nimport numpy as np\nimport pytest\nimport ray\nimport torch\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\nfrom transformers import (\n    AutoConfig,\n    AutoModelForCausalLM,\n    AutoModelForTokenClassification,\n    AutoTokenizer,\n    Qwen3Config,\n    Qwen3MoeConfig,\n)\n\nfrom verl import DataProto\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.model import compute_position_id_with_mask, create_random_mask\nfrom verl.utils.torch_functional import logprobs_from_logits_naive\nfrom verl.workers.config import (\n    ActorConfig,\n    CriticConfig,\n    FSDPEngineConfig,\n    FSDPOptimizerConfig,\n    HFModelConfig,\n    McoreEngineConfig,\n    McoreOptimizerConfig,\n)\nfrom verl.workers.engine_workers import TrainingWorker, TrainingWorkerConfig\nfrom verl.workers.utils.losses import ppo_loss, sft_loss, value_loss\nfrom verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding\n\n\ndef get_test_language_model(device_count):\n    if device_count == 1:\n        model = \"~/models/HuggingFaceTB/SmolLM2-135M-Instruct\"\n    else:\n        model = \"~/models/Qwen/Qwen2.5-0.5B\"\n    model = os.path.expanduser(model)\n    return model\n\n\ndef create_training_config(model_type, strategy, device_count, model):\n    if device_count == 1:\n        tp = pp = cp = fsdp_size = 1\n    else:\n        tp = pp = cp = 2\n        fsdp_size = 4\n\n    path = os.path.expanduser(model)\n    model_config = HFModelConfig(path=path, use_remove_padding=True)\n\n    kwargs = dict(\n        param_offload=True,\n        optimizer_offload=True,\n        grad_offload=True,\n        use_dynamic_bsz=True,\n        use_remove_padding=True,\n        max_token_len_per_gpu=500,\n        infer_max_token_len_per_gpu=1000,\n    )\n\n    if strategy == \"megatron\":\n        engine_config = McoreEngineConfig(\n            forward_only=False,\n            use_mbridge=True,\n            tensor_model_parallel_size=tp,\n            pipeline_model_parallel_size=pp,\n            context_parallel_size=cp,\n            **kwargs,\n        )\n        optimizer_config = McoreOptimizerConfig(lr_decay_steps=10)\n    elif strategy in [\"fsdp\", \"fsdp2\"]:\n        engine_config = FSDPEngineConfig(\n            forward_only=False, fsdp_size=fsdp_size, strategy=strategy, ulysses_sequence_parallel_size=cp, **kwargs\n        )\n        optimizer_config = FSDPOptimizerConfig()\n    else:\n        raise NotImplementedError(f\"strategy {strategy} is not supported\")\n\n    config = TrainingWorkerConfig(\n        model_type=model_type,\n        model_config=model_config,\n        engine_config=engine_config,\n        optimizer_config=optimizer_config,\n        checkpoint_config=None,\n    )\n    return config\n\n\n@pytest.mark.parametrize(\"strategy\", [\"fsdp\", \"fsdp2\", \"megatron\"])\ndef test_actor_engine(strategy):\n    ray.init()\n    device_count = torch.cuda.device_count()\n    config = create_training_config(\n        model_type=\"language_model\",\n        strategy=strategy,\n        device_count=device_count,\n        model=get_test_language_model(device_count),\n    )\n    ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config)\n    resource_pool = RayResourcePool(process_on_nodes=[device_count])\n    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init)\n    # init model\n    wg.reset()\n\n    sft_loss_ = partial(sft_loss, config=config)\n\n    wg.set_loss_fn(sft_loss_)\n\n    batch_size = 8\n    seqlen = 32\n\n    response_length = seqlen // 2\n\n    torch.manual_seed(1)\n    np.random.seed(1)\n\n    input_ids = torch.randint(0, config.model_config.hf_config.vocab_size, (batch_size, seqlen))\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6\n    )\n    position_ids = compute_position_id_with_mask(attention_mask)\n\n    global_token_num = torch.sum(attention_mask, dim=-1).tolist()\n\n    print(input_ids.float().mean(), attention_mask.float().mean())\n\n    responses = input_ids[:, response_length:]\n    response_mask = attention_mask[:, response_length:]\n\n    assert torch.all(response_mask[:, 0] == 1)\n\n    data = DataProto.from_single_dict(\n        {\n            \"input_ids\": input_ids,\n            \"prompts\": input_ids[:, :response_length],\n            \"attention_mask\": attention_mask,\n            \"position_ids\": position_ids,\n            \"responses\": responses,\n            \"response_mask\": response_mask,\n        },\n        meta_info={\"temperature\": 1.0, \"global_token_num\": global_token_num, \"compute_loss\": False},\n    )\n\n    data_td = data.to_tensordict()\n    data_td = left_right_2_no_padding(data_td)\n\n    # eval\n    output = wg.infer_batch(data_td)\n    output = output.get()\n    logprobs_unpad = tu.get(output, \"log_probs\").cpu()\n    logprobs = no_padding_2_padding(logprobs_unpad, data_td)\n\n    output = DataProto.from_single_dict({\"old_log_probs\": logprobs})\n\n    # load hf model and compare results with hf model\n    path = config.model_config.path\n    hf_model = AutoModelForCausalLM.from_pretrained(path, torch_dtype=torch.bfloat16)\n    hf_output = hf_model(input_ids, attention_mask=attention_mask)\n    hf_logprobs = logprobs_from_logits_naive(\n        hf_output.logits[:, -response_length - 1 : -1, :].float(), input_ids[:, -response_length:]\n    )\n    hf_logprobs_mean = torch.mean(hf_logprobs * response_mask)\n    mcore_logprobs_mean = torch.mean(output.batch[\"old_log_probs\"] * response_mask)\n\n    torch.testing.assert_close(hf_logprobs_mean, mcore_logprobs_mean, atol=1e-3, rtol=1e-2)\n\n    data = data.union(output)\n\n    # TODO: sft_loss_ is not compatible with ActorWorker until we replace DataProto with torch.jagged TensorDict\n    # wg.set_loss_fn(sft_loss_)\n\n    # train for one step\n    # metrics = wg.update_actor(data)\n    # print(metrics)\n\n    # add ppo data\n    data.batch[\"advantages\"] = torch.rand_like(responses, dtype=torch.float32)\n    data.batch[\"ref_log_prob\"] = torch.rand_like(responses, dtype=torch.float32)\n\n    # construct actor config\n    actor_config = ActorConfig(strategy=strategy, rollout_n=1, ppo_micro_batch_size_per_gpu=-1)\n\n    # set ppo loss\n    ppo_loss_ = partial(ppo_loss, config=actor_config)\n    wg.set_loss_fn(ppo_loss_)\n\n    # update again\n    data_td = data.to_tensordict()\n    data_td = left_right_2_no_padding(data_td)\n\n    # auto load/offload\n    tu.assign_non_tensor(data_td, global_batch_size=data_td.shape[0])\n    ppo_metrics = wg.train_batch(data_td)\n    ppo_metrics = ppo_metrics.get()\n    ppo_metrics = tu.get(ppo_metrics, \"metrics\")\n    print(ppo_metrics)\n\n    # test manual load/offload\n    tu.assign_non_tensor(data_td, disable_auto_offload=True)\n    wg.to(\"device\")\n    ppo_metrics = wg.train_batch(data_td)\n    ppo_metrics = ppo_metrics.get()\n    ppo_metrics = tu.get(ppo_metrics, \"metrics\")\n    print(ppo_metrics)\n    wg.to(\"cpu\")\n\n    ray.shutdown()\n\n\ndef create_value_model(language_model_path, output_path):\n    config = AutoConfig.from_pretrained(language_model_path)\n    config.num_labels = 1\n    config.classifier_dropout = 0\n    config.tie_word_embeddings = False\n    model = AutoModelForTokenClassification.from_config(config)\n    tokenizer = AutoTokenizer.from_pretrained(os.path.expanduser(language_model_path))\n    assert model.config.num_labels == 1\n    path = os.path.expanduser(output_path)\n    model.save_pretrained(path)\n    tokenizer.save_pretrained(path)\n    config.save_pretrained(path)\n    return path\n\n\n@pytest.mark.parametrize(\"strategy\", [\"fsdp\", \"fsdp2\"])\ndef test_critic_engine(strategy):\n    device_count = torch.cuda.device_count()\n    value_model_path = os.path.expanduser(\"~/models/test_model\")\n    language_model_path = get_test_language_model(device_count=device_count)\n    create_value_model(language_model_path, value_model_path)\n\n    torch.manual_seed(1)\n    np.random.seed(1)\n\n    ray.init()\n\n    config = create_training_config(\n        model_type=\"value_model\", strategy=strategy, device_count=device_count, model=value_model_path\n    )\n    ray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(TrainingWorker), config=config)\n    resource_pool = RayResourcePool(process_on_nodes=[device_count])\n    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init)\n    # init model\n    wg.reset()\n\n    batch_size = 8\n    seqlen = 32\n\n    response_length = seqlen // 2\n    input_ids = torch.randint(0, config.model_config.hf_config.vocab_size, (batch_size, seqlen))\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_valid_token=0.8, max_ratio_of_left_padding=0.2, min_ratio_of_valid_token=0.6\n    )\n    position_ids = compute_position_id_with_mask(attention_mask)\n\n    global_token_num = torch.sum(attention_mask, dim=-1).tolist()\n\n    print(input_ids.float().mean(), attention_mask.float().mean())\n\n    responses = input_ids[:, response_length:]\n    response_mask = attention_mask[:, response_length:]\n\n    assert torch.all(response_mask[:, 0] == 1)\n\n    data = DataProto.from_single_dict(\n        {\n            \"input_ids\": input_ids,\n            \"prompts\": input_ids[:, :response_length],\n            \"attention_mask\": attention_mask,\n            \"position_ids\": position_ids,\n            \"responses\": responses,\n            \"response_mask\": response_mask,\n        },\n        meta_info={\"temperature\": 1.0, \"global_token_num\": global_token_num, \"compute_loss\": False},\n    )\n\n    data_td = data.to_tensordict()\n    data_td = left_right_2_no_padding(data_td)\n\n    # eval\n    output = wg.infer_batch(data_td)\n    output = output.get()\n\n    values_unpad = tu.get(output, \"values\").float().cpu()\n    values = no_padding_2_padding(values_unpad, data_td)\n\n    output = DataProto.from_single_dict({\"values\": values})\n\n    # load hf model and compare results with hf model\n    with torch.device(\"cuda\"), torch.autocast(device_type=\"cuda\", dtype=torch.bfloat16):\n        hf_model = AutoModelForTokenClassification.from_pretrained(\n            value_model_path, torch_dtype=torch.float32, attn_implementation=\"flash_attention_2\"\n        )\n        hf_output = hf_model(input_ids.cuda(), attention_mask=attention_mask.cuda())\n        hf_values = hf_output.logits[:, -response_length - 1 : -1, :].float().squeeze(-1).cpu()\n\n    hf_values_mean = torch.mean(hf_values * response_mask)\n    engine_values = torch.mean(output.batch[\"values\"] * response_mask)\n\n    torch.testing.assert_close(hf_values_mean, engine_values, atol=1e-2, rtol=1e-2)\n\n    data = data.union(output)\n\n    # add ppo data\n    data.batch[\"returns\"] = torch.rand_like(responses, dtype=torch.float32)\n\n    # update again\n    # create critic config\n    critic_config = CriticConfig(\n        strategy=strategy, rollout_n=1, ppo_micro_batch_size_per_gpu=-1, model_config=config.model_config\n    )\n    value_loss_ = partial(value_loss, config=critic_config)\n    wg.set_loss_fn(value_loss_)\n\n    # update again\n    data_td = data.to_tensordict()\n    data_td = left_right_2_no_padding(data_td)\n\n    # auto load/offload\n    tu.assign_non_tensor(data_td, global_batch_size=data_td.shape[0])\n    ppo_metrics = wg.train_batch(data_td)\n    ppo_metrics = ppo_metrics.get()\n    ppo_metrics = tu.get(ppo_metrics, \"metrics\")\n    print(ppo_metrics)\n\n    ray.shutdown()\n\n\ndef create_actor_model(tmp_path, config):\n    model = AutoModelForCausalLM.from_config(config)\n    path = os.path.join(tmp_path, \"test_model\")\n    model.save_pretrained(path)\n    config.save_pretrained(path)\n    return path\n\n\ndef _worker(rank: int, world_size: int, rendezvous_file: str, strategy: str, model_path: str):\n    torch.cuda.set_device(rank)\n    dist.init_process_group(\n        backend=\"nccl\",\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n\n    ref_model_config = AutoConfig.from_pretrained(model_path)\n    with torch.device(\"meta\"):\n        ref_model = AutoModelForCausalLM.from_config(ref_model_config)\n\n    from verl.workers.engine import BaseEngine, EngineRegistry\n\n    # construct configs\n    model_config = HFModelConfig(path=model_path, load_tokenizer=False)\n\n    if strategy == \"megatron\":\n        engine_config = McoreEngineConfig(\n            forward_only=False,\n            use_mbridge=True,\n            tensor_model_parallel_size=2,\n            pipeline_model_parallel_size=2,\n            context_parallel_size=1,\n        )\n        optimizer_config = McoreOptimizerConfig(lr_decay_steps=10)\n    elif strategy in [\"fsdp\", \"fsdp2\"]:\n        engine_config = FSDPEngineConfig(\n            forward_only=False, fsdp_size=4, strategy=strategy, ulysses_sequence_parallel_size=2\n        )\n        optimizer_config = FSDPOptimizerConfig()\n    else:\n        raise NotImplementedError(f\"strategy {strategy} is not supported\")\n\n    checkpoint_config = CheckpointConfig()\n\n    # build model engine\n    engine: BaseEngine = EngineRegistry.new(\n        model_type=\"language_model\",\n        backend=engine_config.strategy,\n        model_config=model_config,\n        engine_config=engine_config,\n        optimizer_config=optimizer_config,\n        checkpoint_config=checkpoint_config,\n    )\n\n    engine.initialize()\n\n    # get per tensor parameter\n    per_tensor_params, _ = engine.get_per_tensor_param()\n\n    ref_state_dict = ref_model.state_dict()\n\n    # load ground truth and compare\n    for key, value in per_tensor_params:\n        assert key in ref_state_dict, f\"{key} not in ref_state_dict\"\n        assert value.shape == ref_state_dict[key].shape, (\n            f\"{key} shape not equal, {value.shape} != {ref_state_dict[key].shape}\"\n        )\n        if rank == 0:\n            print(key, value.shape)\n\n    dist.barrier()\n    dist.destroy_process_group()\n\n\n@pytest.mark.parametrize(\"world_size\", [8])\n@pytest.mark.parametrize(\"config\", [Qwen3Config(num_hidden_layers=2), Qwen3MoeConfig(num_hidden_layers=2)])\n@pytest.mark.parametrize(\"strategy\", [\"megatron\", \"fsdp\", \"fsdp2\"])\ndef test_per_tensor_generator(world_size, tmp_path, config, strategy):\n    rendezvous_file = str(tmp_path / \"rdzv_mask\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n    # create a model\n    model_path = create_actor_model(tmp_path, config)\n    # spawn workers\n    mp.spawn(\n        fn=_worker,\n        args=(world_size, rendezvous_file, strategy, model_path),\n        nprocs=world_size,\n        join=True,\n    )\n"
  },
  {
    "path": "tests/models/test_tiled_mlp_accuracy.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTest script to verify TiledMLP accuracy by comparing logits and gradients\nbetween regular MLP and TiledMLP under FSDP2.\nRun with: torchrun --nproc_per_node=2 tests/test_tiled_mlp_accuracy.py\n\"\"\"\n\nimport torch\nimport torch.distributed as dist\nfrom torch.distributed.device_mesh import init_device_mesh\nfrom torch.distributed.fsdp import fully_shard\n\n\ndef setup_distributed():\n    dist.init_process_group(backend=\"nccl\")\n    rank = dist.get_rank()\n    world_size = dist.get_world_size()\n    torch.cuda.set_device(rank)\n    return rank, world_size\n\n\ndef create_model(model_name=\"Qwen/Qwen3-1.7B\", num_layers=2):\n    \"\"\"Load a Qwen3-1.7B model with only 2 layers from pretrained weights.\"\"\"\n    from transformers import AutoConfig, AutoModelForCausalLM\n\n    config = AutoConfig.from_pretrained(model_name, trust_remote_code=True)\n    config.num_hidden_layers = num_layers\n\n    model = AutoModelForCausalLM.from_pretrained(\n        model_name,\n        config=config,\n        torch_dtype=torch.bfloat16,\n        trust_remote_code=True,\n        attn_implementation=\"flash_attention_2\",\n    )\n    return model\n\n\ndef apply_fsdp2(model, device_mesh):\n    \"\"\"Apply FSDP2 sharding to model.\"\"\"\n    for layer in model.model.layers:\n        fully_shard(layer, mesh=device_mesh)\n    fully_shard(model, mesh=device_mesh)\n    return model\n\n\ndef run_forward_backward(model, input_ids, labels):\n    \"\"\"Run forward and backward pass, return logits and gradients.\"\"\"\n    model.zero_grad()\n\n    outputs = model(input_ids=input_ids, labels=labels)\n    logits = outputs.logits.clone().detach()\n    loss = outputs.loss\n\n    loss.backward()\n\n    # Collect MLP gradients\n    gradients = {}\n    for name, param in model.named_parameters():\n        if \"mlp\" in name and param.grad is not None:\n            gradients[name] = param.grad.clone().detach()\n\n    return logits, gradients, loss.item()\n\n\ndef compare_results(logits1, grads1, logits2, grads2, rank):\n    \"\"\"Compare logits and gradients between two runs.\"\"\"\n    # Compare logits\n    logits_diff = (logits1 - logits2).abs()\n    logits_max_diff = logits_diff.max().item()\n    logits_mean_diff = logits_diff.mean().item()\n\n    # Compare gradients (only for params that exist on this rank due to FSDP sharding)\n    all_pass = True\n    grad_results = []\n    for name in sorted(grads1.keys()):\n        if name in grads2:\n            g1, g2 = grads1[name], grads2[name]\n            diff = (g1 - g2).abs()\n            max_diff = diff.max().item()\n            mean_diff = diff.mean().item()\n\n            # Check if within tolerance (1e-2 for bf16)\n            passed = max_diff < 1e-2\n            if not passed:\n                all_pass = False\n            grad_results.append((name, max_diff, mean_diff, passed))\n\n    # Only print on rank 0 to avoid duplicate output\n    if rank == 0:\n        print(\"\\n=== Comparison Results ===\")\n        print(\"\\nLogits:\")\n        print(f\"  Max diff: {logits_max_diff:.2e}\")\n        print(f\"  Mean diff: {logits_mean_diff:.2e}\")\n\n        print(\"\\nMLP Parameter Gradients:\")\n        if grad_results:\n            for name, max_diff, mean_diff, passed in grad_results:\n                status = \"✓\" if passed else \"✗\"\n                print(f\"  {name}: max={max_diff:.2e}, mean={mean_diff:.2e} {status}\")\n        else:\n            print(\"  (Gradients sharded to other ranks under FSDP2)\")\n\n    return all_pass\n\n\ndef main():\n    rank, world_size = setup_distributed()\n    device_mesh = init_device_mesh(\"cuda\", (world_size,))\n\n    model_name = \"Qwen/Qwen3-1.7B\"\n    num_layers = 2\n\n    if rank == 0:\n        print(f\"Running TiledMLP accuracy test with {world_size} GPUs\")\n        print(f\"Model: {model_name} ({num_layers} layers, from pretrained)\")\n\n    dist.barrier()\n\n    # ========== Create Model 1: WITHOUT TiledMLP ==========\n    if rank == 0:\n        print(\"\\n\" + \"=\" * 60)\n        print(\"Creating Model 1 (without TiledMLP)\")\n        print(\"=\" * 60)\n\n    model1 = create_model(model_name, num_layers)\n    model1 = apply_fsdp2(model1, device_mesh)\n    model1 = model1.cuda()\n\n    # Create deterministic input\n    torch.manual_seed(42)\n    batch_size, seq_len = 2, 256\n    vocab_size = model1.config.vocab_size\n    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=\"cuda\")\n    labels = input_ids.clone()\n\n    # ========== Run Model 1: WITHOUT TiledMLP ==========\n    if rank == 0:\n        print(\"\\n\" + \"=\" * 60)\n        print(\"Running forward/backward on Model 1 (without TiledMLP)\")\n        print(\"=\" * 60)\n\n    logits1, grads1, loss1 = run_forward_backward(model1, input_ids, labels)\n    if rank == 0:\n        print(f\"Loss: {loss1:.4f}\")\n\n    # Free model1 memory before creating model2\n    del model1\n    torch.cuda.empty_cache()\n\n    dist.barrier()\n\n    # ========== Create Model 2, apply TiledMLP patch, then FSDP2 ==========\n    if rank == 0:\n        print(\"\\n\" + \"=\" * 60)\n        print(\"Creating Model 2 (with TiledMLP, patch before FSDP2)\")\n        print(\"=\" * 60)\n\n    model2 = create_model(model_name, num_layers)\n\n    # Apply TiledMLP patch AFTER model instantiation but BEFORE FSDP2 wrap\n    if rank == 0:\n        print(\"Applying TiledMLP monkey patch before FSDP2...\")\n\n    from verl.models.transformers.tiled_mlp import apply_tiled_mlp_monkey_patch\n\n    apply_tiled_mlp_monkey_patch(num_shards=4, model_type=\"qwen3\")\n\n    model2 = apply_fsdp2(model2, device_mesh)\n    model2 = model2.cuda()\n\n    dist.barrier()\n\n    # ========== Run Model 2: WITH TiledMLP ==========\n    if rank == 0:\n        print(\"\\n\" + \"=\" * 60)\n        print(\"Running forward/backward on Model 2 (with TiledMLP)\")\n        print(\"=\" * 60)\n\n    logits2, grads2, loss2 = run_forward_backward(model2, input_ids, labels)\n    if rank == 0:\n        print(f\"Loss: {loss2:.4f}\")\n\n    dist.barrier()\n\n    # ========== Compare Results ==========\n    all_pass = compare_results(logits1, grads1, logits2, grads2, rank)\n\n    dist.barrier()\n\n    if rank == 0:\n        print(\"\\n\" + \"=\" * 60)\n        print(\"SUMMARY\")\n        print(\"=\" * 60)\n        print(f\"Loss diff: {abs(loss1 - loss2):.2e}\")\n        print(f\"All gradient checks: {'PASS' if all_pass else 'FAIL'}\")\n\n    # Cleanup\n    del model2\n    torch.cuda.empty_cache()\n\n    dist.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/models/test_transformer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom transformers import (\n    ApertusConfig,\n    AutoModelForCausalLM,\n    AutoModelForTokenClassification,\n    GemmaConfig,\n    LlamaConfig,\n    MistralConfig,\n    Qwen2Config,\n)\n\nfrom verl.utils.device import get_device_name\n\nif get_device_name() == \"cuda\":\n    from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input\nelif get_device_name() == \"npu\":\n    from verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input\n\nfrom verl.utils.model import compute_position_id_with_mask, create_random_mask\nfrom verl.utils.torch_functional import log_probs_from_logits_all_rmpad, masked_mean\n\n# TODO(sgm): add more models for test\n# we only need one scale for each model\ntest_configs = [\n    LlamaConfig(num_hidden_layers=1),\n    MistralConfig(num_hidden_layers=1),\n    GemmaConfig(num_hidden_layers=1),\n    Qwen2Config(num_hidden_layers=1),\n    ApertusConfig(num_hidden_layers=1),\n]\n\n\ndef test_hf_casual_models():\n    batch_size = 4\n    seqlen = 128\n    response_length = 127\n\n    for config in test_configs:\n        # config = AutoConfig.from_pretrained(test_case)\n        with torch.device(get_device_name()):\n            model = AutoModelForCausalLM.from_config(\n                config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n            )\n            model = model.to(device=get_device_name())\n        input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name())\n        attention_mask = create_random_mask(\n            input_ids=input_ids,\n            max_ratio_of_left_padding=0.1,\n            max_ratio_of_valid_token=0.8,\n            min_ratio_of_valid_token=0.5,\n        )\n        position_ids = compute_position_id_with_mask(\n            attention_mask\n        )  # TODO(sgm): we can construct the position_ids_rmpad here\n\n        input_ids_rmpad, indices, *_ = unpad_input(\n            input_ids.unsqueeze(-1), attention_mask\n        )  # input_ids_rmpad (total_nnz, ...)\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n\n        # unpad the position_ids to align the rotary\n        position_ids_rmpad = index_first_axis(\n            rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n        ).transpose(0, 1)\n\n        # input with input_ids_rmpad and postition_ids to enable flash attention varlen\n        logits_rmpad = model(\n            input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False\n        ).logits  # (1, total_nnz, vocab_size)\n\n        origin_logits = model(\n            input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False\n        ).logits\n        origin_logits_rmpad, origin_logits_indices, *_ = unpad_input(origin_logits, attention_mask)\n\n        logits_rmpad = logits_rmpad.squeeze(0)\n        log_probs = log_probs_from_logits_all_rmpad(\n            input_ids_rmpad=input_ids_rmpad,\n            logits_rmpad=logits_rmpad,\n            indices=indices,\n            batch_size=batch_size,\n            seqlen=seqlen,\n            response_length=response_length,\n        )  # (batch, seqlen)\n        origin_log_probs = log_probs_from_logits_all_rmpad(\n            input_ids_rmpad=input_ids_rmpad,\n            logits_rmpad=origin_logits_rmpad,\n            indices=origin_logits_indices,\n            batch_size=batch_size,\n            seqlen=seqlen,\n            response_length=response_length,\n        )  # (batch, seqlen)\n\n        torch.testing.assert_close(\n            masked_mean(log_probs, attention_mask[:, -response_length - 1 : -1]),\n            masked_mean(origin_log_probs, attention_mask[:, -response_length - 1 : -1]),\n            atol=1e-2,\n            rtol=1e-5,\n        )\n    print(\"Check pass\")\n\n\ndef test_hf_value_models():\n    batch_size = 4\n    seqlen = 128\n\n    for config in test_configs:\n        # config = AutoConfig.from_pretrained(test_case)\n        config.num_labels = 1\n        config.classifier_dropout = 0\n        config.hidden_dropout = 0\n        with torch.device(get_device_name()):\n            model = AutoModelForTokenClassification.from_config(\n                config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n            )\n            model = model.to(device=get_device_name())\n        input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name())\n        attention_mask = create_random_mask(\n            input_ids=input_ids,\n            max_ratio_of_left_padding=0.1,\n            max_ratio_of_valid_token=0.8,\n            min_ratio_of_valid_token=0.5,\n        )\n        position_ids = compute_position_id_with_mask(\n            attention_mask\n        )  # TODO(sgm): we can construct the position_ids_rmpad here\n\n        input_ids_rmpad, indices, *_ = unpad_input(\n            input_ids.unsqueeze(-1), attention_mask\n        )  # input_ids_rmpad (total_nnz, ...)\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n\n        # unpad the position_ids to align the rotary\n        position_ids_rmpad = index_first_axis(\n            rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n        ).transpose(0, 1)\n\n        origin_logits = model(\n            input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids, use_cache=False\n        ).logits\n\n        # input with input_ids_rmpad and postition_ids to enable flash attention varlen\n        rmpad_logits = model(\n            input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False\n        ).logits  # (1, total_nnz, 1)\n        rmpad_logits = rmpad_logits.squeeze(0)\n        pad_logits = pad_input(rmpad_logits, indices, batch_size, seqlen=seqlen)\n\n        torch.testing.assert_close(\n            masked_mean(pad_logits, attention_mask[:, :, None]),\n            masked_mean(origin_logits, attention_mask[:, :, None]),\n            atol=1e-2,\n            rtol=1e-5,\n        )\n    print(\"Value model check pass\")\n\n\ndef test_attn_implementation_override():\n    \"\"\"Test that attn_implementation override config is properly respected.\"\"\"\n    # Test case 1: Test the actual extraction logic (no network required)\n    test_cases = [\n        ({}, \"flash_attention_2\"),  # Default case\n        ({\"attn_implementation\": \"eager\"}, \"eager\"),  # Override case\n        ({\"attn_implementation\": \"sdpa\"}, \"sdpa\"),  # Another override\n        ({\"other_config\": \"value\"}, \"flash_attention_2\"),  # No attn_implementation key\n    ]\n\n    for override_config, expected in test_cases:\n        actual = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert actual == expected, f\"Expected {expected}, got {actual} for config {override_config}\"\n\n    # Test case 2: Test with local config creation (simulate FSDP worker behavior)\n    # Test default behavior\n    override_config_default = {}\n    attn_implementation_default = override_config_default.get(\"attn_implementation\", \"flash_attention_2\")\n    assert attn_implementation_default == \"flash_attention_2\"\n\n    # Test override behavior\n    override_config_eager = {\"attn_implementation\": \"eager\"}\n    attn_implementation_eager = override_config_eager.get(\"attn_implementation\", \"flash_attention_2\")\n    assert attn_implementation_eager == \"eager\"\n\n    # Test that we can create a config with specific attn_implementation\n    config_with_eager = LlamaConfig(num_hidden_layers=1, _attn_implementation=\"eager\")\n    assert config_with_eager._attn_implementation == \"eager\"\n\n    config_with_flash = LlamaConfig(num_hidden_layers=1, _attn_implementation=\"flash_attention_2\")\n    assert config_with_flash._attn_implementation == \"flash_attention_2\"\n\n    print(\"✓ All attn_implementation override config tests passed\")\n\n\ndef test_fsdp_worker_attn_implementation_integration():\n    \"\"\"Test integration of attn_implementation with FSDP worker logic.\"\"\"\n\n    # Mock the FSDP worker configuration scenario\n    mock_override_config = {\"attn_implementation\": \"eager\"}\n\n    # Test the exact logic used in FSDP workers\n    attn_implementation = mock_override_config.get(\"attn_implementation\", \"flash_attention_2\")\n    assert attn_implementation == \"eager\"\n\n    # Test with empty config (should default)\n    mock_override_config_empty = {}\n    attn_implementation_default = mock_override_config_empty.get(\"attn_implementation\", \"flash_attention_2\")\n    assert attn_implementation_default == \"flash_attention_2\"\n\n    # Test that the parameter would be passed correctly to both AutoConfig and Model\n    expected_calls = [\n        (\"AutoConfig.from_pretrained\", {\"attn_implementation\": attn_implementation}),\n        (\"AutoModel.from_pretrained\", {\"attn_implementation\": attn_implementation}),\n    ]\n\n    # Verify the parameter extraction works as expected\n    for call_name, expected_params in expected_calls:\n        assert expected_params[\"attn_implementation\"] == \"eager\"\n\n    print(\"✓ FSDP worker integration test passed\")\n\n\nif __name__ == \"__main__\":\n    test_hf_casual_models()\n    test_hf_value_models()\n    test_attn_implementation_override()\n    test_fsdp_worker_attn_implementation_integration()\n"
  },
  {
    "path": "tests/models/test_transformers_ulysses.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport contextlib\nimport copy\nfrom dataclasses import dataclass\n\nimport pytest\nimport torch\nimport torch.distributed\nimport transformers\nfrom packaging import version\nfrom torch.distributed import init_device_mesh\nfrom transformers import AutoModelForCausalLM, LlamaConfig, PretrainedConfig, Qwen2Config\n\nfrom verl.models.transformers.monkey_patch import apply_monkey_patch\nfrom verl.protocol import DataProto\nfrom verl.utils.device import get_device_name, get_torch_device\nfrom verl.utils.distributed import initialize_global_process_group\nfrom verl.utils.model import compute_position_id_with_mask, create_random_mask\nfrom verl.utils.ulysses import (\n    gather_outputs_and_unpad,\n    get_ulysses_sequence_parallel_world_size,\n    set_ulysses_sequence_parallel_group,\n    ulysses_pad_and_slice_inputs,\n)\nfrom verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager\n\nif get_device_name() == \"cuda\":\n    from flash_attn.bert_padding import index_first_axis, rearrange, unpad_input\nelif get_device_name() == \"npu\":\n    from verl.utils.attention_utils import index_first_axis, rearrange, unpad_input\n\n# TODO(sgm): add more models for test\n# we only need one scale for each model\n\n\n@dataclass\nclass SequenceParallelConfig:\n    config: PretrainedConfig\n    sp_size: int\n    is_valid: bool\n\n\ndef test_configs():\n    configs = [\n        SequenceParallelConfig(\n            LlamaConfig(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=32), sp_size=8, is_valid=True\n        ),\n        SequenceParallelConfig(\n            Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584),\n            sp_size=4,\n            is_valid=True,\n        ),\n        SequenceParallelConfig(\n            Qwen2Config(num_hidden_layers=2, num_attention_heads=28, num_key_value_heads=4, hidden_size=3584),\n            sp_size=8,\n            is_valid=False,\n        ),\n        SequenceParallelConfig(\n            Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=4, is_valid=True\n        ),\n        SequenceParallelConfig(\n            Qwen2Config(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=4), sp_size=8, is_valid=True\n        ),\n    ]\n\n    if version.parse(transformers.__version__) >= version.parse(\"4.56.0\"):\n        from transformers import ApertusConfig\n\n        configs.append(\n            SequenceParallelConfig(\n                ApertusConfig(num_hidden_layers=2, num_attention_heads=32, num_key_value_heads=32, hidden_size=4096),\n                sp_size=8,\n                is_valid=True,\n            )\n        )\n\n    return configs\n\n\ndef sync_model_parameters_global(layer):\n    # synchronize weights\n    for p in layer.parameters():\n        torch.distributed.broadcast(tensor=p.data, src=0)\n\n\n@pytest.mark.parametrize(\"test_config\", test_configs())\ndef test_hf_casual_fwd_bwd(test_config):\n    if not torch.distributed.is_initialized():\n        initialize_global_process_group()\n\n    context = contextlib.nullcontext() if test_config.is_valid else pytest.raises(AssertionError)\n    with context:\n        world_size = torch.distributed.get_world_size()\n        _hf_casual_fwd_bwd(test_config.config, test_config.sp_size, world_size // test_config.sp_size)\n\n    # TODO: seems not work, will cause `socketStartConnect: Connect to xxx failed : Software caused connection abort`\n    # torch.distributed.destroy_process_group()\n\n\ndef _hf_casual_fwd(config, sp_size, dp_size):\n    assert get_torch_device().device_count() >= 2, \"need at least 2 gpus for test\"\n\n    ulysses_device_mesh = init_device_mesh(\n        device_type=get_device_name(), mesh_shape=(dp_size, sp_size), mesh_dim_names=(\"dp\", \"sp\")\n    )\n    sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh)\n\n    batch_size = 1\n    seqlen = 128\n    # response_length = 127\n\n    # patch before load\n    with torch.device(get_device_name()):\n        model = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        apply_monkey_patch(model, sp_size)\n        model = model.to(device=get_device_name())\n        sync_model_parameters_global(model)\n\n    # different rank will generate different input_ids following fsdp\n    input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name())\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8\n    )\n    position_ids = compute_position_id_with_mask(\n        attention_mask\n    )  # TODO(sgm): we can construct the position_ids_rmpad here\n\n    model_inputs = {\n        \"input_ids\": input_ids.to(get_device_name()),\n        \"attention_mask\": attention_mask.to(get_device_name()),\n        \"position_ids\": position_ids.int().to(get_device_name()),\n    }\n\n    model_inputs = DataProto.from_dict(model_inputs)\n\n    # 1. perform ulysses forward\n    with sharding_manager:\n        model_inputs = sharding_manager.preprocess_data(model_inputs)\n        input_ids = model_inputs.batch[\"input_ids\"]\n        attention_mask = model_inputs.batch[\"attention_mask\"]\n        position_ids = model_inputs.batch[\"position_ids\"]\n        input_ids_rmpad, indices, *_ = unpad_input(\n            input_ids.unsqueeze(-1), attention_mask\n        )  # input_ids_rmpad (total_nnz, ...)\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n        # unpad the position_ids to align the rotary\n        position_ids_rmpad = index_first_axis(\n            rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n        ).transpose(0, 1)\n\n        # slice input tensor for ulysses\n        # input_ids are padded and sliced\n        # postition_ids are only padded but not sliced\n        input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs(\n            input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size()\n        )\n\n        # input with input_ids_rmpad and postition_ids to enable flash attention varlen\n        logits_split_in_seq = model(\n            input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False\n        ).logits  # (1, total_nnz/n, vocab_size)\n\n        # all_gather output\n        logits_full = gather_outputs_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size)\n\n    # 2. perform normal forward\n    set_ulysses_sequence_parallel_group(None)\n    logits_rmpad_local = model(\n        input_ids_rmpad, position_ids=position_ids_rmpad, use_cache=False\n    ).logits  # (1, total_nnz, vocab_size)\n\n    mean_local = logits_rmpad_local.mean()\n    mean_full = logits_full.mean()\n    torch.testing.assert_close(mean_local, mean_full, rtol=1e-2, atol=1e-5)\n\n\ndef _hf_casual_fwd_bwd(config, sp_size, dp_size):\n    assert get_torch_device().device_count() >= 2, \"need at least 2 gpus for test\"\n\n    ulysses_device_mesh = init_device_mesh(\n        device_type=get_device_name(), mesh_shape=(dp_size, sp_size), mesh_dim_names=(\"dp\", \"sp\")\n    )\n    sharding_manager = FSDPUlyssesShardingManager(ulysses_device_mesh)\n\n    batch_size = 1\n    seqlen = 128\n    # response_length = 127\n\n    # patch before load\n    with torch.device(get_device_name()):\n        model = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        apply_monkey_patch(model, sp_size)\n        model = model.to(device=get_device_name())\n        sync_model_parameters_global(model)\n\n    # different rank will generate different input_ids following fsdp\n    input_ids = torch.randint(low=0, high=config.vocab_size, size=(batch_size, seqlen), device=get_device_name())\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_left_padding=0, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.8\n    )\n    position_ids = compute_position_id_with_mask(\n        attention_mask\n    )  # TODO(sgm): we can construct the position_ids_rmpad here\n\n    model_inputs = {\n        \"input_ids\": input_ids.to(get_device_name()),\n        \"attention_mask\": attention_mask.to(get_device_name()),\n        \"position_ids\": position_ids.int().to(get_device_name()),\n    }\n\n    model_inputs = DataProto.from_dict(model_inputs)\n\n    # 1. perform ulysses forward\n    with sharding_manager:\n        model_inputs = sharding_manager.preprocess_data(model_inputs)\n        input_ids = model_inputs.batch[\"input_ids\"]\n        attention_mask = model_inputs.batch[\"attention_mask\"]\n        position_ids = model_inputs.batch[\"position_ids\"]\n        input_ids_rmpad, indices, *_ = unpad_input(\n            input_ids.unsqueeze(-1), attention_mask\n        )  # input_ids_rmpad (total_nnz, ...)\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n        # unpad the position_ids to align the rotary\n        position_ids_rmpad = index_first_axis(\n            rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n        ).transpose(0, 1)\n\n        # slice input tensor for ulysses\n        # input_ids are padded and sliced\n        # postition_ids are only padded but not sliced\n        input_ids_rmpad_sliced, position_ids_rmpad_padded, pad_size = ulysses_pad_and_slice_inputs(\n            input_ids_rmpad, position_ids_rmpad, sp_size=get_ulysses_sequence_parallel_world_size()\n        )\n\n        # input with input_ids_rmpad and postition_ids to enable flash attention varlen\n        logits_split_in_seq = model(\n            input_ids_rmpad_sliced, position_ids=position_ids_rmpad_padded, use_cache=False\n        ).logits  # (1, total_nnz/n, vocab_size)\n\n        # all_gather output\n        logits_full = gather_outputs_and_unpad(logits_split_in_seq, gather_dim=1, unpad_dim=1, padding_size=pad_size)\n\n    # 2. perform normal forward\n    set_ulysses_sequence_parallel_group(None)\n    input_ids_full = copy.deepcopy(input_ids_rmpad)\n    position_ids_full = copy.deepcopy(position_ids_rmpad)\n    model_no_sp = copy.deepcopy(model)\n    logits_rmpad_local = model_no_sp(\n        input_ids_full, position_ids=position_ids_full, use_cache=False\n    ).logits  # (1, total_nnz, vocab_size)\n\n    mean_local = logits_rmpad_local.mean()\n    mean_full = logits_full.mean()\n\n    mean_full.backward()\n    mean_local.backward()\n\n    # 3. check the gradients\n    grad = model.model.layers[0].self_attn.q_proj.weight.grad\n    grad_full = model_no_sp.model.layers[0].self_attn.q_proj.weight.grad\n    torch.testing.assert_close(mean_local, mean_full, rtol=1e-2, atol=3e-5)\n    # The check should be less strict because the gradient is not an averaged value.\n    torch.testing.assert_close(grad, grad_full, rtol=1e-2, atol=1e-3)\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-svv\"])\n"
  },
  {
    "path": "tests/single_controller/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/single_controller/base/test_decorator.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nimport verl.single_controller.base.decorator as decorator_module\nfrom verl.single_controller.base.decorator import (\n    DISPATCH_MODE_FN_REGISTRY,\n    Dispatch,\n    _check_dispatch_mode,\n    get_predefined_dispatch_fn,\n    register_dispatch_mode,\n    update_dispatch_mode,\n)\n\n\n@pytest.fixture\ndef reset_dispatch_registry():\n    # Store original state\n    original_registry = DISPATCH_MODE_FN_REGISTRY.copy()\n    yield\n    # Reset registry after test\n    decorator_module.DISPATCH_MODE_FN_REGISTRY.clear()\n    decorator_module.DISPATCH_MODE_FN_REGISTRY.update(original_registry)\n\n\ndef test_register_new_dispatch_mode(reset_dispatch_registry):\n    # Test registration\n    def dummy_dispatch(worker_group, *args, **kwargs):\n        return args, kwargs\n\n    def dummy_collect(worker_group, output):\n        return output\n\n    register_dispatch_mode(\"TEST_MODE\", dummy_dispatch, dummy_collect)\n\n    # Verify enum extension\n    _check_dispatch_mode(Dispatch.TEST_MODE)\n\n    # Verify registry update\n    assert get_predefined_dispatch_fn(Dispatch.TEST_MODE) == {\n        \"dispatch_fn\": dummy_dispatch,\n        \"collect_fn\": dummy_collect,\n    }\n    # Clean up\n    Dispatch.remove(\"TEST_MODE\")\n\n\ndef test_update_existing_dispatch_mode(reset_dispatch_registry):\n    # Store original implementation\n    original_mode = Dispatch.ONE_TO_ALL\n\n    # New implementations\n    def new_dispatch(worker_group, *args, **kwargs):\n        return args, kwargs\n\n    def new_collect(worker_group, output):\n        return output\n\n    # Test update=\n    update_dispatch_mode(original_mode, new_dispatch, new_collect)\n\n    # Verify update\n    assert get_predefined_dispatch_fn(original_mode)[\"dispatch_fn\"] == new_dispatch\n    assert get_predefined_dispatch_fn(original_mode)[\"collect_fn\"] == new_collect\n"
  },
  {
    "path": "tests/single_controller/check_worker_alive/main.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport sys\nimport time\n\nimport ray\n\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n\n@ray.remote\nclass TestActor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    def foo(self, wait_time):\n        time.sleep(wait_time)\n        sys.exit(1)\n\n\nif __name__ == \"__main__\":\n    wait_time = int(os.getenv(\"WAIT_TIME\", \"10\"))\n\n    ray.init()\n\n    # test single-node-no-partition\n    print(\"test single-node-no-partition\")\n    resource_pool = RayResourcePool([2], use_gpu=False)\n    class_with_args = RayClassWithInitArgs(cls=TestActor)\n\n    print(\"create worker group\")\n    wg = RayWorkerGroup(resource_pool, class_with_args, name_prefix=\"test\")\n\n    wg.start_worker_aliveness_check(1)\n    time.sleep(1)\n\n    print(time.time(), \"start foo\")\n\n    _ = wg.foo(wait_time)\n    print(\"foo started\")\n\n    print(\n        time.time(),\n        f\"wait 6x wait time {wait_time * 6} to let signal returned to process but still not exceed process wait time\",\n    )\n    time.sleep(wait_time * 6)\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/detached_worker/README.md",
    "content": "# Detached Worker\n## How to run (Only on a single node)\n- Start a local ray cluster: \n```bash\nray start --head --port=6379\n```\n- Run the server\n```bash\npython3 server.py\n```\n- On another terminal, Run the client\n```bash\npython3 client.py\n```\n"
  },
  {
    "path": "tests/single_controller/detached_worker/client.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nIn client, we can get the server handler and send RPC request\n\"\"\"\n\nimport ray\nimport torch\nfrom server import Trainer\nfrom tensordict import TensorDict\n\nfrom verl import DataProto\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\n\n\ndef compute_position_id_with_mask(mask):\n    return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None)\n\n\nif __name__ == \"__main__\":\n    ray.init(address=\"auto\", namespace=\"verl\")\n    # get the worker group using names\n    worker_names = [\"trainerTrainer_0:0\", \"trainerTrainer_0:1\"]\n    cls_with_init_args = RayClassWithInitArgs(cls=Trainer)\n    worker_group = RayWorkerGroup.from_detached(worker_names=worker_names, ray_cls_with_init=cls_with_init_args)\n\n    batch_size = 16\n    sequence_length = 1024\n\n    # give Trainer some data to train\n    input_ids = torch.randint(low=0, high=256, size=(batch_size, sequence_length), dtype=torch.int64, device=\"cuda\")\n    attention_mask = torch.ones_like(input_ids)\n    position_ids = compute_position_id_with_mask(attention_mask)\n\n    data = DataProto(\n        batch=TensorDict(\n            {\"input_ids\": input_ids, \"attention_mask\": attention_mask, \"position_ids\": position_ids},\n            batch_size=batch_size,\n        ),\n        meta_info={},\n    )\n\n    output = worker_group.train_model(data)\n\n    print(output)\n"
  },
  {
    "path": "tests/single_controller/detached_worker/run.sh",
    "content": "#!/bin/bash\nray start --head --port=6379\npython3 server.py\npython3 client.py\nray stop --force"
  },
  {
    "path": "tests/single_controller/detached_worker/server.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nServer starts a Trainer. Client sends data to the server to train.\n\"\"\"\n\nimport os\n\nos.environ[\"MEGATRON_USE_CUDA_TIMER\"] = \"0\"\nos.environ[\"MEGATRON_START_PROCESS_TIMER\"] = \"False\"\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\n\nimport ray\nimport torch\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core import tensor_parallel\nfrom megatron.core.models.gpt.gpt_model import ModelType\nfrom omegaconf import OmegaConf\nfrom tensordict import TensorDict\nfrom torch import nn\nfrom transformers import LlamaConfig\n\nfrom verl import DataProto\nfrom verl.models.llama.megatron import ParallelLlamaForCausalLMRmPadPP\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.megatron.optimizer import get_megatron_optimizer, init_megatron_optim_config\nfrom verl.utils.megatron_utils import get_model, mcore_model_parallel_config\n\n\n@ray.remote\nclass Trainer(Worker):\n    def __init__(self):\n        super().__init__()\n\n        if not torch.distributed.is_initialized():\n            rank = int(os.environ[\"LOCAL_RANK\"])\n            torch.distributed.init_process_group(backend=\"nccl\")\n            torch.cuda.set_device(rank)\n\n            mpu.initialize_model_parallel(\n                tensor_model_parallel_size=2,\n                pipeline_model_parallel_size=1,\n                virtual_pipeline_model_parallel_size=None,\n                use_sharp=False,\n                context_parallel_size=1,\n                expert_model_parallel_size=1,\n                nccl_communicator_config_path=None,\n            )\n            tensor_parallel.model_parallel_cuda_manual_seed(10)\n\n            is_collect = (\n                mpu.get_tensor_model_parallel_rank() == 0\n                and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1\n                and mpu.get_context_parallel_rank() == 0\n            )\n            self._register_dispatch_collect_info(\n                mesh_name=\"train\", dp_rank=mpu.get_data_parallel_rank(), is_collect=is_collect\n            )\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        actor_model_config = LlamaConfig(\n            vocab_size=256,\n            hidden_size=2048,\n            intermediate_size=5504,\n            num_hidden_layers=24,\n            num_attention_heads=16,\n            num_key_value_heads=16,\n        )\n\n        megatron_config = mcore_model_parallel_config(sequence_parallel=True, params_dtype=torch.bfloat16)\n        self.megatron_config = megatron_config\n\n        def megatron_actor_model_provider(pre_process, post_process):\n            # vpp is not supported yet because it will hang for some reason. Need debugging\n            # this_megatron_config = copy.deepcopy(megatron_config)\n            # this_megatron_config.virtual_pipeline_model_parallel_rank = vpp_rank\n            parallel_model = ParallelLlamaForCausalLMRmPadPP(\n                config=actor_model_config,\n                megatron_config=megatron_config,\n                pre_process=pre_process,\n                post_process=post_process,\n            )\n            parallel_model.cuda()\n            return parallel_model\n\n        actor_module = get_model(\n            model_provider_func=megatron_actor_model_provider,\n            model_type=ModelType.encoder_or_decoder,\n            wrap_with_ddp=True,\n        )\n        actor_module = nn.ModuleList(actor_module)\n\n        optim_config = OmegaConf.create({\"lr\": 1e-6, \"clip_grad\": 1.0})\n\n        optim_config = init_megatron_optim_config(optim_config)\n        self.optimizer_config = optim_config\n        actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config)\n\n        self.model = actor_module[0]\n        self.optimizer = actor_optimizer\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"))\n    def train_model(self, data: DataProto) -> DataProto:\n        input_ids = data.batch[\"input_ids\"]\n        attention_mask = data.batch[\"attention_mask\"]\n        position_ids = data.batch[\"position_ids\"]\n\n        self.optimizer.zero_grad()\n        self.model.zero_grad_buffer(\n            zero_buffer=(not self.optimizer_config.use_distributed_optimizer)\n        )  # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm\n        # update for 1 iteration\n        output = self.model(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids).logits\n        output.mean().backward()\n\n        update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step(\n            self.megatron_config, self.megatron_config.timers\n        )\n\n        return DataProto(batch=TensorDict({\"loss\": output.detach()}, batch_size=output.shape[0]))\n\n\nif __name__ == \"__main__\":\n    ray.init(address=\"auto\", namespace=\"verl\")\n\n    resource_pool = RayResourcePool(process_on_nodes=[2], detached=True)\n    cls_with_init_args = RayClassWithInitArgs(cls=Trainer)\n    worker_group = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=cls_with_init_args,\n        name_prefix=\"trainer\",\n        detached=True,\n    )\n\n    worker_group.init_model()\n\n    worker_names = worker_group.worker_names\n    print(worker_names)\n"
  },
  {
    "path": "tests/single_controller/test_auto_padding_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport numpy as np\nimport ray\nimport torch\n\nfrom verl import DataProto\nfrom verl.protocol import DataProtoConfig\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n# or set env var VERL_AUTO_PADDING = \"1\" / \"true\"\nDataProtoConfig.auto_padding = True\n\n\n@ray.remote\nclass Actor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def add(self, data: DataProto):\n        data.batch[\"a\"] += self.rank\n        return data\n\n\ndef test_auto_padding():\n    ray.init(num_cpus=100)\n\n    chunk_size = 4\n    actor_cls = RayClassWithInitArgs(cls=Actor)\n    resource_pool = RayResourcePool(process_on_nodes=[chunk_size], use_gpu=False)\n    actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls)\n\n    # test locally first\n    for test_size in range(4, 20):\n        local_data = DataProto.from_dict({\"a\": torch.zeros(test_size)}, {\"na\": np.zeros(test_size, dtype=object)})\n        # print(f\"before padding, local_data = {local_data}\")\n        padding_size = (chunk_size - (test_size % chunk_size)) if (test_size % chunk_size > 0) else 0\n        local_data.padding(padding_size)\n        # print(f\"after padding, local_data = {local_data}\")\n        assert len(local_data) == len(local_data) + len(local_data) % chunk_size, (\n            f\"expecting padded length to be {len(local_data) + len(local_data) % chunk_size}, but got {len(local_data)}\"\n        )\n        chunked = local_data.chunk(chunk_size)\n        assert len(chunked) == chunk_size, f\"during test_size = {test_size}, expecting {chunk_size}, got {chunked}\"\n        for dp in chunked:\n            assert len(dp) == test_size // chunk_size + bool(test_size % chunk_size), (\n                f\"test size = {test_size}, expecting dp to be length of \"\n                f\"{test_size // chunk_size + bool(test_size % chunk_size)}, but got {len(dp)}: {dp} {chunked}\"\n            )\n\n    # test with RayWorkerGroup method decorated as dispatch_mode=Dispatch.DP_COMPUTE_PROTO\n    data = DataProto.from_dict({\"a\": torch.zeros(10)}, {\"na\": np.array([str(i) for i in range(10)], dtype=object)})\n    output = actor_wg.add(data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 10, \"Failed in args split and padding.\"\n\n    data = DataProto.from_dict({\"a\": torch.zeros(10)}, {\"na\": np.array([str(i) for i in range(10)], dtype=object)})\n    output = actor_wg.add(data=data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 10, \"Failed in kwargs split and padding.\"\n\n    data = DataProto.from_dict({\"a\": torch.zeros(1)}, {\"na\": np.array([str(i) for i in range(1)], dtype=object)})\n    output = actor_wg.add(data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 1, \"Failed in args split and padding.\"\n\n    data = DataProto.from_dict({\"a\": torch.zeros(1)}, {\"na\": np.array([str(i) for i in range(1)], dtype=object)})\n    output = actor_wg.add(data=data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 1, \"Failed in kwargs split and padding.\"\n\n    data = DataProto.from_dict({\"a\": torch.zeros(8)}, {\"na\": np.array([str(i) for i in range(8)], dtype=object)})\n    output = actor_wg.add(data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 8, \"Failed in args split and padding.\"\n\n    data = DataProto.from_dict({\"a\": torch.zeros(8)}, {\"na\": np.array([str(i) for i in range(8)], dtype=object)})\n    output = actor_wg.add(data=data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 8, \"Failed in kwargs split and padding.\"\n\n    # test data proto specific config\n    DataProtoConfig.auto_padding = False\n\n    data = DataProto.from_dict(\n        {\"a\": torch.zeros(10)}, {\"na\": np.array([str(i) for i in range(10)], dtype=object)}, auto_padding=True\n    )\n    output = actor_wg.add(data)\n    print(output.batch[\"a\"])\n    assert len(output) == 10, \"Failed in args split and padding.\"\n\n    data = DataProto.from_dict(\n        {\"a\": torch.zeros(10)}, {\"na\": np.array([str(i) for i in range(10)], dtype=object)}, auto_padding=True\n    )\n    output = actor_wg.add(data=data)\n    print(output.batch[\"a\"])\n    assert len(output) == 10, \"Failed in kwargs split and padding.\"\n\n    data = DataProto.from_single_dict(\n        {\"a\": torch.zeros(1), \"na\": np.array([str(i) for i in range(1)], dtype=object)}, auto_padding=True\n    )\n    output = actor_wg.add(data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 1, \"Failed in args split and padding.\"\n\n    data = DataProto.from_single_dict(\n        {\"a\": torch.zeros(1), \"na\": np.array([str(i) for i in range(1)], dtype=object)}, auto_padding=True\n    )\n    output = actor_wg.add(data=data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 1, \"Failed in kwargs split and padding.\"\n\n    data = DataProto.from_single_dict({\"a\": torch.zeros(8), \"na\": np.array([str(i) for i in range(8)], dtype=object)})\n    output = actor_wg.add(data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 8, \"Failed in args split and padding.\"\n\n    data = DataProto.from_single_dict({\"a\": torch.zeros(8), \"na\": np.array([str(i) for i in range(8)], dtype=object)})\n    output = actor_wg.add(data=data)\n\n    print(output.batch[\"a\"])\n    assert len(output) == 8, \"Failed in kwargs split and padding.\"\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_auto_padding()\n"
  },
  {
    "path": "tests/single_controller/test_colocated_workers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ray\n\nfrom verl import DataProto\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray.base import (\n    RayClassWithInitArgs,\n    RayResourcePool,\n    RayWorkerGroup,\n    create_colocated_worker_cls,\n)\nfrom verl.utils.device import get_device_name\n\n\n@ray.remote\nclass Actor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def add(self, data: DataProto):\n        data.batch[\"a\"] += self.rank\n        return data\n\n\n@ray.remote\nclass Critic(Worker):\n    def __init__(self, config) -> None:\n        super().__init__()\n        self.config = config\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    async def sub(self, data: DataProto):\n        data.batch[\"a\"] -= self.config[\"b\"]\n        return data\n\n\ndef test_colocated_workers():\n    ray.init()\n\n    import torch\n\n    data = DataProto.from_dict({\"a\": torch.zeros(10)})\n    # create separate workers on the same resource pool\n    actor_cls = RayClassWithInitArgs(cls=Actor)\n    critic_cls = RayClassWithInitArgs(cls=Critic, config={\"b\": 10})\n    resource_pool = RayResourcePool(process_on_nodes=[2])\n\n    actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls, device_name=get_device_name())\n    critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls, device_name=get_device_name())\n\n    expected_actor_output = actor_wg.add(data)\n    expected_critic_output = critic_wg.sub(data)\n\n    # create colocated workers\n    cls_dict = {\"actor\": actor_cls, \"critic\": critic_cls}\n    ray_cls_with_init = create_colocated_worker_cls(cls_dict)\n    wg_dict = RayWorkerGroup(\n        resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name()\n    )\n    spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys())\n\n    colocated_actor_wg = spawn_wg[\"actor\"]\n    colocated_critic_wg = spawn_wg[\"critic\"]\n\n    actor_output = colocated_actor_wg.add(data)\n    critic_output = colocated_critic_wg.sub(data)\n\n    torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0)\n    torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0)\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_colocated_workers_fused.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ray\n\nfrom verl import DataProto\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray.base import (\n    RayClassWithInitArgs,\n    RayResourcePool,\n    RayWorkerGroup,\n    create_colocated_worker_cls_fused,\n)\nfrom verl.utils.device import get_device_name\n\n\n@ray.remote\nclass Actor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def add(self, data: DataProto):\n        data.batch[\"a\"] += self.rank\n        return data\n\n\n@ray.remote\nclass Critic(Worker):\n    def __init__(self, config) -> None:\n        super().__init__()\n        self.config = config\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def sub(self, data: DataProto):\n        data.batch[\"a\"] -= self.config[\"b\"]\n        return data\n\n\ndef test_colocated_workers_fused():\n    ray.init()\n\n    import torch\n\n    data = DataProto.from_dict({\"a\": torch.zeros(10)})\n    # create separate workers on the same resource pool\n    actor_cls = RayClassWithInitArgs(cls=Actor)\n    critic_cls = RayClassWithInitArgs(cls=Critic, config={\"b\": 10})\n    resource_pool = RayResourcePool(process_on_nodes=[2])\n\n    actor_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=actor_cls, device_name=get_device_name())\n    critic_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=critic_cls, device_name=get_device_name())\n\n    expected_actor_output = actor_wg.add(data)\n    expected_critic_output = critic_wg.sub(data)\n\n    # create colocated workers\n    cls_dict = {\"actor\": actor_cls, \"critic\": critic_cls}\n    ray_cls_with_init = create_colocated_worker_cls_fused(cls_dict)\n    wg_dict = RayWorkerGroup(\n        resource_pool=resource_pool, ray_cls_with_init=ray_cls_with_init, device_name=get_device_name()\n    )\n    spawn_wg = wg_dict.spawn(prefix_set=cls_dict.keys())\n\n    colocated_actor_wg = spawn_wg[\"actor\"]\n    colocated_critic_wg = spawn_wg[\"critic\"]\n\n    actor_output = colocated_actor_wg.add(data)\n    critic_output = colocated_critic_wg.sub(data)\n\n    torch.testing.assert_close(expected_actor_output.batch, actor_output.batch, atol=0, rtol=0)\n    torch.testing.assert_close(expected_critic_output.batch, critic_output.batch, atol=0, rtol=0)\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_data_transfer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nIn this test, we instantiate a data parallel worker with 8 GPUs\n\"\"\"\n\nimport ray\nimport tensordict\nimport torch\nfrom codetiming import Timer\nfrom packaging import version\nfrom torch import distributed as dist\n\nfrom verl import DataProto\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.device import get_device_name\nfrom verl.utils.ray_utils import parallel_put\n\n\n@ray.remote\nclass DummyWorker(Worker):\n    def __init__(self):\n        super().__init__()\n        dist.init_process_group()\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False)\n    def do_nothing(self, data):\n        for key in data.batch.keys():\n            data.batch[key] += 1\n        if version.parse(tensordict.__version__) >= version.parse(\"0.5.0\"):\n            data.batch = data.batch.consolidate()\n        return data\n\n\ndef test_data_transfer():\n    ray.init()\n    # construct resource pool\n    resource_pool = RayResourcePool([8])\n    cls_with_init = RayClassWithInitArgs(cls=DummyWorker)\n    # construct worker group\n    wg = RayWorkerGroup(resource_pool, cls_with_init, device_name=get_device_name())\n\n    # this is real dataset size\n    batch_size = 4096\n    seqlen = 32768\n\n    data_dict = {}\n\n    for i in range(2):\n        data_dict[str(i)] = torch.randint(0, 10000, (batch_size, seqlen))\n\n    data = DataProto.from_dict(tensors=data_dict)\n\n    print(data)\n\n    # we manually split data here and send to each worker\n    data_list = data.chunk(wg.world_size)\n\n    for i in range(wg.world_size):\n        # consolidate is necessary\n        if version.parse(tensordict.__version__) >= version.parse(\"0.5.0\"):\n            data_list[i].batch = data_list[i].batch.consolidate()\n\n    with Timer(name=\"ray.pickle\", initial_text=True):\n        for i in range(wg.world_size):\n            ray.cloudpickle.pickle.dumps(data_list[i])\n\n    with Timer(name=\"raw.pickle\", initial_text=True):\n        import pickle\n\n        for i in range(wg.world_size):\n            pickle.dumps(data_list[i])\n\n    # we put in advance\n    with Timer(name=\"put\", initial_text=True):\n        # takes around 40 seconds\n        data_list_ref = parallel_put(data_list)\n        # for i in range(wg.world_size):\n        #     data_list[i] = ray.put(data_list[i])\n\n    with Timer(name=\"launch\", initial_text=True):\n        output_ref = wg.do_nothing(data_list_ref)\n\n    with Timer(name=\"get\", initial_text=True):\n        # takes around 40 seconds\n        output_lst = ray.get(output_ref)\n\n    for input_data, output_data in zip(data_list, output_lst, strict=True):\n        for key in input_data.batch.keys():\n            assert torch.all(torch.eq(input_data.batch[key] + 1, output_data.batch[key])), (\n                input_data.batch[key],\n                output_data.batch[key],\n                key,\n            )\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_decorator_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport time\n\nimport pytest\nimport ray\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl.protocol import DataProto, DataProtoFuture\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils import tensordict_utils as tu\n\n\n# Pytest fixture for Ray setup/teardown\n@pytest.fixture\ndef ray_init_shutdown():\n    ray.init(num_cpus=100)\n    yield\n    ray.shutdown()\n\n\n# Define a simple worker for testing\n@ray.remote\nclass DecoratorTestWorker(Worker):\n    def __init__(self, initial_value=0):\n        super().__init__()\n        self.value = initial_value\n        # Simulate some setup if needed\n        time.sleep(0.1)  # Ensure worker init completes\n\n        self._register_dispatch_collect_info(mesh_name=\"train\", dp_rank=self.rank, is_collect=True)\n\n    # Test method for synchronous DP compute (default behavior)\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def dp_compute(self, data: DataProto) -> DataProto:\n        time.sleep(0.1)  # Simulate work\n        rank_value = torch.tensor(self.rank, device=data.batch[\"input\"].device, dtype=data.batch[\"input\"].dtype)\n        data.batch[\"output\"] = data.batch[\"input\"] + self.value + rank_value\n        return data\n\n    # Test async def method with DP compute (default behavior)\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO, blocking=False)\n    async def async_dp_compute(self, data: DataProto) -> DataProto:\n        # Simulate async work\n        await asyncio.sleep(0.1)  # Simulate async work\n        rank_value = torch.tensor(self.rank, device=data.batch[\"input\"].device, dtype=data.batch[\"input\"].dtype)\n        data.batch[\"output_async\"] = data.batch[\"input\"] * 2 + self.value + rank_value\n        return data\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"), blocking=False)\n    def dp_compute_td(self, data: TensorDict) -> TensorDict:\n        # note that we have to call contiguous so that we can modify data in plac\n        data = tu.contiguous(data)\n        rank_value = torch.tensor(self.rank, device=data[\"input\"].device, dtype=data[\"input\"].dtype)\n        data[\"output\"] = data[\"input\"] + self.value + rank_value\n        position_ids = data.pop(\"position_ids\")\n        position_ids._ragged_idx = 2\n\n        for i, position_id in enumerate(position_ids.unbind(dim=0)):\n            assert (position_id == torch.arange(4 + rank_value * 2 + i).expand(position_id.shape)).all()\n\n        return data\n\n\n# Test function for synchronous DP compute\ndef test_decorator_dp_compute(ray_init_shutdown):\n    \"\"\"\n    Tests the default behavior of a synchronous decorated method with DP_COMPUTE_PROTO.\n    Verifies the result correctness.\n    \"\"\"\n    num_workers = 2\n    resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1)  # Use CPU for simplicity\n    cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=10)\n    worker_group = RayWorkerGroup(\n        resource_pool, cls_with_args, name_prefix=f\"decorator_test_sync_dp_{int(time.time())}\"\n    )\n\n    # Prepare input data (size 4, for 2 workers)\n    input_tensor = torch.arange(4, dtype=torch.float32)\n    data = DataProto(batch=TensorDict({\"input\": input_tensor}, batch_size=[4]))\n\n    # Call the decorated method\n    output = worker_group.dp_compute(data)\n\n    # Assert the result correctness\n    assert isinstance(output, DataProto), \"Expected DataProto result\"\n    assert \"output\" in output.batch.keys()\n    assert len(output) == len(data), \"Output length should match input length\"\n\n    # Expected output calculation for DP_COMPUTE_PROTO with 2 workers\n    # Worker 0 gets data[0:2], Worker 1 gets data[2:4]\n    # Worker 0 adds initial_value(10) + rank(0) = 10\n    # Worker 1 adds initial_value(10) + rank(1) = 11\n    expected_output_part1 = torch.tensor([0, 1], dtype=torch.float32) + 10 + 0\n    expected_output_part2 = torch.tensor([2, 3], dtype=torch.float32) + 10 + 1\n    expected_output = torch.cat([expected_output_part1, expected_output_part2])\n\n    torch.testing.assert_close(output.batch[\"output\"], expected_output, msg=\"Sync DP compute output data mismatch\")\n\n\n# Test function for async def method with DP compute\ndef test_decorator_async_function(ray_init_shutdown):\n    \"\"\"\n    Tests the decorator with an `async def` method using DP_COMPUTE_PROTO.\n    Verifies that the call returns a future and the result is correct after .get().\n    \"\"\"\n    num_workers = 2\n    resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1)\n    cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=5)\n    worker_group = RayWorkerGroup(\n        resource_pool, cls_with_args, name_prefix=f\"decorator_test_async_dp_{int(time.time())}\"\n    )\n\n    # Prepare input data (size 4, for 2 workers)\n    input_tensor = torch.arange(4, dtype=torch.float32)\n    data = DataProto(batch=TensorDict({\"input\": input_tensor}, batch_size=[4]))\n\n    # Call the async decorated method - this should return a future\n    future_output: DataProtoFuture = worker_group.async_dp_compute(data)\n\n    # Assert that the call returned a future\n    assert isinstance(future_output, DataProtoFuture), \"Expected DataProtoFuture for async def call\"\n\n    # Get the result (this should block)\n    result_data = future_output.get()\n\n    # Assert the result correctness\n    assert isinstance(result_data, DataProto)\n    assert \"output_async\" in result_data.batch.keys()\n    assert len(result_data) == len(data), \"Output length should match input length\"\n\n    # Expected output calculation for DP_COMPUTE_PROTO with 2 workers\n    # Worker 0 gets data[0:2], Worker 1 gets data[2:4]\n    # Worker 0 calculates: input * 2 + initial_value(5) + rank(0)\n    # Worker 1 calculates: input * 2 + initial_value(5) + rank(1)\n    expected_output_part1 = (torch.tensor([0, 1], dtype=torch.float32) * 2) + 5 + 0\n    expected_output_part2 = (torch.tensor([2, 3], dtype=torch.float32) * 2) + 5 + 1\n    expected_output = torch.cat([expected_output_part1, expected_output_part2])\n\n    torch.testing.assert_close(\n        result_data.batch[\"output_async\"], expected_output, msg=\"Async DP compute output data mismatch\"\n    )\n\n\ndef test_decorator_dp_compute_td(ray_init_shutdown):\n    num_workers = 2\n    resource_pool = RayResourcePool([num_workers], use_gpu=False, max_colocate_count=1)  # Use CPU for simplicity\n    cls_with_args = RayClassWithInitArgs(cls=DecoratorTestWorker, initial_value=10)\n    worker_group = RayWorkerGroup(\n        resource_pool, cls_with_args, name_prefix=f\"decorator_test_sync_dp_{int(time.time())}\"\n    )\n\n    # Prepare input data (size 4, for 2 workers)\n    input_tensor = torch.arange(4, dtype=torch.float32)\n    position_ids = torch.nested.as_nested_tensor(\n        [\n            torch.arange(4).expand(4, 4).contiguous(),\n            torch.arange(5).expand(4, 5).contiguous(),\n            torch.arange(6).expand(4, 6).contiguous(),\n            torch.arange(7).expand(4, 7).contiguous(),\n        ],\n        layout=torch.jagged,\n    )\n    data = TensorDict({\"input\": input_tensor, \"position_ids\": position_ids}, batch_size=[4])\n\n    # Call the decorated method\n    output = worker_group.dp_compute_td(data)\n\n    output = output.get()\n\n    # Assert the result correctness\n    assert isinstance(output, TensorDict), \"Expected DataProto result\"\n    assert \"output\" in output.keys()\n    assert len(output) == len(data), \"Output length should match input length\"\n\n    # Expected output calculation for DP_COMPUTE_PROTO with 2 workers\n    # Worker 0 gets data[0:2], Worker 1 gets data[2:4]\n    # Worker 0 adds initial_value(10) + rank(0) = 10\n    # Worker 1 adds initial_value(10) + rank(1) = 11\n    expected_output_part1 = torch.tensor([0, 1], dtype=torch.float32) + 10 + 0\n    expected_output_part2 = torch.tensor([2, 3], dtype=torch.float32) + 10 + 1\n    expected_output = torch.cat([expected_output_part1, expected_output_part2])\n\n    torch.testing.assert_close(output[\"output\"], expected_output, msg=\"Sync DP compute output data mismatch\")\n"
  },
  {
    "path": "tests/single_controller/test_device_mesh_register.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport numpy as np\nimport ray\nimport torch\nfrom tensordict import TensorDict\n\nimport verl.utils.tensordict_utils as tu\nfrom verl import DataProto\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils.device import get_device_name, get_nccl_backend\n\n\n@ray.remote\nclass TestActor(Worker):\n    def __init__(self):\n        super().__init__()\n\n        import torch.distributed\n\n        torch.distributed.init_process_group(backend=get_nccl_backend())\n        self.infer_device_mesh = torch.distributed.device_mesh.init_device_mesh(\n            device_type=get_device_name(), mesh_shape=[2, 4], mesh_dim_names=[\"dp\", \"tp\"]\n        )\n        self.train_device_mesh = torch.distributed.device_mesh.init_device_mesh(\n            device_type=get_device_name(), mesh_shape=[2, 2, 2], mesh_dim_names=[\"pp\", \"dp\", \"tp\"]\n        )\n\n        self._register_dispatch_collect_info(\n            \"infer\",\n            dp_rank=self.infer_device_mesh[\"dp\"].get_local_rank(),\n            is_collect=self.infer_device_mesh[\"tp\"].get_local_rank() == 0,\n        )\n        self._register_dispatch_collect_info(\n            \"train\",\n            dp_rank=self.train_device_mesh[\"dp\"].get_local_rank(),\n            is_collect=self.train_device_mesh[\"tp\"].get_local_rank() == 0\n            and self.train_device_mesh[\"pp\"].get_local_rank() == 1,\n        )\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"infer\"))\n    def generate_data_proto(self, data: DataProto):\n        tp_rank = self.infer_device_mesh[\"tp\"].get_local_rank()\n        dp_rank = self.infer_device_mesh[\"dp\"].get_local_rank()\n        data.batch[\"a\"] += (tp_rank + 1) * dp_rank\n        return data\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"infer\"))\n    def generate_tensordict(self, data: TensorDict):\n        tp_rank = self.infer_device_mesh[\"tp\"].get_local_rank()\n        dp_rank = self.infer_device_mesh[\"dp\"].get_local_rank()\n        data[\"a\"] += (tp_rank + 1) * dp_rank\n        return data\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"))\n    def train_data_proto(self, data: DataProto):\n        tp_rank = self.train_device_mesh[\"tp\"].get_local_rank()\n        dp_rank = self.train_device_mesh[\"dp\"].get_local_rank()\n        pp_rank = self.train_device_mesh[\"pp\"].get_local_rank()\n        data.batch[\"a\"] += (tp_rank + 1) * (dp_rank + 2) * (pp_rank + 3)\n        # tp rank 0, pp rank 1, dp rank 0, output data added: 8 + 3 = 11\n        # tp rank 0, pp rank 1, dp rank 1, output data added: 12 + 4 = 16\n        return data\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"))\n    def train_tensordict(self, data: TensorDict):\n        tp_rank = self.train_device_mesh[\"tp\"].get_local_rank()\n        dp_rank = self.train_device_mesh[\"dp\"].get_local_rank()\n        pp_rank = self.train_device_mesh[\"pp\"].get_local_rank()\n        data[\"a\"] += (tp_rank + 1) * (dp_rank + 2) * (pp_rank + 3)\n        # tp rank 0, pp rank 1, dp rank 0, output data added: 8 + 3 = 11\n        # tp rank 0, pp rank 1, dp rank 1, output data added: 12 + 4 = 16\n        return data\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"infer\"))\n    def generate_nested_tensor(self, data: TensorDict):\n        tp_rank = self.infer_device_mesh[\"tp\"].get_local_rank()\n        dp_rank = self.infer_device_mesh[\"dp\"].get_local_rank()\n        assert data.shape[0] == 8\n        data[\"input_ids\"] += tp_rank + dp_rank\n\n        print(data)\n        return data\n\n\ndef test_dist_global_info_wg():\n    # create a worker group with size 8\n    # register a infer dist info with tp=4, dp=2\n    # register a train dist info with tp=2, dp=2, pp=2\n    # test the correctness of data dispatch and computation\n    from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n    ray.init()\n\n    ray_cls = RayClassWithInitArgs(TestActor)\n    resource_pool = RayResourcePool(process_on_nodes=[8])\n    wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=ray_cls, device_name=get_device_name())\n\n    infer_input_data_proto = DataProto.from_single_dict(data={\"a\": torch.tensor([1, 2])})\n    infer_output_data_proto = wg.generate_data_proto(infer_input_data_proto)\n\n    assert wg._dispatch_info[\"infer\"] == [0, 0, 0, 0, 1, 1, 1, 1]\n\n    assert torch.all(torch.eq(infer_output_data_proto.batch[\"a\"], torch.tensor([1, 3])))\n\n    infer_input_tensordict = infer_input_data_proto.to_tensordict()\n    infer_output_tensordict = wg.generate_tensordict(infer_input_tensordict)\n    assert torch.all(torch.eq(infer_output_tensordict[\"a\"], torch.tensor([1, 3])))\n\n    train_input_data_proto = DataProto.from_single_dict(data={\"a\": torch.tensor([3, 4])})\n    train_output_data_proto = wg.train_data_proto(train_input_data_proto)\n\n    assert wg._dispatch_info[\"train\"] == [0, 0, 1, 1, 0, 0, 1, 1]\n\n    assert torch.all(torch.eq(train_output_data_proto.batch[\"a\"], torch.tensor([11, 16])))\n\n    train_input_tensordict = train_input_data_proto.to_tensordict()\n    train_output_tensordict = wg.train_tensordict(train_input_tensordict)\n    assert torch.all(torch.eq(train_output_tensordict[\"a\"], torch.tensor([11, 16])))\n\n    # create a batch size of input_ids\n    input_ids = [\n        torch.randint(low=0, high=128, size=(np.random.randint(low=1, high=10, dtype=np.int64),)) for _ in range(16)\n    ]\n    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)\n    data = tu.get_tensordict(tensor_dict={\"input_ids\": input_ids})\n    output = wg.generate_nested_tensor(data)\n\n    input_ids_chunked = list(input_ids.chunk(2))\n\n    print(input_ids_chunked)\n\n    input_ids_chunked[0] += 0\n    input_ids_chunked[1] += 1\n\n    expected = tu.concat_nested_tensors(input_ids_chunked)\n\n    assert torch.all(torch.eq(output[\"input_ids\"].values(), expected.values()))\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_dist_global_info_wg()\n"
  },
  {
    "path": "tests/single_controller/test_driverfunc_to_worker.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport ray\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl import DataProto\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray import RayWorkerGroup\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool\nfrom verl.utils.device import get_device_name\n\nos.environ[\"RAY_DEDUP_LOGS\"] = \"0\"\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\n\n\n@ray.remote\nclass ModelActor(Worker):\n    def __init__(self):\n        pass\n\n\nclass HackSelf:\n    def __init__(self):\n        pass\n\n\ndef get_aux_metrics(self, test_proto):\n    sequence_ids = test_proto.batch[\"sequence_ids\"]\n    decode_count = []\n    for i in range(sequence_ids.size(0)):\n        decode_count.append(len(sequence_ids[i].tolist()))\n    ret_proto = DataProto(\n        batch=TensorDict(\n            {\"sequence_ids\": sequence_ids, \"decode_count\": torch.tensor(decode_count)}, batch_size=sequence_ids.size(0)\n        )\n    )\n    return ret_proto\n\n\ndef test():\n    # construct model\n    ray.init()\n\n    # create 2 workers, each hold a GPU\n    resource_pool = RayResourcePool([2], use_gpu=True, name_prefix=\"a\")\n\n    class_with_args = RayClassWithInitArgs(cls=ModelActor)\n    shard_wg = RayWorkerGroup(resource_pool, class_with_args, device_name=get_device_name())\n\n    test_bs = 8\n    test_proto = DataProto(\n        TensorDict(\n            {\n                \"sequence_ids\": torch.ones([test_bs, 2048], dtype=torch.int64),\n            },\n            batch_size=test_bs,\n        ),\n        meta_info={\"query_length\": 1536},\n    )\n\n    # Sharding among different ranks\n    ret_proto1 = shard_wg.execute_with_func_generator(get_aux_metrics, test_proto)\n\n    # compare execute on driver\n    hs = HackSelf()\n    ret_proto2 = get_aux_metrics(hs, test_proto)\n\n    torch.testing.assert_close(ret_proto1.batch[\"decode_count\"], ret_proto2.batch[\"decode_count\"])\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_fused_workers_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ray\n\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray.base import (\n    RayClassWithInitArgs,\n    RayResourcePool,\n    RayWorkerGroup,\n    create_colocated_worker_raw_cls,\n)\n\n\n@ray.remote\nclass Actor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def add(self, x):\n        x += self.rank\n        return x\n\n\n@ray.remote\nclass Critic(Worker):\n    def __init__(self, val) -> None:\n        super().__init__()\n        self.val = val\n\n    @register(dispatch_mode=Dispatch.ALL_TO_ALL)\n    def sub(self, x):\n        x -= self.val\n        return x\n\n\nactor_cls = RayClassWithInitArgs(cls=Actor)\ncritic_cls = RayClassWithInitArgs(cls=Critic, val=10)\ncls_dict = {\"actor\": actor_cls, \"critic\": critic_cls}\nFusedBaseClass = create_colocated_worker_raw_cls(cls_dict)\n\n\n@ray.remote\nclass HybridWorker(FusedBaseClass):\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def foo(self, x):\n        return self.critic.sub(self.actor.add(x))\n\n\ndef test_fused_workers():\n    ray.init(num_cpus=100)\n\n    # create separate workers on the same resource pool\n    process_on_nodes = [2]\n    resource_pool = RayResourcePool(process_on_nodes=process_on_nodes, use_gpu=False)\n\n    # create colocated workers\n    hybrid_cls_with_init = RayClassWithInitArgs(cls=HybridWorker)\n    hybrid_cls_with_init.fused_worker_used = True\n\n    fused_wg = RayWorkerGroup(resource_pool=resource_pool, ray_cls_with_init=hybrid_cls_with_init)\n    fused_wg.fuse(cls_dict.keys())\n\n    x = fused_wg.actor.add(0.1)\n    print(x)\n    y = fused_wg.critic.sub(x)\n    print(y)\n    z = fused_wg.foo(0.1)\n    print(z)\n    for i, j in zip(y, z, strict=True):\n        assert i == j\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_fused_workers()\n"
  },
  {
    "path": "tests/single_controller/test_get_set_dispatch_collect_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport pytest\n\nfrom verl.single_controller.base import Worker\n\n\ndef test_get_set_dispatch_collect_cpu():\n    os.environ[\"RANK\"] = \"0\"\n    os.environ[\"LOCAL_RANK\"] = \"0\"\n    os.environ[\"WORLD_SIZE\"] = \"2\"\n    os.environ[\"MASTER_ADDR\"] = \"localhost\"\n    os.environ[\"MASTER_PORT\"] = \"12345\"\n\n    ref = Worker()\n    ref._register_dispatch_collect_info(mesh_name=\"actor\", dp_rank=0, is_collect=True)\n\n    actor = Worker()\n    actor._register_dispatch_collect_info(mesh_name=\"actor\", dp_rank=1, is_collect=False)\n\n    actor_rollout_ref = Worker()\n    actor_rollout_ref.set_dispatch_collect(mesh_name=\"ref\", **ref.get_dispatch_collect())\n    actor_rollout_ref.set_dispatch_collect(mesh_name=\"actor\", **actor.get_dispatch_collect())\n\n    assert actor_rollout_ref._query_dispatch_info(\"ref\") == 0\n    assert actor_rollout_ref._query_collect_info(\"ref\")\n    assert actor_rollout_ref._query_dispatch_info(\"actor\") == 1\n    assert not actor_rollout_ref._query_collect_info(\"actor\")\n\n    # test conflict mesh_name\n    actor2 = Worker()\n    actor2._register_dispatch_collect_info(mesh_name=\"actor\", dp_rank=1, is_collect=False)\n    with pytest.raises(AssertionError):\n        actor_rollout_ref.set_dispatch_collect(mesh_name=\"actor\", **actor2.get_dispatch_collect())\n"
  },
  {
    "path": "tests/single_controller/test_high_level_scheduling_api.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport gc\nimport time\n\nimport ray\n\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, merge_resource_pool\nfrom verl.utils.device import get_device_name\n\n\n@ray.remote\nclass TestActor(Worker):\n    # TODO: pass *args and **kwargs is bug prone and not very convincing\n    def __init__(self, cuda_visible_devices=None) -> None:\n        super().__init__(cuda_visible_devices)\n\n    def get_node_id(self):\n        return ray.get_runtime_context().get_node_id()\n\n\ndef test():\n    ray.init()\n\n    # test single-node-no-partition\n    print(\"test single-node-no-partition\")\n    resource_pool = RayResourcePool([8], use_gpu=True)\n\n    class_with_args = RayClassWithInitArgs(cls=TestActor)\n\n    print(\"create actor worker group\")\n    actor_wg = RayWorkerGroup(\n        resource_pool, class_with_args, name_prefix=\"high_level_api_actor\", device_name=get_device_name()\n    )\n    print(\"create critic worker group\")\n    critic_wg = RayWorkerGroup(\n        resource_pool, class_with_args, name_prefix=\"hight_level_api_critic\", device_name=get_device_name()\n    )\n    print(\"create ref worker group\")\n    ref_wg = RayWorkerGroup(\n        resource_pool, class_with_args, name_prefix=\"high_level_api_ref\", device_name=get_device_name()\n    )\n\n    assert actor_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(8)]\n    assert critic_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(8)]\n    assert ref_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(8)]\n\n    del actor_wg\n    del critic_wg\n    del ref_wg\n    gc.collect()  # make sure ray actors are deleted\n\n    [ray.util.remove_placement_group(pg) for pg in resource_pool.get_placement_groups()]\n    print(\"wait 5s to remove placemeng_group\")\n    time.sleep(5)\n    # test single-node-multi-partition\n\n    print(\"test single-node-multi-partition\")\n    rm_resource_pool = RayResourcePool([4], use_gpu=True, name_prefix=\"rm\")\n    ref_resource_pool = RayResourcePool([4], use_gpu=True, name_prefix=\"ref\")\n    total_resource_pool = merge_resource_pool(rm_resource_pool, ref_resource_pool)\n\n    assert rm_resource_pool.world_size == 4\n    assert ref_resource_pool.world_size == 4\n    assert total_resource_pool.world_size == 8\n\n    actor_wg = RayWorkerGroup(\n        total_resource_pool, class_with_args, name_prefix=\"high_level_api_actor\", device_name=get_device_name()\n    )\n    critic_wg = RayWorkerGroup(\n        total_resource_pool, class_with_args, name_prefix=\"high_level_api_critic\", device_name=get_device_name()\n    )\n    ref_wg = RayWorkerGroup(\n        ref_resource_pool, class_with_args, name_prefix=\"high_level_api_ref\", device_name=get_device_name()\n    )\n\n    assert actor_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(8)]\n    assert critic_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(8)]\n    assert ref_wg.execute_all_sync(\"get_cuda_visible_devices\") == [str(i) for i in range(4, 8)]\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_nested_worker.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport ray\n\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.device import get_device_name\n\n\nclass TestActor(Worker):\n    # TODO: pass *args and **kwargs is bug prone and not very convincing\n    def __init__(self, x) -> None:\n        super().__init__()\n        self.a = x\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def get(self):\n        return self.a + self.rank\n\n\nclass TestHighLevelActor(Worker):\n    def __init__(self, x=None) -> None:\n        super().__init__()\n        self.test_actor = TestActor(x=x)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def get(self):\n        return self.test_actor.get()\n\n\ndef test_nested_worker():\n    ray.init(num_cpus=100)\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=True)\n    class_with_args = RayClassWithInitArgs(cls=ray.remote(TestActor), x=2)\n\n    worker_group = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=class_with_args,\n        name_prefix=\"worker_group_basic\",\n        device_name=get_device_name(),\n    )\n\n    output = worker_group.get()\n\n    assert output == [2, 3, 4, 5]\n\n    class_with_args = RayClassWithInitArgs(cls=ray.remote(TestHighLevelActor), x=2)\n    high_level_worker_group = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=class_with_args,\n        name_prefix=\"worker_group_basic_2\",\n        device_name=get_device_name(),\n    )\n\n    output_1 = high_level_worker_group.get()\n\n    assert output_1 == [2, 3, 4, 5]\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_ray_collectives.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTest for using ray collective group.\nSuppose we Actor and Rollout. Actor contains 4 workers and Rollout contains 2 workers. We established a Worker to\nRollout relationship by using collective groups\nActor: rank 0, 1 - Rollout rank 0\nRollout rank 2, 3 - Rollout rank 1\nThen, we initiate 4 p2p comms from actor to rollout\n\"\"\"\n\nimport ray\nimport ray.util.collective as collective\nimport torch\n\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n\n@ray.remote\nclass Actor(Worker):\n    @register(Dispatch.ONE_TO_ALL)\n    def init(self):\n        remote_rank = self.rank // 2\n        self.group_name = f\"A{self.rank}_R{remote_rank}\"\n        collective.init_collective_group(world_size=2, rank=0, backend=\"nccl\", group_name=self.group_name)\n\n    @register(Dispatch.ONE_TO_ALL, blocking=False)\n    def send_tensors(self):\n        tensor = torch.ones(size=(4,), dtype=torch.float32, device=\"cuda\") * self.rank\n        collective.send(tensor=tensor, dst_rank=1, group_name=self.group_name)\n\n\n@ray.remote\nclass Rollout(Worker):\n    @register(Dispatch.ONE_TO_ALL)\n    def init(self):\n        self.remote_first_rank = self.rank * 2\n        self.remote_second_rank = self.remote_first_rank + 1\n        self.first_group_name = f\"A{self.remote_first_rank}_R{self.rank}\"\n        self.second_group_name = f\"A{self.remote_second_rank}_R{self.rank}\"\n\n        collective.init_collective_group(world_size=2, rank=1, backend=\"nccl\", group_name=self.first_group_name)\n        collective.init_collective_group(world_size=2, rank=1, backend=\"nccl\", group_name=self.second_group_name)\n\n    @register(Dispatch.ONE_TO_ALL, blocking=False)\n    def receive_tensors(self):\n        self.tensor1 = torch.randn(size=(4,), dtype=torch.float32, device=\"cuda\")\n        self.tensor2 = torch.randn(size=(4,), dtype=torch.float32, device=\"cuda\")\n\n        collective.recv(self.tensor1, src_rank=0, group_name=self.first_group_name)\n        collective.recv(self.tensor2, src_rank=0, group_name=self.second_group_name)\n\n    @register(Dispatch.ONE_TO_ALL)\n    def get_tensors(self):\n        return {f\"src_{self.remote_first_rank}\": self.tensor1, f\"src_{self.remote_second_rank}\": self.tensor2}\n\n\ndef test_ray_collective_group():\n    ray.init()\n\n    actor_resource_pool = RayResourcePool([4])\n    rollout_resource_pool = RayResourcePool([2])\n\n    actor_cls = RayClassWithInitArgs(cls=Actor)\n    rollout_cls = RayClassWithInitArgs(cls=Rollout)\n\n    actor_wg = RayWorkerGroup(\n        resource_pool=actor_resource_pool, ray_cls_with_init=actor_cls, name_prefix=\"collective_group_actor\"\n    )\n    rollout_wg = RayWorkerGroup(\n        resource_pool=rollout_resource_pool, ray_cls_with_init=rollout_cls, name_prefix=\"collective_group_rollout\"\n    )\n\n    actor_wg.init()\n    rollout_wg.init()\n\n    out1 = actor_wg.send_tensors()\n    out2 = rollout_wg.receive_tensors()\n\n    # block to wait\n    ray.get(out1)\n    ray.get(out2)\n\n    output = rollout_wg.get_tensors()\n\n    rollout_0_output = output[0]\n    rollout_1_output = output[1]\n\n    output = rollout_0_output | rollout_1_output\n\n    print(output)\n\n    for i in range(4):\n        assert torch.sum(output[f\"src_{i}\"]).item() == 4 * i\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_ray_collective_group()\n"
  },
  {
    "path": "tests/single_controller/test_ray_local_envs_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\ne2e test verl.single_controller.ray\n\"\"\"\n\nimport os\n\nimport ray\n\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n\n@ray.remote\nclass TestActor(Worker):\n    def __init__(self) -> None:\n        super().__init__()\n\n    def getenv(self, key):\n        val = os.getenv(key, f\"{key} not set\")\n        return val\n\n\ndef test_basics():\n    ray.init(num_cpus=100)\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=False)\n    class_with_args = RayClassWithInitArgs(cls=TestActor)\n\n    worker_group = RayWorkerGroup(\n        resource_pool=resource_pool, ray_cls_with_init=class_with_args, name_prefix=\"worker_group_basic\"\n    )\n\n    output = worker_group.execute_all_sync(\"getenv\", key=\"RAY_LOCAL_WORLD_SIZE\")\n    assert output == [\"4\", \"4\", \"4\", \"4\"]\n\n    ray.shutdown()\n\n\ndef test_customized_worker_env():\n    ray.init(num_cpus=100)\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=False)\n    class_with_args = RayClassWithInitArgs(cls=TestActor)\n\n    worker_group = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=class_with_args,\n        name_prefix=\"worker_group_customized\",\n        worker_env={\n            \"test_key\": \"test_value\",  # new key will be appended\n        },\n    )\n\n    output = worker_group.execute_all_sync(\"getenv\", key=\"test_key\")\n    assert output == [\"test_value\", \"test_value\", \"test_value\", \"test_value\"]\n\n    try:\n        worker_group = RayWorkerGroup(\n            resource_pool=resource_pool,\n            ray_cls_with_init=class_with_args,\n            name_prefix=\"worker_group_error\",\n            worker_env={\n                \"WORLD_SIZE\": \"100\",  # override system env will result in error\n            },\n        )\n    except ValueError as e:\n        assert \"WORLD_SIZE\" in str(e)\n    else:\n        raise ValueError(\"test failed\")\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_basics()\n    test_customized_worker_env()\n"
  },
  {
    "path": "tests/single_controller/test_ray_utils_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\nimport ray\n\nfrom verl.utils.ray_utils import parallel_put\n\n\n# Initialize Ray for testing if not already done globally\n@pytest.fixture()\ndef init_ray():\n    ray.init(num_cpus=4)\n    yield\n    ray.shutdown()\n\n\ndef test_parallel_put_basic(init_ray):\n    data = [1, \"hello\", {\"a\": 2}, [3, 4]]\n    refs = parallel_put(data)\n    assert len(refs) == len(data)\n    retrieved_data = [ray.get(ref) for ref in refs]\n    assert retrieved_data == data\n\n\ndef test_parallel_put_empty(init_ray):\n    data = []\n    with pytest.raises(AssertionError):\n        _ = parallel_put(data)\n\n\ndef test_parallel_put_workers(init_ray):\n    data = list(range(20))\n    # Test with specific number of workers\n    refs = parallel_put(data, max_workers=4)\n    assert len(refs) == len(data)\n    retrieved_data = [ray.get(ref) for ref in refs]\n    assert retrieved_data == data\n    # Test with default workers (should cap)\n    refs_default = parallel_put(data)\n    assert len(refs_default) == len(data)\n    retrieved_data_default = [ray.get(ref) for ref in refs_default]\n    assert retrieved_data_default == data\n"
  },
  {
    "path": "tests/single_controller/test_rvdz.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ray\n\n\n@ray.remote\nclass TestWorker:\n    def __init__(self, rank, world_size, group_name):\n        self.rank = rank\n        self.world_size = world_size\n        self.group_name = group_name\n        self.communicator = None\n\n    def init(self):\n        from verl.utils.rendezvous.ray_backend import create_nccl_communicator_in_ray\n\n        self.communicator = create_nccl_communicator_in_ray(self.rank, self.world_size, self.group_name)\n\n    def test(self):\n        if self.communicator is None:\n            return None\n        return self.communicator.rank_id()\n\n\ndef test_rvdz():\n    ray.init()\n\n    group_name = \"test_group\"\n    world_size = 2\n\n    workers = [TestWorker.options(num_gpus=1).remote(rank, world_size, group_name) for rank in range(world_size)]\n\n    ray.get([worker.init.remote() for worker in workers])\n\n    ranks = ray.get([worker.test.remote() for worker in workers])\n\n    assert ranks == [0, 1], f\"expecting [0, 1], got {ranks}\"\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_split_resource_pool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport ray\nimport torch\n\nfrom verl import DataProto\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray.base import (\n    RayClassWithInitArgs,\n    RayResourcePool,\n    RayWorkerGroup,\n    split_resource_pool,\n)\nfrom verl.utils.device import get_device_name, get_nccl_backend\n\n\n@ray.remote\nclass Actor(Worker):\n    def __init__(self, worker_id) -> None:\n        super().__init__()\n        self.worker_id = worker_id\n        self.temp_tensor = torch.rand(4096, 4096).to(get_device_name())\n\n        if not torch.distributed.is_initialized():\n            rank = int(os.environ.get(\"RANK\", 0))\n            world_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n            torch.distributed.init_process_group(backend=get_nccl_backend(), world_size=world_size, rank=rank)\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO)\n    def add(self, data: DataProto):\n        data.batch[\"a\"] += self.rank + self.worker_id\n        return data\n\n\ndef test_split_resource_pool_with_split_size():\n    ray.init()\n    # assume we have 2 nodes, with 4 GPUs each\n    global_resource_pool = RayResourcePool(process_on_nodes=[4, 4])\n    global_resource_pool.get_placement_groups(device_name=get_device_name())\n\n    # first 4 gpus for actor_1, last 4 gpus for actor_2\n    actor_1_resource_pool, actor_2_resource_pool = split_resource_pool(resource_pool=global_resource_pool, split_size=4)\n    actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0)\n    actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100)\n    actor_worker_1 = RayWorkerGroup(\n        resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name()\n    )\n    actor_worker_2 = RayWorkerGroup(\n        resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name()\n    )\n    assert actor_worker_1.world_size == 4\n    assert actor_worker_2.world_size == 4\n\n    data = DataProto.from_dict({\"a\": torch.zeros(8)})\n    actor_output_1 = actor_worker_1.add(data)\n    actor_output_2 = actor_worker_2.add(data)\n    assert actor_output_1.batch[\"a\"].tolist() == [0, 0, 1, 1, 2, 2, 3, 3]\n    assert actor_output_2.batch[\"a\"].tolist() == [100, 100, 101, 101, 102, 102, 103, 103]\n\n    ray.shutdown()\n\n\ndef test_split_resource_pool_with_split_size_list():\n    ray.init()\n    # assume we have 4 nodes, with 2 GPUs each\n    global_resource_pool = RayResourcePool(process_on_nodes=[2, 2, 2, 2])\n    global_resource_pool.get_placement_groups(device_name=get_device_name())\n\n    # first 2 gpus for actor_1, last 6 gpus for actor_2\n    actor_1_resource_pool, actor_2_resource_pool = split_resource_pool(\n        resource_pool=global_resource_pool,\n        split_size=[2, 6],\n    )\n    actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0)\n    actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100)\n    actor_worker_1 = RayWorkerGroup(\n        resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name()\n    )\n    actor_worker_2 = RayWorkerGroup(\n        resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name()\n    )\n    assert actor_worker_1.world_size == 2\n    assert actor_worker_2.world_size == 6\n\n    data_1 = DataProto.from_dict({\"a\": torch.zeros(4)})\n    data_2 = DataProto.from_dict({\"a\": torch.zeros(6)})\n    actor_output_1 = actor_worker_1.add(data_1)\n    actor_output_2 = actor_worker_2.add(data_2)\n    print(actor_output_1.batch[\"a\"].tolist())\n    print(actor_output_2.batch[\"a\"].tolist())\n    assert actor_output_1.batch[\"a\"].tolist() == [0, 0, 1, 1]\n    assert actor_output_2.batch[\"a\"].tolist() == [100, 101, 102, 103, 104, 105]\n\n    ray.shutdown()\n\n\ndef test_split_resource_pool_with_split_size_list_cross_nodes():\n    ray.init()\n    # assume we have 4 nodes, with 2 GPUs each\n    global_resource_pool = RayResourcePool(process_on_nodes=[4, 4])\n    global_resource_pool.get_placement_groups(device_name=get_device_name())\n\n    # first 2 gpus for actor_1, last 6 gpus for actor_2\n    actor_1_resource_pool, actor_2_resource_pool = split_resource_pool(\n        resource_pool=global_resource_pool,\n        split_size=[2, 6],\n    )\n    actor_cls_1 = RayClassWithInitArgs(cls=Actor, worker_id=0)\n    actor_cls_2 = RayClassWithInitArgs(cls=Actor, worker_id=100)\n    actor_worker_1 = RayWorkerGroup(\n        resource_pool=actor_1_resource_pool, ray_cls_with_init=actor_cls_1, device_name=get_device_name()\n    )\n    actor_worker_2 = RayWorkerGroup(\n        resource_pool=actor_2_resource_pool, ray_cls_with_init=actor_cls_2, device_name=get_device_name()\n    )\n\n    assert actor_worker_1.world_size == 2\n    assert actor_worker_2.world_size == 6\n\n    data_1 = DataProto.from_dict({\"a\": torch.zeros(4)})\n    data_2 = DataProto.from_dict({\"a\": torch.zeros(6)})\n    actor_output_1 = actor_worker_1.add(data_1)\n    actor_output_2 = actor_worker_2.add(data_2)\n    print(actor_output_1.batch[\"a\"].tolist())\n    print(actor_output_2.batch[\"a\"].tolist())\n    assert actor_output_1.batch[\"a\"].tolist() == [0, 0, 1, 1]\n    assert actor_output_2.batch[\"a\"].tolist() == [100, 101, 102, 103, 104, 105]\n\n    ray.shutdown()\n\n\ndef test_split_resource_pool_with_split_twice():\n    ray.init()\n\n    # assume we have 4 nodes, with 2 GPUs each\n    global_resource_pool = RayResourcePool(process_on_nodes=[2, 2, 2, 2])\n    global_resource_pool.get_placement_groups(device_name=get_device_name())\n\n    # actors with [2, 1, 1, 1, 1, 2] (split twice)\n    rp_1, rp_2, rp_3 = split_resource_pool(\n        resource_pool=global_resource_pool,\n        split_size=[2, 4, 2],\n    )\n    rp_2_1, rp_2_2, rp_2_3, rp_2_4 = split_resource_pool(\n        resource_pool=rp_2,\n        split_size=1,\n    )\n    fp_list = [rp_1, rp_2_1, rp_2_2, rp_2_3, rp_2_4, rp_3]\n    correct_world_size = [2, 1, 1, 1, 1, 2]\n    correct_output = [\n        [0.0, 0.0, 1.0, 1.0],  # 2 worker\n        [100.0, 100.0, 100.0, 100.0],  # 1 worker\n        [200.0, 200.0, 200.0, 200.0],  # 1 worker\n        [300.0, 300.0, 300.0, 300.0],  # 1 worker\n        [400.0, 400.0, 400.0, 400.0],  # 1 worker\n        [500.0, 500.0, 501.0, 501.0],  # 2 worker\n    ]\n    for idx, rp in enumerate(fp_list):\n        actor_cls = RayClassWithInitArgs(cls=Actor, worker_id=idx * 100)\n        actor_worker = RayWorkerGroup(resource_pool=rp, ray_cls_with_init=actor_cls, device_name=get_device_name())\n        data = DataProto.from_dict({\"a\": torch.zeros(4)})\n        actor_output = actor_worker.add(data)\n        assert actor_worker.world_size == correct_world_size[idx]\n        assert actor_output.batch[\"a\"].tolist() == correct_output[idx]\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/single_controller/test_worker_group_basics.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\ne2e test verl.single_controller.ray\n\"\"\"\n\nimport ray\nimport torch\n\nfrom verl.single_controller.base.decorator import Dispatch, Execute, collect_all_to_all, register\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.device import get_device_name\n\n\ndef two_to_all_dispatch_fn(worker_group, *args, **kwargs):\n    \"\"\"\n    Assume the input is a list of 2. Duplicate the input interleaved and pass to each worker.\n    \"\"\"\n    for arg in args:\n        assert len(arg) == 2\n        for i in range(worker_group.world_size - 2):\n            arg.append(arg[i % 2])\n    for k, v in kwargs.items():\n        assert len(v) == 2\n        for i in range(worker_group.world_size - 2):\n            v.append(v[i % 2])\n    return args, kwargs\n\n\ndef get_ray_remote_options() -> str:\n    \"\"\"Function that gets the torch.device based on the current machine.\n    This currently only supports CPU, CUDA, NPU.\n    Returns:\n        device\n    \"\"\"\n    if get_device_name() == \"cuda\":\n        return dict(num_gpus=0.1)\n    elif get_device_name() == \"npu\":\n        return dict(resources={\"NPU\": 0.1})\n    return dict(num_cpus=0.1)\n\n\n@ray.remote\nclass TestActor(Worker):\n    # TODO: pass *args and **kwargs is bug prone and not very convincing\n    def __init__(self, x) -> None:\n        super().__init__()\n        self._x = x\n\n    def foo(self, y):\n        return self._x + y\n\n    @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)\n    def foo_rank_zero(self, x, y):\n        return self._x + y + x\n\n    @register(Dispatch.ONE_TO_ALL, blocking=False)\n    def foo_one_to_all(self, x, y):\n        return self._x + y + x\n\n    @register(Dispatch.ALL_TO_ALL, blocking=False)\n    def foo_all_to_all(self, x, y):\n        return self._x + y + x\n\n    @register(dispatch_mode={\"dispatch_fn\": two_to_all_dispatch_fn, \"collect_fn\": collect_all_to_all})\n    def foo_custom(self, x, y):\n        return self._x + y + x\n\n\n@ray.remote(**get_ray_remote_options())\ndef remote_call_wg(worker_names):\n    class_with_args = RayClassWithInitArgs(cls=TestActor, x=2)\n    worker_group = RayWorkerGroup.from_detached(\n        worker_names=worker_names, ray_cls_with_init=class_with_args, name_prefix=None\n    )\n    print(worker_group.worker_names)\n\n    output_ref = worker_group.foo_custom(x=[1, 2], y=[5, 6])\n    assert output_ref == [8, 10, 8, 10]\n\n    output_ref = worker_group.foo_rank_zero(x=1, y=2)\n    assert output_ref == 5\n\n    return worker_group.worker_names\n\n\ndef add_one(data):\n    data = data.to(get_device_name())\n    data += 1\n    data = data.to(\"cpu\")\n    return data\n\n\ndef test_basics():\n    ray.init(num_cpus=100)\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=True)\n    class_with_args = RayClassWithInitArgs(cls=TestActor, x=2)\n\n    worker_group = RayWorkerGroup(\n        resource_pool=resource_pool,\n        ray_cls_with_init=class_with_args,\n        name_prefix=\"worker_group_basic\",\n        device_name=get_device_name(),\n    )\n\n    print(worker_group.worker_names)\n\n    # this will wait for all the results\n    output = worker_group.execute_all_sync(\"foo\", y=3)\n    assert output == [5, 5, 5, 5]\n\n    # this is a list of object reference. It won't block.\n    output_ref = worker_group.execute_all_async(\"foo\", y=4)\n    print(output_ref)\n\n    assert ray.get(output_ref) == [6, 6, 6, 6]\n\n    output_ref = worker_group.foo_one_to_all(x=1, y=2)\n    assert ray.get(output_ref) == [5, 5, 5, 5]\n\n    output_ref = worker_group.foo_all_to_all(x=[1, 2, 3, 4], y=[5, 6, 7, 8])\n    assert ray.get(output_ref) == [8, 10, 12, 14]\n\n    print(ray.get(remote_call_wg.remote(worker_group.worker_names)))\n\n    output = worker_group.execute_func_rank_zero(add_one, torch.ones(2, 2))\n    torch.testing.assert_close(output, torch.ones(2, 2) + 1)\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    test_basics()\n"
  },
  {
    "path": "tests/single_controller/test_worker_group_torch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nos.environ[\"RAY_DEDUP_LOGS\"] = \"0\"\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\n\nimport ray\nimport torch\nimport torch.distributed\n\nfrom verl.single_controller.base.worker import Worker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\nfrom verl.utils.device import get_device_name\n\n\n@ray.remote\nclass TestAllGatherActor(Worker):\n    def __init__(self, size) -> None:\n        super().__init__()\n        self.size = size\n\n    def init(self):\n        torch.distributed.init_process_group()\n        self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device=get_device_name())\n        self.tensor += self.rank\n\n    def all_gather(self):\n        world_size = self._world_size\n        output = torch.zeros(\n            size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device\n        )\n        torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False)\n        return output\n\n\n@ray.remote\nclass TestAllGatherActorV2(Worker):\n    def __init__(self, size) -> None:\n        super().__init__()\n        self.size = size\n\n        torch.distributed.init_process_group()\n        self.tensor = torch.zeros(size=(self.size,), dtype=torch.int64, device=get_device_name())\n        self.tensor += self.rank\n\n    def all_gather(self):\n        world_size = self._world_size\n        output = torch.zeros(\n            size=(self.tensor.shape[0] * world_size,), dtype=self.tensor.dtype, device=self.tensor.device\n        )\n        torch.distributed.all_gather_into_tensor(output, self.tensor, async_op=False)\n        return output\n\n\ndef test_all_gather_torch():\n    \"\"\"\n    In this test, we instantiate 4 GPUs in a group and test the all_gather\n    \"\"\"\n    ray.init()\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=True)\n    class_with_args = RayClassWithInitArgs(cls=TestAllGatherActor, size=2)\n\n    worker_group = RayWorkerGroup(\n        resource_pool, class_with_args, name_prefix=\"worker_group_torch\", device_name=get_device_name()\n    )\n\n    worker_group.execute_all_sync(\"init\")\n    output = worker_group.execute_all_sync(\"all_gather\")\n    for i in range(1, len(output)):\n        assert torch.all(output[i] == output[0])\n\n    output = output[0].cpu()\n    print(output)\n    assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64))\n\n    ray.shutdown()\n\n\ndef test_all_gather_torch_v2():\n    \"\"\"\n    In this test, we instantiate 4 GPUs in a group and test the all_gather\n    \"\"\"\n    ray.init()\n\n    # create 4 workers, each hold a GPU\n    resource_pool = RayResourcePool([4], use_gpu=True)\n    class_with_args = RayClassWithInitArgs(cls=TestAllGatherActorV2, size=2)\n\n    worker_group = RayWorkerGroup(\n        resource_pool, class_with_args, name_prefix=\"worker_group_torch\", device_name=get_device_name()\n    )\n\n    output = worker_group.execute_all_sync(\"all_gather\")\n    for i in range(1, len(output)):\n        assert torch.all(output[i] == output[0])\n\n    output = output[0].cpu()\n    print(output)\n    assert torch.all(output == torch.tensor([0, 0, 1, 1, 2, 2, 3, 3], dtype=torch.int64))\n\n    ray.shutdown()\n"
  },
  {
    "path": "tests/special_distributed/README.md",
    "content": "This folder is reserved for unit tests (instead of end-to-end tests) that require multiple GPUs.\n"
  },
  {
    "path": "tests/special_distributed/run_all.sh",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n#!/usr/bin/env bash\n\nset -e -x\ntorchrun --nproc-per-node=4 --standalone tests/special_distributed/test_tensor_dict.py\ntorchrun --nproc-per-node=4 --standalone tests/special_distributed/test_torch_functional.py\n"
  },
  {
    "path": "tests/special_distributed/test_fsdp_ckpt.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport shutil\nimport tempfile\n\nimport torch\nimport torch.distributed\nfrom torch.distributed import init_device_mesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision, ShardingStrategy\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config\n\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.device import get_device_name, get_torch_device\nfrom verl.utils.distributed import initialize_global_process_group\nfrom verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2\n\n\ndef create_random_input_ids(batch_size, seq_len, vocab_size):\n    if get_device_name() == \"cuda\":\n        from flash_attn.bert_padding import unpad_input\n    elif get_device_name() == \"npu\":\n        from verl.utils.attention_utils import unpad_input\n    from verl.utils.model import compute_position_id_with_mask, create_random_mask\n\n    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name())\n\n    attention_mask = create_random_mask(\n        input_ids, max_ratio_of_left_padding=0.1, min_ratio_of_valid_token=0.5, max_ratio_of_valid_token=0.7\n    )\n    position_ids = compute_position_id_with_mask(attention_mask)\n\n    input_ids = unpad_input(input_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1)\n    position_ids = unpad_input(position_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1)\n    return input_ids, position_ids\n\n\ndef test_fsdp_ckpt(strategy=\"fsdp\"):\n    assert get_torch_device().device_count() >= 2, \"need at least 2 gpus for test\"\n    local_rank, rank, world_size = initialize_global_process_group()\n    device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=(\"dp\",))\n\n    model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    config = Qwen2Config(num_hidden_layers=1)\n\n    with torch.device(get_device_name()):\n        model = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        model = model.to(device=get_device_name())\n\n    # Wrap model with FSDP\n    if strategy == \"fsdp\":\n        mixed_precision = MixedPrecision(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32\n        )\n\n        model = FSDP(\n            model,\n            use_orig_params=False,\n            device_id=get_torch_device().current_device(),\n            sharding_strategy=ShardingStrategy.FULL_SHARD,\n            mixed_precision=mixed_precision,\n            device_mesh=device_mesh,\n        )\n    else:\n        mp_policy = MixedPrecisionPolicy(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True\n        )\n        fsdp_kwargs = {\n            \"mesh\": device_mesh,\n            \"mp_policy\": mp_policy,\n        }\n        apply_fsdp2(model, fsdp_kwargs, {})\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\n    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)\n\n    # Create checkpoint manager\n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n    checkpoint_manager = FSDPCheckpointManager(\n        model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer\n    )\n\n    # Generate sample input\n    batch_size = 10\n    seq_len = 1024\n    vocab_size = config.vocab_size\n    # First input for initial update\n    input_ids1, position_ids1 = create_random_input_ids(batch_size, seq_len, vocab_size)\n\n    # Second input for verification\n    input_ids2, position_ids2 = create_random_input_ids(batch_size, seq_len, vocab_size)\n\n    # Step 1: Initial update and save checkpoint\n    outputs1 = model(input_ids=input_ids1, position_ids=position_ids1)\n    loss1 = outputs1.logits.mean()\n    loss1.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Save checkpoint after first update\n    temp_dir = tempfile.mkdtemp()\n    checkpoint_path = os.path.join(temp_dir, \"checkpoint\")\n    checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0)\n    saved_state_dict = model.state_dict()\n\n    # Step 2: Second update and forward pass\n    outputs2 = model(input_ids=input_ids2, position_ids=position_ids2)\n    loss2 = outputs2.logits.mean()\n    loss2.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Record logits after second update\n    with torch.no_grad():\n        logits_before_load = model(input_ids=input_ids2, position_ids=position_ids2).logits\n\n    # Step 3: Load checkpoint and repeat second update\n    checkpoint_manager.load_checkpoint(checkpoint_path)\n    loaded_state_dict = model.state_dict()\n    for key in loaded_state_dict:\n        assert key in saved_state_dict, f\"Key {key} not found in saved state dict\"\n        torch.testing.assert_close(loaded_state_dict[key], saved_state_dict[key], atol=0.0, rtol=0.0)\n\n    # Repeat the second update with same input\n    outputs3 = model(input_ids=input_ids2, position_ids=position_ids2)\n    loss3 = outputs3.logits.mean()\n    loss3.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Record logits after loaded checkpoint and update\n    with torch.no_grad():\n        logits_after_load = model(input_ids=input_ids2, position_ids=position_ids2).logits\n\n    # Step 4: Verify outputs match\n    torch.testing.assert_close(logits_before_load, logits_after_load, atol=0.0, rtol=0.0)\n    print(\"Checkpoint save/load test passed!\")\n\n    # Cleanup\n    shutil.rmtree(temp_dir)\n    torch.distributed.barrier()\n    torch.distributed.destroy_process_group()\n\n\nif __name__ == \"__main__\":\n    strategy = os.environ.get(\"STRATEGY\", \"fsdp\")\n    os.environ[\"FLASH_ATTENTION_DETERMINISTIC\"] = \"1\"\n    test_fsdp_ckpt(strategy=strategy)\n"
  },
  {
    "path": "tests/special_distributed/test_mcore_config_converter.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport megatron.core.parallel_state as mpu\nimport torch\nfrom megatron.core.transformer import MLATransformerConfig, TransformerConfig\nfrom transformers import AutoConfig, PretrainedConfig\n\nfrom verl.models.mcore import hf_to_mcore_config\nfrom verl.utils.distributed import destroy_global_process_group, initialize_global_process_group\n\nTEST_MODELS = [\n    \"Qwen/Qwen2.5-7B\",  # Qwen2 dense\n    \"Qwen/Qwen3-8B\",  # Qwen3 dense\n    \"deepseek-ai/deepseek-coder-1.3b-instruct\",  # deepseek dense\n    \"Qwen/Qwen2-57B-A14B\",  # Qwen2 moe\n    \"Qwen/Qwen3-30B-A3B\",  # Qwen3 moe\n    # \"mistralai/Mixtral-8x7B-v0.1\",  # Mixtral # require authentication\n    \"deepseek-ai/DeepSeek-V3-Base\",  # Deepseek V3\n]\n\n\ndef check_config_converter_results(tf_config: TransformerConfig | MLATransformerConfig, hf_config: PretrainedConfig):\n    assert tf_config.num_layers == hf_config.num_hidden_layers, (\n        f\"Number of layers mismatch: {tf_config.num_layers} != {hf_config.num_hidden_layers}\"\n    )\n    assert tf_config.hidden_size == hf_config.hidden_size, (\n        f\"Hidden size mismatch: {tf_config.hidden_size} != {hf_config.hidden_size}\"\n    )\n    assert tf_config.num_attention_heads == hf_config.num_attention_heads, (\n        f\"Number of attention heads mismatch: {tf_config.num_attention_heads} != {hf_config.num_attention_heads}\"\n    )\n    assert tf_config.num_query_groups == hf_config.num_key_value_heads, (\n        f\"Number of query groups mismatch: {tf_config.num_query_groups} != {hf_config.num_key_value_heads}\"\n    )\n    assert tf_config.ffn_hidden_size == hf_config.intermediate_size, (\n        f\"FFN hidden size mismatch: {tf_config.ffn_hidden_size} != {hf_config.intermediate_size}\"\n    )\n    assert tf_config.attention_dropout == hf_config.attention_dropout, (\n        f\"Attention dropout mismatch: {tf_config.attention_dropout} != {hf_config.attention_dropout}\"\n    )\n    assert tf_config.hidden_dropout == getattr(hf_config, \"hidden_dropout\", 0.0), (\n        f\"Hidden dropout mismatch: {tf_config.hidden_dropout} != {getattr(hf_config, 'hidden_dropout', 0.0)}\"\n    )\n    if getattr(hf_config, \"head_dim\", None) is not None:\n        assert tf_config.kv_channels == getattr(hf_config, \"head_dim\", None), (\n            f\"Head dim mismatch: {tf_config.kv_channels} != {getattr(hf_config, 'head_dim', None)}\"\n        )\n    assert tf_config.layernorm_epsilon == hf_config.rms_norm_eps, (\n        f\"Layernorm epsilon mismatch: {tf_config.layernorm_epsilon} != {hf_config.rms_norm_eps}\"\n    )\n\n\ndef modify_hf_config(name: str, hf_config: PretrainedConfig):\n    if name == \"deepseek-ai/DeepSeek-V3-Base\":\n        hf_config.num_nextn_predict_layers = 0\n        hf_config.quantization_config = None\n    return hf_config\n\n\ndef test_mcore_config_converter():\n    \"\"\"\n    Test the conversion of Hugging Face model configurations to MCore configurations.\n    \"\"\"\n    local_rank, rank, world_size = initialize_global_process_group()\n    mpu.initialize_model_parallel(\n        tensor_model_parallel_size=2,\n        pipeline_model_parallel_size=2,\n        virtual_pipeline_model_parallel_size=None,\n        use_sharp=False,\n        context_parallel_size=2,\n        expert_model_parallel_size=1,\n        expert_tensor_parallel_size=None,\n        nccl_communicator_config_path=None,\n    )\n    for model_name in TEST_MODELS:\n        print(f\"testing {model_name}\")\n        hf_config = AutoConfig.from_pretrained(os.path.expanduser(f\"~/models/configs/{model_name}/config.json\"))\n        hf_config = modify_hf_config(model_name, hf_config)\n        tf_config = hf_to_mcore_config(hf_config, torch.bfloat16)\n        check_config_converter_results(tf_config, hf_config)\n\n    destroy_global_process_group()\n\n\nif __name__ == \"__main__\":\n    test_mcore_config_converter()\n"
  },
  {
    "path": "tests/special_distributed/test_tensor_dict.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\n\nimport numpy as np\nimport torch\nimport torch.distributed\n\nfrom verl.protocol import DataProto, all_gather_data_proto\nfrom verl.utils.device import get_device_name\nfrom verl.utils.distributed import initialize_global_process_group\n\n\ndef test_all_gather_data_proto():\n    device_mesh = torch.distributed.device_mesh.init_device_mesh(\n        get_device_name(), mesh_shape=[2, 2], mesh_dim_names=[\"dp\", \"tp\"]\n    )\n\n    global_rank = torch.distributed.get_rank()\n\n    obs = torch.tensor([[1 * global_rank, 2 * global_rank + 1], [3 * global_rank, 4 * global_rank + 1]])\n\n    labels = [\"a\", \"b\"] if global_rank % 2 == 0 else [\"b\", \"a\"]\n    labels = np.array(labels, dtype=object)\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    all_gather_data_proto(data=data, process_group=device_mesh.get_group(\"dp\"))\n\n    if global_rank == 0:\n        expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device=get_device_name())\n        expected_labels = [\"a\", \"b\", \"a\", \"b\"]\n    elif global_rank == 1:\n        expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device=get_device_name())\n        expected_labels = [\"b\", \"a\", \"b\", \"a\"]\n    elif global_rank == 2:\n        expected_obs = torch.tensor([[0, 1], [0, 1], [2, 5], [6, 9]], device=get_device_name())\n        expected_labels = [\"a\", \"b\", \"a\", \"b\"]\n    elif global_rank == 3:\n        expected_obs = torch.tensor([[1, 3], [3, 5], [3, 7], [9, 13]], device=get_device_name())\n        expected_labels = [\"b\", \"a\", \"b\", \"a\"]\n\n    torch.testing.assert_close(data.batch[\"obs\"], expected_obs, atol=0, rtol=0)\n    assert (data.non_tensor_batch[\"labels\"] == expected_labels).all()\n    assert data.meta_info == {\"info\": \"test_info\"}\n\n\ndef test_vocab_parallel_entropy():\n    from megatron.core import parallel_state as mpu\n\n    from verl.utils.megatron.tensor_parallel import vocab_parallel_entropy\n    from verl.utils.profiler import log_gpu_memory_usage\n    from verl.utils.torch_functional import entropy_from_logits\n\n    mpu.initialize_model_parallel(\n        tensor_model_parallel_size=2, pipeline_model_parallel_size=1, virtual_pipeline_model_parallel_size=None\n    )\n\n    batch_size = 2\n    seqlen = 128\n    vocab_size = 155136\n\n    logits = torch.randn(batch_size * seqlen, vocab_size, device=get_device_name(), requires_grad=True)\n    target = torch.randint(\n        low=0, high=vocab_size, size=(batch_size * seqlen,), device=get_device_name(), dtype=torch.int64\n    )\n\n    # broadcast across tp\n    torch.distributed.broadcast(\n        logits, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()\n    )\n    torch.distributed.broadcast(\n        target, mpu.get_tensor_model_parallel_src_rank(), group=mpu.get_tensor_model_parallel_group()\n    )\n\n    tp_rank = mpu.get_tensor_model_parallel_rank()\n    vocab_size_per_tp = vocab_size // mpu.get_tensor_model_parallel_world_size()\n\n    # get the local logits of each tp\n    vocab_parallel_logits = (\n        logits.clone().detach()[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp].requires_grad_()\n    )\n    logits.grad = None\n    vocab_parallel_logits.grad = None\n\n    log_gpu_memory_usage(\"begin\")\n    output_entropy = vocab_parallel_entropy(vocab_parallel_logits)\n    log_gpu_memory_usage(\"after forward\")\n    grad_output = torch.randn_like(output_entropy)\n    output_entropy.backward(grad_output)\n    log_gpu_memory_usage(\"after backward\")\n\n    target_entropy = entropy_from_logits(logits)\n    torch.testing.assert_close(output_entropy, target_entropy)\n    target_entropy.backward(grad_output)\n    torch.testing.assert_close(\n        logits.grad[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits.grad\n    )\n    # make sure logits is not altered\n    torch.testing.assert_close(\n        logits[:, tp_rank * vocab_size_per_tp : (tp_rank + 1) * vocab_size_per_tp], vocab_parallel_logits\n    )\n\n    if mpu.get_tensor_model_parallel_rank() == 0:\n        print(\"test_vocab_parallel_entropy passes\")\n\n    mpu.destroy_model_parallel()\n\n\nif __name__ == \"__main__\":\n    local_rank, rank, world_size = initialize_global_process_group()\n    test_all_gather_data_proto()\n    test_vocab_parallel_entropy()\n"
  },
  {
    "path": "tests/special_distributed/test_torch_functional.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport torch\n\nfrom verl.utils.torch_functional import allgather_dict_into_dict\n\nif __name__ == \"__main__\":\n    torch.distributed.init_process_group(backend=\"gloo\")\n\n    local_rank = int(os.environ[\"LOCAL_RANK\"])\n    rank = int(os.environ[\"RANK\"])\n    world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    metrics_dict = {\"loss\": [0 + rank, 1 + rank, 2 + rank], \"grad_norm\": rank}\n\n    result = allgather_dict_into_dict(data=metrics_dict, group=None)\n\n    assert result[\"loss\"] == [[0, 1, 2], [1, 2, 3], [2, 3, 4], [3, 4, 5]]\n    assert result[\"grad_norm\"] == [0, 1, 2, 3]\n\n    print(result)\n"
  },
  {
    "path": "tests/special_e2e/README.md",
    "content": "This folder is reserved for end-to-end tests that typically require multiple GPUs.\n"
  },
  {
    "path": "tests/special_e2e/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/special_e2e/check_custom_rwd_fn.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\n\n\ndef check_congratulations_in_file(output_file):\n    with open(output_file) as f:\n        output = f.read()\n\n    success_message = \"Congratulations!!! You have called my_reward_function successfully!!!\"\n    assert success_message in output, f\"Success message of my_reward_function not found in {output_file}\"\n    print(\"Check passes\")\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--output_file\", required=True, type=str)\n\n    args = parser.parse_args()\n\n    check_congratulations_in_file(args.output_file)\n"
  },
  {
    "path": "tests/special_e2e/check_results.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\n\nimport numpy as np\n\n\ndef extract_reward_from_line(line):\n    # TODO: this function needs error handling\n    try:\n        key_vals = line.split(\" - \")\n        for key_val in key_vals:\n            key, val = key_val.split(\":\")\n            if key == \"critic/rewards/mean\":\n                reward = float(val)\n                return reward\n        return -np.inf\n    except Exception:\n        return -np.inf\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\"--output_file\", required=True, type=str)\n    parser.add_argument(\"--target\", type=float, default=0.2, help=\"target reward score\")\n\n    args = parser.parse_args()\n\n    with open(args.output_file) as f:\n        output = f.read().split(\"\\n\")\n\n    best_reward = -np.inf\n    for line in output:\n        if line.startswith(\"step\"):\n            reward = extract_reward_from_line(line)\n            if reward > best_reward:\n                best_reward = reward\n\n    print(f\"Best reward is {best_reward}\")\n    assert best_reward > args.target, f\"Best reward must be greater than {args.target}. best_reward: {best_reward}\"\n    print(\"Check passes\")\n"
  },
  {
    "path": "tests/special_e2e/envs/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .digit_completion import DigitCompletion\n\n__all__ = [\"DigitCompletion\"]\n"
  },
  {
    "path": "tests/special_e2e/envs/digit_completion/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom transformers import AutoTokenizer, LlamaConfig\n\nfrom .task import DigitCompletion, generate_ground_truth_response\nfrom .tokenizer import CharTokenizer\n\nAutoTokenizer.register(LlamaConfig, CharTokenizer, exist_ok=True)\n\n__all__ = [\"DigitCompletion\", \"generate_ground_truth_response\", \"CharTokenizer\"]\n"
  },
  {
    "path": "tests/special_e2e/envs/digit_completion/task.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Task and environment definition for digit completion.\"\"\"\n\nimport numpy as np\n\n\nclass DigitCompletion:\n    \"\"\"\n    The implementation of a simple digit completion task.\n    The prompt is a sequence of numbers with fixed difference. The task is to complete the next N numbers.\n    If the max number is reached, the next number should be modulo with max number.\n\n    For example,\n    - prompt = [1, 2, 3]\n    - N = 5\n    - max_number = 6\n\n    the response should be [4, 5, 6, 7%6, 8%6] = [4, 5, 6, 0, 1]\n\n    Note that the tokenizer is char-level to increase the difficulty.\n    \"\"\"\n\n    def __init__(self, max_number: int, max_diff: int, max_num_in_response: int, seed=0):\n        \"\"\"\n\n        Args:\n            max_number: the maximum number allowed in the arithmetic sequence\n            max_diff: the maximum diff. The actual common diff will be sampled from [0, max_diff]\n            max_num_in_response: the maximum number in the response\n        \"\"\"\n        super().__init__()\n        self.max_number = max_number\n        self.max_diff = max_diff\n        self.max_num_in_response = max_num_in_response\n        assert self.max_num_in_response < 10\n        assert self.max_number > 0\n        assert self.max_diff > 0\n        self.max_number_length = len(str(max_number))\n        # {num1},{num2}:{max_num_in_response},{max_number}\n        self._prompt_length = self.max_number_length * 2 + 4 + self.max_number_length  # no negative is allowed\n\n        self.np_rng = np.random.default_rng(seed=seed)\n\n    def __str__(self):\n        return (\n            f\"Prompt length: {self.prompt_length}. Response length: {self.response_length}, \"\n            f\"Max number: {self.max_number}. Max diff: {self.max_diff}, \"\n            f\"Max number in response: {self.max_num_in_response}\"\n        )\n\n    def get_state(self):\n        return {\"rng\": self.np_rng}\n\n    def set_state(self, state):\n        assert \"rng\" in state, \"rng must be inside state\"\n        self.np_rng = state[\"rng\"]\n\n    @property\n    def prompt_length(self):\n        return self._prompt_length\n\n    @property\n    def response_length(self):\n        # number length + comma length + [EOS]\n        # The actual number times 1.5 to allow 'U'\n        return (self.max_num_in_response * self.max_number_length + (self.max_num_in_response - 1) + 1) * 2\n\n    def add(self, a, b):\n        return (a + b) % self.max_number\n\n    def get_all_prompts(self):\n        all_prompts = []\n        for first_num in range(self.max_number + 1):\n            for diff in range(0, self.max_diff + 1):\n                second_num = self.add(first_num, diff)\n                for num_to_complete in range(self.max_num_in_response + 1):\n                    prompt = str(first_num) + \",\" + str(second_num) + f\":{self.max_number},{num_to_complete}\"\n                    all_prompts.append(prompt)\n        return all_prompts\n\n    def sample_str_prompts(self):\n        # step 1: sample initial numbers\n        first_num = self.np_rng.integers(self.max_number + 1)\n        diff = self.np_rng.integers(self.max_diff + 1)\n        second_num = self.add(first_num, diff)\n        num_to_complete = self.np_rng.integers(self.max_num_in_response + 1)\n        prompt = str(first_num) + \",\" + str(second_num) + f\":{self.max_number},{num_to_complete}\"\n        return prompt\n\n    def sample_batch_str_prompts(self, batch_size):\n        str_prompts = []\n        for _ in range(batch_size):\n            str_prompts.append(self.sample_str_prompts())\n        return str_prompts\n\n\ndef compute_attention_mask(prompts, pad_token_id):\n    mask = np.ones_like(prompts)\n    mask[prompts == pad_token_id] = 0\n    return mask\n\n\ndef compute_position_id_with_mask(mask):\n    return np.clip(np.cumsum(mask, axis=-1) - 1, a_min=0, a_max=None)\n\n\ndef generate_ground_truth_response(prompt: str):\n    \"\"\"Generate ground truth response given a prompt.\"\"\"\n    num, info = prompt.split(\":\")\n    num1, num2 = num.split(\",\")\n    max_number, num_to_gen = info.split(\",\")\n    num1 = int(num1)\n    num2 = int(num2)\n    max_number = int(max_number)\n    num_to_gen = int(num_to_gen)\n    diff = (num2 - num1) % max_number\n    results = []\n    last_num = num2\n    for _ in range(num_to_gen):\n        curr = (last_num + diff) % max_number\n        results.append(str(curr))\n        last_num = curr\n    response = \",\".join(results)\n    return response\n\n\ndef compute_reward(prompt: str, response: str, sequence_reward=1.0):\n    \"\"\"We compute dense reward here so that we can directly train RL without SFT\"\"\"\n    response_length = len(response)\n    ground_truth_response = generate_ground_truth_response(prompt)\n    per_token_reward = sequence_reward / (len(ground_truth_response) + 1)  # including [EOS]\n\n    # pad\n    reward = np.zeros(response_length, dtype=np.float32)  # this assumes that each char is a token\n    # assign reward until mismatches\n    ground_truth_idx = 0\n    for i in range(response_length):\n        if ground_truth_idx == len(ground_truth_response):\n            break\n\n        ground_truth_response_token = ground_truth_response[ground_truth_idx]\n        response_token = response[i]\n        if ground_truth_response_token == response_token:\n            reward[i] = per_token_reward\n            ground_truth_idx += 1\n        else:\n            # no matches\n            break\n\n    return reward, {\"ground_truth_response\": ground_truth_response}\n\n\nif __name__ == \"__main__\":\n    task = DigitCompletion(max_number=20, max_diff=3, max_num_in_response=5)\n    print(task.sample_str_prompts())\n\n    prompt = \"7,8:20,0\"\n    response = \"\"\n    print(compute_reward(prompt, response))\n\n    prompt = \"7,8:20,0\"\n    response = \"E000\"\n    print(compute_reward(prompt, response))\n\n    prompt = \"9,10:20,2\"\n    response = \"11,12,13\"\n    print(compute_reward(prompt, response))\n"
  },
  {
    "path": "tests/special_e2e/envs/digit_completion/tokenizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Copied from https://github.com/dariush-bahrami/character-tokenizer/blob/master/charactertokenizer/core.py\n\nCharacterTokenzier for Hugging Face Transformers.\n\nThis is heavily inspired from CanineTokenizer in transformers package.\n\"\"\"\n\nimport json\nimport os\nfrom pathlib import Path\nfrom typing import Optional, Sequence\n\nfrom transformers.tokenization_utils import AddedToken, PreTrainedTokenizer\n\n\nclass CharTokenizer(PreTrainedTokenizer):\n    def __init__(self, characters: Sequence[str], model_max_length: int, chat_template, **kwargs):\n        \"\"\"Character tokenizer for Hugging Face transformers.\n\n        Args:\n            characters (Sequence[str]): List of desired characters. Any character which\n                is not included in this list will be replaced by a special token called\n                [UNK] with id=6. Following are list of all of the special tokens with\n                their corresponding ids:\n                    \"[CLS]\": 0\n                    \"[SEP]\": 1\n                    \"[BOS]\": 2\n                    \"[MASK]\": 3\n                    \"[PAD]\": 4\n                    \"[RESERVED]\": 5\n                    \"[UNK]\": 6\n                an id (starting at 7) will be assigned to each character.\n\n            model_max_length (int): Model maximum sequence length.\n        \"\"\"\n        eos_token_str = \"E\"\n        sep_token_str = \"S\"\n        pad_token_str = \"P\"\n        unk_token_str = \"U\"\n\n        self.characters = characters\n        self.model_max_length = model_max_length\n        eos_token = AddedToken(eos_token_str, lstrip=False, rstrip=False)\n        sep_token = AddedToken(sep_token_str, lstrip=False, rstrip=False)\n        pad_token = AddedToken(pad_token_str, lstrip=False, rstrip=False)\n        unk_token = AddedToken(unk_token_str, lstrip=False, rstrip=False)\n\n        self._vocab_str_to_int = {\n            sep_token_str: 0,\n            eos_token_str: 1,\n            pad_token_str: 2,\n            unk_token_str: 3,\n            **{ch: i + 4 for i, ch in enumerate(characters)},\n        }\n        self._vocab_int_to_str = {v: k for k, v in self._vocab_str_to_int.items()}\n\n        super().__init__(\n            eos_token=eos_token,\n            sep_token=sep_token,\n            pad_token=pad_token,\n            unk_token=unk_token,\n            add_prefix_space=False,\n            model_max_length=model_max_length,\n            **kwargs,\n        )\n\n        self.chat_template = chat_template\n\n    @property\n    def vocab_size(self) -> int:\n        return len(self._vocab_str_to_int)\n\n    def get_vocab(self):\n        return self._vocab_str_to_int\n\n    def _tokenize(self, text: str) -> list[str]:\n        return list(text)\n\n    def _convert_token_to_id(self, token: str) -> int:\n        return self._vocab_str_to_int.get(token, self._vocab_str_to_int[\"U\"])\n\n    def _convert_id_to_token(self, index: int) -> str:\n        return self._vocab_int_to_str[index]\n\n    def convert_tokens_to_string(self, tokens):\n        return \"\".join(tokens)\n\n    def build_inputs_with_special_tokens(\n        self, token_ids_0: list[int], token_ids_1: Optional[list[int]] = None\n    ) -> list[int]:\n        sep = [self.sep_token_id]\n        cls = [self.cls_token_id]\n        result = cls + token_ids_0 + sep\n        if token_ids_1 is not None:\n            result += token_ids_1 + sep\n        return result\n\n    def get_special_tokens_mask(\n        self,\n        token_ids_0: list[int],\n        token_ids_1: Optional[list[int]] = None,\n        already_has_special_tokens: bool = False,\n    ) -> list[int]:\n        if already_has_special_tokens:\n            return super().get_special_tokens_mask(\n                token_ids_0=token_ids_0,\n                token_ids_1=token_ids_1,\n                already_has_special_tokens=True,\n            )\n\n        result = [1] + ([0] * len(token_ids_0)) + [1]\n        if token_ids_1 is not None:\n            result += ([0] * len(token_ids_1)) + [1]\n        return result\n\n    def get_config(self) -> dict:\n        return {\n            \"char_ords\": [ord(ch) for ch in self.characters],\n            \"model_max_length\": self.model_max_length,\n            \"chat_template\": self.chat_template,\n        }\n\n    @classmethod\n    def from_config(cls, config: dict):\n        cfg = {}\n        cfg[\"characters\"] = [chr(i) for i in config[\"char_ords\"]]\n        cfg[\"model_max_length\"] = config[\"model_max_length\"]\n        cfg[\"chat_template\"] = config[\"chat_template\"]\n        return cls(**cfg)\n\n    def save_pretrained(self, save_directory: str | os.PathLike, **kwargs):\n        cfg_file = Path(save_directory) / \"tokenizer_config.json\"\n        cfg = self.get_config()\n        with open(cfg_file, \"w\") as f:\n            json.dump(cfg, f, indent=4)\n\n    @classmethod\n    def from_pretrained(cls, save_directory: str | os.PathLike, **kwargs):\n        cfg_file = Path(save_directory) / \"tokenizer_config.json\"\n        with open(cfg_file) as f:\n            cfg = json.load(f)\n        return cls.from_config(cfg)\n"
  },
  {
    "path": "tests/special_e2e/generation/run_gen_qwen05.sh",
    "content": "#!/usr/bin/env bash\n# Tested with 1 & 4 GPUs\nset -xeuo pipefail\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\n\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-4}\nOUTPUT_PATH=${OUTPUT_PATH:-$HOME/data/gen/qwen_05_gen_test.parquet}\nGEN_TP=${GEN_TP:-2}  # Default tensor parallel size to 2\n\npython3 -m verl.trainer.main_generation \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    data.path=\"${HOME}/data/gsm8k/test.parquet\" \\\n    data.prompt_key=prompt \\\n    data.n_samples=1 \\\n    data.output_path=\"${OUTPUT_PATH}\" \\\n    model.path=\"${MODEL_ID}\" \\\n    +model.trust_remote_code=True \\\n    rollout.temperature=1.0 \\\n    rollout.top_k=50 \\\n    rollout.top_p=0.7 \\\n    rollout.prompt_length=2048 \\\n    rollout.response_length=1024 \\\n    rollout.tensor_model_parallel_size=\"${GEN_TP}\" \\\n    rollout.gpu_memory_utilization=0.8\n"
  },
  {
    "path": "tests/special_e2e/generation/run_gen_qwen05_server.sh",
    "content": "#!/usr/bin/env bash\n# Tested with 1 & 4 GPUs\nset -xeuo pipefail\n\nMODEL_ID=${MODEL_ID:-$HOME/models/Qwen/Qwen2.5-0.5B-Instruct}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\nOUTPUT_PATH=${OUTPUT_PATH:-$HOME/data/gen/qwen_05_gen_test.parquet}\nGEN_TP=${GEN_TP:-2}  # Default tensor parallel size to 2\n\npython3 -m verl.trainer.main_generation_server \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    actor_rollout_ref.model.path=\"${MODEL_ID}\" \\\n    actor_rollout_ref.model.trust_remote_code=True \\\n    actor_rollout_ref.rollout.temperature=1.0 \\\n    actor_rollout_ref.rollout.top_k=50 \\\n    actor_rollout_ref.rollout.top_p=0.7 \\\n    actor_rollout_ref.rollout.prompt_length=2048 \\\n    actor_rollout_ref.rollout.response_length=1024 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=\"${GEN_TP}\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.n=4 \\\n    data.train_files=\"${HOME}/data/gsm8k/test.parquet\" \\\n    data.prompt_key=prompt \\\n    +data.output_path=\"${OUTPUT_PATH}\" \\\n"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/expert_parallel/qwen2moe_minimal.json",
    "content": "{\n    \"num_hidden_layers\": 2,\n    \"max_window_layers\": 2\n}"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/expert_parallel/qwen3moe_minimal.json",
    "content": "{\n    \"num_hidden_layers\": 2,\n    \"max_window_layers\": 2\n}"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/run_function_reward.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nNUM_GPUS=${NUM_GPUS:-8}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nTRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k/train.parquet}\nVAL_FILES=${VAL_FILES:-$HOME/data/gsm8k/test.parquet}\nMAX_PROMPT_LEN=${MAX_PROMPT_LEN:-512}\nMAX_RESPONSE_LEN=${MAX_RESPONSE_LEN:-512}\n\nENGINE=${ENGINE:-vllm}\nif [ \"$ENGINE\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\nROLLOUT_MODE=\"async\"\n\nRETURN_RAW_CHAT=\"True\"\nSKIP_TOKENIZER_INIT=\"True\"\n\nGPU_MEMORY_UTILIZATION=${GPU_MEMORY_UTILIZATION:-0.7}\nACTOR_FSDP_PARAM_OFFLOAD=${ACTOR_FSDP_PARAM_OFFLOAD:-False}\nACTOR_FSDP_OPTIMIZER_OFFLOAD=${ACTOR_FSDP_OPTIMIZER_OFFLOAD:-False}\nREF_FSDP_PARAM_OFFLOAD=${REF_FSDP_PARAM_OFFLOAD:-True}\nRM_PAD=${RM_PAD:-True}\nFUSED_KERNELS=${FUSED_KERNELS:-False}\nFUSED_KERNEL_BACKEND=${FUSED_KERNEL_BACKEND:-torch} # or 'triton' for triton backend\nADV_ESTIMATOR=${ADV_ESTIMATOR:-gae}\nLOSS_MODE=${LOSS_MODE:-vanilla}\nUSE_KL=${USE_KL:-False}\nCUSTOM_REWARD_FN=${CUSTOM_REWARD_FN:-False}\nENABLE_CHUNKED_PREFILL=${ENABLE_CHUNKED_PREFILL:-True} # For vLLM VLM placeholder issue: https://github.com/vllm-project/vllm/issues/15185\nSTRATEGY=${STRATEGY:-fsdp}\n# LoRA config\nLORA_RANK=${LORA_RANK:-0}\nLORA_ALPHA=${LORA_ALPHA:-${LORA_RANK}}\nLORA_TARGET=${LORA_TARGET:-\"all-linear\"}\nLORA_EXCLUDE=${LORA_EXCLUDE:-\"DONT_EXCLUDE\"}\nUSE_SHM=${USE_SHM:-False}\nLOAD_FORMAT=${LOAD_FORMAT:-dummy}\nLAYERED_SUMMON=${LAYERED_SUMMON:-False}\n# Validation\nVAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False}\nTEST_FREQ=${TEST_FREQ:--1}\n# Save & Resume\nRESUME_MODE=${RESUME_MODE:-disable}\nSAVE_FREQ=${SAVE_FREQ:--1}\nTOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1}\n\n# whether to save hf_model\nSAVE_HF_MODEL=${SAVE_HF_MODEL:-False}\nFSDP_SIZE=${FSDP_SIZE:--1}\nSP_SIZE=${SP_SIZE:-1}\n\nif [ \"${SAVE_HF_MODEL}\" = \"True\" ]; then\n    CHECKPOINT_CONTENTS=\"['model','hf_model','optimizer','extra']\"\nelse\n    CHECKPOINT_CONTENTS=\"['model','optimizer','extra']\"\nfi\n\ntrain_traj_micro_bsz_per_gpu=2 # b\nn_resp_per_prompt=4 # g\n\ntrain_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * NUM_GPUS)) # b * n\ntrain_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n\ntrain_prompt_mini_bsz=$((train_traj_mini_bsz * n_resp_per_prompt)) # 2 * b * n / g\ntrain_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g\n\nreward_fn_name=null\nreward_fn_file_path=null\noutput_file=\"$(pwd)/output.txt\"\nif [ \"${CUSTOM_REWARD_FN}\" = \"True\" ]; then\n    reward_fn_name=\"my_reward_function\"\n    reward_fn_file_path=\"$(pwd)/my_reward_function.py\"\n    rm -rf \"${reward_fn_file_path}\"\n    cat <<EOF > \"$reward_fn_file_path\"\ndef ${reward_fn_name}(data_source, solution_str, ground_truth, extra_info=None):\n    print(f\"Congratulations!!! You have called ${reward_fn_name} successfully!!!\")\n    return 0.1\nEOF\n\n    rm -rf \"${output_file}\"\nfi\n\nexp_name=\"${VERL_EXP_NAME:-$(basename \"${MODEL_ID,,}\")-function-reward-minimal}\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=\"${ADV_ESTIMATOR}\" \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=\"${train_prompt_bsz}\" \\\n    data.max_prompt_length=\"${MAX_PROMPT_LEN}\" \\\n    data.max_response_length=\"${MAX_RESPONSE_LEN}\" \\\n    data.return_raw_chat=${RETURN_RAW_CHAT} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_shm=${USE_SHM} \\\n    actor_rollout_ref.model.lora_rank=${LORA_RANK} \\\n    actor_rollout_ref.model.lora_alpha=${LORA_ALPHA} \\\n    actor_rollout_ref.model.target_modules=${LORA_TARGET} \\\n    actor_rollout_ref.model.exclude_modules=${LORA_EXCLUDE} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=\"${RM_PAD}\" \\\n    actor_rollout_ref.model.use_fused_kernels=${FUSED_KERNELS} \\\n    actor_rollout_ref.model.fused_kernel_options.impl_backend=${FUSED_KERNEL_BACKEND} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.actor.strategy=${STRATEGY} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${ACTOR_FSDP_PARAM_OFFLOAD} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${ACTOR_FSDP_OPTIMIZER_OFFLOAD} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${FSDP_SIZE} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=\"${SP_SIZE}\" \\\n    actor_rollout_ref.actor.checkpoint.save_contents=${CHECKPOINT_CONTENTS} \\\n    actor_rollout_ref.actor.use_kl_loss=\"${USE_KL}\" \\\n    actor_rollout_ref.actor.policy_loss.loss_mode=\"${LOSS_MODE}\" \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=\"${ENGINE}\" \\\n    actor_rollout_ref.rollout.mode=\"${ROLLOUT_MODE}\" \\\n    actor_rollout_ref.rollout.load_format=${LOAD_FORMAT} \\\n    actor_rollout_ref.rollout.layered_summon=${LAYERED_SUMMON} \\\n    actor_rollout_ref.rollout.skip_tokenizer_init=\"${SKIP_TOKENIZER_INIT}\" \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=\"${GPU_MEMORY_UTILIZATION}\" \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=\"${ENABLE_CHUNKED_PREFILL}\" \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=\"${REF_FSDP_PARAM_OFFLOAD}\" \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=\"${RM_PAD}\" \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.model.enable_gradient_checkpointing=False \\\n    critic.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.custom_reward_function.path=\"${reward_fn_file_path}\"\\\n    reward.custom_reward_function.name=\"${reward_fn_name}\"\\\n    algorithm.use_kl_in_reward=\"${USE_KL}\" \\\n    algorithm.kl_penalty=kl \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl-test' \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=\"${NUM_GPUS}\" \\\n    trainer.val_before_train=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.test_freq=\"${TEST_FREQ}\" \\\n    trainer.save_freq=\"${SAVE_FREQ}\" \\\n    trainer.resume_mode=\"${RESUME_MODE}\" \\\n    trainer.total_epochs=2 \\\n    trainer.device=cuda \\\n    trainer.total_training_steps=\"${TOTAL_TRAIN_STEPS}\" $@ \\\n    | tee \"${output_file}\"\n\nif [ \"${CUSTOM_REWARD_FN}\" = \"True\" ]; then\n    python3 tests/special_e2e/check_custom_rwd_fn.py --output_file=\"${output_file}\"\n    check_exit_code=$?\n    rm -rf \"${reward_fn_file_path}\"\n    rm -rf \"${output_file}\"\n    # Return the exit code of check_custom_rwd_fn.py if it fails\n    if [ $check_exit_code -ne 0 ]; then\n        exit $check_exit_code\n    fi\nfi\n"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/run_model_reward.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nNUM_GPUS=${NUM_GPUS:-8}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nTRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k/train.parquet}\nVAL_FILES=${VAL_FILES:-$HOME/data/gsm8k/test.parquet}\n\nRM_PAD=${RM_PAD:-True}\nFUSED_KERNELS=${FUSED_KERNELS:-False}\nFUSED_KERNEL_BACKEND=${FUSED_KERNEL_BACKEND:-torch} # or 'triton' for triton backend\nSP_SIZE=${SP_SIZE:-1}\nSEQ_BALANCE=${SEQ_BALANCE:-False}\nLIGER=${LIGER:-False}\n# Validation\nVAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False}\nTEST_FREQ=${TEST_FREQ:--1}\n# Save & Resume\nRESUME_MODE=${RESUME_MODE:-disable}\nSAVE_FREQ=${SAVE_FREQ:--1}\nTOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1}\n\ntrain_traj_micro_bsz_per_gpu=2 # b\nn_resp_per_prompt=4 # g\n\ntrain_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * NUM_GPUS)) # b * n\ntrain_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n\ntrain_prompt_mini_bsz=$((train_traj_mini_bsz * n_resp_per_prompt)) # 2 * b * n / g\ntrain_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g\n\ntrain_max_token_num_per_gpu=32768\ninfer_max_token_num_per_gpu=32768\n\nexp_name=\"$(basename \"${MODEL_ID,,}\")-model-reward-minimal\"\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=512 \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_liger=\"${LIGER}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=\"${RM_PAD}\" \\\n    actor_rollout_ref.model.use_fused_kernels=${FUSED_KERNELS} \\\n    actor_rollout_ref.model.fused_kernel_options.impl_backend=${FUSED_KERNEL_BACKEND} \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps_ratio=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.use_dynamic_bsz=\"${SEQ_BALANCE}\" \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${train_max_token_num_per_gpu} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=\"${SP_SIZE}\" \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_max_token_num_per_gpu} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_max_token_num_per_gpu} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    critic.optim.lr=1e-5 \\\n    critic.ulysses_sequence_parallel_size=\"${SP_SIZE}\" \\\n    critic.model.use_remove_padding=\"${RM_PAD}\" \\\n    critic.optim.lr_warmup_steps_ratio=0.05 \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.model.enable_gradient_checkpointing=False \\\n    critic.use_dynamic_bsz=\"${SEQ_BALANCE}\" \\\n    critic.ppo_max_token_len_per_gpu=${train_max_token_num_per_gpu} \\\n    critic.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    critic.model.fsdp_config.param_offload=False \\\n    critic.model.fsdp_config.optimizer_offload=False \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=\"${MODEL_PATH}\" \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.8 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=1 \\\n    reward.reward_model.rollout.prompt_length=1024 \\\n    reward.reward_model.rollout.response_length=512 \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl-test' \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=\"${NUM_GPUS}\" \\\n    trainer.val_before_train=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.test_freq=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.save_freq=\"${SAVE_FREQ}\" \\\n    trainer.resume_mode=\"${RESUME_MODE}\" \\\n    trainer.total_epochs=2 \\\n    trainer.total_training_steps=\"${TOTAL_TRAIN_STEPS}\" $@\n"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/run_single_gpu.sh",
    "content": "PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n  data.train_files=$HOME/data/gsm8k/train.parquet \\\n  data.val_files=$HOME/data/gsm8k/test.parquet \\\n  data.train_batch_size=256  \\\n  data.max_prompt_length=512 \\\n  data.max_response_length=256  \\\n  actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  actor_rollout_ref.actor.optim.lr=1e-6 \\\n  actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n  actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4  \\\n  actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n  actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n  actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n  actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n  critic.optim.lr=1e-5 \\\n  critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  critic.ppo_micro_batch_size_per_gpu=4 \\\n  algorithm.kl_ctrl.kl_coef=0.001 \\\n  trainer.logger=console \\\n  trainer.val_before_train=False \\\n  trainer.n_gpus_per_node=1 \\\n  trainer.nnodes=1 \\\n  actor_rollout_ref.rollout.name=hf \\\n  trainer.total_training_steps=2"
  },
  {
    "path": "tests/special_e2e/ppo_trainer/run_single_gpu_with_engine.sh",
    "content": "PYTHONUNBUFFERED=1 python3 -m verl.trainer.main_ppo \\\n  data.train_files=$HOME/data/gsm8k/train.parquet \\\n  data.val_files=$HOME/data/gsm8k/test.parquet \\\n  data.train_batch_size=256  \\\n  data.max_prompt_length=512 \\\n  data.max_response_length=256  \\\n  actor_rollout_ref.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  actor_rollout_ref.actor.optim.lr=1e-6 \\\n  actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n  actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=4  \\\n  actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n  actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n  actor_rollout_ref.rollout.gpu_memory_utilization=0.4 \\\n  actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n  critic.optim.lr=1e-5 \\\n  critic.model.path=Qwen/Qwen2.5-0.5B-Instruct \\\n  critic.ppo_micro_batch_size_per_gpu=4 \\\n  algorithm.kl_ctrl.kl_coef=0.001 \\\n  trainer.logger=['console'] \\\n  trainer.val_before_train=False \\\n  trainer.n_gpus_per_node=1 \\\n  trainer.nnodes=1 \\\n  actor_rollout_ref.rollout.name=hf \\\n  trainer.use_legacy_worker_impl=disable \\\n  trainer.total_training_steps=2\n"
  },
  {
    "path": "tests/special_e2e/run_dapo.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nNUM_GPUS=${NUM_GPUS:-8}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nadv_estimator=grpo\n\nkl_coef=0.0\nuse_kl_in_reward=False\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=1024\nmax_response_length=2048\nenable_overlong_buffer=True\noverlong_buffer_len=128\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\nenable_filter_groups=True\nfilter_groups_metric=seq_reward\nmax_num_gen_batches=10\n\ntrain_traj_micro_bsz_per_gpu=2 # b\nn_resp_per_prompt=4 # g\n\ntrain_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * NUM_GPUS)) # b * n\ntrain_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n\ntrain_prompt_mini_bsz=$((train_traj_mini_bsz * n_resp_per_prompt)) # 2 * b * n / g\ntrain_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g\n\ngen_prompt_bsz=$((train_prompt_bsz * 4))\n\nexp_name=\"$(basename \"${MODEL_ID,,}\")-dapo-minimal\"\n\npython3 -m recipe.dapo.main_dapo \\\n    data.train_files=\"${HOME}/data/gsm8k/train.parquet\" \\\n    data.val_files=\"${HOME}/data/gsm8k/test.parquet\" \\\n    reward.reward_manager.name=dapo \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    reward.overlong_buffer.enable=${enable_overlong_buffer} \\\n    reward.overlong_buffer.len=${overlong_buffer_len} \\\n    reward.overlong_buffer.penalty_factor=${overlong_penalty_factor} \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    algorithm.filter_groups.enable=${enable_filter_groups} \\\n    algorithm.filter_groups.metric=${filter_groups_metric} \\\n    algorithm.filter_groups.max_num_gen_batches=${max_num_gen_batches} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_fused_kernels=True \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    trainer.logger=console \\\n    trainer.project_name='verl-test' \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=${NUM_GPUS} \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=2 \\\n    trainer.resume_mode=disable \\\n    trainer.val_before_train=False \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "tests/special_e2e/run_fully_async_policy.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n# Test script for fully_async_policy E2E regression testing\n# This script runs fully async PPO training with both FSDP2 and Megatron backends\n# to ensure the asynchronous training mechanism works correctly\n\nNUM_GPUS=${NUM_GPUS:-8}\nACTOR_STRATEGY=${ACTOR_STRATEGY:-\"fsdp2\"}  # fsdp2 or megatron\n\n# Download model if not exists\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n# hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=1024\nmax_response_length=2048\nenable_overlong_buffer=True\noverlong_buffer_len=128\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Temperature parameters\ntemperature=1.0\ntop_p=1.0\ntop_k=-1\nval_top_p=0.7\n\n# Fully async specific parameters\nn_gpus_rollout=4\nn_gpus_training=4\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=16\ntotal_rollout_steps=$(((128)))\ntest_freq=-1\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\npartial_rollout=True\nuse_trainer_do_validate=False\n\nexp_name=\"$(basename \"${MODEL_ID,,}\")-fully-async-policy-${ACTOR_STRATEGY}-minimal\"\n\necho \"Running fully_async_policy with ${ACTOR_STRATEGY} strategy\"\necho \"Total GPUs: ${NUM_GPUS}, Rollout GPUs: ${n_gpus_rollout}, Training GPUs: ${n_gpus_training}\"\n\n# Common parameters for both FSDP2 and Megatron\ncommon_params=(\n    data.train_files=\"${HOME}/data/gsm8k/train.parquet\"\n    data.val_files=\"${HOME}/data/gsm8k/test.parquet\"\n    data.prompt_key=prompt\n    data.truncation='left'\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.train_batch_size=${train_prompt_bsz}\n    data.gen_batch_size=${gen_prompt_bsz}\n    data.return_raw_chat=${return_raw_chat}\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    actor_rollout_ref.rollout.calculate_log_probs=True\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n    actor_rollout_ref.hybrid_engine=False\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low}\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high}\n    actor_rollout_ref.actor.clip_ratio_c=10.0\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.optim.lr_warmup_steps=-1\n    actor_rollout_ref.actor.optim.weight_decay=0.1\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80\n    actor_rollout_ref.rollout.temperature=${temperature}\n    actor_rollout_ref.rollout.top_p=${top_p}\n    actor_rollout_ref.rollout.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature}\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p}\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.name=${rollout_name}\n    actor_rollout_ref.rollout.mode=${rollout_mode}\n    actor_rollout_ref.rollout.disable_log_stats=False\n    reward.reward_manager.name=dapo\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer}\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len}\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor}\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False\n    +reward.reward_kwargs.max_resp_len=${max_response_length}\n    trainer.logger=['console']\n    trainer.project_name='verl-test-fully-async'\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.val_before_train=True\n    trainer.save_freq=-1\n    trainer.resume_mode=disable\n    trainer.nnodes=1\n    trainer.n_gpus_per_node=${n_gpus_training}\n    trainer.log_val_generations=10\n    rollout.nnodes=1\n    rollout.n_gpus_per_node=${n_gpus_rollout}\n    rollout.total_rollout_steps=${total_rollout_steps}\n    trainer.total_epochs=2\n    trainer.test_freq=${test_freq}\n    # Fully async specific configurations\n    async_training.staleness_threshold=${staleness_threshold}\n    async_training.partial_rollout=\"${partial_rollout}\"\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\"\n    async_training.use_trainer_do_validate=${use_trainer_do_validate}\n    actor_rollout_ref.rollout.checkpoint_engine.backend='nccl'\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024\n)\n\n    # Detect device\n    device_name=$(python3 - <<'EOF'\nfrom verl.utils.device import get_device_name\nprint(get_device_name())\nEOF\n)\n\nif [ \"${ACTOR_STRATEGY}\" == \"fsdp2\" ]; then\n    echo \"Running fully async training with FSDP2 strategy...\"\n    # FSDP2 specific parameters\n    gen_tp=1\n    sp_size=1\n    fsdp_size=1\n    ref_offload=True\n    actor_offload=False\n\n    if [ -n \"$device_name\" ] && [ \"$device_name\" == \"npu\" ]; then\n        common_params+=(\n            # Todo The checkpoint_engine.backend should be unified to nccl\n            # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl'\n            actor_rollout_ref.rollout.gpu_memory_utilization=0.70\n        )\n        actor_offload=True\n    fi\n    python3 -m verl.experimental.fully_async_policy.fully_async_main \\\n        \"${common_params[@]}\" \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n        critic.strategy=fsdp2 \\\n        actor_rollout_ref.actor.grad_clip=1.0 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.actor.use_dynamic_bsz=True \\\n        actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n        actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n        actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n        actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n        actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} $@\n\nelif [ \"${ACTOR_STRATEGY}\" == \"megatron\" ]; then\n    echo \"Running fully async training with Megatron strategy...\"\n    # Megatron specific parameters\n    gen_tp=2\n    train_tp=1\n    train_pp=2\n    ref_offload=True\n    actor_offload=False\n\n    if [ -n \"$device_name\" ] && [ \"$device_name\" == \"npu\" ]; then\n        train_tp=2\n        actor_offload=True\n        common_params+=(\n            # Todo The checkpoint_engine.backend should be unified to nccl\n            # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl'\n            actor_rollout_ref.rollout.gpu_memory_utilization=0.60\n        )\n    fi\n    python3 -m verl.experimental.fully_async_policy.fully_async_main \\\n        --config-path=config \\\n        --config-name='fully_async_ppo_megatron_trainer.yaml' \\\n        \"${common_params[@]}\" \\\n        actor_rollout_ref.actor.strategy=megatron \\\n        critic.strategy=megatron \\\n        actor_rollout_ref.actor.optim.lr_decay_steps=10000000 \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n        actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n        actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n        actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n        actor_rollout_ref.ref.megatron.param_offload=${ref_offload} $@\nelse\n    echo \"Error: Unknown strategy ${ACTOR_STRATEGY}. Please use 'fsdp2' or 'megatron'\"\n    exit 1\nfi\n\necho \"Fully async policy E2E test completed successfully with ${ACTOR_STRATEGY} strategy\"\n\n"
  },
  {
    "path": "tests/special_e2e/run_geo3k_fsdp_sgl_multiturn_w_tool.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\n#hf download Qwen/Qwen2.5-VL-3B-Instruct --local-dir $HOME/models/Qwen/Qwen2.5-VL-3B-Instruct\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\nFSDP_STRATEGY=${FSDP_STRATEGY:-fsdp}\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='geo3k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=64 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=$HOME/models/Qwen/Qwen2.5-VL-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=64 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.strategy=$FSDP_STRATEGY \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=8 \\\n    actor_rollout_ref.ref.strategy=$FSDP_STRATEGY \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='geo3k_async_rl' \\\n    trainer.experiment_name=qwen2.5-vl-3b_function_rm-geo3k-sgl-multi-w-tool-$FSDP_STRATEGY-rebased-0619-verify-n8 \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    data.train_files=$HOME/data/geo3k_verl_sgl_multi_turn_preprocessed/train.parquet \\\n    data.val_files=$HOME/data/geo3k_verl_sgl_multi_turn_preprocessed/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/geo3k_tool_config.yaml\" \\\n    trainer.val_before_train=False \\\n    trainer.total_training_steps=1 $@"
  },
  {
    "path": "tests/special_e2e/run_grpo_lora_with_merge.sh",
    "content": "#!/usr/bin/env bash\n#\n#  An e2e test script for testing the GRPO LoRA training process \n#  and processing the generated checkpoint using the merge_model.py script.  \n\nset -xeuo pipefail\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\nif [ ! -d \"$MODEL_PATH\" ]; then\n    echo \"Downloading model to ${MODEL_PATH}...\"\n#    hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\nelse\n    echo \"Model directory ${MODEL_PATH} already exists, skip downloading.\"\nfi\n\n\nBATCH_SIZE=16\nEXP_NAME=\"qwen2.5_0.5b_grpo_lora\"\n# step 1. train model with grpo-lora for 1 step\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=${BATCH_SIZE} \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.model.use_shm=True \\\n    actor_rollout_ref.model.lora_rank=64 \\\n    actor_rollout_ref.model.lora_alpha=32 \\\n    actor_rollout_ref.actor.optim.lr=3e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${BATCH_SIZE} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=40 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name=${EXP_NAME} \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.total_training_steps=1 \\\n    trainer.save_freq=1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=1 $@\n\n# step 2. merge model\npython3 -m verl.model_merger merge \\\n    --backend fsdp \\\n    --local_dir checkpoints/verl_grpo_example_gsm8k/${EXP_NAME}/global_step_1/actor/ \\\n    --target_dir checkpoints/verl_grpo_example_gsm8k/${EXP_NAME}/global_step_1/actor/hf\n\n# step 3. assert\n# make sure adapter_model.safetensors exists and its size is larger than 1MB\nfile_path=\"checkpoints/verl_grpo_example_gsm8k/${EXP_NAME}/global_step_1/actor/hf/lora_adapter/adapter_model.safetensors\"\n\nif [ ! -f \"$file_path\" ]; then\n    echo \"Error: File $file_path does not exist!\"\n    exit 1\nfi\n\nfile_size=$(stat -c %s \"$file_path\")\n\nmin_size_mb=1\nmin_size=$((min_size_mb * 1024 * 1024))  # 1MB = 1048576 bytes\n\nif [ \"$file_size\" -lt \"$min_size\" ]; then\n    echo \"Error: File $file_path is too small! Current size: $((file_size/1024))KB, Required: ${min_size_mb}MB\"\n    exit 1\nfi\n\necho \"Check passed: File exists and size is $(($file_size/1024/1024))MB\"\nexit 0\n"
  },
  {
    "path": "tests/special_e2e/run_gsm8k_fsdp_sgl_multiturn_sf_tool.sh",
    "content": "# run on 8xH20\n# make sure your current working directory is the root of the project\n\nset -x\n\n\nexport PYTHONUNBUFFERED=1\nexport RAY_DEDUP_LOGS=0\nexport RUST_BACKTRACE=1\nexport HYDRA_FULL_ERROR=1\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_sf_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=128 \\\n    data.max_prompt_length=2048 \\\n    data.max_response_length=16384 \\\n    data.filter_overlong_prompts=False \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    data.train_files=$HOME/data/retool_dapo/train.parquet \\\n    data.val_files=$HOME/data/retool_aime2024/train.parquet \\\n    actor_rollout_ref.model.path=Qwen/Qwen3-4B \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.use_liger=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    +actor_rollout_ref.model.enable_activation_offload=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=1 \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=32768 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.0 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/sandbox_fusion_tool_config.yaml\" \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='retool_async_rl' \\\n    trainer.experiment_name='qwen3-4b_function_rm-retool-async-sgl-no-sft-n8-v2505271300' \\\n    trainer.val_before_train=False \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=100 \\\n    trainer.test_freq=20 \\\n    trainer.total_training_steps=1000 \\\n    trainer.total_epochs=1 $@"
  },
  {
    "path": "tests/special_e2e/run_gsm8k_fsdp_sgl_multiturn_w_tool.sh",
    "content": "# run on 8xH100\n# make sure your current working directory is the root of the project\n\nset -x\n\n#hf download Qwen/Qwen2.5-3B-Instruct --local-dir $HOME/models/Qwen/Qwen2.5-3B-Instruct\n\nulimit -n 65535\n\nPROJECT_DIR=\"$(pwd)\"\nCONFIG_PATH=\"$PROJECT_DIR/examples/sglang_multiturn/config\"\nFSDP_STRATEGY=${FSDP_STRATEGY:-fsdp}\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=\"$CONFIG_PATH\" \\\n    --config-name='gsm8k_multiturn_grpo' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_batch_size=256 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.return_raw_chat=True \\\n    actor_rollout_ref.model.path=$HOME/models/Qwen/Qwen2.5-3B-Instruct \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=256 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=16 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.strategy=$FSDP_STRATEGY \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.5 \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.strategy=$FSDP_STRATEGY \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='gsm8k_async_rl' \\\n    trainer.experiment_name=qwen2.5-3b_function_rm-gsm8k-sgl-multi-w-tool-$FSDP_STRATEGY-rebased-0427-verify-n16 \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    data.train_files=$HOME/data/gsm8k_verl_sgl_multi_turn_preprocessed/train.parquet \\\n    data.val_files=$HOME/data/gsm8k_verl_sgl_multi_turn_preprocessed/test.parquet \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=\"$PROJECT_DIR/examples/sglang_multiturn/config/tool_config/gsm8k_tool_config.yaml\" \\\n    trainer.val_before_train=False \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "tests/special_e2e/run_one_step_off_policy.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n# Test script for one_step_off_policy E2E regression testing\n# This script runs one_step_off_policy with both FSDP2 and Megatron backends\n# to ensure the asynchronous training mechanism works correctly\n\nNUM_GPUS=${NUM_GPUS:-8}\nACTOR_STRATEGY=${ACTOR_STRATEGY:-\"fsdp2\"}  # fsdp2 or megatron\n\n# Download model if not exists\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=1024\nmax_response_length=2048\nenable_overlong_buffer=True\noverlong_buffer_len=128\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\ntrain_prompt_bsz=8\nn_resp_per_prompt=3\ntrain_prompt_mini_bsz=4\n\n# Temperature parameters\ntemperature=1.0\ntop_p=1.0\ntop_k=-1\nval_top_p=0.7\n\n# One-step-off-policy specific parameters\n# Allocate 2 GPUs for rollout, remaining for training\nn_gpus_rollout=2\nn_gpus_training=$((NUM_GPUS - n_gpus_rollout))\n\nexp_name=\"$(basename \"${MODEL_ID,,}\")-one-step-off-policy-${ACTOR_STRATEGY}-minimal\"\n\necho \"Running one_step_off_policy with ${ACTOR_STRATEGY} strategy\"\necho \"Total GPUs: ${NUM_GPUS}, Rollout GPUs: ${n_gpus_rollout}, Training GPUs: ${n_gpus_training}\"\n\n# Common parameters for both FSDP2 and Megatron\ncommon_params=(\n    data.train_files=\"${HOME}/data/gsm8k/train.parquet\"\n    data.val_files=\"${HOME}/data/gsm8k/test.parquet\"\n    data.prompt_key=prompt\n    data.truncation='left'\n    data.max_prompt_length=${max_prompt_length}\n    data.max_response_length=${max_response_length}\n    data.train_batch_size=${train_prompt_bsz}\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt}\n    algorithm.adv_estimator=${adv_estimator}\n    algorithm.use_kl_in_reward=${use_kl_in_reward}\n    algorithm.kl_ctrl.kl_coef=${kl_coef}\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss}\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef}\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low}\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high}\n    actor_rollout_ref.actor.clip_ratio_c=10.0\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\"\n    actor_rollout_ref.actor.optim.lr=1e-6\n    actor_rollout_ref.actor.optim.lr_warmup_steps=-1\n    actor_rollout_ref.actor.optim.weight_decay=0.1\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz}\n    actor_rollout_ref.actor.entropy_coeff=0\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode}\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80\n    actor_rollout_ref.rollout.temperature=${temperature}\n    actor_rollout_ref.rollout.top_p=${top_p}\n    actor_rollout_ref.rollout.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature}\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p}\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k}\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True\n    actor_rollout_ref.rollout.val_kwargs.n=1\n    actor_rollout_ref.rollout.enable_chunked_prefill=True\n    actor_rollout_ref.rollout.name=vllm\n    actor_rollout_ref.rollout.checkpoint_engine.backend='nccl'\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=1024\n    reward.reward_manager.name=dapo\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer}\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len}\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor}\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False\n    +reward.reward_kwargs.max_resp_len=${max_response_length}\n    trainer.logger=['console']\n    trainer.project_name='verl-test'\n    trainer.experiment_name=\"${exp_name}\"\n    trainer.val_before_train=True\n    trainer.test_freq=-1\n    trainer.save_freq=-1\n    trainer.total_epochs=2\n    trainer.total_training_steps=2\n    trainer.resume_mode=disable\n    trainer.nnodes=1\n    trainer.n_gpus_per_node=${n_gpus_training}\n    rollout.nnodes=1\n    rollout.n_gpus_per_node=${n_gpus_rollout}\n\n)\n\n    # Detect device\n    device_name=$(python3 - <<'EOF'\nfrom verl.utils.device import get_device_name\nprint(get_device_name())\nEOF\n)\n\nif [ \"${ACTOR_STRATEGY}\" == \"fsdp2\" ]; then\n    echo \"Running with FSDP2 strategy...\"\n    # FSDP2 specific parameters\n    gen_tp=2\n    sp_size=2\n    fsdp_size=2\n    ref_offload=True\n    actor_offload=False\n\n    if [ \"$device_name\" ] && [ \"$device_name\" == \"npu\" ]; then\n        common_params+=(\n            # Todo The checkpoint_engine.backend should be unified to nccl\n            # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl'\n            actor_rollout_ref.rollout.gpu_memory_utilization=0.60\n        )\n        actor_offload=True\n    fi\n\n    python3 -m verl.experimental.one_step_off_policy.main_ppo \\\n        \"${common_params[@]}\" \\\n        actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n        critic.strategy=fsdp2 \\\n        actor_rollout_ref.actor.grad_clip=1.0 \\\n        actor_rollout_ref.model.use_remove_padding=True \\\n        actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n        actor_rollout_ref.actor.use_dynamic_bsz=True \\\n        actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n        actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n        actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n        actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n        actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n        actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n        actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n        actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} $@\n\nelif [ \"${ACTOR_STRATEGY}\" == \"megatron\" ]; then\n    echo \"Running with Megatron strategy...\"\n    # Megatron specific parameters\n    gen_tp=2\n    train_tp=1\n    train_pp=2\n    ref_offload=True\n    actor_offload=False\n\n    if [ \"$device_name\" ] && [ \"$device_name\" == \"npu\" ]; then\n        common_params+=(\n            # Todo The checkpoint_engine.backend should be unified to nccl\n            # actor_rollout_ref.rollout.checkpoint_engine.backend='hccl'\n            actor_rollout_ref.rollout.gpu_memory_utilization=0.70\n            trainer.n_gpus_per_node=4\n            rollout.n_gpus_per_node=4\n            actor_rollout_ref.model.use_remove_padding=True \\\n            actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n            actor_rollout_ref.actor.use_dynamic_bsz=True \\\n            actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n            actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n        )\n        train_tp=2\n        actor_offload=True\n    fi\n\n    python3 -m verl.experimental.one_step_off_policy.main_ppo \\\n        --config-path=config \\\n        --config-name='one_step_off_ppo_megatron_trainer.yaml' \\\n        \"${common_params[@]}\" \\\n        actor_rollout_ref.actor.strategy=megatron \\\n        critic.strategy=megatron \\\n        actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n        actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n        actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \\\n        actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n        actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n        actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n        actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n        actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n        actor_rollout_ref.ref.megatron.param_offload=${ref_offload} $@\nelse\n    echo \"Error: Unknown strategy ${ACTOR_STRATEGY}. Please use 'fsdp2' or 'megatron'\"\n    exit 1\nfi\n\necho \"One-step-off-policy E2E test completed successfully with ${ACTOR_STRATEGY} strategy\""
  },
  {
    "path": "tests/special_e2e/run_ppo_trainer_megatron.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\nexport VERL_LOGGING_LEVEL=INFO\nexport VERL_PPO_LOGGING_LEVEL=INFO\n\nNUM_GPUS=${NUM_GPUS:-8}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\nRM_MODEL_PATH=${RM_MODEL_PATH:-${HOME}/models/Skywork/Skywork-Reward-V2-Llama-3.2-1B}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nUSE_DUMMY_MODEL=${USE_DUMMY_MODEL:-False}\nDUMMY_MODEL_PATH=${DUMMY_MODEL_PATH:-${HOME}/dummy_models/${MODEL_ID}}\nif [ \"$USE_DUMMY_MODEL\" = \"True\" ]; then\n    if [ -z \"${DUMMY_MODEL_CONFIG_PATH}\"  ]; then\n        echo \"[ERROR] DUMMY_MODEL_CONFIG_PATH not set\"\n        exit 1\n    fi\n\n    python scripts/init_random_model.py \\\n        --hf_model_path \"${MODEL_PATH}\" \\\n        --new_config_path \"${DUMMY_MODEL_CONFIG_PATH}\" \\\n        --output_path \"${DUMMY_MODEL_PATH}\"\n\n    MODEL_PATH=\"${DUMMY_MODEL_PATH}\"\nfi\n\nTRAIN_FILES=${TRAIN_FILES:-${HOME}/data/gsm8k/train.parquet}\nVAL_FILES=${VAL_FILES:-${HOME}/data/gsm8k/test.parquet}\n\nADV_ESTIMATOR=${ADV_ESTIMATOR:-gae}\n# Validation\nVAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False}\nTEST_FREQ=${TEST_FREQ:--1}\n# Save & Resume\nRESUME_MODE=${RESUME_MODE:-disable}\nSAVE_FREQ=${SAVE_FREQ:--1}\nTOTAL_TRAIN_STEPS=${TOTAL_TRAIN_STEPS:-1}\n\nUSE_DYNAMIC_BSZ=${USE_DYNAMIC_BSZ:-True}\nppo_max_token_len_per_gpu=${PPO_MAX_TOKEN_LEN:-2400}\nforward_max_token_len_per_gpu=${FWD_MAX_TOKEN_LEN:-4800}\ntrain_traj_micro_bsz_per_gpu=${MICRO_BSZ:-2} # b\nn_resp_per_prompt=4 # g\n\ntrain_traj_micro_bsz=$((train_traj_micro_bsz_per_gpu * NUM_GPUS)) # b * n\ntrain_traj_mini_bsz=$((train_traj_micro_bsz * 2)) # 2 * b * n\ntrain_prompt_mini_bsz=$((train_traj_mini_bsz * n_resp_per_prompt)) # 2 * b * n / g\ntrain_prompt_bsz=$((train_prompt_mini_bsz * 2)) # 4 * b * n / g\n\nLORA_RANK=${LORA_RANK:-0}\nCRITIC_LORA_RANK=${CRITIC_LORA_RANK:-$LORA_RANK}\nLORA_ALPHA=${LORA_ALPHA:-${LORA_RANK}}\nLORA_TARGET_MODULES=${LORA_TARGET_MODULES:-\"['linear_qkv','linear_proj','linear_fc1','linear_fc2']\"}\nLORA_MERGE=${LORA_MERGE:-False}\n\nMAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512}\nMAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-512}\nMAX_RM_LENGTH=$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))\n\nCOMMON_PP=${COMMON_PP:-2}\nCOMMON_VPP=${COMMON_VPP:-2}\nCOMMON_CP=${COMMON_CP:-2}\nCOMMON_TP=${COMMON_TP:-2}\nCOMMON_EP=${COMMON_EP:-1}\nCOMMON_ETP=${COMMON_ETP:-1}\n\nTRAIN_TP=${TRAIN_TP:-$COMMON_TP}\nINFER_TP=${INFER_TP:-$COMMON_TP}\n\nACTOR_PP=${ACTOR_PP:-$COMMON_PP}\nACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}\nACTOR_CP=${ACTOR_CP:-$COMMON_CP}\nACTOR_TP=${ACTOR_TP:-$TRAIN_TP}\nACTOR_EP=${ACTOR_EP:-$COMMON_EP}\nACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}\nROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}\nREF_PP=${REF_PP:-$COMMON_PP}\nREF_VPP=${REF_VPP:-$COMMON_VPP}\nREF_CP=${REF_CP:-$COMMON_CP}\nREF_TP=${REF_TP:-$TRAIN_TP}\nREF_EP=${REF_EP:-$COMMON_EP}\nREF_ETP=${REF_ETP:-$COMMON_ETP}\nCRITIC_PP=${CRITIC_PP:-$COMMON_PP}\nCRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}\nCRITIC_CP=${CRITIC_CP:-$COMMON_CP}\nCRITIC_TP=${CRITIC_TP:-$TRAIN_TP}\nCRITIC_EP=${CRITIC_EP:-$COMMON_EP}\nCRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}\n\nALL_OFFLOAD=${ALL_OFFLOAD:-False}\nCOMMON_PARAM_OFFLOAD=${COMMON_PARAM_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_GRAD_OFFLOAD=${COMMON_GRAD_OFFLOAD:-$ALL_OFFLOAD}\nCOMMON_OPTIMIZER_OFFLOAD=${COMMON_OPTIMIZER_OFFLOAD:-$ALL_OFFLOAD}\n\nACTOR_PARAM_OFFLOAD=${ACTOR_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nACTOR_GRAD_OFFLOAD=${ACTOR_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nACTOR_OPTIMIZER_OFFLOAD=${ACTOR_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nREF_PARAM_OFFLOAD=${REF_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_PARAM_OFFLOAD=${CRITIC_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nCRITIC_GRAD_OFFLOAD=${CRITIC_GRAD_OFFLOAD:-$COMMON_GRAD_OFFLOAD}\nCRITIC_OPTIMIZER_OFFLOAD=${CRITIC_OPTIMIZER_OFFLOAD:-$COMMON_OPTIMIZER_OFFLOAD}\nRM_PARAM_OFFLOAD=${RM_PARAM_OFFLOAD:-$COMMON_PARAM_OFFLOAD}\nUSE_MBRIDGE=${USE_MBRIDGE:-False}\nVANILLA_MBRIDGE=${VANILLA_MBRIDGE:-True}\nVALUE_VANILLA_MBRIDGE=${VALUE_VANILLA_MBRIDGE:-$VANILLA_MBRIDGE}\nUSE_FUSED_KERNELS=${USE_FUSED_KERNELS:-False}\n\nLR_WARMUP_STEPS=${LR_WARMUP_STEPS:-null}\n\nCHECKPOINT_CONTENTS=['model','hf_model','optimizer','extra']\nSKIP_SAVE_HF_MODEL=${SKIP_SAVE_HF_MODEL:-0}\nif [ $SKIP_SAVE_HF_MODEL -eq 1 ]; then\n    CHECKPOINT_CONTENTS=['model','optimizer','extra']\nfi\n\nUSE_DIST_CKPT=${USE_DIST_CKPT:-False}\nDIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/${MODEL_ID}}\nif [ \"$USE_DIST_CKPT\" = \"True\" ]; then\n    if [ \"$USE_DUMMY_MODEL\" = \"True\" ]; then\n        DIST_CKPT_PATH=${HOME}/dist_ckpt_dummy/${MODEL_ID}\n    fi\n    python scripts/converter_hf_to_mcore.py \\\n        --hf_model_path \"${MODEL_PATH}\" \\\n        --output_path \"${DIST_CKPT_PATH}\"\nfi\n\nENGINE=${ENGINE:-\"vllm\"}\nif [ \"$ENGINE\" = \"vllm\" ]; then\n    export VLLM_USE_V1=1\nfi\n\nexp_name=\"$(basename \"${MODEL_ID,,}\")-megatron-gsm8k-minimal\"\nROLLOUT_MODE=\"async\"\nROLLOUT_QUANTIZATION=${ROLLOUT_QUANTIZATION:-null}\n\nRETURN_RAW_CHAT=\"True\"\nSKIP_TOKENIZER_INIT=\"True\"\n\nOPTIM_MEMORY_EFFICIENT=${OPTIM_MEMORY_EFFICIENT:-False}\n\nPROFILE_ENABLE=${PROFILE_ENABLE:-False}\nPROFILE_STEPS=${PROFILE_STEPS:-[1]}\nPROFILE_RANKS_ALL=${PROFILE_RANKS_ALL:-True}\nPROFILE_RANKS=${PROFILE_RANKS:-[0,1,2,3]}\nDISCRETE=${DISCRETE:-True}  # or True\n\nUSE_LEGACY_WORKER_IMPL=${USE_LEGACY_WORKER_IMPL:-\"enable\"}\nUSE_REMOVE_PADDING=${USE_REMOVE_PADDING:-False}\nROUTING_REPLAY_MODE=${ROUTING_REPLAY_MODE:-\"disabled\"}\n\nif [ \"$ROUTING_REPLAY_MODE\" = \"R3\" ]; then\n    ENABLE_ROLLOUT_ROUTING_REPLAY=True\nelse\n    ENABLE_ROLLOUT_ROUTING_REPLAY=False\nfi\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=\"${ADV_ESTIMATOR}\" \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.max_prompt_length=${MAX_PROMPT_LENGTH} \\\n    data.max_response_length=${MAX_RESPONSE_LENGTH} \\\n    data.return_raw_chat=${RETURN_RAW_CHAT} \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_fused_kernels=${USE_FUSED_KERNELS} \\\n    actor_rollout_ref.model.use_remove_padding=${USE_REMOVE_PADDING} \\\n    actor_rollout_ref.model.lora.rank=${LORA_RANK} \\\n    actor_rollout_ref.model.lora.alpha=${LORA_ALPHA} \\\n    actor_rollout_ref.model.lora.target_modules=${LORA_TARGET_MODULES} \\\n    actor_rollout_ref.model.lora.merge=${LORA_MERGE} \\\n    +actor_rollout_ref.model.lora.fully_sharded_loras=True \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=$LR_WARMUP_STEPS \\\n    actor_rollout_ref.actor.megatron.router_replay.mode=${ROUTING_REPLAY_MODE} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=$OPTIM_MEMORY_EFFICIENT \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=$OPTIM_MEMORY_EFFICIENT \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=$OPTIM_MEMORY_EFFICIENT \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${USE_DYNAMIC_BSZ} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${ppo_max_token_len_per_gpu} \\\n    actor_rollout_ref.actor.megatron.use_mbridge=${USE_MBRIDGE} \\\n    actor_rollout_ref.actor.megatron.vanilla_mbridge=${VANILLA_MBRIDGE} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=$ACTOR_PP \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=$ACTOR_VPP \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=$ACTOR_CP \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=$ACTOR_TP \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=$ACTOR_EP \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=$ACTOR_ETP \\\n    actor_rollout_ref.actor.megatron.param_offload=${ACTOR_PARAM_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${ACTOR_OPTIMIZER_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${ACTOR_GRAD_OFFLOAD} \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.checkpoint.save_contents=$CHECKPOINT_CONTENTS \\\n    actor_rollout_ref.actor.profiler.enable=$PROFILE_ENABLE \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.rollout.name=\"${ENGINE}\" \\\n    actor_rollout_ref.rollout.mode=\"${ROLLOUT_MODE}\" \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$ROLLOUT_TP \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    ++actor_rollout_ref.rollout.quantization=${ROLLOUT_QUANTIZATION} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.rollout.enable_rollout_routing_replay=${ENABLE_ROLLOUT_ROUTING_REPLAY} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    actor_rollout_ref.ref.megatron.use_mbridge=${USE_MBRIDGE} \\\n    actor_rollout_ref.ref.megatron.vanilla_mbridge=${VANILLA_MBRIDGE} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=$REF_PP \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=$REF_VPP \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=$REF_CP \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=$REF_TP \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=$REF_EP \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=$REF_ETP \\\n    actor_rollout_ref.ref.megatron.param_offload=${REF_PARAM_OFFLOAD} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    critic.optim.lr=2e-5 \\\n    critic.optim.lr_warmup_steps=$LR_WARMUP_STEPS \\\n    +critic.optim.override_optimizer_config.optimizer_cpu_offload=$OPTIM_MEMORY_EFFICIENT \\\n    +critic.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=$OPTIM_MEMORY_EFFICIENT \\\n    +critic.optim.override_optimizer_config.use_precision_aware_optimizer=$OPTIM_MEMORY_EFFICIENT \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.model.lora.rank=${CRITIC_LORA_RANK} \\\n    critic.model.lora.alpha=${LORA_ALPHA} \\\n    critic.model.lora.target_modules=${LORA_TARGET_MODULES} \\\n    critic.ppo_micro_batch_size_per_gpu=${train_traj_micro_bsz_per_gpu} \\\n    critic.ppo_max_token_len_per_gpu=${forward_max_token_len_per_gpu} \\\n    critic.megatron.use_mbridge=${USE_MBRIDGE} \\\n    critic.megatron.vanilla_mbridge=${VALUE_VANILLA_MBRIDGE} \\\n    critic.megatron.pipeline_model_parallel_size=$CRITIC_PP \\\n    critic.megatron.virtual_pipeline_model_parallel_size=$CRITIC_VPP \\\n    critic.megatron.context_parallel_size=$CRITIC_CP \\\n    critic.megatron.tensor_model_parallel_size=$CRITIC_TP \\\n    critic.megatron.expert_model_parallel_size=$CRITIC_EP \\\n    critic.megatron.expert_tensor_parallel_size=$CRITIC_ETP \\\n    critic.megatron.param_offload=${CRITIC_PARAM_OFFLOAD} \\\n    critic.megatron.optimizer_offload=${CRITIC_OPTIMIZER_OFFLOAD} \\\n    critic.megatron.grad_offload=${CRITIC_GRAD_OFFLOAD} \\\n    critic.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    critic.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    critic.checkpoint.save_contents=$CHECKPOINT_CONTENTS \\\n    critic.profiler.enable=$PROFILE_ENABLE \\\n    critic.profiler.ranks=$PROFILE_RANKS \\\n    critic.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    reward.num_workers=8 \\\n    reward.reward_model.enable=True \\\n    reward.reward_model.model_path=\"${RM_MODEL_PATH}\" \\\n    reward.reward_model.rollout.name=${ENGINE} \\\n    reward.reward_model.rollout.gpu_memory_utilization=0.6 \\\n    reward.reward_model.rollout.tensor_model_parallel_size=${INFER_TP} \\\n    reward.reward_model.rollout.prompt_length=${MAX_RM_LENGTH} \\\n    reward.reward_model.rollout.response_length=${MAX_RESPONSE_LENGTH} \\\n    algorithm.use_kl_in_reward=False \\\n    algorithm.kl_penalty=kl \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl-test' \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=${NUM_GPUS} \\\n    trainer.val_before_train=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.test_freq=\"${TEST_FREQ}\" \\\n    trainer.save_freq=\"${SAVE_FREQ}\" \\\n    trainer.resume_mode=\"${RESUME_MODE}\" \\\n    trainer.total_epochs=2 \\\n    trainer.total_training_steps=\"${TOTAL_TRAIN_STEPS}\" \\\n    trainer.use_legacy_worker_impl=${USE_LEGACY_WORKER_IMPL} \\\n    global_profiler.profile_continuous_steps=True \\\n    global_profiler.tool=nsys \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.global_tool_config.nsys.discrete=$DISCRETE $@\n"
  },
  {
    "path": "tests/special_e2e/run_ppo_trainer_torchtitan.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n# Download model if not exists\nMODEL_ID=${MODEL_ID:-Qwen/Qwen3-0.6B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#huggingface-cli download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nVAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-False}\nNUM_GPUS=${NUM_GPUS:-1}\nFSDP_SIZE=${FSDP_SIZE:-1}\nTP_SIZE=${TP_SIZE:-1}\nCP_SIZE=${CP_SIZE:-1}\nEP_SIZE=${EP_SIZE:-1}\nVERL_EXP_NAME=${VERL_EXP_NAME:-Titan_Qwen3_30B_A3B_DP8_EP8}\nMAX_PROMPT_LENGTH=${MAX_PROMPT_LENGTH:-512}\nMAX_RESPONSE_LENGTH=${MAX_RESPONSE_LENGTH:-2048}\nMAX_SEQ_LEN=${MAX_SEQ_LEN:-$((MAX_PROMPT_LENGTH + MAX_RESPONSE_LENGTH))}\n\npython3 -m verl.trainer.main_ppo \\\n    model_engine=torchtitan \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=8  \\\n    data.max_prompt_length=\"${MAX_PROMPT_LENGTH}\" \\\n    data.max_response_length=\"${MAX_RESPONSE_LENGTH}\"  \\\n    data.seed=42 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.min_lr_factor=1.0 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=4 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1  \\\n    actor_rollout_ref.actor.torchtitan.data_parallel_shard_size=\"${FSDP_SIZE}\" \\\n    actor_rollout_ref.actor.torchtitan.tensor_parallel_size=\"${TP_SIZE}\" \\\n    actor_rollout_ref.actor.torchtitan.context_parallel_size=\"${CP_SIZE}\" \\\n    actor_rollout_ref.actor.torchtitan.expert_parallel_size=\"${EP_SIZE}\" \\\n    actor_rollout_ref.actor.torchtitan.attn_type=flex \\\n    actor_rollout_ref.actor.torchtitan.use_torch_compile=False \\\n    actor_rollout_ref.actor.torchtitan.param_offload=True \\\n    actor_rollout_ref.actor.torchtitan.optimizer_offload=True \\\n    actor_rollout_ref.actor.torchtitan.max_seq_len=\"${MAX_SEQ_LEN}\" \\\n    actor_rollout_ref.ref.torchtitan.max_seq_len=\"${MAX_SEQ_LEN}\" \\\n    actor_rollout_ref.ref.torchtitan.use_torch_compile=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=8 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.35 \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.max_model_len=\"${MAX_SEQ_LEN}\" \\\n    critic.optim.lr=1e-5 \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.ppo_micro_batch_size_per_gpu=2 \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.logger=['console','file','wandb'] \\\n    trainer.project_name='verl_grpo_example_gsm8k_0302' \\\n    trainer.experiment_name=\"${VERL_EXP_NAME}\" \\\n    trainer.val_before_train=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.log_val_generations=1 \\\n    trainer.test_freq=1 \\\n    trainer.n_gpus_per_node=\"${NUM_GPUS}\" \\\n    trainer.nnodes=1 \\\n    trainer.total_training_steps=100 $@\n"
  },
  {
    "path": "tests/special_e2e/run_ppo_trainer_veomni.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n\nSAVE_PATH=tests/utils/ci/profiler_data\nrm -rf \"$SAVE_PATH\"\n\nCONTENTS=['cuda']\nPROFILE_STEPS=[1]\nPROFILE_RANKS_ALL=False\nPROFILE_RANKS=[0]\nDISCRETE=True\n\n# Download model if not exists\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#huggingface-cli download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nTRAIN_FILES=${TRAIN_FILES:-${HOME}/data/gsm8k/train.parquet}\nVAL_FILES=${VAL_FILES:-${HOME}/data/gsm8k/test.parquet}\nVAL_BEFORE_TRAIN=${VAL_BEFORE_TRAIN:-True}\nNUM_GPUS=${NUM_GPUS:-8}\nFSDP_SIZE=${FSDP_SIZE:-4}\nSP_SIZE=${SP_SIZE:-2}\nEP_SIZE=${EP_SIZE:-1}\nMODEL_NAME_ONLY=${MODEL_ID##*/}\nVERL_EXP_NAME=${VERL_EXP_NAME:-${MODEL_NAME_ONLY}-function-reward-minimal-fsdp-size${FSDP_SIZE}}\n\npython3 -m verl.trainer.main_ppo \\\n    model_engine=veomni \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=128 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=5e-7 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.veomni.param_offload=True \\\n    actor_rollout_ref.actor.veomni.optimizer_offload=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.actor.veomni.fsdp_size=\"${FSDP_SIZE}\" \\\n    actor_rollout_ref.actor.veomni.ulysses_parallel_size=\"${SP_SIZE}\" \\\n    actor_rollout_ref.actor.veomni.expert_parallel_size=\"${EP_SIZE}\" \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.veomni.param_offload=True \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.ref.veomni.optimizer_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_veomni_test' \\\n    trainer.experiment_name=\"${VERL_EXP_NAME}\" \\\n    trainer.n_gpus_per_node=\"${NUM_GPUS}\" \\\n    trainer.val_before_train=\"${VAL_BEFORE_TRAIN}\" \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.tool_config.torch.discrete=$DISCRETE \\\n    actor_rollout_ref.actor.profiler.tool_config.torch.contents=$CONTENTS \\\n    actor_rollout_ref.ref.profiler.enable=True \\\n    actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.ref.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.ref.profiler.tool_config.torch.discrete=$DISCRETE \\\n    actor_rollout_ref.ref.profiler.tool_config.torch.contents=$CONTENTS \\\n    global_profiler.tool=torch \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.save_path=\"$SAVE_PATH\" $@\n\npython3 \"tests/utils/test_check_profiler_output.py\" --profiler_dir=\"$SAVE_PATH\" --device=\"gpu\"\nrm -rf \"$SAVE_PATH\"\n"
  },
  {
    "path": "tests/special_e2e/run_test.sh",
    "content": "#!/bin/bash\nset -xeuo pipefail\n\n# Get the configuration name and engine name from arguments\nCONFIG_NAME=\"$1\"\nENGINE=\"${2:-vllm}\"\n\n# Download model if needed\n#hf download Qwen/Qwen2.5-0.5B --local-dir \"$HOME/models/Qwen/Qwen2.5-0.5B\"\n\n# Run the training with the specified configuration\npython3 -m verl.trainer.main_ppo \\\n    --config-name \"$CONFIG_NAME\" \"$@\" "
  },
  {
    "path": "tests/special_e2e/sft/compare_sft_engine_results.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\n\nimport torch\n\n\ndef get_result(file):\n    file = os.path.expanduser(file)\n    result = []\n    with open(file) as f:\n        lines = f.readlines()\n        for line in lines:\n            result.append(json.loads(line))\n    return result\n\n\ndef compare_results(golden_results, other_result):\n    golden_loss = golden_results[0][\"data\"][\"train/loss\"]\n    golden_grad_norm = golden_results[0][\"data\"][\"train/grad_norm\"]\n\n    loss = other_result[0][\"data\"][\"train/loss\"]\n    grad_norm = other_result[0][\"data\"][\"train/grad_norm\"]\n\n    torch.testing.assert_close(golden_loss, loss, atol=1e-2, rtol=1e-2)\n    torch.testing.assert_close(golden_grad_norm, grad_norm, atol=1e-4, rtol=3e-2)\n\n\nif __name__ == \"__main__\":\n    golden_results = get_result(\"~/verl/test/log/golden.jsonl\")\n\n    # get all other results\n    other_results = {}\n    # walk through all files in ~/verl/test/log\n    for file in os.listdir(os.path.expanduser(\"~/verl/test/log/verl_sft_test\")):\n        if file.endswith(\".jsonl\"):\n            other_results[file] = get_result(os.path.join(os.path.expanduser(\"~/verl/test/log/verl_sft_test\"), file))\n\n    # # compare results\n    for file, other_result in other_results.items():\n        print(f\"compare results {file}\")\n        compare_results(golden_results, other_result)\n        print(f\"compare results {file} done\")\n\n    print(\"All results are close to golden results\")\n"
  },
  {
    "path": "tests/special_e2e/sft/run_sft.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n\nNUM_GPUS=${NUM_GPUS:-8}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nTRAIN_FILES=${TRAIN_FILES:-$HOME/data/gsm8k_sft/train.parquet}\nVAL_FILES=${VAL_FILES:-$HOME/data/gsm8k_sft/test.parquet}\n\nSP_SIZE=${SP_SIZE:-1}\nLIGER=${LIGER:-False}\nMULTITURN=${MULTITURN:-False}\nLORA_RANK=${LORA_RANK:-0}\nRM_PAD=${RM_PAD:-True}\n\nTOTAL_TRAIN_STEP=${TOTAL_TRAIN_STEP:-1}\nRESUME_MODE=${RESUME_MODE:-disable}\nSAVE_FREQ=${SAVE_FREQ:-1}\n\nmicro_bsz=2\nNUM_GPUS=8\n\nproject_name=\"verl-test\"\nexp_name=\"$(basename \"${MODEL_ID,,}\")-sft-minimal\"\nckpts_home=${ckpts_home:-$HOME/${project_name}/${exp_name}}\n\nmkdir -p \"${ckpts_home}\"\n\ntorchrun --standalone --nnodes=1 --nproc_per_node=${NUM_GPUS} ${ENTRYPOINT} \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.messages_key=messages \\\n    data.micro_batch_size_per_gpu=${micro_bsz} \\\n    optim.lr=1e-4 \\\n    engine=fsdp \\\n    engine.ulysses_sequence_parallel_size=\"${SP_SIZE}\" \\\n    model.path=\"${MODEL_PATH}\" \\\n    model.lora_rank=\"${LORA_RANK}\" \\\n    model.lora_alpha=16 \\\n    model.target_modules=all-linear \\\n    model.use_liger=\"${LIGER}\" \\\n    model.use_remove_padding=\"${RM_PAD}\" \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_training_steps=${TOTAL_TRAIN_STEP} \\\n    trainer.save_freq=${SAVE_FREQ} \\\n    checkpoint.save_contents=[model,optimizer,extra,hf_model] \\\n    trainer.max_ckpt_to_keep=1 \\\n    trainer.resume_mode=${RESUME_MODE} \\\n    trainer.logger=['console'] $@\n\nrm -rf \"${ckpts_home:?}/*\""
  },
  {
    "path": "tests/special_e2e/sft/run_sft_engine.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nNUM_GPUS=${NUM_GPUS:-1}\n\nmode=${mode:-spmd}\n\nif [ \"$mode\" = \"spmd\" ]; then\n  ENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n  COMMAND=\"torchrun --standalone --nnodes=${NNODES:-1} --nproc-per-node=${NUM_GPUS:-1} ${ENTRYPOINT}\"\nelse\n  ENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer_ray\"}\n  COMMAND=\"python ${ENTRYPOINT} trainer.nnodes=${NNODES:-1} trainer.n_gpus_per_node=${NUM_GPUS:-1}\"\nfi\n\nDATASET_DIR=${DATASET_DIR:-~/data/gsm8k_sft}\nTRAIN_FILES=${DATASET_DIR}/train.parquet\nVAL_FILES=${DATASET_DIR}/test.parquet\n\nbackend=${BACKEND:-fsdp}\n\nproject_name=verl_sft_test\n\nRESUME_MODE=disable\n\nckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-${backend}}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/models/${MODEL_ID}}\n#hf download \"${MODEL_ID}\" --local-dir \"${MODEL_PATH}\"\n\nSP_SIZE=${SP_SIZE:-1}\nFSDP_SIZE=${FSDP_SIZE:-1}\nFSDP_STRATEGY=${FSDP_STRATEGY:-\"fsdp\"}\n\nTP_SIZE=${TP_SIZE:-1}\nPP_SIZE=${PP_SIZE:-1}\nVPP_SIZE=${VPP_SIZE:-null}\nCP_SIZE=${CP_SIZE:-1}\n\nPAD_MODE=${PAD_MODE:-no_padding}\n\nUSE_REMOVE_PADDING=${USE_REMOVE_PADDING:-True}\n\nFSDP_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    model=hf_model \\\n    model.path=$MODEL_PATH \\\n    optim=${backend} \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.min_lr_ratio=0.1 \\\n    optim.lr_scheduler_type=cosine \\\n    engine.ulysses_sequence_parallel_size=${SP_SIZE} \\\n    engine.strategy=${FSDP_STRATEGY} \\\n    engine.fsdp_size=${FSDP_SIZE}\"\n\nVEOMNI_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    model=hf_model \\\n    model.path=$MODEL_PATH \\\n    optim=${backend} \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.lr_min=1e-6 \\\n    optim.lr_scheduler_type=cosine \\\n    engine.ulysses_parallel_size=${SP_SIZE} \\\n    engine.fsdp_size=${FSDP_SIZE}\"\n\nMEGATRON_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    model=hf_model \\\n    model.path=$MODEL_PATH \\\n    optim=${backend} \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.lr_warmup_init=0 \\\n    optim.lr_decay_style=cosine \\\n    optim.min_lr=1e-6 \\\n    engine.tensor_model_parallel_size=${TP_SIZE} \\\n    engine.pipeline_model_parallel_size=${PP_SIZE} \\\n    engine.virtual_pipeline_model_parallel_size=${VPP_SIZE} \\\n    engine.context_parallel_size=${CP_SIZE} \\\n    +engine.override_transformer_config.context_parallel_size=${CP_SIZE} \\\n    engine.use_mbridge=True\"\n\nTORCHTITAN_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    model=hf_model \\\n    model.path=${MODEL_PATH} \\\n    optim=${backend} \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.min_lr_factor=0.1 \\\n    optim.decay_type=cosine \\\n    optim.total_training_steps=1000 \\\n    engine.tensor_parallel_size=${TP_SIZE} \\\n    engine.pipeline_parallel_size=${PP_SIZE} \\\n    engine.context_parallel_size=${CP_SIZE} \\\n    engine.data_parallel_shard_size=${FSDP_SIZE} \\\n    engine.use_torch_compile=False\"\n\nAUTOMODEL_ENGINE_CONFIG=\"\\\n    engine=${backend} \\\n    model=hf_model \\\n    model.path=${MODEL_PATH} \\\n    optim=${backend} \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.min_lr_ratio=0.1 \\\n    optim.lr_scheduler_type=cosine \\\n    engine.tp_size=${TP_SIZE} \\\n    engine.cp_size=${CP_SIZE} \\\n    engine.use_torch_compile=False\"\n\n\nif [ \"$backend\" = \"fsdp\" ]; then\n    ENGINE_CONFIG=\"$FSDP_ENGINE_CONFIG\"\n    echo \"Using fsdp engine\"\n    exp_name=gsm8k-${backend}-${FSDP_STRATEGY}-sp${SP_SIZE}-fsdp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode}\nelif [ \"$backend\" = \"veomni\" ]; then\n    ENGINE_CONFIG=\"$VEOMNI_ENGINE_CONFIG\"\n    echo \"Using veomni engine\"\n    exp_name=gsm8k-${backend}-sp${SP_SIZE}-fsdp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode}\nelif [ \"$backend\" = \"torchtitan\" ]; then\n    ENGINE_CONFIG=\"$TORCHTITAN_ENGINE_CONFIG\"\n    echo \"Using torchtitan engine\"\n    exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-dp${FSDP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode}\nelif [ \"$backend\" = \"automodel\" ]; then\n    ENGINE_CONFIG=\"$AUTOMODEL_ENGINE_CONFIG\"\n    echo \"Using automodel engine\"\n    exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-cp${CP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode}\nelse\n    ENGINE_CONFIG=\"$MEGATRON_ENGINE_CONFIG\"\n    echo \"Using megatron engine\"\n    exp_name=gsm8k-${backend}-tp${TP_SIZE}-pp${PP_SIZE}-vpp${VPP_SIZE}-cp${CP_SIZE}-pad-${PAD_MODE}-use_remove_padding-${USE_REMOVE_PADDING}-mode-${mode}\nfi\n\nmkdir -p \"${ckpts_home}\"\n\n$COMMAND \\\n    data.train_files=\"${TRAIN_FILES}\" \\\n    data.val_files=\"${VAL_FILES}\" \\\n    data.train_batch_size=128 \\\n    data.pad_mode=${PAD_MODE} \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=2048 \\\n    data.messages_key=messages \\\n    model.use_remove_padding=${USE_REMOVE_PADDING} \\\n    data.ignore_input_ids_mismatch=True \\\n    ${ENGINE_CONFIG} \\\n    trainer.test_freq=after_each_epoch \\\n    trainer.save_freq=-1 \\\n    trainer.logger=['console','file'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.total_epochs=2 \\\n    trainer.total_training_steps=2 \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.resume_mode=${RESUME_MODE} \\\n\n    # trainer.total_training_steps=${TOTAL_TRAIN_STEP} \\\n    # trainer.checkpoint.save_contents=[model,optimizer,extra,hf_model] \\\n    # trainer.max_ckpt_to_keep=1 \\\n\nrm -rf \"${ckpts_home:?}/*\"\n"
  },
  {
    "path": "tests/special_e2e/sft/test_sft_engine_all.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nrm -rf ~/verl/test/log\nmkdir -p ~/verl/test/log\n\nexport VERL_FILE_LOGGER_ROOT=~/verl/test/log\nVPP_SIZE=${VPP_SIZE:-2}\n\n# test with single gpu as golden\necho \"run with single gpu as golden\"\nBACKEND=fsdp SP_SIZE=1 FSDP_SIZE=1 NUM_GPUS=1 FSDP_STRATEGY=fsdp VERL_FILE_LOGGER_PATH=~/verl/test/log/golden.jsonl bash tests/special_e2e/sft/run_sft_engine.sh\n\n# test with fsdp 1\necho \"run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp pad_mode no_padding\"\nBACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding bash tests/special_e2e/sft/run_sft_engine.sh\n\n# test with fsdp 1 use_remove_padding and pad_mode no_padding\necho \"run with sp4 fsdp_size4 num_gpus8 fsdp_strategy fsdp pad_mode no_padding use_remove_padding False\"\nBACKEND=fsdp SP_SIZE=1 FSDP_SIZE=-1 NUM_GPUS=8 FSDP_STRATEGY=fsdp PAD_MODE=no_padding USE_REMOVE_PADDING=False bash tests/special_e2e/sft/run_sft_engine.sh\n\n\n# test with fsdp 2\necho \"run with sp2 fsdp_size2 num_gpus8 fsdp_strategy fsdp2\"\nBACKEND=fsdp SP_SIZE=2 FSDP_SIZE=2 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh\n\n# test with veomni\necho \"run with sp2 fsdp_size4 num_gpus8 fsdp_strategy fsdp2\"\nBACKEND=veomni SP_SIZE=2 FSDP_SIZE=4 NUM_GPUS=8 FSDP_STRATEGY=fsdp2 bash tests/special_e2e/sft/run_sft_engine.sh\n\n\n# test with megatron\necho \"run with tp2 pp2 vpp2 cp2 num_gpus8\"\nBACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=${VPP_SIZE} CP_SIZE=2 NUM_GPUS=8 bash tests/special_e2e/sft/run_sft_engine.sh\n\n# test with cp in ray\necho \"run with tp2 pp2 vpp2 cp2 num_gpus8 mode=ray\"\nBACKEND=megatron TP_SIZE=2 PP_SIZE=2 VPP_SIZE=${VPP_SIZE} CP_SIZE=2 NUM_GPUS=8 mode=ray bash tests/special_e2e/sft/run_sft_engine.sh\n\n# TODO: Will add back torchtitan CI once everything is ready\n# # test with torchtitan fsdp=2\n# echo \"run with tp1 pp1 cp1 fsdp2 num_gpus2\"\n# BACKEND=torchtitan TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=2 bash tests/special_e2e/sft/run_sft_engine.sh\n\n# # test with torchtitan tp2 fsdp=2\n# echo \"run with tp2 pp1 cp1 fsdp2 num_gpus4\"\n# BACKEND=torchtitan TP_SIZE=2 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=4 bash tests/special_e2e/sft/run_sft_engine.sh\n\n# # test with automodel dp=2\n# echo \"run with automodel tp1 pp1 cp1 dp2 num_gpus2\"\n# BACKEND=automodel TP_SIZE=1 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=2 bash tests/special_e2e/sft/run_sft_engine.sh\n\n# # test with automodel tp2 dp=2\n# echo \"run with automodel tp2 pp1 cp1 dp2 num_gpus4\"\n# BACKEND=automodel TP_SIZE=2 PP_SIZE=1 CP_SIZE=1 FSDP_SIZE=2 NUM_GPUS=4 bash tests/special_e2e/sft/run_sft_engine.sh\n\npython3 tests/special_e2e/sft/compare_sft_engine_results.py\n\nrm -rf ~/verl/test/log\n"
  },
  {
    "path": "tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-7b-instruct_fsdp_npu.sh",
    "content": "set -x\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-7B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=5e-8 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_5_7b_instruct_fsdp' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_training_steps=15 2>&1 | tee /root/.cache/grpo_qwen25-7b-instruct_fsdp_npu.log"
  },
  {
    "path": "tests/special_npu/nightly_ci_ascend/run_grpo_qwen25-vl-3b-instruct_fsdp_npu.sh",
    "content": "set -x\nENGINE=${1:-vllm}\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-VL-3B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=16 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.use_legacy_worker_impl=disable \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_3b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_training_steps=15  2>&1 | tee /root/.cache/grpo_qwen25-vl-3b-instruct_fsdp_npu.log"
  },
  {
    "path": "tests/special_npu/nightly_ci_ascend/run_ppo_qwen3-8b_fsdp_npu.sh",
    "content": "set -x\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen3-8B}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=2000 \\\n    data.max_response_length=2000 \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=4000 \\\n    actor_rollout_ref.rollout.max_num_seqs=64 \\\n    actor_rollout_ref.rollout.checkpoint_engine.update_weights_bucket_megabytes=4096 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=1 \\\n    critic.ulysses_sequence_parallel_size=2 \\\n    critic.model.fsdp_config.param_offload=True \\\n    critic.model.fsdp_config.optimizer_offload=True \\\n    critic.use_dynamic_bsz=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_example_ppo_gsm8k' \\\n    trainer.experiment_name='qwen3_8b_fsdp' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.val_before_train=False \\\n    trainer.max_actor_ckpt_to_keep=1 \\\n    trainer.max_critic_ckpt_to_keep=1 \\\n    trainer.total_training_steps=15 2>&1 | tee /root/.cache/ppo_qwen3-8b_fsdp_npu.log"
  },
  {
    "path": "tests/special_npu/run_qwen2_5_05b_grpo.sh",
    "content": "set -x\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\nSAVE_PATH=tests/utils/ci/profiler_data\nrm -rf \"$SAVE_PATH\"\n\nLEVEL=\"level0\"\nCONTENTS=['npu','cpu']\nANALYSIS=False\nPROFILE_STEPS=[1]\nPROFILE_RANKS_ALL=False\nPROFILE_RANKS=[0]\nDISCRETE=True\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=128 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=5e-7 \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_AND_PIECEWISE\" \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 \\\n    actor_rollout_ref.actor.profiler.enable=True \\\n    actor_rollout_ref.actor.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.actor.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.actor.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    actor_rollout_ref.ref.profiler.enable=True \\\n    actor_rollout_ref.ref.profiler.all_ranks=$PROFILE_RANKS_ALL \\\n    actor_rollout_ref.ref.profiler.ranks=$PROFILE_RANKS \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.discrete=$DISCRETE \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.contents=$CONTENTS \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.level=$LEVEL \\\n    actor_rollout_ref.ref.profiler.tool_config.npu.analysis=$ANALYSIS \\\n    global_profiler.tool=npu \\\n    global_profiler.steps=$PROFILE_STEPS \\\n    global_profiler.save_path=\"$SAVE_PATH\" $@\n\npython3 \"tests/utils/test_check_profiler_output.py\" --profiler_dir=\"$SAVE_PATH\" --device=\"npu\"\nrm -rf \"$SAVE_PATH\"\n"
  },
  {
    "path": "tests/special_npu/run_qwen2_5_05b_grpo_mindspeed.sh",
    "content": "set -x\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\nUSE_DIST_CKPT=${USE_DIST_CKPT:-False}\nDIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/qwen2_5_05b_grpo_mindspeed}\nif [ \"$USE_DIST_CKPT\" = \"True\" ]; then\n    if [ \"$USE_DUMMY_MODEL\" = \"True\" ]; then\n        DIST_CKPT_PATH=${HOME}/dist_ckpt_dummy/${MODEL_ID}\n    fi\n    python scripts/converter_hf_to_mcore.py \\\n        --hf_model_path \"${MODEL_PATH}\" \\\n        --output_path \"${DIST_CKPT_PATH}\"\nfi\n\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=128 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.model.path=${MODEL_PATH} \\\n    actor_rollout_ref.actor.optim.lr=5e-7 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=1 \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_AND_PIECEWISE\" \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.strategy=megatron \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=1 \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=True \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    algorithm.kl_ctrl.kl_coef=0.001 \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_gsm8k' \\\n    trainer.experiment_name='qwen2_7b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True $@\n"
  },
  {
    "path": "tests/special_npu/run_qwen2_5_05b_sft_peft_sp2.sh",
    "content": "set -x\n\nNUM_GPUS=${NUM_GPUS:-4}\n\nmode=${mode:-spmd}\n\nif [ \"$mode\" = \"spmd\" ]; then\n  ENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer\"}\n  COMMAND=\"torchrun --standalone --nnodes=${NNODES:-1} --nproc-per-node=${NUM_GPUS:-1} ${ENTRYPOINT}\"\nelse\n  ENTRYPOINT=${ENTRYPOINT:-\"-m verl.trainer.sft_trainer_ray\"}\n  COMMAND=\"python ${ENTRYPOINT} trainer.nnodes=${NNODES:-1} trainer.n_gpus_per_node=${NUM_GPUS:-1}\"\nfi\n\nRESUME_MODE=disable\n\nckpts_home=${ckpts_home:-~/verl/test/gsm8k-sft-fsdp}\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\nDATASET_DIR=${DATASET_DIR:-$HOME/data/gsm8k_sft}\nTRAIN_FILES=${DATASET_DIR}/train.parquet\nVAL_FILES=${DATASET_DIR}/test.parquet\n\nexp_name=gsm8k-sft-qwen-2.5-0.5b-instruct-mode-${mode}\n\nmkdir -p \"${ckpts_home}\"\n\n$COMMAND \\\n    data.train_files=$TRAIN_FILES \\\n    data.val_files=$VAL_FILES \\\n    data.pad_mode=no_padding \\\n    data.truncation=error \\\n    data.use_dynamic_bsz=True \\\n    data.max_token_len_per_gpu=2048 \\\n    data.messages_key=messages \\\n    model.path=$MODEL_PATH \\\n    model.use_remove_padding=True \\\n    model.lora_rank=32 \\\n    model.lora_alpha=16 \\\n    model.target_modules=all-linear \\\n    engine=fsdp \\\n    optim=fsdp \\\n    optim.lr=1e-5 \\\n    optim.lr_warmup_steps_ratio=0.2 \\\n    optim.weight_decay=0.1 \\\n    optim.betas=\"[0.9,0.95]\" \\\n    optim.clip_grad=1.0 \\\n    optim.min_lr_ratio=0.1 \\\n    optim.lr_scheduler_type=cosine \\\n    engine.ulysses_sequence_parallel_size=2 \\\n    engine.strategy=fsdp2 \\\n    engine.fsdp_size=2 \\\n    trainer.test_freq=after_each_epoch \\\n    trainer.save_freq=-1 \\\n    trainer.logger=['console','file'] \\\n    trainer.project_name=gsm8k-sft \\\n    trainer.experiment_name=gsm8k-sft-qwen-2.5-0.5b-instruct \\\n    trainer.total_epochs=2 \\\n    trainer.total_training_steps=2 \\\n    trainer.default_local_dir=\"${ckpts_home}\" \\\n    trainer.resume_mode=${RESUME_MODE} \\\n\nrm -rf \"${ckpts_home:?}/*\"\n"
  },
  {
    "path": "tests/special_npu/run_qwen2_5_vl_3b_npu.sh",
    "content": "set -x\n\nENGINE=${1:-vllm}\n\n# Some models are optimized by vllm ascend. While in some case, e.g. rlhf training, \n# the optimized model may not be suitable. In this case, set this value to 0 to disable the optimized model.\nexport USE_OPTIMIZED_MODEL=0\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-VL-3B-Instruct}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=$HOME/data/geo3k/train.parquet \\\n    data.val_files=$HOME/data/geo3k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.image_key=images \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_AND_PIECEWISE\" \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger=console \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_3b_function_rm' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 $@"
  },
  {
    "path": "tests/special_npu/run_qwen3_06b_ppo.sh",
    "content": "set -x\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen2.5-0.5B-Instruct}  # TODO: change to Qwen3-0.6B when CI server is ready\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\n\npython3 -m verl.trainer.main_ppo \\\n    algorithm.adv_estimator=gae \\\n    data.train_files=$HOME/data/gsm8k/train.parquet \\\n    data.val_files=$HOME/data/gsm8k/test.parquet \\\n    data.train_batch_size=16 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=128 \\\n    data.shuffle=False \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.enforce_eager=False \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_AND_PIECEWISE\" \\\n    critic.optim.lr=1e-5 \\\n    critic.model.use_remove_padding=True \\\n    critic.model.path=\"${MODEL_PATH}\" \\\n    critic.model.enable_gradient_checkpointing=True \\\n    critic.ppo_micro_batch_size_per_gpu=1 \\\n    critic.ulysses_sequence_parallel_size=2 \\\n    critic.model.fsdp_config.param_offload=True \\\n    critic.model.fsdp_config.optimizer_offload=True \\\n    critic.use_dynamic_bsz=True \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\"]' \\\n    trainer.project_name='verl_ppo_example_gsm8k_qwen3' \\\n    trainer.experiment_name='qwen3_06b_fsdp' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 $@\n"
  },
  {
    "path": "tests/special_npu/run_qwen3_30b_grpo_mindspeed.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\n\nMODEL_ID=${MODEL_ID:-Qwen/Qwen3-30B-A3B-Instruct-2507}\nMODEL_PATH=${MODEL_PATH:-${HOME}/.cache/models/${MODEL_ID}}\nUSE_DIST_CKPT=${USE_DIST_CKPT:-False}\nDIST_CKPT_PATH=${DIST_CKPT_PATH:-${HOME}/dist_ckpt/qwen3_30b_grpo_mindspeed}\n\n# use dummy model\nif [[ \"$USE_DUMMY_MODEL\" == \"True\" ]]; then\n    DUMMY_MODEL_PATH=${DUMMY_MODEL_PATH:-${HOME}/models_dummy/${MODEL_ID}}\n    if [ -z \"${DUMMY_MODEL_CONFIG_PATH}\" ]; then\n        echo \"[ERROR] DUMMY_MODEL_CONFIG_PATH not set\"\n        exit 1\n    fi\n\n    # make sure the path is empty\n    if [[ -d $DUMMY_MODEL_PATH && $DUMMY_MODEL_PATH != \"/\" ]]; then\n        rm -rf $DUMMY_MODEL_PATH\n    fi\n\n    # init model\n    python scripts/init_random_model.py \\\n        --hf_model_path \"${MODEL_PATH}\" \\\n        --new_config_path \"${DUMMY_MODEL_CONFIG_PATH}\" \\\n        --output_path \"${DUMMY_MODEL_PATH}\"\n\n    # replace model path\n    MODEL_PATH=$DUMMY_MODEL_PATH\nfi\n\n# convert to megatron\nif [[ \"$USE_DIST_CKPT\" == \"True\" ]]; then\n\n    if [[ \"$USE_DUMMY_MODEL\" == \"True\" ]]; then\n        DIST_CKPT_PATH=${HOME}/dist_ckpt/qwen3_30b_grpo_mindspeed_dummy\n\n        if [[ -d $DIST_CKPT_PATH && $DIST_CKPT_PATH != \"/\" ]];then\n            rm -rf $DIST_CKPT_PATH\n        fi\n    fi\n\n    torchrun --nproc_per_node 2 --nnodes 1 scripts/converter_hf_to_mcore.py \\\n        --hf_model_path \"${MODEL_PATH}\" \\\n        --output_path \"${DIST_CKPT_PATH}\"\nfi\n\nexp_name='Qwen3-30B-A3B-GRPO-MindSpeed'\n\nmax_prompt_length=512\nmax_response_length=1024\n\ntrain_prompt_bsz=16\n\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\n\npython3 -m verl.trainer.main_ppo --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=${HOME}/data/gsm8k/train.parquet \\\n    data.val_files=${HOME}/data/gsm8k/test.parquet \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.filter_overlong_prompts=True \\\n    data.shuffle=False \\\n    data.truncation='left' \\\n    algorithm.adv_estimator=grpo \\\n    algorithm.use_kl_in_reward=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=False \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=2 \\\n    actor_rollout_ref.rollout.temperature=1.0 \\\n    actor_rollout_ref.rollout.top_p=1.0 \\\n    actor_rollout_ref.rollout.top_k=-1 \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    actor_rollout_ref.actor.kl_loss_coef=0.0 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.ppo_epochs=1 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=8 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.compilation_config.cudagraph_mode=\"FULL_AND_PIECEWISE\" \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=2 \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.actor.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.loss_agg_mode=\"token-mean\" \\\n    actor_rollout_ref.ref.strategy=megatron \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=2 \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.ref.megatron.dist_checkpointing_path=${DIST_CKPT_PATH} \\\n    reward.reward_manager.name=naive \\\n    algorithm.kl_ctrl.kl_coef=0.0 \\\n    trainer.logger=['console'] \\\n    trainer.project_name='verl_gsm8k_example' \\\n    trainer.experiment_name='qwen3_30b_a3b_cut_gsm8k_mindspeed' \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=1 \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=1 \\\n    trainer.total_training_steps=1 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.use_flash_attn=True $@\n\n# clean up\nif [[ \"$USE_DUMMY_MODEL\" == \"True\" ]]; then\n    rm -rf $DUMMY_MODEL_PATH\n    if [[ \"$USE_DIST_CKPT\" == \"True\" ]]; then\n        rm -rf $DIST_CKPT_PATH\n    fi\nfi\n"
  },
  {
    "path": "tests/special_sanity/check_api_docs.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nFail CI if any function or class that is publicly exported via\n``__all__`` lacks a docstring.\n\nUsage\n-----\n  # Check specific modules or packages\n  python check_docstrings.py mypkg.core mypkg.utils\n\n  # Check an entire source tree (all top-level packages under cwd)\n  python check_docstrings.py\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport importlib\nimport inspect\nimport pkgutil\nimport sys\nfrom pathlib import Path\nfrom types import ModuleType\nfrom typing import Iterable\n\n_ALLOW_LIST = [\n    \"verl.third_party.vllm.LLM\",\n    \"verl.third_party.vllm.parallel_state\",\n    \"verl.utils.profiler.WorkerProfiler\",\n    \"verl.utils.profiler.WorkerProfilerExtension\",\n    \"verl.utils.profiler.log_gpu_memory_usage\",\n    \"verl.utils.profiler.log_print\",\n    \"verl.utils.profiler.mark_annotate\",\n    \"verl.utils.profiler.mark_end_range\",\n    \"verl.utils.profiler.mark_start_range\",\n    \"verl.models.mcore.qwen2_5_vl.get_vision_model_config\",\n    \"verl.models.mcore.qwen2_5_vl.get_vision_projection_config\",\n    \"verl.models.mcore.mbridge.freeze_moe_router\",\n    \"verl.models.mcore.mbridge.make_value_model\",\n    \"verl.utils.transformers_compat.flash_attn_supports_top_left_mask\",\n]\n\n\ndef iter_submodules(root: ModuleType) -> Iterable[ModuleType]:\n    \"\"\"Yield *root* and every sub-module inside it.\"\"\"\n    yield root\n\n    def print_pkg_error(pkg_name):\n        print(f\"[warn] Skipping {pkg_name!r}\", file=sys.stderr)\n\n    if getattr(root, \"__path__\", None):  # only packages have __path__\n        for mod_info in pkgutil.walk_packages(root.__path__, prefix=f\"{root.__name__}.\", onerror=print_pkg_error):\n            try:\n                yield importlib.import_module(mod_info.name)\n            except Exception as exc:\n                print(f\"[warn] Skipping {mod_info.name!r}: {exc}\", file=sys.stderr)\n\n\ndef names_missing_doc(mod: ModuleType) -> list[str]:\n    \"\"\"Return fully-qualified names that need docstrings.\"\"\"\n    missing: list[str] = []\n    public = getattr(mod, \"__all__\", [])\n    for name in public:\n        obj = getattr(mod, name, None)\n        if f\"{mod.__name__}.{name}\" in _ALLOW_LIST:\n            continue\n        if obj is None:\n            # Exported but not found in the module: flag it anyway.\n            missing.append(f\"{mod.__name__}.{name}  (not found)\")\n            continue\n\n        if inspect.isfunction(obj) or inspect.isclass(obj):\n            doc = inspect.getdoc(obj)\n            if not doc or not doc.strip():\n                missing.append(f\"{mod.__name__}.{name}\")\n    return missing\n\n\ndef check_module(qualname: str) -> list[str]:\n    \"\"\"Import *qualname* and check it (and sub-modules).\"\"\"\n    try:\n        module = importlib.import_module(qualname)\n    except ModuleNotFoundError as exc:\n        print(f\"[error] Cannot import '{qualname}': {exc}\", file=sys.stderr)\n        return [qualname]\n\n    missing: list[str] = []\n    for submod in iter_submodules(module):\n        missing.extend(names_missing_doc(submod))\n    return missing\n\n\ndef autodiscover_packages() -> list[str]:\n    \"\"\"Detect top-level packages under CWD when no argument is given.\"\"\"\n    pkgs: list[str] = []\n    for p in Path.cwd().iterdir():\n        if p.is_dir() and (p / \"__init__.py\").exists():\n            pkgs.append(p.name)\n    return pkgs\n\n\ndef main() -> None:\n    parser = argparse.ArgumentParser(description=__doc__)\n    parser.add_argument(\n        \"modules\",\n        nargs=\"*\",\n        help=\"Fully-qualified module or package names (defaults to every top-level package found in CWD).\",\n    )\n    args = parser.parse_args()\n\n    targets = args.modules or autodiscover_packages()\n    if not targets:\n        raise ValueError(\"[error] No modules specified and none detected automatically.\")\n\n    all_missing: list[str] = []\n    for modname in targets:\n        all_missing.extend(check_module(modname))\n\n    if all_missing:\n        print(\"\\nMissing docstrings:\")\n        for name in sorted(all_missing):\n            print(f\"  - {name}\")\n        raise ValueError(\"Missing docstrings detected. Please enhance them with docs accordingly.\")\n\n    print(\"✅ All exported functions/classes have docstrings.\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/check_dataproto_usage.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis CI test is used for checking whether DataProto is used in the code of some directory\n\"\"\"\n\nimport os\nfrom argparse import ArgumentParser\nfrom pathlib import Path\n\nSEARCH_WHITELIST = []\n\nSEARCH_KEYWORDS = [\"DataProto\"]\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\"--directory\", \"-d\", required=True, type=str)\n    args = parser.parse_args()\n    directory_in_str = args.directory\n\n    pathlist = Path(directory_in_str).glob(\"**/*.py\")\n    for path in pathlist:\n        path_in_str = str(path.absolute())\n\n        # judge whether current path is in pre-defined search whitelist or not.\n        path_in_whitelist = False\n\n        for sw in SEARCH_WHITELIST:\n            # for easy debugging in non-linux system\n            sw = sw.replace(\"/\", os.sep)\n            if sw in path_in_str:\n                print(f\"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.\")\n                path_in_whitelist = True\n                break\n\n        if path_in_whitelist:\n            continue\n\n        with open(path_in_str, encoding=\"utf-8\") as f:\n            file_content = f.read()\n\n            find_invalid_device_management = False\n\n            for sk in SEARCH_KEYWORDS:\n                if sk in file_content:\n                    find_invalid_device_management = True\n                    break\n\n            print(\n                f\"[CHECK] File {path_in_str} is detected for DataProto usage check, check result: \"\n                f\"{'success' if not find_invalid_device_management else f'failed, because detect {sk}'}.\"\n            )\n\n            assert not find_invalid_device_management, (\n                f\"file {path_in_str} contains DataProto usage, please use TensorDict directly!\"\n            )\n"
  },
  {
    "path": "tests/special_sanity/check_device_api_usage.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis CI test is used for checking whether device api usage is irregular, suggest using api in `verl/utils/device.py`.\nSearch targets include .py files in verl/recipe and verl/verl.\nSome files that must contain \".cuda\", \"cuda\" or \"nccl\" keyword is pre-defined in whitelist below.\n\"\"\"\n\nimport os\nfrom argparse import ArgumentParser\nfrom pathlib import Path\n\n# directory or file path must contain keyword \".cuda\" or \"cuda\"\nCUDA_KEYWORD_CHECK_WHITELIST = [\n    \"verl/utils/device.py\",\n    \"verl/utils/torch_functional.py\",  # import flash_attn only on cuda\n    \"verl/utils/profiler/nvtx_profile.py\",  # appear in NsightSystemsProfiler\n    \"verl/utils/profiler/torch_profile.py\",  # appear in TorchProfiler\n    \"verl/utils/profiler/config.py\",  # appear in TorchProfilerToolConfig\n    \"verl/utils/kernel/linear_cross_entropy.py\",  # appear in nvidia nvtx\n    \"verl/utils/rendezvous/ray_backend.py\",  # appear in cupy importance\n    \"verl/single_controller/ray/base.py\",  # appear in default device_name\n    \"verl/trainer/ppo/ray_trainer.py\",  # appear in default device_name\n    \"verl/experimental/transfer_queue/ray_trainer.py\",  # appear in docstring as default device_name\n    \"verl/experimental/one_step_off_policy/ray_trainer.py\",  # appear in docstring as default device_name\n    \"verl/utils/reward_score/sandbox_fusion/utils.py\",  # appear in sandbox language type\n    \"verl/third_party/torch/distributed/_state_dict_utils.py\",  # torch monkey patch fixes\n    \"verl/third_party/torch/distributed/checkpoint/state_dict.py\",  # torch monkey patch fixes\n    \"verl/workers/engine/base.py\",  # appear in default device_name\n    \"verl/workers/engine/utils.py\",  # appear in enable_full_determinism\n    \"verl/workers/engine/fsdp/transformer_impl.py\",  # appear in default device_name\n    \"verl/workers/engine/veomni/transformer_impl.py\",  # appear in default device_name\n    \"verl/workers/engine/torchtitan/transformer_impl.py\",  # appear in default device_name\n    \"verl/workers/engine/torchtitan/utils.py\",  # appear in torch.cuda.empty_cache()\n    \"verl/workers/engine/automodel/transformer_impl.py\",  # appear in default device_name\n    \"verl/workers/rollout/vllm_rollout/vllm_async_server.py\",  # appear in config.cudagraph_capture_sizes\n    \"verl/workers/rollout/sglang_rollout/async_sglang_server.py\",  # manually set CUDA_VISIBLE_DEVICES\n    \"verl/workers/rollout/trtllm_rollout/trtllm_async_server.py\",  # appear in config.cudagraph_capture_sizes\n    \"verl/workers/rollout/replica.py\",  # appear in default device_name\n    \"verl/checkpoint_engine\",  # checkpoint engine backend are device specific\n]\n\n# directory or file path must contain keyword \"nccl\"\nNCCL_KEYWORD_CHECK_WHITELIST = [\n    \"verl/utils/device.py\",\n    \"verl/third_party/sglang/parallel_state.py\",  # appear in default backend\n]\n\nSEARCH_WHITELIST = CUDA_KEYWORD_CHECK_WHITELIST + NCCL_KEYWORD_CHECK_WHITELIST\n\nSEARCH_KEYWORDS = [\".cuda\", '\"cuda\"', '\"nccl\"']\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\"--directory\", \"-d\", required=True, type=str)\n    args = parser.parse_args()\n    directory_in_str = args.directory\n\n    pathlist = Path(directory_in_str).glob(\"**/*.py\")\n    for path in pathlist:\n        path_in_str = str(path.absolute())\n\n        # judge whether current path is in pre-defined search whitelist or not.\n        path_in_whitelist = False\n\n        for sw in SEARCH_WHITELIST:\n            # for easy debugging in non-linux system\n            sw = sw.replace(\"/\", os.sep)\n            if sw in path_in_str:\n                print(f\"[SKIP] File {path_in_str} is in device api usage check whitelist, checking is skipped.\")\n                path_in_whitelist = True\n                break\n\n        if path_in_whitelist:\n            continue\n\n        with open(path_in_str, encoding=\"utf-8\") as f:\n            file_content = f.read()\n\n            find_invalid_device_management = False\n\n            for sk in SEARCH_KEYWORDS:\n                if sk in file_content:\n                    find_invalid_device_management = True\n                    break\n\n            print(\n                f\"[CHECK] File {path_in_str} is detected for device api usage check, check result: \"\n                f\"{'success' if not find_invalid_device_management else f'failed, because detect {sk}'}.\"\n            )\n\n            assert not find_invalid_device_management, (\n                f'file {path_in_str} contains .cuda/\"cuda\"/\"nccl\" usage, please use api in '\n                f\"verl/utils/device.py directly.\"\n            )\n"
  },
  {
    "path": "tests/special_sanity/check_docs_time_info.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCheck that every .md and .rst file under docs/ contains the substring \"Last updated\",\nwith an allow-list for exceptions.\n\"\"\"\n\nimport sys\nfrom pathlib import Path\n\n# === CONFIGURATION ===\n\n# Relative paths (to docs/) or glob patterns to skip checking\nALLOW_LIST = {\n    \"docs/README.md\",  # you can list individual files\n    \"docs/legacy/*.rst\",  # or glob patterns\n    \"docs/index.rst\",\n    \"docs/start/install.rst\",\n    \"docs/start/quickstart.rst\",\n    \"docs/README_vllm0.7.md\",\n}\n\n# The folder to scan\nDOCS_DIR = Path(\"docs\")\n\n# === SCRIPT ===\n\n\ndef is_allowed(path: Path) -> bool:\n    \"\"\"\n    Return True if `path` matches any entry in ALLOW_LIST.\n    \"\"\"\n    rel = str(path)\n    for pattern in ALLOW_LIST:\n        if Path(rel).match(pattern):\n            return True\n    return False\n\n\ndef main():\n    if not DOCS_DIR.exists():\n        print(f\"Error: Documentation directory '{DOCS_DIR}' does not exist.\", file=sys.stderr)\n        sys.exit(1)\n\n    missing = []\n\n    # Gather all .md and .rst files under docs/\n    for ext in (\"*.md\", \"*.rst\"):\n        for path in DOCS_DIR.rglob(ext):\n            if is_allowed(path):\n                continue\n\n            text = path.read_text(encoding=\"utf-8\", errors=\"ignore\")\n            if \"Last updated\" not in text:\n                missing.append(path)\n\n    # Report\n    if missing:\n        print(\"\\nThe following files are missing the 'Last updated' string:\\n\")\n        for p in missing:\n            print(f\"  - {p}\")\n        print(f\"\\nTotal missing: {len(missing)}\\n\", file=sys.stderr)\n        raise AssertionError(\n            \"Some documentation files lack a 'Last updated' line. Please include info such as \"\n            \"'Last updated: mm/dd/yyyy' to indicate the last update time of the document.\"\n        )\n    else:\n        print(\"✅ All checked files contain 'Last updated'.\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/check_docstrings.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nPython script to check docstrings for functions and classes in specified files.\nChecks that every public function and class has proper docstring documentation.\n\"\"\"\n\nimport ast\nimport os\nimport sys\n\n\nclass DocstringChecker(ast.NodeVisitor):\n    \"\"\"AST visitor to check for missing docstrings in functions and classes.\"\"\"\n\n    def __init__(self, filename: str):\n        self.filename = filename\n        self.missing_docstrings: list[tuple[str, str, int]] = []\n        self.current_class = None\n        self.function_nesting_level = 0\n\n    def visit_FunctionDef(self, node: ast.FunctionDef):\n        \"\"\"Visit function definitions and check for docstrings.\"\"\"\n        if not node.name.startswith(\"_\") and self.function_nesting_level == 0:\n            if not self._has_docstring(node):\n                func_name = f\"{self.current_class}.{node.name}\" if self.current_class else node.name\n                self.missing_docstrings.append((func_name, self.filename, node.lineno))\n\n        self.function_nesting_level += 1\n        self.generic_visit(node)\n        self.function_nesting_level -= 1\n\n    def visit_AsyncFunctionDef(self, node: ast.AsyncFunctionDef):\n        \"\"\"Visit async function definitions and check for docstrings.\"\"\"\n        if not node.name.startswith(\"_\") and self.function_nesting_level == 0:\n            if not self._has_docstring(node):\n                func_name = f\"{self.current_class}.{node.name}\" if self.current_class else node.name\n                self.missing_docstrings.append((func_name, self.filename, node.lineno))\n\n        self.function_nesting_level += 1\n        self.generic_visit(node)\n        self.function_nesting_level -= 1\n\n    def visit_ClassDef(self, node: ast.ClassDef):\n        \"\"\"Visit class definitions and check for docstrings.\"\"\"\n        if not node.name.startswith(\"_\"):\n            if not self._has_docstring(node):\n                self.missing_docstrings.append((node.name, self.filename, node.lineno))\n\n        old_class = self.current_class\n        self.current_class = node.name\n        self.generic_visit(node)\n        self.current_class = old_class\n\n    def _has_docstring(self, node) -> bool:\n        \"\"\"Check if a node has a docstring.\"\"\"\n        return ast.get_docstring(node) is not None\n\n\ndef check_file_docstrings(filepath: str) -> list[tuple[str, str, int]]:\n    \"\"\"Check docstrings in a single file.\"\"\"\n    try:\n        with open(filepath, encoding=\"utf-8\") as f:\n            content = f.read()\n\n        tree = ast.parse(content, filename=filepath)\n        checker = DocstringChecker(filepath)\n        checker.visit(tree)\n        return checker.missing_docstrings\n\n    except Exception as e:\n        print(f\"Error processing {filepath}: {e}\")\n        return []\n\n\ndef main():\n    \"\"\"Main function to check docstrings in specified files.\"\"\"\n\n    files_to_check = [\n        \"verl/trainer/ppo/ray_trainer.py\",\n        \"verl/trainer/main_ppo.py\",\n        \"verl/trainer/ppo/reward.py\",\n        \"verl/utils/reward_score/__init__.py\",\n        \"verl/trainer/ppo/core_algos.py\",\n        \"verl/experimental/agent_loop/agent_loop.py\",\n        \"verl/workers/sharding_manager/fsdp_vllm.py\",\n        \"verl/workers/sharding_manager/fsdp_ulysses.py\",\n    ]\n\n    script_dir = os.path.dirname(os.path.abspath(__file__))\n    repo_path = os.path.dirname(os.path.dirname(script_dir))\n\n    if not os.path.exists(repo_path):\n        print(f\"Repository path {repo_path} does not exist!\")\n        sys.exit(1)\n\n    os.chdir(repo_path)\n\n    all_missing_docstrings = []\n\n    print(\"Checking docstrings in specified files...\")\n    print(\"=\" * 60)\n\n    for file_path in files_to_check:\n        if not os.path.exists(file_path):\n            print(f\"Warning: File {file_path} does not exist!\")\n            continue\n\n        print(f\"Checking {file_path}...\")\n        missing = check_file_docstrings(file_path)\n        all_missing_docstrings.extend(missing)\n\n        if missing:\n            print(f\"  Found {len(missing)} missing docstrings\")\n        else:\n            print(\"  All functions and classes have docstrings [OK]\")\n\n    print(\"=\" * 60)\n\n    if all_missing_docstrings:\n        print(f\"\\nSUMMARY: Found {len(all_missing_docstrings)} functions/classes missing docstrings:\")\n        print(\"-\" * 60)\n\n        by_file = {}\n        for name, filepath, lineno in all_missing_docstrings:\n            if filepath not in by_file:\n                by_file[filepath] = []\n            by_file[filepath].append((name, lineno))\n\n        for filepath in sorted(by_file.keys()):\n            print(f\"\\n{filepath}:\")\n            for name, lineno in sorted(by_file[filepath], key=lambda x: x[1]):\n                print(f\"  - {name} (line {lineno})\")\n\n        print(f\"\\nTotal missing docstrings: {len(all_missing_docstrings)}\")\n\n        raise Exception(f\"Found {len(all_missing_docstrings)} functions/classes without proper docstrings!\")\n\n    else:\n        print(\"\\n[OK] All functions and classes have proper docstrings!\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/check_license.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom argparse import ArgumentParser\nfrom pathlib import Path\nfrom typing import Iterable\n\nlicense_head_bytedance = \"Copyright 2024 Bytedance Ltd. and/or its affiliates\"\nlicense_head_bytedance_25 = \"Copyright 2025 Bytedance Ltd. and/or its affiliates\"\nlicense_head_bytedance_26 = \"Copyright 2026 Bytedance Ltd. and/or its affiliates\"\n# Add custom license headers below\nlicense_head_prime = \"Copyright 2024 PRIME team and/or its affiliates\"\nlicense_head_individual = \"Copyright 2025 Individual Contributor:\"\nlicense_head_sglang = \"Copyright 2023-2024 SGLang Team\"\nlicense_head_modelbest = \"Copyright 2025 ModelBest Inc. and/or its affiliates\"\nlicense_head_amazon = \"Copyright 2025 Amazon.com Inc and/or its affiliates\"\nlicense_head_amazon_26 = \"Copyright 2026 Amazon.com Inc and/or its affiliates\"\nlicense_head_facebook = \"Copyright (c) 2016-     Facebook, Inc\"\nlicense_head_meituan = \"Copyright 2025 Meituan Ltd. and/or its affiliates\"\nlicense_head_huawei = \"Copyright (c) 2025 Huawei Technologies Co., Ltd. All Rights Reserved.\"\nlicense_head_nvidia = \"Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\"\nlicense_headers = [\n    license_head_bytedance,\n    license_head_bytedance_25,\n    license_head_bytedance_26,\n    license_head_prime,\n    license_head_individual,\n    license_head_sglang,\n    license_head_modelbest,\n    license_head_amazon,\n    license_head_amazon_26,\n    license_head_facebook,\n    license_head_meituan,\n    license_head_huawei,\n    license_head_nvidia,\n]\n\n\ndef get_py_files(path_arg: Path) -> Iterable[Path]:\n    \"\"\"get py files under a dir. if already py file return it\n\n    Args:\n        path_arg (Path): path to scan for py files\n\n    Returns:\n        Iterable[Path]: list of py files\n    \"\"\"\n    if path_arg.is_dir():\n        return path_arg.glob(\"**/*.py\")\n    elif path_arg.is_file() and path_arg.suffix == \".py\":\n        return [path_arg]\n    return []\n\n\nif __name__ == \"__main__\":\n    parser = ArgumentParser()\n    parser.add_argument(\n        \"--directories\",\n        \"-d\",\n        required=True,\n        type=Path,\n        nargs=\"+\",\n        help=\"List of directories to check for license headers\",\n    )\n    args = parser.parse_args()\n\n    # Collect all Python files from specified directories\n    pathlist = set(path for path_arg in args.directories for path in get_py_files(path_arg))\n\n    for path in pathlist:\n        # because path is object not string\n        path_in_str = str(path.absolute())\n        print(path_in_str)\n        with open(path_in_str, encoding=\"utf-8\") as f:\n            file_content = f.read()\n\n            has_license = False\n            for lh in license_headers:\n                if lh in file_content:\n                    has_license = True\n                    break\n            assert has_license, f\"file {path_in_str} does not contain license\"\n"
  },
  {
    "path": "tests/special_sanity/check_pr_description.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n#!/usr/bin/env python3\nimport json\nimport os\n\n# Number of lines to check\nNUM_LINES = 5\n\n\n# Custom exception types for clear error handling\nclass TemplateFileError(Exception):\n    pass\n\n\nclass PRBodyLoadError(Exception):\n    pass\n\n\nclass PRDescriptionError(Exception):\n    pass\n\n\n# Path to the PR template file\ntemplate_file = os.path.join(os.getenv(\"GITHUB_WORKSPACE\", \".\"), \".github\", \"PULL_REQUEST_TEMPLATE.md\")\n\n\ndef load_template(path):\n    \"\"\"\n    Load only the first NUM_LINES of the PR template file as a list of lines,\n    without stripping any characters.\n    \"\"\"\n    lines = []\n    try:\n        with open(path, encoding=\"utf-8\") as f:\n            for _ in range(NUM_LINES):\n                line = f.readline()\n                if not line:\n                    break\n                lines.append(line.strip())\n        return lines\n    except Exception as e:\n        raise TemplateFileError(f\"Failed to read PR template (first {NUM_LINES} lines) at {path}: {e}\") from e\n\n\ndef load_pr_body(event_path):\n    try:\n        with open(event_path, encoding=\"utf-8\") as f:\n            payload = json.load(f)\n        return payload.get(\"pull_request\", {}).get(\"body\", \"\") or \"\"\n    except Exception as e:\n        raise PRBodyLoadError(f\"Failed to read PR body from {event_path}: {e}\") from e\n\n\ndef check_pr_description(body, template_lines):\n    \"\"\"\n    Compare the first NUM_LINES lines of the PR body to the template lines.\n    If they match exactly, the placeholder was not modified.\n    \"\"\"\n    pr_lines = body.splitlines(keepends=True)\n    pr_first = [x.strip() for x in pr_lines[:NUM_LINES]]\n    if pr_first == template_lines:\n        raise PRDescriptionError(\n            \"It looks like you haven't updated the '### What does this PR do?' section. Please replace \"\n            \"the placeholder text with a concise description of what your PR does.\"\n        )\n    else:\n        print(pr_first)\n        print(template_lines)\n\n\ndef main():\n    event_path = os.getenv(\"GITHUB_EVENT_PATH\")\n    if not event_path:\n        raise OSError(\"GITHUB_EVENT_PATH is not set.\")\n\n    template_lines = load_template(template_file)\n    pr_body = load_pr_body(event_path)\n    check_pr_description(pr_body, template_lines)\n\n    print(\"✅ '### What does this PR do?' section has been filled out.\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/check_pr_title.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport re\n\n# Get PR title from environment\npr_title = os.environ.get(\"PR_TITLE\", \"\").strip()\n\n# Define rules\nallowed_modules = [\"fsdp\", \"megatron\", \"veomni\", \"sglang\", \"vllm\", \"trtllm\", \"rollout\", \"trainer\"]\nallowed_modules += [\"tests\", \"training_utils\", \"recipe\", \"hardware\", \"deployment\"]\nallowed_modules += [\"ray\", \"worker\", \"single_controller\", \"misc\", \"docker\", \"ci\"]\nallowed_modules += [\"perf\", \"model\", \"algo\", \"env\", \"tool\", \"ckpt\", \"doc\", \"data\", \"cfg\", \"reward\"]\nallowed_modules += [\"fully_async\", \"one_step_off\"]\nallowed_types = [\"feat\", \"fix\", \"refactor\", \"chore\", \"test\"]\n\n# Check for [1/N] prefix and extract the rest of the title\nprogress_match = re.match(r\"^\\[\\d/[\\dNn]\\]\\s*(.+)$\", pr_title, re.IGNORECASE)\nif progress_match:\n    pr_title = progress_match.group(1).strip()\n\n# Check for [BREAKING] prefix and extract the rest of the title\nbreaking_match = re.match(r\"^\\[BREAKING\\]\\s*(.+)$\", pr_title, re.IGNORECASE)\nif breaking_match:\n    core_pr_title = breaking_match.group(1).strip()\n    is_breaking = True\nelse:\n    core_pr_title = pr_title\n    is_breaking = False\n\n# Build dynamic regex pattern for modules (now working on core_pr_title)\nre_modules_pattern = re.compile(r\"^\\[([a-z_,\\s]+)\\]\", re.IGNORECASE)\nre_modules = re_modules_pattern.match(core_pr_title)\nif not re_modules:\n    print(f\"❌ Invalid PR title: '{pr_title}'\")\n    print(\"Expected format: [BREAKING][module] type: description\")\n    print(f\"Allowed modules: {', '.join(allowed_modules)}\")\n    raise Exception(\"Invalid PR title\")\nelse:\n    modules = re.findall(r\"[a-z_]+\", re_modules.group(1).lower())\n    if not all(module in allowed_modules for module in modules):\n        invalid_modules = [module for module in modules if module not in allowed_modules]\n        print(f\"❌ Invalid modules: {', '.join(invalid_modules)}\")\n        print(f\"Allowed modules: {', '.join(allowed_modules)}\")\n        raise Exception(\"Invalid PR title\")\n\ntypes_pattern = \"|\".join(re.escape(t) for t in allowed_types)\nre_types_pattern = re.compile(rf\"^\\[[a-z_,\\s]+\\]\\s+({types_pattern}):\\s+.+$\", re.IGNORECASE)\nmatch = re_types_pattern.match(core_pr_title)\n\nif not match:\n    print(f\"❌ Invalid PR title: '{pr_title}'\")\n    print(\"Expected format: [BREAKING][module] type: description\")\n    print(f\"Allowed types: {', '.join(allowed_types)}\")\n    raise Exception(\"Invalid PR title\")\n\nchange_type = match.group(1).lower()\n\n# Build the success message\nbreaking_info = \" (BREAKING CHANGE)\" if is_breaking else \"\"\nprint(f\"✅ PR title is valid: {pr_title}, modules: {modules}, type: {change_type}{breaking_info}\")\n"
  },
  {
    "path": "tests/special_sanity/test_config_docs.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\nfrom pathlib import Path\n\n\ndef validate_yaml_format(yaml_lines):\n    errors = []\n    i = 0\n\n    while i < len(yaml_lines):\n        line = yaml_lines[i]\n        stripped = line.strip()\n\n        # Skip empty lines\n        if stripped == \"\":\n            i += 1\n            continue\n\n        # Match YAML keys like \"field:\" or \"field: value\"\n        key_match = re.match(r\"^(\\s*)([a-zA-Z0-9_]+):\", line)\n        if key_match:\n            # Check if there's a comment above\n            if i == 0 or not yaml_lines[i - 1].strip().startswith(\"#\"):\n                errors.append(f\"Missing comment above line {i + 1}: {line.strip()}\")\n\n            # Check for inline comment\n            if \"#\" in line and not stripped.startswith(\"#\"):\n                comment_index = line.index(\"#\")\n                colon_index = line.index(\":\")\n                if comment_index > colon_index:\n                    errors.append(f\"Inline comment found on line {i + 1}: {line.strip()}\")\n\n            # Check for blank line after this key line (unless next is a deeper indent)\n            if i + 1 < len(yaml_lines):\n                next_line = yaml_lines[i + 1]\n                next_stripped = next_line.strip()\n\n                # If next is not empty and not a deeper nested line, enforce blank line\n                if next_stripped != \"\":\n                    errors.append(f\"Missing blank line after line {i + 1}: {line.strip()}\")\n\n        i += 1\n\n    return errors\n\n\ndef test_trainer_config_doc():\n    yamls_to_inspect = [\n        \"verl/trainer/config/ppo_trainer.yaml\",\n        \"verl/trainer/config/actor/actor.yaml\",\n        \"verl/trainer/config/actor/dp_actor.yaml\",\n        \"verl/trainer/config/critic/critic.yaml\",\n        \"verl/trainer/config/critic/dp_critic.yaml\",\n        \"verl/trainer/config/ref/ref.yaml\",\n        \"verl/trainer/config/ref/dp_ref.yaml\",\n        \"verl/trainer/config/rollout/rollout.yaml\",\n    ]\n    success = True\n    for yaml_to_inspect in yamls_to_inspect:\n        yaml_path = Path(yaml_to_inspect)  # path to your YAML file\n        with open(yaml_path) as f:\n            lines = f.readlines()\n\n        validation_errors = validate_yaml_format(lines)\n        if validation_errors:\n            success = False\n            print(\"YAML documentation format check failed:\")\n            print(f\"Please read the top block of {yaml_to_inspect} to see format rules:\\n\")\n            for err in validation_errors:\n                print(\" -\", err)\n\n    if not success:\n        raise Exception(\"Please fix documentation format.\")\n    else:\n        print(\"YAML format check passed ✅\")\n"
  },
  {
    "path": "tests/special_sanity/test_import.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\ndef test_import():\n    import verl\n\n    print(verl.__version__)\n\n\ndef test_single_controller_import():\n    import verl.single_controller\n\n    print(verl.single_controller.__version__)\n"
  },
  {
    "path": "tests/special_sanity/type_coverage_check.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Custom type annotation check tool.\nTo inspect the type annotation for functions in the entire codebase, please run:\nfind verl -type f -name \"*.py\" | xargs -n 1 python3 tests/special_sanity/type_coverage_check.py --all-lines\n--debug --target-file\n\"\"\"\n\nimport argparse\nimport ast\nimport linecache\nimport subprocess\nfrom pathlib import Path\n\n\ndef get_changed_files() -> list[Path]:\n    result = subprocess.run(\n        [\"git\", \"diff\", \"--name-only\", \"--diff-filter=AM\", \"origin/main...HEAD\"], stdout=subprocess.PIPE, text=True\n    )\n    return [Path(f) for f in result.stdout.splitlines() if f.endswith(\".py\")]\n\n\ndef get_changed_lines(file_path: Path) -> set[int]:\n    result = subprocess.run(\n        [\"git\", \"diff\", \"-U0\", \"origin/main...HEAD\", \"--\", str(file_path)],\n        stdout=subprocess.PIPE,\n        text=True,\n    )\n    lines: set[int] = set()\n    for line in result.stdout.splitlines():\n        if line.startswith(\"@@\"):\n            for part in line.split():\n                try:\n                    if part.startswith(\"+\") and \",\" in part:\n                        start, count = map(int, part[1:].split(\",\"))\n                        lines.update(range(start, start + count))\n                    elif part.startswith(\"+\") and \",\" not in part:\n                        lines.add(int(part[1:]))\n                except Exception:\n                    # (vermouth1992) There are many edge cases here because + can be in the changed program\n                    pass\n    return lines\n\n\nCHECK_SUCCESS = 0\nCHECK_WARNING = 1\nCHECK_FAILURE = -1\n\n\ndef should_check_type(arg_name: str) -> bool:\n    if arg_name in (\"self\", \"cls\"):\n        return False\n    if arg_name.startswith(\"*\"):\n        return False\n    return True\n\n\ndef has_type_annotations(node: ast.AST, debug: bool = False) -> int:\n    if isinstance(node, ast.FunctionDef):\n        is_private = node.name.startswith(\"_\")\n        if node.args.vararg is not None or node.args.kwarg is not None:\n            return CHECK_SUCCESS\n        has_ann = (\n            all(arg.annotation is not None for arg in node.args.args if should_check_type(arg.arg))\n            and node.returns is not None\n        )\n        if has_ann or is_private:\n            return CHECK_SUCCESS\n        else:\n            if debug:\n                print(node, [(arg.annotation, arg.arg) for arg in node.args.args if should_check_type(arg.arg)])\n            return CHECK_FAILURE\n    return CHECK_SUCCESS\n\n\ndef check_file(\n    file_path: Path, changed_lines: set[int], debug: bool = False\n) -> tuple[int, int, list[tuple[Path, int, str]], list[tuple[Path, int, str]]]:\n    with open(file_path) as f:\n        source: str = f.read()\n    tree = ast.parse(source, filename=str(file_path))\n    annotated = 0\n    total = 0\n    warning_lines: list[tuple[Path, int, str]] = []\n    failure_lines: list[tuple[Path, int, str]] = []\n\n    for node in ast.walk(tree):\n        if hasattr(node, \"lineno\") and node.lineno in changed_lines:\n            if isinstance(node, ast.FunctionDef | ast.Assign | ast.AnnAssign):\n                total += 1\n                result = has_type_annotations(node, debug)\n                if result == CHECK_SUCCESS or result == CHECK_WARNING:\n                    annotated += 1\n                    if result == CHECK_WARNING:\n                        warning_lines.append(\n                            (file_path, node.lineno, linecache.getline(str(file_path), node.lineno).strip())\n                        )\n                else:\n                    source_line = linecache.getline(str(file_path), node.lineno).strip()\n                    failure_lines.append((file_path, node.lineno, source_line))\n\n    return annotated, total, warning_lines, failure_lines\n\n\ndef main() -> None:\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--threshold\", type=float, default=0.3, help=\"Minimum ratio of annotated lines required (0.0 - 1.0)\"\n    )\n    parser.add_argument(\"--target-file\", type=str, default=None, help=\"Path to the Python source file to analyse\")\n    parser.add_argument(\n        \"--all-lines\",\n        action=\"store_true\",\n        help=\"Check all lines in the file instead of only changed lines based on git\",\n    )\n    parser.add_argument(\"--debug\", action=\"store_true\", help=\"Add debugging logs\")\n    args = parser.parse_args()\n\n    total_changed = 0\n    total_annotated = 0\n    all_warnings: list[tuple[Path, int, str]] = []\n    all_failures: list[tuple[Path, int, str]] = []\n\n    target_files = [args.target_file] if args.target_file is not None else get_changed_files()\n    for fpath in target_files:\n        if \"tests/\" in str(fpath):\n            continue\n        if args.all_lines:\n            changed_lines = [i + 1 for i in range(len(open(fpath).readlines()))]\n        else:\n            changed_lines = get_changed_lines(fpath)\n        annotated, total, warning_lines, failure_lines = check_file(fpath, changed_lines, args.debug)\n        total_annotated += annotated\n        total_changed += total\n        all_warnings.extend(warning_lines)\n        all_failures.extend(failure_lines)\n\n    ratio = (total_annotated / total_changed) if total_changed else 1.0\n\n    print(\n        f\"🔍 Type coverage on {'all' if args.all_lines else 'changed'} lines: \"\n        f\"{total_annotated}/{total_changed} = {ratio:.2%}. Files inspected: {target_files}\"\n    )\n\n    if all_warnings:\n        print(\"\\n⚠️ Suggest Improve: Lines missing type annotations for inputs and outputs:\\n\")\n        for fname, lineno, line in all_warnings:\n            print(f\"{fname}:{lineno}: {line}\")\n\n    if all_failures:\n        print(\"⚠️ [ERROR] Lines missing type annotations for inputs and outputs:\\n\")\n        for fname, lineno, line in all_failures:\n            print(f\"{fname}:{lineno}: {line}\")\n\n    if ratio < args.threshold:\n        print(\n            f\"Please add type annotations for inputs and outputs to meet threshold {args.threshold}. \"\n            f\"Cases exempt from checking:\"\n        )\n        print(\"1. Private methods.\")\n        print(\"2. Args with name in ('self', 'cls'), or *args / **kwargs\")\n        print(\"3. Files under tests/\")\n        raise Exception(f\"\\n❌ Type coverage below threshold ({args.threshold:.0%}).\")\n    else:\n        if all_warnings or all_failures:\n            print(\"\")\n        print(\"✅ Type annotation coverage acceptable.\\n\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/validate_imported_docs.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nverify_imported_docs.py\n\nAssert that every function or class *explicitly imported* (via\n`from <module> import <name>`) in a given Python file has a docstring.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport ast\nimport importlib\nimport inspect\nimport pathlib\nimport sys\n\n\ndef _parse_args() -> argparse.Namespace:\n    p = argparse.ArgumentParser(description=\"Verify that imported functions/classes have docstrings.\")\n    p.add_argument(\n        \"--target-file\",\n        default=\"verl/trainer/ppo/ray_trainer.py\",\n        help=\"Path to the Python source file to analyse (e.g. verl/trainer/ppo/ray_trainer.py)\",\n    )\n    p.add_argument(\n        \"--allow-list\",\n        default=[\"omegaconf.open_dict\"],\n        help=\"a list of third_party dependencies that do not have proper docs :(\",\n    )\n    p.add_argument(\n        \"--project-root\",\n        default=\".\",\n        help=\"Directory to prepend to PYTHONPATH so local packages resolve (default: .)\",\n    )\n    p.add_argument(\n        \"--quiet\",\n        action=\"store_true\",\n        help=\"Suppress success message (still prints errors).\",\n    )\n    return p.parse_args()\n\n\ndef _import_attr(module_name: str, attr_name: str):\n    \"\"\"Import `module_name` then return `getattr(module, attr_name)`.\"\"\"\n    module = importlib.import_module(module_name)\n    return getattr(module, attr_name)\n\n\ndef _check_file(py_file: pathlib.Path, project_root: pathlib.Path, allow_list: list[str]) -> list[str]:\n    \"\"\"Return a list of error strings (empty == success).\"\"\"\n    # Ensure local packages resolve\n    sys.path.insert(0, str(project_root.resolve()))\n\n    tree = ast.parse(py_file.read_text(), filename=str(py_file))\n    problems: list[str] = []\n\n    for node in ast.walk(tree):\n        if not isinstance(node, ast.ImportFrom):\n            continue\n\n        # Relative imports (level > 0) get the leading dots stripped\n        module_name = \".\" * node.level + (node.module or \"\")\n        for alias in node.names:\n            if alias.name == \"*\":\n                problems.append(\n                    f\"{py_file}:{node.lineno} - wildcard import `from {module_name} import *` cannot be verified.\"\n                )\n                continue\n\n            imported_name = alias.name\n\n            try:\n                obj = _import_attr(module_name, imported_name)\n            except Exception:  # pragma: no cover – wide net for import quirks\n                pass\n                # For some reason the module cannot be imported, skip for now\n                # problems.append(\n                #     f\"{py_file}:{node.lineno} - could not resolve \"\n                #     f\"`{imported_name}` from `{module_name}` ({exc})\"\n                # )\n                continue\n\n            if f\"{module_name}.{imported_name}\" in allow_list:\n                continue\n            if inspect.isfunction(obj) or inspect.isclass(obj):\n                doc = inspect.getdoc(obj)\n                if not (doc and doc.strip()):\n                    kind = \"class\" if inspect.isclass(obj) else \"function\"\n                    problems.append(\n                        f\"{py_file}:{node.lineno} - {kind} `{module_name}.{imported_name}` is missing a docstring.\"\n                    )\n\n    return problems\n\n\ndef main() -> None:\n    args = _parse_args()\n    target_path = pathlib.Path(args.target_file).resolve()\n    project_root = pathlib.Path(args.project_root).resolve()\n\n    if not target_path.is_file():\n        raise Exception(f\"❌ Target file not found: {target_path}\")\n\n    errors = _check_file(target_path, project_root, args.allow_list)\n\n    if errors:\n        print(\"Docstring verification failed:\\n\")\n        print(\"\\n\".join(f\" • {e}\" for e in errors))\n        raise Exception(\"❌ Docstring verification failed.\")\n\n    if not args.quiet:\n        print(f\"✅ All explicitly imported functions/classes in {target_path} have docstrings.\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_sanity/validate_structure.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n#!/usr/bin/env python3\n\"\"\"\nValidate that test file subfolders mirror the top-level package layout.\n\nUsage examples\n--------------\n\n# Typical run (defaults: impl_root=my_project, tests_root=tests)\npython check_tests_structure.py\n\n# Custom layout and extra allowed folders\npython check_tests_structure.py \\\n    --impl-root verl \\\n    --tests-root tests \\\n    --allow-dirs special_e2e special_sanity special_standalone special_distributed\n\"\"\"\n\nfrom __future__ import annotations\n\nimport argparse\nimport sys\nfrom pathlib import Path\n\n\ndef discover_allowed_modules(impl_root: Path, extra: list[str]) -> set[str]:\n    \"\"\"Return the set of first-level directories that tests may live under.\"\"\"\n    allowed = {p.name for p in impl_root.iterdir() if p.is_dir()}\n    allowed.update(extra)\n    return allowed\n\n\ndef find_violations(tests_root: Path, allowed: set[str], allowed_files: list[str]) -> list[str]:\n    \"\"\"Return a list of error strings for test files in the wrong place.\"\"\"\n    errors: list[str] = []\n    for test_file in tests_root.rglob(\"test*.py\"):\n        if str(test_file) in allowed_files:\n            continue\n        rel_parts = test_file.relative_to(tests_root).parts\n        if len(rel_parts) < 2:\n            errors.append(f\"{test_file}: must be inside one of {sorted(allowed)} (not at tests root)\")\n            continue\n\n        first_folder = rel_parts[0]\n        if first_folder not in allowed:\n            errors.append(\n                f\"{test_file}: subfolder '{first_folder}' under tests/ is not an allowed module. \"\n                f\"The valid ones are: {sorted(allowed)}\"\n            )\n    return errors\n\n\ndef main() -> None:\n    parser = argparse.ArgumentParser(description=\"Check that test files follow tests/<module>/… layout.\")\n    parser.add_argument(\n        \"--impl-root\",\n        type=Path,\n        default=\"verl\",\n        help=\"Implementation root (default: my_project)\",\n    )\n    parser.add_argument(\n        \"--tests-root\",\n        type=Path,\n        default=\"tests\",\n        help=\"Root of test tree (default: tests)\",\n    )\n    parser.add_argument(\n        \"--allow-dirs\",\n        nargs=\"*\",\n        default=[\"special_e2e\", \"special_sanity\", \"special_standalone\", \"special_distributed\"],\n        help=\"Extra top-level test folders that are exempt from the rule\",\n    )\n    parser.add_argument(\n        \"--allow-files\",\n        nargs=\"*\",\n        default=[\n            \"tests/test_protocol_on_cpu.py\",\n            \"tests/test_base_config_on_cpu.py\",\n            \"tests/test_protocol_v2_on_cpu.py\",\n        ],\n        help=\"Extra top-level test folders that are exempt from the rule\",\n    )\n    args = parser.parse_args()\n\n    if not args.impl_root.is_dir():\n        raise Exception(f\"Implementation root '{args.impl_root}' does not exist.\")\n    if not args.tests_root.is_dir():\n        raise Exception(f\"Tests root '{args.tests_root}' does not exist.\")\n\n    allowed = discover_allowed_modules(args.impl_root, args.allow_dirs)\n    violations = find_violations(args.tests_root, allowed, args.allow_files)\n\n    if violations:\n        print(\"❌  Test layout violations found:\\n\", file=sys.stderr)\n        for err in violations:\n            print(\"  -\", err, file=sys.stderr)\n\n        print(\n            f\"\\nGuideline:\\n  Place each test file under   tests/<module_name>/…\\n  where <module_name> is \"\n            f\"one of the top-level packages inside '{args.impl_root}', or is explicitly listed via --allow-dirs.\\n\",\n            file=sys.stderr,\n        )\n        raise Exception(\"❌  Test layout violations found.\")\n\n    print(\"✅  Tests folder structure looks good.\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/special_standalone/README.md",
    "content": "The standalone test folder is reserved for tests that require dedicated environment (e.g. memory stress tests)\n"
  },
  {
    "path": "tests/special_standalone/test_memory_buffers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTest memory buffers\n- We start with two models with the same weights\n- We use Memory buffer to make one of the models and then compare the parameters\n\"\"\"\n\nimport gc\n\nimport torch\nfrom transformers import LlamaConfig, LlamaModel\n\n\ndef test_memory_buffers():\n    llama_config = LlamaConfig(\n        vocab_size=256,\n        hidden_size=4096,\n        intermediate_size=11008,\n        num_hidden_layers=2,\n        num_attention_heads=16,\n        num_key_value_heads=16,\n    )\n\n    model = LlamaModel(config=llama_config).cuda()\n    model_copy = LlamaModel(config=llama_config).cuda()\n    model_copy.load_state_dict(model.state_dict())\n\n    norm_factor = 1024**3\n\n    t_before = torch.cuda.get_device_properties(0).total_memory / norm_factor\n    r_before = torch.cuda.memory_reserved(0) / norm_factor\n    a_before = torch.cuda.memory_allocated(0) / norm_factor\n\n    print(f\"Before Total memory: {t_before} GB, reserved: {r_before} GB, allocated: {a_before} GB\")\n\n    t = torch.cuda.get_device_properties(0).total_memory / norm_factor\n    r = torch.cuda.memory_reserved(0) / norm_factor\n    a = torch.cuda.memory_allocated(0) / norm_factor\n\n    gc.collect()\n    torch.cuda.empty_cache()\n\n    print(f\"After Total memory: {t} GB, reserved: {r} GB, allocated: {a} GB\")\n\n    change_ratio = (a - a_before) / a_before\n    assert change_ratio < 0.01, f\"make sure the allocated change is less than 1%, Got {change_ratio}\"\n\n    for (name1, param1), (name2, param2) in zip(model.named_parameters(), model_copy.named_parameters(), strict=True):\n        assert name1 == name2\n        assert torch.eq(param1.data, param2.data).all(), f\"{param1.data}, {param2.data}, {name1}\"\n\n\nif __name__ == \"__main__\":\n    test_memory_buffers()\n"
  },
  {
    "path": "tests/test_base_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nfrom verl.base_config import BaseConfig\n\n\n@pytest.fixture\ndef base_config_mock():\n    \"\"\"Fixture to create a mock BaseConfig instance with test attributes.\"\"\"\n    mock_config = BaseConfig()\n    mock_config.test_attr = \"test_value\"\n    return mock_config\n\n\ndef test_getitem_success(base_config_mock):\n    \"\"\"Test __getitem__ with existing attribute (happy path).\"\"\"\n    assert base_config_mock[\"test_attr\"] == \"test_value\"\n\n\ndef test_getitem_nonexistent_attribute(base_config_mock):\n    \"\"\"Test __getitem__ with non-existent attribute (exception path 1).\"\"\"\n    with pytest.raises(AttributeError):\n        _ = base_config_mock[\"nonexistent_attr\"]\n\n\ndef test_getitem_invalid_key_type(base_config_mock):\n    \"\"\"Test __getitem__ with invalid key type (exception path 2).\"\"\"\n    with pytest.raises(TypeError):\n        _ = base_config_mock[123]  # type: ignore\n"
  },
  {
    "path": "tests/test_protocol_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport random\n\nimport numpy as np\nimport pytest\nimport tensordict\nimport torch\nfrom packaging.version import parse as parse_version\nfrom tensordict import TensorDict\n\nfrom verl import DataProto\nfrom verl.protocol import (\n    deserialize_single_tensor,\n    deserialize_tensordict,\n    serialize_single_tensor,\n    serialize_tensordict,\n    union_numpy_dict,\n    union_tensor_dict,\n)\nfrom verl.utils import tensordict_utils as tu\n\n\ndef test_union_tensor_dict():\n    obs = torch.randn(100, 10)\n\n    data1 = TensorDict({\"obs\": obs, \"act\": torch.randn(100, 3)}, batch_size=[100])\n    data2 = TensorDict({\"obs\": obs, \"next_obs\": torch.randn(100, 10), \"rew\": torch.randn(100)}, batch_size=[100])\n\n    data_with_copied_obs = TensorDict(\n        {\"obs\": obs.clone(), \"next_obs\": torch.randn(100, 10), \"rew\": torch.randn(100)}, batch_size=[100]\n    )\n\n    union_tensor_dict(data1, data2)\n    with pytest.raises(AssertionError):\n        union_tensor_dict(data1, data_with_copied_obs)\n\n\ndef test_union_numpy_dict():\n    \"\"\"\n    A comprehensive test suite for union_numpy_dict, covering standard use\n    cases, N-dimensional arrays, object-dtype arrays, and NaN value handling.\n    \"\"\"\n    arr_3d = np.arange(8).reshape((2, 2, 2))\n    union_numpy_dict({\"a\": arr_3d}, {\"a\": arr_3d})\n    arr1 = np.array([1, \"hello\", np.array([2, 3])], dtype=object)\n    arr2 = np.array([1, \"hello\", np.array([2, 3])], dtype=object)\n    union_numpy_dict({\"a\": arr1}, {\"a\": arr2})\n    # --- Test Case 1: The original test with mixed object/float types ---\n    # This test case from the original test file is preserved.\n    data = np.random.random(100)\n    # This array intentionally mixes float('nan') and the string 'nan'\n    nan_data = [float(\"nan\") for _ in range(99)]\n    nan_data.append(\"nan\")\n    nan_data_arr = np.array(nan_data, dtype=object)\n\n    dict1 = {\"a\": data, \"b\": nan_data_arr}\n    dict2_same = {\"a\": data.copy(), \"b\": nan_data_arr.copy()}\n    dict3_different = {\"a\": np.random.random(100)}\n\n    union_numpy_dict(dict1, dict2_same)  # Should pass\n    with pytest.raises(AssertionError):\n        union_numpy_dict(dict1, dict3_different)\n\n    # --- Test Case 2: Standard 3D arrays (fixes the core bug) ---\n    arr_3d = np.arange(24, dtype=np.int32).reshape((2, 3, 4))\n    dict_3d_1 = {\"nd_array\": arr_3d}\n    dict_3d_2_same = {\"nd_array\": arr_3d.copy()}\n    dict_3d_3_different = {\"nd_array\": arr_3d + 1}\n\n    union_numpy_dict(dict_3d_1, dict_3d_2_same)  # Should pass\n    with pytest.raises(AssertionError, match=\"`nd_array` in tensor_dict1 and tensor_dict2 are not the same object.\"):\n        union_numpy_dict(dict_3d_1, dict_3d_3_different)\n\n    # --- Test Case 3: Nested 2D and 4D object-dtype arrays ---\n    sub_arr1 = np.array([1, 2])\n    sub_arr2 = np.array([3.0, 4.0])\n    # 2D object array\n    arr_2d_obj = np.array([[sub_arr1, \"text\"], [sub_arr2, None]], dtype=object)\n    arr_2d_obj_diff = np.array([[sub_arr1, \"text\"], [sub_arr2, \"other\"]], dtype=object)\n\n    union_numpy_dict({\"data\": arr_2d_obj}, {\"data\": arr_2d_obj.copy()})  # Should pass\n    with pytest.raises(AssertionError):\n        union_numpy_dict({\"data\": arr_2d_obj}, {\"data\": arr_2d_obj_diff})\n\n    # 4D object array to ensure deep recursion is robust\n    arr_4d_obj = np.array([[[[sub_arr1]]], [[[sub_arr2]]]], dtype=object)\n    arr_4d_obj_diff = np.array([[[[sub_arr1]]], [[[np.array([9, 9])]]]], dtype=object)\n\n    union_numpy_dict({\"data\": arr_4d_obj}, {\"data\": arr_4d_obj.copy()})  # Should pass\n    with pytest.raises(AssertionError):\n        union_numpy_dict({\"data\": arr_4d_obj}, {\"data\": arr_4d_obj_diff})\n\n    # --- Test Case 4: Explicit NaN value comparison ---\n    # This verifies that our new _deep_equal logic correctly handles NaNs.\n    nan_arr = np.array([1.0, np.nan, 3.0])\n    dict_nan_1 = {\"data\": nan_arr}\n    dict_nan_2_same = {\"data\": np.array([1.0, np.nan, 3.0])}  # A new array with same values\n    dict_nan_3_different_val = {\"data\": np.array([1.0, 2.0, 3.0])}\n    dict_nan_4_different_pos = {\"data\": np.array([np.nan, 1.0, 3.0])}\n\n    # NaNs in the same position should be considered equal for merging.\n    union_numpy_dict(dict_nan_1, dict_nan_2_same)  # Should pass\n\n    with pytest.raises(AssertionError):\n        union_numpy_dict(dict_nan_1, dict_nan_3_different_val)\n    with pytest.raises(AssertionError):\n        union_numpy_dict(dict_nan_1, dict_nan_4_different_pos)\n\n    # --- Test Case 5: Circular reference handling ---\n    # Create two separate, but structurally identical, circular references.\n    # This should pass without a RecursionError.\n    circ_arr_1 = np.array([None], dtype=object)\n    circ_arr_1[0] = circ_arr_1\n\n    circ_arr_2 = np.array([None], dtype=object)\n    circ_arr_2[0] = circ_arr_2\n\n    union_numpy_dict({\"data\": circ_arr_1}, {\"data\": circ_arr_2})  # Should pass\n\n    # Create a circular reference and a non-circular one.\n    # This should fail with an AssertionError because they are different.\n    non_circ_arr = np.array([None], dtype=object)\n\n    with pytest.raises(AssertionError):\n        union_numpy_dict({\"data\": circ_arr_1}, {\"data\": non_circ_arr})\n\n\ndef test_tensor_dict_constructor():\n    obs = torch.randn(100, 10)\n    act = torch.randn(100, 10, 3)\n    data = DataProto.from_dict(tensors={\"obs\": obs, \"act\": act})\n\n    assert data.batch.batch_size == torch.Size([100])\n\n    with pytest.raises(AssertionError):\n        data = DataProto.from_dict(tensors={\"obs\": obs, \"act\": act}, num_batch_dims=2)\n\n    with pytest.raises(AssertionError):\n        data = DataProto.from_dict(tensors={\"obs\": obs, \"act\": act}, num_batch_dims=3)\n\n\ndef test_tensor_dict_make_iterator():\n    obs = torch.randn(100, 10)\n    labels = [random.choice([\"abc\", \"cde\"]) for _ in range(100)]\n    dataset = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels})\n\n    data_iter_1 = dataset.make_iterator(mini_batch_size=10, epochs=2, seed=1)\n    data_list_1 = []\n    for data in data_iter_1:\n        data_list_1.append(data)\n\n    data_iter_2 = dataset.make_iterator(mini_batch_size=10, epochs=2, seed=1)\n    data_list_2 = []\n    for data in data_iter_2:\n        data_list_2.append(data)\n\n    for data1, data2 in zip(data_list_1, data_list_2, strict=True):\n        assert isinstance(data1, DataProto)\n        assert isinstance(data2, DataProto)\n        result = torch.all(torch.eq(data1.batch[\"obs\"], data2.batch[\"obs\"]))\n        if not result.item():\n            print(data1.batch[\"obs\"])\n            print(data2.batch[\"obs\"])\n            raise AssertionError()\n        non_tensor_result = np.all(np.equal(data1.non_tensor_batch[\"labels\"], data2.non_tensor_batch[\"labels\"]))\n        if not non_tensor_result.item():\n            print(data1.non_tensor_batch[\"labels\"])\n            print(data2.non_tensor_batch[\"labels\"])\n\n\ndef test_reorder():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    labels = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abdce\"})\n    data.reorder(torch.tensor([3, 4, 2, 0, 1, 5]))\n\n    assert torch.all(torch.eq(data.batch[\"obs\"], torch.tensor([4, 5, 3, 1, 2, 6])))\n    assert np.all(data.non_tensor_batch[\"labels\"] == np.array([\"d\", \"e\", \"c\", \"a\", \"b\", \"f\"]))\n    assert data.meta_info == {\"name\": \"abdce\"}\n\n\ndef test_chunk_concat():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    labels = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abdce\"})\n\n    with pytest.raises(AssertionError):\n        data.chunk(5)\n\n    data_split = data.chunk(2)\n    assert len(data_split) == 2\n    assert torch.all(torch.eq(data_split[0].batch[\"obs\"], torch.tensor([1, 2, 3])))\n    assert np.all(data_split[0].non_tensor_batch[\"labels\"] == np.array([\"a\", \"b\", \"c\"]))\n    assert data_split[0].meta_info == {\"name\": \"abdce\"}\n\n    assert torch.all(torch.eq(data_split[1].batch[\"obs\"], torch.tensor([4, 5, 6])))\n    assert np.all(data_split[1].non_tensor_batch[\"labels\"] == np.array([\"d\", \"e\", \"f\"]))\n    assert data_split[1].meta_info == {\"name\": \"abdce\"}\n\n    concat_data = DataProto.concat(data_split)\n    assert torch.all(torch.eq(concat_data.batch[\"obs\"], data.batch[\"obs\"]))\n    assert np.all(concat_data.non_tensor_batch[\"labels\"] == data.non_tensor_batch[\"labels\"])\n    assert concat_data.meta_info == data.meta_info\n\n\ndef test_concat_metrics_from_multiple_workers():\n    \"\"\"Test that concat() properly merges metrics from all workers in distributed training.\"\"\"\n    # Simulate 3 workers each with their own metrics\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n    obs3 = torch.tensor([5, 6])\n\n    # Each worker has different metrics (as list of dict format)\n    worker1_metrics = [{\"loss\": 0.5, \"accuracy\": 0.9}]\n    worker2_metrics = [{\"loss\": 0.6, \"accuracy\": 0.85}]\n    worker3_metrics = [{\"loss\": 0.55, \"accuracy\": 0.88}]\n\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"metrics\": worker1_metrics, \"config_flag\": True})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"metrics\": worker2_metrics, \"config_flag\": True})\n    data3 = DataProto.from_dict(tensors={\"obs\": obs3}, meta_info={\"metrics\": worker3_metrics, \"config_flag\": True})\n\n    # Concat all workers' data\n    concat_data = DataProto.concat([data1, data2, data3])\n\n    # Verify tensors are concatenated\n    assert torch.all(torch.eq(concat_data.batch[\"obs\"], torch.tensor([1, 2, 3, 4, 5, 6])))\n\n    # Verify ALL workers' metrics are flattened to dict of lists\n    expected_metrics = {\"loss\": [0.5, 0.6, 0.55], \"accuracy\": [0.9, 0.85, 0.88]}\n    assert concat_data.meta_info[\"metrics\"] == expected_metrics\n\n    # Verify config flags are preserved from first worker\n    assert concat_data.meta_info[\"config_flag\"] is True\n\n\ndef test_concat_with_empty_and_non_list_meta_info():\n    \"\"\"Test concat() handles edge cases: empty meta_info, non-list values, and None.\"\"\"\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n\n    # Worker 1 has metrics, worker 2 doesn't\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"metrics\": [{\"loss\": 0.5}], \"flag\": True})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"flag\": True})\n\n    concat_data = DataProto.concat([data1, data2])\n\n    # Should flatten worker1's metrics to dict of lists\n    assert concat_data.meta_info[\"metrics\"] == {\"loss\": [0.5]}\n    assert concat_data.meta_info[\"flag\"] is True\n\n    # Test with non-list meta_info value\n    data3 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"single_value\": 42})\n    data4 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"single_value\": 42})\n\n    concat_data2 = DataProto.concat([data3, data4])\n    assert concat_data2.meta_info[\"single_value\"] == 42\n\n\ndef test_concat_first_worker_missing_metrics():\n    \"\"\"Test that metrics from other workers are preserved even when first worker has no metrics.\n\n    This is a critical edge case - the old buggy implementation only checked data[0].meta_info\n    and would lose all metrics if the first worker didn't have any.\n    \"\"\"\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n    obs3 = torch.tensor([5, 6])\n\n    # First worker has NO metrics, but workers 2 and 3 do\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"config_flag\": True})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"metrics\": {\"loss\": 0.6}, \"config_flag\": True})\n    data3 = DataProto.from_dict(tensors={\"obs\": obs3}, meta_info={\"metrics\": {\"loss\": 0.55}, \"config_flag\": True})\n\n    concat_data = DataProto.concat([data1, data2, data3])\n\n    # Should flatten metrics from workers 2 and 3 into dict of lists\n    expected_metrics = {\"loss\": [0.6, 0.55]}\n    assert concat_data.meta_info[\"metrics\"] == expected_metrics\n    assert concat_data.meta_info[\"config_flag\"] is True\n\n\ndef test_concat_non_list_metrics():\n    \"\"\"Test that concat() handles non-list metrics (single dict) correctly.\n\n    In some cases, metrics might be a single dict instead of a list.\n    The implementation should flatten them into a dict of lists.\n    \"\"\"\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n\n    # Metrics as single dict (not wrapped in list)\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"metrics\": {\"loss\": 0.5, \"accuracy\": 0.9}})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"metrics\": {\"loss\": 0.6, \"accuracy\": 0.85}})\n\n    concat_data = DataProto.concat([data1, data2])\n\n    # Should flatten to dict of lists\n    expected_metrics = {\"loss\": [0.5, 0.6], \"accuracy\": [0.9, 0.85]}\n    assert concat_data.meta_info[\"metrics\"] == expected_metrics\n\n\ndef test_concat_merge_different_non_metric_keys():\n    \"\"\"Test that concat() merges non-metric meta_info keys from all workers.\n\n    When different workers have different non-metric keys, all keys should be preserved.\n    This prevents silent data loss and aligns with the docstring stating meta_info is \"merged\".\n    \"\"\"\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n    obs3 = torch.tensor([5, 6])\n\n    # Each worker has some unique non-metric keys\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"config\": \"A\", \"shared_key\": \"X\"})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"extra_key\": \"B\", \"shared_key\": \"X\"})\n    data3 = DataProto.from_dict(tensors={\"obs\": obs3}, meta_info={\"another_key\": \"C\", \"shared_key\": \"X\"})\n\n    concat_data = DataProto.concat([data1, data2, data3])\n\n    # All unique keys should be preserved\n    assert concat_data.meta_info[\"config\"] == \"A\"\n    assert concat_data.meta_info[\"extra_key\"] == \"B\"\n    assert concat_data.meta_info[\"another_key\"] == \"C\"\n    assert concat_data.meta_info[\"shared_key\"] == \"X\"\n\n\ndef test_concat_conflicting_non_metric_keys():\n    \"\"\"Test that concat() raises an assertion error when non-metric keys have conflicting values.\n\n    This ensures data integrity by catching cases where workers have different values\n    for what should be the same configuration parameter.\n    \"\"\"\n    obs1 = torch.tensor([1, 2])\n    obs2 = torch.tensor([3, 4])\n\n    # Same key \"config\" but different values\n    data1 = DataProto.from_dict(tensors={\"obs\": obs1}, meta_info={\"config\": \"A\"})\n    data2 = DataProto.from_dict(tensors={\"obs\": obs2}, meta_info={\"config\": \"B\"})\n\n    # Should raise an assertion error due to conflicting values\n    with pytest.raises(AssertionError, match=\"Conflicting values for meta_info key 'config'\"):\n        DataProto.concat([data1, data2])\n\n\ndef test_pop():\n    obs = torch.randn(100, 10)\n    act = torch.randn(100, 3)\n    dataset = DataProto.from_dict({\"obs\": obs, \"act\": act}, meta_info={\"2\": 2, \"1\": 1})\n    poped_dataset = dataset.pop(batch_keys=[\"obs\"], meta_info_keys=[\"2\"])\n\n    assert poped_dataset.batch.keys() == {\"obs\"}\n    assert poped_dataset.meta_info.keys() == {\"2\"}\n\n    assert dataset.batch.keys() == {\"act\"}\n    assert dataset.meta_info.keys() == {\"1\"}\n\n\ndef test_repeat():\n    # Create a DataProto object with some batch and non-tensor data\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    # Test interleave=True\n    repeated_data_interleave = data.repeat(repeat_times=2, interleave=True)\n    expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]])\n    expected_labels_interleave = [\"a\", \"a\", \"b\", \"b\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_interleave.batch[\"obs\"], expected_obs_interleave))\n    assert (repeated_data_interleave.non_tensor_batch[\"labels\"] == expected_labels_interleave).all()\n    assert repeated_data_interleave.meta_info == {\"info\": \"test_info\"}\n\n    # Test interleave=False\n    repeated_data_no_interleave = data.repeat(repeat_times=2, interleave=False)\n    expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]])\n    expected_labels_no_interleave = [\"a\", \"b\", \"c\", \"a\", \"b\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_no_interleave.batch[\"obs\"], expected_obs_no_interleave))\n    assert (repeated_data_no_interleave.non_tensor_batch[\"labels\"] == expected_labels_no_interleave).all()\n    assert repeated_data_no_interleave.meta_info == {\"info\": \"test_info\"}\n\n\ndef test_dataproto_pad_unpad():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    from verl.protocol import pad_dataproto_to_divisor, unpad_dataproto\n\n    padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=2)\n    assert pad_size == 1\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2]])\n    expected_labels = [\"a\", \"b\", \"c\", \"a\"]\n\n    assert torch.all(torch.eq(padded_data.batch[\"obs\"], expected_obs))\n    assert (padded_data.non_tensor_batch[\"labels\"] == expected_labels).all()\n    assert padded_data.meta_info == {\"info\": \"test_info\"}\n\n    unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data.batch[\"obs\"], obs))\n    assert (unpadd_data.non_tensor_batch[\"labels\"] == labels).all()\n    assert unpadd_data.meta_info == {\"info\": \"test_info\"}\n\n    padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=3)\n    assert pad_size == 0\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    expected_labels = [\"a\", \"b\", \"c\"]\n\n    assert torch.all(torch.eq(padded_data.batch[\"obs\"], expected_obs))\n    assert (padded_data.non_tensor_batch[\"labels\"] == expected_labels).all()\n    assert padded_data.meta_info == {\"info\": \"test_info\"}\n\n    unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data.batch[\"obs\"], obs))\n    assert (unpadd_data.non_tensor_batch[\"labels\"] == labels).all()\n    assert unpadd_data.meta_info == {\"info\": \"test_info\"}\n\n    padded_data, pad_size = pad_dataproto_to_divisor(data, size_divisor=7)\n    assert pad_size == 4\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]])\n    expected_labels = [\"a\", \"b\", \"c\", \"a\", \"b\", \"c\", \"a\"]\n    assert torch.all(torch.eq(padded_data.batch[\"obs\"], expected_obs))\n    assert (padded_data.non_tensor_batch[\"labels\"] == expected_labels).all()\n    assert padded_data.meta_info == {\"info\": \"test_info\"}\n\n    unpadd_data = unpad_dataproto(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data.batch[\"obs\"], obs))\n    assert (unpadd_data.non_tensor_batch[\"labels\"] == labels).all()\n    assert unpadd_data.meta_info == {\"info\": \"test_info\"}\n\n\ndef test_dataproto_fold_unfold():\n    from verl.protocol import DataProto, fold_batch_dim, unfold_batch_dim\n\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    data1 = data.repeat(repeat_times=2, interleave=True)\n\n    data2 = fold_batch_dim(data1, new_batch_size=3)\n\n    torch.testing.assert_close(data2.batch[\"obs\"], torch.tensor([[[1, 2], [1, 2]], [[3, 4], [3, 4]], [[5, 6], [5, 6]]]))\n    assert (data2.non_tensor_batch[\"labels\"] == [[\"a\", \"a\"], [\"b\", \"b\"], [\"c\", \"c\"]]).all()\n\n    data2.reorder(indices=torch.tensor([1, 2, 0]))\n\n    data3 = unfold_batch_dim(data2, batch_dims=2)\n\n    torch.testing.assert_close(data3.batch[\"obs\"], torch.tensor([[3, 4], [3, 4], [5, 6], [5, 6], [1, 2], [1, 2]]))\n    assert (data3.non_tensor_batch[\"labels\"] == [\"b\", \"b\", \"c\", \"c\", \"a\", \"a\"]).all()\n    assert data3.meta_info == {\"info\": \"test_info\"}\n\n\ndef test_torch_save_data_proto():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n    data.save_to_disk(\"test_data.pt\")\n    loaded_data = DataProto.load_from_disk(\"test_data.pt\")\n\n    assert torch.all(torch.eq(loaded_data.batch[\"obs\"], data.batch[\"obs\"]))\n    assert (loaded_data.non_tensor_batch[\"labels\"] == data.non_tensor_batch[\"labels\"]).all()\n    assert loaded_data.meta_info == data.meta_info\n\n    import os\n\n    os.remove(\"test_data.pt\")\n\n\ndef test_len():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = np.array([\"a\", \"b\", \"c\"], dtype=object)\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    assert len(data) == 3\n\n    data = DataProto(batch=None, non_tensor_batch={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    assert len(data) == 3\n\n    data = DataProto(batch=None, non_tensor_batch={}, meta_info={\"info\": \"test_info\"})\n\n    assert len(data) == 0\n\n    data = DataProto(batch=None, non_tensor_batch=None, meta_info={\"info\": \"test_info\"})\n\n    assert len(data) == 0\n\n\ndef test_dataproto_index():\n    data_len = 100\n    idx_num = 10\n\n    obs = torch.randn(data_len, 10)\n    labels = [random.choice([\"abc\", \"cde\"]) for _ in range(data_len)]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels})\n    labels_np = np.array(labels)\n\n    idx_np_int = np.random.randint(0, data_len, size=(idx_num,))\n    result_np_int = data[idx_np_int]\n    assert result_np_int.batch.keys() == data.batch.keys()\n    assert result_np_int.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_np_int.batch[\"obs\"].shape[0] == idx_num\n    assert result_np_int.non_tensor_batch[\"labels\"].shape[0] == idx_num\n    assert np.array_equal(result_np_int.batch[\"obs\"].cpu().numpy(), obs[idx_np_int].numpy())\n    assert np.array_equal(result_np_int.non_tensor_batch[\"labels\"], labels_np[idx_np_int])\n\n    idx_torch_int = torch.randint(0, data_len, size=(idx_num,))\n    result_torch_int = data[idx_torch_int]\n    assert result_torch_int.batch.keys() == data.batch.keys()\n    assert result_torch_int.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_torch_int.batch[\"obs\"].shape[0] == idx_num\n    assert result_torch_int.non_tensor_batch[\"labels\"].shape[0] == idx_num\n    assert np.array_equal(result_torch_int.batch[\"obs\"].cpu().numpy(), obs[idx_torch_int].cpu().numpy())\n    assert np.array_equal(result_torch_int.non_tensor_batch[\"labels\"], labels_np[idx_torch_int.cpu().numpy()])\n\n    idx_list_int = [np.random.randint(0, data_len) for _ in range(idx_num)]\n    result_list_int = data[idx_list_int]\n    assert result_list_int.batch.keys() == data.batch.keys()\n    assert result_list_int.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_list_int.batch[\"obs\"].shape[0] == idx_num\n    assert result_list_int.non_tensor_batch[\"labels\"].shape[0] == idx_num\n    assert np.array_equal(result_list_int.batch[\"obs\"].cpu().numpy(), obs[idx_list_int].cpu().numpy())\n    assert np.array_equal(result_list_int.non_tensor_batch[\"labels\"], labels_np[idx_list_int])\n\n    idx_np_bool = np.random.randint(0, 2, size=(data_len,), dtype=bool)\n    result_np_bool = data[idx_np_bool]\n    assert result_np_bool.batch.keys() == data.batch.keys()\n    assert result_np_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_np_bool.batch[\"obs\"].shape[0] == idx_np_bool.sum()\n    assert result_np_bool.non_tensor_batch[\"labels\"].shape[0] == idx_np_bool.sum()\n    assert np.array_equal(result_np_bool.batch[\"obs\"].cpu().numpy(), obs[idx_np_bool].cpu().numpy())\n    assert np.array_equal(result_np_bool.non_tensor_batch[\"labels\"], labels_np[idx_np_bool])\n\n    idx_torch_bool = torch.randint(0, 2, size=(data_len,), dtype=torch.bool)\n    result_torch_bool = data[idx_torch_bool]\n    assert result_torch_bool.batch.keys() == data.batch.keys()\n    assert result_torch_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_torch_bool.batch[\"obs\"].shape[0] == idx_torch_bool.sum().item()\n    assert result_torch_bool.non_tensor_batch[\"labels\"].shape[0] == idx_torch_bool.sum().item()\n    assert np.array_equal(result_torch_bool.batch[\"obs\"].cpu().numpy(), obs[idx_torch_bool].cpu().numpy())\n    assert np.array_equal(result_torch_bool.non_tensor_batch[\"labels\"], labels_np[idx_torch_bool])\n\n    idx_list_bool = [np.random.randint(0, 2, dtype=bool) for _ in range(data_len)]\n    result_list_bool = data[idx_list_bool]\n    assert result_list_bool.batch.keys() == data.batch.keys()\n    assert result_list_bool.non_tensor_batch.keys() == data.non_tensor_batch.keys()\n    assert result_list_bool.batch[\"obs\"].shape[0] == sum(idx_list_bool)\n    assert result_list_bool.non_tensor_batch[\"labels\"].shape[0] == sum(idx_list_bool)\n    assert np.array_equal(result_list_bool.batch[\"obs\"].cpu().numpy(), obs[idx_list_bool].cpu().numpy())\n    assert np.array_equal(result_list_bool.non_tensor_batch[\"labels\"], labels_np[idx_list_bool])\n\n\ndef test_old_vs_new_from_single_dict():\n    class CustomProto(DataProto):\n        \"\"\"Uses the new, fixed from_single_dict.\"\"\"\n\n        pass\n\n    class OriginProto(DataProto):\n        \"\"\"Mimics the *old* from_single_dict (always returns a DataProto).\"\"\"\n\n        @classmethod\n        def from_single_dict(cls, data, meta_info=None, auto_padding=False):\n            tensors, non_tensors = {}, {}\n            for k, v in data.items():\n                if torch.is_tensor(v):\n                    tensors[k] = v\n                else:\n                    non_tensors[k] = v\n            # always calls DataProto.from_dict, ignoring `cls`\n            return DataProto.from_dict(\n                tensors=tensors,\n                non_tensors=non_tensors,\n                meta_info=meta_info,\n                auto_padding=auto_padding,\n            )\n\n    sample = {\"x\": torch.tensor([0])}\n\n    orig = OriginProto.from_single_dict(sample)\n    # old behavior: always DataProto, not a CustomOriginProto\n    assert type(orig) is DataProto\n    assert type(orig) is not OriginProto\n\n    cust = CustomProto.from_single_dict(sample)\n    # new behavior: respects subclass\n    assert type(cust) is CustomProto\n\n\ndef test_dataproto_no_batch():\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n    selected = data.select(non_tensor_batch_keys=[\"labels\"])\n    assert (selected.non_tensor_batch[\"labels\"] == labels).all()\n    pop_data = data.pop(non_tensor_batch_keys=[\"labels\"])\n    assert (pop_data.non_tensor_batch[\"labels\"] == labels).all()\n    assert data.non_tensor_batch == {}\n\n\ndef test_sample_level_repeat():\n    # Create a DataProto object with some batch and non-tensor data\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"info\": \"test_info\"})\n\n    # list\n    repeated_data_interleave = data.sample_level_repeat(repeat_times=[3, 1, 2])\n    expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]])\n    expected_labels_interleave = [\"a\", \"a\", \"a\", \"b\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_interleave.batch[\"obs\"], expected_obs_interleave))\n    assert (repeated_data_interleave.non_tensor_batch[\"labels\"] == expected_labels_interleave).all()\n    assert repeated_data_interleave.meta_info == {\"info\": \"test_info\"}\n\n    # torch.tensor\n    repeated_data_no_interleave = data.sample_level_repeat(repeat_times=torch.tensor([1, 2, 3]))\n    expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]])\n    expected_labels_no_interleave = [\"a\", \"b\", \"b\", \"c\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_no_interleave.batch[\"obs\"], expected_obs_no_interleave))\n    assert (repeated_data_no_interleave.non_tensor_batch[\"labels\"] == expected_labels_no_interleave).all()\n    assert repeated_data_no_interleave.meta_info == {\"info\": \"test_info\"}\n\n\ndef test_dataproto_unfold_column_chunks():\n    obs1 = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])\n    obs2 = torch.tensor([[1, 2], [5, 6], [9, 10]])\n\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(\n        tensors={\"obs1\": obs1, \"obs2\": obs2}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abc\"}\n    )\n    ret = data.unfold_column_chunks(2, split_keys=[\"obs1\"])\n\n    expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])\n    expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]])\n    expect_labels = [\"a\", \"a\", \"b\", \"b\", \"c\", \"c\"]\n    assert torch.all(torch.eq(ret.batch[\"obs1\"], expect_obs1))\n    assert torch.all(torch.eq(ret.batch[\"obs2\"], expect_obs2))\n    assert (ret.non_tensor_batch[\"labels\"] == expect_labels).all()\n    assert ret.meta_info == {\"name\": \"abc\"}\n\n    obs1 = torch.tensor([[1, 2, 3, 4], [5, 6, 7, 8], [9, 10, 11, 12]])\n    obs2 = torch.tensor([[1, 2], [5, 6], [9, 10]])\n\n    labels = [[\"a1\", \"a2\"], [\"b1\", \"b2\"], [\"c1\", \"c2\"]]\n    data = DataProto.from_dict(\n        tensors={\"obs1\": obs1, \"obs2\": obs2}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abc\"}\n    )\n    ret = data.unfold_column_chunks(2, split_keys=[\"obs1\", \"labels\"])\n\n    expect_obs1 = torch.tensor([[1, 2], [3, 4], [5, 6], [7, 8], [9, 10], [11, 12]])\n    expect_obs2 = torch.tensor([[1, 2], [1, 2], [5, 6], [5, 6], [9, 10], [9, 10]])\n    expect_labels = [[\"a1\"], [\"a2\"], [\"b1\"], [\"b2\"], [\"c1\"], [\"c2\"]]\n    assert torch.all(torch.eq(ret.batch[\"obs1\"], expect_obs1))\n    assert torch.all(torch.eq(ret.batch[\"obs2\"], expect_obs2))\n    assert (ret.non_tensor_batch[\"labels\"] == expect_labels).all()\n    assert ret.meta_info == {\"name\": \"abc\"}\n\n    obs1 = torch.tensor(\n        [[[1, 1], [2, 2], [3, 3], [4, 4]], [[5, 5], [6, 6], [7, 7], [8, 8]], [[9, 9], [10, 10], [11, 11], [12, 12]]]\n    )\n    obs2 = torch.tensor([[[1, 1], [2, 2]], [[5, 5], [6, 6]], [[9, 9], [10, 10]]])\n\n    labels = [\"a\", \"b\", \"c\"]\n    data = DataProto.from_dict(\n        tensors={\"obs1\": obs1, \"obs2\": obs2}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abc\"}\n    )\n    ret = data.unfold_column_chunks(2, split_keys=[\"obs1\"])\n\n    expect_obs1 = torch.tensor(\n        [\n            [[1, 1], [2, 2]],\n            [[3, 3], [4, 4]],\n            [[5, 5], [6, 6]],\n            [[7, 7], [8, 8]],\n            [[9, 9], [10, 10]],\n            [[11, 11], [12, 12]],\n        ]\n    )\n    expect_obs2 = torch.tensor(\n        [[[1, 1], [2, 2]], [[1, 1], [2, 2]], [[5, 5], [6, 6]], [[5, 5], [6, 6]], [[9, 9], [10, 10]], [[9, 9], [10, 10]]]\n    )\n    expect_labels = [\"a\", \"a\", \"b\", \"b\", \"c\", \"c\"]\n    assert torch.all(torch.eq(ret.batch[\"obs1\"], expect_obs1))\n    assert torch.all(torch.eq(ret.batch[\"obs2\"], expect_obs2))\n    assert (ret.non_tensor_batch[\"labels\"] == expect_labels).all()\n    assert ret.meta_info == {\"name\": \"abc\"}\n\n\ndef test_dataproto_chunk_after_index():\n    data_len = 4\n    obs = torch.randn(data_len, 4)\n    labels = [f\"label_{i}\" for i in range(data_len)]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abc\"})\n\n    # Test with boolean numpy array\n    bool_mask = np.array([True, False, True, False])\n    selected = data[bool_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)  # int or List[int]\n\n    # Test with integer numpy array\n    int_mask = np.array([0, 2])\n    selected = data[int_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)\n\n    # Test with boolean list\n    list_mask = [True, False, True, False]\n    selected = data[list_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)\n\n    # Test with list\n    list_mask = [0, 2]\n    selected = data[list_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)\n\n    # Test with torch tensor (bool)\n    torch_bool_mask = torch.tensor([True, False, True, False])\n    selected = data[torch_bool_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)\n\n    # Test with torch tensor (int)\n    torch_int_mask = torch.tensor([0, 2])\n    selected = data[torch_int_mask]\n    assert isinstance(selected.batch.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch.batch_size)\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    labels = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"labels\": labels}, meta_info={\"name\": \"abdce\"})\n    output = data.to_tensordict()\n\n    assert torch.all(torch.eq(output[\"obs\"], obs)).item()\n    assert output[\"labels\"] == labels\n    assert output[\"name\"] == \"abdce\"\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_from_tensordict():\n    tensor_dict = {\n        \"obs\": torch.tensor([1, 2, 3, 4, 5, 6]),\n        \"labels\": [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"],\n    }\n    non_tensor_dict = {\"name\": \"abdce\"}\n    tensordict = tu.get_tensordict(tensor_dict, non_tensor_dict)\n    data = DataProto.from_tensordict(tensordict)\n\n    assert data.non_tensor_batch[\"labels\"].tolist() == tensor_dict[\"labels\"]\n    assert torch.all(torch.eq(data.batch[\"obs\"], tensor_dict[\"obs\"])).item()\n    assert data.meta_info[\"name\"] == \"abdce\"\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict_with_nested_lists():\n    \"\"\"Test converting DataProto with nested lists to TensorDict (lists of lists).\"\"\"\n    obs = torch.tensor([1, 2, 3])\n    # Simulate turn_scores or tool_rewards: array of lists with varying lengths\n    turn_scores = [[], [0.5, 0.8], [0.9]]\n\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"turn_scores\": turn_scores})\n\n    # This should not raise an error\n    tensordict_output = data.to_tensordict()\n\n    # Verify the data is preserved\n    assert torch.all(torch.eq(tensordict_output[\"obs\"], obs)).item()\n    # Verify nested structure is accessible (TensorDict wraps NonTensorStack as LinkedList)\n    retrieved_scores = tensordict_output[\"turn_scores\"]\n    assert len(retrieved_scores) == len(turn_scores)\n    # Verify content matches\n    assert list(retrieved_scores[0]) == []\n    assert list(retrieved_scores[1]) == [0.5, 0.8]\n    assert list(retrieved_scores[2]) == [0.9]\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict_with_nested_dicts():\n    \"\"\"Test converting DataProto with lists of dicts to TensorDict.\"\"\"\n    obs = torch.tensor([1, 2, 3])\n    # Simulate reward_extra_info: array of dicts\n    reward_extra_info = [{\"acc\": 1.0}, {\"acc\": 0.0}, {\"acc\": 1.0}]\n\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"reward_extra_info\": reward_extra_info})\n\n    # This should not raise an error - this was the original bug\n    tensordict_output = data.to_tensordict()\n\n    # Verify the data is preserved\n    assert torch.all(torch.eq(tensordict_output[\"obs\"], obs)).item()\n    # Verify nested dicts are accessible\n    retrieved_info = tensordict_output[\"reward_extra_info\"]\n    assert len(retrieved_info) == len(reward_extra_info)\n    # Verify content matches\n    for i, expected_dict in enumerate(reward_extra_info):\n        assert dict(retrieved_info[i]) == expected_dict\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict_with_complex_nested_structures():\n    \"\"\"Test converting DataProto with complex nested structures (lists of lists of dicts).\"\"\"\n    obs = torch.tensor([1, 2, 3])\n    # Simulate raw_prompt: array of lists containing dicts\n    raw_prompt = [\n        [{\"content\": \"Question 1\", \"role\": \"user\"}],\n        [{\"content\": \"Question 2\", \"role\": \"user\"}, {\"content\": \"Answer 2\", \"role\": \"assistant\"}],\n        [{\"content\": \"Question 3\", \"role\": \"user\"}],\n    ]\n\n    data = DataProto.from_dict(tensors={\"obs\": obs}, non_tensors={\"raw_prompt\": raw_prompt})\n\n    # This should not raise an error\n    tensordict_output = data.to_tensordict()\n\n    # Verify the data is preserved\n    assert torch.all(torch.eq(tensordict_output[\"obs\"], obs)).item()\n    # Verify complex nested structure is accessible\n    retrieved_prompt = tensordict_output[\"raw_prompt\"]\n    assert len(retrieved_prompt) == len(raw_prompt)\n    # Spot check: verify first prompt has correct structure\n    assert len(retrieved_prompt[0]) == 1\n    assert dict(retrieved_prompt[0][0]) == {\"content\": \"Question 1\", \"role\": \"user\"}\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict_and_back_with_nested_data():\n    \"\"\"Test round-trip conversion: DataProto → TensorDict → DataProto with nested structures.\"\"\"\n    obs = torch.tensor([1, 2, 3, 4])\n    labels = [\"a\", \"b\", \"c\", \"d\"]\n\n    # Multiple types of nested structures\n    turn_scores = [[], [0.5], [0.8, 0.9], [0.7]]\n    reward_extra_info = [\n        {\"acc\": 1.0, \"loss\": 0.1},\n        {\"acc\": 0.5, \"loss\": 0.3},\n        {\"acc\": 1.0, \"loss\": 0.05},\n        {\"acc\": 0.0, \"loss\": 0.9},\n    ]\n    raw_prompt = [\n        [{\"content\": \"Q1\", \"role\": \"user\"}],\n        [{\"content\": \"Q2\", \"role\": \"user\"}],\n        [{\"content\": \"Q3\", \"role\": \"user\"}, {\"content\": \"A3\", \"role\": \"assistant\"}],\n        [{\"content\": \"Q4\", \"role\": \"user\"}],\n    ]\n\n    # Create original DataProto\n    original_data = DataProto.from_dict(\n        tensors={\"obs\": obs},\n        non_tensors={\n            \"labels\": labels,\n            \"turn_scores\": turn_scores,\n            \"reward_extra_info\": reward_extra_info,\n            \"raw_prompt\": raw_prompt,\n        },\n        meta_info={\"experiment\": \"test_nested\"},\n    )\n\n    # Convert to TensorDict\n    tensordict_output = original_data.to_tensordict()\n\n    # Convert back to DataProto\n    reconstructed_data = DataProto.from_tensordict(tensordict_output)\n\n    # Verify tensors are preserved\n    assert torch.all(torch.eq(reconstructed_data.batch[\"obs\"], obs)).item()\n\n    # Verify non-tensor data is preserved\n    assert reconstructed_data.non_tensor_batch[\"labels\"].tolist() == labels\n\n    # Verify nested structures are preserved\n    assert len(reconstructed_data.non_tensor_batch[\"turn_scores\"]) == len(turn_scores)\n    for orig, recon in zip(turn_scores, reconstructed_data.non_tensor_batch[\"turn_scores\"], strict=True):\n        assert list(orig) == list(recon)\n\n    assert len(reconstructed_data.non_tensor_batch[\"reward_extra_info\"]) == len(reward_extra_info)\n    for orig, recon in zip(reward_extra_info, reconstructed_data.non_tensor_batch[\"reward_extra_info\"], strict=True):\n        assert orig == recon\n\n    assert len(reconstructed_data.non_tensor_batch[\"raw_prompt\"]) == len(raw_prompt)\n    for orig, recon in zip(raw_prompt, reconstructed_data.non_tensor_batch[\"raw_prompt\"], strict=True):\n        assert orig == list(recon)\n\n    # Verify meta_info is preserved\n    assert reconstructed_data.meta_info[\"experiment\"] == \"test_nested\"\n\n\n@pytest.mark.skipif(\n    parse_version(tensordict.__version__) < parse_version(\"0.10\"), reason=\"requires at least tensordict 0.10\"\n)\ndef test_to_tensordict_agent_loop_scenario():\n    \"\"\"Test the exact scenario from agent loop: DataProto with tool rewards, acc, etc.\n\n    This test reproduces the exact error from the agent loop where nested structures\n    (lists of lists, lists of dicts) failed to convert to TensorDict.\n    \"\"\"\n    # Simulate real agent loop data structure\n    prompts = torch.tensor([[1, 2, 3], [4, 5, 6]])\n    responses = torch.tensor([[7, 8], [9, 10]])\n\n    # Non-tensor data with nested structures from agent loop\n    data_source = [\"lighteval/MATH\", \"lighteval/MATH\"]\n    uid = [\"uuid-1\", \"uuid-2\"]\n    num_turns = np.array([2, 4], dtype=np.int32)\n    acc = np.array([1.0, 0.0])\n    turn_scores = [[], [0.5, 0.8]]  # Lists of varying lengths\n    reward_extra_info = [{\"acc\": 1.0}, {\"acc\": 0.0}]  # List of dicts\n    raw_prompt = [\n        [{\"content\": \"Compute 4 @ 2\", \"role\": \"user\"}],\n        [{\"content\": \"Compute 8 @ 7\", \"role\": \"user\"}],\n    ]\n    tool_rewards = [[0.0], []]  # List of lists\n\n    data = DataProto.from_dict(\n        tensors={\"prompts\": prompts, \"responses\": responses},\n        non_tensors={\n            \"data_source\": data_source,\n            \"uid\": uid,\n            \"num_turns\": num_turns,\n            \"acc\": acc,\n            \"turn_scores\": turn_scores,\n            \"reward_extra_info\": reward_extra_info,\n            \"raw_prompt\": raw_prompt,\n            \"tool_rewards\": tool_rewards,\n        },\n        meta_info={\"global_steps\": 42},\n    )\n\n    # THE KEY TEST: This should not raise ValueError about TensorDict conversion\n    tensordict_output = data.to_tensordict()\n\n    # Verify tensors are accessible\n    assert torch.all(torch.eq(tensordict_output[\"prompts\"], prompts)).item()\n    assert torch.all(torch.eq(tensordict_output[\"responses\"], responses)).item()\n\n    # Verify all nested structures are accessible (content check, not type check)\n    assert len(tensordict_output[\"turn_scores\"]) == 2\n    assert list(tensordict_output[\"turn_scores\"][0]) == []\n    assert list(tensordict_output[\"turn_scores\"][1]) == [0.5, 0.8]\n\n    assert len(tensordict_output[\"reward_extra_info\"]) == 2\n    assert dict(tensordict_output[\"reward_extra_info\"][0]) == {\"acc\": 1.0}\n\n    assert len(tensordict_output[\"raw_prompt\"]) == 2\n    assert dict(tensordict_output[\"raw_prompt\"][0][0]) == {\"content\": \"Compute 4 @ 2\", \"role\": \"user\"}\n\n    assert len(tensordict_output[\"tool_rewards\"]) == 2\n    assert list(tensordict_output[\"tool_rewards\"][0]) == [0.0]\n    assert list(tensordict_output[\"tool_rewards\"][1]) == []\n\n    # Verify round-trip conversion works perfectly\n    reconstructed = DataProto.from_tensordict(tensordict_output)\n    assert len(reconstructed) == 2\n    assert reconstructed.meta_info[\"global_steps\"] == 42\n    assert torch.all(torch.eq(reconstructed.batch[\"prompts\"], prompts)).item()\n\n\ndef test_serialize_deserialize_single_tensor():\n    \"\"\"Test serialization and deserialization of a single tensor\"\"\"\n    # Create test tensor\n    original_tensor = torch.randn(3, 4, 5)\n\n    # Serialize\n    dtype, shape, data = serialize_single_tensor(original_tensor)\n\n    # Deserialize\n    reconstructed_tensor = deserialize_single_tensor((dtype, shape, data))\n\n    # Verify results\n    assert torch.allclose(original_tensor, reconstructed_tensor)\n    assert original_tensor.shape == reconstructed_tensor.shape\n    assert original_tensor.dtype == reconstructed_tensor.dtype\n\n\ndef test_serialize_deserialize_tensordict_regular_tensors():\n    \"\"\"Test serialization and deserialization of TensorDict with regular tensors\"\"\"\n    # Create test data\n    batch_size = (5, 3)\n    tensor1 = torch.randn(*batch_size, 4)\n    tensor2 = torch.randint(0, 10, (*batch_size, 2))\n\n    # Create TensorDict\n    original_tensordict = TensorDict({\"tensor1\": tensor1, \"tensor2\": tensor2}, batch_size=batch_size)\n\n    # Serialize\n    batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict)\n\n    # Deserialize\n    reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items))\n\n    # Verify results\n    assert original_tensordict.batch_size == reconstructed_tensordict.batch_size\n    assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys())\n\n    for key in original_tensordict.keys():\n        original_tensor = original_tensordict[key]\n        reconstructed_tensor = reconstructed_tensordict[key]\n\n        assert torch.allclose(original_tensor, reconstructed_tensor)\n        assert original_tensor.shape == reconstructed_tensor.shape\n        assert original_tensor.dtype == reconstructed_tensor.dtype\n\n\ndef test_serialize_deserialize_tensordict_nested_tensors():\n    \"\"\"Test serialization and deserialization of TensorDict with nested tensors\"\"\"\n    # Create nested tensor\n    tensor_list = [torch.randn(2, 3), torch.randn(3, 4), torch.randn(1, 5)]\n    nested_tensor = torch.nested.as_nested_tensor(tensor_list)\n\n    # Create regular tensor for comparison\n    regular_tensor = torch.randn(3, 4, 5)\n\n    # Create TensorDict\n    original_tensordict = TensorDict({\"nested\": nested_tensor, \"regular\": regular_tensor}, batch_size=(3,))\n\n    # Serialize\n    batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict)\n\n    # Deserialize\n    reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items))\n\n    # Verify results\n    assert original_tensordict.batch_size == reconstructed_tensordict.batch_size\n    assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys())\n\n    # Verify regular tensor\n    original_regular = original_tensordict[\"regular\"]\n    reconstructed_regular = reconstructed_tensordict[\"regular\"]\n\n    assert torch.allclose(original_regular, reconstructed_regular)\n    assert original_regular.shape == reconstructed_regular.shape\n    assert original_regular.dtype == reconstructed_regular.dtype\n\n    # Verify nested tensor\n    original_nested = original_tensordict[\"nested\"]\n    reconstructed_nested = reconstructed_tensordict[\"nested\"]\n\n    # Check if it's a nested tensor\n    assert original_nested.is_nested\n    assert reconstructed_nested.is_nested\n\n    # Check layout\n    assert original_nested.layout == reconstructed_nested.layout\n\n    # Check each tensor after unbinding\n    original_unbind = original_nested.unbind()\n    reconstructed_unbind = reconstructed_nested.unbind()\n\n    assert len(original_unbind) == len(reconstructed_unbind)\n\n    for orig, recon in zip(original_unbind, reconstructed_unbind, strict=False):\n        assert torch.allclose(orig, recon)\n        assert orig.shape == recon.shape\n        assert orig.dtype == recon.dtype\n\n\ndef test_serialize_deserialize_tensordict_mixed_types():\n    \"\"\"Test serialization and deserialization of TensorDict with mixed tensor types\"\"\"\n    # Create tensors with different data types\n    float_tensor = torch.randn(2, 3).float()\n    double_tensor = torch.randn(2, 3).double()\n    int_tensor = torch.randint(0, 10, (2, 3)).int()\n    long_tensor = torch.randint(0, 10, (2, 3)).long()\n    bool_tensor = torch.tensor([[True, False], [False, True]])\n    bfloat16_tensor = torch.randn(2, 3).bfloat16()\n\n    # Add fp8 tensor (if available)\n    # Note: FP8 is not natively supported in all PyTorch versions\n    # We'll check if it's available and conditionally include it\n    has_fp8 = hasattr(torch, \"float8_e5m2\") or hasattr(torch, \"float8_e4m3fn\")\n    if has_fp8:\n        try:\n            # Try to create an FP8 tensor (implementation may vary)\n            # This is a placeholder - actual FP8 support might require specific hardware\n            fp8_tensor = torch.randn(2, 3)\n            if hasattr(torch, \"float8_e5m2\"):\n                fp8_tensor = fp8_tensor.to(torch.float8_e5m2)\n            elif hasattr(torch, \"float8_e4m3fn\"):\n                fp8_tensor = fp8_tensor.to(torch.float8_e4m3fn)\n        except Exception:\n            has_fp8 = False\n\n    # Create nested tensor\n    tensor_list = [\n        torch.randn(2, 3),\n        torch.randn(3, 4),\n    ]\n    nested_tensor = torch.nested.as_nested_tensor(tensor_list)\n\n    # Create TensorDict with all available types\n    tensordict_data = {\n        \"float\": float_tensor,\n        \"double\": double_tensor,\n        \"int\": int_tensor,\n        \"long\": long_tensor,\n        \"bool\": bool_tensor,\n        \"bfloat16\": bfloat16_tensor,\n        \"nested\": nested_tensor,\n    }\n\n    # Conditionally add fp8 tensor if available\n    if has_fp8:\n        tensordict_data[\"fp8\"] = fp8_tensor\n\n    original_tensordict = TensorDict(\n        tensordict_data,\n        batch_size=(2,),\n    )\n\n    # Serialize\n    batch_size_serialized, device, encoded_items = serialize_tensordict(original_tensordict)\n\n    # Deserialize\n    reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device, encoded_items))\n\n    # Verify results\n    assert original_tensordict.batch_size == reconstructed_tensordict.batch_size\n    assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys())\n\n    for key in original_tensordict.keys():\n        original_tensor = original_tensordict[key]\n        reconstructed_tensor = reconstructed_tensordict[key]\n\n        if original_tensor.is_nested:\n            # For nested tensors, check each tensor after unbinding\n            original_unbind = original_tensor.unbind()\n            reconstructed_unbind = reconstructed_tensor.unbind()\n\n            assert len(original_unbind) == len(reconstructed_unbind)\n\n            for orig, recon in zip(original_unbind, reconstructed_unbind, strict=False):\n                assert torch.allclose(orig, recon, equal_nan=True)\n                assert orig.shape == recon.shape\n                assert orig.dtype == recon.dtype\n        else:\n            # For regular tensors, compare directly\n            assert torch.all(original_tensor == reconstructed_tensor)\n            assert original_tensor.shape == reconstructed_tensor.shape\n            assert original_tensor.dtype == reconstructed_tensor.dtype\n\n\ndef test_serialize_deserialize_tensordict_with_device():\n    \"\"\"Test serialization and deserialization of TensorDict with device information\"\"\"\n    # Create test data\n    batch_size = (2, 3)\n    tensor1 = torch.randn(*batch_size, 4)\n    tensor2 = torch.randint(0, 10, (*batch_size, 2))\n\n    # Create TensorDict with device information\n    device = \"cpu\"\n    original_tensordict = TensorDict({\"tensor1\": tensor1, \"tensor2\": tensor2}, batch_size=batch_size, device=device)\n\n    # Serialize\n    batch_size_serialized, device_serialized, encoded_items = serialize_tensordict(original_tensordict)\n\n    # Deserialize\n    reconstructed_tensordict = deserialize_tensordict((batch_size_serialized, device_serialized, encoded_items))\n\n    # Verify results\n    assert original_tensordict.batch_size == reconstructed_tensordict.batch_size\n    assert str(original_tensordict.device) == str(reconstructed_tensordict.device)\n    assert set(original_tensordict.keys()) == set(reconstructed_tensordict.keys())\n\n    for key in original_tensordict.keys():\n        original_tensor = original_tensordict[key]\n        reconstructed_tensor = reconstructed_tensordict[key]\n\n        assert torch.allclose(original_tensor.cpu(), reconstructed_tensor.cpu())\n        assert original_tensor.shape == reconstructed_tensor.shape\n        assert original_tensor.dtype == reconstructed_tensor.dtype\n\n\ndef test_serialize_dataproto_with_empty_tensordict():\n    \"\"\"Tests that serializing a DataProto with an empty TensorDict does not crash.\n\n    This test verifies the fix for the torch.cat error that occurs when calling\n    consolidate() on an empty TensorDict during serialization.\n    \"\"\"\n    import pickle\n\n    # This test requires tensordict >= 0.5.0 to trigger the code path\n    if parse_version(tensordict.__version__) < parse_version(\"0.5.0\"):\n        pytest.skip(\"Test requires tensordict>=0.5.0\")\n\n    # Create a DataProto with an empty TensorDict but with a batch size\n    empty_td = TensorDict({}, batch_size=[10])\n    data = DataProto(batch=empty_td)\n\n    # This would crash before the fix with:\n    # RuntimeError: torch.cat(): expected a non-empty list of Tensors\n    try:\n        serialized_data = pickle.dumps(data)\n    except Exception as e:\n        pytest.fail(f\"Serializing DataProto with empty TensorDict failed with: {e}\")\n\n    # Verify deserialization works as expected\n    deserialized_data = pickle.loads(serialized_data)\n    assert len(deserialized_data.batch.keys()) == 0\n    assert deserialized_data.batch.batch_size == torch.Size([10])\n"
  },
  {
    "path": "tests/test_protocol_v2_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nReplace DataProto with raw TensorDict\n\"\"\"\n\nimport copy\nimport random\n\nimport numpy as np\nimport pytest\nimport torch\nfrom tensordict.tensorclass import NonTensorData, NonTensorStack\n\nfrom verl.utils import tensordict_utils as tu\n\n\ndef test_union_tensor_dict():\n    obs = torch.randn(100, 10)\n\n    meta_info1 = {\"top_p\": 0.8}\n    meta_info2 = {\"top_p\": 0.9}\n    data1 = {\"obs\": obs, \"act\": torch.randn(100, 3), \"data_sources\": [\"gsm8k\"] * 100}\n    data2 = {\"obs\": obs, \"next_obs\": torch.randn(100, 10), \"rew\": torch.randn(100), \"data_sources\": [\"gsm8k\"] * 100}\n\n    data_with_copied_obs = {\"obs\": obs.clone(), \"next_obs\": torch.randn(100, 10), \"rew\": torch.randn(100)}\n\n    data1 = tu.get_tensordict(tensor_dict=data1)\n    data2 = tu.get_tensordict(tensor_dict=data2)\n    data_with_copied_obs = tu.get_tensordict(data_with_copied_obs)\n\n    tu.union_tensor_dict(data1, data2)\n    with pytest.raises(AssertionError):\n        # conflict in tensor values\n        tu.union_tensor_dict(data1, data_with_copied_obs)\n\n    data1 = tu.assign_non_tensor(data1, **meta_info1)\n    tu.union_tensor_dict(data1, data2)  # works ok\n\n    data2 = tu.assign_non_tensor(data2, **meta_info2)\n\n    with pytest.raises(AssertionError):\n        # conflict in NonTensorData\n        tu.union_tensor_dict(data1, data2)\n\n    data1.pop(\"top_p\")\n    data2.pop(\"top_p\")\n\n    data2[\"data_sources\"][0] = \"math\"\n    with pytest.raises(AssertionError):\n        # conflict in NonTensorData\n        tu.union_tensor_dict(data1, data2)\n\n\ndef test_tensor_dict_constructor():\n    obs = torch.ones(100, 10)\n    act = torch.zeros(100, 10, 3)\n    data_source = [\"gsm8k\"] * 100\n    non_tensor_dict = {\"name\": \"abdce\"}\n\n    data = tu.get_tensordict(\n        tensor_dict={\"obs\": obs, \"act\": act, \"data_source\": data_source}, non_tensor_dict=non_tensor_dict\n    )\n\n    assert data.batch_size == torch.Size([100])\n\n    # test slicing\n    assert torch.all(torch.eq(data[0][\"obs\"], torch.ones(10))).item()\n    assert torch.all(torch.eq(data[0][\"act\"], torch.zeros(10, 3))).item()\n    assert data[0][\"data_source\"] == \"gsm8k\"\n\n    assert torch.all(torch.eq(data[0:2][\"obs\"], torch.ones(2, 10))).item()\n    assert torch.all(torch.eq(data[0:2][\"act\"], torch.zeros(2, 10, 3))).item()\n    assert data[0:2][\"data_source\"] == [\"gsm8k\"] * 2\n\n    # test non tensor data\n    assert data[\"name\"] == \"abdce\"\n\n\ndef test_index_select_tensor_dict():\n    vocab_size = 128\n    a = torch.randint(low=0, high=vocab_size, size=(11,))\n    b = torch.randint(low=0, high=vocab_size, size=(13,))\n    c = torch.randint(low=0, high=vocab_size, size=(12,))\n    d = torch.randint(low=0, high=vocab_size, size=(15,))\n    input_ids = [a, b, c, d]\n    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)\n\n    padded_tensor = torch.randn(4, 10)\n    non_tensor_dict = {\"global_batch_size\": \"4\"}\n\n    data = tu.get_tensordict(\n        tensor_dict={\n            \"input_ids\": input_ids,\n            \"padded_tensor\": padded_tensor,\n        },\n        non_tensor_dict=non_tensor_dict,\n    )\n\n    assert data.batch_size == torch.Size([4])\n\n    # test index select\n    indices = torch.tensor([1, 3])\n    selected_data = tu.index_select_tensor_dict(data, indices)\n\n    assert selected_data.batch_size == torch.Size([2])\n\n    target_input_ids = torch.nested.as_nested_tensor([input_ids[idx] for idx in indices], layout=torch.jagged)\n    target_select_data = tu.get_tensordict(\n        tensor_dict={\n            \"input_ids\": target_input_ids,\n            \"padded_tensor\": padded_tensor[indices],\n        },\n        non_tensor_dict=non_tensor_dict,\n    )\n    tu.assert_tensordict_eq(selected_data, target_select_data)\n\n\ndef test_tensordict_with_images():\n    # each sample contains a sequence with multiple images of different sizes\n    vocab_size = 128\n    a = torch.randint(low=0, high=vocab_size, size=(11,))\n    b = torch.randint(low=0, high=vocab_size, size=(13,))\n    input_ids = [a, b]\n    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)\n\n    # must be numpy\n    # TODO(vermouth1992). We may use nested tensor too. But this requires nested over nested\n    a_images = [\n        torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(),\n        torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(),\n    ]\n    b_images = [\n        torch.randint(low=0, high=255, size=(3, 256, 256), dtype=torch.uint8).numpy(),\n        torch.randint(low=0, high=255, size=(3, 128, 128), dtype=torch.uint8).numpy(),\n        torch.randint(low=0, high=255, size=(3, 64, 64), dtype=torch.uint8).numpy(),\n    ]\n\n    images = [a_images, b_images]\n\n    data = tu.get_tensordict({\"input_ids\": input_ids, \"images\": images})\n\n    assert np.all(np.equal(data[0][\"images\"][0], a_images[0]))\n    assert torch.all(torch.eq(data[0][\"input_ids\"], a))\n\n\ndef test_tensordict_with_packing():\n    vocab_size = 128\n    a = torch.randint(low=0, high=vocab_size, size=(11,))\n    b = torch.randint(low=0, high=vocab_size, size=(13,))\n    input_ids = [a, b]\n    input_ids = torch.nested.as_nested_tensor(input_ids, layout=torch.jagged)\n\n    data = tu.get_tensordict({\"input_ids\": input_ids})\n\n    # test cu_seqlens\n    cu_seqlens = torch.tensor([0, 11, 24])\n    assert torch.all(torch.eq(cu_seqlens, data[\"input_ids\"].offsets()))\n\n    # test index\n    assert torch.all(torch.eq(data[\"input_ids\"][0], a))\n    assert torch.all(torch.eq(data[\"input_ids\"][1], b))\n\n    assert torch.all(torch.eq(data[0][\"input_ids\"], a))\n    assert torch.all(torch.eq(data[1][\"input_ids\"], b))\n\n    data_lst = data.chunk(2)\n\n    assert torch.all(torch.eq(data_lst[0][\"input_ids\"][0], a))\n    assert torch.all(torch.eq(data_lst[1][\"input_ids\"][0], b))\n\n\ndef test_tensordict_eq():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    data_sources = [\"abc\", \"def\", \"abc\", \"def\", \"pol\", \"klj\"]\n    non_tensor_dict = {\"train_sample_kwargs\": {\"top_p\": 1.0}, \"val_sample_kwargs\": {\"top_p\": 0.7}}\n    data = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    data_sources = [\"abc\", \"def\", \"abc\", \"def\", \"pol\", \"klj\"]\n    non_tensor_dict = {\"train_sample_kwargs\": {\"top_p\": 1.0}, \"val_sample_kwargs\": {\"top_p\": 0.7}}\n    data1 = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n\n    tu.assert_tensordict_eq(data, data1)\n\n    data2 = copy.deepcopy(data1)\n    data2[\"obs\"][0] += 1\n\n    with pytest.raises(AssertionError):\n        tu.assert_tensordict_eq(data, data2)\n\n    data2 = copy.deepcopy(data1)\n    data2[\"data_sources\"][0] = \"math\"\n\n    with pytest.raises(AssertionError):\n        tu.assert_tensordict_eq(data, data2)\n\n    data2 = copy.deepcopy(data1)\n    data2[\"train_sample_kwargs\"][\"top_p\"] = 0.9\n\n    with pytest.raises(AssertionError):\n        tu.assert_tensordict_eq(data, data2)\n\n    tensor_list = [\n        torch.tensor([1, 2, 3, 3, 2]),\n        torch.tensor([4, 5]),\n        torch.tensor([7, 8, 10, 14]),\n        torch.tensor([10, 11, 12]),\n        torch.tensor([13, 14, 15, 18]),\n        torch.tensor([16, 17]),\n    ]\n    obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged)\n    data_sources = [\"abc\", \"def\", \"abc\", \"def\", \"pol\", \"klj\"]\n    non_tensor_dict = {\"train_sample_kwargs\": {\"top_p\": 1.0}, \"val_sample_kwargs\": {\"top_p\": 0.7}}\n    data3 = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n\n    tensor_list[0] = torch.tensor([1, 2, 3, 3, 2])\n    obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged)\n    data4 = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n    tu.assert_tensordict_eq(data3, data4)\n\n    tensor_list[0] = torch.tensor([1, 2, 4])\n    obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged)\n    data5 = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n    with pytest.raises(AssertionError):\n        tu.assert_tensordict_eq(data3, data5)\n\n    tensor_list[0] = torch.tensor([4, 5])\n    tensor_list[1] = torch.tensor([1, 2, 3, 3, 2])\n    obs = torch.nested.as_nested_tensor(tensor_list, layout=torch.jagged)\n    data6 = tu.get_tensordict({\"obs\": obs, \"data_sources\": data_sources}, non_tensor_dict=non_tensor_dict)\n    with pytest.raises(AssertionError):\n        tu.assert_tensordict_eq(data3, data6)\n\n\ndef test_tensor_dict_make_iterator():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    input_ids = torch.nested.as_nested_tensor(\n        [\n            torch.tensor([0, 1]),\n            torch.tensor([2]),\n            torch.tensor([3, 4]),\n            torch.tensor([5]),\n            torch.tensor([6, 7, 8]),\n            torch.tensor([9]),\n        ],\n        layout=torch.jagged,\n    )\n    data_sources = [\"abc\", \"def\", \"abc\", \"def\", \"pol\", \"klj\"]\n    non_tensor_dict = {\"train_sample_kwargs\": {\"top_p\": 1.0}, \"val_sample_kwargs\": {\"top_p\": 0.7}}\n    dataset = tu.get_tensordict(\n        {\"obs\": obs, \"data_sources\": data_sources, \"input_ids\": input_ids}, non_tensor_dict=non_tensor_dict\n    )\n\n    dataloader = tu.make_iterator(\n        dataset, mini_batch_size=2, epochs=2, seed=0, dataloader_kwargs={\"shuffle\": False, \"drop_last\": False}\n    )\n\n    expected_tensor_dict = [\n        tu.index_select_tensor_dict(dataset, indices=list(range(0, 2))),\n        tu.index_select_tensor_dict(dataset, indices=list(range(2, 4))),\n        tu.index_select_tensor_dict(dataset, indices=list(range(4, 6))),\n        tu.index_select_tensor_dict(dataset, indices=list(range(0, 2))),\n        tu.index_select_tensor_dict(dataset, indices=list(range(2, 4))),\n        tu.index_select_tensor_dict(dataset, indices=list(range(4, 6))),\n    ]\n\n    i = 0\n\n    for d in dataloader:\n        tu.assert_tensordict_eq(d, expected_tensor_dict[i])\n        i += 1\n\n    data_iter_1 = tu.make_iterator(dataset, mini_batch_size=3, epochs=1, seed=1, dataloader_kwargs={\"shuffle\": True})\n    data_list_1 = []\n    for data in data_iter_1:\n        data_list_1.append(data)\n\n    data_iter_2 = tu.make_iterator(dataset, mini_batch_size=3, epochs=1, seed=1, dataloader_kwargs={\"shuffle\": True})\n    data_list_2 = []\n    for data in data_iter_2:\n        data_list_2.append(data)\n\n    for data1, data2 in zip(data_list_1, data_list_2, strict=True):\n        tu.assert_tensordict_eq(data1, data2)\n\n\ndef test_reorder():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    labels = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n    non_tensor_dict = {\"name\": \"abdce\"}\n\n    data = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict=non_tensor_dict)\n    data = data[torch.tensor([3, 4, 2, 0, 1, 5])]\n\n    assert torch.all(torch.eq(data[\"obs\"], torch.tensor([4, 5, 3, 1, 2, 6])))\n    assert np.all(data[\"labels\"] == np.array([\"d\", \"e\", \"c\", \"a\", \"b\", \"f\"]))\n    assert data[\"name\"] == \"abdce\"\n\n\ndef test_chunk_concat():\n    obs = torch.tensor([1, 2, 3, 4, 5, 6])\n    labels = [\"a\", \"b\", \"c\", \"d\", \"e\", \"f\"]\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"name\": \"abcde\"})\n\n    data_split = data.tensor_split(indices_or_sections=5, dim=0)\n\n    expected_idx_lst = [[0, 1], [2], [3], [4], [5]]\n\n    for d, expected_idx in zip(data_split, expected_idx_lst, strict=False):\n        tu.assert_tensordict_eq(d, data[expected_idx])\n\n    data_split = data.chunk(2)\n    assert len(data_split) == 2\n    assert torch.all(torch.eq(data_split[0][\"obs\"], torch.tensor([1, 2, 3])))\n    assert np.all(data_split[0][\"labels\"] == np.array([\"a\", \"b\", \"c\"]))\n    assert data_split[0][\"name\"] == \"abcde\"\n\n    assert torch.all(torch.eq(data_split[1][\"obs\"], torch.tensor([4, 5, 6])))\n    assert np.all(data_split[1][\"labels\"] == np.array([\"d\", \"e\", \"f\"]))\n    assert data_split[1][\"name\"] == \"abcde\"\n\n    concat_data = torch.cat(data_split, dim=0)\n    assert torch.all(torch.eq(concat_data[\"obs\"], data[\"obs\"]))\n    assert np.all(concat_data[\"labels\"] == data[\"labels\"])\n    assert concat_data[\"name\"] == data[\"name\"]\n\n    data1 = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"name\": \"abcde\"})\n    data2 = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"name\": \"def\"})\n    data3 = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"name\": \"cfg\"})\n\n    output = torch.cat([data1, data2, data3], dim=0)\n\n    # concat NonTensorData will keep the first one.\n    assert output[\"name\"] == \"abcde\"\n\n\ndef test_pop():\n    obs = torch.randn(3, 10)\n    act = torch.randn(3, 3)\n    labels = [\"a\", [\"b\"], []]\n    dataset = tu.get_tensordict({\"obs\": obs, \"act\": act, \"labels\": labels}, non_tensor_dict={\"2\": 2, \"1\": 1})\n\n    dataset1 = copy.deepcopy(dataset)\n\n    # test pop keys\n    popped_dataset = tu.pop_keys(dataset, keys=[\"obs\", \"2\"])\n\n    assert popped_dataset.batch_size[0] == 3\n\n    assert popped_dataset.keys() == {\"obs\", \"2\"}\n    assert torch.all(torch.eq(popped_dataset[\"obs\"], obs)).item()\n    assert popped_dataset[\"2\"] == 2\n\n    assert dataset.keys() == {\"act\", \"1\", \"labels\"}\n\n    # test pop non-exist key\n    with pytest.raises(KeyError):\n        tu.pop_keys(dataset, keys=[\"obs\", \"2\"])\n\n    # test single pop\n    # NonTensorData\n    assert tu.pop(dataset1, key=\"2\") == 2\n    # NonTensorStack\n    assert tu.pop(dataset1, key=\"labels\") == [\"a\", [\"b\"], []]\n    # Tensor\n    assert torch.all(torch.eq(tu.pop(dataset1, key=\"obs\"), obs)).item()\n\n\ndef test_get():\n    obs = torch.randn(3, 10)\n    act = torch.randn(3, 3)\n    labels = [\"a\", [\"b\"], []]\n    dataset = tu.get_tensordict({\"obs\": obs, \"act\": act, \"labels\": labels}, non_tensor_dict={\"2\": 2, \"1\": 1})\n\n    # test pop keys\n    popped_dataset = tu.get_keys(dataset, keys=[\"obs\", \"2\"])\n\n    assert popped_dataset.batch_size[0] == 3\n\n    assert torch.all(torch.eq(popped_dataset[\"obs\"], dataset[\"obs\"])).item()\n\n    assert popped_dataset[\"2\"] == dataset[\"2\"]\n\n    # test pop non-exist key\n    with pytest.raises(KeyError):\n        tu.get_keys(dataset, keys=[\"obs\", \"3\"])\n\n    # test single pop\n    # NonTensorData\n    assert tu.get(dataset, key=\"2\") == 2\n    # NonTensorStack\n    assert tu.get(dataset, key=\"labels\") == [\"a\", [\"b\"], []]\n    # Tensor\n    assert torch.all(torch.eq(tu.get(dataset, key=\"obs\"), obs)).item()\n    # Non-exist key\n    assert tu.get(dataset, key=\"3\", default=3) == 3\n\n\ndef test_repeat():\n    # Create a DataProto object with some batch and non-tensor data\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"info\": \"test_info\"})\n\n    # Test interleave=True\n    repeated_data_interleave = data.repeat_interleave(repeats=2)\n    expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [3, 4], [3, 4], [5, 6], [5, 6]])\n    expected_labels_interleave = [\"a\", \"a\", \"b\", \"b\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_interleave[\"obs\"], expected_obs_interleave))\n    assert repeated_data_interleave[\"labels\"] == expected_labels_interleave\n    assert repeated_data_interleave[\"info\"] == \"test_info\"\n\n    # Test interleave=False\n    repeated_data_no_interleave = data.repeat(2)\n    expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6]])\n    expected_labels_no_interleave = [\"a\", \"b\", \"c\", \"a\", \"b\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_no_interleave[\"obs\"], expected_obs_no_interleave))\n    assert repeated_data_no_interleave[\"labels\"] == expected_labels_no_interleave\n    assert repeated_data_no_interleave[\"info\"] == \"test_info\"\n\n\ndef test_dataproto_pad_unpad():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"info\": \"test_info\"})\n\n    padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=2)\n\n    assert pad_size == 1\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2]])\n    expected_labels = [\"a\", \"b\", \"c\", \"a\"]\n\n    assert torch.all(torch.eq(padded_data[\"obs\"], expected_obs))\n    assert padded_data[\"labels\"] == expected_labels\n    assert padded_data[\"info\"] == \"test_info\"\n\n    unpadd_data = tu.unpad(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data[\"obs\"], obs))\n    assert unpadd_data[\"labels\"] == labels\n    assert unpadd_data[\"info\"] == \"test_info\"\n\n    padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=3)\n    assert pad_size == 0\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    expected_labels = [\"a\", \"b\", \"c\"]\n\n    assert torch.all(torch.eq(padded_data[\"obs\"], expected_obs))\n    assert padded_data[\"labels\"] == expected_labels\n    assert padded_data[\"info\"] == \"test_info\"\n\n    unpadd_data = tu.unpad(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data[\"obs\"], obs))\n    assert unpadd_data[\"labels\"] == labels\n    assert unpadd_data[\"info\"] == \"test_info\"\n\n    padded_data, pad_size = tu.pad_to_divisor(data, size_divisor=7)\n    assert pad_size == 4\n\n    expected_obs = torch.tensor([[1, 2], [3, 4], [5, 6], [1, 2], [3, 4], [5, 6], [1, 2]])\n    expected_labels = [\"a\", \"b\", \"c\", \"a\", \"b\", \"c\", \"a\"]\n    assert torch.all(torch.eq(padded_data[\"obs\"], expected_obs))\n    assert padded_data[\"labels\"] == expected_labels\n    assert padded_data[\"info\"] == \"test_info\"\n\n    unpadd_data = tu.unpad(padded_data, pad_size=pad_size)\n    assert torch.all(torch.eq(unpadd_data[\"obs\"], obs))\n    assert unpadd_data[\"labels\"] == labels\n    assert unpadd_data[\"info\"] == \"test_info\"\n\n\ndef test_torch_save_data_proto():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"info\": \"test_info\"})\n\n    filename = \"test_data.pt\"\n    torch.save(data, filename)\n    loaded_data = torch.load(filename, weights_only=False)\n\n    assert torch.all(torch.eq(loaded_data[\"obs\"], data[\"obs\"]))\n    assert loaded_data[\"labels\"] == data[\"labels\"]\n    assert loaded_data[\"info\"] == data[\"info\"]\n\n    import os\n\n    os.remove(filename)\n\n\ndef test_len():\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = np.array([\"a\", \"b\", \"c\"], dtype=object)\n\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels.tolist()}, non_tensor_dict={\"info\": \"test_info\"})\n    assert len(data) == 3\n\n    data = tu.get_tensordict({\"labels\": labels.tolist()}, non_tensor_dict={\"info\": \"test_info\"})\n    assert len(data) == 3\n\n    data_item = data[0]\n    assert len(data_item) == 0\n\n    data = tu.get_tensordict({}, non_tensor_dict={\"info\": \"test_info\"})\n    assert len(data) == 0\n\n\ndef test_dataproto_index():\n    data_len = 100\n    idx_num = 10\n\n    obs = torch.randn(data_len, 10)\n    labels = [random.choice([\"abc\", \"cde\"]) for _ in range(data_len)]\n\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels})\n\n    labels_np = np.array(labels)\n\n    idx_np_int = np.random.randint(0, data_len, size=(idx_num,))\n    result_np_int = data[idx_np_int]\n    assert result_np_int.keys() == data.keys()\n    assert result_np_int[\"obs\"].shape[0] == idx_num\n    assert len(result_np_int[\"labels\"]) == idx_num\n    assert np.array_equal(result_np_int[\"obs\"].cpu().numpy(), obs[idx_np_int].numpy())\n    assert np.array_equal(result_np_int[\"labels\"], labels_np[idx_np_int])\n\n    idx_torch_int = torch.randint(0, data_len, size=(idx_num,))\n    result_torch_int = data[idx_torch_int]\n    assert result_torch_int.keys() == data.keys()\n    assert result_torch_int[\"obs\"].shape[0] == idx_num\n    assert len(result_torch_int[\"labels\"]) == idx_num\n    assert np.array_equal(result_torch_int[\"obs\"].cpu().numpy(), obs[idx_torch_int].cpu().numpy())\n    assert np.array_equal(result_torch_int[\"labels\"], labels_np[idx_torch_int.cpu().numpy()])\n\n    idx_list_int = [np.random.randint(0, data_len) for _ in range(idx_num)]\n    result_list_int = data[idx_list_int]\n    assert result_list_int.keys() == data.keys()\n    assert result_list_int[\"obs\"].shape[0] == idx_num\n    assert len(result_list_int[\"labels\"]) == idx_num\n    assert np.array_equal(result_list_int[\"obs\"].cpu().numpy(), obs[idx_list_int].cpu().numpy())\n    assert np.array_equal(result_list_int[\"labels\"], labels_np[idx_list_int])\n\n    # idx_np_bool = np.random.randint(0, 2, size=(data_len,), dtype=bool)\n    # result_np_bool = data[idx_np_bool]\n    # assert result_np_bool.keys() == data.keys()\n    # assert result_np_bool[\"obs\"].shape[0] == idx_np_bool.sum()\n    # assert len(result_np_bool[\"labels\"]) == idx_np_bool.sum()\n    # assert np.array_equal(result_np_bool[\"obs\"].cpu().numpy(), obs[idx_np_bool].cpu().numpy())\n    # assert np.array_equal(result_np_bool[\"labels\"], labels_np[idx_np_bool])\n\n    idx_torch_bool = torch.randint(0, 2, size=(data_len,), dtype=torch.bool)\n    result_torch_bool = data[idx_torch_bool]\n    assert result_torch_bool.keys() == data.keys()\n    assert result_torch_bool[\"obs\"].shape[0] == idx_torch_bool.sum().item()\n    assert len(result_torch_bool[\"labels\"]) == idx_torch_bool.sum().item()\n    assert np.array_equal(result_torch_bool[\"obs\"].cpu().numpy(), obs[idx_torch_bool].cpu().numpy())\n    assert np.array_equal(result_torch_bool[\"labels\"], labels_np[idx_torch_bool])\n\n    # idx_list_bool = [np.random.randint(0, 2, dtype=bool) for _ in range(data_len)]\n    # result_list_bool = data[idx_list_bool]\n    # assert result_list_bool.keys() == data.keys()\n    # assert result_list_bool[\"obs\"].shape[0] == sum(idx_list_bool)\n    # assert len(result_list_bool[\"labels\"]) == sum(idx_list_bool)\n    # assert np.array_equal(result_list_bool[\"obs\"].cpu().numpy(), obs[idx_list_bool].cpu().numpy())\n    # assert np.array_equal(result_list_bool[\"labels\"], labels_np[idx_list_bool])\n\n\ndef test_select():\n    obs = torch.randn(100, 10)\n    act = torch.randn(100, 3)\n    dataset = tu.get_tensordict({\"obs\": obs, \"act\": act}, non_tensor_dict={\"2\": 2, \"1\": 1})\n\n    subset = dataset.select(\"obs\", \"2\")\n\n    assert torch.all(torch.eq(subset[\"obs\"], dataset[\"obs\"]))\n    assert subset[\"2\"] == dataset[\"2\"]\n    assert \"act\" not in subset.keys()\n    assert \"1\" not in subset.keys()\n\n\ndef test_dataproto_no_batch():\n    labels = [\"a\", \"b\", \"c\"]\n    data = tu.get_tensordict(tensor_dict={\"labels\": labels}, non_tensor_dict={\"info\": \"test_info\"})\n    selected = data.select(\"labels\")\n\n    assert selected[\"labels\"] == labels\n    pop_data = tu.pop_keys(data, keys=[\"labels\"])\n    assert pop_data[\"labels\"] == labels\n    assert \"labels\" not in data\n\n\ndef test_sample_level_repeat():\n    # Create a DataProto object with some batch and non-tensor data\n    obs = torch.tensor([[1, 2], [3, 4], [5, 6]])\n    labels = [\"a\", \"b\", \"c\"]\n\n    data = tu.get_tensordict({\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"info\": \"test_info\"})\n\n    # list\n    repeated_data_interleave = data.repeat_interleave(repeats=torch.tensor([3, 1, 2]))\n    expected_obs_interleave = torch.tensor([[1, 2], [1, 2], [1, 2], [3, 4], [5, 6], [5, 6]])\n    expected_labels_interleave = [\"a\", \"a\", \"a\", \"b\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_interleave[\"obs\"], expected_obs_interleave))\n    assert repeated_data_interleave[\"labels\"] == expected_labels_interleave\n    assert repeated_data_interleave[\"info\"] == \"test_info\"\n\n    # torch.tensor\n    repeated_data_no_interleave = data.repeat_interleave(repeats=torch.tensor([1, 2, 3]))\n    expected_obs_no_interleave = torch.tensor([[1, 2], [3, 4], [3, 4], [5, 6], [5, 6], [5, 6]])\n    expected_labels_no_interleave = [\"a\", \"b\", \"b\", \"c\", \"c\", \"c\"]\n\n    assert torch.all(torch.eq(repeated_data_no_interleave[\"obs\"], expected_obs_no_interleave))\n    assert repeated_data_no_interleave[\"labels\"] == expected_labels_no_interleave\n    assert repeated_data_no_interleave[\"info\"] == \"test_info\"\n\n\ndef test_dataproto_chunk_after_index():\n    data_len = 4\n    obs = torch.randn(data_len, 4)\n    labels = [f\"label_{i}\" for i in range(data_len)]\n\n    data = tu.get_tensordict(tensor_dict={\"obs\": obs, \"labels\": labels}, non_tensor_dict={\"name\": \"abc\"})\n    # Test with boolean numpy array\n    bool_mask = torch.tensor([True, False, True, False])\n    selected = data[bool_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)  # int or List[int]\n\n    # Test with integer numpy array\n    int_mask = torch.tensor([0, 2])\n    selected = data[int_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)\n\n    # Test with boolean list\n    list_mask = [True, False, True, False]\n    selected = data[list_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)\n\n    # Test with list\n    list_mask = [0, 2]\n    selected = data[list_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)\n\n    # Test with torch tensor (bool)\n    torch_bool_mask = torch.tensor([True, False, True, False])\n    selected = data[torch_bool_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)\n\n    # Test with torch tensor (int)\n    torch_int_mask = torch.tensor([0, 2])\n    selected = data[torch_int_mask]\n    assert isinstance(selected.batch_size, torch.Size)\n    assert all(isinstance(d, int) for d in selected.batch_size)\n\n\ndef test_concat_nested_tensor():\n    # Test 2D nested tensors\n    vocab_size = 128\n    a = torch.randint(low=0, high=vocab_size, size=(11,))\n    b = torch.randint(low=0, high=vocab_size, size=(13,))\n    c = torch.randint(low=0, high=vocab_size, size=(12,))\n    d = torch.randint(low=0, high=vocab_size, size=(15,))\n\n    nested_a_b = torch.nested.as_nested_tensor([a, b], layout=torch.jagged)\n    nested_c_d = torch.nested.as_nested_tensor([c, d], layout=torch.jagged)\n\n    output = tu.concat_nested_tensors([nested_a_b, nested_c_d])\n\n    output_values = output.values()\n    expected = torch.cat([a, b, c, d], dim=0)\n\n    assert torch.all(torch.eq(output_values, expected)).item()\n\n    # Test 3D nested tensors\n    a_3d = torch.randint(low=0, high=vocab_size, size=(4, 4))\n    b_3d = torch.randint(low=0, high=vocab_size, size=(4, 5))\n    c_3d = torch.randint(low=0, high=vocab_size, size=(4, 6))\n    d_3d = torch.randint(low=0, high=vocab_size, size=(4, 7))\n\n    nested_a_b_3d = torch.nested.as_nested_tensor([a_3d, b_3d], layout=torch.jagged)\n    nested_c_d_3d = torch.nested.as_nested_tensor([c_3d, d_3d], layout=torch.jagged)\n\n    output_3d = tu.concat_nested_tensors([nested_a_b_3d, nested_c_d_3d])\n\n    assert output_3d.shape[0] == 4\n    output_3d_unbind = output_3d.unbind(0)\n    assert torch.all(torch.eq(output_3d_unbind[0], a_3d)).item()\n    assert torch.all(torch.eq(output_3d_unbind[1], b_3d)).item()\n    assert torch.all(torch.eq(output_3d_unbind[2], c_3d)).item()\n    assert torch.all(torch.eq(output_3d_unbind[3], d_3d)).item()\n\n    # Test 4D nested tensors\n    a_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 4))\n    b_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 5))\n    c_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 3))\n    d_4d = torch.randint(low=0, high=vocab_size, size=(2, 3, 6))\n\n    nested_a_b_4d = torch.nested.as_nested_tensor([a_4d, b_4d], layout=torch.jagged)\n    nested_c_d_4d = torch.nested.as_nested_tensor([c_4d, d_4d], layout=torch.jagged)\n\n    output_4d = tu.concat_nested_tensors([nested_a_b_4d, nested_c_d_4d])\n\n    assert output_4d.shape[0] == 4\n    output_4d_unbind = output_4d.unbind(0)\n    assert torch.all(torch.eq(output_4d_unbind[0], a_4d)).item()\n    assert torch.all(torch.eq(output_4d_unbind[1], b_4d)).item()\n    assert torch.all(torch.eq(output_4d_unbind[2], c_4d)).item()\n    assert torch.all(torch.eq(output_4d_unbind[3], d_4d)).item()\n\n\ndef test_concat_tensordict():\n    vocab_size = 128\n    a = torch.randint(low=0, high=vocab_size, size=(11,))\n    b = torch.randint(low=0, high=vocab_size, size=(13,))\n    c = torch.randint(low=0, high=vocab_size, size=(12,))\n    d = torch.randint(low=0, high=vocab_size, size=(15,))\n\n    nested_a_b = torch.nested.as_nested_tensor([a, b], layout=torch.jagged)\n    nested_c_d = torch.nested.as_nested_tensor([c, d], layout=torch.jagged)\n\n    tensordict1 = tu.get_tensordict(\n        tensor_dict={\"input_ids\": nested_a_b, \"labels\": [\"a\", \"b\"]}, non_tensor_dict={\"temp\": 1.0}\n    )\n    tensordict2 = tu.get_tensordict(\n        tensor_dict={\"input_ids\": nested_c_d, \"labels\": [\"c\", \"d\"]}, non_tensor_dict={\"temp\": 2.0}\n    )\n\n    tensordict1_copy = copy.deepcopy(tensordict1)\n    tensordict2_copy = copy.deepcopy(tensordict2)\n\n    output = tu.concat_tensordict([tensordict1, tensordict2])\n\n    assert torch.all(torch.eq(output[\"input_ids\"].values(), torch.cat([a, b, c, d]))).item()\n    assert output[\"labels\"] == [\"a\", \"b\", \"c\", \"d\"]\n    assert output[\"temp\"] == 1.0\n\n    # make sure tensordict1 and tensordict2 is untouched\n    tu.assert_tensordict_eq(tensordict1, tensordict1_copy)\n    tu.assert_tensordict_eq(tensordict2, tensordict2_copy)\n\n    # test concat tensordict with only NonTensorStack and NonTensorData\n    tensordict1 = tu.get_tensordict(tensor_dict={\"labels\": [\"a\", \"b\"]}, non_tensor_dict={\"temp\": 1.0})\n    tensordict2 = tu.get_tensordict(tensor_dict={\"labels\": [\"c\", \"d\"]}, non_tensor_dict={\"temp\": 2.0})\n\n    output = tu.concat_tensordict([tensordict1, tensordict2])\n\n    assert output[\"labels\"] == [\"a\", \"b\", \"c\", \"d\"]\n    assert output[\"temp\"] == 1.0\n\n    assert output.batch_size[0] == 4\n\n    # test concat tensordict with only NonTensorData\n    tensordict1 = tu.get_tensordict(tensor_dict={}, non_tensor_dict={\"temp\": 1.0})\n    tensordict2 = tu.get_tensordict(tensor_dict={}, non_tensor_dict={\"temp\": 2.0})\n\n    output = tu.concat_tensordict([tensordict1, tensordict2])\n    assert len(output.batch_size) == 0\n    assert output[\"temp\"] == 1.0\n\n\ndef test_chunk_tensordict():\n    # Qwen-VL 3d position_ids\n    position_ids = torch.nested.as_nested_tensor(\n        [\n            torch.arange(4).expand(4, 4),\n            torch.arange(5).expand(4, 5),\n            torch.arange(6).expand(4, 6),\n            torch.arange(7).expand(4, 7),\n        ],\n        layout=torch.jagged,\n    )\n    input_ids = torch.nested.as_nested_tensor(\n        [torch.arange(4), torch.arange(5), torch.arange(6), torch.arange(7)], layout=torch.jagged\n    )\n    attention_mask = torch.nested.as_nested_tensor(\n        [\n            torch.randint(low=0, high=2, size=[3, 4]),\n            torch.randint(low=0, high=2, size=[3, 5]),\n            torch.randint(low=0, high=2, size=[3, 6]),\n            torch.randint(low=0, high=2, size=[3, 7]),\n        ],\n        layout=torch.jagged,\n    )\n\n    multi_modal_inputs = torch.stack(\n        [\n            NonTensorData({\"pixel_values\": torch.randn(3, 224, 224)}),\n            NonTensorData(None),\n            NonTensorData({\"pixel_values\": torch.randn(3, 128, 128)}),\n            NonTensorData({\"pixel_values\": torch.randn(3, 128, 128)}),\n        ]\n    )\n    td = tu.get_tensordict(\n        {\n            \"input_ids\": input_ids,\n            \"position_ids\": position_ids,\n            \"attention_mask\": attention_mask,\n            \"multi_modal_inputs\": multi_modal_inputs,\n        },\n    )\n    assert len(td) == 4\n    chunks = tu.chunk_tensordict(td, chunks=2)\n\n    for i, chunk in enumerate(chunks):\n        assert len(chunk) == 2\n        for key, val in chunk.items():\n            if isinstance(val, torch.Tensor) and val.is_nested:\n                tensors = td[key].unbind(dim=0)\n                expected = torch.nested.as_nested_tensor(tensors[i * 2 : (i + 1) * 2], layout=torch.jagged)\n                assert torch.all(torch.eq(val.values(), expected.values())).item()\n            else:\n                expected = td[key][i * 2 : (i + 1) * 2]\n                for tensor, expect in zip(val, expected, strict=False):\n                    if tensor.data is None:\n                        assert expect is None\n                    else:\n                        assert torch.all(torch.eq(tensor.data[\"pixel_values\"], expect[\"pixel_values\"])).item()\n\n\ndef test_assign_non_tensor_stack_with_nested_lists():\n    \"\"\"Test assign_non_tensor_stack with lists of lists.\"\"\"\n    td = tu.get_tensordict({\"obs\": torch.randn(3, 4)}, non_tensor_dict={})\n\n    # Lists of varying lengths (like turn_scores or tool_rewards)\n    turn_scores = [[], [0.5, 0.8], [0.9]]\n    tu.assign_non_tensor_stack(td, \"turn_scores\", turn_scores)\n\n    # Verify data is accessible\n    assert len(td[\"turn_scores\"]) == 3\n    assert list(td[\"turn_scores\"][0]) == []\n    assert list(td[\"turn_scores\"][1]) == [0.5, 0.8]\n    assert list(td[\"turn_scores\"][2]) == [0.9]\n\n\ndef test_assign_non_tensor_stack_with_nested_dicts():\n    \"\"\"Test assign_non_tensor_stack with lists of dicts.\"\"\"\n    td = tu.get_tensordict({\"obs\": torch.randn(3, 4)}, non_tensor_dict={})\n\n    # Lists of dicts (like reward_extra_info)\n    reward_extra_info = [{\"acc\": 1.0, \"loss\": 0.1}, {\"acc\": 0.0, \"loss\": 0.9}, {\"acc\": 1.0, \"loss\": 0.05}]\n    tu.assign_non_tensor_stack(td, \"reward_extra_info\", reward_extra_info)\n\n    # Verify data is accessible\n    assert len(td[\"reward_extra_info\"]) == 3\n    assert dict(td[\"reward_extra_info\"][0]) == {\"acc\": 1.0, \"loss\": 0.1}\n    assert dict(td[\"reward_extra_info\"][1]) == {\"acc\": 0.0, \"loss\": 0.9}\n    assert dict(td[\"reward_extra_info\"][2]) == {\"acc\": 1.0, \"loss\": 0.05}\n\n\ndef test_assign_non_tensor_stack_with_complex_nested():\n    \"\"\"Test assign_non_tensor_stack with lists of lists of dicts.\"\"\"\n    td = tu.get_tensordict({\"obs\": torch.randn(2, 4)}, non_tensor_dict={})\n\n    # Lists of lists of dicts (like raw_prompt)\n    raw_prompt = [\n        [{\"content\": \"Question 1\", \"role\": \"user\"}],\n        [{\"content\": \"Question 2\", \"role\": \"user\"}, {\"content\": \"Answer 2\", \"role\": \"assistant\"}],\n    ]\n    tu.assign_non_tensor_stack(td, \"raw_prompt\", raw_prompt)\n\n    # Verify data is accessible\n    assert len(td[\"raw_prompt\"]) == 2\n    assert len(td[\"raw_prompt\"][0]) == 1\n    assert dict(td[\"raw_prompt\"][0][0]) == {\"content\": \"Question 1\", \"role\": \"user\"}\n    assert len(td[\"raw_prompt\"][1]) == 2\n    assert dict(td[\"raw_prompt\"][1][0]) == {\"content\": \"Question 2\", \"role\": \"user\"}\n\n\ndef test_assign_non_tensor_handles_wrappers():\n    td = tu.get_tensordict({\"obs\": torch.randn(3, 4)}, non_tensor_dict={})\n\n    meta = {\"top_p\": 0.8}\n    tu.assign_non_tensor(td, **meta)\n    assert td[\"top_p\"] == 0.8\n\n    wrapped = NonTensorData(0.3)\n    stack = NonTensorStack.from_list([NonTensorData(1.0), NonTensorData(2.0), NonTensorData(3.0)])\n    tu.assign_non_tensor(td, wrapped=wrapped, stack=stack)\n\n    assert td[\"wrapped\"] == 0.3\n    assert td[\"stack\"] == [1.0, 2.0, 3.0]\n\n\ndef test_assign_non_tensor_stack_batch_size_check():\n    td = tu.get_tensordict({\"obs\": torch.randn(3, 4)}, non_tensor_dict={})\n    stack = NonTensorStack.from_list([NonTensorData(1.0), NonTensorData(2.0)])\n\n    with pytest.raises(RuntimeError):\n        tu.assign_non_tensor(td, stack=stack)\n\n\ndef test_assign_non_tensor_with_auto_detection():\n    \"\"\"Test assign_non_tensor automatically detects and handles nested structures.\"\"\"\n    td = tu.get_tensordict({\"obs\": torch.randn(3, 4)}, non_tensor_dict={})\n\n    # Mix of simple and nested data\n    tu.assign_non_tensor(\n        td,\n        metadata=\"experiment_1\",  # Simple value\n        turn_scores=[[], [0.5, 0.8], [0.9]],  # Nested list\n        reward_extra_info=[{\"acc\": 1.0}, {\"acc\": 0.0}, {\"acc\": 1.0}],  # List of dicts\n        simple_list=[\"a\", \"b\", \"c\"],  # Simple list (also uses NonTensorStack for consistency)\n    )\n\n    # Verify all data is accessible\n    assert td[\"metadata\"] == \"experiment_1\"\n    assert len(td[\"turn_scores\"]) == 3\n    assert list(td[\"turn_scores\"][1]) == [0.5, 0.8]\n    assert len(td[\"reward_extra_info\"]) == 3\n    assert dict(td[\"reward_extra_info\"][0]) == {\"acc\": 1.0}\n    assert len(td[\"simple_list\"]) == 3\n    assert td[\"simple_list\"][0] == \"a\"\n\n\ndef test_get_tensordict_with_nested_lists():\n    \"\"\"Test get_tensordict automatically handles nested lists.\"\"\"\n    obs = torch.randn(3, 4)\n    turn_scores = [[], [0.5, 0.8], [0.9]]\n\n    # This should automatically convert turn_scores to NonTensorStack\n    td = tu.get_tensordict({\"obs\": obs, \"turn_scores\": turn_scores})\n\n    # Verify tensors and nested data are both accessible\n    assert torch.all(torch.eq(td[\"obs\"], obs))\n    assert len(td[\"turn_scores\"]) == 3\n    assert list(td[\"turn_scores\"][0]) == []\n    assert list(td[\"turn_scores\"][1]) == [0.5, 0.8]\n\n\ndef test_get_tensordict_with_nested_dicts():\n    \"\"\"Test get_tensordict automatically handles lists of dicts.\"\"\"\n    obs = torch.randn(3, 4)\n    reward_extra_info = [{\"acc\": 1.0}, {\"acc\": 0.0}, {\"acc\": 1.0}]\n\n    td = tu.get_tensordict({\"obs\": obs, \"reward_extra_info\": reward_extra_info})\n\n    assert torch.all(torch.eq(td[\"obs\"], obs))\n    assert len(td[\"reward_extra_info\"]) == 3\n    assert dict(td[\"reward_extra_info\"][0]) == {\"acc\": 1.0}\n\n\ndef test_get_tensordict_with_complex_nested_structures():\n    \"\"\"Test get_tensordict with lists of lists of dicts.\"\"\"\n    obs = torch.randn(2, 4)\n    raw_prompt = [\n        [{\"content\": \"Q1\", \"role\": \"user\"}],\n        [{\"content\": \"Q2\", \"role\": \"user\"}, {\"content\": \"A2\", \"role\": \"assistant\"}],\n    ]\n\n    td = tu.get_tensordict({\"obs\": obs, \"raw_prompt\": raw_prompt})\n\n    assert torch.all(torch.eq(td[\"obs\"], obs))\n    assert len(td[\"raw_prompt\"]) == 2\n    assert dict(td[\"raw_prompt\"][0][0]) == {\"content\": \"Q1\", \"role\": \"user\"}\n\n\ndef test_get_tensordict_agent_loop_scenario():\n    \"\"\"Test the complete agent loop scenario with all nested types.\n\n    This simulates the exact use case from agent loops with:\n    - turn_scores: lists of lists\n    - reward_extra_info: lists of dicts\n    - raw_prompt: lists of lists of dicts\n    - tool_rewards: lists of lists\n    \"\"\"\n    prompts = torch.randn(2, 10)\n    responses = torch.randn(2, 5)\n\n    # Nested structures from agent loop\n    data_source = [\"lighteval/MATH\", \"lighteval/MATH\"]\n    uid = [\"uuid-1\", \"uuid-2\"]\n    turn_scores = [[], [0.5, 0.8]]  # Lists of varying lengths\n    reward_extra_info = [{\"acc\": 1.0, \"loss\": 0.1}, {\"acc\": 0.0, \"loss\": 0.9}]\n    raw_prompt = [\n        [{\"content\": \"Compute 4 @ 2\", \"role\": \"user\"}],\n        [{\"content\": \"Compute 8 @ 7\", \"role\": \"user\"}],\n    ]\n    tool_rewards = [[0.0], []]  # List of lists\n\n    # This should handle all nested structures automatically\n    td = tu.get_tensordict(\n        tensor_dict={\n            \"prompts\": prompts,\n            \"responses\": responses,\n            \"data_source\": data_source,\n            \"uid\": uid,\n            \"turn_scores\": turn_scores,\n            \"reward_extra_info\": reward_extra_info,\n            \"raw_prompt\": raw_prompt,\n            \"tool_rewards\": tool_rewards,\n        },\n        non_tensor_dict={\"global_steps\": 42},\n    )\n\n    # Verify all data types are accessible\n    assert torch.all(torch.eq(td[\"prompts\"], prompts))\n    assert torch.all(torch.eq(td[\"responses\"], responses))\n    assert td[\"data_source\"] == data_source\n    assert td[\"uid\"] == uid\n\n    # Verify nested structures\n    assert len(td[\"turn_scores\"]) == 2\n    assert list(td[\"turn_scores\"][0]) == []\n    assert list(td[\"turn_scores\"][1]) == [0.5, 0.8]\n\n    assert len(td[\"reward_extra_info\"]) == 2\n    assert dict(td[\"reward_extra_info\"][0]) == {\"acc\": 1.0, \"loss\": 0.1}\n\n    assert len(td[\"raw_prompt\"]) == 2\n    assert dict(td[\"raw_prompt\"][0][0]) == {\"content\": \"Compute 4 @ 2\", \"role\": \"user\"}\n\n    assert len(td[\"tool_rewards\"]) == 2\n    assert list(td[\"tool_rewards\"][0]) == [0.0]\n    assert list(td[\"tool_rewards\"][1]) == []\n\n    # Verify metadata\n    assert td[\"global_steps\"] == 42\n\n\ndef test_contiguous():\n    # create a tensordict that contains normal tensor, nested tensor,\n    # nontensorstack with numpy, nontensorstack with tensor, NonTensorData with numpy and NonTensorData with tensor\n\n    a = torch.randn(3, 4)  # contiguous tensor\n    b = torch.randn(3, 4)[:, :-1]  # non contiguous tensor\n    c = torch.nested.as_nested_tensor([torch.randn(3), torch.randn(4), torch.randn(5)], layout=torch.jagged)\n\n    d = torch.randn(10, 12)\n    e = torch.randn(11, 12)\n    f = torch.randn(13, 12)\n\n    data = tu.get_tensordict(\n        tensor_dict={\"a\": a, \"b\": b, \"c\": c, \"nt\": [{\"pixel\": d}, {\"pixel\": e}, {\"pixel\": f}]},\n        non_tensor_dict={\"ntd\": a.clone()},\n    )\n\n    with pytest.raises(RuntimeError):\n        # b is not contiguous\n        data.consolidate()\n\n    data1 = copy.deepcopy(data)\n    data_cont = tu.contiguous(data1)\n\n    tu.assert_tensordict_eq(data_cont, data)\n\n    data_cont.consolidate()\n\n    tu.assert_tensordict_eq(data_cont, data)\n"
  },
  {
    "path": "tests/trainer/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTests for the trainer module.\n\"\"\"\n"
  },
  {
    "path": "tests/trainer/config/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/trainer/config/legacy_ppo_megatron_trainer.yaml",
    "content": "data:\n  tokenizer: null\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1  # set to -1 to use full dataset\n  val_max_samples: -1  # set to -1 to use full dataset\n  prompt_key: prompt\n  reward_fn_key: data_source\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null # DEPRECATED: Validation datasets are sent to inference engines as a whole batch, which will schedule the memory themselves\n  return_raw_input_ids: False  # This should be set to true when the tokenizer between policy and rm differs\n  return_raw_chat: True\n  return_full_prompt: False\n  shuffle: True\n  seed: null # An integer seed to use when shuffling the data. If not set or set to `null`, the data shuffling will not be seeded, resulting in a different data order on each run.\n  filter_overlong_prompts: False # for large-scale dataset, filtering overlong prompts could be timeconsuming. You cat set the filter_overlong_prompts_workers to use multiprocessing to speed up.\n  filter_overlong_prompts_workers: 1\n  truncation: error\n  trust_remote_code: False  # main_ppo will check this config to determine whether to use remote code for tokenizer\n  custom_cls:\n      path: null\n      name: null\n  sampler:\n    class_path: null\n    class_name: null\n  dataloader_num_workers: 8\n  return_multi_modal_inputs: True\n\nactor_rollout_ref:\n  hybrid_engine: True\n  nccl_timeout: 600 # seconds, default is 10 minutes for torch, you can set it to a larger value if you have long-running operations like 32B or 72B model using megatron\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    custom_chat_template: null\n    external_lib: null\n    override_config:\n      model_config: {}\n      moe_config:\n        freeze_moe_router: False\n    enable_gradient_checkpointing: True\n    gradient_checkpointing_kwargs:\n      ## Activation Checkpointing\n      activations_checkpoint_method: null # 'uniform', 'block'; not used with 'selective'\n      # 'uniform' divides the total number of transformer layers and checkpoints the input activation of each chunk\n      # 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity\n      activations_checkpoint_granularity: null # 'selective' or 'full'\n      # 'full' will checkpoint the entire transformer layer and 'selective' only checkpoints memory intensive part of attention\n      activations_checkpoint_num_layers: null # not used with 'selective'\n    trust_remote_code: False\n  actor:\n    strategy: megatron  # This is for backward-compatibility\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: False\n    ppo_max_token_len_per_gpu: 16384 # n * ${data.max_prompt_length} + ${data.max_response_length}\n    use_torch_compile: True # False to disable torch compile\n    # pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high)\n    clip_ratio: 0.2 # default value if clip_ratio_low and clip_ratio_high are not specified\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.2\n    clip_ratio_c: 3.0 # lower bound of the value for Dual-clip PPO from https://arxiv.org/pdf/1912.09729\n    loss_agg_mode: \"token-mean\" # / \"seq-mean-token-sum\" / \"seq-mean-token-mean\" / \"seq-mean-token-sum-norm\"\n    # NOTE: \"token-mean\" is the default behavior\n    loss_scale_factor: null  # Scale factor for \"seq-mean-token-sum-norm\" mode. If null, uses response_length.\n    entropy_coeff: 0\n    use_kl_loss: False # True for GRPO\n    kl_loss_coef: 0.001 # for grpo\n    kl_loss_type: low_var_kl # for grpo\n    ppo_epochs: 1\n    data_loader_seed: 42\n    shuffle: False\n    policy_loss:   # policy loss config\n      loss_mode: \"vanilla\" # Loss function mode: vanilla / clip-cov / kl-cov / gpg from https://arxiv.org/abs/2505.22617,\n      clip_cov_ratio: 0.0002 # Ratio of tokens to be clipped for clip-cov loss\n      clip_cov_lb: 1.0 # Lower bound for clip-cov loss\n      clip_cov_ub: 5.0 # Upper bound for clip-cov loss\n      kl_cov_ratio: 0.0002 # Ratio of tokens to be applied kl penalty for kl-cov loss\n      ppo_kl_coef: 0.1 # KL divergence penalty coefficient\n    optim:\n      optimizer: adam\n      lr: 1e-6\n      clip_grad: 1.0\n      total_training_steps: -1  # must be override by program\n      lr_warmup_init: 0.0  # initial learning rate for warmup, default to 0.0\n      lr_warmup_steps: null # Prioritized. None, 0 or Negative values mean delegating to lr_warmup_steps_ratio.\n      lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n      lr_decay_steps: null\n      lr_decay_style: constant # select from constant/linear/cosine/inverse_square_root\n      min_lr: 0.0 # minimum learning rate, default to 0.0\n      weight_decay: 0.01\n      weight_decay_incr_style: constant # select from constant/linear/cosine\n      lr_wsd_decay_style: exponential # select from constant/exponential/cosine\n      lr_wsd_decay_steps: null\n      use_checkpoint_opt_param_scheduler: False # use checkpoint optimizer parameter scheduler\n    megatron:\n      param_offload: False\n      grad_offload: False\n      optimizer_offload: False\n      tensor_model_parallel_size: 1\n      expert_model_parallel_size: 1\n      expert_tensor_parallel_size: null\n      pipeline_model_parallel_size: 1\n      virtual_pipeline_model_parallel_size: null # change VPP interface for parallelism tests\n      context_parallel_size: 1\n      sequence_parallel: True\n      use_distributed_optimizer: True\n      use_dist_checkpointing: False\n      dist_checkpointing_path: null\n      seed: 42\n      override_transformer_config: {} # additional transformer config like: num_layers_in_first(/last)_pipeline_stage\n      use_mbridge: True\n      vanilla_mbridge: True\n    profile: # profile the actor model in `update_policy`\n      use_profile: False # open it when you want to profile the actor model\n      profile_ranks: null # list, you can specify the ranks to profile\n      step_start: -1 # start step in update_policy\n      step_end: -1 # end step\n      save_path: null # the path to save the profile result\n    load_weight: True\n    checkpoint:\n      async_save: False # save checkpoint asynchronously\n      # What to include in saved checkpoints\n      # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n      save_contents: ['model', 'optimizer', 'extra']\n      # For more flexibility, you can specify the contents to load from the checkpoint.\n      load_contents: ${actor_rollout_ref.actor.checkpoint.save_contents}\n  ref:\n    strategy: ${actor_rollout_ref.actor.strategy}\n    use_torch_compile: ${actor_rollout_ref.actor.use_torch_compile}\n    megatron:\n      param_offload: False\n      tensor_model_parallel_size: 1\n      expert_model_parallel_size: 1\n      expert_tensor_parallel_size: null\n      pipeline_model_parallel_size: 1\n      virtual_pipeline_model_parallel_size: null # change VPP interface for parallelism tests\n      context_parallel_size: 1\n      sequence_parallel: True\n      use_distributed_optimizer: True\n      use_dist_checkpointing: False\n      dist_checkpointing_path: null\n      seed: ${actor_rollout_ref.actor.megatron.seed}\n      override_transformer_config: ${actor_rollout_ref.actor.megatron.override_transformer_config}\n      use_mbridge: ${actor_rollout_ref.actor.megatron.use_mbridge}\n      vanilla_mbridge: ${actor_rollout_ref.actor.megatron.vanilla_mbridge}\n    profile:\n      use_profile: False\n      profile_ranks: null\n      step_start: -1\n      step_end: -1\n      save_path: null\n    load_weight: True\n    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n  rollout:\n    name: vllm\n    mode: async # sync: LLM, async: AsyncLLM\n    temperature: 1.0\n    top_k: -1 # 0 for hf rollout, -1 for vllm rollout\n    top_p: 1\n    prompt_length: ${data.max_prompt_length}  # for xperf_gpt\n    response_length: ${data.max_response_length}\n    # for vllm rollout\n    dtype: bfloat16 # should align with FSDP\n    gpu_memory_utilization: 0.5\n    ignore_eos: False\n    enforce_eager: False\n    free_cache_engine: True\n    load_format: dummy\n    tensor_model_parallel_size: 2\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    log_prob_micro_batch_size: null # will be deprecated, use log_prob_micro_batch_size_per_gpu\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n    disable_log_stats: True\n    enable_chunked_prefill: True # could get higher throughput\n    # for hf rollout\n    do_sample: True\n    layer_name_map:\n      qkv_layer_name: qkv\n      gate_proj_layer_name: gate_up\n    # number of responses (i.e. num sample times)\n    n: 1\n    engine_kwargs: # inference engine parameters, please refer vllm/sglang official doc for detail\n      vllm: {}\n      sglang: {}\n    val_kwargs:\n      # sampling parameters for validation\n      top_k: -1 # 0 for hf rollout, -1 for vllm rollout\n      top_p: 1.0\n      temperature: 0\n      n: 1\n      do_sample: False # default eager for validation\n\n    # Multi-turn interaction config for tools or chat.\n    multi_turn:\n      # set to True for multi-turn tool interaction tasks; should set rollout.name to sglang as well\n      enable: False\n\n      # null for no limit (default max_length // 3)\n      max_assistant_turns: null\n\n      # null for no tool\n      tool_config_path: null\n\n      # null for no limit (default max_length // 3)\n      max_user_turns: null\n\n      # max parallel call for tools in single turn\n      max_parallel_calls: 1\n\n      # max length of tool response\n      max_tool_response_length: 256\n\n      # truncate side of tool response: left, middle, right\n      tool_response_truncate_side: middle\n\n      # null for no interaction\n      interaction_config_path: null\n\n      # - When set to True, the model's default chat template is used for multi-turn rollout, which typically matches production behavior.\n      # - When set to False, the token ids recorded for training are used instead; unlike the default chat template, these always include the model's full output,\n      #   which may contain additional content such as reasoning content. This maintains the consistency between training and rollout, but it will lead to longer prompts.\n      use_inference_chat_template: False\n\n      # Tokenization is performed turn by turn and the resulting token ids are concatenated to form the full conversation.\n      # To ensure this matches the result of tokenizing the entire conversation at once, a sanity check is run at the end of each multi-turn rollout to compare the two sets of token ids.\n      # Some models are known to produce different tokenization results when tokenizing turn by turn vs. all at once. aThis behavior has already been validated for them.\n      # To reduce excessive warnings, you can turn off the sanity check for these models if you are using their default chat template:\n      # Qwen/QwQ-32B, Qwen/Qwen3-xxB\n      # - disable: disable tokenization sanity check\n      # - strict: enable strict tokenization sanity check (default)\n      # - ignore_strippable: ignore strippable tokens when checking tokenization sanity\n      tokenization_sanity_check_mode: strict\n\n      # Format of the multi-turn interaction. Options: hermes, llama3_json, ...\n      format: hermes\n\n    # [Experimental] agent loop based rollout configs\n    agent:\n\n      # Number of agent loop workers\n      num_workers: 8\n\n      custom_async_server:\n        path: null\n        name: null\n\n    # support logging rollout prob for debugging purpose\n    calculate_log_probs: False\n    # Nsight system profiler configs\n  profiler:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n    discrete: False\n    all_ranks: False\n    ranks: []\n\ncritic:\n  rollout_n: ${actor_rollout_ref.rollout.n}\n  strategy: ${actor_rollout_ref.actor.strategy}\n  nccl_timeout: 600 # seconds, default is 10 minutes for torch, you can set it to a larger value if you have long-running operations like 32B or 72B model using megatron\n  optim:\n    optimizer: adam\n    lr: 1e-6\n    clip_grad: 1.0\n    total_training_steps: -1  # must be override by program\n    lr_warmup_init: 0.0  # initial learning rate for warmup, default to 0.0\n    lr_warmup_steps: null # Prioritized. None, 0 or Negative values mean delegating to lr_warmup_steps_ratio.\n    lr_warmup_steps_ratio: 0.  # the total steps will be injected during runtime\n    lr_decay_steps: null\n    lr_decay_style: constant # select from constant/linear/cosine/inverse_square_root\n    min_lr: 0.0 # minimum learning rate, default to 0.0\n    weight_decay: 0.01\n    weight_decay_incr_style: constant # select from constant/linear/cosine\n    lr_wsd_decay_style: exponential # select from constant/exponential/cosine\n    lr_wsd_decay_steps: null\n    use_checkpoint_opt_param_scheduler: False # use checkpoint optimizer parameter scheduler\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${actor_rollout_ref.model.path}\n    override_config:\n      model_config: {}\n      moe_config:\n        freeze_moe_router: False\n    external_lib: ${actor_rollout_ref.model.external_lib}\n    trust_remote_code: False\n    enable_gradient_checkpointing: True\n    gradient_checkpointing_kwargs:\n      ## Activation Checkpointing\n      activations_checkpoint_method: null\n      activations_checkpoint_granularity: null\n      activations_checkpoint_num_layers: null\n  megatron:\n    param_offload: False\n    grad_offload: False\n    optimizer_offload: False\n    tensor_model_parallel_size: 1\n    expert_model_parallel_size: 1\n    expert_tensor_parallel_size: null\n    pipeline_model_parallel_size: 1\n    virtual_pipeline_model_parallel_size: null # change VPP interface for parallelism tests\n    context_parallel_size: 1\n    sequence_parallel: True\n    use_distributed_optimizer: True\n    use_dist_checkpointing: False\n    dist_checkpointing_path: null\n    seed: ${actor_rollout_ref.actor.megatron.seed}\n    override_transformer_config: ${actor_rollout_ref.actor.megatron.override_transformer_config}\n    use_mbridge: ${actor_rollout_ref.actor.megatron.use_mbridge}\n    vanilla_mbridge: ${actor_rollout_ref.actor.megatron.vanilla_mbridge}\n  load_weight: True\n  ppo_mini_batch_size: ${actor_rollout_ref.actor.ppo_mini_batch_size}\n  ppo_micro_batch_size: null # will be deprecated, use ppo_micro_batch_size_per_gpu\n  ppo_micro_batch_size_per_gpu: null\n  use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n  ppo_max_token_len_per_gpu: 32768 # (${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}) * 2\n  forward_max_token_len_per_gpu: ${critic.ppo_max_token_len_per_gpu}\n  ppo_epochs: ${actor_rollout_ref.actor.ppo_epochs}\n  data_loader_seed: ${actor_rollout_ref.actor.data_loader_seed}\n  shuffle: ${actor_rollout_ref.actor.shuffle}\n  cliprange_value: 0.5\n  loss_agg_mode: ${actor_rollout_ref.actor.loss_agg_mode}\n  checkpoint:\n    async_save: False # save checkpoint asynchronously\n    # What to include in saved checkpoints\n    # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n    save_contents: ['model', 'optimizer', 'extra']\n    load_contents: ${critic.checkpoint.save_contents}\n  # Nsight system profiler configs\n  profiler:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n    discrete: False\n    all_ranks: False\n    ranks: []\nreward_model:\n  enable: False\n  strategy: ${actor_rollout_ref.actor.strategy}\n  nccl_timeout: 600 # seconds, default is 10 minutes for torch, you can set it to a larger value if you have long-running operations like 32B or 72B model using megatron\n  megatron:\n    param_offload: False\n    tensor_model_parallel_size: 1\n    expert_model_parallel_size: 1\n    expert_tensor_parallel_size: null\n    pipeline_model_parallel_size: 1\n    virtual_pipeline_model_parallel_size: null # change VPP interface for parallelism tests\n    context_parallel_size: 1\n    sequence_parallel: True\n    use_distributed_optimizer: False\n    use_dist_checkpointing: False\n    dist_checkpointing_path: null\n    seed: ${actor_rollout_ref.actor.megatron.seed}\n    override_transformer_config: {}\n    use_mbridge: ${actor_rollout_ref.actor.megatron.use_mbridge}\n    vanilla_mbridge: ${actor_rollout_ref.actor.megatron.vanilla_mbridge}\n  model:\n    input_tokenizer: ${actor_rollout_ref.model.path}  # set this to null if the chat template is identical\n    path: ~/models/FsfairX-LLaMA3-RM-v0.1\n    trust_remote_code: False\n    external_lib: ${actor_rollout_ref.model.external_lib}\n  load_weight: True\n  micro_batch_size: null # will be deprecated, use micro_batch_size_per_gpu\n  micro_batch_size_per_gpu: null\n  use_dynamic_bsz: ${critic.use_dynamic_bsz}\n  forward_max_token_len_per_gpu: ${critic.forward_max_token_len_per_gpu}\n  max_length: null\n  reward_manager: naive\n  launch_reward_fn_async: False # custom reward function executed async on CPU, during log_prob\n  sandbox_fusion:\n    url: null # faas url to run code in cloud sandbox\n    max_concurrent: 64 # max concurrent requests to sandbox\n    memory_limit_mb: 1024 # Max memory limit for each sandbox process in MB\n  # Nsight system profiler configs\n  profiler:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n    discrete: False\n    all_ranks: False\n    ranks: []\n\ncustom_reward_function:\n  path: null\n  name: compute_score\n\nalgorithm:\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: True\n  use_kl_in_reward: False\n  kl_penalty: kl  # how to estimate kl divergence\n  kl_ctrl:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: False\n  pf_ppo:\n    reweight_method: pow  # [\"pow\", \"max_min\", \"max_random\"]\n    weight_pow: 2.0\n\ntrainer:\n  balance_batch: True\n  total_epochs: 30\n  total_training_steps: null\n  profile_steps: null # [1,2,5] or [] or null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger: ['console', 'wandb']\n  log_val_generations: 0\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n\n  # auto: find the last ckpt to resume. If can't find, start from scratch\n  resume_mode: auto # or disable or resume_path if resume_from_path is set\n  resume_from_path: null\n  del_local_ckpt_after_load: False\n  val_before_train: True\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  # The timeout for ray worker group to wait for the register center to be ready\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  # see ppo_trainer.yaml for more details\n  controller_nsight_options:\n    trace: \"cuda,nvtx,cublas,ucx\"\n    cuda-memory-usage: \"true\"\n    cuda-graph-trace: \"graph\"\n  worker_nsight_options:\n    trace: \"cuda,nvtx,cublas,ucx\"\n    cuda-memory-usage: \"true\"\n    cuda-graph-trace: \"graph\"\n    capture-range: \"cudaProfilerApi\"\n    capture-range-end: null\n    kill: none\n  npu_profile:\n    options:\n      save_path: ./profiler_data\n      roles: [\"all\"]\n      level: level0\n      with_memory: False\n      record_shapes: False\n      with_npu: True\n      with_cpu: True\n      with_module: False\n      with_stack: False\n      analysis: True\n\nray_kwargs:\n  ray_init:\n    num_cpus: null # `None` means using all CPUs, which might cause hang if limited in systems like SLURM. Please set to a number allowed then.\n  timeline_json_file: null\n"
  },
  {
    "path": "tests/trainer/config/legacy_ppo_trainer.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# dataset config\ndata:\n\n  # Tokenizer class or path. If null, it will be inferred from the model.\n  tokenizer: null\n\n  # Whether to use shared memory for data loading.\n  use_shm: False\n\n  # Training set parquet. Can be a list or a single file.\n  # The program will read all files into memory, so it can't be too large (< 100GB).\n  # The path can be either a local path or an HDFS path.\n  # For HDFS path, we provide utils to download it to DRAM and convert it to a local path.\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n\n  # Validation parquet. Can be a list or a single file.\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n\n  # Maximum sample length to be used.\n  # Set to -1 to use full dataset, otherwise, randomly\n  # select the specified number of samples from train dataset\n  train_max_samples: -1\n\n  # Maximum sample length to be used.\n  # Set to -1 to use full dataset, otherwise, randomly\n  # select the specified number of samples from val dataset\n  val_max_samples: -1\n\n  # The field in the dataset where the prompt is located. Default is 'prompt'.\n  prompt_key: prompt\n\n  # The field used to select the reward function (if using different ones per example).\n  reward_fn_key: data_source\n\n  # Maximum prompt length. All prompts will be left-padded to this length.\n  # An error will be reported if the length is too long.\n  max_prompt_length: 512\n\n  # Maximum response length. Rollout in RL algorithms (e.g. PPO) generates up to this length.\n  max_response_length: 512\n\n  # Batch size sampled for one training iteration of different RL algorithms.\n  train_batch_size: 1024\n\n  # Batch size used during validation. Can be null.\n  val_batch_size: null\n\n  # Whether to return the original input_ids without adding chat template.\n  # This is used when the reward model's chat template differs from the policy.\n  # If using a model-based RM with different templates, this should be True.\n  return_raw_input_ids: False\n\n  # Whether to return the original chat (prompt) without applying chat template.\n  return_raw_chat: True\n\n  # Whether to return the full prompt with chat template.\n  return_full_prompt: False\n\n  # Whether to shuffle the data in the dataloader.\n  shuffle: True\n\n  # An integer seed to use when shuffling the data. If not set or set to\n  # `null`, the data shuffling will not be seeded, resulting in a different data order on each run.\n  seed: null\n\n  # num dataloader workers\n  dataloader_num_workers: 8\n\n  # Whether to shuffle the validation set.\n  validation_shuffle: False\n\n  # Whether to filter overlong prompts.\n  filter_overlong_prompts: False\n\n  # Number of workers for filtering overlong prompts.\n  # For large-scale datasets, filtering can be time-consuming.\n  # Use multiprocessing to speed up. Default is 1.\n  filter_overlong_prompts_workers: 1\n\n  # Truncate the input_ids or prompt if they exceed max_prompt_length.\n  # Options: 'error', 'left', or 'right'. Default is 'error'.\n  truncation: error\n\n  # The field in the multi-modal dataset where the image is located. Default is 'images'.\n  image_key: images\n\n  # The field in the multi-modal dataset where the video is located.\n  video_key: videos\n\n  # If the remote tokenizer has a Python file, this flag determines whether to allow using it.\n  trust_remote_code: False\n\n  # Optional: specify a custom dataset class path and name if overriding default loading behavior.\n  custom_cls:\n\n    # The path to the file containing your customized dataset class. If not specified, pre-implemented dataset will be used.\n    path: null\n\n    # The name of the dataset class within the specified file.\n    name: null\n\n  # Whether to return multi-modal inputs in the dataset. Set to False if rollout generates new multi-modal inputs.\n  return_multi_modal_inputs: True\n\n  # Data generation configuration for augmenting the dataset.\n  datagen:\n\n    # The path to the file containing your customized data generation class.\n    # E.g. 'pkg://verl.experimental.dynamic_dataset.dynamicgen_dataset'\n    path: null\n\n    # The class name of the data generation class within the specified file.\n    # E.g. 'MockDataGenerator'\n    name: null\n\n  # settings related to data sampler\n  sampler:\n\n    # the path to the module containing a curriculum class which implements the\n    # AbstractSampler interface\n    class_path: null\n\n    # the name of the curriculum class like `MySampler`\n    class_name: null\n\n  # Additional kwargs when calling tokenizer.apply_chat_template\n  apply_chat_template_kwargs: {}\n\n# config for actor, rollout and reference model\nactor_rollout_ref:\n\n  # Whether it's a hybrid engine, currently only supports hybrid engine\n  hybrid_engine: true\n\n  # common configs for the model\n  model:\n\n    _target_: verl.workers.config.HFModelConfig\n\n    # Huggingface model path. This can be either local path or HDFS path.\n    path: ~/models/deepseek-llm-7b-chat\n\n    # Custom chat template for the model.\n    custom_chat_template: null\n\n    # Whether to use shared memory (SHM) for accelerating the loading of model weights\n    use_shm: false\n\n    # Additional Python packages to register huggingface models/tokenizers.\n    external_lib: null\n\n    # Used to override model's original configurations, mainly dropout\n    override_config: {}\n\n    # Enable gradient checkpointing for actor\n    enable_gradient_checkpointing: true\n\n    # Enable activation offloading for actor\n    enable_activation_offload: false\n\n    # Whether to remove padding tokens in inputs during training\n    use_remove_padding: true\n\n    # Set to positive value to enable LoRA (e.g., 32)\n    lora_rank: 0\n\n    # LoRA scaling factor\n    lora_alpha: 16\n\n    # Target modules to apply LoRA. Options: \"all-linear\" (not recommended for VLMs) or\n    # [q_proj,k_proj,v_proj,o_proj,gate_proj,up_proj,down_proj]\n    target_modules: all-linear\n\n    # Exclude modules from applying Lora. Similar usage to target_modules and Peft.\n    # Example: '.*visual.*' for excluding the ViT in Qwen2.5-VL, as currently vllm does not support ViT Lora.\n    exclude_modules: null\n\n    # Whether to use Liger for linear layer fusion\n    use_liger: false\n\n    # Whether to use custom fused kernels (e.g., FlashAttention, fused MLP)\n    use_fused_kernels: false\n\n    # Options for fused kernels. If use_fused_kernels is true, this will be used.\n    fused_kernel_options:\n\n      # Implementation backend for fused kernels. Options: \"triton\" or \"torch\".\n      impl_backend: torch\n\n    # Whether to enable loading a remote code model\n    trust_remote_code: false\n\n  # actor configs\n  actor:\n\n    # fsdp, fsdp2 or megatron. fsdp backend used here.\n    strategy: fsdp\n\n    # Split each sample into sub-batches of this size for PPO\n    ppo_mini_batch_size: 256\n\n    # [Deprecated] Global micro batch size\n    ppo_micro_batch_size: null\n\n    # Local per-GPU micro batch size\n    ppo_micro_batch_size_per_gpu: null\n\n    # Whether to automatically adjust batch size at runtime\n    use_dynamic_bsz: false\n\n    # Max tokens per GPU in one PPO batch; affects gradient accumulation\n    # Typically it should be: n * ${data.max_prompt_length} + ${data.max_response_length}\n    ppo_max_token_len_per_gpu: 16384\n\n    # Gradient clipping for actor updates\n    grad_clip: 1.0\n\n    # PPO clip ratio\n    clip_ratio: 0.2\n\n    # Lower bound for asymmetric clipping (used in dual-clip PPO)\n    clip_ratio_low: 0.2\n\n    # Upper bound for asymmetric clipping (used in dual-clip PPO)\n    clip_ratio_high: 0.2\n\n    # policy loss config\n    policy_loss:\n\n      # Loss function mode: vanilla / clip-cov / kl-cov /gpg from https://arxiv.org/abs/2505.22617\n      loss_mode: \"vanilla\"\n\n      # Ratio of tokens to be clipped for clip-cov loss\n      clip_cov_ratio: 0.0002\n\n      # Lower bound for clip-cov loss\n      clip_cov_lb: 1.0\n\n      # Upper bound for clip-cov loss\n      clip_cov_ub: 5.0\n\n      # Ratio of tokens to be applied kl penalty for kl-cov loss\n      kl_cov_ratio: 0.0002\n\n      # KL divergence penalty coefficient\n      ppo_kl_coef: 0.1\n\n    # Constant C in Dual-clip PPO; clips when advantage < 0 and ratio > C\n    clip_ratio_c: 3.0\n\n    # Loss aggregation mode: \"token-mean\", \"seq-mean-token-sum\", \"seq-mean-token-mean\", or \"seq-mean-token-sum-norm\"\n    loss_agg_mode: token-mean\n\n    # Scale factor for \"seq-mean-token-sum-norm\" loss aggregation mode.\n    # If null, uses response_length. Set to a constant to ensure consistent normalization.\n    loss_scale_factor: null\n\n    # Entropy regularization coefficient in PPO loss\n    entropy_coeff: 0\n\n    # Whether to use KL loss instead of KL reward penalty. True for GRPO\n    use_kl_loss: false\n\n    # Whether to use torch.compile()\n    use_torch_compile: true\n\n    # KL loss coefficient when use_kl_loss is enabled. For GRPO\n    kl_loss_coef: 0.001\n\n    # Type of KL divergence loss. Options: \"kl\"(k1), \"abs\", \"mse\"(k2), \"low_var_kl\"(k3), \"full\"\n    kl_loss_type: low_var_kl\n\n    # Number of PPO epochs per batch\n    ppo_epochs: 1\n\n    # Shuffle training data across PPO epochs\n    shuffle: false\n\n    # Sequence parallelism size for Ulysses-style model parallelism\n    ulysses_sequence_parallel_size: 1\n\n    # calculate entropy with chunking to reduce memory peak\n    entropy_from_logits_with_chunking: False\n\n    # recompute entropy\n    entropy_checkpointing: False\n\n    # checkpoint configs\n    checkpoint:\n\n      # What to include in saved checkpoints\n      # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n      save_contents: ['model', 'optimizer', 'extra']\n\n      # For more flexibility, you can specify the contents to load from the checkpoint.\n      load_contents: ${actor_rollout_ref.actor.checkpoint.save_contents}\n\n    # optimizer configs\n    optim:\n\n      # Learning rate\n      lr: 1e-6\n\n      # Warmup steps; negative value delegates to lr_warmup_steps_ratio\n      lr_warmup_steps: -1\n\n      # Warmup steps ratio (used if lr_warmup_steps is negative)\n      lr_warmup_steps_ratio: 0.0\n\n      # Minimum LR ratio for cosine schedule\n      min_lr_ratio: 0.0\n\n      # Number of cosine cycles in LR schedule\n      num_cycles: 0.5\n\n      # LR scheduler type: \"constant\" or \"cosine\"\n      lr_scheduler_type: constant\n\n      # Total training steps (must be overridden at runtime)\n      total_training_steps: -1\n\n      # Weight decay\n      weight_decay: 0.01\n\n    # configs for FSDP\n    fsdp_config:\n\n      # policy for wrapping the model\n      wrap_policy:\n\n        # Minimum number of parameters to trigger wrapping a layer with FSDP\n        min_num_params: 0\n\n      # Whether to offload model parameters to CPU (trades speed for memory)\n      param_offload: false\n\n      # Whether to offload optimizer state to CPU\n      optimizer_offload: false\n\n      # Only for FSDP2: offload param/grad/optimizer during train\n      offload_policy: false\n\n      # Only for FSDP2: Reshard after forward pass to reduce memory footprint\n      reshard_after_forward: true\n\n      # Number of GPUs in each FSDP shard group; -1 means auto\n      fsdp_size: -1\n\n      # Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n      # before the current forward computation.\n      forward_prefetch: False\n\n  # Reference model config.\n  # Reference model will be enabled when actor.use_kl_loss or/and algorithm.use_kl_in_reward is/are True.\n  ref:\n\n    # actor_rollout_ref.ref: FSDP config same as actor. For models larger than 7B, it’s recommended to turn on offload for ref by default\n    strategy: ${actor_rollout_ref.actor.strategy}\n\n    # config for FSDP strategy\n    fsdp_config:\n\n      # whether to offload parameters in FSDP\n      param_offload: False\n\n      # whether to perform reshard after model forward to save memory.\n      # only for fsdp2, [True, False, int between 1 and fsdp_size]\n      reshard_after_forward: True\n\n      # Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n      # before the current forward computation.\n      forward_prefetch: False\n\n      # the wrap policy for FSDP model\n      wrap_policy:\n\n        # minimum number of params in a wrapped module\n        min_num_params: 0\n\n    # whether to enable torch.compile\n    use_torch_compile: ${actor_rollout_ref.actor.use_torch_compile}\n\n    # [Will be deprecated, use log_prob_micro_batch_size_per_gpu]\n    # The batch size for one forward pass in the computation of log_prob. Global batch size.\n    log_prob_micro_batch_size: null\n\n    # The batch size for one forward pass in the computation of log_prob. Local batch size per GPU.\n    log_prob_micro_batch_size_per_gpu: null\n\n    # enable dynamic batch size (sequence packing) for log_prob computation\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n\n    # the max token length per GPU\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n\n    # sequence parallel size\n    ulysses_sequence_parallel_size: ${actor_rollout_ref.actor.ulysses_sequence_parallel_size}\n\n    # calculate entropy with chunking to reduce memory peak\n    entropy_from_logits_with_chunking: False\n\n    # recompute entropy\n    entropy_checkpointing: False\n\n  # Rollout model config.\n  rollout:\n\n    # actor_rollout_ref.rollout.name: hf/vllm/sglang.\n    name: vllm\n\n    # sync: LLM, async: AsyncLLM\n    mode: async\n\n    # Sampling temperature for rollout.\n    temperature: 1.0\n\n    # Top-k sampling parameter. -1 for vLLM rollout, 0 for HF rollout.\n    top_k: -1\n\n    # Top-p sampling parameter. Default 1.0.\n    top_p: 1\n\n\n    # typically the same as data max prompt length\n    prompt_length: ${data.max_prompt_length}\n\n    # typically the same as data max response length\n    response_length: ${data.max_response_length}\n\n    # for vllm rollout\n    # Rollout model parameters type. Align with actor model's FSDP/Megatron type.\n    dtype: bfloat16\n\n    # Fraction of GPU memory used by vLLM/SGLang for KV cache.\n    gpu_memory_utilization: 0.5\n\n    # Whether to ignore EOS and continue generating after EOS is hit.\n    ignore_eos: False\n\n    # Whether to disable CUDA graph. Default True to allow cache freeing.\n    enforce_eager: False\n\n    # Whether to free engine KVCache after generation. Set enforce_eager=True when enabled.\n    free_cache_engine: True\n\n    # Which loader to use for rollout model weights: dummy_dtensor, hf, megatron, etc.\n    # safetensors (for huge model, and set use_shm=True); dummy_dtensor: randomly init model weight\n    load_format: dummy\n\n    # for huge model, layered summon can save memory (prevent OOM) but make it slower\n    layered_summon: False\n\n    # TP size for rollout. Only effective for vLLM.\n    tensor_model_parallel_size: 2\n\n    # max number of tokens in a batch\n    max_num_batched_tokens: 8192\n\n    # max length for rollout\n    max_model_len: null\n\n    # max length of sequences\n    max_num_seqs: 1024\n\n    # [Will be deprecated, use log_prob_micro_batch_size_per_gpu] The batch size for one forward pass in the computation of log_prob. Global batch size.\n    log_prob_micro_batch_size: null\n\n    # The batch size for one forward pass in the computation of log_prob. Local batch size per GPU.\n    log_prob_micro_batch_size_per_gpu: null\n\n    # enable dynamic batch size (sequence packing) for log_prob computation\n    log_prob_use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n\n    # max token length for log_prob computation\n    log_prob_max_token_len_per_gpu: ${actor_rollout_ref.actor.ppo_max_token_len_per_gpu}\n\n    # disable logging statistics\n    disable_log_stats: True\n\n    # may get higher throughput when set to True. When activated, Please increase max_num_batched_tokens or decrease max_model_len.\n    enable_chunked_prefill: True\n\n    # for hf rollout\n    # Whether to sample during training rollout. False uses greedy sampling.\n    do_sample: True\n\n    # number of responses (i.e. num sample times). > 1 for grpo\n    n: 1\n\n    # Whether to wake up inference engine in multi-stage to reduce peak memory during training-rollout transition.\n    multi_stage_wake_up: false\n\n    # Extra inference engine arguments, please refer vllm/sglang official doc for detail\n    engine_kwargs:\n\n      # vllm engine config\n      vllm: {}\n\n      # sglang engine config\n      sglang: {}\n\n    # Sampling parameters used during validation.\n    val_kwargs:\n\n      # sampling parameters for validation\n      # Top-k sampling parameter. -1 for vLLM rollout, 0 for HF rollout.\n      top_k: -1\n\n      # Top-p sampling parameter. Default 1.0.\n      top_p: 1.0\n\n      # Sampling temperature for rollout.\n      temperature: 0\n\n      # whether to repeat n times for validation\n      n: 1\n\n      # Whether to sample during training rollout. False uses greedy sampling.\n      do_sample: False\n\n    # Multi-turn interaction config for tools or chat.\n    multi_turn:\n\n      # set to True for multi-turn tool interaction tasks; should set rollout.name to sglang as well\n      enable: False\n\n      # null for no limit (default max_length // 3)\n      max_assistant_turns: null\n\n      # null for no tool\n      tool_config_path: null\n\n      # null for no limit (default max_length // 3)\n      max_user_turns: null\n\n      # max parallel call for tools in single turn\n      max_parallel_calls: 1\n\n      # max length of tool response\n      max_tool_response_length: 256\n\n      # truncate side of tool response: left, middle, right\n      tool_response_truncate_side: middle\n\n      # null for no interaction\n      interaction_config_path: null\n\n      # - When set to True, the model's default chat template is used for multi-turn rollout, which typically matches production behavior.\n      # - When set to False, the token ids recorded for training are used instead; unlike the default chat template, these always include the model's full output,\n      #   which may contain additional content such as reasoning content. This maintains the consistency between training and rollout, but it will lead to longer prompts.\n      use_inference_chat_template: False\n\n      # Tokenization is performed turn by turn and the resulting token ids are concatenated to form the full conversation.\n      # To ensure this matches the result of tokenizing the entire conversation at once, a sanity check is run at the end of each multi-turn rollout to compare the two sets of token ids.\n      # Some models are known to produce different tokenization results when tokenizing turn by turn vs. all at once. aThis behavior has already been validated for them.\n      # To reduce excessive warnings, you can turn off the sanity check for these models if you are using their default chat template:\n      # Qwen/QwQ-32B, Qwen/Qwen3-xxB\n      # - disable: disable tokenization sanity check\n      # - strict: enable strict tokenization sanity check (default)\n      # - ignore_strippable: ignore strippable tokens when checking tokenization sanity\n      tokenization_sanity_check_mode: strict\n\n      # Format of the multi-turn interaction. Options: hermes, llama3_json, ...\n      format: hermes\n\n    # support logging rollout prob for debugging purpose\n    calculate_log_probs: False\n\n    # [Experimental] agent loop based rollout configs\n    agent:\n\n      # Number of agent loop workers\n      num_workers: 8\n\n      # custom async server configs\n      custom_async_server:\n\n        # Path to the custom async server implementation\n        path: null\n\n        # Class name of the custom async server class (e.g. AsyncvLLMServer)\n        name: null\n\n  # profiler configs\n  profiler:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: False\n\n    # Whether to profile all ranks.\n    all_ranks: False\n\n    # The ranks that will be profiled. [] or [0,1,...]\n    ranks: []\n\n# configs for the critic\ncritic:\n\n  # Number of rollouts per update (mirrors actor rollout_n)\n  rollout_n: ${actor_rollout_ref.rollout.n}\n\n  # fsdp or fsdp2 strategy used for critic model training\n  strategy: ${actor_rollout_ref.actor.strategy}\n\n  # optimizer configs\n  optim:\n\n    # Learning rate\n    lr: 1e-5\n\n    # Warmup steps ratio; total steps will be injected at runtime\n    lr_warmup_steps_ratio: 0.\n\n    # Minimum LR ratio for cosine schedule\n    min_lr_ratio: 0.0\n\n    # LR scheduler type: \"constant\" or \"cosine\"\n    lr_scheduler_type: constant\n\n    # Total training steps (must be overridden at runtime)\n    total_training_steps: -1\n\n    # Weight decay\n    weight_decay: 0.01\n\n  # model config for the critic\n  model:\n\n    # Path to pretrained model weights\n    path: ~/models/deepseek-llm-7b-chat\n\n    # Whether to use shared memory for loading the model\n    use_shm: False\n\n    # Tokenizer path (defaults to actor's model path)\n    tokenizer_path: ${actor_rollout_ref.model.path}\n\n    # Hugging Face config override\n    override_config: { }\n\n    # External model implementation (optional)\n    external_lib: ${actor_rollout_ref.model.external_lib}\n\n    # Enable gradient checkpointing to save memory\n    enable_gradient_checkpointing: True\n\n    # Offload activations to CPU to reduce GPU memory usage\n    enable_activation_offload: False\n\n    # Use remove padding optimization (saves compute)\n    use_remove_padding: False\n\n    # Whether to trust remote code from Hugging Face models\n    trust_remote_code: ${actor_rollout_ref.model.trust_remote_code}\n\n    # FSDP-specific config\n    fsdp_config:\n\n      # Whether to offload model parameters to CPU\n      param_offload: False\n\n      # Whether to offload optimizer state to CPU\n      optimizer_offload: False\n\n      # Only for FSDP2: offload param/grad/optimizer during train\n      offload_policy: False\n\n      # Only for FSDP2: Reshard after forward pass to reduce memory footprint\n      reshard_after_forward: True\n\n      # Policy for wrapping layers with FSDP\n      wrap_policy:\n\n        # Minimum number of parameters to trigger wrapping\n        min_num_params: 0\n\n      # Number of GPUs in each FSDP shard group; -1 means auto\n      fsdp_size: -1\n\n      # Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n      # before the current forward computation.\n      forward_prefetch: False\n\n    # Set to positive value to enable LoRA (e.g., 32)\n    lora_rank: 0\n\n    # LoRA scaling factor\n    lora_alpha: 16\n\n    # LoRA target modules: \"all-linear\" or list of linear projection layers\n    target_modules: all-linear\n\n  # PPO mini-batch size per update\n  ppo_mini_batch_size: ${actor_rollout_ref.actor.ppo_mini_batch_size}\n\n  # [Deprecated] Global micro batch size\n  ppo_micro_batch_size: null\n\n  # Local per-GPU micro batch size\n  ppo_micro_batch_size_per_gpu: null\n\n  # Forward-only batch size (global)\n  forward_micro_batch_size: ${critic.ppo_micro_batch_size}\n\n  # Forward-only batch size (per GPU)\n  forward_micro_batch_size_per_gpu: ${critic.ppo_micro_batch_size_per_gpu}\n\n  # Whether to automatically adjust batch size at runtime\n  use_dynamic_bsz: ${actor_rollout_ref.actor.use_dynamic_bsz}\n\n  # Max tokens per GPU in one PPO batch (doubled for critic)\n  ppo_max_token_len_per_gpu: 32768\n\n  # Max token length per GPU in forward pass\n  forward_max_token_len_per_gpu: ${critic.ppo_max_token_len_per_gpu}\n\n  # Sequence parallelism size for Ulysses-style model parallelism\n  ulysses_sequence_parallel_size: 1\n\n  # Number of PPO epochs per batch\n  ppo_epochs: ${actor_rollout_ref.actor.ppo_epochs}\n\n  # Shuffle training data across PPO epochs\n  shuffle: ${actor_rollout_ref.actor.shuffle}\n\n  # Gradient clipping for critic updates\n  grad_clip: 1.0\n\n  # PPO value function clipping range\n  cliprange_value: 0.5\n\n  # Loss aggregation mode: \"token-mean\", \"seq-mean-token-sum\", or \"seq-mean-token-mean\"\n  loss_agg_mode: ${actor_rollout_ref.actor.loss_agg_mode}\n\n  # checkpoint configs\n  checkpoint:\n\n    # What to include in saved checkpoints\n    # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n    save_contents: ['model', 'optimizer', 'extra']\n\n    # What to include when loading checkpoints\n    load_contents: ${critic.checkpoint.save_contents}\n\n  # profiler configs\n  # the corresponding dataclass is verl.utils.profiler.ProfilerConfig.\n  profiler:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: False\n\n    # Whether to profile all ranks.\n    all_ranks: False\n\n    # The ranks that will be profiled. [] or [0,1,...]\n    ranks: []\n\n# configs for the reward model\nreward_model:\n\n  # Whether to enable reward model. If False, we compute the reward only with the user-defined reward functions.\n  # In GSM8K and Math examples, we disable reward model.\n  # For RLHF alignment example using full_hh_rlhf, we utilize reward model to assess the responses.\n  # If False, the following parameters are not effective\n  enable: False\n\n  # FSDP strategy: \"fsdp\" or \"fsdp2\"\n  strategy: ${actor_rollout_ref.actor.strategy}\n\n  # model config for reward scoring\n  model:\n\n    # Input tokenizer. If the reward model’s chat template is inconsistent with the policy,\n    # we need to first decode to plaintext, then apply the rm’s chat_template.\n    # Then score with RM. If chat_templates are consistent, it can be set to null.\n    input_tokenizer: ${actor_rollout_ref.model.path}\n\n    # RM’s HDFS path or local path. Note that RM only supports AutoModelForSequenceClassification.\n    # Other model types need to define their own RewardModelWorker and pass it from the code.\n    path: ~/models/FsfairX-LLaMA3-RM-v0.1\n\n    # Whether to use shared memory for loading the model\n    use_shm: False\n\n    # External model implementation (optional)\n    external_lib: ${actor_rollout_ref.model.external_lib}\n\n    # Use remove padding optimization (saves compute)\n    use_remove_padding: False\n\n    # Whether to use fused reward kernels for speedup\n    use_fused_kernels: ${actor_rollout_ref.model.use_fused_kernels}\n\n    # Whether to enable loading a remote code model, default to False\n    trust_remote_code: False\n\n    # FSDP-specific config\n    fsdp_config:\n\n      # Policy for wrapping layers with FSDP\n      wrap_policy:\n\n        # Minimum number of parameters to trigger wrapping\n        min_num_params: 0\n\n      # Whether to offload model parameters to CPU\n      param_offload: False\n\n      # Only for FSDP2: Reshard after forward pass to reduce memory footprint\n      reshard_after_forward: True\n\n      # Number of GPUs in each FSDP shard group; -1 means auto\n      fsdp_size: -1\n\n      # Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n      # before the current forward computation.\n      forward_prefetch: False\n\n  # [Deprecated] Global micro batch size\n  micro_batch_size: null\n\n  # Local per-GPU micro batch size\n  micro_batch_size_per_gpu: null\n\n  # Maximum sequence length to process for scoring\n  max_length: null\n\n  # Sequence parallelism size for Ulysses-style model parallelism\n  ulysses_sequence_parallel_size: 1\n\n  # Whether to dynamically adjust batch size at runtime\n  use_dynamic_bsz: ${critic.use_dynamic_bsz}\n\n  # Maximum number of tokens per GPU in one forward pass\n  forward_max_token_len_per_gpu: ${critic.forward_max_token_len_per_gpu}\n\n  # Reward Manager. This defines the mechanism of computing rule-based reward and handling different reward sources.\n  # Default is naive. If all verification functions are multiprocessing-safe,\n  # the reward manager can be set to prime for parallel verification.\n  reward_manager: naive\n\n  # Whether to launch custom reward function asynchronously during log_prob\n  launch_reward_fn_async: False\n\n  # Cloud/local sandbox fusion configuration for custom reward logic\n  sandbox_fusion:\n\n    # Cloud/local function URL for sandbox execution\n    url: null\n\n    # Max concurrent requests allowed to sandbox\n    max_concurrent: 64\n\n    # Max memory limit for each sandbox process in MB\n    memory_limit_mb: 1024\n\n  # profiler configs\n  profiler:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.ProfilerConfig\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: False\n\n    # Whether to profile all ranks.\n    all_ranks: False\n\n    # The ranks that will be profiled. [] or [0,1,...]\n    ranks: []\n\n# custom reward function definition\ncustom_reward_function:\n\n  # The path to the file containing your customized reward function.\n  # If not specified, pre-implemented reward functions will be used.\n  path: null\n\n  # The name of the reward function within the specified file. Default is 'compute_score'.\n  name: compute_score\n\n# config for the algorithm\nalgorithm:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.AlgoConfig\n\n  # Discount factor for future rewards\n  gamma: 1.0\n\n  # Trade-off between bias and variance in the GAE estimator\n  lam: 1.0\n\n  # Advantage estimator type: \"gae\", \"grpo\", \"reinforce_plus_plus\", etc.\n  adv_estimator: gae\n\n  # Whether to normalize advantages by std (specific to GRPO)\n  norm_adv_by_std_in_grpo: True\n\n  # Whether to enable in-reward KL penalty\n  use_kl_in_reward: False\n\n  # How to estimate KL divergence: \"kl\", \"abs\", \"mse\", \"low_var_kl\", or \"full\"\n  kl_penalty: kl\n\n  # KL control configuration\n  kl_ctrl:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.trainer.config.KLControlConfig\n\n    # KL control type: \"fixed\" or \"adaptive\"\n    type: fixed\n\n    # Initial coefficient for KL penalty\n    kl_coef: 0.001\n\n    # Horizon value for adaptive controller (if enabled)\n    horizon: 10000\n\n    # Target KL divergence (used for adaptive controller)\n    target_kl: 0.1\n\n  # Whether to enable preference feedback PPO\n  use_pf_ppo: False\n\n  # Preference feedback PPO settings\n  pf_ppo:\n\n    # Method for reweighting samples: \"pow\", \"max_min\", or \"max_random\"\n    reweight_method: pow\n\n    # Power used for weight scaling in \"pow\" method\n    weight_pow: 2.0\n\n# config for the trainer\ntrainer:\n\n  # Whether to balance batch sizes across distributed workers\n  balance_batch: True\n\n  # Number of epochs in training\n  total_epochs: 30\n\n  # Total training steps (can be set explicitly or derived from epochs)\n  total_training_steps: null\n\n  # The steps that will be profiled. null means no profiling. null or [1,2,5,...]\n  profile_steps: null\n\n  # controller Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n  ## reference https://docs.nvidia.com/nsight-systems/UserGuide/index.html\n  ## reference https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html\n  controller_nsight_options:\n\n    # Select the API(s) to be traced.\n    trace: \"cuda,nvtx,cublas,ucx\"\n\n    # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n    cuda-memory-usage: \"true\"\n\n    # CUDA graphs will be traced as a whole\n    cuda-graph-trace: \"graph\"\n\n  # worker Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n  worker_nsight_options:\n\n    # Select the API(s) to be traced.\n    trace: \"cuda,nvtx,cublas,ucx\"\n\n    # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n    cuda-memory-usage: \"true\"\n\n    # CUDA graphs will be traced as a whole\n    cuda-graph-trace: \"graph\"\n\n    # Profiling only in a range of torch.cuda.profiler.start and stop. Do not change this config.\n    capture-range: \"cudaProfilerApi\"\n\n    # Specify the desired behavior when a capture range ends.\n    # In verl we need the orch.cuda.profiler.start/stop pair to repeats n times.\n    # valid values are \"repeat-shutdown:n\" or null.\n    # For normal whole step profiling, n = len(profile_steps);\n    # but for discrete profiling, n = len(profile_steps) * Number(subtasks).\n    # Or you can just leave it null and the program will use n = len(profile_steps) * 6;\n    capture-range-end: null\n\n    # Send signal to the target application's process group. We let the program to exit by itself.\n    kill: none\n\n  # Config for npu profiler. Must set when profile_steps is not None and torch_npu is available.\n  npu_profile:\n\n    # Options for the npu profiler\n    options:\n\n      # Storage path of collected data.\n      save_path: ./profiler_data\n\n      # The roles that will be profiled. Only takes effect in discrete mode.\n      # optional values: all, rollout_generate, actor_compute_log_prob, actor_update and ref_compute_log_prob.\n      # \"all\" means all roles will be profiled.\n      roles: [\"all\"]\n\n      # Collection level, optional values: level_none, level0, level1, level2.\n      level: level0\n\n      # Whether to enable memory analysis.\n      with_memory: False\n\n      # Whether to record tensor shape.\n      record_shapes: False\n\n      # Whether to record Device-side performance data.\n      with_npu: True\n\n      # Whether to record Host-side performance data.\n      with_cpu: True\n\n      # Whether to record Python call stack information.\n      with_module: False\n\n      # Whether to record operator call stack information.\n      with_stack: False\n\n      # Whether to automatically parse the data.\n      analysis: True\n\n  # Project name for experiment tracking (e.g., wandb)\n  project_name: verl_examples\n\n  # Experiment name for run identification in tracking tools\n  experiment_name: gsm8k\n\n  # Logging backends to use: \"console\", \"wandb\", etc.\n  logger: [ 'console', 'wandb' ]\n\n  # Number of generations to log during validation\n  log_val_generations: 0\n\n  # Directory for logging rollout data; no dump if null\n  rollout_data_dir: null\n\n  # Directory for logging validation data; no dump if null\n  validation_data_dir: null\n\n  # Number of nodes used in the training\n  nnodes: 1\n\n  # Number of GPUs per node\n  n_gpus_per_node: 8\n\n  # Save frequency (by iteration) for model checkpoints\n  save_freq: -1\n\n  # ESI refers to the elastic server instance used during training, similar to the training plan. For example,\n  # if you purchase 10 hours of computing power, the ESI will automatically shut down after 10 hours of training.\n  # To ensure a checkpoint is saved before ESI shuts down, the system will start saving a checkpoint in advance.\n  # The advance time is calculated as: Advance Time = Longest historical step duration + Checkpoint save duration + esi_redundant_time.\n  # Here, esi_redundant_time is a user-defined value that further extends the advance time for added safety.\n  esi_redundant_time: 0\n\n  # Resume mode: \"auto\", \"disable\", or \"resume_path\"\n  # \"auto\": resume from last checkpoint if available\n  # \"disable\": start from scratch\n  # \"resume_path\": resume from a user-defined path\n  resume_mode: auto\n\n  # Path to resume training from (only used when resume_mode is \"resume_path\")\n  resume_from_path: null\n\n  # Whether to run validation before training begins\n  val_before_train: True\n\n  # Whether to run validation only\n  val_only: False\n\n  # Validation frequency (in training iterations)\n  test_freq: -1\n\n  # Number of iterations to warm up the critic before updating policy\n  critic_warmup: 0\n\n  # Default path to distributed filesystem for saving checkpoints\n  default_hdfs_dir: null\n\n  # Whether to delete local checkpoints after loading\n  del_local_ckpt_after_load: False\n\n  # Default local directory for saving checkpoints\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n\n  # Maximum number of actor checkpoints to keep\n  max_actor_ckpt_to_keep: null\n\n  # Maximum number of critic checkpoints to keep\n  max_critic_ckpt_to_keep: null\n\n  # Timeout (in seconds) for Ray worker to wait for registration\n  ray_wait_register_center_timeout: 300\n\n  # Device to run training on (e.g., \"cuda\", \"cpu\")\n  device: cuda\n\n# configs related to ray\nray_kwargs:\n  # configs related to ray initialization\n  ray_init:\n\n    # Number of CPUs for Ray. Use a fixed number instead of null when using SLURM.\n    num_cpus: null\n\n  # Path to save Ray timeline JSON for performance profiling\n  timeline_json_file: null\n"
  },
  {
    "path": "tests/trainer/config/test_algo_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nimport numpy as np\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom verl.trainer.config import AlgoConfig, KLControlConfig\nfrom verl.trainer.ppo.core_algos import (\n    compute_gae_advantage_return,\n    compute_grpo_outcome_advantage,\n    get_adv_estimator_fn,\n)\nfrom verl.utils.config import omega_conf_to_dataclass\n\n\nclass TestAlgoConfig(unittest.TestCase):\n    \"\"\"Test the AlgoConfig dataclass and its integration with core algorithms.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up test fixtures.\"\"\"\n        # Create a sample algorithm config as DictConfig (similar to what comes from YAML)\n        self.config_dict = {\n            \"_target_\": \"verl.trainer.config.AlgoConfig\",\n            \"gamma\": 0.99,\n            \"lam\": 0.95,\n            \"adv_estimator\": \"gae\",\n            \"norm_adv_by_std_in_grpo\": True,\n            \"use_kl_in_reward\": True,\n            \"kl_penalty\": \"kl\",\n            \"kl_ctrl\": {\n                \"_target_\": \"verl.trainer.config.KLControlConfig\",\n                \"type\": \"adaptive\",\n                \"kl_coef\": 0.002,\n                \"horizon\": 5000,\n                \"target_kl\": 0.05,\n            },\n            \"use_pf_ppo\": True,\n            \"pf_ppo\": {\"reweight_method\": \"max_min\", \"weight_pow\": 3.0},\n        }\n        self.omega_config = OmegaConf.create(self.config_dict)\n\n    def test_dataclass_creation_from_dict(self):\n        \"\"\"Test creating AlgoConfig from dictionary.\"\"\"\n        config = omega_conf_to_dataclass(self.config_dict)\n\n        self.assertIsInstance(config, AlgoConfig)\n        self.assertEqual(config.gamma, 0.99)\n        self.assertEqual(config.lam, 0.95)\n        self.assertEqual(config.adv_estimator, \"gae\")\n        self.assertTrue(config.norm_adv_by_std_in_grpo)\n        self.assertTrue(config.use_kl_in_reward)\n        self.assertEqual(config.kl_penalty, \"kl\")\n        self.assertTrue(config.use_pf_ppo)\n\n    def test_dataclass_creation_from_omega_config(self):\n        \"\"\"Test creating AlgoConfig from OmegaConf DictConfig.\"\"\"\n        config = omega_conf_to_dataclass(self.omega_config)\n\n        self.assertIsInstance(config, AlgoConfig)\n        self.assertEqual(config.gamma, 0.99)\n        self.assertEqual(config.lam, 0.95)\n\n    def test_nested_configs(self):\n        \"\"\"Test that nested configurations are properly converted.\"\"\"\n        config = omega_conf_to_dataclass(self.omega_config)\n\n        # Test KL control config\n        self.assertIsInstance(config.kl_ctrl, KLControlConfig)\n        self.assertEqual(config.kl_ctrl.type, \"adaptive\")\n        self.assertEqual(config.kl_ctrl.kl_coef, 0.002)\n        self.assertEqual(config.kl_ctrl.horizon, 5000)\n        self.assertEqual(config.kl_ctrl.target_kl, 0.05)\n\n        # Test PF PPO config\n        self.assertEqual(config.pf_ppo.get(\"reweight_method\"), \"max_min\")\n        self.assertEqual(config.pf_ppo.get(\"weight_pow\"), 3.0)\n\n    def test_default_values(self):\n        \"\"\"Test that default values are properly set.\"\"\"\n        minimal_config = {\"gamma\": 0.8}\n        config = omega_conf_to_dataclass(minimal_config, AlgoConfig)\n\n        self.assertEqual(config.gamma, 0.8)\n        self.assertEqual(config.lam, 1.0)  # default value\n        self.assertEqual(config.adv_estimator, \"gae\")  # default value\n        self.assertTrue(config.norm_adv_by_std_in_grpo)  # default value\n        self.assertFalse(config.use_kl_in_reward)  # default value\n        self.assertEqual(config.kl_penalty, \"kl\")  # default value\n        self.assertFalse(config.use_pf_ppo)  # default value\n\n    def test_get_method_backward_compatibility(self):\n        \"\"\"Test the get method for backward compatibility.\"\"\"\n        config = omega_conf_to_dataclass(self.omega_config)\n\n        # Test existing attribute\n        self.assertEqual(config.get(\"gamma\"), 0.99)\n        self.assertEqual(config.get(\"gamma\", 1.0), 0.99)\n\n        # Test non-existing attribute\n        self.assertIsNone(config.get(\"non_existing\"))\n        self.assertEqual(config.get(\"non_existing\", \"default\"), \"default\")\n\n    def test_post_init_nested_configs(self):\n        \"\"\"Test that __post_init__ properly initializes nested configs when None.\"\"\"\n        # Create config without nested configs\n        minimal_config = AlgoConfig(gamma=0.9)\n\n        # Check that nested configs are initialized\n        self.assertIsNotNone(minimal_config.kl_ctrl)\n        self.assertIsInstance(minimal_config.kl_ctrl, KLControlConfig)\n        assert not minimal_config.pf_ppo\n\n    def test_config_init_from_yaml(self):\n        import os\n\n        from hydra import compose, initialize_config_dir\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n            cfg = compose(config_name=\"ppo_trainer\")\n        algo_config = omega_conf_to_dataclass(cfg.algorithm)\n        from verl.trainer.config import AlgoConfig\n\n        assert isinstance(algo_config, AlgoConfig)\n\n\nclass TestAlgoCompute(unittest.TestCase):\n    \"\"\"Test the AlgoConfig dataclass and its integration with core algorithms.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up test fixtures.\"\"\"\n        self.algo_config = AlgoConfig(\n            gamma=0.99,\n            lam=0.95,\n            adv_estimator=\"gae\",\n            norm_adv_by_std_in_grpo=True,\n            use_kl_in_reward=True,\n            kl_penalty=\"kl\",\n            kl_ctrl=KLControlConfig(type=\"adaptive\", kl_coef=0.002, horizon=5000, target_kl=0.05),\n            use_pf_ppo=True,\n            pf_ppo={\"reweight_method\": \"max_min\", \"weight_pow\": 3.0},\n        )\n\n    def test_advantage_estimator_with_cfg(self):\n        \"\"\"Test integration with advantage estimators from core_algos.\"\"\"\n        config = self.algo_config\n\n        # Test GAE advantage estimator\n        adv_fn = get_adv_estimator_fn(config.adv_estimator)\n        self.assertIsNotNone(adv_fn)\n\n        # Test with actual GAE computation\n        batch_size, seq_len = 2, 5\n        token_level_rewards = torch.randn(batch_size, seq_len)\n        values = torch.randn(batch_size, seq_len)\n        response_mask = torch.ones(batch_size, seq_len)\n\n        advantages, returns = compute_gae_advantage_return(\n            token_level_rewards=token_level_rewards,\n            values=values,\n            response_mask=response_mask,\n            gamma=config.gamma,\n            lam=config.lam,\n        )\n\n        self.assertEqual(advantages.shape, (batch_size, seq_len))\n        self.assertEqual(returns.shape, (batch_size, seq_len))\n\n    def test_grpo_advantage_estimator_with_cfg(self):\n        \"\"\"Test integration with GRPO advantage estimator.\"\"\"\n        grpo_config = AlgoConfig(adv_estimator=\"grpo\", norm_adv_by_std_in_grpo=True)\n\n        # Test GRPO advantage computation\n        batch_size, seq_len = 4, 3\n        token_level_rewards = torch.tensor([[1.0, 0.5, 0.0], [2.0, 1.0, 0.0], [0.5, 0.2, 0.0], [1.5, 0.8, 0.0]])\n        response_mask = torch.ones(batch_size, seq_len)\n        index = np.array([0, 0, 1, 1])  # Two groups\n\n        advantages, returns = compute_grpo_outcome_advantage(\n            token_level_rewards=token_level_rewards,\n            response_mask=response_mask,\n            index=index,\n            norm_adv_by_std_in_grpo=grpo_config.norm_adv_by_std_in_grpo,\n        )\n\n        self.assertEqual(advantages.shape, (batch_size, seq_len))\n        self.assertEqual(returns.shape, (batch_size, seq_len))\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/trainer/config/test_legacy_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport unittest\nimport warnings\n\nfrom hydra import compose, initialize_config_dir\nfrom hydra.core.global_hydra import GlobalHydra\nfrom omegaconf import OmegaConf\n\n_BREAKING_CHANGES = [\n    \"critic.optim.lr\",  # mcore critic lr init value 1e-6 -> 1e-5\n    \"actor_rollout_ref.actor.optim.lr_warmup_steps\",  # None -> -1\n    \"critic.optim.lr_warmup_steps\",  # None -> -1\n    \"actor_rollout_ref.rollout.name\",  # vllm -> ???\n    \"actor_rollout_ref.actor.megatron.expert_tensor_parallel_size\",\n    \"actor_rollout_ref.ref.megatron.expert_tensor_parallel_size\",\n    \"critic.megatron.expert_tensor_parallel_size\",\n    \"reward_model.megatron.expert_tensor_parallel_size\",\n]\n\n\nclass TestConfigComparison(unittest.TestCase):\n    \"\"\"Test that current configs match their legacy counterparts exactly.\"\"\"\n\n    ignored_keys = [\n        \"enable_gradient_checkpointing\",\n        \"gradient_checkpointing_kwargs\",\n        \"activations_checkpoint_method\",\n        \"activations_checkpoint_granularity\",\n        \"activations_checkpoint_num_layers\",\n        \"discrete\",\n        \"profiler\",\n        \"profile\",\n        \"use_profile\",\n        \"npu_profile\",\n        \"profile_steps\",\n        \"worker_nsight_options\",\n        \"controller_nsight_options\",\n    ]\n    ignored_paths = [\"reward_model\", \"custom_reward_function\"]\n\n    def _compare_configs_recursively(\n        self, current_config, legacy_config, path=\"\", legacy_allow_missing=True, current_allow_missing=False\n    ):\n        \"\"\"Recursively compare two OmegaConf configs and assert they are identical.\n\n        Args:\n            legacy_allow_missing (bool): sometimes the legacy megatron config contains fewer keys and\n              we allow that to happen\n        \"\"\"\n        if path in self.ignored_paths:\n            return\n\n        if isinstance(current_config, dict) and isinstance(legacy_config, dict):\n            current_keys = set(current_config.keys())\n            legacy_keys = set(legacy_config.keys())\n\n            missing_in_current = legacy_keys - current_keys\n            missing_in_legacy = current_keys - legacy_keys\n\n            # Ignore specific keys that are allowed to be missing\n            for key in self.ignored_keys:\n                if key in missing_in_current:\n                    missing_in_current.remove(key)\n                if key in missing_in_legacy:\n                    missing_in_legacy.remove(key)\n\n            if missing_in_current:\n                msg = f\"Keys missing in current config at {path}: {missing_in_current}\"\n                if current_allow_missing:\n                    warnings.warn(msg, stacklevel=1)\n                else:\n                    self.fail(f\"Keys missing in current config at {path}: {missing_in_current}\")\n            if missing_in_legacy:\n                # if the legacy\n                msg = f\"Keys missing in legacy config at {path}: {missing_in_legacy}\"\n                if legacy_allow_missing:\n                    warnings.warn(msg, stacklevel=1)\n                else:\n                    self.fail(msg)\n\n            for key in current_keys:\n                current_path = f\"{path}.{key}\" if path else key\n                if key in legacy_config:\n                    self._compare_configs_recursively(current_config[key], legacy_config[key], current_path)\n        elif isinstance(current_config, list) and isinstance(legacy_config, list):\n            self.assertEqual(\n                len(current_config),\n                len(legacy_config),\n                f\"List lengths differ at {path}: current={len(current_config)}, legacy={len(legacy_config)}\",\n            )\n            for i, (current_item, legacy_item) in enumerate(zip(current_config, legacy_config, strict=True)):\n                self._compare_configs_recursively(current_item, legacy_item, f\"{path}[{i}]\")\n        elif path not in _BREAKING_CHANGES:\n            self.assertEqual(\n                current_config,\n                legacy_config,\n                f\"Values differ at {path}: current={current_config}, legacy={legacy_config}\",\n            )\n\n    def test_ppo_trainer_config_matches_legacy(self):\n        \"\"\"Test that ppo_trainer.yaml matches legacy_ppo_trainer.yaml exactly.\"\"\"\n        import os\n\n        from hydra import compose, initialize_config_dir\n        from hydra.core.global_hydra import GlobalHydra\n\n        GlobalHydra.instance().clear()\n\n        try:\n            with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n                current_config = compose(config_name=\"ppo_trainer\")\n\n            legacy_config = OmegaConf.load(\"tests/trainer/config/legacy_ppo_trainer.yaml\")\n            current_dict = OmegaConf.to_container(current_config, resolve=True)\n            legacy_dict = OmegaConf.to_container(legacy_config, resolve=True)\n\n            if \"defaults\" in current_dict:\n                del current_dict[\"defaults\"]\n\n            self._compare_configs_recursively(current_dict, legacy_dict)\n        finally:\n            GlobalHydra.instance().clear()\n\n    def test_ppo_megatron_trainer_config_matches_legacy(self):\n        \"\"\"Test that ppo_megatron_trainer.yaml matches legacy_ppo_megatron_trainer.yaml exactly.\"\"\"\n\n        GlobalHydra.instance().clear()\n\n        try:\n            with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n                current_config = compose(config_name=\"ppo_megatron_trainer\")\n\n            legacy_config = OmegaConf.load(\"tests/trainer/config/legacy_ppo_megatron_trainer.yaml\")\n            current_dict = OmegaConf.to_container(current_config, resolve=True)\n            legacy_dict = OmegaConf.to_container(legacy_config, resolve=True)\n\n            if \"defaults\" in current_dict:\n                del current_dict[\"defaults\"]\n\n            self._compare_configs_recursively(\n                current_dict, legacy_dict, legacy_allow_missing=True, current_allow_missing=False\n            )\n        finally:\n            GlobalHydra.instance().clear()\n\n    def test_load_component(self):\n        \"\"\"Test that ppo_megatron_trainer.yaml matches legacy_ppo_megatron_trainer.yaml exactly.\"\"\"\n\n        GlobalHydra.instance().clear()\n        configs_to_load = [\n            (\"verl/trainer/config/actor\", \"dp_actor\"),\n            (\"verl/trainer/config/actor\", \"megatron_actor\"),\n            (\"verl/trainer/config/ref\", \"dp_ref\"),\n            (\"verl/trainer/config/ref\", \"megatron_ref\"),\n            (\"verl/trainer/config/rollout\", \"rollout\"),\n        ]\n        for config_dir, config_file in configs_to_load:\n            try:\n                with initialize_config_dir(config_dir=os.path.abspath(config_dir)):\n                    compose(config_name=config_file)\n            finally:\n                GlobalHydra.instance().clear()\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/trainer/ppo/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTests for the PPO trainer module.\n\"\"\"\n"
  },
  {
    "path": "tests/trainer/ppo/test_core_algos_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport random\nimport unittest\n\nimport numpy as np\nimport pytest\nimport torch\n\nimport verl.trainer.ppo.core_algos\nfrom verl.trainer.ppo.core_algos import (\n    compute_gae_advantage_return,\n    compute_grpo_outcome_advantage,\n    compute_grpo_vectorized_outcome_advantage,\n    compute_rloo_outcome_advantage,\n    compute_rloo_vectorized_outcome_advantage,\n    get_adv_estimator_fn,\n    register_adv_est,\n)\n\n\ndef mock_test_fn():\n    pass\n\n\nclass TestRegisterAdvEst(unittest.TestCase):\n    def setUp(self):\n        \"\"\"Clear the registry before each test\"\"\"\n        verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY.clear()\n        verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY = {\n            \"gae\": lambda x: x * 2,\n            \"vtrace\": lambda x: x + 1,\n        }\n        self.ADV_ESTIMATOR_REGISTRY = verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY\n\n    def tearDown(self) -> None:\n        verl.trainer.ppo.core_algos.ADV_ESTIMATOR_REGISTRY.clear()\n        return super().tearDown()\n\n    def test_register_new_function(self):\n        \"\"\"Test registering a new function with a string name\"\"\"\n\n        @register_adv_est(\"test_estimator\")\n        def test_fn():\n            pass\n\n        self.assertIn(\"test_estimator\", self.ADV_ESTIMATOR_REGISTRY)\n        self.assertEqual(self.ADV_ESTIMATOR_REGISTRY[\"test_estimator\"], test_fn)\n\n    def test_register_with_enum(self):\n        \"\"\"Test registering with an enum value (assuming AdvantageEstimator exists)\"\"\"\n        from enum import Enum\n\n        class AdvantageEstimator(Enum):\n            TEST = \"test_enum_estimator\"\n\n        @register_adv_est(AdvantageEstimator.TEST)\n        def test_fn():\n            pass\n\n        self.assertIn(\"test_enum_estimator\", self.ADV_ESTIMATOR_REGISTRY)\n        self.assertEqual(self.ADV_ESTIMATOR_REGISTRY[\"test_enum_estimator\"], test_fn)\n\n    def test_duplicate_registration_same_function(self):\n        \"\"\"Test that registering the same function twice doesn't raise an error\"\"\"\n        register_adv_est(\"duplicate_test\")(mock_test_fn)\n        register_adv_est(\"duplicate_test\")(mock_test_fn)\n\n        self.assertEqual(self.ADV_ESTIMATOR_REGISTRY[\"duplicate_test\"], mock_test_fn)\n\n    def test_duplicate_registration_different_function(self):\n        \"\"\"Test that registering different functions with same name raises ValueError\"\"\"\n\n        @register_adv_est(\"conflict_test\")\n        def test_fn1():\n            pass\n\n        with self.assertRaises(ValueError):\n\n            @register_adv_est(\"conflict_test\")\n            def test_fn2():\n                pass\n\n    def test_decorator_preserves_function(self):\n        \"\"\"Test that the decorator returns the original function\"\"\"\n\n        def test_fn():\n            return \"original\"\n\n        decorated = register_adv_est(\"preserve_test\")(test_fn)\n        self.assertEqual(decorated(), \"original\")\n\n    def test_multiple_registrations(self):\n        \"\"\"Test registering multiple different functions\"\"\"\n        init_adv_count = len(self.ADV_ESTIMATOR_REGISTRY)\n\n        @register_adv_est(\"estimator1\")\n        def fn1():\n            pass\n\n        @register_adv_est(\"estimator2\")\n        def fn2():\n            pass\n\n        self.assertEqual(len(self.ADV_ESTIMATOR_REGISTRY), 2 + init_adv_count)\n        self.assertEqual(self.ADV_ESTIMATOR_REGISTRY[\"estimator1\"], fn1)\n        self.assertEqual(self.ADV_ESTIMATOR_REGISTRY[\"estimator2\"], fn2)\n\n    def test_get_adv_estimator_fn_valid_names(self):\n        \"\"\"Test that valid names return the correct function from registry.\"\"\"\n        # Test GAE\n        gae_fn = get_adv_estimator_fn(\"gae\")\n        assert gae_fn(5) == 10  # 5 * 2 = 10\n\n        # Test Vtrace\n        vtrace_fn = get_adv_estimator_fn(\"vtrace\")\n        assert vtrace_fn(5) == 6  # 5 + 1 = 6\n\n    def test_get_adv_estimator_fn_invalid_name(self):\n        \"\"\"Test that invalid names raise ValueError.\"\"\"\n        with pytest.raises(ValueError) as excinfo:\n            get_adv_estimator_fn(\"invalid_name\")\n        assert \"Unknown advantage estimator simply: invalid_name\" in str(excinfo.value)\n\n    def test_get_adv_estimator_fn_case_sensitive(self):\n        \"\"\"Test that name lookup is case-sensitive.\"\"\"\n        with pytest.raises(ValueError):\n            get_adv_estimator_fn(\"GAE\")  # Different case\n\n\ndef test_multi_turn_compute_gae_advantage_return():\n    \"\"\"Test multi-turn GAE skip observation tokens.\"\"\"\n    gamma = random.uniform(0.0, 1.0)\n    lam = random.uniform(0.0, 1.0)\n\n    rewards = torch.tensor([[0.0, 0.0, 0.1, 0.1, 0.1, 0.0, 0.0, 0.1, 1.0, 0.0, 0.0]], dtype=torch.float)\n\n    values1 = torch.tensor(\n        [\n            [\n                random.uniform(-100.0, 100.0),\n                random.random(),\n                4.0,\n                5.0,\n                6.0,\n                random.uniform(-100.0, 0),\n                random.random(),\n                7.0,\n                9.0,\n                0.0,\n                0.0,\n            ]\n        ],\n        dtype=torch.float,\n    )\n\n    values2 = torch.tensor(\n        [\n            [\n                random.random(),\n                random.uniform(-100.0, 100.0),\n                4.0,\n                5.0,\n                6.0,\n                random.random(),\n                random.uniform(0.0, 100.0),\n                7.0,\n                9.0,\n                0.0,\n                0.0,\n            ]\n        ],\n        dtype=torch.float,\n    )\n\n    response_mask = torch.tensor([[0, 0, 1, 1, 1, 0, 0, 1, 1, 0, 0]], dtype=torch.float)\n\n    adv1, ret1 = compute_gae_advantage_return(rewards, values1, response_mask, gamma, lam)\n    adv2, ret2 = compute_gae_advantage_return(rewards, values2, response_mask, gamma, lam)\n\n    ret1 *= response_mask\n    ret2 *= response_mask\n    assert torch.equal(adv1, adv2), f\"{adv1=}, {adv2=}\"\n    assert torch.equal(ret1, ret2), f\"{ret1=}, {ret2=}\"\n    print(f\" [CORRECT] \\n\\n{adv1=}, \\n\\n{ret1=}\")\n\n\ndef _make_group_index(batch_size: int, num_groups: int) -> np.ndarray:\n    \"\"\"Create a numpy index array ensuring each group has at least 2 samples.\"\"\"\n    assert num_groups * 2 <= batch_size, \"batch_size must allow >=2 samples per group\"\n    counts: list[int] = [2] * num_groups\n    remaining = batch_size - 2 * num_groups\n    for _ in range(remaining):\n        counts[random.randrange(num_groups)] += 1\n    index = []\n    for gid, c in enumerate(counts):\n        index.extend([gid] * c)\n    random.shuffle(index)\n    return np.asarray(index, dtype=np.int64)\n\n\ndef _rand_mask(batch_size: int, seq_len: int) -> torch.Tensor:\n    mask = torch.randint(0, 2, (batch_size, seq_len), dtype=torch.int64).float()\n    rows_without_one = (mask.sum(dim=-1) == 0).nonzero(as_tuple=True)[0]\n    if len(rows_without_one) > 0:\n        mask[rows_without_one, -1] = 1.0\n    return mask\n\n\n@pytest.mark.parametrize(\n    \"batch_size,seq_len,num_groups,seed\",\n    [\n        (64, 128, 5, 0),\n        (128, 256, 8, 1),\n        (512, 512, 10, 2),\n    ],\n)\ndef test_rloo_and_vectorized_equivalence(batch_size: int, seq_len: int, num_groups: int, seed: int):\n    torch.manual_seed(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n    index = _make_group_index(batch_size, num_groups)\n    response_mask = _rand_mask(batch_size, seq_len)\n    base_rewards = torch.randn(batch_size, seq_len, dtype=torch.float32)\n    token_level_rewards = base_rewards * response_mask\n    adv1, ret1 = compute_rloo_outcome_advantage(\n        token_level_rewards=token_level_rewards,\n        response_mask=response_mask,\n        index=index,\n    )\n    adv2, ret2 = compute_rloo_vectorized_outcome_advantage(\n        token_level_rewards=token_level_rewards,\n        response_mask=response_mask,\n        index=index,\n    )\n    # Print concise diagnostics for visibility during test runs\n    adv_max_diff = (adv1 - adv2).abs().max().item()\n    ret_max_diff = (ret1 - ret2).abs().max().item()\n    total_mask_tokens = int(response_mask.sum().item())\n    print(\n        f\"[RLOO] seed={seed} groups={num_groups} shape={adv1.shape} \"\n        f\"mask_tokens={total_mask_tokens} adv_max_diff={adv_max_diff:.3e} ret_max_diff={ret_max_diff:.3e}\"\n    )\n    assert adv1.shape == adv2.shape == (batch_size, seq_len)\n    assert ret1.shape == ret2.shape == (batch_size, seq_len)\n    assert torch.allclose(adv1, adv2, rtol=1e-5, atol=1e-6)\n    assert torch.allclose(ret1, ret2, rtol=1e-5, atol=1e-6)\n\n\n@pytest.mark.parametrize(\n    \"batch_size,seq_len,num_groups,seed\",\n    [\n        (64, 128, 5, 0),\n        (128, 256, 8, 1),\n        (512, 512, 10, 2),\n    ],\n)\ndef test_grpo_and_vectorized_equivalence(batch_size: int, seq_len: int, num_groups: int, seed: int):\n    # Set seeds for reproducibility\n    torch.manual_seed(seed)\n    random.seed(seed)\n    np.random.seed(seed)\n\n    # Generate group indices (numpy array of shape [batch_size])\n    index = _make_group_index(batch_size, num_groups)\n\n    # Generate binary response mask (at least one valid token per row)\n    response_mask = _rand_mask(batch_size, seq_len)\n\n    # Generate token-level rewards and apply mask\n    base_rewards = torch.randn(batch_size, seq_len, dtype=torch.float32)\n    token_level_rewards = base_rewards * response_mask\n\n    # Compute GRPO outcome advantage (original implementation)\n    adv1, ret1 = compute_grpo_outcome_advantage(\n        token_level_rewards=token_level_rewards,\n        response_mask=response_mask,\n        index=index,\n    )\n\n    # Compute GRPO outcome advantage (vectorized implementation)\n    adv2, ret2 = compute_grpo_vectorized_outcome_advantage(\n        token_level_rewards=token_level_rewards,\n        response_mask=response_mask,\n        index=index,\n    )\n\n    # Diagnostic info for visibility (same style as RLOO test)\n    adv_max_diff = (adv1 - adv2).abs().max().item()\n    ret_max_diff = (ret1 - ret2).abs().max().item()\n    total_mask_tokens = int(response_mask.sum().item())\n    print(\n        f\"[GRPO] seed={seed} groups={num_groups} shape={adv1.shape} \"\n        f\"mask_tokens={total_mask_tokens} adv_max_diff={adv_max_diff:.3e} ret_max_diff={ret_max_diff:.3e}\"\n    )\n\n    # Assert shape and numerical equivalence\n    assert adv1.shape == adv2.shape == (batch_size, seq_len)\n    assert ret1.shape == ret2.shape == (batch_size, seq_len)\n    assert torch.allclose(adv1, adv2, rtol=1e-5, atol=1e-6)\n    assert torch.allclose(ret1, ret2, rtol=1e-5, atol=1e-6)\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/trainer/ppo/test_metric_utils_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTests for the metric utilities in verl.trainer.ppo.metric_utils.\n\"\"\"\n\nimport unittest\nfrom unittest.mock import MagicMock, patch\n\nimport numpy as np\nimport torch\n\nfrom verl.trainer.ppo.metric_utils import (\n    bootstrap_metric,\n    calc_maj_val,\n    compute_data_metrics,\n    compute_throughout_metrics,\n    compute_timing_metrics,\n    process_validation_metrics,\n)\nfrom verl.utils.metric import (\n    reduce_metrics,\n)\nfrom verl.utils.metric.utils import (\n    AggregationType,\n    Metric,\n)\n\n\nclass TestReduceMetrics(unittest.TestCase):\n    \"\"\"Tests for the reduce_metrics function.\"\"\"\n\n    def test_reduce_metrics_basic(self):\n        \"\"\"Test that reduce_metrics correctly computes means.\"\"\"\n        metrics = {\n            \"loss\": [1.0, 2.0, 3.0],\n            \"accuracy\": [0.0, 0.5, 1.0],\n        }\n        result = reduce_metrics(metrics)\n\n        self.assertEqual(result[\"loss\"], 2.0)\n        self.assertEqual(result[\"accuracy\"], 0.5)\n\n    def test_reduce_metrics_empty(self):\n        \"\"\"Test that reduce_metrics handles empty lists.\"\"\"\n        metrics = {\n            \"empty\": [],\n        }\n        result = reduce_metrics(metrics)\n\n        self.assertTrue(np.isnan(result[\"empty\"]))\n\n    def test_reduce_metrics_single_value(self):\n        \"\"\"Test that reduce_metrics works with single values.\"\"\"\n        metrics = {\n            \"single\": [5.0],\n        }\n        result = reduce_metrics(metrics)\n\n        self.assertEqual(result[\"single\"], 5.0)\n\n\nclass TestMetric(unittest.TestCase):\n    \"\"\"Tests for the Metric class.\"\"\"\n\n    def test_init_with_string_aggregation(self):\n        \"\"\"Test Metric initialization with string aggregation type.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        self.assertEqual(metric.aggregation, AggregationType.MEAN)\n        self.assertEqual(metric.values, [])\n\n    def test_init_with_enum_aggregation(self):\n        \"\"\"Test Metric initialization with AggregationType enum.\"\"\"\n        metric = Metric(aggregation=AggregationType.SUM)\n        self.assertEqual(metric.aggregation, AggregationType.SUM)\n        self.assertEqual(metric.values, [])\n\n    def test_init_with_value(self):\n        \"\"\"Test Metric initialization with an initial value.\"\"\"\n        metric = Metric(aggregation=\"mean\", value=5.0)\n        self.assertEqual(metric.values, [5.0])\n\n    def test_init_with_invalid_aggregation(self):\n        \"\"\"Test Metric initialization with invalid aggregation type.\"\"\"\n        with self.assertRaises(ValueError):\n            Metric(aggregation=\"invalid\")\n\n    def test_append_float(self):\n        \"\"\"Test appending float values.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.append(1.0)\n        metric.append(2.0)\n        self.assertEqual(metric.values, [1.0, 2.0])\n\n    def test_append_int(self):\n        \"\"\"Test appending int values.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.append(1)\n        metric.append(2)\n        self.assertEqual(metric.values, [1, 2])\n\n    def test_append_tensor(self):\n        \"\"\"Test appending scalar tensor values.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.append(torch.tensor(3.0))\n        metric.append(torch.tensor(4.0))\n        self.assertEqual(metric.values, [3.0, 4.0])\n\n    def test_append_non_scalar_tensor_raises(self):\n        \"\"\"Test that appending non-scalar tensor raises ValueError.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        with self.assertRaises(ValueError):\n            metric.append(torch.tensor([1.0, 2.0]))\n\n    def test_append_metric(self):\n        \"\"\"Test appending another Metric extends values.\"\"\"\n        metric1 = Metric(aggregation=\"mean\", value=1.0)\n        metric1.append(2.0)\n\n        metric2 = Metric(aggregation=\"mean\", value=3.0)\n        metric2.append(metric1)\n\n        self.assertEqual(metric2.values, [3.0, 1.0, 2.0])\n\n    def test_extend_with_list(self):\n        \"\"\"Test extending with a list of values.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.extend([1.0, 2.0, 3.0])\n        self.assertEqual(metric.values, [1.0, 2.0, 3.0])\n\n    def test_extend_with_metric(self):\n        \"\"\"Test extending with another Metric.\"\"\"\n        metric1 = Metric(aggregation=\"mean\")\n        metric1.extend([1.0, 2.0])\n\n        metric2 = Metric(aggregation=\"mean\")\n        metric2.extend([3.0, 4.0])\n        metric2.extend(metric1)\n\n        self.assertEqual(metric2.values, [3.0, 4.0, 1.0, 2.0])\n\n    def test_extend_aggregation_mismatch_raises(self):\n        \"\"\"Test that extending with mismatched aggregation raises ValueError.\"\"\"\n        metric1 = Metric(aggregation=\"mean\")\n        metric2 = Metric(aggregation=\"sum\")\n\n        with self.assertRaises(ValueError):\n            metric1.extend(metric2)\n\n    def test_aggregate_mean(self):\n        \"\"\"Test aggregation with mean.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.extend([1.0, 2.0, 3.0, 4.0])\n        self.assertEqual(metric.aggregate(), 2.5)\n\n    def test_aggregate_sum(self):\n        \"\"\"Test aggregation with sum.\"\"\"\n        metric = Metric(aggregation=\"sum\")\n        metric.extend([1.0, 2.0, 3.0, 4.0])\n        self.assertEqual(metric.aggregate(), 10.0)\n\n    def test_aggregate_min(self):\n        \"\"\"Test aggregation with min.\"\"\"\n        metric = Metric(aggregation=\"min\")\n        metric.extend([3.0, 1.0, 4.0, 2.0])\n        self.assertEqual(metric.aggregate(), 1.0)\n\n    def test_aggregate_max(self):\n        \"\"\"Test aggregation with max.\"\"\"\n        metric = Metric(aggregation=\"max\")\n        metric.extend([3.0, 1.0, 4.0, 2.0])\n        self.assertEqual(metric.aggregate(), 4.0)\n\n    def test_aggregate_dp_sum_mean(self):\n        \"\"\"Test aggregate_dp with SUM and MEAN aggregations.\"\"\"\n        # Test with SUM: mean over DP ranks, then sum\n        metric1 = Metric(aggregation=\"sum\")\n        metric1.extend([1.0, 2.0])\n\n        metric2 = Metric(aggregation=\"sum\")\n        metric2.extend([3.0, 4.0])\n\n        result = Metric.aggregate_dp([metric1, metric2])\n\n        # value_arrays = [[1.0, 2.0], [3.0, 4.0]]\n        # mean over axis 0 = [2.0, 3.0]\n        # sum = 5.0\n        self.assertEqual(result, 5.0)\n\n        # Test with MEAN: mean over DP ranks, then mean\n        metric4 = Metric(aggregation=\"mean\")\n        metric4.extend([1.0, 2.0])\n\n        metric5 = Metric(aggregation=\"mean\")\n        metric5.extend([3.0, 4.0])\n\n        result = Metric.aggregate_dp([metric4, metric5])\n\n        # value_arrays = [[1.0, 2.0], [3.0, 4.0]]\n        # mean over axis 0 = [2.0, 3.0]\n        # mean = 2.5\n        self.assertEqual(result, 2.5)\n\n    def test_aggregate_dp_min_max(self):\n        \"\"\"Test aggregate_dp with MIN and MAX aggregations.\"\"\"\n        # Test with MAX: flatten, then max\n        metric1 = Metric(aggregation=\"max\")\n        metric1.extend([1.0, 2.0])\n\n        metric2 = Metric(aggregation=\"max\")\n        metric2.extend([3.0, 4.0])\n\n        result = Metric.aggregate_dp([metric1, metric2])\n\n        # value_arrays = [[1.0, 2.0], [3.0, 4.0]]\n        # flatten = [1.0, 2.0, 3.0, 4.0]\n        # max = 4.0\n        self.assertEqual(result, 4.0)\n\n        # Test with MIN: flatten, then min\n        metric4 = Metric(aggregation=\"min\")\n        metric4.extend([1.0, 2.0])\n\n        metric5 = Metric(aggregation=\"min\")\n        metric5.extend([3.0, 4.0])\n\n        result = Metric.aggregate_dp([metric4, metric5])\n\n        # value_arrays = [[1.0, 2.0], [3.0, 4.0]]\n        # flatten = [1.0, 2.0, 3.0, 4.0]\n        # min = 1.0\n        self.assertEqual(result, 1.0)\n\n    def test_aggregate_dp_mismatched_lengths(self):\n        \"\"\"Test aggregate_dp raises error with mismatched value lengths.\"\"\"\n        metric1 = Metric(aggregation=\"sum\")\n        metric1.extend([1.0, 2.0])\n\n        metric2 = Metric(aggregation=\"sum\")\n        metric2.extend([3.0, 4.0, 5.0])  # Different length\n\n        with self.assertRaises(ValueError):\n            Metric.aggregate_dp([metric1, metric2])\n\n    def test_from_dict(self):\n        \"\"\"Test from_dict creates Metrics from dictionary.\"\"\"\n        data = {\"loss\": 1.0, \"accuracy\": 0.9}\n        metrics = Metric.from_dict(data, aggregation=\"mean\")\n\n        self.assertIn(\"loss\", metrics)\n        self.assertIn(\"accuracy\", metrics)\n        self.assertEqual(metrics[\"loss\"].values, [1.0])\n        self.assertEqual(metrics[\"accuracy\"].values, [0.9])\n        self.assertEqual(metrics[\"loss\"].aggregation, AggregationType.MEAN)\n\n    def test_init_list(self):\n        \"\"\"Test init_list creates new empty Metric with same aggregation.\"\"\"\n        metric = Metric(aggregation=\"max\")\n        metric.extend([1.0, 2.0])\n\n        new_metric = metric.init_list()\n\n        self.assertEqual(new_metric.aggregation, AggregationType.MAX)\n        self.assertEqual(new_metric.values, [])\n\n    def test_reduce_metrics_with_metric(self):\n        \"\"\"Test reduce_metrics correctly handles Metric objects.\"\"\"\n        metric = Metric(aggregation=\"mean\")\n        metric.extend([1.0, 2.0, 3.0])\n\n        metrics = {\n            \"custom_metric\": metric,\n            \"list_metric\": [4.0, 5.0, 6.0],\n        }\n        result = reduce_metrics(metrics)\n\n        self.assertEqual(result[\"custom_metric\"], 2.0)\n        self.assertEqual(result[\"list_metric\"], 5.0)\n\n\nclass TestComputeDataMetrics(unittest.TestCase):\n    \"\"\"Tests for the compute_data_metrics function.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up common test data.\"\"\"\n        # Create a mock DataProto object\n        self.batch = MagicMock()\n        self.batch.batch = {\n            \"token_level_scores\": torch.tensor([[1.0, 2.0], [3.0, 4.0]]),\n            \"token_level_rewards\": torch.tensor([[0.5, 1.0], [1.5, 2.0]]),\n            \"advantages\": torch.tensor([[0.1, 0.2], [0.3, 0.4]]),\n            \"returns\": torch.tensor([[1.1, 1.2], [1.3, 1.4]]),\n            \"responses\": torch.zeros((2, 2)),  # 2 samples, 2 tokens each\n            \"attention_mask\": torch.tensor(\n                [\n                    [1, 1, 1, 1],  # 2 prompt tokens, 2 response tokens\n                    [1, 1, 1, 1],\n                ]\n            ),\n            \"response_mask\": torch.tensor(\n                [\n                    [1, 1],  # 2 response tokens\n                    [1, 1],\n                ]\n            ),\n            \"values\": torch.tensor([[0.9, 1.0], [1.1, 1.2]]),\n        }\n\n    def test_compute_data_metrics_with_critic(self):\n        \"\"\"Test compute_data_metrics with critic enabled.\"\"\"\n        metrics = compute_data_metrics(self.batch, use_critic=True)\n\n        # Check that all expected metrics are present\n        self.assertIn(\"critic/score/mean\", metrics)\n        self.assertIn(\"critic/rewards/mean\", metrics)\n        self.assertIn(\"critic/advantages/mean\", metrics)\n        self.assertIn(\"critic/returns/mean\", metrics)\n        self.assertIn(\"critic/values/mean\", metrics)\n        self.assertIn(\"critic/vf_explained_var\", metrics)\n        self.assertIn(\"response_length/mean\", metrics)\n        self.assertIn(\"prompt_length/mean\", metrics)\n\n        # Check some specific values\n        self.assertAlmostEqual(metrics[\"critic/score/mean\"], 5.0)  # Sum of token_level_scores\n        self.assertAlmostEqual(metrics[\"critic/rewards/mean\"], 2.5)  # Sum of token_level_rewards\n\n    def test_compute_data_metrics_without_critic(self):\n        \"\"\"Test compute_data_metrics with critic disabled.\"\"\"\n        metrics = compute_data_metrics(self.batch, use_critic=False)\n\n        # Check that critic-specific metrics are not present\n        self.assertNotIn(\"critic/values/mean\", metrics)\n        self.assertNotIn(\"critic/vf_explained_var\", metrics)\n\n        # Check that other metrics are still present\n        self.assertIn(\"critic/score/mean\", metrics)\n        self.assertIn(\"critic/rewards/mean\", metrics)\n        self.assertIn(\"response_length/mean\", metrics)\n\n\nclass TestComputeTimingMetrics(unittest.TestCase):\n    \"\"\"Tests for the compute_timing_metrics function.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up common test data.\"\"\"\n        # Create a mock DataProto object\n        self.batch = MagicMock()\n        self.batch.batch = {\n            \"responses\": torch.zeros((2, 3)),  # 2 samples, 3 response tokens each\n            \"attention_mask\": torch.tensor(\n                [\n                    [1, 1, 1, 1, 1, 1],  # 3 prompt tokens, 3 response tokens\n                    [1, 1, 1, 1, 1, 1],\n                ]\n            ),\n        }\n\n        # Mock the _compute_response_info function to return known values\n        self.response_info = {\n            \"prompt_length\": torch.tensor([3.0, 3.0]),\n            \"response_length\": torch.tensor([3.0, 3.0]),\n            \"response_mask\": torch.ones((2, 3)),\n        }\n\n    @patch(\"verl.trainer.ppo.metric_utils._compute_response_info\")\n    def test_compute_timing_metrics(self, mock_compute_response_info):\n        \"\"\"Test compute_timing_metrics with various timing data.\"\"\"\n        mock_compute_response_info.return_value = self.response_info\n\n        timing_raw = {\n            \"gen\": 0.5,  # 500ms\n            \"ref\": 0.3,  # 300ms\n            \"values\": 0.2,  # 200ms\n        }\n\n        metrics = compute_timing_metrics(self.batch, timing_raw)\n\n        # Check raw timing metrics\n        self.assertEqual(metrics[\"timing_s/gen\"], 0.5)\n        self.assertEqual(metrics[\"timing_s/ref\"], 0.3)\n        self.assertEqual(metrics[\"timing_s/values\"], 0.2)\n\n        # Check per-token timing metrics\n        # gen uses only response tokens (6 tokens)\n        self.assertAlmostEqual(metrics[\"timing_per_token_ms/gen\"], 0.5 * 1000 / 6, places=5)\n\n        # ref and values use all tokens (12 tokens)\n        self.assertAlmostEqual(metrics[\"timing_per_token_ms/ref\"], 0.3 * 1000 / 12, places=5)\n        self.assertAlmostEqual(metrics[\"timing_per_token_ms/values\"], 0.2 * 1000 / 12, places=5)\n\n\nclass TestComputeThroughputMetrics(unittest.TestCase):\n    \"\"\"Tests for the compute_throughout_metrics function.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up common test data.\"\"\"\n        # Create a mock DataProto object\n        self.batch = MagicMock()\n        self.batch.meta_info = {\n            \"global_token_num\": [100, 200, 300],  # 600 tokens total\n        }\n\n    def test_compute_throughout_metrics(self):\n        \"\"\"Test compute_throughout_metrics with various timing data.\"\"\"\n        timing_raw = {\n            \"step\": 2.0,  # 2 seconds per step\n        }\n\n        # Test with 1 GPU\n        metrics = compute_throughout_metrics(self.batch, timing_raw, n_gpus=1)\n\n        self.assertEqual(metrics[\"perf/total_num_tokens\"], 600)\n        self.assertEqual(metrics[\"perf/time_per_step\"], 2.0)\n        self.assertEqual(metrics[\"perf/throughput\"], 600 / 2.0)  # 300 tokens/sec\n\n        # Test with 2 GPUs\n        metrics = compute_throughout_metrics(self.batch, timing_raw, n_gpus=2)\n\n        self.assertEqual(metrics[\"perf/total_num_tokens\"], 600)\n        self.assertEqual(metrics[\"perf/time_per_step\"], 2.0)\n        self.assertEqual(metrics[\"perf/throughput\"], 600 / (2.0 * 2))  # 150 tokens/sec/GPU\n\n\nclass TestBootstrapMetric(unittest.TestCase):\n    \"\"\"Tests for the bootstrap_metric function.\"\"\"\n\n    def test_bootstrap_metric_basic(self):\n        \"\"\"Test bootstrap_metric with simple data and functions.\"\"\"\n        data = [1, 2, 3, 4, 5]\n        reduce_fns = [np.mean, np.max]\n\n        # Use a fixed seed for reproducibility\n        result = bootstrap_metric(data, subset_size=3, reduce_fns=reduce_fns, n_bootstrap=100, seed=42)\n\n        # Check that we get two results (one for each reduce_fn)\n        self.assertEqual(len(result), 2)\n\n        # Each result should be a tuple of (mean, std)\n        mean_result, max_result = result\n        self.assertEqual(len(mean_result), 2)\n        self.assertEqual(len(max_result), 2)\n\n        # The mean of means should be close to the true mean (3.0)\n        self.assertAlmostEqual(mean_result[0], 3.0, delta=0.3)\n\n        # The mean of maxes should be close to the expected value for samples of size 3\n        # For samples of size 3 from [1,2,3,4,5], the expected max is around 4.0-4.5\n        self.assertGreater(max_result[0], 3.5)\n        self.assertLess(max_result[0], 5.0)\n\n    def test_bootstrap_metric_empty(self):\n        \"\"\"Test bootstrap_metric with empty data.\"\"\"\n        with self.assertRaises(ValueError):\n            bootstrap_metric([], subset_size=1, reduce_fns=[np.mean])\n\n\nclass TestCalcMajVal(unittest.TestCase):\n    \"\"\"Tests for the calc_maj_val function.\"\"\"\n\n    def test_calc_maj_val_basic(self):\n        \"\"\"Test calc_maj_val with simple data.\"\"\"\n        data = [\n            {\"pred\": \"A\", \"val\": 0.9},\n            {\"pred\": \"B\", \"val\": 0.8},\n            {\"pred\": \"A\", \"val\": 0.7},\n        ]\n\n        result = calc_maj_val(data, vote_key=\"pred\", val_key=\"val\")\n\n        # \"A\" is the majority vote, so we should get the first \"val\" for \"A\"\n        self.assertEqual(result, 0.9)\n\n    def test_calc_maj_val_tie(self):\n        \"\"\"Test calc_maj_val with tied votes.\"\"\"\n        data = [\n            {\"pred\": \"A\", \"val\": 0.9},\n            {\"pred\": \"B\", \"val\": 0.8},\n            {\"pred\": \"B\", \"val\": 0.7},\n            {\"pred\": \"A\", \"val\": 0.6},\n        ]\n\n        # In case of a tie, the first key in sorted order wins\n        # This depends on Python's dict implementation, but for this test\n        # we just verify that one of the valid values is returned\n        result = calc_maj_val(data, vote_key=\"pred\", val_key=\"val\")\n\n        self.assertTrue(result in [0.9, 0.8])\n\n\nclass TestProcessValidationMetrics(unittest.TestCase):\n    \"\"\"Tests for the process_validation_metrics function.\"\"\"\n\n    def test_process_validation_metrics_basic(self):\n        \"\"\"Test process_validation_metrics with simple data.\"\"\"\n        data_sources = [\"source1\", \"source1\", \"source2\"]\n        sample_inputs = [\"prompt1\", \"prompt1\", \"prompt2\"]\n        infos_dict = {\n            \"score\": [0.8, 0.9, 0.7],\n        }\n\n        result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42)\n\n        # Check the structure of the result\n        self.assertIn(\"source1\", result)\n        self.assertIn(\"source2\", result)\n\n        # Check that source1 has metrics for score\n        self.assertIn(\"score\", result[\"source1\"])\n\n        # Check that mean@2 is present for source1/score\n        self.assertIn(\"mean@2\", result[\"source1\"][\"score\"])\n\n        # Check the value of mean@2 for source1/score\n        self.assertAlmostEqual(result[\"source1\"][\"score\"][\"mean@2\"], 0.85)\n\n    def test_process_validation_metrics_with_pred(self):\n        \"\"\"Test process_validation_metrics with prediction data.\"\"\"\n        data_sources = [\"source1\", \"source1\", \"source1\"]\n        sample_inputs = [\"prompt1\", \"prompt1\", \"prompt1\"]\n        infos_dict = {\n            \"score\": [0.8, 0.9, 0.7],\n            \"pred\": [\"A\", \"B\", \"A\"],\n        }\n\n        result = process_validation_metrics(data_sources, sample_inputs, infos_dict, seed=42)\n\n        # Check that majority voting metrics are present\n        self.assertIn(\"maj@2/mean\", result[\"source1\"][\"score\"])\n\n        # For bootstrap with n=2, the majority vote could be either A or B\n        # depending on the random sampling, so we don't check the exact value\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/trainer/ppo/test_rollout_corr.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nQuick Sanity Test for Rollout Correction\n\nThis is a standalone test script that can be run without pytest to quickly verify\nthe rollout correction implementation is working correctly. For comprehensive integration\ntests, see: tests/trainer/ppo/test_rollout_corr_integration.py\n\nUsage:\n    python test_rollout_corr.py\n\nThis tests:\n- Basic rollout correction functionality (IS weights + rejection sampling)\n- Metrics completeness (IS metrics + rejection metrics + off-policy metrics)\n- Edge cases\n\"\"\"\n\nimport pytest\nimport torch\n\nfrom verl.trainer.ppo.rollout_corr_helper import (\n    SUPPORTED_ROLLOUT_RS_OPTIONS,\n    compute_offpolicy_metrics,\n    compute_rollout_correction_and_rejection_mask,\n)\n\n\ndef test_basic_rollout_correction():\n    \"\"\"Test basic rollout correction functionality.\"\"\"\n    print(\"Testing basic rollout correction functionality...\")\n\n    # Create test data\n    batch_size, seq_length = 4, 10\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Create slightly different log probs (simulating BF16 vs FP32 mismatch)\n    old_log_prob = torch.randn(batch_size, seq_length, device=device)\n    rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.1\n    eos_mask = torch.ones(batch_size, seq_length, device=device)\n\n    # Test token-level truncate mode\n    print(\"\\n1. Testing token-level truncate mode...\")\n    weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=eos_mask,\n        rollout_is=\"token\",  # Compute IS weights at token level\n        rollout_is_threshold=2.0,\n        rollout_rs=None,  # No rejection sampling (truncate mode)\n    )\n\n    weights = weights_proto.batch[\"rollout_is_weights\"]\n    print(f\"   Weights shape: {weights.shape}\")\n    print(f\"   Mean weight: {metrics['rollout_corr/rollout_is_mean']:.4f}\")\n    print(f\"   Max weight: {metrics['rollout_corr/rollout_is_max']:.4f}\")\n    print(f\"   Min weight: {metrics['rollout_corr/rollout_is_min']:.4f}\")\n    assert weights.shape == old_log_prob.shape\n    assert weights.max() <= 2.0, \"Weights should be capped at threshold\"\n    print(\"   ✓ Token-level truncate mode passed\")\n\n    # Test sequence-level mode\n    print(\"\\n2. Testing sequence-level mode...\")\n    weights_seq_proto, _, metrics_seq = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=eos_mask,\n        rollout_is=\"sequence\",  # Compute IS weights at sequence level\n        rollout_is_threshold=5.0,\n        rollout_rs=None,  # No rejection sampling (truncate mode)\n    )\n\n    weights_seq = weights_seq_proto.batch[\"rollout_is_weights\"]\n    print(f\"   Mean weight: {metrics_seq['rollout_corr/rollout_is_mean']:.4f}\")\n    print(f\"   Effective sample size: {metrics_seq['rollout_corr/rollout_is_eff_sample_size']:.4f}\")\n    # Check that all tokens in a sequence have the same weight\n    for i in range(batch_size):\n        seq_weights = weights_seq[i, eos_mask[i].bool()]\n        assert torch.allclose(seq_weights, seq_weights[0]), \"All tokens in sequence should have same weight\"\n    print(\"   ✓ Sequence-level mode passed\")\n\n    # Test K1 sequence mean rejection sampling (mask mode)\n    print(\"\\n3. Testing K1 (sequence mean) rejection sampling...\")\n    weights_geo_proto, modified_mask_geo, metrics_geo = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=eos_mask,\n        rollout_is=None,  # No IS weights (pure mask mode)\n        rollout_rs=\"seq_mean_k1\",  # Rejection sampling with sequence-mean log ratio bounds\n        rollout_rs_threshold=\"0.5_1.5\",\n    )\n\n    print(f\"   Masked fraction: {metrics_geo['rollout_corr/rollout_rs_masked_fraction']:.4f}\")\n    print(\"   ✓ K1 sequence mean rejection sampling passed\")\n\n    # Test disabled IS (rollout_is=None, rollout_rs=None)\n    print(\"\\n4. Testing disabled IS...\")\n    weights_disabled, modified_response_mask_disabled, metrics_disabled = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=eos_mask,\n        rollout_is=None,\n        rollout_rs=None,\n    )\n\n    assert weights_disabled is None, \"Should return None when IS is disabled\"\n    assert torch.equal(modified_response_mask_disabled, eos_mask), \"Should return original mask unchanged\"\n    # Note: off-policy metrics are still computed even when IS/RS are disabled\n    assert \"rollout_corr/kl\" in metrics_disabled, \"Should still compute off-policy metrics\"\n    print(\"   ✓ Disabled IS passed\")\n\n    print(\"\\n✓ All tests passed!\")\n\n\n@pytest.mark.parametrize(\n    (\"option\", \"threshold\"),\n    [\n        (\"token_k1\", \"0.5_1.5\"),\n        (\"token_k2\", 2.0),\n        (\"token_k3\", 2.0),\n        (\"seq_sum_k1\", \"0.6_1.4\"),\n        (\"seq_sum_k2\", 2.5),\n        (\"seq_sum_k3\", 2.5),\n        (\"seq_mean_k1\", \"0.5_1.5\"),\n        (\"seq_mean_k2\", 2.0),\n        (\"seq_mean_k3\", 2.0),\n        (\"seq_max_k2\", 2.0),\n        (\"seq_max_k3\", 2.0),\n    ],\n)\ndef test_each_supported_rollout_rs_option(option: str, threshold):\n    \"\"\"Ensure every supported RS option produces metrics without error.\"\"\"\n    assert option in SUPPORTED_ROLLOUT_RS_OPTIONS\n\n    batch_size, seq_length = 3, 7\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    old_log_prob = torch.randn(batch_size, seq_length, device=device)\n    rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.15\n    response_mask = torch.ones(batch_size, seq_length, device=device)\n\n    _, modified_mask, metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n        rollout_is=None,\n        rollout_rs=option,\n        rollout_rs_threshold=threshold,\n    )\n\n    expected_key = f\"rollout_corr/rollout_rs_{option}_mean\"\n    assert expected_key in metrics, f\"Missing metric for {option}\"\n    assert modified_mask.shape == response_mask.shape\n\n\ndef test_rollout_rs_multiple_options():\n    \"\"\"Verify multiple RS options with mixed threshold formats.\"\"\"\n    batch_size, seq_length = 2, 6\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    old_log_prob = torch.randn(batch_size, seq_length, device=device)\n    rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.2\n    response_mask = torch.ones(batch_size, seq_length, device=device)\n\n    rollout_rs = \"token_k1,seq_max_k3\"\n    rollout_rs_threshold = \"0.4_1.8,3.0\"\n\n    _, _, metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n        rollout_is=None,\n        rollout_rs=rollout_rs,\n        rollout_rs_threshold=rollout_rs_threshold,\n    )\n\n    for option in rollout_rs.split(\",\"):\n        key = f\"rollout_corr/rollout_rs_{option}_mean\"\n        assert key in metrics, f\"Metrics missing for chained option {option}\"\n\n\ndef test_metrics_completeness():\n    \"\"\"Test that all expected metrics are returned.\"\"\"\n    print(\"\\nTesting metrics completeness...\")\n\n    batch_size, seq_length = 3, 8\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    old_log_prob = torch.randn(batch_size, seq_length, device=device)\n    rollout_log_prob = old_log_prob + torch.randn(batch_size, seq_length, device=device) * 0.2\n    eos_mask = torch.ones(batch_size, seq_length, device=device)\n\n    _, _, metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=eos_mask,\n        rollout_is=\"token\",\n        rollout_is_threshold=2.5,\n        rollout_rs=None,\n    )\n\n    # Expected IS metrics\n    expected_is_metrics = [\n        \"rollout_corr/rollout_is_mean\",\n        \"rollout_corr/rollout_is_max\",\n        \"rollout_corr/rollout_is_min\",\n        \"rollout_corr/rollout_is_std\",\n        \"rollout_corr/rollout_is_eff_sample_size\",\n        \"rollout_corr/rollout_is_ratio_fraction_high\",\n        \"rollout_corr/rollout_is_ratio_fraction_low\",\n    ]\n\n    # Expected off-policy diagnostic metrics (also included now)\n    expected_offpolicy_metrics = [\n        \"rollout_corr/training_ppl\",\n        \"rollout_corr/training_log_ppl\",\n        \"rollout_corr/kl\",\n        \"rollout_corr/k3_kl\",\n        \"rollout_corr/rollout_ppl\",\n        \"rollout_corr/rollout_log_ppl\",\n        \"rollout_corr/log_ppl_diff\",\n        \"rollout_corr/log_ppl_abs_diff\",\n        \"rollout_corr/log_ppl_diff_max\",\n        \"rollout_corr/log_ppl_diff_min\",\n        \"rollout_corr/ppl_ratio\",\n        \"rollout_corr/chi2_token\",\n        \"rollout_corr/chi2_seq\",\n    ]\n\n    expected_metrics = expected_is_metrics + expected_offpolicy_metrics\n\n    missing_metrics = [m for m in expected_metrics if m not in metrics]\n    if missing_metrics:\n        print(f\"   ✗ Missing metrics: {missing_metrics}\")\n        return False\n\n    print(f\"   ✓ All {len(expected_metrics)} expected metrics present\")\n    print(f\"   Total metrics returned: {len(metrics)}\")\n    return True\n\n\ndef test_offpolicy_metrics():\n    \"\"\"Test off-policy metrics computation.\"\"\"\n    print(\"\\nTesting off-policy metrics computation...\")\n\n    batch_size, seq_length = 4, 12\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Create test data with some mismatch\n    old_log_prob = torch.randn(batch_size, seq_length, device=device) - 2.0  # training policy\n    rollout_log_prob = torch.randn(batch_size, seq_length, device=device) - 1.5  # rollout policy (more confident)\n    response_mask = torch.ones(batch_size, seq_length, device=device)\n\n    # Test with rollout log probs\n    metrics = compute_offpolicy_metrics(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n    )\n\n    expected_metrics = [\n        \"training_ppl\",\n        \"training_log_ppl\",\n        \"kl\",\n        \"k3_kl\",\n        \"rollout_ppl\",\n        \"rollout_log_ppl\",\n        \"log_ppl_diff\",\n        \"log_ppl_abs_diff\",\n        \"log_ppl_diff_max\",\n        \"log_ppl_diff_min\",\n        \"ppl_ratio\",\n        \"chi2_token\",\n        \"chi2_seq\",\n    ]\n\n    for metric in expected_metrics:\n        assert metric in metrics, f\"Missing metric: {metric}\"\n\n    print(f\"   Training PPL: {metrics['training_ppl']:.4f}\")\n    print(f\"   Rollout PPL: {metrics['rollout_ppl']:.4f}\")\n    print(f\"   KL divergence: {metrics['kl']:.6f}\")\n    print(f\"   K3 KL: {metrics['k3_kl']:.6f}\")\n    print(f\"   PPL ratio: {metrics['ppl_ratio']:.4f}\")\n    print(f\"   ✓ All {len(expected_metrics)} off-policy metrics present\")\n\n    # Test without rollout log probs\n    metrics_no_rollout = compute_offpolicy_metrics(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=None,\n        response_mask=response_mask,\n    )\n\n    assert \"training_ppl\" in metrics_no_rollout\n    assert \"rollout_ppl\" not in metrics_no_rollout\n    print(\"   ✓ Off-policy metrics work without rollout log probs\")\n\n\ndef test_mask_mode():\n    \"\"\"Test mask mode applies rejection via response_mask, keeps true IS weights.\"\"\"\n    print(\"\\nTesting mask mode behavior...\")\n\n    batch_size = 2\n    seq_length = 5\n    device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n    # Sequence 0: ratio ≈ 0.37 (below 0.5, should be rejected)\n    # Sequence 1: ratio ≈ 1.65 (in [0.5, 2.0], should be accepted)\n    old_log_prob = torch.tensor([[-2.0] * seq_length, [-2.0] * seq_length], device=device)\n    rollout_log_prob = torch.tensor(\n        [\n            [-1.0] * seq_length,  # exp(-2.0 - (-1.0)) = exp(-1.0) ≈ 0.37\n            [-2.5] * seq_length,  # exp(-2.0 - (-2.5)) = exp(0.5) ≈ 1.65\n        ],\n        device=device,\n    )\n    response_mask = torch.ones(batch_size, seq_length, device=device)\n\n    weights_proto, modified_response_mask, metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n        rollout_is=\"token\",  # Compute IS weights\n        rollout_is_threshold=2.0,\n        rollout_rs=\"token_k1\",  # Also apply rejection sampling (mask mode)\n        rollout_rs_threshold=\"0.5_2.0\",\n    )\n\n    weights = weights_proto.batch[\"rollout_is_weights\"]\n\n    # KEY FIX: Weights should be safety-bounded ratios (NOT zeroed)\n    assert torch.all(weights[0, :] > 0), \"Weights should remain as safety-bounded ratios (not zeroed)\"\n    assert torch.allclose(weights[0, 0], torch.tensor(0.368, device=device), atol=0.01), (\n        \"First seq ratio should be ≈0.37\"\n    )\n    assert torch.allclose(weights[1, 0], torch.tensor(1.649, device=device), atol=0.01), (\n        \"Second seq ratio should be ≈1.65\"\n    )\n\n    # Rejection should be applied via response_mask\n    assert torch.all(modified_response_mask[0, :] == 0), \"First sequence should be rejected via mask\"\n    assert torch.all(modified_response_mask[1, :] == 1), \"Second sequence should be accepted\"\n\n    # Verify rejection sampling metrics exist\n    assert \"rollout_corr/rollout_rs_masked_fraction\" in metrics, \"Should have rollout_rs_masked_fraction metric\"\n    assert abs(metrics[\"rollout_corr/rollout_rs_masked_fraction\"] - 0.5) < 0.01, \"Should reject 50% of tokens\"\n\n    print(f\"   First seq IS weight: {weights[0, 0]:.4f} (expected ≈0.37)\")\n    print(f\"   Second seq IS weight: {weights[1, 0]:.4f} (expected ≈1.65)\")\n    print(f\"   First seq mask: {modified_response_mask[0, 0]:.0f} (expected 0 - rejected)\")\n    print(f\"   Second seq mask: {modified_response_mask[1, 0]:.0f} (expected 1 - accepted)\")\n    print(f\"   Masked fraction: {metrics['rollout_corr/rollout_rs_masked_fraction']:.2f}\")\n    print(\"   ✓ Mask mode correctly separates IS weights from rejection\")\n\n\nif __name__ == \"__main__\":\n    print(\"=\" * 60)\n    print(\"Rollout Correction Test Suite\")\n    print(\"=\" * 60)\n\n    try:\n        test_basic_rollout_correction()\n        test_metrics_completeness()\n        test_offpolicy_metrics()\n        test_mask_mode()\n        print(\"\\n\" + \"=\" * 60)\n        print(\"ALL TESTS PASSED ✓\")\n        print(\"=\" * 60)\n    except Exception as e:\n        print(f\"\\n✗ Test failed with error: {e}\")\n        import traceback\n\n        traceback.print_exc()\n        exit(1)\n"
  },
  {
    "path": "tests/trainer/ppo/test_rollout_corr_integration.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Integration tests for Rollout Correction.\"\"\"\n\nimport pytest\nimport torch\n\nfrom verl.trainer.config.algorithm import RolloutCorrectionConfig\nfrom verl.trainer.ppo.core_algos import compute_policy_loss_vanilla\nfrom verl.trainer.ppo.rollout_corr_helper import (\n    compute_offpolicy_metrics,\n    compute_rollout_correction_and_rejection_mask,\n)\nfrom verl.workers.config.actor import ActorConfig\n\n\nclass TestRolloutISIntegration:\n    \"\"\"Integration tests for Rollout Correction with PPO.\"\"\"\n\n    @pytest.fixture\n    def sample_data(self):\n        \"\"\"Create sample training data.\"\"\"\n        batch_size, seq_length = 4, 16\n        device = \"cuda\" if torch.cuda.is_available() else \"cpu\"\n\n        return {\n            \"old_log_prob\": torch.randn(batch_size, seq_length, device=device),\n            \"log_prob\": torch.randn(batch_size, seq_length, device=device),\n            \"rollout_log_prob\": torch.randn(batch_size, seq_length, device=device),\n            \"advantages\": torch.randn(batch_size, seq_length, device=device),\n            \"response_mask\": torch.ones(batch_size, seq_length, device=device),\n        }\n\n    @pytest.fixture\n    def config_with_rollout_is(self):\n        \"\"\"Create config for policy loss computation.\n\n        Note: rollout_is config has been moved to algorithm config.\n        This config only needs fields used by policy loss (clip_ratio, etc).\n        \"\"\"\n        config = ActorConfig(\n            strategy=\"fsdp\",\n            rollout_n=1,\n            ppo_micro_batch_size=2,\n            clip_ratio=0.2,\n        )\n        return config\n\n    def test_policy_loss_with_rollout_is(self, sample_data, config_with_rollout_is):\n        \"\"\"Test that policy loss computation works with rollout correction weights.\n\n        Note: In production, IS weights are computed centrally in the trainer\n        (before advantage computation) and passed to policy loss.\n        This test simulates that workflow.\n        \"\"\"\n        # First compute IS weights (as trainer would do centrally)\n        rollout_is_weights_proto, _, _ = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=\"token\",\n            rollout_is_threshold=2.0,\n            rollout_rs=None,\n        )\n\n        rollout_is_weights = rollout_is_weights_proto.batch[\"rollout_is_weights\"]\n\n        # Policy loss function receives pre-computed IS weights\n        pg_loss, _ = compute_policy_loss_vanilla(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            log_prob=sample_data[\"log_prob\"],\n            advantages=sample_data[\"advantages\"],\n            response_mask=sample_data[\"response_mask\"],\n            loss_agg_mode=\"token-mean\",\n            config=config_with_rollout_is,\n            rollout_is_weights=rollout_is_weights,\n        )\n\n        # Check loss is valid\n        assert isinstance(pg_loss, torch.Tensor)\n        assert pg_loss.ndim == 0  # Scalar\n        assert not torch.isnan(pg_loss)\n        assert not torch.isinf(pg_loss)\n\n    def test_rollout_is_weights_computation(self, sample_data):\n        \"\"\"Test rollout correction weights and metrics computation.\"\"\"\n        weights_proto, _, metrics = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=\"token\",\n            rollout_is_threshold=2.0,\n            rollout_rs=None,\n        )\n\n        # Check weights\n        from verl.protocol import DataProto\n\n        assert isinstance(weights_proto, DataProto)\n        weights = weights_proto.batch[\"rollout_is_weights\"]\n        assert isinstance(weights, torch.Tensor)\n        assert weights.shape == sample_data[\"old_log_prob\"].shape\n\n        # Check metrics are returned\n        assert isinstance(metrics, dict)\n        assert len(metrics) > 0\n        assert \"rollout_corr/rollout_is_mean\" in metrics\n\n    def test_all_aggregation_levels(self, sample_data):\n        \"\"\"Test all aggregation levels (token, sequence for IS; K1 for RS).\"\"\"\n        # Test IS weight levels\n        is_levels = [\"token\", \"sequence\"]\n        for level in is_levels:\n            _, _, metrics = compute_rollout_correction_and_rejection_mask(\n                old_log_prob=sample_data[\"old_log_prob\"],\n                rollout_log_prob=sample_data[\"rollout_log_prob\"],\n                response_mask=sample_data[\"response_mask\"],\n                rollout_is=level,\n                rollout_is_threshold=2.0,\n                rollout_rs=None,\n            )\n            assert \"rollout_corr/rollout_is_mean\" in metrics\n\n        # Test rejection sampling with K1 sequence mean level\n        _, _, metrics_geo = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=\"0.999_1.001\",\n        )\n        assert \"rollout_corr/rollout_rs_seq_mean_k1_mean\" in metrics_geo\n\n    def test_both_bounding_modes(self, sample_data):\n        \"\"\"Test both truncate and mask modes.\"\"\"\n        # Test truncate mode (IS weights only)\n        _, _, metrics_truncate = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=\"token\",\n            rollout_is_threshold=2.0,\n            rollout_rs=None,\n        )\n        assert \"rollout_corr/rollout_is_mean\" in metrics_truncate\n\n        # Test mask mode (rejection sampling)\n        _, _, metrics_mask = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=\"token\",  # Can also compute IS weights in mask mode\n            rollout_is_threshold=2.0,\n            rollout_rs=\"token_k1\",  # Enable rejection sampling\n            rollout_rs_threshold=1.3,  # Float upper bound (lower inferred automatically)\n        )\n        assert \"rollout_corr/rollout_is_mean\" in metrics_mask\n        assert \"rollout_corr/rollout_rs_token_k1_mean\" in metrics_mask\n\n    def test_offpolicy_metrics(self, sample_data):\n        \"\"\"Test off-policy diagnostic metrics computation.\"\"\"\n        metrics = compute_offpolicy_metrics(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n        )\n\n        # Check key metrics are present\n        assert \"training_ppl\" in metrics\n        assert \"rollout_ppl\" in metrics\n        assert \"kl\" in metrics\n        assert isinstance(metrics[\"kl\"], float)\n\n    def test_metrics_only_mode(self, sample_data, config_with_rollout_is):\n        \"\"\"Test metrics-only mode: compute IS weights/metrics but don't apply to loss.\n\n        This tests the use case where rollout_is_threshold is set (enables computation)\n        but rollout_is=False (disables weight application to policy loss).\n        \"\"\"\n        # Compute IS weights (as trainer would do)\n        rollout_is_weights_proto, _, is_metrics = compute_rollout_correction_and_rejection_mask(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            rollout_log_prob=sample_data[\"rollout_log_prob\"],\n            response_mask=sample_data[\"response_mask\"],\n            rollout_is=\"token\",\n            rollout_is_threshold=2.0,\n            rollout_rs=None,\n        )\n\n        # Metrics should be computed\n        assert len(is_metrics) > 0\n        assert \"rollout_corr/rollout_is_mean\" in is_metrics\n\n        # In metrics-only mode, we compute loss WITHOUT applying weights\n        # (simulating rollout_is=False)\n        pg_loss_no_weights, _ = compute_policy_loss_vanilla(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            log_prob=sample_data[\"log_prob\"],\n            advantages=sample_data[\"advantages\"],\n            response_mask=sample_data[\"response_mask\"],\n            loss_agg_mode=\"token-mean\",\n            config=config_with_rollout_is,\n            rollout_is_weights=None,  # Don't apply weights\n        )\n\n        # Compare to loss WITH weights (rollout_is=True)\n        rollout_is_weights = rollout_is_weights_proto.batch[\"rollout_is_weights\"]\n        pg_loss_with_weights, _ = compute_policy_loss_vanilla(\n            old_log_prob=sample_data[\"old_log_prob\"],\n            log_prob=sample_data[\"log_prob\"],\n            advantages=sample_data[\"advantages\"],\n            response_mask=sample_data[\"response_mask\"],\n            loss_agg_mode=\"token-mean\",\n            config=config_with_rollout_is,\n            rollout_is_weights=rollout_is_weights,\n        )\n\n        # Losses should be different (weights have an effect)\n        assert not torch.allclose(pg_loss_no_weights, pg_loss_with_weights)\n\n\nclass TestRolloutCorrectionConfigNormalization:\n    \"\"\"Unit tests for RolloutCorrectionConfig canonicalization logic.\"\"\"\n\n    def test_alias_normalization_and_threshold_parsing(self):\n        config = RolloutCorrectionConfig(\n            rollout_is=\"token\",\n            rollout_is_threshold=2.5,\n            rollout_rs=\"seq_mean_k1,seq_max_k3\",\n            rollout_rs_threshold=\"0.8_1.2,3.0\",\n        )\n\n        assert config.rollout_is == \"token\"\n        assert config.rollout_is_threshold == pytest.approx(2.5)\n        assert config.rollout_rs == \"seq_mean_k1,seq_max_k3\"\n        assert config.rollout_rs_threshold == \"0.8_1.2,3.0\"\n\n    def test_missing_threshold_raises(self):\n        config = RolloutCorrectionConfig(rollout_rs=\"token_k1\")\n        assert config.rollout_rs == \"token_k1\"\n        assert config.rollout_rs_threshold is None\n\n    def test_float_threshold_conversion_in_factory(self):\n        config = RolloutCorrectionConfig.decoupled_geo_rs_seq_tis(rs_threshold=1.001)\n        assert config.rollout_rs == \"seq_mean_k1\"\n        assert config.rollout_rs_threshold == 1.001\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\", \"-s\"])\n"
  },
  {
    "path": "tests/utils/_test_module.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n# Test module for import_utils.load_extern_object testing\nclass TestClass:\n    \"\"\"A test class to be imported by load_extern_object\"\"\"\n\n    def __init__(self, value=None):\n        self.value = value or \"default\"\n\n    def get_value(self):\n        return self.value\n\n\nTEST_CONSTANT = \"test_constant_value\"\n\n\ndef test_function():\n    return \"test_function_result\"\n"
  },
  {
    "path": "tests/utils/ckpt/test_checkpoint_cleanup_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport shutil\nimport tempfile\n\nimport pytest\n\n\nclass TestCheckpointCleanupLogic:\n    \"\"\"Tests for checkpoint cleanup methods in BaseCheckpointManager.\"\"\"\n\n    @pytest.fixture(autouse=True)\n    def setup(self):\n        \"\"\"Set up test fixtures.\"\"\"\n        self.test_dir = tempfile.mkdtemp()\n        yield\n        shutil.rmtree(self.test_dir, ignore_errors=True)\n\n    @pytest.fixture\n    def manager(self, monkeypatch):\n        \"\"\"Create a minimal BaseCheckpointManager for testing.\"\"\"\n        import torch.distributed\n\n        monkeypatch.setattr(torch.distributed, \"get_rank\", lambda: 0)\n        monkeypatch.setattr(torch.distributed, \"get_world_size\", lambda: 1)\n\n        from verl.utils.checkpoint.checkpoint_manager import BaseCheckpointManager\n\n        class MockModel:\n            pass\n\n        class MockOptimizer:\n            pass\n\n        return BaseCheckpointManager(\n            model=MockModel(),\n            optimizer=MockOptimizer(),\n            lr_scheduler=None,\n            processing_class=None,\n            checkpoint_config=None,\n        )\n\n    def _create_checkpoint_dir(self, step: int) -> str:\n        \"\"\"Create a mock checkpoint directory.\"\"\"\n        path = os.path.join(self.test_dir, f\"global_step_{step}\")\n        os.makedirs(path, exist_ok=True)\n        with open(os.path.join(path, \"checkpoint.txt\"), \"w\") as f:\n            f.write(f\"step={step}\")\n        return path\n\n    def test_max_ckpt_1_preserves_existing_before_save(self, manager):\n        \"\"\"\n        Regression test: max_ckpt_to_keep=1 must NOT delete existing checkpoint before save.\n        \"\"\"\n        ckpt_100 = self._create_checkpoint_dir(100)\n        manager.previous_saved_paths = [ckpt_100]\n\n        manager.ensure_checkpoint_capacity(max_ckpt_to_keep=1)\n\n        assert os.path.exists(ckpt_100), \"Bug: checkpoint deleted before save!\"\n        assert manager.previous_saved_paths == [ckpt_100]\n\n    def test_max_ckpt_1_deletes_old_after_save(self, manager):\n        \"\"\"After save succeeds, old checkpoint should be deleted.\"\"\"\n        ckpt_100 = self._create_checkpoint_dir(100)\n        manager.previous_saved_paths = [ckpt_100]\n\n        ckpt_200 = self._create_checkpoint_dir(200)\n        manager.register_checkpoint(ckpt_200, max_ckpt_to_keep=1)\n\n        assert not os.path.exists(ckpt_100)\n        assert os.path.exists(ckpt_200)\n        assert manager.previous_saved_paths == [ckpt_200]\n\n    def test_max_ckpt_2_keeps_one_before_save(self, manager):\n        \"\"\"With max_ckpt_to_keep=2, pre-save cleanup keeps 1 checkpoint.\"\"\"\n        ckpt_100 = self._create_checkpoint_dir(100)\n        ckpt_200 = self._create_checkpoint_dir(200)\n        manager.previous_saved_paths = [ckpt_100, ckpt_200]\n\n        manager.ensure_checkpoint_capacity(max_ckpt_to_keep=2)\n\n        assert not os.path.exists(ckpt_100)\n        assert os.path.exists(ckpt_200)\n        assert len(manager.previous_saved_paths) == 1\n\n    def test_max_ckpt_0_keeps_all(self, manager):\n        \"\"\"max_ckpt_to_keep=0 means unlimited - no deletions.\"\"\"\n        ckpt_100 = self._create_checkpoint_dir(100)\n        ckpt_200 = self._create_checkpoint_dir(200)\n        manager.previous_saved_paths = [ckpt_100, ckpt_200]\n\n        manager.ensure_checkpoint_capacity(max_ckpt_to_keep=0)\n        ckpt_300 = self._create_checkpoint_dir(300)\n        manager.register_checkpoint(ckpt_300, max_ckpt_to_keep=0)\n\n        assert os.path.exists(ckpt_100)\n        assert os.path.exists(ckpt_200)\n        assert os.path.exists(ckpt_300)\n        assert len(manager.previous_saved_paths) == 3\n\n    def test_full_save_cycle_max_ckpt_1(self, manager):\n        \"\"\"Simulate multiple save cycles with max_ckpt_to_keep=1.\"\"\"\n        # First save\n        manager.ensure_checkpoint_capacity(1)\n        ckpt_100 = self._create_checkpoint_dir(100)\n        manager.register_checkpoint(ckpt_100, 1)\n        assert manager.previous_saved_paths == [ckpt_100]\n\n        # Second save - existing checkpoint must survive pre-save\n        manager.ensure_checkpoint_capacity(1)\n        assert os.path.exists(ckpt_100), \"Bug: checkpoint deleted before save!\"\n\n        ckpt_200 = self._create_checkpoint_dir(200)\n        manager.register_checkpoint(ckpt_200, 1)\n        assert not os.path.exists(ckpt_100)\n        assert manager.previous_saved_paths == [ckpt_200]\n\n        # Third save\n        manager.ensure_checkpoint_capacity(1)\n        assert os.path.exists(ckpt_200), \"Bug: checkpoint deleted before save!\"\n\n        ckpt_300 = self._create_checkpoint_dir(300)\n        manager.register_checkpoint(ckpt_300, 1)\n        assert not os.path.exists(ckpt_200)\n        assert manager.previous_saved_paths == [ckpt_300]\n"
  },
  {
    "path": "tests/utils/ckpt/test_esi_save_ckpt_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport time\nfrom datetime import datetime, timedelta\nfrom unittest import TestCase\n\nfrom verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi\n\n\nclass TestShouldSaveCkptEsi(TestCase):\n    def test_no_expiration_timestamp(self):\n        \"\"\"Test case when no expiration timestamp is set\"\"\"\n        os.environ.pop(\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\", None)\n        os.environ.pop(\"SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\", None)\n        self.assertFalse(should_save_ckpt_esi(100))\n\n    def test_mlp_expiration_valid(self):\n        \"\"\"Test valid MLP expiration timestamp requiring save\"\"\"\n        current_time = time.time()\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(current_time + 90)\n        self.assertTrue(should_save_ckpt_esi(30))  # max_steps_duration=30 seconds\n\n    def test_mlp_expiration_passed(self):\n        \"\"\"Test expired MLP timestamp\"\"\"\n        current_time = time.time()\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(current_time - 10)\n        self.assertFalse(should_save_ckpt_esi(30))\n\n    def test_mlp_invalid_timestamp(self):\n        \"\"\"Test invalid MLP timestamp format\"\"\"\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = \"invalid\"\n        self.assertFalse(should_save_ckpt_esi(30))\n\n    def test_mlp_expiration_not_reached(self):\n        \"\"\"Test MLP expiration timestamp with insufficient remaining time\"\"\"\n        current_time = time.time()\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(current_time + 200)\n        self.assertFalse(should_save_ckpt_esi(30))  # max_steps_duration=30\n\n    def test_aws_expiration_not_reached(self):\n        \"\"\"Test AWS expiration timestamp with sufficient remaining time\"\"\"\n        now = datetime.now()\n        expiration = now + timedelta(minutes=100)  # Exceeds 90-minute threshold\n        os.environ[\"SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(int(expiration.timestamp()))\n        self.assertFalse(should_save_ckpt_esi(30 * 60))\n\n    def test_redundant_time(self):\n        \"\"\"Test redundant_time parameter effect\"\"\"\n        current_time = time.time()\n        # Total required: 60+30+30=120 seconds\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(current_time + 120)\n        self.assertTrue(should_save_ckpt_esi(30, redundant_time=30))\n\n    def test_zero_max_steps_duration(self):\n        \"\"\"Test zero max_steps_duration\"\"\"\n        current_time = time.time()\n        os.environ[\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\"] = str(current_time + 60)\n        self.assertFalse(should_save_ckpt_esi(0))\n"
  },
  {
    "path": "tests/utils/dataset/test_create_rl_sampler_on_cpu.py",
    "content": "# Copyright 2025 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\ntest create_rl_sampler\n\"\"\"\n\nfrom collections.abc import Sized\n\nimport pytest\nimport torch\nfrom omegaconf import DictConfig, OmegaConf\nfrom torch.utils.data import Dataset, RandomSampler\n\nfrom verl.experimental.dataset.sampler import AbstractCurriculumSampler\nfrom verl.trainer.main_ppo import create_rl_sampler\n\n\nclass RandomCurriculumSampler(AbstractCurriculumSampler):\n    def __init__(\n        self,\n        data_source: Sized,\n        data_config: DictConfig,\n    ):\n        train_dataloader_generator = torch.Generator()\n        train_dataloader_generator.manual_seed(1)\n        sampler = RandomSampler(data_source=data_source)\n        self.sampler = sampler\n\n    def __iter__(self):\n        return self.sampler.__iter__()\n\n    def __len__(self) -> int:\n        return len(self.sampler)\n\n    def update(self, batch) -> None:\n        return\n\n\nclass MockIncorrectSampler:\n    \"\"\"A fake sampler class that does not adhere to the AbstractCurriculumSampler interface.\"\"\"\n\n    def __init__(self, data_source, data_config):\n        pass\n\n\nclass MockChatDataset(Dataset):\n    def __init__(self):\n        self.data = [\n            {\"prompt\": \"What's your name?\", \"response\": \"My name is Assistant.\"},\n            {\"prompt\": \"How are you?\", \"response\": \"I'm doing well, thank you.\"},\n            {\"prompt\": \"What is the capital of France?\", \"response\": \"Paris.\"},\n            {\n                \"prompt\": \"Tell me a joke.\",\n                \"response\": \"Why did the chicken cross the road? To get to the other side!\",\n            },\n            {\"prompt\": \"What is 2+2?\", \"response\": \"4\"},\n        ]\n\n    def __getitem__(self, index):\n        return self.data[index]\n\n    def __len__(self):\n        return len(self.data)\n\n\ndef test_create_custom_curriculum_samper():\n    data_config = OmegaConf.create(\n        {\n            \"dataloader_num_workers\": 0,\n            \"sampler\": {\n                \"class_path\": \"pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu\",\n                \"class_name\": \"RandomCurriculumSampler\",\n            },\n        }\n    )\n\n    dataset = MockChatDataset()\n\n    # doesn't raise\n    create_rl_sampler(data_config, dataset)\n\n\ndef test_create_custom_curriculum_samper_wrong_class():\n    data_config = OmegaConf.create(\n        {\n            \"sampler\": {\n                \"class_path\": \"pkg://tests.utils.dataset.test_create_rl_sampler_on_cpu\",\n                \"class_name\": \"MockIncorrectSampler\",\n            }\n        }\n    )\n\n    dataset = MockChatDataset()\n\n    # MockIncorrectSampler is not an instance of AbstractCurriculumSampler, so raises\n    with pytest.raises(AssertionError):\n        create_rl_sampler(data_config, dataset)\n"
  },
  {
    "path": "tests/utils/dataset/test_multiturn_sft_dataset_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTest the MultiTurnSFTDataset implementation\n\"\"\"\n\nimport os\nfrom io import BytesIO\nfrom pathlib import Path\n\nimport pandas as pd\nimport pytest\nimport torch\nfrom PIL import Image\nfrom tensordict import TensorDict\nfrom torch.utils.data import DistributedSampler\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom transformers.utils import get_json_schema\n\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode, SFTTensorCollator\nfrom verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset\nfrom verl.utils.model import extract_multi_modal_inputs\n\ncustom_model_prefix = Path(\"~/models\").expanduser().resolve()\n\n\n@pytest.mark.parametrize(\n    \"model_path, ignore_input_ids_mismatch\",\n    [\n        (f\"{custom_model_prefix}/Qwen/Qwen2.5-0.5B\", False),\n        (f\"{custom_model_prefix}/Qwen/Qwen3-0.6B\", True),\n        (f\"{custom_model_prefix}/Qwen/Qwen3.5-0.8B\", False),\n    ],\n)\ndef test_multiturn_sft_dataset(model_path: str, ignore_input_ids_mismatch: bool):\n    print(f\"Starting test... model_path={model_path}, ignore_input_ids_mismatch={ignore_input_ids_mismatch}\")\n    # Create a temporary parquet file with test data\n    test_data = {\n        \"messages\": [\n            [\n                {\"role\": \"user\", \"content\": \"What is 2+2?\"},\n                {\"role\": \"assistant\", \"content\": \"2+2 equals 4.\"},\n                {\"role\": \"tool\", \"content\": \"And what is 4+4?\"},\n                {\"role\": \"assistant\", \"content\": \"4+4 equals 8.\"},\n            ],\n            [\n                # {\"role\": \"system\", \"content\": \"You are a powerful assistant.\"},\n                {\"role\": \"user\", \"content\": \"Tell me a joke.\"},\n                {\"role\": \"assistant\", \"content\": \"Why did the chicken cross the road?\"},\n                {\"role\": \"tool\", \"content\": \"Why?\"},\n                {\"role\": \"assistant\", \"content\": \"To get to the other side!\"},\n            ],\n        ]\n    }\n\n    # Create test directory if it doesn't exist\n    os.makedirs(\"test_data\", exist_ok=True)\n    test_file = \"test_data/test.parquet\"\n\n    # Save test data to parquet\n    df = pd.DataFrame(test_data)\n    df.to_parquet(test_file)\n\n    # Initialize tokenizer and dataset\n    tokenizer = hf_tokenizer(model_path)\n    # processor = hf_processor(model_path)\n    processor = None\n    config = {\n        \"max_length\": 512,\n        \"truncation\": \"error\",\n        \"multiturn\": {\"messages_key\": \"messages\"},\n        \"ignore_input_ids_mismatch\": ignore_input_ids_mismatch,\n    }\n    dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=config)\n\n    # Test 1: Dataset Length\n    assert len(dataset) == 2, f\"Expected dataset length 2, got {len(dataset)}\"\n\n    # Get items for testing\n    item0 = dataset[0]  # Math conversation\n    item1 = dataset[1]  # Joke conversation\n\n    # Test 2: Required Keys and Types\n    required_keys = [\"input_ids\", \"attention_mask\", \"position_ids\", \"loss_mask\"]\n    for key in required_keys:\n        assert key in item0, f\"Missing key {key} in dataset item\"\n        assert isinstance(item0[key], torch.Tensor), f\"Expected torch.Tensor for {key}\"\n        assert item0[key].dtype == torch.long, f\"Expected torch.long for {key}, got {item0[key].dtype}\"\n\n    # Test 3: Shape Consistency\n    assert item0[\"loss_mask\"].shape == item0[\"input_ids\"].shape, \"Loss mask shape doesn't match input_ids shape\"\n    assert item0[\"attention_mask\"].shape == item0[\"input_ids\"].shape, (\n        \"Attention mask shape doesn't match input_ids shape\"\n    )\n    assert item0[\"position_ids\"].shape == item0[\"input_ids\"].shape, \"Position IDs shape doesn't match input_ids shape\"\n\n    # Test 4: Loss Mask Pattern - Math Conversation\n    loss_mask0 = item0[\"loss_mask\"]\n    input_ids0 = item0[\"input_ids\"]\n\n    # Find assistant response positions\n    assistant_positions0 = torch.where(loss_mask0 == 1)[0]\n    assert len(assistant_positions0) > 0, \"No assistant positions found in loss mask\"\n\n    # Decode and verify assistant responses\n    assistant_text0 = tokenizer.decode(input_ids0[loss_mask0 == 1])\n    print(f\"Math conversation assistant text: {assistant_text0}\")\n    assert \"2+2 equals 4\" in assistant_text0, \"First assistant response not found\"\n    assert \"4+4 equals 8\" in assistant_text0, \"Second assistant response not found\"\n\n    # Test 5: Loss Mask Pattern - Joke Conversation\n    loss_mask1 = item1[\"loss_mask\"]\n    input_ids1 = item1[\"input_ids\"]\n\n    # Find assistant response positions\n    assistant_positions1 = torch.where(loss_mask1 == 1)[0]\n    assert len(assistant_positions1) > 0, \"No assistant positions found in loss mask\"\n\n    # Decode and verify assistant responses\n    assistant_text1 = tokenizer.decode(input_ids1[loss_mask1 == 1])\n    print(f\"Joke conversation assistant text: {assistant_text1}\")\n    assert \"chicken cross the road\" in assistant_text1, \"First assistant response not found\"\n    assert \"other side\" in assistant_text1, \"Second assistant response not found\"\n\n    # Test 6: Attention Mask Pattern\n    attention_mask0 = item0[\"attention_mask\"]\n    sequence_length = torch.sum(attention_mask0)\n    assert sequence_length > 0, \"No tokens marked as attended in attention mask\"\n    assert torch.all(attention_mask0[:sequence_length] == 1), \"Incorrect attention mask pattern\"\n    if sequence_length < len(attention_mask0):\n        assert torch.all(attention_mask0[sequence_length:] == 0), \"Padding not properly masked\"\n\n    # Test 7: Position IDs Pattern\n    position_ids0 = item0[\"position_ids\"]\n    assert torch.equal(position_ids0[:sequence_length], torch.arange(sequence_length)), (\n        \"Position IDs not sequential for non-padded tokens\"\n    )\n    if sequence_length < len(position_ids0):\n        assert torch.all(position_ids0[sequence_length:] == 0), \"Padding position IDs not zero\"\n\n    # Test 8: Verify loss mask for assistant responses\n    # Get the full conversation text\n    full_text = tokenizer.decode(input_ids0)\n    print(f\"\\nFull conversation text:\\n{full_text}\")\n\n    # Get the assistant responses\n    assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 1])\n    print(f\"\\nAssistant responses (from loss mask):\\n{assistant_text}\")\n\n    # Verify that loss mask is set for all assistant responses\n    for msg in test_data[\"messages\"][0]:  # First conversation\n        if msg[\"role\"] == \"assistant\":\n            # The content should appear in the masked text\n            assert msg[\"content\"] in assistant_text, f\"Assistant message '{msg['content']}' not found in masked text\"\n\n            # The content should NOT appear in the non-masked text\n            non_assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 0])\n            assert msg[\"content\"] not in non_assistant_text, (\n                f\"Assistant message '{msg['content']}' found in non-assistant text\"\n            )\n\n    # Test 9: Verify non-assistant parts have loss_mask=0\n    # Get non-assistant text\n    non_assistant_text = tokenizer.decode(input_ids0[loss_mask0 == 0])\n    print(f\"\\nNon-assistant text (from loss mask):\\n{non_assistant_text}\")\n\n    # Verify that system and user messages are in the non-assistant text\n    for msg in test_data[\"messages\"][0]:  # First conversation\n        if msg[\"role\"] in [\"system\", \"user\"]:\n            assert msg[\"content\"] in non_assistant_text, (\n                f\"{msg['role'].title()} message '{msg['content']}' not found in non-assistant text\"\n            )\n\n            # And verify they're NOT in the assistant text\n            assert msg[\"content\"] not in assistant_text, (\n                f\"{msg['role'].title()} message '{msg['content']}' found in assistant text\"\n            )\n\n    # Test 10: Verify padding behavior\n    padding_config = {\n        \"max_length\": 1024,\n        \"truncation\": \"error\",\n        \"multiturn\": {\"messages_key\": \"messages\"},\n        \"ignore_input_ids_mismatch\": ignore_input_ids_mismatch,\n    }\n    small_dataset = MultiTurnSFTDataset(\n        parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=padding_config\n    )\n    padded_item = small_dataset[0]\n\n    # Get actual sequence length (before padding)\n    actual_length = torch.sum(padded_item[\"attention_mask\"])\n\n    # Verify padding tokens\n    assert torch.all(padded_item[\"input_ids\"][actual_length:] == tokenizer.pad_token_id), (\n        \"Padding tokens not set correctly\"\n    )\n    assert torch.all(padded_item[\"attention_mask\"][actual_length:] == 0), \"Attention mask not set correctly for padding\"\n    assert torch.all(padded_item[\"loss_mask\"][actual_length:] == 0), \"Loss mask not set correctly for padding\"\n\n    # test no-padding\n    config = {\n        \"max_length\": 512,\n        \"truncation\": \"error\",\n        \"multiturn\": {\"messages_key\": \"messages\"},\n        \"pad_mode\": \"no_padding\",\n        \"ignore_input_ids_mismatch\": ignore_input_ids_mismatch,\n    }\n    dataset = MultiTurnSFTDataset(parquet_files=test_file, tokenizer=tokenizer, processor=processor, config=config)\n\n    item0 = dataset[0]\n\n    # Verify that the output contains expected keys for no-padding mode\n    required_keys = [\"input_ids\", \"position_ids\", \"loss_mask\"]\n    for key in required_keys:\n        assert key in item0, f\"Missing key {key} in no-padding mode dataset item\"\n        assert isinstance(item0[key], torch.Tensor), f\"Expected torch.Tensor for {key} in no-padding mode\"\n\n    # make sure assistant_text matches with expected\n    assistant_text = tokenizer.decode(item0[\"input_ids\"][item0[\"loss_mask\"] == 1])\n    assert assistant_text == \"2+2 equals 4.<|im_end|>\\n4+4 equals 8.<|im_end|>\\n\"\n\n    print(\"All tests passed!\")\n    print(\"Starting test...\")\n\n\ndef generate_image(description: str, size: str = \"256x256\"):\n    \"\"\"Generate a simple image based on description.\n\n    Args:\n        description: The description of the image to generate.\n        size: The size of the image. Defaults to \"256x256\". (choices: [\"256x256\", \"512x512\"])\n\n    Returns:\n        A generated image\n    \"\"\"\n    ...\n\n\n@pytest.fixture\ndef vlm_data_file():\n    test_data = [\n        # sample 0: single turn with image input\n        {\n            \"messages\": [\n                {\n                    \"role\": \"user\",\n                    \"content\": \"<image>Describe this image.\",\n                },\n                {\n                    \"role\": \"assistant\",\n                    \"content\": \"The image is a red square.\",\n                },\n            ],\n            \"images\": [Image.new(\"RGB\", (300, 300), color=\"red\")],\n            \"tools\": [],\n        },\n        # sample 1: single turn with multiple images input\n        {\n            \"messages\": [\n                {\n                    \"role\": \"user\",\n                    \"content\": \"<image><image>Compare these images.\",\n                },\n                {\n                    \"role\": \"assistant\",\n                    \"content\": \"The first image is a red square and the second image is a green square.\",\n                },\n            ],\n            \"images\": [Image.new(\"RGB\", (100, 100), color=\"red\"), Image.new(\"RGB\", (100, 300), color=\"green\")],\n            \"tools\": [],\n        },\n        # sample 2: multi turn with image input and tool generated image\n        {\n            \"messages\": [\n                {\n                    \"role\": \"user\",\n                    \"content\": \"<image>Describe this image.\",\n                },\n                {\n                    \"role\": \"assistant\",\n                    \"content\": \"Let's generate a zoom-in image.\",\n                    \"tool_calls\": [\n                        {\n                            \"function\": {\"arguments\": {\"bbox_2d\": \"[0, 1, 2, 4]\"}, \"name\": \"image_zoom_in_tool\"},\n                            \"type\": \"function\",\n                        }\n                    ],\n                },\n                {\n                    \"role\": \"tool\",\n                    \"content\": \"<image>Generated image.\",\n                },\n                {\"role\": \"assistant\", \"content\": \"The zoom-in image is a red square.\"},\n            ],\n            \"images\": [Image.new(\"RGB\", (300, 500), color=\"red\"), Image.new(\"RGB\", (100, 100), color=\"red\")],\n            \"tools\": [get_json_schema(generate_image)],\n        },\n        # sample 3: single turn without image input\n        {\n            \"messages\": [\n                {\"role\": \"user\", \"content\": \"How is the weather today?\"},\n                {\"role\": \"assistant\", \"content\": \"The weather is sunny.\"},\n            ],\n            \"images\": [],\n            \"tools\": [],\n        },\n    ]\n\n    # Create test directory if it doesn't exist\n    os.makedirs(\"test_data\", exist_ok=True)\n    test_file = \"test_data/test_vlm.parquet\"\n\n    # Save test data to parquet\n    df = pd.DataFrame(test_data)\n\n    def serialize_image(img):\n        if isinstance(img, Image.Image):\n            img_byte_arr = BytesIO()\n            img.save(img_byte_arr, format=\"PNG\")\n            return {\"bytes\": img_byte_arr.getvalue()}\n        return img\n\n    df[\"images\"] = df[\"images\"].apply(lambda x: [serialize_image(img) for img in x])\n\n    df.to_parquet(test_file)\n    return test_file\n\n\n@pytest.mark.parametrize(\n    \"model_path\",\n    [\n        f\"{custom_model_prefix}/Qwen/Qwen3-VL-2B-Instruct\",\n        f\"{custom_model_prefix}/Qwen/Qwen3.5-0.8B\",\n    ],\n)\ndef test_multiturn_sft_vlm_dataset_on_cpu(model_path, vlm_data_file):\n    df = pd.read_parquet(vlm_data_file)\n    tokenizer = hf_tokenizer(model_path)\n    processor = hf_processor(model_path)\n    config = {\"max_length\": 1024, \"pad_mode\": \"no_padding\", \"truncation\": \"error\", \"messages_key\": \"messages\"}\n    dataset = MultiTurnSFTDataset(parquet_files=vlm_data_file, tokenizer=tokenizer, processor=processor, config=config)\n    assert dataset.pad_mode == DatasetPadMode.NO_PADDING\n\n    for i in range(len(dataset)):\n        item = dataset[i]\n        input_ids = item[\"input_ids\"]\n        loss_mask = item[\"loss_mask\"]\n        position_ids = item[\"position_ids\"]\n        pixel_values = item.get(\"multi_modal_inputs\", {}).get(\"pixel_values\")\n        image_grid_thw = item.get(\"multi_modal_inputs\", {}).get(\"image_grid_thw\")\n\n        assert input_ids.shape == loss_mask.shape, \"Shapes of input_ids and loss_mask must be equal\"\n        assert position_ids.dim() == 2, \"position_ids must be 2-dimensional\"\n        assert position_ids.shape[0] == 4, f\"position_ids[0] should be 4: {position_ids[0]}\"\n        assert position_ids.shape[1] == input_ids.shape[0]\n\n        # 1. verify input_ids without assistant text\n        text = tokenizer.decode(input_ids[loss_mask == 0], skip_special_tokens=True)\n        print(f\"Text without assistant: {repr(text)}\")\n        for message in df[\"messages\"][i]:\n            if message[\"role\"] != \"assistant\":\n                content = message[\"content\"].replace(\"<image>\", \"\")\n                assert content in text, f\"user/tool text should be in the input_ids: {text}\"\n\n        # 2. verify input_ids with assistant text\n        text = tokenizer.decode(input_ids[loss_mask == 1], skip_special_tokens=True)\n        print(f\"Text with assistant: {repr(text)}\")\n        for message in df[\"messages\"][i]:\n            if message[\"role\"] == \"assistant\":\n                assert message[\"content\"] in text, f\"Assistant text should be in the input_ids: {text}\"\n                assert \"assistant\" not in text, f\"Assistant token should not be in the input_ids: {text}\"\n\n        # 3. verify image token match with image_grid_thw\n        if len(df[\"images\"][i]) > 0:\n            patch_size = processor.image_processor.patch_size\n            temporal_patch_size = processor.image_processor.temporal_patch_size\n            merge_size = processor.image_processor.merge_size\n            num_patches = image_grid_thw.prod(dim=1).sum()\n            assert image_grid_thw.shape == (len(df[\"images\"][i]), 3), (\n                f\"image_grid_thw: {image_grid_thw.shape} should have shape ({len(df['images'][i])}, 3)\"\n            )\n            assert pixel_values.shape == (num_patches, 3 * temporal_patch_size * patch_size * patch_size), (\n                f\"pixel_values: {pixel_values.shape} should have shape ({num_patches}, {3 * patch_size * patch_size})\"\n            )\n            assert (input_ids == processor.image_token_id).sum() == num_patches // (merge_size**2)\n        else:\n            assert pixel_values is None, \"pixel_values should be None when no image is provided\"\n            assert image_grid_thw is None, \"image_grid_thw should be None when no image is provided\"\n\n\n@pytest.mark.parametrize(\n    \"model_path\",\n    [\n        f\"{custom_model_prefix}/Qwen/Qwen3-VL-2B-Instruct\",\n        f\"{custom_model_prefix}/Qwen/Qwen3.5-0.8B\",\n    ],\n)\ndef test_multiturn_sft_vlm_dataloader_on_cpu(model_path, vlm_data_file):\n    df = pd.read_parquet(vlm_data_file)\n    tokenizer = hf_tokenizer(model_path)\n    processor = hf_processor(model_path)\n    config = {\"max_length\": 1024, \"pad_mode\": \"no_padding\", \"truncation\": \"error\", \"messages_key\": \"messages\"}\n    dataset = MultiTurnSFTDataset(parquet_files=vlm_data_file, tokenizer=tokenizer, processor=processor, config=config)\n    assert dataset.pad_mode == DatasetPadMode.NO_PADDING\n\n    collate_fn = SFTTensorCollator(DatasetPadMode.NO_PADDING)\n    sampler = DistributedSampler(dataset, shuffle=False, num_replicas=1, rank=0, drop_last=True)\n    batch_size = 2\n    dataloader = StatefulDataLoader(\n        dataset=dataset,\n        batch_size=batch_size,\n        sampler=sampler,\n        collate_fn=collate_fn,\n        num_workers=0,\n        pin_memory=False,\n        drop_last=True,\n    )\n\n    for i, batch in enumerate(dataloader):\n        # 1. verify input_ids, loss_mask\n        input_ids = batch[\"input_ids\"]\n        loss_mask = batch[\"loss_mask\"]\n        assert input_ids.is_nested, \"input_ids should be a nested tensor\"\n        assert loss_mask.is_nested, \"loss_mask should be a nested tensor\"\n        assert input_ids.shape[0] == loss_mask.shape[0] == batch_size, \"Shapes of input_ids, loss_mask must be equal\"\n\n        # 2. verify position_ids: (bs, 4, seq_len)\n        position_ids = batch[\"position_ids\"]\n        assert position_ids.is_nested, \"position_ids should be a nested tensor\"\n        assert position_ids.dim() == 3, \"position_ids must be 3-dimensional\"\n        assert position_ids.shape[0] == batch_size\n        assert position_ids.shape[1] == 4\n        values = position_ids.values()\n        assert values.shape == (4, len(input_ids.values()))\n\n        # 3. verify multi-modal data\n        td = TensorDict(**batch, batch_size=batch_size)\n        multi_modal_inputs = extract_multi_modal_inputs(td[\"multi_modal_inputs\"])\n        pixel_values = multi_modal_inputs[\"pixel_values\"]\n        image_grid_thw = multi_modal_inputs[\"image_grid_thw\"]\n\n        num_images = sum([len(images) for images in df[\"images\"][i * batch_size : (i + 1) * batch_size]])\n        assert image_grid_thw.shape == (num_images, 3), (\n            f\"image_grid_thw: {image_grid_thw.shape} should have shape ({num_images}, 3)\"\n        )\n        patch_size = processor.image_processor.patch_size\n        temporal_patch_size = processor.image_processor.temporal_patch_size\n        num_patches = image_grid_thw.prod(dim=1).sum()\n        assert pixel_values.shape[0] == num_patches, (\n            f\"pixel_values: {pixel_values.shape} should have shape \"\n            f\"({num_patches}, 3 * {temporal_patch_size} * {patch_size} * {patch_size})\"\n        )\n"
  },
  {
    "path": "tests/utils/dataset/test_rl_collate_fn_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport torch\n\n\ndef test_rl_collate_fn():\n    from verl.utils.dataset.rl_dataset import collate_fn\n\n    max_prompt_length = 5\n\n    test_data = [\n        {\n            # test tensor\n            \"input_ids\": torch.randint(0, 10, (max_prompt_length,)),\n            # test fixed length (1) list within a batch\n            \"messages\": [{\"role\": \"user\", \"content\": \"Hi.\"}],\n            # test variable length list within a batch\n            \"raw_prompt_ids\": [1, 2, 3, 4],\n            # test string\n            \"ability\": \"math\",\n            # test dict\n            \"reward_model\": {\"ground_truth\": 5, \"style\": \"rule\"},\n            # test empty dict\n            \"tools_kwargs\": {},\n        },\n        {\n            \"input_ids\": torch.randint(0, 10, (max_prompt_length,)),\n            \"messages\": [{\"role\": \"user\", \"content\": \"Hello.\"}],\n            \"raw_prompt_ids\": [1, 2, 3],\n            \"ability\": \"toolcall\",\n            \"reward_model\": {\n                \"ground_truth\": '[{\"name\": \"rgb_to_cmyk\", \"arguments\": {\"r\": 0, \"g\": 0, \"b\": 255}}]',\n                \"style\": \"rule\",\n            },\n            \"tools_kwargs\": {},\n        },\n    ]\n\n    batch_size = len(test_data)\n    batch = collate_fn(test_data)\n\n    # Tensor part\n    assert batch[\"input_ids\"].shape == (batch_size, max_prompt_length)\n    assert isinstance(batch[\"input_ids\"], torch.Tensor)\n\n    # Non-tensor parts\n    expected_types = {\n        \"messages\": list,\n        \"raw_prompt_ids\": list,\n        \"ability\": str,\n        \"reward_model\": dict,\n        \"tools_kwargs\": dict,\n    }\n\n    for key, dtype in expected_types.items():\n        assert batch[key].shape == (batch_size,), (\n            f\"Expected shape {(batch_size,)} for '{key}', but got {batch[key].shape}\"\n        )\n        assert isinstance(batch[key][0], dtype), (\n            f\"'{key}' should contain elements of type {dtype}, but got {type(batch[key][0])}\"\n        )\n"
  },
  {
    "path": "tests/utils/dataset/test_rl_dataset_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport os\n\nimport pytest\nimport torch\nfrom omegaconf import OmegaConf\nfrom PIL import Image\nfrom torch.utils.data import DataLoader\n\nfrom verl import DataProto\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.dataset.rl_dataset import RLHFDataset, collate_fn\n\n\ndef get_gsm8k_data():\n    # prepare test dataset\n    local_folder = os.path.expanduser(\"~/data/gsm8k/\")\n    local_path = os.path.join(local_folder, \"train.parquet\")\n    os.makedirs(local_folder, exist_ok=True)\n    return local_path\n\n\ndef test_rl_dataset():\n    tokenizer = hf_tokenizer(os.path.expanduser(\"~/models/deepseek-ai/deepseek-coder-1.3b-instruct\"))\n    local_path = get_gsm8k_data()\n    config = OmegaConf.create(\n        {\n            \"prompt_key\": \"prompt\",\n            \"max_prompt_length\": 256,\n            \"filter_overlong_prompts\": True,\n            \"filter_overlong_prompts_workers\": 2,\n        }\n    )\n    dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config)\n\n    dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn)\n\n    a = next(iter(dataloader))\n\n    tensors = {}\n    non_tensors = {}\n\n    for key, val in a.items():\n        if isinstance(val, torch.Tensor):\n            tensors[key] = val\n        else:\n            non_tensors[key] = val\n\n    data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors)\n    assert len(data_proto) == 16\n    assert \"raw_prompt\" in data_proto.non_tensor_batch\n\n\ndef test_rl_dataset_with_max_samples():\n    tokenizer = hf_tokenizer(os.path.expanduser(\"~/models/deepseek-ai/deepseek-coder-1.3b-instruct\"))\n    local_path = get_gsm8k_data()\n    config = OmegaConf.create(\n        {\n            \"prompt_key\": \"prompt\",\n            \"max_prompt_length\": 256,\n            \"filter_overlong_prompts\": True,\n            \"filter_overlong_prompts_workers\": 2,\n            \"max_samples\": 5,\n        }\n    )\n    dataset = RLHFDataset(data_files=local_path, tokenizer=tokenizer, config=config, max_samples=5)\n    assert len(dataset) == 5\n\n\ndef test_image_rl_data():\n    tokenizer = hf_tokenizer(os.path.expanduser(\"~/models/Qwen/Qwen2-VL-2B-Instruct\"))\n    processor = hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2-VL-2B-Instruct\"))\n    config = OmegaConf.create(\n        {\n            \"prompt_key\": \"prompt\",\n            \"max_prompt_length\": 1024,\n            \"filter_overlong_prompts\": True,\n            \"filter_overlong_prompts_workers\": None,  # num_workers=1 hang in ci\n        }\n    )\n    dataset = RLHFDataset(\n        data_files=os.path.expanduser(\"~/data/geo3k/train.parquet\"),\n        tokenizer=tokenizer,\n        config=config,\n        processor=processor,\n    )\n\n    dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn)\n\n    a = next(iter(dataloader))\n\n    tensors = {}\n    non_tensors = {}\n\n    for key, val in a.items():\n        if isinstance(val, torch.Tensor):\n            tensors[key] = val\n        else:\n            non_tensors[key] = val\n\n    data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors)\n    assert len(data_proto) == 16\n    assert \"images\" not in data_proto.non_tensor_batch\n\n    for prompt in data_proto.non_tensor_batch[\"raw_prompt\"]:\n        assert len(prompt) == 1\n        prompt = prompt[0]\n        role, content = prompt[\"role\"], prompt[\"content\"]\n        assert role == \"user\"\n        assert len(content) == 2\n        assert content[0][\"type\"] == \"image\" and isinstance(content[0][\"image\"], Image.Image)\n        assert content[1][\"type\"] == \"text\" and isinstance(content[1][\"text\"], str)\n\n    print(\"raw_prompt\", data_proto.non_tensor_batch[\"raw_prompt\"][0])\n\n\n@pytest.fixture\ndef video_data_file():\n    data = [\n        {\n            \"problem_id\": 17,\n            \"problem\": \"How does the crowd's excitement change as the match progresses?\",\n            \"data_type\": \"video\",\n            \"prompt\": [\n                {\n                    \"role\": \"user\",\n                    \"content\": [\n                        {\"type\": \"video\", \"video\": \"LLaVA-Video-178K/academic_source/activitynet/v_2g9GrshWQrU.mp4\"},\n                        {\n                            \"type\": \"text\",\n                            \"text\": \"How does the crowd's excitement change as the match progresses? \"\n                            \"A. It fluctuates; B. It decreases; C. It builds up; D. It remains the same. \"\n                            \"Put your answer in <answer></answer>\",\n                        },\n                    ],\n                }\n            ],\n            \"problem_type\": \"multiple choice\",\n            \"solution\": \"C\",\n            \"data_source\": \"LLaVA-Video-178K/2_3_m_academic_v0_1\",\n        }\n    ] * 30\n\n    # Create test directory if it doesn't exist\n    os.makedirs(\"test_data\", exist_ok=True)\n    test_file = \"test_data/test_video.json\"\n    with open(test_file, \"w\") as f:\n        json.dump(data, f, indent=2)\n\n    return test_file\n\n\ndef test_video_rl_data(video_data_file):\n    tokenizer = hf_tokenizer(os.path.expanduser(\"~/models/Qwen/Qwen2-VL-2B-Instruct\"))\n    processor = hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2-VL-2B-Instruct\"))\n    config = OmegaConf.create(\n        {\n            \"prompt_key\": \"prompt\",\n            \"max_prompt_length\": 1024,\n            \"filter_overlong_prompts\": False,\n        }\n    )\n    dataset = RLHFDataset(\n        data_files=video_data_file,\n        tokenizer=tokenizer,\n        config=config,\n        processor=processor,\n    )\n\n    dataloader = DataLoader(dataset=dataset, batch_size=16, shuffle=True, drop_last=True, collate_fn=collate_fn)\n    batch = next(iter(dataloader))\n    tensors = {}\n    non_tensors = {}\n    for key, val in batch.items():\n        if isinstance(val, torch.Tensor):\n            tensors[key] = val\n        else:\n            non_tensors[key] = val\n\n    data_proto = DataProto.from_dict(tensors=tensors, non_tensors=non_tensors)\n    assert len(data_proto) == 16\n    assert \"images\" not in data_proto.non_tensor_batch\n\n    print(\"raw_prompt\", data_proto.non_tensor_batch[\"raw_prompt\"][0])\n"
  },
  {
    "path": "tests/utils/debug/test_metrics.py",
    "content": "# Copyright 2025 Individual Contributor: TomQunChaoA\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nimport torch\n\nfrom verl.protocol import DataProto\nfrom verl.utils.debug.metrics import calculate_debug_metrics\n\n\nclass TestMetrics(unittest.TestCase):\n    def test_calculate_debug_metrics(self):\n        data = DataProto.from_dict(\n            {\n                \"rollout_log_probs\": torch.tensor(\n                    [\n                        [-1.5085, -0.1200, -0.6650, -0.4823, -0.1426, -1.5557, -2.8532, -0.3919, -0.4294, -0.4700],\n                        [-0.0585, -0.0573, -0.4681, -0.5187, -0.7451, -1.2737, -0.0682, -0.4284, -0.5754, -0.0611],\n                    ]\n                ),\n                \"old_log_probs\": torch.tensor(\n                    [\n                        [-1.8636, -0.7863, -0.2136, -0.4376, -2.0257, -0.2579, -1.1547, -0.5203, -0.3802, -0.9872],\n                        [-0.3507, -0.5426, -0.2725, -0.4637, -0.3577, -0.3733, -1.7560, -1.9542, -0.4229, -1.3098],\n                    ]\n                ),\n                \"loss_mask\": torch.tensor([[1, 0, 0, 0, 1, 1, 0, 1, 1, 0], [1, 0, 1, 0, 1, 1, 1, 0, 1, 1]]),\n                \"responses\": torch.zeros((2, 10)),\n            }\n        )\n        metrics = calculate_debug_metrics(data)\n        print(metrics)\n        assert metrics[\"training/rollout_probs_diff_valid\"] == 1\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/megatron/test_pipeline_parallel.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nfrom verl.model_merger.megatron_model_merger import get_dynamic_pipeline_shards\nfrom verl.utils.megatron.pipeline_parallel import make_batch_generator\n\n\ndef test_make_batch_generator_no_vpp():\n    batches = [1, 2, 3]\n    vpp_size = 1\n    generator = make_batch_generator(batches, vpp_size)\n    assert list(generator) == batches\n\n\ndef test_make_batch_generator_with_vpp():\n    batches = [{\"data\": 1}, {\"data\": 2}]\n    vpp_size = 2\n    generators = make_batch_generator(batches, vpp_size)\n    assert isinstance(generators, list)\n    assert len(generators) == vpp_size\n\n    # Check each generator yields the original batches\n    for gen in generators:\n        assert list(gen) == batches\n\n\ndef test_make_batch_generator_empty():\n    batches = []\n    vpp_size = 1\n    generator = make_batch_generator(batches, vpp_size)\n    assert list(generator) == []\n\n    vpp_size = 3\n    generators = make_batch_generator(batches, vpp_size)\n    assert len(generators) == vpp_size\n    for gen in generators:\n        assert list(gen) == []\n\n\n@pytest.mark.parametrize(\n    \"layer_num,pp_size,gt\",\n    [\n        (61, 8, [6, 8, 8, 8, 8, 8, 8, 7]),\n        (61, 7, [8, 9, 9, 9, 9, 9, 8]),\n        (61, 1, [61]),\n        (61, 0, ValueError),\n        (10, 16, ValueError),\n    ],\n)\ndef test_get_dynamic_pipeline_shards(layer_num, pp_size, gt):\n    if isinstance(gt, list):\n        shards = get_dynamic_pipeline_shards(layer_num, pp_size)\n        assert len(shards) == len(gt) == pp_size, f\"Expected {pp_size} shards, got {len(shards)}\"\n        assert all([shard == gt[i] for i, shard in enumerate(shards)]), f\"Expected shards {gt}, got {shards}\"\n    elif issubclass(gt, Exception):\n        with pytest.raises(gt):\n            shards = get_dynamic_pipeline_shards(layer_num, pp_size)\n"
  },
  {
    "path": "tests/utils/reward_score/reward_score/test_sandbox_fusion_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport multiprocessing\nimport os\nimport time\nfrom concurrent.futures import ProcessPoolExecutor\nfrom unittest.mock import patch\n\nimport pytest\n\n# Import the function to be tested\nfrom verl.utils.reward_score.sandbox_fusion.utils import check_correctness\n\n# Get SANDBOX_URL from environment variable\nSANDBOX_URL = os.environ.get(\"SANDBOX_FUSION_URL\")\n# Define skip condition and reason\nskip_reason = \"SANDBOX_FUSION_URL environment variable not set\"\nskip_condition = not SANDBOX_URL\n\n# --- Test code (for real API calls) ---\nCODE_SUCCESS = \"\"\"\nimport sys\ndata = sys.stdin.read()\nif data == 'input1':\n    print('output1\\\\n', end='')\nelif data == 'input2':\n    print('output2\\\\n', end='')\nelse:\n    print('unexpected input', end='')\n\"\"\"\n\nCODE_WRONG_OUTPUT = \"\"\"\nprint('wrong_output\\\\n', end='')\n\"\"\"\n\nCODE_COMPILE_ERROR = \"\"\"\na=b\n\"\"\"\n\nCODE_RUNTIME_ERROR = \"\"\"\nimport sys\nprint(\"About to raise error\", file=sys.stderr)\nraise ValueError(\"This is a runtime error\")\n\"\"\"\n\nCODE_TIMEOUT = \"\"\"\nimport time\nimport sys\nprint(\"Sleeping...\", file=sys.stderr)\ntime.sleep(10) # Sleep time should be longer than the timeout set in the test\nprint(\"Finished sleeping\", file=sys.stderr)\n\"\"\"\n\n# --- Test input/output data ---\nINPUT_OUTPUT_VALID = {\"inputs\": [\"input1\", \"input2\"], \"outputs\": [\"output1\\n\", \"output2\\n\"]}\n\nINPUT_OUTPUT_SINGLE = {\"inputs\": [\"input1\"], \"outputs\": [\"output1\\n\"]}\n\nINPUT_OUTPUT_MISMATCH = {\"inputs\": [\"input1\"], \"outputs\": [\"output1\\n\", \"output2\\n\"]}\n\nINPUT_OUTPUT_INVALID_MISSING_KEY = {\"inputs\": [\"input1\"]}\n\n# --- Integration test cases (calling real API) ---\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_success_correct():\n    \"\"\"Integration test: Code is correct, output is correct\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_SUCCESS)\n    assert results == [True, True]\n    assert metadata_list[0][\"status\"] == \"success\"\n    assert metadata_list[0][\"stdout\"] == \"output1\\n\"\n    assert metadata_list[1][\"status\"] == \"success\"\n    assert metadata_list[1][\"stdout\"] == \"output2\\n\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_success_wrong_output():\n    \"\"\"Integration test: Code runs successfully, but output is wrong\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_WRONG_OUTPUT)\n    assert results == [False, False]\n    assert metadata_list[0][\"status\"] == \"wrong_answer\"\n    assert metadata_list[0][\"stdout\"] == \"wrong_output\\n\"\n    assert metadata_list[1][\"status\"] == \"wrong_answer\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_compile_error():\n    \"\"\"Integration test: Code causes compile error\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_VALID, CODE_COMPILE_ERROR, language=\"cpp\")\n    assert results == [-4, -4]\n    assert metadata_list[0][\"status\"] == \"compile_error\"\n    assert metadata_list[1][\"status\"] == \"compile_error\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_runtime_error():\n    \"\"\"Integration test: Code causes runtime error\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_RUNTIME_ERROR)\n    assert results == [-2]\n    assert metadata_list[0][\"status\"] == \"runtime_error\"\n    # More assertions can be added based on the actual API response, e.g., exit_code, stderr\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_runtime_timeout():\n    \"\"\"Integration test: Code causes runtime timeout\"\"\"\n    test_timeout = 5  # Set a timeout shorter than the sleep time in CODE_TIMEOUT\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_SINGLE, CODE_TIMEOUT, timeout=test_timeout)\n    assert results == [-3]\n    assert metadata_list[0][\"status\"] == \"timeout\"\n    # More assertions can be added based on the actual API response, e.g., run_status\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_concurrency_high_load():\n    \"\"\"Integration test: High concurrency (100 cases) against real API with mixed results (success, wrong\n    answer, timeout)\"\"\"\n    concurrency_level = 100\n    # Indices for different expected outcomes\n    wrong_answer_indices = {10, 25, 50}\n    timeout_indices = {5, 30, 60, 90}  # Indices where we expect a timeout\n\n    # Generate 100 input/output pairs and code\n    high_load_inputs = []\n    high_load_outputs = []\n    expected_results_map = {}  # Store expected result for each index\n\n    for i in range(concurrency_level):\n        if i in timeout_indices:\n            # Use a special input to trigger timeout in the code\n            high_load_inputs.append(f\"input_timeout_{i}\")\n            # Output doesn't matter for timeout, but keep it consistent\n            high_load_outputs.append(f\"output_{i}\\n\")\n            expected_results_map[i] = -3  # Expect timeout\n        elif i in wrong_answer_indices:\n            high_load_inputs.append(f\"input_{i}\")\n            # Intentionally set wrong expected output\n            high_load_outputs.append(f\"wrong_output_{i}\\n\")\n            expected_results_map[i] = False  # Expect wrong answer\n        else:\n            high_load_inputs.append(f\"input_{i}\")\n            # Correct expected output\n            high_load_outputs.append(f\"output_{i}\\n\")\n            expected_results_map[i] = True  # Expect success\n\n    high_load_in_outs = {\"inputs\": high_load_inputs, \"outputs\": high_load_outputs}\n\n    # Code that handles normal inputs, and sleeps on specific \"timeout\" inputs\n    code_mixed_concurrent = \"\"\"\nimport sys\nimport time\ndata = sys.stdin.read()\nif data.startswith('input_timeout_'):\n    time.sleep(20) # Sleep longer than the test timeout\n    print(f\"output_{data.split('_')[-1]}\\\\n\", end='') # Still print something in case it finishes early\nelif data.startswith('input_'):\n    print(f\"output_{data.split('_')[-1]}\\\\n\", end='')\nelse:\n    print(\"unknown_input\\\\n\", end='')\n\"\"\"\n    # Set a reasonable timeout per case (must be less than the sleep time in the code)\n    test_timeout = 15  # Allow slightly more time due to potential API load, but less than 20s sleep\n\n    start_time = time.time()\n    results, metadata_list = check_correctness(\n        SANDBOX_URL,\n        high_load_in_outs,\n        code_mixed_concurrent,  # Use the new code\n        timeout=test_timeout,\n    )\n    end_time = time.time()\n    duration = end_time - start_time\n    print(\n        f\"\\nHigh concurrency test ({concurrency_level} cases with {len(wrong_answer_indices)} wrong answers, \"\n        f\"{len(timeout_indices)} timeouts) duration: {duration:.2f} seconds\"\n    )\n\n    # Verify results against the expected map\n    assert len(results) == concurrency_level, f\"Expected {concurrency_level} results, got {len(results)}\"\n\n    correct_count = 0\n    wrong_count = 0\n    timeout_count = 0\n    unexpected_results = []\n    for i, r in enumerate(results):\n        expected = expected_results_map[i]\n        if r == expected:\n            if expected is True:\n                correct_count += 1\n            elif expected is False:\n                wrong_count += 1\n            elif expected == -3:\n                timeout_count += 1\n        else:\n            unexpected_results.append((i, r, f\"Expected {expected}\"))\n\n    print(\n        f\"Correct results (True): {correct_count}/\"\n        f\"{concurrency_level - len(wrong_answer_indices) - len(timeout_indices)}\"\n    )\n    print(f\"Expected wrong answers (False, correctly identified): {wrong_count}/{len(wrong_answer_indices)}\")\n    print(f\"Expected timeouts (-3, correctly identified): {timeout_count}/{len(timeout_indices)}\")\n\n    if unexpected_results:\n        print(\"Unexpected results found:\")\n        for idx, res, expected_str in unexpected_results[:10]:  # Print first 10 unexpected\n            print(f\"  Index {idx}: Got {res}, {expected_str}. Metadata: {metadata_list[idx]}\")\n        raise AssertionError(f\"Found {len(unexpected_results)} unexpected results.\")\n\n    assert correct_count == concurrency_level - len(wrong_answer_indices) - len(timeout_indices), (\n        \"Incorrect number of successful results\"\n    )\n    assert wrong_count == len(wrong_answer_indices), \"Incorrect number of identified wrong answers\"\n    assert timeout_count == len(timeout_indices), \"Incorrect number of identified timeouts\"\n\n    # Verify metadata count and basic status of one of each type\n    assert len(metadata_list) == concurrency_level\n    # Find the first correct index\n    first_correct_index = next(\n        i for i in range(concurrency_level) if i not in wrong_answer_indices and i not in timeout_indices\n    )\n    assert metadata_list[first_correct_index][\"status\"] == \"success\"\n    assert metadata_list[first_correct_index][\"stdout\"] == f\"output_{first_correct_index}\\n\"\n\n    # Check the status of the first intentionally wrong case\n    first_wrong_index = min(wrong_answer_indices)\n    assert metadata_list[first_wrong_index][\"status\"] == \"wrong_answer\"\n    assert metadata_list[first_wrong_index][\"stdout\"] == f\"output_{first_wrong_index}\\n\"\n    assert metadata_list[first_wrong_index][\"expected_output\"] == f\"wrong_output_{first_wrong_index}\\n\"\n\n    # Check the status of the first intentionally timeout case\n    first_timeout_index = min(timeout_indices)\n    assert metadata_list[first_timeout_index][\"status\"] == \"timeout\"\n    # For timeout, stdout might be None or empty depending on when the timeout occurred\n    # assert metadata_list[first_timeout_index][\"stdout\"] is None or metadata_list[first_timeout_index][\"stdout\"] == \"\"\n\n\n# --- Unit test cases (using mock) ---\n\n\n@patch(\"verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api\")\ndef test_unit_concurrency_order(mock_call_sandbox_api):\n    sandbox_url = \"mock_url\"\n    generation = \"print(input())\"\n    language = \"python\"\n    timeout = 5\n    in_outs = {\"inputs\": [\"input1\", \"input2\", \"input3\"], \"outputs\": [\"output1\", \"output2\", \"output3\"]}\n\n    def side_effect(*args, **kwargs):\n        stdin = kwargs.get(\"stdin\")\n        if stdin == \"input1\":\n            return (\n                {\"status\": \"Success\", \"run_result\": {\"status\": \"Finished\", \"stdout\": \"output1\", \"return_code\": 0}},\n                None,\n            )\n        elif stdin == \"input2\":\n            time.sleep(0.1)\n            return (\n                {\"status\": \"Success\", \"run_result\": {\"status\": \"Finished\", \"stdout\": \"output2\", \"return_code\": 0}},\n                None,\n            )\n        elif stdin == \"input3\":\n            return (\n                {\"status\": \"Success\", \"run_result\": {\"status\": \"Finished\", \"stdout\": \"output3\", \"return_code\": 0}},\n                None,\n            )\n        else:\n            return (None, \"Unknown input in mock\")\n\n    mock_call_sandbox_api.side_effect = side_effect\n\n    results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language)\n\n    assert results == [True, True, True]\n    assert len(metadata_list) == 3\n    assert metadata_list[0][\"case_index\"] == 0\n    assert metadata_list[0][\"status\"] == \"success\"\n    assert metadata_list[1][\"case_index\"] == 1\n    assert metadata_list[1][\"status\"] == \"success\"\n    assert metadata_list[2][\"case_index\"] == 2\n    assert metadata_list[2][\"status\"] == \"success\"\n    assert mock_call_sandbox_api.call_count == 3\n\n\n@patch(\"verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api\")\ndef test_unit_api_timeout_error_concurrent(mock_call_sandbox_api):\n    sandbox_url = \"mock_url\"\n    generation = \"print(input())\"\n    language = \"python\"\n    timeout = 5\n    in_outs = {\"inputs\": [\"input1\", \"input2_timeout\", \"input3\"], \"outputs\": [\"output1\", \"output2\", \"output3\"]}\n\n    api_error_message = \"API Call Failed: Gateway Timeout (504) on attempt 3/3\"\n\n    def side_effect(*args, **kwargs):\n        stdin = kwargs.get(\"stdin\")\n        if stdin == \"input1\":\n            return (\n                {\"status\": \"Success\", \"run_result\": {\"status\": \"Finished\", \"stdout\": \"output1\", \"return_code\": 0}},\n                None,\n            )\n        elif stdin == \"input2_timeout\":\n            return (None, api_error_message)\n        elif stdin == \"input3\":\n            return (\n                {\"status\": \"Success\", \"run_result\": {\"status\": \"Finished\", \"stdout\": \"output3\", \"return_code\": 0}},\n                None,\n            )\n        else:\n            return (None, \"Unknown input in mock\")\n\n    mock_call_sandbox_api.side_effect = side_effect\n\n    results, metadata_list = check_correctness(sandbox_url, in_outs, generation, timeout, language)\n\n    assert results == [True, -1, True]\n    assert len(metadata_list) == 3\n    assert metadata_list[0][\"status\"] == \"success\"\n    assert metadata_list[1][\"status\"] == \"api_error\"\n    assert metadata_list[1][\"api_request_error\"] == api_error_message\n    assert metadata_list[2][\"status\"] == \"success\"\n    assert mock_call_sandbox_api.call_count == 3\n\n\n# --- Constants for the new concurrency test ---\n# Define a low global concurrency limit to test the semaphore's effect\nMAX_GLOBAL_CONCURRENCY_LIMIT_TEST = 5\n# Define the number of processes used in the test\nNUM_PROCESSES_TEST = 4\n# Define the number of tasks processed by check_correctness in each process (i.e., internal\n# ThreadPoolExecutor's concurrency potential)\nNUM_TASKS_PER_PROCESS_TEST = 3\n# Simulate API call duration to ensure calls can overlap\nSIMULATED_API_CALL_DURATION_TEST = 0.2  # seconds\n\n\n# --- Mock API call function for concurrency tracking ---\n# This function will replace the real call_sandbox_api and use shared variables to track concurrency\ndef _mock_api_call_for_concurrency_tracking(\n    active_calls_counter,  # multiprocessing.Value\n    max_calls_tracker,  # multiprocessing.Value\n    call_lock,  # multiprocessing.Lock\n    # Standard call_sandbox_api parameters\n    sandbox_fusion_url,\n    code,\n    stdin,\n    compile_timeout,\n    run_timeout,\n    memory_limit_mb,\n    language,\n):\n    # entry_time = time.time() # For detailed logging\n    with call_lock:\n        active_calls_counter.value += 1\n        if active_calls_counter.value > max_calls_tracker.value:\n            max_calls_tracker.value = active_calls_counter.value\n        # Optional debug log:\n        # print(f\"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call Start. Active: \"\n        #       f\"{active_calls_counter.value}, Max Observed: {max_calls_tracker.value}, Input: {stdin}\")\n\n    time.sleep(SIMULATED_API_CALL_DURATION_TEST)  # Simulate actual work duration\n\n    # exit_time = time.time() # For detailed logging\n    with call_lock:\n        active_calls_counter.value -= 1\n        # Optional debug log:\n        # print(f\"[PID:{os.getpid()}-TID:{threading.get_ident()}] API Call End. Active: \"\n        #       f\"{active_calls_counter.value}, Input: {stdin}, Duration: {exit_time - entry_time:.2f}s\")\n\n    # Return a simulated successful API response\n    return {\n        \"status\": \"Success\",\n        \"run_result\": {\"status\": \"Finished\", \"stdout\": f\"mock_output_for_{stdin}\", \"return_code\": 0},\n    }, None\n\n\n# --- Worker function for ProcessPoolExecutor ---\n# This function runs in each child process of ProcessPoolExecutor\ndef _process_pool_worker_for_concurrency_test(\n    sandbox_url,\n    in_outs,\n    generation,\n    memory_limit_mb,\n    language,\n    timeout,\n    mp_semaphore_for_check_correctness,\n    active_calls_counter,\n    max_calls_tracker,\n    call_lock,\n):\n    # Corrected lambda to accept keyword arguments matching call_sandbox_api's usage\n    curried_mock_api_call = (\n        lambda sandbox_fusion_url, code, stdin, compile_timeout, run_timeout, memory_limit_mb, language: (\n            _mock_api_call_for_concurrency_tracking(\n                active_calls_counter,\n                max_calls_tracker,\n                call_lock,\n                sandbox_fusion_url,\n                code,\n                stdin,\n                compile_timeout,\n                run_timeout,\n                memory_limit_mb,\n                language,\n            )\n        )\n    )\n\n    # ---- START DEBUG PRINTS ----\n    import os\n\n    import verl.utils.reward_score.sandbox_fusion.utils\n\n    print(\n        f\"[Worker PID:{os.getpid()}] Original call_sandbox_api: \"\n        f\"{verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api}\",\n        flush=True,\n    )\n    # ---- END DEBUG PRINTS ----\n\n    with patch(\n        \"verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api\", side_effect=curried_mock_api_call\n    ) as mock_obj:\n        # ---- START DEBUG PRINTS ----\n        print(\n            f\"[Worker PID:{os.getpid()}] Patched call_sandbox_api: \"\n            f\"{verl.utils.reward_score.sandbox_fusion.utils.call_sandbox_api}\",\n            flush=True,\n        )\n        print(f\"[Worker PID:{os.getpid()}] Mock object: {mock_obj}\", flush=True)\n        # ---- END DEBUG PRINTS ----\n        results, metadata_list = check_correctness(\n            sandbox_fusion_url=sandbox_url,\n            in_outs=in_outs,\n            generation=generation,\n            timeout=timeout,\n            memory_limit_mb=memory_limit_mb,\n            language=language,\n            concurrent_semaphore=mp_semaphore_for_check_correctness,  # Pass multiprocessing.Semaphore\n        )\n        # print(f\"Process {os.getpid()} finished check_correctness. Processed {len(results)} tasks.\")\n    return len(results)  # Return the number of processed tasks for basic validation\n\n\n# --- The actual test case for multiprocess concurrency control ---\ndef test_multiprocess_global_concurrency_limit_with_semaphore():\n    \"\"\"\n    Tests that the global concurrent_semaphore (multiprocessing.Semaphore)\n    correctly limits the number of concurrent calls to call_sandbox_api\n    across multiple processes, each potentially running multiple threads\n    via check_correctness's internal ThreadPoolExecutor.\n    \"\"\"\n    manager = multiprocessing.Manager()\n    active_calls_counter = manager.Value(\"i\", 0)  # Current active mock API calls\n    max_calls_tracker = manager.Value(\"i\", 0)  # Observed maximum concurrent mock API calls\n    call_lock = manager.Lock()  # Lock to protect counters\n\n    # Create a multiprocessing.Semaphore instance, this is the global semaphore we are testing.\n    # It will be passed to check_correctness and used by _process_single_case to limit calls to call_sandbox_api.\n    global_mp_semaphore = manager.Semaphore(MAX_GLOBAL_CONCURRENCY_LIMIT_TEST)\n\n    mock_sandbox_url = \"mock_url_for_concurrency_test\"\n    mock_generation = \"pass\"  # Specific code content is not important as API call is mocked\n    mock_memory_limit_mb = 1024\n    mock_language = \"python\"\n    mock_timeout = 5  # Timeout setting, not critical for mock calls\n\n    # Input/output data for each process\n    # NUM_TASKS_PER_PROCESS_TEST tasks will be handled by check_correctness's internal ThreadPoolExecutor\n    process_in_outs = {\n        \"inputs\": [f\"task_input_{i}\" for i in range(NUM_TASKS_PER_PROCESS_TEST)],\n        \"outputs\": [f\"task_output_{i}\" for i in range(NUM_TASKS_PER_PROCESS_TEST)],\n    }\n\n    futures = []\n    total_tasks_expected_to_run = NUM_PROCESSES_TEST * NUM_TASKS_PER_PROCESS_TEST\n\n    test_start_time = time.time()\n\n    with ProcessPoolExecutor(max_workers=NUM_PROCESSES_TEST) as executor:\n        for i in range(NUM_PROCESSES_TEST):\n            future = executor.submit(\n                _process_pool_worker_for_concurrency_test,  # Worker function\n                mock_sandbox_url,\n                process_in_outs,\n                mock_generation,\n                mock_memory_limit_mb,\n                mock_language,\n                mock_timeout,\n                global_mp_semaphore,  # Global semaphore to test\n                active_calls_counter,  # Shared variables for tracking\n                max_calls_tracker,\n                call_lock,\n            )\n            futures.append(future)\n\n    # Wait for all processes to complete and collect results\n    num_tasks_processed_per_worker = [f.result() for f in futures]\n    test_end_time = time.time()\n    total_execution_time = test_end_time - test_start_time\n\n    # Print some test statistics for debugging and validation\n    print(\"\\n--- Global Concurrency Test Stats ---\")\n    print(f\"Semaphore Limit (MAX_GLOBAL_CONCURRENCY_LIMIT_TEST): {MAX_GLOBAL_CONCURRENCY_LIMIT_TEST}\")\n    print(f\"Number of Processes (NUM_PROCESSES_TEST): {NUM_PROCESSES_TEST}\")\n    print(f\"Tasks per Process (NUM_TASKS_PER_PROCESS_TEST): {NUM_TASKS_PER_PROCESS_TEST}\")\n    print(f\"Total Tasks Submitted: {total_tasks_expected_to_run}\")\n    print(f\"Simulated API Call Duration: {SIMULATED_API_CALL_DURATION_TEST}s\")\n    print(f\"Total Test Execution Time: {total_execution_time:.2f}s\")\n    print(f\"Max Concurrent Mock API Calls Observed: {max_calls_tracker.value}\")\n    # print(f\"Tasks processed per worker: {num_tasks_processed_per_worker}\")\n\n    # Verify that all submitted tasks have been processed\n    assert sum(num_tasks_processed_per_worker) == total_tasks_expected_to_run, (\n        \"Mismatch in the number of tasks processed.\"\n    )\n\n    # Verify that the mock API was called at least once\n    assert max_calls_tracker.value > 0, \"The mocked API call_sandbox_api was not called.\"\n\n    # Core assertion: Observed maximum concurrent calls should not exceed the semaphore's limit\n    assert max_calls_tracker.value <= MAX_GLOBAL_CONCURRENCY_LIMIT_TEST, (\n        f\"Observed concurrency ({max_calls_tracker.value}) exceeded semaphore limit \"\n        f\"({MAX_GLOBAL_CONCURRENCY_LIMIT_TEST}).\"\n    )\n\n    # Optional: Rough check on execution time to verify semaphore is working to limit concurrency\n    # Theoretical minimum execution time = (Total tasks / Concurrency limit) * Single task duration\n    # Actual time will be longer due to various overheads\n    min_expected_duration = (\n        total_tasks_expected_to_run * SIMULATED_API_CALL_DURATION_TEST\n    ) / MAX_GLOBAL_CONCURRENCY_LIMIT_TEST\n    # print(f\"Minimum Expected Execution Time (approx): {min_expected_duration:.2f}s\")\n    # Allow some margin, e.g., 80% of theoretical minimum time\n    assert total_execution_time >= min_expected_duration * 0.8, (\n        f\"Total execution time ({total_execution_time:.2f}s) was unexpectedly short, suggesting the \"\n        f\"semaphore might not be effectively limiting concurrency as expected \"\n        f\"(min expected: {min_expected_duration * 0.8:.2f}s).\"\n    )\n\n\n# Ensure there is no more code after this point if these were the last functions.\n# If there was other code, it would follow here.\ndef test_unit_invalid_input_format():\n    \"\"\"Unit test: Invalid in_outs format passed\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, None, CODE_SUCCESS)\n    assert results == [-1]\n    assert metadata_list[0][\"error\"] == \"Invalid input/output data\"\n\n    results, metadata_list = check_correctness(SANDBOX_URL, {}, CODE_SUCCESS)\n    assert results == [-1]\n    assert metadata_list[0][\"error\"] == \"Invalid input/output data\"\n\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_INVALID_MISSING_KEY, CODE_SUCCESS)\n    assert results == [-1]\n    assert metadata_list[0][\"error\"] == \"Invalid input/output data\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_unit_input_output_mismatch():\n    \"\"\"Unit test: Mismatch between the number of inputs and outputs\"\"\"\n    results, metadata_list = check_correctness(SANDBOX_URL, INPUT_OUTPUT_MISMATCH, CODE_SUCCESS)\n    assert results == [-1]\n    assert len(metadata_list) == 1\n    assert metadata_list[0][\"error\"] == \"Input/output count mismatch\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_integration_concurrency_all_timeout():\n    \"\"\"Integration test: High concurrency (100 cases) against real API, all causing timeout\"\"\"\n    concurrency_level = 100\n    code_infinite_loop = \"\"\"\ndef knight_moves(X, Y):\n    MOD = 10**9 + 7\n    dp = [[0] * (Y + 1) for _ in range(X + 1)]\n    dp[0][0] = 1\n    for i in range(1, X + 1):\n        for j in range(1, Y + 1):\n            dp[i][j] = (dp[i - 1][j] + dp[i][j - 1]) % MOD\n    return dp[X][Y]\n\ndef solve():\n    X, Y = map(int, input().split())\n    print(knight_moves(X, Y))\n\nif __name__ == \"__main__\":\n    solve()\n    \"\"\"\n\n    # Generate 100 simple input/output pairs (content doesn't matter)\n    timeout_inputs = [\"324 384429\" for i in range(concurrency_level)]\n    timeout_outputs = [f\"output_{i}\\n\" for i in range(concurrency_level)]\n    timeout_in_outs = {\"inputs\": timeout_inputs, \"outputs\": timeout_outputs}\n\n    # Set a timeout for the test cases\n    test_timeout = 10  # Set a timeout value\n\n    start_time = time.time()\n    results, metadata_list = check_correctness(SANDBOX_URL, timeout_in_outs, code_infinite_loop, timeout=test_timeout)\n    end_time = time.time()\n    duration = end_time - start_time\n    print(f\"\\nHigh concurrency all timeout test ({concurrency_level} cases) duration: {duration:.2f} seconds\")\n\n    # Verify all results are -3 (timeout)\n    assert len(results) == concurrency_level, f\"Expected {concurrency_level} results, got {len(results)}\"\n    all_timed_out = all(r == -3 for r in results)\n    if not all_timed_out:\n        non_timeout_indices = [i for i, r in enumerate(results) if r != -3]\n        print(f\"Indices that did not time out: {non_timeout_indices}\")\n        # Print metadata for the first few non-timeout cases for debugging\n        for i in non_timeout_indices[:5]:\n            print(f\"Metadata for non-timeout case {i}: {metadata_list[i]}\")\n    assert all_timed_out, f\"Not all {concurrency_level} concurrent tests resulted in timeout (-3). Results: {results}\"\n\n    # Verify metadata count and status of the first case\n    assert len(metadata_list) == concurrency_level\n    assert metadata_list[0][\"status\"] == \"timeout\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_fn_name_success_single_case():\n    \"\"\"Tests successful execution for a single test case with fn_name.\n    from livecodebench/code_generation_lite test 510\n    \"\"\"\n    generation_code = \"\"\"\nclass Solution:\n    def occurrencesOfElement(self, nums: List[int], queries: List[int], x: int) -> List[int]:\n        positions = defaultdict(list)\n        for idx, num in enumerate(nums):\n            positions[num].append(idx)\n\n        x_positions = positions[x]\n        answer = []\n        for k in queries:\n            if k > len(x_positions):\n                answer.append(-1)\n            else:\n                answer.append(x_positions[k-1])\n        return answer\n\"\"\"\n    in_outs = {\n        \"fn_name\": \"occurrencesOfElement\",\n        \"inputs\": [\"[1, 3, 1, 7]\\n[1, 3, 2, 4]\\n1\", \"[1, 2, 3]\\n[10]\\n5\"],\n        \"outputs\": [\"[0, -1, 2, -1]\", \"[-1]\"],\n    }\n\n    # Use a short timeout for fast tests\n    results, metadata_list = check_correctness(SANDBOX_URL, in_outs, generation_code, timeout=5)\n    # from verl.utils.reward_score.prime_code import apps_check_correctness\n    # results, metadata_list = apps_check_correctness(in_outs=in_outs, generation=generation_code,\n    #                                                        timeout=50000, debug=True)\n\n    assert results == [True, True]\n    assert \"error\" not in metadata_list[0]\n    assert metadata_list[0].get(\"status\") != \"compile_error\"\n    assert metadata_list[0].get(\"status\") != \"runtime_error\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_none_and_empty_stdin_passed_correctly():\n    \"\"\"\n    Tests that when stdin data is set to an empty string or None, it is still\n    is passed correctly to Sandbox Fusion as an empty string.\n    \"\"\"\n    echo_code = \"\"\"\nimport sys\nprint(f\"You said '{sys.stdin.readline().strip()}'\")\n\"\"\"\n    in_outs = {\n        \"inputs\": [None, \"\", \"hello\"],\n        \"outputs\": [\"You said ''\", \"You said ''\", \"You said 'hello'\"],\n    }\n\n    # Use a short timeout for fast tests\n    results, metadata_list = check_correctness(SANDBOX_URL, in_outs, echo_code, timeout=5)\n\n    assert results == [True, True, True]\n    assert \"error\" not in metadata_list[0]\n    assert metadata_list[0].get(\"status\") != \"compile_error\"\n    assert metadata_list[0].get(\"status\") != \"runtime_error\"\n\n\n@pytest.mark.skipif(skip_condition, reason=skip_reason)\ndef test_assert_case_success():\n    \"\"\"Tests successful execution for assert case.\n    from KodCode\n    \"\"\"\n    generation_code = \"\"\"\nfrom typing import List, Tuple\n\ndef merge_intervals(intervals: List[Tuple[int, int]]) -> List[Tuple[int, int]]:\n    if not intervals:\n        return []\n\n    # Sort intervals by the start time\n    intervals.sort(key=lambda x: x[0])\n\n    merged = [intervals[0]]\n\n    for current in intervals[1:]:\n        last = merged[-1]\n        # If intervals overlap, merge them\n        if current[0] <= last[1]:\n            merged[-1] = (last[0], max(last[1], current[1]))\n        else:\n            merged.append(current)\n\n    return merged\n\"\"\"\n    test_cases = {\n        \"fn_name\": \"merge_intervals\",\n        \"assert_case\": [\n            \"assert merge_intervals([(0, 1), (3, 5), (4, 7), (6, 8), (10, 12),\"\n            \" (12, 14)]) == [(0, 1), (3, 8), (10, 14)]\",\n            \"assert merge_intervals([(1, 2), (2, 3), (3, 4)]) == [(1, 4)]\",\n            \"assert merge_intervals([(1, 2), (3, 4), (5, 6)]) == [(1, 2), (3, 4), (5, 5)]\",\n        ],\n    }\n\n    assert_cases = test_cases.get(\"assert_case\")\n    test_cases.setdefault(\"inputs\", [\"\" for _ in assert_cases])\n    test_cases.setdefault(\"outputs\", [None for _ in assert_cases])\n\n    # Use a short timeout for fast tests\n    results, metadata_list = check_correctness(SANDBOX_URL, test_cases, generation_code, timeout=5)\n    assert results == [True, True, -2]\n    for i in range(2):\n        assert \"error\" not in metadata_list[i]\n        assert metadata_list[i].get(\"status\") == \"success\"\n        assert metadata_list[i].get(\"expected_output\") is None\n        assert metadata_list[i].get(\"status\") != \"runtime_error\"\n    assert \"error\" not in metadata_list[2]\n    assert metadata_list[2].get(\"status\") != \"success\"\n    assert metadata_list[2].get(\"expected_output\") is None\n    assert metadata_list[2].get(\"status\") == \"runtime_error\"\n"
  },
  {
    "path": "tests/utils/reward_score/test_sandbox_on_cpu.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport json\nimport os\n\nimport pytest\n\nfrom verl.utils.reward_score import default_compute_score, sandbox_fusion\nfrom verl.workers.reward_manager.prime import parallel_compute_score_async\n\nprime_math_answers = [\n    \"\"\"\\\\begin{bmatrix}\\n -7 & 6 & -8 \\\\\\\\\\n 11 & -9 & 12 \\\\\\\\\\n 15 & -16 & 19 \\n \\\\end{bmatrix}\"\"\",\n    \"\"\"\\\\frac{\\\\sqrt{505}}{7}\"\"\",\n    \"\"\"x^2 + y^2 + 4x - 6y + 13\"\"\",\n]\nprime_math_gts = [\n    \"\"\"\\\\begin{pmatrix}\\n -7 & 6 & -8 \\\\\\\\\\n 11 & -9 & 12 \\\\\\\\\\n 15 & -16 & 19\\n \\\\end{pmatrix}\"\"\",  # mat test\n    \"\"\"\\\\frac{\\\\sqrt{505}}{7}\"\"\",  # frac test\n    \"\"\"(x + 2)^2 + (y - 3)^2 \"\"\",  # symbolic test\n]\n\nprime_code_answers = [\n    \"\"\"import sys\nfrom collections import deque\n\ndef main():\n    data = sys.stdin.read().split()\n    it = iter(data)\n    \n    # Read start and target positions\n    x0, y0, x1, y1 = int(next(it)), int(next(it)), int(next(it)), int(next(it))\n    \n    n = int(next(it))\n    allowed = set()\n    # The total number of allowed cells is at most 10^5.\n    for _ in range(n):\n        r = int(next(it))\n        a = int(next(it))\n        b = int(next(it))\n        for c in range(a, b + 1):\n            allowed.add((r, c))\n    \n    # Directions for the king (8 neighboring cells)\n    directions = [(-1, -1), (-1, 0), (-1, 1),\n                  (0, -1),           (0, 1),\n                  (1, -1),  (1, 0),  (1, 1)]\n    \n    start = (x0, y0)\n    target = (x1, y1)\n    \n    # BFS initialization\n    queue = deque()\n    queue.append((x0, y0, 0))\n    # Mark the starting cell as visited by removing it from allowed set.\n    allowed.discard(start)\n    \n    while queue:\n        x, y, moves = queue.popleft()\n        if (x, y) == target:\n            print(moves)\n            return\n        for dx, dy in directions:\n            nx, ny = x + dx, y + dy\n            if (nx, ny) in allowed:\n                allowed.remove((nx, ny))\n                queue.append((nx, ny, moves + 1))\n    \n    print(-1)\n\nif __name__ == '__main__':\n    main()\n\"\"\"\n] * 2\nprime_code_gts = [\n    \"\"\"{\\n \\\"inputs\\\": [\\n \\\"5 7 6 11\\\\n3\\\\n5 3 8\\\\n6 7 11\\\\n5 2 5\\\\n\\\",\\n \\\"3 4 3 10\\\\n3\\\\n3 1 4\\\\n4 5 9\\\\n3 10 10\\\\n\\\",\\n \\\"1 1 2 10\\\\n2\\\\n1 1 3\\\\n2 6 10\\\\n\\\",\\n \\\"9 8 7 8\\\\n9\\\\n10 6 6\\\\n10 6 6\\\\n7 7 8\\\\n9 5 6\\\\n8 9 9\\\\n9 5 5\\\\n9 8 8\\\\n8 5 6\\\\n9 10 10\\\\n\\\",\\n \\\"6 15 7 15\\\\n9\\\\n6 15 15\\\\n7 14 14\\\\n6 15 15\\\\n9 14 14\\\\n7 14 16\\\\n6 15 15\\\\n6 15 15\\\\n7 14 14\\\\n8 15 15\\\\n\\\",\\n \\\"13 16 20 10\\\\n18\\\\n13 16 16\\\\n20 10 10\\\\n19 10 10\\\\n12 15 15\\\\n20 10 10\\\\n18 11 11\\\\n19 10 10\\\\n19 10 10\\\\n20 10 10\\\\n19 10 10\\\\n20 10 10\\\\n20 10 10\\\\n19 10 10\\\\n18 11 11\\\\n13 16 16\\\\n12 15 15\\\\n19 10 10\\\\n19 10 10\\\\n\\\",\\n \\\"89 29 88 30\\\\n16\\\\n87 31 31\\\\n14 95 95\\\\n98 88 89\\\\n96 88 88\\\\n14 97 97\\\\n13 97 98\\\\n100 88 88\\\\n88 32 32\\\\n99 88 89\\\\n90 29 29\\\\n87 31 31\\\\n15 94 96\\\\n89 29 29\\\\n88 32 32\\\\n97 89 89\\\\n88 29 30\\\\n\\\",\\n \\\"30 14 39 19\\\\n31\\\\n35 7 11\\\\n37 11 12\\\\n32 13 13\\\\n37 5 6\\\\n46 13 13\\\\n37 14 14\\\\n31 13 13\\\\n43 13 19\\\\n45 15 19\\\\n46 13 13\\\\n32 17 17\\\\n41 14 19\\\\n30 14 14\\\\n43 13 17\\\\n34 16 18\\\\n44 11 19\\\\n38 13 13\\\\n40 12 20\\\\n37 16 18\\\\n46 16 18\\\\n34 10 14\\\\n36 9 10\\\\n36 15 19\\\\n38 15 19\\\\n42 13 19\\\\n33 14 15\\\\n35 15 19\\\\n33 17 18\\\\n39 12 20\\\\n36 5 7\\\\n45 12 12\\\\n\\\",\\n \\\"2 1 1 1\\\\n2\\\\n1 1 2\\\\n2 1 2\\\\n\\\",\\n \\\"1 1 1 2\\\\n5\\\\n1000000000 1 10000\\\\n19920401 1188 5566\\\\n1000000000 1 10000\\\\n1 1 10000\\\\n5 100 200\\\\n\\\",\\n \\\"1 1 1000000000 2\\\\n5\\\\n1000000000 1 10000\\\\n19920401 1188 5566\\\\n1000000000 1 10000\\\\n1 1 10000\\\\n5 100 200\\\\n\\\"\\n ],\\n \\\"outputs\\\": [\\n \\\"4\\\\n\\\",\\n \\\"6\\\\n\\\",\\n \\\"-1\\\\n\\\",\\n \\\"2\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"-1\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"9\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"-1\\\\n\\\"\\n ]\\n}\"\"\",  # A correct sample # noqa: E501\n    \"\"\"{\\n \\\"inputs\\\": [\\n \\\"5 7 6 11\\\\n3\\\\n5 3 8\\\\n6 7 11\\\\n5 2 5\\\\n\\\",\\n \\\"3 4 3 10\\\\n3\\\\n3 1 4\\\\n4 5 9\\\\n3 10 10\\\\n\\\",\\n \\\"1 1 2 10\\\\n2\\\\n1 1 3\\\\n2 6 10\\\\n\\\",\\n \\\"9 8 7 8\\\\n9\\\\n10 6 6\\\\n10 6 6\\\\n7 7 8\\\\n9 5 6\\\\n8 9 9\\\\n9 5 5\\\\n9 8 8\\\\n8 5 6\\\\n9 10 10\\\\n\\\",\\n \\\"6 15 7 15\\\\n9\\\\n6 15 15\\\\n7 14 14\\\\n6 15 15\\\\n9 14 14\\\\n7 14 16\\\\n6 15 15\\\\n6 15 15\\\\n7 14 14\\\\n8 15 15\\\\n\\\",\\n \\\"13 16 20 10\\\\n18\\\\n13 16 16\\\\n20 10 10\\\\n19 10 10\\\\n12 15 15\\\\n20 10 10\\\\n18 11 11\\\\n19 10 10\\\\n19 10 10\\\\n20 10 10\\\\n19 10 10\\\\n20 10 10\\\\n20 10 10\\\\n19 10 10\\\\n18 11 11\\\\n13 16 16\\\\n12 15 15\\\\n19 10 10\\\\n19 10 10\\\\n\\\",\\n \\\"89 29 88 30\\\\n16\\\\n87 31 31\\\\n14 95 95\\\\n98 88 89\\\\n96 88 88\\\\n14 97 97\\\\n13 97 98\\\\n100 88 88\\\\n88 32 32\\\\n99 88 89\\\\n90 29 29\\\\n87 31 31\\\\n15 94 96\\\\n89 29 29\\\\n88 32 32\\\\n97 89 89\\\\n88 29 30\\\\n\\\",\\n \\\"30 14 39 19\\\\n31\\\\n35 7 11\\\\n37 11 12\\\\n32 13 13\\\\n37 5 6\\\\n46 13 13\\\\n37 14 14\\\\n31 13 13\\\\n43 13 19\\\\n45 15 19\\\\n46 13 13\\\\n32 17 17\\\\n41 14 19\\\\n30 14 14\\\\n43 13 17\\\\n34 16 18\\\\n44 11 19\\\\n38 13 13\\\\n40 12 20\\\\n37 16 18\\\\n46 16 18\\\\n34 10 14\\\\n36 9 10\\\\n36 15 19\\\\n38 15 19\\\\n42 13 19\\\\n33 14 15\\\\n35 15 19\\\\n33 17 18\\\\n39 12 20\\\\n36 5 7\\\\n45 12 12\\\\n\\\",\\n \\\"2 1 1 1\\\\n2\\\\n1 1 2\\\\n2 1 2\\\\n\\\",\\n \\\"1 1 1 2\\\\n5\\\\n1000000000 1 10000\\\\n19920401 1188 5566\\\\n1000000000 1 10000\\\\n1 1 10000\\\\n5 100 200\\\\n\\\",\\n \\\"1 1 1000000000 2\\\\n5\\\\n1000000000 1 10000\\\\n19920401 1188 5566\\\\n1000000000 1 10000\\\\n1 1 10000\\\\n5 100 200\\\\n\\\"\\n ],\\n \\\"outputs\\\": [\\n \\\"4\\\\n\\\",\\n \\\"6\\\\n\\\",\\n \\\"-1\\\\n\\\",\\n \\\"-1\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"-1\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"9\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"1\\\\n\\\",\\n \\\"-1\\\\n\\\"\\n ]\\n}\"\"\",  # noqa: E501\n]  # A failed sample with first several in-out passed\n\nprime_code_scores = [1.0, 0.9]\n\n\ndef test_parallelism():\n    \"\"\"\n    Test if process pool works properly\n    \"\"\"\n    sequences_str = []\n    ground_truth = []\n    data_sources = []\n    while len(sequences_str) < 32:\n        sequences_str.extend(prime_code_answers)\n        ground_truth.extend(prime_code_gts)\n        data_sources.extend([\"codecontests\"] * len(prime_code_answers))\n\n        sequences_str.extend(prime_math_answers)\n        ground_truth.extend(prime_math_gts)\n        data_sources.extend([\"numina_aops_forum\"] * len(prime_math_answers))\n\n    scores = asyncio.run(\n        parallel_compute_score_async(default_compute_score, sequences_str, ground_truth, data_sources, num_processes=16)\n    )\n    print(scores)\n\n\n@pytest.mark.skip(\"pyext not compatible with python 3.12\")\ndef test_prime_code():\n    \"\"\"\n    Test PRIME code sandbox.\n    \"\"\"\n    data_source = \"codecontests\"\n    for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True):\n        score = default_compute_score(data_source, completion, ground_truth)\n        assert float(score) == score_\n\n\n# Use the pytest.mark.skipif decorator to skip the test\n@pytest.mark.skipif(not os.environ.get(\"SANDBOX_FUSION_URL\"), reason=\"SANDBOX_FUSION_URL environment variable not set\")\ndef test_prime_code_sandbox_fusion():\n    \"\"\"\n    Test PRIME code on sandbox fusion. Skips if SANDBOX_FUSION_URL is not set.\n    \"\"\"\n    data_source = \"codecontests\"\n    # Get the URL from the environment variable, as skipif ensures it is set at this point\n    sandbox_fusion_url = os.environ.get(\"SANDBOX_FUSION_URL\")\n    # Removed the previous 'if not sandbox_url' check block\n\n    for completion, ground_truth, score_ in zip(prime_code_answers, prime_code_gts, prime_code_scores, strict=True):\n        score = default_compute_score(\n            data_source, completion, ground_truth, extra_info={\"sandbox_fusion_url\": sandbox_fusion_url}\n        )  # <-- Use the URL obtained from the environment variable\n        assert float(score) == score_\n\n\n@pytest.mark.skipif(not os.environ.get(\"SANDBOX_FUSION_URL\"), reason=\"SANDBOX_FUSION_URL environment variable not set\")\ndef test_continuous_score_consistency():\n    \"\"\"\n    Verify that continuous score calculation is consistent between prime_code and sandbox_fusion.\n    Uses a test case where the first 9 out of 11 sub-cases pass (expected score 0.9).\n    \"\"\"\n    from verl.utils.reward_score import prime_code\n\n    completion = prime_code_answers[1]  # Use the second sample\n    ground_truth = prime_code_gts[1]  # Use the second sample (9/11 pass, first 9 pass)\n    expected_continuous_score = 0.9\n\n    # 1. Calculate score using prime_code (default) with continuous=True\n    prime_score, _ = sandbox_fusion.compute_score(\n        os.environ.get(\"SANDBOX_FUSION_URL\"), None, completion, ground_truth, continuous=True\n    )\n\n    # 2. Calculate score using sandbox_fusion with continuous=True\n    # Ensure the extra_info key triggers the sandbox_fusion path in default_compute_score\n    fusion_score, _ = prime_code.compute_score(completion, ground_truth, continuous=True)\n\n    # 3. Assert scores are equal (using pytest.approx for float comparison)\n    assert float(prime_score) == pytest.approx(expected_continuous_score)\n    assert float(fusion_score) == pytest.approx(expected_continuous_score)\n    assert float(prime_score) == pytest.approx(float(fusion_score))\n    print(f\"Continuous Score (Prime Code): {prime_score}\")\n    print(f\"Continuous Score (Sandbox Fusion): {fusion_score}\")\n\n\n@pytest.mark.skip(\"pyext not compatible with python 3.12\")\ndef test_check_correctness():\n    from verl.utils.reward_score.prime_code import apps_check_correctness\n\n    completion = prime_code_answers[0]\n    ground_truth = json.loads(prime_code_gts[0])\n    ground_truth_single = {\"inputs\": ground_truth[\"inputs\"][:1], \"outputs\": ground_truth[\"outputs\"][:1]}\n    res, meta = apps_check_correctness(in_outs=ground_truth_single, generation=completion, timeout=5, debug=False)\n    print(res, meta)\n\n\ndef test_prime_math():\n    data_source = \"numina_aops_forum\"\n    for completion, ground_truth in zip(prime_math_answers, prime_math_gts, strict=True):\n        score = default_compute_score(data_source, completion, ground_truth)\n        assert float(score) == 1.0\n"
  },
  {
    "path": "tests/utils/test_activation_offload.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport shutil\nimport tempfile\n\nimport pytest\nimport torch\nimport torch.distributed\nimport torch.multiprocessing as mp\nfrom torch.distributed import init_device_mesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision, ShardingStrategy\nfrom transformers import AutoModelForCausalLM, AutoTokenizer, Qwen2Config\n\nfrom verl.utils.activation_offload import enable_activation_offloading\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.fsdp_utils import MixedPrecisionPolicy, apply_fsdp2, get_fsdp_wrap_policy\n\n\ndef create_random_input_ids(batch_size, seq_len, vocab_size):\n    if get_device_name() == \"cuda\":\n        from flash_attn.bert_padding import unpad_input\n    elif get_device_name() == \"npu\":\n        from verl.utils.attention_utils import unpad_input\n    from verl.utils.model import compute_position_id_with_mask, create_random_mask\n\n    input_ids = torch.randint(0, vocab_size, (batch_size, seq_len), device=get_device_name())\n\n    attention_mask = create_random_mask(\n        input_ids, max_ratio_of_left_padding=0.1, min_ratio_of_valid_token=0.5, max_ratio_of_valid_token=0.7\n    )\n    position_ids = compute_position_id_with_mask(attention_mask)\n\n    input_ids = unpad_input(input_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1)\n    position_ids = unpad_input(position_ids.unsqueeze(-1), attention_mask)[0].transpose(0, 1)\n    return input_ids, position_ids\n\n\ndef _fsdp_activation_offloading_test(rank, world_size, rendezvous_file, strategy=\"fsdp\"):\n    get_torch_device().set_device(rank)\n    torch.distributed.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n    device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=(\"dp\",))\n\n    model_name = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    config = Qwen2Config(num_hidden_layers=4)\n\n    with torch.device(get_device_name()):\n        model = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        model = model.to(device=get_device_name())\n\n    # Wrap model with FSDP\n    mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32)\n\n    if strategy == \"fsdp\":\n        model = FSDP(\n            model,\n            use_orig_params=False,\n            device_id=get_torch_device().current_device(),\n            sharding_strategy=ShardingStrategy.FULL_SHARD,\n            mixed_precision=mixed_precision,\n            device_mesh=device_mesh,\n            auto_wrap_policy=get_fsdp_wrap_policy(module=model),\n        )\n    else:\n        mp_policy = MixedPrecisionPolicy(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True\n        )\n        fsdp_kwargs = {\n            \"mesh\": device_mesh,\n            \"mp_policy\": mp_policy,\n        }\n        apply_fsdp2(model, fsdp_kwargs, {})\n\n    optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)\n    lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.9)\n\n    # Create checkpoint manager\n    tokenizer = AutoTokenizer.from_pretrained(model_name)\n    checkpoint_manager = FSDPCheckpointManager(\n        model=model, optimizer=optimizer, lr_scheduler=lr_scheduler, tokenizer=tokenizer\n    )\n\n    # Generate sample input\n    batch_size = 2\n    seq_len = 32\n    vocab_size = 32000\n    # First input for initial update\n    input_ids1, position_ids1 = create_random_input_ids(batch_size, seq_len, vocab_size)\n\n    # Second input for verification\n    input_ids2, position_ids2 = create_random_input_ids(batch_size, seq_len, vocab_size)\n\n    # Step 1: Initial update and save checkpoint\n    outputs1 = model(input_ids=input_ids1, position_ids=position_ids1)\n    loss1 = outputs1.logits.mean()\n    loss1.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Save checkpoint after first update\n    temp_dir = tempfile.mkdtemp()\n    checkpoint_path = os.path.join(temp_dir, \"checkpoint\")\n    checkpoint_manager.save_checkpoint(local_path=checkpoint_path, hdfs_path=None, global_step=0)\n\n    # Step 2: Second update and forward pass\n    outputs2 = model(input_ids=input_ids2, position_ids=position_ids2)\n    loss2 = outputs2.logits.mean()\n    loss2.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Record logits after second update\n    with torch.no_grad():\n        logits_without_offloading = model(input_ids=input_ids2, position_ids=position_ids2).logits\n\n    # Step 3: wrap module with activation offloading and load checkpoint\n    enable_activation_offloading(model, strategy=strategy)\n    checkpoint_manager.load_checkpoint(checkpoint_path)\n\n    # Step 4: Repeat the second update with same input\n    outputs3 = model(input_ids=input_ids2, position_ids=position_ids2)\n    loss3 = outputs3.logits.mean()\n    loss3.backward()\n    optimizer.step()\n    lr_scheduler.step()\n    optimizer.zero_grad()\n\n    # Record logits after loaded checkpoint and update\n    with torch.no_grad():\n        logits_with_offloading = model(input_ids=input_ids2, position_ids=position_ids2).logits\n\n    # Step 4: Verify outputs match\n    torch.testing.assert_close(logits_without_offloading, logits_with_offloading, atol=0.0, rtol=0.0)\n    print(f\"Activaiton offloading for {strategy} test passed on {world_size} GPUs!\")\n\n    # Cleanup\n    shutil.rmtree(temp_dir)\n    torch.distributed.barrier()\n    torch.distributed.destroy_process_group()\n\n\n@pytest.mark.parametrize(\"world_size\", (2, 4))\n@pytest.mark.parametrize(\"strategy\", (\"fsdp\", \"fsdp2\"))\ndef test_activation_offloading(world_size, strategy, tmp_path):\n    rendezvous_file = str(tmp_path / \"rdzv_file\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    mp.spawn(\n        fn=_fsdp_activation_offloading_test,\n        args=(world_size, rendezvous_file, strategy),\n        nprocs=world_size,\n        join=True,\n    )\n"
  },
  {
    "path": "tests/utils/test_bucketed_weight_transfer.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Tests for BucketedWeightSender and BucketedWeightReceiver.\n\nSender and receiver run in separate processes to match real-world usage\nand because CUDA IPC requires distinct processes.\n\"\"\"\n\nimport asyncio\nimport multiprocessing as mp\nimport uuid\n\nimport pytest\nimport torch\n\nfrom verl.utils.device import get_device_name, get_torch_device, is_support_ipc\n\nPROCESS_TIMEOUT = 60\n\n# Use string checks to avoid initializing CUDA in the main pytest process,\n# which would make subsequent fork-based multiprocessing in other tests unsafe.\nHAS_ACCELERATOR = get_device_name() != \"cpu\"\nHAS_CUDA = \"cuda\" in get_device_name()\n\n\ndef _unique_zmq_handle():\n    return f\"ipc:///tmp/test-bwt-{uuid.uuid4().hex}.sock\"\n\n\ndef _generate_weights(weight_specs, seed):\n    \"\"\"Deterministically generate weights on the best available device from specs.\n\n    Args:\n        weight_specs: list of (name, shape, dtype) tuples\n        seed: random seed for reproducibility\n    Returns:\n        list of (name, tensor_on_device) tuples\n    \"\"\"\n    device_name = get_device_name()\n    device = torch.device(f\"{device_name}:0\")\n    get_torch_device().manual_seed(seed)\n    weights = []\n    for name, shape, dtype in weight_specs:\n        # Generate in float32 then cast, since torch.randn doesn't support all dtypes\n        t = torch.randn(shape, dtype=torch.float32, device=device).to(dtype)\n        weights.append((name, t))\n    return weights\n\n\n# ---------------------------------------------------------------------------\n# Process entry points (must be module-level for pickling with spawn)\n# ---------------------------------------------------------------------------\ndef _sender_fn(zmq_handle, weight_specs, seed, bucket_size_mb, use_shm):\n    \"\"\"Sender process: generate weights, move to device, send.\"\"\"\n    from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightSender\n\n    weights = _generate_weights(weight_specs, seed)\n    sender = BucketedWeightSender(\n        zmq_handle=zmq_handle,\n        bucket_size_mb=bucket_size_mb,\n        use_shm=use_shm,\n    )\n    asyncio.run(sender.async_send_weights(iter(weights)))\n\n\ndef _receiver_fn(zmq_handle, use_shm, result_queue):\n    \"\"\"Receiver process: receive weights, send back (name, dtype, shape, checksum).\"\"\"\n    from verl.utils.device import get_device_name\n    from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightReceiver\n\n    device = torch.device(f\"{get_device_name()}:0\")\n    receiver = BucketedWeightReceiver(\n        zmq_handle=zmq_handle,\n        device=device,\n        use_shm=use_shm,\n    )\n    received = []\n    receiver.receive_weights(on_bucket_received=lambda w: received.extend(w))\n    # Only send lightweight metadata + checksum back through the queue\n    summaries = [(name, t.dtype, tuple(t.shape), t.float().sum().item()) for name, t in received]\n    result_queue.put(summaries)\n\n\n# ---------------------------------------------------------------------------\n# Test helper\n# ---------------------------------------------------------------------------\ndef _transfer_and_validate(weight_specs, bucket_size_mb, use_shm):\n    \"\"\"Spawn sender + receiver processes, then validate received tensors.\"\"\"\n    zmq_handle = _unique_zmq_handle()\n    seed = 42\n    ctx = mp.get_context(\"spawn\")\n    result_queue = ctx.Queue()\n\n    sender_p = ctx.Process(\n        target=_sender_fn,\n        args=(zmq_handle, weight_specs, seed, bucket_size_mb, use_shm),\n    )\n    receiver_p = ctx.Process(\n        target=_receiver_fn,\n        args=(zmq_handle, use_shm, result_queue),\n    )\n\n    # Start sender first (it binds), then receiver (it connects)\n    sender_p.start()\n    receiver_p.start()\n\n    sender_p.join(timeout=PROCESS_TIMEOUT)\n    receiver_p.join(timeout=PROCESS_TIMEOUT)\n\n    assert sender_p.exitcode == 0, f\"Sender process failed with exit code {sender_p.exitcode}\"\n    assert receiver_p.exitcode == 0, f\"Receiver process failed with exit code {receiver_p.exitcode}\"\n\n    summaries = result_queue.get(timeout=5)\n\n    # Regenerate expected weights on device with the same seed\n    expected = _generate_weights(weight_specs, seed)\n\n    assert len(summaries) == len(expected), f\"Expected {len(expected)} weights, got {len(summaries)}\"\n\n    for (exp_name, exp_tensor), (recv_name, recv_dtype, recv_shape, recv_cksum) in zip(\n        expected, summaries, strict=False\n    ):\n        assert exp_name == recv_name, f\"Name mismatch: expected {exp_name}, got {recv_name}\"\n        assert tuple(exp_tensor.shape) == recv_shape, (\n            f\"Shape mismatch for {exp_name}: expected {tuple(exp_tensor.shape)}, got {recv_shape}\"\n        )\n        assert exp_tensor.dtype == recv_dtype, (\n            f\"Dtype mismatch for {exp_name}: expected {exp_tensor.dtype}, got {recv_dtype}\"\n        )\n        exp_sum = exp_tensor.float().sum().item()\n        assert exp_sum == recv_cksum, f\"Data mismatch for {exp_name}\"\n\n\n# ---------------------------------------------------------------------------\n# Shared memory tests\n# ---------------------------------------------------------------------------\n@pytest.mark.skipif(not (HAS_ACCELERATOR and not HAS_CUDA), reason=\"Requires (shm only tested)\")\nclass TestBucketedWeightTransferSHM:\n    \"\"\"Test BucketedWeightSender/Receiver via shared memory path.\"\"\"\n\n    def test_single_small_weight(self):\n        specs = [(\"layer.weight\", (32, 16), torch.float32)]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=True)\n\n    def test_multiple_weights_single_bucket(self):\n        specs = [\n            (\"layer0.weight\", (16, 16), torch.float32),\n            (\"layer0.bias\", (16,), torch.float32),\n            (\"layer1.weight\", (16, 8), torch.bfloat16),\n        ]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=True)\n\n    def test_multiple_buckets(self):\n        # ~64 KB each x 20 = ~1.25 MB, bucket = 1 MB => spans 2 buckets\n        specs = [(f\"layer{i}.weight\", (128, 128), torch.float32) for i in range(20)]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=True)\n\n    def test_mixed_dtypes(self):\n        specs = [\n            (\"fp32_param\", (64, 64), torch.float32),\n            (\"bf16_param\", (64, 64), torch.bfloat16),\n            (\"fp16_param\", (32, 32), torch.float16),\n        ]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=True)\n\n    def test_empty_weights(self):\n        _transfer_and_validate([], bucket_size_mb=1, use_shm=True)\n\n\n# ---------------------------------------------------------------------------\n# CUDA IPC tests (CUDA only — IPC is not supported on NPU)\n# ---------------------------------------------------------------------------\n@pytest.mark.skipif(not is_support_ipc(), reason=\"Requires IPC support\")\nclass TestBucketedWeightTransferIPC:\n    \"\"\"Test BucketedWeightSender/Receiver via CUDA IPC path.\"\"\"\n\n    def test_single_small_weight(self):\n        specs = [(\"layer.weight\", (32, 16), torch.float32)]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=False)\n\n    def test_multiple_weights_single_bucket(self):\n        specs = [\n            (\"layer0.weight\", (16, 16), torch.float32),\n            (\"layer0.bias\", (16,), torch.float32),\n            (\"layer1.weight\", (16, 8), torch.bfloat16),\n        ]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=False)\n\n    def test_multiple_buckets(self):\n        specs = [(f\"layer{i}.weight\", (128, 128), torch.float32) for i in range(20)]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=False)\n\n    def test_mixed_dtypes(self):\n        specs = [\n            (\"fp32_param\", (64, 64), torch.float32),\n            (\"bf16_param\", (64, 64), torch.bfloat16),\n            (\"fp16_param\", (32, 32), torch.float16),\n        ]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=False)\n\n    def test_empty_weights(self):\n        _transfer_and_validate([], bucket_size_mb=1, use_shm=False)\n\n    def test_exact_bucket_boundary(self):\n        # 1 MB bucket = 1048576 bytes; float32 = 4 bytes => 262144 elements\n        numel = (1 << 20) // 4\n        specs = [(\"exact_fit\", (numel,), torch.float32)]\n        _transfer_and_validate(specs, bucket_size_mb=1, use_shm=False)\n"
  },
  {
    "path": "tests/utils/test_check_ipc_version_support_on_npu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport unittest\nfrom unittest.mock import Mock, mock_open, patch\n\nfrom verl.utils.device import check_ipc_version_support, get_npu_versions\n\n\nclass TestCheckIPCVersionSupport(unittest.TestCase):\n    \"\"\"Test cases for the check_ipc_version_support function.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up test logging to suppress INFO messages.\"\"\"\n        # Suppress INFO log messages during testing\n        logging.disable(logging.INFO)\n\n    def tearDown(self):\n        \"\"\"Restore logging.\"\"\"\n        logging.disable(logging.NOTSET)\n\n    def test_standard_version_with_support(self):\n        \"\"\"Test standard version that meets minimum requirements.\"\"\"\n        # Software 25.5.0 >= 25.3.rc1, CANN 8.3.0 >= 8.3.rc1\n        result = check_ipc_version_support(\"25.5.0\", \"8.3.0\")\n        self.assertTrue(result)\n\n    def test_standard_version_newer(self):\n        \"\"\"Test newer standard versions.\"\"\"\n        # Software 26.0.0 >= 25.3.rc1, CANN 9.0.0 >= 8.3.rc1\n        result = check_ipc_version_support(\"26.0.0\", \"9.0.0\")\n        self.assertTrue(result)\n\n    def test_rc_version_format(self):\n        \"\"\"Test RC version format with additional parts.\"\"\"\n        # Software 25.3.rc1.2 -> 25.3.rc1 >= 25.3.rc1\n        # CANN 8.3.rc1.2 -> 8.3.rc1 >= 8.3.rc1\n        result = check_ipc_version_support(\"25.3.rc1.2\", \"8.3.rc1.2\")\n        self.assertTrue(result)\n\n    def test_exact_rc_version(self):\n        \"\"\"Test exact RC version.\"\"\"\n        # Software 25.3.rc1 >= 25.3.rc1\n        # CANN 8.3.rc1 >= 8.3.rc1\n        result = check_ipc_version_support(\"25.3.rc1\", \"8.3.rc1\")\n        self.assertTrue(result)\n\n    def test_t_suffix_version(self):\n        \"\"\"Test version with lowercase t suffix.\"\"\"\n        # Software 25.5.t3.b001 -> 25.5 >= 25.3.rc1\n        # CANN 8.3.rc1 >= 8.3.rc1\n        result = check_ipc_version_support(\"25.5.t3.b001\", \"8.3.rc1\")\n        self.assertTrue(result)\n\n    def test_t_suffix_version_older(self):\n        \"\"\"Test version with lowercase t suffix that's too old.\"\"\"\n        # Software 25.5.t3.b001 -> 25.5 >= 25.3.rc1 (should pass)\n        # CANN 8.2.rc1 < 8.3.rc1 (should fail)\n        result = check_ipc_version_support(\"25.5.t3.b001\", \"8.2.rc1\")\n        self.assertFalse(result)\n\n    def test_software_version_below_minimum(self):\n        \"\"\"Test software version below minimum requirement.\"\"\"\n        # Software 25.2.0 < 25.3.rc1\n        result = check_ipc_version_support(\"25.2.0\", \"8.3.0\")\n        self.assertFalse(result)\n\n    def test_cann_version_below_minimum(self):\n        \"\"\"Test CANN version below minimum requirement.\"\"\"\n        # Software 25.5.0 >= 25.3.rc1\n        # CANN 8.2.0 < 8.3.rc1\n        result = check_ipc_version_support(\"25.5.0\", \"8.2.0\")\n        self.assertFalse(result)\n\n    def test_both_versions_below_minimum(self):\n        \"\"\"Test both versions below minimum requirement.\"\"\"\n        # Software 25.2.0 < 25.3.rc1\n        # CANN 8.2.0 < 8.3.rc1\n        result = check_ipc_version_support(\"25.2.0\", \"8.2.0\")\n        self.assertFalse(result)\n\n    def test_invalid_software_version(self):\n        \"\"\"Test invalid software version format.\"\"\"\n        with self.assertRaises(RuntimeError) as context:\n            check_ipc_version_support(\"invalid.version\", \"8.3.0\")\n        self.assertIn(\"Invalid software version format\", str(context.exception))\n\n    def test_invalid_cann_version(self):\n        \"\"\"Test invalid CANN version format.\"\"\"\n        with self.assertRaises(RuntimeError) as context:\n            check_ipc_version_support(\"25.5.0\", \"invalid.version\")\n        self.assertIn(\"Invalid CANN version format\", str(context.exception))\n\n    def test_rc_with_more_parts(self):\n        \"\"\"Test RC version with more than 3 parts.\"\"\"\n        # Should extract only first 3 parts: 25.3.rc1\n        result = check_ipc_version_support(\"25.3.rc1.2.3.4\", \"8.3.rc1.2.3.4\")\n        self.assertTrue(result)\n\n    def test_standard_with_more_parts(self):\n        \"\"\"Test standard version with more than 3 parts.\"\"\"\n        # Should extract only first 3 parts: 25.5.0\n        result = check_ipc_version_support(\"25.5.0.1.2.3\", \"8.3.0.1.2.3\")\n        self.assertTrue(result)\n\n    def test_rc_edge_case_versions(self):\n        \"\"\"Test edge case RC versions.\"\"\"\n        # RC1 is the minimum\n        result = check_ipc_version_support(\"25.3.rc1\", \"8.3.rc1\")\n        self.assertTrue(result)\n\n        # RC0 should fail\n        result = check_ipc_version_support(\"25.3.rc0\", \"8.3.rc1\")\n        self.assertFalse(result)\n\n    def test_major_version_differences(self):\n        \"\"\"Test major version number differences.\"\"\"\n        # Much newer major versions\n        result = check_ipc_version_support(\"30.0.0\", \"10.0.0\")\n        self.assertTrue(result)\n\n        # Older major versions\n        result = check_ipc_version_support(\"24.0.0\", \"7.0.0\")\n        self.assertFalse(result)\n\n\nclass TestGetNPUVersions(unittest.TestCase):\n    \"\"\"Test cases for the get_npu_versions function.\"\"\"\n\n    @patch(\"subprocess.run\")\n    @patch(\"platform.machine\")\n    @patch(\"os.path.exists\")\n    @patch(\"builtins.open\", new_callable=mock_open, read_data=\"version=8.3.rc1\\n\")\n    def test_get_npu_versions_success(self, mock_file, mock_exists, mock_machine, mock_run):\n        \"\"\"Test successful retrieval of versions.\"\"\"\n        # Mock npu-smi output\n        mock_run.return_value = Mock(stdout=\"Software Version : 25.5.0\\nOther Info\\n\", check=True)\n\n        # Mock architecture\n        mock_machine.return_value = \"x86_64\"\n\n        # Mock path exists\n        mock_exists.return_value = True\n\n        software_version, cann_version = get_npu_versions()\n\n        self.assertEqual(software_version, \"25.5.0\")\n        self.assertEqual(cann_version, \"8.3.rc1\")\n\n    @patch(\"subprocess.run\")\n    def test_get_npu_versions_missing_software_version(self, mock_run):\n        \"\"\"Test error when Software Version is missing.\"\"\"\n        mock_run.return_value = Mock(stdout=\"Other Info Without Software Version\\n\", check=True)\n\n        with self.assertRaises(RuntimeError) as context:\n            get_npu_versions()\n\n        self.assertIn(\"Could not find Software Version\", str(context.exception))\n\n    @patch(\"subprocess.run\")\n    @patch(\"platform.machine\")\n    @patch(\"os.path.exists\")\n    @patch(\"builtins.open\", new_callable=mock_open, read_data=\"version=8.3.rc1\\n\")\n    def test_get_npu_versions_unsupported_architecture(self, mock_file, mock_exists, mock_machine, mock_run):\n        \"\"\"Test error with unsupported architecture.\"\"\"\n        mock_run.return_value = Mock(stdout=\"Software Version : 25.5.0\\n\", check=True)\n\n        mock_machine.return_value = \"armv7l\"  # Unsupported architecture\n        mock_exists.return_value = True\n\n        with self.assertRaises(RuntimeError) as context:\n            get_npu_versions()\n\n        self.assertIn(\"Unsupported architecture\", str(context.exception))\n\n    @patch(\"subprocess.run\")\n    @patch(\"platform.machine\")\n    @patch(\"os.path.exists\")\n    @patch(\"builtins.open\", new_callable=mock_open, read_data=\"version=8.3.rc1\\n\")\n    def test_get_npu_versions_cann_path_not_exists(self, mock_file, mock_exists, mock_machine, mock_run):\n        \"\"\"Test error when CANN path doesn't exist.\"\"\"\n        mock_run.return_value = Mock(stdout=\"Software Version : 25.5.0\\n\", check=True)\n\n        mock_machine.return_value = \"x86_64\"\n        mock_exists.return_value = False  # Path doesn't exist\n\n        with self.assertRaises(RuntimeError) as context:\n            get_npu_versions()\n\n        self.assertIn(\"CANN toolkit path does not exist\", str(context.exception))\n\n    @patch(\"subprocess.run\")\n    @patch(\"platform.machine\")\n    @patch(\"os.path.exists\")\n    @patch(\"builtins.open\")\n    def test_get_npu_versions_info_file_not_exists(self, mock_file, mock_exists, mock_machine, mock_run):\n        \"\"\"Test error when CANN info file doesn't exist.\"\"\"\n        mock_run.return_value = Mock(stdout=\"Software Version : 25.5.0\\n\", check=True)\n\n        mock_machine.return_value = \"x86_64\"\n\n        # First call is for CANN path exists, second call is for info file exists\n        mock_exists.side_effect = [True, False]\n\n        with self.assertRaises(RuntimeError) as context:\n            get_npu_versions()\n\n        self.assertIn(\"CANN toolkit info file does not exist\", str(context.exception))\n\n    @patch(\"subprocess.run\")\n    @patch(\"platform.machine\")\n    @patch(\"os.path.exists\")\n    @patch(\"builtins.open\", new_callable=mock_open, read_data=\"other_info=no_version\\n\")\n    def test_get_npu_versions_missing_cann_version(self, mock_file, mock_exists, mock_machine, mock_run):\n        \"\"\"Test error when CANN version is missing from info file.\"\"\"\n        mock_run.return_value = Mock(stdout=\"Software Version : 25.5.0\\n\", check=True)\n\n        mock_machine.return_value = \"x86_64\"\n        mock_exists.return_value = True\n\n        with self.assertRaises(RuntimeError) as context:\n            get_npu_versions()\n\n        self.assertIn(\"Could not find version in CANN toolkit info file\", str(context.exception))\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_check_profiler_output.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport glob\nimport logging\nimport os\nimport sys\nfrom dataclasses import dataclass\nfrom typing import Callable\n\n# Initialize logger\nlogger = logging.getLogger(__file__)\nlogging.basicConfig(level=logging.INFO, format=\"%(asctime)s - %(name)s - %(levelname)s - %(message)s\")\n\n\n@dataclass\nclass DeviceCheckConfig:\n    \"\"\"Device check configuration: encapsulates device-specific validation rules\"\"\"\n\n    # Search path pattern\n    search_pattern: str\n    # Directory count validation function: takes stage and dir list, returns bool\n    dir_count_validator: Callable[[str, list[str]], bool]\n    # PROF file/dir validation function: takes directory path, returns bool\n    prof_validator: Callable[[str], bool]\n\n\nclass ProfilerChecker:\n    \"\"\"Unified Profiler checker supporting GPU/NPU devices\"\"\"\n\n    TARGET_STAGES = [\"actor_update\", \"*_rollout_*\", \"ref_*\"]\n\n    def __init__(self, device_type: str, profiler_dir: str):\n        self.device_type = device_type.lower()\n        self.profiler_dir = profiler_dir\n\n        # Validate device type\n        if self.device_type not in [\"gpu\", \"npu\"]:\n            raise ValueError(f\"Unsupported device type: {device_type}, only gpu/npu are supported\")\n\n        # Initialize device-specific configuration\n        self._init_device_config()\n\n    def _init_device_config(self):\n        \"\"\"Initialize validation rules for different devices (core: device differences as config)\"\"\"\n        if self.device_type == \"gpu\":\n            self.config = DeviceCheckConfig(\n                # GPU search pattern: match stage directory directly\n                search_pattern=os.path.join(self.profiler_dir, \"{stage}\"),\n                # GPU: all stages must have exactly 1 directory\n                dir_count_validator=lambda stage, dirs: len(dirs) == 1,\n                # GPU: any file/subdirectory exists under the directory\n                prof_validator=lambda d: len(glob.glob(os.path.join(d, \"*\"))) > 0,\n            )\n        else:  # NPU\n            self.config = DeviceCheckConfig(\n                # NPU search pattern: match ascend subdirectory under stage\n                search_pattern=os.path.join(self.profiler_dir, \"{stage}\", \"*_ascend_*\"),\n                # NPU: rollout requires >1 dir, others require exactly 1 dir\n                dir_count_validator=lambda stage, dirs: (len(dirs) > 1 if stage == \"*_rollout_*\" else len(dirs) == 1),\n                # NPU: PROF_* subdirectory must exist and be a valid directory\n                prof_validator=lambda d: (\n                    len(glob.glob(os.path.join(d, \"PROF_*\"))) > 0\n                    and os.path.isdir(glob.glob(os.path.join(d, \"PROF_*\"))[0])\n                ),\n            )\n\n    def _validate_stage_dirs(self, stage: str) -> bool:\n        \"\"\"Generic stage directory validation: extracted common logic for GPU/NPU\"\"\"\n        # 1. Generate search path and match directories\n        search_pattern = self.config.search_pattern.format(stage=stage)\n        dirs = glob.glob(search_pattern, recursive=True)\n\n        # 2. Log found directories\n        for d in dirs:\n            logger.info(f\"[{stage}] Found: {d}\")\n\n        # 3. Validate directory count\n        if not self.config.dir_count_validator(stage, dirs):\n            expected = \">1\" if stage == \"*_rollout_*\" and self.device_type == \"npu\" else 1\n            logger.error(f\"[{stage}] Expected {expected} directories, found {len(dirs)}\")\n            return False\n\n        # 4. Validate PROF files/directories\n        for target_dir in dirs:\n            if not self.config.prof_validator(target_dir):\n                logger.error(f\"[{stage}] PROF not found in {target_dir}\")\n                return False\n\n        return True\n\n    def check(self) -> bool:\n        \"\"\"Unified check entry point\"\"\"\n        logger.info(f\"Starting profiler deliverables check for {self.device_type.upper()}...\")\n\n        # Validate root directory exists\n        if not os.path.exists(self.profiler_dir):\n            logger.error(f\"Profiler data directory not found: {self.profiler_dir}\")\n            return False\n\n        # Run validation for all target stages\n        for stage in self.TARGET_STAGES:\n            if not self._validate_stage_dirs(stage):\n                return False\n\n        logger.info(f\"All {self.device_type.upper()} validation stages passed\")\n        return True\n\n\ndef parse_args():\n    \"\"\"Parse command line arguments\"\"\"\n    parser = argparse.ArgumentParser(description=\"Check Profiler deliverables (support GPU/NPU)\")\n    parser.add_argument(\n        \"--device\",\n        type=str,\n        required=True,\n        choices=[\"gpu\", \"npu\"],\n        help=\"Device type, available values: gpu/npu (required)\",\n    )\n    parser.add_argument(\n        \"--profiler_dir\",\n        type=str,\n        default=\"./profiler_data\",\n        help=\"Path to profiler data directory (default: ./profiler_data)\",\n    )\n    return parser.parse_args()\n\n\ndef main():\n    args = parse_args()\n\n    try:\n        checker = ProfilerChecker(device_type=args.device, profiler_dir=args.profiler_dir)\n        if checker.check():\n            logger.info(f\"All {args.device.upper()} profiler deliverables check passed!\")\n            sys.exit(0)\n        else:\n            logger.error(f\"{args.device.upper()} profiler check failed!\")\n            sys.exit(1)\n\n    except Exception as e:\n        logger.exception(f\"Check failed with error: {str(e)}\")\n        sys.exit(1)\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/utils/test_config_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom dataclasses import dataclass, field\n\nfrom omegaconf import OmegaConf\n\nfrom verl.base_config import BaseConfig\nfrom verl.utils import omega_conf_to_dataclass\n\n\n@dataclass\nclass TestDataclass(BaseConfig):\n    hidden_size: int = 0\n    activation: str = \"relu\"\n\n\n@dataclass\nclass TestTrainConfig(BaseConfig):\n    batch_size: int = 0\n    model: TestDataclass = field(default_factory=TestDataclass)\n    override_config: dict = field(default_factory=dict)\n\n\n_cfg_str = \"\"\"train_config:\n  _target_: tests.utils.test_config_on_cpu.TestTrainConfig\n  batch_size: 32\n  model:\n    hidden_size: 768\n    activation: relu\n  override_config: {}\"\"\"\n\n\nclass TestConfigOnCPU(unittest.TestCase):\n    \"\"\"Test cases for configuration utilities on CPU.\n\n    Test Plan:\n    1. Test basic OmegaConf to dataclass conversion for simple nested structures\n    2. Test nested OmegaConf to dataclass conversion for complex hierarchical configurations\n    3. Verify all configuration values are correctly converted and accessible\n    \"\"\"\n\n    def setUp(self):\n        self.config = OmegaConf.create(_cfg_str)\n\n    def test_omega_conf_to_dataclass(self):\n        sub_cfg = self.config.train_config.model\n        cfg = omega_conf_to_dataclass(sub_cfg, TestDataclass)\n        self.assertEqual(cfg.hidden_size, 768)\n        self.assertEqual(cfg.activation, \"relu\")\n        assert isinstance(cfg, TestDataclass)\n\n    def test_nested_omega_conf_to_dataclass(self):\n        cfg = omega_conf_to_dataclass(self.config.train_config, TestTrainConfig)\n        self.assertEqual(cfg.batch_size, 32)\n        self.assertEqual(cfg.model.hidden_size, 768)\n        self.assertEqual(cfg.model.activation, \"relu\")\n        assert isinstance(cfg, TestTrainConfig)\n        assert isinstance(cfg.model, TestDataclass)\n\n\nclass TestPrintCfgCommand(unittest.TestCase):\n    \"\"\"Test suite for the print_cfg.py command-line tool.\"\"\"\n\n    def test_command_with_override(self):\n        \"\"\"Test that the command runs without error when overriding config values.\"\"\"\n        import subprocess\n\n        # Run the command\n        result = subprocess.run(\n            [\"python3\", \"scripts/print_cfg.py\"],\n            capture_output=True,\n            text=True,\n        )\n\n        # Verify the command exited successfully\n        self.assertEqual(result.returncode, 0, f\"Command failed with stderr: {result.stderr}\")\n\n        # Verify the output contains expected config information\n        self.assertIn(\"critic\", result.stdout)\n        self.assertIn(\"profiler\", result.stdout)\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_flops_counter.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport math\n\nimport pytest\n\nfrom verl.utils.flops_counter import FlopsCounter\n\nVALID_CONFIG_TYPE = {\"llama\", \"qwen2\", \"qwen3\", \"qwen3_moe\", \"deepseek_v3\", \"mistral\", \"gemma3_text\", \"apertus\"}\n\n\nclass Config:\n    def __init__(self, config_dict):\n        for key, value in config_dict.items():\n            if isinstance(value, dict):\n                value = Config(value)\n            setattr(self, key, value)\n\n\nCONFIG = {\n    \"llama\": {\n        \"config\": {  # llama2-7B\n            \"model_type\": \"llama\",\n            \"vocab_size\": 32000,\n            \"hidden_size\": 4096,\n            \"intermediate_size\": 11008,\n            \"num_hidden_layers\": 32,\n            \"num_attention_heads\": 32,\n            \"num_key_value_heads\": 32,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # 6*(vocab*hidden*2+layer*(hidden*(q+k+v+head*head_dim)+ hidden*inter*3))*token_sum +\n        # 6*sum(seqlen^2)*layer*head*head_dim\n        # 6*(32000*4096*2+32*(4096*4096*4+4096*11008*3))*(512+1024+2048) +\n        # 6*(512*512+1024*1024+2048*2048)*32*4096\n        # 6*(32000*4096*2+32*(4096*4096*4+4096*11008*3))*(4096+4096+4096) +\n        # 6*(4096*4096+4096*4096+4096*4096)*32*4096\n        \"expected_flops_tuple\": (149226491215872 / 1e12, 536372695793664 / 1e12),\n    },\n    \"qwen2\": {\n        \"config\": {  # Qwen/Qwen2.5-7B-Instruct\n            \"model_type\": \"qwen2\",\n            \"vocab_size\": 152064,\n            \"hidden_size\": 3584,\n            \"intermediate_size\": 18944,\n            \"num_hidden_layers\": 28,\n            \"num_attention_heads\": 28,\n            \"num_key_value_heads\": 4,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # 6*(vocab*hidden*2+layer*(hidden*(q+k+v+head*head_dim)+ hidden*inter*3))*token_sum +\n        # 6*sum(seqlen^2)*layer*head*head_dim\n        # 6*(152064*3584*2+28*(3584*(3584+512+512+3584)+3584*18944*3))*(512+1024+2048) +\n        # 6*(512*512+1024*1024+2048*2048)*28*3584\n        # 6*(152064*3584*2+28*(3584*(3584+512+512+3584)+3584*18944*3))*(4096+4096+4096) +\n        # 6*(4096*4096+4096*4096+4096*4096)*28*3584\n        \"expected_flops_tuple\": (167073690943488 / 1e12, 591764889010176 / 1e12),\n    },\n    \"qwen3\": {\n        \"config\": {  # Qwen/Qwen3-8B\n            \"model_type\": \"qwen3\",\n            \"vocab_size\": 151936,\n            \"hidden_size\": 4096,\n            \"intermediate_size\": 12288,\n            \"num_hidden_layers\": 36,\n            \"num_attention_heads\": 32,\n            \"num_key_value_heads\": 8,\n            \"head_dim\": 128,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # 6*(vocab*hidden*2+layer*(hidden*(q+k+v+head*head_dim)+ hidden*inter*3))*token_sum +\n        # 6*sum(seqlen^2)*layer*head*head_dim\n        # 6*(151936*4096*2+36*(4096*(128*32+128*8*2+128*32)+4096*12288*3))*(512+1024+2048) +\n        # 6*(512*512+1024*1024+2048*2048)*36*128*32\n        # 6*(151936*4096*2+36*(4096*(128*32+128*8*2+128*32)+4096*12288*3))*(4096+4096+4096) +\n        # 6*(4096*4096+4096*4096+4096*4096)*36*128*32\n        \"expected_flops_tuple\": (180997438046208 / 1e12, 648394032807936 / 1e12),\n    },\n    \"qwen3_moe\": {\n        \"config\": {  # Qwen/Qwen3-30B-A3B-Base\n            \"model_type\": \"qwen3_moe\",\n            \"hidden_size\": 2048,\n            \"vocab_size\": 151936,\n            \"num_hidden_layers\": 48,\n            \"num_key_value_heads\": 4,\n            \"num_attention_heads\": 32,\n            \"head_dim\": 128,\n            \"moe_intermediate_size\": 768,\n            \"num_experts_per_tok\": 8,\n            \"num_experts\": 128,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # 6*(vocab*hidden*2+layer*(hidden*(q+k+v+head*head_dim)+hidden*inter*top_k_exp*3 +\n        # hidden*num_experts))*token_sum + 6*sum(seqlen^2)*layer*head*head_dim\n        # 6*(151936*2048*2+48*(2048*(128*32+128*4*2+128*32)+2048*768*8*3+2048*128))*(512+1024+2048) +\n        # 6*(512*512+1024*1024+2048*2048)*48*128*32\n        # 6*(151936*2048*2+48*(2048*(128*32+128*4*2+128*32)+2048*768*8*3+2048*128))*(4096+4096+4096) +\n        # 6*(4096*4096+4096*4096+4096*4096)*48*128*32\n        \"expected_flops_tuple\": (78593069678592 / 1e12, 306570470621184 / 1e12),\n    },\n    \"deepseek_v3\": {\n        \"config\": {  # deepseek-ai/DeepSeek-Prover-V2-671B\n            \"model_type\": \"deepseek_v3\",\n            \"hidden_size\": 7168,\n            \"vocab_size\": 129280,\n            \"moe_intermediate_size\": 2048,\n            \"num_hidden_layers\": 61,\n            \"first_k_dense_replace\": 3,\n            \"num_attention_heads\": 128,\n            \"n_routed_experts\": 256,\n            \"num_experts_per_tok\": 8,\n            \"n_shared_experts\": 1,\n            \"kv_lora_rank\": 512,\n            \"qk_rope_head_dim\": 64,\n            \"v_head_dim\": 128,\n            \"intermediate_size\": 18432,\n            \"qk_nope_head_dim\": 128,\n            \"q_lora_rank\": 1536,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # (1536*7168+128*192*1536+7168*(512+64)+128*(128+128)*512+128*128*7168) = 187105280\n        # 6*(129280*7168*2+ 3*(7168*18432*3+187105280)+ 58*(187105280+7168*256+7168*2048*9*3))*(512+1024+2048) +\n        # 3*(512*512+1024*1024+2048*2048)*61*(192+128)*128\n        # 6*(129280*7168*2+ 3*(7168*18432*3+187105280)+ 58*(187105280+7168*256+7168*2048*9*3))*(4096+4096+4096) +\n        # 3*(4096*4096+4096*4096+4096*4096)*61*(192+128)*128\n        \"expected_flops_tuple\": (848766538088448 / 1e12, 3145850406567936 / 1e12),\n    },\n    \"mistral\": {\n        \"config\": {  # mistralai/Mistral-Small-24B-Instruct-2501\n            \"model_type\": \"mistral\",\n            \"vocab_size\": 131072,\n            \"hidden_size\": 5120,\n            \"intermediate_size\": 32768,\n            \"num_hidden_layers\": 40,\n            \"num_attention_heads\": 32,\n            \"num_key_value_heads\": 8,\n            \"head_dim\": 128,\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # Mistral uses same architecture as Llama, with GQA\n        # 6*(vocab*hidden*2+layer*(hidden*(q+k+v+head*head_dim)+ hidden*inter*3))*token_sum +\n        # 12*sum(seqlen^2)*layer*head*head_dim\n        # vocab part: 131072*5120*2 = 1342177280\n        # attn part per layer: 5120*(128*32+128*8+128*8+128*32) = 5120*10240 = 52428800\n        # mlp part per layer: 5120*32768*3 = 503316480\n        # total per layer: 52428800 + 503316480 = 555745280\n        # all layers: 1342177280 + 40*555745280 = 23571988480\n        # For batch [512, 1024, 2048], tokens_sum = 3584:\n        # dense flops: 6 * 23571988480 * 3584 = 506892040273920\n        # attn flops: 6 * 5505024 * 40 * 128 * 32 = 10823317585920\n        # total: 517715357859840 / 1e12 = 517.71535785984\n        # For batch [4096, 4096, 4096], tokens_sum = 12288:\n        # dense flops: 6 * 23571988480 * 12288 = 1737915566653440\n        # attn flops: 6 * 50331648 * 40 * 128 * 32 = 98956046499840\n        # total: 1836871613153280 / 1e12 = 1836.87161315328\n        \"expected_flops_tuple\": (512303699066880 / 1e12, 1787393589903360 / 1e12),\n    },\n    \"gemma3_text\": {\n        \"config\": {  # Gemma3-12B-IT-TextOnly\n            \"model_type\": \"gemma3_text\",\n            \"vocab_size\": 262208,\n            \"hidden_size\": 3840,\n            \"intermediate_size\": 15360,\n            \"num_hidden_layers\": 48,\n            \"num_attention_heads\": 16,\n            \"num_key_value_heads\": 8,\n            \"head_dim\": 256,\n            \"sliding_window\": 1024,\n            \"layer_types\": None,\n            # Will be auto-generated based on sliding_window_pattern\n            \"sliding_window_pattern\": 6,\n            # Every 6th layer is full attention\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # Gemma3 has alternating sliding window attention\n        # With sliding_window_pattern=6: layers 5,11,17,23,29,35,41,47 use full attention (8 layers)\n        # Other 40 layers use sliding window attention with window_size=1024\n        #\n        # Non-attention FLOPs:\n        # vocab part: 262208*3840*2 = 2013757440\n        # attn part per layer: 3840*(256*16+256*8+256*8+256*16) = 3840*12288 = 47185920\n        # mlp part per layer: 3840*15360*3 = 176947200\n        # total per layer: 47185920 + 176947200 = 224133120\n        # all layers: 2013757440 + 48*224133120 = 12772147200\n        #\n        # For batch [512, 1024, 2048], tokens_sum = 3584:\n        # dense flops: 6 * 12772147200 * 3584 = 274652253388800\n        # seqlen_square_sum: 180355072 (calculated with sliding window logic)\n        # attn flops: 6 * 180355072 * 256 * 16 = 8864812498944\n        # total: 283517065887744 / 1e12 = 283.517065887744\n        #\n        # For batch [4096, 4096, 4096], tokens_sum = 12288:\n        # dense flops: 6 * 12772147200 * 12288 = 941664868761600\n        # seqlen_square_sum: 905969664 (calculated with sliding window logic)\n        # attn flops: 6 * 905969664 * 256 * 16 = 44530220924928\n        # total: 986195089686528 / 1e12 = 986.195089686528\n        \"expected_flops_tuple\": (279084659638272 / 1e12, 963929979224064 / 1e12),\n    },\n    \"gpt_oss\": {\n        \"config\": {\n            \"model_type\": \"gpt_oss\",\n            \"vocab_size\": 201088,\n            \"hidden_size\": 2880,\n            \"num_hidden_layers\": 24,\n            \"num_attention_heads\": 64,\n            \"num_key_value_heads\": 8,\n            \"head_dim\": 64,\n            \"intermediate_size\": 2880,\n            \"num_local_experts\": 32,\n            \"num_experts_per_tok\": 4,\n            \"sliding_window\": 128,\n            \"layer_types\": [\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n                \"sliding_attention\",\n                \"full_attention\",\n            ],\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # GPT-OSS has alternating sliding / full attention\n        # Even layers (12 layers) use sliding window attention with window_size = 128\n        # Odd layers  (12 layers) use full attention\n        #\n        # Non-attention FLOPs:\n        # vocab part: 201088 * 2880 * 2 = 1158266880\n        # attn linear part per layer:\n        #   Q: 2880 * (64 * 64) = 11796480\n        #   K: 2880 * (8  * 64) = 1474560\n        #   V: 2880 * (8  * 64) = 1474560\n        #   O: (64 * 64) * 2880 = 11796480\n        #   attn linear total = 26542080\n        # mlp (MoE, SwiGLU) part per layer:\n        #   gate: 2880 * 32 = 92160\n        #   active experts: 3 * 2880 * 2880 * 4 = 99532800\n        #   mlp total = 99624960\n        # total per layer: 26542080 + 99624960 = 126167040\n        # all layers:\n        #   126167040 * 24 = 3028008960\n        # total dense params:\n        #   3028008960 + 1158266880 = 4186275840\n        #\n        # For batch [512, 1024, 2048], tokens_sum = 3584:\n        # dense flops: 6 * 4186275840 * 3584 = 90021675663360\n        # seqlen_square_sum: 71565312 (calculated with sliding window logic)\n        # attn flops: 6 * 71565312 * 64 * 64 = 3517578215424\n        # total: 93539253878784 / 1e12 = 93.539253878784\n        #\n        # For batch [4096, 4096, 4096], tokens_sum = 12288:\n        # dense flops: 6 * 4186275840 * 12288 = 308646629068800\n        # seqlen_square_sum: 622854144 (calculated with sliding window logic)\n        # attn flops: 6 * 622854144 * 64 * 64 = 30613642948608\n        # total: 339260272017408 / 1e12 = 339.260272017408\n        \"expected_flops_tuple\": (91780464771072 / 1e12, 323953008574464 / 1e12),\n    },\n    \"apertus\": {\n        \"config\": {  # swiss-ai/Apertus-8B\n            \"model_type\": \"apertus\",\n            \"vocab_size\": 131072,\n            \"hidden_size\": 4096,\n            \"intermediate_size\": 21504,\n            \"num_hidden_layers\": 32,\n            \"num_attention_heads\": 32,\n            \"num_key_value_heads\": 32,\n            \"hidden_act\": \"xielu\",\n            # head_dim will be derived as 4096 / 32 = 128\n        },\n        \"batch_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # Calculation for Apertus (hidden_act=\"xielu\" -> MLP uses [k_mlp=2]*H*I params; qk_norm=True -> [k_qkn=2]*H):\n        # V=131072, H=4096, I=21504, L=32, k_mlp=2 (XIELU), k_qkn=2 (QK norm), S=6\n        # S*(2*V*H + L*(4*H**2 + k_mlp*H*I + k_qkn*H)) * (SUM[seqlen]) + 6*SUM[seqlen**2]*L*H\n        \"expected_flops_tuple\": (194825353691136 / 1e12, 692711652851712 / 1e12),\n    },\n    \"qwen3_vl\": {\n        \"config\": {  # Qwen/Qwen3-VL-8B\n            \"model_type\": \"qwen3_vl\",\n            # -------- Text config --------\n            \"text_config\": {\n                \"vocab_size\": 151936,\n                \"hidden_size\": 4096,\n                \"intermediate_size\": 12288,\n                \"num_hidden_layers\": 36,\n                \"num_attention_heads\": 32,\n                \"num_key_value_heads\": 8,\n                \"head_dim\": 128,\n            },\n            # -------- Vision config (ViT) --------\n            \"vision_config\": {\n                \"deepstack_visual_indexes\": [8, 16, 24],\n                \"num_heads\": 16,\n                \"depth\": 27,\n                \"hidden_size\": 1152,\n                \"intermediate_size\": 4304,\n                \"out_hidden_size\": 4096,\n                \"spatial_merge_size\": 2,\n                \"temporal_patch_size\": 2,\n                \"in_channels\": 3,\n                \"patch_size\": 16,\n            },\n        },\n        \"batch_seqlens_tuple\": (\n            [512, 1024, 2048],\n            [4096, 4096, 4096],\n        ),\n        \"images_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # -----Text-----\n        # 6*(vocab*hidden*2\n        #   + layer*(hidden*(q+k+v+o) + hidden*inter*3)\n        # )*token_sum\n        # + 6*sum(seqlen^2)*layer*hidden\n        #\n        # -----ViT-----\n        # patch_embed_N =hidden*temporal_patch_size*in_channels* patch_size^2\n        # attn_linear_N =hidden*(4*hidden)\n        # mlp_N =hidden*inter*2\n        # merger_N =((o+hidden*spatial_merge_size^2) * (hidden*spatial_merge_size^2))\n        # deepstack_merger_N =merger_N * 3\n        # dense_N =patch_embed_N + (attn_linear_N + mlp_N) * 27 + deepstack_merger_N + merger_N\n        #\n        # 6*(151936*4096*2\n        #   + 36*(4096*(4096+1024+1024+4096) + 4096*12288*3)\n        # )*(512+1024+2048)\n        # + 12*(512*512+1024*1024+2048*2048)*36*4096\n        # + 6 * dense_N * (512 + 1024 + 2048)\n        # + 12 * (512**2 + 1024**2 + 2048**2) * 27 * 16 * 72\n        #\n        # 6*(151936*4096*2\n        #   + 36*(4096*(4096+1024+1024+4096) + 4096*12288*3)\n        # )*(4096+4096+4096)\n        # + 12*(4096*4096+4096*4096+4096*4096)*36*4096\n        # + 6 * dense_N * (4096 + 4096 + 2048)\n        # + 12 * (4096**2 + 4096**2 + 4096**2) * 27 * 16 * 72\n        \"expected_flops_tuple\": (\n            195379819708416 / 1e12,\n            709446422495232 / 1e12,\n        ),\n    },\n    \"qwen3_vl_moe\": {\n        \"config\": {  # Qwen/Qwen3-VL-30B-A3B\n            \"model_type\": \"qwen3_vl_moe\",\n            # -------- Text config --------\n            \"text_config\": {\n                \"vocab_size\": 151936,\n                \"hidden_size\": 2048,\n                \"num_hidden_layers\": 48,\n                \"num_attention_heads\": 32,\n                \"num_key_value_heads\": 4,\n                \"head_dim\": 128,\n                \"moe_intermediate_size\": 768,\n                \"num_experts\": 128,\n                \"num_experts_per_tok\": 8,\n            },\n            # -------- Vision config (ViT) --------\n            \"vision_config\": {\n                \"deepstack_visual_indexes\": [8, 16, 24],\n                \"num_heads\": 16,\n                \"depth\": 27,\n                \"hidden_size\": 1152,\n                \"intermediate_size\": 4304,\n                \"out_hidden_size\": 4096,\n                \"spatial_merge_size\": 2,\n                \"temporal_patch_size\": 2,\n                \"in_channels\": 3,\n                \"patch_size\": 16,\n            },\n        },\n        \"batch_seqlens_tuple\": (\n            [512, 1024, 2048],\n            [4096, 4096, 4096],\n        ),\n        \"images_seqlens_tuple\": ([512, 1024, 2048], [4096, 4096, 4096]),\n        # -----Text-----\n        # 6*(vocab*hidden*2\n        #   + layer*(hidden*(q+k+v+head*head_dim)+hidden*inter*top_k_exp*3+hidden*num_experts)\n        # )*token_sum\n        # + 6*sum(seqlen^2)*layer*hidden\n        #\n        # -----ViT-----\n        # patch_embed_N =hidden*temporal_patch_size*in_channels* patch_size^2\n        # attn_linear_N =hidden*(4*hidden)\n        # mlp_N =hidden*inter*2\n        # merger_N =((o+hidden*spatial_merge_size^2) * (hidden*spatial_merge_size^2))\n        # deepstack_merger_N =merger_N * 3\n        # dense_N =patch_embed_N + (attn_linear_N + mlp_N) * 27 + deepstack_merger_N + merger_N\n        #\n        # 6*(151936*2048*2\n        #   + 48*(2048*(128*32+128*4*2+128*32)+2048*768*8*3+2048*128)\n        # )*(512+1024+2048)\n        # + 12*(512*512+1024*1024+2048*2048)*48*4096\n        # + 6 * dense_N * (512 + 1024 + 2048)\n        # + 12 * (512**2 + 1024**2 + 2048**2) * 27 * 16 * 72\n        #\n        # 6*(151936*2048*2\n        #   48*(2048*(128*32+128*4*2+128*32)+2048*768*8*3+2048*128)\n        # )*(4096+4096+4096)\n        # + 12*(4096*4096+4096*4096+4096*4096)*48*4096\n        # + 6 * dense_N * (4096 + 4096 + 2048)\n        # + 12 * (4096**2 + 4096**2 + 4096**2) * 27 * 16 * 72\n        \"expected_flops_tuple\": (\n            92975451340800 / 1e12,\n            367622860308480 / 1e12,\n        ),\n    },\n}\n\n\n@pytest.mark.parametrize(\n    \"config_type\",\n    [\n        \"llama\",\n        \"qwen2\",\n        \"qwen3\",\n        \"qwen3_moe\",\n        \"deepseek_v3\",\n        \"mistral\",\n        \"gemma3_text\",\n        \"apertus\",\n        \"gpt_oss\",\n        \"qwen3_vl\",\n        \"qwen3_vl_moe\",\n    ],\n)\ndef test_flops_counter(config_type: str):\n    test_config = CONFIG[config_type]\n    config = Config(test_config[\"config\"])\n    flops_counter = FlopsCounter(config)\n    if \"images_seqlens_tuple\" in test_config:\n        for batch_seqlens, images_seqlens, expected_flops in zip(\n            test_config[\"batch_seqlens_tuple\"],\n            test_config[\"images_seqlens_tuple\"],\n            test_config[\"expected_flops_tuple\"],\n            strict=True,\n        ):\n            # set delta time to 1 to get the flops\n            counted_flops, _ = flops_counter.estimate_flops(batch_seqlens, 1, images_seqlens=images_seqlens)\n            print(f\"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}\")\n            assert math.isclose(counted_flops, expected_flops), (\n                f\"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}\"\n            )\n    else:\n        for batch_seqlens, expected_flops in zip(\n            test_config[\"batch_seqlens_tuple\"], test_config[\"expected_flops_tuple\"], strict=True\n        ):\n            # set delta time to 1 to get the flops\n            counted_flops, _ = flops_counter.estimate_flops(batch_seqlens, 1)\n            print(f\"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}\")\n            assert math.isclose(counted_flops, expected_flops), (\n                f\"Expect flops for {test_config['config']} is {expected_flops}, but get {counted_flops}\"\n            )\n"
  },
  {
    "path": "tests/utils/test_fs_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom pathlib import Path\n\nimport verl.utils.fs as fs\n\n\ndef test_record_and_check_directory_structure(tmp_path):\n    # Create test directory structure\n    test_dir = tmp_path / \"test_dir\"\n    test_dir.mkdir()\n    (test_dir / \"file1.txt\").write_text(\"test\")\n    (test_dir / \"subdir\").mkdir()\n    (test_dir / \"subdir\" / \"file2.txt\").write_text(\"test\")\n\n    # Create structure record\n    record_file = fs._record_directory_structure(test_dir)\n\n    # Verify record file exists\n    assert os.path.exists(record_file)\n\n    # Initial check should pass\n    assert fs._check_directory_structure(test_dir, record_file) is True\n\n    # Modify structure and verify check fails\n    (test_dir / \"new_file.txt\").write_text(\"test\")\n    assert fs._check_directory_structure(test_dir, record_file) is False\n\n\ndef test_copy_from_hdfs_with_mocks(tmp_path, monkeypatch):\n    # Mock HDFS dependencies\n    monkeypatch.setattr(fs, \"is_non_local\", lambda path: True)\n\n    # side_effect will simulate the copy by creating parent dirs + empty file\n    def fake_copy(src: str, dst: str, *args, **kwargs):\n        dst_path = Path(dst)\n        dst_path.parent.mkdir(parents=True, exist_ok=True)\n        dst_path.write_bytes(b\"\")  # touch an empty file\n\n    monkeypatch.setattr(fs, \"copy\", fake_copy)  # Mock actual HDFS copy\n\n    # Test parameters\n    test_cache = tmp_path / \"cache\"\n    hdfs_path = \"hdfs://test/path/file.txt\"\n\n    # Test initial copy\n    local_path = fs.copy_to_local(hdfs_path, cache_dir=test_cache)\n    expected_path = os.path.join(test_cache, fs.md5_encode(hdfs_path), os.path.basename(hdfs_path))\n    assert local_path == expected_path\n    assert os.path.exists(local_path)\n\n\ndef test_always_recopy_flag(tmp_path, monkeypatch):\n    # Mock HDFS dependencies\n    monkeypatch.setattr(fs, \"is_non_local\", lambda path: True)\n\n    copy_call_count = 0\n\n    def fake_copy(src: str, dst: str, *args, **kwargs):\n        nonlocal copy_call_count\n        copy_call_count += 1\n        dst_path = Path(dst)\n        dst_path.parent.mkdir(parents=True, exist_ok=True)\n        dst_path.write_bytes(b\"\")\n\n    monkeypatch.setattr(fs, \"copy\", fake_copy)  # Mock actual HDFS copy\n\n    test_cache = tmp_path / \"cache\"\n    hdfs_path = \"hdfs://test/path/file.txt\"\n\n    # Initial copy (always_recopy=False)\n    fs.copy_to_local(hdfs_path, cache_dir=test_cache)\n    assert copy_call_count == 1\n\n    # Force recopy (always_recopy=True)\n    fs.copy_to_local(hdfs_path, cache_dir=test_cache, always_recopy=True)\n    assert copy_call_count == 2\n\n    # Subsequent normal call (always_recopy=False)\n    fs.copy_to_local(hdfs_path, cache_dir=test_cache)\n    assert copy_call_count == 2  # Should not increment\n"
  },
  {
    "path": "tests/utils/test_fsdp2_peft_wrapping.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Test that apply_fsdp2's module selection handles peft-wrapped models.\n\npeft wraps embed_tokens in a ModulesToSaveWrapper, so isinstance(module, nn.Embedding)\nfails. Without name-based matching, embed_tokens + lm_head land in the root FSDP unit,\ncausing OOM from oversized allgather. These tests verify the module selection logic\nworks for: (1) vanilla models, (2) peft-wrapped models, (3) tied embeddings.\n\"\"\"\n\nimport unittest\nfrom types import SimpleNamespace\n\nimport torch.nn as nn\n\nfrom verl.utils.fsdp_utils import _select_fsdp2_wrap_targets\n\n\nclass MockDecoderLayer(nn.Module):\n    \"\"\"Simulates a transformer decoder layer (e.g. Qwen3DecoderLayer).\"\"\"\n\n    def __init__(self, hidden_size=64):\n        super().__init__()\n        self.self_attn = nn.Linear(hidden_size, hidden_size)\n        self.mlp = nn.Linear(hidden_size, hidden_size)\n\n\nclass MockModulesToSaveWrapper(nn.Module):\n    \"\"\"Simulates peft's ModulesToSaveWrapper around nn.Embedding.\n\n    peft wraps modules listed in modules_to_save (like embed_tokens) in this wrapper,\n    which breaks isinstance(module, nn.Embedding) checks.\n    \"\"\"\n\n    def __init__(self, original_module):\n        super().__init__()\n        self.original_module = original_module\n        self.weight = original_module.weight  # peft exposes weight\n\n\nclass MockCausalLM(nn.Module):\n    \"\"\"Simulates a causal LM with embed_tokens, decoder layers, and lm_head.\"\"\"\n\n    _no_split_modules = [\"MockDecoderLayer\"]\n\n    def __init__(self, vocab_size=1000, hidden_size=64, num_layers=2, tie_word_embeddings=False):\n        super().__init__()\n        self.config = SimpleNamespace(tie_word_embeddings=tie_word_embeddings)\n        self.model = nn.Module()\n        self.model.embed_tokens = nn.Embedding(vocab_size, hidden_size)\n        self.model.layers = nn.ModuleList([MockDecoderLayer(hidden_size) for _ in range(num_layers)])\n        self.lm_head = nn.Linear(hidden_size, vocab_size, bias=False)\n\n        if tie_word_embeddings:\n            self.lm_head.weight = self.model.embed_tokens.weight\n\n\nclass TestFSDP2PeftWrapping(unittest.TestCase):\n    \"\"\"Test module selection in apply_fsdp2 for vanilla and peft-wrapped models.\"\"\"\n\n    def _get_wrapped_names(self, model, cls_names):\n        \"\"\"Return names of modules selected for wrapping.\"\"\"\n        selected = _select_fsdp2_wrap_targets(model, cls_names)\n        # _select_fsdp2_wrap_targets returns module objects; map back to names\n        module_to_name = {id(m): n for n, m in model.named_modules()}\n        return [module_to_name[id(m)] for m in selected]\n\n    def test_vanilla_model_wraps_layers_and_embedding(self):\n        \"\"\"Vanilla model (no peft): embed_tokens matched by isinstance, layers by class name.\"\"\"\n        model = MockCausalLM(tie_word_embeddings=False)\n        names = self._get_wrapped_names(model, [\"MockDecoderLayer\"])\n\n        self.assertIn(\"model.embed_tokens\", names)\n        self.assertIn(\"lm_head\", names)\n        self.assertTrue(any(\"layers.0\" in n for n in names))\n        self.assertTrue(any(\"layers.1\" in n for n in names))\n\n    def test_peft_wrapped_model_wraps_embed_tokens_by_name(self):\n        \"\"\"peft-wrapped model: embed_tokens fails isinstance but is matched by name.\"\"\"\n        model = MockCausalLM(tie_word_embeddings=False)\n        original_embed = model.model.embed_tokens\n        model.model.embed_tokens = MockModulesToSaveWrapper(original_embed)\n\n        names = self._get_wrapped_names(model, [\"MockDecoderLayer\"])\n\n        self.assertIn(\"model.embed_tokens\", names)\n        self.assertIn(\"lm_head\", names)\n        self.assertTrue(any(\"layers.0\" in n for n in names))\n\n    def test_tied_embeddings_skips_name_based_wrapping(self):\n        \"\"\"With tie_word_embeddings=True, embed_tokens/lm_head are NOT wrapped separately.\"\"\"\n        model = MockCausalLM(tie_word_embeddings=True)\n        names = self._get_wrapped_names(model, [\"MockDecoderLayer\"])\n\n        self.assertNotIn(\"model.embed_tokens\", names)\n        self.assertNotIn(\"lm_head\", names)\n        self.assertTrue(any(\"layers.0\" in n for n in names))\n\n    def test_peft_wrapped_tied_embeddings_skips_wrapping(self):\n        \"\"\"peft + tied embeddings: name-based matching is disabled, no wrapping.\"\"\"\n        model = MockCausalLM(tie_word_embeddings=True)\n        original_embed = model.model.embed_tokens\n        model.model.embed_tokens = MockModulesToSaveWrapper(original_embed)\n\n        names = self._get_wrapped_names(model, [\"MockDecoderLayer\"])\n\n        self.assertNotIn(\"model.embed_tokens\", names)\n        self.assertNotIn(\"lm_head\", names)\n\n    def test_no_duplicate_wrapping_for_vanilla_embedding(self):\n        \"\"\"Vanilla nn.Embedding should not be wrapped twice (by isinstance AND by name).\"\"\"\n        model = MockCausalLM(tie_word_embeddings=False)\n        names = self._get_wrapped_names(model, [\"MockDecoderLayer\"])\n\n        embed_count = sum(1 for n in names if n == \"model.embed_tokens\")\n        self.assertEqual(embed_count, 1, f\"embed_tokens wrapped {embed_count} times, expected 1\")\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_fsdp_lora_merge.py",
    "content": "# Copyright 2026 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport pytest\nimport torch\nimport torch.distributed\nimport torch.multiprocessing as mp\nfrom peft import LoraConfig, get_peft_model\nfrom torch.distributed import init_device_mesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision, ShardingStrategy\nfrom transformers import AutoModelForCausalLM, GptOssConfig, Qwen2Config\n\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.fsdp_utils import (\n    MixedPrecisionPolicy,\n    apply_fsdp2,\n    get_fsdp_wrap_policy,\n    merged_lora_context,\n)\n\n\ndef _test_merged_lora_context_worker(\n    rank, world_size, rendezvous_file, strategy, model_config, lora_config_dict, backup_adapters\n):\n    \"\"\"Worker function for testing merged_lora_context with FSDP.\n\n    Args:\n        rank: Process rank\n        world_size: Total number of processes\n        rendezvous_file: Path to rendezvous file for distributed init\n        strategy: FSDP strategy (\"fsdp\" or \"fsdp2\")\n        model_config: Model configuration object (Qwen2Config, GptOssConfig, etc.)\n        lora_config_dict: Dictionary of LoRA configuration parameters\n        backup_adapters: Whether to backup adapter weights before merging\n    \"\"\"\n    get_torch_device().set_device(rank)\n    torch.distributed.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n    device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=(\"dp\",))\n\n    # Create model from provided config\n    with torch.device(get_device_name()):\n        model = AutoModelForCausalLM.from_config(\n            config=model_config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        model = model.to(device=get_device_name())\n\n    # Add LoRA with provided config\n    lora_config = LoraConfig(**lora_config_dict)\n    model = get_peft_model(model, lora_config)\n\n    # Initialize LoRA adapter weights to non-zero values for testing\n    from peft.tuners.lora import LoraLayer\n\n    with torch.no_grad():\n        for name, module in model.named_modules():\n            if isinstance(module, LoraLayer):\n                for adapter_name in module.lora_A.keys():\n                    if adapter_name in module.lora_A:\n                        # Initialize lora_A with values around 1.0\n                        module.lora_A[adapter_name].weight.data.uniform_(0.5, 1.5)\n                    if adapter_name in module.lora_B:\n                        # Initialize lora_B with values around 2.0\n                        module.lora_B[adapter_name].weight.data.uniform_(1.5, 2.5)\n\n    # Wrap model with FSDP\n    if strategy == \"fsdp\":\n        mixed_precision = MixedPrecision(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32\n        )\n        model = FSDP(\n            model,\n            use_orig_params=True,\n            device_id=get_torch_device().current_device(),\n            sharding_strategy=ShardingStrategy.FULL_SHARD,\n            mixed_precision=mixed_precision,\n            device_mesh=device_mesh,\n            auto_wrap_policy=get_fsdp_wrap_policy(module=model, is_lora=True),\n        )\n    else:\n        mp_policy = MixedPrecisionPolicy(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True\n        )\n        fsdp_kwargs = {\n            \"mesh\": device_mesh,\n            \"mp_policy\": mp_policy,\n        }\n        apply_fsdp2(model, fsdp_kwargs, {})\n\n    # Test: backup adapter weights, merge, restore\n    from peft.tuners.lora import LoraLayer\n\n    lora_layers = [m for m in model.modules() if isinstance(m, LoraLayer)]\n\n    # Verify LoRA layers exist\n    assert len(lora_layers) > 0, \"Model should have LoRA layers\"\n\n    # Initially not merged\n    for layer in lora_layers:\n        assert not getattr(layer, \"merged\", False), \"LoRA should not be merged initially\"\n\n    # Backup adapter weights before merge\n    from peft.utils.save_and_load import get_peft_model_state_dict\n\n    original_adapter_weights = get_peft_model_state_dict(model)\n\n    # Use merged_lora_context with the specified backup_adapters flag\n    for _ in range(3):\n        with merged_lora_context(model, backup_adapters=backup_adapters):\n            # Inside context, LoRA should be merged\n            for layer in lora_layers:\n                assert getattr(layer, \"merged\", False), \"LoRA should be merged inside context\"\n\n    # After context, check the state based on backup_adapters flag\n    for layer in lora_layers:\n        assert not getattr(layer, \"merged\", False), \"LoRA should be unmerged after context\"\n\n    restored_adapter_weights = get_peft_model_state_dict(model)\n\n    # Verify adapter weights are restored exactly\n    for key in original_adapter_weights.keys():\n        assert key in restored_adapter_weights, f\"Key {key} should be in restored weights\"\n        torch.testing.assert_close(\n            original_adapter_weights[key].cpu(),\n            restored_adapter_weights[key].cpu(),\n            rtol=1e-5,\n            atol=1e-6,\n            msg=f\"Adapter weight {key} should be restored to original value\",\n        )\n\n    if rank == 0:\n        model_name = model_config.__class__.__name__\n        backup_mode = \"with backup\" if backup_adapters else \"without backup\"\n        print(f\"merged_lora_context test with {model_name} {strategy} {backup_mode} passed on {world_size} GPUs!\")\n\n    torch.distributed.barrier()\n    torch.distributed.destroy_process_group()\n\n\n@pytest.mark.parametrize(\"world_size\", (2,))\n@pytest.mark.parametrize(\"strategy\", (\"fsdp\", \"fsdp2\"))\n@pytest.mark.parametrize(\"backup_adapters\", (True, False))\ndef test_merged_lora_context_qwen2(world_size, strategy, backup_adapters, tmp_path):\n    \"\"\"Test merged_lora_context with FSDP on Qwen2 model.\"\"\"\n    rendezvous_file = str(tmp_path / f\"rdzv_file_qwen2_{backup_adapters}\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    # Create Qwen2 model config\n    model_config = Qwen2Config(num_hidden_layers=2, num_attention_heads=2, hidden_size=128)\n\n    # Create LoRA config for Qwen2\n    lora_config_dict = {\n        \"r\": 8,\n        \"lora_alpha\": 16,\n        \"target_modules\": [\"q_proj\", \"v_proj\"],\n        \"lora_dropout\": 0.0,\n        \"bias\": \"none\",\n        \"task_type\": \"CAUSAL_LM\",\n    }\n\n    mp.spawn(\n        fn=_test_merged_lora_context_worker,\n        args=(world_size, rendezvous_file, strategy, model_config, lora_config_dict, backup_adapters),\n        nprocs=world_size,\n        join=True,\n    )\n\n\n@pytest.mark.parametrize(\"world_size\", (2,))\n@pytest.mark.parametrize(\"strategy\", (\"fsdp\", \"fsdp2\"))\n@pytest.mark.parametrize(\"backup_adapters\", (True, False))\ndef test_merged_lora_context_gptoss(world_size, strategy, backup_adapters, tmp_path):\n    \"\"\"Test merged_lora_context with FSDP on GPT-OSS model.\"\"\"\n    rendezvous_file = str(tmp_path / f\"rdzv_file_gptoss_{backup_adapters}\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    # Create GPT-OSS model config\n    model_config = GptOssConfig(\n        num_hidden_layers=2,\n        num_attention_heads=2,\n        num_key_value_heads=2,\n        hidden_size=128,\n        intermediate_size=256,\n    )\n\n    # Create LoRA config for GPT-OSS\n    lora_config_dict = {\n        \"r\": 8,\n        \"lora_alpha\": 16,\n        \"target_modules\": \"all-linear\",\n        \"target_parameters\": [\"mlp.experts.gate_up_proj\", \"mlp.experts.down_proj\"],\n        \"exclude_modules\": [\"mlp.router\"],\n        \"lora_dropout\": 0.0,\n        \"bias\": \"none\",\n        \"task_type\": \"CAUSAL_LM\",\n    }\n\n    mp.spawn(\n        fn=_test_merged_lora_context_worker,\n        args=(world_size, rendezvous_file, strategy, model_config, lora_config_dict, backup_adapters),\n        nprocs=world_size,\n        join=True,\n    )\n"
  },
  {
    "path": "tests/utils/test_groupwise.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nos.environ.setdefault(\"VERL_FORCE_DEVICE\", \"cpu\")  # ensure CPU for tests\n\nimport numpy as np\nimport pytest\nimport torch\n\nfrom verl.utils import as_torch_index, group_mean_std\n\n\ndef test_as_torch_index_basic_integers():\n    g = as_torch_index([2, 2, 5, 7, 5, 2])\n    assert g.dtype == torch.long\n    assert g.device.type == \"cpu\"\n    # Values should be contiguous 0..G-1, keeping equal labels equal\n    assert g.tolist()[0] == g.tolist()[1]\n    assert len(torch.unique(g)) == 3  # {2,5,7} -> 3 groups\n\n\ndef test_as_torch_index_near_integer_floats():\n    arr = np.array([1.0000001, 2.0, 1.0, 3.0000000001], dtype=np.float64)\n    g = as_torch_index(arr)  # should round to integers then factorize\n    assert g.dtype == torch.long\n    assert len(torch.unique(g)) == 3  # {1,2,3}\n\n\ndef test_as_torch_index_factorization_mixed():\n    labels = [\"a\", \"b\", \"a\", \"c\", \"0042\", 42]\n    g = as_torch_index(labels)\n    # \"0042\" and 42 should NOT be the same group (strings are not coerced here)\n    assert g.tolist()[4] != g.tolist()[5]\n    assert len(torch.unique(g)) == 5\n\n\ndef test_group_mean_std_simple():\n    # groups: 0 -> [1, 3], 1 -> [2]\n    scores = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)\n    gidx = as_torch_index([0, 1, 0])\n\n    mean_g, std_g, cnt_g = group_mean_std(scores, gidx)\n    # group 0: mean = (1+3)/2 = 2\n    # sample std (unbiased) = sqrt( (sum(x^2) - (sum(x)^2)/n) / (n-1) )\n    # = sqrt( (1^2+3^2) - (1+3)^2/2 ) / (2-1) = sqrt(10 - 16/2) = sqrt(2)\n    assert torch.allclose(mean_g, torch.tensor([2.0, 0.0]))\n    assert torch.allclose(cnt_g, torch.tensor([2.0, 1.0]))\n    # singleton group -> std = 1.0\n    assert mean_g[1].item() == 0.0\n    assert std_g[1].item() == 1.0\n    assert pytest.approx(std_g[0].item(), rel=1e-6) == (2.0**0.5)\n\n\ndef test_group_mean_std_empty():\n    scores = torch.tensor([], dtype=torch.float32)\n    gidx = torch.tensor([], dtype=torch.long)\n    mean_g, std_g, cnt_g = group_mean_std(scores, gidx)\n    assert mean_g.numel() == 0 and std_g.numel() == 0 and cnt_g.numel() == 0\n\n\ndef test_group_mean_std_default_device_no_force_env(monkeypatch):\n    \"\"\"\n    Regression test:\n    - group_mean_std(device=None) must not pass a device *module* (e.g., torch.cuda)\n      into Tensor.to(device=...), which crashes with:\n      TypeError: to() received an invalid combination of arguments - got (..., device=module, ...)\n    \"\"\"\n    # Simulate a non-pytest environment (training code path) while keeping the test CPU-only.\n    monkeypatch.delenv(\"VERL_FORCE_DEVICE\", raising=False)\n    monkeypatch.delenv(\"PYTEST_CURRENT_TEST\", raising=False)\n\n    # Force device selection to CPU even if CUDA is available on the test machine.\n    import verl.utils.device as device_mod\n\n    monkeypatch.setattr(device_mod, \"is_cuda_available\", False)\n    monkeypatch.setattr(device_mod, \"is_npu_available\", False)\n\n    scores = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float32)\n    gidx = torch.tensor([0, 1, 0], dtype=torch.long)\n\n    mean_g, std_g, cnt_g = group_mean_std(scores, gidx)\n    assert mean_g.device.type == \"cpu\"\n    assert std_g.device.type == \"cpu\"\n    assert cnt_g.device.type == \"cpu\"\n"
  },
  {
    "path": "tests/utils/test_import_utils_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport pytest\n\nfrom verl.utils.import_utils import load_extern_object\n\n# Path to the test module\nTEST_MODULE_PATH = os.path.join(os.path.dirname(__file__), \"_test_module.py\")\n\n\ndef test_load_extern_object_class():\n    \"\"\"Test loading a class from an external file\"\"\"\n    TestClass = load_extern_object(TEST_MODULE_PATH, \"TestClass\")\n\n    # Verify the class was loaded correctly\n    assert TestClass is not None\n    assert TestClass.__name__ == \"TestClass\"\n\n    # Test instantiation and functionality\n    instance = TestClass()\n    assert instance.value == \"default\"\n\n    # Test with a custom value\n    custom_instance = TestClass(\"custom\")\n    assert custom_instance.get_value() == \"custom\"\n\n\ndef test_load_extern_object_function():\n    \"\"\"Test loading a function from an external file\"\"\"\n    test_function = load_extern_object(TEST_MODULE_PATH, \"test_function\")\n\n    # Verify the function was loaded correctly\n    assert test_function is not None\n    assert callable(test_function)\n\n    # Test function execution\n    result = test_function()\n    assert result == \"test_function_result\"\n\n\ndef test_load_extern_object_constant():\n    \"\"\"Test loading a constant from an external file\"\"\"\n    constant = load_extern_object(TEST_MODULE_PATH, \"TEST_CONSTANT\")\n\n    # Verify the constant was loaded correctly\n    assert constant is not None\n    assert constant == \"test_constant_value\"\n\n\ndef test_load_extern_object_nonexistent_file():\n    \"\"\"Test behavior when file doesn't exist\"\"\"\n    with pytest.raises(FileNotFoundError):\n        load_extern_object(\"/nonexistent/path.py\", \"SomeType\")\n\n\ndef test_load_extern_object_nonexistent_type():\n    \"\"\"Test behavior when type doesn't exist in the file\"\"\"\n    with pytest.raises(AttributeError):\n        load_extern_object(TEST_MODULE_PATH, \"NonExistentType\")\n\n\ndef test_load_extern_object_none_path():\n    \"\"\"Test behavior when file path is None\"\"\"\n    with pytest.raises(AttributeError):\n        load_extern_object(None, \"SomeType\")\n\n\ndef test_load_extern_object_invalid_module():\n    \"\"\"Test behavior when module has syntax errors\"\"\"\n    # Create a temporary file with syntax errors\n    import tempfile\n\n    with tempfile.NamedTemporaryFile(suffix=\".py\", mode=\"w+\", delete=False) as temp_file:\n        temp_file.write(\"This is not valid Python syntax :\")\n        temp_path = temp_file.name\n\n    try:\n        with pytest.raises(RuntimeError):\n            load_extern_object(temp_path, \"SomeType\")\n    finally:\n        # Clean up the temporary file\n        if os.path.exists(temp_path):\n            os.remove(temp_path)\n"
  },
  {
    "path": "tests/utils/test_linear_cross_entropy.py",
    "content": "#\n# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport torch\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.utils.device import is_torch_npu_available\nfrom verl.utils.experimental.torch_functional import FusedLinearForPPO\nfrom verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\nfrom verl.utils.torch_functional import logprobs_from_logits\n\ncompute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True)\nfused_linear_for_ppo = FusedLinearForPPO()\nfused_linear_for_ppo.compile(dynamic=True)\n\nMAX_TEST_CASES = os.environ.get(\"MAX_TEST_CASES\", 5)\n\n\ndef run_torch_entropy(\n    hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, temperature: float, reduction=\"none\"\n) -> list[torch.Tensor]:\n    hidden = hidden.squeeze(0).to(torch.float32)\n    weight = weight.transpose(0, 1).to(torch.float32)\n    logits = torch.matmul(hidden, weight)  # [num_tokens, vocab_size]\n    logits /= temperature\n    pd = torch.nn.functional.softmax(logits, dim=-1)  # [num_tokens, vocab_size]\n    entropy_a = torch.logsumexp(logits, dim=-1)  # [num_tokens]\n    entropy_b = torch.sum(pd * logits, dim=-1)  # [num_tokens]\n    entropy = entropy_a - entropy_b\n    logprobs = torch.nn.functional.cross_entropy(logits, labels.squeeze(0), reduction=reduction)  # [num_tokens]\n    logprobs = torch.neg(logprobs)\n    return logprobs, entropy\n\n\ndef run_verl_original_entropy(\n    hidden: torch.Tensor,\n    weight: torch.Tensor,\n    labels: torch.Tensor,\n    temperature: float,\n) -> list[torch.Tensor]:\n    hidden = hidden.squeeze(0).to(torch.float32)\n    weight = weight.transpose(0, 1).to(torch.float32)\n    logits = torch.matmul(hidden, weight)  # [num_tokens, vocab_size]\n    logits /= temperature\n    # compute entropy\n    entropy = compute_entropy_from_logits(logits)  # ((total_nnz / sp) + pad)\n    # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen)\n    logprobs = logprobs_from_logits(logits=logits, labels=labels, inplace_backward=False)\n    return logprobs, entropy\n\n\n# To be tested\ndef run_verl_torch_fused_entropy(\n    hidden: torch.Tensor,\n    weight: torch.Tensor,\n    labels: torch.Tensor,\n    temperature: float,\n):\n    hidden = hidden.to(torch.float32)\n    weight = weight.to(torch.float32)\n    logprobs, entropy = fused_linear_for_ppo(\n        hidden,\n        weight,\n        labels,\n        temperature=temperature,\n    )\n    return logprobs.squeeze(0), entropy.squeeze(0)\n\n\nclass TestLinearCrossEntropy:\n    def __init__(self, test_case_idx: int, temperature: float = 1.5) -> None:\n        self.test_case_idx = test_case_idx\n        self.temperature = temperature\n\n    def cleanup(self):\n        torch.cuda.empty_cache()\n        torch.cuda.reset_peak_memory_stats()\n        import gc\n\n        gc.collect()\n        torch.cuda.synchronize()\n\n    def generate_hyper(self):\n        global MAX_TEST_CASES\n\n        self.dtype = torch.bfloat16\n        if self.test_case_idx == 0:\n            self.batch_size = 1\n            self.num_tokens = 1937\n            self.hidden_size = 3584\n            self.vocab_size = 152064\n        elif self.test_case_idx == 1:\n            self.batch_size = 1\n            self.num_tokens = 2169\n            self.hidden_size = 896\n            self.vocab_size = 151936\n        elif self.test_case_idx == 2:\n            self.batch_size = 1\n            self.num_tokens = 1530\n            self.hidden_size = 2048\n            self.vocab_size = 32256\n        elif self.test_case_idx == 3:\n            self.batch_size = 1\n            self.num_tokens = 1388\n            self.hidden_size = 4096\n            self.vocab_size = 102400\n        elif self.test_case_idx == 4:\n            self.batch_size = 1\n            self.num_tokens = 8192\n            self.hidden_size = 4096\n            self.vocab_size = 102400\n        else:\n            raise ValueError(f\"Invalid test case index: {self.test_case_idx}\")\n        assert MAX_TEST_CASES <= 5, \"MAX_TEST_CASES should be less than or equal to 5.\"\n\n    def generate_forward_inputs(self):\n        hidden = (\n            torch.empty((self.batch_size, self.num_tokens, self.hidden_size), dtype=self.dtype, device=\"cuda\")\n            .uniform_(-0.5, 0.5)\n            .requires_grad_()\n        )\n        weight = (\n            torch.empty((self.vocab_size, self.hidden_size), dtype=self.dtype, device=\"cuda\")\n            .uniform_(-0.5, 0.5)\n            .requires_grad_()\n        )\n        labels = torch.randint(0, self.vocab_size, (self.batch_size, self.num_tokens), device=\"cuda\")\n        return hidden, weight, labels\n\n    def generate_backward_inputs(self):\n        g_entropy = torch.empty((self.num_tokens,), dtype=self.dtype, device=\"cuda\").uniform_(-0.5, 0.5)\n        g_logprobs = torch.empty((self.num_tokens,), dtype=self.dtype, device=\"cuda\").uniform_(-1, 1)\n        return g_entropy, g_logprobs\n\n    def verify_correctness(self, iterations=5):\n        self.cleanup()\n        self.generate_hyper()\n\n        torch_forward_latency = list()\n        torch_backward_latency = list()\n        verl_forward_latency = list()\n        verl_backward_latency = list()\n        verl_fused_forward_latency = list()\n        verl_fused_backward_latency = list()\n        kernel_forward_latency = list()\n        kernel_backward_latency = list()\n\n        start_event = torch.cuda.Event(enable_timing=True)\n        end_event = torch.cuda.Event(enable_timing=True)\n\n        for i in range(iterations):\n            print(f\"[INFO]: Iteration {i + 1} / {iterations}...\", end=\"\\r\")\n            hidden, weight, labels = self.generate_forward_inputs()\n\n            start_event.record()\n            (torch_logprobs, torch_entropy) = run_torch_entropy(hidden, weight, labels, self.temperature)\n            end_event.record()\n            torch.cuda.synchronize()\n            torch_forward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (verl_logprobs, verl_entropy) = run_verl_original_entropy(hidden, weight, labels, self.temperature)\n            end_event.record()\n            torch.cuda.synchronize()\n            verl_forward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (verl_fused_logprobs, verl_fused_entropy) = run_verl_torch_fused_entropy(\n                hidden, weight, labels, self.temperature\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            verl_fused_forward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (kernel_logprobs, kernel_entropy) = linear_cross_entropy(hidden, weight, labels, self.temperature)\n            end_event.record()\n            torch.cuda.synchronize()\n            kernel_forward_latency.append(start_event.elapsed_time(end_event))\n\n            torch.testing.assert_close(torch_logprobs, verl_logprobs, atol=1e-4, rtol=1e-4)\n            torch.testing.assert_close(torch_entropy, verl_entropy, atol=1e-4, rtol=1e-4)\n\n            torch.testing.assert_close(torch_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4)\n            torch.testing.assert_close(torch_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4)\n            torch.testing.assert_close(verl_logprobs, verl_fused_logprobs, atol=1e-4, rtol=1e-4)\n            torch.testing.assert_close(verl_entropy, verl_fused_entropy, atol=1e-4, rtol=1e-4)\n\n            torch.testing.assert_close(torch_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4)\n            torch.testing.assert_close(torch_entropy, kernel_entropy, atol=5e-3, rtol=5e-4)\n            torch.testing.assert_close(verl_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4)\n            torch.testing.assert_close(verl_entropy, kernel_entropy, atol=5e-3, rtol=5e-4)\n            torch.testing.assert_close(verl_fused_logprobs, kernel_logprobs, atol=1e-3, rtol=2e-4)\n            torch.testing.assert_close(verl_fused_entropy, kernel_entropy, atol=5e-3, rtol=5e-4)\n\n            # backward\n            g_entropy, g_logprobs = self.generate_backward_inputs()\n\n            start_event.record()\n            (d_torch_hidden, d_torch_weight) = torch.autograd.grad(\n                (torch_entropy, torch_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            torch_backward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (d_verl_hidden, d_verl_weight) = torch.autograd.grad(\n                (verl_entropy, verl_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            verl_backward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (d_verl_fused_hidden, d_verl_fused_weight) = torch.autograd.grad(\n                (verl_fused_entropy, verl_fused_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            verl_fused_backward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (d_kernel_hidden, d_kernel_weight) = torch.autograd.grad(\n                (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            kernel_backward_latency.append(start_event.elapsed_time(end_event))\n\n            torch.testing.assert_close(d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4)\n\n            torch.testing.assert_close(d_torch_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_torch_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_verl_hidden, d_verl_fused_hidden, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_verl_weight, d_verl_fused_weight, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_torch_hidden, d_verl_hidden, atol=1e-2, rtol=1e-4)\n            torch.testing.assert_close(d_torch_weight, d_verl_weight, atol=1e-2, rtol=1e-4)\n\n            torch.testing.assert_close(d_torch_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(d_torch_weight, d_kernel_weight, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(d_verl_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(d_verl_weight, d_kernel_weight, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(d_verl_fused_hidden, d_kernel_hidden, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(d_verl_fused_weight, d_kernel_weight, atol=2e-2, rtol=4e-2)\n\n        # remove first latency\n        torch_forward_latency = torch_forward_latency[1:]\n        torch_backward_latency = torch_backward_latency[1:]\n        verl_forward_latency = verl_forward_latency[1:]\n        verl_backward_latency = verl_backward_latency[1:]\n        verl_fused_forward_latency = verl_fused_forward_latency[1:]\n        verl_fused_backward_latency = verl_fused_backward_latency[1:]\n        kernel_forward_latency = kernel_forward_latency[1:]\n        kernel_backward_latency = kernel_backward_latency[1:]\n\n        print(\"\\n[INFO]: Verified forward & backward correctness.\")\n\n        print(\n            f\"[INFO]: Forward pass: Torch implementation average time: \"\n            f\"{sum(torch_forward_latency) / len(torch_forward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Backward pass: torch implementation average time: \"\n            f\"{sum(torch_backward_latency) / len(torch_backward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Forward pass: VeRL implementation average time: \"\n            f\"{sum(verl_forward_latency) / len(verl_forward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Backward pass: VeRL implementation average time: \"\n            f\"{sum(verl_backward_latency) / len(verl_backward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Forward pass: VeRL Fused Entropy implementation average time: \"\n            f\"{sum(verl_fused_forward_latency) / len(verl_fused_forward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Backward pass: VeRL Fused Entropy implementation average time: \"\n            f\"{sum(verl_fused_backward_latency) / len(verl_fused_backward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Forward pass: Kernel implementation average time: \"\n            f\"{sum(kernel_forward_latency) / len(kernel_forward_latency):.2f} ms\"\n        )\n        print(\n            f\"[INFO]: Backward pass: kernel implementation average time: \"\n            f\"{sum(kernel_backward_latency) / len(kernel_backward_latency):.2f} ms\"\n        )\n\n    def check_storage(self, method_name, run_forward):\n        self.cleanup()\n        self.generate_hyper()\n\n        hidden, weight, labels = self.generate_forward_inputs()\n\n        torch.cuda.reset_peak_memory_stats()\n        (logprobs, entropy) = run_forward(hidden, weight, labels, self.temperature)\n        torch.cuda.synchronize()\n        torch_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n        print(f\"[INFO]: {method_name} Forward pass peak memory: {torch_max_memory:.2f} MB\")\n\n        g_entropy, g_logprobs = self.generate_backward_inputs()\n\n        torch.cuda.reset_peak_memory_stats()\n        (d_torch_hidden, d_torch_weight) = torch.autograd.grad(\n            (entropy, logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n        )\n        torch.cuda.synchronize()\n        torch_backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n        print(f\"[INFO]: {method_name} Backward pass peak memory: {torch_backward_max_memory:.2f} MB\")\n\n    def check_storage_all(self):\n        self.check_storage(\"Torch\", run_torch_entropy)\n        self.check_storage(\"VeRL\", run_verl_original_entropy)\n        self.check_storage(\"VeRL Torch Fused\", run_verl_torch_fused_entropy)\n        self.check_storage(\"Kernel\", linear_cross_entropy)\n\n\ndef test_lce_non_divisible_vocab_padding():\n    \"\"\"Regression test for the logsumexp padding bug.\n\n    When vocab_size % BLOCK_SIZE_N != 0 the last tile has fewer than\n    BLOCK_SIZE_N valid entries. Without the fix, out-of-bounds positions\n    are loaded as weight=0 → logit=0 → exp(0)=1, adding phantom probability\n    mass to the logsumexp denominator. For peaked softmax distributions\n    (small denominator) this causes large log-prob errors.\n\n    Reproducing construction: one token-logit at +3, all others at -15\n    → denominator ≈ 20, phantom adds ≈ 25 → error ≈ 0.82 per token.\n    \"\"\"\n    if not torch.cuda.is_available() or is_torch_npu_available(check_device=False):\n        return\n\n    torch.manual_seed(0)\n\n    V = 152064  # vocab_size % 1024 == 512 (triggers bug)\n    V_div = 149 * 1024  # vocab_size % 1024 == 0 (control)\n    D = 3584\n    N = 512\n    T = 1.5\n\n    def reference(hidden, weight, labels):\n        h = hidden.squeeze(0).float()\n        logits = torch.matmul(h, weight.float().T) / T\n        lp = -torch.nn.functional.cross_entropy(logits, labels.squeeze(0), reduction=\"none\")\n        pd = torch.nn.functional.softmax(logits, dim=-1)\n        ent = torch.logsumexp(logits, dim=-1) - (pd * logits).sum(-1)\n        return lp, ent\n\n    for vocab_size, desc in [(V, \"non-divisible vocab (mod1024=512)\"), (V_div, \"divisible vocab (mod1024=0)\")]:\n        w = torch.zeros(vocab_size, D, dtype=torch.bfloat16, device=\"cuda\")\n        w[:, 0] = -15.0 * T\n        w[0, 0] = 3.0 * T\n        h = torch.zeros(1, N, D, dtype=torch.bfloat16, device=\"cuda\")\n        h[:, :, 0] = 1.0\n        labels = torch.zeros(1, N, dtype=torch.long, device=\"cuda\")\n\n        ref_lp, ref_ent = reference(h, w, labels)\n        ker_lp, ker_ent = linear_cross_entropy(h, w, labels, T)\n\n        torch.testing.assert_close(ref_lp, ker_lp, atol=1e-3, rtol=1e-3, msg=f\"logprob mismatch: {desc}\")\n        torch.testing.assert_close(ref_ent, ker_ent, atol=1e-3, rtol=1e-3, msg=f\"entropy mismatch: {desc}\")\n\n\nif __name__ == \"__main__\":\n    # torch.cuda.memory._record_memory_history()\n\n    for test_case_idx in range(MAX_TEST_CASES):\n        print(f\"[INFO] Running test case {test_case_idx}\")\n        test = TestLinearCrossEntropy(test_case_idx)\n\n        test.verify_correctness()\n        test.check_storage_all()\n\n    test_lce_non_divisible_vocab_padding()\n\n    # torch.cuda.memory._dump_snapshot(\"test_linear_cross_entropy.pkl\")\n"
  },
  {
    "path": "tests/utils/test_mlflow_key_sanitization.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom unittest.mock import patch\n\nfrom verl.utils.tracking import _MlflowLoggingAdapter\n\n\nclass TestMlflowLoggingAdapter(unittest.TestCase):\n    def test_sanitize_key_and_warning(self):\n        \"\"\"Test key sanitization for invalid characters and consecutive slashes with warnings.\"\"\"\n        adapter = _MlflowLoggingAdapter()\n        data = {\n            \"valid_key\": 1.0,\n            \"invalid@key!\": 2.0,\n            \"another/valid-key\": 3.0,\n            \"bad key#\": 4.0,\n            \"val-aux//reward/mean_at_1\": 5.0,\n            \"val-core///acc/best_at_5\": 6.0,\n            \"metric////with/many////slashes\": 7.0,\n        }\n        # Patch mlflow.log_metrics to capture the metrics actually sent\n        with (\n            patch(\"mlflow.log_metrics\") as mock_log_metrics,\n            patch.object(adapter, \"logger\") as mock_logger,\n        ):\n            adapter.log(data, step=5)\n            # Check that invalid characters are sanitized\n            sent_metrics = mock_log_metrics.call_args[1][\"metrics\"]\n            self.assertIn(\"invalid_at_key_\", sent_metrics)  # @ becomes _at_, ! becomes _\n            self.assertIn(\"bad key_\", sent_metrics)  # # becomes _, space remains\n            self.assertNotIn(\"invalid@key!\", sent_metrics)\n            self.assertNotIn(\"bad key#\", sent_metrics)\n            # Check that consecutive slashes are collapsed to single slashes\n            self.assertIn(\"val-aux/reward/mean_at_1\", sent_metrics)\n            self.assertIn(\"val-core/acc/best_at_5\", sent_metrics)\n            self.assertIn(\"metric/with/many/slashes\", sent_metrics)\n            self.assertNotIn(\"val-aux//reward/mean_at_1\", sent_metrics)\n            self.assertNotIn(\"val-core///acc/best_at_5\", sent_metrics)\n            # Check that warnings were logged for all sanitized keys\n            warning_msgs = [str(call) for call in mock_logger.warning.call_args_list]\n            # Warnings for invalid characters\n            self.assertTrue(any(\"invalid@key!\" in msg and \"invalid_at_key_\" in msg for msg in warning_msgs))\n            self.assertTrue(any(\"bad key#\" in msg and \"bad key_\" in msg for msg in warning_msgs))\n            # Warnings for consecutive slashes\n            self.assertTrue(any(\"val-aux//reward/mean_at_1\" in msg for msg in warning_msgs))\n            self.assertTrue(any(\"val-core///acc/best_at_5\" in msg for msg in warning_msgs))\n            self.assertTrue(any(\"metric////with/many////slashes\" in msg for msg in warning_msgs))\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_model_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom types import SimpleNamespace  # Or use a mock object library\n\nimport pytest\n\nfrom verl.utils.model import update_model_config\n\n\n# Parametrize with different override scenarios\n@pytest.mark.parametrize(\n    \"override_kwargs\",\n    [\n        {\"param_a\": 5, \"new_param\": \"plain_added\"},\n        {\"param_a\": 2, \"nested_params\": {\"sub_param_x\": \"updated_x\", \"sub_param_z\": True}},\n    ],\n)\ndef test_update_model_config(override_kwargs):\n    \"\"\"\n    Tests that update_model_config correctly updates attributes,\n    handling both plain and nested overrides via parametrization.\n    \"\"\"\n    # Create a fresh mock config object for each test case\n    mock_config = SimpleNamespace(\n        param_a=1, nested_params=SimpleNamespace(sub_param_x=\"original_x\", sub_param_y=100), other_param=\"keep_me\"\n    )\n    # Apply the updates using the parametrized override_kwargs\n    update_model_config(mock_config, override_kwargs)\n\n    # Assertions to check if the config was updated correctly\n    if \"nested_params\" in override_kwargs:  # Case 2: Nested override\n        override_nested = override_kwargs[\"nested_params\"]\n        assert mock_config.nested_params.sub_param_x == override_nested[\"sub_param_x\"], \"Nested sub_param_x mismatch\"\n        assert mock_config.nested_params.sub_param_y == 100, \"Nested sub_param_y should be unchanged\"\n        assert hasattr(mock_config.nested_params, \"sub_param_z\"), \"Expected nested sub_param_z to be added\"\n        assert mock_config.nested_params.sub_param_z == override_nested[\"sub_param_z\"], \"Value of sub_param_z mismatch\"\n    else:  # Case 1: Plain override (nested params untouched)\n        assert mock_config.nested_params.sub_param_x == \"original_x\", \"Nested sub_param_x should be unchanged\"\n        assert mock_config.nested_params.sub_param_y == 100, \"Nested sub_param_y should be unchanged\"\n        assert not hasattr(mock_config.nested_params, \"sub_param_z\"), \"Nested sub_param_z should not exist\"\n"
  },
  {
    "path": "tests/utils/test_normalize_peft_param_name.py",
    "content": "# Copyright 2026 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport pytest\nimport torch\nimport torch.distributed\nimport torch.multiprocessing as mp\nfrom peft import LoraConfig, get_peft_model\nfrom torch.distributed import init_device_mesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import MixedPrecision, ShardingStrategy, StateDictType\nfrom transformers import AutoModelForCausalLM, Qwen3Config\n\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.fsdp_utils import (\n    MixedPrecisionPolicy,\n    apply_fsdp2,\n    get_fsdp_wrap_policy,\n    normalize_peft_param_name,\n)\nfrom verl.utils.model import convert_weight_keys\n\n\ndef _test_normalize_peft_with_fsdp_worker(rank, world_size, rendezvous_file, strategy):\n    \"\"\"Worker function for testing normalize_peft_param_name with FSDP-wrapped models.\n\n    Args:\n        rank: Process rank\n        world_size: Total number of processes\n        rendezvous_file: Path to rendezvous file for distributed init\n        strategy: FSDP strategy (\"fsdp\" or \"fsdp2\")\n    \"\"\"\n    get_torch_device().set_device(rank)\n    torch.distributed.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n    device_mesh = init_device_mesh(get_device_name(), mesh_shape=(world_size,), mesh_dim_names=(\"dp\",))\n\n    # Create model config\n    config = Qwen3Config(\n        num_hidden_layers=2,\n        num_attention_heads=2,\n        num_key_value_heads=2,\n        hidden_size=128,\n        intermediate_size=256,\n    )\n\n    # Create base model\n    with torch.device(get_device_name()):\n        base_model = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        base_model = base_model.to(device=get_device_name())\n\n    # Create PEFT model with LoRA\n    lora_config = LoraConfig(\n        r=8, lora_alpha=16, target_modules=\"all-linear\", lora_dropout=0.0, bias=\"none\", task_type=\"CAUSAL_LM\"\n    )\n    peft_model = get_peft_model(base_model, lora_config)\n\n    # Wrap base model with FSDP (create a fresh copy for base model)\n    with torch.device(get_device_name()):\n        base_model_for_fsdp = AutoModelForCausalLM.from_config(\n            config=config, torch_dtype=torch.bfloat16, attn_implementation=\"flash_attention_2\"\n        )\n        base_model_for_fsdp = base_model_for_fsdp.to(device=get_device_name())\n\n    if strategy == \"fsdp\":\n        mixed_precision = MixedPrecision(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32\n        )\n\n        # Wrap base model with FSDP\n        fsdp_base_model = FSDP(\n            base_model_for_fsdp,\n            use_orig_params=True,\n            device_id=get_torch_device().current_device(),\n            sharding_strategy=ShardingStrategy.FULL_SHARD,\n            mixed_precision=mixed_precision,\n            device_mesh=device_mesh,\n            auto_wrap_policy=get_fsdp_wrap_policy(module=base_model_for_fsdp, is_lora=False),\n        )\n\n        # Wrap PEFT model with FSDP\n        fsdp_peft_model = FSDP(\n            peft_model,\n            use_orig_params=True,\n            device_id=get_torch_device().current_device(),\n            sharding_strategy=ShardingStrategy.FULL_SHARD,\n            mixed_precision=mixed_precision,\n            device_mesh=device_mesh,\n            auto_wrap_policy=get_fsdp_wrap_policy(module=peft_model, is_lora=True),\n        )\n    else:\n        # FSDP2\n        mp_policy = MixedPrecisionPolicy(\n            param_dtype=torch.bfloat16, reduce_dtype=torch.float32, cast_forward_inputs=True\n        )\n        fsdp_kwargs = {\n            \"mesh\": device_mesh,\n            \"mp_policy\": mp_policy,\n        }\n\n        # Wrap base model with FSDP2\n        apply_fsdp2(base_model_for_fsdp, fsdp_kwargs, {})\n        fsdp_base_model = base_model_for_fsdp\n\n        # Wrap PEFT model with FSDP2\n        apply_fsdp2(peft_model, fsdp_kwargs, {})\n        fsdp_peft_model = peft_model\n\n    # Get state dicts from FSDP models\n    if strategy == \"fsdp\":\n        # FSDP v1: Use full_state_dict context\n        with FSDP.state_dict_type(fsdp_base_model, StateDictType.FULL_STATE_DICT):\n            base_state_dict = fsdp_base_model.state_dict()\n\n        with FSDP.state_dict_type(fsdp_peft_model, StateDictType.FULL_STATE_DICT):\n            peft_state_dict = fsdp_peft_model.state_dict()\n    else:\n        # FSDP2: Direct state_dict call\n        base_state_dict = fsdp_base_model.state_dict()\n        peft_state_dict = fsdp_peft_model.state_dict()\n\n    # Normalize PEFT model state dict\n    normalized_peft_state_dict = normalize_peft_param_name(peft_state_dict)\n\n    base_state_dict = convert_weight_keys(\n        base_state_dict, getattr(fsdp_base_model, \"_fsdp_wrapped_module\", fsdp_base_model)\n    )\n    normalized_peft_state_dict = convert_weight_keys(\n        normalized_peft_state_dict, getattr(fsdp_peft_model, \"_fsdp_wrapped_module\", fsdp_peft_model)\n    )\n\n    # Get key sets\n    base_keys = set(base_state_dict.keys())\n    normalized_peft_keys = set(normalized_peft_state_dict.keys())\n\n    # if rank == 0:\n    print(f\"\\n=== FSDP {strategy} Test Results ===\")\n    print(f\"Base model keys: {base_keys=}\")\n    print(f\"Normalized PEFT keys: {normalized_peft_keys=}\")\n\n    # Check for missing keys\n    missing_keys = base_keys - normalized_peft_keys\n    if missing_keys:\n        print(f\"Missing keys from base model: {missing_keys}\")\n\n    # Check for extra keys\n    extra_keys = normalized_peft_keys - base_keys\n    if extra_keys:\n        print(f\"Extra keys not in base model: {extra_keys}\")\n\n    # Verify that all base model keys are in the normalized PEFT keys\n    missing_keys = base_keys - normalized_peft_keys\n    assert len(missing_keys) == 0, f\"Missing keys from base model: {missing_keys}\"\n\n    # Verify that all normalized PEFT keys are in the base model\n    extra_keys = normalized_peft_keys - base_keys\n    assert len(extra_keys) == 0, f\"Extra keys not in base model: {extra_keys}\"\n\n    # Verify exact match\n    assert base_keys == normalized_peft_keys, \"Normalized PEFT keys should exactly match FSDP base model keys\"\n\n    # Verify tensor shapes match\n    for key in base_keys:\n        base_shape = base_state_dict[key].shape\n        peft_shape = normalized_peft_state_dict[key].shape\n        assert base_shape == peft_shape, f\"Shape mismatch for {key}: base={base_shape}, peft={peft_shape}\"\n\n    # Verify no LoRA keys remain in normalized state dict\n    lora_keys = [k for k in normalized_peft_keys if \"lora_\" in k or \"adapter_\" in k]\n    assert len(lora_keys) == 0, f\"Normalized state dict should not contain LoRA keys, but found: {lora_keys}\"\n\n    if rank == 0:\n        print(f\"✓ All tests passed for FSDP {strategy}\")\n\n    torch.distributed.barrier()\n    torch.distributed.destroy_process_group()\n\n\n@pytest.mark.parametrize(\"world_size\", (2,))\n@pytest.mark.parametrize(\"strategy\", (\"fsdp\", \"fsdp2\"))\ndef test_normalize_peft_param_name_with_fsdp(world_size, strategy, tmp_path):\n    \"\"\"Test normalize_peft_param_name with FSDP-wrapped models.\n\n    This test verifies that after applying FSDP to both base and PEFT models,\n    the normalized PEFT model keys match the FSDP base model keys.\n    \"\"\"\n    rendezvous_file = str(tmp_path / f\"rdzv_file_normalize_{strategy}\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    mp.spawn(\n        fn=_test_normalize_peft_with_fsdp_worker,\n        args=(world_size, rendezvous_file, strategy),\n        nprocs=world_size,\n        join=True,\n    )\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\"])\n"
  },
  {
    "path": "tests/utils/test_normalize_peft_param_name_on_cpu.py",
    "content": "# Copyright 2026 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\nimport torch\nfrom peft import LoraConfig, get_peft_model\nfrom transformers import AutoModelForCausalLM, Qwen3Config\n\nfrom verl.utils.fsdp_utils import normalize_peft_param_name\n\n\ndef create_base_model():\n    \"\"\"Create a simple base model for testing.\"\"\"\n    config = Qwen3Config(\n        num_hidden_layers=2,\n        num_attention_heads=2,\n        num_key_value_heads=2,\n        hidden_size=128,\n        intermediate_size=256,\n    )\n    model = AutoModelForCausalLM.from_config(config)\n    return model\n\n\ndef create_peft_model():\n    lora_config = LoraConfig(\n        r=8, lora_alpha=16, target_modules=\"all-linear\", lora_dropout=0.0, bias=\"none\", task_type=\"CAUSAL_LM\"\n    )\n    model = create_base_model()\n    model = get_peft_model(model, lora_config)\n    return model\n\n\n@pytest.fixture\ndef base_model():\n    \"\"\"Create a simple base model for testing.\"\"\"\n    return create_base_model()\n\n\n@pytest.fixture\ndef peft_model():\n    \"\"\"Create a PEFT model with LoRA adapters.\"\"\"\n    return create_peft_model()\n\n\ndef test_normalize_peft_param_name_keys_match_base_model():\n    \"\"\"Test that normalized PEFT model keys match base model keys.\"\"\"\n    # Get state dicts\n    base_model = create_base_model()\n    peft_model = create_peft_model()\n    base_state_dict = base_model.state_dict()\n    peft_state_dict = peft_model.state_dict()\n\n    # Normalize PEFT model keys\n    normalized_peft_state_dict = normalize_peft_param_name(peft_state_dict)\n\n    # Get key sets\n    base_keys = set(base_state_dict.keys())\n    normalized_peft_keys = set(normalized_peft_state_dict.keys())\n    print(f\"{base_keys=}\")\n    print(f\"{normalized_peft_keys=}\")\n\n    # Verify that all base model keys are in the normalized PEFT keys\n    missing_keys = base_keys - normalized_peft_keys\n    assert len(missing_keys) == 0, f\"Missing keys from base model: {missing_keys}\"\n\n    # Verify that all normalized PEFT keys are in the base model\n    extra_keys = normalized_peft_keys - base_keys\n    assert len(extra_keys) == 0, f\"Extra keys not in base model: {extra_keys}\"\n\n    # Verify exact match\n    assert base_keys == normalized_peft_keys, \"Normalized PEFT keys should exactly match base model keys\"\n\n\ndef test_normalize_peft_param_name_removes_lora_keys(peft_model):\n    \"\"\"Test that LoRA-specific parameters are removed after normalization.\"\"\"\n    peft_state_dict = peft_model.state_dict()\n\n    # Before normalization, should have lora_A and lora_B keys\n    lora_keys_before = [k for k in peft_state_dict.keys() if \"lora_\" in k]\n    assert len(lora_keys_before) > 0, \"PEFT model should have LoRA parameters\"\n\n    # After normalization, should not have any lora keys\n    normalized_state_dict = normalize_peft_param_name(peft_state_dict)\n    lora_keys_after = [k for k in normalized_state_dict.keys() if \"lora_\" in k]\n    assert len(lora_keys_after) == 0, (\n        f\"Normalized state dict should not contain LoRA keys, but found: {lora_keys_after}\"\n    )\n\n\ndef test_normalize_peft_param_name_removes_base_model_prefix(peft_model):\n    \"\"\"Test that base_model prefix is removed from parameter names.\"\"\"\n    peft_state_dict = peft_model.state_dict()\n\n    # Before normalization, should have base_model prefix\n    base_model_keys = [k for k in peft_state_dict.keys() if \"base_model\" in k]\n    assert len(base_model_keys) > 0, \"PEFT model should have base_model prefix\"\n\n    # After normalization, should not have base_model prefix\n    normalized_state_dict = normalize_peft_param_name(peft_state_dict)\n    base_model_keys_after = [k for k in normalized_state_dict.keys() if \"base_model\" in k]\n    assert len(base_model_keys_after) == 0, (\n        f\"Normalized keys should not contain base_model prefix, but found: {base_model_keys_after}\"\n    )\n\n\ndef test_normalize_peft_param_name_removes_base_layer_suffix(peft_model):\n    \"\"\"Test that .base_layer suffix is removed from parameter names.\"\"\"\n    peft_state_dict = peft_model.state_dict()\n\n    # Before normalization, should have .base_layer suffix\n    base_layer_keys = [k for k in peft_state_dict.keys() if \".base_layer\" in k]\n    assert len(base_layer_keys) > 0, \"PEFT model should have .base_layer suffix\"\n\n    # After normalization, should not have .base_layer suffix\n    normalized_state_dict = normalize_peft_param_name(peft_state_dict)\n    base_layer_keys_after = [k for k in normalized_state_dict.keys() if \".base_layer\" in k]\n    assert len(base_layer_keys_after) == 0, (\n        f\"Normalized keys should not contain .base_layer suffix, but found: {base_layer_keys_after}\"\n    )\n\n\ndef test_normalize_peft_param_name_tensor_shapes_match(base_model, peft_model):\n    \"\"\"Test that tensor shapes match between base model and normalized PEFT model.\"\"\"\n    base_state_dict = base_model.state_dict()\n    peft_state_dict = peft_model.state_dict()\n\n    # Normalize PEFT model keys\n    normalized_peft_state_dict = normalize_peft_param_name(peft_state_dict)\n\n    # Check that shapes match for all common keys\n    for key in base_state_dict.keys():\n        assert key in normalized_peft_state_dict, f\"Key {key} not found in normalized PEFT state dict\"\n        base_shape = base_state_dict[key].shape\n        peft_shape = normalized_peft_state_dict[key].shape\n        assert base_shape == peft_shape, f\"Shape mismatch for {key}: base={base_shape}, peft={peft_shape}\"\n\n\ndef test_normalize_peft_param_name_empty_dict():\n    \"\"\"Test that normalize_peft_param_name handles empty dict.\"\"\"\n    result = normalize_peft_param_name({})\n    assert result == {}, \"Empty dict should return empty dict\"\n\n\n@pytest.mark.parametrize(\n    \"lora_key_pattern\",\n    [\n        \"model.layers.0.self_attn.q_proj.lora_A.default.weight\",\n        \"model.layers.0.self_attn.q_proj.lora_B.default.weight\",\n        \"model.layers.0.adapter_layer.weight\",\n        \"base_model.model.layers.0.lora_embedding_A\",\n    ],\n)\ndef test_normalize_peft_param_name_filters_lora_patterns(lora_key_pattern):\n    \"\"\"Test that various LoRA key patterns are filtered out.\"\"\"\n    test_dict = {\n        lora_key_pattern: torch.randn(10, 10),\n        \"model.layers.0.weight\": torch.randn(10, 10),\n    }\n\n    normalized = normalize_peft_param_name(test_dict)\n\n    # LoRA key should be filtered out\n    assert lora_key_pattern not in normalized, f\"LoRA key {lora_key_pattern} should be filtered out\"\n\n    # Regular key should remain\n    assert len(normalized) == 1, \"Should have exactly one key remaining\"\n    assert \"model.layers.0.weight\" in normalized, \"Regular weight should remain\"\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\"])\n"
  },
  {
    "path": "tests/utils/test_nvtx_profile.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom unittest.mock import MagicMock, patch\n\nfrom verl.utils import omega_conf_to_dataclass\nfrom verl.utils.profiler.config import NsightToolConfig, ProfilerConfig\nfrom verl.utils.profiler.profile import DistProfiler\n\n\nclass TestProfilerConfig(unittest.TestCase):\n    def test_config_init(self):\n        import os\n\n        from hydra import compose, initialize_config_dir\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n            cfg = compose(config_name=\"ppo_trainer\")\n        for config in [\n            cfg.actor_rollout_ref.actor.profiler,\n            cfg.actor_rollout_ref.rollout.profiler,\n            cfg.actor_rollout_ref.ref.profiler,\n            cfg.critic.profiler,\n        ]:\n            profiler_config = omega_conf_to_dataclass(config)\n            self.assertEqual(profiler_config.tool, config.tool)\n            self.assertEqual(profiler_config.enable, config.enable)\n            self.assertEqual(profiler_config.all_ranks, config.all_ranks)\n            self.assertEqual(profiler_config.ranks, config.ranks)\n            self.assertEqual(profiler_config.save_path, config.save_path)\n            self.assertEqual(profiler_config.ranks, config.ranks)\n            assert isinstance(profiler_config, ProfilerConfig)\n            with self.assertRaises(AttributeError):\n                _ = profiler_config.non_existing_key\n            assert config.get(\"non_existing_key\") == profiler_config.get(\"non_existing_key\")\n            assert config.get(\"non_existing_key\", 1) == profiler_config.get(\"non_existing_key\", 1)\n\n    def test_frozen_config(self):\n        \"\"\"Test that modifying frozen keys in ProfilerConfig raises exceptions.\"\"\"\n        from dataclasses import FrozenInstanceError\n\n        from verl.utils.profiler.config import ProfilerConfig\n\n        # Create a new ProfilerConfig instance\n        config = ProfilerConfig(all_ranks=False, ranks=[0])\n\n        with self.assertRaises(FrozenInstanceError):\n            config.all_ranks = True\n\n        with self.assertRaises(FrozenInstanceError):\n            config.ranks = [1, 2, 3]\n\n        with self.assertRaises(TypeError):\n            config[\"all_ranks\"] = True\n\n        with self.assertRaises(TypeError):\n            config[\"ranks\"] = [1, 2, 3]\n\n\nclass TestNsightSystemsProfiler(unittest.TestCase):\n    \"\"\"Test suite for NsightSystemsProfiler functionality.\n\n    Test Plan:\n    1. Initialization: Verify profiler state after creation\n    2. Basic Profiling: Test start/stop functionality\n    3. Discrete Mode: TODO: Test discrete profiling behavior\n    4. Annotation: Test the annotate decorator in both normal and discrete modes\n    5. Config Validation: Verify proper config initialization from OmegaConf\n    \"\"\"\n\n    def setUp(self):\n        self.config = ProfilerConfig(tool=\"nsys\", enable=True, all_ranks=True)\n        self.rank = 0\n        self.profiler = DistProfiler(self.rank, self.config, tool_config=NsightToolConfig(discrete=False))\n\n    def test_initialization(self):\n        self.assertEqual(self.profiler.check_this_rank(), True)\n        self.assertEqual(self.profiler.check_this_step(), False)\n\n    def test_start_stop_profiling(self):\n        with patch(\"torch.cuda.profiler.start\") as mock_start, patch(\"torch.cuda.profiler.stop\") as mock_stop:\n            # Test start\n            self.profiler.start()\n            self.assertTrue(self.profiler.check_this_step())\n            mock_start.assert_called_once()\n\n            # Test stop\n            self.profiler.stop()\n            self.assertFalse(self.profiler.check_this_step())\n            mock_stop.assert_called_once()\n\n    # def test_discrete_profiling(self):\n    #     discrete_config = ProfilerConfig(discrete=True, all_ranks=True)\n    #     profiler = NsightSystemsProfiler(self.rank, discrete_config)\n\n    #     with patch(\"torch.cuda.profiler.start\") as mock_start, patch(\"torch.cuda.profiler.stop\") as mock_stop:\n    #         profiler.start()\n    #         self.assertTrue(profiler.this_step)\n    #         mock_start.assert_not_called()  # Shouldn't start immediately in discrete mode\n\n    #         profiler.stop()\n    #         self.assertFalse(profiler.this_step)\n    #         mock_stop.assert_not_called()  # Shouldn't stop immediately in discrete mode\n\n    def test_annotate_decorator(self):\n        mock_self = MagicMock()\n        mock_self.profiler = self.profiler\n        mock_self.profiler.start()\n        decorator = mock_self.profiler.annotate(message=\"test\")\n\n        @decorator\n        def test_func(self, *args, **kwargs):\n            return \"result\"\n\n        with (\n            patch(\"torch.cuda.profiler.start\") as mock_start,\n            patch(\"torch.cuda.profiler.stop\") as mock_stop,\n            patch(\"verl.utils.profiler.nvtx_profile.mark_start_range\") as mock_start_range,\n            patch(\"verl.utils.profiler.nvtx_profile.mark_end_range\") as mock_end_range,\n        ):\n            result = test_func(mock_self)\n            self.assertEqual(result, \"result\")\n            mock_start_range.assert_called_once()\n            mock_end_range.assert_called_once()\n            mock_start.assert_not_called()  # Not discrete mode\n            mock_stop.assert_not_called()  # Not discrete mode\n\n    # def test_annotate_discrete_mode(self):\n    #     discrete_config = ProfilerConfig(discrete=True, all_ranks=True)\n    #     profiler = NsightSystemsProfiler(self.rank, discrete_config)\n    #     mock_self = MagicMock()\n    #     mock_self.profiler = profiler\n    #     mock_self.profiler.this_step = True\n\n    #     @NsightSystemsProfiler.annotate(message=\"test\")\n    #     def test_func(self, *args, **kwargs):\n    #         return \"result\"\n\n    #     with (\n    #         patch(\"torch.cuda.profiler.start\") as mock_start,\n    #         patch(\"torch.cuda.profiler.stop\") as mock_stop,\n    #         patch(\"verl.utils.profiler.nvtx_profile.mark_start_range\") as mock_start_range,\n    #         patch(\"verl.utils.profiler.nvtx_profile.mark_end_range\") as mock_end_range,\n    #     ):\n    #         result = test_func(mock_self)\n    #         self.assertEqual(result, \"result\")\n    #         mock_start_range.assert_called_once()\n    #         mock_end_range.assert_called_once()\n    #         mock_start.assert_called_once()  # Should start in discrete mode\n    #         mock_stop.assert_called_once()  # Should stop in discrete mode\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_padding_on_cpu.py",
    "content": "# Copyright 2026 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding\n\n\ndef test_padding_conversion_with_log_probs():\n    \"\"\"Test that log probability tensors remain in padded format after conversion\n\n    This test verifies the fix for the bug where ratio values were ~451,728 instead of ~1.0.\n    The key insight is that old_log_probs should STAY in padded format and be sliced\n    in the loss computation to match log_prob from model output, rather than being\n    converted to nested format.\n    \"\"\"\n    batch_size = 4\n    max_seq_len = 128\n    max_response_len = 64\n\n    # Create test data with varying sequence lengths\n    input_ids = torch.randint(0, 1000, (batch_size, max_seq_len))\n\n    # Create attention masks with different valid lengths per sample\n    attention_mask = torch.zeros(batch_size, max_seq_len)\n    valid_lens = [100, 120, 90, 128]  # Different lengths for each batch item\n    for i, vlen in enumerate(valid_lens):\n        attention_mask[i, :vlen] = 1\n\n    # Create response masks aligned with the end of each sequence\n    response_mask = torch.zeros(batch_size, max_response_len)\n    response_lens = [50, 60, 45, 64]  # Different response lengths\n    for i, rlen in enumerate(response_lens):\n        response_mask[i, :rlen] = 1\n\n    # Create position IDs\n    position_ids = torch.arange(max_seq_len).unsqueeze(0).expand(batch_size, -1)\n\n    # Add log probability tensors in padded format\n    old_log_probs = torch.randn(batch_size, max_seq_len)\n    ref_log_prob = torch.randn(batch_size, max_seq_len)\n    advantages = torch.randn(batch_size, max_response_len)\n    rollout_log_probs = torch.randn(batch_size, max_seq_len)\n\n    data = TensorDict(\n        {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n            \"response_mask\": response_mask,\n            \"position_ids\": position_ids,\n            \"old_log_probs\": old_log_probs,\n            \"ref_log_prob\": ref_log_prob,\n            \"advantages\": advantages,\n            \"rollout_log_probs\": rollout_log_probs,\n        }\n    )\n\n    # Convert to no-padding format\n    data_converted = left_right_2_no_padding(data)\n\n    # Verify input_ids and position_ids are nested tensors\n    assert isinstance(data_converted[\"input_ids\"], torch.Tensor)\n    assert data_converted[\"input_ids\"].is_nested\n    assert data_converted[\"position_ids\"].is_nested\n\n    # Verify log probs REMAIN in padded format (NOT converted to nested)\n    # They will be sliced in the loss computation to match log_prob format\n    assert isinstance(data_converted[\"old_log_probs\"], torch.Tensor)\n    assert not data_converted[\"old_log_probs\"].is_nested, \"old_log_probs should remain in padded format\"\n    assert not data_converted[\"ref_log_prob\"].is_nested, \"ref_log_prob should remain in padded format\"\n    assert not data_converted[\"advantages\"].is_nested, \"advantages should remain in padded format\"\n    assert not data_converted[\"rollout_log_probs\"].is_nested, \"rollout_log_probs should remain in padded format\"\n\n    # Verify they maintain their original shapes\n    assert data_converted[\"old_log_probs\"].shape == (batch_size, max_seq_len)\n    assert data_converted[\"ref_log_prob\"].shape == (batch_size, max_seq_len)\n    assert data_converted[\"advantages\"].shape == (batch_size, max_response_len)\n    assert data_converted[\"rollout_log_probs\"].shape == (batch_size, max_seq_len)\n\n    # Verify that nested tensors (input_ids, position_ids) have correct number of elements per batch item\n    for i, vlen in enumerate(valid_lens):\n        assert data_converted[\"input_ids\"][i].numel() == vlen, (\n            f\"Batch {i}: input_ids should have {vlen} elements, got {data_converted['input_ids'][i].numel()}\"\n        )\n\n\ndef test_padding_conversion_without_log_probs():\n    \"\"\"Test that padding conversion works correctly when log prob tensors are not present\"\"\"\n    batch_size = 4\n    max_seq_len = 128\n    max_response_len = 64\n\n    # Create minimal test data\n    input_ids = torch.randint(0, 1000, (batch_size, max_seq_len))\n    attention_mask = torch.ones(batch_size, max_seq_len)\n    response_mask = torch.ones(batch_size, max_response_len)\n    position_ids = torch.arange(max_seq_len).unsqueeze(0).expand(batch_size, -1)\n\n    data = TensorDict(\n        {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n            \"response_mask\": response_mask,\n            \"position_ids\": position_ids,\n        }\n    )\n\n    # Convert to no-padding format\n    data_converted = left_right_2_no_padding(data)\n\n    # Verify basic conversion works\n    assert data_converted[\"input_ids\"].is_nested\n    assert data_converted[\"position_ids\"].is_nested\n    assert \"old_log_probs\" not in data_converted\n    assert \"ref_log_prob\" not in data_converted\n\n\ndef test_padding_roundtrip():\n    \"\"\"Test that converting from padding to nested and back preserves values in the response region\"\"\"\n    batch_size = 2\n    max_seq_len = 64\n    max_response_len = 32\n    prompt_len = max_seq_len - max_response_len  # 32\n\n    # Create simple test data with known values\n    input_ids = torch.arange(1, max_seq_len + 1).unsqueeze(0).expand(batch_size, -1).clone()\n    attention_mask = torch.ones(batch_size, max_seq_len)\n    response_mask = torch.ones(batch_size, max_response_len)\n    position_ids = torch.arange(max_seq_len).unsqueeze(0).expand(batch_size, -1)\n\n    # Create nested prompts and responses (required by no_padding_2_padding)\n    prompt_list = [input_ids[i, :prompt_len] for i in range(batch_size)]\n    response_list = [input_ids[i, prompt_len:] for i in range(batch_size)]\n    prompts_nested = torch.nested.as_nested_tensor(prompt_list, layout=torch.jagged)\n    responses_nested = torch.nested.as_nested_tensor(response_list, layout=torch.jagged)\n\n    data = TensorDict(\n        {\n            \"input_ids\": input_ids,\n            \"prompts\": prompts_nested,\n            \"responses\": responses_nested,\n            \"attention_mask\": attention_mask,\n            \"response_mask\": response_mask,\n            \"position_ids\": position_ids,\n        }\n    )\n\n    # Convert to nested format\n    data_nested = left_right_2_no_padding(data)\n\n    # Verify input_ids is nested\n    assert data_nested[\"input_ids\"].is_nested\n\n    # Convert back to padding format\n    recovered = no_padding_2_padding(data_nested[\"input_ids\"], data_nested)\n\n    # Verify the shape is correct (response region only)\n    assert recovered.shape == (batch_size, max_response_len)\n\n    # Verify values are correct (left-shifted by 1 for log_probs alignment)\n    # Response tokens are 33,34,...,64 -> left-shifted: 32,33,...,63\n    expected = torch.arange(prompt_len, max_seq_len, dtype=torch.long).unsqueeze(0).expand(batch_size, -1)\n    torch.testing.assert_close(recovered, expected)\n\n\ndef test_no_padding_2_padding_varying_lengths():\n    \"\"\"Test no_padding_2_padding with varied prompt/response lengths.\"\"\"\n    batch_size = 4\n    max_seq_len = 100\n    max_response_len = 50\n\n    prompt_lens = [10, 30, 5, 40]\n    response_lens = [40, 20, 45, 10]\n\n    input_ids = torch.zeros(batch_size, max_seq_len, dtype=torch.long)\n    for i in range(batch_size):\n        total_len = prompt_lens[i] + response_lens[i]\n        input_ids[i, :total_len] = torch.arange(1, total_len + 1)\n\n    attention_mask = torch.zeros(batch_size, max_seq_len)\n    for i in range(batch_size):\n        attention_mask[i, : prompt_lens[i] + response_lens[i]] = 1\n\n    response_mask = torch.zeros(batch_size, max_response_len)\n    for i in range(batch_size):\n        response_mask[i, : response_lens[i]] = 1\n\n    position_ids = torch.arange(max_seq_len).unsqueeze(0).expand(batch_size, -1).clone()\n\n    prompt_list = [input_ids[i, : prompt_lens[i]] for i in range(batch_size)]\n    response_list = [input_ids[i, prompt_lens[i] : prompt_lens[i] + response_lens[i]] for i in range(batch_size)]\n\n    prompts_nested = torch.nested.as_nested_tensor(prompt_list, layout=torch.jagged)\n    responses_nested = torch.nested.as_nested_tensor(response_list, layout=torch.jagged)\n\n    data = TensorDict(\n        {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n            \"response_mask\": response_mask,\n            \"position_ids\": position_ids,\n            \"prompts\": prompts_nested,\n            \"responses\": responses_nested,\n        }\n    )\n\n    data_nested = left_right_2_no_padding(data)\n    input_ids_nested = data_nested[\"input_ids\"]\n    log_probs_values = input_ids_nested.values().float()\n    log_probs_nested = torch.nested.nested_tensor_from_jagged(log_probs_values, offsets=input_ids_nested.offsets())\n\n    result_slice_response = no_padding_2_padding(log_probs_nested, data_nested)\n\n    # Verify no_padding_2_padding produces correct values (left-shifted by 1)\n    for i in range(batch_size):\n        resp_len = response_lens[i]\n        expected_start = prompt_lens[i]\n        expected_values = torch.arange(expected_start, expected_start + resp_len, dtype=torch.float)\n        torch.testing.assert_close(\n            result_slice_response[i, :resp_len],\n            expected_values,\n            rtol=1e-5,\n            atol=1e-6,\n            msg=f\"Batch {i} (prompt_len={prompt_lens[i]}, resp_len={resp_len}): values incorrect\",\n        )\n    print(\"All varied length tests passed\")\n\n\nif __name__ == \"__main__\":\n    test_padding_conversion_with_log_probs()\n    test_padding_conversion_without_log_probs()\n    test_padding_roundtrip()\n    test_no_padding_2_padding_varying_lengths()\n    print(\"All padding conversion tests passed!\")\n"
  },
  {
    "path": "tests/utils/test_prepare_micro_batches_with_group_size.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nTests for prepare_micro_batches with force_group_size > 1 and use_dynamic_bsz=True.\n\nFocuses on verifying that:\n1. Samples within the same group (consecutive force_group_size samples) always\n   end up in the same micro-batch.\n2. All original samples are covered exactly once across all micro-batches.\n3. The returned batch_idx_list correctly maps micro-batch positions back to\n   original batch positions.\n4. Token budget (max_token_len) is respected per micro-batch.\n\"\"\"\n\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl.utils import tensordict_utils as tu\nfrom verl.workers.engine.utils import prepare_micro_batches\n\n\ndef _make_batch(seq_lens: list[int], force_group_size: int, max_token_len_per_gpu: int) -> TensorDict:\n    \"\"\"Build a minimal TensorDict accepted by prepare_micro_batches.\n\n    Args:\n        seq_lens: Effective sequence length for each sample.\n        force_group_size: Group size constraint to embed in the batch.\n        max_token_len_per_gpu: Token budget per GPU to embed in the batch.\n\n    Returns:\n        A TensorDict with ``input_ids``, ``attention_mask``, and the required\n        non-tensor metadata fields.\n    \"\"\"\n    batch_size = len(seq_lens)\n    max_len = max(seq_lens)\n\n    # Build padded attention_mask: each row has seq_lens[i] ones followed by zeros.\n    attention_mask = torch.zeros(batch_size, max_len, dtype=torch.long)\n    for i, sl in enumerate(seq_lens):\n        attention_mask[i, :sl] = 1\n\n    input_ids = torch.randint(1, 100, (batch_size, max_len))\n\n    batch = TensorDict(\n        {\"input_ids\": input_ids, \"attention_mask\": attention_mask},\n        batch_size=[batch_size],\n    )\n\n    # Embed metadata that prepare_micro_batches reads via get_non_tensor_data.\n    tu.assign_non_tensor_data(batch, \"use_dynamic_bsz\", True)\n    tu.assign_non_tensor_data(batch, \"sp_size\", 1)\n    tu.assign_non_tensor_data(batch, \"force_group_size\", force_group_size)\n    tu.assign_non_tensor_data(batch, \"max_token_len_per_gpu\", max_token_len_per_gpu)\n\n    return batch\n\n\ndef _verify_group_integrity(batch_idx_list: list[list[int]], force_group_size: int, batch_size: int):\n    \"\"\"Assert that every group of force_group_size consecutive samples stays together.\n\n    Args:\n        batch_idx_list: Index lists returned by prepare_micro_batches.\n        force_group_size: Expected group size.\n        batch_size: Total number of samples in the original batch.\n    \"\"\"\n    # Build a mapping: original_sample_idx -> micro_batch_id\n    sample_to_mb = {}\n    for mb_id, indices in enumerate(batch_idx_list):\n        for idx in indices:\n            assert idx not in sample_to_mb, f\"Sample {idx} appears in multiple micro-batches\"\n            sample_to_mb[idx] = mb_id\n\n    # Every sample must be assigned.\n    assert set(sample_to_mb.keys()) == set(range(batch_size)), (\n        f\"Not all samples covered. Missing: {set(range(batch_size)) - set(sample_to_mb.keys())}\"\n    )\n\n    # Samples within the same group must share the same micro-batch.\n    num_groups = batch_size // force_group_size\n    for g in range(num_groups):\n        start = g * force_group_size\n        group_indices = list(range(start, start + force_group_size))\n        mb_ids = {sample_to_mb[i] for i in group_indices}\n        assert len(mb_ids) == 1, f\"Group {g} (samples {group_indices}) was split across micro-batches {mb_ids}\"\n\n\ndef test_force_group_size_2_basic():\n    \"\"\"Basic test: batch_size=8, force_group_size=2, dynamic bsz enabled.\"\"\"\n    # 4 groups of 2; alternating short/long sequences within each group.\n    seq_lens = [50, 60, 80, 90, 40, 45, 100, 110]\n    force_group_size = 2\n    batch_size = len(seq_lens)\n    max_token_len_per_gpu = 200\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None, \"batch_idx_list must not be None when use_dynamic_bsz=True\"\n    assert len(micro_batches) > 0\n\n    _verify_group_integrity(batch_idx_list, force_group_size, batch_size)\n\n\ndef test_force_group_size_4_basic():\n    \"\"\"Test with force_group_size=4 (e.g., 4 responses per prompt in RM training).\"\"\"\n    # 4 groups of 4 samples each.\n    seq_lens = [\n        100,\n        110,\n        90,\n        95,  # group 0\n        200,\n        210,\n        190,\n        205,  # group 1\n        50,\n        55,\n        45,\n        60,  # group 2\n        150,\n        160,\n        140,\n        155,  # group 3\n    ]\n    force_group_size = 4\n    batch_size = len(seq_lens)\n    max_token_len_per_gpu = 500\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None\n    assert len(micro_batches) > 0\n\n    _verify_group_integrity(batch_idx_list, force_group_size, batch_size)\n\n\ndef test_force_group_size_reconstruction():\n    \"\"\"Verify that micro-batches can be reconstructed back to the original batch order.\"\"\"\n    seq_lens = [80, 85, 120, 130, 60, 65, 200, 210]\n    force_group_size = 2\n    max_token_len_per_gpu = 300\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None\n\n    # Flatten micro-batches and index lists.\n    flat_input_ids = torch.cat([mb[\"input_ids\"] for mb in micro_batches], dim=0)\n    flat_indices = [idx for indices in batch_idx_list for idx in indices]\n\n    # Build reverse mapping and reconstruct.\n    reverse_idx = [0] * len(flat_indices)\n    for new_pos, orig_pos in enumerate(flat_indices):\n        reverse_idx[orig_pos] = new_pos\n\n    reconstructed = flat_input_ids[torch.tensor(reverse_idx)]\n    torch.testing.assert_close(reconstructed, batch[\"input_ids\"])\n\n\ndef test_force_group_size_single_micro_batch():\n    \"\"\"When all samples fit in one micro-batch, grouping constraint is trivially satisfied.\"\"\"\n    seq_lens = [10, 12, 15, 11, 8, 9, 14, 13]\n    force_group_size = 2\n    max_token_len_per_gpu = 10000  # very large budget\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None\n    # All samples should be in a single micro-batch.\n    assert len(micro_batches) == 1\n    assert len(batch_idx_list[0]) == len(seq_lens)\n\n    _verify_group_integrity(batch_idx_list, force_group_size, len(seq_lens))\n\n\ndef test_force_group_size_large_group():\n    \"\"\"Test with a larger batch and force_group_size=3.\"\"\"\n    # 6 groups of 3 samples each.\n    seq_lens = [\n        100,\n        105,\n        95,  # group 0\n        200,\n        205,\n        195,  # group 1\n        50,\n        55,\n        45,  # group 2\n        150,\n        155,\n        145,  # group 3\n        80,\n        85,\n        75,  # group 4\n        120,\n        125,\n        115,  # group 5\n    ]\n    force_group_size = 3\n    batch_size = len(seq_lens)\n    max_token_len_per_gpu = 400\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None\n    assert len(micro_batches) > 0\n\n    _verify_group_integrity(batch_idx_list, force_group_size, batch_size)\n\n\ndef test_force_group_size_1_unchanged():\n    \"\"\"force_group_size=1 should behave identically to the default (no grouping constraint).\"\"\"\n    seq_lens = [100, 200, 50, 150, 80, 120]\n    force_group_size = 1\n    max_token_len_per_gpu = 300\n\n    batch = _make_batch(seq_lens, force_group_size, max_token_len_per_gpu)\n    micro_batches, batch_idx_list = prepare_micro_batches(batch)\n\n    assert batch_idx_list is not None\n    assert len(micro_batches) > 0\n\n    # All samples covered exactly once.\n    all_indices = [idx for indices in batch_idx_list for idx in indices]\n    assert sorted(all_indices) == list(range(len(seq_lens)))\n"
  },
  {
    "path": "tests/utils/test_rollout_skip_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport shutil\nimport tempfile\nfrom pathlib import Path\nfrom unittest.mock import MagicMock\n\nimport pytest\nimport torch\n\nfrom verl.utils.rollout_skip import DataProto, RolloutSkip\n\nlen_prompt = 50\nlen_response = 100\n\n\ndef temp_dir():\n    # Create a temporary directory\n    temp_dir = Path(tempfile.mkdtemp())\n    yield temp_dir\n    # Cleanup\n    shutil.rmtree(temp_dir)\n\n\ndef build_generate_fn(gen_bs, n):\n    len_tokenizer = 1024\n\n    def iterate():\n        while True:\n            prompt = torch.randint(len_tokenizer, size=(gen_bs, len_prompt)).repeat_interleave(n, dim=0)\n            generate = torch.randint(len_tokenizer, size=(gen_bs * n, len_response))\n            data = DataProto.from_dict(tensors={\"prompt\": prompt, \"response\": generate})\n            yield data\n\n    mock_infer_engine = iterate()\n\n    def fn(batch, **kwargs):\n        # Simulate the inference engine returning the next batch\n        return next(mock_infer_engine)\n\n    return fn\n\n\n@pytest.fixture(params=[(32, 4), (64, 4), (64, 8)])\ndef mock_rollout_wg(request):\n    gen_bs, n = request.param\n    rollout_wg = MagicMock()\n\n    config = MagicMock()\n    config.actor_rollout_ref.rollout = {\n        \"n\": n,\n        \"skip_dump_dir\": next(temp_dir()),\n    }\n    config.data = {\"gen_batch_size\": gen_bs}\n\n    rollout_wg.generate_sequences = build_generate_fn(gen_bs, n)\n\n    yield config, rollout_wg\n    # Cleanup\n    shutil.rmtree(next(temp_dir()))\n\n\nclass TestRolloutSkip:\n    def test_initialization(self, capsys):\n        \"\"\"Test that RolloutSkip initializes correctly\"\"\"\n        config = MagicMock()\n        config.actor_rollout_ref.rollout = {\n            \"n\": 16,\n            \"skip_dump_dir\": \"tmp/rollout_dump\",\n        }\n        config.data = {\"gen_batch_size\": 128}\n        mock_rollout_wg = MagicMock()\n        skip = RolloutSkip(config, mock_rollout_wg)\n\n        assert skip.n == 16\n        assert skip.gbs == 128\n        assert str(skip.dumped_dir) == \"tmp/rollout_dump\"\n\n        assert skip._rollout_wg == mock_rollout_wg\n        skip.wrap_generate_sequences()\n        captured = capsys.readouterr()\n        assert \"Successfully patched\" in captured.out\n\n    def test_generate_without_wrap(self, mock_rollout_wg):\n        \"\"\"Test that generate_sequences works without wrapping\"\"\"\n\n        config, rollout_wg = mock_rollout_wg\n        _ = RolloutSkip(config, rollout_wg)\n\n        _result = rollout_wg.generate_sequences(MagicMock())\n        for _ in range(10):\n            result = rollout_wg.generate_sequences(MagicMock())\n            assert isinstance(result, DataProto)\n            # * make sure the data is different\n            assert torch.abs(_result.batch[\"prompt\"] - result.batch[\"prompt\"]).sum() > 0\n            assert torch.abs(_result.batch[\"response\"] - result.batch[\"response\"]).sum() > 0\n            _result = result\n\n    def test_dump(self, mock_rollout_wg, capsys):\n        config, rollout_wg = mock_rollout_wg\n        skip = RolloutSkip(config, rollout_wg)\n        skip.wrap_generate_sequences()\n\n        result = rollout_wg.generate_sequences(MagicMock())\n        # * check if dump is OK\n        assert skip.curr_path_dump.exists()\n        captured = capsys.readouterr()\n        assert \"Successfully dump data in\" in captured.out\n        # * get file size, estimate file size\n        file_size = skip.curr_path_dump.stat().st_size\n        est_file_size = (len_prompt + len_response) * skip.gbs * skip.n * result.batch[\"prompt\"].dtype.itemsize\n        assert file_size >= est_file_size, \"Dumped file size is smaller than expected\"\n\n    def test_generate_with_wrap(self, mock_rollout_wg, capsys):\n        \"\"\"Test that generate_sequences works without wrapping\"\"\"\n\n        config, rollout_wg = mock_rollout_wg\n        skip = RolloutSkip(config, rollout_wg)\n        skip.wrap_generate_sequences()\n\n        _result = rollout_wg.generate_sequences(MagicMock())\n\n        for _ in range(10):\n            result = rollout_wg.generate_sequences(MagicMock())\n            assert isinstance(result, DataProto)\n            # * make sure the data is different\n            assert torch.abs(_result.batch[\"prompt\"] - result.batch[\"prompt\"]).sum() == 0\n            assert torch.abs(_result.batch[\"response\"] - result.batch[\"response\"]).sum() == 0\n            captured = capsys.readouterr()\n            assert \"Successfully load pre-generated data from\" in captured.out\n            _result = result\n"
  },
  {
    "path": "tests/utils/test_rollout_trace_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport sys\nfrom unittest.mock import MagicMock, patch\n\nimport pytest\n\nfrom verl.utils.rollout_trace import RolloutTraceConfig, rollout_trace_attr, rollout_trace_op\n\n\n@pytest.fixture(autouse=True)\ndef reset_rollout_trace_config_singleton():\n    \"\"\"Fixture to reset the RolloutTraceConfig singleton before each test.\"\"\"\n    RolloutTraceConfig.reset()\n\n\n@pytest.fixture\ndef mock_weave_client():\n    \"\"\"Mocks the weave module and its client, yielding the mock client.\"\"\"\n    mock_weave = MagicMock()\n    mock_client = MagicMock()\n    mock_call = MagicMock()\n    mock_client.create_call.return_value = mock_call\n    mock_weave.init.return_value = mock_client\n\n    # Also mock the call_context if it's used internally by the decorator\n    mock_weave.trace.context.call_context.return_value = MagicMock()\n\n    with patch.dict(sys.modules, {\"weave\": mock_weave, \"weave.trace.context\": mock_weave.trace.context}):\n        yield mock_client\n\n\nclass TracedClass:\n    @rollout_trace_op\n    # @weave.op\n    # @mlflow.trace\n    async def my_method(self, a, b=\"default\"):\n        return f\"result: {a}, {b}\"\n\n    @rollout_trace_op\n    # @weave.op\n    # @mlflow.trace\n    async def middle_method(self, a, b=\"default\"):\n        await self.my_method(\"test_a1\", b=\"test_b1\")\n        return f\"result: {a}, {b}\"\n\n    @rollout_trace_op\n    # @mlflow.trace\n    async def my_method_with_exception(self):\n        raise ValueError(\"Test Exception\")\n\n    async def upper_method(self):\n        await self.my_method(\"test_a0\", b=\"test_b0\")\n        await self.middle_method(\"test_a2\", b=\"test_b2\")\n        return True\n\n\nclass UntracedClass:\n    @rollout_trace_op\n    async def my_method(self, x):\n        return x * 2\n\n\nasync def test_rollout_trace_on_untraced_class():\n    \"\"\"Tests that the decorator works correctly when no backend is configured.\"\"\"\n    instance = UntracedClass()\n    assert await instance.my_method(10) == 20\n\n\nasync def test_rollout_trace_with_tracer(mock_weave_client):\n    \"\"\"Tests that the decorator calls the tracer's methods correctly.\"\"\"\n    RolloutTraceConfig.init(project_name=\"my-project\", experiment_name=\"my-experiment\", backend=\"weave\")\n    instance = TracedClass()\n    assert RolloutTraceConfig.get_client() is mock_weave_client\n\n    result = await instance.my_method(\"test_a\", b=\"test_b\")\n\n    assert result == \"result: test_a, test_b\"\n    mock_weave_client.create_call.assert_called_once()\n    call_kwargs = mock_weave_client.create_call.call_args.kwargs\n    assert call_kwargs[\"op\"] == \"TracedClass.my_method\"\n    expected_inputs = {\"a\": \"test_a\", \"b\": \"test_b\"}\n    assert call_kwargs[\"inputs\"] == expected_inputs\n\n    mock_call = mock_weave_client.create_call.return_value\n    mock_weave_client.finish_call.assert_called_once_with(mock_call, output=result)\n\n\nasync def test_rollout_trace_with_exception(mock_weave_client):\n    \"\"\"Tests that `finish` is called with the exception when one is raised.\"\"\"\n    RolloutTraceConfig.init(project_name=\"my-project\", experiment_name=\"my-experiment\", backend=\"weave\")\n    instance = TracedClass()\n\n    with pytest.raises(ValueError, match=\"Test Exception\"):\n        await instance.my_method_with_exception()\n\n    mock_weave_client.create_call.assert_called_once()\n    mock_call = mock_weave_client.create_call.return_value\n    mock_weave_client.finish_call.assert_called_once()\n\n    # Check that finish_call was called with the exception\n    args, kwargs = mock_weave_client.finish_call.call_args\n    assert args[0] == mock_call\n    assert \"exception\" in kwargs\n    assert isinstance(kwargs[\"exception\"], ValueError)\n\n\nasync def test_rollout_trace_with_dummy_backend(mock_weave_client):\n    \"\"\"Tests that the tracer is not called when the backend is 'dummy'.\"\"\"\n    RolloutTraceConfig.init(project_name=\"my-project\", experiment_name=\"my-experiment\", backend=\"dummy\")\n    instance = TracedClass()\n\n    await instance.my_method(\"test_a\")\n\n    mock_weave_client.create_call.assert_not_called()\n\n\nasync def test_trace_disabled_with_trace_false(mock_weave_client):\n    \"\"\"Tests that tracing is disabled when trace=False.\"\"\"\n    RolloutTraceConfig.init(\n        project_name=\"my-project\",\n        experiment_name=\"my-experiment\",\n        backend=\"weave\",\n    )\n    instance = TracedClass()\n\n    assert RolloutTraceConfig.get_backend() == \"weave\"\n\n    with rollout_trace_attr(step=1, sample_index=0, rollout_n=0, trace=False):\n        result = await instance.my_method(\"test_a\", b=\"test_b\")\n        assert result == \"result: test_a, test_b\"\n\n    # No tracing should have occurred\n    mock_weave_client.create_call.assert_not_called()\n\n    # Verify that tracing works again with trace=True (default)\n    with rollout_trace_attr(step=1, sample_index=0, rollout_n=0):\n        result = await instance.my_method(\"test_a\", b=\"test_b\")\n        assert result == \"result: test_a, test_b\"\n\n    assert mock_weave_client.create_call.call_count == 1\n\n\nasync def test_trace_false_disables_nested_trace_ops(mock_weave_client):\n    \"\"\"Tests that trace=False disables all nested @rollout_trace_op calls.\"\"\"\n    RolloutTraceConfig.init(\n        project_name=\"my-project\",\n        experiment_name=\"my-experiment\",\n        backend=\"weave\",\n    )\n    instance = TracedClass()\n\n    with rollout_trace_attr(step=1, sample_index=0, rollout_n=0, trace=False):\n        # Call upper_method which internally calls my_method and middle_method\n        # All of these are decorated with @rollout_trace_op\n        result = await instance.upper_method()\n        assert result is True\n\n    # No tracing should have occurred for any of the nested calls\n    mock_weave_client.create_call.assert_not_called()\n\n    with rollout_trace_attr(step=1, sample_index=0, rollout_n=0):\n        result = await instance.my_method(\"test_a\", b=\"test_b\")\n        assert result == \"result: test_a, test_b\"\n\n    assert mock_weave_client.create_call.call_count == 1\n\n\nasync def test_trace_enabled_restored_after_exception(mock_weave_client):\n    \"\"\"Tests that trace state is restored even if an exception occurs when trace=False.\"\"\"\n    RolloutTraceConfig.init(\n        project_name=\"my-project\",\n        experiment_name=\"my-experiment\",\n        backend=\"weave\",\n    )\n    instance = TracedClass()\n\n    assert RolloutTraceConfig.get_backend() == \"weave\"\n\n    # Use trace=False and raise an exception\n    try:\n        with rollout_trace_attr(step=1, sample_index=0, rollout_n=0, trace=False):\n            raise RuntimeError(\"Test exception with trace disabled\")\n    except RuntimeError:\n        pass\n\n    with rollout_trace_attr(step=1, sample_index=0, rollout_n=0):\n        result = await instance.my_method(\"test_a\", b=\"test_b\")\n        assert result == \"result: test_a, test_b\"\n\n    assert mock_weave_client.create_call.call_count == 1\n\n\n@pytest.mark.skipif(\n    os.environ.get(\"RUN_WEAVE_INTEGRATION_TESTS\", \"false\").lower() != \"true\",\n    reason=\"Skipping weave integration test. Set RUN_WEAVE_INTEGRATION_TESTS=true to run.\",\n)\nasync def test_rollout_trace_with_real_weave_backend():\n    \"\"\"Integration test with a real weave backend.\"\"\"\n\n    # This assumes that the weave environment (e.g., project) is configured\n    RolloutTraceConfig.init(project_name=\"my-project\", experiment_name=\"my-experiment\", backend=\"weave\")\n\n    instance = TracedClass()\n\n    with rollout_trace_attr(step=1, sample_index=2, rollout_n=3):\n        await instance.upper_method()\n\n    with pytest.raises(ValueError, match=\"Test Exception\"):\n        await instance.my_method_with_exception()\n\n    print(\"\\nWeave integration test ran successfully. Check your weave project for the trace.\")\n\n\n@pytest.mark.skipif(\n    os.environ.get(\"RUN_MLFLOW_INTEGRATION_TESTS\", \"false\").lower() != \"true\",\n    reason=\"Skipping mlflow integration test. Set RUN_MLFLOW_INTEGRATION_TESTS=true to run.\",\n)\nasync def test_rollout_trace_with_real_mlflow_backend():\n    \"\"\"Integration test with a real mlflow backend.\"\"\"\n\n    # This assumes that the mlflow environment (e.g., project) is configured\n    RolloutTraceConfig.init(project_name=\"my-project\", experiment_name=\"my-experiment\", backend=\"mlflow\")\n\n    instance = TracedClass()\n\n    with rollout_trace_attr(step=1, sample_index=2, rollout_n=3, name=\"agent_run\"):\n        assert await instance.upper_method()\n\n    # with pytest.raises(ValueError, match=\"Test Exception\"):\n    #     await instance.my_method_with_exception()\n\n    print(\"\\nWeave integration test ran successfully. Check your weave project for the trace.\")\n"
  },
  {
    "path": "tests/utils/test_seqlen_balancing.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\n\nfrom verl import DataProto\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.model import create_random_mask\nfrom verl.utils.seqlen_balancing import (\n    ceildiv,\n    get_reverse_idx,\n    prepare_dynamic_batch,\n    rearrange_micro_batches,\n    restore_dynamic_batch,\n)\n\n\ndef test_seqlen_balancing():\n    input_ids = torch.randint(low=0, high=10, size=(20, 100))\n\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_left_padding=0.1, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.5\n    )\n    data = {\"input_ids\": input_ids, \"attention_mask\": attention_mask}\n    dataproto = DataProto.from_single_dict(data)\n    micro_batches, micro_bsz_idx_lst = rearrange_micro_batches(dataproto.batch, max_token_len=300)\n    batch = torch.cat(micro_batches)\n    micro_bsz_idx = []\n    for idx in micro_bsz_idx_lst:\n        micro_bsz_idx.extend(idx)\n    reverse_idx_map = get_reverse_idx(micro_bsz_idx)\n    reverse_idx_map = torch.tensor(reverse_idx_map)\n    new_batch = batch[reverse_idx_map]\n    torch.testing.assert_close(new_batch, dataproto.batch)\n\n\ndef test_dynamic_batch():\n    input_ids = torch.randint(low=0, high=10, size=(20, 100))\n\n    attention_mask = create_random_mask(\n        input_ids=input_ids, max_ratio_of_left_padding=0.1, max_ratio_of_valid_token=0.9, min_ratio_of_valid_token=0.5\n    )\n    data = {\"input_ids\": input_ids, \"attention_mask\": attention_mask}\n    dataproto = DataProto.from_single_dict(data)\n    micro_batches, micro_bsz_idx_lst = prepare_dynamic_batch(dataproto, max_token_len=300)\n    input_ids = torch.cat([micro_batch.batch[\"input_ids\"] for micro_batch in micro_batches], dim=0)\n    input_ids = restore_dynamic_batch(input_ids, micro_bsz_idx_lst)\n    torch.testing.assert_close(input_ids, dataproto.batch[\"input_ids\"])\n\n\ndef _worker(rank, world_size, init_method, max_token_len, use_same_dp, min_mb):\n    # 1) init process group & CUDA\n    get_torch_device().set_device(rank)\n    dist.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=init_method,\n        world_size=world_size,\n        rank=rank,\n    )\n\n    # 2) build a small random batch (each rank different length to force mismatch)\n    torch.manual_seed(42 + rank)\n    input_ids = torch.randint(0, 10, (20 + rank * 5, 100), device=f\"{get_device_name()}:{rank}\")\n    attention_mask = create_random_mask(\n        input_ids=input_ids,\n        max_ratio_of_left_padding=0.1,\n        max_ratio_of_valid_token=0.9,\n        min_ratio_of_valid_token=0.5,\n    )\n    dp = {\"input_ids\": input_ids, \"attention_mask\": attention_mask}\n    proto = DataProto.from_single_dict(dp)\n    batch = proto.batch\n\n    # 3) call rearrange_micro_batches with one of the two params under test\n    micros, idx_lst = rearrange_micro_batches(\n        batch,\n        max_token_len=max_token_len,\n        dp_group=dist.group.WORLD,\n        same_micro_num_in_dp=use_same_dp,\n        min_num_micro_batch=min_mb,\n    )\n\n    # 4) check the enforced counts\n    seq_len_effective: torch.Tensor = batch[\"attention_mask\"].sum(dim=1)\n    total_seqlen = seq_len_effective.sum().item()\n    local = min(len(seq_len_effective), ceildiv(total_seqlen, max_token_len))\n\n    if min_mb is not None:\n        expected = max(local, min_mb)\n        assert len(micros) == expected\n    if use_same_dp:\n        # gather all local_counts\n        counts = [torch.zeros(1, device=f\"{get_device_name()}:{rank}\") for _ in range(world_size)]\n        counts[rank].fill_(local)\n        dist.all_gather(counts, counts[rank])\n        expected = max(int(c.item()) for c in counts)\n        assert len(micros) == expected\n    else:\n        # if neither, we get the local natural count\n        assert len(micros) == local\n\n    # 5) reconstruction sanity: concat→reverse_idx→orig\n    flat = torch.cat(micros, dim=0)\n    idx = []\n    for sub in idx_lst:\n        idx.extend(sub)\n    inv = get_reverse_idx(idx)\n    inv = torch.tensor(inv, device=flat.device)\n    reconstructed = flat[inv]\n    torch.testing.assert_close(reconstructed, batch)\n\n    dist.destroy_process_group()\n\n\ndef test_dataproto_split_uneven():\n    \"\"\"Test DataProto.split with uneven splits\"\"\"\n    # Create test data with 10 items\n    input_ids = torch.randint(low=0, high=10, size=(10, 5))\n    attention_mask = torch.ones(10, 5)\n    data = {\"input_ids\": input_ids, \"attention_mask\": attention_mask}\n    dataproto = DataProto.from_single_dict(data)\n\n    # Test split with size 3 (should create chunks of [3, 3, 3, 1])\n    splits = dataproto.split(3)\n    assert len(splits) == 4\n    assert len(splits[0]) == 3\n    assert len(splits[1]) == 3\n    assert len(splits[2]) == 3\n    assert len(splits[3]) == 1\n\n    reconstructed = DataProto.concat(splits)\n    torch.testing.assert_close(reconstructed.batch[\"input_ids\"], dataproto.batch[\"input_ids\"])\n    torch.testing.assert_close(reconstructed.batch[\"attention_mask\"], dataproto.batch[\"attention_mask\"])\n\n    # Test split with size equal to length (should create one chunk)\n    splits = dataproto.split(10)\n    assert len(splits) == 1\n    assert len(splits[0]) == 10\n\n    # Test split with size larger than length (should create one chunk with all data)\n    splits = dataproto.split(15)\n    assert len(splits) == 1\n    assert len(splits[0]) == 10\n\n    # Test with non-tensor batch data\n    import numpy as np\n\n    data_with_non_tensor = {\n        \"input_ids\": input_ids,\n        \"attention_mask\": attention_mask,\n        \"labels\": np.array([f\"label_{i}\" for i in range(10)], dtype=object),\n    }\n    dataproto_with_non_tensor = DataProto.from_single_dict(data_with_non_tensor)\n\n    splits = dataproto_with_non_tensor.split(3)\n    assert len(splits) == 4\n    assert len(splits[0]) == 3\n    assert len(splits[1]) == 3\n    assert len(splits[2]) == 3\n    assert len(splits[3]) == 1\n\n    # Verify non-tensor data integrity\n    reconstructed = DataProto.concat(splits)\n    np.testing.assert_array_equal(\n        reconstructed.non_tensor_batch[\"labels\"], dataproto_with_non_tensor.non_tensor_batch[\"labels\"]\n    )\n\n\ndef test_seqlen_balancing_distributed_params(tmp_path):\n    world_size = 2\n    init_file = tmp_path / \"dist_init\"\n    init_file.write_text(\"\")  # empty file\n    init_method = f\"file://{init_file}\"\n\n    # test min_num_micro_batch only\n    mp.spawn(\n        _worker,\n        args=(world_size, init_method, 300, False, 4),\n        nprocs=world_size,\n        join=True,\n    )\n\n    # test same_micro_num_in_dp only\n    mp.spawn(\n        _worker,\n        args=(world_size, init_method, 300, True, None),\n        nprocs=world_size,\n        join=True,\n    )\n\n\ndef test_group_balanced_partitions():\n    \"\"\"Test group-level balancing keeps same-uid samples together.\"\"\"\n    from verl.utils.seqlen_balancing import get_group_balanced_partitions\n\n    # Create test data: 4 groups with different sizes\n    # Group 0 (uid=0): indices 0,1,2,3 with seqlens [100, 100, 100, 100]\n    # Group 1 (uid=1): indices 4,5,6,7 with seqlens [200, 200, 200, 200]\n    # Group 2 (uid=2): indices 8,9,10,11 with seqlens [150, 150, 150, 150]\n    # Group 3 (uid=3): indices 12,13,14,15 with seqlens [50, 50, 50, 50]\n    seqlen_list = [100] * 4 + [200] * 4 + [150] * 4 + [50] * 4\n    uid_list = [0] * 4 + [1] * 4 + [2] * 4 + [3] * 4\n\n    # Partition into 2 groups\n    partitions = get_group_balanced_partitions(seqlen_list, uid_list, k_partitions=2)\n\n    assert len(partitions) == 2\n\n    # Verify all indices are covered\n    all_indices = set()\n    for partition in partitions:\n        all_indices.update(partition)\n    assert all_indices == set(range(16))\n\n    # Verify same-uid samples stay together\n    for partition in partitions:\n        uids_in_partition = set(uid_list[i] for i in partition)\n        for uid in uids_in_partition:\n            # All samples with this uid should be in this partition\n            uid_indices = [i for i, u in enumerate(uid_list) if u == uid]\n            assert all(i in partition for i in uid_indices), f\"uid {uid} samples split across partitions\"\n\n\ndef test_group_balanced_partitions_single_sample_groups():\n    \"\"\"Test group balancing with single-sample groups (n=1).\"\"\"\n    from verl.utils.seqlen_balancing import get_group_balanced_partitions\n\n    # Each sample is its own group\n    seqlen_list = [100, 200, 150, 50, 300, 250]\n    uid_list = [0, 1, 2, 3, 4, 5]\n\n    partitions = get_group_balanced_partitions(seqlen_list, uid_list, k_partitions=2)\n\n    assert len(partitions) == 2\n    all_indices = set()\n    for partition in partitions:\n        all_indices.update(partition)\n    assert all_indices == set(range(6))\n\n\ndef test_group_balanced_partitions_equal_size():\n    \"\"\"Test group balancing with equal_size constraint simulation.\"\"\"\n    from verl.utils.seqlen_balancing import get_group_balanced_partitions\n\n    # 8 groups, partition into 4 (simulating world_size=4)\n    # Each group has 2 samples\n    seqlen_list = [100, 100, 200, 200, 150, 150, 50, 50, 300, 300, 250, 250, 180, 180, 120, 120]\n    uid_list = [0, 0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7, 7]\n\n    partitions = get_group_balanced_partitions(seqlen_list, uid_list, k_partitions=4)\n\n    assert len(partitions) == 4\n\n    # Verify all indices are covered\n    all_indices = set()\n    for partition in partitions:\n        all_indices.update(partition)\n    assert all_indices == set(range(16))\n\n    # Verify same-uid samples stay together\n    for partition in partitions:\n        uids_in_partition = set(uid_list[i] for i in partition)\n        for uid in uids_in_partition:\n            uid_indices = [i for i, u in enumerate(uid_list) if u == uid]\n            assert all(i in partition for i in uid_indices)\n"
  },
  {
    "path": "tests/utils/test_server_profiler.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nimport unittest\nfrom unittest.mock import AsyncMock, MagicMock, patch\n\nfrom verl.utils.profiler.config import (\n    ProfilerConfig,\n    TorchProfilerToolConfig,\n    build_sglang_profiler_args,\n    build_vllm_profiler_args,\n)\n\n\nclass TestServerProfilerArgs(unittest.TestCase):\n    def test_build_vllm_profiler_args(self):\n        # Case 1: All features enabled\n        tool_config = TorchProfilerToolConfig(contents=[\"stack\", \"shapes\", \"memory\"])\n        config = ProfilerConfig(save_path=\"/tmp/test\", tool_config=tool_config)\n\n        # Patch environ to avoid side effects and verify calls\n        with patch.dict(os.environ, {}, clear=True):\n            args = build_vllm_profiler_args(config, tool_config, rank=0)\n\n            # Check Env vars (backward compatibility)\n            self.assertEqual(os.environ.get(\"VLLM_TORCH_PROFILER_DIR\"), \"/tmp/test/agent_loop_rollout_replica_0\")\n            self.assertEqual(os.environ.get(\"VLLM_TORCH_PROFILER_WITH_STACK\"), \"1\")\n            self.assertEqual(os.environ.get(\"VLLM_TORCH_PROFILER_RECORD_SHAPES\"), \"1\")\n            self.assertEqual(os.environ.get(\"VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY\"), \"1\")\n\n            # Check Args (new API)\n            self.assertIn(\"profiler_config\", args)\n            profiler_config_dict = json.loads(args[\"profiler_config\"])\n            self.assertEqual(profiler_config_dict[\"torch_profiler_dir\"], \"/tmp/test/agent_loop_rollout_replica_0\")\n            self.assertTrue(profiler_config_dict[\"torch_profiler_with_stack\"])\n            self.assertTrue(profiler_config_dict[\"torch_profiler_record_shapes\"])\n            self.assertTrue(profiler_config_dict[\"torch_profiler_with_memory\"])\n\n    def test_build_sglang_profiler_args(self):\n        # Case 1: Basic features\n        tool_config = TorchProfilerToolConfig(contents=[\"stack\", \"shapes\", \"memory\"])\n        config = ProfilerConfig(save_path=\"/tmp/test\", tool_config=tool_config)\n        with self.assertWarns(UserWarning):\n            args = build_sglang_profiler_args(config, tool_config, rank=0)\n        self.assertEqual(args[\"output_dir\"], \"/tmp/test/agent_loop_rollout_replica_0\")\n        self.assertTrue(args[\"with_stack\"])\n        self.assertTrue(args[\"record_shapes\"])\n\n\nclass TestServerProfilerFunctionality(unittest.IsolatedAsyncioTestCase):\n    async def test_vllm_start_stop_profile(self):\n        try:\n            # Import strictly inside test to avoid import errors if dependencies missing\n            from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMHttpServer\n        except ImportError:\n            self.skipTest(\"vllm or dependencies not installed\")\n            return\n\n        # Mock dependencies\n        mock_profiler = MagicMock()\n        mock_profiler.check_enable.return_value = True\n        mock_profiler.check_this_rank.return_value = True\n        mock_profiler.is_discrete_mode.return_value = True\n\n        mock_engine = AsyncMock()\n\n        # Mock self object\n        mock_self = MagicMock()\n        mock_self.profiler_controller = mock_profiler\n        mock_self.engine = mock_engine\n\n        # Test start_profile using the unbound method\n        await vLLMHttpServer.start_profile(mock_self)\n        mock_engine.start_profile.assert_called_once()\n\n        # Test stop_profile\n        await vLLMHttpServer.stop_profile(mock_self)\n        mock_engine.stop_profile.assert_called_once()\n\n    async def test_sglang_start_stop_profile(self):\n        try:\n            # Import strictly inside test to avoid import errors if dependencies missing\n            from verl.workers.rollout.sglang_rollout.async_sglang_server import SGLangHttpServer\n        except ImportError:\n            self.skipTest(\"sglang or dependencies not installed\")\n            return\n\n        # Mock dependencies\n        mock_profiler = MagicMock()\n        mock_profiler.check_enable.return_value = True\n        mock_profiler.check_this_rank.return_value = True\n        mock_profiler.is_discrete_mode.return_value = True\n        mock_profiler.config = MagicMock()\n        mock_profiler.tool_config = MagicMock()\n\n        mock_tokenizer_manager = AsyncMock()\n\n        mock_self = MagicMock()\n        mock_self.profiler_controller = mock_profiler\n        mock_self.tokenizer_manager = mock_tokenizer_manager\n        mock_self.replica_rank = 0\n\n        # Mock build_sglang_profiler_args to return known dict\n        with patch(\"verl.workers.rollout.sglang_rollout.async_sglang_server.build_sglang_profiler_args\") as mock_build:\n            mock_args = {\"arg1\": \"val1\"}\n            mock_build.return_value = mock_args\n\n            # Test start_profile\n            await SGLangHttpServer.start_profile(mock_self)\n\n            mock_build.assert_called_once_with(mock_profiler.config, mock_profiler.tool_config, mock_self.replica_rank)\n            mock_tokenizer_manager.start_profile.assert_called_once_with(**mock_args)\n\n            # Test stop_profile\n            await SGLangHttpServer.stop_profile(mock_self)\n            mock_tokenizer_manager.stop_profile.assert_called_once()\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_shared_memory.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport multiprocessing\nimport unittest\nfrom multiprocessing import shared_memory\n\nimport torch\n\nfrom verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import create_shared_memory, rebuild_shared_memory\n\n\nclass TestSharedMemory(unittest.TestCase):\n    \"\"\"Test cases for shared memory utility functions.\"\"\"\n\n    def setUp(self):\n        \"\"\"Set up test fixtures before each test method.\"\"\"\n        # Use short unique names to avoid POSIX shared memory name length limits\n        import uuid\n\n        short_id = uuid.uuid4().hex[:8]\n        self.test_name = f\"shm_{short_id}\"\n\n    def tearDown(self):\n        \"\"\"Clean up shared memory after each test method.\"\"\"\n        # Note: We're relying on the OS to clean up shared memory\n        # as we properly delete all references in the tests\n        pass\n\n    def test_create_shared_memory_new(self):\n        \"\"\"Test creating new shared memory with unique name.\"\"\"\n        size = 1024\n\n        shm = create_shared_memory(size, self.test_name)\n\n        # Verify shared memory object is created correctly\n        self.assertIsNotNone(shm)\n        # Note: shared memory may have system-dependent size rounding\n        self.assertGreaterEqual(shm.size, size)\n        self.assertEqual(shm.name, self.test_name)\n\n        # Clean up - delete tensor references first\n        del shm\n\n    def test_create_shared_memory_attach_existing(self):\n        \"\"\"Test that create_shared_memory attaches to existing shared memory when FileExistsError occurs.\"\"\"\n        size = 2048\n\n        # First, create shared memory\n        shm1 = create_shared_memory(size, self.test_name)\n        self.assertGreaterEqual(shm1.size, size)\n\n        # Second call should attach to existing memory\n        shm2 = create_shared_memory(size, self.test_name)\n\n        # Verify we attached to the same shared memory\n        self.assertIsNotNone(shm2)\n        self.assertGreaterEqual(shm2.size, size)\n        self.assertEqual(shm2.name, self.test_name)\n\n        # Both should reference the same shared memory\n        self.assertEqual(shm1.name, shm2.name)\n\n        # Clean up\n        del shm1, shm2\n\n    def test_rebuild_shared_memory_default_dtype(self):\n        \"\"\"Test rebuilding tensor from shared memory with default dtype (uint8).\"\"\"\n        size = 1024\n\n        # Create and write to shared memory\n        shm = create_shared_memory(size, self.test_name)\n        test_data = torch.arange(size, dtype=torch.uint8)\n        shm.buf[:size] = test_data.numpy().tobytes()\n\n        # Rebuild tensor from shared memory\n        tensor, _ = rebuild_shared_memory(self.test_name, size)\n\n        # Verify tensor properties\n        self.assertEqual(tensor.dtype, torch.uint8)\n        self.assertEqual(len(tensor), size)\n\n        # Verify data integrity\n        reconstructed = torch.frombuffer(shm.buf[:size], dtype=torch.uint8)\n        self.assertTrue(torch.equal(tensor, reconstructed))\n\n        # Clean up - delete references before closing\n        del tensor, reconstructed\n\n    def test_rebuild_shared_memory_custom_dtype(self):\n        \"\"\"Test rebuilding tensor from shared memory with custom dtype.\"\"\"\n        size = 256  # 256 bytes = 64 float32 values\n\n        # Create and write to shared memory\n        shm = create_shared_memory(size, self.test_name)\n        test_data = torch.arange(64, dtype=torch.float32)\n        shm.buf[:size] = test_data.numpy().tobytes()\n\n        # Rebuild tensor with custom dtype\n        tensor, _ = rebuild_shared_memory(self.test_name, size, dtype=torch.float32)\n\n        # Verify tensor properties\n        self.assertEqual(tensor.dtype, torch.float32)\n        self.assertEqual(len(tensor), 64)\n\n        # Verify data integrity\n        reconstructed = torch.frombuffer(shm.buf[:size], dtype=torch.float32)\n        self.assertTrue(torch.equal(tensor, reconstructed))\n\n        # Clean up - delete references before closing\n        del tensor, reconstructed\n\n    def test_shared_memory_data_integrity(self):\n        \"\"\"Test that data remains intact between create and rebuild operations.\"\"\"\n        size = 512\n\n        # Create test data with various patterns\n        test_data = torch.randint(0, 256, (size,), dtype=torch.uint8)\n\n        # Create shared memory and write data\n        shm = create_shared_memory(size, self.test_name)\n        shm.buf[:size] = test_data.numpy().tobytes()\n\n        # Rebuild tensor\n        tensor, _ = rebuild_shared_memory(self.test_name, size)\n\n        # Verify data integrity\n        reconstructed = torch.frombuffer(shm.buf[:size], dtype=torch.uint8)\n        self.assertTrue(torch.equal(test_data, reconstructed))\n\n        # Clean up - delete references before closing\n        del tensor, reconstructed\n\n    def test_shared_memory_different_dtypes(self):\n        \"\"\"Test shared memory operations with different tensor dtypes.\"\"\"\n        test_cases = [\n            (torch.float32, 256, 64),  # 256 bytes / 4 bytes = 64 values\n            (torch.float64, 256, 32),  # 256 bytes / 8 bytes = 32 values\n            (torch.int32, 256, 64),  # 256 bytes / 4 bytes = 64 values\n            (torch.int64, 256, 32),  # 256 bytes / 8 bytes = 32 values\n            (torch.uint8, 256, 256),  # 256 bytes / 1 byte = 256 values\n        ]\n\n        for dtype, size, expected_len in test_cases:\n            # Create test data\n            test_data = torch.arange(expected_len, dtype=dtype)\n\n            # Create shared memory and write data\n            shm = create_shared_memory(size, self.test_name)\n            shm.buf[:size] = test_data.numpy().tobytes()\n\n            # Rebuild tensor\n            tensor, _ = rebuild_shared_memory(self.test_name, size, dtype=dtype)\n\n            # Verify properties and data\n            self.assertEqual(tensor.dtype, dtype)\n            self.assertEqual(len(tensor), expected_len)\n\n            reconstructed = torch.frombuffer(shm.buf[:size], dtype=dtype)\n            self.assertTrue(torch.equal(test_data, reconstructed))\n\n            # Clean up - delete references before closing\n            del tensor, reconstructed\n\n    def test_shared_memory_multiple_operations(self):\n        \"\"\"Test multiple create/rebuild operations with the same name.\"\"\"\n        size = 512\n\n        # First iteration\n        test_data1 = torch.arange(size, dtype=torch.uint8)\n        shm1 = create_shared_memory(size, self.test_name)\n        shm1.buf[:size] = test_data1.numpy().tobytes()\n        tensor1, _ = rebuild_shared_memory(self.test_name, size)\n        reconstructed1 = torch.frombuffer(shm1.buf[:size], dtype=torch.uint8)\n        self.assertTrue(torch.equal(test_data1, reconstructed1))\n        del tensor1, reconstructed1, shm1\n\n        # Second iteration with different data\n        test_data2 = torch.arange(size, dtype=torch.uint8) * 2\n        shm2 = create_shared_memory(size, self.test_name)\n        shm2.buf[:size] = test_data2.numpy().tobytes()\n        tensor2, _ = rebuild_shared_memory(self.test_name, size)\n        reconstructed2 = torch.frombuffer(shm2.buf[:size], dtype=torch.uint8)\n        self.assertTrue(torch.equal(test_data2, reconstructed2))\n        del tensor2, reconstructed2, shm2\n\n\n# Module-level function for cross-process testing\ndef child_process_function(name, size, test_data_bytes):\n    \"\"\"Child process function to rebuild and verify tensor.\"\"\"\n    shm = None\n    tensor = None\n    test_data = None\n    try:\n        # Convert bytes back to tensor\n        test_data = torch.frombuffer(test_data_bytes, dtype=torch.uint8)\n\n        # Attach to shared memory\n        shm = shared_memory.SharedMemory(name=name)\n\n        # Rebuild tensor from shared memory\n        tensor = torch.frombuffer(shm.buf[:size], dtype=torch.uint8)\n\n        # Verify data integrity\n        assert torch.equal(test_data, tensor), \"Data mismatch in child process\"\n        return True\n    except Exception as e:\n        print(f\"Error in child process: {e}\")\n        return False\n    finally:\n        # Clean up shared memory in child process\n        # Delete all references first\n        del tensor, test_data\n        if shm is not None:\n            shm.close()\n            # Note: Don't unlink in child process, parent will clean up\n\n\nclass TestSharedMemoryIntegration(unittest.TestCase):\n    \"\"\"Integration tests for shared memory operations across process boundaries.\"\"\"\n\n    def test_cross_process_shared_memory(self):\n        \"\"\"Test shared memory can be created in one process and accessed in another.\"\"\"\n        size = 1024\n        test_data = torch.arange(size, dtype=torch.uint8)\n\n        # Create shared memory in parent process\n        shm = create_shared_memory(size, \"test_cross_proc\")\n        shm.buf[:size] = test_data.numpy().tobytes()\n\n        # Convert tensor to bytes for passing to child process\n        test_data_bytes = test_data.numpy().tobytes()\n\n        # Start child process\n        process = multiprocessing.Process(\n            target=child_process_function, args=(\"test_cross_proc\", size, test_data_bytes)\n        )\n        process.start()\n        process.join(timeout=5)\n\n        # Verify child process completed successfully\n        self.assertEqual(process.exitcode, 0, \"Child process failed\")\n\n        # Clean up\n        del shm\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_special_linear_cross_entropy_tp.py",
    "content": "#\n# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport torch\nimport torch.distributed as dist\n\ntry:\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\nexcept ImportError:\n    # FIXME: remove these manually included paths\n    import sys\n\n    sys.path.append(os.path.abspath(os.path.join(os.path.dirname(os.path.abspath(__file__)), \"../../\")))\nfinally:\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\n\nimport verl.utils.torch_functional as verl_F\n\ncompute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True)\n\nMAX_TEST_CASES = os.environ.get(\"MAX_TEST_CASES\", 5)\nVERIFY_TORCH_SELF = os.environ.get(\"VERIFY_TORCH_SELF\", False)\nLOW_MEMORY = os.environ.get(\"LOW_MEMORY\", False)\nLOW_MEMORY_DIV_FACTOR = os.environ.get(\"LOW_MEMORY_DIV_FACTOR\", 16)\n\n\ndef run_torch_entropy(\n    hidden: torch.Tensor, weight: torch.Tensor, labels: torch.Tensor, temperature: float, reduction=\"none\"\n) -> list[torch.Tensor]:\n    # [num_tokens, vocab_size]\n    if len(hidden.shape) > 2:\n        hidden = hidden.view(-1, hidden.shape[-1])  # [num_tokens, hidden_size]\n    if len(labels.shape) > 1:\n        labels = labels.view(-1)\n    logits = torch.matmul(\n        hidden.to(torch.float32),\n        weight.to(torch.float32) if weight.size(0) == hidden.size(1) else weight.T.to(torch.float32),\n    )\n    logits /= temperature\n    pd = torch.nn.functional.softmax(logits, dim=-1)  # [num_tokens, vocab_size]\n    entropy_a = torch.logsumexp(logits, dim=-1)  # [num_tokens]\n    entropy_b = torch.sum(pd * logits, dim=-1)  # [num_tokens]\n    entropy = entropy_a - entropy_b\n    logprobs = torch.nn.functional.cross_entropy(logits, labels, reduction=reduction)  # [num_tokens]\n    logprobs = torch.neg(logprobs)\n    return logprobs, entropy\n\n\nclass TorchEntropyTP(torch.autograd.Function):\n    \"\"\"\n    it is used for testing the correctness of the kernel\n    it is not efficient and is not recommended to use in practice\n    \"\"\"\n\n    @staticmethod\n    def forward(\n        ctx,\n        hidden: torch.Tensor,\n        weight: torch.Tensor,\n        labels: torch.Tensor,\n        temperature: float,\n        dist_process_group: torch.distributed.ProcessGroup,\n    ):\n        # weight has shape [vocab_size, hidden_size], hidden has shape [num_tokens, hidden_size]\n        ctx.original_hidden_shape = hidden.shape\n        if len(hidden.shape) > 2:\n            hidden = hidden.view(-1, hidden.shape[-1])  # [num_tokens, hidden_size]\n        if len(labels.shape) > 1:\n            labels = labels.view(-1)\n\n        logits = torch.matmul(hidden.to(torch.float32), weight.to(torch.float32).T)  # [num_tokens, vocab_size]\n        logits /= temperature\n        whole_logits = torch.empty(\n            (logits.shape[0], logits.shape[1] * dist.get_world_size(dist_process_group)),\n            dtype=logits.dtype,\n            device=logits.device,\n        )\n        whole_logits_ref = [\n            whole_logits[:, i * logits.shape[1] : (i + 1) * logits.shape[1]]\n            for i in range(dist.get_world_size(dist_process_group))\n        ]\n        dist.all_gather(whole_logits_ref, logits, group=dist_process_group)\n\n        pd = torch.nn.functional.softmax(whole_logits, dim=-1)\n        entropy_a = torch.logsumexp(whole_logits, dim=-1)  # [num_tokens]\n        entropy_b = torch.sum(pd * whole_logits, dim=-1)  # [num_tokens]\n        entropy = entropy_a - entropy_b\n\n        logprobs = torch.nn.functional.cross_entropy(whole_logits, labels, reduction=\"none\")\n        logprobs = torch.neg(logprobs)\n\n        ctx.save_for_backward(hidden, weight, labels, whole_logits, entropy_b)\n        ctx.dist_process_group = dist_process_group\n        ctx.temperature = temperature\n        return logprobs, entropy\n\n    @staticmethod\n    def backward(ctx, g_logprobs: torch.Tensor, g_entropy: torch.Tensor):\n        hidden, weight, labels, whole_logits, entropy_b = ctx.saved_tensors\n        dist_process_group = ctx.dist_process_group\n        temperature = ctx.temperature\n        batch_size, hidden_size = hidden.shape\n        vocab_size, hidden_size = weight.shape\n        rank = dist.get_rank(dist_process_group)\n\n        # Compute softmax probabilities\n        maximum, _ = torch.max(whole_logits, dim=-1, keepdim=True)\n        exp_logits = torch.exp(whole_logits - maximum)\n        accumulate = exp_logits.sum(dim=-1, keepdim=True)\n        pd = exp_logits / accumulate\n\n        # Gradient for entropy\n        # entropy = entropy_a - entropy_b\n        # entropy_a = log(sum(exp(logits)))\n        # entropy_b = sum(pd * logits)\n        # d_entropy_a/d_logits = pd\n        # d_entropy_b/d_logits = pd * (logits - b.unsqueeze(1) + 1)\n        # d_entropy/d_logits = d_entropy_a - d_entropy_b\n        # d_entropy/d_logits = pd - pd * (logits - b.unsqueeze(1) + 1)\n        # d_entropy/d_logits = -pd * (logits - b.unsqueeze(1))\n        d_logits_entropy = g_entropy.unsqueeze(1) * (-pd * (whole_logits - entropy_b.unsqueeze(1)))\n\n        # Gradient for logprobs\n        # logprobs = -cross_entropy = -log(pd[labels])\n        # d_logprobs/d_logits = (pd - one_hot(labels))\n        one_hot = torch.zeros_like(whole_logits)\n        one_hot.scatter_(1, labels.unsqueeze(1), 1)\n        g_logprobs = torch.neg(g_logprobs)\n        d_logits_logprobs = g_logprobs.unsqueeze(1) * (pd - one_hot)\n        # NOTE: This will lead to wrong result\n        # d_logits_logprobs = g_logprobs.unsqueeze(1) * (pd - 1) * one_hot\n\n        # Combine gradients\n        d_logits = d_logits_entropy + d_logits_logprobs\n        d_logits /= temperature\n\n        # Get local slice of gradients\n        local_d_logits = d_logits[:, rank * vocab_size : (rank + 1) * vocab_size]\n\n        # Compute gradients for hidden and weight\n        d_hidden = torch.matmul(local_d_logits, weight.to(torch.float32))\n        d_weight = torch.matmul(local_d_logits.T, hidden.to(torch.float32))\n        d_hidden = d_hidden.view(ctx.original_hidden_shape)\n\n        return d_hidden, d_weight, None, None, None\n\n\nrun_torch_entropy_tp = TorchEntropyTP.apply\n\n\nclass TestLinearCrossEntropy_TensorParallel:\n    def __init__(self):\n        dist.init_process_group(backend=\"nccl\")\n        self.group = dist.group.WORLD\n\n        self.local_rank = dist.get_rank(self.group)\n        self.world_size = dist.get_world_size(self.group)\n        device = torch.device(f\"cuda:{self.local_rank}\")\n        torch.cuda.set_device(device)\n        print(f\"[INFO]: Local rank: {self.local_rank}, World size: {self.world_size}\")\n\n    def initialize(self, test_case_idx: int, temperature: float = 1.5):\n        self.test_case_idx = test_case_idx\n        self.temperature = temperature\n\n    def shutdown(self):\n        dist.destroy_process_group()\n\n    def cleanup(self):\n        torch.cuda.empty_cache()\n        torch.cuda.reset_peak_memory_stats()\n        import gc\n\n        gc.collect()\n        torch.cuda.synchronize()\n\n    def generate_hyper(self):\n        global LOW_MEMORY, LOW_MEMORY_DIV_FACTOR, MAX_TEST_CASES\n\n        self.dtype = torch.bfloat16\n        if self.test_case_idx == 0:\n            self.batch_size = 1\n            self.num_tokens = 1937\n            self.hidden_size = 3584\n            self.vocab_size = 152064\n        elif self.test_case_idx == 1:\n            self.batch_size = 1\n            self.num_tokens = 2169\n            self.hidden_size = 896\n            self.vocab_size = 151936\n        elif self.test_case_idx == 2:\n            self.batch_size = 1\n            self.num_tokens = 1530\n            self.hidden_size = 2048\n            self.vocab_size = 32256\n        elif self.test_case_idx == 3:\n            self.batch_size = 1\n            self.num_tokens = 1388\n            self.hidden_size = 4096\n            self.vocab_size = 102400\n        elif self.test_case_idx == 4:\n            self.batch_size = 1\n            self.num_tokens = 8192\n            self.hidden_size = 4096\n            self.vocab_size = 102400\n        else:\n            raise ValueError(f\"Invalid test case index: {self.test_case_idx}\")\n        if LOW_MEMORY:\n            self.vocab_size = int(self.vocab_size / LOW_MEMORY_DIV_FACTOR)\n        assert MAX_TEST_CASES <= 5, \"MAX_TEST_CASES should be less than or equal to 5.\"\n\n    def generate_forward_inputs(self):\n        hidden = (\n            torch.empty((self.batch_size, self.num_tokens, self.hidden_size), dtype=self.dtype, device=\"cuda\")\n            .uniform_(-0.5, 0.5)\n            .requires_grad_()\n        )\n        weight = (\n            torch.empty((self.vocab_size, self.hidden_size), dtype=self.dtype, device=\"cuda\")\n            .uniform_(-0.5, 0.5)\n            .requires_grad_()\n        )\n        labels = torch.randint(0, self.vocab_size, (self.batch_size, self.num_tokens), device=\"cuda\")\n        return hidden, weight, labels\n\n    def generate_backward_inputs(self):\n        g_entropy = torch.empty((self.num_tokens,), dtype=self.dtype, device=\"cuda\").uniform_(-0.5, 0.5)\n        g_logprobs = torch.empty((self.num_tokens,), dtype=self.dtype, device=\"cuda\").uniform_(-1, 1)\n        return g_entropy, g_logprobs\n\n    def verify_torch_itself(self, iterations: int = 5):\n        self.cleanup()\n        self.generate_hyper()\n\n        for i in range(iterations):\n            hidden, weight, labels = self.generate_forward_inputs()\n\n            # NOTE: we need to manually synchronize hidden and labels among Process Group\n            dist.broadcast(hidden, src=0, group=self.group)\n            dist.broadcast(labels, src=0, group=self.group)\n\n            # forward pass\n            # Create a tensor to hold the gathered weights from all ranks\n            # weight has shape [vocab_size, hidden_size]\n            # We want to gather along the first dimension to get [vocab_size * world_size, hidden_size]\n\n            # Create a single contiguous tensor to hold all gathered weights\n            whole_weight = torch.empty(\n                (self.vocab_size * self.world_size, self.hidden_size), dtype=weight.dtype, device=weight.device\n            )\n\n            # Create views into the tensor for each rank's portion\n            whole_weight_views = [\n                whole_weight[i * self.vocab_size : (i + 1) * self.vocab_size] for i in range(self.world_size)\n            ]\n\n            # Perform all_gather operation using the views\n            dist.all_gather(whole_weight_views, weight, group=self.group)\n\n            # Set requires_grad for autograd\n            whole_weight.requires_grad_()\n\n            (single_logprobs, single_entropy) = run_torch_entropy(hidden, whole_weight, labels, self.temperature)\n\n            (tp_logprobs, tp_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group)\n\n            torch.testing.assert_close(single_logprobs, tp_logprobs, atol=1e-4, rtol=1e-4)\n            torch.testing.assert_close(single_entropy, tp_entropy, atol=1e-4, rtol=1e-4)\n\n            # backward pass\n            g_entropy, g_logprobs = self.generate_backward_inputs()\n            # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group\n            dist.broadcast(g_entropy, src=0, group=self.group)\n            dist.broadcast(g_logprobs, src=0, group=self.group)\n\n            (single_d_hidden, single_d_weight) = torch.autograd.grad(\n                (single_entropy, single_logprobs), (hidden, whole_weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n\n            (tp_d_hidden, tp_d_weight) = torch.autograd.grad(\n                (tp_entropy, tp_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            # NOTE: all-reduce on hidden is conducted outside the kernel\n            dist.all_reduce(tp_d_hidden, op=dist.ReduceOp.SUM, group=self.group)\n\n            torch.testing.assert_close(tp_d_hidden, single_d_hidden, atol=1e-2, rtol=1e-4)\n            # Extract the corresponding slice from single_d_weight for comparison\n            # tp_d_weight has shape [vocab_size, hidden_size]\n            # single_d_weight has shape [vocab_size * world_size, hidden_size]\n            torch.testing.assert_close(\n                tp_d_weight,\n                single_d_weight[self.local_rank * self.vocab_size : (self.local_rank + 1) * self.vocab_size],\n                atol=1e-2,\n                rtol=1e-4,\n            )\n\n            # atol=1e-3, rtol=1e-4)\n        if self.local_rank == 0:\n            print(\"[PASS] torch TP correctness is verified\")\n\n    def check_torch_storage(self):\n        self.cleanup()\n        self.generate_hyper()\n\n        hidden, weight, labels = self.generate_forward_inputs()\n\n        # NOTE: we need to manually synchronize hidden and labels among Process Group\n        dist.broadcast(hidden, src=0, group=self.group)\n        dist.broadcast(labels, src=0, group=self.group)\n\n        torch.cuda.reset_peak_memory_stats()\n        (tp_logprobs, tp_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group)\n        torch.cuda.synchronize()\n        forward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n\n        g_entropy, g_logprobs = self.generate_backward_inputs()\n        # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group\n        dist.broadcast(g_entropy, src=0, group=self.group)\n        dist.broadcast(g_logprobs, src=0, group=self.group)\n\n        torch.cuda.reset_peak_memory_stats()\n        (d_tp_hidden, d_tp_weight) = torch.autograd.grad(\n            (tp_entropy, tp_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n        )\n        torch.cuda.synchronize()\n        backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n        # NOTE: all-reduce on hidden is conducted outside the kernel\n        dist.all_reduce(d_tp_hidden, op=dist.ReduceOp.SUM, group=self.group)\n\n        if self.local_rank == 0:\n            print(f\"[INFO]: Torch Forward pass peak memory: {forward_max_memory:.2f} MB\")\n            print(f\"[INFO]: Torch Backward pass peak memory: {backward_max_memory:.2f} MB\")\n\n    def verify_kernel_correctness(self, iterations: int = 5):\n        self.cleanup()\n        self.generate_hyper()\n\n        torch_forward_latency = list()\n        torch_backward_latency = list()\n        kernel_forward_latency = list()\n        kernel_backward_latency = list()\n\n        start_event = torch.cuda.Event(enable_timing=True)\n        end_event = torch.cuda.Event(enable_timing=True)\n\n        for i in range(iterations):\n            hidden, weight, labels = self.generate_forward_inputs()\n\n            # NOTE: we need to manually synchronize hidden and labels among Process Group\n            dist.broadcast(hidden, src=0, group=self.group)\n            dist.broadcast(labels, src=0, group=self.group)\n\n            start_event.record()\n            (torch_logprobs, torch_entropy) = run_torch_entropy_tp(hidden, weight, labels, self.temperature, self.group)\n            end_event.record()\n            torch.cuda.synchronize()\n            torch_forward_latency.append(start_event.elapsed_time(end_event))\n\n            start_event.record()\n            (kernel_logprobs, kernel_entropy) = linear_cross_entropy(\n                hidden, weight, labels, self.temperature, \"none\", self.group\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            kernel_forward_latency.append(start_event.elapsed_time(end_event))\n\n            torch.testing.assert_close(torch_logprobs, kernel_logprobs, atol=1e-1, rtol=1e-2)\n            torch.testing.assert_close(torch_entropy, kernel_entropy, atol=1e-1, rtol=1e-2)\n\n            # backward pass\n            g_entropy, g_logprobs = self.generate_backward_inputs()\n            # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group\n            dist.broadcast(g_entropy, src=0, group=self.group)\n            dist.broadcast(g_logprobs, src=0, group=self.group)\n\n            start_event.record()\n            (torch_d_hidden, torch_d_weight) = torch.autograd.grad(\n                (torch_entropy, torch_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            torch_backward_latency.append(start_event.elapsed_time(end_event))\n            # NOTE: all-reduce on hidden is conducted outside the kernel\n            dist.all_reduce(torch_d_hidden, op=dist.ReduceOp.SUM, group=self.group)\n\n            start_event.record()\n            (kernel_d_hidden, kernel_d_weight) = torch.autograd.grad(\n                (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n            )\n            end_event.record()\n            torch.cuda.synchronize()\n            kernel_backward_latency.append(start_event.elapsed_time(end_event))\n            # NOTE: all-reduce on hidden is conducted outside the kernel\n            dist.all_reduce(kernel_d_hidden, op=dist.ReduceOp.SUM, group=self.group)\n\n            torch.testing.assert_close(torch_d_hidden, kernel_d_hidden, atol=2e-2, rtol=4e-2)\n            torch.testing.assert_close(torch_d_weight, kernel_d_weight, atol=2e-2, rtol=4e-2)\n\n        # remove first latency\n        torch_forward_latency = torch_forward_latency[1:]\n        torch_backward_latency = torch_backward_latency[1:]\n        kernel_forward_latency = kernel_forward_latency[1:]\n        kernel_backward_latency = kernel_backward_latency[1:]\n\n        if self.local_rank == 0:\n            print(\"\\n[PASS]: Verified kernel forward & backward correctness.\")\n\n            print(\n                f\"[INFO]: Forward pass: Torch implementation average time: \"\n                f\"{sum(torch_forward_latency) / len(torch_forward_latency):.2f} ms\"\n            )\n            print(\n                f\"[INFO]: Backward pass: torch implementation average time: \"\n                f\"{sum(torch_backward_latency) / len(torch_backward_latency):.2f} ms\"\n            )\n            print(\n                f\"[INFO]: Forward pass: Kernel implementation average time: \"\n                f\"{sum(kernel_forward_latency) / len(kernel_forward_latency):.2f} ms\"\n            )\n            print(\n                f\"[INFO]: Backward pass: kernel implementation average time: \"\n                f\"{sum(kernel_backward_latency) / len(kernel_backward_latency):.2f} ms\"\n            )\n\n    def check_kernel_storage(self):\n        self.cleanup()\n        self.generate_hyper()\n\n        hidden, weight, labels = self.generate_forward_inputs()\n\n        # NOTE: we need to manually synchronize hidden and labels among Process Group\n        dist.broadcast(hidden, src=0, group=self.group)\n        dist.broadcast(labels, src=0, group=self.group)\n\n        torch.cuda.reset_peak_memory_stats()\n        (kernel_logprobs, kernel_entropy) = linear_cross_entropy(\n            hidden, weight, labels, self.temperature, \"none\", self.group\n        )\n        torch.cuda.synchronize()\n        kernel_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n\n        g_entropy, g_logprobs = self.generate_backward_inputs()\n        # NOTE: we need to manually synchronize g_entropy and g_logprobs among Process Group\n        dist.broadcast(g_entropy, src=0, group=self.group)\n        dist.broadcast(g_logprobs, src=0, group=self.group)\n\n        torch.cuda.reset_peak_memory_stats()\n        (d_kernel_hidden, d_kernel_weight) = torch.autograd.grad(\n            (kernel_entropy, kernel_logprobs), (hidden, weight), (g_entropy, g_logprobs), retain_graph=False\n        )\n        torch.cuda.synchronize()\n        kernel_backward_max_memory = torch.cuda.max_memory_allocated() / 1024 / 1024\n        # NOTE: all-reduce on hidden is conducted outside the kernel\n        dist.all_reduce(d_kernel_hidden, op=dist.ReduceOp.SUM, group=self.group)\n\n        if self.local_rank == 0:\n            print(f\"[INFO]: Kernel Forward pass peak memory: {kernel_max_memory:.2f} MB\")\n            print(f\"[INFO]: Kernel Backward pass peak memory: {kernel_backward_max_memory:.2f} MB\")\n\n\nif __name__ == \"__main__\":\n    # TP command: torchrun --standalone --nnodes=1 --nproc-per-node=2 tests/kernels/test_linear_cross_entropy_tp.py\n\n    # Check if running with torchrun (distributed mode)\n    assert int(os.environ[\"WORLD_SIZE\"]) > 1, (\n        \"[ERROR]: This test is designed to run in distributed mode with torchrun. Please use torchrun to \"\n        \"execute this script.\"\n    )\n    torch.manual_seed(233376 + int(os.environ.get(\"RANK\", 0)))\n\n    # set_backward_method(BackwardEnum._Total_Fuse_MN)\n    # set_backward_method(BackwardEnum._Split_Dlogits_N)\n\n    test = TestLinearCrossEntropy_TensorParallel()\n    for test_case_idx in range(MAX_TEST_CASES):\n        print(f\"[INFO] Running test case {test_case_idx}\")\n        test.initialize(test_case_idx)\n        if VERIFY_TORCH_SELF:\n            test.verify_torch_itself()\n        test.check_torch_storage()\n        test.verify_kernel_correctness()\n        test.check_kernel_storage()\n\n    test.shutdown()\n"
  },
  {
    "path": "tests/utils/test_special_mstx_profile.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom unittest.mock import MagicMock, patch\n\nfrom verl.utils.profiler.config import NPUToolConfig, ProfilerConfig\nfrom verl.utils.profiler.mstx_profile import NPUProfiler\nfrom verl.utils.profiler.profile import DistProfiler\n\n\nclass TestNPUProfilerInitialization(unittest.TestCase):\n    def setUp(self):\n        NPUProfiler._define_count = 0\n\n    def test_init_with_default_config(self):\n        tool_config = NPUToolConfig()\n        config = ProfilerConfig(tool=\"npu\")\n        profiler = DistProfiler(rank=0, config=config, tool_config=tool_config)\n        self.assertFalse(profiler.check_enable())\n\n    def test_init_with_disabled_config(self):\n        config = ProfilerConfig(enable=False, tool=\"npu\")\n        tool_config = NPUToolConfig()\n        profiler = DistProfiler(rank=0, config=config, tool_config=tool_config)\n        self.assertFalse(profiler.check_enable())\n\n    def test_init_with_all_ranks_true(self):\n        config = ProfilerConfig(enable=True, all_ranks=True, tool=\"npu\")\n        tool_config = NPUToolConfig()\n        profiler = DistProfiler(rank=0, config=config, tool_config=tool_config)\n        self.assertTrue(profiler.check_this_rank())\n\n    def test_init_with_ranks_list(self):\n        config = ProfilerConfig(enable=True, ranks=[1, 2], tool=\"npu\")\n        tool_config = NPUToolConfig()\n        profiler = DistProfiler(rank=1, config=config, tool_config=tool_config)\n        self.assertTrue(profiler.check_this_rank())\n\n    def test_init_with_rank_not_in_ranks(self):\n        config = ProfilerConfig(enable=True, ranks=[1, 2], tool=\"npu\")\n        tool_config = NPUToolConfig()\n        profiler = DistProfiler(rank=3, config=config, tool_config=tool_config)\n        self.assertFalse(profiler.check_this_rank())\n\n\nclass TestNPUProfilerStart(unittest.TestCase):\n    def setUp(self):\n        NPUProfiler._define_count = 0\n        self.config = ProfilerConfig(enable=True, ranks=[0], tool=\"npu\")\n        self.tool_config = NPUToolConfig(discrete=False)\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_start_when_enabled_and_this_rank(self, mock_get_profiler):\n        profiler = DistProfiler(rank=0, config=self.config, tool_config=self.tool_config)\n        profiler.start(role=\"worker\", profile_step=\"1\")\n        self.assertTrue(profiler.check_this_step())\n        self.assertEqual(NPUProfiler._define_count, 1)\n        mock_get_profiler.assert_called_once()\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_start_when_not_this_rank(self, mock_get_profiler):\n        profiler = DistProfiler(rank=1, config=self.config, tool_config=self.tool_config)\n        profiler.start()\n        self.assertFalse(profiler.check_this_step())\n        self.assertEqual(NPUProfiler._define_count, 0)\n        mock_get_profiler.assert_not_called()\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_start_discrete_mode_does_not_increase_count(self, mock_get_profiler):\n        tool_config = NPUToolConfig(discrete=True)\n        profiler = DistProfiler(rank=0, config=self.config, tool_config=tool_config)\n        profiler.start()\n        self.assertEqual(NPUProfiler._define_count, 0)\n        mock_get_profiler.assert_not_called()\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_multiple_start_calls_do_not_increase_count(self, mock_get_profiler):\n        profiler = DistProfiler(rank=0, config=self.config, tool_config=self.tool_config)\n        profiler.start()\n        profiler.start()\n        self.assertEqual(NPUProfiler._define_count, 1)\n        mock_get_profiler.assert_called_once()\n\n\nclass TestNPUProfilerStartStopInteraction(unittest.TestCase):\n    def setUp(self):\n        NPUProfiler._define_count = 0\n        self.config = ProfilerConfig(enable=True, ranks=[0], tool=\"npu\")\n        self.tool_config = NPUToolConfig(discrete=False)\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_start_stop_cycle(self, mock_get_profiler):\n        mock_profile_npu = MagicMock()\n        mock_get_profiler.return_value = mock_profile_npu\n\n        profiler = DistProfiler(rank=0, config=self.config, tool_config=self.tool_config)\n        profiler.start()\n        self.assertEqual(NPUProfiler._define_count, 1)\n        self.assertEqual(mock_profile_npu.start.call_count, 1)\n        profiler.stop()\n        self.assertEqual(NPUProfiler._define_count, 0)\n        self.assertEqual(mock_profile_npu.step.call_count, 1)\n        self.assertEqual(mock_profile_npu.stop.call_count, 1)\n\n    @patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\")\n    def test_multiple_instances_share_define_count(self, mock_get_profiler):\n        mock_profile_npu = MagicMock()\n        mock_get_profiler.return_value = mock_profile_npu\n\n        profiler1 = DistProfiler(rank=0, config=self.config, tool_config=self.tool_config)\n        profiler2 = DistProfiler(rank=0, config=self.config, tool_config=self.tool_config)\n        profiler1.start()\n        profiler2.start()\n        self.assertEqual(NPUProfiler._define_count, 1)\n        self.assertEqual(mock_profile_npu.start.call_count, 1)\n        profiler1.stop()\n        self.assertEqual(NPUProfiler._define_count, 0)\n\n\nclass TestNPUProfilerAnnotate(unittest.TestCase):\n    def setUp(self):\n        self.config = ProfilerConfig(enable=True, all_ranks=True, tool=\"npu\")\n        self.tool_config = NPUToolConfig(discrete=False)\n        self.rank = 0\n\n    def test_annotate_decorator_applied_correctly(self):\n        mock_worker = MagicMock()\n        mock_worker.profiler = DistProfiler(rank=self.rank, config=self.config, tool_config=self.tool_config)\n        # Manually set private attribute for testing annotation in active step\n        mock_worker.profiler._this_step = True\n\n        mock_mark_range = \"mocked_range_handle\"\n\n        with (\n            patch(\"verl.utils.profiler.mstx_profile.mark_start_range\") as mock_start_patch,\n            patch(\"verl.utils.profiler.mstx_profile.mark_end_range\") as mock_end_patch,\n        ):\n            mock_start_patch.return_value = mock_mark_range\n\n            with patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\") as mock_get_profiler:\n                decorator = mock_worker.profiler.annotate(message=\"test\")\n\n                @decorator\n                def test_func(self, *args, **kwargs):\n                    return \"result\"\n\n                result = test_func(mock_worker)\n\n                self.assertEqual(result, \"result\")\n                mock_start_patch.assert_called_once_with(message=\"test\")\n                mock_end_patch.assert_called_once_with(mock_mark_range)\n                mock_get_profiler.assert_not_called()\n\n    def test_annotate_when_profiler_disabled(self):\n        disabled_config = ProfilerConfig(enable=False, tool=\"npu\")\n        mock_worker = MagicMock()\n        mock_worker.profiler = DistProfiler(rank=self.rank, config=disabled_config, tool_config=self.tool_config)\n\n        with (\n            patch(\"verl.utils.profiler.mstx_profile.mark_start_range\") as mock_start_patch,\n            patch(\"verl.utils.profiler.mstx_profile.mark_end_range\") as mock_end_patch,\n            patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\") as mock_get_profiler,\n        ):\n            decorator = mock_worker.profiler.annotate(message=\"test\")\n\n            @decorator\n            def test_func(self, *args, **kwargs):\n                return \"result\"\n\n            result = test_func(mock_worker)\n\n            self.assertEqual(result, \"result\")\n            mock_start_patch.assert_not_called()\n            mock_end_patch.assert_not_called()\n            mock_get_profiler.assert_not_called()\n\n    def test_annotate_when_this_step_disabled(self):\n        mock_worker = MagicMock()\n        mock_worker.profiler = DistProfiler(rank=self.rank, config=self.config, tool_config=self.tool_config)\n        mock_worker.profiler._this_step = False\n\n        with (\n            patch(\"verl.utils.profiler.mstx_profile.mark_start_range\") as mock_start_patch,\n            patch(\"verl.utils.profiler.mstx_profile.mark_end_range\") as mock_end_patch,\n            patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\") as mock_get_profiler,\n        ):\n            decorator = mock_worker.profiler.annotate(message=\"test\")\n\n            @decorator\n            def test_func(self, *args, **kwargs):\n                return \"result\"\n\n            result = test_func(mock_worker)\n\n            self.assertEqual(result, \"result\")\n            mock_start_patch.assert_not_called()\n            mock_end_patch.assert_not_called()\n            mock_get_profiler.assert_not_called()\n\n    def test_annotate_discrete_mode_enabled(self):\n        discrete_tool_config = NPUToolConfig(discrete=True)\n        mock_worker = MagicMock()\n        mock_worker.profiler = DistProfiler(rank=self.rank, config=self.config, tool_config=discrete_tool_config)\n        mock_worker.profiler._this_step = True\n\n        mock_mark_range = \"mocked_range_handle\"\n        mock_profile_npu = MagicMock()\n\n        with (\n            patch(\"verl.utils.profiler.mstx_profile.mark_start_range\") as mock_start_patch,\n            patch(\"verl.utils.profiler.mstx_profile.mark_end_range\") as mock_end_patch,\n            patch(\"verl.utils.profiler.mstx_profile.get_npu_profiler\") as mock_get_profiler,\n        ):\n            mock_start_patch.return_value = mock_mark_range\n            mock_get_profiler.return_value = mock_profile_npu\n            decorator = mock_worker.profiler.annotate(message=\"test\", role=\"test_role\")\n\n            @decorator\n            def test_func(self, *args, **kwargs):\n                return \"result\"\n\n            result = test_func(mock_worker)\n\n            self.assertEqual(result, \"result\")\n            mock_start_patch.assert_called_once_with(message=\"test\")\n            mock_end_patch.assert_called_once_with(mock_mark_range)\n            mock_get_profiler.assert_called_once_with(\n                contents=mock_worker.profiler._impl.profile_contents,\n                profile_level=mock_worker.profiler._impl.profile_level,\n                profile_save_path=mock_worker.profiler._impl.profile_save_path,\n                analysis=mock_worker.profiler._impl.analysis,\n                role=\"test_role\",\n            )\n            mock_profile_npu.start.assert_called_once()\n            mock_profile_npu.step.assert_called_once()\n            mock_profile_npu.stop.assert_called_once()\n\n    def test_annotate_with_default_message(self):\n        mock_worker = MagicMock()\n        mock_worker.profiler = DistProfiler(rank=self.rank, config=self.config, tool_config=self.tool_config)\n        mock_worker.profiler._this_step = True\n\n        mock_mark_range = \"mocked_range_handle\"\n        with (\n            patch(\"verl.utils.profiler.mstx_profile.mark_start_range\") as mock_start_patch,\n            patch(\"verl.utils.profiler.mstx_profile.mark_end_range\") as mock_end_patch,\n        ):\n            mock_start_patch.return_value = mock_mark_range\n            decorator = mock_worker.profiler.annotate()\n\n            @decorator\n            def test_func(self, *args, **kwargs):\n                return \"result\"\n\n            test_func(mock_worker)\n\n            mock_start_patch.assert_called_once_with(message=\"test_func\")\n            mock_end_patch.assert_called_once_with(mock_mark_range)\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/utils/test_temp_env_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport pytest\n\nfrom verl.utils.py_functional import temp_env_var\n\n\n@pytest.fixture(autouse=True)\ndef clean_env():\n    \"\"\"Fixture to clean up environment variables before and after each test.\"\"\"\n    # Store original environment state\n    original_env = dict(os.environ)\n\n    # Clean up any test variables that might exist\n    test_vars = [\"TEST_VAR\", \"TEST_VAR_2\", \"EXISTING_VAR\"]\n    for var in test_vars:\n        if var in os.environ:\n            del os.environ[var]\n\n    # Yield control to the test function\n    yield\n\n    # Restore original environment state after test\n    os.environ.clear()\n    os.environ.update(original_env)\n\n\ndef test_set_new_env_var():\n    \"\"\"Test setting a new environment variable that didn't exist before.\"\"\"\n    # Ensure variable doesn't exist\n    assert \"TEST_VAR\" not in os.environ\n\n    with temp_env_var(\"TEST_VAR\", \"test_value\"):\n        # Variable should be set inside context\n        assert os.environ[\"TEST_VAR\"] == \"test_value\"\n        assert \"TEST_VAR\" in os.environ\n\n    # Variable should be removed after context\n    assert \"TEST_VAR\" not in os.environ\n\n\ndef test_restore_existing_env_var():\n    \"\"\"Test restoring an environment variable that already existed.\"\"\"\n    # Set up existing variable\n    os.environ[\"EXISTING_VAR\"] = \"original_value\"\n\n    with temp_env_var(\"EXISTING_VAR\", \"temporary_value\"):\n        # Variable should be temporarily changed\n        assert os.environ[\"EXISTING_VAR\"] == \"temporary_value\"\n\n    # Variable should be restored to original value\n    assert os.environ[\"EXISTING_VAR\"] == \"original_value\"\n\n\ndef test_env_var_restored_on_exception():\n    \"\"\"Test that environment variables are restored even when exceptions occur.\"\"\"\n    # Set up existing variable\n    os.environ[\"EXISTING_VAR\"] = \"original_value\"\n\n    with pytest.raises(ValueError):\n        with temp_env_var(\"EXISTING_VAR\", \"temporary_value\"):\n            # Verify variable is set\n            assert os.environ[\"EXISTING_VAR\"] == \"temporary_value\"\n            # Raise exception\n            raise ValueError(\"Test exception\")\n\n    # Variable should still be restored despite exception\n    assert os.environ[\"EXISTING_VAR\"] == \"original_value\"\n\n\ndef test_nested_context_managers():\n    \"\"\"Test nested temp_env_var context managers.\"\"\"\n    # Set up original variable\n    os.environ[\"TEST_VAR\"] = \"original\"\n\n    with temp_env_var(\"TEST_VAR\", \"level1\"):\n        assert os.environ[\"TEST_VAR\"] == \"level1\"\n\n        with temp_env_var(\"TEST_VAR\", \"level2\"):\n            assert os.environ[\"TEST_VAR\"] == \"level2\"\n\n        # Should restore to level1\n        assert os.environ[\"TEST_VAR\"] == \"level1\"\n\n    # Should restore to original\n    assert os.environ[\"TEST_VAR\"] == \"original\"\n\n\ndef test_multiple_different_vars():\n    \"\"\"Test setting multiple different environment variables.\"\"\"\n    # Set up one existing variable\n    os.environ[\"EXISTING_VAR\"] = \"existing_value\"\n\n    with temp_env_var(\"EXISTING_VAR\", \"modified\"):\n        with temp_env_var(\"TEST_VAR\", \"new_value\"):\n            assert os.environ[\"EXISTING_VAR\"] == \"modified\"\n            assert os.environ[\"TEST_VAR\"] == \"new_value\"\n\n    # Check restoration\n    assert os.environ[\"EXISTING_VAR\"] == \"existing_value\"\n    assert \"TEST_VAR\" not in os.environ\n\n\ndef test_empty_string_value():\n    \"\"\"Test setting environment variable to empty string.\"\"\"\n    with temp_env_var(\"TEST_VAR\", \"\"):\n        assert os.environ[\"TEST_VAR\"] == \"\"\n        assert \"TEST_VAR\" in os.environ\n\n    # Should be removed after context\n    assert \"TEST_VAR\" not in os.environ\n\n\ndef test_overwrite_with_empty_string():\n    \"\"\"Test overwriting existing variable with empty string.\"\"\"\n    os.environ[\"EXISTING_VAR\"] = \"original\"\n\n    with temp_env_var(\"EXISTING_VAR\", \"\"):\n        assert os.environ[\"EXISTING_VAR\"] == \"\"\n\n    # Should restore original value\n    assert os.environ[\"EXISTING_VAR\"] == \"original\"\n\n\ndef test_context_manager_returns_none():\n    \"\"\"Test that context manager yields None.\"\"\"\n    with temp_env_var(\"TEST_VAR\", \"value\") as result:\n        assert result is None\n        assert os.environ[\"TEST_VAR\"] == \"value\"\n"
  },
  {
    "path": "tests/utils/test_timeout_decorator_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport multiprocessing\nimport sys\nimport threading\nimport time\n\nimport pytest  # Import pytest\n\nfrom verl.utils.py_functional import timeout_limit as timeout\n\n# --- Test Task Functions ---\nTEST_TIMEOUT_SECONDS = 1.5  # Timeout duration for tests\nLONG_TASK_DURATION = TEST_TIMEOUT_SECONDS + 0.5  # Duration slightly longer than timeout\n\n\n@timeout(seconds=TEST_TIMEOUT_SECONDS)  # Keep global decorator for mp tests\ndef quick_task(x):\n    \"\"\"A task that completes quickly.\"\"\"\n    time.sleep(0.1)\n    return \"quick_ok\"\n\n\n@timeout(seconds=TEST_TIMEOUT_SECONDS)  # Keep global decorator for mp tests\ndef slow_task(x):\n    \"\"\"A task that takes longer than the timeout.\"\"\"\n    time.sleep(LONG_TASK_DURATION)\n    return \"slow_finished\"  # This return value indicates it didn't time out\n\n\n# REMOVE global decorator here\ndef task_raises_value_error():  # Now truly not globally decorated\n    \"\"\"A task that intentionally raises a ValueError.\"\"\"\n    raise ValueError(\"Specific value error from task\")\n\n\n# --- Top-level function for signal test in subprocess ---\n# Keep this decorated globally for the specific subprocess test case\n@timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)\ndef top_level_decorated_quick_task_signal():\n    \"\"\"A pickleable top-level function decorated with signal timeout.\"\"\"\n    # Assuming this calls the logic of quick_task directly for the test purpose\n    time.sleep(0.1)\n    return \"quick_ok_signal_subprocess\"  # Different return for clarity if needed\n\n\n# --- Top-level function for signal test in subprocess ---\n# Keep this decorated globally for the specific subprocess test case\n@timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)\ndef top_level_decorated_slow_task_signal():\n    \"\"\"A pickleable top-level function decorated with signal timeout.\"\"\"\n    time.sleep(LONG_TASK_DURATION)\n    return \"slow_finished\"\n\n\n# --- NEW: Top-level helper function to run target in process ---\ndef run_target_and_put_in_queue(target_func, q):\n    \"\"\"\n    Top-level helper function to run a target function and put its result or exception into a queue.\n    This function is pickleable and can be used as the target for multiprocessing.Process.\n    \"\"\"\n    try:\n        result = target_func()\n        q.put((\"success\", result))\n    except Exception as e:\n        q.put((\"error\", e))\n\n\n# Use a module-level fixture to set the start method on macOS\n@pytest.fixture(scope=\"module\", autouse=True)  # Changed scope to module\ndef set_macos_start_method():\n    if sys.platform == \"darwin\":\n        # Force fork method on macOS to avoid pickling issues with globally decorated functions\n        # when running tests via pytest discovery.\n        current_method = multiprocessing.get_start_method(allow_none=True)\n        # Only set if not already set or if set to something else (less likely in test run)\n        if current_method is None or current_method != \"fork\":\n            try:\n                multiprocessing.set_start_method(\"fork\", force=True)\n            except RuntimeError:\n                # Might fail if context is already started, ignore in that case.\n                pass\n\n\ndef test_quick_task():  # Renamed from test_multiprocessing_quick_task\n    \"\"\"Tests timeout handles a quick task correctly.\"\"\"\n    # Call the globally decorated function directly\n    result = quick_task(1)\n    assert result == \"quick_ok\"  # Use pytest assert\n\n\ndef test_slow_task_timeout():  # Renamed from test_multiprocessing_slow_task_timeout\n    \"\"\"Tests timeout correctly raises TimeoutError for a slow task.\"\"\"\n    # Call the globally decorated function directly within pytest.raises\n    with pytest.raises(TimeoutError) as excinfo:  # Use pytest.raises\n        slow_task(1)\n    # Check the error message from the multiprocessing implementation\n    assert f\"timed out after {TEST_TIMEOUT_SECONDS} seconds\" in str(excinfo.value)  # Use pytest assert\n\n\ndef test_internal_exception():  # Renamed from test_multiprocessing_internal_exception\n    \"\"\"Tests timeout correctly propagates internal exceptions.\"\"\"\n    # Apply the default timeout decorator dynamically to the undecorated function\n    decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS)(task_raises_value_error)  # Apply decorator dynamically\n    with pytest.raises(ValueError) as excinfo:  # Use pytest.raises\n        decorated_task()  # Call the dynamically decorated function\n    assert str(excinfo.value) == \"Specific value error from task\"  # Use pytest assert\n\n\n# --- Test the signal implementation (use_signals=True) ---\n# Note: As per py_functional.py, use_signals=True currently falls back to\n# multiprocessing on POSIX. These tests verify that behavior.\n\n\ndef test_signal_quick_task_main_process():  # Removed self\n    \"\"\"Tests signal timeout handles a quick task correctly in the main process.\"\"\"\n\n    # Apply the signal decorator dynamically\n    def plain_quick_task_logic():\n        time.sleep(0.1)\n        return \"quick_ok_signal\"\n\n    decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)(plain_quick_task_logic)\n    assert decorated_task() == \"quick_ok_signal\"  # Use pytest assert\n\n\ndef test_signal_slow_task_main_process_timeout():  # Removed self\n    \"\"\"Tests signal timeout correctly raises TimeoutError for a slow task in the main process.\"\"\"\n\n    # Apply the signal decorator dynamically\n    def plain_slow_task_logic():\n        time.sleep(LONG_TASK_DURATION)\n        return \"slow_finished_signal\"\n\n    decorated_task = timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)(plain_slow_task_logic)\n    with pytest.raises(TimeoutError) as excinfo:  # Use pytest.raises\n        decorated_task()\n    # Check the error message (falls back to multiprocessing message on POSIX)\n    assert f\"timed out after {TEST_TIMEOUT_SECONDS} seconds\" in str(excinfo.value)  # Use pytest assert\n\n\n@pytest.mark.skip(reason=\"this test won't pass. Just to show why use_signals should not be used\")\ndef test_signal_in_thread_does_not_timeout():\n    \"\"\"\n    Tests that signal-based timeout does NOT work reliably in a child thread.\n    The TimeoutError from the signal handler is not expected to be raised.\n    \"\"\"\n    result_container = []  # Use a list to store result from thread\n    exception_container = []  # Use a list to store exception from thread\n\n    @timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=True)\n    def slow_task_in_thread():\n        try:\n            print(\"Thread: Starting slow task...\")\n            time.sleep(LONG_TASK_DURATION)\n            print(\"Thread: Slow task finished.\")\n            return \"slow_finished_in_thread\"\n        except Exception as e:\n            # Catch any exception within the thread's target function\n            print(f\"Thread: Caught exception: {e}\")\n            exception_container.append(e)\n            return None  # Indicate failure\n\n    def thread_target():\n        try:\n            # Run the decorated function inside the thread\n            res = slow_task_in_thread()\n            if res is not None:\n                result_container.append(res)\n        except Exception as e:\n            # This might catch exceptions happening *outside* the decorated function\n            # but still within the thread target, though less likely here.\n            print(f\"Thread Target: Caught exception: {e}\")\n            exception_container.append(e)\n\n    thread = threading.Thread(target=thread_target)\n    print(\"Main: Starting thread...\")\n    thread.start()\n    # Wait longer than the timeout + task duration to ensure the thread finishes\n    # regardless of whether timeout worked or not.\n    thread.join(timeout=LONG_TASK_DURATION + 1)\n\n    assert len(exception_container) == 1\n    assert isinstance(exception_container[0], TimeoutError)\n    assert not result_container\n\n\ndef test_in_thread_timeout():\n    result_container = []  # Use a list to store result from thread\n    exception_container = []  # Use a list to store exception from thread\n\n    @timeout(seconds=TEST_TIMEOUT_SECONDS, use_signals=False)\n    def slow_task_in_thread():\n        try:\n            print(\"Thread: Starting slow task...\")\n            time.sleep(LONG_TASK_DURATION)\n            print(\"Thread: Slow task finished.\")\n            return \"slow_finished_in_thread\"\n        except Exception as e:\n            # Catch any exception within the thread's target function\n            print(f\"Thread: Caught exception: {e}\")\n            exception_container.append(e)\n            return None  # Indicate failure\n\n    def thread_target():\n        try:\n            # Run the decorated function inside the thread\n            res = slow_task_in_thread()\n            if res is not None:\n                result_container.append(res)\n        except Exception as e:\n            # This might catch exceptions happening *outside* the decorated function\n            # but still within the thread target, though less likely here.\n            print(f\"Thread Target: Caught exception: {e}\")\n            exception_container.append(e)\n\n    thread = threading.Thread(target=thread_target)\n    print(\"Main: Starting thread...\")\n    thread.start()\n    # Wait longer than the timeout + task duration to ensure the thread finishes\n    # regardless of whether timeout worked or not.\n    thread.join(timeout=LONG_TASK_DURATION + 1)\n\n    assert len(exception_container) == 1\n    assert isinstance(exception_container[0], TimeoutError)\n    assert not result_container\n"
  },
  {
    "path": "tests/utils/test_tokenizer_normalize_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport numpy as np\nimport pytest\n\nfrom verl.utils.tokenizer import normalize_token_ids\n\n\nclass DummyBatchEncoding:\n    def __init__(self, input_ids):\n        self.input_ids = input_ids\n\n\nclass DummyToList:\n    def __init__(self, data):\n        self._data = data\n\n    def tolist(self):\n        return self._data\n\n\n@pytest.mark.parametrize(\n    (\"tokenized_output\", \"expected\"),\n    [\n        # transformers v4-style direct token ids\n        ([1, 2, 3], [1, 2, 3]),\n        ((1, 2, 3), [1, 2, 3]),\n        # common list-like outputs with tolist()/ndarray paths\n        (DummyToList([1, 2, 3]), [1, 2, 3]),\n        (np.array([1, 2, 3], dtype=np.int64), [1, 2, 3]),\n        # transformers v5-like mapping / BatchEncoding-style outputs\n        ({\"input_ids\": [1, 2, 3]}, [1, 2, 3]),\n        ({\"input_ids\": DummyToList([1, 2, 3])}, [1, 2, 3]),\n        ({\"input_ids\": [[1, 2, 3]]}, [1, 2, 3]),\n        (DummyBatchEncoding([1, 2, 3]), [1, 2, 3]),\n        (DummyBatchEncoding(DummyToList([[1, 2, 3]])), [1, 2, 3]),\n        # scalar item() support\n        ([np.int64(1), np.int32(2), np.int16(3)], [1, 2, 3]),\n    ],\n)\ndef test_normalize_token_ids_valid_outputs(tokenized_output, expected):\n    assert normalize_token_ids(tokenized_output) == expected\n\n\n@pytest.mark.parametrize(\n    \"tokenized_output\",\n    [\n        \"not-token-ids\",\n        {\"attention_mask\": [1, 1, 1]},\n        [[1, 2], [3, 4]],  # ambiguous batched ids should fail fast\n        [1, object(), 3],\n    ],\n)\ndef test_normalize_token_ids_invalid_outputs(tokenized_output):\n    with pytest.raises(TypeError):\n        normalize_token_ids(tokenized_output)\n"
  },
  {
    "path": "tests/utils/test_torch_functional.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport pytest\nimport torch\nimport torch.distributed as dist\nimport torch.multiprocessing as mp\n\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.torch_functional import (\n    distributed_masked_mean,\n    distributed_mean_max_min_std,\n    expand_as_nested,\n    masked_mean,\n)\n\n\ndef _worker_mean(rank: int, world_size: int, rendezvous_file: str):\n    # 1) set GPU and init NCCL\n    get_torch_device().set_device(rank)\n    dist.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n    # each rank holds tensor [rank+1]\n    local = torch.tensor([float(rank + 1)], device=f\"{get_device_name()}:{rank}\")\n    mean, gmax, gmin, gstd = distributed_mean_max_min_std(local, True, True, True)\n\n    values = [float(i + 1) for i in range(world_size)]\n    exp_mean = sum(values) / len(values)\n    exp_max = max(values)\n    exp_min = min(values)\n    var = sum((x - exp_mean) ** 2 for x in values) / (len(values) - 1)\n    exp_std = var**0.5\n\n    # all ranks should see the same result\n    assert torch.allclose(mean.cpu(), torch.tensor(exp_mean)), f\"mean@{rank}\"\n    assert torch.allclose(gmax.cpu(), torch.tensor(exp_max)), f\"max@{rank}\"\n    assert torch.allclose(gmin.cpu(), torch.tensor(exp_min)), f\"min@{rank}\"\n    assert torch.allclose(gstd.cpu(), torch.tensor(exp_std)), f\"std@{rank}\"\n\n    dist.destroy_process_group()\n\n\n@pytest.mark.parametrize(\n    \"value,mask,gt\",\n    [\n        ([1.0, 2.0, 3.0, 4.0], [1, 0, 0, 1], 2.5),\n        ([1.0, 2.0, float(\"nan\"), 4.0], [1, 0, 0, 1], 2.5),\n        ([1.0, 2.0, float(\"nan\"), 4.0], [1, 0, 1, 0], float(\"nan\")),\n    ],\n)\ndef test_masked_mean(value, mask, gt):\n    res = masked_mean(torch.tensor(value), torch.tensor(mask))\n    gt = torch.tensor(gt)\n    assert torch.allclose(res, gt) or (torch.isnan(res) and torch.isnan(gt))\n\n\n@pytest.mark.parametrize(\"world_size\", [2, 4])\ndef test_distributed_mean_max_min_std(world_size, tmp_path):\n    rendezvous_file = str(tmp_path / \"rdzv_mean\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    mp.spawn(\n        fn=_worker_mean,\n        args=(world_size, rendezvous_file),\n        nprocs=world_size,\n        join=True,\n    )\n\n\ndef _worker_mask(rank: int, world_size: int, rendezvous_file: str):\n    get_torch_device().set_device(rank)\n    dist.init_process_group(\n        backend=get_nccl_backend(),\n        init_method=f\"file://{rendezvous_file}\",\n        rank=rank,\n        world_size=world_size,\n    )\n\n    # build per‐rank tensor and mask\n    local_tensor = torch.tensor([rank * 2 + 1.0, rank * 2 + 2.0], device=f\"{get_device_name()}:{rank}\")\n    if rank == 0:\n        mask = torch.tensor([1, 0], device=f\"{get_device_name()}:{rank}\", dtype=torch.float32)\n    else:\n        mask = torch.tensor([0, 1], device=f\"{get_device_name()}:{rank}\", dtype=torch.float32)\n\n    gmean = distributed_masked_mean(local_tensor, mask)\n\n    valid_values = [1.0] + [2 * i + 2.0 for i in range(1, world_size)]\n    expected_mean = sum(valid_values) / len(valid_values)\n    assert torch.allclose(gmean.cpu(), torch.tensor(expected_mean)), f\"masked_mean@{rank}\"\n\n    dist.destroy_process_group()\n\n\n@pytest.mark.parametrize(\"world_size\", [2, 4])\ndef test_distributed_masked_mean(world_size, tmp_path):\n    rendezvous_file = str(tmp_path / \"rdzv_mask\")\n    os.makedirs(os.path.dirname(rendezvous_file), exist_ok=True)\n\n    mp.spawn(\n        fn=_worker_mask,\n        args=(world_size, rendezvous_file),\n        nprocs=world_size,\n        join=True,\n    )\n\n\ndef test_expand_as_nested():\n    a = torch.randn(2)\n    b = torch.randn(3)\n    c = torch.randn(4)\n    nested_tensor = torch.nested.as_nested_tensor([a, b, c], layout=torch.jagged)\n    tensor = torch.tensor([1, 2, 3])\n\n    output = expand_as_nested(tensor, nested_tensor)\n\n    assert output.values().tolist() == [1, 1, 2, 2, 2, 3, 3, 3, 3]\n    assert torch.all(output.offsets() == nested_tensor.offsets()).item()\n\n    # test exceptions\n    with pytest.raises(AssertionError):\n        expand_as_nested(tensor, tensor)\n\n    other_tensor = torch.tensor([1, 2, 3, 4])\n\n    with pytest.raises(AssertionError):\n        expand_as_nested(other_tensor, nested_tensor)\n\n    other_tensor = torch.tensor([[1, 2, 3]])\n\n    with pytest.raises(AssertionError):\n        expand_as_nested(other_tensor, nested_tensor)\n\n    with pytest.raises(AssertionError):\n        expand_as_nested(tensor, nested_tensor.unsqueeze(-1))\n"
  },
  {
    "path": "tests/utils/test_torch_profile.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\nfrom unittest.mock import MagicMock, patch\n\nimport torch\n\nfrom verl.utils.profiler.config import ProfilerConfig, TorchProfilerToolConfig\nfrom verl.utils.profiler.torch_profile import Profiler, get_torch_profiler\n\n\nclass TestTorchProfile(unittest.TestCase):\n    def setUp(self):\n        # Reset Profiler class state\n        Profiler._define_count = 0\n\n    @patch(\"torch.profiler.profile\")\n    def test_get_torch_profiler(self, mock_profile):\n        # Test wrapper function\n        get_torch_profiler(contents=[\"cpu\", \"cuda\", \"stack\"], save_path=\"/tmp/test\", rank=0)\n        mock_profile.assert_called_once()\n        _, kwargs = mock_profile.call_args\n\n        # Verify activities\n        activities = kwargs[\"activities\"]\n        self.assertIn(torch.profiler.ProfilerActivity.CPU, activities)\n        self.assertIn(torch.profiler.ProfilerActivity.CUDA, activities)\n\n        # Verify options\n        self.assertTrue(kwargs[\"with_stack\"])\n        self.assertFalse(kwargs[\"record_shapes\"])\n        self.assertFalse(kwargs[\"profile_memory\"])\n\n    @patch(\"verl.utils.profiler.torch_profile.get_torch_profiler\")\n    def test_profiler_lifecycle(self, mock_get_profiler):\n        # Mock the underlying torch profiler object\n        mock_prof_instance = MagicMock()\n        mock_get_profiler.return_value = mock_prof_instance\n\n        # Initialize\n        tool_config = TorchProfilerToolConfig(contents=[\"cpu\"], discrete=False)\n        config = ProfilerConfig(save_path=\"/tmp/test\", enable=True, tool_config=tool_config)\n        profiler = Profiler(rank=0, config=config, tool_config=tool_config)\n\n        # Test Start\n        profiler.start()\n        mock_get_profiler.assert_called_once()\n        mock_prof_instance.start.assert_called_once()\n\n        # Test Step\n        profiler.step()\n        mock_prof_instance.step.assert_called_once()\n\n        # Test Stop\n        profiler.stop()\n        mock_prof_instance.stop.assert_called_once()\n\n    @patch(\"verl.utils.profiler.torch_profile.get_torch_profiler\")\n    def test_discrete_mode(self, mock_get_profiler):\n        # Mock for discrete mode\n        mock_prof_instance = MagicMock()\n        mock_get_profiler.return_value = mock_prof_instance\n\n        tool_config = TorchProfilerToolConfig(contents=[\"cpu\"], discrete=True)\n        config = ProfilerConfig(save_path=\"/tmp/test\", enable=True, tool_config=tool_config)\n        profiler = Profiler(rank=0, config=config, tool_config=tool_config)\n\n        # In discrete mode, start/stop shouldn't trigger global profiler immediately\n        profiler.start()\n        mock_get_profiler.assert_not_called()\n\n        profiler.stop()\n        mock_prof_instance.stop.assert_not_called()\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/workers/actor/test_special_dp_actor.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport unittest\n\nimport torch\nimport torch.nn as nn\nfrom tensordict import TensorDict\nfrom transformers import AutoModelForCausalLM, Qwen3Config\n\nfrom verl import DataProto\nfrom verl.utils.device import get_device_name\nfrom verl.workers.actor.dp_actor import DataParallelPPOActor\nfrom verl.workers.config import FSDPActorConfig, OptimizerConfig\n\n\nclass MockTransformerModel(nn.Module):\n    \"\"\"Mock transformer model for testing DataParallelPPOActor\"\"\"\n\n    def __init__(self, vocab_size=1000, hidden_size=64):\n        super().__init__()\n        self.vocab_size = vocab_size\n        self.hidden_size = hidden_size\n        self.embedding = nn.Embedding(vocab_size, hidden_size)\n        self.transformer = nn.TransformerEncoder(\n            nn.TransformerEncoderLayer(d_model=hidden_size, nhead=4, batch_first=True), num_layers=2\n        )\n        self.lm_head = nn.Linear(hidden_size, vocab_size)\n\n    def forward(self, input_ids, attention_mask=None, position_ids=None, use_cache=False, **kwargs):\n        batch_size, seq_len = input_ids.shape\n\n        embeddings = self.embedding(input_ids)\n        hidden_states = self.transformer(embeddings)\n        logits = self.lm_head(hidden_states)\n\n        class MockOutput:\n            def __init__(self, logits):\n                self.logits = logits\n\n        return MockOutput(logits)\n\n\nclass TestDataParallelPPOActor(unittest.TestCase):\n    \"\"\"Test DataParallelPPOActor compute_log_prob and update_policy methods\"\"\"\n\n    @classmethod\n    def setUpClass(cls):\n        \"\"\"Set up distributed environment\"\"\"\n        if get_device_name() == \"cuda\":\n            backend_name = \"nccl\"\n        elif get_device_name() == \"npu\":\n            backend_name = \"hccl\"\n        else:\n            backend_name = \"gloo\"\n\n        if not torch.distributed.is_initialized():\n            torch.distributed.init_process_group(backend=backend_name, init_method=\"env://\")\n\n        cls.rank = torch.distributed.get_rank()\n        cls.world_size = torch.distributed.get_world_size()\n\n        if get_device_name() == \"cuda\":\n            torch.cuda.set_device(cls.rank)\n            cls.device = torch.device(f\"cuda:{cls.rank}\")\n        elif get_device_name() == \"npu\":\n            torch.npu.set_device(cls.rank)\n            cls.device = torch.device(f\"npu:{cls.rank}\")\n        else:\n            cls.device = torch.device(\"cpu\")\n\n    def setUp(self):\n        \"\"\"Set up test fixtures\"\"\"\n        self.config = FSDPActorConfig(\n            strategy=\"fsdp2\",\n            ppo_mini_batch_size=4,\n            ppo_micro_batch_size_per_gpu=2,\n            ppo_epochs=1,\n            clip_ratio=0.2,\n            entropy_coeff=0.01,\n            grad_clip=1.0,\n            use_dynamic_bsz=False,\n            use_torch_compile=False,  # Disable torch.compile for testing\n            ulysses_sequence_parallel_size=1,\n            optim=OptimizerConfig(lr=1e-6),\n            rollout_n=1,\n        )\n\n        self.mock_model = MockTransformerModel(vocab_size=1000, hidden_size=64).to(self.device)\n        self.mock_optimizer = torch.optim.Adam(self.mock_model.parameters(), lr=1e-4)\n\n        self.actor = DataParallelPPOActor(\n            config=self.config, actor_module=self.mock_model, actor_optimizer=self.mock_optimizer\n        )\n\n    @classmethod\n    def tearDownClass(cls):\n        \"\"\"Clean up distributed environment\"\"\"\n        if torch.distributed.is_initialized():\n            torch.distributed.destroy_process_group()\n\n    def _create_test_data_for_compute_log_prob(self):\n        \"\"\"Create test DataProto for compute_log_prob method\"\"\"\n        batch_size = 2\n        prompt_length = 8\n        response_length = 4\n        total_length = prompt_length + response_length\n        vocab_size = 1000\n\n        input_ids = torch.randint(0, vocab_size, (batch_size, total_length)).to(self.device)\n        attention_mask = torch.ones(batch_size, total_length).to(self.device)\n        position_ids = torch.arange(total_length).unsqueeze(0).expand(batch_size, -1).to(self.device)\n        responses = input_ids[:, -response_length:]  # Last part is the response\n\n        tensor_dict = TensorDict(\n            {\n                \"input_ids\": input_ids,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n                \"responses\": responses,\n            },\n            batch_size=[batch_size],\n        )\n\n        meta_info = {\"micro_batch_size\": batch_size, \"temperature\": 1.0, \"use_dynamic_bsz\": False}\n\n        return DataProto(batch=tensor_dict, meta_info=meta_info)\n\n    def _create_test_data_for_update_policy(self):\n        \"\"\"Create test DataProto for update_policy method\"\"\"\n        batch_size = 4  # Must match ppo_mini_batch_size\n        prompt_length = 8\n        response_length = 4\n        total_length = prompt_length + response_length\n        vocab_size = 1000\n\n        input_ids = torch.randint(0, vocab_size, (batch_size, total_length)).to(self.device)\n        attention_mask = torch.ones(batch_size, total_length).to(self.device)\n        position_ids = torch.arange(total_length).unsqueeze(0).expand(batch_size, -1).to(self.device)\n        responses = input_ids[:, -response_length:]\n        response_mask = torch.ones(batch_size, response_length).to(self.device)\n        old_log_probs = torch.randn(batch_size, response_length).to(self.device) * 0.1  # Small values\n        advantages = torch.randn(batch_size, response_length).to(self.device) * 0.5\n\n        tensor_dict = TensorDict(\n            {\n                \"input_ids\": input_ids,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n                \"responses\": responses,\n                \"response_mask\": response_mask,\n                \"old_log_probs\": old_log_probs,\n                \"advantages\": advantages,\n            },\n            batch_size=[batch_size],\n        )\n\n        meta_info = {\"temperature\": 1.0}\n\n        return DataProto(batch=tensor_dict, meta_info=meta_info)\n\n    def test_compute_log_prob(self):\n        \"\"\"Test compute_log_prob method\"\"\"\n        data = self._create_test_data_for_compute_log_prob()\n\n        outputs = self.actor.compute_log_prob(data, calculate_entropy=True)\n        log_probs = outputs[\"log_probs\"]\n        entropys = outputs[\"entropys\"]\n\n        batch_size = data.batch[\"responses\"].shape[0]\n        response_length = data.batch[\"responses\"].shape[1]\n\n        self.assertIsInstance(log_probs, torch.Tensor)\n        self.assertEqual(log_probs.shape, (batch_size, response_length))\n        self.assertTrue(torch.all(torch.isfinite(log_probs)))\n\n        self.assertIsInstance(entropys, torch.Tensor)\n        self.assertEqual(entropys.shape, (batch_size, response_length))\n        self.assertTrue(torch.all(torch.isfinite(entropys)))\n        self.assertTrue(torch.all(entropys >= 0))  # Entropy should be non-negative\n\n    def test_compute_log_prob_without_entropy(self):\n        \"\"\"Test compute_log_prob method without entropy calculation\"\"\"\n        data = self._create_test_data_for_compute_log_prob()\n\n        outputs = self.actor.compute_log_prob(data, calculate_entropy=False)\n        log_probs = outputs[\"log_probs\"]\n        entropys = outputs.get(\"entropys\", None)\n\n        batch_size = data.batch[\"responses\"].shape[0]\n        response_length = data.batch[\"responses\"].shape[1]\n\n        self.assertIsInstance(log_probs, torch.Tensor)\n        self.assertEqual(log_probs.shape, (batch_size, response_length))\n        self.assertTrue(torch.all(torch.isfinite(log_probs)))\n        self.assertIsNone(entropys)\n\n    def test_update_policy(self):\n        \"\"\"Test update_policy method\"\"\"\n        data = self._create_test_data_for_update_policy()\n\n        metrics = self.actor.update_policy(data)\n\n        self.assertIsInstance(metrics, dict)\n\n        expected_metric_keys = [\n            \"actor/pg_loss\",\n            \"actor/pg_clipfrac\",\n            \"actor/ppo_kl\",\n            \"actor/pg_clipfrac_lower\",\n            \"actor/grad_norm\",\n        ]\n\n        for key in expected_metric_keys:\n            self.assertIn(key, metrics)\n            if isinstance(metrics[key], list):\n                self.assertTrue(all(torch.isfinite(torch.tensor(v)) for v in metrics[key]))\n            else:\n                self.assertIsInstance(metrics[key], (float, int))\n                self.assertTrue(torch.isfinite(torch.tensor(metrics[key])))\n\n    def test_dataparallelppoactor_initialization(self):\n        \"\"\"Test DataParallelPPOActor initialization\"\"\"\n        self.assertIsNotNone(self.actor.actor_module)\n        self.assertIsNotNone(self.actor.actor_optimizer)\n        self.assertEqual(self.actor.config, self.config)\n\n        self.assertEqual(self.actor.config.strategy, \"fsdp2\")\n        self.assertEqual(self.actor.config.ppo_mini_batch_size, 4)\n        self.assertEqual(self.actor.config.clip_ratio, 0.2)\n\n    def test_dataparallelppoactor_with_qwen3_model(self):\n        \"\"\"Test DataParallelPPOActor with real Qwen3ForCausalLM model\"\"\"\n        qwen_config = Qwen3Config(\n            vocab_size=1000,\n            hidden_size=64,\n            intermediate_size=128,\n            num_hidden_layers=2,\n            num_attention_heads=4,\n            num_key_value_heads=2,\n            max_position_embeddings=512,\n            torch_dtype=torch.float32,\n            use_cache=False,\n        )\n\n        with torch.device(self.device):\n            qwen_model = AutoModelForCausalLM.from_config(config=qwen_config, torch_dtype=torch.float32).to(self.device)\n\n        qwen_optimizer = torch.optim.Adam(qwen_model.parameters(), lr=1e-4)\n\n        qwen_actor = DataParallelPPOActor(config=self.config, actor_module=qwen_model, actor_optimizer=qwen_optimizer)\n\n        data = self._create_test_data_for_compute_log_prob()\n        outputs = qwen_actor.compute_log_prob(data, calculate_entropy=True)\n        log_probs = outputs[\"log_probs\"]\n        entropys = outputs[\"entropys\"]\n\n        batch_size = data.batch[\"responses\"].shape[0]\n        response_length = data.batch[\"responses\"].shape[1]\n\n        self.assertIsInstance(log_probs, torch.Tensor)\n        self.assertEqual(log_probs.shape, (batch_size, response_length))\n        self.assertTrue(torch.all(torch.isfinite(log_probs)))\n\n        self.assertIsInstance(entropys, torch.Tensor)\n        self.assertEqual(entropys.shape, (batch_size, response_length))\n        self.assertTrue(torch.all(torch.isfinite(entropys)))\n        self.assertTrue(torch.all(entropys >= 0))\n\n        policy_data = self._create_test_data_for_update_policy()\n        metrics = qwen_actor.update_policy(policy_data)\n\n        self.assertIsInstance(metrics, dict)\n\n        expected_metric_keys = [\n            \"actor/pg_loss\",\n            \"actor/pg_clipfrac\",\n            \"actor/ppo_kl\",\n            \"actor/pg_clipfrac_lower\",\n            \"actor/grad_norm\",\n        ]\n\n        for key in expected_metric_keys:\n            self.assertIn(key, metrics)\n            if isinstance(metrics[key], list):\n                self.assertTrue(all(torch.isfinite(torch.tensor(v)) for v in metrics[key]))\n            else:\n                self.assertIsInstance(metrics[key], (float, int))\n                self.assertTrue(torch.isfinite(torch.tensor(metrics[key])))\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/workers/config/test_actor_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport unittest\n\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.workers.config import (\n    ActorConfig,\n    FSDPActorConfig,\n    McoreActorConfig,\n    OptimizerConfig,\n)\n\n\nclass TestActorConfig(unittest.TestCase):\n    \"\"\"Test the ActorConfig dataclass and its variants.\"\"\"\n\n    def test_config_inheritance(self):\n        \"\"\"Test that the inheritance hierarchy works correctly.\"\"\"\n        megatron_dict = {\n            \"_target_\": \"verl.workers.config.McoreActorConfig\",\n            \"strategy\": \"megatron\",\n            \"ppo_mini_batch_size\": 256,\n            \"ppo_micro_batch_size_per_gpu\": 256,\n            \"clip_ratio\": 0.2,\n            \"optim\": {\n                \"_target_\": \"verl.workers.config.McoreOptimizerConfig\",\n                \"lr\": 0.1,\n            },\n            \"rollout_n\": 1,\n        }\n        fsdp_dict = {\n            \"_target_\": \"verl.workers.config.FSDPActorConfig\",\n            \"strategy\": \"fsdp\",\n            \"ppo_mini_batch_size\": 256,\n            \"ppo_micro_batch_size_per_gpu\": 256,\n            \"clip_ratio\": 0.2,\n            \"optim\": {\n                \"_target_\": \"verl.workers.config.FSDPOptimizerConfig\",\n                \"lr\": 0.1,\n            },\n            \"rollout_n\": 1,\n        }\n\n        megatron_config = omega_conf_to_dataclass(megatron_dict)\n        fsdp_config = omega_conf_to_dataclass(fsdp_dict)\n\n        self.assertIsInstance(megatron_config, ActorConfig)\n        self.assertIsInstance(fsdp_config, ActorConfig)\n\n        self.assertEqual(megatron_config.ppo_mini_batch_size, fsdp_config.ppo_mini_batch_size)\n        self.assertEqual(megatron_config.clip_ratio, fsdp_config.clip_ratio)\n\n    def test_actor_config_from_yaml(self):\n        \"\"\"Test creating ActorConfig from YAML file.\"\"\"\n        from hydra import compose, initialize_config_dir\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config/actor\")):\n            cfg = compose(config_name=\"actor\", overrides=[\"strategy=fsdp\", \"ppo_micro_batch_size_per_gpu=128\"])\n\n        config = omega_conf_to_dataclass(cfg)\n\n        self.assertIsInstance(config, ActorConfig)\n        self.assertEqual(config.strategy, \"fsdp\")\n\n    def test_fsdp_actor_config_from_yaml(self):\n        \"\"\"Test creating FSDPActorConfig from YAML file.\"\"\"\n        from hydra import compose, initialize_config_dir\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config/actor\")):\n            cfg = compose(config_name=\"dp_actor\", overrides=[\"strategy=fsdp2\", \"ppo_micro_batch_size_per_gpu=128\"])\n\n        config = omega_conf_to_dataclass(cfg)\n\n        self.assertIsInstance(config, FSDPActorConfig)\n        self.assertEqual(config.strategy, \"fsdp2\")\n\n    def test_megatron_actor_config_from_yaml(self):\n        \"\"\"Test creating McoreActorConfig from YAML file.\"\"\"\n        from hydra import compose, initialize_config_dir\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config/actor\")):\n            cfg = compose(config_name=\"megatron_actor\", overrides=[\"ppo_micro_batch_size_per_gpu=128\"])\n\n        config = omega_conf_to_dataclass(cfg)\n\n        self.assertIsInstance(config, McoreActorConfig)\n        self.assertEqual(config.strategy, \"megatron\")\n\n    def test_config_get_method(self):\n        \"\"\"Test the get method for backward compatibility.\"\"\"\n        config_dict = {\n            \"_target_\": \"verl.workers.config.ActorConfig\",\n            \"strategy\": \"fsdp\",\n            \"ppo_mini_batch_size\": 256,\n            \"ppo_micro_batch_size_per_gpu\": 256,\n            \"optim\": {\n                \"_target_\": \"verl.workers.config.OptimizerConfig\",\n                \"lr\": 0.1,\n            },\n            \"rollout_n\": 1,\n        }\n        config = omega_conf_to_dataclass(config_dict)\n\n        self.assertEqual(config.get(\"strategy\"), \"fsdp\")\n        self.assertEqual(config.get(\"ppo_mini_batch_size\"), 256)\n\n        self.assertIsNone(config.get(\"non_existing\"))\n        self.assertEqual(config.get(\"non_existing\", \"default\"), \"default\")\n\n    def test_config_dict_like_access(self):\n        \"\"\"Test dictionary-like access to config fields.\"\"\"\n        config_dict = {\n            \"_target_\": \"verl.workers.config.ActorConfig\",\n            \"strategy\": \"fsdp\",\n            \"ppo_mini_batch_size\": 256,\n            \"ppo_micro_batch_size_per_gpu\": 256,\n            \"optim\": {\n                \"_target_\": \"verl.workers.config.OptimizerConfig\",\n                \"lr\": 0.1,\n            },\n            \"rollout_n\": 1,\n        }\n        config = omega_conf_to_dataclass(config_dict)\n\n        self.assertEqual(config[\"strategy\"], \"fsdp\")\n        self.assertEqual(config[\"ppo_mini_batch_size\"], 256)\n\n        field_names = list(config)\n        self.assertIn(\"strategy\", field_names)\n        self.assertIn(\"ppo_mini_batch_size\", field_names)\n\n        self.assertGreater(len(config), 0)\n\n    def test_frozen_fields_modification_raises_exception(self):\n        \"\"\"Test that modifying frozen fields raises an exception.\"\"\"\n        config_dict = {\n            \"_target_\": \"verl.workers.config.ActorConfig\",\n            \"strategy\": \"fsdp\",\n            \"ppo_mini_batch_size\": 256,\n            \"ppo_micro_batch_size_per_gpu\": 256,\n            \"optim\": {\n                \"_target_\": \"verl.workers.config.OptimizerConfig\",\n                \"lr\": 0.1,\n            },\n            \"rollout_n\": 1,\n        }\n        config = omega_conf_to_dataclass(config_dict)\n\n        with self.assertRaises(AttributeError):\n            config.strategy = \"megatron\"\n\n        with self.assertRaises(AttributeError):\n            config.clip_ratio = 0.5\n\n        config.ppo_mini_batch_size = 512  # This should work since it's not in frozen fields anymore\n        self.assertEqual(config.ppo_mini_batch_size, 512)\n\n    def test_actor_config_validation_exceptions(self):\n        \"\"\"Test that ActorConfig.__post_init__ raises appropriate validation exceptions.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        with self.assertRaises((ValueError, AssertionError)) as cm:\n            ActorConfig(\n                strategy=\"fsdp\",\n                loss_agg_mode=\"invalid-mode\",\n                use_dynamic_bsz=True,\n                optim=optim,\n                ppo_micro_batch_size_per_gpu=4,\n                rollout_n=1,\n            )\n        self.assertIn(\"Invalid loss_agg_mode\", str(cm.exception))\n\n        with self.assertRaises((ValueError, AssertionError)) as cm:\n            ActorConfig(\n                strategy=\"fsdp\",\n                use_dynamic_bsz=False,\n                ppo_micro_batch_size=4,\n                ppo_micro_batch_size_per_gpu=2,\n                optim=optim,\n                rollout_n=1,\n            )\n        self.assertIn(\"You have set both\", str(cm.exception))\n\n        with self.assertRaises((ValueError, AssertionError)) as cm:\n            ActorConfig(\n                strategy=\"fsdp\",\n                use_dynamic_bsz=False,\n                ppo_micro_batch_size=None,\n                ppo_micro_batch_size_per_gpu=None,\n                optim=optim,\n                rollout_n=1,\n            )\n        self.assertIn(\"Please set at least one\", str(cm.exception))\n\n        config = ActorConfig(\n            strategy=\"fsdp\",\n            use_dynamic_bsz=True,\n            ppo_micro_batch_size=None,\n            ppo_micro_batch_size_per_gpu=None,\n            optim=optim,\n            rollout_n=1,\n        )\n        self.assertIsNotNone(config)  # Should not raise an exception\n\n    def test_fsdp_actor_config_validation_exceptions(self):\n        \"\"\"Test that FSDPActorConfig.validate() raises appropriate validation exceptions.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        config = FSDPActorConfig(\n            strategy=\"fsdp\",\n            ulysses_sequence_parallel_size=2,\n            use_dynamic_bsz=True,  # Skip batch size validation to focus on FSDP validation\n            optim=optim,\n            rollout_n=1,\n        )\n\n        model_config = {\"use_remove_padding\": False}\n        with self.assertRaises(ValueError) as cm:\n            config.validate(n_gpus=8, train_batch_size=256, model_config=model_config)\n        self.assertIn(\"you must enable `use_remove_padding`\", str(cm.exception))\n\n    def test_actor_config_validate_method_exceptions(self):\n        \"\"\"Test that ActorConfig.validate() raises appropriate validation exceptions.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        config = ActorConfig(\n            strategy=\"fsdp\",\n            use_dynamic_bsz=False,\n            ppo_mini_batch_size=256,\n            ppo_micro_batch_size=8,\n            ppo_micro_batch_size_per_gpu=None,  # Ensure only one batch size setting is used\n            optim=optim,\n            rollout_n=1,\n        )\n\n        with self.assertRaises(ValueError) as cm:\n            config.validate(n_gpus=8, train_batch_size=128)\n        self.assertIn(\"train_batch_size\", str(cm.exception))\n\n        with self.assertRaises(ValueError) as cm:\n            config.validate(n_gpus=16, train_batch_size=512)\n        self.assertIn(\"must be >= n_gpus\", str(cm.exception))\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/workers/config/test_critic_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom pathlib import Path\n\nimport pytest\nfrom hydra import compose, initialize_config_dir\n\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.profiler import ProfilerConfig\nfrom verl.workers.config import (\n    CriticConfig,\n    FSDPCriticConfig,\n    FSDPOptimizerConfig,\n    McoreCriticConfig,\n    McoreOptimizerConfig,\n    OptimizerConfig,\n)\n\n\n@pytest.mark.skip(reason=\"This test is flaky when we actively load model config\")\nclass TestCriticConfig:\n    \"\"\"Test suite for critic configuration dataclasses.\"\"\"\n\n    @pytest.fixture\n    def config_dir(self):\n        \"\"\"Get the path to the config directory.\"\"\"\n        return Path(__file__).parent.parent.parent.parent / \"verl\" / \"trainer\" / \"config\" / \"critic\"\n\n    def test_megatron_critic_config_instantiation_from_yaml(self, config_dir):\n        \"\"\"Test that McoreCriticConfig can be instantiated from megatron_critic.yaml.\"\"\"\n        yaml_path = config_dir / \"megatron_critic.yaml\"\n        assert yaml_path.exists(), f\"Config file not found: {yaml_path}\"\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config/critic\")):\n            test_config = compose(config_name=\"megatron_critic\", overrides=[\"ppo_micro_batch_size_per_gpu=1\"])\n\n        megatron_config_obj = omega_conf_to_dataclass(test_config)\n\n        assert isinstance(megatron_config_obj, McoreCriticConfig)\n        assert isinstance(megatron_config_obj, CriticConfig)\n\n        expected_attrs = [\n            \"strategy\",\n            \"rollout_n\",\n            \"optim\",\n            \"model\",\n            \"ppo_mini_batch_size\",\n            \"ppo_max_token_len_per_gpu\",\n            \"cliprange_value\",\n            \"get\",\n            \"nccl_timeout\",\n            \"megatron\",\n            \"load_weight\",\n        ]\n        for attr in expected_attrs:\n            assert hasattr(megatron_config_obj, attr), f\"Missing attribute: {attr}\"\n\n        assert callable(megatron_config_obj.get)\n        assert megatron_config_obj.strategy == \"megatron\"\n\n    def test_fsdp_critic_config_instantiation_from_yaml(self, config_dir):\n        \"\"\"Test that FSDPCriticConfig can be instantiated from dp_critic.yaml.\"\"\"\n        yaml_path = config_dir / \"dp_critic.yaml\"\n        assert yaml_path.exists(), f\"Config file not found: {yaml_path}\"\n\n        with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config/critic\")):\n            test_config = compose(config_name=\"dp_critic\", overrides=[\"ppo_micro_batch_size_per_gpu=1\"])\n\n        fsdp_config_obj = omega_conf_to_dataclass(test_config)\n\n        assert isinstance(fsdp_config_obj, FSDPCriticConfig)\n        assert isinstance(fsdp_config_obj, CriticConfig)\n\n        expected_attrs = [\n            \"strategy\",\n            \"rollout_n\",\n            \"optim\",\n            \"model\",\n            \"ppo_mini_batch_size\",\n            \"ppo_max_token_len_per_gpu\",\n            \"cliprange_value\",\n            \"get\",\n            \"forward_micro_batch_size\",\n            \"forward_micro_batch_size_per_gpu\",\n            \"ulysses_sequence_parallel_size\",\n            \"grad_clip\",\n        ]\n        for attr in expected_attrs:\n            assert hasattr(fsdp_config_obj, attr), f\"Missing attribute: {attr}\"\n\n        assert callable(fsdp_config_obj.get)\n        assert fsdp_config_obj.strategy == \"fsdp\"\n\n    def test_config_inheritance_hierarchy(self):\n        \"\"\"Test that the inheritance hierarchy is correct.\"\"\"\n        megatron_config = McoreCriticConfig(ppo_micro_batch_size_per_gpu=1, optim=McoreOptimizerConfig(lr=0.1))\n        assert isinstance(megatron_config, CriticConfig)\n        assert isinstance(megatron_config, McoreCriticConfig)\n\n        fsdp_config = FSDPCriticConfig(ppo_micro_batch_size_per_gpu=1, optim=FSDPOptimizerConfig(lr=0.1))\n        assert isinstance(fsdp_config, CriticConfig)\n        assert isinstance(fsdp_config, FSDPCriticConfig)\n\n        critic_config = CriticConfig(ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=OptimizerConfig(lr=0.1))\n        assert isinstance(critic_config, CriticConfig)\n        assert not isinstance(critic_config, McoreCriticConfig)\n        assert not isinstance(critic_config, FSDPCriticConfig)\n\n    def test_config_dict_interface(self):\n        \"\"\"Test that configs provide dict-like interface from BaseConfig.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        config = CriticConfig(ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=optim)\n\n        assert \"strategy\" in config\n        assert config[\"strategy\"] == \"fsdp2\"\n\n        assert config.get(\"strategy\") == \"fsdp2\"\n        assert config.get(\"nonexistent_key\", \"default\") == \"default\"\n\n        keys = list(config)\n        assert \"strategy\" in keys\n        assert \"rollout_n\" in keys\n\n        assert len(config) > 0\n\n    def test_frozen_fields_immutability(self):\n        \"\"\"Test that frozen fields raise exceptions when modified after creation.\"\"\"\n        critic_config = CriticConfig(ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=OptimizerConfig(lr=0.1))\n        frozen_fields = [\"rollout_n\", \"strategy\", \"cliprange_value\"]\n\n        for field_name in frozen_fields:\n            with pytest.raises((AttributeError, TypeError, ValueError)):\n                setattr(critic_config, field_name, \"modified_value\")\n\n        megatron_config = McoreCriticConfig(ppo_micro_batch_size_per_gpu=1, optim=McoreOptimizerConfig(lr=0.1))\n        megatron_frozen_fields = [\"nccl_timeout\", \"load_weight\", \"data_loader_seed\"]\n\n        for field_name in megatron_frozen_fields:\n            with pytest.raises((AttributeError, TypeError, ValueError)):\n                setattr(megatron_config, field_name, \"modified_value\")\n\n        fsdp_config = FSDPCriticConfig(ppo_micro_batch_size_per_gpu=1, optim=FSDPOptimizerConfig(lr=0.1))\n        fsdp_frozen_fields = [\"ulysses_sequence_parallel_size\", \"grad_clip\"]\n\n        for field_name in fsdp_frozen_fields:\n            with pytest.raises((AttributeError, TypeError, ValueError)):\n                setattr(fsdp_config, field_name, \"modified_value\")\n\n    def test_batch_size_fields_modifiable(self):\n        \"\"\"Test that batch size fields can be modified after creation.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        critic_config = CriticConfig(ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=optim)\n\n        critic_config.ppo_mini_batch_size = 8\n        critic_config.ppo_micro_batch_size = 4\n        critic_config.ppo_micro_batch_size_per_gpu = 2\n\n        assert critic_config.ppo_mini_batch_size == 8\n        assert critic_config.ppo_micro_batch_size == 4\n        assert critic_config.ppo_micro_batch_size_per_gpu == 2\n\n        fsdp_config = FSDPCriticConfig(ppo_micro_batch_size_per_gpu=1, optim=FSDPOptimizerConfig(lr=0.1))\n\n        fsdp_config.forward_micro_batch_size = 16\n        fsdp_config.forward_micro_batch_size_per_gpu = 8\n\n        assert fsdp_config.forward_micro_batch_size == 16\n        assert fsdp_config.forward_micro_batch_size_per_gpu == 8\n\n    def test_profiler_config_type_validation(self):\n        \"\"\"Test that profiler field has correct type and validation.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        critic_config = CriticConfig(ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=optim)\n        assert isinstance(critic_config.profiler, ProfilerConfig)\n        assert critic_config.profiler.all_ranks is False\n        assert critic_config.profiler.ranks == []\n\n        custom_profiler = ProfilerConfig(all_ranks=True, ranks=[0, 1])\n        critic_config_custom = CriticConfig(\n            profiler=custom_profiler, ppo_micro_batch_size_per_gpu=1, strategy=\"fsdp2\", optim=optim\n        )\n        assert isinstance(critic_config_custom.profiler, ProfilerConfig)\n        assert critic_config_custom.profiler.all_ranks is True\n        assert critic_config_custom.profiler.ranks == [0, 1]\n\n        profiler1 = ProfilerConfig(enable=True, ranks=[0, 1])\n        profiler2 = ProfilerConfig(all_ranks=True, ranks=[1, 2])\n\n        union_result = profiler1.union(profiler2)\n        assert union_result.enable is True\n        assert union_result.all_ranks is True\n        assert set(union_result.ranks) == {0, 1, 2}\n\n        intersect_result = profiler1.intersect(profiler2)\n        assert intersect_result.all_ranks is False\n        assert intersect_result.ranks == [1]\n\n    def test_critic_config_validation_logic(self):\n        \"\"\"Test the __post_init__ validation logic for CriticConfig.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        valid_config = CriticConfig(\n            strategy=\"fsdp2\", ppo_micro_batch_size_per_gpu=2, use_dynamic_bsz=False, optim=optim\n        )\n        assert valid_config.ppo_micro_batch_size_per_gpu == 2\n\n        valid_config2 = CriticConfig(\n            strategy=\"fsdp2\",\n            ppo_micro_batch_size_per_gpu=None,\n            ppo_micro_batch_size=4,\n            ppo_mini_batch_size=8,\n            use_dynamic_bsz=False,\n            optim=optim,\n        )\n        assert valid_config2.ppo_micro_batch_size == 4\n\n        dynamic_config = CriticConfig(\n            strategy=\"fsdp2\", ppo_micro_batch_size_per_gpu=2, use_dynamic_bsz=True, optim=optim\n        )\n        assert dynamic_config.use_dynamic_bsz is True\n\n        with pytest.raises(ValueError, match=\"You have set both.*micro_batch_size.*AND.*micro_batch_size_per_gpu\"):\n            CriticConfig(\n                strategy=\"fsdp2\",\n                ppo_micro_batch_size=4,\n                ppo_micro_batch_size_per_gpu=2,\n                use_dynamic_bsz=False,\n                optim=optim,\n            )\n\n        with pytest.raises(\n            ValueError, match=\"Please set at least one of.*micro_batch_size.*or.*micro_batch_size_per_gpu\"\n        ):\n            CriticConfig(\n                strategy=\"fsdp2\",\n                ppo_micro_batch_size=None,\n                ppo_micro_batch_size_per_gpu=None,\n                use_dynamic_bsz=False,\n                optim=optim,\n            )\n\n    def test_micro_batch_size_divisibility_validation(self):\n        \"\"\"Test micro batch size divisibility validation in __post_init__.\"\"\"\n        optim = OptimizerConfig(lr=0.1)\n        valid_config = CriticConfig(\n            strategy=\"fsdp2\", ppo_micro_batch_size_per_gpu=2, ppo_mini_batch_size=8, use_dynamic_bsz=False, optim=optim\n        )\n        assert valid_config.ppo_mini_batch_size == 8\n        assert valid_config.ppo_micro_batch_size_per_gpu == 2\n\n        valid_config_with_mbs = CriticConfig(\n            strategy=\"fsdp2\", ppo_mini_batch_size=8, ppo_micro_batch_size=4, use_dynamic_bsz=False, optim=optim\n        )\n        assert valid_config_with_mbs.ppo_mini_batch_size == 8\n        assert valid_config_with_mbs.ppo_micro_batch_size == 4\n\n        with pytest.raises(ValueError, match=\"ppo_mini_batch_size.*must be divisible by.*ppo_micro_batch_size\"):\n            CriticConfig(\n                strategy=\"fsdp2\", ppo_mini_batch_size=7, ppo_micro_batch_size=4, use_dynamic_bsz=False, optim=optim\n            )\n\n        dynamic_config = CriticConfig(\n            strategy=\"fsdp2\", ppo_mini_batch_size=7, ppo_micro_batch_size=4, use_dynamic_bsz=True, optim=optim\n        )\n        assert dynamic_config.use_dynamic_bsz is True\n\n    def test_fsdp_sequence_parallelism_validation(self):\n        \"\"\"Test FSDP sequence parallelism validation in FSDPCriticConfig.__post_init__.\"\"\"\n        valid_config = FSDPCriticConfig(\n            ppo_micro_batch_size_per_gpu=2,\n            ulysses_sequence_parallel_size=2,\n            model={\"use_remove_padding\": True},\n            optim=FSDPOptimizerConfig(lr=0.1),\n        )\n        assert valid_config.ulysses_sequence_parallel_size == 2\n\n        with pytest.raises(\n            ValueError, match=\"When using sequence parallelism for critic, you must enable.*use_remove_padding\"\n        ):\n            FSDPCriticConfig(\n                ppo_micro_batch_size_per_gpu=2,\n                ulysses_sequence_parallel_size=2,\n                model={\"use_remove_padding\": False},\n                optim=FSDPOptimizerConfig(lr=0.1),\n            )\n\n        valid_config_no_sp = FSDPCriticConfig(\n            ppo_micro_batch_size_per_gpu=2,\n            ulysses_sequence_parallel_size=1,\n            model={\"use_remove_padding\": False},\n            optim=FSDPOptimizerConfig(lr=0.1),\n        )\n        assert valid_config_no_sp.ulysses_sequence_parallel_size == 1\n"
  },
  {
    "path": "tests/workers/config/test_engine_config_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nfrom verl.workers.config.engine import FSDPEngineConfig, McoreEngineConfig\n\n\nclass TestMcoreEngineConfig:\n    def test_default_values(self):\n        config = McoreEngineConfig()\n        assert config.tensor_model_parallel_size == 1\n        assert config.sequence_parallel is False  # Should be auto-corrected\n        assert config.seed == 42\n\n    def test_post_init_validation(self):\n        # Test TP size 1 forces sequence_parallel=False\n        config = McoreEngineConfig(tensor_model_parallel_size=1)\n        assert config.sequence_parallel is False\n\n        # Test TP >1 keeps sequence_parallel=True\n        config = McoreEngineConfig(tensor_model_parallel_size=2)\n        assert config.sequence_parallel is True\n\n    def test_mutable_fields(self):\n        config = McoreEngineConfig()\n        config.sequence_parallel = True  # Should be mutable\n        with pytest.raises(AttributeError):\n            config.tensor_model_parallel_size = 2  # Frozen field\n\n    @pytest.mark.parametrize(\"offload_field\", [\"param_offload\", \"grad_offload\", \"optimizer_offload\"])\n    def test_offload_flags(self, offload_field):\n        config = McoreEngineConfig(**{offload_field: True})\n        assert getattr(config, offload_field) is True\n\n\nclass TestFSDPEngineConfigCPU:\n    def test_default_values(self):\n        config = FSDPEngineConfig()\n        assert config.param_offload is False\n        assert config.optimizer_offload is False\n        assert config.fsdp_size == -1\n\n    @pytest.mark.parametrize(\n        \"offload_params\",\n        [{\"param_offload\": True}, {\"optimizer_offload\": True}, {\"param_offload\": True, \"optimizer_offload\": True}],\n    )\n    def test_offload_combinations(self, offload_params):\n        config = FSDPEngineConfig(**offload_params)\n        assert config.param_offload == offload_params.get(\"param_offload\", False)\n        assert config.optimizer_offload == offload_params.get(\"optimizer_offload\", False)\n\n    def test_wrap_policy_configuration(self):\n        test_policy = {\"layer_class\": \"TransformerBlock\"}\n        config = FSDPEngineConfig(wrap_policy=test_policy)\n        assert config.wrap_policy == test_policy\n"
  },
  {
    "path": "tests/workers/config/test_model_config_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright Amazon.com and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nimport pytest\nfrom omegaconf import OmegaConf\n\nfrom verl.workers.config.model import HFModelConfig\n\n\nclass TestHFModelConfigCPU:\n    model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B\")  # Just a path string, not loaded\n\n    def test_target_modules_accepts_list_via_omegaconf(self):\n        \"\"\"\n        Test that target_modules field accepts both string and list values\n        when merging OmegaConf configs (simulates CLI override behavior).\n\n        The purpose is to ensure we can pass\n        actor_rollout_ref.model.target_modules='[\"k_proj\",\"o_proj\",\"down_proj\",\"q_proj\"]'\n        \"\"\"\n\n        # Create structured config from the dataclass defaults\n        # This is what omega_conf_to_dataclass does internally\n        cfg_from_dataclass = OmegaConf.structured(HFModelConfig)\n\n        # Simulate CLI override with target_modules as a list\n        cli_config = OmegaConf.create(\n            {\n                \"path\": self.model_path,\n                \"target_modules\": [\"k_proj\", \"o_proj\", \"q_proj\", \"v_proj\"],\n            }\n        )\n\n        # This merge should NOT raise ValidationError\n        # Before the fix (target_modules: str), this would fail with:\n        # \"Cannot convert 'ListConfig' to string\"\n        merged = OmegaConf.merge(cfg_from_dataclass, cli_config)\n\n        # Verify the list was merged correctly\n        assert list(merged.target_modules) == [\"k_proj\", \"o_proj\", \"q_proj\", \"v_proj\"]\n\n    def test_target_modules_accepts_none_via_omegaconf(self):\n        \"\"\"Test that target_modules still accepts None values.\"\"\"\n\n        cfg_from_dataclass = OmegaConf.structured(HFModelConfig)\n\n        cli_config = OmegaConf.create(\n            {\n                \"path\": self.model_path,\n                \"target_modules\": None,\n            }\n        )\n\n        merged = OmegaConf.merge(cfg_from_dataclass, cli_config)\n        assert merged.target_modules is None\n\n    def test_target_modules_accepts_string_via_omegaconf(self):\n        \"\"\"Test that target_modules still accepts string values.\"\"\"\n\n        cfg_from_dataclass = OmegaConf.structured(HFModelConfig)\n\n        cli_config = OmegaConf.create(\n            {\n                \"path\": self.model_path,\n                \"target_modules\": \"all-linear\",\n            }\n        )\n\n        merged = OmegaConf.merge(cfg_from_dataclass, cli_config)\n        assert merged.target_modules == \"all-linear\"\n\n    def test_target_modules_raises_on_invalid_type(self):\n        \"\"\"Test that __post_init__ raises TypeError for invalid target_modules types.\"\"\"\n        base_config = OmegaConf.structured(HFModelConfig)\n        invalid_cli_config = OmegaConf.create(\n            {\n                \"path\": self.model_path,\n                \"target_modules\": [1, 2, 3],  # list of ints instead of strings\n            }\n        )\n        merged_config = OmegaConf.merge(base_config, invalid_cli_config)\n        with pytest.raises(TypeError):\n            OmegaConf.to_object(merged_config)\n"
  },
  {
    "path": "tests/workers/config/test_optim_config_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\nfrom verl.workers.config.optimizer import FSDPOptimizerConfig\n\n\nclass TestFSDPOptimizerConfigCPU:\n    def test_default_configuration(self):\n        config = FSDPOptimizerConfig(lr=0.1)\n        assert config.min_lr_ratio is None\n        assert config.lr_scheduler_type == \"constant\"\n        assert config.num_cycles == 0.5\n\n    @pytest.mark.parametrize(\"lr_scheduler_type\", [\"constant\", \"cosine\"])\n    def test_valid_lr_scheduler_types(self, lr_scheduler_type):\n        config = FSDPOptimizerConfig(lr_scheduler_type=lr_scheduler_type, lr=0.1)\n        assert config.lr_scheduler_type == lr_scheduler_type\n\n    @pytest.mark.parametrize(\"warmup_style\", [\"constant\", \"cosine\"])\n    def test_valid_warmup_style_types(self, warmup_style):\n        config = FSDPOptimizerConfig(warmup_style=warmup_style, lr=0.1)\n        assert config.lr_scheduler_type == warmup_style\n\n    def test_invalid_lr_scheduler_type(self):\n        with pytest.raises((ValueError, AssertionError)):\n            FSDPOptimizerConfig(lr_scheduler_type=\"invalid_style\", lr=0.1)\n\n    def test_invalid_warmup_style_type(self):\n        with pytest.raises((ValueError, AssertionError)):\n            FSDPOptimizerConfig(warmup_style=\"invalid_style\", lr=0.1)\n\n    @pytest.mark.parametrize(\"num_cycles\", [0.1, 1.0, 2.5])\n    def test_num_cycles_configuration(self, num_cycles):\n        config = FSDPOptimizerConfig(num_cycles=num_cycles, lr=0.1)\n        assert config.num_cycles == num_cycles\n"
  },
  {
    "path": "tests/workers/critic/test_special_dp_critic.py",
    "content": "#!/usr/bin/env python3\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nimport tempfile\nimport unittest\nfrom unittest.mock import Mock, patch\n\nimport torch\nimport torch.distributed\nfrom omegaconf import OmegaConf\nfrom tensordict import TensorDict\nfrom transformers import AutoConfig\n\nfrom verl import DataProto\nfrom verl.workers.config import FSDPCriticConfig, FSDPOptimizerConfig\nfrom verl.workers.config.critic import FSDPCriticModelCfg\nfrom verl.workers.config.engine import FSDPEngineConfig\nfrom verl.workers.fsdp_workers import CriticWorker\n\n\nclass TestCriticWorker(unittest.TestCase):\n    @classmethod\n    def setUpClass(cls):\n        \"\"\"Set up distributed environment\"\"\"\n        if not torch.distributed.is_initialized():\n            torch.distributed.init_process_group(\n                backend=\"nccl\" if torch.cuda.is_available() else \"gloo\", init_method=\"env://\"\n            )\n\n        cls.rank = torch.distributed.get_rank()\n        cls.world_size = torch.distributed.get_world_size()\n\n        if torch.cuda.is_available():\n            torch.cuda.set_device(cls.rank)\n            cls.device = torch.device(f\"cuda:{cls.rank}\")\n        else:\n            cls.device = torch.device(\"cpu\")\n\n    @classmethod\n    def tearDownClass(cls):\n        \"\"\"Clean up distributed environment\"\"\"\n        if torch.distributed.is_initialized():\n            torch.distributed.destroy_process_group()\n\n    def setUp(self):\n        \"\"\"Set up test fixtures\"\"\"\n\n        self.device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n        self.temp_dir = tempfile.mkdtemp()\n\n        model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n        config = AutoConfig.from_pretrained(model_path)\n        config.save_pretrained(self.temp_dir)\n\n        self.config = FSDPCriticConfig(\n            strategy=\"fsdp2\",\n            ppo_mini_batch_size=4,\n            ppo_micro_batch_size_per_gpu=2,\n            forward_micro_batch_size_per_gpu=2,\n            ppo_epochs=1,\n            cliprange_value=0.5,\n            grad_clip=1.0,\n            use_dynamic_bsz=False,\n            ulysses_sequence_parallel_size=1,\n            rollout_n=1,\n            optim=FSDPOptimizerConfig(lr=1e-6),\n            model=FSDPCriticModelCfg(\n                path=model_path,\n                tokenizer_path=model_path,\n                fsdp_config=FSDPEngineConfig(fsdp_size=-1),\n                use_remove_padding=False,\n            ),\n        )\n        assert self.world_size <= 4 // 2\n\n    def tearDown(self):\n        \"\"\"Clean up test fixtures\"\"\"\n        import shutil\n\n        shutil.rmtree(self.temp_dir, ignore_errors=True)\n\n    def _create_test_data_for_compute_values(self, batch_size=2, seq_len=10, response_len=5):\n        \"\"\"Create test data for compute_values method\"\"\"\n        input_ids = torch.randint(0, 1000, (batch_size, seq_len), dtype=torch.long)\n        attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)\n        position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)\n        responses = torch.randint(0, 1000, (batch_size, response_len), dtype=torch.long)\n        response_mask = torch.ones(batch_size, response_len, dtype=torch.float)\n\n        batch = TensorDict(\n            {\n                \"input_ids\": input_ids,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n                \"responses\": responses,\n                \"response_mask\": response_mask,\n            },\n            batch_size=[batch_size],\n        )\n\n        data = DataProto(\n            batch=batch, meta_info={\"micro_batch_size\": 2, \"max_token_len\": seq_len, \"use_dynamic_bsz\": False}\n        )\n\n        return data\n\n    def _create_test_data_for_update_critic(self, batch_size=2, seq_len=10, response_len=5):\n        \"\"\"Create test data for update_critic method\"\"\"\n        input_ids = torch.randint(0, 1000, (batch_size, seq_len), dtype=torch.long)\n        attention_mask = torch.ones(batch_size, seq_len, dtype=torch.long)\n        position_ids = torch.arange(seq_len).unsqueeze(0).expand(batch_size, -1)\n        responses = torch.randint(0, 1000, (batch_size, response_len), dtype=torch.long)\n        response_mask = torch.ones(batch_size, response_len, dtype=torch.float)\n        values = torch.randn(batch_size, response_len, dtype=torch.float)\n        returns = torch.randn(batch_size, response_len, dtype=torch.float)\n\n        batch = TensorDict(\n            {\n                \"input_ids\": input_ids,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n                \"responses\": responses,\n                \"response_mask\": response_mask,\n                \"values\": values,\n                \"returns\": returns,\n            },\n            batch_size=[batch_size],\n        )\n\n        data = DataProto(\n            batch=batch,\n            meta_info={\"global_token_num\": [response_len] * batch_size, \"batch_seqlens\": [response_len] * batch_size},\n        )\n\n        return data\n\n    def test_init_model(self):\n        \"\"\"Test CriticWorker.init_model() method\"\"\"\n        worker = CriticWorker(self.config)\n        worker.init_model()\n\n        self.assertIsNotNone(worker.critic_module)\n        self.assertIsNotNone(worker.critic_optimizer)\n        self.assertIsNotNone(worker.critic)\n        self.assertIsNotNone(worker.checkpoint_manager)\n\n    def test_compute_values(self):\n        \"\"\"Test CriticWorker.compute_values() method\"\"\"\n        worker = CriticWorker(self.config)\n        worker.init_model()\n\n        data = self._create_test_data_for_compute_values()\n\n        result = worker.compute_values(data)\n\n        self.assertIsInstance(result, DataProto)\n        self.assertIn(\"values\", result.batch)\n        values = result.batch[\"values\"]\n\n        batch_size, response_len = 2, 5\n        self.assertEqual(values.shape, (batch_size, response_len))\n\n        self.assertTrue(torch.isfinite(values).all())\n\n    def test_update_critic(self):\n        \"\"\"Test CriticWorker.update_critic() method\"\"\"\n        worker = CriticWorker(self.config)\n        worker.init_model()\n\n        data = self._create_test_data_for_update_critic()\n\n        result = worker.update_critic(data)\n\n        self.assertIsInstance(result, DataProto)\n        self.assertIn(\"metrics\", result.meta_info)\n        metrics = result.meta_info[\"metrics\"]\n\n        expected_keys = [\"critic/vf_loss\", \"critic/vf_clipfrac\", \"critic/vpred_mean\", \"critic/grad_norm\"]\n        for key in expected_keys:\n            self.assertIn(key, metrics)\n\n        for key, value in metrics.items():\n            if isinstance(value, list | tuple):\n                for v in value:\n                    self.assertTrue(torch.isfinite(torch.tensor(v)).all())\n            else:\n                self.assertTrue(torch.isfinite(torch.tensor(value)).all())\n\n    @patch(\"transformers.AutoConfig.from_pretrained\")\n    def test_critic_attn_implementation_override_functionality(self, mock_config_from_pretrained):\n        \"\"\"Test that CriticWorker correctly uses attn_implementation from override_config\"\"\"\n\n        # Mock the AutoConfig return value\n        mock_config = Mock()\n        mock_config.tie_word_embeddings = False\n        mock_config.architectures = [\"LlamaForCausalLM\"]\n        mock_config.num_labels = 1\n        mock_config_from_pretrained.return_value = mock_config\n\n        # Test different attn_implementation values\n        test_cases = [\n            (\"eager\", \"eager\"),\n            (\"sdpa\", \"sdpa\"),\n            (\"flash_attention_2\", \"flash_attention_2\"),\n            (None, \"flash_attention_2\"),  # Default case\n        ]\n\n        for override_value, expected_value in test_cases:\n            mock_config_from_pretrained.reset_mock()\n\n            # Create config with override_config\n            config_dict = {\n                \"model\": {\n                    \"path\": \"/test/model/path\",\n                    \"tokenizer_path\": \"/test/tokenizer/path\",\n                    \"fsdp_config\": {\n                        \"fsdp_size\": 1,\n                        \"param_offload\": False,\n                        \"optimizer_offload\": False,\n                    },\n                },\n                \"optim\": {\"lr\": 1e-4, \"type\": \"AdamW\"},\n                \"strategy\": \"fsdp\",\n                \"ppo_mini_batch_size\": 1,\n                \"ppo_epochs\": 1,\n                \"rollout_n\": 1,\n                \"checkpoint\": {\"save_contents\": [], \"load_contents\": []},\n            }\n\n            # Add override_config with attn_implementation if specified\n            if override_value is not None:\n                config_dict[\"model\"][\"override_config\"] = {\"attn_implementation\": override_value}\n\n            # Convert to OmegaConf\n            test_config = OmegaConf.create(config_dict)\n\n            # Test the extraction logic that should happen in CriticWorker._build_critic_model_optimizer\n            override_config = OmegaConf.to_container(OmegaConf.create(test_config.model.get(\"override_config\", {})))\n            extracted_attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n            # Verify the extraction works correctly\n            self.assertEqual(\n                extracted_attn_implementation,\n                expected_value,\n                f\"Expected {expected_value}, got {extracted_attn_implementation} for override_value {override_value}\",\n            )\n\n    def test_critic_model_config_structure(self):\n        \"\"\"Test that critic model config properly incorporates override settings\"\"\"\n\n        # Test configuration scenarios\n        test_scenarios = [\n            {\"name\": \"default_flash_attention\", \"override_config\": {}, \"expected_attn\": \"flash_attention_2\"},\n            {\"name\": \"eager_override\", \"override_config\": {\"attn_implementation\": \"eager\"}, \"expected_attn\": \"eager\"},\n            {\"name\": \"sdpa_override\", \"override_config\": {\"attn_implementation\": \"sdpa\"}, \"expected_attn\": \"sdpa\"},\n            {\n                \"name\": \"mixed_config\",\n                \"override_config\": {\"attn_implementation\": \"eager\", \"dropout\": 0.1, \"num_labels\": 1},\n                \"expected_attn\": \"eager\",\n            },\n        ]\n\n        for scenario in test_scenarios:\n            with self.subTest(scenario=scenario[\"name\"]):\n                # Simulate the config processing logic from CriticWorker\n                override_config = scenario[\"override_config\"]\n\n                # Test the extraction logic\n                extracted_attn = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n                # Verify correct extraction\n                self.assertEqual(extracted_attn, scenario[\"expected_attn\"], f\"Failed for scenario {scenario['name']}\")\n\n                # Verify other configs are preserved\n                if \"dropout\" in override_config:\n                    self.assertEqual(override_config[\"dropout\"], 0.1)\n\n    def test_critic_hydra_config_compatibility(self):\n        \"\"\"Test that Hydra +prefix configurations work correctly for CriticWorker\"\"\"\n\n        # Simulate Hydra configuration with +prefix for critic\n        # This would come from: +critic.model.override_config.attn_implementation=eager\n        hydra_config_dict = {\n            \"critic\": {\"model\": {\"path\": \"/test/model/path\", \"override_config\": {\"attn_implementation\": \"eager\"}}}\n        }\n\n        omegaconf = OmegaConf.create(hydra_config_dict)\n\n        # Extract override config as would be done in CriticWorker\n        override_model_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {}))\n        )\n\n        # Test extraction\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        self.assertEqual(attn_implementation, \"eager\")\n\n    def test_critic_backward_compatibility(self):\n        \"\"\"Test that CriticWorker maintains backward compatibility with existing configurations\"\"\"\n\n        # Test cases for backward compatibility\n        compatibility_tests = [\n            {\"name\": \"no_override_config\", \"config\": {}, \"expected\": \"flash_attention_2\"},\n            {\"name\": \"empty_override_config\", \"config\": {\"override_config\": {}}, \"expected\": \"flash_attention_2\"},\n            {\n                \"name\": \"other_overrides_only\",\n                \"config\": {\"override_config\": {\"dropout\": 0.1, \"hidden_size\": 768}},\n                \"expected\": \"flash_attention_2\",\n            },\n        ]\n\n        for test in compatibility_tests:\n            with self.subTest(test=test[\"name\"]):\n                override_config = test[\"config\"].get(\"override_config\", {})\n                attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n                self.assertEqual(\n                    attn_implementation, test[\"expected\"], f\"Backward compatibility failed for {test['name']}\"\n                )\n\n    def test_critic_and_actor_independent_configuration(self):\n        \"\"\"Test that critic and actor can have independent attention implementation configurations\"\"\"\n\n        # Simulate a complete training configuration with both actor and critic\n        complete_config = {\n            \"actor_rollout_ref\": {\"model\": {\"override_config\": {\"attn_implementation\": \"eager\"}}},\n            \"critic\": {\"model\": {\"override_config\": {\"attn_implementation\": \"sdpa\"}}},\n        }\n\n        omegaconf = OmegaConf.create(complete_config)\n\n        # Extract actor config\n        actor_override = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.actor_rollout_ref.model.get(\"override_config\", {}))\n        )\n        actor_attn = actor_override.get(\"attn_implementation\", \"flash_attention_2\")\n\n        # Extract critic config\n        critic_override = OmegaConf.to_container(OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {})))\n        critic_attn = critic_override.get(\"attn_implementation\", \"flash_attention_2\")\n\n        # Verify independent configuration\n        self.assertEqual(actor_attn, \"eager\")\n        self.assertEqual(critic_attn, \"sdpa\")\n        self.assertNotEqual(actor_attn, critic_attn)  # Ensure they are indeed different\n\n\nif __name__ == \"__main__\":\n    unittest.main()\n"
  },
  {
    "path": "tests/workers/reward_manager/test_registry_on_cpu.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\n\n# Assuming REWARD_MANAGER_REGISTRY is defined somewhere in the module\nfrom verl.workers.reward_manager.registry import REWARD_MANAGER_REGISTRY, get_reward_manager_cls, register\n\n\n@pytest.fixture\ndef setup():\n    \"\"\"Setup test cases with a mock registry.\"\"\"\n    REWARD_MANAGER_REGISTRY.clear()\n    REWARD_MANAGER_REGISTRY.update({\"manager1\": \"Manager1Class\", \"manager2\": \"Manager2Class\"})\n    return REWARD_MANAGER_REGISTRY\n\n\ndef test_get_existing_manager(setup):\n    \"\"\"Test getting an existing reward manager class.\"\"\"\n    assert get_reward_manager_cls(\"manager1\") == \"Manager1Class\"\n    assert get_reward_manager_cls(\"manager2\") == \"Manager2Class\"\n\n\ndef test_get_nonexistent_manager(setup):\n    \"\"\"Test getting a non-existent reward manager raises ValueError.\"\"\"\n    with pytest.raises(ValueError) as excinfo:\n        get_reward_manager_cls(\"unknown_manager\")\n    assert \"Unknown reward manager: unknown_manager\" in str(excinfo.value)\n\n\ndef test_case_sensitivity(setup):\n    \"\"\"Test that manager names are case-sensitive.\"\"\"\n    with pytest.raises(ValueError):\n        get_reward_manager_cls(\"MANAGER1\")\n    with pytest.raises(ValueError):\n        get_reward_manager_cls(\"Manager1\")\n\n\ndef test_empty_registry(setup):\n    \"\"\"Test behavior when registry is empty.\"\"\"\n    REWARD_MANAGER_REGISTRY.clear()\n    with pytest.raises(ValueError) as excinfo:\n        get_reward_manager_cls(\"any_manager\")\n    assert \"Unknown reward manager: any_manager\" in str(excinfo.value)\n\n\ndef test_register_new_class(setup):\n    \"\"\"Test registering a new class with the decorator.\"\"\"\n\n    @register(\"test_manager\")\n    class TestManager:\n        pass\n\n    assert \"test_manager\" in REWARD_MANAGER_REGISTRY\n    assert REWARD_MANAGER_REGISTRY[\"test_manager\"] == TestManager\n\n\ndef test_register_different_classes_same_name(setup):\n    \"\"\"Test that registering different classes with same name raises ValueError.\"\"\"\n\n    @register(\"conflict_manager\")\n    class Manager1:\n        pass\n\n    with pytest.raises(ValueError):\n\n        @register(\"conflict_manager\")\n        class Manager2:\n            pass\n\n    assert REWARD_MANAGER_REGISTRY[\"conflict_manager\"] == Manager1\n\n\ndef test_decorator_returns_original_class(setup):\n    \"\"\"Test that the decorator returns the original class unchanged.\"\"\"\n\n    @register(\"return_test\")\n    class OriginalClass:\n        def method(setup):\n            return 42\n\n    assert OriginalClass().method() == 42\n    assert REWARD_MANAGER_REGISTRY[\"return_test\"] == OriginalClass\n"
  },
  {
    "path": "tests/workers/rollout/perf/vllm_async_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nCompare vLLM AsyncLLM backend: ExternalRayDistributedExecutor(remote call) vs RayDistributedExecutor(compiled graph)\n\n1. Prepare openai/gsm8k dataset\npython3 examples/data_preprocess/gsm8k.py\n\n2. Run perf test\npython3 tests/workers/rollout/perf/vllm_async_rollout.py >perf.log 2>&1\n\nhardware: Nvidia 8*H20\npackages:\n- torch==2.6.0\n- vllm==0.8.5\n\n[DEBUG] backend: sync, n_gpus_per_node: 8, batch_size: 2048, step: 0, step_time: 21.27 secs\n[DEBUG] backend: zeromq, n_gpus_per_node: 8, batch_size: 2048, step: 0, step_time: 23.40 secs\n[DEBUG] backend: ray, n_gpus_per_node: 8, batch_size: 2048, step: 0, step_time: 25.33 secs\n\"\"\"\n\nimport os\nimport time\n\nimport ray\nfrom omegaconf import DictConfig\nfrom torch.utils.data import SequentialSampler\nfrom torchdata.stateful_dataloader import StatefulDataLoader\n\nfrom tests.experimental.agent_loop.agent_utils import AgentLoopManager, RayWorkerGroup, init_agent_loop_manager\nfrom verl.protocol import DataProto\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.dataset import RLHFDataset\nfrom verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn\n\n\ndef init_config(n_gpus_per_node) -> DictConfig:\n    import os\n\n    from hydra import compose, initialize_config_dir\n\n    with initialize_config_dir(config_dir=os.path.abspath(\"verl/trainer/config\")):\n        config = compose(\n            config_name=\"ppo_trainer\",\n            overrides=[\n                \"actor_rollout_ref.actor.use_dynamic_bsz=true\",\n                \"actor_rollout_ref.actor.fsdp_config.param_offload=True\",\n                \"actor_rollout_ref.actor.fsdp_config.optimizer_offload=True\",\n            ],\n        )\n    config.trainer.n_gpus_per_node = n_gpus_per_node\n    config.data.train_batch_size = 128\n    config.data.return_raw_chat = True\n    config.actor_rollout_ref.model.path = \"Qwen/Qwen2.5-7B-Instruct\"\n    config.actor_rollout_ref.rollout.mode = \"async\"\n    config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n    config.actor_rollout_ref.rollout.gpu_memory_utilization = 0.9\n    config.actor_rollout_ref.rollout.multi_turn.format = \"hermes\"\n    config.actor_rollout_ref.rollout.prompt_length = 4096\n    config.actor_rollout_ref.rollout.response_length = 4096\n    config.actor_rollout_ref.rollout.n = 16\n\n    return config\n\n\ndef initialize(config, backend) -> tuple[AgentLoopManager | RayWorkerGroup, StatefulDataLoader]:\n    env_vars = {\n        \"NCCL_DEBUG\": \"WARN\",\n        \"VLLM_USE_V1\": \"1\",\n        \"VERL_VLLM_DISTRIBUTED_BACKEND\": backend,\n    }\n    ray.init(runtime_env={\"env_vars\": env_vars})\n\n    # STEP 1: init async llm server\n    server = init_agent_loop_manager(config)\n\n    # STEP 2: create dataloader\n    tokenizer = hf_tokenizer(config.actor_rollout_ref.model.path)\n    dataset = RLHFDataset(\n        data_files=os.path.expanduser(\"~/data/gsm8k/train.parquet\"),\n        tokenizer=tokenizer,\n        config=config.data,\n    )\n    dataloader = StatefulDataLoader(\n        dataset=dataset,\n        batch_size=config.data.get(\"gen_batch_size\", config.data.train_batch_size),\n        num_workers=config.data.get(\"dataloader_num_workers\", 8),\n        drop_last=True,\n        collate_fn=default_collate_fn,\n        sampler=SequentialSampler(dataset),\n    )\n\n    return server, dataloader\n\n\ndef perf_rollout(mode, backend, n_gpus_per_node, num_steps):\n    config = init_config(n_gpus_per_node)\n    config.actor_rollout_ref.rollout.mode = mode\n    agent_loop_manager, dataloader = initialize(config, backend)\n\n    for step, batch in enumerate(dataloader):\n        batch: DataProto = DataProto.from_single_dict(batch)\n        batch = batch.pop(\n            batch_keys=[\"input_ids\", \"attention_mask\", \"position_ids\"],\n            non_tensor_batch_keys=[\"raw_prompt_ids\", \"raw_prompt\"],\n        )\n        t_start = time.time()\n        gen_batch = agent_loop_manager.generate_sequences(batch)\n        t_end = time.time()\n        print(\n            f\"[DEBUG] backend: {backend}, n_gpus_per_node: {n_gpus_per_node}, batch_size: {len(gen_batch)}, \"\n            f\"step: {step}, step_time: {t_end - t_start:.2f} secs\"\n        )\n        if step + 1 >= num_steps:\n            break\n\n    ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    num_steps = 1\n    n_gpus_per_node = 8\n\n    # test_cases = [(\"sync\", \"sync\"), (\"async\", \"zeromq\"), (\"async\", \"ray\")]\n    test_cases = [(\"async\", \"zeromq\"), (\"async\", \"ray\")]\n    for mode, backend in test_cases:\n        perf_rollout(mode=mode, backend=backend, n_gpus_per_node=n_gpus_per_node, num_steps=num_steps)\n"
  },
  {
    "path": "tests/workers/rollout/resource/tool_configs/mcp_server.json",
    "content": "{\n    \"mcpServers\": {\n        \"Tavily Expert\": {\n            \"url\": \"https://tavily.api.tadata.com/mcp/tavily/your_expert\",\n            \"auth_token\": \"your_tavily_token\"\n        }\n    }\n}"
  },
  {
    "path": "tests/workers/rollout/resource/tool_configs/mcp_tool_config",
    "content": "tools:\n  - class_name: verl.tools.mcp_search_tool.MCPSearchTool\n    config:\n      rate_limit: 120\n      timeout: 120\n      type: mcp\n    mcp:\n      mcp_servers_config_path: ./resource/tool_configs/mcp_server.json\n      # optional\n      tool_selected_list: \n        - tavily_search_tool"
  },
  {
    "path": "tests/workers/rollout/resource/tool_configs/sandbox_fusion_tool_config",
    "content": "tools:\n  - class_name: \"verl.tools.sandbox_fusion_tools.SandboxFusionTool\"\n    config: \n      sandbox_fusion_url: \"https://xxx.apigateway-cn-beijing.volceapi.com/run_code\"\n      type: native\n    tool_schema:\n      type: \"function\"\n      function:\n        name: \"code_interpreter\"\n        description: \"A tool for executing code.\"\n        parameters:\n          type: \"object\"\n          properties:\n            code:\n              type: \"string\"\n              description: \"The code to execute.\"\n          required: [\"code\"]"
  },
  {
    "path": "tests/workers/rollout/resource/tool_configs/search_tool_config",
    "content": "tools:\n  - class_name: verl.tools.search_tool.SearchTool\n    config:\n      retrieval_service_url: http://127.0.0.1:8000/retrieve\n      num_workers: 120\n      rate_limit: 120\n      timeout: 30\n      type: native\n    tool_schema:\n      type: function\n      function:\n        name: search\n        description: Searches the web for relevant information based on the given query.\n        parameters:\n          type: object\n          properties:\n            query_list:\n              type: array\n              item:\n                type: string\n              description: A list of fully-formed semantic queries. The tool will return search results for each query.\n          required: \n            - query_list"
  },
  {
    "path": "tests/workers/rollout/rollout_sglang/test_http_server_engine.py",
    "content": "# Copyright 2025 z.ai\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# This file is adapted from multiple sources:\n# 1. THUDM/slime project\n#    Original source: https://github.com/THUDM/slime/blob/main/slime/backends/sglang_utils/http_server_engine.py\n#    Copyright 2025 z.ai\n#    Licensed under the Apache License, Version 2.0\n# 2. SGLang project\n#    Original source: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server_engine.py\n#    Copyright 2023-2024 SGLang Team\n#    Licensed under the Apache License, Version 2.0\n#\n# Modifications made by z.ai and ModelBest Inc. include but are not limited to:\n# - Enhanced error handling and retry logic\n# - Added async support with connection pooling\n# - Extended functionality for distributed weight updates\n# - Improved logging and monitoring capabilities\n# - Additional configuration options and optimizations\n\n\"\"\"Complete unit tests for HTTP Server Engine Adapters.\n\nThis module contains comprehensive unit tests for both HttpServerEngineAdapter\nand AsyncHttpServerEngineAdapter classes, covering all public methods,\nerror handling scenarios, edge cases, and boundary conditions using pytest and mock frameworks.\n\nTests use real SGLang modules for integration testing while mocking external dependencies.\n\"\"\"\n\nimport asyncio\nfrom unittest.mock import AsyncMock, Mock, patch\n\nimport aiohttp\nimport pytest\nimport requests\nfrom sglang.srt.managers.io_struct import (\n    UpdateWeightsFromTensorReqInput,\n)\nfrom sglang.srt.utils import MultiprocessingSerializer\n\n# Import the module under test\nfrom verl.workers.rollout.sglang_rollout.http_server_engine import (\n    AsyncHttpServerAdapter,\n    HttpServerAdapter,\n    launch_server_process,\n)\n\n\n@pytest.fixture(scope=\"session\")\ndef event_loop():\n    \"\"\"Create an event loop for the entire test session.\"\"\"\n    loop = asyncio.new_event_loop()\n    yield loop\n    loop.close()\n\n\n@pytest.fixture\ndef basic_adapter_kwargs():\n    \"\"\"Provide basic kwargs for creating HTTP server adapters.\"\"\"\n    return {\n        \"host\": \"localhost\",\n        \"port\": 8000,\n        \"node_rank\": 0,\n        \"model_path\": \"/tmp/test_model\",\n    }\n\n\n@pytest.fixture\ndef router_adapter_kwargs():\n    \"\"\"Provide kwargs for creating adapters with router configuration.\"\"\"\n    return {\n        \"router_ip\": \"192.168.1.1\",\n        \"router_port\": 8080,\n        \"host\": \"localhost\",\n        \"port\": 8000,\n        \"node_rank\": 0,\n        \"model_path\": \"/tmp/test_model\",\n    }\n\n\n@pytest.fixture\ndef non_master_adapter_kwargs():\n    \"\"\"Provide kwargs for creating non-master node adapters.\"\"\"\n    return {\n        \"host\": \"localhost\",\n        \"port\": 8000,\n        \"node_rank\": 1,  # Non-master\n        \"model_path\": \"/tmp/test_model\",\n    }\n\n\n@pytest.fixture\ndef mock_launch_server_process():\n    \"\"\"Mock the launch_server_process function for testing without actual server startup.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.launch_server_process\") as mock_launch:\n        mock_process = Mock()\n        mock_process.is_alive.return_value = True\n        mock_process.pid = 12345\n        mock_launch.return_value = mock_process\n        yield mock_launch\n\n\n@pytest.fixture\ndef mock_multiprocessing_process():\n    \"\"\"Create mock multiprocessing.Process for testing without actual process creation.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.multiprocessing.Process\") as mock_process_class:\n        mock_process = Mock()\n        mock_process.is_alive.return_value = True\n        mock_process.pid = 12345\n        mock_process_class.return_value = mock_process\n        yield mock_process\n\n\n@pytest.fixture\ndef mock_requests_session():\n    \"\"\"Create mock requests.Session for testing HTTP interactions.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.Session\") as mock_session_class:\n        mock_session = Mock()\n        mock_response = Mock()\n        mock_response.status_code = 200\n        mock_response.json.return_value = {\"status\": \"success\"}\n        mock_session.get.return_value = mock_response\n        mock_session.post.return_value = mock_response\n        mock_session_class.return_value.__enter__.return_value = mock_session\n        yield mock_session\n\n\n@pytest.fixture\ndef mock_requests_post():\n    \"\"\"Mock requests.post for testing HTTP POST requests.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n        mock_response = Mock()\n        mock_response.status_code = 200\n        mock_response.json.return_value = {\"status\": \"success\"}\n        mock_post.return_value = mock_response\n        yield mock_post\n\n\n@pytest.fixture\ndef mock_requests_get():\n    \"\"\"Mock requests.get for testing HTTP GET requests.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.get\") as mock_get:\n        mock_response = Mock()\n        mock_response.status_code = 200\n        mock_response.json.return_value = {\"status\": \"success\"}\n        mock_get.return_value = mock_response\n        yield mock_get\n\n\n@pytest.fixture\ndef mock_aiohttp_session():\n    \"\"\"Create mock aiohttp.ClientSession for testing async HTTP interactions.\"\"\"\n    mock_session = AsyncMock()\n    mock_session.closed = False\n\n    # Mock response\n    mock_response = AsyncMock()\n    mock_response.status = 200\n    mock_response.json = AsyncMock(return_value={\"status\": \"success\"})\n    mock_response.raise_for_status = Mock()\n\n    # Mock context managers\n    mock_session.get.return_value.__aenter__.return_value = mock_response\n    mock_session.post.return_value.__aenter__.return_value = mock_response\n\n    return mock_session\n\n\n@pytest.fixture\ndef mock_kill_process_tree():\n    \"\"\"Mock kill_process_tree function for testing cleanup without actual process termination.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.kill_process_tree\") as mock_kill:\n        yield mock_kill\n\n\n# Test environment fixtures for real SGLang testing\n@pytest.fixture(scope=\"session\")\ndef sglang_test_model_path():\n    \"\"\"Provide a test model path for SGLang tests.\n\n    This can be overridden by environment variable SGLANG_TEST_MODEL_PATH\n    for tests that need a real model.\n    \"\"\"\n    import os\n\n    return os.getenv(\"SGLANG_TEST_MODEL_PATH\", \"/tmp/test_model\")\n\n\n@pytest.fixture\ndef real_adapter_kwargs(sglang_test_model_path):\n    \"\"\"Provide kwargs for creating adapters with real SGLang integration.\"\"\"\n    return {\n        \"host\": \"localhost\",\n        \"port\": 8000,\n        \"node_rank\": 0,\n        \"model_path\": sglang_test_model_path,\n    }\n\n\n@pytest.fixture(autouse=True)\ndef mock_server_args_post_init():\n    \"\"\"Mock ServerArgs.__post_init__ to skip model path validation.\"\"\"\n    from unittest.mock import patch\n\n    with patch(\n        \"verl.workers.rollout.sglang_rollout.http_server_engine.ServerArgs.__post_init__\", return_value=None\n    ) as mock_post_init:\n        yield mock_post_init\n\n\nclass TestLaunchServerProcess:\n    \"\"\"Test cases for launch_server_process function.\"\"\"\n\n    def test_launch_server_process_success(\n        self, mock_multiprocessing_process, mock_requests_session, real_adapter_kwargs\n    ):\n        \"\"\"Test successful server process launch and health check.\"\"\"\n        # Import real SGLang ServerArgs\n        from sglang.srt.server_args import ServerArgs\n\n        # Create server args using real ServerArgs\n        server_args = ServerArgs(**real_adapter_kwargs)\n\n        # Test\n        with patch(\n            \"verl.workers.rollout.sglang_rollout.http_server_engine.multiprocessing.Process\"\n        ) as mock_process_class:\n            mock_process_class.return_value = mock_multiprocessing_process\n            with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.Session\") as mock_session_class:\n                mock_session_class.return_value.__enter__.return_value = mock_requests_session\n\n                result = launch_server_process(server_args, first_rank_in_node=True)\n\n                # Assertions\n                assert result == mock_multiprocessing_process\n                mock_multiprocessing_process.start.assert_called_once()\n                assert mock_requests_session.get.call_count >= 2  # health_generate and flush_cache\n\n    def test_launch_server_process_non_master(self, mock_multiprocessing_process, non_master_adapter_kwargs):\n        \"\"\"Test server launch for non-master nodes (should return immediately).\"\"\"\n        from sglang.srt.server_args import ServerArgs\n\n        server_args = ServerArgs(**non_master_adapter_kwargs)\n\n        with patch(\n            \"verl.workers.rollout.sglang_rollout.http_server_engine.multiprocessing.Process\"\n        ) as mock_process_class:\n            mock_process_class.return_value = mock_multiprocessing_process\n            result = launch_server_process(server_args, first_rank_in_node=True)\n\n            assert result == mock_multiprocessing_process\n            mock_multiprocessing_process.start.assert_not_called()\n\n    def test_launch_server_process_timeout(self, mock_multiprocessing_process, real_adapter_kwargs):\n        \"\"\"Test timeout during server health check.\"\"\"\n        from sglang.srt.server_args import ServerArgs\n\n        server_args = ServerArgs(**real_adapter_kwargs)\n\n        with patch(\n            \"verl.workers.rollout.sglang_rollout.http_server_engine.multiprocessing.Process\"\n        ) as mock_process_class:\n            mock_process_class.return_value = mock_multiprocessing_process\n            with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.Session\") as mock_session_class:\n                mock_session = Mock()\n                mock_session.get.side_effect = requests.RequestException(\"Connection failed\")\n                mock_session_class.return_value.__enter__.return_value = mock_session\n\n            import itertools\n\n            with patch(\n                \"verl.workers.rollout.sglang_rollout.http_server_engine.time.time\",\n                side_effect=itertools.chain([0], itertools.repeat(400)),  # 第一次返回0，之后一直返回400\n            ):\n                with pytest.raises(TimeoutError):\n                    launch_server_process(server_args, first_rank_in_node=True)\n\n                mock_multiprocessing_process.terminate.assert_called_once()\n\n    def test_launch_server_process_died(self, real_adapter_kwargs):\n        \"\"\"Test server process dies during startup.\"\"\"\n        from sglang.srt.server_args import ServerArgs\n\n        server_args = ServerArgs(**real_adapter_kwargs)\n\n        with patch(\n            \"verl.workers.rollout.sglang_rollout.http_server_engine.multiprocessing.Process\"\n        ) as mock_process_class:\n            mock_process = Mock()\n            mock_process.is_alive.return_value = False\n            mock_process_class.return_value = mock_process\n\n            with pytest.raises(RuntimeError, match=\"Server process terminated unexpectedly\"):\n                launch_server_process(server_args, first_rank_in_node=True)\n\n\nclass TestHttpServerEngineAdapter:\n    \"\"\"Test cases for HttpServerEngineAdapter class.\"\"\"\n\n    def test_init_with_router_registration(self, mock_launch_server_process, mock_requests_post, router_adapter_kwargs):\n        \"\"\"Test initialization with router registration.\"\"\"\n        adapter = HttpServerAdapter(**router_adapter_kwargs)\n\n        assert adapter.router_ip == \"192.168.1.1\"\n        assert adapter.router_port == 8080\n        assert adapter.process == mock_launch_server_process.return_value\n        mock_requests_post.assert_called_once()\n\n    def test_init_without_router(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test initialization without router registration.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        assert adapter.router_ip is None\n        assert adapter.router_port is None\n        assert adapter.process == mock_launch_server_process.return_value\n\n    def test_register_with_router_failure(self, mock_launch_server_process, router_adapter_kwargs):\n        \"\"\"Test router registration failure handling.\"\"\"\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            mock_post.side_effect = requests.RequestException(\"Connection failed\")\n\n            # Should not raise exception, just log error\n            adapter = HttpServerAdapter(**router_adapter_kwargs)\n\n            assert adapter.router_ip == \"192.168.1.1\"\n            mock_post.assert_called_once()\n\n    def test_make_request_success(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test successful HTTP request.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            mock_response = Mock()\n            mock_response.status_code = 200\n            mock_response.json.return_value = {\"status\": \"success\"}\n            mock_post.return_value = mock_response\n\n            result = adapter._make_request(\"test_endpoint\", {\"param\": \"value\"})\n\n            assert result == {\"status\": \"success\"}\n            mock_post.assert_called_with(\n                \"http://localhost:8000/test_endpoint\",\n                json={\"param\": \"value\"},\n                timeout=adapter.timeout,\n            )\n\n    def test_make_request_get_method(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test HTTP GET request.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.get\") as mock_get:\n            mock_response = Mock()\n            mock_response.status_code = 200\n            mock_response.json.return_value = {\"data\": \"test\"}\n            mock_get.return_value = mock_response\n\n            result = adapter._make_request(\"test_endpoint\", method=\"GET\")\n\n            assert result == {\"data\": \"test\"}\n            mock_get.assert_called_with(\"http://localhost:8000/test_endpoint\", timeout=adapter.timeout)\n\n    def test_make_request_non_master(self, mock_launch_server_process):\n        \"\"\"Test request from non-master node returns empty dict.\"\"\"\n        kwargs = {\"host\": \"localhost\", \"port\": 8000, \"node_rank\": 1, \"model_path\": \"/tmp/test_model\"}\n        adapter = HttpServerAdapter(**kwargs)\n        result = adapter._make_request(\"test_endpoint\")\n\n        assert result == {}\n\n    def test_make_request_retry_logic(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test retry logic for failed requests.\"\"\"\n        adapter = HttpServerAdapter(max_attempts=3, **basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            with patch(\"time.sleep\") as mock_sleep:\n                # First two calls fail, third succeeds\n                mock_post.side_effect = [\n                    requests.exceptions.Timeout(),\n                    requests.exceptions.ConnectionError(),\n                    Mock(status_code=200, json=lambda: {\"success\": True}),\n                ]\n\n                result = adapter._make_request(\"test_endpoint\")\n\n                assert result == {\"success\": True}\n                assert mock_post.call_count == 3\n                assert mock_sleep.call_count == 2\n\n    def test_make_request_http_error(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test HTTP error handling.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            mock_response = Mock()\n            mock_response.raise_for_status.side_effect = requests.exceptions.HTTPError(\"404 Not Found\")\n            mock_post.return_value = mock_response\n\n            with pytest.raises(requests.exceptions.HTTPError):\n                adapter._make_request(\"test_endpoint\")\n\n    def test_make_request_max_attempts_exceeded(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test max retries exceeded.\"\"\"\n        adapter = HttpServerAdapter(max_attempts=1, **basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            with patch(\"time.sleep\"):\n                mock_post.side_effect = requests.exceptions.Timeout()\n\n                with pytest.raises(RuntimeError, match=\"Failed to complete request\"):\n                    adapter._make_request(\"test_endpoint\")\n\n                assert mock_post.call_count == 1  # Initial retry\n\n    def test_update_weights_from_tensor_strict(self, mock_launch_server_process, basic_adapter_kwargs):\n        import base64\n\n        from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput\n\n        from verl.workers.rollout.sglang_rollout.http_server_engine import HttpServerAdapter\n\n        basic_adapter_kwargs.setdefault(\"node_rank\", 0)\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"updated\"}\n\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=[b\"tensor1\", b\"tensor2\"],\n                load_format=\"safetensors\",\n                flush_cache=True,\n            )\n            result = adapter.update_weights_from_tensor(req)\n\n            assert result == {\"status\": \"updated\"}\n\n            expected_b64_1 = base64.b64encode(b\"tensor1\").decode(\"utf-8\")\n            expected_b64_2 = base64.b64encode(b\"tensor2\").decode(\"utf-8\")\n\n            mock_request.assert_called_once_with(\n                \"update_weights_from_tensor\",\n                {\n                    \"serialized_named_tensors\": [expected_b64_1, expected_b64_2],\n                    \"load_format\": \"safetensors\",\n                    \"flush_cache\": True,\n                },\n            )\n\n    def test_update_weights_from_tensor_empty(self, mock_launch_server_process, basic_adapter_kwargs):\n        from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput\n\n        from verl.workers.rollout.sglang_rollout.http_server_engine import HttpServerAdapter\n\n        basic_adapter_kwargs.setdefault(\"node_rank\", 0)\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"updated\"}\n\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=[],\n                load_format=\"safetensors\",\n                flush_cache=True,\n            )\n            result = adapter.update_weights_from_tensor(req)\n\n            assert result == {\"status\": \"updated\"}\n\n            mock_request.assert_called_once_with(\n                \"update_weights_from_tensor\",\n                {\n                    \"serialized_named_tensors\": [],\n                    \"load_format\": \"safetensors\",\n                    \"flush_cache\": True,\n                },\n            )\n\n    def test_update_weights_from_tensor_none(self, mock_launch_server_process, basic_adapter_kwargs):\n        from sglang.srt.managers.io_struct import UpdateWeightsFromTensorReqInput\n\n        from verl.workers.rollout.sglang_rollout.http_server_engine import HttpServerAdapter\n\n        basic_adapter_kwargs.setdefault(\"node_rank\", 0)\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"updated\"}\n\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=None,\n                load_format=\"safetensors\",\n                flush_cache=True,\n            )\n            result = adapter.update_weights_from_tensor(req)\n\n            assert result == {\"status\": \"updated\"}\n\n            mock_request.assert_called_once_with(\n                \"update_weights_from_tensor\",\n                {\n                    \"serialized_named_tensors\": [],\n                    \"load_format\": \"safetensors\",\n                    \"flush_cache\": True,\n                },\n            )\n\n    def test_generate(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test generate method.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"text\": \"Generated text\"}\n\n            result = adapter.generate(\n                prompt=\"Hello world\",\n                sampling_params={\"temperature\": 0.7},\n                return_logprob=True,\n            )\n\n            assert result == {\"text\": \"Generated text\"}\n            mock_request.assert_called_once_with(\n                \"generate\",\n                {\n                    \"text\": \"Hello world\",\n                    \"sampling_params\": {\"temperature\": 0.7},\n                    \"return_logprob\": True,\n                },\n                only_master=False,\n            )\n\n    def test_flush_cache(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test flush_cache method.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.get\") as mock_get:\n            with patch(\"time.sleep\") as mock_sleep:\n                # First call fails, second succeeds\n                mock_responses = [\n                    Mock(status_code=503),  # Service unavailable\n                    Mock(status_code=200, json=lambda: {\"cache_flushed\": True}),\n                ]\n                mock_get.side_effect = mock_responses\n\n                result = adapter.flush_cache()\n\n                assert result == {\"cache_flushed\": True}\n                assert mock_get.call_count == 2\n                mock_sleep.assert_called_once()\n\n    def test_flush_cache_non_master(self, mock_launch_server_process):\n        \"\"\"Test flush_cache for non-master node.\"\"\"\n        kwargs = {\"host\": \"localhost\", \"port\": 8000, \"node_rank\": 1, \"model_path\": \"/tmp/test_model\"}\n        adapter = HttpServerAdapter(**kwargs)\n        result = adapter.flush_cache()\n\n        assert result == {}\n\n    def test_memory_management_methods(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test memory release and resume methods.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n\n            # Test release_memory_occupation\n            result = adapter.release_memory_occupation([\"weights\", \"kv_cache\"])\n            assert result == {\"status\": \"success\"}\n            mock_request.assert_called_with(\"release_memory_occupation\", {\"tags\": [\"weights\", \"kv_cache\"]})\n\n            # Test resume_memory_occupation\n            result = adapter.resume_memory_occupation([\"weights\"])\n            assert result == {\"status\": \"success\"}\n            mock_request.assert_called_with(\"resume_memory_occupation\", {\"tags\": [\"weights\"]})\n\n    def test_generation_control_methods(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test generation control methods.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n\n    def test_shutdown(self, mock_launch_server_process, mock_kill_process_tree, router_adapter_kwargs):\n        \"\"\"Test shutdown method.\"\"\"\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            mock_response = Mock()\n            mock_response.status_code = 200\n            mock_post.return_value = mock_response\n\n            adapter = HttpServerAdapter(**router_adapter_kwargs)\n\n            adapter.shutdown()\n\n            # Should unregister from router\n            assert mock_post.call_count == 2  # Once for registration, once for unregistration\n            # Should kill process\n            mock_kill_process_tree.assert_called_once_with(mock_launch_server_process.return_value.pid)\n\n    def test_shutdown_with_errors(self, mock_launch_server_process, mock_kill_process_tree, router_adapter_kwargs):\n        \"\"\"Test shutdown method with errors.\"\"\"\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            # Mock registration success but unregistration failure\n            mock_post.side_effect = [\n                Mock(status_code=200),  # Registration success\n                requests.RequestException(\"Unregistration failed\"),  # Unregistration failure\n            ]\n\n            # Mock process kill failure\n            mock_kill_process_tree.side_effect = Exception(\"Kill failed\")\n\n            adapter = HttpServerAdapter(**router_adapter_kwargs)\n\n            # Should not raise exceptions\n            adapter.shutdown()\n\n            assert mock_post.call_count == 2\n            mock_kill_process_tree.assert_called_once_with(mock_launch_server_process.return_value.pid)\n\n    # Edge cases for HttpServerEngineAdapter\n    def test_empty_and_none_parameters(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test handling of empty and None parameters.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=None,\n                load_format=None,\n                flush_cache=None,\n            )\n\n            # Test generate with all None parameters\n            result = adapter.generate()\n            assert result == {\"status\": \"success\"}\n\n            # Test with empty lists\n            result = adapter.update_weights_from_tensor(req)\n            assert result == {\"status\": \"success\"}\n\n            # Test with empty tags\n            result = adapter.release_memory_occupation(req)\n            assert result == {\"status\": \"success\"}\n\n    def test_large_payload_handling(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test handling of large payloads.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n\n            # Test with large tensor list\n            large_tensor_list = [MultiprocessingSerializer.serialize(f\"tensor_{i}\") for i in range(1000)]\n\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=large_tensor_list,\n                load_format=\"safetensors\",\n                flush_cache=True,\n            )\n            result = adapter.update_weights_from_tensor(req)\n            assert result == {\"status\": \"success\"}\n\n            # Test with large prompt\n            large_prompt = \"A\" * 10000\n            result = adapter.generate(prompt=large_prompt)\n            assert result == {\"status\": \"success\"}\n\n    def test_timeout_edge_cases(self, mock_launch_server_process):\n        \"\"\"Test various timeout scenarios.\"\"\"\n        # Test with very small timeout\n        kwargs = {\"host\": \"localhost\", \"port\": 8000, \"node_rank\": 0, \"model_path\": \"/tmp/test_model\", \"timeout\": 0.001}\n        adapter = HttpServerAdapter(**kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            mock_post.side_effect = requests.exceptions.Timeout()\n\n            with pytest.raises(RuntimeError, match=\"Failed to complete request\"):\n                adapter._make_request(\"test_endpoint\")\n\n    def test_extreme_configuration_values(self, mock_launch_server_process):\n        \"\"\"Test extreme configuration values.\"\"\"\n        # Test with extreme values\n        kwargs = {\n            \"host\": \"localhost\",\n            \"port\": 8000,\n            \"node_rank\": 0,\n            \"model_path\": \"/tmp/test_model\",\n            \"timeout\": 0.001,  # Very small\n            \"max_attempts\": 100,  # Very large\n            \"retry_delay\": 0.001,  # Very small\n        }\n        adapter = HttpServerAdapter(**kwargs)\n\n        assert adapter.timeout == 0.001\n        assert adapter.max_attempts == 100\n        assert adapter.retry_delay == 0.001\n\n\nclass TestAsyncHttpServerEngineAdapter:\n    \"\"\"Test cases for AsyncHttpServerEngineAdapter class.\"\"\"\n\n    def test_init(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test async adapter initialization.\"\"\"\n        adapter = AsyncHttpServerAdapter(max_connections=50, **basic_adapter_kwargs)\n\n        assert adapter.max_connections == 50\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_success(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test successful async HTTP request.\"\"\"\n\n        # Instantiate adapter\n        adapter = AsyncHttpServerAdapter(**basic_adapter_kwargs)\n\n        mock_response = AsyncMock()\n        mock_response.status = 200\n        mock_response.json = AsyncMock(return_value={\"status\": \"success\"})\n        mock_response.raise_for_status = Mock()\n\n        mock_post_context_manager = AsyncMock()\n        mock_post_context_manager.__aenter__.return_value = mock_response\n\n        mock_session = AsyncMock(spec=aiohttp.ClientSession)\n        mock_session.closed = False\n        mock_session.post.return_value = mock_post_context_manager\n\n        mock_session_cm = AsyncMock()\n        mock_session_cm.__aenter__.return_value = mock_session\n\n        with patch.object(adapter, \"_get_session\", return_value=mock_session_cm):\n            result = await adapter._make_async_request(\"test_endpoint\", {\"param\": \"value\"})\n\n            # Assert result is correct\n            assert result == {\"status\": \"success\"}\n\n            # Verify post was called\n            mock_session.post.assert_called_once_with(\n                \"http://localhost:8000/test_endpoint\", json={\"param\": \"value\"}, timeout=adapter.timeout\n            )\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_get_method(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test async GET request using aiohttp and proper context mocking.\"\"\"\n\n        # Instantiate the async adapter\n        adapter = AsyncHttpServerAdapter(**basic_adapter_kwargs)\n\n        mock_response = AsyncMock()\n        mock_response.status = 200\n        mock_response.json = AsyncMock(return_value={\"data\": \"test\"})\n        mock_response.raise_for_status = Mock()\n\n        mock_get_context_manager = AsyncMock()\n        mock_get_context_manager.__aenter__.return_value = mock_response\n\n        mock_session = AsyncMock(spec=aiohttp.ClientSession)\n        mock_session.closed = False\n        mock_session.get.return_value = mock_get_context_manager\n\n        mock_session_cm = AsyncMock()\n        mock_session_cm.__aenter__.return_value = mock_session\n\n        with patch.object(adapter, \"_get_session\", return_value=mock_session_cm):\n            result = await adapter._make_async_request(\"test_endpoint\", method=\"GET\")\n\n            # Validate\n            assert result == {\"data\": \"test\"}\n            mock_session.get.assert_called_once_with(\"http://localhost:8000/test_endpoint\", timeout=adapter.timeout)\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_non_master(self, mock_launch_server_process):\n        \"\"\"Test async request from non-master node.\"\"\"\n        kwargs = {\"host\": \"localhost\", \"port\": 8000, \"node_rank\": 1, \"model_path\": \"/tmp/test_model\"}\n        adapter = AsyncHttpServerAdapter(**kwargs)\n        result = await adapter._make_async_request(\"test_endpoint\")\n\n        assert result == {}\n\n    @pytest.mark.asyncio\n    async def test_async_generate(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test async generate method.\"\"\"\n        adapter = AsyncHttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_async_request\", new_callable=AsyncMock) as mock_request:\n            mock_request.return_value = {\"text\": \"Generated text\"}\n\n            result = await adapter.generate(\n                prompt=\"Hello world\",\n                sampling_params={\"temperature\": 0.7},\n                return_logprob=True,\n            )\n\n            assert result == {\"text\": \"Generated text\"}\n            mock_request.assert_called_once()\n\n    @pytest.mark.asyncio\n    async def test_async_memory_management(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test async memory management methods.\"\"\"\n        adapter = AsyncHttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_async_request\", new_callable=AsyncMock) as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n\n            # Test release_memory_occupation\n            result = await adapter.release_memory_occupation([\"weights\"])\n            assert result == {\"status\": \"success\"}\n            mock_request.assert_called_with(\"release_memory_occupation\", {\"tags\": [\"weights\"]})\n\n            # Test resume_memory_occupation\n            result = await adapter.resume_memory_occupation([\"weights\"])\n            assert result == {\"status\": \"success\"}\n            mock_request.assert_called_with(\"resume_memory_occupation\", {\"tags\": [\"weights\"]})\n            assert (\n                mock_request.call_count == 2\n            )  # resume memory occupation will also call release memory occupation once\n\n\nclass TestErrorRecovery:\n    \"\"\"Test error recovery mechanisms.\"\"\"\n\n    def test_flush_cache_recovery(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test flush cache recovery from failures.\"\"\"\n        adapter = HttpServerAdapter(max_attempts=2, **basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.get\") as mock_get:\n            # Simulate multiple failures then success\n            mock_get.side_effect = [\n                requests.exceptions.ConnectionError(),\n                requests.exceptions.Timeout(),\n                Mock(status_code=503),  # Service unavailable\n                Mock(status_code=200, json=lambda: {\"cache_flushed\": True}),\n            ]\n\n            with patch(\"time.sleep\"):\n                result = adapter.flush_cache()\n                assert result == {\"cache_flushed\": True}\n\n    def test_flush_cache_max_attempts(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test flush cache max retries exceeded.\"\"\"\n        adapter = HttpServerAdapter(max_attempts=1, **basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.get\") as mock_get:\n            # All attempts fail\n            mock_get.side_effect = requests.exceptions.ConnectionError()\n\n            with patch(\"time.sleep\"):\n                result = adapter.flush_cache()\n                assert result == {}  # Should return empty dict on failure\n\n    def test_network_partition_recovery(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test recovery from network partition scenarios.\"\"\"\n        adapter = HttpServerAdapter(max_attempts=3, **basic_adapter_kwargs)\n\n        with patch(\"verl.workers.rollout.sglang_rollout.http_server_engine.requests.post\") as mock_post:\n            # Simulate network partition then recovery\n            mock_post.side_effect = [\n                requests.exceptions.ConnectionError(\"Network unreachable\"),\n                requests.exceptions.ConnectionError(\"Network unreachable\"),\n                Mock(status_code=200, json=lambda: {\"recovered\": True}),\n            ]\n\n            with patch(\"time.sleep\"):\n                result = adapter._make_request(\"test_endpoint\")\n                assert result == {\"recovered\": True}\n\n\nclass TestResourceManagement:\n    \"\"\"Test resource management and cleanup.\"\"\"\n\n    def test_resource_cleanup_on_exception(\n        self, mock_launch_server_process, mock_kill_process_tree, basic_adapter_kwargs\n    ):\n        \"\"\"Test resource cleanup when exceptions occur.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        # Simulate exception during operation\n        with patch.object(adapter, \"_make_request\", side_effect=Exception(\"Test error\")):\n            try:\n                adapter.generate(prompt=\"test\")\n            except Exception:\n                pass\n\n        # Cleanup should still work\n        adapter.shutdown()\n        mock_kill_process_tree.assert_called_once_with(mock_launch_server_process.return_value.pid)\n\n    def test_multiple_shutdown_calls(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test multiple shutdown calls are safe.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        # Multiple shutdown calls should be safe\n        adapter.shutdown()\n        adapter.shutdown()\n        adapter.shutdown()\n\n\nclass TestDataTypeHandling:\n    \"\"\"Test handling of various data types.\"\"\"\n\n    def test_complex_data_structures(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test handling of complex data structures.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {\"status\": \"success\"}\n\n            # Test with complex sampling params\n            complex_sampling_params = {\n                \"temperature\": 0.7,\n                \"top_p\": 0.9,\n                \"top_k\": 50,\n                \"repetition_penalty\": 1.1,\n                \"stop_sequences\": [\"</s>\", \"\\n\\n\"],\n                \"max_tokens\": 100,\n                \"logit_bias\": {\"token_123\": 0.5, \"token_456\": -0.5},\n                \"nested_config\": {\n                    \"beam_search\": True,\n                    \"num_beams\": 4,\n                    \"early_stopping\": True,\n                },\n            }\n\n            result = adapter.generate(\n                prompt=\"Test prompt\",\n                sampling_params=complex_sampling_params,\n            )\n\n            assert result == {\"status\": \"success\"}\n            # Verify the complex structure was passed through\n            call_args = mock_request.call_args[0][1]\n            assert call_args[\"sampling_params\"] == complex_sampling_params\n\n\nclass TestIntegration:\n    \"\"\"Integration tests for both adapters.\"\"\"\n\n    def test_error_scenarios(self, mock_launch_server_process, basic_adapter_kwargs):\n        \"\"\"Test various error scenarios.\"\"\"\n        adapter = HttpServerAdapter(**basic_adapter_kwargs)\n\n        # Test with None payload\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {}\n            result = adapter.generate()\n            assert result == {}\n\n        # Test with empty parameters\n        with patch.object(adapter, \"_make_request\") as mock_request:\n            mock_request.return_value = {}\n            req = UpdateWeightsFromTensorReqInput(\n                serialized_named_tensors=None,\n                load_format=None,\n                flush_cache=None,\n            )\n            result = adapter.update_weights_from_tensor(req)\n            assert result == {}\n"
  },
  {
    "path": "tests/workers/rollout/rollout_trtllm/__init__.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "tests/workers/rollout/rollout_trtllm/test_adapter.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport os\nimport subprocess\nfrom unittest.mock import AsyncMock, Mock, patch\n\nimport aiohttp\nimport pytest\nimport ray\n\nfrom verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica\nfrom verl.workers.rollout.trtllm_rollout.trtllm_rollout import AsyncTRTLLMHttpAdapter\n\n\nclass TestAsyncTRTLLMHttpAdapter:\n    def _build_async_session(\n        self,\n        *,\n        adapter: AsyncTRTLLMHttpAdapter,\n        method: str,\n        response: AsyncMock | None = None,\n    ) -> tuple[AsyncMock, AsyncMock]:\n        mock_session = AsyncMock(spec=aiohttp.ClientSession)\n        mock_session.closed = False\n\n        if response is not None:\n            mock_context_manager = AsyncMock()\n            mock_context_manager.__aenter__.return_value = response\n            getattr(mock_session, method).return_value = mock_context_manager\n\n        mock_session_cm = AsyncMock()\n        mock_session_cm.__aenter__.return_value = mock_session\n        return mock_session_cm, mock_session\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_get_method(self):\n        \"\"\"Test HTTP GET request.\"\"\"\n        adapter = AsyncTRTLLMHttpAdapter(host=\"localhost\", port=8000)\n\n        get_response = AsyncMock()\n        get_response.status = 200\n        get_response.headers = {\"Content-Type\": \"application/json\"}\n        get_response.raise_for_status = Mock()\n        get_response.json = AsyncMock(return_value={\"data\": \"test\"})\n\n        get_session_cm, get_session = self._build_async_session(\n            adapter=adapter,\n            method=\"get\",\n            response=get_response,\n        )\n        with patch.object(adapter, \"_get_session\", return_value=get_session_cm):\n            get_result = await adapter._make_async_request(\"test_endpoint\", method=\"GET\")\n\n        assert get_result == {\"data\": \"test\"}\n        get_session.get.assert_called_once_with(\"http://localhost:8000/test_endpoint\", timeout=adapter.timeout)\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_post_method(self):\n        \"\"\"Test HTTP POST request.\"\"\"\n        adapter = AsyncTRTLLMHttpAdapter(host=\"localhost\", port=8000)\n\n        post_response = AsyncMock()\n        post_response.status = 200\n        post_response.headers = {\"Content-Type\": \"application/json\"}\n        post_response.raise_for_status = Mock()\n        post_response.json = AsyncMock(return_value={\"status\": \"ok\"})\n\n        post_session_cm, post_session = self._build_async_session(\n            adapter=adapter,\n            method=\"post\",\n            response=post_response,\n        )\n        with patch.object(adapter, \"_get_session\", return_value=post_session_cm):\n            post_result = await adapter._make_async_request(\"test_endpoint\", {\"param\": \"value\"})\n\n        assert post_result == {\"status\": \"ok\"}\n        post_session.post.assert_called_once_with(\n            \"http://localhost:8000/test_endpoint\", json={\"param\": \"value\"}, timeout=adapter.timeout\n        )\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_http_error(self):\n        \"\"\"Test HTTP error handling.\"\"\"\n        adapter = AsyncTRTLLMHttpAdapter(host=\"localhost\", port=8000)\n\n        mock_response = AsyncMock()\n        mock_response.status = 500\n        mock_response.headers = {\"Content-Type\": \"application/json\"}\n        mock_response.raise_for_status = Mock(\n            side_effect=aiohttp.ClientResponseError(\n                request_info=Mock(real_url=\"http://localhost:8000/test_endpoint\"),\n                history=(),\n                status=500,\n                message=\"server error\",\n            )\n        )\n\n        mock_session_cm, _mock_session = self._build_async_session(\n            adapter=adapter,\n            method=\"post\",\n            response=mock_response,\n        )\n        with patch.object(adapter, \"_get_session\", return_value=mock_session_cm):\n            with pytest.raises(aiohttp.ClientResponseError):\n                await adapter._make_async_request(\"test_endpoint\", {\"param\": \"value\"})\n\n    @pytest.mark.asyncio\n    async def test_make_async_request_max_attempts_exceeded(self):\n        \"\"\"Test max retries exceeded.\"\"\"\n        adapter = AsyncTRTLLMHttpAdapter(host=\"localhost\", port=8000, max_attempts=1)\n\n        mock_session_cm, mock_session = self._build_async_session(\n            adapter=adapter,\n            method=\"post\",\n            response=None,\n        )\n        mock_session.post.side_effect = asyncio.TimeoutError()\n        with patch.object(adapter, \"_get_session\", return_value=mock_session_cm):\n            with pytest.raises(RuntimeError, match=\"Failed to complete async request\"):\n                await adapter._make_async_request(\"test_endpoint\", {\"param\": \"value\"})\n\n\nclass TestTRTLLMServerAdapter:\n    def test_init_without_device_mesh(self):\n        \"\"\"Test ServerAdapter init path without device mesh.\"\"\"\n        from hydra import compose, initialize_config_dir\n\n        prev_rank = os.environ.get(\"RANK\")\n        os.environ[\"RANK\"] = \"0\"\n\n        try:\n            os.environ.setdefault(\"TLLM_RAY_FORCE_LOCAL_CLUSTER\", \"1\")\n            ray.init(address=\"local\", ignore_reinit_error=True, include_dashboard=False)\n\n            config_dir = os.path.abspath(\"verl/verl/trainer/config\")\n            if not os.path.exists(config_dir):\n                config_dir = os.path.abspath(\"verl/trainer/config\")\n\n            with initialize_config_dir(config_dir=config_dir, version_base=None):\n                config = compose(config_name=\"ppo_trainer\")\n\n            config.trainer.n_gpus_per_node = 2\n            config.trainer.nnodes = 1\n            model_root = os.path.expanduser(os.getenv(\"TRTLLM_TEST_MODEL_PATH_ROOT\", \"~/models\"))\n            config.actor_rollout_ref.model.path = os.path.join(model_root, \"Qwen/Qwen2.5-1.5B-Instruct\")\n            config.actor_rollout_ref.rollout.name = \"trtllm\"\n            config.actor_rollout_ref.rollout.mode = \"async\"\n            config.actor_rollout_ref.rollout.tensor_model_parallel_size = 2\n\n            rollout_config = config.actor_rollout_ref.rollout\n            model_config = config.actor_rollout_ref.model\n\n            replica = TRTLLMReplica(\n                replica_rank=0,\n                config=rollout_config,\n                model_config=model_config,\n                gpus_per_node=2,\n            )\n\n            asyncio.run(replica.init_standalone())\n\n            assert len(replica.workers) == 2\n\n            worker0 = replica.workers[0]\n            worker1 = replica.workers[1]\n            replica_rank = ray.get(worker0.get_replica_rank.remote())\n            is_leader_rank_0 = ray.get(worker0.is_leader_rank.remote())\n            is_leader_rank_1 = ray.get(worker1.is_leader_rank.remote())\n\n            assert replica_rank == 0\n            assert is_leader_rank_0 is True\n            assert is_leader_rank_1 is False\n        finally:\n            if prev_rank is None:\n                os.environ.pop(\"RANK\", None)\n            else:\n                os.environ[\"RANK\"] = prev_rank\n            ray.shutdown()\n            subprocess.run([\"ray\", \"stop\"], capture_output=True)\n"
  },
  {
    "path": "tests/workers/rollout/rollout_trtllm/test_async_server.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport subprocess\nimport time\nfrom unittest.mock import MagicMock, patch\n\nimport ray\nimport torch\nfrom ray.util import placement_group_table\nfrom ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy\n\nfrom verl.single_controller.ray import RayResourcePool, SubRayResourcePool\nfrom verl.workers.rollout.replica import RolloutMode\nfrom verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMHttpServer, TRTLLMReplica\n\n\nclass TestTRTLLMReplica:\n    def test_placement_group_with_sub_ray_resource_pool(self):\n        \"\"\"\n        Scenario: SubRayResourcePool, 1 node, 8 GPUs, TP=4, replica_rank=1\n        SubRayResourcePool pre-assigns start_bundle_index=4 for replica 1.\n        Expected: Replica 1 gets bundles [4, 5, 6, 7]\n        \"\"\"\n        with patch(\"verl.workers.rollout.trtllm_rollout.trtllm_async_server.ray\"):\n            mock_config = MagicMock()\n            mock_config.tensor_model_parallel_size = 4\n            mock_config.data_parallel_size = 1\n            mock_config.pipeline_model_parallel_size = 1\n\n            replica = TRTLLMReplica(\n                replica_rank=1,\n                config=mock_config,\n                model_config=MagicMock(),\n                gpus_per_node=8,\n            )\n\n            mock_pg = MagicMock()\n            mock_pg.bundle_count = 8\n\n            resource_pool = SubRayResourcePool(\n                placement_groups=[mock_pg],\n                start_bundle_index=4,\n                subgroup_world_size=4,\n            )\n\n            replica.resource_pool = resource_pool\n            replica.world_size = 4  # TP=4\n\n            pgs, bundle_indices = replica.get_pgs_and_bundle_indices()\n\n            assert len(pgs) == 1\n            assert pgs[0] == mock_pg\n            assert len(bundle_indices) == 1\n            assert bundle_indices[0] == [4, 5, 6, 7]\n\n    def test_placement_group_with_ray_resource_pool(self):\n        \"\"\"\n        Scenario: RayResourcePool, 1 node, 8 GPUs, TP=2, replica_rank=1\n        RayResourcePool calculates: local_bundle_index = world_size * replica_rank = 2 * 1 = 2\n        Expected: Replica 1 gets bundles [2, 3]\n        \"\"\"\n        with patch(\"verl.workers.rollout.trtllm_rollout.trtllm_async_server.ray\"):\n            mock_config = MagicMock()\n            mock_config.tensor_model_parallel_size = 2\n            mock_config.data_parallel_size = 1\n            mock_config.pipeline_model_parallel_size = 1\n\n            replica = TRTLLMReplica(\n                replica_rank=1,\n                config=mock_config,\n                model_config=MagicMock(),\n                gpus_per_node=8,\n            )\n\n            mock_pg = MagicMock()\n            mock_pg.bundle_count = 8\n\n            resource_pool = RayResourcePool(\n                process_on_nodes=[8],\n                use_gpu=True,\n                max_colocate_count=1,\n                name_prefix=\"test_rollout\",\n            )\n            resource_pool.pgs = [mock_pg]\n\n            replica.resource_pool = resource_pool\n            replica.world_size = 2  # TP=2\n\n            pgs, bundle_indices = replica.get_pgs_and_bundle_indices()\n\n            assert len(pgs) == 1\n            assert pgs[0] == mock_pg\n            assert len(bundle_indices) == 1\n            assert bundle_indices[0] == [2, 3]\n\n\nclass TestTRTLLMHttpServer:\n    @staticmethod\n    def _build_rollout_config(*, response_length: int | None = None, free_cache_engine: bool = False):\n        from hydra import compose, initialize_config_dir\n\n        config_dir = os.path.abspath(\"verl/verl/trainer/config\")\n        if not os.path.exists(config_dir):\n            config_dir = os.path.abspath(\"verl/trainer/config\")\n\n        with initialize_config_dir(config_dir=config_dir, version_base=None):\n            config = compose(config_name=\"ppo_trainer\")\n\n        config.trainer.n_gpus_per_node = 1\n        config.trainer.nnodes = 1\n        model_root = os.path.expanduser(os.getenv(\"TRTLLM_TEST_MODEL_PATH_ROOT\", \"~/models\"))\n        config.actor_rollout_ref.model.path = os.path.join(model_root, \"Qwen/Qwen2.5-0.5B-Instruct\")\n        config.actor_rollout_ref.rollout.name = \"trtllm\"\n        config.actor_rollout_ref.rollout.mode = \"async\"\n        config.actor_rollout_ref.rollout.tensor_model_parallel_size = 1\n        if response_length is not None:\n            config.actor_rollout_ref.rollout.response_length = response_length\n        if free_cache_engine:\n            config.actor_rollout_ref.rollout.free_cache_engine = True\n\n        return config.actor_rollout_ref.rollout, config.actor_rollout_ref.model\n\n    @staticmethod\n    def _create_server(rollout_config, model_config, *, name: str):\n        resource_pool = RayResourcePool(\n            process_on_nodes=[1],\n            use_gpu=True,\n            max_colocate_count=1,\n            name_prefix=\"test_rollout\",\n        )\n        pgs = resource_pool.get_placement_groups()\n        bundle_indices = [[0]]\n\n        pg_data = placement_group_table(pgs[0])\n        node_id = pg_data[\"bundles_to_node_id\"][bundle_indices[0][0]]\n\n        return TRTLLMHttpServer.options(\n            scheduling_strategy=NodeAffinitySchedulingStrategy(\n                node_id=node_id,\n                soft=False,\n            ),\n            runtime_env={\"env_vars\": {\"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES\": \"1\"}},\n            name=name,\n        ).remote(\n            config=rollout_config,\n            model_config=model_config,\n            is_reward_model=False,\n            rollout_mode=RolloutMode.COLOCATED,\n            workers=[],\n            replica_rank=0,\n            max_colocate_count=1,\n            pgs=pgs,\n            bundle_indices=bundle_indices,\n        )\n\n    def test_async_generate(self):\n        \"\"\"Test TRT-LLM generate method with real model.\"\"\"\n        try:\n            os.environ.setdefault(\"TLLM_RAY_FORCE_LOCAL_CLUSTER\", \"1\")\n            ray.init(address=\"local\", ignore_reinit_error=True, include_dashboard=False)\n\n            rollout_config, model_config = self._build_rollout_config(response_length=50)\n\n            server = self._create_server(\n                rollout_config,\n                model_config,\n                name=\"trtllm_server_test_generate\",\n            )\n\n            ray.get(server.launch_server.remote())\n\n            # Test generate with a simple prompt\n            prompt_ids = [1, 2, 3, 4, 5]  # Simple test prompt\n            sampling_params = {\n                \"temperature\": 1.0,\n                \"top_k\": 0,\n                \"logprobs\": 1,\n            }\n            request_id = \"test_request_1\"\n\n            result = ray.get(server.generate.remote(prompt_ids, sampling_params, request_id))\n\n            print(f\"Result: {result}\")\n            # Verify the result structure\n            assert hasattr(result, \"token_ids\"), \"Result should have token_ids attribute\"\n            assert hasattr(result, \"log_probs\"), \"Result should have log_probs attribute\"\n            assert isinstance(result.token_ids, list), \"token_ids should be a list\"\n            assert len(result.token_ids) > 0, \"Generated tokens should not be empty\"\n\n            # Verify logprobs are returned when requested\n            assert result.log_probs is not None, \"log_probs should not be None when requested\"\n            assert len(result.log_probs) == len(result.token_ids), \"log_probs length should match token_ids\"\n\n            print(f\"Generated {len(result.token_ids)} tokens\")\n            print(f\"Token IDs: {result.token_ids[:10]}...\")  # Print first 10 tokens\n            print(f\"Log probs: {result.log_probs[:10]}...\")  # Print first 10 log probs\n\n        finally:\n            ray.shutdown()\n            subprocess.run([\"ray\", \"stop\"], capture_output=True)\n\n    def test_async_memory_management(self):\n        \"\"\"Test TRT-LLM async memory management (sleep) reduces memory usage.\"\"\"\n        try:\n            os.environ.setdefault(\"TLLM_RAY_FORCE_LOCAL_CLUSTER\", \"1\")\n            ray.init(address=\"local\", ignore_reinit_error=True, include_dashboard=False)\n\n            rollout_config, model_config = self._build_rollout_config(free_cache_engine=True)\n\n            server = self._create_server(\n                rollout_config,\n                model_config,\n                name=\"trtllm_server_test_0\",\n            )\n\n            ray.get(server.launch_server.remote())\n            device_ids = ray.get(server.report_device_ids.remote())\n            print(f\"TRTLLM device UUIDs: {device_ids}\")\n\n            def _uuid_to_device_index(device_uuid: str) -> int | None:\n                for idx in range(torch.cuda.device_count()):\n                    props = torch.cuda.get_device_properties(idx)\n                    uuid = getattr(props, \"uuid\", None)\n                    if uuid is None:\n                        # fall back to rank 0\n                        return 0\n                    if isinstance(uuid, bytes):\n                        uuid_str = uuid.decode(\"utf-8\", errors=\"ignore\")\n                    else:\n                        uuid_str = str(uuid)\n                    if uuid_str == device_uuid or uuid_str in device_uuid:\n                        print(f\"Mapped device UUID {device_uuid} to torch device index {idx}\")\n                        return idx\n                return 0\n\n            def get_gpu_memory_mb_for_device(device_uuid: str) -> float:\n                device_index = _uuid_to_device_index(device_uuid)\n                prev_device = torch.cuda.current_device()\n                torch.cuda.set_device(device_index)\n                mem_free, mem_total = torch.cuda.mem_get_info()\n                torch.cuda.set_device(prev_device)\n                return (mem_total - mem_free) / (1024**2)\n\n            baseline_memory_mb = get_gpu_memory_mb_for_device(device_ids[0])\n            print(f\"   Baseline memory: {baseline_memory_mb:.2f} MB\")\n\n            ray.get(server.sleep.remote())\n            time.sleep(2)\n\n            sleep_memory_mb = get_gpu_memory_mb_for_device(device_ids[0])\n            memory_freed_mb = baseline_memory_mb - sleep_memory_mb\n            print(f\"   Memory after sleep: {sleep_memory_mb:.2f} MB\")\n            print(f\"   Memory freed: {memory_freed_mb:.2f} MB\")\n\n            assert memory_freed_mb >= baseline_memory_mb * 0.6, (\n                f\"Expected sleep() to free >=60% of baseline memory. \"\n                f\"Baseline: {baseline_memory_mb:.2f} MB, freed: {memory_freed_mb:.2f} MB.\"\n            )\n\n        finally:\n            ray.shutdown()\n            subprocess.run([\"ray\", \"stop\"], capture_output=True)\n"
  },
  {
    "path": "tests/workers/rollout/rollout_trtllm/test_trtllm_rollout_utils.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport uuid\n\nimport numpy as np\nimport pytest\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\nfrom PIL import Image\nfrom transformers import AutoTokenizer\n\nUNIMODAL_MODEL_PATH = \"Qwen/Qwen2.5-Math-7B\"\nMULTIMODAL_MODEL_PATH = \"Qwen/Qwen2.5-VL-7B-Instruct\"\n\nMAX_MODEL_LEN = 4096\nRESPONSE_LENGTH = 256\nMAX_NUM_SEQS = 16\nGPU_MEMORY_UTILIZATION = 0.8\nTENSOR_PARALLEL_SIZE = 1\n\n\ndef create_test_image(width: int = 224, height: int = 224) -> Image.Image:\n    img_array = np.zeros((height, width, 3), dtype=np.uint8)\n    for i in range(height):\n        for j in range(width):\n            img_array[i, j] = [\n                int(255 * i / height),\n                int(255 * j / width),\n                int(255 * (i + j) / (height + width)),\n            ]\n    return Image.fromarray(img_array)\n\n\ndef create_rollout_config_dict():\n    config_dict = {\n        \"_target_\": \"verl.workers.config.RolloutConfig\",\n        \"name\": \"trtllm\",\n        \"mode\": \"async\",\n        \"temperature\": 0.7,\n        \"top_k\": 50,\n        \"top_p\": 0.9,\n        \"do_sample\": True,\n        \"n\": 1,\n        \"prompt_length\": 512,\n        \"response_length\": RESPONSE_LENGTH,\n        \"dtype\": \"bfloat16\",\n        \"gpu_memory_utilization\": GPU_MEMORY_UTILIZATION,\n        \"ignore_eos\": False,\n        \"enforce_eager\": True,\n        \"free_cache_engine\": False,\n        \"data_parallel_size\": 1,\n        \"tensor_model_parallel_size\": TENSOR_PARALLEL_SIZE,\n        \"pipeline_model_parallel_size\": 1,\n        \"max_num_batched_tokens\": 8192,\n        \"max_model_len\": MAX_MODEL_LEN,\n        \"max_num_seqs\": MAX_NUM_SEQS,\n        \"load_format\": \"auto\",\n        \"enable_chunked_prefill\": True,\n        \"enable_prefix_caching\": True,\n    }\n    return OmegaConf.create(config_dict)\n\n\ndef create_model_config_dict(model_path: str):\n    config_dict = {\n        \"_target_\": \"verl.workers.config.HFModelConfig\",\n        \"path\": model_path,\n        \"trust_remote_code\": True,\n        \"load_tokenizer\": True,\n    }\n    return OmegaConf.create(config_dict)\n\n\ndef get_tokenizer(model_path: str):\n    return AutoTokenizer.from_pretrained(model_path, trust_remote_code=True)\n\n\ndef get_processor(model_path: str):\n    from transformers import AutoProcessor\n\n    return AutoProcessor.from_pretrained(model_path, trust_remote_code=True)\n\n\n@pytest.mark.skipif(\n    not torch.cuda.is_available(),\n    reason=\"CUDA not available\",\n)\nclass TestUnimodalTRTLLMRollout:\n    @pytest.fixture(scope=\"class\")\n    def ray_context(self):\n        if ray.is_initialized():\n            ray.shutdown()\n        ray.init(ignore_reinit_error=True)\n        yield\n        ray.shutdown()\n\n    @pytest.fixture(scope=\"class\")\n    def trtllm_replica(self, ray_context):\n        from verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica\n\n        rollout_config = create_rollout_config_dict()\n        model_config = create_model_config_dict(UNIMODAL_MODEL_PATH)\n\n        replica = TRTLLMReplica(\n            replica_rank=0,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=torch.cuda.device_count(),\n            is_reward_model=False,\n        )\n\n        loop = asyncio.new_event_loop()\n        asyncio.set_event_loop(loop)\n        loop.run_until_complete(replica.init_standalone())\n\n        yield replica\n\n        loop.close()\n\n    @pytest.fixture(scope=\"class\")\n    def tokenizer(self):\n        return get_tokenizer(UNIMODAL_MODEL_PATH)\n\n    @pytest.mark.parametrize(\n        \"prompt\",\n        [\n            \"What is 2 + 2?\",\n            \"Solve for x: 3x + 5 = 20\",\n            \"Calculate the derivative of x^2 + 3x + 1\",\n        ],\n    )\n    def test_unimodal_generate(self, trtllm_replica, tokenizer, prompt):\n        replica = trtllm_replica\n\n        messages = [\n            {\"role\": \"system\", \"content\": \"You are a helpful math assistant.\"},\n            {\"role\": \"user\", \"content\": prompt},\n        ]\n\n        text = tokenizer.apply_chat_template(\n            messages,\n            tokenize=False,\n            add_generation_prompt=True,\n        )\n        input_ids = tokenizer.encode(text, return_tensors=\"pt\")[0].tolist()\n\n        sampling_params = {\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 50,\n            \"logprobs\": True,\n        }\n\n        request_id = str(uuid.uuid4())\n        output = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=request_id,\n            )\n        )\n\n        assert output is not None\n        assert hasattr(output, \"token_ids\")\n        assert len(output.token_ids) > 0\n\n        generated_text = tokenizer.decode(output.token_ids, skip_special_tokens=True)\n        print(\"\\n[Unimodal Test]\")\n        print(f\"Prompt: {prompt}\")\n        print(f\"Generated ({len(output.token_ids)} tokens): {generated_text[:300]}...\")\n\n    def test_unimodal_batch_generate(self, trtllm_replica, tokenizer):\n        replica = trtllm_replica\n\n        prompts = [\n            \"What is 1 + 1?\",\n            \"What is 2 * 3?\",\n            \"What is 10 / 2?\",\n        ]\n\n        sampling_params = {\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 50,\n            \"logprobs\": False,\n        }\n\n        results = []\n\n        for i, prompt in enumerate(prompts):\n            messages = [{\"role\": \"user\", \"content\": prompt}]\n            text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n            input_ids = tokenizer.encode(text, return_tensors=\"pt\")[0].tolist()\n\n            output = ray.get(\n                replica.server_handle.generate.remote(\n                    prompt_ids=input_ids,\n                    sampling_params=sampling_params,\n                    request_id=str(uuid.uuid4()),\n                )\n            )\n            results.append(output)\n\n        assert len(results) == len(prompts)\n        for i, (prompt, result) in enumerate(zip(prompts, results, strict=False)):\n            assert result is not None\n            assert len(result.token_ids) > 0\n            generated = tokenizer.decode(result.token_ids, skip_special_tokens=True)\n            print(f\"\\n[Batch {i}] Prompt: {prompt}\")\n            print(f\"Generated: {generated[:100]}...\")\n\n\n@pytest.mark.skipif(\n    not torch.cuda.is_available(),\n    reason=\"CUDA not available\",\n)\nclass TestMultimodalTRTLLMRollout:\n    @pytest.fixture(scope=\"class\")\n    def ray_context(self):\n        if ray.is_initialized():\n            ray.shutdown()\n        ray.init(ignore_reinit_error=True)\n        yield\n        ray.shutdown()\n\n    @pytest.fixture(scope=\"class\")\n    def trtllm_vlm_replica(self, ray_context):\n        from verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica\n\n        rollout_config = create_rollout_config_dict()\n        model_config = create_model_config_dict(MULTIMODAL_MODEL_PATH)\n\n        replica = TRTLLMReplica(\n            replica_rank=0,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=torch.cuda.device_count(),\n            is_reward_model=False,\n        )\n\n        loop = asyncio.new_event_loop()\n        asyncio.set_event_loop(loop)\n        loop.run_until_complete(replica.init_standalone())\n\n        yield replica\n\n        loop.close()\n\n    @pytest.fixture(scope=\"class\")\n    def tokenizer(self):\n        return get_tokenizer(MULTIMODAL_MODEL_PATH)\n\n    @pytest.fixture(scope=\"class\")\n    def processor(self):\n        return get_processor(MULTIMODAL_MODEL_PATH)\n\n    @pytest.mark.parametrize(\n        \"prompt\",\n        [\n            \"Describe this image in detail.\",\n            \"What colors do you see in this image?\",\n            \"What patterns are visible in this image?\",\n        ],\n    )\n    def test_multimodal_generate_with_image(self, trtllm_vlm_replica, processor, tokenizer, prompt):\n        replica = trtllm_vlm_replica\n\n        test_image = create_test_image(224, 224)\n\n        messages = [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"image\"},\n                    {\"type\": \"text\", \"text\": prompt},\n                ],\n            }\n        ]\n\n        text = processor.apply_chat_template(\n            messages,\n            tokenize=False,\n            add_generation_prompt=True,\n        )\n        print(\"text: \", text)\n        input_ids = processor.tokenizer(text, return_tensors=\"pt\", padding=True)[\"input_ids\"][0].tolist()\n\n        print(\n            \"input_ids decoded: \",\n            processor.tokenizer.decode(input_ids, skip_special_tokens=False, add_special_tokens=False),\n        )\n\n        sampling_params = {\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 50,\n            \"logprobs\": False,\n        }\n\n        output = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=str(uuid.uuid4()),\n                image_data=[test_image],\n            )\n        )\n\n        assert output is not None\n        assert hasattr(output, \"token_ids\")\n        assert len(output.token_ids) > 0\n\n        generated_text = tokenizer.decode(output.token_ids, skip_special_tokens=True)\n        print(\"\\n[Multimodal Test]\")\n        print(f\"Prompt: {prompt}\")\n        print(f\"Image size: {test_image.size}\")\n        print(f\"Generated ({len(output.token_ids)} tokens): {generated_text[:300]}...\")\n\n    @pytest.mark.parametrize(\n        \"image_size\",\n        [(224, 224), (384, 384), (512, 512)],\n    )\n    def test_multimodal_different_image_sizes(self, trtllm_vlm_replica, processor, tokenizer, image_size):\n        replica = trtllm_vlm_replica\n\n        width, height = image_size\n        test_image = create_test_image(width, height)\n\n        prompt = \"What is shown in this image?\"\n        messages = [\n            {\n                \"role\": \"user\",\n                \"content\": [\n                    {\"type\": \"image\"},\n                    {\"type\": \"text\", \"text\": prompt},\n                ],\n            }\n        ]\n\n        text = processor.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n        input_ids = processor.tokenizer(text, return_tensors=\"pt\", padding=True)[\"input_ids\"][0].tolist()\n\n        sampling_params = {\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 50,\n            \"logprobs\": False,\n        }\n\n        output = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=str(uuid.uuid4()),\n                image_data=[test_image],\n            )\n        )\n\n        assert output is not None\n        assert len(output.token_ids) > 0\n        print(f\"\\n[Image Size {image_size}] Generated {len(output.token_ids)} tokens\")\n\n    def test_multimodal_text_only_fallback(self, trtllm_vlm_replica, tokenizer):\n        replica = trtllm_vlm_replica\n\n        prompt = \"What is the capital of China?\"\n        messages = [{\"role\": \"user\", \"content\": prompt}]\n\n        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n        input_ids = tokenizer.encode(text, return_tensors=\"pt\")[0].tolist()\n\n        sampling_params = {\n            \"temperature\": 0.7,\n            \"top_p\": 0.9,\n            \"top_k\": 50,\n            \"logprobs\": False,\n        }\n\n        output = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=str(uuid.uuid4()),\n            )\n        )\n\n        assert output is not None\n        assert len(output.token_ids) > 0\n\n        generated_text = tokenizer.decode(output.token_ids, skip_special_tokens=True)\n        print(\"\\n[Text-only on VLM]\")\n        print(f\"Prompt: {prompt}\")\n        print(f\"Generated: {generated_text}\")\n\n\n@pytest.mark.skipif(\n    not torch.cuda.is_available(),\n    reason=\"CUDA not available\",\n)\nclass TestTRTLLMServerLifecycle:\n    @pytest.fixture(scope=\"class\")\n    def ray_context(self):\n        if ray.is_initialized():\n            ray.shutdown()\n        ray.init(ignore_reinit_error=True)\n        yield\n        ray.shutdown()\n\n    @pytest.fixture(scope=\"class\")\n    def trtllm_replica_lifecycle(self, ray_context):\n        from verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica\n\n        rollout_config = create_rollout_config_dict()\n        model_config = create_model_config_dict(UNIMODAL_MODEL_PATH)\n\n        replica = TRTLLMReplica(\n            replica_rank=0,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=torch.cuda.device_count(),\n            is_reward_model=False,\n        )\n\n        loop = asyncio.new_event_loop()\n        asyncio.set_event_loop(loop)\n        loop.run_until_complete(replica.init_standalone())\n\n        yield replica, loop\n\n        loop.close()\n\n    @pytest.fixture(scope=\"class\")\n    def tokenizer(self):\n        return get_tokenizer(UNIMODAL_MODEL_PATH)\n\n    def test_wake_sleep_cycle(self, trtllm_replica_lifecycle, tokenizer):\n        replica, loop = trtllm_replica_lifecycle\n\n        prompt = \"Hello, world!\"\n        messages = [{\"role\": \"user\", \"content\": prompt}]\n        text = tokenizer.apply_chat_template(messages, tokenize=False, add_generation_prompt=True)\n        input_ids = tokenizer.encode(text, return_tensors=\"pt\")[0].tolist()\n\n        sampling_params = {\"temperature\": 0.7, \"top_p\": 0.9, \"top_k\": 50, \"logprobs\": False}\n\n        output1 = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=str(uuid.uuid4()),\n            )\n        )\n        assert output1 is not None\n        assert len(output1.token_ids) > 0\n        print(f\"\\n[Before Sleep] Generated {len(output1.token_ids)} tokens\")\n\n        loop.run_until_complete(replica.sleep())\n        print(\"[Sleep] Server put to sleep\")\n\n        loop.run_until_complete(replica.wake_up())\n        print(\"[Wake Up] Server woken up\")\n\n        output2 = ray.get(\n            replica.server_handle.generate.remote(\n                prompt_ids=input_ids,\n                sampling_params=sampling_params,\n                request_id=str(uuid.uuid4()),\n            )\n        )\n        assert output2 is not None\n        assert len(output2.token_ids) > 0\n        print(f\"[After Wake Up] Generated {len(output2.token_ids)} tokens\")\n"
  },
  {
    "path": "tests/workers/rollout/rollout_vllm/run_fsdp_vllm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport time\n\nimport torch\nimport torch.distributed as dist\nfrom torch.distributed.fsdp import CPUOffload, MixedPrecision\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp.api import ShardedStateDictConfig, ShardingStrategy, StateDictType\nfrom transformers import AutoConfig, AutoModelForCausalLM, AutoTokenizer\nfrom vllm import SamplingParams\n\nfrom verl.third_party.vllm import LLM\nfrom verl.utils.distributed import initialize_global_process_group\n\n\ndef _pre_process_inputs(pad_token_id, prompt_token_ids: torch.Tensor) -> list[int]:\n    \"\"\"Remove left padding tokens before feeding prompts to vLLM.\"\"\"\n    non_pad_index = torch.nonzero(prompt_token_ids != pad_token_id, as_tuple=False)[0][0]\n    return prompt_token_ids[non_pad_index:].tolist()\n\n\ndef main():\n    assert torch.cuda.is_available(), \"CUDA must be present to run FSDP vLLM example\"\n    local_rank, rank, world_size = initialize_global_process_group()\n\n    local_cache_path = \"~/.cache/verl/rlhf\"\n    local_cache_path = os.path.expanduser(local_cache_path)\n    hdfs_path = \"Qwen/Qwen2-7B-Instruct\"\n\n    from verl.utils.fs import copy_to_local\n\n    local_model_path = copy_to_local(src=hdfs_path, cache_dir=local_cache_path)\n    tokenizer = AutoTokenizer.from_pretrained(local_model_path, trust_remote_code=True)\n    actor_model_config = AutoConfig.from_pretrained(local_model_path, trust_remote_code=True)\n    with torch.device(\"cuda\"):\n        actor_model = AutoModelForCausalLM.from_pretrained(local_model_path, trust_remote_code=True)\n        actor_model.to(torch.bfloat16)\n\n    max_prompt_length = 16\n    response_length = 32\n    preencode_prompts = [\n        \"The president of the United States is\",\n        \"The capital of France is\",\n        \"The future of AI is\",\n    ]\n    tokenizer.pad_token = tokenizer.eos_token\n    prompts = tokenizer(preencode_prompts, return_tensors=\"pt\", padding=True)\n    input_ids = prompts[\"input_ids\"]\n    attention_mask = prompts[\"attention_mask\"]\n    from verl.utils.torch_functional import pad_sequence_to_length\n\n    input_ids = pad_sequence_to_length(input_ids, max_prompt_length, tokenizer.pad_token_id, left_pad=True).cuda()\n    attention_mask = pad_sequence_to_length(attention_mask, max_prompt_length, 0, left_pad=True).cuda()\n\n    from transformers import GenerationConfig\n\n    generation_config = GenerationConfig(do_sample=False)\n    actor_model.cuda()\n    output = actor_model.generate(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        max_new_tokens=32,\n        # max_length=max_length,\n        eos_token_id=tokenizer.eos_token_id,\n        pad_token_id=tokenizer.pad_token_id,\n        generation_config=generation_config,\n        # renormalize_logits=True,\n        output_scores=False,  # this is potentially very large\n        return_dict_in_generate=True,\n        use_cache=False,\n    )  # may OOM when use_cache = True\n    seq = output.sequences\n    response = seq[:, max_prompt_length:]\n\n    print(f\"hf response: {tokenizer.batch_decode(response)}\")\n\n    tensor_model_parallel_size = 4\n    from torch.distributed.device_mesh import init_device_mesh\n\n    device_mesh = init_device_mesh(\"cuda\", mesh_shape=(world_size,), mesh_dim_names=[\"fsdp\"])\n\n    mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32)\n    fsdp_model = FSDP(\n        actor_model,\n        use_orig_params=True,\n        auto_wrap_policy=None,\n        device_id=torch.cuda.current_device(),\n        sharding_strategy=ShardingStrategy.FULL_SHARD,\n        mixed_precision=mixed_precision,\n        cpu_offload=CPUOffload(offload_params=False),\n        sync_module_states=False,\n        device_mesh=device_mesh,\n    )\n\n    FSDP.set_state_dict_type(\n        fsdp_model, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig()\n    )\n\n    state_dict = fsdp_model.state_dict()\n\n    sampling_params = SamplingParams(\n        temperature=0, top_p=1, n=1, max_tokens=response_length, logprobs=1, ignore_eos=True, detokenize=False\n    )\n\n    print(actor_model_config)\n    llm = LLM(\n        model=None,\n        tokenizer=tokenizer,\n        model_hf_config=actor_model_config,\n        tensor_parallel_size=tensor_model_parallel_size,\n        enforce_eager=True,\n        dtype=\"bfloat16\",\n        load_format=\"dummy_dtensor\",\n        gpu_memory_utilization=0.8,\n        trust_remote_code=True,\n    )\n\n    # Warmup iterations\n    for _ in range(10):\n        torch.cuda.synchronize()\n        llm.sync_model_weights(actor_weights=state_dict, load_format=\"dtensor\")\n        torch.cuda.synchronize()\n        dist.barrier()\n\n    start_time = time.time()\n    llm.sync_model_weights(actor_weights=state_dict, load_format=\"dtensor\")\n    torch.cuda.synchronize()\n    dist.barrier()\n    end_time = time.time()\n\n    # Calculate elapsed time\n    elapsed_time = end_time - start_time\n    print(f\"Time taken: {elapsed_time:.6f} seconds\")\n\n    input_ids = input_ids.cuda()\n    attention_mask = attention_mask.cuda()\n    idx_list = []\n    batch_size = input_ids.shape[0]\n\n    pad_token_id = tokenizer.pad_token_id if tokenizer.pad_token_id is not None else tokenizer.eos_token_id\n    for i in range(batch_size):\n        idx_list.append(_pre_process_inputs(pad_token_id, input_ids[i]))\n    print(\"start generation\")\n    outputs = llm.generate(prompt_token_ids=idx_list, sampling_params=sampling_params, use_tqdm=False)\n    vllm_output = outputs[0].cuda()\n    if torch.distributed.get_rank() == 0:\n        print(f\"hf response: {tokenizer.batch_decode(response)}\")\n        print(f\"vllm response: {tokenizer.batch_decode(vllm_output)}\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "tests/workers/rollout/rollout_vllm/test_vllm_abort.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTest vLLM abort functionality.\n\nUsage:\n    pytest tests/workers/rollout/rollout_vllm/test_vllm_abort.py -v -s\n    or\n    python tests/workers/rollout/rollout_vllm/test_vllm_abort.py\n\"\"\"\n\nimport asyncio\nimport os\nimport time\nfrom uuid import uuid4\n\n\ndef test_vllm_abort():\n    # ==================== Configuration ====================\n    MODEL_PATH = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")  # /root/models/Qwen/Qwen2.5-1.5B-Instruct\n    GPUS_PER_NODE = 2\n    TP_SIZE = 1\n    ROLLOUT_NAME = \"vllm\"\n    ABORT_DELAY = 0.5  # seconds to wait before aborting\n\n    print(\"=\" * 60)\n    print(\"vLLM Abort Test\")\n    print(\"=\" * 60)\n    print(f\"Model: {MODEL_PATH}\")\n    print(f\"GPUs: {GPUS_PER_NODE}, TP Size: {TP_SIZE}\")\n    print(f\"Abort Delay: {ABORT_DELAY}s\")\n    print(\"=\" * 60)\n\n    # ==================== Initialize Ray ====================\n    print(\"\\n[1] Initializing Ray...\")\n    import ray\n\n    ray.init(\n        runtime_env={\n            \"env_vars\": {\n                \"TOKENIZERS_PARALLELISM\": \"true\",\n                \"NCCL_DEBUG\": \"WARN\",\n                \"VLLM_LOGGING_LEVEL\": \"INFO\",\n                \"VLLM_USE_V1\": \"1\",\n            }\n        },\n        ignore_reinit_error=True,\n    )\n\n    try:\n        # ==================== Create Config ====================\n        print(\"\\n[2] Creating config...\")\n        from hydra import compose, initialize_config_dir\n\n        from verl.utils.tokenizer import normalize_token_ids\n\n        config_dir = os.path.abspath(\"verl/verl/trainer/config\")\n        if not os.path.exists(config_dir):\n            config_dir = os.path.abspath(\"verl/trainer/config\")\n\n        with initialize_config_dir(config_dir=config_dir, version_base=None):\n            config = compose(config_name=\"ppo_trainer\")\n\n        config.trainer.n_gpus_per_node = GPUS_PER_NODE\n        config.trainer.nnodes = 1\n        config.actor_rollout_ref.model.path = MODEL_PATH\n        config.actor_rollout_ref.rollout.name = ROLLOUT_NAME\n        config.actor_rollout_ref.rollout.mode = \"async\"\n        config.actor_rollout_ref.rollout.tensor_model_parallel_size = TP_SIZE\n        config.actor_rollout_ref.rollout.prompt_length = 512\n        config.actor_rollout_ref.rollout.response_length = 512  # Longer for abort test\n\n        # ==================== Create Rollout Server ====================\n        print(\"\\n[3] Creating rollout server (this may take a while)...\")\n        from verl.workers.rollout.replica import get_rollout_replica_class\n\n        rollout_config = config.actor_rollout_ref.rollout\n        model_config = config.actor_rollout_ref.model\n\n        rollout_server_class = get_rollout_replica_class(ROLLOUT_NAME)\n        server = rollout_server_class(\n            replica_rank=0,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=GPUS_PER_NODE,\n        )\n\n        asyncio.run(server.init_standalone())\n        server_handle = server._server_handle\n        print(f\"Server address: {server._server_address}\")\n\n        # ==================== Load Tokenizer ====================\n        print(\"\\n[4] Loading tokenizer...\")\n        from transformers import AutoTokenizer\n\n        tokenizer = AutoTokenizer.from_pretrained(MODEL_PATH, trust_remote_code=True)\n\n        # ==================== Prepare Prompts ====================\n        print(\"\\n[5] Preparing prompts (to ensure generation takes time)...\")\n        NUM_PROMPTS = 8\n        prompts = [\n            \"Write a very long story about a brave knight and dragon.\",\n            \"Explain the history of the Roman Empire in great detail.\",\n            \"Describe quantum computing and its applications thoroughly.\",\n            \"Write an essay about climate change and its global effects.\",\n            \"Who won the Champions League in 2019?\",\n            \"Write a detailed analysis of Shakespeare's Hamlet.\",\n            \"Describe the process of photosynthesis in plants.\",\n            \"Write about the French Revolution and its consequences.\",\n        ]\n\n        all_prompt_ids = []\n        for prompt in prompts[:NUM_PROMPTS]:\n            messages = [{\"role\": \"user\", \"content\": prompt}]\n            prompt_ids = normalize_token_ids(\n                tokenizer.apply_chat_template(messages, add_generation_prompt=True, tokenize=True)\n            )\n            all_prompt_ids.append(prompt_ids)\n        print(f\"Prepared {NUM_PROMPTS} prompts\")\n\n        # ==================== Start Generations and Abort ====================\n        print(\"\\n[6] Starting generations and then aborting...\")\n\n        sampling_params = {\n            \"temperature\": 1.0,\n            \"top_p\": 1.0,\n            \"logprobs\": False,\n        }\n\n        # Start all generations concurrently\n        print(f\"\\n   Starting {NUM_PROMPTS} generations...\")\n        generate_refs = []\n        for i, prompt_ids in enumerate(all_prompt_ids):\n            request_id = f\"abort_test_{i}_{uuid4().hex[:8]}\"\n            ref = server_handle.generate.remote(\n                request_id=request_id,\n                prompt_ids=prompt_ids,\n                sampling_params=sampling_params,\n                image_data=None,\n            )\n            generate_refs.append((i, request_id, ref))\n            print(f\"      Started request {i}: {request_id}\")\n\n        # Wait before aborting\n        print(f\"\\n   Waiting {ABORT_DELAY}s before abort...\")\n        time.sleep(ABORT_DELAY)\n\n        # Call abort\n        print(\"   Calling abort_all_requests...\")\n        abort_start = time.perf_counter()\n        abort_result = ray.get(server_handle.abort_all_requests.remote())\n        abort_time = time.perf_counter() - abort_start\n\n        print(f\"   Abort took: {abort_time * 1000:.2f}ms\")\n        print(f\"   Abort result: {abort_result}\")\n\n        # Wait for all generations to finish\n        print(\"\\n   Waiting for all generations to complete...\")\n        outputs = []\n        for i, request_id, ref in generate_refs:\n            try:\n                output = ray.get(ref, timeout=10.0)\n                outputs.append((i, request_id, output))\n            except ray.exceptions.GetTimeoutError:\n                print(f\"      Request {i} timed out!\")\n                outputs.append((i, request_id, None))\n\n        # ==================== Print Results ====================\n        print(\"\\n\" + \"=\" * 60)\n        print(\"RESULTS\")\n        print(\"=\" * 60)\n\n        aborted_count = 0\n        completed_count = 0\n        timeout_count = 0\n\n        for i, request_id, output in outputs:\n            if output is None:\n                timeout_count += 1\n                print(f\"[{i}] {request_id}: TIMEOUT\")\n            elif output.stop_reason == \"aborted\":\n                aborted_count += 1\n                print(f\"[{i}] {request_id}: ABORTED ({len(output.token_ids)} tokens)\")\n                print(f\"Partial Output: {tokenizer.decode(output.token_ids)}\")\n            else:\n                completed_count += 1\n                print(f\"[{i}] {request_id}: COMPLETED ({output.stop_reason}, {len(output.token_ids)} tokens)\")\n                print(f\"Full Output: {tokenizer.decode(output.token_ids)}\")\n\n        print(f\"\\nSummary: {aborted_count} aborted, {completed_count} completed, {timeout_count} timeout\")\n\n        print(\"\\n\" + \"=\" * 60)\n        print(f\"Abort result: {abort_result}\")\n        print(\"=\" * 60)\n        print(\"Abort test completed!\")\n\n        # Assertions for pytest\n        assert timeout_count == 0, \"No requests should timeout\"\n        assert aborted_count + completed_count == NUM_PROMPTS, \"All requests should finish\"\n        assert \"aborted_count\" in abort_result, \"Abort result should contain aborted_count\"\n        assert abort_time < 1.0, \"Abort should be fast (< 1 second)\"\n\n    finally:\n        print(\"\\nShutting down Ray...\")\n        ray.shutdown()\n\n\nif __name__ == \"__main__\":\n    # Can still run as standalone script\n    test_vllm_abort()\n"
  },
  {
    "path": "tests/workers/rollout/test_hf_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\n\nimport torch\nfrom omegaconf import OmegaConf\nfrom torch.distributed.fsdp import CPUOffload, MixedPrecision\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp.api import ShardedStateDictConfig, ShardingStrategy, StateDictType\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\nfrom verl import DataProto\nfrom verl.utils.distributed import initialize_global_process_group\nfrom verl.utils.fs import copy_to_local\nfrom verl.utils.model import compute_position_id_with_mask\nfrom verl.workers.rollout.hf_rollout import HFRollout\n\nBASE_HF_ROLLOUT_CONFIG = {\n    \"temperature\": 1.0,\n    \"top_k\": -1,\n    \"top_p\": 1,\n    \"prompt_length\": 64,\n    \"response_length\": 64,\n    \"do_sample\": True,\n    \"n\": 1,\n    \"val_kwargs\": {\n        \"top_k\": -1,\n        \"top_p\": 1.0,\n        \"temperature\": 0,\n        \"n\": 1,\n        \"do_sample\": False,\n    },\n}\n\n\ndef prepare_input_dataproto(tokenizer, config, validate):\n    preencode_prompts = [\n        [{\"role\": \"user\", \"content\": \"Who won the Champions League in 2019?\"}],\n        [{\"role\": \"user\", \"content\": \"The founder of Apple is\"}],\n        [{\"role\": \"user\", \"content\": \"What's your name\"}],\n    ]\n    formatted_prompts = [\n        tokenizer.apply_chat_template(conversation, tokenize=False, add_generation_prompt=True)\n        for conversation in preencode_prompts\n    ]\n    prompts = tokenizer(formatted_prompts, return_tensors=\"pt\", padding=\"max_length\", max_length=config.prompt_length)\n    input_dataproto = DataProto.from_dict(\n        {\n            \"input_ids\": prompts[\"input_ids\"],\n            \"attention_mask\": prompts[\"attention_mask\"],\n            \"position_ids\": compute_position_id_with_mask(prompts[\"attention_mask\"]),\n        },\n        meta_info={\n            \"bos_token_id\": tokenizer.bos_token_id,\n            \"eos_token_id\": tokenizer.eos_token_id,\n            \"pad_token_id\": tokenizer.pad_token_id,\n            \"validate\": validate,\n        },\n    )\n    return input_dataproto\n\n\ndef prepare_fsdp_model(model, world_size):\n    from torch.distributed.device_mesh import init_device_mesh\n\n    device_mesh = init_device_mesh(\"cuda\", mesh_shape=(world_size,), mesh_dim_names=[\"fsdp\"])\n\n    mixed_precision = MixedPrecision(param_dtype=torch.bfloat16, reduce_dtype=torch.float32, buffer_dtype=torch.float32)\n\n    fsdp_model = FSDP(\n        model,\n        use_orig_params=True,\n        auto_wrap_policy=None,\n        device_id=torch.cuda.current_device(),\n        sharding_strategy=ShardingStrategy.FULL_SHARD,\n        mixed_precision=mixed_precision,\n        cpu_offload=CPUOffload(offload_params=False),\n        sync_module_states=False,\n        device_mesh=device_mesh,\n    )\n\n    FSDP.set_state_dict_type(\n        fsdp_model, state_dict_type=StateDictType.SHARDED_STATE_DICT, state_dict_config=ShardedStateDictConfig()\n    )\n    return fsdp_model\n\n\ndef test_hf_rollout(n: int = 1, do_sample: bool = True, validate: bool = False):\n    config = OmegaConf.create(BASE_HF_ROLLOUT_CONFIG)\n    config.update({\"n\": n, \"do_sample\": do_sample})\n\n    assert torch.cuda.device_count() >= 2, \"At least 2 GPUs is required to run tp+dp tests.\"\n    local_rank, rank, world_size = initialize_global_process_group()\n\n    # Initialize model and tokenizer\n    local_cache_path = \"~/.cache/verl/rlhf\"\n    local_cache_path = os.path.expanduser(local_cache_path)\n    hdfs_path = \"Qwen/Qwen2-7B-Instruct\"\n    local_model_path = copy_to_local(src=hdfs_path, cache_dir=local_cache_path)\n    tokenizer = AutoTokenizer.from_pretrained(local_model_path, padding_side=\"left\", trust_remote_code=True)\n    tokenizer.pad_token = tokenizer.eos_token\n\n    # Initialize FSDP model\n    actor_model = AutoModelForCausalLM.from_pretrained(local_model_path, trust_remote_code=True)\n    actor_model.to(torch.bfloat16)\n    fsdp_model = prepare_fsdp_model(actor_model, world_size)\n\n    # Initialize HFRollout and start generate\n    hf_rollout = HFRollout(fsdp_model, OmegaConf.create(config))\n    input = prepare_input_dataproto(tokenizer, config, validate).to(torch.cuda.current_device())\n    outputs = hf_rollout.generate_sequences(input)\n\n    # check generated batch size is expected\n    generated_batch_size = outputs.batch.batch_size[0]\n    assert generated_batch_size == input.batch.batch_size[0] * config.n\n\n    for i in range(generated_batch_size):\n        prompt_tokens = outputs.batch[\"prompts\"][i]\n        prompt_mask = prompt_tokens != tokenizer.pad_token_id\n        prompt_tokens = prompt_tokens[prompt_mask]\n        decoded_prompt = tokenizer.decode(prompt_tokens, skip_special_tokens=False)\n\n        response_tokens = outputs.batch[\"responses\"][i]\n        response_mask = response_tokens != tokenizer.pad_token_id\n        response_tokens = response_tokens[response_mask]\n        decoded_response = tokenizer.decode(response_tokens, skip_special_tokens=False)\n\n        attention_mask = outputs.batch[\"attention_mask\"][i]\n        position_ids = outputs.batch[\"position_ids\"][i]\n        prompt_length = outputs.batch[\"prompts\"].size(1)\n        response_length = outputs.batch[\"responses\"].size(1)\n\n        assert attention_mask.size(0) == prompt_length + response_length\n        assert position_ids.size(0) == prompt_length + response_length\n\n        # check response attention mask is expected\n        response_attention = attention_mask[prompt_length:]\n        eos_positions = (outputs.batch[\"responses\"][i] == tokenizer.pad_token_id).nonzero(as_tuple=True)[0]\n        if len(eos_positions) > 0:\n            first_eos_pos = eos_positions[0].item()\n            assert response_attention[: first_eos_pos + 1].all(), \"Response attention mask should be 1 until EOS\"\n            if first_eos_pos + 1 < response_length:\n                assert not response_attention[first_eos_pos + 1 :].any(), (\n                    \"Response attention mask should be 0 after EOS\"\n                )\n        else:\n            assert response_attention.all(), \"Response attention mask should be all 1 if no EOS token\"\n\n        # check response position ids is expected\n        prompt_positions = position_ids[:prompt_length]\n        response_positions = position_ids[prompt_length:]\n        valid_response_length = min(len(response_tokens), response_length)\n        if valid_response_length > 0:\n            assert response_positions[0] == prompt_positions[-1] + 1\n            for j in range(1, valid_response_length):\n                assert response_positions[j] == response_positions[j - 1] + 1\n\n        # print generated text for inspection\n        if torch.distributed.get_rank() == 0:\n            print(f\"prompt: {decoded_prompt}\")\n            print(f\"response: {decoded_response}\")\n            print(\"=\" * 30)\n\n\nif __name__ == \"__main__\":\n    test_hf_rollout(n=2, do_sample=True, validate=False)\n    # test_hf_rollout(n=1, do_sample=False, validate=True)\n    # test_hf_rollout(n=1, do_sample=True, validate=False)\n"
  },
  {
    "path": "tests/workers/rollout/test_sglang_async_rollout_multimodal_delta.py",
    "content": "# Copyright 2025 Amazon.com, Inc. or its affiliates\n# Copyright 2023-2024 SGLang Team\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport os\n\nimport pytest\n\nfrom verl.tools.schemas import ToolResponse\nfrom verl.utils.dataset.vision_utils import process_image\nfrom verl.utils.tokenizer import hf_processor\nfrom verl.workers.rollout.schemas import (\n    AsyncRolloutRequest,\n    AsyncRolloutRequestStateEnum,\n    TokenizationSanityCheckModeEnum,\n)\n\n\ndef _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=False):\n    assert len(image_list) == len(description_list)\n    # Get the smallest dimensions across all images\n    processed_images = []\n    for img_url in image_list:\n        img = process_image(img_url)\n        processed_images.append(img)\n\n    min_width = min(img.size[0] for img in processed_images)\n    min_height = min(img.size[1] for img in processed_images)\n    min_size = (min_width, min_height)\n\n    if resize_image:\n        processed_images_resized = []\n        for img in processed_images:\n            img = img.resize(min_size)\n            processed_images_resized.append(img)\n        processed_images = processed_images_resized\n\n    # Initial message history\n    system_prompt = (\n        \"You will be provided with an image. Describe this image and then generate a new image for the next round\"\n    )\n    messages = [\n        {\n            \"role\": \"system\",\n            \"content\": system_prompt,\n        },\n        {\n            \"role\": \"user\",\n            \"content\": [\n                {\"type\": \"text\", \"text\": \"Here is the first image provided: \"},\n                {\"type\": \"image\", \"image\": [processed_images[0]]},\n            ],\n        },\n    ]\n\n    # Initial multi_modal_data with one image\n    multi_modal_data = {\"image\": [processed_images[0]], \"video\": []}\n    # Minimal required fields for AsyncRolloutRequest\n\n    req = AsyncRolloutRequest(\n        batch_data_id=0,\n        request_id=\"test-req-1\",\n        state=AsyncRolloutRequestStateEnum.PENDING,\n        messages=messages,\n        multi_modal_keys=[\"image\", \"video\"],\n        multi_modal_data=multi_modal_data.copy(),\n        tool_schemas=[],\n        tools_kwargs={},\n        interaction_kwargs={},\n        input_ids=None,\n        prompt_ids=None,\n        response_ids=None,\n        attention_mask=None,\n        prompt_attention_mask=None,\n        response_attention_mask=None,\n        position_ids=None,\n        prompt_position_ids=None,\n        response_position_ids=None,\n        loss_mask=None,\n        prompt_loss_mask=None,\n        response_loss_mask=None,\n        reward_scores={},\n        max_prompt_len=8192,\n        max_response_len=8192,\n        max_model_len=16384,\n        metrics={},\n        use_inference_chat_template=True,\n        tokenization_sanity_check_mode=TokenizationSanityCheckModeEnum.STRICT,\n        generation_prompt_ids=None,\n        base_conv_wo_gen_prompt_end_pos=0,\n        base_conv_with_gen_prompt_end_pos=0,\n        processing_class=processor,\n    )\n\n    prev_generated_len = 0\n    # Add First Assistant Message and first tool response message(image)\n    for idx, img in enumerate(processed_images):\n        if idx == 0:\n            continue\n        _ = req.get_generation_prompt_ids(processor)\n        req.add_assistant_message(processor, content=description_list[idx - 1])\n        before_tool_call_len = req.input_ids.shape[-1]\n        req.add_tool_response_messages(\n            processor, [ToolResponse(image=[img], text=\"Here is the new image you requested: \")]\n        )\n        after_tool_call_len = req.input_ids.shape[-1]\n        if prev_generated_len == 0:\n            prev_generated_len = after_tool_call_len - before_tool_call_len\n        else:\n            if resize_image:\n                assert after_tool_call_len - before_tool_call_len == prev_generated_len\n        assert req.multi_modal_data[\"image\"] == processed_images[: idx + 1]\n\n    _ = req.get_generation_prompt_ids(processor)\n    req.add_assistant_message(processor, content=description_list[-1])\n\n    messages = [msg.model_dump() for msg in req.messages]\n    tools = [tool.model_dump() for tool in req.tool_schemas] if req.tool_schemas else None\n    full_prompt_info = req._handle_apply_chat_template(\n        processor,\n        messages,\n        multi_modal_data=req.multi_modal_data,\n        tools=tools,\n        add_generation_prompt=False,\n        tokenize=True,\n        return_dict=True,\n    )\n    full_prompt_ids = full_prompt_info[\"input_ids\"]\n    assert full_prompt_ids.eq(req.input_ids).all()\n\n    # We must use dict(full_prompt_info) to convert BatchFeature values to a new dict\n    # because np.array() only keeps the keys for BatchFeature.\n    full_prompt_multi_modal_inputs = full_prompt_info.copy()\n    full_prompt_multi_modal_inputs.pop(\"input_ids\", None)\n    full_prompt_multi_modal_inputs.pop(\"attention_mask\", None)\n\n    for key in full_prompt_multi_modal_inputs:\n        assert full_prompt_multi_modal_inputs[key].eq(req.multi_modal_inputs[key]).all()\n\n\n@pytest.mark.skipif(\n    hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2.5-VL-3B-Instruct\")) is None,\n    reason=\"Processor not available for Qwen/Qwen2.5-VL-B-Instruct\",\n)\ndef test_add_tool_response_messages_image_delta():\n    processor = hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2.5-VL-3B-Instruct\"))\n\n    # From Qwen2.5-VL-3B-Instruct HF example\n    img_1_url = {\"image\": \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg\"}\n    img_1_description = \"A woman sits on the beach at sunset, smiling as she shares a high five with her large dog.\"\n    # GitHub Logo\n    img_2_url = {\"image\": \"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\"}\n    img_2_description = \"A GitHub Logo image\"\n    # Octocat\n    img_3_url = {\"image\": \"https://octodex.github.com/images/orderedlistocat.png\"}\n    img_3_description = \"An Octocat image\"\n\n    image_list = [img_1_url, img_2_url, img_3_url]\n    description_list = [img_1_description, img_2_description, img_3_description]\n    _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=False)\n\n\n@pytest.mark.skipif(\n    hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2.5-VL-3B-Instruct\")) is None,\n    reason=\"Processor not available for Qwen/Qwen2.5-VL-B-Instruct\",\n)\ndef test_add_tool_response_messages_image_delta_resize_image():\n    processor = hf_processor(os.path.expanduser(\"~/models/Qwen/Qwen2.5-VL-3B-Instruct\"))\n\n    # From Qwen2.5-VL-3B-Instruct HF example\n    img_1_url = {\"image\": \"https://qianwen-res.oss-cn-beijing.aliyuncs.com/Qwen-VL/assets/demo.jpeg\"}\n    img_1_description = \"A woman sits on the beach at sunset, smiling as she shares a high five with her large dog.\"\n    # GitHub Logo\n    img_2_url = {\"image\": \"https://github.githubassets.com/assets/GitHub-Mark-ea2971cee799.png\"}\n    img_2_description = \"A GitHub Logo image\"\n    # Octocat\n    img_3_url = {\"image\": \"https://octodex.github.com/images/orderedlistocat.png\"}\n    img_3_description = \"An Octocat image\"\n\n    image_list = [img_1_url, img_2_url, img_3_url]\n    description_list = [img_1_description, img_2_description, img_3_description]\n    _test_add_tool_response_messages_image_delta(processor, image_list, description_list, resize_image=True)\n"
  },
  {
    "path": "tests/workers/rollout/test_sglang_rollout_sharding_manager.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pytest\nimport torch\n\nfrom verl.workers.rollout.sglang_rollout.utils import get_named_tensor_buckets\n\n_TENSOR_1MB = torch.zeros(512, 512)\n_BYTES_1MB = 1 << 20\n\n\n@pytest.mark.parametrize(\n    \"named_tensors, bucket_size_mb, gt_groups\",\n    [\n        (\n            [(\"a\", _TENSOR_1MB), (\"b\", _TENSOR_1MB)],\n            0.5 * _BYTES_1MB,\n            [[\"a\"], [\"b\"]],\n        ),\n        (\n            [(\"a\", _TENSOR_1MB), (\"b\", _TENSOR_1MB)],\n            1 * _BYTES_1MB,\n            [[\"a\"], [\"b\"]],\n        ),\n        (\n            [(\"a\", _TENSOR_1MB), (\"b\", _TENSOR_1MB)],\n            1.5 * _BYTES_1MB,\n            [[\"a\"], [\"b\"]],\n        ),\n        (\n            [(\"a\", _TENSOR_1MB), (\"b\", _TENSOR_1MB)],\n            2 * _BYTES_1MB,\n            [[\"a\", \"b\"]],\n        ),\n    ],\n)\ndef test_get_named_tensor_buckets(named_tensors, bucket_size_mb, gt_groups: list[list[str]]):\n    named_tensors_iter = iter(named_tensors)\n    groups = list(get_named_tensor_buckets(named_tensors_iter, bucket_size_mb))\n    assert len(groups) == len(gt_groups)\n    for group, gt_group in zip(groups, gt_groups, strict=True):\n        assert len(group) == len(gt_group)\n        for (name, _), (gt_name) in zip(group, gt_group, strict=True):\n            assert name == gt_name\n"
  },
  {
    "path": "tests/workers/rollout/test_vllm_cli_args_on_cpu.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\n\nimport pytest\n\nfrom verl.workers.rollout.vllm_rollout.utils import build_cli_args_from_config\n\n\nclass TestBuildCliArgsFromConfig:\n    \"\"\"Tests for CLI argument serialization from config dictionaries.\"\"\"\n\n    def test_string_value(self):\n        \"\"\"String values become '--key value'.\"\"\"\n        config = {\"model\": \"gpt2\"}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--model\", \"gpt2\"]\n\n    def test_integer_value(self):\n        \"\"\"Integer values are converted to strings.\"\"\"\n        config = {\"tensor-parallel-size\": 4}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--tensor-parallel-size\", \"4\"]\n\n    def test_float_value(self):\n        \"\"\"Float values are converted to strings.\"\"\"\n        config = {\"temperature\": 0.7}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--temperature\", \"0.7\"]\n\n    def test_bool_true(self):\n        \"\"\"Bool True adds flag without value.\"\"\"\n        config = {\"enable-prefix-caching\": True}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--enable-prefix-caching\"]\n\n    def test_bool_false(self):\n        \"\"\"Bool False is skipped entirely.\"\"\"\n        config = {\"enable-prefix-caching\": False}\n        result = build_cli_args_from_config(config)\n        assert result == []\n\n    def test_none_value(self):\n        \"\"\"None values are skipped.\"\"\"\n        config = {\"lora-path\": None}\n        result = build_cli_args_from_config(config)\n        assert result == []\n\n    def test_list_values(self):\n        \"\"\"List values are expanded into multiple arguments.\"\"\"\n        config = {\"cudagraph-capture-sizes\": [1, 2, 4, 8]}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--cudagraph-capture-sizes\", \"1\", \"2\", \"4\", \"8\"]\n\n    def test_empty_list(self):\n        \"\"\"Empty lists are skipped (vLLM nargs='+' requires at least one value).\"\"\"\n        config = {\"cudagraph-capture-sizes\": []}\n        result = build_cli_args_from_config(config)\n        assert result == []\n\n    def test_list_with_strings(self):\n        \"\"\"List of strings is properly expanded.\"\"\"\n        config = {\"allowed-origins\": [\"http://localhost\", \"http://example.com\"]}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--allowed-origins\", \"http://localhost\", \"http://example.com\"]\n\n    def test_dict_value(self):\n        \"\"\"Dict values are JSON serialized.\"\"\"\n        config = {\"extra-config\": {\"key\": \"value\", \"nested\": True}}\n        result = build_cli_args_from_config(config)\n        assert result[0] == \"--extra-config\"\n        # JSON output may have different key ordering, so parse and compare\n        assert json.loads(result[1]) == {\"key\": \"value\", \"nested\": True}\n\n    def test_mixed_config(self):\n        \"\"\"Test a realistic mixed configuration.\"\"\"\n        config = {\n            \"tensor-parallel-size\": 4,\n            \"enable-prefix-caching\": True,\n            \"disable-log-requests\": False,\n            \"lora-path\": None,\n            \"cudagraph-capture-sizes\": [1, 2, 4, 8],\n            \"max-model-len\": 2048,\n        }\n        result = build_cli_args_from_config(config)\n\n        # Check expected args are present\n        assert \"--tensor-parallel-size\" in result\n        assert \"4\" in result\n        assert \"--enable-prefix-caching\" in result\n        assert \"--cudagraph-capture-sizes\" in result\n        assert \"1\" in result\n        assert \"8\" in result\n        assert \"--max-model-len\" in result\n        assert \"2048\" in result\n\n        # Check skipped values are not present\n        assert \"--disable-log-requests\" not in result\n        assert \"--lora-path\" not in result\n\n    def test_preserves_order(self):\n        \"\"\"Arguments should preserve dictionary order (Python 3.7+).\"\"\"\n        config = {\"first\": \"a\", \"second\": \"b\", \"third\": \"c\"}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--first\", \"a\", \"--second\", \"b\", \"--third\", \"c\"]\n\n    def test_empty_config(self):\n        \"\"\"Empty config returns empty list.\"\"\"\n        config = {}\n        result = build_cli_args_from_config(config)\n        assert result == []\n\n    def test_single_element_list(self):\n        \"\"\"Single element list works correctly.\"\"\"\n        config = {\"sizes\": [42]}\n        result = build_cli_args_from_config(config)\n        assert result == [\"--sizes\", \"42\"]\n\n\nif __name__ == \"__main__\":\n    pytest.main([__file__, \"-v\"])\n"
  },
  {
    "path": "tests/workers/test_fsdp_attn_implementation.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nTest for attn_implementation override configuration in FSDP workers.\n\nThis test verifies that the fix for honoring attn_implementation override config\nworks correctly in the ActorRolloutRefWorker._build_model_optimizer method.\n\"\"\"\n\nfrom unittest.mock import Mock, patch\n\nimport pytest\nimport torch\nfrom omegaconf import OmegaConf\nfrom transformers import AutoConfig, AutoModelForCausalLM\n\n# Only run these tests if we can import verl components\ntry:\n    from verl.workers.config import FSDPEngineConfig  # noqa: F401\n    from verl.workers.fsdp_workers import (\n        ActorRolloutRefWorker,  # noqa: F401\n        CriticWorker,  # noqa: F401\n    )\n\n    VERL_AVAILABLE = True\nexcept ImportError:\n    VERL_AVAILABLE = False\n\n\n@pytest.mark.skipif(not VERL_AVAILABLE, reason=\"VERL components not available\")\nclass TestFSDPAttnImplementation:\n    \"\"\"Test cases for attn_implementation override in FSDP workers.\"\"\"\n\n    def test_attn_implementation_extraction_logic(self):\n        \"\"\"Test the core logic for extracting attn_implementation from override config.\"\"\"\n\n        # Test case 1: Default behavior\n        override_config = {}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"flash_attention_2\"\n\n        # Test case 2: Override to eager\n        override_config = {\"attn_implementation\": \"eager\"}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"eager\"\n\n        # Test case 3: Override to sdpa\n        override_config = {\"attn_implementation\": \"sdpa\"}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"sdpa\"\n\n        # Test case 4: Other configs don't affect attn_implementation\n        override_config = {\"other_setting\": \"value\", \"dropout\": 0.1}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"flash_attention_2\"\n\n    @patch(\"transformers.AutoConfig.from_pretrained\")\n    @patch(\"transformers.AutoModelForCausalLM.from_pretrained\")\n    def test_attn_implementation_passed_to_autoconfig(self, mock_model_from_pretrained, mock_config_from_pretrained):\n        \"\"\"Test that attn_implementation is correctly passed to AutoConfig.from_pretrained.\"\"\"\n\n        # Mock the AutoConfig return value\n        mock_config = Mock()\n        mock_config.tie_word_embeddings = False\n        mock_config.architectures = [\"LlamaForCausalLM\"]\n        mock_config_from_pretrained.return_value = mock_config\n\n        # Mock the model return value\n        mock_model = Mock()\n        mock_model_from_pretrained.return_value = mock_model\n\n        # Test data\n        test_cases = [\n            ({}, \"flash_attention_2\"),  # Default\n            ({\"attn_implementation\": \"eager\"}, \"eager\"),  # Override to eager\n            ({\"attn_implementation\": \"sdpa\"}, \"sdpa\"),  # Override to sdpa\n        ]\n\n        for override_config, expected_attn_impl in test_cases:\n            # Reset mocks\n            mock_config_from_pretrained.reset_mock()\n            mock_model_from_pretrained.reset_mock()\n\n            # Simulate the logic from FSDP workers\n            attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n            # This simulates what happens in _build_model_optimizer\n            AutoConfig.from_pretrained(\"test_path\", trust_remote_code=False, attn_implementation=attn_implementation)\n\n            # Verify AutoConfig.from_pretrained was called with correct attn_implementation\n            mock_config_from_pretrained.assert_called_once_with(\n                \"test_path\", trust_remote_code=False, attn_implementation=expected_attn_impl\n            )\n\n    @patch(\"transformers.AutoConfig.from_pretrained\")\n    @patch(\"transformers.AutoModelForCausalLM.from_pretrained\")\n    def test_attn_implementation_passed_to_model(self, mock_model_from_pretrained, mock_config_from_pretrained):\n        \"\"\"Test that attn_implementation is correctly passed to model.from_pretrained.\"\"\"\n\n        # Mock the AutoConfig return value\n        mock_config = Mock()\n        mock_config.tie_word_embeddings = False\n        mock_config.architectures = [\"LlamaForCausalLM\"]\n        mock_config_from_pretrained.return_value = mock_config\n\n        # Mock the model return value\n        mock_model = Mock()\n        mock_model_from_pretrained.return_value = mock_model\n\n        # Test with override config\n        override_config = {\"attn_implementation\": \"eager\"}\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n        # This simulates what happens in _build_model_optimizer\n        AutoModelForCausalLM.from_pretrained(\n            pretrained_model_name_or_path=\"test_path\",\n            torch_dtype=torch.bfloat16,\n            config=mock_config,\n            trust_remote_code=False,\n            attn_implementation=attn_implementation,\n        )\n\n        # Verify AutoModelForCausalLM.from_pretrained was called with correct attn_implementation\n        mock_model_from_pretrained.assert_called_once_with(\n            pretrained_model_name_or_path=\"test_path\",\n            torch_dtype=torch.bfloat16,\n            config=mock_config,\n            trust_remote_code=False,\n            attn_implementation=\"eager\",\n        )\n\n    def test_override_config_integration(self):\n        \"\"\"Test that override_config from Hydra configuration works correctly.\"\"\"\n\n        # Simulate the OmegaConf configuration structure used in VERL\n        config_dict = {\n            \"model\": {\"path\": \"/test/path\", \"override_config\": {\"attn_implementation\": \"eager\", \"dropout\": 0.1}}\n        }\n\n        # Convert to OmegaConf structure\n        omegaconf = OmegaConf.create(config_dict)\n\n        # Simulate what happens in the FSDP worker\n        override_model_config = OmegaConf.to_container(OmegaConf.create(omegaconf.model.get(\"override_config\", {})))\n\n        # Test extraction\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"eager\"\n\n        # Test that other configs are preserved\n        assert override_model_config.get(\"dropout\") == 0.1\n\n    def test_hydra_plus_prefix_config(self):\n        \"\"\"Test that Hydra +prefix configurations work correctly.\"\"\"\n\n        # This simulates the configuration when user specifies:\n        # +actor_rollout_ref.model.override_config.attn_implementation=eager\n\n        # The + prefix in Hydra adds new keys to the config\n        config_dict = {\n            \"actor_rollout_ref\": {\n                \"model\": {\n                    \"path\": \"/test/path\",\n                    \"override_config\": {\n                        \"attn_implementation\": \"eager\"  # This gets added via +prefix\n                    },\n                }\n            }\n        }\n\n        omegaconf = OmegaConf.create(config_dict)\n\n        # Extract override config as done in FSDP workers\n        override_model_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.actor_rollout_ref.model.get(\"override_config\", {}))\n        )\n\n        # Verify extraction works\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"eager\"\n\n    def test_backward_compatibility(self):\n        \"\"\"Test that the fix maintains backward compatibility.\"\"\"\n\n        # Test case 1: No override_config at all (old behavior)\n        config_without_override = {}\n        attn_implementation = config_without_override.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n        # Test case 2: Empty override_config\n        config_with_empty_override = {\"override_config\": {}}\n        override_config = config_with_empty_override.get(\"override_config\", {})\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n        # Test case 3: override_config with other settings but no attn_implementation\n        config_with_other_overrides = {\"override_config\": {\"dropout\": 0.1, \"hidden_size\": 1024}}\n        override_config = config_with_other_overrides.get(\"override_config\", {})\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n    def test_critic_attn_implementation_extraction_logic(self):\n        \"\"\"Test the core logic for extracting attn_implementation from override config for CriticWorker.\"\"\"\n\n        # Test case 1: Default behavior for critic\n        override_config = {}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"flash_attention_2\"\n\n        # Test case 2: Override to eager for critic\n        override_config = {\"attn_implementation\": \"eager\"}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"eager\"\n\n        # Test case 3: Override to sdpa for critic\n        override_config = {\"attn_implementation\": \"sdpa\"}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"sdpa\"\n\n        # Test case 4: Other configs don't affect attn_implementation for critic\n        override_config = {\"other_setting\": \"value\", \"dropout\": 0.1}\n        attn_impl = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_impl == \"flash_attention_2\"\n\n    @patch(\"transformers.AutoConfig.from_pretrained\")\n    def test_critic_attn_implementation_passed_to_autoconfig(self, mock_config_from_pretrained):\n        \"\"\"Test that attn_implementation is correctly passed to AutoConfig.from_pretrained in CriticWorker.\"\"\"\n\n        # Mock the AutoConfig return value\n        mock_config = Mock()\n        mock_config.tie_word_embeddings = False\n        mock_config.architectures = [\"LlamaForCausalLM\"]\n        mock_config.num_labels = 1\n        mock_config_from_pretrained.return_value = mock_config\n\n        # Test data for critic model\n        test_cases = [\n            ({}, \"flash_attention_2\"),  # Default\n            ({\"attn_implementation\": \"eager\"}, \"eager\"),  # Override to eager\n            ({\"attn_implementation\": \"sdpa\"}, \"sdpa\"),  # Override to sdpa\n        ]\n\n        for override_config, expected_attn_impl in test_cases:\n            # Reset mocks\n            mock_config_from_pretrained.reset_mock()\n\n            # Simulate the logic from CriticWorker _build_critic_model_optimizer\n            attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n            # This simulates what should happen in CriticWorker._build_critic_model_optimizer\n            # (This is where the fix needs to be applied in the actual implementation)\n            AutoConfig.from_pretrained(\n                \"test_path\",\n                attn_implementation=attn_implementation,\n                trust_remote_code=False,\n            )\n\n            # Verify AutoConfig.from_pretrained was called with correct attn_implementation\n            mock_config_from_pretrained.assert_called_once_with(\n                \"test_path\",\n                attn_implementation=expected_attn_impl,\n                trust_remote_code=False,\n            )\n\n    def test_critic_override_config_integration(self):\n        \"\"\"Test that override_config from Hydra configuration works correctly for CriticWorker.\"\"\"\n\n        # Simulate the OmegaConf configuration structure used in VERL for critic\n        config_dict = {\n            \"critic\": {\n                \"model\": {\"path\": \"/test/path\", \"override_config\": {\"attn_implementation\": \"eager\", \"dropout\": 0.1}}\n            }\n        }\n\n        # Convert to OmegaConf structure\n        omegaconf = OmegaConf.create(config_dict)\n\n        # Simulate what happens in the CriticWorker\n        override_model_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {}))\n        )\n\n        # Test extraction for critic\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"eager\"\n\n        # Test that other configs are preserved for critic\n        assert override_model_config.get(\"dropout\") == 0.1\n\n    def test_critic_hydra_plus_prefix_config(self):\n        \"\"\"Test that Hydra +prefix configurations work correctly for CriticWorker.\"\"\"\n\n        # This simulates the configuration when user specifies:\n        # +critic.model.override_config.attn_implementation=eager\n\n        # The + prefix in Hydra adds new keys to the config\n        config_dict = {\n            \"critic\": {\n                \"model\": {\n                    \"path\": \"/test/path\",\n                    \"override_config\": {\n                        \"attn_implementation\": \"eager\"  # This gets added via +prefix for critic\n                    },\n                }\n            }\n        }\n\n        omegaconf = OmegaConf.create(config_dict)\n\n        # Extract override config as done in CriticWorker\n        override_model_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {}))\n        )\n\n        # Verify extraction works for critic\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"eager\"\n\n    def test_both_actor_and_critic_configuration(self):\n        \"\"\"Test that both actor and critic can have different attn_implementation overrides simultaneously.\"\"\"\n\n        # This simulates a complete training configuration with both actor and critic overrides\n        config_dict = {\n            \"actor_rollout_ref\": {\"model\": {\"override_config\": {\"attn_implementation\": \"eager\"}}},\n            \"critic\": {\"model\": {\"override_config\": {\"attn_implementation\": \"sdpa\"}}},\n        }\n\n        omegaconf = OmegaConf.create(config_dict)\n\n        # Extract actor override config\n        actor_override_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.actor_rollout_ref.model.get(\"override_config\", {}))\n        )\n        actor_attn_implementation = actor_override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n        # Extract critic override config\n        critic_override_config = OmegaConf.to_container(\n            OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {}))\n        )\n        critic_attn_implementation = critic_override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n        # Verify both can be configured independently\n        assert actor_attn_implementation == \"eager\"\n        assert critic_attn_implementation == \"sdpa\"\n\n    def test_critic_backward_compatibility(self):\n        \"\"\"Test that the CriticWorker fix maintains backward compatibility.\"\"\"\n\n        # Test case 1: No override_config at all for critic (old behavior)\n        config_without_override = {}\n        attn_implementation = config_without_override.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n        # Test case 2: Empty override_config for critic\n        config_with_empty_override = {\"override_config\": {}}\n        override_config = config_with_empty_override.get(\"override_config\", {})\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n        # Test case 3: override_config with other settings but no attn_implementation for critic\n        config_with_other_overrides = {\"override_config\": {\"dropout\": 0.1, \"num_labels\": 1}}\n        override_config = config_with_other_overrides.get(\"override_config\", {})\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        assert attn_implementation == \"flash_attention_2\"\n\n\ndef test_attn_implementation_fix_integration():\n    \"\"\"Integration test to verify the entire fix works as expected.\"\"\"\n\n    # This test simulates the complete flow from configuration to model creation\n\n    # Step 1: Simulate Hydra configuration with +prefix\n    # user_config = \"+actor_rollout_ref.model.override_config.attn_implementation=eager\"\n\n    # This would result in a config structure like:\n    config_dict = {\"actor_rollout_ref\": {\"model\": {\"override_config\": {\"attn_implementation\": \"eager\"}}}}\n\n    # Step 2: Extract override_model_config as done in FSDP workers\n    omegaconf = OmegaConf.create(config_dict)\n    override_model_config = OmegaConf.to_container(\n        OmegaConf.create(omegaconf.actor_rollout_ref.model.get(\"override_config\", {}))\n    )\n\n    # Step 3: Apply the fix logic\n    attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n    # Step 4: Verify the fix works\n    assert attn_implementation == \"eager\"\n\n    # Step 5: Verify this would be passed to both AutoConfig and Model creation\n    # (This would normally be done with mocks, but we can test the parameter preparation)\n    config_params = {\"attn_implementation\": attn_implementation}\n    model_params = {\"attn_implementation\": attn_implementation}\n\n    assert config_params[\"attn_implementation\"] == \"eager\"\n    assert model_params[\"attn_implementation\"] == \"eager\"\n\n\ndef test_critic_attn_implementation_fix_integration():\n    \"\"\"Integration test to verify the entire fix works as expected for CriticWorker.\"\"\"\n\n    # This test simulates the complete flow from configuration to model creation for critic\n\n    # Step 1: Simulate Hydra configuration with +prefix for critic\n    # user_config = \"+critic.model.override_config.attn_implementation=sdpa\"\n\n    # This would result in a config structure like:\n    config_dict = {\"critic\": {\"model\": {\"override_config\": {\"attn_implementation\": \"sdpa\"}}}}\n\n    # Step 2: Extract override_model_config as should be done in CriticWorker\n    omegaconf = OmegaConf.create(config_dict)\n    override_model_config = OmegaConf.to_container(OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {})))\n\n    # Step 3: Apply the fix logic (what needs to be implemented in CriticWorker)\n    attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n    # Step 4: Verify the fix works for critic\n    assert attn_implementation == \"sdpa\"\n\n    # Step 5: Verify this would be passed to AutoConfig creation for critic\n    config_params = {\"attn_implementation\": attn_implementation}\n\n    assert config_params[\"attn_implementation\"] == \"sdpa\"\n\n\ndef test_complete_training_configuration():\n    \"\"\"Integration test for a complete training configuration with both actor and critic overrides.\"\"\"\n\n    # This test simulates a realistic training configuration where both\n    # actor and critic have different attention implementations\n    config_dict = {\n        \"actor_rollout_ref\": {\n            \"model\": {\n                \"path\": \"/shared/models/llama-7b\",\n                \"override_config\": {\"attn_implementation\": \"eager\", \"torch_dtype\": \"bfloat16\"},\n            }\n        },\n        \"critic\": {\n            \"model\": {\n                \"path\": \"/shared/models/llama-7b\",\n                \"override_config\": {\"attn_implementation\": \"sdpa\", \"num_labels\": 1},\n            }\n        },\n    }\n\n    omegaconf = OmegaConf.create(config_dict)\n\n    # Extract configurations as would be done in the workers\n    actor_override_config = OmegaConf.to_container(\n        OmegaConf.create(omegaconf.actor_rollout_ref.model.get(\"override_config\", {}))\n    )\n    critic_override_config = OmegaConf.to_container(OmegaConf.create(omegaconf.critic.model.get(\"override_config\", {})))\n\n    # Apply the fix logic for both\n    actor_attn_implementation = actor_override_config.get(\"attn_implementation\", \"flash_attention_2\")\n    critic_attn_implementation = critic_override_config.get(\"attn_implementation\", \"flash_attention_2\")\n\n    # Verify both configurations work independently\n    assert actor_attn_implementation == \"eager\"\n    assert critic_attn_implementation == \"sdpa\"\n\n    # Verify other configs are preserved\n    assert actor_override_config.get(\"torch_dtype\") == \"bfloat16\"\n    assert critic_override_config.get(\"num_labels\") == 1\n\n\nif __name__ == \"__main__\":\n    # Run basic tests\n    test_attn_implementation_fix_integration()\n    test_critic_attn_implementation_fix_integration()\n    test_complete_training_configuration()\n\n    if VERL_AVAILABLE:\n        # Run class-based tests\n        test_class = TestFSDPAttnImplementation()\n        test_class.test_attn_implementation_extraction_logic()\n        test_class.test_override_config_integration()\n        test_class.test_hydra_plus_prefix_config()\n        test_class.test_backward_compatibility()\n\n        # Run new critic tests\n        test_class.test_critic_attn_implementation_extraction_logic()\n        test_class.test_critic_override_config_integration()\n        test_class.test_critic_hydra_plus_prefix_config()\n        test_class.test_both_actor_and_critic_configuration()\n        test_class.test_critic_backward_compatibility()\n\n        print(\"✓ All FSDP attn_implementation tests passed!\")\n        print(\"✓ All CriticWorker attn_implementation tests passed!\")\n    else:\n        print(\"⚠ VERL components not available, skipping VERL-specific tests\")\n\n    print(\"✓ Integration tests passed!\")\n    print(\"✓ Critic integration tests passed!\")\n"
  },
  {
    "path": "tests/workers/test_fsdp_workers.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nfrom omegaconf import OmegaConf\n\nfrom verl.workers.fsdp_workers import ActorRolloutRefWorker\n\n\ndef test_actor_rollout_ref_worker_actor_ref_model():\n    \"\"\"Test specifying different reference/actor model\"\"\"\n    os.environ[\"RANK\"] = \"0\"\n    os.environ[\"WORLD_SIZE\"] = \"1\"\n    os.environ[\"MASTER_ADDR\"] = \"127.0.0.1\"\n    os.environ[\"MASTER_PORT\"] = \"8888\"\n\n    actor_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-0.5B-Instruct\")\n    ref_model_path = os.path.expanduser(\"~/models/Qwen/Qwen2.5-1.5B-Instruct\")\n    config_str = f\"\"\"\n    model:\n      path: {actor_model_path}\n    actor:\n      _target_: verl.workers.config.FSDPActorConfig\n      strategy: fsdp\n      fsdp_config:\n        _target_: verl.workers.config.FSDPEngineConfig\n        fsdp_size: -1\n        forward_prefetch: false\n      profiler:\n        tool: torch_memory\n        save_path: ./mem_snapshots\n        tool_config:\n          torch_memory:\n            _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n            trace_alloc_max_entries: 100000\n            stack_depth: 32\n    ref:\n      model:\n        path: {ref_model_path}\n      fsdp_config:\n        _target_: verl.workers.config.FSDPEngineConfig\n        fsdp_size: -1\n      profiler:\n        tool: torch_memory\n        save_path: ./mem_snapshots\n        tool_config:\n          torch_memory:\n            _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n            trace_alloc_max_entries: 100000\n            stack_depth: 32\n      log_prob_micro_batch_size: 1\n      ulysses_sequence_parallel_size: 1\n      entropy_from_logits_with_chunking: false\n    \"\"\"\n    dict_conf = OmegaConf.create(config_str)\n    actor_rollout_ref_worker = ActorRolloutRefWorker(dict_conf, role=\"ref\")\n    actor_rollout_ref_worker.init_model()\n\n    model_config = actor_rollout_ref_worker.ref_module_fsdp._fsdp_wrapped_module.config\n    assert model_config.hidden_size == 1536\n\n    # set ref.model to null, fallback to default case where actor is the same as reference\n    dict_conf[\"ref\"][\"model\"] = None\n    actor_rollout_ref_worker = ActorRolloutRefWorker(dict_conf, role=\"ref\")\n    actor_rollout_ref_worker.init_model()\n\n    model_config = actor_rollout_ref_worker.ref_module_fsdp._fsdp_wrapped_module.config\n    assert model_config.hidden_size == 896\n"
  },
  {
    "path": "verl/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport importlib\nimport logging\nimport os\n\nfrom packaging.version import parse as parse_version\n\nfrom .protocol import DataProto\nfrom .utils.device import is_npu_available\nfrom .utils.import_utils import import_external_libs\nfrom .utils.logging_utils import set_basic_config\n\nversion_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))\n\nwith open(os.path.join(version_folder, \"version/version\")) as f:\n    __version__ = f.read().strip()\n\n\nset_basic_config(level=logging.WARNING)\n\n\n__all__ = [\"DataProto\", \"__version__\"]\n\n\nmodules = os.getenv(\"VERL_USE_EXTERNAL_MODULES\", \"\")\nif modules:\n    modules = modules.split(\",\")\n    import_external_libs(modules)\n\n\nif os.getenv(\"VERL_USE_MODELSCOPE\", \"False\").lower() == \"true\":\n    if importlib.util.find_spec(\"modelscope\") is None:\n        raise ImportError(\"You are using the modelscope hub, please install modelscope by `pip install modelscope -U`\")\n    # Patch hub to download models from modelscope to speed up.\n    from modelscope.utils.hf_util import patch_hub\n\n    patch_hub()\n\n\nif is_npu_available:\n    # Workaround for torch-npu's lack of support for creating nested tensors from NPU tensors.\n    #\n    # ```\n    # >>> a, b = torch.arange(3).npu(), torch.arange(5).npu() + 3\n    # >>> nt = torch.nested.nested_tensor([a, b], layout=torch.jagged)\n    # ```\n    # throws \"not supported in npu\" on Ascend NPU.\n    # See https://github.com/Ascend/pytorch/blob/294cdf5335439b359991cecc042957458a8d38ae/torch_npu/utils/npu_intercept.py#L109\n    # for details.\n\n    import torch\n\n    try:\n        if hasattr(torch.nested.nested_tensor, \"__wrapped__\"):\n            torch.nested.nested_tensor = torch.nested.nested_tensor.__wrapped__\n        if hasattr(torch.nested.as_nested_tensor, \"__wrapped__\"):\n            torch.nested.as_nested_tensor = torch.nested.as_nested_tensor.__wrapped__\n    except AttributeError:\n        pass\n\n    # In verl, the driver process aggregates the computation results of workers via Ray.\n    # Therefore, after a worker completes its computation job, it will package the output\n    # using tensordict and transfer it to the CPU. Since the `to` operation of tensordict\n    # is non-blocking, when transferring data from a device to the CPU, it is necessary to\n    # ensure that a batch of data has been completely transferred before being used on the\n    # host; otherwise, unexpected precision issues may arise. Tensordict has already noticed\n    # this problem and fixed it. Ref: https://github.com/pytorch/tensordict/issues/725\n    # However, the relevant modifications only cover CUDA and MPS devices and do not take effect\n    # for third-party devices such as NPUs. This patch fixes this issue, and the relevant\n    # modifications can be removed once the fix is merged into tensordict.\n\n    import tensordict\n\n    if parse_version(tensordict.__version__) < parse_version(\"0.10.0\"):\n        from tensordict.base import TensorDictBase\n\n        def _sync_all_patch(self):\n            from torch._utils import _get_available_device_type, _get_device_module\n\n            device_type = _get_available_device_type()\n            if device_type is None:\n                return\n\n            device_module = _get_device_module(device_type)\n            device_module.synchronize()\n\n        TensorDictBase._sync_all = _sync_all_patch\n"
  },
  {
    "path": "verl/base_config.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport collections\nfrom dataclasses import FrozenInstanceError, dataclass, fields\nfrom typing import Any\n\n\n# BaseConfig class inherits from collections.abc.Mapping, which means it can act like a dictionary\n@dataclass\nclass BaseConfig(collections.abc.Mapping):\n    \"\"\"The BaseConfig provides dict-like interface for a dataclass config.\n\n    By default all fields in the config is not mutable, unless specified in\n    \"_mutable_fields\". The BaseConfig class implements the Mapping Abstract Base Class.\n    This allows instances of this class to be used like dictionaries.\n    \"\"\"\n\n    _mutable_fields = set()\n    _target_: str = \"\"\n\n    def __setattr__(self, name: str, value):\n        \"\"\"Set the value of an attribute. Check if the attr is mutable before setting the value.\"\"\"\n        # If the field already exists, it's considered frozen unless it's in _mutable_fields\n        if name in self.__dict__ and name not in getattr(self, \"_mutable_fields\", set()):\n            raise FrozenInstanceError(f\"Field '{name}' is frozen and cannot be modified\")\n        super().__setattr__(name, value)\n\n    def get(self, key: str, default: Any = None) -> Any:\n        \"\"\"Get the value associated with the given key. If the key does not exist, return the default value.\n\n        Args:\n            key (str): The attribute name to retrieve.\n            default (Any, optional): The value to return if the attribute does not exist. Defaults to None.\n\n        Returns:\n            Any: The value of the attribute or the default value.\n        \"\"\"\n        try:\n            return getattr(self, key)\n        except AttributeError:\n            return default\n\n    def __getitem__(self, key: str):\n        \"\"\"Implement the [] operator for the class. Allows accessing attributes like dictionary items.\n\n        Args:\n            key (str): The attribute name to retrieve.\n\n        Returns:\n            Any: The value of the attribute.\n\n        Raises:\n            AttributeError: If the attribute does not exist.\n            TypeError: If the key type is not string\n        \"\"\"\n        return getattr(self, key)\n\n    def __iter__(self):\n        \"\"\"Implement the iterator protocol. Allows iterating over the attribute names of the instance.\n\n        Yields:\n            str: The name of each field in the dataclass.\n        \"\"\"\n        for f in fields(self):\n            yield f.name\n\n    def __len__(self):\n        \"\"\"\n        Return the number of fields in the dataclass.\n\n        Returns:\n            int: The number of fields in the dataclass.\n        \"\"\"\n        return len(fields(self))\n"
  },
  {
    "path": "verl/checkpoint_engine/README.md",
    "content": "Checkpoint Engine\n---\n\n### Overview\n\nCheckpoint Engine is an unified abstract layer to synchronize weights between various training backends and inference backends. It provides three unified APIs:\n- send_weights: get named tensors from generator and send them in streaming manner.\n- receive_weights: return a tensor generator that yield named tensors in streaming manner.\n- get_weights: return a tensor generator that yield named tensors in streaming manner, used for each inference instance update weight independently from local cache (e.g share memory, disk).\n\n![checkpoint-engine](https://github.com/wuxibin89/verl/blob/wuxibin/doc_images/docs/_static/checkpoint_engine.png?raw=true)\n\n### Supported Backends\n\n||Comm Library|Topology|Hardware|Performance|Elastic|Use case|\n|----|----|----|----|----|----|----|\n|naive|torch.distributed|all_gather|NVIDIA/AMD/Ascend|Very High|NA|On-policy training<br>- Trainer/rollout colocated\n|nccl|NCCL|all_gather+broadcast|NVIDIA GPU & NCCL|Very High|Low: rebuild nccl group|Off-policy training<br>- Trainer/rollout disaggregated<br>- Fixed clusters\n|hccl|HCCL|all_gather+broadcast|Ascend NPU & HCCL| High|Low: rebuild hccl group|Off-policy training<br>- Trainer/rollout disaggregated<br>- Fixed clusters\n|nixl|NIXL|all_gather+ring p2p|Various transport backends (D2D, H2H, H2D, etc)<br>- UCX<br>- UCCL<br>- Mooncacke|Medium/High|High: dynamic adjust ring topology|Off-policy training<br>- Trainer/rollout disaggregated<br>- Elastic rollout<br>- Rollout fault tolerance<br>- Heterogeneous hardware rollout\n|kimi_ckpt_engine|MOONCAKE+NCCL/HCCL|p2p+broadcast|NVIDIA/Ascend|High|Low: rebuild communication group|Off-policy training<br>- Trainer/rollout disaggregated<br>- Save checkpoint each time\n|mooncake|Mooncake Transfer Engine|all_gather+ring p2p|NVIDIA/Ascend|High|High: dynamic adjust ring topology|Off-policy training<br>- Trainer/rollout disaggregated<br>- Fixed clusters\n\n##### kimi_ckpt_engine detail:\n\nIn the kimi_ckpt_engine workflow, the trainer first offloads the weights to the CPU, and the rollout creates a sub communication group that includes all the cards for the rollout. Then, using Mooncake transfer engine, these weights are transmitted via P2P to a specific worker in the rollout, followed by a broadcast to all other rollout workers.\n\n<img src=\"https://github.com/kip-cxj/verl/blob/cxj/doc_imgs/docs/_static/kimi_ckpt_engine.png?raw=true\" alt=\"kimi-ckpt-engine\" width=\"50%\">\n\nThis mode requires the P2P feature of checkpoint_engine. Please ensure you have installed it via pip install 'checkpoint-engine[p2p]' and that your version is 0.4.0 or higher.\n\nIn addition, during the installation of checkpoint-engine[p2p], the transfer engine will be installed. However, This library has no prebuilt packages for Ascend devices and must be compiled from source. For detailed compilation instructions, see: [transfer-engine: ascend direct](https://github.com/kvcache-ai/Mooncake/blob/main/docs/source/design/transfer-engine/ascend_direct_transport.md)\n\nNote: Important Configuration for Ascend Devices\nIf you are using CANN version >= 8.5.0 on Ascend devices, you must set the following environment variable to enable intra-node ROCE:\n\n```bash\nexport HCCL_INTRA_ROCE_ENABLE=1\n```\n\n### Benchmark\n1. benchmark setup\n- model: Qwen/Qwen3-30B-A3B-Base\n- trainer: fsdp world_size=2 (since Ascend 910C has 64GB of HBM, we set world_size=4)\n- rollout: num_rollout=30 (only receive weight without cuda ipc to vllm/sglang)\n```bash\npytest tests/checkpoint_engine/test_correctness_on_gpu.py\npytest tests/checkpoint_engine/test_correctness_on_npu.py\npytest tests/checkpoint_engine/test_special_server_adapter.py\n```\n\n2. benchmark result\n\n| hardware | backend | time cost (s) | Bandwidth(GB/s) |\n|----|----|----|----|\n|4*8 H100, ConnectX-7 400 Gbps (InfiniBand)| NCCL | ~7 | 8.25|\n|4*8 H100, ConnectX-7 400 Gbps (InfiniBand)| NIXL | ~7 | 8.25|\n|2*16 Ascend 910C, inner suppernode| HCCL | ~11 | 5.3|\n|2*16 Ascend 910C, inner suppernode| kimi_ckpt_engine | offload: 7 update: 3.5 | 16.5|\n|2*8 H100, ConnectX-7 400 Gbps (InfiniBand)| mooncake | 5.93 | 9.44|\n"
  },
  {
    "path": "verl/checkpoint_engine/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .base import (\n    CheckpointEngine,\n    CheckpointEngineManager,\n    CheckpointEngineRegistry,\n    CheckpointEngineWorker,\n    ColocatedCheckpointEngine,\n    TensorMeta,\n)\n\n__all__ = [\n    \"CheckpointEngine\",\n    \"CheckpointEngineRegistry\",\n    \"TensorMeta\",\n    \"ColocatedCheckpointEngine\",\n    \"CheckpointEngineManager\",\n    \"CheckpointEngineWorker\",\n]\n\ntry:\n    from .nccl_checkpoint_engine import NCCLCheckpointEngine\n\n    __all__ += [\"NCCLCheckpointEngine\"]\nexcept ImportError:\n    NCCLCheckpointEngine = None\n\ntry:\n    from .hccl_checkpoint_engine import HCCLCheckpointEngine\n\n    __all__ += [\"HCCLCheckpointEngine\"]\nexcept ImportError:\n    HCCLCheckpointEngine = None\n\ntry:\n    from .nixl_checkpoint_engine import NIXLCheckpointEngine\n\n    __all__ += [\"NIXLCheckpointEngine\"]\nexcept ImportError:\n    NIXLCheckpointEngine = None\n\ntry:\n    from .kimi_checkpoint_engine import KIMICheckpointEngine\n\n    __all__ += [\"KIMICheckpointEngine\"]\nexcept ImportError:\n    KIMICheckpointEngine = None\n\ntry:\n    from .mooncake_checkpoint_engine import MooncakeCheckpointEngine\n\n    __all__ += [\"MooncakeCheckpointEngine\"]\nexcept ImportError:\n    MooncakeCheckpointEngine = None\n"
  },
  {
    "path": "verl/checkpoint_engine/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Generator, TypedDict\n\nimport ray\nimport torch\n\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\nfrom verl.utils.distributed import initialize_global_process_group_ray\nfrom verl.utils.ray_utils import auto_await\nfrom verl.workers.config import CheckpointEngineConfig, HFModelConfig, RolloutConfig\nfrom verl.workers.rollout import BaseRollout, RolloutReplica, get_rollout_class\n\n\nclass TensorMeta(TypedDict):\n    name: str\n    shape: torch.Size\n    dtype: torch.dtype\n    offset: int\n\n\nclass CheckpointEngineRegistry:\n    \"\"\"Checkpoint engine registry.\"\"\"\n\n    _registry: dict[str, type[\"CheckpointEngine\"]] = {}\n\n    def register(backend: str):\n        \"\"\"Register a checkpoint engine.\n\n        Args:\n            backend: The backend of the checkpoint engine.\n        \"\"\"\n\n        def wrapper(cls: type[\"CheckpointEngine\"]):\n            CheckpointEngineRegistry._registry[backend] = cls\n            return cls\n\n        return wrapper\n\n    @classmethod\n    def get(cls, backend: str) -> type[\"CheckpointEngine\"]:\n        \"\"\"Get the checkpoint engine class.\n\n        Args:\n            backend: The backend of the checkpoint engine.\n\n        Returns:\n            The checkpoint engine class.\n        \"\"\"\n        return cls._registry[backend]\n\n    @classmethod\n    def new(cls, backend: str, *args, **kwargs) -> \"CheckpointEngine\":\n        \"\"\"Create a new checkpoint engine instance.\n\n        Args:\n            backend: The backend of the checkpoint engine.\n            *args: Variable length argument pass to the checkpoint engine constructor.\n            **kwargs: Arbitrary keyword arguments pass to the checkpoint engine constructor.\n\n        Returns:\n            A new checkpoint engine instance.\n        \"\"\"\n        if backend not in cls._registry:\n            raise ValueError(f\"Checkpoint engine {backend} not registered\")\n        return cls._registry[backend](*args, **kwargs)\n\n\nclass CheckpointEngine(ABC):\n    \"\"\"CheckpointEngine is an abstraction to transfer weights from trainer to rollout.\n\n    In trainer process:\n    >>> trainer = EngineRegistry.new(...) # FSDP, Megatron, VeOmini, TorchTitan, ...\n    >>> engine = CheckpointEngine.new(...) # NCCLCheckpointEngine, NIXLCheckpointEngine, ...\n    >>> await engine.send_weights(trainer.get_per_tensor_param())\n\n    In rollout process:\n    >>> engine = CheckpointEngine.new(...)\n    >>> server_adapter = ServerAdapter()\n    >>> await server_adapter.update_weights(engine.get_weights()) # update weights via cuda ipc\n    \"\"\"\n\n    @abstractmethod\n    def prepare(self) -> dict[str, Any]:\n        \"\"\"Prepare checkpoint engine before each step send_weights/receive_weights.\n\n        1. Allocate weight bucket.\n        2. [Optional] Register weight bucket for RDMA.\n        3. Return metadata to build communication topology: master ip:port, register RDMA description, etc.\n\n        Args:\n            worker_group: The worker group that the checkpoint engine will be used.\n\n        Returns:\n            A dictionary that contains the metadata of the worker group.\n        \"\"\"\n        raise NotImplementedError\n\n    @classmethod\n    @abstractmethod\n    def build_topology(\n        cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]\n    ) -> tuple[dict[str, list[Any]], dict[str, list[Any]]]:\n        \"\"\"Build communication topology between all workers.\n\n        Args:\n            trainer_world_size: The world size of the trainer worker group.\n            rollout_world_size: The world size of the rollout replica.\n            metadata: A list of metadata `prepare` from all workers.\n\n        Returns:\n            A tuple of two dictionaries that contains the communication topology for trainer and rollout worker group.\n            Each dict value should be a list argument equal to the world size of the worker group to dispatch to\n            `init_process_group`.\n\n            ```\n            world_size = rollout.world_size + trainer.world_size\n            kwargs = {\n                \"rank\": list(range(world_size)),\n                \"world_size\": [world_size] * world_size,\n                \"master_metadata\": [metadata[0]] * world_size,\n            }\n            ```\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    def init_process_group(self, **kwargs):\n        \"\"\"Init process group for checkpoint engine.\n\n        Args:\n            **kwargs: Keyword arguments from `build_topology`.\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    def finalize(self):\n        \"\"\"Finalize checkpoint engine after each step send_weights/receive_weights.\n\n        1. Free weight bucket.\n        1. [Optional] Deregister weight bucket for RDMA.\n        2. [Optional] Destroy process group.\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        raise NotImplementedError\n\n    @abstractmethod\n    async def receive_weights(self) -> Generator[tuple[str, torch.Tensor], None, None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        raise NotImplementedError\n\n\nclass CheckpointEngineWithCache(CheckpointEngine):\n    \"\"\"Checkpoint engine with local cache: shm, disk, etc. This allow to synchronize weights without interrupting\n    rollout ongoing requests (partial rollout). After requests exhausted, rollout can get weights from local cache.\n\n    Laminar: https://arxiv.org/abs/2510.12633\n    \"\"\"\n\n    @abstractmethod\n    async def get_weights(self) -> Generator[tuple[str, torch.Tensor], None, None]:\n        \"\"\"Get the weights of the model from local cache.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        raise NotImplementedError\n\n\n@CheckpointEngineRegistry.register(\"naive\")\nclass ColocatedCheckpointEngine(CheckpointEngine):\n    \"\"\"Checkpoint engine for trainer and rollout colocated on same GPU.\n\n    In trainer process:\n    >>> engine = ColocatedCheckpointEngine()\n    >>> trainer = Trainer()\n    >>> server_adapter = ServerAdapter()\n    >>> engine.send_weights(trainer.get_per_tensor_param())\n    >>> server_adapter.update_weights(engine.receive_weights())\n    \"\"\"\n\n    def __init__(self, bucket_size: int, is_master: bool = False) -> None:\n        self.bucket_size = bucket_size\n        self.is_master = is_master\n\n    def prepare(self):\n        raise NotImplementedError\n\n    def init_process_group(self, **kwargs):\n        raise NotImplementedError\n\n    def finalize(self):\n        raise NotImplementedError\n\n    @classmethod\n    def build_topology(cls, *args, **kwargs):\n        raise NotImplementedError\n\n    def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        self.weights = weights\n\n    def receive_weights(self) -> Generator[tuple[str, torch.Tensor], None, None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        yield from self.weights\n        self.weights = None\n\n\nclass CheckpointEngineWorker(Worker):\n    \"\"\"CheckpointEngineWorker colocated with inference engine's WorkerProc on same GPU.\n\n    Args:\n        rollout_config: The rollout configuration.\n        model_config: The model configuration.\n        server_adapter: The server adapter to update weights.\n    \"\"\"\n\n    def __init__(\n        self,\n        rollout_config: RolloutConfig,\n        model_config: HFModelConfig,\n        server_adapter: BaseRollout = None,\n        *args,\n        **kwargs,\n    ) -> None:\n        super().__init__()\n        self.rollout_config = rollout_config\n        self.model_config = model_config\n\n        self.server_adapter: BaseRollout = server_adapter\n        backend = self.rollout_config.checkpoint_engine.backend\n        bucket_size = self.rollout_config.checkpoint_engine.update_weights_bucket_megabytes << 20\n        engine_kwargs = self.rollout_config.checkpoint_engine.engine_kwargs.get(backend, {})\n        self.checkpoint_engine: CheckpointEngine = CheckpointEngineRegistry.new(\n            backend, bucket_size=bucket_size, **engine_kwargs\n        )\n        self.extra_rollout_args = args\n        self.extra_rollout_kwargs = kwargs\n        if self.server_adapter is None:\n            self.server_adapter = get_rollout_class(self.rollout_config.name, self.rollout_config.mode)(\n                *self.extra_rollout_args,\n                config=self.rollout_config,\n                model_config=self.model_config,\n                device_mesh=None,\n                **self.extra_rollout_kwargs,\n            )\n        # sglang and trt-llm need device_mesh for internal communication\n        initialize_global_process_group_ray(timeout_second=None, backend=\"cpu:gloo\")\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    async def update_weights(self, global_steps: int = None):\n        weights = self.checkpoint_engine.receive_weights()\n        await self.server_adapter.update_weights(weights, global_steps=global_steps)\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False)\n    def execute_checkpoint_engine(self, method: str, *args, **kwargs):\n        return getattr(self.checkpoint_engine, method)(*args, **kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def get_replica_rank(self) -> int:\n        \"\"\"Get replica rank from the underlying rollout server adapter.\"\"\"\n        return self.server_adapter.replica_rank\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def is_leader_rank(self) -> bool:\n        \"\"\"Get leader rank flag from the underlying rollout server adapter.\"\"\"\n        return self.server_adapter.is_leader_rank\n\n\n_worker_cls = ray.remote(CheckpointEngineWorker)\n\n\nclass CheckpointEngineManager:\n    \"\"\"Checkpoint engine manager to coordinate weight synchronization between trainer and rollout replicas.\n\n    - ME: model engine, FSDP, MCore, VeOmni, export full tensor generator `get_per_tensor_param`\n    - CE: checkpoint engine, NCCL, NIXL, etc\n\n    In trainer, model engine and checkpoint engine are in same process.\n    In rollout, checkpoint engine and rollout worker are in separate process, update weights via cuda ipc.\n\n    ```\n    ┌────────┬────────┬─────┬────────┐         ┌───────────────────┬───────────────────┐\n    │ ┌────┐ │ ┌────┐ │     │ ┌────┐ │         │     Replica 0     │     Replica 1     │\n    │ │ ME0│ │ │ ME1│ │     │ │ MEn│ │         ├────┬────┬────┬────┼────┬────┬────┬────┤\n    │ └──┬─┘ │ └────┘ │ ... │ └────┘ │         │ 0  │ 1  │ 2  │ 3  │ 0  │ 1  │ 2  │ 3  │\n    │    v   |        |     |        |         └──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┘\n    | ┌──┴─┐ │ ┌────┐ │     │ ┌────┐ │            ^    ^    ^   cuda ipc   ^    ^    ^\n    │ │ CE │ │ │ CE │ │     │ │ CE │ │         ┌──┴─┬──┴─┬──┴─┬──┴─┬──┴─┬──┴─┬──┴─┬──┴─┐\n    │ └──┬─┘ │ └────┘ │     │ └────┘ │         │ CE │ CE │ CE │ CE │ CE │ CE │ CE │ CE |\n    └────┼───┴────────┴─────┴────────┘         └──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┴──┬─┘\n         v                                        |    |    |    |    |    |    |    |\n         └─────────────(nccl/nixl/..)─────────────┴────┴────┴────┴────┴────┴────┴────┘\n    ```\n\n    Args:\n        config: The checkpoint engine config.\n        trainer: The trainer worker group.\n        replicas: The list of rollout replicas.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: CheckpointEngineConfig,\n        trainer: RayWorkerGroup,\n        replicas: list[RolloutReplica],\n    ) -> None:\n        self.config = config\n        self.backend = config.backend\n        self.backend_cls = CheckpointEngineRegistry.get(config.backend)\n        self.trainer = trainer\n        self.replicas = replicas\n\n    def build_process_group(self, rollout: RayWorkerGroup):\n        \"\"\"Build process group for trainer and rollout replicas.\"\"\"\n        trainer = self.trainer\n\n        # 1. prepare all workers\n        metadata = ray.get(\n            trainer.execute_checkpoint_engine([\"prepare\"] * trainer.world_size)\n            + rollout.execute_checkpoint_engine([\"prepare\"] * rollout.world_size)\n        )\n\n        # 2. build communication topology between all workers\n        trainer_kwargs, rollout_kwargs = self.backend_cls.build_topology(\n            trainer.world_size, rollout.world_size, metadata\n        )\n        for k, v in trainer_kwargs.items():\n            assert len(v) == trainer.world_size, f\"trainer_kwargs[{k}] must have length of {trainer.world_size}\"\n        for k, v in rollout_kwargs.items():\n            assert len(v) == rollout.world_size, f\"rollout_kwargs[{k}] must have length of {rollout.world_size}\"\n\n        trainer_kwargs[\"method\"] = [\"init_process_group\"] * trainer.world_size\n        rollout_kwargs[\"method\"] = [\"init_process_group\"] * rollout.world_size\n\n        # 3. init process group between all workers\n        ray.get(\n            trainer.execute_checkpoint_engine(**trainer_kwargs) + rollout.execute_checkpoint_engine(**rollout_kwargs)\n        )\n\n    def add_replicas(self, replicas: list[RolloutReplica]):\n        \"\"\"Add rollout replicas to the manager for elastic scale up, will rebuild process group.\n\n        Args:\n            replicas: The list of rollout replicas to add.\n        \"\"\"\n        self.replicas.extend(replicas)\n\n    def remove_replicas(self, replicas: list[RolloutReplica]):\n        \"\"\"Remove rollout replicas from the manager for elastic scale down, will rebuild process group.\n\n        Args:\n            replicas: The list of rollout replicas to remove.\n        \"\"\"\n        replicas_set = set(replicas)\n        self.replicas = [r for r in self.replicas if r not in replicas_set]\n\n    @auto_await\n    async def sleep_replicas(self):\n        \"\"\"Sleep all rollout replicas: free weight and kv_cache device memory.\"\"\"\n        await asyncio.gather(*[r.sleep() for r in self.replicas])\n\n    @auto_await\n    async def wake_up_replicas(self):\n        \"\"\"Resume all rollout replicas: recover kv_cache and weights device memory.\"\"\"\n        await asyncio.gather(*[r.wake_up() for r in self.replicas])\n\n    @auto_await\n    async def update_weights(self, global_steps: int = None):\n        \"\"\"Update weights from trainer to rollout replicas.\n\n        Args:\n            global_steps: The global steps of the trainer.\n        \"\"\"\n\n        # 0. update weights for sync training with colocated trainer and rollout\n        if self.backend == \"naive\":\n            ray.get(self.trainer.update_weights(global_steps=global_steps))\n            return\n\n        # 1. abort and save all unfinished requests for partial rollout\n        await asyncio.gather(*[r.abort_all_requests() for r in self.replicas])\n\n        # 2. create a temporay worker group for all replicas\n        workers = []\n        for replica in self.replicas:\n            workers.extend(replica.workers)\n        rollout = RayWorkerGroup(worker_handles=workers, ray_cls_with_init=RayClassWithInitArgs(cls=_worker_cls))\n        trainer = self.trainer\n\n        # 3. sleep replicas to free kv_cache before weight sync (if free_cache_engine is enabled)\n        await self.sleep_replicas()\n\n        # 4. build process group\n        self.build_process_group(rollout)\n\n        # 5. update weights of all workers\n        ray.get(trainer.update_weights(global_steps=global_steps) + rollout.update_weights(global_steps=global_steps))\n\n        # 6. finalize all workers\n        ray.get(\n            trainer.execute_checkpoint_engine([\"finalize\"] * trainer.world_size)\n            + rollout.execute_checkpoint_engine([\"finalize\"] * rollout.world_size)\n        )\n\n        # 7. resume replicas to recover kv_cache (for free_cache_engine scenarios)\n        await self.wake_up_replicas()\n\n        # 8. resume all unfinished requests for partial rollout\n        await asyncio.gather(*[r.resume_generation() for r in self.replicas])\n"
  },
  {
    "path": "verl/checkpoint_engine/hccl_checkpoint_engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\nimport time\nfrom dataclasses import dataclass\nfrom typing import AsyncGenerator, Generator\n\nimport ray\nimport torch\nimport zmq\nfrom vllm.distributed.utils import StatelessProcessGroup\n\nfrom verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry, TensorMeta\nfrom verl.utils.device import is_torch_npu_available\nfrom verl.utils.distributed import stateless_init_process_group\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\n\nif not is_torch_npu_available(check_device=False):\n    raise ImportError(\"HCCLCheckpointEngine is unavailable because the torch.npu module is not available.\")\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@dataclass\nclass MasterMetadata:\n    zmq_ip: str\n    zmq_port: int\n    dist_ip: str\n    dist_port: int\n\n\nclass BroadcastOperation:\n    \"\"\"Async broadcast operation with HCCL in separate thread.\n\n    Args:\n        rank (int): The rank of the current process.\n        group_name (str): The name of the HCCL process group.\n        bucket (torch.Tensor): The tensor to broadcast.\n        metadata (dict[str, TensorMeta]): The metadata of the tensor.\n        socket (zmq.Socket): The zeromq socket to communicate with master.\n        topic (str): The topic to subscribe.\n    \"\"\"\n\n    def __init__(\n        self,\n        rank: int,\n        process_group: StatelessProcessGroup | str,\n        bucket: torch.Tensor,\n        metadata: dict[str, TensorMeta],\n        socket: zmq.Socket,\n        topic: str,\n    ) -> None:\n        self.rank = rank\n        self.pyhccl = process_group\n        self.bucket = bucket\n        self.metadata = metadata\n        self.socket = socket\n        self.topic = topic\n\n        self._run()\n\n    def _run(self):\n        # broadcast tensor meta via zeromq PUB/SUB\n        if self.rank == 0:\n            self.socket.send_string(self.topic, flags=zmq.SNDMORE)\n            self.socket.send_pyobj(self.metadata)\n        else:\n            self.socket.recv_string()\n            self.metadata = self.socket.recv_pyobj()\n\n        # broadcast tensor via HCCL\n        self.pyhccl.broadcast(self.bucket, src=0)\n\n    async def wait_for_complete(self) -> dict[str, TensorMeta]:\n        \"\"\"Wait for the broadcast operation to complete.\n\n        Returns:\n            dict[str, TensorMeta]: The bucket meta after broadcast.\n        \"\"\"\n        return self.metadata\n\n\n@CheckpointEngineRegistry.register(\"nccl\")\nclass HCCLCheckpointEngine(CheckpointEngine):\n    \"\"\"HCCL checkpoint engine with collective communication.\n\n    Args:\n        bucket_size (int): Bucket size in bytes to transfer multiple weights at one time. Note that we use\n            two buffer to send and recv weights at same time, so the device memory overhead is 2 * bucket_size.\n        group_name (str): The name of the HCCL process group. Defaults to \"default\".\n        rebuild_group (bool): Whether to rebuild the HCCL process group in each update. Defaults to False.\n        is_master (bool): Whether the current process is the master process. Defaults to False.\n        rollout_dtype (torch.dtype): The dtype of the weights received from rollout workers. Defaults to torch.bfloat16.\n    \"\"\"\n\n    def __init__(\n        self,\n        bucket_size: int,\n        group_name: str = \"default\",\n        rebuild_group: bool = False,\n        is_master: bool = False,\n        rollout_dtype: torch.dtype = torch.bfloat16,\n    ) -> None:\n        self.bucket_size = bucket_size\n        self.group_name = group_name\n        self.rebuild_group = rebuild_group\n        self.rollout_dtype = rollout_dtype\n        self.pyhccl = None\n        self.device = torch.npu.current_device()\n\n        # start zeromq server for broadcasting bucket tensor metadata\n        self.is_master = is_master\n        self.topic = \"bucket_metadata\"\n        if self.is_master:\n            self._start_zmq_server()\n            self.dist_port, _ = get_free_port(self.ip)\n\n    def prepare(self) -> MasterMetadata:\n        self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=\"npu\")\n        self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=\"npu\")\n\n        return (\n            MasterMetadata(zmq_ip=self.ip, zmq_port=self.zmq_port, dist_ip=self.ip, dist_port=self.dist_port)\n            if self.is_master\n            else None\n        )\n\n    def finalize(self):\n        \"\"\"Destroy the HCCL process group if rebuild_group is True.\"\"\"\n        if self.rebuild_group:\n            if self.rank >= 0:\n                self.pyhccl.destroyComm(self.pyhccl.comm)\n                self.pyhccl = None\n            self.rank = None\n            self.world_size = None\n\n        self.send_buf = None\n        self.recv_buf = None\n        torch.npu.empty_cache()\n\n    @classmethod\n    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]):\n        trainer_kwargs = {\n            \"rank\": [0] + [-1] * (trainer_world_size - 1),\n            \"world_size\": [rollout_world_size + 1] * trainer_world_size,\n            \"master_metadata\": [metadata[0]] * trainer_world_size,\n        }\n        rollout_kwargs = {\n            \"rank\": list(range(1, rollout_world_size + 1)),\n            \"world_size\": [rollout_world_size + 1] * rollout_world_size,\n            \"master_metadata\": [metadata[0]] * rollout_world_size,\n        }\n        return trainer_kwargs, rollout_kwargs\n\n    def _start_zmq_server(self):\n        self.ip = ray.util.get_node_ip_address().strip(\"[]\")\n        self.zmq_port, _ = get_free_port(self.ip)\n\n        context = zmq.Context()\n        self.socket = context.socket(zmq.PUB)\n        if is_valid_ipv6_address(self.ip):\n            address = f\"tcp://[{self.ip}]:{self.zmq_port}\"\n            self.socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{self.ip}:{self.zmq_port}\"\n\n        self.socket.bind(address)\n\n    def _connect_zmq_client(self, metadata: MasterMetadata):\n        assert not self.is_master, \"Master process should not connect to other processes.\"\n        context = zmq.Context()\n        self.socket = context.socket(zmq.SUB)\n        if is_valid_ipv6_address(metadata.zmq_ip):\n            address = f\"tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}\"\n            self.socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{metadata.zmq_ip}:{metadata.zmq_port}\"\n\n        self.socket.connect(address)\n        self.socket.setsockopt_string(zmq.SUBSCRIBE, self.topic)\n\n    def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata):\n        \"\"\"Initialize the HCCL process group.\n\n        Args:\n            rank (int): The rank of the current process.\n            world_size (int): The total number of processes.\n        \"\"\"\n        # For trainer workers other than rank 0, their rank should be -1.\n        if rank < 0:\n            self.rank = rank\n            self.world_size = world_size\n            return\n\n        if self.rebuild_group or self.pyhccl is None:\n            self.pyhccl = stateless_init_process_group(\n                master_metadata.dist_ip, master_metadata.dist_port, rank, world_size, self.device\n            )\n            self.rank = rank\n            self.world_size = world_size\n        else:\n            assert self.rank == rank, f\"rank {rank} is not equal to self.rank {self.rank}\"\n            assert self.world_size == world_size, (\n                f\"world_size {world_size} is not equal to self.world_size {self.world_size}\"\n            )\n\n        if self.rank > 0:\n            self._connect_zmq_client(master_metadata)\n\n        # barrier\n        signal = torch.tensor([1], dtype=torch.int8, device=torch.npu.current_device())\n        self.pyhccl.all_reduce(signal)\n\n        logger.info(f\"init_process_group rank: {self.rank}, world_size: {self.world_size}\")\n\n    @torch.no_grad()\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.rank <= 0, \"Trainer workers other than rank 0 should not send weights.\"\n\n        # For trainer rank other than 0, consume weights without sending.\n        if self.rank < 0:\n            for name, weight in weights:\n                pass\n            return\n\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        broadcast_op = None\n\n        start_time = time.time()\n        bucket_meta: dict[str, TensorMeta] = {}\n        offset = 0\n        for name, weight in weights:\n            # fill the tensor bucket\n            if offset + weight.nbytes > self.bucket_size:\n                torch.npu.synchronize()\n\n                # wait previous broadcast op finish\n                if broadcast_op is not None:\n                    await broadcast_op.wait_for_complete()\n\n                broadcast_op = BroadcastOperation(\n                    rank=self.rank,\n                    process_group=self.pyhccl,\n                    bucket=send_buf,\n                    metadata={\"bucket_meta\": bucket_meta, \"is_last\": False},\n                    socket=self.socket,\n                    topic=self.topic,\n                )\n\n                # swap send_buf and recv_buf\n                send_buf, recv_buf = recv_buf, send_buf\n                bucket_meta = {}\n                offset = 0\n\n            assert offset + weight.nbytes <= self.bucket_size, (\n                f\"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket.\"\n            )\n\n            bucket_meta[name] = {\n                \"name\": name,\n                \"shape\": weight.shape,\n                \"dtype\": weight.dtype,\n                \"offset\": offset,\n            }\n            send_buf[offset : offset + weight.nbytes] = weight.view(-1).view(torch.uint8)\n            offset += weight.nbytes\n\n        # broadcast last bucket\n        torch.npu.synchronize()\n        if broadcast_op is not None:\n            await broadcast_op.wait_for_complete()\n\n        broadcast_op = BroadcastOperation(\n            rank=self.rank,\n            process_group=self.pyhccl,\n            bucket=send_buf,\n            metadata={\"bucket_meta\": bucket_meta, \"is_last\": True},\n            socket=self.socket,\n            topic=self.topic,\n        )\n        await broadcast_op.wait_for_complete()\n        logger.info(f\"Rank {self.rank} send weights done, time cost: {time.time() - start_time:.2f}s\")\n\n    @torch.no_grad()\n    async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.rank > 0, \"Rank 0 should not receive weights.\"\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        total_bytes, total_params = 0, 0\n\n        # receive first bucket\n        start_time = time.time()\n        broadcast_op = BroadcastOperation(\n            rank=self.rank,\n            process_group=self.pyhccl,\n            bucket=recv_buf,\n            metadata=None,\n            socket=self.socket,\n            topic=self.topic,\n        )\n        metadata = await broadcast_op.wait_for_complete()\n        total_bytes += self.bucket_size\n        total_params += len(metadata[\"bucket_meta\"])\n\n        # swap send_buf and recv_buf\n        send_buf, recv_buf = recv_buf, send_buf\n        while not metadata[\"is_last\"]:\n            # 1. receive next bucket\n            broadcast_op = BroadcastOperation(\n                rank=self.rank,\n                process_group=self.pyhccl,\n                bucket=recv_buf,\n                metadata=None,\n                socket=self.socket,\n                topic=self.topic,\n            )\n\n            # 2. yield tensor from send_buf\n            for name, meta in metadata[\"bucket_meta\"].items():\n                dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n                size = dtype.itemsize * shape.numel()\n                tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n                yield name, tensor\n\n            # 3. wait for next bucket broadcast finish\n            metadata = await broadcast_op.wait_for_complete()\n            total_bytes += self.bucket_size\n            total_params += len(metadata[\"bucket_meta\"])\n\n            # 4. swap send_buf and recv_buf\n            torch.npu.synchronize()  # sync non-blocking copy\n            send_buf, recv_buf = recv_buf, send_buf\n\n        # yield tensor from send_buf\n        for name, meta in metadata[\"bucket_meta\"].items():\n            dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n            size = dtype.itemsize * shape.numel()\n            tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n            yield name, tensor\n\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} receive weights done, total_params: {total_params}, \"\n            f\"time cost: {time_cost:.2f}s, bandwidth: {bandwidth:.2f} GB/s\"\n        )\n"
  },
  {
    "path": "verl/checkpoint_engine/kimi_checkpoint_engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport concurrent.futures\nimport logging\nimport os\nimport time\nimport types\nfrom collections import defaultdict\nfrom dataclasses import dataclass\nfrom typing import AsyncGenerator, Generator\n\nimport checkpoint_engine.distributed as dist\nimport ray\nimport torch\nfrom checkpoint_engine.ps import H2DBucket, ParameterMeta, ParameterServer, _gen_h2d_buckets, _to_named_tensor\n\nfrom verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry\nfrom verl.utils.device import get_nccl_backend, get_torch_device\nfrom verl.utils.net_utils import get_free_port\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef ckpt_get_named_tensor_buckets(\n    iterable: Generator[tuple[str, torch.Tensor], None, None],\n    bucket_bytes: int,\n    world_size: int,\n    rank_id: int,\n    rollout_dtype: torch.dtype = torch.bfloat16,\n) -> dict[str, torch.Tensor]:\n    if bucket_bytes <= 0:\n        raise ValueError(f\"bucket_bytes must be greater than 0, got {bucket_bytes}\")\n\n    current_bucket = {}\n    current_size = 0\n    for tensor_idx, (name, tensor) in enumerate(iterable):\n        tensor = tensor.to(rollout_dtype)\n        if tensor_idx % world_size == rank_id:\n            tensor_size = tensor.element_size() * tensor.numel()\n            if current_size + tensor_size > bucket_bytes:\n                if current_bucket:\n                    yield current_bucket\n                    current_bucket = {}\n                    current_size = 0\n\n            current_bucket[name] = tensor\n            current_size += tensor_size\n\n    if current_bucket:\n        yield current_bucket\n\n\nasync def receive_tensor(\n    self,\n    checkpoint_name: str,\n    ranks_group: int,\n    ranks: list[int] | None = None,\n    bucket_size: int = 2 << 30,\n    disable_h2d_buffer: bool = False,\n) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n    assert len(self._current_global_parameter_metas) != 0, \"parameter metas is empty\"\n    assert dist.is_initialized(), \"process group is not initialized\"\n    assert self._p2p_store is not None, \"p2p store is not initialized\"\n    assert ranks, \"ranks should be set\"\n\n    # first execute a barrier to avoid subsequent device oom\n    dist.barrier(group=ranks_group)\n    buckets = _gen_h2d_buckets(\n        self._current_global_parameter_metas,\n        bucket_size,\n        self._local_rdma_devices,\n        self._remote_rdma_devices,\n        ranks,\n    )\n    h2d_buffer: torch.Tensor | None = (\n        None\n        if disable_h2d_buffer\n        else torch.empty(bucket_size, dtype=torch.uint8, device=self.device_manager.device_type)\n    )\n    # p2p store need to register h2d_buffer to let other ranks read\n    if ranks:\n        h2d_buffer_name = \"__h2d_buffer__\"\n        if h2d_buffer is not None and self._p2p_store is not None:\n            self._p2p_store.register_named_tensors({h2d_buffer_name: h2d_buffer})\n    receiver_rank_buckets: list[tuple[int, H2DBucket]] = []\n    for receiver_rank, owner_rank, bucket in buckets:\n        if receiver_rank != self._rank:\n            continue\n        receiver_rank_buckets.append((owner_rank, bucket))\n    buffer = torch.empty(bucket_size * 2, dtype=torch.uint8, device=self.device_manager.device_type)\n    buckets_by_receiver_rank: dict[int, list[H2DBucket]] = defaultdict(list)\n\n    max_len = 0\n    for receiver_rank, _, bucket in buckets:\n        buckets_by_receiver_rank[receiver_rank].append(bucket)\n        if len(buckets_by_receiver_rank[receiver_rank]) > max_len:\n            max_len = len(buckets_by_receiver_rank[receiver_rank])\n    gidx = 0\n    metadata: list[ParameterMeta]\n    try:\n        for i in range(max_len):\n            if i < len(receiver_rank_buckets) and not disable_h2d_buffer:\n                self._copy_to_buffer(\n                    checkpoint_name,\n                    receiver_rank_buckets[i][1],\n                    h2d_buffer,\n                    receiver_rank_buckets[i][0] if ranks else None,\n                )\n            for receiver_rank, _buckets in buckets_by_receiver_rank.items():\n                if i >= len(_buckets):\n                    continue\n                bucket = _buckets[i]\n                start = gidx % 2 * bucket_size\n                buffer_b: torch.Tensor = buffer[start : start + bucket.size]\n                if receiver_rank == self._rank:\n                    if disable_h2d_buffer:\n                        self._copy_to_buffer(checkpoint_name, bucket, buffer_b)\n                    else:\n                        buffer_b.data.copy_(h2d_buffer[: bucket.size])\n                broadcast_op = BroadcastOperation(\n                    rank=receiver_rank,\n                    ranks_group=ranks_group,\n                    bucket=buffer_b,\n                    metadata=bucket.items,\n                )\n                if gidx == 0:\n                    metadata = await broadcast_op.wait_for_complete()\n                    gidx += 1\n                    continue\n                meta_list = _to_named_tensor(metadata, (gidx - 1) % 2 * bucket_size)\n                for item in meta_list:\n                    shape = item[\"shape\"]\n                    if isinstance(shape, list | tuple):\n                        shape = torch.Size(shape)\n                    assert isinstance(shape, torch.Size)\n                    dtype, offset = item[\"dtype\"], item[\"offset\"]\n                    size = dtype.itemsize * shape.numel()\n                    tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape)\n                    yield item[\"name\"], tensor\n                metadata = await broadcast_op.wait_for_complete()\n                self.device_manager.device_module.synchronize()\n                gidx += 1\n\n        meta_list = _to_named_tensor(metadata, (gidx - 1) % 2 * bucket_size)\n        for item in meta_list:\n            shape = item[\"shape\"]\n            if isinstance(shape, list | tuple):\n                shape = torch.Size(shape)\n            assert isinstance(shape, torch.Size)\n            dtype, offset = item[\"dtype\"], item[\"offset\"]\n            size = dtype.itemsize * shape.numel()\n            tensor = buffer[offset : offset + size].view(dtype=dtype).view(shape)\n            yield item[\"name\"], tensor\n\n    finally:\n        dist.barrier(group=ranks_group)\n        if ranks and h2d_buffer is not None:\n            self._p2p_store.unregister_named_tensors([h2d_buffer_name])\n        self.device_manager.device_module.empty_cache()\n\n\n@dataclass\nclass MasterMetadata:\n    zmq_ip: str\n    zmq_port: int\n    dist_ip: str\n    dist_port: int\n\n\nclass BroadcastOperation:\n    \"\"\"Async broadcast operation in separate thread.\n\n    Args:\n        rank (int): The rank of the current process.\n        ranks_group (int): The process group's value.\n        bucket (torch.Tensor): The tensor to broadcast.\n        metadata (list[ParameterMeta]): The metadata of the tensor.\n    \"\"\"\n\n    def __init__(\n        self,\n        rank: int,\n        ranks_group: int,\n        bucket: torch.Tensor,\n        metadata: list[ParameterMeta],\n    ) -> None:\n        self.rank = rank\n        self.ranks_group = ranks_group\n        self.bucket = bucket\n        self.metadata = metadata\n\n        loop = asyncio.get_running_loop()\n        self._task = loop.run_in_executor(None, self._run)\n\n    def _run(self):\n        # broadcast tensor\n        dist.broadcast(self.bucket, src=self.rank, group=self.ranks_group)\n\n    async def wait_for_complete(self) -> list[ParameterMeta]:\n        \"\"\"Wait for the broadcast operation to complete.\n\n        Returns:\n            list[ParameterMeta]: The bucket meta after broadcast.\n        \"\"\"\n        await self._task\n        return self.metadata\n\n\n@CheckpointEngineRegistry.register(\"kimi_ckpt_engine\")\nclass KIMICheckpointEngine(CheckpointEngine):\n    \"\"\"kimi checkpoint engine with collective communication.\n\n    Args:\n        bucket_size (int): Bucket size in bytes to transfer multiple weights at one time. Note that we use\n            two buffer to send and recv weights at same time, so the device memory overhead is 2 * bucket_size.\n        rebuild_group (bool): Whether to rebuild the process group in each update. Defaults to False.\n        is_master (bool): Whether the current process is the master process. Defaults to False.\n        rollout_dtype (torch.dtype): The dtype of the weights received from rollout workers. Defaults to torch.bfloat16.\n    \"\"\"\n\n    def __init__(\n        self,\n        bucket_size: int,\n        rebuild_group: bool = False,\n        is_master: bool = False,\n        rollout_dtype: torch.dtype = torch.bfloat16,\n    ) -> None:\n        self.bucket_size = bucket_size\n        self.rebuild_group = rebuild_group\n        self.rollout_dtype = rollout_dtype\n        self.is_master = is_master\n        self.initialized = False\n        self.checkpoint_name = \"kimi_checkpoint_engine\"\n\n    def prepare(self) -> MasterMetadata:\n        if self.is_master:\n            self.ip = ray.util.get_node_ip_address().strip(\"[]\")\n            self.listen_port, _ = get_free_port(self.ip)\n\n        return (\n            MasterMetadata(zmq_ip=None, zmq_port=None, dist_ip=self.ip, dist_port=self.listen_port)\n            if self.is_master\n            else None\n        )\n\n    def finalize(self):\n        \"\"\"Destroy the ckpt engine process group if rebuild_group is True.\"\"\"\n        if self.rebuild_group:\n            dist.destroy_process_group()\n            self.rank = None\n            self.world_size = None\n            self.initialized = False\n\n    @classmethod\n    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]):\n        trainer_kwargs = {\n            \"method\": [\"init_process_group\"] * trainer_world_size,\n            \"rank\": list(range(0, trainer_world_size)),\n            \"trainer_world_size\": [trainer_world_size] * trainer_world_size,\n            \"rollout_world_size\": [rollout_world_size] * trainer_world_size,\n            \"master_metadata\": [metadata[0]] * trainer_world_size,\n        }\n        rollout_kwargs = {\n            \"method\": [\"init_process_group\"] * rollout_world_size,\n            \"rank\": list(range(trainer_world_size, trainer_world_size + rollout_world_size)),\n            \"trainer_world_size\": [trainer_world_size] * rollout_world_size,\n            \"rollout_world_size\": [rollout_world_size] * rollout_world_size,\n            \"master_metadata\": [metadata[0]] * rollout_world_size,\n        }\n        return trainer_kwargs, rollout_kwargs\n\n    def init_process_group(\n        self,\n        rank: int,\n        trainer_world_size: int,\n        rollout_world_size: int,\n        master_metadata: MasterMetadata,\n    ):\n        \"\"\"Initialize the ckpt engine process group.\n\n        Args:\n            rank (int): The rank of the current process.\n            world_size (int): The total number of processes.\n        \"\"\"\n        self.rank = rank\n        self.trainer_world_size = trainer_world_size\n        self.rollout_world_size = rollout_world_size\n        self.world_size = trainer_world_size + rollout_world_size\n\n        if not self.initialized:\n            self.parameter_server = ParameterServer(\n                rank=rank,\n                world_size=self.world_size,\n                auto_pg=False,\n                master_addr=master_metadata.dist_ip,\n                master_port=master_metadata.dist_port,\n            )\n            self.parameter_server.receive_tensor = types.MethodType(receive_tensor, self.parameter_server)\n\n            dist.use_backend(f\"vllm_{get_nccl_backend()}\")\n            self.parameter_server.init_process_group()\n\n            self.rollout_ranks = list(range(self.trainer_world_size, self.world_size))\n            self.rollout_group = dist.new_group(self.rollout_ranks)\n            self.initialized = True\n\n    @torch.no_grad()\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n\n        def offload_cpu(name: str, tensor: torch.Tensor) -> tuple[str, torch.Tensor]:\n            return name, tensor.to(\"cpu\", non_blocking=True)\n\n        start_time = time.time()\n        named_tensors = {}\n        for named_tensors_gpu in ckpt_get_named_tensor_buckets(\n            weights, self.bucket_size, self.trainer_world_size, self.rank, self.rollout_dtype\n        ):\n            with concurrent.futures.ThreadPoolExecutor(max_workers=32) as executor:\n                futures = [\n                    executor.submit(\n                        offload_cpu,\n                        name,\n                        tensor,\n                    )\n                    for name, tensor in named_tensors_gpu.items()\n                ]\n            for future in concurrent.futures.as_completed(futures):\n                name, tensor_cpu = future.result()\n                named_tensors[name] = tensor_cpu\n\n        get_torch_device().synchronize()\n\n        self.parameter_server.register_checkpoint(self.checkpoint_name, named_tensors=named_tensors)\n        named_tensors = {}\n        get_torch_device().empty_cache()\n        logger.info(f\"Rank {self.rank} offload and register, time cost: {time.time() - start_time:.2f}s\")\n\n        self.parameter_server.gather_metas(self.checkpoint_name)\n        dist.barrier()\n        self.parameter_server.unregister_checkpoint(self.checkpoint_name)\n        logger.info(f\"Rank {self.rank} send weights done, time cost: {time.time() - start_time:.2f}s\")\n\n    @torch.no_grad()\n    async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        self.parameter_server.gather_metas(self.checkpoint_name)\n\n        start_time = time.time()\n        total_bytes, total_params = 0, 0\n        async for name, tensor in self.parameter_server.receive_tensor(\n            self.checkpoint_name, self.rollout_group, self.rollout_ranks, self.bucket_size\n        ):\n            total_bytes += tensor.element_size() * tensor.nelement()\n            total_params += 1\n            yield name, tensor\n        dist.barrier()\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} receive weights done, total_params: {total_params}, \"\n            f\"time cost: {time_cost:.2f}s, bandwidth: {bandwidth:.2f} GB/s\"\n        )\n"
  },
  {
    "path": "verl/checkpoint_engine/mooncake_checkpoint_engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport gc\nimport logging\nimport os\nimport time\nfrom typing import Any, AsyncGenerator, Generator\n\nimport ray\nimport torch\nfrom mooncake.engine import TransferEngine\nfrom vllm.distributed.utils import StatelessProcessGroup\n\nfrom verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry, TensorMeta\nfrom verl.utils.device import get_torch_device\nfrom verl.utils.net_utils import get_free_port\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\n\n\n@CheckpointEngineRegistry.register(\"mooncake\")\nclass MooncakeCheckpointEngine(CheckpointEngine):\n    \"\"\"Mooncake checkpoint engine with p2p communication using TransferEngine\n\n    Args:\n        bucket_size (int): Bucket size in bytes to transfer multiple weights at one time.\n        device (str): The device to use for the checkpoint engine, \"cpu\" or \"cuda\".\n        rollout_dtype (torch.dtype): The dtype of the weights received from rollout workers.\n        device_name (str): Mooncake device name filter.\n    \"\"\"\n\n    def __init__(\n        self,\n        bucket_size: int,\n        device: str = \"cuda\",\n        rollout_dtype: torch.dtype = torch.bfloat16,\n        device_name: str = \"\",\n        is_master: bool = False,\n        rebuild_group: bool = False,\n    ):\n        self.bucket_size = bucket_size\n        self.device = device\n        self.rollout_dtype = rollout_dtype\n        self.is_master = is_master\n        self.rebuild_group = rebuild_group\n\n        rank = int(os.environ[\"RANK\"])\n        device_count = get_torch_device().device_count()\n        local_rank = rank % device_count\n        get_torch_device().set_device(local_rank)\n\n        self.engine = TransferEngine()\n        hostname = ray.util.get_node_ip_address().strip(\"[]\")\n        ret = self.engine.initialize(\n            hostname,\n            \"P2PHANDSHAKE\",\n            \"ascend_direct\" if self.device == \"npu\" else \"rdma\",\n            device_name,\n        )\n        assert ret == 0, f\"TransferEngine initialize failed ret={ret}\"\n\n        rpc_port = self.engine.get_rpc_port()\n        self.session_id = f\"{hostname}:{rpc_port}\"\n        self.hostname = hostname\n\n        self.buf = torch.empty(2 * self.bucket_size, dtype=torch.uint8, device=self.device)\n        self.magic_buf = torch.empty(4 * 1024, dtype=torch.uint8, device=self.device)\n        ret = self.engine.batch_register_memory(\n            [self.buf.data_ptr(), self.magic_buf.data_ptr()],\n            [2 * self.bucket_size, 4 * 1024],\n        )\n        assert ret == 0, f\"batch_register_memory failed ret={ret}\"\n        logger.info(f\"__init__ session_id={self.session_id}\")\n\n    def prepare(self) -> dict[str, Any]:\n        \"\"\"Prepare send and recv buckets\"\"\"\n        logger.info(\n            f\"prepare ptr={self.buf.data_ptr():#x} len={2 * self.bucket_size} \"\n            f\"magic_buf_ptr={self.magic_buf.data_ptr():#x}\"\n        )\n        port, _ = get_free_port(self.hostname)\n        return {\"addr\": self.hostname, \"port\": port}\n\n    @classmethod\n    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadatas: list[dict]):\n        trainer_kwargs = {\n            \"rank\": [0] + [-1] * (trainer_world_size - 1),\n            \"world_size\": [rollout_world_size + 1] * trainer_world_size,\n            \"metadata\": [metadatas[0]] * trainer_world_size,\n        }\n        rollout_kwargs = {\n            \"rank\": list(range(1, rollout_world_size + 1)),\n            \"world_size\": [rollout_world_size + 1] * rollout_world_size,\n            \"metadata\": [metadatas[0]] * rollout_world_size,\n        }\n        return trainer_kwargs, rollout_kwargs\n\n    def init_process_group(self, rank: int, world_size: int, metadata: dict[str, Any]):\n        self.rank = rank\n        self.world_size = world_size\n        if rank < 0:\n            logger.info(f\"init_process_group rank={rank}\")\n            return\n\n        self.store = StatelessProcessGroup.create(\n            host=metadata[\"addr\"],\n            port=metadata[\"port\"],\n            rank=rank,\n            world_size=world_size,\n        )\n\n        info = {\n            \"session_id\": self.session_id,\n            \"ptr\": self.buf.data_ptr(),\n        }\n\n        info_list = self.store.all_gather_obj(info)\n        self.buffer_info = None if rank == 0 else info_list[rank - 1]\n\n        logger.info(f\"init_process_group rank={rank} world_size={world_size} buffer_info={self.buffer_info}\")\n\n    def finalize(self):\n        \"\"\"Cleanup communication and deregister memory\"\"\"\n        self.store = None\n        get_torch_device().empty_cache()\n        gc.collect()\n        logger.info(f\"finalize rank={self.rank}\")\n\n    async def wait_for_complete(self, buf: torch.Tensor):\n        magic = torch.tensor([0xAB, 0xDC, 0xEF, 0x88], dtype=torch.uint8, device=self.device)\n        while True:\n            if torch.equal(buf[:4], magic):\n                break\n            await asyncio.sleep(0)\n\n    @torch.no_grad()\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send weights using Mooncake TransferEngine\"\"\"\n        if self.rank < 0:\n            for name, weight in weights:\n                pass\n            logger.info(f\"send_weights rank={self.rank}\")\n            return\n\n        total_bytes = 0\n        start_time = time.time()\n        bucket_meta: dict[str, TensorMeta] = {}\n        offset = 0\n        should_wait = False\n        bufs = [self.buf[: self.bucket_size], self.buf[self.bucket_size :]]\n        idx = 0\n        current = bufs[idx]\n\n        for name, weight in weights:\n            weight = weight.to(self.rollout_dtype)\n\n            if offset + weight.nbytes > self.bucket_size:\n                total_bytes += offset\n                get_torch_device().synchronize()\n                info = {\n                    \"bucket_meta\": bucket_meta,\n                    \"ptr\": current.data_ptr(),\n                    \"len\": offset,\n                    \"is_last\": False,\n                }\n                # send to rank 1\n                self.store.send_obj(info, 1)\n\n                idx ^= 1\n                current = bufs[idx]\n                bucket_meta = {}\n                offset = 0\n\n                if should_wait:\n                    await self.wait_for_complete(current)\n                should_wait = True\n\n            assert offset + weight.nbytes <= self.bucket_size, (\n                f\"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket.\"\n            )\n\n            bucket_meta[name] = {\n                \"name\": name,\n                \"shape\": weight.shape,\n                \"dtype\": weight.dtype,\n                \"offset\": offset,\n            }\n            current[offset : offset + weight.nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)\n            offset += weight.nbytes\n\n        get_torch_device().synchronize()\n        info = {\n            \"bucket_meta\": bucket_meta,\n            \"ptr\": current.data_ptr(),\n            \"len\": offset,\n            \"is_last\": True,\n        }\n        self.store.send_obj(info, 1)\n        await self.wait_for_complete(current)\n\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} send weights done, \"\n            f\"total bytes: {total_bytes} time cost: {time_cost:.2f}s bandwidth: {bandwidth:.2f} GB/s\"\n        )\n\n    @torch.no_grad()\n    async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n        \"\"\"Receive weights using Mooncake TransferEngine\"\"\"\n        start_time = time.time()\n        total_bytes = 0\n        bufs = [self.buf[: self.bucket_size], self.buf[self.bucket_size :]]\n        idx = 0\n        current = bufs[idx]\n        self.magic_buf[:4] = torch.tensor([0xAB, 0xDC, 0xEF, 0x88], dtype=torch.uint8, device=self.device)\n\n        while True:\n            # 1 receive info from previous rank\n            info = self.store.recv_obj(self.rank - 1)\n            if idx >= 2 and self.rank < self.world_size - 1:\n                await self.wait_for_complete(current)\n\n            ptr = info[\"ptr\"]\n            ret = self.engine.transfer_sync_read(\n                self.buffer_info[\"session_id\"],\n                current.data_ptr(),\n                ptr,\n                info[\"len\"],\n            )\n            assert ret == 0, f\"transfer_sync_read failed {ret}\"\n            total_bytes += info[\"len\"]\n\n            # 2 send info to next rank\n            info[\"ptr\"] = current.data_ptr()\n            if self.rank < self.world_size - 1:\n                self.store.send_obj(info, self.rank + 1)\n\n            # 3 yield tensor from current buffer\n            for name, meta in info[\"bucket_meta\"].items():\n                dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n                size = dtype.itemsize * shape.numel()\n                tensor = current[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n                yield name, tensor\n\n            # 4 write magic data to previous rank\n            ret = self.engine.transfer_sync_write(\n                self.buffer_info[\"session_id\"],\n                self.magic_buf.data_ptr(),\n                ptr,\n                4,\n            )\n            assert ret == 0, f\"transfer_sync_write failed {ret}\"\n\n            # 5 swap buffer\n            idx += 1\n            current = bufs[idx % 2]\n            get_torch_device().synchronize()\n\n            if info[\"is_last\"]:\n                break\n\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} receive weights done, time cost: {time_cost:.2f}s, bandwidth: {bandwidth:.2f} GB/s\"\n        )\n"
  },
  {
    "path": "verl/checkpoint_engine/nccl_checkpoint_engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nimport time\nfrom dataclasses import dataclass\nfrom typing import AsyncGenerator, Generator\nfrom unittest.mock import patch\n\nwith patch(\"importlib.metadata.distributions\", return_value=[]):\n    import cupy as cp\n\nimport ray\nimport ray.util.collective as collective\nimport torch\nimport zmq\n\nfrom verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry, TensorMeta\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@dataclass\nclass MasterMetadata:\n    zmq_ip: str\n    zmq_port: int\n\n\nclass BroadcastOperation:\n    \"\"\"Async broadcast operation with NCCL in separate thread.\n\n    Args:\n        rank (int): The rank of the current process.\n        group_name (str): The name of the NCCL process group.\n        bucket (cp.ndarray | torch.Tensor): The tensor to broadcast.\n        metadata (dict[str, TensorMeta]): The metadata of the tensor.\n        socket (zmq.Socket): The zeromq socket to communicate with master.\n        topic (str): The topic to subscribe.\n    \"\"\"\n\n    def __init__(\n        self,\n        rank: int,\n        group_name: str,\n        bucket: cp.ndarray | torch.Tensor,\n        metadata: dict[str, TensorMeta],\n        socket: zmq.Socket,\n        topic: str,\n    ) -> None:\n        self.rank = rank\n        self.group_name = group_name\n        self.bucket = bucket\n        self.metadata = metadata\n        self.socket = socket\n        self.topic = topic\n\n        loop = asyncio.get_running_loop()\n        self._task = loop.run_in_executor(None, self._run)\n\n    def _run(self):\n        # broadcast tensor meta via zeromq PUB/SUB\n        if self.rank == 0:\n            self.socket.send_string(self.topic, flags=zmq.SNDMORE)\n            self.socket.send_pyobj(self.metadata)\n        else:\n            self.socket.recv_string()\n            self.metadata = self.socket.recv_pyobj()\n\n        # broadcast tensor via NCCL\n        collective.broadcast(self.bucket, src_rank=0, group_name=self.group_name)\n\n    async def wait_for_complete(self) -> dict[str, TensorMeta]:\n        \"\"\"Wait for the broadcast operation to complete.\n\n        Returns:\n            dict[str, TensorMeta]: The bucket meta after broadcast.\n        \"\"\"\n        await self._task\n        return self.metadata\n\n\n@CheckpointEngineRegistry.register(\"nccl\")\nclass NCCLCheckpointEngine(CheckpointEngine):\n    \"\"\"NCCL checkpoint engine with collective communication.\n\n    Args:\n        bucket_size (int): Bucket size in bytes to transfer multiple weights at one time. Note that we use\n            two buffer to send and recv weights at same time, so the device memory overhead is 2 * bucket_size.\n        group_name (str): The name of the NCCL process group. Defaults to \"default\".\n        rebuild_group (bool): Whether to rebuild the NCCL process group in each update. Defaults to False.\n        is_master (bool): Whether the current process is the master process. Defaults to False.\n        rollout_dtype (torch.dtype): The dtype of the weights received from rollout workers. Defaults to torch.bfloat16.\n    \"\"\"\n\n    def __init__(\n        self,\n        bucket_size: int,\n        group_name: str = \"default\",\n        rebuild_group: bool = False,\n        is_master: bool = False,\n        rollout_dtype: torch.dtype = torch.bfloat16,\n    ) -> None:\n        self.bucket_size = bucket_size\n        self.group_name = group_name\n        self.rebuild_group = rebuild_group\n        self.rollout_dtype = rollout_dtype\n\n        # start zeromq server for broadcasting bucket tensor metadata\n        self.is_master = is_master\n        self.topic = \"bucket_metadata\"\n        if self.is_master:\n            self._start_zmq_server()\n\n    def prepare(self) -> MasterMetadata:\n        # For master process, use cupy instead of torch to avoid memory register error\n        # when `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.\n        if self.is_master:\n            self.send_buf = cp.zeros(self.bucket_size, dtype=cp.uint8)\n            self.recv_buf = cp.zeros(self.bucket_size, dtype=cp.uint8)\n        else:\n            self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=\"cuda\")\n            self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=\"cuda\")\n\n        return MasterMetadata(zmq_ip=self.ip, zmq_port=self.listen_port) if self.is_master else None\n\n    def finalize(self):\n        \"\"\"Destroy the NCCL process group if rebuild_group is True.\"\"\"\n        if self.rebuild_group:\n            if self.rank >= 0:\n                collective.destroy_collective_group(self.group_name)\n            self.rank = None\n            self.world_size = None\n\n        self.send_buf = None\n        self.recv_buf = None\n\n        torch.cuda.empty_cache()\n\n    @classmethod\n    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]):\n        trainer_kwargs = {\n            \"rank\": [0] + [-1] * (trainer_world_size - 1),\n            \"world_size\": [rollout_world_size + 1] * trainer_world_size,\n            \"master_metadata\": [metadata[0]] * trainer_world_size,\n        }\n        rollout_kwargs = {\n            \"rank\": list(range(1, rollout_world_size + 1)),\n            \"world_size\": [rollout_world_size + 1] * rollout_world_size,\n            \"master_metadata\": [metadata[0]] * rollout_world_size,\n        }\n        return trainer_kwargs, rollout_kwargs\n\n    def _start_zmq_server(self):\n        self.ip = ray.util.get_node_ip_address().strip(\"[]\")\n        self.listen_port, _ = get_free_port(self.ip)\n\n        context = zmq.Context()\n        self.socket = context.socket(zmq.PUB)\n        if is_valid_ipv6_address(self.ip):\n            address = f\"tcp://[{self.ip}]:{self.listen_port}\"\n            self.socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{self.ip}:{self.listen_port}\"\n\n        self.socket.bind(address)\n\n    def _connect_zmq_client(self, metadata: MasterMetadata):\n        assert not self.is_master, \"Master process should not connect to other processes.\"\n        context = zmq.Context()\n        self.socket = context.socket(zmq.SUB)\n        if is_valid_ipv6_address(metadata.zmq_ip):\n            address = f\"tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}\"\n            self.socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{metadata.zmq_ip}:{metadata.zmq_port}\"\n\n        self.socket.connect(address)\n        self.socket.setsockopt_string(zmq.SUBSCRIBE, self.topic)\n\n    def init_process_group(self, rank: int, world_size: int, master_metadata: MasterMetadata):\n        \"\"\"Initialize the NCCL process group.\n\n        Args:\n            rank (int): The rank of the current process.\n            world_size (int): The total number of processes.\n        \"\"\"\n        # For trainer workers other than rank 0, their rank should be -1.\n        if rank < 0:\n            self.rank = rank\n            self.world_size = world_size\n            return\n\n        if self.rebuild_group or not collective.is_group_initialized(self.group_name):\n            collective.init_collective_group(world_size, rank, \"nccl\", self.group_name)\n            self.rank = rank\n            self.world_size = world_size\n        else:\n            assert self.rank == rank, f\"rank {rank} is not equal to self.rank {self.rank}\"\n            assert self.world_size == world_size, (\n                f\"world_size {world_size} is not equal to self.world_size {self.world_size}\"\n            )\n\n        if self.rank > 0:\n            self._connect_zmq_client(master_metadata)\n        collective.barrier(self.group_name)\n\n        logger.info(f\"init_process_group rank: {self.rank}, world_size: {self.world_size}\")\n\n    @torch.no_grad()\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.rank <= 0, \"Trainer workers other than rank 0 should not send weights.\"\n\n        # For trainer rank other than 0, consume weights without sending.\n        if self.rank < 0:\n            for name, weight in weights:\n                pass\n            return\n\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        broadcast_op = None\n\n        start_time = time.time()\n        bucket_meta: dict[str, TensorMeta] = {}\n        offset = 0\n        for name, weight in weights:\n            # fill the tensor bucket\n            if offset + weight.nbytes > self.bucket_size:\n                torch.cuda.synchronize()\n\n                # wait previous broadcast op finish\n                if broadcast_op is not None:\n                    await broadcast_op.wait_for_complete()\n\n                broadcast_op = BroadcastOperation(\n                    rank=self.rank,\n                    group_name=self.group_name,\n                    bucket=send_buf,\n                    metadata={\"bucket_meta\": bucket_meta, \"is_last\": False},\n                    socket=self.socket,\n                    topic=self.topic,\n                )\n\n                # swap send_buf and recv_buf\n                send_buf, recv_buf = recv_buf, send_buf\n                bucket_meta = {}\n                offset = 0\n\n            assert offset + weight.nbytes <= self.bucket_size, (\n                f\"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket.\"\n            )\n\n            bucket_meta[name] = {\n                \"name\": name,\n                \"shape\": weight.shape,\n                \"dtype\": weight.dtype,\n                \"offset\": offset,\n            }\n            send_buf[offset : offset + weight.nbytes] = cp.asarray(weight.view(-1).view(torch.uint8))\n            offset += weight.nbytes\n\n        # broadcast last bucket\n        torch.cuda.synchronize()\n        if broadcast_op is not None:\n            await broadcast_op.wait_for_complete()\n\n        broadcast_op = BroadcastOperation(\n            rank=self.rank,\n            group_name=self.group_name,\n            bucket=send_buf,\n            metadata={\"bucket_meta\": bucket_meta, \"is_last\": True},\n            socket=self.socket,\n            topic=self.topic,\n        )\n        await broadcast_op.wait_for_complete()\n        logger.info(f\"Rank {self.rank} send weights done, time cost: {time.time() - start_time:.2f}s\")\n\n    @torch.no_grad()\n    async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.rank > 0, \"Rank 0 should not receive weights.\"\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        total_bytes, total_params = 0, 0\n\n        # receive first bucket\n        start_time = time.time()\n        broadcast_op = BroadcastOperation(\n            rank=self.rank,\n            group_name=self.group_name,\n            bucket=recv_buf,\n            metadata=None,\n            socket=self.socket,\n            topic=self.topic,\n        )\n        metadata = await broadcast_op.wait_for_complete()\n        total_bytes += self.bucket_size\n        total_params += len(metadata[\"bucket_meta\"])\n\n        # swap send_buf and recv_buf\n        send_buf, recv_buf = recv_buf, send_buf\n        while not metadata[\"is_last\"]:\n            # 1. receive next bucket\n            broadcast_op = BroadcastOperation(\n                rank=self.rank,\n                group_name=self.group_name,\n                bucket=recv_buf,\n                metadata=None,\n                socket=self.socket,\n                topic=self.topic,\n            )\n\n            # 2. yield tensor from send_buf\n            for name, meta in metadata[\"bucket_meta\"].items():\n                dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n                size = dtype.itemsize * shape.numel()\n                tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n                yield name, tensor\n\n            # 3. wait for next bucket broadcast finish\n            metadata = await broadcast_op.wait_for_complete()\n            total_bytes += self.bucket_size\n            total_params += len(metadata[\"bucket_meta\"])\n\n            # 4. swap send_buf and recv_buf\n            torch.cuda.synchronize()  # sync non-blocking copy\n            send_buf, recv_buf = recv_buf, send_buf\n\n        # yield tensor from send_buf\n        for name, meta in metadata[\"bucket_meta\"].items():\n            dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n            size = dtype.itemsize * shape.numel()\n            tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n            yield name, tensor\n\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} receive weights done, total_params: {total_params}, \"\n            f\"time cost: {time_cost:.2f}s, bandwidth: {bandwidth:.2f} GB/s\"\n        )\n"
  },
  {
    "path": "verl/checkpoint_engine/nixl_checkpoint_engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nimport time\nimport uuid\nfrom collections import defaultdict, deque\nfrom dataclasses import dataclass\nfrom typing import AsyncGenerator, Generator\nfrom unittest.mock import patch\n\nwith patch(\"importlib.metadata.distributions\", return_value=[]):\n    import cupy as cp\n\nimport nixl._api as nixl_api\nimport nixl._bindings as nixl_bindings\nimport ray\nimport torch\nimport zmq\nimport zmq.asyncio\n\nfrom verl.checkpoint_engine.base import CheckpointEngine, CheckpointEngineRegistry, TensorMeta\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@dataclass\nclass NixlAgentMetadata:\n    agent_name: str\n    agent_metadata: bytes\n    zmq_ip: str\n    zmq_port: int\n\n\nclass NixlAgent:\n    \"\"\"This is a wrapper class for nixl_agent, the main purpose is to use ZeroMQ instead of\n    `nixl_agent.send_notif` to send bucket tensor metadata.\n    \"\"\"\n\n    def __init__(self):\n        self.agent_name = str(uuid.uuid4())\n        self.agent = nixl_api.nixl_agent(self.agent_name)\n        self.notifications: dict[str, deque[bytes]] = defaultdict(deque)\n\n        self.start_zmq_server()\n        self.zmq_clients: dict[str, zmq.Socket] = {}\n        self.messages: dict[str, deque[bytes]] = defaultdict(deque)\n\n    def __getattr__(self, name):\n        attr = getattr(self.agent, name)\n\n        if callable(attr):\n\n            def wrapper(*args, **kwargs):\n                return attr(*args, **kwargs)\n\n            return wrapper\n        else:\n            return attr\n\n    def get_agent_metadata(self) -> NixlAgentMetadata:\n        return NixlAgentMetadata(\n            agent_name=self.agent_name,\n            agent_metadata=self.agent.get_agent_metadata(),\n            zmq_ip=self.ip,\n            zmq_port=self.listen_port,\n        )\n\n    def start_zmq_server(self):\n        self.ip = ray.util.get_node_ip_address().strip(\"[]\")\n        self.listen_port, _ = get_free_port(self.ip)\n\n        context = zmq.asyncio.Context()\n        self.socket = context.socket(zmq.PULL)\n        if is_valid_ipv6_address(self.ip):\n            address = f\"tcp://[{self.ip}]:{self.listen_port}\"\n            self.socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{self.ip}:{self.listen_port}\"\n\n        self.socket.bind(address)\n\n    def add_remote_agent(self, metadata: NixlAgentMetadata) -> str:\n        agent_name = self.agent.add_remote_agent(metadata.agent_metadata).decode(\"utf-8\")\n        assert agent_name == metadata.agent_name, f\"Agent name {agent_name} not equal to {metadata.agent_name}\"\n\n        context = zmq.Context()\n        socket = context.socket(zmq.PUSH)\n        if is_valid_ipv6_address(metadata.zmq_ip):\n            address = f\"tcp://[{metadata.zmq_ip}]:{metadata.zmq_port}\"\n            socket.setsockopt(zmq.IPV6, 1)\n        else:\n            address = f\"tcp://{metadata.zmq_ip}:{metadata.zmq_port}\"\n\n        socket.connect(address)\n        self.zmq_clients[agent_name] = socket\n        return agent_name\n\n    def remove_remote_agent(self, agent_name: str):\n        self.agent.remove_remote_agent(agent_name)\n        socket = self.zmq_clients.pop(agent_name)\n        socket.close()\n\n    def send_message(self, agent_name, message: dict):\n        socket = self.zmq_clients[agent_name]\n        socket.send_pyobj((self.agent_name, message), zmq.DONTWAIT)\n\n    async def read_message(self, agent_name: str) -> dict:\n        while len(self.messages[agent_name]) == 0:\n            recv_agent_name, message = await self.socket.recv_pyobj()\n            self.messages[recv_agent_name].append(message)\n        return self.messages[agent_name].popleft()\n\n    async def get_notification(self, remote_name: str) -> bytes:\n        while len(self.notifications[remote_name]) == 0:\n            notifs = self.agent.get_new_notifs()\n            for remote_name, notif in notifs.items():\n                self.notifications[remote_name].extend(notif)\n            await asyncio.sleep(0)\n        return self.notifications[remote_name].popleft()\n\n\nclass ReadableOperation:\n    \"\"\"Encapsulates a readable operation to remote agent.\n       1. send metadata to remote agent\n       2. wait until remote agent read complete.\n\n    Args:\n        agent (NixlAgent): The Nixl agent.\n        remote_agent (str): The name of the remote agent.\n        local_descs (nixl_bindings.nixlXferDList): The local transfer descriptors.\n        metadata (dict): Metadata for the read operation.\n        bucket_size (int): The size of the bucket in bytes.\n    \"\"\"\n\n    def __init__(\n        self,\n        agent: NixlAgent,\n        remote_agent: str,\n        local_descs: nixl_bindings.nixlXferDList,\n        metadata: dict,\n    ):\n        self.agent = agent\n        self.remote_agent = remote_agent\n        self.local_descs = local_descs\n        self.notify_key = uuid.uuid4().bytes\n        message = {\"notify_key\": self.notify_key, \"remote_descs\": self.local_descs, **metadata}\n        self.agent.send_message(self.remote_agent, message)\n\n    async def wait_for_complete(self):\n        \"\"\"Block until remote agent read complete.\"\"\"\n        notification = await self.agent.get_notification(self.remote_agent)\n        assert self.notify_key == notification, f\"Notify key {self.notify_key} not equal to {notification}\"\n        logger.debug(f\"ReadableOperation to {self.remote_agent} complete\")\n\n\nclass ReadOperation:\n    \"\"\"Encapsulates a read operation from remote agent.\n    1. read medata from remote agent\n    2. start read transfer operation\n    3. wait until read complete\n\n    Args:\n        agent (NixlAgent): The Nixl agent.\n        remote_agent (str): The name of the remote agent.\n        local_descs (nixl_bindings.nixlXferDList): The local transfer descriptors.\n        bucket_size (int): The size of the bucket in bytes.\n    \"\"\"\n\n    def __init__(self, agent: NixlAgent, remote_agent: str, local_descs: nixl_bindings.nixlXferDList, bucket_size: int):\n        self.agent = agent\n        self.remote_agent = remote_agent\n        self.local_descs = local_descs\n        self.remote_descs = None\n        self.xfer_handle = None\n        self.notify_key = None\n        self.bucket_size = bucket_size\n        self.start_time = None\n\n    async def read_metadata(self) -> dict:\n        \"\"\"Block until the remote agent sends the metadata.\n\n        Returns:\n            dict: Metadata from the remote agent.\n        \"\"\"\n        metadata = await self.agent.read_message(self.remote_agent)\n        self.remote_descs = metadata.pop(\"remote_descs\")\n        self.notify_key = metadata.pop(\"notify_key\")\n        return metadata\n\n    def begin_read(self):\n        \"\"\"Start the read operation.\"\"\"\n        assert self.remote_descs is not None and self.notify_key is not None\n        self.xfer_handle = self.agent.initialize_xfer(\n            \"READ\", self.local_descs, self.remote_descs, self.remote_agent, self.notify_key\n        )\n        state = self.agent.transfer(self.xfer_handle)\n        assert state != \"ERR\", f\"Read from {self.remote_agent} got to {state} state.\"\n        self.start_time = time.time()\n\n    async def wait_for_complete(self):\n        \"\"\"Block until the read operation complete.\"\"\"\n        while True:\n            state = self.agent.check_xfer_state(self.xfer_handle)\n            if state == \"ERR\":\n                logger.error(f\"Read from {self.remote_agent} got to {state} state.\")\n                exit(-1)\n            elif state == \"DONE\":\n                break\n            else:\n                await asyncio.sleep(0)\n        self.agent.release_xfer_handle(self.xfer_handle)\n        end_time = time.time()\n        bandwidth = self.bucket_size / (end_time - self.start_time) / (1024 * 1024 * 1024)\n        logger.debug(f\"ReadOperation read data from {self.remote_agent} complete, bandwidth: {bandwidth:.2f} GB/s\")\n\n\n@CheckpointEngineRegistry.register(\"nixl\")\nclass NIXLCheckpointEngine(CheckpointEngine):\n    \"\"\"NIXL checkpoint engine with p2p communication, support various backends: ucx, uccl, mooncacke, etc.\n\n    For UCX backend, some environment variables need to be set: UCX_TLS, UCX_IB_GID_INDEX, UCX_IB_DEVICES, etc.\n    Please refer to: https://openucx.readthedocs.io/en/master/faq.html\n\n    Args:\n        bucket_size (int): Bucket size in bytes to transfer multiple weights at one time. Note that we use\n            two buffer to send and recv weights at same time, so the device memory overhead is 2 * bucket_size.\n        device (str): The device to use for the checkpoint engine, \"cpu\" or \"cuda\".\n        rollout_dtype (torch.dtype): The dtype of the weights received from rollout workers. Defaults to torch.bfloat16.\n    \"\"\"\n\n    def __init__(\n        self,\n        bucket_size: int,\n        device: str = \"cuda\",\n        rollout_dtype: torch.dtype = torch.bfloat16,\n        is_master: bool = False,\n    ):\n        self.bucket_size = bucket_size\n        self.device = device\n        self.rollout_dtype = rollout_dtype\n        self.agent = NixlAgent()\n        self.is_master = is_master\n\n    def prepare(self) -> NixlAgentMetadata:\n        \"\"\"Prepare send and recv bucket.\n\n        Returns:\n            NixlAgentMetadata: The metadata of the current nixl agent.\n        \"\"\"\n        # For master process, use cupy instead of torch to avoid memory register error\n        # when `PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True`.\n        if self.device == \"cuda\":\n            send_buf = cp.zeros(self.bucket_size, dtype=cp.uint8)\n            recv_buf = cp.zeros(self.bucket_size, dtype=cp.uint8)\n            self.send_buf = torch.as_tensor(send_buf, dtype=torch.uint8)\n            self.recv_buf = torch.as_tensor(recv_buf, dtype=torch.uint8)\n        else:\n            self.send_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=self.device, pin_memory=True)\n            self.recv_buf = torch.zeros(self.bucket_size, dtype=torch.uint8, device=self.device, pin_memory=True)\n        self.send_reg_descs = self.agent.register_memory(self.send_buf)\n        self.recv_reg_descs = self.agent.register_memory(self.recv_buf)\n        self.send_descs = self.agent.get_xfer_descs(self.send_buf)\n        self.recv_descs = self.agent.get_xfer_descs(self.recv_buf)\n\n        return self.agent.get_agent_metadata()\n\n    @classmethod\n    def build_topology(cls, trainer_world_size: int, rollout_world_size: int, metadata: list[dict]):\n        trainer_kwargs = {\n            \"method\": [\"init_process_group\"] * trainer_world_size,\n            \"rank\": [0] + [-1] * (trainer_world_size - 1),\n            \"world_size\": [rollout_world_size + 1] * trainer_world_size,\n            \"prev_agent_metadata\": [None] * trainer_world_size,\n            \"next_agent_metadata\": [metadata[-rollout_world_size]] + [None] * (trainer_world_size - 1),\n        }\n\n        rollout_kwargs = {\n            \"method\": [\"init_process_group\"] * rollout_world_size,\n            \"rank\": list(range(1, rollout_world_size + 1)),\n            \"world_size\": [rollout_world_size + 1] * rollout_world_size,\n            \"prev_agent_metadata\": [metadata[0]] + metadata[-rollout_world_size:-1],\n            \"next_agent_metadata\": metadata[-rollout_world_size + 1 :] + [None],\n        }\n        return trainer_kwargs, rollout_kwargs\n\n    def init_process_group(\n        self, rank: int, world_size: int, prev_agent_metadata: NixlAgentMetadata, next_agent_metadata: NixlAgentMetadata\n    ):\n        \"\"\"Setup the communication with the previous and next agent.\n\n        Args:\n            rank (int): The rank of the current process.\n            world_size (int): The total number of processes.\n            prev_agent_metadata (NixlAgentMetadata): The metadata of the previous nixl agent.\n            next_agent_metadata (NixlAgentMetadata): The metadata of the next nixl agent.\n        \"\"\"\n        if rank < 0:\n            assert not prev_agent_metadata and not next_agent_metadata, (\n                f\"rank {rank} should not have prev_agent_metadata or next_agent_metadata\"\n            )\n        elif rank == 0:\n            assert not prev_agent_metadata and next_agent_metadata, f\"rank {rank} should have next_agent_metadata\"\n        elif 0 < rank < world_size - 1:\n            assert prev_agent_metadata and next_agent_metadata, (\n                f\"rank {rank} should have prev_agent_metadata and next_agent_metadata\"\n            )\n        elif rank == world_size - 1:\n            assert prev_agent_metadata and not next_agent_metadata, (\n                f\"rank {rank} should have prev_agent_metadata and not next_agent_metadata\"\n            )\n\n        self.rank = rank\n        self.world_size = world_size\n        self.prev_agent = None\n        self.next_agent = None\n\n        if prev_agent_metadata is not None:\n            self.prev_agent = self.agent.add_remote_agent(prev_agent_metadata)\n\n        if next_agent_metadata is not None:\n            self.next_agent = self.agent.add_remote_agent(next_agent_metadata)\n\n        logger.info(\n            f\"init_process_group rank: {self.rank}, world_size: {self.world_size}, \"\n            f\"prev_agent: {self.prev_agent}, next_agent: {self.next_agent}\"\n        )\n\n    def finalize(self):\n        \"\"\"Cleanup communication with the previous and next agent, and deregister the memory.\"\"\"\n        if self.prev_agent:\n            self.agent.remove_remote_agent(self.prev_agent)\n        if self.next_agent:\n            self.agent.remove_remote_agent(self.next_agent)\n\n        self.agent.deregister_memory(self.send_reg_descs)\n        self.agent.deregister_memory(self.recv_reg_descs)\n        self.send_buf = None\n        self.recv_buf = None\n        self.send_reg_descs = None\n        self.recv_reg_descs = None\n        self.send_descs = None\n        self.recv_descs = None\n\n        self.rank = None\n        self.world_size = None\n        self.prev_agent = None\n        self.next_agent = None\n\n    @torch.no_grad()\n    async def send_weights(self, weights: Generator[tuple[str, torch.Tensor], None, None]):\n        \"\"\"Send the weights of the model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.rank <= 0, \"Trainer workers other than rank 0 should not send weights.\"\n\n        # For trainer workers other than rank 0, just consume weights and do nothing.\n        if self.rank < 0:\n            for name, weight in weights:\n                pass\n            return\n\n        assert self.next_agent is not None, \"Next agent is not set.\"\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        send_descs, recv_descs = self.send_descs, self.recv_descs\n        readable_op = None\n\n        start_time = time.time()\n        bucket_meta: dict[str, TensorMeta] = {}\n        offset = 0\n        for name, weight in weights:\n            # fill the tensor bucket\n            if offset + weight.nbytes > self.bucket_size:\n                torch.cuda.synchronize()\n\n                # wait previous bucket to be received\n                if readable_op is not None:\n                    await readable_op.wait_for_complete()\n\n                # send bucket meta to next agent\n                readable_op = ReadableOperation(\n                    self.agent,\n                    self.next_agent,\n                    send_descs,\n                    {\"bucket_meta\": bucket_meta, \"is_last\": False},\n                )\n\n                # swap send and recv buf\n                send_buf, recv_buf = recv_buf, send_buf\n                send_descs, recv_descs = recv_descs, send_descs\n                bucket_meta = {}\n                offset = 0\n\n            assert offset + weight.nbytes <= self.bucket_size, (\n                f\"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket.\"\n            )\n\n            bucket_meta[name] = {\n                \"name\": name,\n                \"shape\": weight.shape,\n                \"dtype\": weight.dtype,\n                \"offset\": offset,\n            }\n            send_buf[offset : offset + weight.nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)\n            offset += weight.nbytes\n\n        # send last bucket meta to next agent\n        torch.cuda.synchronize()\n        if readable_op is not None:\n            await readable_op.wait_for_complete()\n\n        readable_op = ReadableOperation(\n            self.agent, self.next_agent, send_descs, {\"bucket_meta\": bucket_meta, \"is_last\": True}\n        )\n        await readable_op.wait_for_complete()\n        logger.info(f\"Rank {self.rank} send weights done, time cost: {time.time() - start_time:.2f}s\")\n\n    @torch.no_grad()\n    async def receive_weights(self) -> AsyncGenerator[tuple[str, torch.Tensor], None]:\n        \"\"\"Receive the weights of the model.\n\n        Yields:\n            A tuple of the name of the weight tensor and the tensor itself.\n        \"\"\"\n        assert self.prev_agent is not None, \"Previous agent is not set.\"\n        send_buf, recv_buf = self.send_buf, self.recv_buf\n        send_descs, recv_descs = self.send_descs, self.recv_descs\n        total_bytes, total_params = 0, 0\n\n        # receive first bucket from previous agent\n        start_time = time.time()\n        read_op = ReadOperation(self.agent, self.prev_agent, recv_descs, self.bucket_size)\n        metadata = await read_op.read_metadata()\n        read_op.begin_read()\n        await read_op.wait_for_complete()\n        total_bytes += self.bucket_size\n        total_params += len(metadata[\"bucket_meta\"])\n\n        # swap send and recv buf\n        send_buf, recv_buf = recv_buf, send_buf\n        send_descs, recv_descs = recv_descs, send_descs\n        while not metadata[\"is_last\"]:\n            # 1. send bucket to next agent\n            readable_op = None\n            if self.next_agent is not None:\n                readable_op = ReadableOperation(\n                    self.agent,\n                    self.next_agent,\n                    send_descs,\n                    metadata,\n                )\n\n            # 2. receive bucket from previous agent\n            read_op = ReadOperation(self.agent, self.prev_agent, recv_descs, self.bucket_size)\n            next_metadata = await read_op.read_metadata()\n            read_op.begin_read()\n\n            # 3. yield tensor from send_buf\n            for name, meta in metadata[\"bucket_meta\"].items():\n                dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n                size = dtype.itemsize * shape.numel()\n                tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n                yield name, tensor\n\n            # 4. wait for next agent read complete and read from previous agent complete\n            if readable_op is not None:\n                await readable_op.wait_for_complete()\n            await read_op.wait_for_complete()\n            total_bytes += self.bucket_size\n            total_params += len(next_metadata[\"bucket_meta\"])\n\n            # 5. swap send and recv buf\n            torch.cuda.synchronize()  # sync non-blocking copy\n            metadata = next_metadata\n            send_buf, recv_buf = recv_buf, send_buf\n            send_descs, recv_descs = recv_descs, send_descs\n\n        # send last bucket to next agent\n        readable_op = None\n        if self.next_agent is not None:\n            readable_op = ReadableOperation(\n                self.agent,\n                self.next_agent,\n                send_descs,\n                metadata,\n            )\n\n        # yield tensor from send_buf\n        for name, meta in metadata[\"bucket_meta\"].items():\n            dtype, shape = meta[\"dtype\"], meta[\"shape\"]\n            size = dtype.itemsize * shape.numel()\n            tensor = send_buf[meta[\"offset\"] : meta[\"offset\"] + size].view(dtype=dtype).view(shape)\n            yield name, tensor\n\n        # wait for next agent read complete\n        if readable_op is not None:\n            await readable_op.wait_for_complete()\n        time_cost = time.time() - start_time\n        bandwidth = total_bytes / time_cost / (1024 * 1024 * 1024)\n        logger.info(\n            f\"Rank {self.rank} receive weights done, total_params: {total_params}, \"\n            f\"time cost: {time_cost:.2f}s, bandwidth: {bandwidth:.2f} GB/s\"\n        )\n"
  },
  {
    "path": "verl/experimental/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/agent_loop/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .agent_loop import AgentLoopBase, AgentLoopManager, AgentLoopWorker, AsyncLLMServerManager\nfrom .single_turn_agent_loop import SingleTurnAgentLoop\nfrom .tool_agent_loop import ToolAgentLoop\n\n_ = [SingleTurnAgentLoop, ToolAgentLoop]\n\n__all__ = [\"AgentLoopBase\", \"AgentLoopManager\", \"AsyncLLMServerManager\", \"AgentLoopWorker\"]\n"
  },
  {
    "path": "verl/experimental/agent_loop/agent_loop.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nimport random\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nimport hydra\nimport numpy as np\nimport ray\nimport torch\nfrom cachetools import LRUCache\nfrom omegaconf import DictConfig, OmegaConf\nfrom PIL import Image\nfrom pydantic import BaseModel, ConfigDict\nfrom tensordict import TensorDict\nfrom transformers import AutoProcessor, AutoTokenizer\n\nfrom verl.experimental.agent_loop.prometheus_utils import update_prometheus_config\nfrom verl.experimental.agent_loop.utils import resolve_config_path\nfrom verl.protocol import DataProto\nfrom verl.single_controller.ray.base import RayResourcePool, RayWorkerGroup\nfrom verl.utils.chat_template import apply_chat_template, initialize_system_prompt\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.dataset.rl_dataset import RLHFDataset, get_dataset_class\nfrom verl.utils.model import compute_position_id_with_mask\nfrom verl.utils.ray_utils import auto_await, get_event_loop\nfrom verl.utils.rollout_trace import (\n    RolloutTraceConfig,\n    rollout_trace_attr,\n    rollout_trace_op,\n)\nfrom verl.utils.tokenizer import normalize_token_ids\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.replica import TokenOutput, get_rollout_replica_class\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\nDEFAULT_ROUTING_CACHE_SIZE = 10000\n\n\n@ray.remote\nclass GlobalRequestLoadBalancer:\n    \"\"\"Global sticky-session + in-flight load balancer shared by all AgentLoopWorkers.\"\"\"\n\n    def __init__(self, server_actor_ids: list[str], max_cache_size: int = DEFAULT_ROUTING_CACHE_SIZE):\n        if not server_actor_ids:\n            raise ValueError(\"server_actor_ids must be non-empty\")\n\n        self._inflight_requests: dict[str, int] = {sid: 0 for sid in server_actor_ids}\n        self._request_id_to_server: LRUCache = LRUCache(maxsize=max_cache_size)\n\n    def acquire_server(self, request_id: str) -> str:\n        \"\"\"Acquire a server for the given request, reusing the same server for multi-turn conversations.\"\"\"\n        # request-level sticky (multi-turn: same conversation -> same server)\n        if request_id in self._request_id_to_server:\n            server_id = self._request_id_to_server[request_id]\n            self._inflight_requests[server_id] += 1\n            return server_id\n\n        # new request: route to least loaded server\n        server_id = min(self._inflight_requests, key=self._inflight_requests.get)\n        self._request_id_to_server[request_id] = server_id\n        self._inflight_requests[server_id] += 1\n        return server_id\n\n    def release_server(self, server_id: str) -> None:\n        \"\"\"Release a server after a request completes, decrementing its inflight count.\"\"\"\n        if server_id not in self._inflight_requests:\n            raise ValueError(f\"Invalid server_id for release: {server_id}\")\n        if self._inflight_requests[server_id] <= 0:\n            raise ValueError(f\"Release called with no inflight requests on server {server_id}\")\n        self._inflight_requests[server_id] -= 1\n\n\ndef _get_rollout_and_model_config(config: DictConfig) -> tuple[DictConfig, DictConfig]:\n    # TODO: backward compatibility, remove this once we switch to new trainer.\n    if config.get(\"actor_rollout_ref\"):\n        return config.actor_rollout_ref.rollout, config.actor_rollout_ref.model\n    else:\n        return config.rollout, config.model\n\n\nclass AsyncLLMServerManager:\n    \"\"\"\n    A class to manage multiple OpenAI compatible LLM servers. This class provides\n    - Load balance: least in-flight requests load balancing via global coordination\n    - Sticky session: send multi-turn chat completions to same server for automatic prefix caching\n    \"\"\"\n\n    def __init__(\n        self,\n        config: DictConfig,\n        servers: list[tuple[str, ray.actor.ActorHandle]],\n        load_balancer_handle: ray.actor.ActorHandle,\n    ):\n        \"\"\"Initialize the AsyncLLMServerManager.\n\n        Args:\n            config (DictConfig): whole config for main entrypoint.\n            servers (list[tuple[str, ray.actor.ActorHandle]]): (address, handle) pairs for each LLM server.\n            load_balancer_handle (ray.actor.ActorHandle): shared global load balancer actor.\n        \"\"\"\n        self.config = config\n        self._load_balancer = load_balancer_handle\n        self._server_id_to_handle: dict[str, ray.actor.ActorHandle] = dict(servers)\n\n    async def _acquire_server(self, request_id: str) -> tuple[str, ray.actor.ActorHandle]:\n        server_id = await self._load_balancer.acquire_server.remote(request_id=request_id)\n        handle = self._server_id_to_handle.get(server_id)\n        if handle is None:\n            raise RuntimeError(f\"Unknown server_id returned by load balancer: {server_id}\")\n        return server_id, handle\n\n    def _release_server(self, server_id: str) -> None:\n        # Fire-and-forget: release is just a counter decrement, no need to await.\n        # Awaiting here risks blocking the finally clause if the LB actor is unresponsive.\n        self._load_balancer.release_server.remote(server_id=server_id)\n\n    @rollout_trace_op\n    async def generate(\n        self,\n        request_id,\n        *,\n        prompt_ids: list[int],\n        sampling_params: dict[str, Any],\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> TokenOutput:\n        \"\"\"Generate tokens from prompt ids.\n\n        Args:\n            request_id (str): request id for sticky session.\n            prompt_ids (List[int]): List of prompt token ids.\n            sampling_params (Dict[str, Any]): Sampling parameters for the chat completion.\n\n        Returns:\n            TokenOutput: token output\n        \"\"\"\n        server_id, server = await self._acquire_server(request_id)\n        try:\n            output: TokenOutput = await server.generate.remote(\n                request_id=uuid4().hex,  # use new request_id for each turn\n                prompt_ids=prompt_ids,\n                sampling_params=sampling_params,\n                image_data=image_data,\n                video_data=video_data,\n            )\n            return output\n        finally:\n            self._release_server(server_id)\n\n\nclass AgentLoopMetrics(BaseModel):\n    \"\"\"Agent loop performance metrics.\"\"\"\n\n    generate_sequences: float = 0.0\n    tool_calls: float = 0.0\n    num_preempted: int = -1  # -1 means not available\n\n\nclass AgentLoopOutput(BaseModel):\n    \"\"\"Agent loop output.\"\"\"\n\n    prompt_ids: list[int]\n    \"\"\"Prompt token ids.\"\"\"\n    response_ids: list[int]\n    \"\"\"Response token ids including LLM generated token, tool response token.\"\"\"\n    response_mask: list[int]\n    \"\"\"Response mask, 1 for LLM generated token, 0 for tool response token.\"\"\"\n    response_logprobs: Optional[list[float]] = None\n    \"\"\"Log probabilities for the response tokens.\"\"\"\n    routed_experts: Optional[Any] = None\n    \"\"\"Routed experts for the total tokens.\"\"\"\n    multi_modal_data: Optional[dict[str, Any]] = None\n    \"\"\"Multi-modal data for multi-modal tools.\"\"\"\n    reward_score: Optional[float] = None\n    \"\"\"Reward score for the trajectory.\"\"\"\n    num_turns: int = 0\n    \"\"\"Number of chat turns, including user, assistant, tool.\"\"\"\n    metrics: AgentLoopMetrics\n    \"\"\"Auxiliary performance metrics\"\"\"\n    extra_fields: dict[str, Any] = {}\n    \"\"\"Extra fields for dynamic addition.\"\"\"\n\n\nclass _InternalAgentLoopOutput(AgentLoopOutput):\n    \"\"\"Internal agent loop output with padded sequences.\"\"\"\n\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n\n    prompt_ids: torch.Tensor\n    \"\"\"Padded prompt token ids.\"\"\"\n    response_ids: torch.Tensor\n    \"\"\"Padded response token ids.\"\"\"\n    input_ids: torch.Tensor\n    \"\"\"Padded input ids(prompt_ids + response_ids).\"\"\"\n    position_ids: torch.Tensor\n    \"\"\"Padded position ids.\"\"\"\n    response_mask: torch.Tensor\n    \"\"\"Padded response mask.\"\"\"\n    attention_mask: torch.Tensor\n    \"\"\"Padded attention mask.\"\"\"\n    response_logprobs: Optional[torch.Tensor] = None\n    \"\"\"Padded log probabilities for the response tokens.\"\"\"\n    routed_experts: Optional[torch.Tensor] = None\n    \"\"\"Padded routed experts for the total tokens.\"\"\"\n    multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None\n    \"\"\"Multi-modal inputs for processors (e.g., pixel_values, image_grid_thw).\"\"\"\n    extra_fields: dict[str, Any] = {}\n    \"\"\"Extra fields for dynamic addition.\"\"\"\n\n\nclass DictConfigWrap:\n    \"\"\"Wrapper for DictConfig to avoid hydra.utils.instantiate recursive resolve.\"\"\"\n\n    def __init__(self, config: DictConfig):\n        self.config = config\n\n\nclass AgentLoopBase(ABC):\n    \"\"\"An agent loop takes an input message, chat with OpenAI compatible LLM server and interact with various\n    environments.\n\n    Args:\n        trainer_config (DictConfig): whole config for main entrypoint.\n        server_manager (AsyncLLMServerManager): OpenAI compatible LLM server manager.\n        tokenizer (AutoTokenizer): Tokenizer for tokenize messages.\n        processor (AutoProcessor): Processor for process messages.\n        dataset_cls (type[Dataset]): Dataset class for creating dataset, Defaults to RLHFDataset.\n        data_config (DictConfigWrap): Dataset config.\n    \"\"\"\n\n    def __init__(\n        self,\n        trainer_config: DictConfigWrap,\n        server_manager: AsyncLLMServerManager,\n        tokenizer: AutoTokenizer,\n        processor: AutoProcessor,\n        dataset_cls: type[RLHFDataset],\n        data_config: DictConfigWrap,\n        **kwargs,\n    ):\n        self.config = trainer_config.config\n        self.rollout_config, _ = _get_rollout_and_model_config(self.config)\n        self.server_manager = server_manager\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.dataset_cls = dataset_cls\n        self.data_config = data_config.config\n        self.apply_chat_template_kwargs = self.data_config.get(\"apply_chat_template_kwargs\", {})\n        self.system_prompt = initialize_system_prompt(self.tokenizer, **self.apply_chat_template_kwargs)\n        self.loop = get_event_loop()\n\n    async def process_vision_info(self, messages: list[dict]) -> dict:\n        \"\"\"Extract images and videos from messages.\n\n        Args:\n            messages (list[dict]): Input messages.\n\n        Returns:\n            dict: Multi-modal data with keys \"images\" and \"videos\".\n        \"\"\"\n        multi_modal_data = {}\n        if self.processor is not None:\n            images, videos = await self.dataset_cls.process_vision_info(\n                messages, image_patch_size=self.processor.image_processor.patch_size, config=self.data_config\n            )\n            if images is not None:\n                multi_modal_data[\"images\"] = images\n            if videos is not None:\n                multi_modal_data[\"videos\"] = videos\n\n        return multi_modal_data\n\n    async def apply_chat_template(\n        self,\n        messages: list[dict],\n        tools: list[dict] = None,\n        images: list[Image.Image] = None,\n        videos: list[tuple[torch.Tensor, dict]] = None,\n        remove_system_prompt: bool = False,\n    ):\n        \"\"\"Apply chat template to messages with optional tools, images, and videos.\n\n        Args:\n            messages (list[dict]): Input messages.\n            tools (list[dict], optional): Tools schemas. Defaults to None.\n            images (list[Image.Image], optional): Input images. Defaults to None.\n            videos (list[tuple[torch.Tensor, dict]], optional): Input videos. Defaults to None.\n            remove_system_prompt (bool, optional): Whether to remove system prompt. Defaults to False.\n\n        Returns:\n            list[int]: Prompt token ids.\n        \"\"\"\n        if self.processor is not None:\n            raw_prompt = await self.loop.run_in_executor(\n                None,\n                lambda: apply_chat_template(\n                    self.processor,\n                    messages,\n                    tools=tools,\n                    add_generation_prompt=True,\n                    tokenize=False,\n                    **self.apply_chat_template_kwargs,\n                ),\n            )\n\n            # split the videos and according metadatas\n            if videos is not None:\n                videos, video_metadatas = zip(*videos, strict=False)\n                videos, video_metadatas = list(videos), list(video_metadatas)\n            else:\n                video_metadatas = None\n\n            model_inputs = self.processor(\n                text=[raw_prompt],\n                images=images,\n                videos=videos,\n                video_metadata=video_metadatas,\n                return_tensors=\"pt\",\n                do_sample_frames=False,\n            )\n            prompt_ids = normalize_token_ids(model_inputs.pop(\"input_ids\"))\n        else:\n            tokenized_prompt = await self.loop.run_in_executor(\n                None,\n                lambda: apply_chat_template(\n                    self.tokenizer,\n                    messages,\n                    tools=tools,\n                    add_generation_prompt=True,\n                    tokenize=True,\n                    **self.apply_chat_template_kwargs,\n                ),\n            )\n            prompt_ids = normalize_token_ids(tokenized_prompt)\n\n        if remove_system_prompt:\n            prompt_ids = prompt_ids[len(self.system_prompt) :]\n\n        return prompt_ids\n\n    @abstractmethod\n    async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:\n        \"\"\"Run agent loop to interact with LLM server and environment.\n\n        Args:\n            sampling_params (Dict[str, Any]): LLM sampling params.\n            **kwargs: dataset fields from `verl.utils.dataset.RLHFDataset`.\n\n        Returns:\n            AgentLoopOutput: Agent loop output.\n        \"\"\"\n        raise NotImplementedError\n\n\n\"\"\"Agent loop registry: key is agent_name, value is a dict of agent loop config\nused by hydra.utils.instantiate to initialize agent loop instance.\n\nhttps://hydra.cc/docs/advanced/instantiate_objects/overview/\n\"\"\"\n_agent_loop_registry: dict[str, dict] = {}\n\n\ndef register(agent_name: str):\n    \"\"\"Register agent loop class.\"\"\"\n\n    def decorator(subclass: type[AgentLoopBase]) -> type[AgentLoopBase]:\n        fqdn = f\"{subclass.__module__}.{subclass.__qualname__}\"\n        _agent_loop_registry[agent_name] = {\"_target_\": fqdn}\n        return subclass\n\n    return decorator\n\n\nclass AgentLoopWorker:\n    \"\"\"Agent loop worker takes a batch of messages and run each message in an agent loop.\n\n    Args:\n        config (DictConfig): whole config for main entrypoint.\n        servers (list[tuple[str, ray.actor.ActorHandle]]): (address, handle) pairs for each LLM server.\n        reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: DictConfig,\n        servers: list[tuple[str, ray.actor.ActorHandle]],\n        load_balancer_handle: ray.actor.ActorHandle,\n        reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,\n    ):\n        \"\"\"Initialize agent loop manager.\n        Args:\n            config (DictConfig): YAML config.\n            servers (list[tuple[str, ray.actor.ActorHandle]]): (address, handle) pairs for each LLM server.\n            load_balancer_handle (ray.actor.ActorHandle): shared global load balancer actor.\n            reward_loop_worker_handles (list[ray.actor.ActorHandle]): Actor handles for streaming reward computation.\n        \"\"\"\n        self.config = config\n        rollout_config, model_config = _get_rollout_and_model_config(config)\n        self.rollout_config: RolloutConfig = omega_conf_to_dataclass(rollout_config)\n        self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config)\n\n        # for recipe to change\n        if not hasattr(self, \"server_manager\"):\n            self.server_manager = AsyncLLMServerManager(\n                config,\n                servers,\n                load_balancer_handle=load_balancer_handle,\n            )\n\n        self.dataset_cls = get_dataset_class(config.data)\n        self.reward_loop_worker_handles = reward_loop_worker_handles\n\n        self.tokenizer = self.model_config.tokenizer\n        self.processor = self.model_config.processor\n\n        agent_loop_config_path = self.rollout_config.agent.agent_loop_config_path\n        if agent_loop_config_path:\n            resolved_path = resolve_config_path(agent_loop_config_path)\n            agent_loop_configs = OmegaConf.load(resolved_path)\n            for agent_loop_config in agent_loop_configs:\n                _agent_loop_registry[agent_loop_config.name] = agent_loop_config\n        if self.model_config.get(\"custom_chat_template\", None) is not None:\n            if self.model_config.processor is not None:\n                self.model_config.processor.chat_template = self.model_config.custom_chat_template\n            self.model_config.tokenizer.chat_template = self.model_config.custom_chat_template\n\n        trace_config = self.rollout_config.trace\n        RolloutTraceConfig.init(\n            self.rollout_config.trace.project_name,\n            self.rollout_config.trace.experiment_name,\n            trace_config.get(\"backend\"),\n            trace_config.get(\"token2text\", False),\n            trace_config.get(\"max_samples_per_step_per_worker\", None),\n        )\n\n    async def generate_sequences(self, batch: DataProto) -> DataProto:\n        \"\"\"Generate sequences from agent loop.\n\n        Args:\n            batch (DataProto): Input batch.\n\n        Returns:\n            DataProto: Output batch.\n            - prompts: [bsz, prompt_length], prompt token ids from dataset.\n            - responses: [bsz, response_length], output token ids include response tokens\n              from LLM generation and observation tokens from tool_calls.\n            - response_mask: [bsz, response_length], 1 for LLM generated tokens, 0 for observation/padding tokens.\n            - input_ids: [bsz, prompt_length + response_length], whole sequence token ids, including prompt tokens\n              and response tokens.\n            - attention_mask: [bsz, prompt_length + response_length], 0 for padding tokens, 1 for other tokens.\n            - position_ids: [bsz, prompt_length + response_length], incremental position ids.\n\n            For multi-turn conversations:\n            responses:     |<- LLM generation ->|<- tool_calls ->|<- LLM generation ->|<- padding ->|\n            response_mask: | 1, 1, 1, ..., 1, 1 | 0, 0, .., 0, 0 | 1, 1, 1, ..., 1, 1 | 0, 0, ..., 0|\n        \"\"\"\n        config = self.rollout_config\n        sampling_params = dict(\n            temperature=config.temperature,\n            top_p=config.top_p,\n            top_k=config.top_k,\n            repetition_penalty=1.0,\n            logprobs=config.calculate_log_probs,\n        )\n\n        # override sampling params for validation\n        if batch.meta_info.get(\"validate\", False):\n            sampling_params[\"top_p\"] = config.val_kwargs.top_p\n            sampling_params[\"top_k\"] = config.val_kwargs.top_k\n            sampling_params[\"temperature\"] = config.val_kwargs.temperature\n\n        # by default, we assume it's a single turn agent\n        if \"agent_name\" not in batch.non_tensor_batch:\n            default_agent_loop = config.agent.default_agent_loop\n            batch.non_tensor_batch[\"agent_name\"] = np.array([default_agent_loop] * len(batch), dtype=object)\n\n        if \"index\" in batch.non_tensor_batch:\n            index = batch.non_tensor_batch[\"index\"]\n        else:\n            index = np.arange(len(batch))\n\n        max_samples_per_worker = RolloutTraceConfig.get_instance().max_samples_per_step_per_worker\n\n        # For n rollouts per sample, we trace all n rollouts for selected samples\n        # Note: This sampling happens per-worker, so total traces = max_samples_per_worker * num_workers * n\n        if max_samples_per_worker is not None:\n            unique_sample_indices = np.unique(index)\n            if max_samples_per_worker < len(unique_sample_indices):\n                selected_samples = set(\n                    np.random.choice(unique_sample_indices, max_samples_per_worker, replace=False).tolist()\n                )\n                traced_indices = set(i for i in range(len(batch)) if index[i] in selected_samples)\n            else:\n                traced_indices = set(range(len(batch)))\n        else:\n            traced_indices = set(range(len(batch)))\n\n        trajectory_info = await get_trajectory_info(\n            batch.meta_info.get(\"global_steps\", -1), index.tolist(), batch.meta_info.get(\"validate\", False)\n        )\n\n        tasks = []\n        for i in range(len(batch)):\n            trace_this_sample = i in traced_indices\n            kwargs = {k: v[i] for k, v in batch.non_tensor_batch.items()}\n            tasks.append(\n                asyncio.create_task(\n                    self._run_agent_loop(sampling_params, trajectory_info[i], trace=trace_this_sample, **kwargs)\n                )\n            )\n        outputs = await asyncio.gather(*tasks)\n\n        output = self._postprocess(outputs, input_non_tensor_batch=batch.non_tensor_batch)\n\n        return output\n\n    async def _run_agent_loop(\n        self,\n        sampling_params: dict[str, Any],\n        trajectory: dict[str, Any],\n        *,\n        agent_name: str,\n        trace: bool = True,\n        **kwargs,\n    ) -> _InternalAgentLoopOutput:\n        with rollout_trace_attr(\n            step=trajectory[\"step\"],\n            sample_index=trajectory[\"sample_index\"],\n            rollout_n=trajectory[\"rollout_n\"],\n            validate=trajectory[\"validate\"],\n            name=\"agent_loop\",\n            trace=trace,\n        ):\n            assert agent_name in _agent_loop_registry, (\n                f\"Agent loop {agent_name} not registered, registered agent loops: {_agent_loop_registry.keys()}\"\n            )\n\n            agent_loop_config = _agent_loop_registry[agent_name]\n            agent_loop = hydra.utils.instantiate(\n                config=agent_loop_config,\n                trainer_config=DictConfigWrap(config=self.config),\n                server_manager=self.server_manager,\n                tokenizer=self.tokenizer,\n                processor=self.processor,\n                dataset_cls=self.dataset_cls,\n                data_config=DictConfigWrap(self.config.data),\n            )\n            output: AgentLoopOutput = await agent_loop.run(sampling_params, **kwargs)\n            return await self._agent_loop_postprocess(output, **kwargs)\n\n    async def _agent_loop_postprocess(self, output, **kwargs) -> _InternalAgentLoopOutput:\n        \"\"\"Perform post-processing operations on the output of each individual agent loop.\"\"\"\n        output.extra_fields[\"raw_prompt\"] = kwargs[\"raw_prompt\"]\n\n        # Some AgentLoop may have already computed the reward score, e.g SWE-agent.\n\n        # NOTE: consistent with the legacy batch version of generate_sequences that existed in the\n        # deprecated vLLM SPMD rollout implementation.\n        # prompt_ids: left padded with zeros (e.g., [0,0,0,0,1,2,3,4])\n        # response_ids: right padded with zeros (e.g., [5,6,7,8,0,0,0,0])\n        # input_ids: concatenation of prompt + response\n        # Mask:\n        # For example, if the prompt is [1,2,3,4] and the response is [5,6,7,(tool start)8,9(tool end),10,11,12]\n        # - prompt_attention_mask: 0s for padding, 1s for tokens\n        #   e.g., [0,0,0,0,1,1,1,1]\n        # - response_attention_mask: 0s for padding, 1s for tokens\n        #   e.g., [1,1,1,1,1,1,1,1,1,1,1,0,0,0,0]\n        # attention_mask: concatenation of prompt_attention_mask and response_attention_mask\n        #   e.g., [0,0,0,0,1,1,1,1(prompt),1,1,1,1,1,1,1,1,1,1,1,0,0,0,0(response)]\n        # - response_mask: 1s for LLM generated tokens, 0 for tool response/padding tokens\n        #   e.g., [1,1,1,1,1,1,1,(tool start),0,0(tool end),1,1,0,0,0,0]\n        # - position_ids: sequential positions for tokens, starting at 0\n        #   e.g., [0,0,0,0,0,1,2,3,4,5,6,7,8,9,10,11,12,13,14,0,0,0,0]\n\n        # TODO(wuxibin): remove padding and use tensordict.\n        self.tokenizer.padding_side = \"left\"\n        prompt_output = self.tokenizer.pad(\n            {\"input_ids\": output.prompt_ids},\n            padding=\"max_length\",\n            max_length=self.rollout_config.prompt_length,\n            return_tensors=\"pt\",\n            return_attention_mask=True,\n        )\n        if prompt_output[\"input_ids\"].dim() == 1:\n            prompt_output[\"input_ids\"] = prompt_output[\"input_ids\"].unsqueeze(0)\n            prompt_output[\"attention_mask\"] = prompt_output[\"attention_mask\"].unsqueeze(0)\n\n        self.tokenizer.padding_side = \"right\"\n        response_output = self.tokenizer.pad(\n            {\"input_ids\": output.response_ids},\n            padding=\"max_length\",\n            max_length=self.rollout_config.response_length,\n            return_tensors=\"pt\",\n            return_attention_mask=True,\n        )\n        if response_output[\"input_ids\"].dim() == 1:\n            response_output[\"input_ids\"] = response_output[\"input_ids\"].unsqueeze(0)\n            response_output[\"attention_mask\"] = response_output[\"attention_mask\"].unsqueeze(0)\n\n        response_mask_output = self.tokenizer.pad(\n            {\"input_ids\": output.response_mask},\n            padding=\"max_length\",\n            max_length=self.rollout_config.response_length,\n            return_tensors=\"pt\",\n            return_attention_mask=False,\n        )\n        if response_mask_output[\"input_ids\"].dim() == 1:\n            response_mask_output[\"input_ids\"] = response_mask_output[\"input_ids\"].unsqueeze(0)\n\n        response_logprobs = None\n        if output.response_logprobs is not None:\n            pad_size = self.rollout_config.response_length - len(output.response_logprobs)\n            response_logprobs = torch.tensor(output.response_logprobs + [0.0] * pad_size).unsqueeze(0)\n\n        response_mask = response_mask_output[\"input_ids\"] * response_output[\"attention_mask\"]\n        attention_mask = torch.cat([prompt_output[\"attention_mask\"], response_output[\"attention_mask\"]], dim=1)\n        input_ids = torch.cat([prompt_output[\"input_ids\"], response_output[\"input_ids\"]], dim=1)\n\n        routed_experts = None\n        if output.routed_experts is not None:\n            total_length = input_ids.shape[1]\n            length, layer_num, topk_num = output.routed_experts.shape\n            if isinstance(output.routed_experts, np.ndarray):\n                routed_experts_array = output.routed_experts\n                if not routed_experts_array.flags.writeable:\n                    routed_experts_array = routed_experts_array.copy()\n                experts_tensor = torch.from_numpy(routed_experts_array)\n            elif isinstance(output.routed_experts, torch.Tensor):\n                experts_tensor = output.routed_experts\n            else:\n                raise TypeError(f\"Unsupported type for routed_experts: {type(output.routed_experts)}\")\n            routed_experts = torch.zeros(1, total_length, layer_num, topk_num, dtype=experts_tensor.dtype)\n\n            # Calculate start position: left padding means original prompt starts at the end\n            start_pos = prompt_output[\"input_ids\"].shape[1] - len(output.prompt_ids)\n            end_pos = min(start_pos + length, total_length)\n\n            # Add boundary checks for robustness\n            if start_pos < 0 or end_pos > total_length:\n                raise ValueError(\n                    f\"Invalid position range: start_pos={start_pos}, end_pos={end_pos}, total_length={total_length}\"\n                )\n\n            routed_experts[:, start_pos:end_pos] = experts_tensor.unsqueeze(0)\n\n        multi_modal_inputs = self._compute_multi_modal_inputs(output, input_ids)\n        position_ids = self._compute_position_ids(input_ids, attention_mask, multi_modal_inputs)\n        await self._compute_score(\n            output,\n            prompts=prompt_output[\"input_ids\"],\n            responses=response_output[\"input_ids\"],\n            attention_mask=attention_mask,\n            input_ids=input_ids,\n            position_ids=position_ids,\n            kwargs=kwargs,\n        )\n\n        return _InternalAgentLoopOutput(\n            prompt_ids=prompt_output[\"input_ids\"],\n            response_ids=response_output[\"input_ids\"],\n            input_ids=input_ids,\n            position_ids=position_ids,\n            response_mask=response_mask,\n            attention_mask=attention_mask,\n            response_logprobs=response_logprobs,\n            routed_experts=routed_experts,\n            multi_modal_inputs=multi_modal_inputs,\n            multi_modal_data=output.multi_modal_data,\n            reward_score=output.reward_score,\n            num_turns=output.num_turns,\n            metrics=output.metrics,\n            extra_fields=output.extra_fields,\n        )\n\n    def _compute_multi_modal_inputs(self, output, input_ids) -> dict[str, torch.Tensor]:\n        \"\"\"Compute multi-modal inputs with image and video.\"\"\"\n        multi_modal_inputs = {}\n        if self.processor is None:\n            return multi_modal_inputs\n\n        images = output.multi_modal_data.get(\"images\")\n        videos = output.multi_modal_data.get(\"videos\")\n        # split the videos and according metadatas\n        if videos is not None:\n            videos, video_metadatas = zip(*videos, strict=False)\n            videos, video_metadatas = list(videos), list(video_metadatas)\n        else:\n            video_metadatas = None\n        current_text = self.tokenizer.decode(input_ids.squeeze(0), skip_special_tokens=True)\n        multi_modal_inputs = self.processor(\n            text=[current_text],\n            images=images,\n            videos=videos,\n            video_metadata=video_metadatas,\n            return_tensors=\"pt\",\n            do_sample_frames=False,\n        )\n        multi_modal_inputs.pop(\"input_ids\", None)\n        multi_modal_inputs.pop(\"attention_mask\", None)\n\n        # We must use dict(multi_modal_inputs) to convert BatchFeature values to a new dict\n        # because np.array() only keeps the keys for BatchFeature.\n        multi_modal_inputs = dict(multi_modal_inputs.convert_to_tensors(\"pt\"))\n        image_grid_thw = multi_modal_inputs.get(\"image_grid_thw\")\n        if image_grid_thw is not None:\n            images_seqlens = torch.repeat_interleave(image_grid_thw[:, 1] * image_grid_thw[:, 2], image_grid_thw[:, 0])\n            multi_modal_inputs[\"images_seqlens\"] = images_seqlens\n        return multi_modal_inputs\n\n    def _compute_position_ids(self, input_ids, attention_mask, multi_modal_inputs) -> torch.Tensor:\n        \"\"\"Compute position ids for multi-modal inputs.\"\"\"\n        if self.processor is None:\n            return compute_position_id_with_mask(attention_mask)  # (1, seq_len)\n\n        multi_modal_kwargs = {\n            \"image_grid_thw\": multi_modal_inputs.get(\"image_grid_thw\"),\n            \"video_grid_thw\": multi_modal_inputs.get(\"video_grid_thw\"),\n        }\n        # For transformers>=5.3.0, mm_token_type_ids is only used to calculate position ids.\n        if multi_modal_inputs.pop(\"mm_token_type_ids\", None) is not None:\n            mm_token_type_ids = torch.zeros_like(input_ids)\n            mm_token_type_ids[0][input_ids[0] == self.processor.image_token_id] = 1\n            mm_token_type_ids[0][input_ids[0] == self.processor.video_token_id] = 2\n            multi_modal_kwargs[\"mm_token_type_ids\"] = mm_token_type_ids\n\n        # Model's get_rope_index has been dynamically bind to the processor.\n        vision_position_ids, _ = self.processor.get_rope_index(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            **multi_modal_kwargs,\n        )\n        vision_position_ids = vision_position_ids.transpose(0, 1)  # (3, 1, seq_len) => (1, 3, seq_len)\n\n        valid_mask = attention_mask[0].bool()\n        text_position_ids = torch.ones((1, len(input_ids[0])), dtype=torch.long)\n        text_position_ids[0, valid_mask] = torch.arange(valid_mask.sum().item())\n        text_position_ids = text_position_ids.unsqueeze(0)\n        position_ids = torch.cat((text_position_ids, vision_position_ids), dim=1)  # (1, 4, seq_length)\n        return position_ids\n\n    async def _compute_score(self, output, prompts, responses, attention_mask, input_ids, position_ids, kwargs):\n        \"\"\"Compute reward score for single sample.\"\"\"\n        enable_async_reward = self.reward_loop_worker_handles is not None\n\n        if output.reward_score is None and enable_async_reward:\n            batch = TensorDict(\n                {\n                    \"prompts\": prompts,  # [1, prompt_length]\n                    \"responses\": responses,  # [1, response_length]\n                    \"attention_mask\": attention_mask,  # [1, prompt_length + response_length]\n                    \"input_ids\": input_ids,  # [1, prompt_length + response_length]\n                    \"position_ids\": position_ids,\n                },\n                batch_size=1,\n            )\n            non_tensor_batch = {\n                **{k: np.array([v]) for k, v in kwargs.items()},\n                \"__num_turns__\": np.array([output.num_turns]),\n                \"tool_extra_fields\": np.array([output.extra_fields], dtype=object),\n            }\n\n            data = DataProto(\n                batch=batch,\n                non_tensor_batch=non_tensor_batch,\n            )\n            selected_reward_loop_worker_handle = random.choice(self.reward_loop_worker_handles)\n            result = await selected_reward_loop_worker_handle.compute_score.remote(data)\n            output.reward_score = result[\"reward_score\"]\n            output.extra_fields[\"reward_extra_info\"] = result[\"reward_extra_info\"]\n\n    def _postprocess(\n        self,\n        inputs: list[_InternalAgentLoopOutput],\n        input_non_tensor_batch: dict | None = None,\n    ) -> DataProto:\n        \"\"\"Process the padded outputs from _run_agent_loop and combine them into a batch.\"\"\"\n        # Convert lists back to tensors and stack them to create a batch.\n        prompt_ids = torch.cat([input.prompt_ids for input in inputs], dim=0)\n        response_ids = torch.cat([input.response_ids for input in inputs], dim=0)\n        response_mask = torch.cat([input.response_mask for input in inputs], dim=0)\n        attention_mask = torch.cat([input.attention_mask for input in inputs], dim=0)\n        input_ids = torch.cat([input.input_ids for input in inputs], dim=0)\n        position_ids = torch.cat([input.position_ids for input in inputs], dim=0)\n        optional_outputs = {}\n        if inputs[0].response_logprobs is not None:\n            optional_outputs[\"rollout_log_probs\"] = torch.cat([input.response_logprobs for input in inputs], dim=0)\n        if inputs[0].routed_experts is not None:\n            optional_outputs[\"routed_experts\"] = torch.cat([input.routed_experts for input in inputs], dim=0)\n\n        batch = TensorDict(\n            {\n                \"prompts\": prompt_ids,  # [bsz, prompt_length]\n                \"responses\": response_ids,  # [bsz, response_length]\n                \"response_mask\": response_mask,  # [bsz, response_length]\n                \"input_ids\": input_ids,  # [bsz, prompt_length + response_length]\n                \"attention_mask\": attention_mask,  # [bsz, prompt_length + response_length]\n                # position_ids: [bsz, 3, prompt_length + response_length] or [bsz, prompt_length + response_length]\n                \"position_ids\": position_ids,\n                **optional_outputs,\n            },\n            batch_size=len(inputs),\n        )\n\n        scores = [input.reward_score for input in inputs]\n        if all(score is not None for score in scores):\n            prompt_length = prompt_ids.size(1)\n            response_length = attention_mask[:, prompt_length:].sum(dim=1) - 1\n            rm_scores = torch.zeros_like(response_mask, dtype=torch.float32)\n            rm_scores[torch.arange(response_mask.size(0)), response_length] = torch.tensor(scores, dtype=torch.float32)\n            batch[\"rm_scores\"] = rm_scores\n\n        non_tensor_batch = {\n            \"__num_turns__\": np.array([input.num_turns for input in inputs], dtype=np.int32),\n        }\n        if self.reward_loop_worker_handles is None and input_non_tensor_batch:\n            non_tensor_batch.update(input_non_tensor_batch)\n\n        # add reward_extra_info to non_tensor_batch\n        reward_extra_infos = [input.extra_fields.get(\"reward_extra_info\", {}) for input in inputs]\n        reward_extra_keys = list(reward_extra_infos[0].keys())\n        for key in reward_extra_keys:\n            non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos])\n\n        # Add multi_modal_inputs to non_tensor_batch if any samples have them\n        multi_modal_inputs_list = [input.multi_modal_inputs for input in inputs]\n        if any(mmi is not None for mmi in multi_modal_inputs_list):\n            non_tensor_batch[\"multi_modal_inputs\"] = np.array(multi_modal_inputs_list, dtype=object)\n\n        metrics = [input.metrics.model_dump() for input in inputs]\n        # Collect extra fields from all inputs and convert them to np.ndarray\n        # Keep a stable set of keys so downstream batch concat stays consistent across agent loops.\n        extra_fields = {}\n        default_extra_keys = {\n            \"turn_scores\",\n            \"tool_rewards\",\n            \"min_global_steps\",\n            \"max_global_steps\",\n            \"extras\",\n        }\n        all_keys = set(key for input_item in inputs for key in input_item.extra_fields) | default_extra_keys\n        for key in all_keys:\n            temp_arr = np.empty(len(inputs), dtype=object)\n            temp_arr[:] = [input.extra_fields.get(key) for input in inputs]\n            extra_fields[key] = temp_arr\n\n        non_tensor_batch.update(extra_fields)\n\n        # Only include reward_extra_keys in meta_info if rm_scores is in batch\n        # This avoids conflicts when reward_tensor is merged later in ray_trainer.py\n        if \"rm_scores\" in batch.keys():\n            meta_info = {\"metrics\": metrics, \"reward_extra_keys\": reward_extra_keys}\n        else:\n            meta_info = {\"metrics\": metrics}\n\n        return DataProto(\n            batch=batch,\n            non_tensor_batch=non_tensor_batch,\n            meta_info=meta_info,\n        )\n\n\nasync def get_trajectory_info(step, index, validate):\n    \"\"\"Get trajectory info.\n\n    Args:\n        step (int): global steps in the trainer.\n        index (list): form datastore extra_info.index column.\n        validate (bool): whether is a validate step.\n\n    Returns:\n        list: trajectory.\n    \"\"\"\n    trajectory_info = []\n    rollout_n = 0\n    for i in range(len(index)):\n        if i > 0 and index[i - 1] == index[i]:\n            rollout_n += 1\n        else:\n            rollout_n = 0\n        trajectory_info.append({\"step\": step, \"sample_index\": index[i], \"rollout_n\": rollout_n, \"validate\": validate})\n    return trajectory_info\n\n\nclass AgentLoopManager:\n    \"\"\"Agent loop manager that manages a group of agent loop workers.\n\n    - if worker_group is not None, rollout server is in hybrid mode, share GPUs with training engine.\n    - otherwise, rollout server is in standalone mode, use separate GPUs, e.g., one-step-off/fully async training.\n\n    Args:\n        config (DictConfig): whole config for main entrypoint.\n        worker_group (RayWorkerGroup): ActorRolloutRef worker group for hybrid mode; None for standalone mode.\n        rollout_resource_pool (RayResourcePool): Resource pool for hybrid mode, only used by TensorRT-LLM.\n        reward_loop_worker_handles (List[ray.actor.ActorHandle]): Actor handles for streaming reward computation.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: DictConfig,\n        worker_group: RayWorkerGroup = None,\n        rollout_resource_pool: RayResourcePool = None,\n        reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,\n    ):\n        self.config = config\n        self.rollout_config, self.model_config = _get_rollout_and_model_config(config)\n        self.worker_group = worker_group\n        self.rollout_resource_pool = rollout_resource_pool\n        self.reward_loop_worker_handles = reward_loop_worker_handles\n\n        assert worker_group is not None or self.rollout_config.nnodes > 0, \"nnodes must be > 0 in standalone mode\"\n\n        # for recipe to change\n        if not hasattr(self, \"rollout_replica_class\"):\n            self.rollout_replica_class = get_rollout_replica_class(self.rollout_config.name)\n        if not hasattr(self, \"agent_loop_workers_class\"):\n            self.agent_loop_workers_class = ray.remote(AgentLoopWorker)\n\n    @classmethod\n    @auto_await\n    async def create(\n        cls,\n        config: DictConfig,\n        worker_group: RayWorkerGroup = None,\n        rollout_resource_pool: RayResourcePool = None,\n        reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,\n    ):\n        \"\"\"Create agent loop manager.\"\"\"\n        instance = cls(config, worker_group, rollout_resource_pool, reward_loop_worker_handles)\n        await instance._initialize_llm_servers()\n        await instance._init_global_load_balancer()\n        await instance._init_agent_loop_workers()\n        return instance\n\n    async def _initialize_llm_servers(self):\n        rollout_world_size = (\n            self.rollout_config.tensor_model_parallel_size\n            * self.rollout_config.data_parallel_size\n            * self.rollout_config.pipeline_model_parallel_size\n        )\n        world_size = (\n            self.worker_group.world_size\n            if self.worker_group\n            else self.rollout_config.n_gpus_per_node * self.rollout_config.nnodes\n        )\n        num_replicas = world_size // rollout_world_size\n\n        self.rollout_replicas = [\n            self.rollout_replica_class(\n                replica_rank=replica_rank,\n                config=self.rollout_config,\n                model_config=self.model_config,\n                gpus_per_node=self.rollout_config.n_gpus_per_node,\n            )\n            for replica_rank in range(num_replicas)\n        ]\n\n        if self.worker_group and self.rollout_config.name != \"trtllm\":\n            await asyncio.gather(*[server.init_hybrid(self.worker_group) for server in self.rollout_replicas])\n        # TODO: unify trtllm to init_hybrid\n        elif self.worker_group and self.rollout_config.name == \"trtllm\":\n            await asyncio.gather(\n                *[\n                    server.init_hybrid_colocated(self.worker_group, self.rollout_resource_pool)\n                    for server in self.rollout_replicas\n                ]\n            )\n        else:\n            await asyncio.gather(*[server.init_standalone() for server in self.rollout_replicas])\n\n        self.server_handles = [server._server_handle for server in self.rollout_replicas]\n        self.server_addresses = [server._server_address for server in self.rollout_replicas]\n\n        print(f\"AgentLoopManager: {self.server_addresses}\")\n\n        # Update Prometheus configuration with server addresses\n        if self.rollout_config.prometheus.enable:\n            if self.rollout_config.disable_log_stats:\n                raise ValueError(\"PROMETHEUS needs disable_log_stats==False, but it is currently True.\")\n            update_prometheus_config(self.rollout_config.prometheus, self.server_addresses, self.rollout_config.name)\n\n    async def _init_agent_loop_workers(self):\n        self.agent_loop_workers = []\n        num_workers = self.rollout_config.agent.num_workers\n        load_balancer_handle = self.global_load_balancer\n        servers = list(zip(self.server_addresses, self.server_handles, strict=True))\n\n        node_ids = [node[\"NodeID\"] for node in ray.nodes() if node[\"Alive\"] and node[\"Resources\"].get(\"CPU\", 0) > 0]\n        for i in range(num_workers):\n            # Round-robin scheduling over the all nodes\n            node_id = node_ids[i % len(node_ids)]\n            self.agent_loop_workers.append(\n                self.agent_loop_workers_class.options(\n                    name=f\"agent_loop_worker_{i}\" + f\"_{uuid4().hex[:8]}\",\n                    scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                        node_id=node_id, soft=True\n                    ),\n                ).remote(\n                    self.config,\n                    servers,\n                    load_balancer_handle,\n                    self.reward_loop_worker_handles,\n                )\n            )\n\n    async def _init_global_load_balancer(self) -> None:\n        self.global_load_balancer = GlobalRequestLoadBalancer.remote(\n            server_actor_ids=self.server_addresses,\n            max_cache_size=DEFAULT_ROUTING_CACHE_SIZE,\n        )\n\n    @auto_await\n    async def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Split input batch and dispatch to agent loop workers.\n\n        Args:\n            prompts (DataProto): Input batch.\n\n        Returns:\n            DataProto: Output batch.\n        \"\"\"\n\n        chunkes = prompts.chunk(len(self.agent_loop_workers))\n        outputs = await asyncio.gather(\n            *[\n                worker.generate_sequences.remote(chunk)\n                for worker, chunk in zip(self.agent_loop_workers, chunkes, strict=True)\n            ]\n        )\n        output = DataProto.concat(outputs)\n\n        # calculate performance metrics\n        metrics = [output.meta_info.pop(\"metrics\") for output in outputs]  # List[List[Dict[str, str]]]\n        timing = self._performance_metrics(metrics, output)\n\n        output.meta_info = {\"timing\": timing, **outputs[0].meta_info}\n        return output\n\n    def _performance_metrics(self, metrics: list[list[dict[str, str]]], output: DataProto) -> dict[str, float]:\n        timing = {}\n        t_generate_sequences = np.array([metric[\"generate_sequences\"] for chunk in metrics for metric in chunk])\n        t_tool_calls = np.array([metric[\"tool_calls\"] for chunk in metrics for metric in chunk])\n        num_preempted = np.array([metric[\"num_preempted\"] for chunk in metrics for metric in chunk])\n        timing[\"agent_loop/num_preempted/min\"] = num_preempted.min()\n        timing[\"agent_loop/num_preempted/max\"] = num_preempted.max()\n        timing[\"agent_loop/num_preempted/mean\"] = num_preempted.mean()\n        timing[\"agent_loop/generate_sequences/min\"] = t_generate_sequences.min()\n        timing[\"agent_loop/generate_sequences/max\"] = t_generate_sequences.max()\n        timing[\"agent_loop/generate_sequences/mean\"] = t_generate_sequences.mean()\n        timing[\"agent_loop/tool_calls/min\"] = t_tool_calls.min()\n        timing[\"agent_loop/tool_calls/max\"] = t_tool_calls.max()\n        timing[\"agent_loop/tool_calls/mean\"] = t_tool_calls.mean()\n\n        # batch sequence generation is bounded by the slowest sample\n        slowest = np.argmax(t_generate_sequences + t_tool_calls)\n        attention_mask = output.batch[\"attention_mask\"][slowest]\n        prompt_length = output.batch[\"prompts\"].shape[1]\n        timing[\"agent_loop/slowest/generate_sequences\"] = t_generate_sequences[slowest]\n        timing[\"agent_loop/slowest/tool_calls\"] = t_tool_calls[slowest]\n        timing[\"agent_loop/slowest/prompt_length\"] = attention_mask[:prompt_length].sum().item()\n        timing[\"agent_loop/slowest/response_length\"] = attention_mask[prompt_length:].sum().item()\n        timing[\"agent_loop/slowest/num_preempted\"] = num_preempted[slowest]\n\n        return timing\n\n    @auto_await\n    async def clear_kv_cache(self):\n        \"\"\"Clear all rollout kv cache, but don`t sleep.\"\"\"\n        await asyncio.gather(*[replica.clear_kv_cache() for replica in self.rollout_replicas])\n\n    @auto_await\n    async def start_profile(self, **kwargs):\n        \"\"\"Start profiling on all rollout replicas.\"\"\"\n        await asyncio.gather(*[replica.start_profile(**kwargs) for replica in self.rollout_replicas])\n\n    @auto_await\n    async def stop_profile(self):\n        \"\"\"Stop profiling on all rollout replicas.\"\"\"\n        await asyncio.gather(*[replica.stop_profile() for replica in self.rollout_replicas])\n"
  },
  {
    "path": "verl/experimental/agent_loop/prometheus_utils.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\nimport os\n\nimport ray\nimport yaml\n\nfrom verl.workers.config.rollout import PrometheusConfig\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef update_prometheus_config(config: PrometheusConfig, server_addresses: list[str], rollout_name: str | None = None):\n    \"\"\"\n    Update Prometheus configuration file with server addresses and reload on first node.\n\n    server_addresses: vllm or sglang server addresses\n\n    rollout_name: name of the rollout backend (e.g., \"vllm\", \"sglang\")\n    \"\"\"\n\n    if not server_addresses:\n        logger.warning(\"No server addresses available to update Prometheus config\")\n        return\n\n    try:\n        # Get Prometheus config file path from environment or use default\n        prometheus_config_json = {\n            \"global\": {\"scrape_interval\": \"10s\", \"evaluation_interval\": \"10s\"},\n            \"scrape_configs\": [\n                {\n                    \"job_name\": \"ray\",\n                    \"file_sd_configs\": [{\"files\": [\"/tmp/ray/prom_metrics_service_discovery.json\"]}],\n                },\n                {\"job_name\": \"rollout\", \"static_configs\": [{\"targets\": server_addresses}]},\n            ],\n        }\n\n        # Write configuration file to all nodes\n        @ray.remote(num_cpus=0)\n        def write_config_file(config_data, config_path):\n            os.makedirs(os.path.dirname(config_path), exist_ok=True)\n            with open(config_path, \"w\") as f:\n                yaml.dump(config_data, f, default_flow_style=False, indent=2)\n            return True\n\n        # Reload Prometheus on all nodes. Only master node should succeed, skip errors on other nodes.\n        @ray.remote(num_cpus=0)\n        def reload_prometheus(port):\n            import socket\n            import subprocess\n\n            hostname = socket.gethostname()\n            ip_address = socket.gethostbyname(hostname)\n\n            reload_url = f\"http://{ip_address}:{port}/-/reload\"\n\n            try:\n                subprocess.run([\"curl\", \"-X\", \"POST\", reload_url], capture_output=True, text=True, timeout=10)\n                print(f\"Reloading Prometheus on node: {reload_url}\")\n            except Exception:\n                # Skip errors on non-master nodes\n                pass\n\n        # Get all available nodes and schedule tasks on each node\n        nodes = ray.nodes()\n        alive_nodes = [node for node in nodes if node[\"Alive\"]]\n\n        # Write config files on all nodes\n        write_tasks = []\n        for node in alive_nodes:\n            node_ip = node[\"NodeManagerAddress\"]\n            task = write_config_file.options(\n                resources={\"node:\" + node_ip: 0.001}  # Schedule to specific node\n            ).remote(prometheus_config_json, config.file)\n            write_tasks.append(task)\n\n        ray.get(write_tasks)\n\n        server_type = rollout_name.upper() if rollout_name else \"rollout\"\n        print(f\"Updated Prometheus configuration at {config.file} with {len(server_addresses)} {server_type} servers\")\n\n        # Reload Prometheus on all nodes\n        reload_tasks = []\n        for node in alive_nodes:\n            node_ip = node[\"NodeManagerAddress\"]\n            task = reload_prometheus.options(\n                resources={\"node:\" + node_ip: 0.001}  # Schedule to specific node\n            ).remote(config.port)\n            reload_tasks.append(task)\n\n        ray.get(reload_tasks)\n\n    except Exception as e:\n        logger.error(f\"Failed to update Prometheus configuration: {e}\")\n"
  },
  {
    "path": "verl/experimental/agent_loop/single_turn_agent_loop.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\nfrom typing import Any\nfrom uuid import uuid4\n\nfrom verl.experimental.agent_loop.agent_loop import AgentLoopBase, AgentLoopOutput, register\nfrom verl.utils.profiler import simple_timer\nfrom verl.workers.rollout.replica import TokenOutput\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@register(\"single_turn_agent\")\nclass SingleTurnAgentLoop(AgentLoopBase):\n    \"\"\"Naive agent loop that only do single turn chat completion.\"\"\"\n\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n        self.prompt_length = self.rollout_config.prompt_length\n        self.response_length = self.rollout_config.response_length\n\n    async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:\n        messages = list(kwargs[\"raw_prompt\"])\n\n        # 1. extract images and videos from messages\n        multi_modal_data = await self.process_vision_info(messages)\n        images = multi_modal_data.get(\"images\")\n        videos = multi_modal_data.get(\"videos\")\n\n        # 2. apply chat template and tokenize\n        prompt_ids = await self.apply_chat_template(\n            messages,\n            images=images,\n            videos=videos,\n        )\n\n        # 3. generate sequences\n        metrics = {}\n        with simple_timer(\"generate_sequences\", metrics):\n            output: TokenOutput = await self.server_manager.generate(\n                request_id=uuid4().hex,\n                prompt_ids=prompt_ids,\n                sampling_params=sampling_params,\n                image_data=images,\n                video_data=videos,\n            )\n        if metrics.get(\"num_preempted\") is None:\n            metrics[\"num_preempted\"] = output.num_preempted if output.num_preempted is not None else -1\n        response_mask = [1] * len(output.token_ids)\n\n        output: AgentLoopOutput = AgentLoopOutput(\n            prompt_ids=prompt_ids,\n            response_ids=output.token_ids[: self.response_length],\n            response_mask=response_mask[: self.response_length],\n            response_logprobs=output.log_probs[: self.response_length] if output.log_probs else None,\n            routed_experts=(\n                output.routed_experts[: len(prompt_ids) + self.response_length]\n                if output.routed_experts is not None\n                else None\n            ),\n            multi_modal_data=multi_modal_data,\n            num_turns=2,\n            metrics=metrics,\n            extra_fields=output.extra_fields,\n        )\n\n        # keeping the schema consistent with tool_agent_loop\n        output.extra_fields.update({\"turn_scores\": [], \"tool_rewards\": []})\n\n        return output\n"
  },
  {
    "path": "verl/experimental/agent_loop/tool_agent_loop.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport json\nimport logging\nimport os\nfrom enum import Enum\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nimport torch\nfrom PIL import Image\n\nfrom verl.experimental.agent_loop.agent_loop import (\n    AgentLoopBase,\n    AgentLoopOutput,\n    register,\n)\nfrom verl.experimental.agent_loop.tool_parser import FunctionCall, ToolParser\nfrom verl.experimental.agent_loop.utils import build_gpt_oss_tool_response_text\nfrom verl.interactions.base import BaseInteraction\nfrom verl.interactions.utils.interaction_registry import initialize_interactions_from_config\nfrom verl.tools.schemas import ToolResponse\nfrom verl.tools.utils.tool_registry import initialize_tools_from_config\nfrom verl.utils.profiler import simple_timer\nfrom verl.utils.rollout_trace import rollout_trace_op\nfrom verl.workers.rollout.replica import TokenOutput\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass AgentState(Enum):\n    PENDING = \"pending\"\n    GENERATING = \"generating\"\n    PROCESSING_TOOLS = \"processing_tools\"\n    TERMINATED = \"terminated\"\n    INTERACTING = \"interacting\"\n\n\nclass AgentData:\n    \"\"\"Encapsulates all state variables for the agent loop. AgentData is passed to tool calling in case that\n    tool may need to access full history state. User can store any tool session data in `extra_fields`.\"\"\"\n\n    def __init__(\n        self,\n        messages: list[dict[str, Any]],\n        image_data: list[Image.Image],\n        video_data: list[tuple[torch.Tensor, dict[str, Any]]],\n        metrics: dict[str, Any],\n        request_id: str,\n        tools_kwargs: dict[str, Any],\n        interaction: Optional[BaseInteraction] = None,\n        interaction_kwargs: Optional[dict[str, Any]] = None,\n    ):\n        self.messages = messages\n        self.image_data = image_data\n        self.video_data = video_data\n        self.metrics = metrics\n        self.request_id = request_id\n        self.tools_kwargs = tools_kwargs\n        self.interaction = interaction\n        self.interaction_kwargs = interaction_kwargs or {}\n\n        # State variables\n        self.prompt_ids: list[int] = []\n        self.response_ids: list[int] = []\n        self.response_mask: list[int] = []\n        self.response_logprobs: list[float] = []\n        self.turn_scores: list[float] = []\n        self.tool_rewards: list[float] = []\n        self.user_turns = 0\n        self.assistant_turns = 0\n\n        # Temporary state for tool calls\n        self.tool_calls: list[FunctionCall] = []\n\n        self.routed_experts = None\n\n        # Extra fields for dynamic addition, e.g., tool session data\n        self.extra_fields: dict[str, Any] = {}\n\n\n@register(\"tool_agent\")\nclass ToolAgentLoop(AgentLoopBase):\n    def __init__(self, *args, **kwargs):\n        super().__init__(*args, **kwargs)\n\n        # Initialize tools from config file\n        self.max_user_turns = self.rollout_config.multi_turn.max_user_turns\n        self.max_assistant_turns = self.rollout_config.multi_turn.max_assistant_turns\n        self.max_parallel_calls = self.rollout_config.multi_turn.max_parallel_calls\n        self.max_tool_response_length = self.rollout_config.multi_turn.max_tool_response_length\n        self.tool_response_truncate_side = self.rollout_config.multi_turn.tool_response_truncate_side\n        tool_config_path = self.rollout_config.multi_turn.tool_config_path\n        tool_list = initialize_tools_from_config(tool_config_path) if tool_config_path else []\n        self.tools = {tool.name: tool for tool in tool_list}\n        self.tool_schemas = [tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list]\n        self.tool_parser = ToolParser.get_tool_parser(self.rollout_config.multi_turn.format, self.tokenizer)\n        self.tool_parser_name = self.rollout_config.multi_turn.format\n\n        self.prompt_length = self.rollout_config.prompt_length\n        self.response_length = self.rollout_config.response_length\n\n        # Initialize interactions from config file\n        self.interaction_config_file = self.rollout_config.multi_turn.interaction_config_path\n        if self.interaction_config_file:\n            self.interaction_map: dict[str, BaseInteraction] = self._initialize_interactions(\n                self.interaction_config_file\n            )\n\n    @rollout_trace_op\n    async def run(self, sampling_params: dict[str, Any], **kwargs) -> AgentLoopOutput:\n        messages = list(kwargs[\"raw_prompt\"])\n\n        # extract images and videos from messages\n        multi_modal_data = await self.process_vision_info(messages)\n        images = multi_modal_data.get(\"images\")\n        videos = multi_modal_data.get(\"videos\")\n\n        metrics = {}\n        request_id = uuid4().hex\n        tools_kwargs = kwargs.get(\"tools_kwargs\", {})\n\n        # Initialize interaction if needed\n        interaction = None\n        interaction_kwargs = {}\n        if self.interaction_config_file:\n            interaction_kwargs = kwargs[\"extra_info\"][\"interaction_kwargs\"]\n            if \"name\" not in interaction_kwargs:\n                raise ValueError(\"'name' key is required in interaction_kwargs\")\n            interaction_name = interaction_kwargs[\"name\"]\n            if interaction_name not in self.interaction_map:\n                raise ValueError(\n                    f\"Interaction '{interaction_name}' not found in interaction_map. Available interactions: \"\n                    f\"{list(self.interaction_map.keys())}\"\n                )\n            interaction = self.interaction_map[interaction_name]\n            await interaction.start_interaction(request_id, **interaction_kwargs)\n        # Create AgentData instance to encapsulate all state\n        agent_data = AgentData(\n            messages=messages,\n            image_data=images,\n            video_data=videos,\n            metrics=metrics,\n            request_id=request_id,\n            tools_kwargs=tools_kwargs,\n            interaction=interaction,\n            interaction_kwargs=interaction_kwargs,\n        )\n\n        # State machine loop\n        state = AgentState.PENDING\n        while state != AgentState.TERMINATED:\n            if state == AgentState.PENDING:\n                state = await self._handle_pending_state(agent_data, sampling_params)\n            elif state == AgentState.GENERATING:\n                state = await self._handle_generating_state(agent_data, sampling_params)\n            elif state == AgentState.PROCESSING_TOOLS:\n                state = await self._handle_processing_tools_state(agent_data)\n            elif state == AgentState.INTERACTING:\n                state = await self._handle_interacting_state(agent_data)\n            else:\n                logger.error(f\"Invalid state: {state}\")\n                state = AgentState.TERMINATED\n\n        # Finalize output\n        response_ids = agent_data.prompt_ids[-len(agent_data.response_mask) :]\n        prompt_ids = agent_data.prompt_ids[: len(agent_data.prompt_ids) - len(agent_data.response_mask)]\n        multi_modal_data = {}\n        if agent_data.image_data is not None:\n            multi_modal_data[\"images\"] = agent_data.image_data\n        if agent_data.video_data is not None:\n            multi_modal_data[\"videos\"] = agent_data.video_data\n\n        output: AgentLoopOutput = AgentLoopOutput(\n            prompt_ids=prompt_ids,\n            response_ids=response_ids[: self.response_length],\n            response_mask=agent_data.response_mask[: self.response_length],\n            multi_modal_data=multi_modal_data,\n            response_logprobs=agent_data.response_logprobs[: self.response_length]\n            if agent_data.response_logprobs\n            else None,\n            num_turns=agent_data.user_turns + agent_data.assistant_turns + 1,\n            metrics=agent_data.metrics,\n            routed_experts=agent_data.routed_experts,\n            extra_fields=agent_data.extra_fields,\n        )\n        output.extra_fields.update({\"turn_scores\": agent_data.turn_scores, \"tool_rewards\": agent_data.tool_rewards})\n        return output\n\n    async def _handle_pending_state(self, agent_data: AgentData, sampling_params: dict[str, Any]) -> AgentState:\n        \"\"\"Handle the pending state: prepare the prompt and start generation.\"\"\"\n        prompt_ids = await self.apply_chat_template(\n            agent_data.messages,\n            tools=self.tool_schemas,\n            images=agent_data.image_data,\n            videos=agent_data.video_data,\n        )\n        agent_data.prompt_ids = prompt_ids\n        return AgentState.GENERATING\n\n    async def _handle_generating_state(\n        self, agent_data: AgentData, sampling_params: dict[str, Any], ignore_termination: bool = False\n    ) -> AgentState:\n        \"\"\"Handle the generating state: generate model response and check for tool calls.\"\"\"\n        add_messages: list[dict[str, Any]] = []\n\n        with simple_timer(\"generate_sequences\", agent_data.metrics):\n            output: TokenOutput = await self.server_manager.generate(\n                request_id=agent_data.request_id,\n                prompt_ids=agent_data.prompt_ids,\n                sampling_params=sampling_params,\n                image_data=agent_data.image_data,\n                video_data=agent_data.video_data,\n            )\n        # first time to set num_preempted\n        if agent_data.metrics.get(\"num_preempted\") is None:\n            agent_data.metrics[\"num_preempted\"] = output.num_preempted if output.num_preempted is not None else -1\n        # then add num_preempted to the metrics\n        else:\n            agent_data.metrics[\"num_preempted\"] += output.num_preempted if output.num_preempted is not None else 0\n\n        if not agent_data.extra_fields:\n            agent_data.extra_fields.update(output.extra_fields)\n        else:\n            # Multi-round calls, only update the maximum max_global_steps.\n            max_global_steps = output.extra_fields.get(\"max_global_steps\", None)\n            if max_global_steps:\n                agent_data.extra_fields[\"max_global_steps\"] = max_global_steps\n\n        agent_data.assistant_turns += 1\n        agent_data.response_ids = output.token_ids\n        agent_data.prompt_ids += agent_data.response_ids\n        agent_data.response_mask += [1] * len(agent_data.response_ids)\n        if output.log_probs:\n            agent_data.response_logprobs += output.log_probs\n\n        if output.routed_experts is not None:\n            agent_data.routed_experts = output.routed_experts\n\n        # Check termination conditions\n        if not ignore_termination and len(agent_data.response_mask) >= self.response_length:\n            return AgentState.TERMINATED\n        if self.max_assistant_turns and agent_data.assistant_turns >= self.max_assistant_turns:\n            return AgentState.TERMINATED\n        if self.max_user_turns and agent_data.user_turns >= self.max_user_turns:\n            return AgentState.TERMINATED\n\n        # Extract tool calls\n        tools = [tool.tool_schema for tool in self.tools.values()]\n        _, agent_data.tool_calls = await self.tool_parser.extract_tool_calls(agent_data.response_ids, tools)\n\n        # Handle interaction if needed\n        if self.interaction_config_file:\n            assistant_message = await self.loop.run_in_executor(\n                None, lambda: self.tokenizer.decode(agent_data.response_ids, skip_special_tokens=True)\n            )\n            add_messages.append({\"role\": \"assistant\", \"content\": assistant_message})\n            agent_data.messages.extend(add_messages)\n\n        # Determine next state\n        if agent_data.tool_calls:\n            return AgentState.PROCESSING_TOOLS\n        elif self.interaction_config_file:\n            return AgentState.INTERACTING\n        else:\n            return AgentState.TERMINATED\n\n    async def _handle_processing_tools_state(self, agent_data: AgentData) -> AgentState:\n        \"\"\"Handle the processing tools state: execute tool calls and prepare tool responses.\"\"\"\n        add_messages: list[dict[str, Any]] = []\n        new_images_this_turn: list[Any] = []  # Local variable instead of agent_data attribute\n\n        tasks = []\n        tool_call_names = []\n        for tool_call in agent_data.tool_calls[: self.max_parallel_calls]:\n            tasks.append(self._call_tool(tool_call, agent_data.tools_kwargs, agent_data))\n            tool_call_names.append(tool_call.name)\n\n        with simple_timer(\"tool_calls\", agent_data.metrics):\n            responses = await asyncio.gather(*tasks)\n\n        # Process tool responses and update multi_modal_data\n        # Removed: agent_data.new_images_this_turn = []\n        for tool_response, tool_reward, _ in responses:\n            # Create message from tool response\n            if tool_response.image or tool_response.video:\n                # Multi-modal content with structured format\n                if not getattr(self.processor, \"image_processor\", None):\n                    raise ValueError(\n                        \"Multimedia data can only be processed by `processor`, but the processor is None. \"\n                        \"This error is often caused if you are using a LLM model but your tool returns multimodal \"\n                        \"data. Plase use a vlm as the base model.\"\n                    )\n                content = []\n                if tool_response.image:\n                    content.append({\"type\": \"image\"})\n                if tool_response.video:\n                    content.append({\"type\": \"video\"})\n                if tool_response.text:\n                    content.append({\"type\": \"text\", \"text\": tool_response.text})\n                message = {\"role\": \"tool\", \"content\": content}\n            else:\n                # Text-only content\n                message = {\"role\": \"tool\", \"content\": tool_response.text or \"\"}\n\n            add_messages.append(message)\n\n            # Handle image data\n            if tool_response.image:\n                # Add new image data\n                if isinstance(tool_response.image, list):\n                    # Ensure all elements in the list are valid image objects\n                    for img in tool_response.image:\n                        if img is not None:  # Add a check to ensure the image is not None\n                            new_images_this_turn.append(img)  # Using local variable\n                else:\n                    # Ensure the image is not None\n                    if tool_response.image is not None:\n                        new_images_this_turn.append(tool_response.image)  # Using local variable\n\n            # Handle video data\n            if tool_response.video:\n                # Currently not supported, raise informative error\n                logger.warning(\"Multimedia type 'video' is not currently supported. Only 'image' is supported.\")\n                raise NotImplementedError(\n                    \"Multimedia type 'video' is not currently supported. Only 'image' is supported.\"\n                )\n\n            if tool_reward is not None:\n                agent_data.tool_rewards.append(tool_reward)\n\n        agent_data.messages.extend(add_messages)\n\n        if self.tool_parser_name == \"gpt-oss\":\n            logger.info(\"manually format tool responses for gpt-oss\")\n            tool_response_text = build_gpt_oss_tool_response_text(add_messages, tool_call_names)\n            response_ids = await self.loop.run_in_executor(\n                None, lambda: self.tokenizer.encode(tool_response_text, add_special_tokens=False)\n            )\n        else:\n            # Note that we have to pass None to the images and videos if there are no new images / videos\n            # to stay compatible with downstream image processing logic!\n            images = new_images_this_turn if new_images_this_turn else None\n            videos = None\n            response_ids = await self.apply_chat_template(\n                add_messages,\n                images=images,\n                videos=videos,\n                remove_system_prompt=True,\n            )\n\n        if len(agent_data.response_mask) + len(response_ids) >= self.response_length:\n            return AgentState.TERMINATED\n        # Update prompt_ids and response_mask\n\n        if new_images_this_turn:\n            if agent_data.image_data is None:\n                agent_data.image_data = []\n            elif not isinstance(agent_data.image_data, list):\n                agent_data.image_data = [agent_data.image_data]\n            for img in new_images_this_turn:\n                agent_data.image_data.append(img)\n\n        agent_data.prompt_ids += response_ids\n        agent_data.response_mask += [0] * len(response_ids)\n        if agent_data.response_logprobs:\n            agent_data.response_logprobs += [0.0] * len(response_ids)\n        agent_data.user_turns += 1\n        return AgentState.GENERATING\n\n    async def _handle_interacting_state(self, agent_data: AgentData) -> AgentState:\n        \"\"\"Handle the interacting state: get user input from interaction.\"\"\"\n        (\n            should_terminate_sequence,\n            interaction_responses,\n            reward,\n            metrics,\n        ) = await agent_data.interaction.generate_response(\n            agent_data.request_id, agent_data.messages, **agent_data.interaction_kwargs\n        )\n        agent_data.user_turns += 1\n\n        add_messages: list[dict[str, Any]] = [{\"role\": \"user\", \"content\": interaction_responses}]\n        agent_data.messages.extend(add_messages)\n\n        if reward is not None:\n            agent_data.turn_scores.append(reward)\n\n        # Update prompt with user responses (similar to _handle_processing_tools_state)\n        response_ids = await self.apply_chat_template(\n            add_messages,\n            remove_system_prompt=True,\n        )\n\n        # Update prompt_ids and response_mask\n        agent_data.prompt_ids += response_ids\n        agent_data.response_mask += [0] * len(response_ids)\n        if agent_data.response_logprobs:\n            agent_data.response_logprobs += [0.0] * len(response_ids)\n\n        # double check prompt\n        # Check termination condition\n        if should_terminate_sequence:\n            return AgentState.TERMINATED\n        else:\n            return AgentState.GENERATING\n\n    async def _call_tool(\n        self, tool_call: FunctionCall, tools_kwargs: dict[str, Any], agent_data: AgentData\n    ) -> tuple[ToolResponse, float, dict]:\n        \"\"\"Call tool and return tool response.\"\"\"\n        tool, instance_id = None, None\n        try:\n            # TODO: append malformed tool_call to the prompt: invalid function name or arguments\n            tool_name = tool_call.name\n            tool_args = json.loads(tool_call.arguments)\n            tool = self.tools[tool_name]\n            kwargs = tools_kwargs.get(tool_name, {})\n            instance_id, _ = await tool.create(create_kwargs=kwargs.get(\"create_kwargs\", {}))\n            tool_execution_response, tool_reward, res = await tool.execute(\n                instance_id, tool_args, agent_data=agent_data\n            )\n        except Exception as e:\n            logger.warning(f\"Error when executing tool: {e}\")\n            return (\n                ToolResponse(\n                    text=f\"Error when executing tool: {e}\",\n                ),\n                0.0,\n                {},\n            )\n        finally:\n            if tool and instance_id:\n                await tool.release(instance_id)\n\n        tool_response_text = tool_execution_response.text\n        if tool_response_text and len(tool_response_text) > self.max_tool_response_length:\n            if self.tool_response_truncate_side == \"left\":\n                tool_response_text = tool_response_text[: self.max_tool_response_length] + \"...(truncated)\"\n            elif self.tool_response_truncate_side == \"right\":\n                tool_response_text = \"(truncated)...\" + tool_response_text[-self.max_tool_response_length :]\n            else:\n                length = self.max_tool_response_length // 2\n                tool_response_text = tool_response_text[:length] + \"...(truncated)...\" + tool_response_text[-length:]\n\n        # Create ToolResponse from tool execution result\n        tool_response_kwargs = {\"text\": tool_response_text}\n\n        # Add multimedia data if present\n        for attr_name in [\"image\", \"video\"]:\n            if hasattr(tool_execution_response, attr_name):\n                attr_value = getattr(tool_execution_response, attr_name)\n                if attr_value is not None:\n                    tool_response_kwargs[attr_name] = attr_value\n\n        return ToolResponse(**tool_response_kwargs), tool_reward, res\n\n    def _initialize_interactions(self, interaction_config_file):\n        \"\"\"Initialize interactions from configuration.\n        Returns:\n            dict[str, BaseInteraction]: A dictionary mapping interaction names to interaction instances.\n        \"\"\"\n        if interaction_config_file is None:\n            return {}\n\n        interaction_map = initialize_interactions_from_config(interaction_config_file)\n        return interaction_map\n"
  },
  {
    "path": "verl/experimental/agent_loop/tool_parser.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport logging\nimport os\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Optional\n\nimport regex\nfrom pydantic import BaseModel\n\nfrom verl.tools.schemas import OpenAIFunctionToolSchema\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass FunctionCall(BaseModel):\n    arguments: str\n    \"\"\"\n    The arguments to call the function with, as generated by the model in JSON\n    format. Note that the model does not always generate valid JSON, and may\n    hallucinate parameters not defined by your function schema. Validate the\n    arguments in your code before calling your function.\n    \"\"\"\n\n    name: str\n    \"\"\"The name of the function to call.\"\"\"\n\n\nclass ToolParser(ABC):\n    _registry: dict[str, type[\"ToolParser\"]] = {}\n\n    def __init__(self, tokenizer) -> None:\n        self.tokenizer = tokenizer\n\n    @abstractmethod\n    async def extract_tool_calls(\n        self, responses_ids: list[int], tools: list[OpenAIFunctionToolSchema] = None\n    ) -> tuple[str, list[FunctionCall]]:\n        \"\"\"Extract tool calls from the responses.\n\n        Args:\n            responses_ids (List[int]): The ids of the responses.\n            tools (List[OpenAIFunctionToolSchema], optional): OpenAI function tool schema.\n\n        Returns:\n            Tuple[str, List[FunctionCall]]: Content and extracted tool calls.\n        \"\"\"\n        raise NotImplementedError\n\n    @classmethod\n    def get_tool_parser(cls, name: str, tokenizer):\n        if name not in cls._registry:\n            raise ValueError(f\"Unknown tool parser: {name}\")\n        return cls._registry[name](tokenizer)\n\n    @classmethod\n    def register(cls, name: str):\n        def decorator(subclass: type[ToolParser]) -> type[ToolParser]:\n            cls._registry[name] = subclass\n            return subclass\n\n        return decorator\n\n\n@ToolParser.register(\"hermes\")\nclass HermesToolParser(ToolParser):\n    \"\"\"Adapted from https://github.com/vllm-project/vllm/blob/v0.9.1/vllm/entrypoints/openai/tool_parsers/hermes_tool_parser.py\"\"\"\n\n    def __init__(self, tokenizer) -> None:\n        super().__init__(tokenizer)\n\n        self.tool_call_start_token: str = \"<tool_call>\"\n        self.tool_call_end_token: str = \"</tool_call>\"\n        self.tool_call_regex = regex.compile(r\"<tool_call>(.*?)</tool_call>\", regex.DOTALL)\n\n    @rollout_trace_op\n    async def extract_tool_calls(\n        self, responses_ids: list[int], tools: list[OpenAIFunctionToolSchema] = None\n    ) -> tuple[str, list[FunctionCall]]:\n        loop = get_event_loop()\n        text = await loop.run_in_executor(None, self.tokenizer.decode, responses_ids)\n        if self.tool_call_start_token not in text or self.tool_call_end_token not in text:\n            return text, []\n\n        matches = self.tool_call_regex.findall(text)\n        function_calls = []\n        for match in matches:\n            try:\n                function_call = json.loads(match)\n                name, arguments = function_call[\"name\"], function_call[\"arguments\"]\n                function_calls.append(FunctionCall(name=name, arguments=json.dumps(arguments, ensure_ascii=False)))\n            except Exception as e:\n                logger.error(f\"Failed to decode tool call: {e}\")\n\n        # remaing text exclude tool call tokens\n        content = self.tool_call_regex.sub(\"\", text)\n\n        return content, function_calls\n\n\n@ToolParser.register(\"gpt-oss\")\nclass GptOssToolParser(ToolParser):\n    \"\"\"\n    Tool parser for gpt-oss model.\n    Adapted from https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/function_call/gpt_oss_detector.py\n\n    Args:\n        tokenizer: The tokenizer to use.\n    \"\"\"\n\n    def __init__(self, tokenizer) -> None:\n        super().__init__(tokenizer)\n        # check https://cookbook.openai.com/articles/openai-harmony for more details.\n        self.cot_pattern = regex.compile(\n            r\"<\\|start\\|>assistant<\\|channel\\|>analysis<\\|message\\|>.*?<\\|end\\|>\", regex.DOTALL\n        )\n        # <|start|>assistant may be pre-appended in prompts, so we need to remove it.\n        self.partial_cot_pattern = regex.compile(r\"<\\|channel\\|>analysis<\\|message\\|>(.*?)<\\|end\\|>\", regex.DOTALL)\n        self.tool_call_pattern = regex.compile(\n            r\"<\\|start\\|>assistant<\\|channel\\|>[^<]* to=functions\\.([^<]+) \"\n            r\"<\\|constrain\\|>json<\\|message\\|>(.*?)<\\|call\\|>\",\n            regex.DOTALL,\n        )\n\n    @rollout_trace_op\n    async def extract_tool_calls(\n        self, responses_ids: list[int], tools: list[OpenAIFunctionToolSchema] = None\n    ) -> tuple[str, list[FunctionCall]]:\n        loop = get_event_loop()\n        # We need to keep special tokens for gpt-oss model for better tool call extraction.\n        text = await loop.run_in_executor(None, lambda: self.tokenizer.decode(responses_ids, skip_special_tokens=False))\n        # Need to remove padding tokens for better tool call extraction.\n        text = text.replace(self.tokenizer.pad_token, \"\")\n        # Need to reomve COT since COT may contain tool call tokens.But they are not valid tool calls.\n        text = regex.sub(self.cot_pattern, \"\", text)\n        text = regex.sub(self.partial_cot_pattern, \"\", text)\n\n        # check if there are tool calls in the text by re.findall\n        matches = regex.findall(self.tool_call_pattern, text)\n        if not matches:\n            return text, []\n\n        function_calls = []\n        for match in matches:\n            try:\n                name, arguments = match[0], match[1]\n                # don't check if arguments is valid JSON and leave it to client\n                function_calls.append(FunctionCall(name=name, arguments=arguments))\n            except Exception as e:\n                logger.error(f\"Failed to decode tool call: {e}\")\n\n        # remaing text exclude tool call tokens\n        content = regex.sub(self.tool_call_pattern, \"\", text)\n\n        return content, function_calls\n\n\n@ToolParser.register(\"qwen3_coder\")\nclass Qwen3XMLToolParser(ToolParser):\n    \"\"\"\n    Tool parser for qwen3_coder/qwen3.5 model.\n    Adapted from https://huggingface.co/Qwen/Qwen3-Coder-30B-A3B-Instruct/blob/main/qwen3coder_tool_parser.py\n\n    Args:\n        tokenizer: The tokenizer to use.\n    \"\"\"\n\n    def __init__(self, tokenizer):\n        super().__init__(tokenizer)\n\n        self.tool_call_start_token: str = \"<tool_call>\"\n        self.tool_call_end_token: str = \"</tool_call>\"\n        self.tool_call_prefix: str = \"<function=\"\n\n        self.tool_call_complete_regex = regex.compile(r\"<tool_call>(.*?)</tool_call>\", regex.DOTALL)\n        self.tool_call_regex = regex.compile(r\"<tool_call>(.*?)</tool_call>|<tool_call>(.*?)$\", regex.DOTALL)\n        self.tool_call_function_regex = regex.compile(r\"<function=(.*?)</function>|<function=(.*)$\", regex.DOTALL)\n        self.tool_call_parameter_regex = regex.compile(r\"<parameter=(.*?)</parameter>|<parameter=(.*?)$\", regex.DOTALL)\n\n    def _parse_xml_function_call(\n        self, function_call_str: str, tools: Optional[list[OpenAIFunctionToolSchema]]\n    ) -> FunctionCall:\n        def get_arguments_config(func_name: str) -> dict:\n            for config in tools:\n                if config.type == \"function\" and config.function.name == func_name:\n                    properties = config.function.parameters.properties\n                    return {k: v.model_dump() for k, v in properties.items()}\n            logger.warning(f\"Tool '{func_name}' is not defined in the tools list.\")\n            return {}\n\n        def convert_param_value(param_value: str, param_name: str, param_config: dict, func_name: str) -> Any:\n            # Handle null value for any type\n            if param_value.lower() == \"null\":\n                return None\n\n            if param_name not in param_config:\n                if param_config != {}:\n                    logger.warning(\n                        f\"Parsed parameter '{param_name}' is not defined in the tool \"\n                        f\"parameters for tool '{func_name}', directly returning the string value.\"\n                    )\n                return param_value\n\n            if isinstance(param_config[param_name], dict) and \"type\" in param_config[param_name]:\n                param_type = str(param_config[param_name][\"type\"]).strip().lower()\n            else:\n                param_type = \"string\"\n            if param_type in [\"string\", \"str\", \"text\", \"varchar\", \"char\", \"enum\"]:\n                return param_value\n            elif (\n                param_type.startswith(\"int\")\n                or param_type.startswith(\"uint\")\n                or param_type.startswith(\"long\")\n                or param_type.startswith(\"short\")\n                or param_type.startswith(\"unsigned\")\n            ):\n                try:\n                    param_value = int(param_value)\n                except Exception:\n                    logger.warning(\n                        f\"Parsed value '{param_value}' of parameter '{param_name}' is not an integer in tool \"\n                        f\"'{func_name}', degenerating to string.\"\n                    )\n                return param_value\n            elif param_type.startswith(\"num\") or param_type.startswith(\"float\"):\n                try:\n                    float_param_value = float(param_value)\n                    param_value = (\n                        float_param_value if float_param_value - int(float_param_value) != 0 else int(float_param_value)\n                    )\n                except Exception:\n                    logger.warning(\n                        f\"Parsed value '{param_value}' of parameter '{param_name}' is not a float in tool \"\n                        f\"'{func_name}', degenerating to string.\"\n                    )\n                return param_value\n            elif param_type in [\"boolean\", \"bool\", \"binary\"]:\n                param_value = param_value.lower()\n                if param_value not in [\"true\", \"false\"]:\n                    logger.warning(\n                        f\"Parsed value '{param_value}' of parameter '{param_name}' is not a \"\n                        f\"boolean (`true` of `false`) in tool '{func_name}', degenerating to false.\"\n                    )\n                return param_value == \"true\"\n            else:\n                if param_type == \"object\" or param_type.startswith(\"dict\"):\n                    try:\n                        param_value = json.loads(param_value)\n                        return param_value\n                    except Exception:\n                        logger.warning(\n                            f\"Parsed value '{param_value}' of parameter '{param_name}' is not a valid \"\n                            f\"JSON object in tool '{func_name}', will try other methods to parse it.\"\n                        )\n                try:\n                    param_value = eval(param_value)\n                except Exception:\n                    logger.warning(\n                        f\"Parsed value '{param_value}' of parameter '{param_name}' cannot be converted \"\n                        f\"via Python `eval()` in tool '{func_name}', degenerating to string.\"\n                    )\n                return param_value\n\n        # Extract function name\n        end_index = function_call_str.index(\">\")\n        function_name = function_call_str[:end_index]\n        param_config = get_arguments_config(function_name)\n        parameters = function_call_str[end_index + 1 :]\n        param_dict = {}\n        for match in self.tool_call_parameter_regex.findall(parameters):\n            match_text = match[0] if match[0] else match[1]\n            idx = match_text.index(\">\")\n            param_name = match_text[:idx]\n            param_value = str(match_text[idx + 1 :])\n            # Remove prefix and trailing \\n\n            if param_value.startswith(\"\\n\"):\n                param_value = param_value[1:]\n            if param_value.endswith(\"\\n\"):\n                param_value = param_value[:-1]\n\n            param_dict[param_name] = convert_param_value(param_value, param_name, param_config, function_name)\n        return FunctionCall(name=function_name, arguments=json.dumps(param_dict, ensure_ascii=False))\n\n    def _get_function_calls(self, model_output: str) -> list[str]:\n        # Find all tool calls\n        matched_ranges = self.tool_call_regex.findall(model_output)\n        raw_tool_calls = [match[0] if match[0] else match[1] for match in matched_ranges]\n\n        # Back-off strategy if no tool_call tags found\n        if len(raw_tool_calls) == 0:\n            raw_tool_calls = [model_output]\n\n        raw_function_calls = []\n        for tool_call in raw_tool_calls:\n            raw_function_calls.extend(self.tool_call_function_regex.findall(tool_call))\n\n        function_calls = [match[0] if match[0] else match[1] for match in raw_function_calls]\n        return function_calls\n\n    @rollout_trace_op\n    async def extract_tool_calls(\n        self, responses_ids: list[int], tools: list[OpenAIFunctionToolSchema] = None\n    ) -> tuple[str, list[FunctionCall]]:\n        loop = get_event_loop()\n        text = await loop.run_in_executor(None, self.tokenizer.decode, responses_ids)\n        if self.tool_call_start_token not in text:\n            return text, []\n\n        try:\n            function_calls = self._get_function_calls(text)\n            if len(function_calls) == 0:\n                return text, []\n\n            tool_calls = [\n                self._parse_xml_function_call(function_call_str, tools) for function_call_str in function_calls\n            ]\n\n            # Extract content before tool calls\n            content_index = text.find(self.tool_call_start_token)\n            content_index = content_index if content_index >= 0 else text.find(self.tool_call_prefix)\n            content = text[:content_index]  # .rstrip()\n\n            return content, tool_calls\n        except Exception as e:\n            logger.exception(f\"Error in extracting tool call from response: {e}\")\n            return text, []\n"
  },
  {
    "path": "verl/experimental/agent_loop/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom typing import Any\n\n\ndef resolve_config_path(config_path: str) -> str:\n    \"\"\"Resolve agent loop configuration file path.\n\n    In multi-node Ray training, relative paths may not resolve correctly\n    because the working directory on remote nodes can differ from the driver node.\n    This function resolves relative paths by checking multiple locations in order:\n    1. If already absolute, return as-is\n    2. Try current working directory\n    3. Try relative to verl package installation (project root)\n\n    Args:\n        config_path: Configuration file path (relative or absolute)\n\n    Returns:\n        Absolute path to the configuration file\n\n    Raises:\n        FileNotFoundError: If the configuration file cannot be found\n    \"\"\"\n    # Return absolute paths unchanged\n    if os.path.isabs(config_path):\n        return config_path\n\n    # Try current working directory first\n    cwd = os.path.abspath(os.getcwd())\n    cwd_path = os.path.abspath(os.path.join(cwd, config_path))\n    if (cwd_path == cwd or cwd_path.startswith(cwd + os.sep)) and os.path.exists(cwd_path):\n        return cwd_path\n\n    # Try relative to verl project root (where verl package is installed)\n    try:\n        import verl\n\n        verl_package_dir = os.path.abspath(os.path.dirname(verl.__file__))\n\n        # Strategy 1: For development/editable installs.\n        project_root = os.path.dirname(verl_package_dir)\n        dev_path = os.path.abspath(os.path.join(project_root, config_path))\n        if (dev_path == project_root or dev_path.startswith(project_root + os.sep)) and os.path.exists(dev_path):\n            return dev_path\n\n        # Strategy 2: For standard package installations.\n        install_path = os.path.abspath(os.path.join(verl_package_dir, config_path))\n        if (install_path == verl_package_dir or install_path.startswith(verl_package_dir + os.sep)) and os.path.exists(\n            install_path\n        ):\n            return install_path\n    except (ImportError, AttributeError):\n        pass  # verl not installed or __file__ not available\n\n    # File not found - raise clear error\n    raise FileNotFoundError(\n        f\"Agent loop configuration file not found: {config_path}. Tried current directory and verl project root.\"\n    )\n\n\n# tokenizer.apply_chat_template is not working properly for gpt-oss model.\n# Because the chat template requires tool call messages to parse tool response messages\n# so we need to format the tool response manually.\ndef format_gpt_oss_tool_response_manually(tool_response: str, tool_call_name: str) -> str:\n    \"\"\"Format tool response for gpt-oss model.\n    Args:\n        tool_response: Tool response string\n        tool_call_name: Name of the tool that was called\n\n    Returns:\n        Formatted tool response string\n    \"\"\"\n    return f\"<|start|>functions.{tool_call_name} to=assistant<|channel|>commentary<|message|>{tool_response}<|end|>\"\n\n\ndef add_generation_prompt_for_gpt_oss(message_content: str) -> str:\n    \"\"\"Add generation prompt for gpt-oss model.\n    Args:\n        message_content: Message content string\n\n    Returns:\n        Message content string with generation prompt\n    \"\"\"\n    return message_content + \"<|start|>assistant\"\n\n\ndef build_gpt_oss_tool_response_text(messages: list[dict[str, Any]], tool_call_names: list[str]) -> str:\n    \"\"\"Build gpt-oss tool response text (manual formatting + generation prompt).\"\"\"\n    tool_response_texts: list[str] = []\n    for i, tool_msg in enumerate(messages):\n        actual_tool_name = tool_call_names[i]\n        formatted = format_gpt_oss_tool_response_manually(tool_msg[\"content\"], actual_tool_name)\n        tool_response_texts.append(formatted)\n    return add_generation_prompt_for_gpt_oss(\"\".join(tool_response_texts))\n"
  },
  {
    "path": "verl/experimental/dataset/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/dataset/sampler.py",
    "content": "# Copyright 2025 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom abc import abstractmethod\nfrom collections.abc import Sized\n\nfrom omegaconf import DictConfig\nfrom torch.utils.data import Sampler\n\nfrom verl import DataProto\n\n\nclass AbstractSampler(Sampler[int]):\n    \"\"\"Abstract interface for custom samplers.\"\"\"\n\n    @abstractmethod\n    def __init__(\n        self,\n        data_source: Sized,\n        data_config: DictConfig,\n    ):\n        pass\n\n\nclass AbstractCurriculumSampler(AbstractSampler):\n    \"\"\"Experimental interface for curriculum learning samplers.\"\"\"\n\n    @abstractmethod\n    def update(self, batch: DataProto) -> None:\n        pass\n"
  },
  {
    "path": "verl/experimental/dynamic_dataset/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/dynamic_dataset/dynamicgen_dataset.py",
    "content": "# Copyright 2025 Amazon.com Inc and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nDataset class that enables dynamic data generation strategies between iterations of training.\nThis class extends RLHFDataset and uses an AbstractDataGen instance to generate data.\n\nThis is especially useful in settings where proposer model generates new tasks based\non rollout data.\n\"\"\"\n\nimport logging\nfrom abc import ABC, abstractmethod\nfrom typing import Optional\n\nimport datasets\nfrom omegaconf import DictConfig\nfrom torch.utils.data import Dataset\nfrom transformers import PreTrainedTokenizer, ProcessorMixin\n\nfrom verl import DataProto\nfrom verl.utils.dataset import RLHFDataset\nfrom verl.utils.import_utils import load_extern_object\n\nlogger = logging.getLogger(__name__)\n\n\nclass AbstractDataGenerator(ABC):\n    def __init__(self, config: DictConfig):\n        self.config = config\n\n    @abstractmethod\n    def generate(self, dataset: Dataset) -> datasets.Dataset:\n        \"\"\"\n        Generate method must be implemented by subclasses.\n        Args:\n            dataset: The dataset to generate from.\n        Returns:\n            Processed data or result as implemented by the subclass.\n        \"\"\"\n        pass\n\n\nclass MockDataGenerator(AbstractDataGenerator):\n    \"\"\"\n    A noop data gen class that only reappends the first datapoint.\n    This class is useful as a placeholder and testing.\n    \"\"\"\n\n    def __init__(self, config: DictConfig = None):\n        super().__init__(config)\n\n    def generate(self, dataset: Dataset) -> datasets.Dataset:\n        print(\"MockDataGenerator: No operation performed on the dataset.\")\n        return dataset.dataframe.select([0])\n\n\nclass DynamicGenDataset(RLHFDataset):\n    \"\"\"\n    A dataset class that uses a data generation strategy to process data.\n    This class extends RLHFDataset and uses an AbstractDataGen instance to generate data.\n    \"\"\"\n\n    def __init__(\n        self,\n        data_files: str | list[str],\n        tokenizer: PreTrainedTokenizer,\n        config: DictConfig,\n        processor: Optional[ProcessorMixin] = None,\n    ):\n        super().__init__(data_files, tokenizer, config, processor)\n        self.datagen: AbstractDataGenerator = config.datagen\n        assert \"datagen\" in config and config.datagen.get(\"path\", None) is not None, (\n            f\"datagen path is not set in config: {config}\"\n        )\n        # Dynamically load the custom datagen class\n        datagen_cls = load_extern_object(config.datagen.path, config.datagen.name)\n\n        # Verify that the custom datagen class inherits from AbstractDataGenerator\n        abs_cls = AbstractDataGenerator\n        if not issubclass(datagen_cls, abs_cls):\n            raise TypeError(\n                f\"The custom datagen class '{config.datagen.name}' from '{config.datagen.path}'\"\n                + \" must inherit from {abs_cls}\"\n            )\n\n        self.data_generator = datagen_cls(config.datagen)\n        self.on_batch_end()\n\n    def append_dataframe(self, new_dataframe: datasets.Dataset):\n        new_dataframe = self.maybe_filter_out_long_prompts(new_dataframe)\n        self.dataframe = datasets.concatenate_datasets([self.dataframe, new_dataframe])\n\n        logger.info(f\"new dataset len: {len(self.dataframe)}\")\n\n    def on_batch_end(self, batch: DataProto) -> None:\n        \"\"\"\n        Generate data using the provided data generation strategy.\n        Note: This method is intended to change the dataset after each training batch.\n        \"\"\"\n        new_data = self.data_generator.generate(self)\n        self.append_dataframe(new_data)\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/README.md",
    "content": "# Recipe: Fully Async Policy Trainer\n\n**Author:** `https://github.com/meituan-search`\n\nLast updated: 02/05/2026.\n\nThis document introduces a fully asynchronous PPO training system that completely decouples the Trainer and Rollouter,\nsupporting asynchronous sample generation and training.\nUnder this system, we achieved a 2.35x-2.67x performance improvement when training the Qwen2.5-7B model with 128 GPUs,\nwithout significantly affecting the results.\n\n## Introduction\n\n### Background\n\nThe separated rollout and train architecture, compared to the colocate architecture, can allocate resources more\nflexibly and design more flexible training logic, thereby addressing issues such as low GPU utilization and training\nefficiency caused by long-tail problems.\nThe one_step_off_policy alleviates the problem of long rollout times and achieves some gains in training efficiency by\ndesigning a separated architecture and performing asynchronous training between rollout and train for one round.\nHowever, it forcibly uses data from one round of asynchronous training, which is not flexible enough and cannot\ncompletely eliminate the impact of long-tail on training efficiency.\nIn other frameworks such as AReaL, Magistral, StreamRL, and AsyncFlow, asynchronous training and streaming training have\nbeen implemented based on the separated architecture and have achieved gains.\nWe borrow from their methods and implemented them in VERL. The fully_async_policy supports asynchronous, streaming, and\npartial\nrollout training.\nBy reasonably setting parameters such as resource allocation and parameter synchronization frequency, fully_async_policy\ncan significantly improve training efficiency.\n\n> Magistral https://arxiv.org/abs/2506.10910\n>\n> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language\n> Reasoning https://arxiv.org/abs/2505.24298\n>\n> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream\n> Generation https://arxiv.org/abs/2504.15930\n>\n> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663\n>\n\n### Core Contributions\n\n* **Resource Isolation**: Unlike using hybrid_engine, Rollouter and Trainer use separate computing resources and need to\n  specify the resources they occupy separately.\n* **Parallel Generation and Training**: While the Trainer is training, the Rollouter is generating new samples.\n* **Multi-step Asynchronous**: Compared to one step off policy, it supports asynchronous settings from 0.x steps to\n  multiple steps, making the asynchronous solution more flexible.\n* **NCCL Parameter Synchronization**: Based on the nccl communication primitive, refer to [checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine) to\n  achieve efficient parameter synchronization between Rollouter and Trainer.\n* **Stream Inference and Training**: Rollouter generates data sample by sample, and data transmission uses a single\n  sample as the minimum transmission unit.\n* **Asynchronous Training and Freshness Control**: By setting the parameter async_training.staleness_threshold, it\n  supports training with samples generated by old parameters.\n* **PartialRollout**: The Rollouter's inference process supports partial rollout logic. During parameter\n  synchronization, by adding `sleep() and resume()` logic, it\n  saves samples from ongoing rollouts and continues using them in the next rollout, reducing the time spent waiting for\n  ongoing tasks to finish during parameter synchronization.\n\nCurrently, the supported usage mode is megatron/fsdp+vllm. vllm must use the server mode based on AgentLoop.\n\n## Design\n\nThe overall architecture of fully_async_policy is shown in the figure below. fully_async_policy mainly consists of four\nparts: Rollouter, MessageQueue, Trainer, and ParameterSynchronizer.\n\n![fully_async_policy_structure](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)\n\n1. Rollouter generates sequences sample by sample and puts the generated samples into the MessageQueue, with the\n   production speed controlled by freshness.\n2. MessageQueue is used to temporarily store samples generated by Rollouter.\n3. Trainer fetches samples from MessageQueue sample by sample. After fetching `require_batches*ppo_mini_batch_size`\n   samples, it will perform training. After training for async_training.trigger_parameter_sync_step rounds, it triggers\n   a parameter synchronization with Rollouter.\n4. ParameterSynchronizer implements the NCCL synchronous parameter synchronization capability.\n\nThe source of benefits compared to the base scheme lies in the fact that in the colocate case, using more resources for\nrollout cannot solve the idleness caused by long-tail samples.\nAfter we perform resource isolation, the time for rollout and train may be longer than before (because fewer resources\nare used),\nbut the overlap in their time consumption reduces the end-to-end time consumption.\n\n![fully_async_policy_revenue](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)\n\n## Usage\n\n### Parameter Description\n\n| super params                                                     | implication                                                                                    |\n|------------------------------------------------------------------|------------------------------------------------------------------------------------------------|\n| `trainer.nnodes`                                                 | Number of nodes for Trainer                                                                    |\n| `trainer.n_gpus_per_node`                                        | Number of GPUs per node for Trainer                                                            |\n| `rollout.nnodes`                                                 | Number of nodes for Rollouter                                                                  |\n| `rollout.n_gpus_per_node`                                        | Number of GPUs per node for Rollouter                                                          |\n| `data.train_batch_size`                                          | In the fully async strategy, this value is not effective (default is 0)                        |\n| `data.gen_batch_size`                                            | In the fully async strategy, uses streaming sample production logic (default is 1)             |\n| `rollout.total_rollout_steps`                                    | Total number of rollout samples                                                                |\n| `rollout.test_freq`                                              | How many times Rollouter updates parameters before performing a validation                     |\n| `actor_rollout_ref.actor.ppo_mini_batch_size`                    | The ppo_mini_batch_size is a global num across all workers/gpus                                |\n| `actor_rollout_ref.actor.use_rollout_log_probs=True`             | Use log_probs generated by rollout                                                             |\n| `algorithm.rollout_correction.bypass_mode`                       | Whether to compute log_prob using the training model's parameters during the training phase.   |\n| `async_training.require_batches`                                 | Number of ppo_mini_batch_size that FullyAsyncTrainer fetches at once                           |\n| `async_training.trigger_parameter_sync_step`                     | Indicates how many local updates FullyAsyncTrainer performs before a parameter synchronization |\n| `async_training.staleness_threshold`                             | Freshness control                                                                              |\n| `async_training.partial_rollout`                                 | Whether to perform partial_rollout                                                             |\n| `async_training.use_trainer_do_validate`                         | Whether use trainer node to do validate process, default `False`                               |\n\n**Further Explanation:**\n\n* `rollout.total_rollout_steps`\n\n  Compared to colocate, the quantity can be aligned by multiplying train_batch_size and step:\n  `rollout.total_rollout_steps = data.train_batch_size * step`.\n\n* `async_training.trigger_parameter_sync_step`\n\n  In the fully async strategy, it indicates how many local updates the Trainer performs (i.e., how many times it fetches\n  `require_batches * ppo_mini_batch_size` samples) before a parameter synchronization with Rollouter.\n  Between every two parameter synchronizations between Rollouter and Trainer, the Trainer will process\n  `trigger_parameter_sync_step* require_batches*ppo_mini_batch_size` samples.\n  To fairly compare speed with colocate, `trigger_parameter_sync_step` should be set to\n  `data.train_batch_size / (require_batches * ppo_mini_batch_size)`.\n\n* `async_training.staleness_threshold`\n\n  In the fully async strategy, it indicates the maximum proportion of stale samples allowed to be used.\n\n    * `staleness_threshold`=0, indicates synchronous training.\n      Rollouter will generate a fixed number of samples between two parameter updates, the sample count is:\n      \n      `rollout_num = (trigger_parameter_sync_step*require_batches*ppo_mini_batch_size)`\n    * `staleness_threshold`>0, indicates asynchronous training, can be set to a decimal for more flexible asynchronous\n      calls.\n      Rollouter will generate at most the following number of samples between two parameter updates:\n\n      `rollout_num = (1+staleness_threshold)*(trigger_parameter_sync_step*require_batches*ppo_mini_batch_size) - num_staleness_sample`\n\n  `num_staleness_sample` represents the number of stale samples generated in excess during the last rollout.\n\n  Since it's a streaming system, rollout continues to generate and trainer continues to consume. If rollouter is slower,\n  trainer will trigger parameter synchronization earlier, and rollouter will not actually produce rollout_num samples.\n  When rollout is fast enough, setting `staleness_threshold` to 1 is basically equivalent to one_step_off policy.\n  To avoid too many expired samples affecting training accuracy, it is recommended to set this value to less than 1.\n\n* `async_training.partial_rollout`\n\n  partial_rollout only actually takes effect when staleness_threshold>0.\n\n* `async_training.require_batches`\n\n  In streaming training, require_batches should be set to 1, indicating that training is performed after producing\n  enough ppo_mini_batch_size samples.\n  In actual testing, we found that if fewer samples are issued at once, due to the order of data distribution, it can\n  cause training instability and longer response lengths.\n  Here, we additionally provide require_batches for streaming distribution and control the number of samples\n  participating in training at once.\n\n* `actor_rollout_ref.actor.use_rollout_log_probs=True`\n\n  In reinforcement learning algorithms, log_probs have implicit correlations with parameter versions and tokens. Due to\n  the settings of algorithms like PPO/GRPO/DAPO, when calculating importance sampling,\n  old_log_prob must use the log_probs corresponding to the rollout parameters and tokens to ensure algorithm\n  correctness. In the fully\n  async strategy, we default to old_log_prob being calculated by rollout rather than by trainer.\n\n* `algorithm.rollout_correction.bypass_mode`\n\n  > algorithm.rollout_correction.bypass_mode default is True, using rollout log prob.\n\n  During the training process, we observed that metrics and response lengths may become unstable in the later\n  stages of training. To mitigate this issue, we can use\n  the [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html)\n  technique for importance sampling. To utilize Rollout Importance Sampling, we need to compute log_prob using\n  the training engine, which requires enabling this switch.\n  Additionally, when `algorithm.rollout_correction.bypass_mode=False` and Rollout Importance Sampling are enabled under\n  mode d\n  (async stream pipeline with partial rollout), our implementation approximates `Areal's Decoupled PPO`.\n\n* `async_training.use_trainer_do_validate`\n\n  It controls whether to use the trainer's `do_validate` method for validation.\n  If set to True, the trainer will perform validation after each parameter update. It can reduce the validation time\n  overhead and trainer node idle time.\n  If set to False, the trainer will not perform validation.\n\n### Supported Modes\n\n1. on policy pipeline:\n    1. **trigger_parameter_sync_step=1, staleness_threshold=0**\n    2. Rollouter produces `require_batches*ppo_mini_batch_size` samples at once, Trainer fetches these samples for\n       training, and after training completes, Trainer and Rollouter perform a parameter synchronization;\n    3. During the rollout phase, if there are long-tail samples but few rollout samples, shorter samples cannot fill\n       idle resources, causing some resource waste.\n    4. As shown in figure a;\n\n2. stream off policy pipeline:\n    1. **trigger_parameter_sync_step>1, staleness_threshold=0**\n    2. Synchronous streaming training will be performed. Rollouter produces\n       `require_batches*ppo_mini_batch_size*trigger_parameter_sync_step` samples at once, Trainer performs a local\n       training every time it fetches `require_batches*ppo_mini_batch_size` samples, and after training\n       trigger_parameter_sync_step times, Trainer and Rollouter perform a parameter synchronization;\n    3. Compared to a, since more samples are generated at once, resource idleness will be lower.\n    4. In one step training, there will be two periods of resource idleness: when fetching the first batch of samples,\n       train waits for `require_batches*ppo_mini_batch_size` samples to be produced, and during the last parameter\n       update, rollout waits for training to complete.\n    5. As shown in figure b;\n\n3. async stream pipeline with stale samples:\n    1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=False**\n    2. After each parameter update, Rollouter will plan to produce at most rollout_num samples (in practice, the number\n       of samples generated may be less than this value depending on rollout speed).\n    3. If the rollout process is relatively fast, Rollouter will generate some additional samples num_stale_samples\n       before parameter synchronization for immediate use by Trainer after synchronization.\n       When triggering parameter synchronization, if Rollouter has ongoing tasks, it will wait for the tasks to complete\n       and not add new tasks;\n    4. Compared to b, except for the first step training, subsequent training will not have the time to wait for the\n       first batch rollout to finish, but will have the time to wait for active tasks to finish.\n    5. As shown in figure c;\n\n4. async stream pipeline with partial rollout:\n    1. **trigger_parameter_sync_step>=1, staleness_threshold>0, partial_rollout=True**\n    2. Compared to c, when triggering parameter synchronization, if Rollouter has samples being produced, it will\n       interrupt the rollout process and perform parameter synchronization. The interrupted samples will continue to be\n       generated after synchronization. This reduces the time to wait for active tasks to finish.\n    3. As shown in figure d;\n\n![fully_async_policy_mode](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)\n\n### Key Metrics\n\n| metrics                                        | implication                                                                                            |\n|------------------------------------------------|--------------------------------------------------------------------------------------------------------|\n| `trainer/idle_ratio`                           | Trainer idle rate                                                                                      |\n| `rollouter/idle_ratio`                         | Rollouter idle rate                                                                                    |\n| `fully_async/count/stale_samples_processed`    | Total number of old samples used in training                                                           |\n| `fully_async/count/stale_trajectory_processed` | Total number of old trajectories used in training (one sample produces rollout.n trajectories)         |\n| `fully_async/partial/total_partial_num`        | Number of partial samples processed by Trainer between two trigger_parameter_sync_step                 |\n| `fully_async/partial/partial_ratio`            | Ratio of partial samples processed by Trainer between two trigger_parameter_sync_step                  |\n| `fully_async/partial/max_partial_span`         | Maximum parameter span of partial samples processed by Trainer between two trigger_parameter_sync_step |\n\n### Parameter Tuning Recommendations\n\n* Resource Allocation and Adjustment:\n    * Reasonable resource allocation is the prerequisite for achieving good training efficiency. The ideal resource\n      allocation should make the rollout time and train time close, thereby minimizing pipeline bubbles in the entire\n      training process,\n      avoiding resource idleness, and ensuring Trainer does not use old samples. In real training scenarios, resource\n      allocation can be adjusted based on the idle time of rollout and train during actual training,\n      which can be obtained from rollouter/idle_ratio and trainer/idle_ratio. If rollouter/idle_ratio is high and\n      trainer/idle_ratio is low,\n      Trainer resources should be increased and Rollouter resources should be reduced, and vice versa.\n\n* Key Parameters:\n    * staleness_threshold: Setting it too high will cause more old samples to be used, affecting model performance. It\n      is recommended to set it to less than 1.\n    * require_batches: The closer to 1, the closer to a pure streaming process, the smaller the training bubbles, and\n      the faster the acceleration effect that can be achieved in terms of speed, but it will affect the order of sample\n      processing;\n    * trigger_parameter_sync_step: The smaller the setting, the closer to on policy, but it will cause frequent\n      parameter synchronization. Long-tail samples waste resources that cannot be filled by short samples, resulting in\n      low resource utilization.\n      The larger the setting, the higher the computational efficiency, but the accuracy will be affected by off policy.\n    * rollout.test_freq: It will occupy Rollouter resources and is not recommended to be set too small.\n\n* Mode Selection: By adjusting different parameters, the Fully Async architecture supports optimization acceleration at\n  different levels, suitable for tasks in different scenarios.\n    * For small-scale tasks that need to ensure training stability and on-policy nature, and have low speed\n      requirements, the on policy pipeline mode (Mode 1) can be tried.\n    * For scenarios that need to improve training throughput but are sensitive to staleness, the stream off policy\n      pipeline mode can be tried. That is, by\n      setting trigger_parameter_sync_step>1 to improve training efficiency, but still maintaining the synchronization\n      mechanism (staleness_threshold=0) (Mode 2).\n    * For large-scale tasks with high training speed requirements and can tolerate a certain degree of off-policy and\n      staleness, setting staleness_threshold>\n      0 and partial_rollout=True can improve training efficiency, using the async stream pipeline mode (Mode 3 or 4).\n\n### Quick Start\n\n```shell\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=10\nstaleness_threshold=0\ntrigger_parameter_sync_step=16\npartial_rollout=False\n\n\npython -m recipe.fully_async_policy.fully_async_main \\\n\ttrain_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    rollout.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n```\n\n## Experiments\n\n### Asynchronous Training on 7B Model\n\nWe used Qwen2.5-Math-7B to verify the benefits of the fully async strategy under long candidates and multiple resources.\nUsing the `async stream pipeline with stale samples` strategy, we achieved about 2x performance improvement on 32 cards,\n64 cards, and 128 cards without significantly affecting experimental results.\n\n* Machine: H20\n* Model: Qwen2.5-Math-7B\n* Rollout length: max_response_length FSDP2: 28K tokens;\n* Algorithm: DAPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 20\n\n* colocate sync:\n    * step: 400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|  training mode   \t   | resource allocation \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t  | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1          \t      |\n|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:---------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:|\n| colocate sync      \t | 32                  \t | 790.10 \t | 357.41 \t | 107.71       \t | 269.80        \t | 13h 44m                \t | 1d 3h 43m              \t | 2d 9h 22m              \t | 3d 17h 5m              \t | max: 0.3313<br>last: 0.2448  \t  |\n| fully_async_policy \t | 16:16               \t |  294.77  |  21.26   | \\            \t |     313.81      |    7h 58m<br>(1.72x)     |    16h 21m<br>(1.70x)    |   1d 0h 53m<br>(2.31x)   |   1d 9h 26m<br>(2.66x)   | max: 0.3302<br>last: 0.2333   \t |\n| colocate sync      \t | 64                  \t | 365.28 \t | 150.72 \t | 70.26        \t | 133.41       \t  | 10h 22m                \t | 20h 45m                \t | 1d 7h 6m               \t | 1d 17h 32m             \t | max: 0.3365<br>last:  0.2333 \t  |\n| fully_async_policy \t | 32:32               \t | 189.26 \t | 28.46  \t | \\            \t | 156.98       \t  | 4h 57m<br>(2.09x)      \t | 10h 14m<br>(2.03x)     \t | 16h 58m<br>(1.83x)     \t | 21h 40m<br>(1.92x)     \t | max: 0.3677<br>last: 0.3406  \t  |\n| colocate sync      \t | 128                 \t | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t  | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t  |\n| fully_async_policy \t | 64:64               \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t  | 3h 13m<br>(2.67x)      \t | 6h 46m<br>(2.65x)      \t | 10h 53m<br>(2.67x)     \t | 17h 22m<br>(2.35x)     \t | max: 0.3521<br>last: 0.3094  \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg\n\n### 128-card 7B Asynchronous Mode Experiment\n\nWe used Qwen2.5-Math-7B to verify the effects of various modes supported by fully async.\nWe can see that the benefit brought by streaming is approximately 1.6x, and after combining staleness and\npartial_rollout, the benefit reaches 2.35x.\n\n|                             mode                                         \t                              | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1         \t      |\n|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:|\n|                                          colocate sync      \t                                           | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t |\n| `stream off policy pipeline`<br>(+fully async: trigger_parameter_sync_step= 4,<br>require_batches= 4) \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t  |\n|          `async stream pipeline with stale samples`<br>(+staleness_threshold=0.5)            \t          |    \t     |    \t     |       \t        |       \t        |            \t             |            \t             |            \t             |            \t             |               \t                |\n|        `async stream pipeline with partial rollout`<br>(+partial_rollout=True)                 \t        | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg\n\n### 128-card Stale Ablation Experiment\n\nUnder the `async stream pipeline with partial rollout` mode, we verified the impact of staleness settings on training\nefficiency.\nWe found that the larger the staleness, the more obvious the final gains.\nWe also noticed that the times for staleness values of 0.3 and 0.5 are quite close, because as the training steps\nincrease, the response length changes significantly, causing training instability.\nFurther analysis and optimization are needed for this issue.\n\n| staleness_threshold \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |     acc/mean@1         \t      |\n|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 0                   \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t |\n| 0.1                 \t | 171.30 \t | 58.17  \t | \\            \t | 109.12       \t | 3h 53m                 \t | 8h 37m                 \t | 14h 25m                \t | 19h 59m                \t | max: 0.3542<br>last: 0.2979 \t |\n| 0.3                 \t | 146.11 \t | 38.88  \t | \\            \t | 103.22       \t | 3h 18m                 \t | 6h 49m                 \t | 11h 40m                \t | 17h 20m                \t | max: 0.3469<br>last: 0.2865 \t |\n| 0.5                 \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg\n\n### 128-card 7B require_batches Ablation Experiment\n\nIn multiple tests, we found that the number of samples issued each time in streaming affects the response length during\ntraining, which in turn affects training time. We verified the impact on results by modifying\n`async_training.require_batches`.\n\n| require_batches \t | step  \t  | gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t |     acc/mean@1         \t      |\n|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 1               \t | 203.47 \t | 30.88 \t | \\            \t | 181.08       \t | 3h 31m                 \t | 8h 29m                 \t | 17h 36m                \t | max: 0.349<br>last: 0.326   \t |\n| 2               \t | 158.72 \t | 26.32 \t | \\            \t | 128.08       \t | 3h 35m                 \t | 7h 38m                 \t | 13h 57m                \t | max: 0.351<br>last: 0.3406  \t |\n| 4               \t | 124.64 \t | 25.62 \t | \\            \t | 95.06        \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | max: 0.3521<br>last: 0.3521 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg\n\n### 30B Model Mode Experiment\n\nWe achieved a 1.7x performance improvement with `async stream pipeline with staleness samples` strategy on the\nQwen3-30B-A3B-Base model compared to the colocate setup. It is worth noting that this is far from the upper limit of\nperformance gains achievable through asynchrony. Firstly, the comparative experiments used a maximum response length of\nonly 8k, which is much shorter than the 20k sequence length in previous experiments, resulting in a less pronounced\nrollout tail effect. Secondly, we adopted a highly skewed resource allocation, with rollout using 96 GPUs and trainer\nusing 32 GPUs, which is not an optimal configuration. During the experiments, we observed that the current verl\nimplementation imposes certain constraints, such as requiring data to be evenly divisible by the number of GPUs, making\nresource adjustment less flexible. Additionally, as asynchronous training and deployment accelerate, the performance gap\nis gradually narrowing. Therefore, enabling more flexible resource allocation and dynamic resource adjustment in the\nfuture will be our next focus.\n\n* Machine: H20\n* Model: Qwen3-30B-A3B-Base\n* Rollout length: max_response_length : 8K tokens;\n* Algorithm: GRPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+Megatron\n* rollout.n: 16\n* ppo_mini_batch_size: 128\n* test_freq: 20\n\n* colocate sync:\n    * step:400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * trigger_parameter_sync_step: 512/128 = 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n| Training Mode      | Resource Allocation | Step   | Gen    | Old Log Prob | Ref   | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1                  |\n|--------------------|---------------------|--------|--------|--------------|-------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------|\n| Colocate Sync      | 128                 | 497.89 | 348.05 | 28.73        | 20.86 | 86.27        | 13h 36m             | 1d 3h 48m           | 1d 19h 4m           | 2d 11h 39m          | max: 0.3500<br>last: 0.3208 |\n| Fully Async Policy | 96:32               | 282.75 | 22.06  | \\            | 50.05 | 206.63       | 6h 45m (2.01x)      | 14h 48m (1.88x)     | 1d 0h 9m (1.78x)    | 1d 10h 41m (1.72x)  | max: 0.3813<br>last: 0.3448 |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg           | | |\n\n### checkpoint-engine Ablation Experiment\nWe tested the single-step parameter synchronization time of the checkpoint-engine on three models: Qwen2.5-Math-7B, Qwen3-30B-A3B, and Qwen3-235B-A22B, using default checkpoint-engine configurations. All experiments were performed on H20 machines, and the Megatron engine was used for training.\n\n|      model      | trainer rank \t | rollout rank\t | checkpoint-engine \t | total sync time \t |\n|:---------------:|:--------------:|:-------------:|:-------------------:|:-----------------:|\n| Qwen2.5-Math-7B |       4        |       4       |        False        |       0.12s       |\n| Qwen2.5-Math-7B |       4        |       4       |        True         |       0.02s       |\n|  Qwen3-30B-A3B  |       16       |      16       |        False        |      15.76s       |\n|  Qwen3-30B-A3B  |       16       |      16       |        True         |       4.38s       |\n| Qwen3-235B-A22B |       64       |      64       |        False        |      58.57s       |\n| Qwen3-235B-A22B |       64       |      64       |        True         |      23.70s       |\n\n\n### use_trainer_do_validate Experiment\nWe tested the effect of setting `use_trainer_do_validate=True` on the training process. The results show that setting\nthis parameter to True can reduce the validation time overhead and trainer node idle time.\nWe used Qwen2.5-Math-7B to verify the benefits of `use_trainer_do_validate=True` on the training process, we achieved about 2x performance improvement on validation time, and the trainer node idle time is reduced by about 40%.\n\n* Machine: H20\n* Model: Qwen2.5-Math-7B\n* Rollout length: max_response_length FSDP2: 10K tokens;\n* Algorithm: DAPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 10\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|   training mode    | resource allocation |  step   |   gen   | old_log_prob | update_actor | validate time | total time<br>50 step | acc/mean@2 |\n|:------------------:|:-------------------:|:-------:|:-------:|:------------:|:------------:|:-------------:|:---------------------:|:----------:|\n|   colocate sync    |         16          | 484.623 | 52.939\t |      0\t      |   430.263    |  205.080  \t   |        7h9m  \t        |    22.6    |\n| fully_async_policy |         8:8         | 489.953 | 52.622\t |      0\t      |   435.874    |   95.699  \t   |        7h2m  \t        |    21.0    |\n\n\n## Multi-Turn Tool Calling\n\nReferencing **recipe/retool** and **ToolAgentLoop**, we implemented **AsyncPartialToolAgentLoop**, a multi-turn\ntool-calling loop that supports partial_rollout for **fully_async_policy**.\n\n### Core Design\n\n`AsyncPartialToolAgentLoop` inherits from `ToolAgentLoop` and is adapted for the asynchronous training mode of\n`fully_async_policy`. When `partial_rollout=True`, the Rollouter interrupts ongoing generation tasks before\nsynchronizing parameters with the Trainer. `AsyncPartialToolAgentLoop` is capable of:\n\n1. **Interrupting Tasks**: Responding to an interrupt signal to save the current state. Currently, interruptions occur\n   during the `GENERATING` process or after other states have completed.\n2. **Resuming Tasks**: Resuming execution from the saved state after parameter synchronization is complete, rather than\n   starting over.\n\n### How to Use\n\nRL training with multi-turn tool calling in `fully_async_policy` is similar to `recipe/retool`. It is enabled by\nspecifying `multi_turn` configurations in the config file.\n\n1. **SFT Stage**: First, the model should undergo SFT to learn how to follow tool-calling format instructions.\n2. **Multi-turn Configuration**: In the `fully_async_policy` training configuration, set the following parameters:\n   ```yaml\n   actor_rollout_ref:\n     rollout:\n       multi_turn:\n         enable: True # AsyncPartialToolAgentLoop will be used by default in fully_async_policy mode\n         # Other multi_turn related configurations\n   ```\n3. **Async Parameters**: To improve efficiency, enable `partial_rollout` and `staleness_threshold` when using multi-turn\n   tool calling:\n   ```yaml\n   async_training:\n     partial_rollout: True\n     staleness_threshold: 0.5\n     # Other async parameters\n   ```\n4. **Example**: See `recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`.\n\n### Experimental Results\n\nTo validate the performance of `fully_async_policy` on multi-turn tool-calling tasks, we compared it with the standard\n`colocate` synchronous mode. Key parameter settings are as follows.\n\n* **SFT Model**: Based on `Qwen2.5-7B-Instruct`, trained for 6 epochs on the `ReTool-SFT` dataset\n* **RL Algorithm**: DAPO\n* **Dataset**:\n    * Train: `DAPO-Math-17k`\n    * Test: `aime_2025`\n* **Resource and Mode Comparison**:\n    * `colocate sync`: 32 H20 gpus\n    * `fully_async_policy`: 16 gpus for Trainer + 16 gpus for Rollouter\n* **Key Configurations**:\n    1. **Tool Calling Configuration**:\n        * `multi_turn.enable: True`\n        * `multi_turn.max_user_turns: 16`\n        * `multi_turn.max_assistant_turns: 16`\n        * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml`\n    2. **`colocate sync` Configuration**:\n        * `ppo_mini_batch_size: 16`\n        * `train_batch_size: 64`\n    3. **`fully_async_policy` Configuration**:\n        * `ppo_mini_batch_size: 16`\n        * `trigger_parameter_sync_step: 4`\n        * `require_batches: 1`\n        * `staleness_threshold: 1`\n        * `partial_rollout: True`\n\n|  training mode   \t   | Resource allocation \t |  step  \t  |  gen   \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t |   aime_2025<br>acc/mean@30  \t   |\n|:--------------------:|:---------------------:|:---------:|:---------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:-------------------------------:|\n| colocate           \t | 32                  \t | 375.47  \t | 228.03 \t  | 35.19        \t | 111.84       \t | 9h 46m                 \t | 22h 28m                \t | start:0.1078<br>last:0.2056   \t |\n| fully_async_policy \t | 16: 16              \t | 221.36 \t  | 40.59   \t | \\            \t | 179.58       \t | 6h 19m<br>(1.55x)      \t | 14h 4m<br>(1.60x)      \t |   start:0.11<br>last:0.2044 \t   |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg\n\n## Future Plans\n\n* GRPO experiments\n* Megatron adaptation\n* SGLang integration\n* Transfer queue integration\n* Asynchronous parameter synchronization\n* AReaL asynchronous algorithm implementation\n* TPPO algorithm implementation\n* Multi-turn and Tool support\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/README_zh.md",
    "content": "# Recipe: Fully Async Policy Trainer\n\n**Author:**  `https://github.com/meituan-search`\n\nLast updated: 02/05/2026.\n\n本文档介绍了完全异步PPO训练系统，该系统实现了 Trainer 和 Rollouter 的完全解耦，支持异步样本生成和训练。\n在该系统下，我们使用128卡训练qwen2.5-7B模型取得了2.35x-2.67x的性能提升,同时效果没有显著受到影响。\n\n## Introduction\n\n### Background\n\nrollout和train分离架构相较于colocate的架构能够更加灵活地分配资源，设计更加灵活的训练逻辑，从而处理长尾等问题带来的GPU利用率低，训练效率低的问题。\none_step_off_policy通过分离架构的设计并进行rollout和train一轮异步的训练方法，缓解了rollout时间过长的问题，并在训练效率上取得了一些收益，\n但其强制使用一轮异步的数据，存在不够灵活等问题，而且并不能完全去除长尾对训练效率带来的的影响；在其他框架如areal、Magistral、streamrl、asyncflow上，\n已经基于分离架构实现了异步训练、流式训练，并取得了收益；我们借鉴其方法，在verl上进行了实现。fully_async_policy支持异步、流式、partial\nrollout的训练， 通过合理设置资源分配情况、参数同步频率等参数，fully_async_policy能够显著提高训练效率。\n\n> Magistral https://arxiv.org/abs/2506.10910\n>\n> AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language\n> Reasoning https://arxiv.org/abs/2505.24298\n>\n> StreamRL: Scalable, Heterogeneous, and Elastic RL for LLMs with Disaggregated Stream\n> Generation https://arxiv.org/abs/2504.15930\n>\n> AsyncFlow: An Asynchronous Streaming RL Framework for Efficient LLM Post-Training https://arxiv.org/abs/2507.01663\n>\n\n### 核心贡献\n\n* **资源隔离**：与使用hybrid_engine不同，Rollouter和Trainer使用分离的计算资源，需要分别指定所占用的资源。\n* **生成与训练并行**：Trainer在训练的同时，Rollouter在生成新的样本。\n* **多步异步**: 相比 one step off policy 支持0.x步到多步的异步设定，异步方案更加灵活。\n* **nccl参数同步**：基于nccl通信原语，参考[checkpoint-engine](https://github.com/MoonshotAI/checkpoint-engine)实现Rollouter与Trainer间的高效参数同步。\n* **Stream推理与训练**：Rollouter逐样本生成数据，同时数据传输以单个sample为最小传输单位。\n* **异步训练与新鲜度控制**：通过设置参数async_training.staleness_threshold，支持使用旧参数生成的样本进行训练。\n* **PartialRollout**: Rollouter推理过程支持partial rollout逻辑，通过参数同步时，添加`sleep()`和`resume()`\n  逻辑，保存进行中的rollout的样本，并在下一次rollout中继续使用，减少参数同步等待进行中的任务结束时间。\n\n目前支持使用模式为 megatron/fsdp+vllm。vllm必须使用基于AgentLoop的server模式。\n\n## 设计\n\nfully_async_policy的整体架构如下图所示，fully_async_policy主要由Rollouter、MessageQueue、Trainer、ParameterSynchronizer四部分组成。\n\n![fully_async_policy_structure](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_structure.svg?raw=true)\n\n1. Rollouter逐样本生成序列，并将生成的sample放入MessageQueue中，生产的速度受新鲜度控制。\n2. MessageQueue用于暂存Rollouter生成的sample。\n3. Trainer逐样本从MessageQueue中获取，获取到`require_batches*ppo_mini_batch_size`\n   数量的样本后，就会进行训练，训练async_training.trigger_parameter_sync_step轮后，触发与Rollouter的一次参数同步。\n4. ParameterSynchronizer 实现了Nccl的同步参数同步能力。\n\n当前方案对比base的收益来源，在于colocate情况下，rollout使用更多的资源无法解决长尾样本带来的空闲，\n当我们进行资源隔离后，rollout的时间和train的时间都可能相较于之前更长（因为使用的资源变少了），\n但是相互之间的耗时overlap，端到端的耗时反而有所缩减。\n\n![fully_async_policy_revenue](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_revenue.svg?raw=true)\n\n## 使用方式\n\n### 参数说明\n\n| super params                                                     | implication                                                     |\n|------------------------------------------------------------------|-----------------------------------------------------------------|\n| `trainer.nnodes`                                                 | Trainer的node数量                                                  |\n| `trainer.n_gpus_per_node`                                        | Trainer每个node上gpu的数量                                            |\n| `rollout.nnodes`                                                 | Rollouter的node数量                                                |\n| `rollout.n_gpus_per_node`                                        | Rollouter每个node上gpu的数量                                          |\n| `data.train_batch_size`                                          | 在fully async策略中，该值不生效（默认设置为0）                                   |\n| `data.gen_batch_size`                                            | 在fully async策略中，使用流式的样本生产逻辑（默认设置为1)                             |\n| `rollout.total_rollout_steps`                                    | 总的rollout的sample数量                                              |\n| `rollout.test_freq`                                              | Rollouter每更新多少次参数，进行一次validation                                |\n| `actor_rollout_ref.actor.ppo_mini_batch_size`                    | The ppo_mini_batch_size is a global num across all workers/gpus |\n| `actor_rollout_ref.actor.use_rollout_log_probs=True`             | 使用rollout产生的log_probs                                           |\n| `algorithm.rollout_correction.bypass_mode`                       | 是否在train阶段，使用train模型的参数计算token的 log_prob                        |\n| `async_training.require_batches`                                 | FullyAsyncTrainer一次性获取的ppo_mini_batch_size的数量                   |\n| `async_training.trigger_parameter_sync_step`                     | 表示FullyAsyncTrainer进行多少次本地更新后,进行一次参数同步                          |\n| `async_training.staleness_threshold`                             | 新鲜度控制                                                           |\n| `async_training.partial_rollout`                                 | 是否进行partial_rollout                                             |\n| `async_training.use_trainer_do_validate`                         | 是否使用Trainer的do_validate方法进行validation，默认值False                  |\n\n**进一步的解释：**\n\n* `rollout.total_rollout_steps`\n\n  与 colocate 相比，数量可以通过 train_batch_size 与 step 相乘对齐:\n  `rollout.total_rollout_steps = data.train_batch_size * step`。\n\n* `async_training.trigger_parameter_sync_step`\n\n  在fully async策略中，表示Trainer进行多少次本地更新后（也就是获取多少次`require_batches * ppo_mini_batch_size`数量样本），\n  与Rollouter之间进行一次参数同步。\n  每两次Rollouter和Trainer参数同步之间，Trainer将会处理`trigger_parameter_sync_step* require_batches\\\n  ppo_mini_batch_size`份sample。\n  如果为了与colocate在公平的情况下对比速度，trigger_parameter_sync_step应该设置为 `data.train_batch_size / (\n  require_batches * ppo_mini_batch_size)`。\n\n* `async_training.staleness_threshold`\n\n  在fully async策略中，表示最大允许使用的staleness样本的比例。\n\n    * staleness_threshold=0，表示同步训练。\n      Rollouter两次参数更新之间将会生成固定数量的样本，样本数为：\n      $$rollout\\_num = (trigger\\_parameter\\_sync\\_step*require\\_batches*ppo\\_mini\\_batch\\_size)$$\n    * staleness_threshold>0，表示异步训练， 可以设置为小数，支持更灵活的异步调用。\n      Rollouter两次参数更新之间将会最多生成的样本数为：\n      $$rollout\\_num = (1+staleness\\_threshold)*(trigger\\_parameter\\_sync\\_step*require\\_batches*ppo\\_mini\\_batch\\_size) - num\\_staleness\\_sample $$\n\n  num_staleness_sample 表示上一次rollout多生成的陈旧样本数。\n\n  由于是流式系统，rollout持续生成，trainer持续消费。如果rollouter较慢，trainer会更早触发参数同步，rollouter并不会实际生产rollout_num个样本。\n  当rollout 足够快时，staleness_threshold设置为1，基本上等价于one_step_off policy。\n  为了避免过期样本太多影响训练精度，建议该值设置小于1。\n\n* `async_training.partial_rollout`\n\n  partial_rollout只会在staleness_threshold>0时才实际上起作用。\n\n* `actor_rollout_ref.actor.use_rollout_log_probs=True`\n\n  在强化学习算法中，log_probs与参数版本，token都存在隐性的相关性。由于PPO/GRPO/DAPO等算法的设定，我们在计算重要性采样时，\n  即 old_log_prob必须使用rollout参数及token所对应log_probs，才能保证算法的正确性。在fully\n  async策略中，我们默认old_log_prob是由rollout所计算的，而不是由trainer所计算。\n\n* `algorithm.rollout_correction.bypass_mode`\n  algorithm.rollout_correction.bypass_mode 默认为 True, 直接使用rollout log prob。\n\n  我们在训练过程中，观测到随着训练的进行，训练后期指标和response长度可能会出现不稳定的情况，\n  这里我们可以使用 [Rollout Importance Sampling](https://verl.readthedocs.io/en/latest/advance/rollout_is.html) 的技术进行\n  重要性采样，缓解这一问题。为了使用 `Rollout Importance Sampling` 我们需要使用训练引擎使用当前的参数版本计算old_log_prob，此开关需要打开。\n  此外，在 mode d (async stream pipeline with partial rollout) 的情况下 `algorithm.rollout_correction.bypass_mode=False`\n  以及\n  `Rollout Importance Sampling` 后，我们的实现已近似Areal的 `Decoupled PPO`。\n\n* `async_training.require_batches`\n\n  在流式训练中，require_batches 应该设置为1，表示生产够ppo_mini_batch_size样本后，就进行训练。\n  在实际测试中，我们发现，如果单次下发的样本较少，由于数据分发的顺序，会导致训练不稳定，response 长度变长。\n  在这里，我们额外提供 require_batches 进行流式分发，单次参与训练的样本数量控制。\n\n* `async_training.use_trainer_do_validate`\n\n  控制是否使用trainer的`do_validate`方法进行validation。\n  如果设置为True，trainer会在每次参数更新后，调用`do_validate`方法进行validation。\n  如果设置为False，trainer不会调用`do_validate`方法。\n\n### 模式支持\n\n1. on policy pipeline:\n    1. **trigger_parameter_sync_step=1，staleness_threshold=0**\n    2. Rollouter一次生产`require_batches*ppo_mini_batch_size`\n       的samples，Trainer获取这些samples后进行训练，训练完后Trainer和Rollouter之间进行一次参数同步;\n    3. 在rollout阶段，如果存在长尾的样本，但是rollout样本数较少时，较短的样本无法填充到空闲的资源中，会造成一定的资源浪费。\n    4. 如图a所示；\n\n2. stream off policy pipeline:\n    1. **trigger_parameter_sync_step>1，staleness_threshold=0**\n    2. 将会进行同步的流式训练，Rollouter一次生产`require_batches*ppo_mini_batch_size*trigger_parameter_sync_step`\n       的samples，Trainer每获取`require_batches*ppo_mini_batch_size`\n       就进行一次本地训练，训练trigger_parameter_sync_step次后，Trainer和Rollouter之间进行一次参数同步;\n    3. 相较于a，由于一次生成的样本更多，资源的空闲会更低。\n    4. 在一次step训练中，会存在两次资源闲置的时间，分别是在第一次获取样本时，train等待`require_batches*ppo_mini_batch_size`\n       个样本生产，以及最后一次参数更新时，rollout等待训练完成。\n    5. 如图b所示；\n\n3. async stream pipeline with staleness samples:\n    1. **trigger_parameter_sync_step>=1，staleness_threshold>0，partial_rollout=Flase**\n    2. Rollouter在每次参数更新后将计划最多生产rollout_num个样本（实际根据rollout速度，生成的样本可能会少与这个值）。\n    3. 如果rollout过程比较快，Rollouter将会在参数同步前额外生成一部分样本num_stale_samples，用于参数同步后立即给Trainer使用。\n       触发参数同步时，如果Rollouter有正在生产的任务，将会等待任务完成，同时不会添加新的任务；\n    4. 相较于b，除第一次step训练外，后续的训练都不会有wait first batch rollout finish的时间，但是会有wait active task\n       finish的时间。\n    5. 如图c所示；\n\n4. async stream pipeline with partial rollout:\n    1. **trigger_parameter_sync_step>=1，staleness_threshold>0，partial_rollout=True**\n    2. 相较于c，触发参数同步时，Rollouter如果有正在生产的sample，会打断rollout过程并进行参数同步，被中断的sample会在参数同步后继续生成。减少了wait\n       active task finish的时间。\n    3. 如图d所示；\n\n![fully_async_policy_mode](\nhttps://github.com/ArronHZG/verl-community/blob/main/docs/fully_async_policy_mode.svg?raw=true)\n\n### 关键指标\n\n| metrics                                        | implication                                               |\n|------------------------------------------------|-----------------------------------------------------------|\n| `trainer/idle_ratio`                           | Trainer闲置率                                                |\n| `rollouter/idle_ratio`                         | Rollouter闲置率                                              |\n| `fully_async/count/stale_samples_processed`    | 训练使用的旧sample总数                                            |\n| `fully_async/count/stale_trajectory_processed` | 训练使用的旧trajectory总数(一个sample会生产rollout.n条trajectory)       |\n| `fully_async/partial/total_partial_num`        | 两次trigger_parameter_sync_step之间Trainer处理的partial样本数       |\n| `fully_async/partial/partial_ratio`            | 两次trigger_parameter_sync_step之间Trainer处理的partial样本的比例     |\n| `fully_async/partial/max_partial_span`         | 两次trigger_parameter_sync_step之间Trainer处理的partial样本的最大参数跨度 |\n\n### 调参建议\n\n* 资源分配与调整:\n    * 合理的资源分配是获得好的训练效率的前提。理想的资源分配情况应该是使得Rollout的时间和Train的时间接近，从而使得整个训练过程流水气泡最小，\n      避免资源闲置，同时Trainer不会使用旧样本。在真实训练场景下，可以根据实际训练过程中rollout和train的空闲时间调整资源分配，\n      可从rollouter/idle_ratio和trainer/idle_ratio获得，如果rollouter/idle_ratio较高trainer/idle_ratio较低，\n      应该增多Trainer的资源减少Rollouter的资源，反之亦然。\n\n* 关键参数：\n    * staleness_threshold: 设置太大会导致较多的旧样本使用，影响模型效果，建议设置小于1。\n    * require_batches：越接近1，越接近纯流式过程，训练过程中bubble越小，能够在速度上获得更快的加速效果，但会对样本的处理顺序产生影响；\n    * trigger_parameter_sync_step: 设置的越小越接近on policy，但会导致频繁的参数同步，长尾样本浪费的资源无法被短样本填充，资源利用率低。\n      设置的越大有更高的计算效率，但是精度上会受到off policy的影响。\n    * rollout.test_freq: 会占用Rollouter资源，不建议设置太小。\n\n* 模式选择：通过调整不同的参数，Fully Async架构支持不同程度上的优化加速，适用于不同场景的任务。\n    * 对于小规模任务，需要保证训练的稳定性和 on-policy 性，对速度要求不高的场景，可以尝试使用on policy pipeline的模式（模式1）。\n    * 对于需要提高训练吞吐量，但对 staleness 敏感的场景，可以尝试使用 stream off policy pipeline 的模式。即通过\n      设置trigger_parameter_sync_step>1 ，提高 训练效率，但仍保持同步机制 (staleness_threshold=0 )（模式2）。\n    * 对于大规模任务，对训练速度有较高要求，且可以容忍一定 off-policy 程度、staleness的场景，可以设置staleness_threshold>\n      0、partial_rollout=True提高训练效率，使用 async stream pipeline 模式（模式 3 或 4）。\n\n### 快速开始\n\n```shell\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=10\nstaleness_threshold=0\ntrigger_parameter_sync_step=16\npartial_rollout=False\n\n\npython -m recipe.fully_async_policy.fully_async_main \\\n\ttrain_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    rollout.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n```\n\n## 实验\n\n### 在7B模型上进行异步训练\n\n我们使用 Qwen2.5-Math-7B 验证 fully async 策略在长候选下，多种资源下的收益情况。\n使用`async stream pipeline with staleness samples` 策略，我们在32卡，64卡，128卡都取得2x左右的性能提升，同时没有显著影响实验效果。\n\n* 机器：H20\n* 模型：Qwen2.5-Math-7B\n* rollout长度：max_response_length FSDP2: 28K tokens;\n* 算法：DAPO\n* 数据集： TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 20\n\n* colocate sync:\n    * step: 400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|  training mode   \t   | resource allocation \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1          \t      |\n|:--------------------:|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-------------------------------:|\n| colocate sync      \t | 32                  \t | 790.10 \t | 357.41 \t | 107.71       \t | 269.80      \t | 13h 44m                \t | 1d 3h 43m              \t | 2d 9h 22m              \t | 3d 17h 5m              \t | max: 0.3313<br>last: 0.2448  \t  |\n| fully_async_policy \t | 16:16               \t |  294.77  |  21.26   | \\            \t |      313.81     |    7h 58m<br>(1.72x)     |    16h 21m<br>(1.70x)    |   1d 0h 53m<br>(2.31x)   |   1d 9h 26m<br>(2.66x)   | max: 0.3302<br>last: 0.2333   \t |\n| colocate sync      \t | 64                  \t | 365.28 \t | 150.72 \t | 70.26        \t | 133.41       \t | 10h 22m                \t | 20h 45m                \t | 1d 7h 6m               \t | 1d 17h 32m             \t | max: 0.3365<br>last:  0.2333 \t  |\n| fully_async_policy \t | 32:32               \t | 189.26 \t | 28.46  \t | \\            \t | 156.98       \t | 4h 57m<br>(2.09x)      \t | 10h 14m<br>(2.03x)     \t | 16h 58m<br>(1.83x)     \t | 21h 40m<br>(1.92x)     \t | max: 0.3677<br>last: 0.3406  \t  |\n| colocate sync      \t | 128                 \t | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t  |\n| fully_async_policy \t | 64:64               \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m<br>(2.67x)      \t | 6h 46m<br>(2.65x)      \t | 10h 53m<br>(2.67x)     \t | 17h 22m<br>(2.35x)     \t | max: 0.3521<br>last: 0.3094  \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-colocate_async?nw=nwuserhouzg\n\n### 128卡 7B 异步模式实验\n\n我们使用 Qwen2.5-Math-7B 验证 fully async 所支持的各个模式的效果。\n我们可以看到 stream 带来的收益大约1.6x，叠加 staleness 和 partial_rollout 后，收益为2.35x。\n\n|                             mode                                         \t                              | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |      acc/mean@1         \t      |\n|:-------------------------------------------------------------------------------------------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------------:|\n|                                          colocate sync      \t                                           | 356.30 \t | 177.85 \t | 53.92        \t | 113.81       \t | 8h 36m                 \t | 17h 56m                \t | 1d 5h 6m               \t | 1d 16h 48m             \t | max: 0.3573<br>last: 0.2958  \t |\n| `stream off policy pipeline`<br>(+fully async: trigger_parameter_sync_step= 4,<br>require_batches= 4) \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t  |\n|        `async stream pipeline with staleness samples`<br>(+staleness_threshold=0.5)            \t        |    \t     |    \t     |       \t        |       \t        |            \t             |            \t             |            \t             |            \t             |               \t                |\n|        `async stream pipeline with partial rollout`<br>(+partial_rollout=True)                 \t        | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t  |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-stream_stale_partial?nw=nwuserhouzg\n\n### 128卡 stale 消融实验\n\n在 `async stream pipeline with partial rollout` 模式下，我们验证 staleness 的设置对于训练效率的影响。\n我们可以发现，staleness 越大，最终取得的收益越明显。\n同时我们也注意到 staleness 取 0.3 和 0.5 的时间比较接近，原因是随着训练步数的增量，response 长度变化较大，训练出现了不稳定的问题。\n后续还需要针对该问题进行进一步的分析和优化。\n\n| staleness_threshold \t | step  \t  |  gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t | total time<br>400 step \t |     acc/mean@1         \t      |\n|:---------------------:|:--------:|:--------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 0                   \t | 231.34 \t | 128.47 \t | \\            \t | 98.77        \t | 4h 25m                 \t | 9h 41m                 \t | 15h 2m                 \t | 1d 1h 53m              \t | max: 0.2844<br>last: 0.2604 \t |\n| 0.1                 \t | 171.30 \t | 58.17  \t | \\            \t | 109.12       \t | 3h 53m                 \t | 8h 37m                 \t | 14h 25m                \t | 19h 59m                \t | max: 0.3542<br>last: 0.2979 \t |\n| 0.3                 \t | 146.11 \t | 38.88  \t | \\            \t | 103.22       \t | 3h 18m                 \t | 6h 49m                 \t | 11h 40m                \t | 17h 20m                \t | max: 0.3469<br>last: 0.2865 \t |\n| 0.5                 \t | 150.63 \t | 33.14  \t | \\            \t | 113.16       \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | 17h 22m                \t | max: 0.3521<br>last: 0.3094 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_stale?nw=nwuserhouzg\n\n### 128卡 7B require_batches 消融实验\n\n在多次测试下，我们发现流式每次下发样本的数量会影响训练的response长度，进而影响训练时长，我们通过修改\n`async_training.require_batches` 验证对与结果的影响。\n\n| require_batches \t | step  \t  | gen  \t  | old_log_prob \t | update_actor \t | total time<br>100 step \t | total time<br>200 step \t | total time<br>300 step \t |     acc/mean@1         \t      |\n|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|:------------------------:|:------------------------:|:------------------------:|:-----------------------------:|\n| 1               \t | 203.47 \t | 30.88 \t | \\            \t | 181.08       \t | 3h 31m                 \t | 8h 29m                 \t | 17h 36m                \t | max: 0.349<br>last: 0.326   \t |\n| 2               \t | 158.72 \t | 26.32 \t | \\            \t | 128.08       \t | 3h 35m                 \t | 7h 38m                 \t | 13h 57m                \t | max: 0.351<br>last: 0.3406  \t |\n| 4               \t | 124.64 \t | 25.62 \t | \\            \t | 95.06        \t | 3h 13m                 \t | 6h 46m                 \t | 10h 53m                \t | max: 0.3521<br>last: 0.3521 \t |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-ablation_require_batches?nw=nwuserhouzg\n\n### 30B模型模式实验\n\n我们在 Qwen3-30B-A3B-Base 模型上通过`async stream pipeline with staleness samples` 策略，相比于 colocate 方案取得了 1.7\n倍的性能提升。值得说明的是，这距离异步方式所能带来的性能提升上限还有很大空间。首先，对比实验中使用的最大响应长度仅为\n8k，这远低于此前实验的 20k 序列长度，因此 rollout 的长尾效应并不明显。其次，我们采用了极为倾斜的资源分配方案，rollout 使用了\n96 张 GPU，而 trainer 仅使用了 32 张 GPU，这并不是最优的配置。在实验过程中，我们观察到当前的 verl 实现存在一些限制，比如要求数据必须能被\nGPU 数量整除，这使得资源调整的灵活性受到影响。此外，随着异步训练和部署的加速，性能差距也在逐渐缩小。因此，未来我们将重点关注如何实现更灵活的资源分配和动态调整资源。\n\n* 机器：H20\n* 模型：Qwen3-30B-A3B-Base\n* rollout长度：max_response_length : 8K tokens;\n* 算法： GRPO\n* 数据集： TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+Megatron\n* rollout.n: 16\n* ppo_mini_batch_size: 128\n* test_freq: 20\n\n* colocate sync:\n    * step:400\n    * train_batch_size: 512\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * trigger_parameter_sync_step: 512/128 = 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n| Training Mode        | Resource Allocation | Step    | Gen    | Old Log Prob | Ref    | Update Actor | Total Time 100 Step | Total Time 200 Step | Total Time 300 Step | Total Time 400 Step | Acc/Mean@1                 |\n|----------------------|--------------------|---------|--------|--------------|--------|--------------|---------------------|---------------------|---------------------|---------------------|-----------------------------|\n| Colocate Sync        | 128                | 497.89  | 348.05 | 28.73        | 20.86  | 86.27        | 13h 36m             | 1d 3h 48m           | 1d 19h 4m           | 2d 11h 39m          | max: 0.3500<br>last: 0.3208 |\n| Fully Async Policy   | 96:32              | 282.75  | 22.06  | \\            | 50.05  | 206.63       | 6h 45m (2.01x)      | 14h 48m (1.88x)     | 1d 0h 9m (1.78x)    | 1d 10h 41m (1.72x)  | max: 0.3813<br>last: 0.3448 |\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-30B?nw=nwuserhouzg\n\n### checkpoint-engine参数同步消融实验\n我们在Qwen2.5-Math-7B，Qwen3-30B-A3B和Qwen3-235B-A22B三个模型上测试了checkpoint-engine参数同步的单步参数同步耗时，使用的参数均为默认参数配置。实验均在H20机器上完成，并使用megatron训练引擎。\n| model |  trainer rank \t  | rollout rank\t  | checkpoint-engine \t | total sync time \t |\n|:-----------------:|:--------:|:-------:|:--------------:|:--------------:|\n| Qwen2.5-Math-7B   | 4        | 4       | False      | 0.12s      |\n| Qwen2.5-Math-7B   | 4        | 4       | True      | 0.02s      |\n|  Qwen3-30B-A3B     | 16        | 16       | False      | 15.76s   |\n|  Qwen3-30B-A3B     | 16        | 16       | True      | 4.38s   |\n|  Qwen3-235B-A22B    | 64        | 64       | False      | 58.57s   |\n|  Qwen3-235B-A22B    | 64        | 64       | True      | 23.70s   |\n\n### use_trainer_do_validate 实验测试\n我们在Qwen2.5-Math-7B模型上测试了`use_trainer_do_validate`参数的影响。这个结果展示使用`use_trainer_do_validate=True`可以减少验证时间开销，并且训练器节点的空闲时间也减少了。\n\n* Machine: H20\n* Model: Qwen2.5-Math-7B\n* Rollout length: max_response_length FSDP2: 10K tokens;\n* Algorithm: DAPO\n* Dataset: TRAIN_FILE: dapo-math-17k.parquet TEST_FILE: aime-2024.parquet\n* Engine: vllm+FSDP2\n* rollout.n: 16\n* ppo_mini_batch_size: 32\n* test_freq: 10\n\n* fully_async_policy\n    * total_rollout_steps: 512*400\n    * require_batches: 4\n    * trigger_parameter_sync_step: 4\n    * staleness_threshold: 0.5\n    * partial_rollout: True\n\n|  training mode  | resource allocation | step  |  gen  | old_log_prob | update_actor | validate time | total time<br>50 step | acc/mean@2 |\n|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|:---------------:|\n| colocate sync      | 16  |  484.623  |  52.939\t |   0\t |   430.263   |  205.080  \t |     7h9m  \t |     22.6     |\n| fully_async_policy | 8:8 |  489.953  |  52.622\t |   0\t |   435.874   |  95.699  \t |     7h2m  \t |     21.0    |\n\n\n## 多轮工具调用\n\n参考 **recipe/retool** 和 **ToolAgentLoop**，我们为 **fully_async_policy** 实现了支持partial rollout的多轮工具调用循环 *\n*AsyncPartialToolAgentLoop**。\n\n### 核心设计\n\n`AsyncPartialToolAgentLoop` 继承自 `ToolAgentLoop`，其核心是适配了 `fully_async_policy` 的异步训练模式。当\n`partial_rollout=True` 时，Rollouter 在与 Trainer 同步参数前会中断正在进行的生成任务。`AsyncPartialToolAgentLoop` 能够：\n\n1. **中断任务**: 响应中断信号，保存当前的生成状态。目前，中断会发生在GENERATING过程中，或其他状态结束后；\n2. **恢复任务**: 在参数同步完成后，从保存的状态恢复，继续执行，而不是从头开始。\n\n### 使用方法\n\n`fully_async_policy`多轮与工具调用的RL训练与 `recipe/retool` 类似，通过在配置文件中指定 `multi_turn` 相关配置来启用。\n\n1. **SFT 阶段**: 首先，需要对模型进行 SFT训练，使其具备遵循工具调用格式指令的能力。\n2. **配置启用**: 在 `fully_async_policy` 的训练配置中，设置以下参数:\n   ```yaml\n   actor_rollout_ref:\n     rollout:\n       multi_turn:\n         enable: True # 在fully_async_policy模式下将默认使用AsyncPartialToolAgentLoop\n         # 其他 multi_turn 相关配置\n   ```\n3. **配置async参数**: 为提高效率，在启用多轮工具调用时，同时开启 `partial_rollout`和`staleness_threshold`：\n   ```yaml\n   async_training:\n     partial_rollout: True\n     staleness_threshold: 0.5\n     # 其他async参数\n   ```\n4. **example**: 参考`recipe/fully_async_policy/shell/dapo_7b_async_retool.sh`\n\n### 实验结果\n\n为验证 `fully_async_policy` 在多轮工具调用任务中的性能，我们将其与标准 `colocate` 同步模式进行了对比。实验具体设置如下。\n\n* **SFT模型**: 实验基于 `Qwen2.5-7B-Instruct` 模型，使用`ReTool-SFT`数据集训练6个epoch；\n* **RL算法**: DAPO\n* **数据集**:\n    * 训练集: `DAPO-Math-17k`\n    * 测试集: `aime_2025`\n* **资源与模式对比**:\n    * `colocate sync`: 32卡 H20\n    * `fully_async_policy`: 16卡 Trainer + 16卡 Rollouter\n* **关键配置**:\n    1. **工具调用配置**:\n        * `multi_turn.enable: True`\n        * `multi_turn.max_user_turns: 16`\n        * `multi_turn.max_assistant_turns: 16`\n        * `multi_turn.tool_config_path: recipe/retool/sandbox_fusion_tool_config.yaml`\n    2. **`colocate sync`配置**:\n        * `ppo_mini_batch_size: 16`\n        * `train_batch_size: 64`\n    3. **`fully_async_policy`配置**:\n        * `ppo_mini_batch_size: 16`\n        * `trigger_parameter_sync_step: 4`\n        * `require_batches: 1`\n        * `staleness_threshold: 1`\n        * `partial_rollout: True`\n\n|    training mode   \t| Resource allocation \t|   step  \t|   gen   \t| old_log_prob \t| update_actor \t| total time<br>100 step \t| total time<br>200 step \t|   aime_2025<br>acc/mean@30  \t|\n|:------------------:\t|:-------------------:\t|:-------:\t|:-------:\t|:------------:\t|:------------:\t|:----------------------:\t|:----------------------:\t|:---------------------------:\t|\n| colocate           \t| 32                  \t| 375.47  \t|  228.03 \t| 35.19        \t| 111.84       \t| 9h 46m                 \t| 22h 28m                \t| start:0.1078<br>last:0.2056   \t|\n| fully_async_policy \t| 16: 16              \t|  221.36 \t| 40.59   \t| \\            \t| 179.58       \t| 6h 19m<br>(1.55x)      \t| 14h 4m<br>(1.60x)      \t| start:0.11<br>last:0.2044 \t|\n\n> source data: https://wandb.ai/hou-zg-meituan/fully-async-policy-multiturn-tool?nw=nwuserhouzg\n\n## 后续计划\n\n* GRPO实验\n* megatron 适配\n* sglang 集成\n* transfer queue 集成\n* 异步参数同步\n* Areal异步算法实现\n* TPPO算法实现\n* 多轮及Tool的支持"
  },
  {
    "path": "verl/experimental/fully_async_policy/agent_loop/__init__.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .agent_loop import FullyAsyncAgentLoopManager\n\n__all__ = [FullyAsyncAgentLoopManager]\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/agent_loop/agent_loop.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nfrom typing import Any, Optional\n\nimport ray\nimport torch\nfrom omegaconf import DictConfig\n\nfrom verl.experimental.agent_loop.agent_loop import (\n    AgentLoopManager,\n    AgentLoopWorker,\n    AsyncLLMServerManager,\n    TokenOutput,\n)\nfrom verl.protocol import DataProto\nfrom verl.single_controller.ray import RayResourcePool, RayWorkerGroup\nfrom verl.utils.ray_utils import auto_await\nfrom verl.utils.rollout_trace import (\n    rollout_trace_op,\n)\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass FullyAsyncLLMServerManager(AsyncLLMServerManager):\n    \"\"\"FullyAsyncLLMServerManager supports resume generation on partial rollout, making rollout interruption\n    invisible to the AgentLoop.\n    \"\"\"\n\n    @rollout_trace_op\n    async def generate(\n        self,\n        request_id,\n        *,\n        prompt_ids: list[int],\n        sampling_params: dict[str, Any],\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> TokenOutput:\n        \"\"\"Generate tokens from prompt ids.\n\n        Args:\n            request_id (str): request id for sticky session.\n            prompt_ids (List[int]): List of prompt token ids.\n            sampling_params (Dict[str, Any]): Sampling parameters for the chat completion.\n            image_data (Optional[List[Any]]): Image data for the chat completion.\n            video_data (Optional[List[Any]]): Video data for the chat completion.\n\n        Returns:\n            TokenOutput: token output\n        \"\"\"\n        limit_key = None\n        if \"max_tokens\" in sampling_params:\n            limit_key = \"max_tokens\"\n        elif \"max_new_tokens\" in sampling_params:\n            limit_key = \"max_new_tokens\"\n        original_max_tokens = sampling_params.get(limit_key) if limit_key else None\n\n        final_output = TokenOutput(\n            token_ids=[],\n            log_probs=[],\n            num_preempted=0,\n        )\n        min_global_steps, max_global_steps = None, None\n\n        while True:\n            # 1. generate tokens\n            output = await super().generate(\n                request_id=request_id,\n                prompt_ids=prompt_ids + final_output.token_ids,\n                sampling_params=sampling_params,\n                image_data=image_data,\n                video_data=video_data,\n            )\n\n            # 2. merge output into final_output\n            final_output.token_ids.extend(output.token_ids)\n            if output.log_probs is not None:\n                final_output.log_probs.extend(output.log_probs)\n            if output.routed_experts is not None:\n                if final_output.routed_experts is None:\n                    final_output.routed_experts = output.routed_experts\n                else:\n                    final_output.routed_experts = torch.cat([final_output.routed_experts, output.routed_experts], dim=0)\n            if output.num_preempted is not None:\n                final_output.num_preempted += output.num_preempted\n            final_output.stop_reason = output.stop_reason\n\n            # update model weights version\n            global_steps = output.extra_fields.get(\"global_steps\", None)\n            if min_global_steps is None:\n                min_global_steps = global_steps\n            max_global_steps = global_steps\n\n            # 3. update max_new_tokens\n            if original_max_tokens is not None:\n                sampling_params[limit_key] = original_max_tokens - len(final_output.token_ids)\n                if len(final_output.token_ids) >= original_max_tokens:\n                    final_output.stop_reason = \"length\"\n                    break\n\n            # 4. check stop reason\n            if output.stop_reason not in (\"aborted\", \"abort\") or not self.config.async_training.partial_rollout:\n                break\n        final_output.extra_fields[\"global_steps\"] = global_steps\n        final_output.extra_fields[\"min_global_steps\"] = min_global_steps\n        final_output.extra_fields[\"max_global_steps\"] = max_global_steps\n        return final_output\n\n\n@ray.remote\nclass FullyAsyncAgentLoopWorker(AgentLoopWorker):\n    def __init__(\n        self,\n        config: DictConfig,\n        servers: list[tuple[str, ray.actor.ActorHandle]],\n        load_balancer_handle: ray.actor.ActorHandle,\n        reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,\n    ):\n        self.server_manager = FullyAsyncLLMServerManager(config, servers, load_balancer_handle)\n        super().__init__(config, servers, load_balancer_handle, reward_loop_worker_handles)\n\n\nclass FullyAsyncAgentLoopManager(AgentLoopManager):\n    def __init__(\n        self,\n        config: DictConfig,\n        worker_group: RayWorkerGroup = None,\n        rollout_resource_pool: RayResourcePool = None,\n        reward_loop_worker_handles: list[ray.actor.ActorHandle] = None,\n    ):\n        self.agent_loop_workers_class = FullyAsyncAgentLoopWorker\n        super().__init__(config, worker_group, rollout_resource_pool, reward_loop_worker_handles)\n\n    @auto_await\n    async def generate_sequences_single(self, prompts: DataProto) -> DataProto:\n        \"\"\"Split input batch and dispatch to agent loop workers.\n\n        Args:\n            prompts (DataProto): Input batch. Single sample data\n        Returns:\n            DataProto: Output batch.\n        \"\"\"\n        worker = self._select_best_worker()\n        output_future = worker.generate_sequences.remote(prompts)\n        return await asyncio.wrap_future(output_future.future())\n\n    def _select_best_worker(self):\n        \"\"\"Select the best worker, simple round-robin load balancing\"\"\"\n        if not hasattr(self, \"_worker_index\"):\n            self._worker_index = 0\n\n        worker = self.agent_loop_workers[self._worker_index]\n        self._worker_index = (self._worker_index + 1) % len(self.agent_loop_workers)\n        return worker\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/config/fully_async_ppo_megatron_trainer.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_megatron_trainer\n  - _self_\n\ntrainer:\n  use_legacy_worker_impl: disable\n\nasync_training:\n\n  # Maximum samples staleness threshold\n  staleness_threshold: 0.1\n\n  # Frequency of parameter synchronization between rollouter and trainer, \n  # One step means trainer obtains a batch of required samples\n  trigger_parameter_sync_step: 4\n  \n  # The number of ppo_mini_batches that the FullyAsyncTrainer obtains once\n  require_batches: 1\n\n  # When synchronizing parameters, Whether to resume generation when rollout is interrupted.\n  # If True, AsyncLLMServerManager auto resume generation, making rollout interruption invisible to the AgentLoop.\n  partial_rollout: True\n\n  # whether to use trainer do_validate\n  use_trainer_do_validate: False\n\n# Rollout config\nrollout:\n\n  # Number of nodes used in the rollout\n  nnodes: 1\n\n  # Number of GPUs per node                     \n  n_gpus_per_node: 8\n\n  # number of responses (i.e. num sample times). > 1 for grpo\n  n: 4\n\n  # total rollout samples # TODO rename to total_rollout_samples\n  total_rollout_steps: 100\n\ndata:\n  # Number of samples generated, currently only support 1\n  gen_batch_size: 1\n\nactor_rollout_ref:\n\n  rollout:\n    # Must be enabled! Otherwise, log_probs cannot be calculated.\n    calculate_log_probs: True\n\n    checkpoint_engine:\n      backend: \"nccl\"\n\n  actor:\n    # Must use rollout log probs for training\n    use_rollout_log_probs: True\n\n  model:\n    # To use remove padding (thd)\n    use_remove_padding: True\n\n# Only then will the use of log probs be correct.\n# And it can be used in conjunction with other rollout_correction algorithms.\nalgorithm:\n  rollout_correction:\n    bypass_mode: True"
  },
  {
    "path": "verl/experimental/fully_async_policy/config/fully_async_ppo_trainer.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ntrainer:\n  use_legacy_worker_impl: disable\n\nasync_training:\n\n  # Maximum samples staleness threshold\n  staleness_threshold: 0.1\n\n  # Frequency of parameter synchronization between rollouter and trainer, \n  # One step means trainer obtains a batch of required samples\n  trigger_parameter_sync_step: 4\n  \n  # The number of ppo_mini_batches that the FullyAsyncTrainer obtains once\n  require_batches: 1\n\n  # When synchronizing parameters, Whether to resume generation when rollout is interrupted.\n  # If True, AsyncLLMServerManager auto resume generation, making rollout interruption invisible to the AgentLoop.\n  partial_rollout: True\n\n  # whether to use trainer do_validate\n  use_trainer_do_validate: False\n\n# Rollout config\nrollout:\n\n  # Number of nodes used in the rollout\n  nnodes: 1\n\n  # Number of GPUs per node                     \n  n_gpus_per_node: 8\n\n  # number of responses (i.e. num sample times). > 1 for grpo\n  n: 4\n\n  # total rollout samples # TODO rename to total_rollout_samples\n  total_rollout_steps: 100\n\ndata:\n  # Number of samples generated, currently only support 1\n  gen_batch_size: 1\n\nactor_rollout_ref:\n\n  rollout:\n    # Must be enabled! Otherwise, log_probs cannot be calculated.\n    calculate_log_probs: True\n\n    checkpoint_engine:\n      backend: \"nccl\"\n\n  actor:\n    # Must use rollout log probs for training\n    use_rollout_log_probs: True\n\n  model:\n    # To use remove padding (thd)\n    use_remove_padding: True\n\n\n# Only then will the use of log probs be correct.\n# And it can be used in conjunction with other rollout_correction algorithms.\nalgorithm:\n  rollout_correction:\n    bypass_mode: True\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/detach_utils.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport time\nfrom collections import defaultdict\nfrom dataclasses import dataclass\nfrom typing import Any, Optional\n\nimport numpy as np\nimport torch\n\nfrom verl import DataProto\nfrom verl.trainer.ppo.ray_trainer import compute_response_mask\n\n\n@dataclass\nclass RolloutSample:\n    \"\"\"Enhanced rollout sample containing both original batch info and AgentLoopOutput\"\"\"\n\n    # Original batch information\n    full_batch: Any\n\n    # Metadata\n    sample_id: str\n    epoch: int\n\n    # Processing metadata\n    rollout_status: dict[str, Any]\n\n\n@dataclass\nclass ValidateMetrics:\n    \"\"\"Metrics for validation\"\"\"\n\n    timing_raw: dict[str, Any]\n    metrics: Optional[dict[str, Any]] = None\n\n\ndef prepare_single_generation_data(batch_dict, config) -> DataProto:\n    \"\"\"\n    Similar to the logic of ray_trainer._prepare_generate_batch, but for a single sample.\n    Separate the data used for generation from the original data.\n\n    Returns:\n        tuple: (original_batch_dict, gen_data_for_single_sample)\n    \"\"\"\n\n    full_batch = DataProto.from_single_dict(batch_dict)\n\n    batch_keys_to_pop = []\n    non_tensor_batch_keys_to_pop = []\n\n    existing_batch_keys = [k for k in batch_keys_to_pop if k in full_batch.batch.keys()]\n    existing_non_tensor_keys = [k for k in non_tensor_batch_keys_to_pop if k in full_batch.non_tensor_batch.keys()]\n\n    if existing_batch_keys or existing_non_tensor_keys:\n        full_batch.pop(\n            batch_keys=existing_batch_keys,\n            non_tensor_batch_keys=existing_non_tensor_keys,\n        )\n\n    # Setting selected agent, that supports partial\n    if config.actor_rollout_ref.rollout.multi_turn.enable:\n        full_batch.non_tensor_batch[\"agent_name\"] = np.array([\"tool_agent\"] * len(full_batch), dtype=object)\n    else:\n        full_batch.non_tensor_batch[\"agent_name\"] = np.array([\"single_turn_agent\"] * len(full_batch), dtype=object)\n\n    # Add global step count to generated data\n    full_batch = full_batch.repeat(repeat_times=config.actor_rollout_ref.rollout.n, interleave=True)\n    return full_batch\n\n\ndef addition_process(output: DataProto):\n    \"\"\"collect metirics\"\"\"\n    metrics = output.meta_info.pop(\"metrics\")  # List[Dict[str, str]]\n    processing_times_list = [item[\"generate_sequences\"] for item in metrics]\n    tool_calls_times_list = [item[\"tool_calls\"] for item in metrics]\n    output.non_tensor_batch[\"processing_times\"] = processing_times_list\n    output.non_tensor_batch[\"tool_calls_times\"] = tool_calls_times_list\n    return output\n\n\ndef assemble_batch_from_rollout_samples(\n    rollout_samples: list[RolloutSample], tokenizer, config, balance_batch=None\n) -> DataProto:\n    \"\"\"\n    Assemble gen_batch_output from RolloutSample objects\n    Assembles batches from RolloutSample objects, similar to the _post_generate_batch logic in ray_trainer.\n\n    Args:\n        rollout_samples: List of RolloutSample objects\n        tokenizer: Tokenizer instance\n        config: Configuration object containing trainer settings\n        balance_batch: Whether to balance the batch (simplified version)\n\n    Returns:\n        DataProto: Assembled gen_batch_output\n\n    Raises:\n        ValueError: If rollout_samples is empty\n    \"\"\"\n    start_time = time.time()\n\n    if not rollout_samples:\n        raise ValueError(\"Empty rollout_samples provided for batch assembly\")\n\n    print(f\"[BatchUtils] Assembling batch from {len(rollout_samples)} RolloutSample objects\")\n\n    rollout_samples_batch = []\n    rollout_status = rollout_samples[0].rollout_status\n    # Add a prefix to all rollout_status keys\n    rollout_status = {f\"fully_async/{key}\": value for key, value in rollout_status.items()}\n\n    for rs in rollout_samples:\n        batch = addition_process(rs.full_batch)\n        rollout_samples_batch.append(batch)\n    final_batch = DataProto.concat(rollout_samples_batch)\n\n    # Calculate response_mask (if not present)\n    if \"response_mask\" not in final_batch.batch.keys():\n        final_batch.batch[\"response_mask\"] = compute_response_mask(final_batch)\n\n    if balance_batch:\n        balance_batch(final_batch, metrics={})\n\n    # Calculate the global valid token number\n    if \"attention_mask\" in final_batch.batch:\n        final_batch.meta_info[\"global_token_num\"] = torch.sum(final_batch.batch[\"attention_mask\"], dim=-1).tolist()\n\n    processing_times = final_batch.non_tensor_batch[\"processing_times\"]\n    tool_calls = final_batch.non_tensor_batch[\"tool_calls_times\"]\n    # Collect statistics\n    processing_time_stats = {\n        \"processing_time/avg\": np.mean(processing_times),\n        \"processing_time/max\": np.max(processing_times),\n        \"processing_time/min\": np.min(processing_times),\n        \"processing_time/tp50\": np.percentile(processing_times, 50),\n        \"processing_time/tp99\": np.percentile(processing_times, 99),\n        \"processing_time/tp95\": np.percentile(processing_times, 95),\n    }\n    tool_calls_stats = {}\n    if len(tool_calls) > 0:\n        tool_calls_stats = {\n            \"timing_s/agent_loop/tool_calls/max\": np.max(tool_calls),\n            \"timing_s/agent_loop/tool_calls/min\": np.min(tool_calls),\n            \"timing_s/agent_loop/tool_calls/mean\": np.mean(tool_calls),\n        }\n    processing_time_stats = {f\"fully_async/{key}\": value for key, value in processing_time_stats.items()}\n\n    param_version_start = final_batch.non_tensor_batch[\"min_global_steps\"]\n    param_version_end = final_batch.non_tensor_batch[\"max_global_steps\"]\n    param_version_diff = [abs(a - b) for a, b in zip(param_version_end, param_version_start, strict=False)]\n    num_diff0 = param_version_diff.count(0)\n    partial_stats = {\n        \"fully_async/partial/total_partial_num\": len(param_version_diff) - num_diff0,\n        \"fully_async/partial/partial_ratio\": (len(param_version_diff) - num_diff0) / len(param_version_diff),\n        \"fully_async/partial/max_partial_span\": max(param_version_diff),\n    }\n    # add meta_info\n    trajectory_param_versions = final_batch.non_tensor_batch[\"max_global_steps\"]\n\n    final_batch.meta_info.update(\n        {\n            \"param_version_diversity\": len(set(trajectory_param_versions)),\n            \"trajectory_param_versions\": trajectory_param_versions,\n            **processing_time_stats,\n            **rollout_status,\n            **partial_stats,\n            **tool_calls_stats,\n        }\n    )\n\n    print(f\"[BatchUtils] Batch assembly completed in {time.time() - start_time:.2f}s\")\n\n    return final_batch\n\n\nclass MetricsAggregator:\n    \"\"\"Metrics aggregator, used to combine metrics from multiple training steps\"\"\"\n\n    def __init__(self, total_gpus: int):\n        # Store all values ​​for each metric\n        self.metric_values: dict[str, list[float]] = defaultdict(list)\n        # Store the number of samples at each step for weighted averaging\n        self.sample_counts: list[int] = []\n        # Store the timestamp of each step for time-related calculations\n        self.timestamps: list[float] = []\n        # Step Count\n        self.step_count = 0\n        # total num gpus used\n        self.total_gpus = total_gpus\n\n        # Metric aggregation rule configuration\n        self.aggregation_rules = self._init_aggregation_rules()\n\n    def _init_aggregation_rules(self) -> dict[str, dict[str, list[str]]]:\n        \"\"\"Initialize metrics aggregation rules\"\"\"\n        return {\n            # Time-Based metrics, can add metrics here\n            \"time_sum\": [\"perf/time_per_step\"],\n            \"min\": [\"timing_s/agent_loop/tool_calls/min\"],\n            \"avg\": [\"timing_s/agent_loop/tool_calls/mean\"],\n            \"max\": [\"timing_s/agent_loop/tool_calls/max\"],\n            \"last\": [\n                \"fully_async/count/total_generated_samples\",\n                \"fully_async/count/stale_samples_processed\",\n                \"fully_async/count/stale_trajectory_processed\",\n                \"fully_async/count/current_param_version\",\n                \"fully_async/count/dropped_stale_samples\",\n                \"training/global_step\",  # TODO change name to: total_step\n            ],\n        }\n\n    def add_step_metrics(self, metrics: dict[str, Any], sample_count: int, timestamp: float = None):\n        \"\"\"Adding a single-step metrics\"\"\"\n        if timestamp is None:\n            timestamp = time.time()\n\n        self.sample_counts.append(sample_count)\n        self.timestamps.append(timestamp)\n        self.step_count += 1\n\n        # Store all metrics values\n        for key, value in metrics.items():\n            if isinstance(value, int | float | np.number):\n                self.metric_values[key].append(float(value))\n            elif isinstance(value, torch.Tensor):\n                self.metric_values[key].append(float(value.item()))\n\n    def _get_aggregation_type(self, metric_name: str) -> str:\n        \"\"\"Determine the aggregation type based on the metric name\"\"\"\n        for agg_type, metric_list in self.aggregation_rules.items():\n            if metric_name in metric_list:\n                return agg_type\n\n        metric_lower = metric_name.lower()\n        if any(keyword in metric_lower for keyword in [\"timing_s/\"]):\n            return \"time_sum\"\n        if any(keyword in metric_lower for keyword in [\"mean\", \"avg\", \"average\"]):\n            return \"avg\"\n        if any(keyword in metric_lower for keyword in [\"max\", \"maximum\"]):\n            return \"max\"\n        if any(keyword in metric_lower for keyword in [\"min\", \"minimum\"]):\n            return \"min\"\n        if any(keyword in metric_lower for keyword in [\"sum\", \"total\"]):\n            return \"sum\"\n        if any(keyword in metric_lower for keyword in [\"weighted_avg\"]):\n            return \"weighted_avg\"\n\n        return \"avg\"\n\n    def _aggregate_single_metric(self, metric_name: str, values: list[float]) -> float:\n        \"\"\"Aggregating a single metric\"\"\"\n        if not values:\n            return 0.0\n\n        agg_type = self._get_aggregation_type(metric_name)\n\n        if agg_type == \"last\":\n            return values[-1]\n\n        elif agg_type == \"weighted_avg\":\n            # Weighted average\n            if len(values) != len(self.sample_counts):\n                # If the lengths do not match, use a simple average\n                return sum(values) / len(values)\n\n            total_samples = sum(self.sample_counts)\n            if total_samples == 0:\n                return sum(values) / len(values)\n\n            weighted_sum = sum(v * c for v, c in zip(values, self.sample_counts, strict=False))\n            return weighted_sum / total_samples\n\n        elif agg_type == \"sum\" or agg_type == \"time_sum\":\n            return sum(values)\n\n        elif agg_type == \"avg\":\n            return sum(values) / len(values)\n\n        elif agg_type == \"max\":\n            return max(values)\n\n        elif agg_type == \"min\":\n            return min(values)\n\n        else:\n            # Default average\n            return sum(values) / len(values)\n\n    def get_aggregated_metrics(self) -> dict[str, Any]:\n        \"\"\"aggregated metrics\"\"\"\n        t = time.time()\n        if self.step_count == 0:\n            return {}\n\n        aggregated = {}\n\n        # Aggregate all metrics\n        for metric_name, values in self.metric_values.items():\n            aggregated[metric_name] = self._aggregate_single_metric(metric_name, values)\n\n        # Aggregate special metrics\n        aggregated = self._special_metrics_aggergate(aggregated)\n\n        print(f\"aggregated metrics done. cost {time.time() - t:.4f} seconds.\")\n\n        return aggregated\n\n    def _special_metrics_aggergate(self, aggregated: dict[str, Any]) -> dict[str, Any]:\n        \"\"\"calculate special metrics\"\"\"\n\n        # global_seqlen/minmax_diff\n        if \"global_seqlen/minmax_diff\" in aggregated.keys():\n            aggregated[\"global_seqlen/minmax_diff\"] = aggregated[\"global_seqlen/max\"] - aggregated[\"global_seqlen/min\"]\n\n        # perf/throughput\n        REQUIRED_PERF_KEYS = {\"perf/throughput\", \"perf/total_num_tokens\", \"perf/time_per_step\"}\n        if REQUIRED_PERF_KEYS.issubset(aggregated):\n            aggregated[\"perf/throughput\"] = aggregated[\"perf/total_num_tokens\"] / (\n                aggregated[\"perf/time_per_step\"] * self.total_gpus\n            )\n\n        # trainer/idle_ratio\n        if \"timing_s/gen\" in aggregated.keys() and \"timing_s/step\" in aggregated.keys():\n            aggregated[\"fully_async/trainer/idle_ratio\"] = aggregated[\"timing_s/gen\"] / aggregated[\"timing_s/step\"]\n\n        return aggregated\n\n    def reset(self):\n        \"\"\"Reset Aggregator\"\"\"\n        self.metric_values.clear()\n        self.sample_counts.clear()\n        self.timestamps.clear()\n        self.step_count = 0\n\n    def get_current_stats(self) -> dict[str, Any]:\n        \"\"\"Get statistics about the current aggregation state (for debugging)\"\"\"\n        return {\n            \"step_count\": self.step_count,\n            \"metric_count\": len(self.metric_values),\n            \"total_samples\": sum(self.sample_counts),\n            \"metric_names\": list(self.metric_values.keys()),\n        }\n\n\ndef task_exception_handler(task: asyncio.Task):\n    \"\"\"Handle task exceptions and log them\"\"\"\n    try:\n        task.result()\n    except asyncio.CancelledError:\n        pass  # Task was cancelled, this is expected\n    except Exception as e:\n        print(f\"Task {task.get_name()} failed with exception: {e}\")\n        raise e\n\n\ndef safe_create_task(coro, name: str, task_set: set = None):\n    \"\"\"Safely create a task with exception handling\n\n    Args:\n        coro: The coroutine to run\n        name: Name for the task\n        task_set: Optional set to add the task to\n\n    Returns:\n        The created asyncio.Task\n    \"\"\"\n    task = asyncio.create_task(coro, name=name)\n    task.add_done_callback(task_exception_handler)\n    if task_set is not None:\n        task_set.add(task)\n    return task\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/fully_async_main.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport os\nimport socket\nimport threading\nfrom pprint import pprint\n\nimport hydra\nimport ray\nfrom omegaconf import OmegaConf\n\nfrom verl.experimental.fully_async_policy.fully_async_rollouter import FullyAsyncRollouter\nfrom verl.experimental.fully_async_policy.fully_async_trainer import FullyAsyncTrainer\nfrom verl.experimental.fully_async_policy.message_queue import MessageQueue, MessageQueueClient\nfrom verl.experimental.separation.utils import create_resource_pool_manager, create_role_worker_mapping\nfrom verl.trainer.ppo.utils import Role\nfrom verl.utils.device import auto_set_device\nfrom verl.utils.fs import copy_to_local\n\n\n@ray.remote(num_cpus=1)\nclass FullyAsyncTaskRunner:\n    \"\"\"\n    Ray remote class for executing distributed PPO training tasks.\n    \"\"\"\n\n    def __init__(self):\n        self.running = False\n        self.components = {}\n        self.shutdown_event = threading.Event()\n\n    def run(self, config):\n        print(\"[ASYNC MAIN] Starting fully async PPO training...\")\n        self._initialize_components(config)\n        self._run_training_loop()\n\n    def _initialize_components(self, config) -> None:\n        print(f\"[ASYNC MAIN] TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}\")\n        pprint(OmegaConf.to_container(config, resolve=True))\n        OmegaConf.resolve(config)\n\n        print(\"[ASYNC MAIN] Initializing model and tokenizer...\")\n        local_path = copy_to_local(\n            config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get(\"use_shm\", False)\n        )\n        from verl.utils import hf_processor, hf_tokenizer\n\n        trust_remote_code = config.data.get(\"trust_remote_code\", False)\n        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)\n\n        # Used for multimodal LLM, could be None\n        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)\n\n        self.components[\"tokenizer\"] = tokenizer\n        self.components[\"processor\"] = processor\n        self.components[\"config\"] = config\n\n        print(\"[ASYNC MAIN] Creating worker mapping and resource pools...\")\n        role_worker_mapping, ray_worker_group_cls = create_role_worker_mapping(config)\n        self.components[\"role_worker_mapping\"] = role_worker_mapping\n        self.components[\"ray_worker_group_cls\"] = ray_worker_group_cls\n\n        from concurrent.futures import ThreadPoolExecutor\n\n        print(\"[ASYNC MAIN] Creating FullyAsyncRollouter and FullyAsyncTrainer in parallel...\")\n        with ThreadPoolExecutor(max_workers=2) as executor:\n            # Rollouter does not permit continuous allocation, so we allocate trainer first.\n            trainer_future = executor.submit(self._create_trainer, config)\n            trainer_future.result()\n\n            rollouter_future = executor.submit(self._create_rollouter, config)\n            rollouter_future.result()\n\n        # sync total_train_steps between rollouter and trainer\n        total_train_steps = ray.get(self.components[\"rollouter\"].get_total_train_steps.remote())\n        print(f\"total_train_steps {total_train_steps}\")\n        ray.get(self.components[\"trainer\"].set_total_train_steps.remote(total_train_steps))\n\n        # max_queue_size\n        max_queue_size = ray.get(self.components[\"rollouter\"].get_max_queue_size.remote())\n        print(f\"[ASYNC MAIN] Creating MessageQueue... max_queue_size {max_queue_size}\")\n        message_queue = MessageQueue.remote(config, max_queue_size)\n        message_queue_client = MessageQueueClient(message_queue)\n        self.components[\"message_queue\"] = message_queue\n        self.components[\"message_queue_client\"] = message_queue_client\n\n        ray.get(self.components[\"rollouter\"].set_message_queue_client.remote(self.components[\"message_queue_client\"]))\n        ray.get(self.components[\"trainer\"].set_message_queue_client.remote(self.components[\"message_queue_client\"]))\n\n        # param_version resume from ckpt or default 0\n        ray.get(self.components[\"trainer\"].load_checkpoint.remote())\n        ray.get(self.components[\"rollouter\"].load_checkpoint.remote())\n\n        print(\"[ASYNC MAIN] Setting up parameter synchronization...\")\n        ray.get(self.components[\"trainer\"].set_rollouter.remote(self.components[\"rollouter\"]))\n\n        print(\"[ASYNC MAIN] Param sync before fit..\")\n        ray.get(self.components[\"trainer\"]._fit_update_weights.remote())\n\n        if config.trainer.get(\"val_before_train\", True):\n            ray.get(self.components[\"trainer\"]._fit_validate.remote(True))\n\n        print(\"[ASYNC MAIN] All components initialized successfully\")\n\n    def _create_rollouter(self, config) -> None:\n        print(\"[ASYNC MAIN] Starting create rollouter...\")\n        rollouter = FullyAsyncRollouter.remote(\n            config=config,\n            tokenizer=self.components[\"tokenizer\"],\n            role_worker_mapping=None,\n            resource_pool_manager=create_resource_pool_manager(config, roles=[Role.Rollout]),\n            ray_worker_group_cls=self.components[\"ray_worker_group_cls\"],\n            processor=self.components[\"processor\"],\n            device_name=config.trainer.device,\n        )\n\n        ray.get(rollouter.init_workers.remote())\n        ray.get(rollouter.set_max_required_samples.remote())\n\n        self.components[\"rollouter\"] = rollouter\n        print(\"[ASYNC MAIN] Rollouter created and initialized successfully\")\n\n    def _create_trainer(self, config) -> None:\n        print(\"[ASYNC MAIN] Starting create trainer...\")\n        trainer_role_mapping = {\n            role: worker_cls\n            for role, worker_cls in self.components[\"role_worker_mapping\"].items()\n            if role != Role.Rollout\n        }\n\n        trainer = FullyAsyncTrainer.remote(\n            config=config,\n            tokenizer=self.components[\"tokenizer\"],\n            role_worker_mapping=trainer_role_mapping,\n            resource_pool_manager=create_resource_pool_manager(config, roles=list(trainer_role_mapping.keys())),\n            ray_worker_group_cls=self.components[\"ray_worker_group_cls\"],\n            processor=self.components[\"processor\"],\n            device_name=config.trainer.device,\n        )\n\n        ray.get(trainer.init_workers.remote())\n        self.components[\"trainer\"] = trainer\n        print(\"[ASYNC MAIN] FullyAsyncTrainer created and initialized successfully\")\n\n    def _run_training_loop(self):\n        self.running = True\n\n        print(\"[ASYNC MAIN] Starting Rollouter and Trainer...\")\n        rollouter_future = self.components[\"rollouter\"].fit.remote()\n        trainer_future = self.components[\"trainer\"].fit.remote()\n\n        futures = [rollouter_future, trainer_future]\n\n        try:\n            while futures:\n                # Use ray.wait to monitor all futures and return when any one is completed.\n                done_futures, remaining_futures = ray.wait(futures, num_returns=1, timeout=None)\n\n                for future in done_futures:\n                    try:\n                        ray.get(future)\n                        print(\"[ASYNC MAIN] One component completed successfully\")\n                    except Exception as e:\n                        print(f\"[ASYNC MAIN] Component failed with error: {e}\")\n                        for remaining_future in remaining_futures:\n                            ray.cancel(remaining_future)\n                        raise e\n\n                futures = remaining_futures\n\n        except Exception as e:\n            print(f\"[ASYNC MAIN] Training failed: {e}\")\n            for future in futures:\n                ray.cancel(future)\n            raise\n        finally:\n            asyncio.run(self.components[\"message_queue_client\"].clear_queue())\n            print(\"[ASYNC MAIN] Training completed or interrupted\")\n\n\n@hydra.main(config_path=\"config\", config_name=\"fully_async_ppo_trainer\", version_base=None)\ndef main(config):\n    from verl.trainer.main_ppo import run_ppo\n\n    # Ensure async training config exists\n    if not hasattr(config, \"async_training\"):\n        raise RuntimeError(\"must set async_training config\")\n\n    assert config.async_training.use_trainer_do_validate is False, \"use_trainer_do_validate is not ready to use.\"\n\n    from time import time\n\n    start_time = time()\n    auto_set_device(config)\n    # TODO: unify rollout config with actor_rollout_ref\n    config.actor_rollout_ref.rollout.nnodes = config.rollout.nnodes\n    config.actor_rollout_ref.rollout.n_gpus_per_node = config.rollout.n_gpus_per_node\n    run_ppo(config, task_runner_class=FullyAsyncTaskRunner)\n    print(f\"total time: {time() - start_time:.2f} seconds\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/fully_async_rollouter.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport multiprocessing\nimport os\nimport time\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pprint import pformat\n\nimport numpy as np\nimport ray\nimport torch\n\nfrom verl.experimental.fully_async_policy.detach_utils import (\n    RolloutSample,\n    ValidateMetrics,\n    prepare_single_generation_data,\n    safe_create_task,\n)\nfrom verl.experimental.fully_async_policy.message_queue import MessageQueueClient\nfrom verl.experimental.separation.ray_trainer import SeparateRayPPOTrainer\nfrom verl.single_controller.ray import RayWorkerGroup\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.trainer.ppo.utils import Role, WorkerType\nfrom verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path\nfrom verl.utils.profiler import marked_timer\nfrom verl.utils.tracking import ValidationGenerationsLogger\n\n\n@ray.remote(num_cpus=10, max_concurrency=100)\nclass FullyAsyncRollouter(SeparateRayPPOTrainer):\n    \"\"\"\n    Asynchronous sample generator, responsible for continuously generating training samples\n    and putting them into MessageQueue\n    Based on the mature implementation improvements of OneStepOffRayTrainer\n    \"\"\"\n\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        role_worker_mapping: dict[Role, WorkerType],\n        resource_pool_manager: ResourcePoolManager,\n        ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,\n        processor=None,\n        device_name=None,\n    ):\n        # Store the tokenizer for text processing\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.config = config\n        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine\n\n        assert not self.hybrid_engine\n        assert self.config.data.train_batch_size == 0, \"train_batch_size must be zero\"\n        assert self.config.data.gen_batch_size == 1, \"gen_batch_size must be one\"\n        assert self.config.async_training.staleness_threshold >= 0, \"staleness_threshold must larger than 0\"\n        assert self.config.async_training.trigger_parameter_sync_step >= 1, (\n            \"trigger_parameter_sync_step must larger or equal than 1\"\n        )\n\n        self.role_worker_mapping = role_worker_mapping\n        self.resource_pool_manager = resource_pool_manager\n        self.use_reference_policy = False\n\n        self.use_rm = False\n\n        self.use_critic = False\n        self.ray_worker_group_cls = ray_worker_group_cls\n        self.device_name = device_name if device_name else self.config.trainer.device\n        self.validation_generations_logger = ValidationGenerationsLogger(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n        )\n\n        self.ref_in_actor = False\n        self.kl_ctrl_in_reward = False\n\n        self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get(\"use_prefix_grouper\", False)\n        self.use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n\n        # ==================== fully async config ====================\n\n        print(\"[FullyAsyncRollouter] Creating datasets...\")\n        from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler\n        from verl.utils.dataset.rl_dataset import collate_fn\n\n        train_dataset = create_rl_dataset(\n            config.data.train_files,\n            config.data,\n            tokenizer,\n            processor,\n            max_samples=config.data.get(\"train_max_samples\", -1),\n        )\n        val_dataset = create_rl_dataset(\n            config.data.val_files,\n            config.data,\n            tokenizer,\n            processor,\n            max_samples=config.data.get(\"val_max_samples\", -1),\n        )\n        train_sampler = create_rl_sampler(config.data, train_dataset)\n\n        self._validate_config()\n        if self.config.async_training.use_trainer_do_validate:\n            rollout_gpus = config.rollout.nnodes * config.rollout.n_gpus_per_node\n            train_gpus = config.trainer.nnodes * config.trainer.n_gpus_per_node\n            total_gpus = rollout_gpus + train_gpus\n            print(f\"[FullyAsyncRollouter] split before val_dataset total len: {len(val_dataset)}\")\n            split_dataset = val_dataset.split(total_gpus)\n            rollout_val_dataset0 = split_dataset[:rollout_gpus]\n            from torch.utils.data import ConcatDataset\n\n            val_dataset = ConcatDataset(rollout_val_dataset0)\n            print(f\"[FullyAsyncRollouter] split after val_dataset total len: {len(val_dataset)}\")\n        print(f\"[FullyAsyncRollouter] Rollouter _create_dataloader...\\n{train_dataset}\\n{val_dataset}\")\n\n        self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)\n\n        self.total_rollout_steps = len(self.train_dataloader) * self.config.trainer.total_epochs\n        if self.config.rollout.total_rollout_steps is not None:\n            self.total_rollout_steps = min(self.config.rollout.total_rollout_steps, self.total_rollout_steps)\n        print(f\"[FullyAsyncRollouter] Total rollout steps: {self.total_rollout_steps}\")\n        self.total_train_steps = None\n\n        # Rollouter parameter configuration\n        self.message_queue_client = None\n\n        # Worker groups: rollout_wg is same to actor_rollout_wg\n        self.rollout_wg = None\n        self.actor_rollout_wg = None\n        self.async_rollout_manager = None\n\n        # Config\n        self.staleness_threshold: float = config.async_training.get(\"staleness_threshold\", 1)\n        # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.\n        self.require_batches = config.async_training.require_batches\n        self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches\n        self.max_required_samples = None\n        self.max_concurrent_samples = None\n        # queue size\n        self.max_queue_size = None\n\n        # Statistics\n        self.total_generated_samples = 0\n        self.staleness_samples = 0\n        self.dropped_stale_samples = 0\n        self.processed_sample_count = 0\n        # we start from step 1\n        self.global_steps = 1\n        self.idle_start_time = time.time()\n        self.step_start_time = time.time()\n\n        # Concurrency control\n        # Modified by self.pause() or self._should_pause_generation()\n        self.paused = False\n        self.running = True\n\n        # Add dataloader lock\n        self.dataloader_lock = asyncio.Lock()\n\n        # Initialize async queues\n        self.pending_queue = asyncio.Queue(maxsize=128)\n        self.active_tasks = set()\n\n        cpu_cores = multiprocessing.cpu_count()\n        # cpu case use cpu_cores; io case use cpu_cores*2\n        self.validate_executor = ThreadPoolExecutor(max_workers=cpu_cores)\n        self.validate_task = None\n\n    def _init_async_objects(self):\n        # Initialize asyncio synchronization primitives.\n        # We let asyncio.Condition create the Lock internally to ensure they share the same Event Loop.\n        # This avoids 'ValueError: loop argument must agree with lock' which can occur in Ray environments\n        # where the lock's captured loop (get_running_loop) differs from Condition's default loop check.\n        # Explicitly passing the loop is deprecated/removed in Python 3.10+, so this reverse-initialization\n        # is the most robust workaround.\n        self.condition = asyncio.Condition()\n        self.lock = self.condition._lock\n\n    async def set_message_queue_client(self, message_queue_client: MessageQueueClient):\n        \"\"\"Set message queue client\"\"\"\n        async with self.lock:\n            self.message_queue_client = message_queue_client\n\n    async def set_max_required_samples(self):\n        async with self.lock:\n            self.max_required_samples = int(\n                self.required_samples\n                * (self.staleness_threshold + 1)\n                * self.config.async_training.trigger_parameter_sync_step\n            )\n            self.total_train_steps = int(\n                self.total_rollout_steps\n                / (self.required_samples * self.config.async_training.trigger_parameter_sync_step)\n            )\n\n            self.max_concurrent_samples = len(self.async_rollout_manager.server_handles) * 16\n            self.max_concurrent_samples = min(self.max_concurrent_samples, self.max_required_samples)\n            self.max_queue_size = self.max_required_samples\n\n            print(\n                f\"[FullyAsyncRollouter] required_samples : {self.required_samples} \"\n                f\"max_required_samples: {self.max_required_samples} \"\n                f\"max_queue_size: {self.max_queue_size} \"\n                f\"total_train_steps: {self.total_train_steps} \"\n                f\"total_rollout_steps: {self.total_rollout_steps} \"\n                f\"max_concurrent_samples: {self.max_concurrent_samples} \"\n            )\n\n    def get_rollout_wg(self):\n        \"\"\"Get rollout worker group\"\"\"\n        return self.rollout_wg\n\n    def get_replicas(self):\n        \"\"\"Get rollout worker group\"\"\"\n        return self.async_rollout_manager.rollout_replicas\n\n    def get_max_queue_size(self):\n        return self.max_queue_size\n\n    def get_total_train_steps(self):\n        return self.total_train_steps\n\n    async def reset_staleness(self):\n        \"\"\"\n        Reset staleness samples after parameter update.\n        Returns timing_raw dictionary for metrics.\n        \"\"\"\n        async with self.lock:\n            self.paused = False\n            self.condition.notify_all()\n            # every time param change, reset staleness_samples\n            self.staleness_samples = len(self.active_tasks) + await self.message_queue_client.get_queue_size()\n            timing_raw = {}\n            rollout_active_time = self.idle_start_time - self.step_start_time\n            rollout_version_time = time.time() - self.step_start_time\n            idle_ratio = 1 - rollout_active_time / rollout_version_time\n            timing_raw[\"fully_async/rollouter/active_time\"] = rollout_active_time\n            timing_raw[\"fully_async/rollouter/version_time\"] = rollout_version_time\n            timing_raw[\"fully_async/rollouter/idle_ratio\"] = idle_ratio\n\n            print(\n                f\"[FullyAsyncRollouter][Public][reset_staleness] \"\n                f\"reset staleness_samples to: {self.staleness_samples} \"\n                f\"idle_ratio: {timing_raw['fully_async/rollouter/idle_ratio']:.4f}\"\n            )\n            self.step_start_time = time.time()\n        return timing_raw\n\n    def do_validate(self) -> ValidateMetrics:\n        \"\"\"Run validation and return metrics\"\"\"\n        timing_raw = {}\n        with marked_timer(\"rollouter/validate_time\", timing_raw, color=\"green\"):\n            val_metrics: dict = self._validate()\n        return ValidateMetrics(timing_raw=timing_raw, metrics=val_metrics)\n\n    async def save_checkpoint(self, local_global_step_folder: str):\n        # WARNING!: Due to the asynchronous nature, there are some in-flight samples\n        # (pending/cancel/result queue and message queue).\n        # Therefore, directly saving the state of the dataloader will result in losing these\n        # samples when resuming training.\n        # TODO: Implement dataloader recovery without losing in-flight samples.\n        from verl.utils.fs import local_mkdir_safe\n\n        # save dataloader\n        local_mkdir_safe(local_global_step_folder)\n        dataloader_local_path = os.path.join(local_global_step_folder, \"data.pt\")\n        async with self.dataloader_lock:\n            dataloader_state_dict = self.train_dataloader.state_dict()\n        torch.save(dataloader_state_dict, dataloader_local_path)\n        print(f\"[FullyAsyncRollouter] Saved dataloader checkpoint to {dataloader_local_path}\")\n\n    def load_checkpoint(self):\n        \"\"\"Load checkpoint including dataloader state based on resume mode\"\"\"\n\n        if self.config.trainer.resume_mode == \"disable\":\n            print(\"[FullyAsyncRollouter] Resume mode is disabled, starting from scratch\")\n            return 0\n\n        # Determine checkpoint folder path\n        if self.config.trainer.default_hdfs_dir is not None:\n            raise NotImplementedError(\"[FullyAsyncRollouter] Load from hdfs is not implemented yet\")\n        else:\n            checkpoint_folder = self.config.trainer.default_local_dir\n            if not os.path.isabs(checkpoint_folder):\n                working_dir = os.getcwd()\n                checkpoint_folder = os.path.join(working_dir, checkpoint_folder)\n\n            global_step_folder = find_latest_ckpt_path(checkpoint_folder)\n\n        # Find and validate global_step_folder based on resume mode\n        if self.config.trainer.resume_mode == \"auto\":\n            if global_step_folder is None:\n                print(\"[FullyAsyncRollouter] Training from scratch (no checkpoint found)\")\n                return 0\n        elif self.config.trainer.resume_mode == \"resume_path\":\n            assert isinstance(self.config.trainer.resume_from_path, str), (\n                \"[FullyAsyncRollouter] resume_from_path must be str type\"\n            )\n            assert \"global_step_\" in self.config.trainer.resume_from_path, (\n                \"[FullyAsyncRollouter] resume_from_path must specify the global_steps\"\n            )\n            global_step_folder = self.config.trainer.resume_from_path\n            if not os.path.isabs(global_step_folder):\n                working_dir = os.getcwd()\n                global_step_folder = os.path.join(working_dir, global_step_folder)\n        else:\n            raise ValueError(f\"[FullyAsyncRollouter] Unknown resume_mode: {self.config.trainer.resume_mode}\")\n\n        print(f\"[FullyAsyncRollouter] Loading checkpoint from: {global_step_folder}\")\n\n        # Extract and set global step\n        trainer_global_steps = int(global_step_folder.split(\"global_step_\")[-1])\n        self.global_steps = (\n            trainer_global_steps * self.required_samples * self.config.async_training.trigger_parameter_sync_step + 1\n        )\n        print(f\"[FullyAsyncRollouter] Setting global_steps to {self.global_steps}\")\n\n        # Load dataloader state\n        dataloader_local_path = os.path.join(global_step_folder, \"data.pt\")\n        if os.path.exists(dataloader_local_path):\n            dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False)\n            self.train_dataloader.load_state_dict(dataloader_state_dict)\n            print(f\"[FullyAsyncRollouter] Loaded dataloader state from {dataloader_local_path}\")\n        else:\n            print(\n                f\"[FullyAsyncRollouter] Warning: No dataloader state found at {dataloader_local_path}, \"\n                f\"will start from scratch\"\n            )\n\n    def _validate_config(self):\n        # Validate asynchronous training configuration\n        if not hasattr(self.config, \"async_training\"):\n            raise ValueError(\"[FullyAsyncRollouter] Missing async_training configuration\")\n        assert self.config.actor_rollout_ref.rollout.calculate_log_probs, \"must rollout calculate log_probs\"\n\n    async def init_workers(self):\n        \"\"\"Initialize distributed training workers using Ray backend.\n\n        Creates:\n        1. Ray resource pools from configuration\n        2. Worker groups for each role (actor, critic, etc.)\n        \"\"\"\n        self._init_async_objects()\n        self._create_worker_classes()\n        self._init_reward_loop()\n        await self._init_async_rollout_manager()\n\n    def _create_actor_rollout_classes(self):\n        # Skip rollout creation and let agentloop handle it\n        pass\n\n    def _init_models(self):\n        self.rollout_wg = self.all_wg[str(Role.Rollout)]\n        self.rollout_wg.init_model()\n        self.actor_rollout_wg = self.rollout_wg\n\n    def _create_continuous_iterator(self):\n        \"\"\"\n        Create a continuous data iterator across epoch\n        \"\"\"\n        for epoch in range(self.config.trainer.total_epochs):\n            iterator = iter(self.train_dataloader)\n            for batch_dict in iterator:\n                yield epoch, batch_dict\n\n    async def _init_async_rollout_manager(self):\n        # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design\n        # agent_reward_loop: streaming reward computation with actor rollout\n        # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool\n        enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool\n\n        # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager\n        # to stream reward computation with actor rollout\n        reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None\n\n        # create async rollout manager and request scheduler\n        assert self.config.actor_rollout_ref.rollout.mode == \"async\"\n        from verl.experimental.fully_async_policy.agent_loop import FullyAsyncAgentLoopManager\n\n        self.async_rollout_mode = True\n        self.async_rollout_manager = await FullyAsyncAgentLoopManager.create(\n            config=self.config, worker_group=self.rollout_wg, reward_loop_worker_handles=reward_loop_worker_handles\n        )\n\n    # Add samples to the pending_queue\n    async def _feed_samples(self):\n        continuous_iterator = self._create_continuous_iterator()\n\n        for epoch, batch_dict in continuous_iterator:\n            # Similar to _prepare_generate_batch: Separate data\n            full_batch = prepare_single_generation_data(batch_dict, self.config)\n\n            sample_id = f\"sample_{epoch}_{self.global_steps}\"\n\n            rollout_sample = RolloutSample(\n                full_batch=full_batch,\n                sample_id=sample_id,\n                epoch=epoch,\n                rollout_status={},\n            )\n\n            await self.pending_queue.put(rollout_sample)\n\n            # Check if have reached the last step\n            if self.global_steps >= self.total_rollout_steps:\n                print(\n                    f\"[FullyAsyncRollouter][Feed] \"\n                    f\"Maximum count has been reached, stop adding new samples: \"\n                    f\"{self.global_steps} >= {self.total_rollout_steps}\"\n                )\n                break\n\n            self.global_steps += 1\n\n        # End signal\n        await self.pending_queue.put(None)\n        print(f\"[FullyAsyncRollouter][Feed] Sample addition is complete, {self.global_steps} samples have been added\")\n\n    async def _processor_worker(self):\n        \"\"\"\n        Streaming worker coroutines, a sample is submitted for processing without waiting for batches\n        \"\"\"\n        while True:\n            if self.paused or await self._should_pause_generation():\n                print(\n                    \"[FullyAsyncRollouter][Processor] Received pause signal, waiting for remaining tasks to return...\"\n                )\n                async with self.lock:\n                    self.paused = True\n                while self.active_tasks:\n                    async with self.lock:\n                        # After acquiring the lock, the number of active_tasks may change, need to be verified again\n                        if self.active_tasks:\n                            done_tasks, self.active_tasks = await asyncio.wait(\n                                self.active_tasks, return_when=asyncio.FIRST_COMPLETED\n                            )\n                            for task in done_tasks:\n                                await task\n\n                async with self.lock:\n                    while self.paused:\n                        self.idle_start_time = time.time()\n                        await self.condition.wait()\n                continue\n            # Get sample from appropriate queue and immediately mark task as done\n            rollout_sample = await self.pending_queue.get()\n            self.pending_queue.task_done()\n            self.staleness_samples += 1\n\n            if rollout_sample is None:\n                print(\n                    \"[FullyAsyncRollouter][Processor] Received end signal, waiting for remaining tasks to complete...\"\n                )\n                while self.active_tasks:\n                    async with self.lock:\n                        if self.active_tasks:\n                            done_tasks, self.active_tasks = await asyncio.wait(\n                                self.active_tasks, return_when=asyncio.FIRST_COMPLETED\n                            )\n                            for task in done_tasks:\n                                await task\n                break\n\n            # Check whether the number of concurrent tasks exceeds the limit\n            while len(self.active_tasks) >= self.max_concurrent_samples:\n                async with self.lock:\n                    if self.active_tasks:\n                        done_tasks, self.active_tasks = await asyncio.wait(\n                            self.active_tasks, return_when=asyncio.FIRST_COMPLETED\n                        )\n                        for task in done_tasks:\n                            await task\n\n            # Submit single sample processing\n            async with self.lock:\n                # After the pause is over, the lock is acquired and it is necessary\n                # to determine whether it is the pause phase, otherwise continue to wait\n                while self.paused:\n                    await self.condition.wait()\n                task = safe_create_task(\n                    self._process_single_sample_streaming(rollout_sample),\n                    name=rollout_sample.sample_id,\n                    task_set=self.active_tasks,\n                )\n\n    async def _process_single_sample_streaming(self, rollout_sample: RolloutSample):\n        \"\"\"Process a single sample streamingly\"\"\"\n        # Calling asynchronous generation methods\n        ret = await self.async_rollout_manager.generate_sequences_single(rollout_sample.full_batch)\n        rollout_sample.full_batch = ret\n        rollout_sample.full_batch.non_tensor_batch[\"uid\"] = np.array(\n            [f\"uid_{rollout_sample.sample_id}\"] * len(rollout_sample.full_batch), dtype=object\n        )\n        rollout_sample.rollout_status = await self.get_statistics()\n\n        success = await self.message_queue_client.put_sample(\n            sample=ray.cloudpickle.dumps(rollout_sample),\n        )\n        if success:\n            self.total_generated_samples += 1\n        else:\n            self.dropped_stale_samples += 1\n        self.processed_sample_count += 1\n\n    async def _streaming_generation_main(self):\n        \"\"\"The main entry method for stream processing\"\"\"\n\n        if self.async_rollout_manager is None:\n            await self._init_async_rollout_manager()\n\n        # Start the streaming loop\n        print(f\"[FullyAsyncRollouter] Start streaming mode, maximum concurrent samples: {self.max_concurrent_samples}\")\n\n        # Start sample feed coroutine, streaming process coroutine\n        self.feed_task = safe_create_task(self._feed_samples(), name=\"feed_task\")\n        self.processor_task = safe_create_task(self._processor_worker(), name=\"processor_task\")\n\n        try:\n            # Wait for sample feed to complete\n            # Use asyncio.wait to monitor all tasks. If processor exits early,\n            # detect it instead of blocking on feed_task (it might be stuck on a full queue).\n            done, pending = await asyncio.wait(\n                [self.feed_task, self.processor_task], return_when=asyncio.FIRST_COMPLETED\n            )\n\n            for task in done:\n                if task.exception():\n                    raise task.exception()\n\n            if self.feed_task not in done:\n                raise RuntimeError(\"Processor task exited prematurely\")\n\n            print(\"[FullyAsyncRollouter] Sample feed completed\")\n\n            # Wait for streaming to complete\n            await self.processor_task\n            print(\"[FullyAsyncRollouter] Streaming process completed\")\n\n            await self.pending_queue.join()\n            print(\"[FullyAsyncRollouter] pending_queue joined\")\n\n        except Exception as e:\n            print(f\"[FullyAsyncRollouter] Streaming process exception: {e}\")\n            raise e\n\n        finally:\n            if self.feed_task and not self.feed_task.done():\n                self.feed_task.cancel()\n                await asyncio.gather(self.feed_task, return_exceptions=True)\n\n            if self.processor_task and not self.processor_task.done():\n                self.processor_task.cancel()\n                await asyncio.gather(self.processor_task, return_exceptions=True)\n\n            self.feed_task = None\n            self.processor_task = None\n\n            # Send a finish signal\n            await self.message_queue_client.put_sample(sample=None)\n\n        async with self.lock:\n            self.running = False\n\n    async def fit(self):\n        \"\"\"\n        Start the async rollouter - entry point that sets up and runs async tasks\n        Main async fit method that coordinates all coroutines\n        \"\"\"\n\n        print(\"[FullyAsyncRollouter] Starting FullyAsyncRollouter...\")\n\n        if self.message_queue_client is None:\n            raise ValueError(\"MessageQueue client not set. Call set_message_queue_client() first.\")\n\n        # Set the running status flag\n        async with self.lock:\n            self.paused = False\n            self.running = True\n\n        # Create the main asynchronous task\n        generation_task = safe_create_task(self._streaming_generation_main(), name=\"generation_task\")\n        monitor_task = safe_create_task(self._async_monitor_loop(), name=\"monitor_task\")\n\n        try:\n            # Run build and monitoring tasks concurrently\n            await asyncio.gather(generation_task, monitor_task, return_exceptions=True)\n        except Exception as e:\n            print(f\"[FullyAsyncRollouter] Asynchronous task execution error: {e}\")\n        finally:\n            if not generation_task.done():\n                generation_task.cancel()\n            if not monitor_task.done():\n                monitor_task.cancel()\n\n            # Wait for the task to complete\n            await asyncio.gather(generation_task, monitor_task, return_exceptions=True)\n\n        print(\"[FullyAsyncRollouter] Rollouter fit completed\")\n\n    async def _async_monitor_loop(self):\n        \"\"\"\n        Async coroutine for monitoring:\n        Function 1: Log information output\n        Function 2: Trigger rollout recovery\n        \"\"\"\n        last_stats_time = time.time()\n        stats_interval = 60.0\n        check_interval = 10.0\n\n        while True:\n            async with self.lock:\n                if not self.running:\n                    break\n            await asyncio.sleep(check_interval)\n            # Print statistics periodically\n            current_time = time.time()\n            if current_time - last_stats_time >= stats_interval:\n                stats = await self.get_statistics()\n                print(f\"[FullyAsyncRollouter][MonitorLoop][Statistics] {pformat(stats)}\")\n                last_stats_time = current_time\n\n            # Trigger rollout recovery\n            if self.paused and not await self._should_pause_generation():\n                async with self.lock:\n                    self.paused = False\n                    print(\"[FullyAsyncRollouter][ShouldPause] notify all wait tasks.\")\n                    self.condition.notify_all()\n\n    async def _should_pause_generation(self) -> bool:\n        \"\"\"Determine whether the build should be paused\"\"\"\n        queue_stats = self.message_queue_client.get_statistics_sync()\n        queue_size = queue_stats[\"queue_size\"]\n\n        if queue_size >= self.max_queue_size:\n            if not self.paused:\n                print(\n                    f\"[FullyAsyncRollouter][ShouldPause]  \"\n                    f\"due to full queue: size={queue_size}, max={self.max_queue_size}\"\n                )\n            return True\n\n        if self.staleness_samples >= self.max_required_samples:\n            if not self.paused:\n                print(\n                    \"[FullyAsyncRollouter][ShouldPause] \"\n                    f\"due to \"\n                    f\"staleness_samples {self.staleness_samples} >= max_required_samples {self.max_required_samples} \"\n                )\n            return True\n\n        return False\n\n    async def get_statistics(self) -> dict:\n        queue_stats = self.message_queue_client.get_statistics_sync()\n\n        stats = {\n            # monitor stats\n            \"monitor/active_tasks_size\": len(self.active_tasks),\n            \"monitor/queue/pending_queue_size\": self.pending_queue.qsize(),\n            \"monitor/queue/mq_queue_size\": queue_stats[\"queue_size\"],\n            # counting stats\n            \"count/total_generated_samples\": self.total_generated_samples,\n            \"count/staleness_samples\": self.staleness_samples,\n            \"count/dropped_stale_samples\": self.dropped_stale_samples,\n            # static stats\n            \"static/max_required_samples\": self.max_required_samples,\n            \"static/required_samples\": self.required_samples,\n            \"static/staleness_threshold\": self.staleness_threshold,\n            \"static/max_queue_size\": self.max_queue_size,\n            \"static/max_concurrent_samples\": self.max_concurrent_samples,\n        }\n\n        return stats\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/fully_async_trainer.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nimport time\nfrom datetime import datetime\nfrom pprint import pprint\nfrom typing import Any\n\nimport ray\nfrom omegaconf import OmegaConf, open_dict\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.fully_async_policy.detach_utils import (\n    MetricsAggregator,\n    ValidateMetrics,\n    assemble_batch_from_rollout_samples,\n)\nfrom verl.experimental.fully_async_policy.message_queue import MessageQueueClient\nfrom verl.experimental.separation.ray_trainer import SeparateRayPPOTrainer\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\nfrom verl.trainer.ppo import core_algos\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model\nfrom verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.tracking import Tracking\n\nlogger = logging.getLogger(__name__)\n\n\nclass TrainingStopException(Exception):\n    \"\"\"Exception raised to signal training should stop\"\"\"\n\n    pass\n\n\n@ray.remote(num_cpus=10)\nclass FullyAsyncTrainer(SeparateRayPPOTrainer):\n    \"\"\"\n    A fully asynchronous PPO trainer that obtains samples from a MessageQueue for training.\n    Based on an improved implementation of OneStepOffRayTrainer\n    \"\"\"\n\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        role_worker_mapping: dict[Role, WorkerType],\n        resource_pool_manager: ResourcePoolManager,\n        ray_worker_group_cls: RayWorkerGroup = RayWorkerGroup,\n        processor=None,\n        device_name=None,\n    ):\n        # ==================== RayPPOTrainer config ====================\n\n        # Store the tokenizer for text processing\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.config = config\n\n        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine\n        assert not self.hybrid_engine\n\n        self.role_worker_mapping = role_worker_mapping\n        self.resource_pool_manager = resource_pool_manager\n        self.use_reference_policy = need_reference_policy(self.config)\n\n        self.use_rm = need_reward_model(self.config)\n\n        self.use_critic = need_critic(self.config)\n        self.ray_worker_group_cls = ray_worker_group_cls\n        self.device_name = device_name if device_name else self.config.trainer.device\n\n        # if ref_in_actor is True, the reference policy will be actor without lora applied\n        lora_rank = config.actor_rollout_ref.model.get(\"lora\", {}).get(\"rank\", 0)\n        if lora_rank <= 0:\n            lora_rank = config.actor_rollout_ref.model.get(\"lora_rank\", 0)\n        self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get(\"lora_adapter_path\") is not None\n\n        # define in-reward KL control\n        # kl loss control currently not suppoorted\n        if self.config.algorithm.use_kl_in_reward:\n            self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)\n\n        self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get(\"use_prefix_grouper\", False)\n        self.use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n\n        # ==================== SeparateRayPPOTrainer config ====================\n        self.global_steps = 0\n        self.epoch = 0\n        self.max_steps_duration = 0\n        self.progress_bar = None\n        self.is_last_step = False\n        self.prev_step_profile = False\n        self.curr_step_profile = False\n        self.next_step_profile = False\n        self.last_val_metrics = {}\n        self.metrics = {}\n        self.timing_raw = {}\n        # reward message\n        self.future_reward = None\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n\n        self.logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        # ==================== fully async config ====================\n\n        self.message_queue_client = None\n\n        # Statistics\n        self.local_trigger_step = 1\n        self.processed_samples = 0\n        self.stale_trajectory_processed = 0\n        self.current_param_version = 0\n        self.total_train_steps = None\n        self.progress_bar = None\n        self.trigger_parameter_sync_step = config.async_training.trigger_parameter_sync_step\n        self.last_ckpt_version = 0\n        self.train_role = Role.ActorRollout if config.async_training.use_trainer_do_validate else Role.Actor\n\n        # required_samples use ppo_mini_batch_size*require_batches as the minimum number of samples.\n        self.require_batches = config.async_training.require_batches\n        self.required_samples = config.actor_rollout_ref.actor.ppo_mini_batch_size * self.require_batches\n        total_gpus = (\n            config.trainer.nnodes * config.trainer.n_gpus_per_node\n            + config.rollout.nnodes * config.rollout.n_gpus_per_node\n        )\n        self.metrics_aggregator = MetricsAggregator(total_gpus=total_gpus)\n\n        # use trainer to do validation\n        if self.config.async_training.use_trainer_do_validate:\n            from verl.trainer.main_ppo import create_rl_dataset\n            from verl.utils.dataset.rl_dataset import collate_fn\n\n            val_dataset = create_rl_dataset(config.data.val_files, config.data, tokenizer, processor)\n            rollout_gpus = config.rollout.nnodes * config.rollout.n_gpus_per_node\n            print(f\"[FullyAsyncTrainer] split before val_dataset total len: {len(val_dataset)}\")\n            split_dataset = val_dataset.split(total_gpus)\n            rollout_val_dataset0 = split_dataset[rollout_gpus:]\n            from torch.utils.data import ConcatDataset\n\n            val_dataset = ConcatDataset(rollout_val_dataset0)\n            print(f\"[FullyAsyncTrainer] split after val_dataset total len: {len(val_dataset)}\")\n            self.val_dataset = val_dataset\n            # update val_dataloader\n            val_batch_size = self.config.data.val_batch_size  # Prefer config value if set\n            if val_batch_size is None:\n                val_batch_size = len(val_dataset)\n            from torchdata.stateful_dataloader import StatefulDataLoader\n\n            print(f\"[FullyAsyncTrainer] create val_dataloader with batch_size: {val_batch_size}\")\n            self.val_dataloader = StatefulDataLoader(\n                dataset=val_dataset,\n                batch_size=val_batch_size,\n                num_workers=self.config.data[\"dataloader_num_workers\"],\n                shuffle=self.config.data.get(\"validation_shuffle\", True),\n                drop_last=False,\n                collate_fn=collate_fn,\n            )\n        # Reference to rollouter for parameter synchronization\n        self.rollouter = None\n        self.checkpoint_manager = None\n\n        # when use_trainer_do_validate == Ture, use colocate_checkpoint_manager to sync params\n        self.colocate_checkpoint_manager = None\n\n    def _setup_checkpoint_manager(self, rollouter):\n        \"\"\"Setup checkpoint manager after rollouter is initialized\"\"\"\n        replicas = ray.get(rollouter.get_replicas.remote())\n        checkpoint_engine_config = omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine)\n        self.checkpoint_manager = CheckpointEngineManager(\n            config=checkpoint_engine_config, trainer=self.actor_wg, replicas=replicas\n        )\n        print(\"[FullyAsyncTrainer] Checkpoint manager initialized\")\n\n    def set_message_queue_client(self, message_queue_client: MessageQueueClient):\n        \"\"\"Set message queue client\"\"\"\n        self.message_queue_client = message_queue_client\n\n    def set_rollouter(self, rollouter):\n        \"\"\"Set rollouter reference for parameter synchronization\"\"\"\n        self.rollouter = rollouter\n        # Setup checkpoint manager after rollouter is set\n        self._setup_checkpoint_manager(rollouter)\n\n    def set_total_train_steps(self, total_training_steps):\n        self.total_train_steps = total_training_steps\n\n        try:\n            OmegaConf.set_struct(self.config, True)\n            with open_dict(self.config):\n                if OmegaConf.select(self.config, \"actor_rollout_ref.actor.optim\"):\n                    self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps\n                if OmegaConf.select(self.config, \"critic.optim\"):\n                    self.config.critic.optim.total_training_steps = total_training_steps\n        except Exception as e:\n            print(f\"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}\")\n\n        self.progress_bar = tqdm(total=self.total_train_steps, initial=0, desc=\"Training Progress\")\n\n    def get_actor_wg(self):\n        \"\"\"Get actor worker group\"\"\"\n        return self.actor_wg\n\n    async def _get_samples_from_queue(self) -> tuple[None, None] | tuple[int, Any]:\n        \"\"\"\n        Get samples from message queue and compose gen_batch_output\n        Uses a loop to continuously collect samples until enough are gathered\n\n        Returns:\n            tuple: (epoch, batch_dict, gen_batch_output)\n        \"\"\"\n        print(\n            f\"[FullyAsyncTrainer] Requesting {self.required_samples} samples from queue\",\n            flush=True,\n        )\n\n        # Collect samples using a simple loop calling get_sample\n        consumer_start = time.time()\n        queue_samples = []\n        queue_len = 0\n        while len(queue_samples) < self.required_samples:\n            # Get a single sample and wait until there is a sample or None is received\n            sample, queue_len = self.message_queue_client.get_sample_sync()\n\n            if sample is None:\n                print(\n                    f\"[FullyAsyncTrainer] Detected termination signal (None), stopping sample collection. \"\n                    f\"Collected {len(queue_samples)}/{self.required_samples} samples\"\n                )\n                break\n\n            queue_samples.append(sample)\n\n            if len(queue_samples) % 64 == 0:\n                print(\n                    f\"[FullyAsyncTrainer] Collected {len(queue_samples)}/{self.required_samples} samples. \"\n                    f\"mq_len: {queue_len}\"\n                )\n\n        consumer_end = time.time()\n\n        if not queue_samples or len(queue_samples) < self.required_samples:\n            print(\"[FullyAsyncTrainer] not enough samples collected after loop\")\n            return None, None\n        total_wait_time = consumer_end - consumer_start\n\n        print(\n            f\"[FullyAsyncTrainer] Loop collection completed: {len(queue_samples)}/{self.required_samples} samples, \"\n            f\"total wait time: {total_wait_time:.2f} seconds. \"\n            f\"mq_len: {queue_len}\"\n        )\n\n        queue_samples = [ray.cloudpickle.loads(x) for x in queue_samples]\n        # Assemble batch - now working directly with RolloutSample objects\n        if self.config.trainer.balance_batch:\n            batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, self._balance_batch)\n        else:\n            batch = assemble_batch_from_rollout_samples(queue_samples, self.tokenizer, self.config, None)\n\n        batch.meta_info[\"fully_async/total_wait_time\"] = total_wait_time\n        return 0, batch\n\n    def _create_actor_rollout_classes(self):\n        # create actor\n        for role in [self.train_role]:\n            resource_pool = self.resource_pool_manager.get_resource_pool(role)\n            role_cls = RayClassWithInitArgs(\n                cls=self.role_worker_mapping[role],\n                config=self.config.actor_rollout_ref,\n                role=str(role),\n            )\n            self.resource_pool_to_cls[resource_pool][str(role)] = role_cls\n\n    def _init_models(self):\n        if self.use_critic:\n            self.critic_wg = self.all_wg[str(Role.Critic)]\n            self.critic_wg.init_model()\n\n        if self.use_reference_policy and not self.ref_in_actor:\n            self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]\n            self.ref_policy_wg.init_model()\n\n        if self.use_rm:\n            self.rm_wg = self.all_wg[str(Role.RewardModel)]\n            self.rm_wg.init_model()\n\n        self.actor_wg = self.all_wg[str(self.train_role)]\n        self.actor_wg.init_model()\n        self.actor_rollout_wg = self.actor_wg  # to be compatible with the functions that not be modified\n\n    async def init_workers(self):\n        \"\"\"Initialize distributed training workers using Ray backend.\n        Creates:\n        1. Ray resource pools from configuration\n        2. Worker groups for each role (actor, critic, etc.)\n        \"\"\"\n        self._init_resource_pools()\n        self._create_worker_classes()\n        self._init_worker_groups()\n        self._init_models()\n        self._init_reward_loop()\n        await self._init_async_rollout_manager()\n\n    def _init_reward_loop(self):\n        if self.config.async_training.use_trainer_do_validate:\n            print(\"[FullyAsyncTrainer] Init reward loop\")\n            super()._init_reward_loop()\n\n    async def _init_async_rollout_manager(self):\n        # use async rollout do validate\n        print(f\"[FullyAsyncTrainer] use_trainer_do_validate: {self.config.async_training.use_trainer_do_validate}\")\n        if self.config.async_training.use_trainer_do_validate:\n            print(\"[FullyAsyncTrainer] Init async rollout manager\")\n\n            # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design\n            # agent_reward_loop: streaming reward computation with actor rollout\n            # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool\n            enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool\n\n            # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager\n            # to stream reward computation with actor rollout\n            reward_loop_worker_handles = (\n                self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None\n            )\n\n            # create async rollout manager and request scheduler\n            assert self.config.actor_rollout_ref.rollout.mode == \"async\"\n\n            self.async_rollout_mode = True\n            from verl.experimental.agent_loop import AgentLoopManager\n\n            self.async_rollout_manager = await AgentLoopManager.create(\n                config=self.config,\n                worker_group=self.actor_rollout_wg,\n                reward_loop_worker_handles=reward_loop_worker_handles,\n            )\n            print(\"[FullyAsyncTrainer] async_rollout_manager initialized\")\n\n            # Modify checkpoint_engine config to use naive backend\n            checkpoint_engine_cfg = self.config.actor_rollout_ref.rollout.checkpoint_engine\n            original_backend = checkpoint_engine_cfg.backend\n            with open_dict(checkpoint_engine_cfg):\n                checkpoint_engine_cfg.backend = \"naive\"\n            checkpoint_engine_config = omega_conf_to_dataclass(checkpoint_engine_cfg)\n\n            print(f\"[FullyAsyncTrainer] checkpoint_engine_config: {checkpoint_engine_config}\")\n\n            self.colocate_checkpoint_manager = CheckpointEngineManager(\n                config=checkpoint_engine_config,\n                trainer=self.actor_rollout_wg,\n                replicas=self.async_rollout_manager.rollout_replicas,\n            )\n\n            # sleep all replicas to load checkpoint\n            await self.colocate_checkpoint_manager.sleep_replicas()\n\n            # Restore original backend value\n            with open_dict(checkpoint_engine_cfg):\n                checkpoint_engine_cfg.backend = original_backend\n\n            print(\"[FullyAsyncTrainer] colocate_checkpoint_manager initialized\")\n\n        else:\n            print(\"[FullyAsyncTrainer] Skip async rollout manager (use_trainer_do_validate=False)\")\n\n    async def fit(self):\n        \"\"\"\n        The training loop of PPO.\n        The driver process only need to call the compute functions of the worker group through RPC\n        to construct the PPO dataflow.\n        The light-weight advantage computation is done on the driver process.\n        \"\"\"\n        print(\"[FullyAsyncTrainer] Starting FullyAsyncTrainer...\")\n        if self.message_queue_client is None:\n            raise ValueError(\"MessageQueue client not set. Call set_message_queue_client() first.\")\n        if self.rollouter is None:\n            raise ValueError(\"rollouter not set. Call set_rollouter() first.\")\n\n        self.max_steps_duration = 0\n\n        self.global_steps += 1\n\n        # Use queue mode, no need for traditional dataloader iterator\n        # Initialize to get the first batch of data\n        while True:\n            try:\n                await self.fit_step()\n            except TrainingStopException:\n                print(\"[FullyAsyncTrainer] Training stopped by queue termination signal\")\n                break\n\n        self.progress_bar.close()\n        if self.current_param_version % self.config.trainer.test_freq != 0 or self.local_trigger_step > 1:\n            await self._fit_update_weights()\n            await self._fit_validate()\n        self._fit_save_checkpoint(force=True)\n\n    async def fit_step(self, batch_dict: dict = None):\n        \"\"\"\n        Single-step training template method. Handles all logic for one training step.\n\n        Flow:\n        1. Pre-step processing -> 2. Get batch -> 3. Generate sequences ->\n        4. Compute reward -> 5. Compute log_prob -> 6. Compute reward ->\n        7. Compute advantage -> 8. Update critic -> 9. Update actor -> 10. Post-step processing\n\n        Args:\n            batch_dict: Raw data dictionary\n        \"\"\"\n        self.metrics = {\"training/global_step\": self.global_steps, \"training/epoch\": self.epoch}\n        self.timing_raw = {}\n        # reward message\n        self.future_reward = None\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n\n        self._fit_start_profile()\n\n        with marked_timer(\"step\", self.timing_raw):\n            batch = await self._fit_generate(None)\n            batch = self._fit_compute_reward(batch)\n            batch = self._fit_compute_log_prob(batch)\n            batch = self._fit_compute_ref_log_prob(batch)\n            batch = self._fit_compute_critic(batch)\n            batch = self._fit_compute_advantage(batch)\n            batch = self._fit_update_critic(batch)\n            batch = self._fit_update_actor(batch)\n            self._fit_update_local_step()\n            await self._fit_update_weights()\n            self._fit_dump_data(batch)\n\n        await self._fit_validate()\n        self._fit_save_checkpoint()\n        self._fit_stop_profile()\n        self._fit_collect_metrics(batch)\n        self._fit_torch_memory()\n        self._fit_postprocess_step()\n\n    async def _fit_generate(self, batch: DataProto = None) -> DataProto | None:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        with marked_timer(\"gen\", timing_raw, color=\"red\"):\n            epoch, batch = await self._get_samples_from_queue()\n            if batch is None:\n                raise TrainingStopException(\"Training terminated: queue returned None\")\n            self._collect_metrics_from_samples(batch, metrics)\n        batch.meta_info[\"temperature\"] = self.config.actor_rollout_ref.rollout.temperature\n        return batch\n\n    def _compute_old_log_prob(self, batch: DataProto):\n        \"\"\"\n        If algorithm.rollout_correction.bypass_mode is False,\n        use model engine and first version model params to re-calculate old_log_prob.\n\n        If local_trigger_step == 1, load the training engine's parameters to the CPU\n          and save a copy for subsequent MIS use.\n\n        If local_trigger_step == 2, 3, ..., restore the parameters of version 1 to calculate the old_log_prob,\n        then restore the parameters of the current version.\n        \"\"\"\n        if self.local_trigger_step == 1:\n            self.actor_rollout_wg.save_model_to_cpu(1)\n            old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch)\n        else:\n            self.actor_rollout_wg.save_model_to_cpu(self.local_trigger_step)\n            self.actor_rollout_wg.restore_model_from_cpu(1)\n            old_log_prob, old_log_prob_mfu = super()._compute_old_log_prob(batch)\n            self.actor_rollout_wg.restore_model_from_cpu(self.local_trigger_step)\n            self.actor_rollout_wg.clear_cpu_model(self.local_trigger_step)\n        return old_log_prob, old_log_prob_mfu\n\n    def _fit_update_local_step(self):\n        time_str = datetime.now().strftime(\"%H:%M:%S.%f\")[:-3]\n        print(\n            f\"[FullyAsyncTrainer] global_steps: {self.global_steps} \"\n            f\"local_trigger_step: {self.local_trigger_step} \"\n            f\"trigger_parameter_sync_step: {self.trigger_parameter_sync_step} \"\n            f\"{time_str}\"\n        )\n        if self.local_trigger_step < self.trigger_parameter_sync_step:\n            self.local_trigger_step += 1\n        else:\n            self.current_param_version += 1\n            self.local_trigger_step = 1\n\n    async def _fit_update_weights(self):\n        if self.local_trigger_step != 1:\n            return\n\n        with marked_timer(\"timing_s/param_sync\", self.timing_raw):\n            await self.checkpoint_manager.update_weights(global_steps=self.current_param_version)\n        print(\n            f\"[FullyAsyncTrainer] _fit_update_weights, \"\n            f\"timing_s/param_sync: {self.timing_raw['timing_s/param_sync']:.4f} seconds \"\n            f\"self.current_param_version: {self.current_param_version}\"\n        )\n\n        # Reset staleness in rollouter\n        timing_raw = ray.get(self.rollouter.reset_staleness.remote())\n        self.logger.log(\n            data=timing_raw,\n            step=self.current_param_version,\n        )\n\n        # Log aggregated training metrics\n        self.logger.log(\n            data=self.metrics_aggregator.get_aggregated_metrics(),\n            step=self.current_param_version,\n        )\n        self.metrics_aggregator.reset()\n\n    async def _validate_process(self):\n        \"\"\"Run trainer-side validation using async rollout manager\"\"\"\n        if self.config.async_training.use_trainer_do_validate:\n            print(\"[FullyAsyncTrainer] _validate_process\")\n            from verl.utils.profiler import marked_timer\n\n            # Wake up rollouter replicas and sync weights\n            print(\"[FullyAsyncTrainer] wake up replicas before validation\")\n            await self.colocate_checkpoint_manager.update_weights(global_steps=self.current_param_version)\n\n            with marked_timer(\"trainer/validate_time\", self.timing_raw):\n                train_val_metrics = self._validate(True)\n\n            # Sleep rollouter replicas to free GPU memory for validation\n            print(\"[FullyAsyncTrainer] sleep replicas after validation\")\n            await self.colocate_checkpoint_manager.sleep_replicas()\n\n            print(f\"[FullyAsyncTrainer] validate timing: {self.timing_raw['trainer/validate_time']}\")\n            return train_val_metrics\n        else:\n            print(\"[FullyAsyncTrainer] _validate_process without async_rollout_manager\")\n            return None\n\n    async def _fit_validate(self, val_before_train=False):\n        if self.local_trigger_step != 1:\n            return\n\n        # Check if validation is needed\n        need_validate = (\n            self.config.trainer.test_freq > 0\n            and self.current_param_version % self.config.trainer.test_freq == 0\n            and self.current_param_version > 0\n        )\n        # Skip validation if not needed and not validation before training\n        if not need_validate and not val_before_train:\n            return\n\n        # Trigger rollouter validation and get future\n        val_future = self.rollouter.do_validate.remote()\n\n        # Run trainer-side validation\n        train_val_metrics = await self._validate_process()\n\n        # Wait for rollouter validation result and log\n        val_metrics: ValidateMetrics = ray.get(val_future)\n        if train_val_metrics:\n            # Merge trainer and rollouter validation results\n            with marked_timer(\"timing_s/merge_val\", self.timing_raw):\n                new_metrics = self._merge_validation_results(train_val_metrics, val_metrics.metrics)\n            if new_metrics:\n                self.logger.log(data=new_metrics, step=self.current_param_version)\n                pprint(\n                    f\"[FullyAsyncTrainer] parameter version: {self.current_param_version} \"\n                    f\"Validation metrics: {new_metrics}, timing: {self.timing_raw['timing_s/merge_val']}\"\n                )\n        else:\n            if val_metrics.metrics:\n                self.logger.log(data=val_metrics.metrics, step=self.current_param_version)\n                pprint(\n                    f\"[FullyAsyncTrainer] parameter version: {self.current_param_version} \"\n                    f\"Validation metrics: {val_metrics.metrics}\"\n                )\n        self.logger.log(data=val_metrics.timing_raw, step=self.current_param_version)\n\n    def _fit_save_checkpoint(self, force=False):\n        if self.current_param_version == self.last_ckpt_version:\n            return\n\n        timing_raw = self.timing_raw\n        # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.\n        esi_close_to_expiration = should_save_ckpt_esi(\n            max_steps_duration=self.max_steps_duration,\n            redundant_time=self.config.trainer.esi_redundant_time,\n        )\n        # Check if the conditions for saving a checkpoint are met.\n        # The conditions include a mandatory condition (1) and\n        # one of the following optional conditions (2/3/4):\n        # 1. The save frequency is set to a positive value.\n        # 2. It's the last training step.\n        # 3. The current step number is a multiple of the save frequency.\n        # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.\n        if self.config.trainer.save_freq > 0 and (\n            force or self.current_param_version % self.config.trainer.save_freq == 0 or esi_close_to_expiration\n        ):\n            if esi_close_to_expiration:\n                print(\"Force saving checkpoint: ESI instance expiration approaching.\")\n            with marked_timer(\"save_checkpoint\", timing_raw, color=\"green\"):\n                # sleep replicas to avoid OOM during checkpoint saving\n                self._save_checkpoint()\n                self.last_ckpt_version = self.current_param_version\n\n    def _fit_postprocess_step(self):\n        self.global_steps += 1\n\n        self.metrics_aggregator.add_step_metrics(\n            metrics=self.metrics, sample_count=self.required_samples, timestamp=time.time()\n        )\n\n        if self.local_trigger_step == 1:\n            self.progress_bar.update(1)\n\n    def _save_checkpoint(self):\n        # Warning: Currently, to align the training process and metrics of colocate,\n        # we use current_param_version instead of global step.\n        # This can be logically aligned with the original self.global_steps of colocate\n        # and is used for metrics and ckpt. which means that the parameter synchronization\n        # from trainer to rollouter will increase by 1 each time.\n\n        # path: given_path + `/global_step_{global_steps}` + `/actor`\n        local_global_step_folder = os.path.join(\n            self.config.trainer.default_local_dir, f\"global_step_{self.current_param_version}\"\n        )\n\n        print(f\"[FullyAsyncTrainer] local_global_step_folder: {local_global_step_folder}\")\n        actor_local_path = os.path.join(local_global_step_folder, \"actor\")\n\n        actor_remote_path = (\n            None\n            if self.config.trainer.default_hdfs_dir is None\n            else os.path.join(\n                self.config.trainer.default_hdfs_dir, f\"global_step_{self.current_param_version}\", \"actor\"\n            )\n        )\n\n        remove_previous_ckpt_in_save = self.config.trainer.get(\"remove_previous_ckpt_in_save\", False)\n        if remove_previous_ckpt_in_save:\n            print(\n                \"[FullyAsyncTrainer] Warning: remove_previous_ckpt_in_save is deprecated,\"\n                + \" set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead\"\n            )\n        max_actor_ckpt_to_keep = (\n            self.config.trainer.get(\"max_actor_ckpt_to_keep\", None) if not remove_previous_ckpt_in_save else 1\n        )\n        max_critic_ckpt_to_keep = (\n            self.config.trainer.get(\"max_critic_ckpt_to_keep\", None) if not remove_previous_ckpt_in_save else 1\n        )\n\n        self.actor_rollout_wg.save_checkpoint(\n            actor_local_path, actor_remote_path, self.current_param_version, max_ckpt_to_keep=max_actor_ckpt_to_keep\n        )\n\n        if self.use_critic:\n            critic_local_path = os.path.join(local_global_step_folder, str(Role.Critic))\n            critic_remote_path = (\n                None\n                if self.config.trainer.default_hdfs_dir is None\n                else os.path.join(\n                    self.config.trainer.default_hdfs_dir, f\"global_step_{self.current_param_version}\", str(Role.Critic)\n                )\n            )\n            self.critic_wg.save_checkpoint(\n                critic_local_path,\n                critic_remote_path,\n                self.current_param_version,\n                max_ckpt_to_keep=max_critic_ckpt_to_keep,\n            )\n        ray.get(self.rollouter.save_checkpoint.remote(local_global_step_folder))\n        # latest checkpointed iteration tracker (for atomic usage)\n        local_latest_checkpointed_iteration = os.path.join(\n            self.config.trainer.default_local_dir, \"latest_checkpointed_iteration.txt\"\n        )\n        with open(local_latest_checkpointed_iteration, \"w\") as f:\n            f.write(str(self.current_param_version))\n\n    async def load_checkpoint(self):\n        if self.config.trainer.resume_mode == \"disable\":\n            return 0\n\n        # load from hdfs\n        if self.config.trainer.default_hdfs_dir is not None:\n            raise NotImplementedError(\"load from hdfs is not implemented yet\")\n        else:\n            checkpoint_folder = self.config.trainer.default_local_dir  # TODO: check path\n            if not os.path.isabs(checkpoint_folder):\n                working_dir = os.getcwd()\n                checkpoint_folder = os.path.join(working_dir, checkpoint_folder)\n            global_step_folder = find_latest_ckpt_path(checkpoint_folder)  # None if no latest\n\n        # find global_step_folder\n        if self.config.trainer.resume_mode == \"auto\":\n            if global_step_folder is None:\n                return 0\n        else:\n            if self.config.trainer.resume_mode == \"resume_path\":\n                assert isinstance(self.config.trainer.resume_from_path, str), \"resume ckpt must be str type\"\n                assert \"global_step_\" in self.config.trainer.resume_from_path, (\n                    \"resume ckpt must specify the global_steps\"\n                )\n                global_step_folder = self.config.trainer.resume_from_path\n                if not os.path.isabs(global_step_folder):\n                    working_dir = os.getcwd()\n                    global_step_folder = os.path.join(working_dir, global_step_folder)\n        print(f\"[FullyAsyncTrainer] Load from checkpoint folder: {global_step_folder}\")\n        # set global step\n        self.current_param_version = int(global_step_folder.split(\"global_step_\")[-1])\n        self.global_steps = self.current_param_version * self.trigger_parameter_sync_step + 1\n        self.last_ckpt_version = self.current_param_version\n        print(\n            f\"[FullyAsyncTrainer] Setting global step to {self.global_steps}, \"\n            f\"current_param_version to {self.current_param_version}\"\n        )\n        print(f\"[FullyAsyncTrainer] Resuming from  {global_step_folder}\")\n\n        actor_path = os.path.join(global_step_folder, \"actor\")\n        critic_path = os.path.join(global_step_folder, str(Role.Critic))\n        # load actor\n        self.actor_rollout_wg.load_checkpoint(\n            actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load\n        )\n        # load critic\n        if self.use_critic:\n            self.critic_wg.load_checkpoint(\n                critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load\n            )\n\n        if self.colocate_checkpoint_manager:\n            await self.colocate_checkpoint_manager.update_weights(self.current_param_version)\n            await self.colocate_checkpoint_manager.sleep_replicas()\n\n        return self.current_param_version\n\n    def _collect_metrics_from_samples(self, batch, metrics):\n        \"\"\"\n        Collect metrics from samples\n        \"\"\"\n        if hasattr(batch, \"meta_info\") and batch.meta_info:\n            trajectory_param_versions = batch.meta_info[\"trajectory_param_versions\"]\n            stale_traj_count = sum(1 for v in trajectory_param_versions if self.current_param_version - v >= 1)\n            self.stale_trajectory_processed += stale_traj_count\n            metrics.update(\n                {\n                    \"fully_async/count/stale_trajectory_processed\": self.stale_trajectory_processed,\n                    \"fully_async/count/current_param_version\": self.current_param_version,\n                }\n            )\n            for key, value in batch.meta_info.items():\n                if key.startswith(\"fully_async\") or key.startswith(\"timing_s\"):\n                    metrics[key] = value\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/message_queue.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nfrom collections import deque\nfrom typing import Any\n\nimport ray\nfrom omegaconf import DictConfig\n\nlogger = logging.getLogger(__name__)\n\n\n@ray.remote(num_cpus=2, max_concurrency=20)\nclass MessageQueue:\n    \"\"\"\n    Simplified Ray-based asynchronous message queue for communication between Rollouter and Trainer\n    \"\"\"\n\n    def __init__(self, config: DictConfig, max_queue_size: int = 1000):\n        self.config = config\n        if max_queue_size is None:\n            raise ValueError(f\"max_queue_size cannot be None, got: {max_queue_size}\")\n        self.max_queue_size = int(max_queue_size)\n        self.queue = deque(maxlen=self.max_queue_size)\n\n        self.val_queue = deque()\n\n        # Asyncio for message handling\n        self.running = True\n\n        # async safe\n        self._lock = asyncio.Lock()\n        self._consumer_condition = asyncio.Condition(self._lock)\n\n        # statistic message\n        self.total_produced = 0\n        self.total_consumed = 0\n        self.dropped_samples = 0\n\n        print(f\"[MessageQueue] initialized with max_queue_size={max_queue_size}\")\n\n    async def put_sample(self, sample: Any) -> bool:\n        \"\"\"\n        Put a batch sample into the queue\n\n        Args:\n            sample: Sample data\n\n        Returns:\n            bool: Whether the sample was successfully put into the queue\n        \"\"\"\n        async with self._lock:\n            # If queue is full, remove the oldest sample (rarely happens)\n            is_drop = False\n            if len(self.queue) >= self.max_queue_size:\n                self.queue.popleft()\n                self.dropped_samples += 1\n                is_drop = True\n                logger.warning(\"Queue full, dropped sample\")\n            self.queue.append(sample)\n            self.total_produced += 1\n\n            # Notify waiting consumers\n            self._consumer_condition.notify_all()\n\n            if self.total_produced % 100 == 0:\n                print(f\"MessageQueue stats: produced={self.total_produced}, queue_size={len(self.queue)}\")\n            if is_drop:\n                return False\n            return True\n\n    async def get_sample(self) -> Any | None:\n        \"\"\"\n        Get a single sample from the queue, wait until one is available\n\n        Returns:\n            Any: Single sample data or None if queue is closed\n        \"\"\"\n        async with self._lock:\n            while len(self.queue) == 0 and self.running:\n                await self._consumer_condition.wait()\n\n            # If queue is closed and empty, return None\n            if not self.running and len(self.queue) == 0:\n                return None\n\n            # Get one sample\n            data = self.queue.popleft()\n            self.total_consumed += 1\n            return data, len(self.queue)\n\n    async def get_queue_size(self) -> int:\n        \"\"\"Get current queue length\"\"\"\n        async with self._lock:\n            return len(self.queue)\n\n    async def get_statistics(self) -> dict[str, Any]:\n        \"\"\"Get queue statistics\"\"\"\n        async with self._lock:\n            return {\n                \"queue_size\": len(self.queue),\n                \"total_produced\": self.total_produced,\n                \"total_consumed\": self.total_consumed,\n                \"dropped_samples\": self.dropped_samples,\n                \"max_queue_size\": self.max_queue_size,\n            }\n\n    async def clear_queue(self):\n        \"\"\"Clear the queue\"\"\"\n        async with self._lock:\n            cleared_count = len(self.queue)\n            self.queue.clear()\n            logger.info(f\"Cleared {cleared_count} samples from queue\")\n\n    async def shutdown(self):\n        \"\"\"Shutdown the message queue\"\"\"\n        async with self._lock:\n            self.running = False\n            # Notify all waiting coroutines so they can exit\n            self._consumer_condition.notify_all()\n        logger.info(\"MessageQueue shutdown\")\n\n    async def get_memory_usage(self) -> dict:\n        \"\"\"Get memory usage statistics\"\"\"\n        async with self._lock:\n            # Estimate memory usage of samples in queue\n            import sys\n\n            total_size = 0\n            sample_count = len(self.queue)\n\n            if sample_count > 0:\n                # Estimate size of a single sample (simplified estimation)\n                sample = list(self.queue)[0]\n                try:\n                    sample_size = sys.getsizeof(sample)\n                    # Since we now store RolloutSample directly, estimate based on its components\n                    if hasattr(sample, \"original_batch_dict\") and sample.original_batch_dict:\n                        # Estimate batch data size\n                        batch_data = sample.original_batch_dict.get(\"batch\", {})\n                        sample_size += len(batch_data) * 1000  # Roughly estimate 1KB per batch entry\n                    if hasattr(sample, \"agent_loop_output\"):\n                        # Estimate AgentLoopOutput size\n                        sample_size += 5000  # Roughly estimate 5KB for AgentLoopOutput\n                    total_size = sample_size * sample_count\n                except Exception:\n                    total_size = sample_count * 15000  # Roughly estimate 15KB per RolloutSample\n\n            return {\n                \"queue_samples\": sample_count,\n                \"estimated_memory_bytes\": total_size,\n                \"estimated_memory_mb\": total_size / (1024 * 1024),\n            }\n\n    async def put_validate(self, data):\n        async with self._lock:\n            self.val_queue.append(data)\n\n    async def get_validate(self):\n        async with self._lock:\n            if self.val_queue:\n                return self.val_queue.popleft()\n            else:\n                return None\n\n\nclass MessageQueueClient:\n    \"\"\"Asyncio-compatible MessageQueue client for communicating with MessageQueue Actor\"\"\"\n\n    def __init__(self, queue_actor: Any):\n        self.queue_actor = queue_actor\n\n    async def put_sample(self, sample: Any) -> bool:\n        \"\"\"Put batch into queue (async)\"\"\"\n        future = self.queue_actor.put_sample.remote(sample)\n        return await asyncio.wrap_future(future.future())\n\n    async def put_validate(self, data: Any) -> bool:\n        future = self.queue_actor.put_validate.remote(data)\n        return await asyncio.wrap_future(future.future())\n\n    def get_validate_sync(self) -> Any | None:\n        return ray.get(self.queue_actor.get_validate.remote())\n\n    async def get_sample(self) -> Any | None:\n        \"\"\"Get single sample from queue, wait until one is available (async)\"\"\"\n        future = self.queue_actor.get_sample.remote()\n        return await asyncio.wrap_future(future.future())\n\n    async def get_queue_size(self) -> int:\n        \"\"\"Get queue size (async)\"\"\"\n        future = self.queue_actor.get_queue_size.remote()\n        return await asyncio.wrap_future(future.future())\n\n    async def get_statistics(self) -> dict[str, Any]:\n        \"\"\"Get statistics (async)\"\"\"\n        future = self.queue_actor.get_statistics.remote()\n        return await asyncio.wrap_future(future.future())\n\n    async def clear_queue(self):\n        \"\"\"Clear queue (async)\"\"\"\n        future = self.queue_actor.clear_queue.remote()\n        await asyncio.wrap_future(future.future())\n\n    async def shutdown(self):\n        \"\"\"Shutdown queue (async)\"\"\"\n        future = self.queue_actor.shutdown.remote()\n        await asyncio.wrap_future(future.future())\n\n    async def get_memory_usage(self) -> dict:\n        \"\"\"Get memory usage statistics (async)\"\"\"\n        future = self.queue_actor.get_memory_usage.remote()\n        return await asyncio.wrap_future(future.future())\n\n    def get_sample_sync(self) -> Any | None:\n        \"\"\"Get single sample from queue (sync - deprecated, use get_sample instead)\"\"\"\n        return ray.get(self.queue_actor.get_sample.remote())\n\n    def get_statistics_sync(self) -> dict[str, Any]:\n        \"\"\"Get statistics (sync - deprecated, use get_statistics instead)\"\"\"\n        return ray.get(self.queue_actor.get_statistics.remote())\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_30b_a3b_base_math_fsdp.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO-Qwen3-30B-A3B-Base-Async'\nexp_name='Fsdp2-tp4sp4'\n\n# Ray\nRAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\nWORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\nRUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nDATA_PATH=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nDATA_PATH=${DATA_PATH:-\"/mnt/bn/${BYTENAS}\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${DATA_PATH}/shared/models/Qwen3-30B-A3B-Base\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${DATA_PATH}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${DATA_PATH}/shared/data/dapo-math/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${DATA_PATH}/shared/data/dapo-math/aime-2024.parquet\"}\n\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 20))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\nenable_filter_groups=True\nfilter_groups_metric=acc\nmax_num_gen_batches=10\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n\nNNODES=${NNODES:-4}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\n# Fully async specific parameters\nn_gpus_rollout=8\nn_gpus_training=8 \nn_nodes_rollout=2 \nn_nodes_train=2 # $((NNODES - n_nodes_rollout))\n\ntrain_bsz=512\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((train_bsz * 400)))\ntest_freq=25\nstaleness_threshold=0.6 # 0 0.3 1\nrequire_batches=1\ntotal_train_gpus=$((n_gpus_training * n_nodes_train))\ntotal_rollout_gpus=$((n_gpus_rollout * n_nodes_rollout))\ntrigger_parameter_sync_step=$((train_bsz / ( train_prompt_mini_bsz * require_batches))) # 8 16 32\npartial_rollout=True\nenforce_eager=False\nnccl_timeout=72000\nenable_sleep_mode=False\n\n# Performance Related Parameter\nsp_size=4\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$((max_prompt_length + max_response_length))\ninfer_ppo_max_token_len=$((max_prompt_length + max_response_length))\nref_offload=True\nactor_offload=False\ngen_tp=4\nfsdp_size=-1\n\n\nray job submit --no-wait --runtime-env=\"${RUNTIME_ENV}\" \\\n    --working-dir \"${WORKING_DIR}\" \\\n    --address \"${RAY_ADDRESS}\" \\\n    -- python3 -m verl.experimental.fully_async_policy.fully_async_main \\\n    --config-path=config \\\n    --config-name='fully_async_dapo_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp \\\n    critic.strategy=fsdp \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.nccl_timeout=${nccl_timeout} \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.50 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    +actor_rollout_ref.rollout.enable_sleep_mode=${enable_sleep_mode} \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.enforce_eager=${enforce_eager} \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','wandb'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}-i${total_rollout_gpus}_t${total_train_gpus}_s${staleness_threshold}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=\"${test_freq}\" \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${n_nodes_train}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${n_nodes_rollout}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=${test_freq} \\\n    async_training.require_batches=${require_batches} \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_async_retool.sh",
    "content": "set -x\n\nexport VLLM_USE_V1=1\n\n# ================= data/model/tool =================\nHDFS_ROOT=${HDFS_ROOT:-$PWD}\nDATA_ROOT=${DATA_ROOT:-$PWD}\n\ndapo_math_17k=$DATA_ROOT/dataset/BytedTsinghua-SIA/DAPO-Math-17k\naime_2024=$DATA_ROOT/dataset/Maxwell-Jia/AIME_2024\naime_2025=$DATA_ROOT/dataset/yentinglin/aime_2025\nmodel_path=$HDFS_ROOT/checkpoint/multiturn-sft-qwen-2.5-7b-instruct/global_step_372\n\ntrain_files=\"['$dapo_math_17k']\"\ntest_files=\"['$aime_2025', '$aime_2024']\"\n\n# tool\ntool_config_path=recipe/retool/sandbox_fusion_tool_config.yaml\nretool_path=recipe/retool/retool.py\n\n# wandb / tensorboard\nproject_name=retool\nexperiment_name=qwen2.5-7b_dapo_async_tool\ndefault_local_dir=$DATA_ROOT/checkpoint/$experiment_name\n\n# ================= algorithm =================\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_turns=16\nmax_prompt_length=2048\nmax_response_length=16384\nactor_lr=1e-6\n\n# ================= perfomance =================\ninfer_tp=4 # vllm\ntrain_sp=4 # train\nfsdp_size=4 # train\noffload=False\n\nactor_max_token_len_per_gpu=$(( (max_prompt_length + max_response_length) * 1 ))\nlog_prob_max_token_len_per_gpu=$(( actor_max_token_len_per_gpu * 4 ))\n\n# ================= async policy =================\nrollout_name=\"vllm\"\nrollout_mode=\"async\"\n\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\nn_gpus_rollout=4\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\ntrain_batch_size=0\nppo_mini_batch_size=16\ngen_prompt_bsz=1\nn_resp_per_prompt=16\nn_resp_per_prompt_val=30\ntotal_rollout_steps=$(((64*250)))\ntest_freq=10\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\nrequire_batches=1\npartial_rollout=True\n\npython3 -m verl.experimental.fully_async_policy.fully_async_main \\\n    algorithm.adv_estimator=$adv_estimator \\\n    algorithm.use_kl_in_reward=$use_kl_in_reward \\\n    algorithm.kl_ctrl.kl_coef=$kl_coef \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.return_raw_chat=True \\\n    data.train_batch_size=$train_batch_size \\\n    data.max_prompt_length=$max_prompt_length \\\n    data.max_response_length=$max_response_length \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.custom_cls.path=$retool_path \\\n    data.custom_cls.name=CustomRLHFDataset \\\n    reward.custom_reward_function.path=$retool_path \\\n    reward.custom_reward_function.name=compute_score \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.model.path=$model_path \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.use_kl_loss=$use_kl_loss \\\n    actor_rollout_ref.actor.kl_loss_coef=$kl_loss_coef \\\n    actor_rollout_ref.actor.clip_ratio_low=$clip_ratio_low \\\n    actor_rollout_ref.actor.clip_ratio_high=$clip_ratio_high \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.actor.optim.lr=$actor_lr \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$ppo_mini_batch_size \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=$actor_max_token_len_per_gpu \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=$train_sp \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=$offload \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=$offload \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=$log_prob_max_token_len_per_gpu \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=$infer_tp \\\n    actor_rollout_ref.rollout.multi_turn.enable=True \\\n    actor_rollout_ref.rollout.multi_turn.max_user_turns=$max_turns \\\n    actor_rollout_ref.rollout.multi_turn.max_assistant_turns=$max_turns \\\n    actor_rollout_ref.rollout.multi_turn.tool_config_path=$tool_config_path \\\n    actor_rollout_ref.rollout.multi_turn.format=hermes \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=$n_resp_per_prompt \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=0.6 \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=1.0 \\\n    actor_rollout_ref.rollout.val_kwargs.n=$n_resp_per_prompt_val \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=$project_name \\\n    trainer.experiment_name=$experiment_name \\\n    trainer.val_before_train=True \\\n    trainer.log_val_generations=20 \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=$default_local_dir \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    trainer.nnodes=$NNODES \\\n    trainer.n_gpus_per_node=$n_gpus_training \\\n    rollout.nnodes=$NNODES \\\n    rollout.n_gpus_per_node=$n_gpus_rollout \\\n    rollout.total_rollout_steps=$total_rollout_steps \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=$test_freq \\\n    async_training.staleness_threshold=$staleness_threshold \\\n    async_training.trigger_parameter_sync_step=$trigger_parameter_sync_step \\\n    async_training.require_batches=$require_batches \\\n    async_training.partial_rollout=$partial_rollout"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_16_16.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_16-16'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-2}\nNNODES_TRAIN=${NNODES_TRAIN:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\""
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_32_32.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_32-32'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-4}\nNNODES_TRAIN=${NNODES_TRAIN:-4}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_12.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-4-12'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=1\nsp_size=1\nfsdp_size=2\n\n# Fully async specific parameters\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*100)))\ntest_freq=10\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=\"${test_freq}\" \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.test_freq=\"${test_freq}\" \\\n    trainer.total_epochs=10 \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\""
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_4_4.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-4-4'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=1\nsp_size=1\nfsdp_size=2\n\n# Fully async specific parameters\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=4\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*100)))\ntest_freq=10\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_64-64'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-8}\nNNODES_TRAIN=${NNODES_TRAIN:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_64_64_mis.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_fully-async_64-64'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-8}\nNNODES_TRAIN=${NNODES_TRAIN:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\n# Rollout Correction\nrollout_is=token\nrollout_is_threshold=2.0\nrollout_rs=seq_mean_k1\nrollout_rs_threshold=\"0.99_1.001\"\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\" \\\n    algorithm.rollout_correction.bypass_mode=False \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_is_threshold=${rollout_is_threshold} \\\n    algorithm.rollout_correction.rollout_rs=${rollout_rs} \\\n    algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold}\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/dapo_7b_math_fsdp2_8_8.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-fully-async-8-8'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=1\nsp_size=1\nfsdp_size=2\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-1}\nNNODES_TRAIN=${NNODES_TRAIN:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntotal_rollout_steps=$(((512*100)))\ntest_freq=10\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=4\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\""
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/geo3k_qwen25vl_7b_megatron_4_4.sh",
    "content": "set -x\nENGINE=${1:-vllm}\nexport CUDA_DEVICE_MAX_CONNECTIONS=1 # For megatron communication/computation overlapping\n\n\nHF_MODEL_PATH=${HF_MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-VL-7B-Instruct\"}\n\ntrain_path=$HOME/data/geo3k/train.parquet\ntest_path=$HOME/data/geo3k/test.parquet\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n\n# Fully async specific parameters\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=4\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=4\ntrain_prompt_mini_bsz=128\ntotal_rollout_steps=$(((512*100)))\ntest_freq=5\nstaleness_threshold=0.1\ntrigger_parameter_sync_step=4\nrequire_batches=2\npartial_rollout=True\ntotal_epochs=200\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    --config-path=config \\\n    --config-name='fully_async_ppo_megatron_trainer.yaml'\\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"$train_path\" \\\n    data.val_files=\"$test_path\" \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.max_prompt_length=1024 \\\n    data.max_response_length=2048 \\\n    actor_rollout_ref.rollout.max_model_len=32768 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=32768 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    data.gen_batch_size=${gen_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.model.path=$HF_MODEL_PATH \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_decay_steps=51200 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=1 \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.01 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=True \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=5120 \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=5120 \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=True \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=5120 \\\n    actor_rollout_ref.rollout.name=$ENGINE \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    +actor_rollout_ref.rollout.engine_kwargs.vllm.disable_mm_preprocessor_cache=True \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=1 \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=4 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=1 \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=True \\\n    actor_rollout_ref.actor.megatron.param_offload=True \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=True \\\n    actor_rollout_ref.actor.megatron.grad_offload=True \\\n    actor_rollout_ref.ref.megatron.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.logger='[\"console\",\"wandb\"]' \\\n    trainer.project_name='verl_grpo_example_geo3k' \\\n    trainer.experiment_name='qwen2_5_vl_7b_megatron_async' \\\n    trainer.test_freq=\"${test_freq}\" \\\n    trainer.total_epochs=\"${total_epochs}\" \\\n    trainer.val_before_train=False \\\n    trainer.save_freq=-1 \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=\"${total_epochs}\" \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\""
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='GRPO-Qwen3-30b-Base-MATH'\nexp_name='GRPO-Qwen3-30b-Base-MATH-megatron-fully-async_96-32'\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\nkl_loss_type=low_var_kl\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\noffload=True\ntrain_ppo_micro_batch_size_per_gpu=2\ninfer_ppo_micro_batch_size_per_gpu=2\n\noptimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}\n\nCOMMON_PP=${COMMON_PP:-1}\nCOMMON_VPP=${COMMON_VPP:-null}\nCOMMON_CP=${COMMON_CP:-2}\nCOMMON_TP=${COMMON_TP:-2}\nCOMMON_EP=${COMMON_EP:-8}\nCOMMON_ETP=${COMMON_ETP:-1}\n\nTRAIN_TP=${TRAIN_TP:-$COMMON_TP}\nINFER_TP=${INFER_TP:-4}\n\nACTOR_PP=${ACTOR_PP:-$COMMON_PP}\nACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}\nACTOR_CP=${ACTOR_CP:-$COMMON_CP}\nACTOR_TP=${ACTOR_TP:-$TRAIN_TP}\nACTOR_EP=${ACTOR_EP:-$COMMON_EP}\nACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}\nROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}\nREF_PP=${REF_PP:-$COMMON_PP}\nREF_VPP=${REF_VPP:-$COMMON_VPP}\nREF_CP=${REF_CP:-$COMMON_CP}\nREF_TP=${REF_TP:-$TRAIN_TP}\nREF_EP=${REF_EP:-$COMMON_EP}\nREF_ETP=${REF_ETP:-$COMMON_ETP}\nCRITIC_PP=${CRITIC_PP:-$COMMON_PP}\nCRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}\nCRITIC_CP=${CRITIC_CP:-$COMMON_CP}\nCRITIC_TP=${CRITIC_TP:-$TRAIN_TP}\nCRITIC_EP=${CRITIC_EP:-$COMMON_EP}\nCRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}\nRM_PP=${RM_PP:-$COMMON_PP}\nRM_VPP=${RM_VPP:-$COMMON_VPP}\nRM_CP=${RM_CP:-$COMMON_CP}\nRM_TP=${RM_TP:-$TRAIN_TP}\nRM_EP=${RM_EP:-$COMMON_EP}\nRM_ETP=${RM_ETP:-$COMMON_ETP}\n\n# install mbridge\n# pip3 install git+https://github.com/ISEEKYAN/mbridge\nUSE_MBRIDGE=True\nUSE_DIST_CKPT=False\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-12}\nNNODES_TRAIN=${NNODES_TRAIN:-4}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=128\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\nrequire_batches=1\npartial_rollout=True\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    --config-path=config \\\n    --config-name='fully_async_ppo_megatron_trainer.yaml'\\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    +actor_rollout_ref.model.override_config.model_config.max_position_embeddings=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.model.use_fused_kernels=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.lr_decay_style='constant' \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.optim.lr_decay_steps=${total_rollout_steps} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=\"flex\" \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\"\n\n"
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/grpo_30b_a3b_base_math_megatron_96_32_mis.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='GRPO-Qwen3-30b-Base-MATH'\nexp_name='GRPO-Qwen3-30b-Base-MATH-megatron-fully-async_96-32'\n\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-30B-A3B-Base\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\nrollout_mode=\"async\"\nrollout_name=\"vllm\" # sglang or vllm\nif [ \"$rollout_mode\" = \"async\" ]; then\n    export VLLM_USE_V1=1\n    return_raw_chat=\"True\"\nfi\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=True\nkl_loss_coef=0.001\nkl_loss_type=low_var_kl\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length)))\noffload=True\ntrain_ppo_micro_batch_size_per_gpu=2\ninfer_ppo_micro_batch_size_per_gpu=2\n\noptimizer_offload_fraction=${OFFLOAD_FRACTION:-1.}\n\nCOMMON_PP=${COMMON_PP:-1}\nCOMMON_VPP=${COMMON_VPP:-null}\nCOMMON_CP=${COMMON_CP:-2}\nCOMMON_TP=${COMMON_TP:-2}\nCOMMON_EP=${COMMON_EP:-8}\nCOMMON_ETP=${COMMON_ETP:-1}\n\nTRAIN_TP=${TRAIN_TP:-$COMMON_TP}\nINFER_TP=${INFER_TP:-4}\n\nACTOR_PP=${ACTOR_PP:-$COMMON_PP}\nACTOR_VPP=${ACTOR_VPP:-$COMMON_VPP}\nACTOR_CP=${ACTOR_CP:-$COMMON_CP}\nACTOR_TP=${ACTOR_TP:-$TRAIN_TP}\nACTOR_EP=${ACTOR_EP:-$COMMON_EP}\nACTOR_ETP=${ACTOR_ETP:-$COMMON_ETP}\nROLLOUT_TP=${ROLLOUT_TP:-$INFER_TP}\nREF_PP=${REF_PP:-$COMMON_PP}\nREF_VPP=${REF_VPP:-$COMMON_VPP}\nREF_CP=${REF_CP:-$COMMON_CP}\nREF_TP=${REF_TP:-$TRAIN_TP}\nREF_EP=${REF_EP:-$COMMON_EP}\nREF_ETP=${REF_ETP:-$COMMON_ETP}\nCRITIC_PP=${CRITIC_PP:-$COMMON_PP}\nCRITIC_VPP=${CRITIC_VPP:-$COMMON_VPP}\nCRITIC_CP=${CRITIC_CP:-$COMMON_CP}\nCRITIC_TP=${CRITIC_TP:-$TRAIN_TP}\nCRITIC_EP=${CRITIC_EP:-$COMMON_EP}\nCRITIC_ETP=${CRITIC_ETP:-$COMMON_ETP}\nRM_PP=${RM_PP:-$COMMON_PP}\nRM_VPP=${RM_VPP:-$COMMON_VPP}\nRM_CP=${RM_CP:-$COMMON_CP}\nRM_TP=${RM_TP:-$TRAIN_TP}\nRM_EP=${RM_EP:-$COMMON_EP}\nRM_ETP=${RM_ETP:-$COMMON_ETP}\n\n# install mbridge\n# pip3 install git+https://github.com/ISEEKYAN/mbridge\nUSE_MBRIDGE=True\nUSE_DIST_CKPT=False\n\n# Fully async specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-12}\nNNODES_TRAIN=${NNODES_TRAIN:-4}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=0\ngen_prompt_bsz=1\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=128\ntotal_rollout_steps=$(((512*400)))\ntest_freq=20\nstaleness_threshold=0.5\ntrigger_parameter_sync_step=4\nrequire_batches=1\npartial_rollout=True\n\n# Rollout Importance Sampling\n\nrollout_is=null\nrollout_rs=seq_mean_k1\nrollout_rs_threshold=\"0.999_1.001\"\n\npython -m verl.experimental.fully_async_policy.fully_async_main \\\n    --config-path=config \\\n    --config-name='fully_async_ppo_megatron_trainer.yaml'\\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    data.return_raw_chat=${return_raw_chat} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    algorithm.rollout_correction.rollout_is=${rollout_is} \\\n    algorithm.rollout_correction.rollout_rs=${rollout_rs} \\\n    algorithm.rollout_correction.rollout_rs_threshold=${rollout_rs_threshold} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    +actor_rollout_ref.model.override_config.model_config.max_position_embeddings=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.model.use_fused_kernels=False \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=${train_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.lr_decay_style='constant' \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.optim.lr_decay_steps=${total_rollout_steps} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_offload_fraction=${optimizer_offload_fraction} \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.overlap_cpu_optimizer_d2h_h2d=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.use_precision_aware_optimizer=True \\\n    +actor_rollout_ref.actor.optim.override_optimizer_config.optimizer_cpu_offload=True \\\n    actor_rollout_ref.actor.megatron.use_mbridge=$USE_MBRIDGE \\\n    actor_rollout_ref.actor.megatron.use_dist_checkpointing=$USE_DIST_CKPT \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${ACTOR_TP} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${ACTOR_PP} \\\n    actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size=${ACTOR_VPP} \\\n    actor_rollout_ref.actor.megatron.context_parallel_size=${ACTOR_CP} \\\n    actor_rollout_ref.actor.megatron.expert_model_parallel_size=${ACTOR_EP} \\\n    actor_rollout_ref.actor.megatron.expert_tensor_parallel_size=${ACTOR_ETP} \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.apply_rope_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.masked_softmax_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_activation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.bias_dropout_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.gradient_accumulation_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.deallocate_pipeline_outputs=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.persist_layer_norm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_grouped_gemm=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_permute_fusion=True \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_token_dispatcher_type=\"flex\" \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_router_dtype=fp32 \\\n    +actor_rollout_ref.actor.megatron.override_transformer_config.moe_enable_deepep=True \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.7 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${INFER_TP} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=${rollout_name} \\\n    actor_rollout_ref.rollout.mode=${rollout_mode} \\\n    actor_rollout_ref.rollout.calculate_log_probs=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.free_cache_engine=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=${infer_ppo_micro_batch_size_per_gpu} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.megatron.use_dist_checkpointing=${USE_DIST_CKPT} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${REF_TP} \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${REF_PP} \\\n    actor_rollout_ref.ref.megatron.virtual_pipeline_model_parallel_size=${REF_VPP} \\\n    actor_rollout_ref.ref.megatron.context_parallel_size=${REF_CP} \\\n    actor_rollout_ref.ref.megatron.expert_model_parallel_size=${REF_EP} \\\n    actor_rollout_ref.ref.megatron.expert_tensor_parallel_size=${REF_ETP} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.total_rollout_steps=\"${total_rollout_steps}\" \\\n    trainer.total_epochs=10 \\\n    trainer.test_freq=\"${test_freq}\" \\\n    async_training.staleness_threshold=\"${staleness_threshold}\" \\\n    async_training.trigger_parameter_sync_step=\"${trigger_parameter_sync_step}\" \\\n    async_training.require_batches=\"${require_batches}\" \\\n    async_training.partial_rollout=\"${partial_rollout}\""
  },
  {
    "path": "verl/experimental/fully_async_policy/shell/runtime_env.yaml",
    "content": "env_vars:\n  VLLM_USE_V1: \"1\"\n  NCCL_DEBUG: \"INFO\"\n  HYDRA_FULL_ERROR: \"1\""
  },
  {
    "path": "verl/experimental/fully_async_policy/unittest/simple_streaming_demo.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport random\nimport time\n\n\nclass SimpleStreamingSystem:\n    \"\"\"Simplified streaming system demonstration\"\"\"\n\n    def __init__(self, max_concurrent_tasks: int = 4):\n        self.max_concurrent_tasks = max_concurrent_tasks\n        self.data_queue = asyncio.Queue()\n        self.result_queue = asyncio.Queue()\n        self.consumer_count = 0\n\n    # Data stream coroutine\n    async def data_stream(self):\n        # Add initial data\n        # Prepare test data\n        test_data = [{\"id\": f\"task_{i}\", \"content\": f\"data_{i}\"} for i in range(8)]\n        await self.add_data_stream(test_data)\n\n        # Simulate subsequent data stream\n        await asyncio.sleep(3)\n        print(\"\\nAdding second batch of data...\")\n        extra_data = [{\"id\": f\"extra_{i}\", \"content\": f\"extra_data_{i}\"} for i in range(5)]\n        await self.add_data_stream(extra_data)\n\n        # Send termination signal\n        await asyncio.sleep(1)\n        await self.data_queue.put(\"DONE\")\n        print(\"Sending termination signal\")\n\n    async def add_data_stream(self, data_list: list[dict]):\n        \"\"\"Simulate data stream\"\"\"\n        print(\"Starting to add data stream...\")\n\n        for i, data_item in enumerate(data_list):\n            await self.data_queue.put(data_item)\n            print(f\"Data {data_item['id']} added to pending queue\")\n\n            # Simulate interval between data streams\n            if i < len(data_list) - 1:  # Don't wait after the last item\n                await asyncio.sleep(0.8)\n\n        print(\"Initial data stream added successfully\")\n\n    async def _process_data_async(self, data_item: dict):\n        \"\"\"Asynchronously process a single data item\"\"\"\n        data_id = data_item[\"id\"]\n        content = data_item[\"content\"]\n\n        # Simulate different processing times (1-3 seconds)\n        processing_time = random.uniform(1, 3)\n\n        print(f\"    Starting to process {data_id}, estimated time {processing_time:.1f}s\")\n\n        # Asynchronously wait for processing completion\n        await asyncio.sleep(processing_time)\n\n        result = {\n            \"id\": data_id,\n            \"processed_content\": f\"Processed {content}\",\n            \"processing_time\": round(processing_time, 2),\n            \"completed_at\": time.time(),\n        }\n\n        # Immediately put into result queue\n        await self.result_queue.put(result)\n        print(f\"    {data_id} processing completed! (took {processing_time:.1f}s) -> Added to result queue\")\n\n    async def _submit_worker(self):\n        \"\"\"Stream submission worker coroutine\"\"\"\n        active_tasks = set()\n\n        print(\"Stream submitter started...\")\n\n        while True:\n            # Get data to process\n            data_item = await self.data_queue.get()\n\n            if data_item == \"DONE\":\n                print(\"Received termination signal, waiting for remaining tasks to complete...\")\n                if active_tasks:\n                    await asyncio.gather(*active_tasks, return_exceptions=True)\n                break\n\n            # Check concurrent limit\n            while len(active_tasks) >= self.max_concurrent_tasks:\n                print(f\"Reached maximum concurrency {self.max_concurrent_tasks}, waiting for tasks to complete...\")\n                done_tasks, active_tasks = await asyncio.wait(active_tasks, return_when=asyncio.FIRST_COMPLETED)\n\n                # Clean up completed tasks\n                for task in done_tasks:\n                    try:\n                        await task\n                        print(f\"Task completed {task}\")\n                    except Exception as e:\n                        print(f\"Task execution failed: {e}\")\n\n            # Immediately submit new task\n            task = asyncio.create_task(self._process_data_async(data_item), name=f\"active {data_item}\")\n            active_tasks.add(task)\n\n            print(f\"Submitted task {data_item['id']}, current concurrency: {len(active_tasks)}\")\n\n    async def _consumer_worker(self):\n        \"\"\"Result consumer coroutine\"\"\"\n        print(\"Consumer started...\")\n\n        while True:\n            try:\n                # Get processing result from result queue\n                result = await asyncio.wait_for(self.result_queue.get(), timeout=2.0)\n\n                self.consumer_count += 1\n\n                print(\n                    f\"Consumed #{self.consumer_count}: {result['id']} \"\n                    f\"(processing time {result['processing_time']}s) - {result['processed_content']}\"\n                )\n\n            except asyncio.TimeoutError:\n                print(\"    Consumer waiting...\")\n                await asyncio.sleep(0.5)\n\n    async def run_demo(self):\n        \"\"\"Run demonstration\"\"\"\n        print(\"=\" * 60)\n        print(f\"Maximum concurrency: {self.max_concurrent_tasks}\")\n        print(\"=\" * 60)\n\n        # Start core coroutines\n        stream_task = asyncio.create_task(self.data_stream())\n        submit_task = asyncio.create_task(self._submit_worker())\n        consumer_task = asyncio.create_task(self._consumer_worker())\n\n        try:\n            # Wait for data stream to complete\n            await stream_task\n            print(\"Data stream completed\")\n\n            # Wait for processing to complete\n            await submit_task\n            print(\"All tasks processed\")\n\n        finally:\n            # Cleanup\n            submit_task.cancel()\n            consumer_task.cancel()\n            await asyncio.gather(submit_task, consumer_task, return_exceptions=True)\n\n        print(f\"\\nFinal statistics: Consumed {self.consumer_count} results\")\n\n\nasync def main():\n    \"\"\"Main function\"\"\"\n    system = SimpleStreamingSystem(max_concurrent_tasks=3)\n    await system.run_demo()\n\n\nif __name__ == \"__main__\":\n    asyncio.run(main())\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/README.md",
    "content": "# Recipe: One Step Off Policy Async Trainer\n\n**Author:** `https://github.com/meituan-search`\n\nLast updated: 07/17/2025.\n\n## Introduction\n\n### Background\n\nThe current reinforcement learning training process implemented by verl is synchronous, adhering to the algorithmic\nworkflows of established methods like PPO, GRPO, and DAPO. In each step, training samples are generated by the latest\nmodel, and the model is updated after training completes. While this approach aligns with off-policy reinforcement\nlearning and stabilizes RL training, but it suffers from severe efficiency issues.\nModel updates must wait for the longest output in the generation phase to complete.\nDuring the generation of long-tail samples, GPUs remain idle, resulting in significant underutilization.\nThe more severe the long-tail problem in sample generation, the lower the overall training efficiency.\nFor example, in DAPO 32B training, the Rollout phase accounts for approximately 70% of the total time,\nand increasing resources does not reduce the Rollout duration.\n\n![DAPO 32B Math Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/dapo_32b_math.png)\n\n> source data: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/workspace?nw=nwusertongyuxuan361\n\n### Solution\n\nWe have implemented the **One Step Off Async Trainer** to help alleviate this issue. This approach parallelizes the\ngeneration and training processes, utilizing samples generated in the previous step for current training.\nIt also involves appropriately partitioning resources, allocating dedicated resources for generation while automatically\nassigning the remainder to training. By reducing resources allocated to the generation phase, we mitigate GPU idle time\nduring long-tail sample generation. Throughout this process, generation and training parameters maintain a one-step off\npolicy.\n\n![One Step Off Policy Diagram](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_policy.png)\n\n> reference: [AReaL: A Large-Scale Asynchronous Reinforcement Learning System for Language Reasoning](https://arxiv.org/abs/2505.24298)\n> original work: [Asynchronous RLHF: Faster and More Efficient Off-Policy RL for Language Models](https://arxiv.org/abs/2410.18252)\n\nOur core contributions include:\n\n1. **Parallel Generation and Training**:\n   Samples for the next batch are asynchronously generated while the current batch is being trained.\n\n2. **Resource Isolation**:\n   Unlike `hybrid_engine`, this method requires explicit resource allocation for rollout, with remaining resources\n   automatically assigned to training.\n\n3. **NCCL Parameter Synchronization**:\n   Employs NCCL communication primitives for seamless parameter transfer between generation and training modules.\n\n### Experimental Results\n\n- **Machine Configuration**: 2 nodes with 16 H20 GPUs each\n  - Generation: 4 GPUs\n  - Training: 12 GPUs\n- **Model**: Qwen2.5-Math-7B\n- **Rollout Configuration**:\n- **Max Response Length**: FSDP2: 20,480 tokens; Megatron: 8,192 tokens\n- **Algorithm**: DAPO\n- **Rollout Engine**: vLLM\n\n| training mode          | engine        | step | gen | wait_prev_gen | generate_sequences | old_log_prob | update_actor | total time     | acc/best@32/mean | acc/maj@32/mean |\n| ---------------------- | ------------- | ---- | --- | ------------- | ------------------ | ------------ | ------------ | -------------- | ---------------- | --------------- |\n| colocate sync          | VLLM+FSDP2    | 749  | 321 | -             | 247                | 88           | 286          | 19h18m         | 0.5948           | 0.417           |\n| one-step-overlap async | VLLM+FSDP2    | 520  | -   | 45            | 458                | 108          | 337          | 15h34m（+23%） | 0.6165           | 0.494           |\n| colocate sync          | VLLM+Megatron | 699  | 207 | -             | 162                | 119          | 344          | 18h21m         | 0.605            | 0.4217          |\n| one-step-overlap async | VLLM+Megatron | 566  | -   | 59            | 501                | 120          | 347          | 13h06m (+40%)  | 0.6569           | 0.4038          |\n\n- colocate sync: step ≈ gen + old_log_prob + update_actor\n- one-step-overlap async: step ≈ wait_prev_gen + old_log_prob + update_actor\n\n![One Step Off Megatron Performance](https://raw.githubusercontent.com/eric-haibin-lin/verl-community/refs/heads/main/docs/one_step_off_megatron.png)\n\n> source data: https://wandb.ai/hou-zg-meituan/one-step-off-policy?nw=nwuserhouzg\n\n## Implementation\n\n### One Step Off Policy Async Pipeline\n\nOur implemented **One Step Off Policy Async Pipeline** integrates seamlessly into existing training logic at minimal\ncost,\neliminating the need for additional sample storage management. The core mechanism uses `async_gen_next_batch`\nfor asynchronous rollout generation while maintaining continuous operation during epoch transitions\nvia `create_continuous_iterator`.\n\n```python\n# iterator generator, simplify one-step integration of the training process\ndef _create_continuous_iterator(self):\n    for epoch in range(self.config.trainer.total_epochs):\n        iterator = iter(self.train_dataloader)\n        for batch_dict in iterator:\n            yield epoch, batch_dict\n\n\n# read next batch samples, parameters sync and launch asyn gen_seq\ndef _async_gen_next_batch(self, continuous_iterator):\n    # read train_data\n    try:\n        epoch, batch_dict = next(continuous_iterator)\n    except StopIteration:\n        return None\n    batch = DataProto.from_single_dict(batch_dict)\n    gen_batch = batch_pocess(batch)\n    # sync weights from actor to rollout\n    self.sync_rollout_weights()\n    # async generation\n    gen_batch_output = self.rollout_wg.async_generate_sequences(gen_batch)\n    # future encapsulated\n    return GenerationBatchFuture(epoch, batch, gen_batch_output)\n\n\ncontinuous_iterator = self._create_continuous_iterator()\n# run rollout first to achieve one-step-off\nbatch_data_future = self._async_gen_next_batch(continuous_iterator)\n\nwhile batch_data_future is not None:\n    # wait for the gen_seq result from the previous step\n    batch = batch_data_future.get()\n    # launch the next async call to generate sequences\n    batch_data_future = self._async_gen_next_batch(continuous_iterator)\n\n    # compute advantages\n    batch = critic.compute_values(batch)\n    batch = reference.compute_log_prob(batch)\n    batch = reward.compute_reward(batch)\n    batch = compute_advantages(batch)\n\n    # model update\n    critic_metrics = critic.update_critic(batch)\n    actor_metrics = actor.update_actor(batch)\n```\n\n### Parameter Synchronization\n\nThe exciting point is that our nccl based weights updating for rollout model has great performance.\nAt most of time, the latency is under 300ms, which is negligible for RLHF.\n\n> **sync_rollout_weights**：The time for synchronizing parameters from actor to rollout is extremely fast and can almost\n> be ignored because it is implemented with nccl.\n\n```python\nclass ActorRolloutRefWorker:\n    # actor acquires the meta-info of model parameters for parameter sync\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def get_actor_weights_info(self):\n        params = self._get_actor_params()\n        ret = []\n        for key, tensor in params.items():\n            ret.append((key, tensor.size(), tensor.dtype))\n        self._weights_info = ret\n        return ret\n\n    # rollout sets the meta-info of model parameters for parameter sync\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def set_actor_weights_info(self, weights_info):\n        self._weights_info = weights_info\n\n\nclass AsyncRayPPOTrainer(RayPPOTrainer):\n    def init_workers(self):\n\n\n...\n# rollout obtains the meta-info of model parameters from the actor for parameter sync\nweights_info = self.actor_wg.get_actor_weights_info()[0]\nself.rollout_wg.set_actor_weights_info(weights_info)\n\n# Create an actor-rollout communication group for parameter sync\nactor_rollout_workers = self.actor_wg.workers + self.rollout_wg.workers\ncollective.create_collective_group(\n    actor_rollout_workers,\n    len(actor_rollout_workers),\n    list(range(0, len(actor_rollout_workers))),\n    backend=\"nccl\",\n    group_name=\"actor_rollout\"\n)\n```\n\n```python\n# drive process call the actor and rollout respectively to sync parameters by nccl\ndef sync_rollout_weights(self):\n    self.actor_wg.sync_rollout_weights()\n    ray.get(self.rollout_wg.sync_rollout_weights())\n\n\n# fsdp model parameter sync\n@register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\ndef sync_rollout_weights(self):\n    params = self._get_actor_params() if self._is_actor else None\n    if self._is_rollout:\n        inference_model = (\n            self.rollout.inference_engine.llm_engine.model_executor.driver_worker.worker.model_runner.model\n        )\n        from verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader\n        patch_vllm_moe_model_weight_loader(inference_model)\n    # Model parameters are broadcast tensor-by-tensor from actor to rollout\n    for key, shape, dtype in self._weights_info:\n        tensor = torch.empty(shape, dtype=dtype, device=get_torch_device().current_device())\n        if self._is_actor:\n            assert key in params\n            origin_data = params[key]\n            if hasattr(origin_data, \"full_tensor\"):\n                origin_data = origin_data.full_tensor()\n            if torch.distributed.get_rank() == 0:\n                tensor.copy_(origin_data)\n        from ray.util.collective import collective\n\n        collective.broadcast(tensor, src_rank=0, group_name=\"actor_rollout\")\n        if self._is_rollout:\n            inference_model.load_weights([(key, tensor)])\n```\n\n### PPO Correctness\n\nTo ensure the correctness of the PPO algorithm, we use rollout log_probs for PPO importance sampling.\nFor the related algorithm details, please refer to: https://verl.readthedocs.io/en/latest/algo/rollout_corr_math.html\nThe default mode is `bypass_ppo_clip`, but other modification strategies can also be explored.\n\n### AgentLoop\n\nIn the current implementation, we no longer provide SPMD model rollout mode.\nInstead, we have switched to AgentLoop mode, which also supports multi-turn tool calling.\n\n## Usage\n\n### FSDP2 Configuration Example\n\n```shell\npython3 -m verl.experimental.one_step_off_policy.async_main_ppo \\\n    --config-path=config \\\n    --config-name='one_step_off_ppo_trainer.yaml' \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    # actor and rollout are placed separately\n    actor_rollout_ref.hybrid_engine=False \\\n    # actor and rollout resource\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=6 \\\n    rollout.nnodes=1 \\\n    rollout.n_gpus_per_node=2\n```\n\n### Megatron Configuration Example\n\n```shell\npython3 -m verl.experimental.one_step_off_policy.async_main_ppo \\\n    --config-path=config \\\n    --config-name='one_step_off_ppo_megatron_trainer.yaml' \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    # actor and rollout are placed separately\n    actor_rollout_ref.hybrid_engine=False \\\n    # actor and rollout resource\n    trainer.nnodes=1 \\\n    trainer.n_gpus_per_node=6 \\\n    rollout.nnodes=1 \\\n    rollout.n_gpus_per_node=2\n```\n\n### Configuration Guidelines\n\n1. **Card Number Relationships**\n   Maintain either of these relationships for optimal batch distribution:\n\n   - `actor_rollout_ref.rollout.n` should be an integer divisor of:\n     `trainer.n_gpus_per_node * trainer.nnodes`\n   - `actor_rollout_ref.rollout.n * data.train_batch_size` should be evenly divisible by:\n     `trainer.n_gpus_per_node * trainer.nnodes`\n\n   > Rationale: Ensures training samples can be evenly distributed across training GPUs when using partial resources for\n   > generation.\n\n2. **Dynamic Resource Tuning**\n   Adjust `trainer.nnodes` `trainer.n_gpus_per_node` `rollout.nnodes` `rollout.n_gpus_per_node` based on phase\n   durations:\n   - **Ideal state**: Rollout and training phases have comparable durations\n   - **Diagnostic metrics**:\n     - Monitor `wait_prev_gen` duration\n     - Analyze `sequence_length` distribution\n   - **Adjustment strategy**: - High `wait_prev_gen` + uniform sequence lengths → Increase rollout resources - High `wait_prev_gen` + long-tail sequences → Optimize stopping criteria (resource increase won't help)\n     > **wait_prev_gen**：The time consumed waiting for the previous rollout to end (the part that is not fully\n     > overlapped).\n     > **Resource Configuration Strategies:**\n   - **Resource-constrained scenario**: Optimize resource utilization by adjusting GPU allocation ratios,\n     keeping the number of nodes equal to allow training and rollout to share nodes;\n     - Configure `trainer.nnodes = rollout.nnodes` with\n       `trainer.n_gpus_per_node + rollout.n_gpus_per_node = physical_gpus_per_node`. Control rollout resource\n       allocation by adjusting `n_gpus_per_node`.\n   - **Resource-abundant scenario**: Optimize performance by adjusting the number of nodes,\n     keeping the number of GPUs per node equal to enable independent scaling of training and rollout\n     parallelism. - Configure `trainer.n_gpus_per_node = rollout.n_gpus_per_node` and control rollout resource allocation by\n     adjusting `trainer.nnodes` and `rollout.nnodes`to achieve optimal performance.\n     > **Note**: The total number of nodes required by the system is not simply `trainer.nnodes + rollout.nnodes`. The\n     > actual calculation depends on GPU capacity:\n     >\n     > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node <= physical_gpus_per_node`,\n     >   the required node count is `max(trainer.nnodes, rollout.nnodes)`\n     > - When `trainer.n_gpus_per_node + rollout.n_gpus_per_node > physical_gpus_per_node`,\n     >   the required node count is `trainer.nnodes + rollout.nnodes`\n\n## Functional Support\n\n| Category           | Support Situation                                                                                               |\n| ------------------ | --------------------------------------------------------------------------------------------------------------- |\n| train engine       | FSDP2 <br/> Megatron                                                                                            |\n| rollout engine     | vLLM <br/> SGLang                                                                                               |\n| AdvantageEstimator | GRPO <br/> GRPO_PASSK <br/> REINFORCE_PLUS_PLUS <br/> RLOO <br/> OPO <br/> REINFORCE_PLUS_PLUS_BASELINE<br/>GPG |\n| Reward             | all                                                                                                             |\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/config/one_step_off_ppo_megatron_trainer.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_megatron_trainer\n  - _self_\n\ntrainer:\n  use_legacy_worker_impl: disable\n\n# config for the rollout (only for resource isolation)\nrollout:\n  # Number of nodes used in the rollout\n  nnodes: 1\n  # Number of GPUs per node\n  n_gpus_per_node: 8\n\n# To adapt to the current logic of AgentLoopManager\nactor_rollout_ref:\n  rollout:\n    # Must be turned off! Otherwise, Parameter synchronization cannot be performed.\n    free_cache_engine: False\n    # Must be enabled! Otherwise, log_probs cannot be calculated.\n    calculate_log_probs: True\n\n    checkpoint_engine:\n      backend: \"nccl\"\n\n# Only then will the use of log probs be correct.\n# And it can be used in conjunction with other rollout_correction algorithms.\nalgorithm:\n  rollout_correction:\n    bypass_mode: True"
  },
  {
    "path": "verl/experimental/one_step_off_policy/config/one_step_off_ppo_trainer.yaml",
    "content": "hydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\ntrainer:\n  use_legacy_worker_impl: disable\n\n# config for the rollout (only for resource isolation)\nrollout:\n  # Number of nodes used in the rollout\n  nnodes: 1\n  # Number of GPUs per node\n  n_gpus_per_node: 8\n\n# To adapt to the current logic of AgentLoopManager\nactor_rollout_ref:\n  rollout:\n    # Must be turned off! Otherwise, Parameter synchronization cannot be performed.\n    free_cache_engine: False\n    # Must be enabled! Otherwise, log_probs cannot be calculated.\n    calculate_log_probs: True\n\n    checkpoint_engine:\n      backend: \"nccl\"\n\n# Only then will the use of log probs be correct.\n# And it can be used in conjunction with other rollout_correction algorithms.\nalgorithm:\n  rollout_correction:\n    bypass_mode: True"
  },
  {
    "path": "verl/experimental/one_step_off_policy/main_ppo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nNote that we don't combine the main with ray_trainer as ray_trainer is used by other main.\n\"\"\"\n\nimport asyncio\nimport os\nimport socket\n\nimport hydra\nimport ray\n\nfrom verl.experimental.one_step_off_policy.ray_trainer import OneStepOffRayTrainer\nfrom verl.experimental.separation.utils import create_resource_pool_manager, create_role_worker_mapping\nfrom verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler\nfrom verl.trainer.ppo.utils import need_critic, need_reference_policy\nfrom verl.utils.config import validate_config\nfrom verl.utils.device import auto_set_device\n\n\n@ray.remote(num_cpus=10, max_concurrency=100)  # please make sure main_task is not scheduled on head\nclass OneStepTaskRunner:\n    def run(self, config):\n        # Print the initial configuration. `resolve=True` will evaluate symbolic values.\n        from pprint import pprint\n\n        from omegaconf import OmegaConf\n\n        from verl.utils.fs import copy_to_local\n\n        print(f\"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}\")\n\n        pprint(OmegaConf.to_container(config, resolve=True))\n\n        OmegaConf.resolve(config)\n\n        role_worker_mapping, ray_worker_group_cls = create_role_worker_mapping(config)\n\n        # validate config\n        validate_config(\n            config=config,\n            use_reference_policy=need_reference_policy(config),\n            use_critic=need_critic(config),\n        )\n\n        # Download the checkpoint from HDFS to the local machine.\n        # `use_shm` determines whether to use shared memory, which could lead to faster model loading if turned on\n        local_path = copy_to_local(\n            config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get(\"use_shm\", False)\n        )\n\n        # Instantiate the tokenizer and processor.\n        from verl.utils import hf_processor, hf_tokenizer\n\n        trust_remote_code = config.data.get(\"trust_remote_code\", False)\n        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)\n        # Used for multimodal LLM, could be None\n        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)\n\n        resource_pool_manager = create_resource_pool_manager(config, role_worker_mapping.keys())\n\n        from verl.utils.dataset.rl_dataset import collate_fn\n\n        # Create training and validation datasets.\n        train_dataset = create_rl_dataset(\n            config.data.train_files,\n            config.data,\n            tokenizer,\n            processor,\n            max_samples=config.data.get(\"train_max_samples\", -1),\n        )\n        val_dataset = create_rl_dataset(\n            config.data.val_files, config.data, tokenizer, processor, max_samples=config.data.get(\"val_max_samples\", -1)\n        )\n        train_sampler = create_rl_sampler(config.data, train_dataset)\n\n        # Initialize the PPO trainer.\n        trainer = OneStepOffRayTrainer(\n            config=config,\n            tokenizer=tokenizer,\n            processor=processor,\n            role_worker_mapping=role_worker_mapping,\n            resource_pool_manager=resource_pool_manager,\n            ray_worker_group_cls=ray_worker_group_cls,\n            train_dataset=train_dataset,\n            val_dataset=val_dataset,\n            collate_fn=collate_fn,\n            train_sampler=train_sampler,\n            device_name=config.trainer.device,\n        )\n        # Initialize the workers of the trainer.\n        trainer.init_workers()\n        # Start the training process.\n        asyncio.run(trainer.fit())\n\n\n@hydra.main(config_path=\"config\", config_name=\"one_step_off_ppo_trainer\", version_base=None)\ndef main(config):\n    from time import time\n\n    from verl.trainer.main_ppo import run_ppo\n\n    start_time = time()\n\n    # Automatically set `config.trainer.device = npu` when running on Ascend NPU.\n    auto_set_device(config)\n\n    # TODO: unify rollout config with actor_rollout_ref\n    config.actor_rollout_ref.rollout.nnodes = config.rollout.nnodes\n    config.actor_rollout_ref.rollout.n_gpus_per_node = config.rollout.n_gpus_per_node\n\n    run_ppo(config, task_runner_class=OneStepTaskRunner)\n    print(f\"total time: {time() - start_time:.2f} seconds\")\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/ray_trainer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis trainer supports model-agonistic model initialization with huggingface\n\"\"\"\n\nimport asyncio\nimport uuid\nfrom pprint import pprint\nfrom typing import Optional\n\nimport numpy as np\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import Dataset, Sampler\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.experimental.separation.ray_trainer import SeparateRayPPOTrainer\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup\nfrom verl.trainer.ppo import core_algos\nfrom verl.trainer.ppo.ray_trainer import (\n    ResourcePoolManager,\n    compute_response_mask,\n)\nfrom verl.trainer.ppo.reward import extract_reward\nfrom verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.rollout_skip import RolloutSkip\nfrom verl.utils.tracking import ValidationGenerationsLogger\n\n\nclass OneStepOffRayTrainer(SeparateRayPPOTrainer):\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        role_worker_mapping: dict[Role, WorkerType],\n        resource_pool_manager: ResourcePoolManager,\n        ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup,\n        processor=None,\n        train_dataset: Optional[Dataset] = None,\n        val_dataset: Optional[Dataset] = None,\n        collate_fn=None,\n        train_sampler: Optional[Sampler] = None,\n        device_name=None,\n    ):\n        \"\"\"\n        Initialize distributed PPO trainer with Ray backend.\n        Note that this trainer runs on the driver process on a single CPU/GPU node.\n\n        Args:\n            config: Configuration object containing training parameters.\n            tokenizer: Tokenizer used for encoding and decoding text.\n            role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes.\n            resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools.\n            ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup.\n            processor: Optional data processor, used for multimodal data\n            train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None.\n            val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None.\n            collate_fn: Function to collate data samples into batches.\n            train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None.\n            device_name (str, optional): Device name for training (e.g., \"cuda\", \"cpu\"). Defaults to None.\n        \"\"\"\n\n        # Store the tokenizer for text processing\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.config = config\n\n        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine\n        assert not self.hybrid_engine\n\n        # Skip rollout worker mapping and let agentloop create it.\n        role_worker_mapping.pop(Role.Rollout, None)\n        self.role_worker_mapping = role_worker_mapping\n        self.resource_pool_manager = resource_pool_manager\n        self.use_reference_policy = need_reference_policy(self.config)\n\n        self.use_rm = need_reward_model(self.config)\n\n        self.use_critic = need_critic(self.config)\n\n        self.ray_worker_group_cls = ray_worker_group_cls\n        self.device_name = device_name if device_name else self.config.trainer.device\n        self.validation_generations_logger = ValidationGenerationsLogger(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n        )\n\n        # if ref_in_actor is True, the reference policy will be actor without lora applied\n        lora_rank = config.actor_rollout_ref.model.get(\"lora\", {}).get(\"rank\", 0)\n        if lora_rank <= 0:\n            lora_rank = config.actor_rollout_ref.model.get(\"lora_rank\", 0)\n        self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get(\"lora_adapter_path\") is not None\n\n        # define in-reward KL control\n        # kl loss control currently not suppoorted\n        if self.config.algorithm.use_kl_in_reward:\n            self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)\n\n        self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get(\"use_prefix_grouper\", False)\n        self.use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n\n        self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)\n\n        # ==================== SeparateRayPPOTrainer config ====================\n\n        self.global_steps = 0\n        self.epoch = 0\n        self.max_steps_duration = 0\n        self.progress_bar = None\n        self.logger = None\n        self.is_last_step = False\n        self.prev_step_profile = False\n        self.curr_step_profile = False\n        self.next_step_profile = False\n        self.last_val_metrics = {}\n        self.metrics = {}\n        self.timing_raw = {}\n        # reward message\n        self.future_reward = None\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n\n    def _create_actor_rollout_classes(self):\n        for role in [Role.Actor]:\n            resource_pool = self.resource_pool_manager.get_resource_pool(role)\n            role_cls = RayClassWithInitArgs(\n                cls=self.role_worker_mapping[role],\n                config=self.config.actor_rollout_ref,\n                role=str(role),\n            )\n            self.resource_pool_to_cls[resource_pool][str(role)] = role_cls\n\n    def _init_models(self):\n        if self.use_critic:\n            self.critic_wg = self.all_wg[str(Role.Critic)]\n            self.critic_wg.init_model()\n\n        if self.use_reference_policy and not self.ref_in_actor:\n            self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]\n            self.ref_policy_wg.init_model()\n\n        self.rm_wg = None\n        if self.use_rm:\n            self.rm_wg = self.all_wg[str(Role.RewardModel)]\n            self.rm_wg.init_model()\n\n        self.actor_wg = self.all_wg[str(Role.Actor)]\n        self.actor_wg.init_model()\n        self.actor_rollout_wg = self.actor_wg\n\n    def _init_async_rollout_manager(self):\n        # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design\n        # agent_reward_loop: streaming reward computation with actor rollout\n        # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool\n        enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool\n\n        # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager\n        # to stream reward computation with actor rollout\n        reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None\n\n        # create async rollout manager and request scheduler\n        assert self.config.actor_rollout_ref.rollout.mode == \"async\"\n        from verl.experimental.agent_loop import AgentLoopManager\n\n        self.async_rollout_mode = True\n        self.async_rollout_manager = AgentLoopManager.create(\n            config=self.config, reward_loop_worker_handles=reward_loop_worker_handles\n        )\n\n    def _create_continuous_iterator(self):\n        \"\"\"\n        Create a continuous data iterator across epoch\n        \"\"\"\n        for epoch in range(self.config.trainer.total_epochs):\n            iterator = iter(self.train_dataloader)\n            for batch_dict in iterator:\n                yield epoch, batch_dict\n\n    async def _async_gen_next_batch(self, continuous_iterator):\n        \"\"\"\n        Call parameter synchronization and asynchronous sequence generation.\n        \"\"\"\n        try:\n            epoch, batch_dict = next(continuous_iterator)\n        except StopIteration:\n            return None\n        except Exception as e:\n            print(f\"Error in async_gen_next_batch: {e}\")\n            return None\n\n        metrics = {}\n        timing_raw = {}\n\n        # Create the initial batch from the data loader\n        batch = DataProto.from_single_dict(batch_dict)\n\n        # add uid to batch\n        batch.non_tensor_batch[\"uid\"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)\n\n        gen_batch = self._get_gen_batch(batch)\n\n        # pass global_steps to trace\n        gen_batch.meta_info[\"global_steps\"] = self.global_steps\n        gen_batch_output = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n\n        # async generation\n        with marked_timer(\"generate_async\", timing_raw, color=\"purple\"):\n            gen_batch_output = await self.async_rollout_manager.generate_sequences(gen_batch_output)\n\n        # repeat to align with repeated responses in rollout\n        batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n        batch = batch.union(gen_batch_output)\n\n        if \"response_mask\" not in batch.batch.keys():\n            batch.batch[\"response_mask\"] = compute_response_mask(batch)\n        # Balance the number of valid tokens across DP ranks.\n        # NOTE: This usually changes the order of data in the `batch`,\n        # which won't affect the advantage calculation (since it's based on uid),\n        # but might affect the loss calculation (due to the change of mini-batching).\n        if self.config.trainer.balance_batch:\n            self._balance_batch(batch, metrics=metrics)\n\n        # compute global_valid tokens\n        batch.meta_info[\"global_token_num\"] = torch.sum(batch.batch[\"attention_mask\"], dim=-1).tolist()\n\n        # Launch individual reward computations as each generation completes\n        future_reward = None\n\n        # Return the original, now-modified `batch` and the `future_reward`\n        return metrics, timing_raw, epoch, batch, future_reward\n\n    @staticmethod\n    @ray.remote\n    def _launch_individual_rewards(batch, config, tokenizer):\n        reward_tensor, reward_extra_info = extract_reward(batch)\n        return reward_tensor, reward_extra_info\n\n    async def fit(self):\n        \"\"\"\n        The training loop of PPO.\n        The driver process only need to call the compute functions of the worker group through RPC\n        to construct the PPO dataflow.\n        The light-weight advantage computation is done on the driver process.\n        \"\"\"\n\n        from verl.utils.tracking import Tracking\n\n        self.logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        self.global_steps = 0\n\n        # load checkpoint and update weights before doing anything\n        self._load_checkpoint()\n        self._fit_update_weights()\n\n        # perform validation before training\n        # currently, we only support validation using the reward_function.\n        if self.config.trainer.get(\"val_before_train\", True):\n            val_metrics = self._validate()\n            assert val_metrics, f\"{val_metrics=}\"\n            pprint(f\"Initial validation metrics: {val_metrics}\")\n            self.logger.log(data=val_metrics, step=self.global_steps)\n            if self.config.trainer.get(\"val_only\", False):\n                return\n\n        if self.config.actor_rollout_ref.rollout.get(\"skip_rollout\", False):\n            rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)\n            rollout_skip.wrap_generate_sequences()\n\n        # add tqdm\n        self.progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc=\"Training Progress\")\n\n        # we start from step 1\n        self.global_steps += 1\n        self.last_val_metrics = None\n        self.max_steps_duration = 0\n\n        self.prev_step_profile = False\n        self.curr_step_profile = (\n            self.global_steps in self.config.global_profiler.steps\n            if self.config.global_profiler.steps is not None\n            else False\n        )\n        self.next_step_profile = False\n\n        # across epoch iterator\n        continuous_iterator = self._create_continuous_iterator()\n        # Start the first asynchronous generation task.\n        batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))\n        while batch_data_future is not None:\n            batch_data_future = await self.fit_step(batch_data_future, continuous_iterator)\n            if self.is_last_step:\n                return\n\n    async def fit_step(self, batch_data_future, continuous_iterator):\n        \"\"\"\n        Single-step training template method. Handles all logic for one training step.\n\n        Flow:\n        1. Pre-step processing -> 2. Get batch -> 3. Generate sequences ->\n        4. Compute reward -> 5. Compute log_prob -> 6. Compute reward ->\n        7. Compute advantage -> 8. Update critic -> 9. Update actor -> 10. Post-step processing\n\n        Args:\n            batch_data_future: batch future\n        \"\"\"\n        self.metrics = {\"training/global_step\": self.global_steps, \"training/epoch\": self.epoch}\n        self.timing_raw = {}\n        # reward message\n        self.future_reward = None\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n\n        self._fit_prepare_step()\n        self._fit_start_profile()\n\n        with marked_timer(\"step\", self.timing_raw):\n            batch, batch_data_future = await self._fit_generate(batch_data_future, continuous_iterator)\n\n            # await asyncio.sleep(0) ensures:\n            # Asynchronous tasks can start executing immediately\n            # The event loop can handle other pending coroutines\n            # Prevents computations in a certain phase from blocking the entire asynchronous workflow\n            #\n            # The purpose here is to ensure that after triggering\n            # `self.async_rollout_manager.generate_sequences(gen_batch_output)`,\n            # the subsequent relevant logic can proceed in a timely manner\n            await asyncio.sleep(0)\n            batch = self._fit_compute_reward(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_compute_log_prob(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_compute_ref_log_prob(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_compute_critic(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_compute_advantage(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_update_critic(batch)\n            await asyncio.sleep(0)\n            batch = self._fit_update_actor(batch)\n            await asyncio.sleep(0)\n            self._fit_update_weights()\n            await asyncio.sleep(0)\n            self._fit_dump_data(batch)\n            await asyncio.sleep(0)\n\n        self._fit_validate()\n        await asyncio.sleep(0)\n        self._fit_save_checkpoint()\n        await asyncio.sleep(0)\n        self._fit_stop_profile()\n        self._fit_collect_metrics(batch)\n        self._fit_torch_memory()\n        self._fit_experimental(batch)\n        self._fit_postprocess_step()\n\n        return batch_data_future\n\n    async def _fit_generate(self, batch_data_future, continuous_iterator):\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n\n        with marked_timer(\"gen\", timing_raw, color=\"red\"):\n            _metrics, _timing_raw, epoch, batch, future_reward = await batch_data_future\n            batch.meta_info[\"temperature\"] = self.config.actor_rollout_ref.rollout.temperature\n            timing_raw.update(batch.meta_info[\"timing\"])\n            timing_raw.update(_timing_raw)\n            metrics.update(_metrics)\n            batch.meta_info.pop(\"timing\", None)\n\n        # sync weights from actor to rollout\n        with marked_timer(\"sync_rollout_weights\", timing_raw, color=\"purple\"):\n            self._fit_update_weights()\n            await self.async_rollout_manager.clear_kv_cache()\n\n        # async next generation\n        if not self.is_last_step:\n            batch_data_future = asyncio.create_task(self._async_gen_next_batch(continuous_iterator))\n            await asyncio.sleep(0)\n        else:\n            batch_data_future = None\n\n        return batch, batch_data_future\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_4_12.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-one-step-off-4-12'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=12\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=2\nsp_size=4\nfsdp_size=2\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\"\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_one_step_off_64-64'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n#  one stepa specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-8}\nNNODES_TRAIN=${NNODES_TRAIN:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntest_freq=20\n\npython -m verl.experimental.one_step_off_policy.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=20 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=400 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\""
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_64_64_ris.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='dapo_qwen2-7B-math_28k_fsdp2_one_step_off_64-64-ris'\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm parameters\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\n# Response length parameters\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 28))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\n# Training parameters\nloss_agg_mode=\"token-mean\"\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=4\nsp_size=4\nfsdp_size=8\n\n#  one stepa specific parameters\nNNODES_ROLLOUT=${NNODES_ROLLOUT:-8}\nNNODES_TRAIN=${NNODES_TRAIN:-8}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\ntest_freq=20\n\n# https://github.com/volcengine/verl/blob/main/docs/algo/rollout_corr.md\n# use decoupled_geo_rs\n#algorithm:\n#  rollout_correction:\n#    rollout_is: null\n#    rollout_is_threshold=null\n#    rollout_rs: seq_mean_k1\n#    rollout_rs_threshold: 0.999_1.001\n#    bypass_mode: false  # Decoupled mode\n\npython -m verl.experimental.one_step_off_policy.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=400 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES_TRAIN}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    rollout.nnodes=\"${NNODES_ROLLOUT}\" \\\n    rollout.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    algorithm.rollout_correction.rollout_is=null \\\n    algorithm.rollout_correction.rollout_is_threshold=null \\\n    algorithm.rollout_correction.rollout_rs=seq_mean_k1 \\\n    algorithm.rollout_correction.rollout_rs_threshold=\"0.999_1.001\" \\\n    algorithm.rollout_correction.bypass_mode=false\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_colocate.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-colocate'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=12\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=2\nsp_size=4\nfsdp_size=2\n\n# reference run wandb: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_4_12.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-sglang-one-step-off-4-12'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=12\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=2\nsp_size=4\nfsdp_size=2\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${ref_offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\"\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_fsdp2_sglang_colocate.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-fsdp2-sglang-colocate'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=12\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=2\nsp_size=4\nfsdp_size=2\n\n# reference run wandb: https://wandb.ai/verl-org/DAPO%20Reproduction%20on%20verl/runs/ow47vvon?nw=nwusertongyuxuan361\n\npython3 -m verl.trainer.main_ppo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.grad_clip=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=${offload} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=\"${NGPUS_PER_NODE}\" \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_4_12.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0527a1-megatron-one-step-off-4-12'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=12\ntrain_prompt_mini_bsz=32\n\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\nref_offload=True\nactor_offload=False\ngen_tp=2\ntrain_tp=2\ntrain_pp=2\n\n# TODO: support dynamic_bsz for megatron\n# actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n# actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n# actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    --config-path=config \\\n    --config-name='one_step_off_ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    critic.strategy=megatron \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${actor_offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${actor_offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${actor_offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.param_offload=${ref_offload} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\"\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/dapo_7b_math_megatron_colocate.sh",
    "content": "#!/usr/bin/env bash\nset -xeuo pipefail\n\nproject_name='DAPO'\nexp_name='DAPO-Qwen2.5-7b-MATH-0519a1-megatron-colocate'\n\nadv_estimator=grpo\n\nuse_kl_in_reward=False\nkl_coef=0.0\nuse_kl_loss=False\nkl_loss_coef=0.0\n\nclip_ratio_low=0.2\nclip_ratio_high=0.28\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 8))\nenable_overlong_buffer=True\noverlong_buffer_len=$((1024 * 4))\noverlong_penalty_factor=1.0\n\nloss_agg_mode=\"token-mean\"\n\ntrain_prompt_bsz=512\nn_resp_per_prompt=16\ntrain_prompt_mini_bsz=32\n\n# Ray\n# RAY_ADDRESS=${RAY_ADDRESS:-\"http://localhost:8265\"}\n# WORKING_DIR=${WORKING_DIR:-\"${PWD}\"}\n# RUNTIME_ENV=${RUNTIME_ENV:-\"${WORKING_DIR}/verl/trainer/runtime_env.yaml\"}\nNNODES=${NNODES:-2}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\n# very important! please modify the max_position_embeddings in config.json to 32768 after downloading from huggingface\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen2.5-Math-7B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/dapo-math-17k.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/aime-2024.parquet\"}\n\n# Algorithm\ntemperature=1.0\ntop_p=1.0\ntop_k=-1 # 0 for HF rollout, -1 for vLLM rollout\nval_top_p=0.7\n\n# Performance Related Parameter\nuse_dynamic_bsz=True\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 2))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) * 3))\noffload=True\ngen_tp=2\ntrain_tp=2\ntrain_pp=2\n\n# TODO: support dynamic_bsz for megatron\n# actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n# actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n# actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n# actor_rollout_ref.rollout.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n\npython3 -m verl.trainer.main_ppo \\\n    --config-path=config \\\n    --config-name='ppo_megatron_trainer.yaml' \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.prompt_key=prompt \\\n    data.truncation='left' \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.train_batch_size=${train_prompt_bsz} \\\n    actor_rollout_ref.rollout.n=${n_resp_per_prompt} \\\n    algorithm.adv_estimator=${adv_estimator} \\\n    algorithm.use_kl_in_reward=${use_kl_in_reward} \\\n    algorithm.kl_ctrl.kl_coef=${kl_coef} \\\n    actor_rollout_ref.actor.strategy=megatron \\\n    critic.strategy=megatron \\\n    actor_rollout_ref.actor.use_kl_loss=${use_kl_loss} \\\n    actor_rollout_ref.actor.kl_loss_coef=${kl_loss_coef} \\\n    actor_rollout_ref.actor.clip_ratio_low=${clip_ratio_low} \\\n    actor_rollout_ref.actor.clip_ratio_high=${clip_ratio_high} \\\n    actor_rollout_ref.actor.clip_ratio_c=10.0 \\\n    +actor_rollout_ref.model.override_config.max_position_embeddings=32768 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=4 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.actor.optim.lr_warmup_steps=10 \\\n    actor_rollout_ref.actor.optim.weight_decay=0.1 \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=${train_prompt_mini_bsz} \\\n    actor_rollout_ref.actor.megatron.param_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.optimizer_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.grad_offload=${offload} \\\n    actor_rollout_ref.actor.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.actor.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.optim.clip_grad=1.0 \\\n    actor_rollout_ref.actor.loss_agg_mode=${loss_agg_mode} \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.80 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=${gen_tp} \\\n    actor_rollout_ref.rollout.enable_chunked_prefill=True \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.temperature=${temperature} \\\n    actor_rollout_ref.rollout.top_p=${top_p} \\\n    actor_rollout_ref.rollout.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.temperature=${temperature} \\\n    actor_rollout_ref.rollout.val_kwargs.top_p=${val_top_p} \\\n    actor_rollout_ref.rollout.val_kwargs.top_k=${top_k} \\\n    actor_rollout_ref.rollout.val_kwargs.do_sample=True \\\n    actor_rollout_ref.rollout.val_kwargs.n=1 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.ref.megatron.pipeline_model_parallel_size=${train_pp} \\\n    actor_rollout_ref.ref.megatron.tensor_model_parallel_size=${train_tp} \\\n    actor_rollout_ref.ref.megatron.param_offload=${offload} \\\n    reward.reward_manager.name=dapo \\\n    +reward.reward_kwargs.overlong_buffer_cfg.enable=${enable_overlong_buffer} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.len=${overlong_buffer_len} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.penalty_factor=${overlong_penalty_factor} \\\n    +reward.reward_kwargs.overlong_buffer_cfg.log=False \\\n    +reward.reward_kwargs.max_resp_len=${max_response_length} \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.n_gpus_per_node=8 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.val_before_train=True \\\n    trainer.test_freq=10 \\\n    trainer.save_freq=-1 \\\n    trainer.total_epochs=10 \\\n    trainer.total_training_steps=100 \\\n    trainer.default_local_dir=\"${CKPTS_DIR}\" \\\n    trainer.resume_mode=auto \\\n    trainer.log_val_generations=10\n"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_2_6.sh",
    "content": "set -x\n\nproject_name='GRPO'\nexp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-one-step-off-2-6'\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-0.6B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/test.parquet\"}\n\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=1152 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=192 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.val_before_train=True \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=2 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" $@"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/grpo_0.6b_gsm8k_fsdp2_sglang_2_6.sh",
    "content": "set -x\n\nproject_name='GRPO'\nexp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-sglang-one-step-off-2-6'\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen3-0.6B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/test.parquet\"}\n\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=1152 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=192 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=sglang \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.val_before_train=True \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=2 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" $@"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/grpo_3b_gsm8k_fsdp2_2_6.sh",
    "content": "set -x\n\nproject_name='GRPO'\nexp_name='GRPO-Qwen3-0.6b-gsm8k-fsdp2-one-step-off-2-6'\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen/Qwen2.5-3B-Instruct\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/train.parquet\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/gsm8k/test.parquet\"}\n\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-8}\n\nn_gpus_rollout=2\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=1152 \\\n    data.max_prompt_length=512 \\\n    data.max_response_length=1024 \\\n    data.filter_overlong_prompts=True \\\n    data.truncation='error' \\\n    actor_rollout_ref.actor.fsdp_config.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=192 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.actor.use_kl_loss=True \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=False \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=False \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=2 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.6 \\\n    actor_rollout_ref.rollout.n=5 \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.rollout.layered_summon=True \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=32 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.use_kl_in_reward=False \\\n    trainer.critic_warmup=0 \\\n    trainer.val_before_train=True \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.save_freq=-1 \\\n    trainer.test_freq=5 \\\n    trainer.total_epochs=2 \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" $@"
  },
  {
    "path": "verl/experimental/one_step_off_policy/shell/grpo_qwen3_8b_gsm8k_fsdp2_8_8_npu.sh",
    "content": "# The script has been validated on the Ascend Atlas 800T A3.\nset -x\n\nexport HCCL_EXEC_TIMEOUT=60000\nexport HCCL_CONNECT_TIMEOUT=7200\n\nproject_name='GRPO'\nexp_name='GRPO-Qwen3-8b-gsm8k-fsdp2-one-step-off-8-8-npu'\n\n# Paths\nRAY_DATA_HOME=${RAY_DATA_HOME:-\"${HOME}/verl\"}\nMODEL_PATH=${MODEL_PATH:-\"${RAY_DATA_HOME}/models/Qwen/Qwen3-8B\"}\nCKPTS_DIR=${CKPTS_DIR:-\"${RAY_DATA_HOME}/ckpts/${project_name}/${exp_name}\"}\nTRAIN_FILE=${TRAIN_FILE:-\"${RAY_DATA_HOME}/data/BytedTsinghua-SIA/DAPO-Math-17k\"}\nTEST_FILE=${TEST_FILE:-\"${RAY_DATA_HOME}/data/BytedTsinghua-SIA/DAPO-Math-17k\"}\n\nNNODES=${NNODES:-1}\nNGPUS_PER_NODE=${NGPUS_PER_NODE:-16}\n\nn_gpus_rollout=8\nn_gpus_training=$((NGPUS_PER_NODE - n_gpus_rollout))\n\nmax_prompt_length=$((1024 * 2))\nmax_response_length=$((1024 * 32))\n\nuse_dynamic_bsz=True\nsp_size=8\nfsdp_size=8\nactor_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\ninfer_ppo_max_token_len=$(((max_prompt_length + max_response_length) / sp_size))\n\npython3 -m verl.experimental.one_step_off_policy.main_ppo \\\n    algorithm.adv_estimator=grpo \\\n    data.train_files=\"${TRAIN_FILE}\" \\\n    data.val_files=\"${TEST_FILE}\" \\\n    data.train_batch_size=32 \\\n    data.max_prompt_length=${max_prompt_length} \\\n    data.max_response_length=${max_response_length} \\\n    data.filter_overlong_prompts=True \\\n    data.filter_overlong_prompts_workers=64 \\\n    data.truncation='error' \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.model.path=\"${MODEL_PATH}\" \\\n    actor_rollout_ref.actor.optim.lr=1e-6 \\\n    actor_rollout_ref.hybrid_engine=False \\\n    actor_rollout_ref.model.use_remove_padding=True \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=32 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.actor.use_kl_loss=False \\\n    actor_rollout_ref.actor.kl_loss_coef=0.001 \\\n    actor_rollout_ref.actor.kl_loss_type=low_var_kl \\\n    actor_rollout_ref.actor.entropy_coeff=0 \\\n    actor_rollout_ref.actor.use_torch_compile=False \\\n    actor_rollout_ref.ref.use_torch_compile=False \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=True \\\n    actor_rollout_ref.actor.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.optimizer_offload=True \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=4 \\\n    actor_rollout_ref.rollout.name=vllm \\\n    actor_rollout_ref.rollout.mode=async \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.8 \\\n    actor_rollout_ref.rollout.max_num_batched_tokens=$((max_prompt_length + max_response_length)) \\\n    actor_rollout_ref.rollout.n=8 \\\n    actor_rollout_ref.rollout.enforce_eager=True \\\n    actor_rollout_ref.rollout.load_format=safetensors \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu=1 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    actor_rollout_ref.actor.fsdp_config.fsdp_size=${fsdp_size} \\\n    actor_rollout_ref.actor.use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.actor.ppo_max_token_len_per_gpu=${actor_ppo_max_token_len} \\\n    actor_rollout_ref.actor.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.ref.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    actor_rollout_ref.ref.log_prob_max_token_len_per_gpu=${infer_ppo_max_token_len} \\\n    actor_rollout_ref.ref.ulysses_sequence_parallel_size=${sp_size} \\\n    actor_rollout_ref.rollout.log_prob_use_dynamic_bsz=${use_dynamic_bsz} \\\n    algorithm.use_kl_in_reward=False \\\n    actor_rollout_ref.nccl_timeout=14400 \\\n    trainer.critic_warmup=0 \\\n    trainer.val_before_train=False \\\n    trainer.logger=['console','tensorboard'] \\\n    trainer.project_name=\"${project_name}\" \\\n    trainer.experiment_name=\"${exp_name}\" \\\n    trainer.default_local_dir=${CKPTS_DIR} \\\n    trainer.save_freq=10 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=15 \\\n    trainer.resume_mode=auto \\\n    trainer.nnodes=\"${NNODES}\" \\\n    trainer.n_gpus_per_node=\"${n_gpus_training}\" \\\n    rollout.nnodes=\"${NNODES}\" \\\n    rollout.n_gpus_per_node=\"${n_gpus_rollout}\" $@"
  },
  {
    "path": "verl/experimental/reward_loop/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .reward_loop import RewardLoopManager, RewardLoopWorker, migrate_legacy_reward_impl\nfrom .reward_model import RewardModelManager\n\n__all__ = [\"RewardModelManager\", \"RewardLoopWorker\", \"RewardLoopManager\", \"migrate_legacy_reward_impl\"]\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_loop.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nimport os\n\nimport aiohttp\nimport numpy as np\nimport ray\nimport torch\nfrom omegaconf import DictConfig, open_dict\nfrom tensordict import TensorDict\n\nfrom verl.protocol import DataProto\nfrom verl.single_controller.ray.base import RayResourcePool\nfrom verl.trainer.ppo.reward import load_reward_manager\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.fs import copy_to_local\n\nfrom .reward_model import RewardModelManager\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef migrate_legacy_reward_impl(config):\n    \"\"\"\n    Migrate the legacy reward model implementation to the new one.\n    \"\"\"\n    # 1. reward workers migration\n    # config.reward_model.num_workers -> config.reward.num_workers\n    if config.reward_model.num_workers is not None:\n        config.reward.num_workers = config.reward_model.num_workers\n\n    # 2. reward manager migration\n    # config.reward_model.reward_manager -> config.reward.reward_manager\n    if config.reward_model.reward_manager is not None:\n        config.reward.reward_manager.name = config.reward_model.reward_manager\n    if config.reward_model.reward_loop_source is not None:\n        config.reward.reward_manager.source = config.reward_model.reward_loop_source\n        config.reward.reward_manager.module.path = config.reward_model.reward_loop_module_path\n        config.reward.reward_manager.module.name = config.reward_model.reward_loop_class_name\n\n    # 3. custom reward function migration\n    # config.custom_reward_function -> config.reward.custom_reward_function\n    if not all(v is None for v in config.custom_reward_function.values()):\n        config.reward.custom_reward_function = config.custom_reward_function\n\n    # 4. reward model migration\n    # config.reward_model -> config.reward.reward_model\n    for key in [\"enable\", \"enable_resource_pool\", \"n_gpus_per_node\", \"nnodes\"]:\n        if config.reward_model.get(key) is not None:\n            config.reward.reward_model[key] = config.reward_model[key]\n    if config.reward_model.model.path is not None:\n        config.reward.reward_model.model_path = config.reward_model.model.path\n    # config.reward_model.reward_kwargs -> config.reward.reward_kwargs (for dapo algo)\n    if config.reward_model.get(\"reward_kwargs\") is not None:\n        with open_dict(config.reward):\n            config.reward[\"reward_kwargs\"] = config.reward_model[\"reward_kwargs\"]\n    # config.reward_model.rollout -> config.reward.reward_model.rollout\n    legacy_rollout = config.reward_model.rollout\n    for key in legacy_rollout.keys():\n        if legacy_rollout[key] is not None:\n            config.reward.reward_model.rollout[key] = legacy_rollout[key]\n\n    # 5. sandbox_fusion migration\n    # config.sandbox_fusion -> reward.sandbox_fusion\n    if not all(v is None for v in config.sandbox_fusion.values()):\n        config.reward.sandbox_fusion = config.sandbox_fusion\n\n    # 6. delete legacy config from configs\n    with open_dict(config):\n        del config.reward_model\n        del config.custom_reward_function\n        del config.sandbox_fusion\n\n    return config\n\n\nclass RewardLoopWorker:\n    \"\"\"\n    RewardLoopWork can tackle reward computation:\n    (1) rule-based reward computation\n    (2) reward model-based reward computation (both disrm and genrm)\n    (3) high-flexible user-customized reward function (can access rm by posting requests to reward_model_router)\n\n    Reward Computation Logic:\n    - if user-customized reward function is provided:\n        -> directly use user-customized reward function\n    - if user-customized reward function is not provided:\n        -> rm is not enabled: use default rule-based reward function\n        -> rm is disrm: compute reward score using disrm\n        -> rm is genrm: raise error (user-costomized reward func must be provided)\n    \"\"\"\n\n    def __init__(self, config: DictConfig, reward_router_address: str = None):\n        \"\"\"\n        Args:\n            config: DictConfig, the config for reward loop worker.\n            reward_router_address: str, the address of reward router.\n        \"\"\"\n        self.config = config\n        self.reward_router_address = reward_router_address\n        self._init_reward_fn()\n\n    def _init_reward_fn(self):\n        input_tokenizer_local_path = copy_to_local(self.config.actor_rollout_ref.model.path)\n        self.input_tokenizer = hf_tokenizer(input_tokenizer_local_path, trust_remote_code=True)\n        self.reward_model_tokenizer = None\n        if self.config.reward.reward_model.enable:\n            reward_model_tokenizer_local_path = copy_to_local(self.config.reward.reward_model.model_path)\n            self.reward_model_tokenizer = hf_tokenizer(reward_model_tokenizer_local_path, trust_remote_code=True)\n\n        self.reward_manager = load_reward_manager(\n            self.config,\n            self.input_tokenizer,\n            reward_router_address=self.reward_router_address,\n            reward_model_tokenizer=self.reward_model_tokenizer,\n        )\n\n    async def compute_score_batch(self, data: DataProto) -> list[dict]:\n        tasks = []\n        for i in range(len(data)):\n            tasks.append(asyncio.create_task(self.compute_score(data[i : i + 1])))\n        outputs = await asyncio.gather(*tasks)\n        return outputs\n\n    async def compute_score(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"RewardLoopWorker only support single data item\"\n        if self.config.reward.custom_reward_function.path is not None:\n            # directly use user-customized reward function\n            return await self.reward_manager.run_single(data)\n        else:\n            if self.config.reward.reward_model.enable:\n                # we assume the rm is disrm\n                # genrm must set custom_reward_function\n                return await self.compute_score_disrm(data)\n            else:\n                return await self.reward_manager.run_single(data)\n\n    async def _post_request(self, payload: dict, endpoint: str, max_retries: int = 16):\n        url = f\"http://{self.reward_router_address}/{endpoint}\"\n        last_exception = None\n        for attempt in range(max_retries):\n            try:\n                # It's safer to have a timeout instead of None, which can hang indefinitely.\n                timeout = aiohttp.ClientTimeout(total=None)\n                async with aiohttp.ClientSession(timeout=timeout) as session:\n                    async with session.post(url, json=payload) as resp:\n                        resp.raise_for_status()\n                        return await resp.json()\n            except aiohttp.ClientResponseError as e:\n                # Do not retry on 4xx client errors, but retry on 5xx server errors.\n                if 400 <= e.status < 500:\n                    logger.error(f\"Request to {url} failed with client error HTTP {e.status}: {e}. Not retrying.\")\n                    raise\n                last_exception = e\n                logger.warning(\n                    f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with HTTP {e.status}: {e}. \"\n                    \"Retrying...\"\n                )\n            except (asyncio.TimeoutError, aiohttp.ClientConnectorError) as e:\n                last_exception = e\n                logger.warning(f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed: {e}. Retrying...\")\n            except Exception as e:\n                last_exception = e\n                logger.warning(\n                    f\"[Attempt {attempt + 1}/{max_retries}] Request to {url} failed with unexpected error: {e}. \"\n                    \"Retrying...\"\n                )\n\n            if attempt < max_retries - 1:\n                # Using exponential backoff is generally better than a fixed sleep.\n                backoff_seconds = 2**attempt\n                await asyncio.sleep(min(backoff_seconds, 30))\n\n        logger.error(f\"Max retries ({max_retries}) reached for request to {url}.\")\n        if last_exception:\n            raise last_exception\n\n    async def _preprocess_reward_inputs(self, data: DataProto) -> str:\n        assert len(data) == 1, \"RewardLoopWorker only support single data item\"\n        data_item = data[0]\n        assert \"raw_prompt\" in data_item.non_tensor_batch\n\n        # extract raw prompt\n        chat: list = list(data_item.non_tensor_batch[\"raw_prompt\"])\n\n        # extract response\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        # decode\n        rollout_response = self.input_tokenizer.decode(valid_response_ids)\n        # remove bos and eos\n        rollout_response = rollout_response.replace(self.input_tokenizer.eos_token, \"\")\n\n        chat.append({\"role\": \"assistant\", \"content\": rollout_response})\n\n        rm_prompt = self.reward_model_tokenizer.apply_chat_template(\n            chat,\n            add_generation_prompt=False,\n            tokenize=False,\n        )\n\n        # llama tokenizer will add bos token by default\n        # will be removed in vllm >= 0.11.2, where we can add \"add_special_tokens\" = False\n        if self.reward_model_tokenizer.bos_token is not None and rm_prompt.startswith(\n            self.reward_model_tokenizer.bos_token\n        ):\n            rm_prompt = rm_prompt[len(self.reward_model_tokenizer.bos_token) :]\n\n        return rm_prompt\n\n    async def compute_score_disrm(self, data: DataProto) -> dict:\n        disrm_prompt = await self._preprocess_reward_inputs(data)\n        engine_name = self.config.reward.reward_model.rollout.name\n        model_name = self.config.reward.reward_model.model_path\n        if engine_name == \"vllm\":\n            payloads = {\n                \"model\": model_name,\n                \"input\": disrm_prompt,\n                \"use_activation\": False,\n            }\n            output = await self._post_request(payloads, \"classify\")\n            rm_score = output[\"data\"][-1][\"probs\"][-1]\n        elif engine_name == \"sglang\":\n            payloads = {\n                \"model\": model_name,\n                \"input\": disrm_prompt,\n            }\n            output = await self._post_request(payloads, \"v1/embeddings\")\n            rm_score = output[\"data\"][-1][\"embedding\"][-1]\n        elif engine_name == \"trtllm\":\n            # TODO: remove this once TRT-LLM switches to TorchSampler\n            raise ValueError(\"TensorRT-LLM backend does not support reward models currently.\")\n\n            payloads = {\n                \"model\": model_name,\n                \"prompt\": disrm_prompt,\n                \"return_context_logits\": True,\n            }\n            output = await self._post_request(payloads, \"v1/completions\")\n            rm_score = output[\"choices\"][0][\"context_logits\"]\n            assert isinstance(rm_score, list) and len(rm_score) > 0, (\n                \"TensorRT-LLM OpenAI server response for reward score is not in the expected format.\"\n            )\n\n            rm_score = float(rm_score[0][0])\n            logger.debug(f\"rm score: {rm_score}\")\n        else:\n            raise NotImplementedError(f\"RewardLoopManager does not support {engine_name}\")\n\n        return {\"reward_score\": rm_score}\n\n\nclass RewardLoopManager:\n    \"\"\"\n    RewardLoopManager run in single controller.\n    This class will create reward loop workers and manage them.\n    \"\"\"\n\n    def __init__(self, config: DictConfig, rm_resource_pool: RayResourcePool = None):\n        self.config = config\n        if self.config.reward.reward_model.enable:\n            self.reward_model_manager = RewardModelManager(config.reward.reward_model, rm_resource_pool)\n            self.reward_router_address = self.reward_model_manager.get_router_address()\n        else:\n            self.reward_model_manager = None\n            self.reward_router_address = None\n\n        self.reward_loop_workers_class = ray.remote(RewardLoopWorker)\n        self._init_reward_loop_workers()\n\n    def _init_reward_loop_workers(self):\n        self.reward_loop_workers = []\n        num_workers = self.config.reward.num_workers\n        node_ids = [node[\"NodeID\"] for node in ray.nodes() if node[\"Alive\"] and node[\"Resources\"].get(\"CPU\", 0) > 0]\n\n        for i in range(num_workers):\n            # Round-robin scheduling over the all nodes\n            node_id = node_ids[i % len(node_ids)]\n\n            self.reward_loop_workers.append(\n                self.reward_loop_workers_class.options(\n                    name=f\"reward_loop_worker_{i}\",\n                    scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                        node_id=node_id,\n                        soft=True,\n                    ),\n                ).remote(self.config, self.reward_router_address)\n            )\n\n    def compute_rm_score(self, data: DataProto) -> DataProto:\n        if self.reward_model_manager is not None:\n            self.reward_model_manager.wake_up()\n\n        chunks = data.chunk(len(self.reward_loop_workers))\n        outputs = ray.get(\n            [\n                worker.compute_score_batch.remote(chunk)\n                for worker, chunk in zip(self.reward_loop_workers, chunks, strict=True)\n            ]\n        )\n        outputs_flat = [item for sublist in outputs for item in sublist]\n\n        # compute rm score\n        scores = [item[\"reward_score\"] for item in outputs_flat]\n        prompt_length = data.batch[\"prompts\"].size(1)\n        valid_response_length = data.batch[\"attention_mask\"][:, prompt_length:].sum(dim=1)\n        rm_scores = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n        rm_scores[torch.arange(rm_scores.size(0)), valid_response_length - 1] = torch.tensor(\n            scores, dtype=torch.float32\n        )\n        batch = TensorDict({\"rm_scores\": rm_scores}, batch_size=len(data))\n\n        reward_extra_infos = [output.get(\"reward_extra_info\", {}) for output in outputs_flat]\n        reward_extra_keys = list(reward_extra_infos[0].keys())\n        non_tensor_batch = {}\n        for key in reward_extra_keys:\n            non_tensor_batch[key] = np.array([info[key] for info in reward_extra_infos])\n\n        if self.reward_model_manager is not None:\n            self.reward_model_manager.sleep()\n\n        return DataProto(\n            batch=batch, non_tensor_batch=non_tensor_batch, meta_info={\"reward_extra_keys\": reward_extra_keys}\n        )\n\n    def _run_all(self, tasks: list[asyncio.Task]):\n        async def run_all():\n            return await asyncio.gather(*tasks)\n\n        return asyncio.run(run_all())\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .registry import get_reward_manager_cls, register  # noqa: I001\nfrom .dapo import DAPORewardManager\nfrom .gdpo import GDPORewardManager\nfrom .naive import NaiveRewardManager\nfrom .limited import RateLimitedRewardManager\nfrom .remote import RemoteRewardManager\n\n__all__ = [\n    \"DAPORewardManager\",\n    \"GDPORewardManager\",\n    \"NaiveRewardManager\",\n    \"RateLimitedRewardManager\",\n    \"RemoteRewardManager\",\n    \"register\",\n    \"get_reward_manager_cls\",\n]\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Callable\n\nfrom omegaconf import DictConfig\nfrom transformers import AutoTokenizer\n\nfrom verl import DataProto\nfrom verl.utils.ray_utils import get_event_loop\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nRawRewardFn = Callable[..., Any] | None\n\n\nclass RewardManagerBase(ABC):\n    _class_initialized = False\n\n    def __init__(self, config: DictConfig, tokenizer: AutoTokenizer, compute_score: RawRewardFn):\n        \"\"\"Initialize reward manager.\n\n        Args:\n            config (DictConfig): YAML config.\n            tokenizer (AutoTokenizer): Tokenizer for tokenize messages.\n        \"\"\"\n        self.config = config\n        self.tokenizer = tokenizer\n        self.compute_score = compute_score\n        self.loop = get_event_loop()\n        self.init_class(config, tokenizer)\n\n    @classmethod\n    def init_class(cls, config: DictConfig, tokenizer: AutoTokenizer):\n        \"\"\"Initialize class state shared across all instances.\"\"\"\n        if cls._class_initialized:\n            return\n        cls._class_initialized = True\n\n    @abstractmethod\n    async def run_single(self, data: DataProto):\n        raise NotImplementedError\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/dapo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager import register\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\nfrom verl.utils.reward_score import default_compute_score\n\n\n@register(\"dapo\")\nclass DAPORewardManager(RewardManagerBase):\n    \"\"\"DAPO Reward Manager.\"\"\"\n\n    def __init__(self, config, tokenizer, compute_score, reward_router_address=None, reward_model_tokenizer=None):\n        super().__init__(config, tokenizer, compute_score)\n        self.compute_score = compute_score or default_compute_score\n        self.is_async_reward_score = inspect.iscoroutinefunction(self.compute_score)\n\n        # DAPO Reward Config\n        overlong_buffer_cfg = config.reward.get(\"reward_kwargs\", {}).get(\"overlong_buffer_cfg\", None)\n        self.overlong_buffer_cfg = overlong_buffer_cfg\n        self.max_resp_len = config.reward.get(\"reward_kwargs\", {}).get(\"max_resp_len\", None)\n        self.reward_router_address = reward_router_address\n        self.reward_model_tokenizer = reward_model_tokenizer\n\n        if self.overlong_buffer_cfg is not None:\n            assert self.max_resp_len is not None, (\n                f\"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None\"\n            )\n            assert self.max_resp_len >= self.overlong_buffer_cfg.len, (\n                \"max_resp_len must be larger than overlong_buffer.len\"\n            )\n            assert not self.overlong_buffer_cfg.enable or self.overlong_buffer_cfg.len > 0, (\n                \"overlong_buffer.len must be positive when overlong penalty is enabled,\"\n                f\"but got {self.overlong_buffer_cfg.len}.\"\n                \"To disable the overlong penalty, set overlong_buffer.enable = False\"\n            )\n\n    async def run_single(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"Only support single data item\"\n        data_item = data[0]\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        data_source = data_item.non_tensor_batch[\"data_source\"]\n        ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n        extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n\n        response_str = await self.loop.run_in_executor(\n            None, lambda: self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n        )\n        extra_reward_kwargs = (\n            {\n                \"reward_router_address\": self.reward_router_address,\n                \"reward_model_tokenizer\": self.reward_model_tokenizer,\n            }\n            if self.reward_router_address is not None\n            else {}\n        )\n        if self.is_async_reward_score:\n            result = await self.compute_score(\n                data_source=data_source,\n                solution_str=response_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n                **extra_reward_kwargs,\n            )\n        else:\n            result = await self.loop.run_in_executor(\n                None,\n                lambda: self.compute_score(\n                    data_source=data_source,\n                    solution_str=response_str,\n                    ground_truth=ground_truth,\n                    extra_info=extra_info,\n                    **extra_reward_kwargs,\n                ),\n            )\n\n        reward_extra_info = {}\n\n        score: float\n        if isinstance(result, dict):\n            score = result[\"score\"]\n            for key, value in result.items():\n                reward_extra_info[key] = value\n        else:\n            score = result\n            reward_extra_info[\"acc\"] = score\n\n        reward = score\n\n        if self.overlong_buffer_cfg is not None and self.overlong_buffer_cfg.enable:\n            overlong_buffer_len = self.overlong_buffer_cfg.len\n            expected_len = self.max_resp_len - overlong_buffer_len\n            exceed_len = valid_response_length - expected_len\n            overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor\n            overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)\n            reward += overlong_reward\n            if self.overlong_buffer_cfg.log:\n                reward_extra_info[\"overlong_reward\"] = overlong_reward\n                reward_extra_info[\"overlong\"] = overlong_reward < 0\n\n        return {\"reward_score\": reward, \"reward_extra_info\": reward_extra_info}\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/gdpo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager import register\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\nfrom verl.utils.reward_score import default_compute_score\n\n\n@register(\"gdpo\")\nclass GDPORewardManager(RewardManagerBase):\n    \"\"\"GDPO Reward Manager.\"\"\"\n\n    def __init__(self, config, tokenizer, compute_score, reward_router_address=None, reward_model_tokenizer=None):\n        super().__init__(config, tokenizer, compute_score)\n        self.compute_score = compute_score or default_compute_score\n        self.is_async_reward_score = inspect.iscoroutinefunction(self.compute_score)\n\n        self.reward_router_address = reward_router_address\n        self.reward_model_tokenizer = reward_model_tokenizer\n\n    async def run_single(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"Only support single data item\"\n        data_item = data[0]\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        data_source = data_item.non_tensor_batch[\"data_source\"]\n        ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n        extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n        extra_info[\"experiment_name\"] = self.config.trainer.experiment_name\n\n        response_str = await self.loop.run_in_executor(\n            None, lambda: self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n        )\n        extra_reward_kwargs = (\n            {\n                \"reward_router_address\": self.reward_router_address,\n                \"reward_model_tokenizer\": self.reward_model_tokenizer,\n            }\n            if self.reward_router_address is not None\n            else {}\n        )\n        if self.is_async_reward_score:\n            result = await self.compute_score(\n                data_source=data_source,\n                solution_str=response_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n                **extra_reward_kwargs,\n            )\n        else:\n            result = await self.loop.run_in_executor(\n                None,\n                lambda: self.compute_score(\n                    data_source=data_source,\n                    solution_str=response_str,\n                    ground_truth=ground_truth,\n                    extra_info=extra_info,\n                    **extra_reward_kwargs,\n                ),\n            )\n\n        reward_extra_info = {}\n\n        score: float\n        if isinstance(result, dict):\n            score = result[\"score\"]\n            for key, value in result.items():\n                reward_extra_info[key] = value\n        else:\n            score = result\n            reward_extra_info[\"acc\"] = score\n\n        reward = score\n\n        return {\"reward_score\": reward, \"reward_extra_info\": reward_extra_info}\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/limited.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport inspect\nimport logging\n\nfrom omegaconf import DictConfig\nfrom transformers import AutoTokenizer\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager import register as register_manager\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.reward_score import default_compute_score\nfrom verl.workers.reward_manager import register as register_manager_legacy\n\nlogger = logging.getLogger(__file__)\n\n\nclass AsyncTokenBucket:\n    \"\"\"Async token bucket for rate limiting with variable token consumption.\n\n    The token bucket algorithm is a classic rate limiting technique that allows\n    for burst traffic while maintaining an average rate limit. This implementation\n    is async-first and thread-safe, designed for use in concurrent environments.\n\n    The bucket starts full and refills at a constant rate (rate_limit tokens/second).\n    When tokens are acquired, they are consumed from the bucket. If insufficient\n    tokens are available, the acquire() method will sleep until enough tokens\n    have been refilled.\n\n    This implementation supports variable token consumption, making it suitable\n    for rate limiting based on request size (e.g., API token usage).\n\n    Args:\n        rate_limit (float): The rate at which tokens are added to the bucket,\n            in tokens per second. For example, rate_limit=10.0 means 10 tokens\n            are added per second (or 600 per minute).\n        max_tokens (float, optional): The maximum capacity of the token bucket.\n            Defaults to rate_limit if not specified. This value determines the\n            maximum burst size allowed.\n\n    Attributes:\n        rate_limit (float): Tokens added per second.\n        max_tokens (float): Maximum bucket capacity.\n        tokens (float): Current number of available tokens.\n        last_update (float | None): Timestamp of last token update (from event loop).\n        lock (asyncio.Lock): Async lock for thread-safe token operations.\n\n    Example:\n        >>> # Limit to 60 requests per minute (1 request per second)\n        >>> rpm_limiter = AsyncTokenBucket(rate_limit=1.0, max_tokens=1.0)\n        >>> await rpm_limiter.acquire(1.0)  # Consumes 1 token\n        >>>\n        >>> # Limit to 10000 tokens per minute (~166.67 tokens per second)\n        >>> tpm_limiter = AsyncTokenBucket(rate_limit=166.67, max_tokens=166.67)\n        >>> await tpm_limiter.acquire(100.0)  # Consumes 100 tokens\n\n    Thread Safety:\n        All operations are protected by an asyncio.Lock, making this class safe\n        for concurrent use across multiple coroutines.\n\n    Algorithm Details:\n        1. On each acquire(), calculate elapsed time since last update\n        2. Refill tokens: tokens += elapsed * rate_limit (capped at max_tokens)\n        3. If tokens >= num_tokens: consume tokens and return\n        4. Otherwise: calculate wait_time = tokens_needed / rate_limit, then sleep\n        5. Retry after sleep (loop back to step 1)\n    \"\"\"\n\n    def __init__(self, rate_limit: float, max_tokens: float = None):\n        self.rate_limit = rate_limit\n        self.max_tokens = max_tokens or rate_limit\n        self.tokens = self.max_tokens\n        self.last_update = None\n        self.lock = asyncio.Lock()\n\n    async def acquire(self, num_tokens: float = 1.0) -> None:\n        \"\"\"Acquire tokens from the bucket, waiting if necessary.\n\n        This method will block (using asyncio.sleep) until sufficient tokens\n        are available. It automatically refills tokens based on elapsed time\n        and the configured rate_limit.\n\n        For requests exceeding max_tokens, the method will wait for enough time\n        to accumulate the required tokens at the configured rate_limit, allowing\n        tokens to temporarily go negative.\n\n        Args:\n            num_tokens (float): Number of tokens to consume. Defaults to 1.0.\n                Can be fractional for fine-grained rate limiting.\n\n        Returns:\n            None: Returns when tokens have been successfully acquired.\n\n        Raises:\n            No exceptions are raised. This method will wait indefinitely until\n            tokens become available.\n\n        Example:\n            >>> bucket = AsyncTokenBucket(rate_limit=10.0)\n            >>> await bucket.acquire(5.0)  # Acquire 5 tokens\n            >>> await bucket.acquire(1.0)  # Acquire 1 more token\n\n        Implementation Notes:\n            - Uses event loop's time() for high-precision timestamps\n            - Lock is released during sleep to allow other coroutines to proceed\n            - Tokens are refilled continuously based on elapsed time\n            - For requests > max_tokens, allows temporary negative balance\n        \"\"\"\n        # Handle requests larger than max_tokens separately\n        if num_tokens > self.max_tokens:\n            wait_time = 0.0\n            async with self.lock:\n                loop = get_event_loop()\n                now = loop.time()\n                if self.last_update is None:\n                    self.last_update = now\n\n                elapsed = now - self.last_update\n                new_tokens = elapsed * self.rate_limit\n                self.tokens = min(self.max_tokens, self.tokens + new_tokens)\n\n                tokens_needed = num_tokens - self.tokens\n                if tokens_needed > 0:\n                    wait_time = tokens_needed / self.rate_limit\n\n                self.tokens -= num_tokens\n                self.last_update = now\n\n            if wait_time > 0:\n                await asyncio.sleep(wait_time)\n            return\n\n        # Standard case: request <= max_tokens\n        while True:\n            wait_time = 0.0\n            async with self.lock:\n                loop = get_event_loop()\n                now = loop.time()\n                if self.last_update is None:\n                    self.last_update = now\n\n                elapsed = now - self.last_update\n                new_tokens = elapsed * self.rate_limit\n                self.tokens = min(self.max_tokens, self.tokens + new_tokens)\n                self.last_update = now\n\n                if self.tokens >= num_tokens:\n                    self.tokens -= num_tokens\n                    return\n\n                tokens_needed = num_tokens - self.tokens\n                wait_time = tokens_needed / self.rate_limit\n\n            if wait_time > 0:\n                await asyncio.sleep(wait_time)\n\n\n@register_manager(\"rate_limited\")\n@register_manager_legacy(\"rate_limited\")\nclass RateLimitedRewardManager(RewardManagerBase):\n    \"\"\"Reward manager with rate limiting for API-based reward functions.\n\n    This manager implements a sophisticated three-layer rate limiting system\n    designed for LLM-as-judge scenarios where reward computation involves\n    external API calls (e.g., OpenAI, Anthropic, Claude) that have rate limits.\n\n    The three layers of rate limiting are:\n        1. **Concurrency limiting** (max_concurrent): Limits the number of\n           simultaneous API requests using asyncio.Semaphore. This prevents\n           overwhelming the API with too many parallel connections.\n\n        2. **Request rate limiting** (max_rpm): Limits requests per minute\n           using AsyncTokenBucket. Each request consumes 1 token. Useful for\n           APIs with per-minute request quotas.\n\n        3. **Token rate limiting** (max_tpm): Limits tokens per minute using\n           AsyncTokenBucket. Each request consumes estimated_tokens_per_request\n           tokens. Essential for APIs that bill or limit based on token usage\n           (e.g., GPT-4 API).\n\n    All rate limiters are **global class-level resources**, meaning they are\n    shared across all instances of this manager. This ensures that rate limits\n    are enforced consistently across multiple workers in distributed training.\n\n    Rate Limiting Flow:\n        When processing a reward request, the manager:\n        1. Acquires RPM token (if rpm_limiter enabled)\n        2. Acquires TPM tokens (if tpm_limiter enabled)\n        3. Acquires concurrency semaphore\n        4. Executes reward computation with timeout\n        5. Releases concurrency semaphore\n        6. Tokens are automatically refilled by the token buckets\n\n    Args:\n        config (DictConfig): Configuration object containing reward_model settings:\n            - max_concurrent (int): Max parallel requests. Default: 1\n            - max_rpm (int | None): Max requests per minute. Default: None (unlimited)\n            - max_tpm (int | None): Max tokens per minute. Default: None (unlimited)\n            - estimated_tokens_per_request (int): Estimated tokens per request for\n              TPM limiting. Default: 2000\n            - timeout (float): Timeout for reward computation in seconds. Default: 300\n        tokenizer (AutoTokenizer): HuggingFace tokenizer for decoding responses.\n        compute_score (callable, optional): Custom reward scoring function. Can be\n            sync or async. Defaults to default_compute_score.\n        reward_router_address (str | None): Address for reward router service.\n        reward_model_tokenizer (AutoTokenizer | None): Optional tokenizer for reward model.\n\n    Class Attributes (Global State):\n        _semaphore (asyncio.Semaphore): Global concurrency limiter\n        _max_concurrent (int): Max concurrent requests\n        _rpm_limiter (AsyncTokenBucket | None): Request rate limiter\n        _max_rpm (int | None): Max requests per minute\n        _tpm_limiter (AsyncTokenBucket | None): Token rate limiter\n        _max_tpm (int | None): Max tokens per minute\n        _estimated_tokens_per_request (int): Estimated tokens per request\n        _class_initialized (bool): Whether class has been initialized\n\n    Example Configuration:\n        >>> config = DictConfig({\n        ...     \"reward\": {\n        ...         \"max_concurrent\": 10,      # 10 parallel requests\n        ...         \"max_rpm\": 500,            # 500 requests/minute\n        ...         \"max_tpm\": 100000,         # 100k tokens/minute\n        ...         \"estimated_tokens_per_request\": 2000,\n        ...         \"timeout\": 60.0,\n        ...     }\n        ... })\n        >>> manager = RateLimitedRewardManager(config, tokenizer)\n\n    Thread Safety:\n        This class is designed for concurrent use. All rate limiting resources\n        are protected by asyncio primitives (Lock, Semaphore).\n\n    See Also:\n        - AsyncTokenBucket: Token bucket implementation for rate limiting\n        - RewardManagerBase: Base class for reward managers\n        - verl.utils.reward_score.default_compute_score: Default scoring function\n    \"\"\"\n\n    # Class-level state for global rate limiting\n    _semaphore = None\n    _max_concurrent = None\n    _rpm_limiter = None\n    _max_rpm = None\n    _tpm_limiter = None\n    _max_tpm = None\n    _estimated_tokens_per_request = None\n    _class_initialized = False\n\n    @classmethod\n    def init_class(cls, config: DictConfig, tokenizer: AutoTokenizer):\n        \"\"\"Initialize class state shared across all instances.\"\"\"\n        # Check if already initialized before calling parent.\n        #\n        # NOTE: This class owns a *global*, class-level set of rate limiters. Once the class has been\n        # initialized, subsequent instantiations cannot change the shared limiters. This is by design,\n        # but it can be surprising (and dangerous) when the first initialization happens with default\n        # values (often \"unlimited\") and later code tries to apply limits.\n        if cls._class_initialized:\n            rm_cfg = config.get(\"reward\") or {}\n            incoming_max_rpm = rm_cfg.get(\"max_rpm\", None)\n            incoming_max_tpm = rm_cfg.get(\"max_tpm\", None)\n\n            # Warn when a caller is trying to change the global RPM/TPM limits after initialization.\n            # This commonly happens if the first instance was created without a config (legacy signature),\n            # which initializes the global limiters to their defaults and locks them in.\n            if (incoming_max_rpm != cls._max_rpm) or (incoming_max_tpm != cls._max_tpm):\n                if (\n                    incoming_max_rpm is not None\n                    or incoming_max_tpm is not None\n                    or cls._max_rpm is not None\n                    or cls._max_tpm is not None\n                ):\n                    logger.warning(\n                        \"RateLimitedRewardManager has already been initialized and its rate limiters are shared \"\n                        \"globally across instances. The incoming (max_rpm/max_tpm) settings will be ignored. \"\n                        \"This can lead to unexpected behavior (e.g., exceeding API rate limits) if the first \"\n                        \"initialization used defaults (often unlimited). \"\n                        f\"Existing: max_rpm={cls._max_rpm}, max_tpm={cls._max_tpm}. \"\n                        f\"Incoming: max_rpm={incoming_max_rpm}, max_tpm={incoming_max_tpm}. \"\n                        \"To apply different limits, ensure the first RateLimitedRewardManager created in this \"\n                        \"process uses the desired configuration (or restart/reset the process).\"\n                    )\n            return\n\n        super().init_class(config, tokenizer)\n\n        rm_cfg = config.get(\"reward\") or {}\n\n        # Concurrency limiter\n        cls._max_concurrent = rm_cfg.get(\"max_concurrent\", 1)\n        cls._semaphore = asyncio.Semaphore(cls._max_concurrent)\n\n        # Request rate limiter (RPM)\n        cls._max_rpm = rm_cfg.get(\"max_rpm\", None)\n        if cls._max_rpm is not None:\n            requests_per_second = cls._max_rpm / 60.0\n            cls._rpm_limiter = AsyncTokenBucket(rate_limit=requests_per_second, max_tokens=requests_per_second)\n        else:\n            cls._rpm_limiter = None\n\n        # Token rate limiter (TPM)\n        cls._max_tpm = rm_cfg.get(\"max_tpm\", None)\n        cls._estimated_tokens_per_request = rm_cfg.get(\"estimated_tokens_per_request\", 2000)\n        if cls._max_tpm is not None:\n            tokens_per_second = cls._max_tpm / 60.0\n            cls._tpm_limiter = AsyncTokenBucket(rate_limit=tokens_per_second, max_tokens=tokens_per_second)\n        else:\n            cls._tpm_limiter = None\n\n        log_msg = \"Rate limiting configuration:\\n\"\n        log_msg += f\"  - Concurrency limit: {cls._max_concurrent}\\n\"\n        if cls._max_rpm is not None:\n            log_msg += f\"  - Request rate limit: {cls._max_rpm} RPM ({cls._max_rpm / 60.0:.2f} RPS)\\n\"\n        else:\n            log_msg += \"  - Request rate limit: unlimited\\n\"\n        if cls._max_tpm is not None:\n            log_msg += f\"  - Token rate limit: {cls._max_tpm} TPM ({cls._max_tpm / 60.0:.2f} TPS)\\n\"\n            log_msg += f\"  - Estimated tokens per request: {cls._estimated_tokens_per_request}\\n\"\n        else:\n            log_msg += \"  - Token rate limit: unlimited\\n\"\n        log_msg += \"All limiters are shared globally across all workers.\"\n        logger.info(log_msg)\n\n        cls._class_initialized = True\n\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        compute_score,\n        reward_router_address=None,\n        reward_model_tokenizer=None,\n        # Legacy (AbstractRewardManager) kwargs for compatibility. Not used.\n        num_examine: int | None = None,\n        reward_fn_key: str | None = None,\n        **kwargs,\n    ):\n        # When called via the legacy AbstractRewardManager signature, `config` may be absent.\n        # In that case we fall back to an empty config so training can proceed.\n        if config is None:\n            config = DictConfig({\"reward\": {}})\n        if tokenizer is None:\n            raise TypeError(\"RateLimitedRewardManager requires `tokenizer`.\")\n\n        super().__init__(config, tokenizer, compute_score)\n        self.compute_score = compute_score or default_compute_score\n        self.is_async_reward_score = inspect.iscoroutinefunction(self.compute_score)\n        self.reward_router_address = reward_router_address\n        self.reward_model_tokenizer = reward_model_tokenizer\n        self.timeout = config.reward.get(\"timeout\", 300.0)\n\n    async def _compute_reward(\n        self, data_source: str, solution_str: str, ground_truth: str, extra_info: dict\n    ) -> dict | float:\n        extra_reward_kwargs = (\n            {\n                \"reward_router_address\": self.reward_router_address,\n                \"reward_model_tokenizer\": self.reward_model_tokenizer,\n            }\n            if self.reward_router_address is not None\n            else {}\n        )\n        if self.is_async_reward_score:\n            return await self.compute_score(\n                data_source=data_source,\n                solution_str=solution_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n                **extra_reward_kwargs,\n            )\n        else:\n            return await self.loop.run_in_executor(\n                None,\n                lambda: self.compute_score(\n                    data_source=data_source,\n                    solution_str=solution_str,\n                    ground_truth=ground_truth,\n                    extra_info=extra_info,\n                    **extra_reward_kwargs,\n                ),\n            )\n\n    async def run_single(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"Only support single data item\"\n        data_item = data[0]\n\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        data_source = data_item.non_tensor_batch[\"data_source\"]\n        ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n        extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n        tool_extra_fields = data_item.non_tensor_batch.get(\"tool_extra_fields\", None)\n        if tool_extra_fields is not None:\n            extra_info.update(tool_extra_fields.items())\n\n        response_str = await self.loop.run_in_executor(\n            None, lambda: self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n        )\n\n        reward_extra_info = {}\n\n        # Apply rate limiting layers\n        if self._rpm_limiter is not None:\n            await self._rpm_limiter.acquire(1.0)\n\n        if self._tpm_limiter is not None:\n            estimated_tokens = self._estimated_tokens_per_request\n            await self._tpm_limiter.acquire(estimated_tokens)\n\n        async with self._semaphore:\n            try:\n                result = await asyncio.wait_for(\n                    self._compute_reward(\n                        data_source=data_source,\n                        solution_str=response_str,\n                        ground_truth=ground_truth,\n                        extra_info=extra_info,\n                    ),\n                    timeout=self.timeout,\n                )\n\n                score: float\n                if isinstance(result, dict):\n                    score = result[\"score\"]\n                    for key, value in result.items():\n                        reward_extra_info[key] = value\n                else:\n                    score = result\n                    reward_extra_info[\"acc\"] = score\n\n                reward = score\n\n            except asyncio.TimeoutError:\n                logger.warning(\n                    f\"Reward computation timed out after {self.timeout}s for data_source={data_source}. \"\n                    f\"Response preview: {response_str[:100]}...\"\n                )\n                reward = 0.0\n                reward_extra_info[\"timeout\"] = True\n                reward_extra_info[\"acc\"] = 0.0\n\n            except Exception as e:\n                logger.error(\n                    f\"Reward computation failed for data_source={data_source}: {e}. \"\n                    f\"Response preview: {response_str[:100]}...\"\n                )\n                reward = 0.0\n                reward_extra_info[\"error\"] = str(e)\n                reward_extra_info[\"acc\"] = 0.0\n\n        return {\"reward_score\": reward, \"reward_extra_info\": reward_extra_info}\n\n    def __call__(self, data: DataProto, return_dict: bool = False):\n        \"\"\"Make the manager callable like traditional reward managers.\n\n        This method provides compatibility with the existing reward manager interface\n        by wrapping the async run_single method in a synchronous call.\n\n        Args:\n            data (DataProto): Input data containing prompts and responses.\n            return_dict (bool): If True, return a dict with reward_tensor and reward_extra_info.\n                               If False, return only the reward_tensor. Defaults to False.\n\n        Returns:\n            torch.Tensor | dict: If return_dict is False, returns a tensor of shape [batch_size, response_length]\n                                with rewards. If return_dict is True, returns a dict with:\n                                - reward_tensor: The reward tensor\n                                - reward_extra_info: Dict containing extra information about rewards\n        \"\"\"\n        from collections import defaultdict\n\n        import torch\n\n        # If there are pre-computed rm_scores, return them directly\n        if \"rm_scores\" in data.batch.keys():\n            if return_dict:\n                reward_extra_keys = data.meta_info.get(\"reward_extra_keys\", [])\n                reward_extra_info = {key: data.non_tensor_batch[key] for key in reward_extra_keys}\n                return {\"reward_tensor\": data.batch[\"rm_scores\"], \"reward_extra_info\": reward_extra_info}\n            else:\n                return data.batch[\"rm_scores\"]\n\n        # Initialize reward tensor\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n        reward_extra_info = defaultdict(list)\n\n        # Process each data item through the async event loop\n        async def process_batch():\n            tasks = []\n            for i in range(len(data)):\n                data_item = data[i : i + 1]  # Get single item as DataProto slice\n                tasks.append(self.run_single(data_item))\n\n            results = await asyncio.gather(*tasks)\n            return results\n\n        # Run the async processing using self.loop property which lazily gets/creates event loop\n        # This ensures rate limiters and semaphores work correctly by using the same loop\n        results = self.loop.run_until_complete(process_batch())\n\n        # Aggregate results into reward tensor and extra info\n        for i, result in enumerate(results):\n            data_item = data[i]\n            response_ids = data_item.batch[\"responses\"]\n            response_length = response_ids.shape[-1]\n            valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n\n            reward = result[\"reward_score\"]\n            reward_tensor[i, valid_response_length - 1] = reward\n\n            # Collect extra info\n            if \"reward_extra_info\" in result:\n                for key, value in result[\"reward_extra_info\"].items():\n                    reward_extra_info[key].append(value)\n\n        if return_dict:\n            return {\n                \"reward_tensor\": reward_tensor,\n                \"reward_extra_info\": reward_extra_info,\n            }\n        else:\n            return reward_tensor\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/naive.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager import register\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\nfrom verl.utils.reward_score import default_compute_score\n\n\n@register(\"naive\")\nclass NaiveRewardManager(RewardManagerBase):\n    \"\"\"The reward manager.\"\"\"\n\n    def __init__(self, config, tokenizer, compute_score, reward_router_address=None, reward_model_tokenizer=None):\n        super().__init__(config, tokenizer, compute_score)\n        self.compute_score = compute_score or default_compute_score\n        self.is_async_reward_score = inspect.iscoroutinefunction(self.compute_score)\n        self.reward_router_address = reward_router_address\n        self.reward_model_tokenizer = reward_model_tokenizer\n\n    async def run_single(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"Only support single data item\"\n        data_item = data[0]\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        data_source = data_item.non_tensor_batch[\"data_source\"]\n        ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n        extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n        tool_extra_fields = data_item.non_tensor_batch.get(\"tool_extra_fields\", None)\n        if tool_extra_fields is not None:\n            extra_info.update(tool_extra_fields.items())\n\n        num_turns = data_item.non_tensor_batch.get(\"__num_turns__\", None)\n        rollout_reward_scores = data_item.non_tensor_batch.get(\"reward_scores\", {})\n        extra_info[\"num_turns\"] = num_turns\n        extra_info[\"rollout_reward_scores\"] = rollout_reward_scores\n\n        response_str = await self.loop.run_in_executor(\n            None, lambda: self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n        )\n\n        extra_reward_kwargs = (\n            {\n                \"reward_router_address\": self.reward_router_address,\n                \"reward_model_tokenizer\": self.reward_model_tokenizer,\n            }\n            if self.reward_router_address is not None\n            else {}\n        )\n        if self.is_async_reward_score:\n            result = await self.compute_score(\n                data_source=data_source,\n                solution_str=response_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n                **extra_reward_kwargs,\n            )\n        else:\n            result = await self.loop.run_in_executor(\n                None,\n                lambda: self.compute_score(\n                    data_source=data_source,\n                    solution_str=response_str,\n                    ground_truth=ground_truth,\n                    extra_info=extra_info,\n                    **extra_reward_kwargs,\n                ),\n            )\n\n        reward_extra_info = {}\n\n        score: float\n        if isinstance(result, dict):\n            score = result[\"score\"]\n            for key, value in result.items():\n                reward_extra_info[key] = value\n        else:\n            score = result\n            reward_extra_info[\"acc\"] = score\n\n        reward = score\n\n        return {\"reward_score\": reward, \"reward_extra_info\": reward_extra_info}\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/registry.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable\n\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\n\n__all__ = [\"register\", \"get_reward_manager_cls\"]\n\nREWARD_MANAGER: dict[str, type[RewardManagerBase]] = {}\n\n\ndef register(name: str) -> Callable[[type[RewardManagerBase]], type[RewardManagerBase]]:\n    \"\"\"Decorator to register a reward manager class with a given name.\n\n    Args:\n        name: `(str)`\n            The name of the reward manager.\n    \"\"\"\n\n    def decorator(cls: type[RewardManagerBase]) -> type[RewardManagerBase]:\n        if name in REWARD_MANAGER and REWARD_MANAGER[name] != cls:\n            raise ValueError(f\"reward manager {name} has already been registered: {REWARD_MANAGER[name]} vs {cls}\")\n        REWARD_MANAGER[name] = cls\n        return cls\n\n    return decorator\n\n\ndef get_reward_manager_cls(name: str) -> type[RewardManagerBase]:\n    \"\"\"Get the reward manager class with a given name.\n\n    Args:\n        name: `(str)`\n            The name of the reward manager.\n\n    Returns:\n        `(type)`: The reward manager class.\n    \"\"\"\n    if name not in REWARD_MANAGER:\n        raise ValueError(f\"Unknown reward manager: {name}\")\n    return REWARD_MANAGER[name]\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_manager/remote.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport itertools\n\nimport ray\n\nfrom verl import DataProto\nfrom verl.experimental.reward_loop.reward_manager import register\nfrom verl.experimental.reward_loop.reward_manager.base import RewardManagerBase\nfrom verl.utils.reward_score import default_compute_score\n\n\n@ray.remote(num_cpus=1)\nclass RewardComputeWorker:\n    \"\"\"\n    WARNING: This class cannot have async methods.\n    \"\"\"\n\n    def __init__(self, compute_score_fn):\n        # since the reward function may not be pickleable, we need to init it in the worker\n        self.compute_score_fn = compute_score_fn\n\n    def compute_score(self, **kwargs) -> dict:\n        return self.compute_score_fn(**kwargs)\n\n\n@register(\"remote\")\nclass RemoteRewardManager(RewardManagerBase):\n    \"\"\"\n    The reward manager.\n    Some errors exist when using default thread pool to compute reward score, e.g., math-verify.\n    https://github.com/volcengine/verl/issues/3407\n    To avoid the above issues, we use a separate process to compute reward score.\n    Moreover, process may be more suitable for cpu-intensive requests.\n    \"\"\"\n\n    def __init__(self, config, tokenizer, compute_score, reward_router_address=None, reward_model_tokenizer=None):\n        super().__init__(config, tokenizer, compute_score)\n        self.compute_score = compute_score or default_compute_score\n        self.is_async_reward_score = inspect.iscoroutinefunction(self.compute_score)\n        assert not self.is_async_reward_score, \"Async reward score is not supported in remote reward manager. \"\n        self.reward_router_address = reward_router_address\n        self.reward_model_tokenizer = reward_model_tokenizer\n        num_reward_workers = config.reward.num_workers\n        # in the rollout & reward parallel mode\n        # the sum of final reward workers will be agent_loop_workers * num_reward_workers\n        self.reward_worker = [\n            # register the reward worker in the same node\n            RewardComputeWorker.options(\n                scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                    node_id=ray.get_runtime_context().get_node_id(),\n                    soft=True,\n                ),\n            ).remote(self.compute_score)\n            for _ in range(num_reward_workers)\n        ]\n        self.reward_worker_pool = itertools.cycle(self.reward_worker)\n\n    def choose_reward_worker(self):\n        return next(self.reward_worker_pool)\n\n    async def run_single(self, data: DataProto) -> dict:\n        assert len(data) == 1, \"Only support single data item\"\n        data_item = data[0]\n        response_ids = data_item.batch[\"responses\"]\n        response_length = response_ids.shape[-1]\n        valid_response_length = data_item.batch[\"attention_mask\"][-response_length:].sum()\n        valid_response_ids = response_ids[:valid_response_length]\n\n        data_source = data_item.non_tensor_batch[\"data_source\"]\n        ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n        extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n        tool_extra_fields = data_item.non_tensor_batch.get(\"tool_extra_fields\", None)\n        if tool_extra_fields is not None:\n            extra_info.update(tool_extra_fields.items())\n\n        num_turns = data_item.non_tensor_batch.get(\"__num_turns__\", None)\n        rollout_reward_scores = data_item.non_tensor_batch.get(\"reward_scores\", {})\n        extra_info[\"num_turns\"] = num_turns\n        extra_info[\"rollout_reward_scores\"] = rollout_reward_scores\n\n        response_str = await self.loop.run_in_executor(\n            None, lambda: self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n        )\n\n        extra_reward_kwargs = (\n            {\n                \"reward_router_address\": self.reward_router_address,\n                \"reward_model_tokenizer\": self.reward_model_tokenizer,\n            }\n            if self.reward_router_address is not None\n            else {}\n        )\n\n        reward_worker = self.choose_reward_worker()\n        result = await reward_worker.compute_score.remote(\n            data_source=data_source,\n            solution_str=response_str,\n            ground_truth=ground_truth,\n            extra_info=extra_info,\n            **extra_reward_kwargs,\n        )\n\n        reward_extra_info = {}\n\n        score: float\n        if isinstance(result, dict):\n            score = result[\"score\"]\n            for key, value in result.items():\n                reward_extra_info[key] = value\n        else:\n            score = result\n            reward_extra_info[\"acc\"] = score\n\n        reward = score\n\n        return {\"reward_score\": reward, \"reward_extra_info\": reward_extra_info}\n"
  },
  {
    "path": "verl/experimental/reward_loop/reward_model.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nimport os\n\nfrom verl.single_controller.ray.base import RayResourcePool, split_resource_pool\nfrom verl.workers.config import HFModelConfig, RewardModelConfig\nfrom verl.workers.rollout.replica import get_rollout_replica_class\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass RewardModelManager:\n    \"\"\"Reward model manager.\"\"\"\n\n    def __init__(\n        self,\n        config: RewardModelConfig,\n        resource_pool: RayResourcePool = None,\n    ):\n        \"\"\"\n        Initialize the reward model manager.\n\n        Args:\n            config (RewardModelConfig): Reward model configuration.\n            resource_pool (RayResourcePool, optional): Resource pool. Defaults to None.\n        \"\"\"\n        self.config = config\n        self.resource_pool = resource_pool\n        self._initialize_llm_servers()\n        self._initialize_router()\n        assert self.config.rollout.skip_tokenizer_init is False, \"Reward model should not skip tokenizer init.\"\n        if self.config.rollout.free_cache_engine:\n            self.sleep()\n\n    def _initialize_llm_servers(self):\n        rollout_world_size = self.config.rollout.tensor_model_parallel_size\n        world_size = (\n            self.resource_pool.world_size\n            if self.resource_pool  # colocate mode\n            else self.config.n_gpus_per_node * self.config.nnodes  # standalone mode\n        )\n        num_replicas = world_size // rollout_world_size\n\n        rollout_replica_class = get_rollout_replica_class(self.config.rollout.name)\n        rollout_config = self.config.rollout\n        model_config = HFModelConfig(path=self.config.model_path)\n        self.tokenizer = model_config.get_processor()\n        self.rollout_replicas = [\n            rollout_replica_class(\n                replica_rank=replica_rank,\n                config=rollout_config,\n                model_config=model_config,\n                gpus_per_node=self.config.n_gpus_per_node,\n                is_reward_model=True,\n            )\n            for replica_rank in range(num_replicas)\n        ]\n        if self.resource_pool:\n            split_resource_pools = split_resource_pool(self.resource_pool, split_size=rollout_world_size)\n            assert len(split_resource_pools) == len(self.rollout_replicas)\n            self._run_all(\n                [\n                    server.init_colocated(resource_pool)\n                    for server, resource_pool in zip(self.rollout_replicas, split_resource_pools, strict=True)\n                ]\n            )\n        else:\n            self._run_all([server.init_standalone() for server in self.rollout_replicas])\n        self.server_handles = [server._server_handle for server in self.rollout_replicas]\n        self.server_addresses = [server._server_address for server in self.rollout_replicas]\n\n    def _initialize_router(self):\n        worker_urls = [f\"http://{server_address}\" for server_address in self.server_addresses]\n\n        # TODO (dyy): sglang router is not ready yet.\n        # if self.config.rollout.name == \"sglang\":\n        #     from .router.inner_sglang_router import launch_router_process\n        # else:\n        #     from .router.naive_router import launch_router_process\n\n        from .router.naive_router import launch_router_process\n\n        self.router_address, _ = launch_router_process(worker_urls=worker_urls)\n\n    def get_router_address(self):\n        return self.router_address\n\n    def wake_up(self):\n        \"\"\"Wake up all rollout replica instances.\"\"\"\n        self._run_all([replica.wake_up() for replica in self.rollout_replicas])\n\n    def sleep(self):\n        \"\"\"Sleep all rollout replica instances.\"\"\"\n        self._run_all([replica.sleep() for replica in self.rollout_replicas])\n\n    def _run_all(self, tasks: list[asyncio.Task]):\n        async def run_all():\n            await asyncio.gather(*tasks)\n\n        asyncio.run(run_all())\n"
  },
  {
    "path": "verl/experimental/reward_loop/router/inner_sglang_router.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport multiprocessing\nimport os\nimport time\n\nimport ray\nimport requests\nfrom sglang_router.launch_server import RouterArgs, launch_router\n\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef launch_router_process(\n    worker_urls: list[str],\n    request_timeout: int = 180,\n    max_wait_time: int = 300,\n    timeout: int = 30,\n) -> str:\n    router_ip = ray.util.get_node_ip_address().strip(\"[]\")\n    router_port, _ = get_free_port(router_ip)\n    router_address = (\n        f\"[{router_ip}]:{router_port}\" if is_valid_ipv6_address(router_ip) else f\"{router_ip}:{router_port}\"\n    )\n    router_args = RouterArgs(\n        host=router_ip,\n        port=router_port,\n        worker_urls=worker_urls,\n        balance_abs_threshold=0,\n        log_level=\"warn\",\n        request_timeout_secs=request_timeout,\n    )\n    router_process = multiprocessing.Process(target=launch_router, args=(router_args,))\n    router_process.daemon = True\n    router_process.start()\n    time.sleep(3)\n    assert router_process.is_alive()\n\n    # health check\n    start_time = time.time()\n    url = f\"http://{router_address}/health\"\n    with requests.Session() as session:\n        while time.time() - start_time < max_wait_time:\n            try:\n                response = session.get(url, timeout=timeout)\n                if response.status_code == 200:\n                    break\n            except requests.RequestException as e:\n                logger.debug(f\"Health check failed: {e}\")\n\n            time.sleep(2)\n        else:\n            router_process.terminate()\n            raise RuntimeError(f\"Router health check failed after {max_wait_time} seconds.\")\n\n    logger.info(f\"Router is running on {router_address}\")\n    return router_address, router_process\n"
  },
  {
    "path": "verl/experimental/reward_loop/router/naive_router.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nimport multiprocessing\nimport os\nimport time\nfrom typing import Any\n\nimport aiohttp\nimport ray\nimport uvicorn\nfrom fastapi import FastAPI, Request\nfrom fastapi.responses import JSONResponse\n\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nasync def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]:\n    if resp.status == 204 or (resp.content_length == 0):\n        return {}\n\n    try:\n        return await resp.json(content_type=None)\n    except Exception:\n        try:\n            text = await resp.text()\n        except Exception:\n            return {}\n        return {\n            \"content_type\": (resp.headers.get(\"Content-Type\") or \"\"),\n            \"text\": text,\n        }\n\n\ndef launch_router_process(\n    worker_urls: list[str],\n):\n    router_ip = ray.util.get_node_ip_address().strip(\"[]\")\n    router_port, _ = get_free_port(router_ip)\n    router_address = (\n        f\"[{router_ip}]:{router_port}\" if is_valid_ipv6_address(router_ip) else f\"{router_ip}:{router_port}\"\n    )\n\n    router_process = multiprocessing.Process(\n        target=run_router,\n        args=(\n            router_ip,\n            router_port,\n            worker_urls,\n        ),\n    )\n    router_process.daemon = True\n    router_process.start()\n    time.sleep(3)\n    assert router_process.is_alive()\n\n    logger.info(f\"Router is running on {router_address}\")\n    return router_address, router_process\n\n\ndef run_router(router_ip: str, router_port: int, worker_urls: list[str]):\n    router = NaiveRouter(worker_urls=worker_urls, verbose=False)\n    uvicorn.run(router.app, host=router_ip, port=router_port, log_level=\"warning\")\n\n\nclass NaiveRouter:\n    def __init__(\n        self,\n        worker_urls: list[str],\n        max_connections: int = 1024,\n        timeout: int = 60,\n        max_attempts: int = 3,\n        retry_delay: float = 2.0,\n        verbose: bool = False,\n    ) -> None:\n        \"\"\"A minimal async load-balancing router.\"\"\"\n        self.verbose = verbose\n        self.app = FastAPI()\n        self.worker_urls = worker_urls\n        self.request_counts = {url: 0 for url in worker_urls}\n\n        self.max_connections = max_connections\n        self.timeout = timeout\n        self.max_attempts = max_attempts\n        self.retry_delay = retry_delay\n\n        self.app = FastAPI()\n\n        # Register startup / shutdown hooks\n        self.app.on_event(\"startup\")(self._on_startup)\n        self.app.on_event(\"shutdown\")(self._on_shutdown)\n\n        # Catch-all proxy route\n        self.app.api_route(\"/{endpoint:path}\", methods=[\"GET\", \"POST\"])(self._make_async_request)\n\n        # Placeholder for aiohttp client\n        self.client = None\n\n    async def _on_startup(self):\n        \"\"\"Initialize aiohttp client safely inside the event loop\"\"\"\n        connector = aiohttp.TCPConnector(\n            limit=self.max_connections,\n            limit_per_host=self.max_connections // 4,\n            ttl_dns_cache=300,\n            use_dns_cache=True,\n        )\n        timeout = aiohttp.ClientTimeout(total=None)\n        self.client = aiohttp.ClientSession(connector=connector, timeout=timeout)\n        if self.verbose:\n            logger.info(f\"[router] aiohttp client initialized with max_connections={self.max_connections}\")\n\n    async def _on_shutdown(self):\n        \"\"\"Gracefully close aiohttp client\"\"\"\n        if self.client and not self.client.closed:\n            await self.client.close()\n            if self.verbose:\n                logger.info(\"[router] aiohttp client closed\")\n\n    async def _make_async_request(self, request: Request, endpoint: str):\n        \"\"\"Proxy single request to a worker URL.\"\"\"\n        if not self.worker_urls:\n            return JSONResponse(status_code=503, content={\"error\": \"No available workers\"})\n\n        worker_url = self._select_worker()\n        target_url = f\"{worker_url}/{endpoint}\"\n\n        if self.verbose:\n            logger.debug(f\"[router] Forwarding request → {target_url}\")\n\n        # Copy request data\n        body = await request.body()\n        headers = dict(request.headers)\n\n        for attempt in range(self.max_attempts):\n            # Send request to worker\n            try:\n                async with self.client.request(request.method, target_url, data=body, headers=headers) as response:\n                    response.raise_for_status()\n                    output = await _read_async_response(response)\n                    self._release_worker(worker_url)\n                    return output\n            except asyncio.TimeoutError:\n                logger.warning(f\"Async request to {endpoint} timed out (attempt {attempt + 1})\")\n            except aiohttp.ClientConnectorError:\n                logger.warning(f\"Connection error for {endpoint} (attempt {attempt + 1})\")\n            except aiohttp.ClientResponseError as e:\n                logger.error(f\"HTTP error for {endpoint}: {e}\")\n                raise\n            except Exception as e:\n                logger.error(f\"Unexpected error for {endpoint}: {e}\")\n                if attempt == self.max_attempts - 1:\n                    raise\n\n            if attempt < self.max_attempts - 1:\n                await asyncio.sleep(self.retry_delay * (2**attempt))\n\n        raise RuntimeError(f\"Failed to complete async request to {endpoint} after {self.max_attempts} attempts\")\n\n    def _select_worker(self) -> str:\n        \"\"\"Select the least-loaded worker (simple round-robin by request count).\"\"\"\n        url = min(self.request_counts, key=self.request_counts.get)\n        self.request_counts[url] += 1\n        return url\n\n    def _release_worker(self, url: str) -> None:\n        \"\"\"Mark worker as free after request completes.\"\"\"\n        self.request_counts[url] = max(0, self.request_counts[url] - 1)\n"
  },
  {
    "path": "verl/experimental/separation/__init__.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/separation/engine_workers.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\n\nfrom omegaconf import DictConfig\n\nfrom verl.single_controller.base.decorator import Dispatch, register\nfrom verl.utils.device import (\n    get_device_name,\n)\nfrom verl.workers.engine_workers import ActorRolloutRefWorker\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\ndevice_name = get_device_name()\n\n__all__ = [\"DetachActorWorker\"]\n\n\nclass DetachActorWorker(ActorRolloutRefWorker):\n    \"\"\"\n    A worker class that extends ActorRolloutRefWorker to support detaching and restoring the actor model.\n\n    This worker facilitates saving the model state to CPU and restoring it, enabling efficient\n    resource management and checkpointing in distributed training. It currently supports\n    FSDP, FSDP2, and Megatron strategies.\n    \"\"\"\n\n    def __init__(self, config: DictConfig, role: str):\n        \"\"\"\n        Initialize the DetachActorWorker.\n\n        Args:\n            config: Configuration dictionary.\n            role: The role of the worker (e.g., 'actor', 'rollout', 'ref').\n        \"\"\"\n        ActorRolloutRefWorker.__init__(self, config, role)\n        self._strategy_handlers = None\n        self.copy_handler, self.restore_handler = self._get_strategy_handlers()\n\n    def _get_strategy_handlers(self):\n        \"\"\"\n        Get the strategy-specific handlers for saving and restoring the model.\n\n        Returns:\n            tuple: A tuple containing (save_handler, restore_handler).\n\n        Raises:\n            NotImplementedError: If the strategy is not supported.\n        \"\"\"\n        if self._strategy_handlers is not None:\n            return self._strategy_handlers\n\n        strategy = self.config.actor.strategy\n\n        if strategy in [\"fsdp\", \"fsdp2\"]:\n            from verl.utils.fsdp_utils import (\n                fsdp2_sharded_load_from_cpu,\n                fsdp2_sharded_save_to_cpu,\n            )\n\n            self._strategy_handlers = (fsdp2_sharded_save_to_cpu, fsdp2_sharded_load_from_cpu)\n        elif strategy == \"megatron\":\n            from verl.utils.megatron_utils import (\n                copy_megatron_model_to_cpu,\n                restore_megatron_model_from_cpu,\n            )\n\n            self._strategy_handlers = (copy_megatron_model_to_cpu, restore_megatron_model_from_cpu)\n        else:\n            raise NotImplementedError(f\"Unsupported strategy: {strategy}\")\n\n        return self._strategy_handlers\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_model_to_cpu(self, n):\n        \"\"\"\n        Save the current model state to CPU memory.\n\n        Args:\n            n: Identifier/Key for the saved model state.\n        \"\"\"\n        if not hasattr(self, \"cpu_saved_models\"):\n            self.cpu_saved_models = {}\n\n        self.cpu_saved_models[n] = self.copy_handler(self.actor.engine.module)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def restore_model_from_cpu(self, n):\n        \"\"\"\n        Restore the model state from CPU memory.\n\n        Args:\n            n: Identifier/Key for the saved model state to restore.\n        \"\"\"\n        if n in self.cpu_saved_models:\n            strategy = self.config.actor.strategy\n\n            if strategy in [\"fsdp\", \"fsdp2\"]:\n                cpu_sharded_state, global_spec = self.cpu_saved_models[n]\n                self.restore_handler(self.actor.engine.module, cpu_sharded_state, global_spec)\n            else:\n                self.restore_handler(self.actor.engine.module, self.cpu_saved_models[n])\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def clear_cpu_model(self, n):\n        \"\"\"\n        Clear the saved model state from CPU memory.\n\n        Args:\n            n: Identifier/Key for the saved model state to remove.\n        \"\"\"\n        if n in self.cpu_saved_models:\n            del self.cpu_saved_models[n]\n"
  },
  {
    "path": "verl/experimental/separation/ray_trainer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPPO Trainer with Ray-based single controller.\nThis trainer supports model-agonistic model initialization with huggingface\n\"\"\"\n\nimport uuid\nfrom copy import deepcopy\nfrom pprint import pprint\nfrom typing import Any, Optional\n\nimport numpy as np\nimport torch\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import Dataset, Sampler\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.dataset.sampler import AbstractCurriculumSampler\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager\nfrom verl.single_controller.ray.base import create_colocated_worker_cls\nfrom verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss\nfrom verl.trainer.ppo.metric_utils import (\n    compute_data_metrics,\n    compute_throughout_metrics,\n    compute_timing_metrics,\n    compute_variance_proxy_metrics,\n)\nfrom verl.trainer.ppo.ray_trainer import RayPPOTrainer, apply_kl_penalty, compute_advantage, compute_response_mask\nfrom verl.trainer.ppo.reward import extract_reward\nfrom verl.trainer.ppo.utils import Role, WorkerType\nfrom verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.metric import reduce_metrics\nfrom verl.utils.rollout_skip import RolloutSkip\n\n\nclass SeparateRayPPOTrainer(RayPPOTrainer):\n    \"\"\"\n    Support for the initialization and fit process of Ray Trainer in the resource-separated scenario:\n        - Fully async policy\n        - One-step off-policy\n    \"\"\"\n\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        role_worker_mapping: dict[Role, WorkerType],\n        resource_pool_manager: ResourcePoolManager,\n        ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup,\n        processor=None,\n        reward_fn=None,\n        val_reward_fn=None,\n        train_dataset: Optional[Dataset] = None,\n        val_dataset: Optional[Dataset] = None,\n        collate_fn=None,\n        train_sampler: Optional[Sampler] = None,\n        device_name=None,\n    ):\n        super().__init__(\n            config,\n            tokenizer,\n            role_worker_mapping,\n            resource_pool_manager,\n            ray_worker_group_cls,\n            processor,\n            train_dataset,\n            val_dataset,\n            collate_fn,\n            train_sampler,\n            device_name,\n        )\n        self.global_steps = 0\n        self.epoch = 0\n        self.max_steps_duration = 0\n        self.progress_bar = None\n        self.logger = None\n        self.is_last_step = False\n        self.prev_step_profile = False\n        self.curr_step_profile = False\n        self.next_step_profile = False\n        self.last_val_metrics = {}\n        self.metrics = {}\n        self.timing_raw = {}\n        # reward message\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n        self.checkpoint_manager = None\n\n    def init_workers(self):\n        \"\"\"Initialize distributed training workers using Ray backend.\n\n        Creates:\n        1. Ray resource pools from configuration\n        2. Worker groups for each role (actor, critic, etc.)\n        \"\"\"\n        self._init_resource_pools()\n        self._create_worker_classes()\n        self._init_worker_groups()\n        self._init_models()\n        self._init_reward_loop()\n        self._init_async_rollout_manager()\n\n        self.checkpoint_manager = CheckpointEngineManager(\n            config=omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine),\n            trainer=self.actor_rollout_wg,\n            replicas=self.async_rollout_manager.rollout_replicas,\n        )\n\n    def _init_resource_pools(self):\n        self.resource_pool_manager.create_resource_pool()\n        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}\n\n    def _create_worker_classes(self):\n        self._create_actor_rollout_classes()\n        self._create_critic_class()\n        self._create_reference_policy_class()\n        self._create_reward_model_class()\n\n    def _create_actor_rollout_classes(self):\n        raise NotImplementedError\n\n    def _create_critic_class(self):\n        # create critic\n        if self.use_critic:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)\n            critic_cfg = omega_conf_to_dataclass(self.config.critic)\n\n            if self.use_legacy_worker_impl == \"disable\":\n                # convert critic_cfg into TrainingWorkerConfig\n                from verl.workers.config import FSDPEngineConfig\n                from verl.workers.engine_workers import TrainingWorkerConfig\n\n                self.orig_critic_cfg = critic_cfg\n                if self.orig_critic_cfg.strategy == \"fsdp\":\n                    engine_config: FSDPEngineConfig = self.orig_critic_cfg.model.fsdp_config\n                    engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu\n                    engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu\n                else:\n                    raise NotImplementedError(f\"Unknown strategy {self.orig_critic_cfg.strategy=}\")\n\n                critic_cfg = TrainingWorkerConfig(\n                    model_type=\"value_model\",\n                    model_config=self.orig_critic_cfg.model_config,\n                    engine_config=engine_config,\n                    optimizer_config=self.orig_critic_cfg.optim,\n                    checkpoint_config=self.orig_critic_cfg.checkpoint,\n                )\n\n            critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)\n            self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls\n\n    def _create_reference_policy_class(self):\n        # create reference policy if needed\n        if self.use_reference_policy:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)\n            ref_policy_cls = RayClassWithInitArgs(\n                self.role_worker_mapping[Role.RefPolicy],\n                config=self.config.actor_rollout_ref,\n                role=str(Role.RefPolicy),\n                # profile_option=self.config.trainer.npu_profile.options,\n            )\n            self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls\n\n    def _create_reward_model_class(self):\n        # create a reward model if reward_fn is None\n        if self.use_rm:\n            # we create a RM here\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel)\n            rm_cls = RayClassWithInitArgs(\n                self.role_worker_mapping[Role.RewardModel], config=self.config.reward.reward_model\n            )\n            self.resource_pool_to_cls[resource_pool][str(Role.RewardModel)] = rm_cls\n\n    def _init_worker_groups(self):\n        # initialize WorkerGroup\n        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,\n        # you should not use `create_colocated_worker_cls`.\n        # Instead, directly pass different resource pool to different worker groups.\n        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.\n        all_wg = {}\n        wg_kwargs = {}  # Setting up kwargs for RayWorkerGroup\n        if OmegaConf.select(self.config.trainer, \"ray_wait_register_center_timeout\") is not None:\n            wg_kwargs[\"ray_wait_register_center_timeout\"] = self.config.trainer.ray_wait_register_center_timeout\n        if OmegaConf.select(self.config.global_profiler, \"steps\") is not None:\n            wg_kwargs[\"profile_steps\"] = OmegaConf.select(self.config.global_profiler, \"steps\")\n            # Only require nsight worker options when tool is nsys\n            if OmegaConf.select(self.config.global_profiler, \"tool\") == \"nsys\":\n                assert (\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                    is not None\n                ), \"worker_nsight_options must be set when using nsys with profile_steps\"\n                wg_kwargs[\"worker_nsight_options\"] = OmegaConf.to_container(\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                )\n        wg_kwargs[\"device_name\"] = self.device_name\n\n        for resource_pool, class_dict in self.resource_pool_to_cls.items():\n            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n            wg_dict = self.ray_worker_group_cls(\n                resource_pool=resource_pool,\n                ray_cls_with_init=worker_dict_cls,\n                **wg_kwargs,\n            )\n            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n            all_wg.update(spawn_wg)\n        self.all_wg = all_wg\n\n    def _init_models(self):\n        if self.use_critic:\n            self.critic_wg = self.all_wg[str(Role.Critic)]\n            if self.use_legacy_worker_impl == \"disable\":\n                self.critic_wg.reset()\n                # assign critic loss\n                from functools import partial\n\n                from verl.workers.utils.losses import value_loss\n\n                value_loss_ = partial(value_loss, config=self.orig_critic_cfg)\n                self.critic_wg.set_loss_fn(value_loss_)\n            else:\n                self.critic_wg.init_model()\n\n        if self.use_reference_policy and not self.ref_in_actor:\n            self.ref_policy_wg = self.all_wg[str(Role.RefPolicy)]\n            self.ref_policy_wg.init_model()\n\n        if self.use_rm:\n            self.rm_wg = self.all_wg[str(Role.RewardModel)]\n            self.rm_wg.init_model()\n\n        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory\n        self.actor_rollout_wg = self.all_wg[str(Role.ActorRollout)]\n        self.actor_rollout_wg.init_model()\n\n    def _init_reward_loop(self):\n        from verl.experimental.reward_loop import RewardLoopManager\n\n        # initalize reward loop manager\n        # reward model (colocate or standalone): get resource_pool\n        # no reward model: resource_pool = None\n        resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None\n        self.reward_loop_manager = RewardLoopManager(\n            config=self.config,\n            rm_resource_pool=resource_pool,\n        )\n\n    def _init_async_rollout_manager(self):\n        pass\n\n    def fit(self):\n        \"\"\"\n        The training loop of PPO.\n        The driver process only need to call the compute functions of the worker group through RPC\n        to construct the PPO dataflow.\n        The light-weight advantage computation is done on the driver process.\n\n        !!!\n        The logic of fit is consistent with that of fit_refactor;\n        if any modifications are made, apply them to both methods simultaneously.\n        \"\"\"\n        from omegaconf import OmegaConf\n\n        from verl.utils.tracking import Tracking\n\n        self.logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        self.global_steps = 0\n\n        # load checkpoint and update weights before doing anything\n        self._load_checkpoint()\n        self.checkpoint_manager.update_weights(self.global_steps)\n\n        current_epoch = self.global_steps // len(self.train_dataloader)\n\n        # perform validation before training\n        # currently, we only support validation using the reward_function.\n        if self.config.trainer.get(\"val_before_train\", True):\n            val_metrics = self._validate()\n            assert val_metrics, f\"{val_metrics=}\"\n            pprint(f\"Initial validation metrics: {val_metrics}\")\n            self.logger.log(data=val_metrics, step=self.global_steps)\n            if self.config.trainer.get(\"val_only\", False):\n                return\n\n        if self.config.actor_rollout_ref.rollout.get(\"skip_rollout\", False):\n            rollout_skip = RolloutSkip(self.config, self.actor_rollout_wg)\n            rollout_skip.wrap_generate_sequences()\n\n        # add tqdm\n        self.progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc=\"Training Progress\")\n\n        # we start from step 1\n        self.global_steps += 1\n        self.last_val_metrics = None\n        self.max_steps_duration = 0\n\n        self.prev_step_profile = False\n        self.curr_step_profile = (\n            self.global_steps in self.config.global_profiler.steps\n            if self.config.global_profiler.steps is not None\n            else False\n        )\n        self.next_step_profile = False\n\n        for epoch in range(current_epoch, self.config.trainer.total_epochs):\n            for batch_dict in self.train_dataloader:\n                self.epoch = epoch\n                self.fit_step(batch_dict)\n                if self.is_last_step:\n                    return\n\n    def fit_step(self, batch_dict: Any = None):\n        \"\"\"\n        Single-step training template method. Handles all logic for one training step.\n\n        Flow:\n        1. Pre-step processing -> 2. Get batch -> 3. Generate sequences ->\n        4. Compute reward -> 5. Compute log_prob -> 6. Compute reward ->\n        7. Compute advantage -> 8. Update critic -> 9. Update actor -> 10. Post-step processing\n\n        Args:\n            batch_dict: Raw data dictionary\n        \"\"\"\n        self.metrics = {\"training/global_step\": self.global_steps, \"training/epoch\": self.epoch}\n        self.timing_raw = {}\n        # reward message\n        self.reward_tensor = None\n        self.reward_extra_infos_dict = {}\n\n        self._fit_prepare_step()\n        self._fit_start_profile()\n\n        with marked_timer(\"step\", self.timing_raw):\n            batch = self._fit_get_batch(batch_dict)\n            batch = self._fit_generate(batch)\n            batch = self._fit_compute_reward(batch)\n            batch = self._fit_compute_log_prob(batch)\n            batch = self._fit_compute_ref_log_prob(batch)\n            batch = self._fit_compute_critic(batch)\n            batch = self._fit_compute_advantage(batch)\n            batch = self._fit_update_critic(batch)\n            batch = self._fit_update_actor(batch)\n            self._fit_update_weights()\n            self._fit_dump_data(batch)\n\n        self._fit_validate()\n        self._fit_save_checkpoint()\n        self._fit_stop_profile()\n        self._fit_collect_metrics(batch)\n        self._fit_torch_memory()\n        self._fit_experimental(batch)\n        self._fit_postprocess_step()\n\n    def _fit_prepare_step(self):\n        if hasattr(self.actor_rollout_wg, \"async_calls_finalize_fn_exec\"):\n            self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False)\n        self.is_last_step = self.global_steps >= self.total_training_steps\n\n    def _fit_start_profile(self):\n        timing_raw = self.timing_raw\n        with marked_timer(\"start_profile\", timing_raw):\n            self._start_profiling(\n                not self.prev_step_profile and self.curr_step_profile\n                if self.config.global_profiler.profile_continuous_steps\n                else self.curr_step_profile\n            )\n\n    def _fit_get_batch(self, batch_dict: dict) -> DataProto:\n        batch = DataProto.from_single_dict(batch_dict)\n        batch.meta_info[\"temperature\"] = self.config.actor_rollout_ref.rollout.temperature\n        # add uid\n        batch.non_tensor_batch[\"uid\"] = np.array([str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object)\n        return batch\n\n    def _fit_generate(self, batch: DataProto = None) -> DataProto:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        gen_batch = self._get_gen_batch(batch)\n        # pass global_steps to trace\n        gen_batch.meta_info[\"global_steps\"] = self.global_steps\n        gen_batch_output = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n\n        with marked_timer(\"gen\", timing_raw, color=\"red\"):\n            if self.curr_step_profile:\n                self.async_rollout_manager.start_profile(global_step=self.global_steps)\n            gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch_output)\n            self.checkpoint_manager.sleep_replicas()\n            if self.curr_step_profile:\n                self.async_rollout_manager.stop_profile()\n\n            timing_raw.update(gen_batch_output.meta_info[\"timing\"])\n            gen_batch_output.meta_info.pop(\"timing\", None)\n\n        if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:\n            with marked_timer(\"gen_max\", timing_raw, color=\"purple\"):\n                gen_baseline_batch = deepcopy(gen_batch)\n                gen_baseline_batch.meta_info[\"do_sample\"] = False\n                if self.curr_step_profile:\n                    self.async_rollout_manager.start_profile()\n                gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)\n                self.checkpoint_manager.sleep_replicas()\n                if self.curr_step_profile:\n                    self.async_rollout_manager.stop_profile()\n                batch = batch.union(gen_baseline_output)\n                # compute reward model score on batch\n                rm_scores = None\n                if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                    batch_reward = self._compute_reward_colocate(batch)\n                    batch = batch.union(batch_reward)\n\n                # Compute or extract reward for REMAX baseline\n                reward_baseline_tensor = batch.batch[\"rm_scores\"].sum(dim=-1)\n\n                keys_to_pop = set(gen_baseline_output.batch.keys())\n                if rm_scores is not None:\n                    keys_to_pop.update(rm_scores.batch.keys())\n                batch.pop(batch_keys=list(keys_to_pop))\n\n                batch.batch[\"reward_baselines\"] = reward_baseline_tensor\n\n                del rm_scores, gen_baseline_batch, gen_baseline_output\n        # repeat to align with repeated responses in rollout\n        batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n        batch = batch.union(gen_batch_output)\n\n        if \"response_mask\" not in batch.batch.keys():\n            batch.batch[\"response_mask\"] = compute_response_mask(batch)\n        # Balance the number of valid tokens across DP ranks.\n        # NOTE: This usually changes the order of data in the `batch`,\n        # which won't affect the advantage calculation (since it's based on uid),\n        # but might affect the loss calculation (due to the change of mini-batching).\n        if self.config.trainer.balance_batch:\n            self._balance_batch(batch, metrics=metrics)\n\n        # compute global_valid tokens\n        batch.meta_info[\"global_token_num\"] = torch.sum(batch.batch[\"attention_mask\"], dim=-1).tolist()\n        # get images_seqlens\n        images_seqlens_all = []\n        for multi_modal_input in batch.non_tensor_batch[\"multi_modal_inputs\"]:\n            if \"image_grid_thw\" not in multi_modal_input.keys():\n                continue\n            images_seqlens_all.extend(multi_modal_input[\"images_seqlens\"].tolist())\n        batch.meta_info[\"images_seqlens\"] = images_seqlens_all\n        return batch\n\n    def _fit_compute_reward(self, batch: DataProto) -> DataProto:\n        timing_raw = self.timing_raw\n        with marked_timer(\"reward\", timing_raw, color=\"yellow\"):\n            # compute reward model score\n            if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                batch_reward = self._compute_reward_colocate(batch)\n                batch = batch.union(batch_reward)\n\n            # Compute or extract reward_tensor and reward_extra_infos_dict for training\n            reward_tensor, reward_extra_infos_dict = extract_reward(batch)\n            self.reward_tensor = reward_tensor\n            self.reward_extra_infos_dict = reward_extra_infos_dict\n        return batch\n\n    def _fit_compute_log_prob(self, batch: DataProto) -> DataProto:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        # Operating Mode Selection:\n        # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ)\n        # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ)\n        #   Note: π_old computed once per data batch, serves as stable reference during mini-batch updates\n        rollout_corr_config = self.config.algorithm.get(\"rollout_correction\", None)\n        bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get(\"bypass_mode\", False)\n        if bypass_recomputing_logprobs:  # Use `rollout_log_probs`\n            from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode\n\n            apply_bypass_mode(\n                batch=batch,\n                rollout_corr_config=rollout_corr_config,\n                policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss,\n            )\n        else:  # Recompute old_log_probs\n            with marked_timer(\"old_log_prob\", timing_raw, color=\"blue\"):\n                old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch)\n                entropys = old_log_prob.batch[\"entropys\"]\n                response_masks = batch.batch[\"response_mask\"]\n                actor_config = self.config.actor_rollout_ref.actor\n                entropy_agg = agg_loss(\n                    loss_mat=entropys,\n                    loss_mask=response_masks,\n                    loss_agg_mode=actor_config.loss_agg_mode,\n                    loss_scale_factor=actor_config.loss_scale_factor,\n                )\n                old_log_prob_metrics = {\n                    \"actor/entropy\": entropy_agg.detach().item(),\n                    \"perf/mfu/actor_infer\": old_log_prob_mfu,\n                }\n                metrics.update(old_log_prob_metrics)\n                old_log_prob.batch.pop(\"entropys\")\n                if \"routed_experts\" in batch.batch and \"routed_experts\" in old_log_prob.batch:\n                    router_mode = getattr(self.config.actor_rollout_ref.actor.router_replay, \"mode\", \"disabled\")\n                    if router_mode == \"R2\":\n                        batch.batch.pop(\"routed_experts\")\n                    else:\n                        old_log_prob.batch.pop(\"routed_experts\")\n                batch = batch.union(old_log_prob)\n                if \"rollout_log_probs\" in batch.batch.keys():\n                    # TODO: we may want to add diff of probs too.\n                    from verl.utils.debug.metrics import calculate_debug_metrics\n\n                    metrics.update(calculate_debug_metrics(batch))\n\n        assert \"old_log_probs\" in batch.batch, f'\"old_log_prob\" not in {batch.batch.keys()=}'\n        return batch\n\n    def _fit_compute_ref_log_prob(self, batch: DataProto) -> DataProto:\n        timing_raw = self.timing_raw\n        if self.use_reference_policy:\n            with marked_timer(str(Role.RefPolicy), timing_raw, color=\"olive\"):\n                ref_log_prob = self._compute_ref_log_prob(batch)\n                batch = batch.union(ref_log_prob)\n        return batch\n\n    def _fit_compute_critic(self, batch: DataProto) -> DataProto:\n        timing_raw = self.timing_raw\n        if self.use_critic:\n            with marked_timer(\"values\", timing_raw, color=\"cyan\"):\n                values = self._compute_values(batch)\n                batch = batch.union(values)\n        return batch\n\n    def _fit_compute_advantage(self, batch) -> DataProto:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        reward_tensor = self.reward_tensor\n        reward_extra_infos_dict = self.reward_extra_infos_dict\n\n        with marked_timer(\"adv\", timing_raw, color=\"brown\"):\n            # we combine with rule-based rm\n            reward_extra_infos_dict: dict[str, list]\n            batch.batch[\"token_level_scores\"] = reward_tensor\n\n            if reward_extra_infos_dict:\n                batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})\n\n            # compute rewards. apply_kl_penalty if available\n            if self.config.algorithm.use_kl_in_reward:\n                batch, kl_metrics = apply_kl_penalty(\n                    batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty\n                )\n                metrics.update(kl_metrics)\n            else:\n                batch.batch[\"token_level_rewards\"] = batch.batch[\"token_level_scores\"]\n\n            # Compute rollout correction: IS weights, rejection sampling, and metrics\n            # Only runs in decoupled mode (computes once per batch using stable π_old)\n            # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout\n            rollout_corr_config = self.config.algorithm.get(\"rollout_correction\", None)\n            bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get(\"bypass_mode\", False)\n            if (\n                rollout_corr_config is not None\n                and \"rollout_log_probs\" in batch.batch\n                and not bypass_recomputing_logprobs  # Only in decoupled mode\n            ):\n                from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch\n\n                # Compute IS weights, apply rejection sampling, compute metrics\n                batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)\n                # IS and off-policy metrics already have rollout_corr/ prefix\n                metrics.update(is_metrics)\n\n            # compute advantages, executed on the driver process\n            norm_adv_by_std_in_grpo = self.config.algorithm.get(\n                \"norm_adv_by_std_in_grpo\", True\n            )  # GRPO adv normalization factor\n\n            batch = compute_advantage(\n                batch,\n                adv_estimator=self.config.algorithm.adv_estimator,\n                gamma=self.config.algorithm.gamma,\n                lam=self.config.algorithm.lam,\n                num_repeat=self.config.actor_rollout_ref.rollout.n,\n                norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n                config=self.config.algorithm,\n            )\n        return batch\n\n    def _fit_update_critic(self, batch: DataProto) -> DataProto:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        if self.use_critic:\n            with marked_timer(\"update_critic\", timing_raw, color=\"pink\"):\n                critic_output = self._update_critic(batch)\n            critic_output_metrics = reduce_metrics(critic_output.meta_info[\"metrics\"])\n            metrics.update(critic_output_metrics)\n        return batch\n\n    def _fit_update_actor(self, batch: DataProto) -> DataProto:\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        # implement critic warmup\n        if self.config.trainer.critic_warmup <= self.global_steps:\n            # update actor\n            with marked_timer(\"update_actor\", timing_raw, color=\"red\"):\n                actor_output = self._update_actor(batch)\n\n            actor_output_metrics = reduce_metrics(actor_output.meta_info[\"metrics\"])\n            metrics.update(actor_output_metrics)\n        return batch\n\n    def _fit_update_weights(self):\n        timing_raw = self.timing_raw\n        if self.config.trainer.critic_warmup <= self.global_steps:\n            # update weights from trainer to rollout\n            with marked_timer(\"update_weights\", timing_raw, color=\"red\"):\n                self.checkpoint_manager.update_weights(self.global_steps)\n\n    def _fit_dump_data(self, batch: DataProto):\n        timing_raw = self.timing_raw\n        reward_extra_infos_dict = self.reward_extra_infos_dict\n        # Log rollout generations if enabled\n        rollout_data_dir = self.config.trainer.get(\"rollout_data_dir\", None)\n        if rollout_data_dir:\n            self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir)\n\n    def _fit_validate(self):\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n        if self.config.trainer.test_freq > 0 and (\n            self.is_last_step or self.global_steps % self.config.trainer.test_freq == 0\n        ):\n            with marked_timer(\"testing\", timing_raw, color=\"green\"):\n                val_metrics: dict = self._validate()\n                if self.is_last_step:\n                    self.last_val_metrics = val_metrics\n            metrics.update(val_metrics)\n\n    def _fit_save_checkpoint(self):\n        timing_raw = self.timing_raw\n        # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.\n        esi_close_to_expiration = should_save_ckpt_esi(\n            max_steps_duration=self.max_steps_duration,\n            redundant_time=self.config.trainer.esi_redundant_time,\n        )\n        # Check if the conditions for saving a checkpoint are met.\n        # The conditions include a mandatory condition (1) and\n        # one of the following optional conditions (2/3/4):\n        # 1. The save frequency is set to a positive value.\n        # 2. It's the last training step.\n        # 3. The current step number is a multiple of the save frequency.\n        # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.\n        if self.config.trainer.save_freq > 0 and (\n            self.is_last_step or self.global_steps % self.config.trainer.save_freq == 0 or esi_close_to_expiration\n        ):\n            if esi_close_to_expiration:\n                print(\"Force saving checkpoint: ESI instance expiration approaching.\")\n            with marked_timer(\"save_checkpoint\", timing_raw, color=\"green\"):\n                # sleep replicas to avoid OOM during checkpoint saving\n                # self.checkpoint_manager.sleep_replicas()\n                self._save_checkpoint()\n                # wake replicas to avoid OOM during checkpoint saving\n                # TODO: Check separation is needed.\n                # self.checkpoint_manager.update_weights()\n\n    def _fit_stop_profile(self):\n        timing_raw = self.timing_raw\n        with marked_timer(\"stop_profile\", timing_raw):\n            self.next_step_profile = (\n                self.global_steps + 1 in self.config.global_profiler.steps\n                if self.config.global_profiler.steps is not None\n                else False\n            )\n            self._stop_profiling(\n                self.curr_step_profile and not self.next_step_profile\n                if self.config.global_profiler.profile_continuous_steps\n                else self.curr_step_profile\n            )\n            self.prev_step_profile = self.curr_step_profile\n            self.curr_step_profile = self.next_step_profile\n\n    def _fit_collect_metrics(self, batch):\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n\n        # collect metrics\n        metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))\n        metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))\n        # TODO: implement actual tflpo and theoretical tflpo\n        n_gpus = self.resource_pool_manager.get_n_gpus()\n        metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))\n        # compute variance proxy metrics\n        gradient_norm = metrics.get(\"actor/grad_norm\", None)\n        metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm))\n\n    def _fit_torch_memory(self):\n        if (\n            hasattr(self.config.actor_rollout_ref.actor, \"profiler\")\n            and self.config.actor_rollout_ref.actor.profiler.tool == \"torch_memory\"\n        ):\n            self.actor_rollout_wg.dump_memory_snapshot(\n                tag=f\"post_update_step{self.global_steps}\", sub_dir=f\"step{self.global_steps}\"\n            )\n\n    def _fit_experimental(self, batch):\n        # this is experimental and may be changed/removed in the future in favor of a general-purpose one\n        if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):\n            self.train_dataloader.sampler.update(batch=batch)\n\n        # this is experimental and may be changed/removed in the future\n        # in favor of a general-purpose data buffer pool\n        if hasattr(self.train_dataset, \"on_batch_end\"):\n            # The dataset may be changed after each training batch\n            self.train_dataset.on_batch_end(batch=batch)\n\n    def _fit_postprocess_step(self):\n        metrics = self.metrics\n        timing_raw = self.timing_raw\n\n        steps_duration = timing_raw[\"step\"]\n        self.max_steps_duration = max(self.max_steps_duration, steps_duration)\n\n        # TODO: make a canonical logger that supports various backend\n        self.logger.log(data=metrics, step=self.global_steps)\n        self.progress_bar.update(1)\n        self.global_steps += 1\n        if self.is_last_step:\n            if hasattr(self.actor_rollout_wg, \"async_calls_finalize_fn_exec\"):\n                self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True)\n            pprint(f\"Final validation metrics: {self.last_val_metrics}\")\n            self.progress_bar.close()\n"
  },
  {
    "path": "verl/experimental/separation/utils.py",
    "content": "# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport ray\n\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.trainer.ppo.utils import Role, need_reference_policy\n\n\ndef create_resource_pool_manager(config, roles: list) -> ResourcePoolManager:\n    \"\"\"\n    Create resource pool manager\n\n    Args:\n        config: Configuration object\n        roles: List of roles that need to create resource pools\n\n    Returns:\n        ResourcePoolManager: Resource pool manager\n    \"\"\"\n    resource_pool_spec = {}\n    mapping = {}\n\n    # Actor/Critic resource pool\n    if any(role in roles for role in [Role.Actor, Role.ActorRollout, Role.Critic, Role.RefPolicy, Role.RewardModel]):\n        assert config.trainer.n_gpus_per_node > 0, \"config.trainer.n_gpus_per_node must be greater than 0\"\n        assert config.trainer.nnodes > 0, \"config.trainer.nnodes must be greater than 0\"\n\n        trainer_pool = [config.trainer.n_gpus_per_node] * config.trainer.nnodes\n        resource_pool_spec[\"trainer_pool\"] = trainer_pool\n\n        # Map training-related roles to the same resource pool\n        for role in [Role.Actor, Role.ActorRollout, Role.Critic, Role.RefPolicy, Role.RewardModel]:\n            if role in roles:\n                mapping[role] = \"trainer_pool\"\n\n    # Rollout resource pool\n    if Role.Rollout in roles:\n        assert config.rollout.n_gpus_per_node > 0, \"config.rollout.n_gpus_per_node must be greater than 0\"\n        assert config.rollout.nnodes > 0, \"config.rollout.nnodes must be greater than 0\"\n\n    return ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n\n\ndef create_role_worker_mapping(config):\n    \"\"\"\n    Create mapping from roles to worker classes\n\n    Args:\n        config: Configuration object\n\n    Returns:\n        dict: Mapping from roles to worker classes\n    \"\"\"\n    # Select worker class based on strategy\n    if config.trainer.get(\"use_legacy_worker_impl\", \"auto\") != \"disable\":\n        raise NotImplementedError(\n            \"Fully async policy or One step off policy does not support legacy worker implementation\"\n        )\n\n    from verl.experimental.separation.engine_workers import DetachActorWorker\n    from verl.single_controller.ray import RayWorkerGroup\n    from verl.workers.engine_workers import TrainingWorker\n\n    ray_worker_group_cls = RayWorkerGroup\n\n    train_role = Role.Actor\n    if config.get(\"async_training\", {}).get(\"use_trainer_do_validate\", False):\n        train_role = Role.ActorRollout\n\n    role_worker_mapping = {\n        train_role: ray.remote(DetachActorWorker),\n        Role.Critic: ray.remote(TrainingWorker),\n    }\n\n    # Add reference policy (if KL loss or reward is required)\n    if need_reference_policy(config):\n        role_worker_mapping[Role.RefPolicy] = ray.remote(DetachActorWorker)\n\n    return role_worker_mapping, ray_worker_group_cls\n"
  },
  {
    "path": "verl/experimental/vla/README.md",
    "content": "# [WIP] Experimental VLA RL Support\n\nThis recipe introduces experimental support for training SimpleVLA-OFT, a VLA model.\n\nA key challenge in VLA RL training, which differs from standard LLM RL training, is that the environment/simulation phase has a higher computational overhead than the generation phase. To achieve high efficiency, RL in this context requires an effective environment scheduling mechanism in addition to verl's existing efficient training and inference scheduling. The goal is to reduce the inefficiency caused by the environment and the model's generation process waiting on each other.\n\nThe core computational model of this PR is inspired by the pipeline parallelism design from RLinf. It aims to overlap the environment's execution time with the model's generation time, thereby maximizing environment utilization.\n\nThis PR also proposes a future direction: creating a unified `Env` class. This class would encapsulate functionalities like tool calling, MCP, etc., under a single interface. The environment would manage its state internally, allowing the agent to communicate simply by calling `step(action)` to submit an action and receive an observation.\n\nCurrently, this code is located independently within the `recipes` folder. Much of the design is tightly coupled with the SimpleVLA model and the Libero environment, serving as an initial version for demonstration and discussion.\n\n## Supported Simulators\n\n| Simulator | Env Name |  Difference | Benchmark data source |\n| --- | --- | --- | --- | \n| Mujoco | LiberoEnv | 1. init task from init_states in Libero dataset<br>2. each env can have different tasks | https://github.com/Lifelong-Robot-Learning/LIBERO |\n| IsaacSim | IsaacEnv  | 1. init task from random states, which has more variety than init_states in dataset<br>2. each sim process must using the same task for its envs | https://huggingface.co/datasets/china-sae-robotics/IsaacLabPlayGround_Dataset |\n\n## Hardware Requirements\n\n*   Simulator GPU: NVIDIA L20 or L40 with 48GB memory and RT Cores\n\nNotes: \n1. Mujoco can failback to CPU mode with degraded performance if no RT Cores is available\n2. IsaacSim only support GPU with RT Cores\n3. RTX GPU will be supported in the future release with remote deployment feature, but it can not work with colocated mode because of the limitation of GPU memory capacity.\n\n## Docker image\n\nThe Isaac Lab support for libero dataset depends on RobotLearningLab project from The Isaac Lab Project Developers team. The project is in the process of being public available and is currently build in this image with BSD-3-Clause license. \n\n`recipe/vla/run_simpleVLA_libero_grpo.sh` is the example of training SimpleVLA-OFT with this image:\n\n`vemlp-cn-shanghai.cr.volces.com/preset-images/verl_vla:preview_vla_0.1`\n\n## Disaggregation Mode for Train-Rollout / Simulation\n\nDisaggregate Train-Rollout workers and Simulation workers into different nodes.\n\nTo enable disaggregation mode for Train-Rollout nodes and Simulation nodes, we need to establish ray connection before running verl.\n* On Train-Rollout node (default main node):\n```shell\nray start --head --dashboard-host=0.0.0.0 --resources='{\"train_rollout\": 1}'\n```\n* On Simulation node:\n```shell\nray start --address='<main_node_ip>:6379' --resources='{\"sim\": 1}'\n```\n\nThen run verl on main node **only**. See `run_simpleVLA_isaac_disagg.sh` for example.\n- `env.disagg_sim.enable=True` enable disagg mode\n- `trainer.n_env_gpus_per_node` GPUs for simulaton per node\n- `trainer.n_rollout_gpus_per_node` GPUs for train-rollout node\n- `env.disagg_sim.nnodes` sim node num\n- `trainer.nnodes` train-rollout node num\n\n*Tips: you can run the following command on the sim node to check whether sim workers are scheduled up*\n```shell\npython -c \"import ray; ray.init(address=\\\"<main_node_ip>:6379\\\"); print(ray._private.state.available_resources_per_node())\"\n```\n*If you see output pattern like \"'train_rollout': 0.9992\" and \"'sim': 0.9992\", the sim workers are scheduled up successfully*\n*The actual value depends on your GPUs per node, usually <1 - 1e-4 * num_gpus>*\n\n**References:**\n*   [https://github.com/PRIME-RL/SimpleVLA-RL](https://github.com/PRIME-RL/SimpleVLA-RL)\n*   [https://github.com/RLinf/RLinf](https://github.com/RLinf/RLinf)"
  },
  {
    "path": "verl/experimental/vla/config/rob_ppo_trainer.yaml",
    "content": "# the rob_ppo config will override default ppo_trainer.yaml\n\nhydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\nenv:\n  rollout:\n    pipeline_stage_num: 2\n  actor:\n    model:\n      num_action_chunks: 8\n      action_dim: 7\n  train:\n    simulator_type: libero\n    max_episode_steps: 512\n    reward_coef: 1.0\n    only_eval: False\n    video_cfg:\n      save_video: True\n      video_base_dir: /tmp/videos\n    num_envs: 16\n    seed: 42\n    task_suite_name: libero_10\n    init_params:\n      camera_depths: False\n      camera_heights: 256\n      camera_widths: 256\n      camera_names: \n        - agentview\n        - robot0_eye_in_hand\n    \n    # Profile the env worker\n    profiler:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.ProfilerConfig\n\n      # Profiling tool to use\n      # options: nsys, npu, torch, torch_memory\n      # Defaults to global_profiler.tool if set\n      tool: ${oc.select:global_profiler.tool,null}\n\n      # Whether to enable profiling for env worker\n      enable: False\n\n      # Whether to profile all ranks\n      all_ranks: False\n\n      # List of ranks to profile (empty means no specific ranks)\n      ranks: []\n\n      # Path to save profiling results\n      # Defaults to global_profiler.save_path if set\n      save_path: ${oc.select:global_profiler.save_path,null}\n\n      # Tool-specific configurations\n      tool_config:\n\n        # nsys tool config\n        nsys:\n\n          # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n          _target_: verl.utils.profiler.config.NsightToolConfig\n        \n          # True for each task has its own database, False for all tasks in one training step share one database.\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        \n        # npu config\n        npu:\n\n          # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n          _target_: verl.utils.profiler.config.NPUToolConfig\n\n          # Contents to profile, can be empty\n          # options: npu, cpu, memory, shapes, module, stack\n          contents: []\n\n          # Collection level, optional values: level_none, level0, level1, level2.\n          level: \"level1\"\n\n          # Whether to automatically parse the data.\n          analysis: True\n\n          # True for each task has its own database, False for all tasks in one training step share one database.\n          discrete: False\n        \n        # torch profiler config\n        torch:\n\n          # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n          # Contents to profile, can be empty\n          # options: cuda, cpu, memory, shapes, stack\n          contents: []\n\n          # True for each task has its own database, False for all tasks in one training step share one database.\n          discrete: False\n\n\n        # torch memory profiler config\n        torch_memory:\n\n          # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n\n          # Maximum number of memory allocation entries to track\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n\n          # Stack trace depth for memory allocations\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n  disagg_sim:\n    enable: False\n    nnodes: 1\n\n\nactor_rollout_ref:\n  actor:\n    num_images_in_input: 1\n    traj_mini_batch_size: 16\n    fsdp_config:\n      wrap_policy:\n        transformer_layer_cls_to_wrap: \n          - PrismaticProjector\n          - LlamaDecoderLayer\n        min_num_params: 0\n      param_offload: False\n      optimizer_offload: False\n      forward_prefetch: True\n      fsdp_size: -1\n  rollout:\n    mode: async_envloop\n    prompt_length: 512\n"
  },
  {
    "path": "verl/experimental/vla/config/rob_sac_trainer.yaml",
    "content": "# the rob_ppo config will override default ppo_trainer.yaml\n\nhydra:\n  searchpath:\n    - file://verl/trainer/config\n\ndefaults:\n  - ppo_trainer\n  - _self_\n\nenv:\n  rollout:\n    pipeline_stage_num: 2\n  actor:\n    model:\n      num_action_chunks: 8\n      action_dim: 7\n  train:\n    simulator_type: libero\n    max_episode_steps: 512\n    reward_coef: 1.0\n    step_penalty: 0.001\n    only_eval: False\n    video_cfg:\n      save_video: True\n      video_base_dir: /tmp/videos\n    num_envs: 16\n    seed: 42\n    task_suite_name: libero_spatial\n    init_params:\n      camera_depths: False\n      camera_heights: 256\n      camera_widths: 256\n      camera_names:\n        - agentview\n        - robot0_eye_in_hand\n  disagg_sim:\n    enable: False\n    nnodes: 1\n\n\nactor_rollout_ref:\n  model:\n    override_config:\n      sac_enable: True\n      flow_sde_enable: True\n      flow_sde_noise_level: 0.065\n      flow_sde_rollout_noise_scale: 1.0\n      flow_sde_train_noise_scale: 1.0\n      flow_sde_initial_beta: 1.0\n      flow_sde_beta_min: 0.02\n      flow_sde_beta_schedule_T: 4000\n  actor:\n    sac:\n      gamma: 0.99\n      tau: 1.0\n      bc_loss_coef: 0\n      initial_alpha: 0\n      critic_replay_positive_sample_ratio: 0.5\n      actor_replay_positive_sample_ratio: 0.8\n      auto_entropy: False\n      alpha_type: exp\n      alpha_lr: 0.0003\n      target_entropy: -64.0\n    critic_lr: 0.0001\n    critic_weight_decay: 0\n    warm_rollout_steps: 23\n    critic_warmup_steps: 200\n    actor_update_interval: 1\n    actor_ema_enabled: true\n    actor_ema_decay: 0.95\n    replay_pool_save_interval: 500\n    num_images_in_input: 1\n    traj_mini_batch_size: 16\n    replay_pool_single_size: 2000\n    replay_pool_save_dir: /tmp/replay_pools\n    fsdp_config:\n      wrap_policy:\n        transformer_layer_cls_to_wrap:\n          - PrismaticProjector\n          - LlamaDecoderLayer\n        min_num_params: 0\n      param_offload: False\n      optimizer_offload: False\n      forward_prefetch: True\n      fsdp_size: -1\n  rollout:\n    mode: async_envloop\n    prompt_length: 512\n"
  },
  {
    "path": "verl/experimental/vla/dp_rob.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nSingle Process Actor\n\"\"\"\n\nimport logging\n\nimport torch\nfrom tensordict.base import TensorDictBase\nfrom torch import nn\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.protocol import DataProto\nfrom verl.trainer.ppo import core_algos\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import prepare_dynamic_batch, restore_dynamic_batch\nfrom verl.utils.torch_functional import logprobs_from_logits\nfrom verl.workers.actor import BasePPOActor\n\nlogger = logging.getLogger(__name__)\n\n__all__ = [\"RobDataParallelPPOActor\"]\n\n\nclass RobDataParallelPPOActor(BasePPOActor):\n    def __init__(\n        self,\n        config,\n        actor_module: nn.Module,\n        actor_optimizer: torch.optim.Optimizer = None,\n    ):\n        \"\"\"When optimizer is None, it is Reference Policy\"\"\"\n        super().__init__(config)\n        self.actor_module = actor_module\n        self.actor_optimizer = actor_optimizer\n        self.use_remove_padding = self.config.get(\"use_remove_padding\", False)\n        logger.info(f\"Actor use_remove_padding={self.use_remove_padding}\")\n        logger.info(f\"PRM use dynamic bsz={self.config.get('use_dynamic_bsz', False)}\")\n        self.ulysses_sequence_parallel_size = self.config.ulysses_sequence_parallel_size\n        self.use_ulysses_sp = False  # self.ulysses_sequence_parallel_size > 1\n        self.compute_entropy_from_logits = torch.compile(verl_F.entropy_from_logits, dynamic=True)\n\n    def process_tensor(self, tensor, pad_id):\n        mask = tensor != pad_id\n        if not torch.all(mask == mask[0:1], dim=1).all():\n            raise ValueError(\"Padding error!\")\n        base_mask = mask[0]\n        valid_len = base_mask.sum().item()\n        return tensor[:, base_mask], valid_len\n\n    def generate_traj_mask(self, end_step, traj_len):\n        \"\"\"\n        Args:\n            end_step: (batch_size,),\n            traj_len:\n        Returns:\n            mask: (batch_size, traj_len),\n        \"\"\"\n        steps = torch.arange(traj_len, device=end_step.device)  # (traj_len,)\n        steps_expanded = steps.unsqueeze(0).expand(end_step.size(0), -1)\n        mask = steps_expanded < end_step.unsqueeze(1)  # (batch_size, traj_len)\n        return mask\n\n    def apply_mask_with_grad_control(self, log_probs, entropy, mask):\n        \"\"\"\n        Args:\n            log_probs: (batch_size, 7*8)\n            entropy:   (batch_size, 7*8)\n            # mask:      (batch_size, 8)\n            mask:      (batch_size, 7*8)\n        Returns:\n            log_probs_masked:\n            entropy_masked:\n        \"\"\"\n\n        mask = mask.to(log_probs.device)\n        log_probs_masked = torch.where(mask, log_probs, torch.zeros_like(log_probs, requires_grad=False))\n        entropy_masked = torch.where(mask, entropy, torch.zeros_like(entropy, requires_grad=False))\n        return log_probs_masked, entropy_masked\n\n    def _forward_micro_batch(self, micro_batch, temperature) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"\n        micro_batch:\n\n        Returns:\n            entropy: # (bs, response_len)\n            log_probs: # (bs, response_len)\n        \"\"\"\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            input_ids = micro_batch[\"input_ids\"]\n            attention_mask = micro_batch[\"attention_mask\"]\n            pixel_values = micro_batch[\"pixel_values\"]\n            responses = micro_batch[\"responses\"]\n\n            input_ids_unpad, _ = self.process_tensor(input_ids, self.pad_token_id)\n            attention_mask_unpad, _ = self.process_tensor(attention_mask, 0)\n\n            logits = self.actor_module(\n                input_ids=input_ids_unpad,\n                attention_mask=attention_mask_unpad,\n                pixel_values=pixel_values,\n            )  # prevent model thinks we are generating\n\n            assert self.actor_module.vocab_size == 32000\n            start_index = self.actor_module.vocab_size - 256\n            logits = logits[..., -256 - 64 : -64]  # Shape: [batch_size, seq_len, 256]\n            responses = responses - start_index\n            # assert (0<=responses<=255).all()\n\n            logits = logits.div(temperature)\n\n            log_probs = logprobs_from_logits(logits, responses.to(logits.device))\n            entropy = verl_F.entropy_from_logits(logits)  # (bsz, response_length)\n\n            # assert len(log_probs.shape) == 2 and len(entropy.shape) == 2\n\n            # TODO(caiyunke.astra): check here\n\n            mask = micro_batch[\"response_mask\"]\n            log_probs, entropy = self.apply_mask_with_grad_control(log_probs, entropy, mask)\n\n            return entropy, log_probs\n\n    def _forward_micro_batch_update(\n        self, input_ids, attention_mask, pixel_values, responses, temperature\n    ) -> tuple[torch.Tensor, torch.Tensor]:\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            input_ids_unpad, _ = self.process_tensor(input_ids, self.pad_token_id)\n            attention_mask_unpad, _ = self.process_tensor(attention_mask, 0)\n\n            logits = self.actor_module(\n                input_ids=input_ids_unpad,\n                attention_mask=attention_mask_unpad,\n                pixel_values=pixel_values,\n            )\n\n            assert logits.requires_grad\n\n            assert self.actor_module.vocab_size == 32000\n            start_index = self.actor_module.vocab_size - 256\n            logits = logits[..., -256 - 64 : -64]  # Shape: [batch_size, seq_len, 256]\n            responses = responses - start_index\n\n            logits = logits.div(temperature)\n\n            log_probs = logprobs_from_logits(logits, responses)\n            entropy = verl_F.entropy_from_logits(logits)  # (bsz, response_length)\n            return entropy, log_probs\n\n    def _optimizer_step(self):\n        assert self.config.grad_clip is not None\n\n        if isinstance(self.actor_module, FSDP):\n            grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip)\n        self.actor_optimizer.step()\n        return grad_norm\n\n    def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n        \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids\n\n        Args:\n            data (DataProto): a DataProto containing keys\n\n                ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the\n                concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``.\n\n                ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``responses``:  tensor of shape [batch_size, response_length]. torch.int64.\n\n        Returns:\n            torch.Tensor: the log_prob tensor\n        \"\"\"\n        self.actor_module.eval()\n\n        micro_batch_size = data.meta_info[\"micro_batch_size\"]  # 256\n        temperature = data.meta_info[\n            \"temperature\"\n        ]  # temperature must be in the data.meta_info to avoid slient error # 1\n        use_dynamic_bsz = data.meta_info[\"use_dynamic_bsz\"]  # trues\n        self.pad_token_id = data.meta_info[\"pad_token_id\"]\n\n        select_keys = [\"responses\", \"input_ids\", \"attention_mask\", \"pixel_values\", \"response_mask\"]\n        data = data.select(batch_keys=select_keys).batch\n\n        if use_dynamic_bsz:\n            max_token_len = data.meta_info[\"max_token_len\"] * self.ulysses_sequence_parallel_size\n            micro_batches, batch_idx_list = prepare_dynamic_batch(data, max_token_len=max_token_len)\n        else:\n            micro_batches = data.split(micro_batch_size)\n\n        log_probs_lst = []\n        entropy_lst = []\n        for micro_batch in micro_batches:\n            with torch.no_grad():\n                entropy, log_probs = self._forward_micro_batch(micro_batch, temperature=temperature)\n            log_probs_lst.append(log_probs)\n            if calculate_entropy:\n                entropy_lst.append(entropy)\n        log_probs = torch.concat(log_probs_lst, dim=0)\n        entropys = None\n        if calculate_entropy:\n            entropys = torch.concat(entropy_lst, dim=0)\n\n        if use_dynamic_bsz:\n            log_probs = restore_dynamic_batch(log_probs, batch_idx_list)\n            if calculate_entropy:\n                entropys = restore_dynamic_batch(entropys, batch_idx_list)\n\n        return {\"log_probs\": log_probs, \"entropys\": entropys}\n\n    def update_policy(self, data: DataProto):\n        self.actor_module.train()\n\n        assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size_per_gpu == 0\n        self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu\n        temperature = data.meta_info[\"temperature\"]  # temperature must be in the data.meta_info to avoid slient error\n\n        select_keys = [\n            \"responses\",\n            \"response_mask\",\n            \"input_ids\",\n            \"attention_mask\",\n            \"pixel_values\",\n            \"old_log_probs\",\n            \"advantages\",\n        ]\n        batch = data.select(batch_keys=select_keys).batch\n        self.pad_token_id = data.meta_info[\"pad_token_id\"]\n        # TODO(caiyunke.astra): check here\n        # assert self.config.ppo_micro_batch_size_per_gpu == 1\n\n        # Split to make minibatch iterator for updating the actor\n        # See PPO paper for details. https://arxiv.org/abs/1707.06347\n        mini_batches = batch.split(self.config.ppo_mini_batch_size)\n        metrics = {}\n        for batch_idx, mini_batch in enumerate(mini_batches):\n            if self.config.use_dynamic_bsz:\n                max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size\n                micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len)\n            else:\n                self.gradient_accumulation = self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu\n                micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)\n\n            self.actor_optimizer.zero_grad()\n\n            for _, micro_batch in enumerate[DataProto | TensorDictBase](micro_batches):\n                micro_batch = micro_batch.to(get_device_id())  # actor device is cpu when using offload\n                responses = micro_batch[\"responses\"]\n\n                response_mask = micro_batch[\"response_mask\"]  # (batch_size, traj_len)\n\n                old_log_prob = micro_batch[\"old_log_probs\"]\n                advantages = micro_batch[\"advantages\"]\n\n                # clip_ratio = self.config.clip_ratio\n                clip_ratio_high = self.config.clip_ratio_high\n                clip_ratio_low = self.config.clip_ratio_low\n\n                input_ids = micro_batch[\"input_ids\"]\n                attention_mask = micro_batch[\"attention_mask\"]\n                pixel_values = micro_batch[\"pixel_values\"]\n                responses = micro_batch[\"responses\"]\n\n                loss_info = {\n                    \"actor/pg_loss\": 0,\n                    \"actor/pg_clipfrac\": 0,\n                    \"actor/ppo_kl\": 0,\n                    \"actor/pg_clipfrac_lower\": 0,\n                }\n\n                _, log_prob = self._forward_micro_batch_update(\n                    input_ids=input_ids,\n                    attention_mask=attention_mask,\n                    pixel_values=pixel_values,\n                    responses=responses,\n                    temperature=temperature,\n                )\n\n                pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower = core_algos.compute_policy_loss(\n                    old_log_prob=old_log_prob,\n                    log_prob=log_prob,\n                    advantages=advantages,\n                    response_mask=response_mask,\n                    cliprange_high=clip_ratio_high,\n                    cliprange_low=clip_ratio_low,\n                )\n                loss = pg_loss / self.gradient_accumulation\n\n                loss.backward()\n\n                loss_info[\"actor/pg_loss\"] = loss_info[\"actor/pg_loss\"] + pg_loss.detach().item()\n                loss_info[\"actor/pg_clipfrac\"] = loss_info[\"actor/pg_clipfrac\"] + pg_clipfrac.detach().item()\n                loss_info[\"actor/ppo_kl\"] = loss_info[\"actor/ppo_kl\"] + ppo_kl.detach().item()\n                loss_info[\"actor/pg_clipfrac_lower\"] = (\n                    loss_info[\"actor/pg_clipfrac_lower\"] + pg_clipfrac_lower.detach().item()\n                )\n                append_to_dict(metrics, loss_info)\n\n            grad_norm = self._optimizer_step()\n            mini_batch_metrics = {\"actor/grad_norm\": grad_norm.detach().item()}\n            append_to_dict(metrics, mini_batch_metrics)\n        self.actor_optimizer.zero_grad()\n        return metrics\n"
  },
  {
    "path": "verl/experimental/vla/env_loop.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport logging\nimport os\n\nimport numpy as np\nimport torch\nfrom omegaconf import DictConfig\n\nfrom verl import DataProto\nfrom verl.single_controller.ray import RayWorkerGroup\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass EnvLoop:\n    \"\"\"An env loop manages interactions between models and vectorized environments. It's designed for computationally\n    intensive environments, such as robotics simulators.\"\"\"\n\n    def __init__(self, env_wg: RayWorkerGroup, rollout_wg: RayWorkerGroup, config: DictConfig):\n        \"\"\"\n        Initialize the EnvLoop.\n\n        Args:\n            env_wg (RayWorkerGroup): Environment worker group.\n            rollout_wg (RayWorkerGroup): Rollout worker group for model inference.\n            config (DictConfig): YAML config.\n        \"\"\"\n        self.env_wg = env_wg\n        self.rollout_wg = rollout_wg\n        self.config = config\n        # Extract relevant configuration\n        self.max_interactions = config.env.train.max_episode_steps // config.env.actor.model.num_action_chunks\n        self.stage_num = config.env.rollout.pipeline_stage_num\n        self.num_envs_per_worker = config.env.train.num_envs\n        self.action_dim = config.env.actor.model.action_dim\n        self.num_action_chunks = config.env.actor.model.num_action_chunks\n        # Derived properties\n        self.total_envs = self.env_wg.world_size * self.num_envs_per_worker\n        if self.total_envs % self.stage_num != 0:\n            raise ValueError(f\"Total envs ({self.total_envs}) must be divisible by stage_num ({self.stage_num})\")\n        self.envs_per_stage = self.total_envs // self.stage_num\n\n        self.env_wg.init_worker()\n        self.env_wg.init_simulator()\n\n    def generate_sequences(self, prompts: DataProto, reset_future: asyncio.Future) -> DataProto:\n        \"\"\"Split input batch and dispatch to env loop workers.\n\n        Args:\n            prompts (DataProto): Input batch.\n\n        Returns:\n            DataProto: Output batch.\n        \"\"\"\n\n        reset_results = reset_future.get()\n\n        loop = asyncio.get_event_loop()\n        self.rollout_wg.switch_to_rollout()\n        output = loop.run_until_complete(self.run(prompts, reset_results))\n        self.rollout_wg.switch_to_train()\n        # TODO(caiyunke.astra): add timing metrics\n        return output\n\n    async def run(self, prompts: DataProto, reset_results: DataProto) -> DataProto:\n        \"\"\"\n        Run the environment interaction loop.\n        This method orchestrates a pipelined process:\n        1. Resets environments to specified initial states.\n        2. In a loop, it gets actions from the rollout workers and applies them to the environments.\n        3. Collects all trajectory data (observations, actions, rewards, dones).\n        4. Formats and returns the collected trajectories as a single batch.\n        Args:\n            prompts (DataProto): Contains initial state IDs and other settings.\n                                 - 'non_tensor_batch.state_ids': A numpy array of state IDs to reset envs.\n        Returns:\n            DataProto: A batch containing the complete trajectories.\n        \"\"\"\n        initial_state_ids = prompts.non_tensor_batch[\"state_ids\"]\n\n        staged_obs = self._restructure_obs_data(reset_results)\n        # --- Pipeline state ---\n        trajectories = {i: [] for i in range(self.stage_num)}  # To store (obs, action, rew, done) tuples\n        rollout_futures = {}\n        # is_complete = torch.zeros((self.total_envs,), dtype=torch.bool)\n\n        for stage_id in range(self.stage_num):\n            # trajectories[stage_id].append({'obs': staged_obs[stage_id]})\n            trajectories[stage_id].append({})\n            vla_input = staged_obs[stage_id]\n            vla_input.meta_info = prompts.meta_info  # Pass along rollout config\n            rollout_futures[stage_id] = self.rollout_wg.generate_sequences(vla_input)\n\n        async def _stage_loop(stage_id: int):\n            for step_idx in range(self.max_interactions):\n                if stage_id == 0:\n                    logger.info(f\"[{step_idx}/{self.max_interactions - 1}] rollout step\")\n                action_result: DataProto = await asyncio.to_thread(rollout_futures[stage_id].get)\n\n                trajectories[stage_id][-1][\"action\"] = action_result\n                action_data = DataProto.from_dict(\n                    non_tensors={\n                        \"actions\": action_result.batch[\"action\"].cpu().numpy(),\n                        \"critic_values\": action_result.batch[\"critic_value\"].cpu().numpy(),\n                    },\n                    meta_info={\"stage_id\": stage_id},\n                )\n\n                env_ref = self.env_wg.env_interact_step(action_data)\n                env_result: DataProto = await asyncio.to_thread(env_ref.get)\n\n                trajectories[stage_id][-1][\"rew\"] = env_result.batch[\"rews\"]\n                trajectories[stage_id][-1][\"done\"] = env_result.batch[\"terminations\"]\n\n                next_obs = DataProto(\n                    batch=env_result.batch.select(\"full_image\", \"wrist_image\", \"state\"),\n                    non_tensor_batch={\"task_descriptions\": env_result.non_tensor_batch[\"task_descriptions\"]},\n                )\n\n                if step_idx < self.max_interactions - 1:\n                    trajectories[stage_id].append({})\n                    vla_input = next_obs\n                    vla_input.meta_info = prompts.meta_info\n                    rollout_futures[stage_id] = self.rollout_wg.generate_sequences(vla_input)\n\n        await asyncio.gather(*[asyncio.create_task(_stage_loop(sid)) for sid in range(self.stage_num)])\n        self.env_wg.finish_rollout()\n\n        return self._collate_trajectories(trajectories, initial_state_ids, meta_info=prompts.meta_info)\n\n    def _restructure_obs_data(self, data_proto: DataProto) -> list[DataProto]:\n        \"\"\"Reshapes flat observation data from env_wg into a list of per-stage DataProto objects.\"\"\"\n        # env_wg returns a flat batch ordered by [worker0_stage0, worker0_stage1, ...,\n        # worker1_stage0, worker1_stage1, ...]\n        # First, un-flatten by worker, then by stage\n\n        num_workers = self.env_wg.world_size\n\n        staged_data = [[] for _ in range(self.stage_num)]\n        chunks = data_proto.chunk(num_workers)\n        for worker_chunk in chunks:\n            stage_chunks = worker_chunk.chunk(self.stage_num)\n            for stage_id, data in enumerate(stage_chunks):\n                staged_data[stage_id].append(data)\n\n        # Concatenate data from all workers for each stage\n        return [DataProto.concat(data_list) for data_list in staged_data]\n\n    def _collate_trajectories(self, trajectories: dict, initial_state_ids: np.ndarray, meta_info) -> DataProto:\n        \"\"\"\n        Collates the collected trajectory data into the final batch format.\n        \"\"\"\n        flat_trajs = [{} for _ in range(len(trajectories[0]))]\n        for stage_id in range(self.stage_num):\n            for step_idx, step_data in enumerate(trajectories[stage_id]):\n                if not flat_trajs[step_idx]:  # if dict is empty\n                    flat_trajs[step_idx] = step_data\n                else:\n                    # Concatenate DataProto objects\n                    for key, value in step_data.items():\n                        if isinstance(value, DataProto):\n                            flat_trajs[step_idx][key] = DataProto.concat([flat_trajs[step_idx][key], value])\n                        elif isinstance(value, torch.Tensor):\n                            flat_trajs[step_idx][key] = torch.cat([flat_trajs[step_idx][key], value], dim=0)\n\n        # iterate all action batch keys (e.g., action, images, pixel_values, input_ids, ...)\n        batch_dict = {}\n        action_batch_keys = list(flat_trajs[0][\"action\"].batch.keys())\n        for key in action_batch_keys:\n            per_step_values = [step[\"action\"].batch[key] for step in flat_trajs]\n            batch_dict[key] = torch.stack(per_step_values, dim=1)\n\n        batch_dict[\"complete\"] = torch.stack([step[\"done\"] for step in flat_trajs], dim=1).squeeze(-1)\n        batch_dict[\"env_state_id\"] = torch.from_numpy(initial_state_ids.astype(int))\n\n        return DataProto.from_single_dict(batch_dict, meta_info=meta_info)\n"
  },
  {
    "path": "verl/experimental/vla/envs/__init__.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/vla/envs/action_utils.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom io import BytesIO\nfrom typing import Any, Optional\n\nimport imageio\nimport numpy as np\nimport torch\nimport torchvision.transforms.functional as F\nfrom PIL import Image, ImageDraw, ImageFont\n\n\ndef prepare_actions_simplevla(\n    raw_chunk_actions,\n) -> torch.Tensor:\n    from verl.experimental.vla.envs.libero_env.utils import invert_gripper_action, normalize_gripper_action\n\n    normalized_action = normalize_gripper_action(raw_chunk_actions, binarize=True)\n    inverted_action = invert_gripper_action(normalized_action)\n    return inverted_action\n\n\ndef prepare_actions(\n    simulator_type,\n    raw_chunk_actions,\n    num_action_chunks,\n    action_dim,\n    action_scale: float = 1.0,\n    policy: str = \"widowx_bridge\",\n) -> torch.Tensor:\n    # TODO: prepare_actions according to simulator_type\n    chunk_actions = prepare_actions_simplevla(\n        raw_chunk_actions=raw_chunk_actions,\n    )\n\n    return chunk_actions\n\n\ndef to_tensor(array: dict | torch.Tensor | np.ndarray | list | Any, device: str = \"cpu\") -> dict | torch.Tensor:\n    \"\"\"\n    Copied from ManiSkill!\n    Maps any given sequence to a torch tensor on the CPU/GPU. If physx gpu\n    is not enabled then we use CPU, otherwise GPU, unless specified\n    by the device argument\n\n    Args:\n        array: The data to map to a tensor\n        device: The device to put the tensor on. By default this is None\n        and to_tensor will put the device on the GPU if physx is enabled\n            and CPU otherwise\n\n    \"\"\"\n    if isinstance(array, (dict)):\n        return {k: to_tensor(v, device=device) for k, v in array.items()}\n    elif isinstance(array, torch.Tensor):\n        ret = array.to(device)\n    elif isinstance(array, np.ndarray):\n        if array.dtype == np.uint16:\n            array = array.astype(np.int32)\n        elif array.dtype == np.uint32:\n            array = array.astype(np.int64)\n        ret = torch.tensor(array).to(device)\n    else:\n        if isinstance(array, list) and isinstance(array[0], np.ndarray):\n            array = np.array(array)\n        ret = torch.tensor(array, device=device)\n    if ret.dtype == torch.float64:\n        ret = ret.to(torch.float32)\n    return ret\n\n\ndef tile_images(images: list[np.ndarray | torch.Tensor], nrows: int = 1) -> np.ndarray | torch.Tensor:\n    \"\"\"\n    Copied from maniskill https://github.com/haosulab/ManiSkill\n    Tile multiple images to a single image comprised of nrows and an\n    appropriate number of columns to fit all the images.\n    The images can also be batched (e.g. of shape (B, H, W, C)), but\n    give images must all have the same batch size.\n\n    if nrows is 1, images can be of different sizes. If nrows > 1,\n    they must all be the same size.\n    \"\"\"\n    # Sort images in descending order of vertical height\n    batched = False\n    if len(images[0].shape) == 4:\n        batched = True\n    if nrows == 1:\n        images = sorted(images, key=lambda x: x.shape[0 + batched], reverse=True)\n\n    columns: list[list[np.ndarray | torch.Tensor]] = []\n    if batched:\n        max_h = images[0].shape[1] * nrows\n        cur_h = 0\n        cur_w = images[0].shape[2]\n    else:\n        max_h = images[0].shape[0] * nrows\n        cur_h = 0\n        cur_w = images[0].shape[1]\n\n    # Arrange images in columns from left to right\n    column = []\n    for im in images:\n        if cur_h + im.shape[0 + batched] <= max_h and cur_w == im.shape[1 + batched]:\n            column.append(im)\n            cur_h += im.shape[0 + batched]\n        else:\n            columns.append(column)\n            column = [im]\n            cur_h, cur_w = im.shape[0 + batched : 2 + batched]\n    columns.append(column)\n\n    # Tile columns\n    total_width = sum(x[0].shape[1 + batched] for x in columns)\n\n    is_torch = False\n    if torch is not None:\n        is_torch = isinstance(images[0], torch.Tensor)\n\n    output_shape = (max_h, total_width, 3)\n    if batched:\n        output_shape = (images[0].shape[0], max_h, total_width, 3)\n    if is_torch:\n        output_image = torch.zeros(output_shape, dtype=images[0].dtype)\n    else:\n        output_image = np.zeros(output_shape, dtype=images[0].dtype)\n    cur_x = 0\n    for column in columns:\n        cur_w = column[0].shape[1 + batched]\n        next_x = cur_x + cur_w\n        if is_torch:\n            column_image = torch.concatenate(column, dim=0 + batched)\n        else:\n            column_image = np.concatenate(column, axis=0 + batched)\n        cur_h = column_image.shape[0 + batched]\n        output_image[..., :cur_h, cur_x:next_x, :] = column_image\n        cur_x = next_x\n    return output_image\n\n\ndef put_text_on_image(image: np.ndarray, lines: list[str], max_width: int = 200) -> np.ndarray:\n    \"\"\"\n    Put text lines on an image with automatic line wrapping.\n\n    Args:\n        image: Input image as numpy array\n        lines: List of text lines to add\n        max_width: Maximum width for text wrapping\n    \"\"\"\n    assert image.dtype == np.uint8, image.dtype\n    image = image.copy()\n    image = Image.fromarray(image)\n    draw = ImageDraw.Draw(image)\n    font = ImageFont.load_default(size=20)\n\n    new_lines = []\n    for line in lines:\n        words = line.split()\n        current_line = []\n\n        for word in words:\n            test_line = \" \".join(current_line + [word])\n            test_width = font.getlength(test_line)\n\n            if test_width <= max_width:\n                current_line.append(word)\n            else:\n                new_lines.append(\" \".join(current_line))\n                current_line = [word]\n        if current_line:\n            new_lines.append(\" \".join(current_line))\n\n    y = -10\n    for line in new_lines:\n        bbox = draw.textbbox((0, 0), text=line)\n        textwidth = bbox[2] - bbox[0]\n        textheight = bbox[3] - bbox[1]\n        y += textheight + 10\n        x = 10\n        pad = 2\n        draw.rectangle(\n            [(x - pad, y - pad), (x + textwidth + pad, y + textheight + pad)],\n            fill=(0, 0, 0),\n        )\n        draw.text((x, y), text=line, fill=(255, 255, 255))\n    return np.array(image)\n\n\ndef put_info_on_image(\n    image: np.ndarray,\n    info: dict[str, float],\n    extras: Optional[list[str]] = None,\n    overlay: bool = True,\n) -> np.ndarray:\n    \"\"\"\n    Put information dictionary and extra lines on an image.\n\n    Args:\n        image: Input image\n        info: Dictionary of key-value pairs to display\n        extras: Additional text lines to display\n        overlay: Whether to overlay text on image\n    \"\"\"\n    lines = [f\"{k}: {v:.3f}\" if isinstance(v, float) else f\"{k}: {v}\" for k, v in info.items()]\n    if extras is not None:\n        lines.extend(extras)\n    return put_text_on_image(image, lines)\n\n\ndef list_of_dict_to_dict_of_list(\n    list_of_dict: list[dict[str, Any]],\n) -> dict[str, list[Any]]:\n    \"\"\"\n    Convert a list of dictionaries to a dictionary of lists.\n\n    Args:\n        list_of_dict: List of dictionaries with same keys\n\n    Returns:\n        Dictionary where each key maps to a list of values\n    \"\"\"\n    if len(list_of_dict) == 0:\n        return {}\n    keys = list_of_dict[0].keys()\n    output = {key: [] for key in keys}\n    for data in list_of_dict:\n        for key, item in data.items():\n            assert key in output\n            output[key].append(item)\n    return output\n\n\ndef save_rollout_video(rollout_images: list[np.ndarray], output_dir: str, video_name: str, fps: int = 30) -> None:\n    \"\"\"\n    Saves an MP4 replay of an episode.\n\n    Args:\n        rollout_images: List of images from the episode\n        output_dir: Directory to save the video\n        video_name: Name of the output video file\n        fps: Frames per second for the video\n    \"\"\"\n    os.makedirs(output_dir, exist_ok=True)\n    mp4_path = os.path.join(output_dir, f\"{video_name}.mp4\")\n    video_writer = imageio.get_writer(mp4_path, fps=fps)\n    for img in rollout_images:\n        video_writer.append_data(img)\n    video_writer.close()\n\n\ndef resize_image(img: np.ndarray, resize_size: tuple[int, int]) -> np.ndarray:\n    \"\"\"\n    Takes numpy array corresponding to a single image and returns resized image as numpy array.\n\n    Args:\n        img: Input image as numpy array\n        resize_size: Target size for resizing\n\n    Returns:\n        Resized image as numpy array\n    \"\"\"\n\n    assert isinstance(resize_size, tuple), \"resize_size must be a tuple\"\n    assert isinstance(img, np.ndarray), \"img must be a numpy array\"\n\n    # Convert numpy array to PIL Image\n    pil_img = Image.fromarray(img)\n\n    # Encode as JPEG, as done in RLDS dataset builder\n    buffer = BytesIO()\n    pil_img.save(buffer, format=\"JPEG\")\n    buffer.seek(0)\n\n    # Immediately decode back\n    img = Image.open(buffer)\n\n    img = img.resize(resize_size, Image.Resampling.LANCZOS)\n    img = np.array(img)\n    img = np.clip(np.round(img), 0, 255).astype(np.uint8)\n\n    return img\n\n\ndef center_crop_image(image: Image.Image) -> Image.Image:\n    crop_scale = 0.9\n    orig_w, orig_h = image.size\n    image_tensor = F.to_tensor(image)\n    crop_h = int(orig_h * crop_scale)\n    crop_w = int(orig_w * crop_scale)\n    image_tensor = F.center_crop(image_tensor, (crop_h, crop_w))\n    image_tensor = F.resize(image_tensor, (orig_h, orig_w))\n    final_image = F.to_pil_image(image_tensor)\n\n    final_image = final_image.convert(\"RGB\")\n    return final_image\n"
  },
  {
    "path": "verl/experimental/vla/envs/isaac_env/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .isaac_env import IsaacEnv\n\n__all__ = [\"IsaacEnv\"]\n"
  },
  {
    "path": "verl/experimental/vla/envs/isaac_env/isaac_env.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Optional\n\nimport gymnasium as gym\nimport numpy as np\nimport omni\nimport torch\n\nfrom verl.experimental.vla.envs.action_utils import (\n    put_info_on_image,\n    save_rollout_video,\n    tile_images,\n    to_tensor,\n)\n\nlogger = logging.getLogger(__name__)\n\n\nclass IsaacEnv(gym.Env):\n    def __init__(self, cfg, rank, world_size):\n        self.rank = rank\n        self.cfg = cfg\n        self.world_size = world_size\n        self.seed = self.cfg.seed + rank\n        self.num_envs = self.cfg.num_envs\n        self.action_dim = self.cfg.get(\"action_dim\", 7)\n        self.device = self.cfg.get(\"device\", \"cuda:0\")\n\n        self._generator = np.random.default_rng(seed=self.seed)\n\n        self.task_suite_name = self.cfg.task_suite_name\n\n        self.env = None\n        self.prev_step_reward = np.zeros(self.num_envs)\n        self.use_rel_reward = False\n\n        self._init_metrics()\n        self._elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)\n        self.max_episode_steps = cfg.max_episode_steps\n        self.video_cfg = cfg.video_cfg\n\n        self.render_images = []\n        self.video_cnt = 0\n        self.camera_name = cfg.init_params.camera_names\n\n        # sys env must be set before import isaaclab\n        from isaaclab.app import AppLauncher\n\n        launch_args = {\"headless\": True, \"enable_cameras\": True}\n        app_launcher = AppLauncher(**launch_args)\n        self.app = app_launcher.app\n        # force franka registration\n        import isaaclab_playground.tasks.manipulation.libero.config.franka  # noqa\n\n    def _init_env(self, task_id=0):\n        \"\"\"Initializes the Isaac Sim environment.\"\"\"\n\n        self.task_name = self.cfg.get(\"task_name\")\n        self.task_id = task_id\n        # FIXME since isaac use env to set task id, all env have to use the same task id\n        if self.task_suite_name.startswith(\"libero\"):\n            os.environ[\"LIBERO_TASK_SUITE\"] = self.task_suite_name\n            os.environ[\"LIBERO_TASK_ID\"] = str(task_id)\n            os.environ[\"LIBERO_OSC_TYPE\"] = \"pose_rel\"\n\n            if not self.task_name:\n                self.task_name = \"Isaac-Libero-Franka-OscPose-v0\"\n\n        from isaaclab_tasks.utils import parse_env_cfg\n\n        self.env_cfg = parse_env_cfg(self.task_name, num_envs=self.num_envs)\n        self.env_cfg.env_name = self.cfg.get(\"env_name\", str(self.task_id))\n        self.env_cfg.sim.device = self.device\n        self.env_cfg.sim.physx.enable_ccd = True\n        self.env_cfg.terminations.time_out = None\n        self.env_cfg.observations.policy.concatenate_terms = False\n\n        # create environment from loaded config\n        if self.env:\n            self.env.close()\n            omni.usd.get_context().new_stage()\n        self.env = gym.make(self.task_name, cfg=self.env_cfg).unwrapped\n\n        if self.cfg.video_cfg.save_video:\n            video_dir = os.path.join(self.cfg.video_cfg.video_base_dir, f\"rank_{self.rank}\")\n            os.makedirs(video_dir, exist_ok=True)\n\n        self.action_space = self.env.action_space\n        self.observation_space = self.env.observation_space\n\n        # TODO support other task suite\n        if self.task_suite_name.startswith(\"libero\"):\n            self.task_descriptions = self.env.cfg.libero_config.task_info[\"language_instruction\"]\n            assert self.env_cfg.osc_type == \"pose_rel\", (\n                f\"Only pose_rel osc type is supported for libero. Received: {self.env_cfg.osc_type}\"\n            )\n        else:\n            raise ValueError(f\"Task suite {self.task_suite_name} is not supported.\")\n        logger.info(\"Isaac Sim environment initialized\")\n\n    def _init_metrics(self):\n        self.success_once = np.zeros(self.num_envs, dtype=bool)\n        self.returns = np.zeros(self.num_envs)\n\n    def _reset_metrics(self, env_idx=None):\n        if env_idx is not None:\n            mask = np.zeros(self.num_envs, dtype=bool)\n            mask[env_idx] = True\n            self.prev_step_reward[mask] = 0.0\n            self.success_once[mask] = False\n            self.returns[mask] = 0\n            self._elapsed_steps[env_idx] = 0\n        else:\n            self.prev_step_reward[:] = 0\n            self.success_once[:] = False\n            self.returns[:] = 0.0\n            self._elapsed_steps[:] = 0\n\n    def _record_metrics(self, step_reward, terminations, infos):\n        episode_info = {}\n        self.returns += step_reward\n        # Ensure terminations is a numpy array before the bitwise OR\n        if isinstance(terminations, torch.Tensor):\n            terminations = terminations.cpu().numpy()\n        self.success_once = self.success_once | terminations\n        episode_info[\"success_once\"] = self.success_once.copy()\n        episode_info[\"return\"] = self.returns.copy()\n        episode_info[\"episode_len\"] = self.elapsed_steps.copy()\n        if any(self.elapsed_steps > 0):\n            episode_info[\"reward\"] = episode_info[\"return\"] / self.elapsed_steps\n        else:\n            episode_info[\"reward\"] = 0\n        infos[\"episode\"] = to_tensor(episode_info)\n        return infos\n\n    def reset(self, env_idx: Optional[int | list[int] | np.ndarray] = None, options: Optional[dict] = None):\n        if env_idx is None:\n            env_idx = np.arange(self.num_envs)\n\n        raw_obs, infos = self.env.reset()\n\n        obs = self._wrap_obs(raw_obs)\n\n        self._reset_metrics(env_idx)\n\n        return obs, infos\n\n    def step(self, actions=None, critic_values=None):\n        if actions is None:\n            # isaac should start with reset_envs_to_initial_state\n            # do nothing for None\n            return (None, None, None, None, None)\n\n        truncations = self.elapsed_steps >= self.max_episode_steps\n        # _actions = torch.zeros(self.action_space.shape)\n\n        if isinstance(actions, np.ndarray):\n            actions = torch.from_numpy(actions)\n\n        self._elapsed_steps += 1\n        raw_obs, _reward, terminations, _, infos = self.env.step(actions)\n        self.last_obs = raw_obs\n        self.last_infos = infos\n\n        obs = self._wrap_obs(raw_obs)\n\n        step_reward = self._calc_step_reward(_reward.cpu().numpy())\n\n        if self.video_cfg.save_video:\n            plot_infos = {\n                \"rewards\": step_reward,\n                \"terminations\": terminations,\n                \"task\": self.task_descriptions,\n            }\n            if critic_values is not None:\n                plot_infos[\"critic_value\"] = np.asarray(critic_values, dtype=np.float32)\n            self.add_new_frames(obs, plot_infos)\n\n        infos = self._record_metrics(step_reward, terminations, infos)\n\n        return (\n            obs,\n            to_tensor(step_reward),\n            to_tensor(terminations),\n            to_tensor(truncations),\n            infos,\n        )\n\n    def chunk_step(self, chunk_actions, chunk_values=None):\n        # chunk_actions: [num_envs, chunk_step, action_dim]\n        chunk_size = chunk_actions.shape[1]\n\n        chunk_rewards = []\n\n        raw_chunk_terminations = []\n        raw_chunk_truncations = []\n        for i in range(chunk_size):\n            actions = chunk_actions[:, i]\n            step_values = None\n            if chunk_values is not None:\n                if len(chunk_values.shape) == 1:\n                    step_values = chunk_values\n                elif len(chunk_values.shape) == 2:\n                    step_values = chunk_values[:, i]\n\n            extracted_obs, step_reward, terminations, truncations, infos = self.step(actions, critic_values=step_values)\n\n            chunk_rewards.append(step_reward)\n            raw_chunk_terminations.append(terminations)\n            raw_chunk_truncations.append(truncations)\n\n        chunk_rewards = torch.stack(chunk_rewards, dim=1)  # [num_envs, chunk_steps]\n        raw_chunk_terminations = torch.stack(raw_chunk_terminations, dim=1)  # [num_envs, chunk_steps]\n        raw_chunk_truncations = torch.stack(raw_chunk_truncations, dim=1)  # [num_envs, chunk_steps]\n\n        chunk_terminations = raw_chunk_terminations.clone()\n        chunk_truncations = raw_chunk_truncations.clone()\n        return (\n            extracted_obs,\n            chunk_rewards,\n            chunk_terminations,\n            chunk_truncations,\n            infos,\n        )\n\n    def _calc_step_reward(self, reward):\n        if self.use_rel_reward:\n            reward_diff = reward - self.prev_step_reward\n            self.prev_step_reward = reward\n            return reward_diff\n        else:\n            return reward\n\n    def _wrap_obs(self, raw_obs):\n        images_and_states = self._extract_image_and_state(raw_obs)\n\n        obs = {\n            \"images_and_states\": to_tensor(images_and_states),\n            \"task_descriptions\": [self.task_descriptions] * self.num_envs,\n        }\n        return obs\n\n    def _extract_image_and_state(self, obs):\n        # TODO support multiple camera\n        camera_name = self.camera_name[0]\n        for key in self.env.unwrapped.scene.keys():\n            if key.startswith(camera_name):\n                cam = self.env.unwrapped.scene[key]\n                break\n        assert cam is not None, f\"camera {camera_name} not found in scene\"\n\n        rgb = cam.data.output[\"rgb\"]\n\n        full_image = rgb.cpu().numpy()\n        return {\n            \"full_image\": full_image,\n            \"state\": np.concatenate(\n                [\n                    obs[\"policy\"][\"eef_pose\"].cpu(),\n                    # quat2axisangle(obs[\"robot0_eef_quat\"]), # isaac do not return robot0_eef_quat\n                    # obs[\"policy\"][\"gripper_pos\"].cpu(),\n                ],\n                axis=-1,\n            ),\n        }\n\n    def add_new_frames(self, obs, plot_infos):\n        images = []\n        for env_id, img in enumerate(obs[\"images_and_states\"][\"full_image\"]):\n            info_item = {k: v if np.size(v) == 1 else v[env_id] for k, v in plot_infos.items()}\n            img = put_info_on_image(img.cpu().numpy(), info_item)\n            images.append(img)\n        full_image = tile_images(images, nrows=int(np.sqrt(self.num_envs)))\n        self.render_images.append(full_image)\n\n    def flush_video(self, video_sub_dir: Optional[str] = None):\n        output_dir = os.path.join(self.video_cfg.video_base_dir, f\"rank_{self.rank}\")\n        if video_sub_dir is not None:\n            output_dir = os.path.join(output_dir, f\"{video_sub_dir}\")\n        save_rollout_video(\n            self.render_images,\n            output_dir=output_dir,\n            video_name=f\"{self.video_cnt}\",\n        )\n        self.video_cnt += 1\n        self.render_images = []\n\n    def close(self):\n        if self.env is not None:\n            self.env.close()\n            self.app.close()\n\n    def load_state(self, state_buffer: bytes):\n        self.env.load_state(state_buffer)\n\n    def get_state(self):\n        return None\n\n    def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):\n        logger.info(f\"IsaacEnv reset_envs_to_state_ids task_ids_list: {task_ids_list}\")\n        assert len(set(task_ids_list)) == 1, \"Isaac env only support single task\"\n\n        self._init_env(task_ids_list[0])\n\n        # In Isaac, reset to random status in groups to have more test coverage\n        # TODO support reset in group with options = {\"group\": len(set(state_ids_list))}\n        raw_obs, infos = self.env.reset()\n        env_idx = np.arange(self.num_envs)\n        self._reset_metrics(env_idx)\n\n        self.elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)\n\n        # stablize the environment\n        for _ in range(10):\n            zero_actions = torch.zeros((self.num_envs, self.action_dim), device=self.device)\n            raw_obs, _, _, _, infos = self.env.step(zero_actions)\n\n        obs = self._wrap_obs(raw_obs)\n        return obs, infos\n"
  },
  {
    "path": "verl/experimental/vla/envs/libero_env/__init__.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/vla/envs/libero_env/libero_env.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Optional\n\nimport gymnasium as gym\nimport numpy as np\nimport torch\nfrom libero.libero import get_libero_path\nfrom libero.libero.benchmark import Benchmark, get_benchmark\nfrom libero.libero.envs import OffScreenRenderEnv\nfrom omegaconf.omegaconf import OmegaConf\n\nfrom verl.experimental.vla.envs.action_utils import (\n    list_of_dict_to_dict_of_list,\n    put_info_on_image,\n    save_rollout_video,\n    tile_images,\n    to_tensor,\n)\nfrom verl.experimental.vla.envs.libero_env.utils import get_libero_image, get_libero_wrist_image, quat2axisangle\nfrom verl.experimental.vla.envs.libero_env.venv import ReconfigureSubprocEnv\n\nlogger = logging.getLogger(__name__)\n\n\ndef patched_get_task_init_states(self, i):\n    init_states_path = os.path.join(\n        get_libero_path(\"init_states\"),\n        self.tasks[i].problem_folder,\n        self.tasks[i].init_states_file,\n    )\n    init_states = torch.load(init_states_path, weights_only=False)\n    return init_states\n\n\nBenchmark.get_task_init_states = patched_get_task_init_states\n\n\nclass LiberoEnv(gym.Env):\n    def __init__(self, cfg, rank, world_size, stage_id: int = 0):\n        self.rank = rank\n        self.stage_id = stage_id\n        self.cfg = cfg\n        self.world_size = world_size\n        self.seed = int(self.cfg.seed)\n        self.rollout_id = 0\n        self.num_envs = self.cfg.num_envs\n\n        self.ignore_terminations = False\n\n        self._generator = np.random.default_rng(seed=self._compose_seed(env_id=0, rollout_id=0, stream_id=0))\n        self._generator_ordered = np.random.default_rng(seed=self._compose_seed(env_id=0, rollout_id=0, stream_id=1))\n        self.start_idx = 0\n\n        self.task_suite: Benchmark = get_benchmark(cfg.task_suite_name)()\n\n        self._compute_total_num_group_envs()\n        self.reset_state_ids_all = self.get_reset_state_ids_all()\n        self.reset_state_ids = self._get_ordered_reset_state_ids(self.num_envs)\n        self._init_task_and_trial_ids()\n        self._init_env()\n\n        self.prev_step_reward = np.zeros(self.num_envs)\n        self.use_rel_reward = False\n\n        self._init_metrics()\n        self._elapsed_steps = np.zeros(self.num_envs, dtype=np.int32)\n\n        self.video_cfg = cfg.video_cfg\n        self.video_cnt = 0\n        self.render_images = []\n\n    def _compose_seed(self, env_id: int, rollout_id: Optional[int] = None, stream_id: int = 0) -> int:\n        if rollout_id is None:\n            rollout_id = self.rollout_id\n        mixed_seed = (\n            self.seed * 1000003\n            + self.rank * 10007\n            + self.stage_id * 1009\n            + int(rollout_id) * 97\n            + int(env_id)\n            + int(stream_id) * 53\n        )\n        return int(mixed_seed % (2**31 - 1))\n\n    @property\n    def elapsed_steps(self):\n        return self._elapsed_steps\n\n    def get_all_state_ids(self):\n        \"\"\"Returns all possible state IDs from the entire benchmark.\"\"\"\n        return np.arange(self.total_num_group_envs)  # (total_num_states,)\n\n    def _init_env(self):\n        env_fns = self.get_env_fns()\n        self.env = ReconfigureSubprocEnv(env_fns)\n\n    def get_env_fns(self):\n        env_fn_params = self.get_env_fn_params()\n        env_fns = []\n        for env_fn_param in env_fn_params:\n\n            def env_fn(param=env_fn_param):\n                seed = param.pop(\"seed\")\n                env = OffScreenRenderEnv(**param)\n                env.seed(seed)\n                return env\n\n            env_fns.append(env_fn)\n        return env_fns\n\n    def get_env_fn_params(self, env_idx=None):\n        env_fn_params = []\n        raw_base_env_args = OmegaConf.to_container(self.cfg.init_params, resolve=True)\n        if raw_base_env_args is None:\n            base_env_args = {}\n        elif isinstance(raw_base_env_args, dict):\n            base_env_args = raw_base_env_args\n        else:\n            raise TypeError(f\"Expected init_params to be a mapping, got {type(raw_base_env_args)}\")\n\n        task_descriptions = []\n        if env_idx is None:\n            env_idx = np.arange(self.cfg.num_envs)\n        for env_id in range(self.cfg.num_envs):\n            if env_id not in env_idx:\n                task_descriptions.append(self.task_descriptions[env_id])\n                continue\n            task = self.task_suite.get_task(self.task_ids[env_id])\n            task_bddl_file = os.path.join(get_libero_path(\"bddl_files\"), task.problem_folder, task.bddl_file)\n            env_fn_params.append(\n                {\n                    **base_env_args,\n                    \"bddl_file_name\": task_bddl_file,\n                    \"seed\": self._compose_seed(env_id=env_id),\n                }\n            )\n            task_descriptions.append(task.language)\n        self.task_descriptions = task_descriptions\n        return env_fn_params\n\n    def _compute_total_num_group_envs(self):\n        self.total_num_group_envs = 0\n        self.trial_id_bins = []\n        for task_id in range(self.task_suite.get_num_tasks()):\n            task_num_trials = len(self.task_suite.get_task_init_states(task_id))\n            self.trial_id_bins.append(task_num_trials)\n\n            self.total_num_group_envs += task_num_trials\n\n        self.cumsum_trial_id_bins = np.cumsum(self.trial_id_bins)\n\n    def _init_task_and_trial_ids(self):\n        self.task_ids, self.trial_ids = self._get_task_and_trial_ids_from_reset_state_ids(self.reset_state_ids)\n\n    def _get_random_reset_state_ids(self, num_reset_states):\n        reset_state_ids = self._generator.integers(low=0, high=self.total_num_group_envs, size=(num_reset_states,))\n        return reset_state_ids\n\n    def get_reset_state_ids_all(self):\n        reset_state_ids = np.arange(self.total_num_group_envs)\n        valid_size = len(reset_state_ids) - (len(reset_state_ids) % self.world_size)\n        if not self.cfg.only_eval:\n            self._generator_ordered.shuffle(reset_state_ids)\n        reset_state_ids = reset_state_ids[:valid_size]\n        reset_state_ids = reset_state_ids.reshape(self.world_size, -1)\n        return reset_state_ids\n\n    def _get_ordered_reset_state_ids(self, num_reset_states):\n        reset_state_ids = self.reset_state_ids_all[self.rank][self.start_idx : self.start_idx + num_reset_states]\n        self.start_idx = self.start_idx + num_reset_states\n        if self.start_idx >= len(self.reset_state_ids_all[0]):\n            self.reset_state_ids_all = self.get_reset_state_ids_all()\n            self.start_idx = 0\n        return reset_state_ids\n\n    def _get_task_and_trial_ids_from_reset_state_ids(self, reset_state_ids):\n        task_ids = []\n        trial_ids = []\n        # get task id and trial id from reset state ids\n        for reset_state_id in reset_state_ids:\n            start_pivot = 0\n            for task_id, end_pivot in enumerate(self.cumsum_trial_id_bins):\n                if reset_state_id < end_pivot and reset_state_id >= start_pivot:\n                    task_ids.append(task_id)\n                    trial_ids.append(reset_state_id - start_pivot)\n                    break\n                start_pivot = end_pivot\n        logger.debug(\n            \"get task and trial id\",\n            self.cumsum_trial_id_bins,\n            reset_state_ids,\n            task_ids,\n            trial_ids,\n        )\n        return np.array(task_ids), np.array(trial_ids)\n\n    def _get_reset_states(self, env_idx):\n        if env_idx is None:\n            env_idx = np.arange(self.num_envs)\n        init_state = [\n            self.task_suite.get_task_init_states(self.task_ids[env_id])[self.trial_ids[env_id]] for env_id in env_idx\n        ]\n        return init_state\n\n    def _init_metrics(self):\n        self.success_once = np.zeros(self.num_envs, dtype=bool)\n        self.fail_once = np.zeros(self.num_envs, dtype=bool)\n        self.returns = np.zeros(self.num_envs)\n\n    def _reset_metrics(self, env_idx=None):\n        if env_idx is not None:\n            mask = np.zeros(self.num_envs, dtype=bool)\n            mask[env_idx] = True\n            self.prev_step_reward[mask] = 0.0\n            self.success_once[mask] = False\n            self.fail_once[mask] = False\n            self.returns[mask] = 0\n            self._elapsed_steps[env_idx] = 0\n        else:\n            self.prev_step_reward[:] = 0\n            self.success_once[:] = False\n            self.fail_once[:] = False\n            self.returns[:] = 0.0\n            self._elapsed_steps[:] = 0\n\n    def _record_metrics(self, step_reward, terminations, infos):\n        episode_info = {}\n        self.returns += step_reward\n        self.success_once = self.success_once | terminations\n        episode_info[\"success_once\"] = self.success_once.copy()\n        episode_info[\"return\"] = self.returns.copy()\n        episode_info[\"episode_len\"] = self.elapsed_steps.copy()\n        episode_info[\"reward\"] = episode_info[\"return\"] / episode_info[\"episode_len\"]\n        infos[\"episode\"] = to_tensor(episode_info)\n        return infos\n\n    def _extract_image_and_state(self, obs):\n        return {\n            \"full_image\": get_libero_image(obs),\n            \"wrist_image\": get_libero_wrist_image(obs),\n            \"state\": np.concatenate(\n                [\n                    obs[\"robot0_eef_pos\"],\n                    quat2axisangle(obs[\"robot0_eef_quat\"]),\n                    obs[\"robot0_gripper_qpos\"],\n                ]\n            ),\n        }\n\n    def _wrap_obs(self, obs_list):\n        images_and_states_list = []\n        for obs in obs_list:\n            images_and_states = self._extract_image_and_state(obs)\n            images_and_states_list.append(images_and_states)\n\n        obs = {\n            \"images_and_states\": to_tensor(list_of_dict_to_dict_of_list(images_and_states_list)),\n            \"task_descriptions\": self.task_descriptions,\n        }\n        return obs\n\n    def _reconfigure(self, reset_state_ids, env_idx):\n        reconfig_env_idx = []\n        task_ids, trial_ids = self._get_task_and_trial_ids_from_reset_state_ids(reset_state_ids)\n        for j, env_id in enumerate(env_idx):\n            if self.task_ids[env_id] != task_ids[j]:\n                reconfig_env_idx.append(env_id)\n            self.task_ids[env_id] = task_ids[j]\n            self.trial_ids[env_id] = trial_ids[j]\n        if reconfig_env_idx:\n            env_fn_params = self.get_env_fn_params(reconfig_env_idx)\n            self.env.reconfigure_env_fns(env_fn_params, reconfig_env_idx)\n\n        seed_list = [self._compose_seed(env_id=int(env_id)) for env_id in env_idx]\n        self.env.seed(seed_list)\n        self.env.reset(id=env_idx)\n        init_state = self._get_reset_states(env_idx=env_idx)\n        self.env.set_init_state(init_state=init_state, id=env_idx)\n\n    def reset(\n        self,\n        env_idx: Optional[int | list[int] | np.ndarray] = None,\n        reset_state_ids=None,\n        options: Optional[dict] = None,\n    ):\n        self.rollout_id += 1\n        if env_idx is None:\n            env_idx = np.arange(self.num_envs)\n\n        if reset_state_ids is None:\n            num_reset_states = len(env_idx)\n            reset_state_ids = self._get_random_reset_state_ids(num_reset_states)\n\n        self._reconfigure(reset_state_ids, env_idx)\n\n        for _ in range(10):\n            zero_actions = np.zeros((self.num_envs, 7))\n            raw_obs, _reward, terminations, info_lists = self.env.step(zero_actions)\n\n        obs = self._wrap_obs(raw_obs)\n        if env_idx is not None:\n            self._reset_metrics(env_idx)\n        else:\n            self._reset_metrics()\n        infos = {}\n        return obs, infos\n\n    def step(self, actions=None, critic_values=None):\n        if actions is None:\n            obs, infos = self.reset(reset_state_ids=self.reset_state_ids)\n            terminations = np.zeros(self.num_envs, dtype=bool)\n            truncations = np.zeros(self.num_envs, dtype=bool)\n\n            return obs, None, to_tensor(terminations), to_tensor(truncations), infos\n\n        if isinstance(actions, torch.Tensor):\n            actions = actions.detach().cpu().numpy()\n\n        self._elapsed_steps += 1\n        raw_obs, _reward, terminations, info_lists = self.env.step(actions)\n        infos = list_of_dict_to_dict_of_list(info_lists)\n        truncations = self.elapsed_steps >= self.cfg.max_episode_steps\n\n        obs = self._wrap_obs(raw_obs)\n        step_reward = self._calc_step_reward(terminations)\n\n        if self.video_cfg.save_video:\n            plot_infos = {\n                \"rewards\": step_reward,\n                \"terminations\": terminations,\n                \"critic_value\": np.asarray(critic_values, dtype=np.float32),\n            }\n            plot_infos[\"task\"] = self.task_descriptions\n            self.add_new_frames(raw_obs, plot_infos)\n\n        infos = self._record_metrics(step_reward, terminations, infos)\n\n        return (\n            obs,\n            to_tensor(step_reward),\n            to_tensor(terminations),\n            to_tensor(truncations),\n            infos,\n        )\n\n    def chunk_step(self, chunk_actions, chunk_values=None):\n        # chunk_actions: [num_envs, chunk_step, action_dim]\n        chunk_size = chunk_actions.shape[1]\n\n        chunk_rewards = []\n\n        raw_chunk_terminations = []\n        raw_chunk_truncations = []\n        for i in range(chunk_size):\n            actions = chunk_actions[:, i]\n            if len(chunk_values.shape) == 1:\n                step_values = chunk_values\n            elif len(chunk_values.shape) == 2:\n                step_values = chunk_values[:, i]\n\n            extracted_obs, step_reward, terminations, truncations, infos = self.step(actions, critic_values=step_values)\n\n            chunk_rewards.append(step_reward)\n            raw_chunk_terminations.append(terminations)\n            raw_chunk_truncations.append(truncations)\n\n        chunk_rewards = torch.stack(chunk_rewards, dim=1)  # [num_envs, chunk_steps]\n        raw_chunk_terminations = torch.stack(raw_chunk_terminations, dim=1)  # [num_envs, chunk_steps]\n        raw_chunk_truncations = torch.stack(raw_chunk_truncations, dim=1)  # [num_envs, chunk_steps]\n\n        chunk_terminations = raw_chunk_terminations.clone()\n        chunk_truncations = raw_chunk_truncations.clone()\n        return (\n            extracted_obs,\n            chunk_rewards,\n            chunk_terminations,\n            chunk_truncations,\n            infos,\n        )\n\n    def _calc_step_reward(self, terminations):\n        reward = self.cfg.reward_coef * terminations\n        reward_diff = reward - self.prev_step_reward\n        self.prev_step_reward = reward\n\n        if self.use_rel_reward:\n            return reward_diff\n        else:\n            return reward\n\n    def add_new_frames(self, raw_obs, plot_infos):\n        images = []\n        for env_id, raw_single_obs in enumerate(raw_obs):\n            info_item = {k: v if np.size(v) == 1 else v[env_id] for k, v in plot_infos.items()}\n            img = raw_single_obs[\"agentview_image\"][::-1, ::-1]\n            img = put_info_on_image(img, info_item)\n            images.append(img)\n        full_image = tile_images(images, nrows=int(np.sqrt(self.num_envs)))\n        self.render_images.append(full_image)\n\n    def flush_video(self, video_sub_dir: Optional[str] = None):\n        output_dir = os.path.join(self.video_cfg.video_base_dir, f\"rank_{self.rank}\")\n        if video_sub_dir is not None:\n            output_dir = os.path.join(output_dir, f\"{video_sub_dir}\")\n        save_rollout_video(\n            self.render_images,\n            output_dir=output_dir,\n            video_name=f\"{self.video_cnt}\",\n        )\n        self.video_cnt += 1\n        self.render_images = []\n\n    def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):\n        \"\"\"Reset environments to specified state IDs.\n\n        Args:\n            state_ids_list: List of state IDs to reset environments to\n        \"\"\"\n        env_idx = np.arange(len(state_ids_list))\n        obs, infos = self.reset(env_idx=env_idx, reset_state_ids=state_ids_list)\n        return obs, infos\n\n    def load_state(self, state_buffer: bytes):\n        self.env.load_state(state_buffer)\n"
  },
  {
    "path": "verl/experimental/vla/envs/libero_env/utils.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utils for evaluating policies in LIBERO simulation environments.\"\"\"\n\nimport math\n\nimport numpy as np\n\n\ndef get_libero_image(obs: dict[str, np.ndarray]) -> np.ndarray:\n    \"\"\"\n    Extracts image from observations and preprocesses it.\n\n    Args:\n        obs: Observation dictionary from LIBERO environment\n\n    Returns:\n        Preprocessed image as numpy array\n    \"\"\"\n    img = obs[\"agentview_image\"]\n    img = img[::-1, ::-1]  # IMPORTANT: rotate 180 degrees to match train preprocessing\n    return img\n\n\ndef get_libero_wrist_image(obs: dict[str, np.ndarray]) -> np.ndarray:\n    \"\"\"\n    Extracts wrist camera image from observations and preprocesses it.\n\n    Args:\n        obs: Observation dictionary from LIBERO environment\n\n    Returns:\n        Preprocessed wrist camera image as numpy array\n    \"\"\"\n    img = obs[\"robot0_eye_in_hand_image\"]\n    img = img[::-1, ::-1]  # IMPORTANT: rotate 180 degrees to match train preprocessing\n    return img\n\n\ndef quat2axisangle(quat: np.ndarray) -> np.ndarray:\n    \"\"\"\n    Copied from robosuite: https://github.com/ARISE-Initiative/robosuite/blob/eafb81f54ffc104f905ee48a16bb15f059176ad3/robosuite/utils/transform_utils.py#L490C1-L512C55\n\n    Converts quaternion to axis-angle format.\n    Returns a unit vector direction scaled by its angle in radians.\n\n    Args:\n        quat (np.array): (x,y,z,w) vec4 float angles\n\n    Returns:\n        np.array: (ax,ay,az) axis-angle exponential coordinates\n    \"\"\"\n    # clip quaternion\n    if quat[3] > 1.0:\n        quat[3] = 1.0\n    elif quat[3] < -1.0:\n        quat[3] = -1.0\n\n    den = np.sqrt(1.0 - quat[3] * quat[3])\n    if math.isclose(den, 0.0):\n        # This is (close to) a zero degree rotation, immediately return\n        return np.zeros(3)\n\n    return (quat[:3] * 2.0 * math.acos(quat[3])) / den\n\n\ndef normalize_gripper_action(action: np.ndarray, binarize: bool = True) -> np.ndarray:\n    \"\"\"\n    Normalize gripper action from [0,1] to [-1,+1] range.\n\n    This is necessary for some environments because the dataset wrapper\n    standardizes gripper actions to [0,1]. Note that unlike the other action\n    dimensions, the gripper action is not normalized to [-1,+1] by default.\n\n    Normalization formula: y = 2 * (x - orig_low) / (orig_high - orig_low) - 1\n\n    Args:\n        action: Action array with gripper action in the last dimension\n        binarize: Whether to binarize gripper action to -1 or +1\n\n    Returns:\n        np.ndarray: Action array with normalized gripper action\n    \"\"\"\n    # Create a copy to avoid modifying the original\n    normalized_action = action.copy()\n\n    # Normalize the last action dimension to [-1,+1]\n    orig_low, orig_high = 0.0, 1.0\n    normalized_action[..., -1] = 2 * (normalized_action[..., -1] - orig_low) / (orig_high - orig_low) - 1\n\n    if binarize:\n        # Binarize to -1 or +1\n        normalized_action[..., -1] = np.sign(normalized_action[..., -1])\n\n    return normalized_action\n\n\ndef invert_gripper_action(action: np.ndarray) -> np.ndarray:\n    \"\"\"\n    Flip the sign of the gripper action (last dimension of action vector).\n\n    This is necessary for environments where -1 = open, +1 = close, since\n    the RLDS dataloader aligns gripper actions such that 0 = close, 1 = open.\n\n    Args:\n        action: Action array with gripper action in the last dimension\n\n    Returns:\n        np.ndarray: Action array with inverted gripper action\n    \"\"\"\n    # Create a copy to avoid modifying the original\n    inverted_action = action.copy()\n\n    # Invert the gripper action\n    inverted_action[..., -1] = inverted_action[..., -1] * -1.0\n\n    return inverted_action\n"
  },
  {
    "path": "verl/experimental/vla/envs/libero_env/venv.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom multiprocessing import Pipe, connection\nfrom multiprocessing.context import Process\nfrom typing import Any, Callable, Optional\n\nimport gymnasium as gym\nimport numpy as np\nfrom libero.libero.envs import OffScreenRenderEnv\nfrom libero.libero.envs.venv import (\n    BaseVectorEnv,\n    CloudpickleWrapper,\n    EnvWorker,\n    ShArray,\n    SubprocEnvWorker,\n    SubprocVectorEnv,\n    _setup_buf,\n)\n\n\ndef _worker(\n    parent: connection.Connection,\n    p: connection.Connection,\n    env_fn_wrapper: CloudpickleWrapper,\n    obs_bufs: Optional[dict | tuple | ShArray] = None,\n) -> None:\n    def _encode_obs(obs: dict | tuple | np.ndarray, buffer: dict | tuple | ShArray) -> None:\n        if isinstance(obs, np.ndarray) and isinstance(buffer, ShArray):\n            buffer.save(obs)\n        elif isinstance(obs, tuple) and isinstance(buffer, tuple):\n            for o, b in zip(obs, buffer, strict=False):\n                _encode_obs(o, b)\n        elif isinstance(obs, dict) and isinstance(buffer, dict):\n            for k in obs.keys():\n                _encode_obs(obs[k], buffer[k])\n        return None\n\n    parent.close()\n    env = env_fn_wrapper.data()\n    try:\n        while True:\n            try:\n                cmd, data = p.recv()\n            except EOFError:  # the pipe has been closed\n                p.close()\n                break\n            if cmd == \"step\":\n                env_return = env.step(data)\n                if obs_bufs is not None:\n                    _encode_obs(env_return[0], obs_bufs)\n                    env_return = (None, *env_return[1:])\n                p.send(env_return)\n            elif cmd == \"reset\":\n                retval = env.reset(**data)\n                reset_returns_info = (\n                    isinstance(retval, (tuple | list)) and len(retval) == 2 and isinstance(retval[1], dict)\n                )\n                if reset_returns_info:\n                    obs, info = retval\n                else:\n                    obs = retval\n                if obs_bufs is not None:\n                    _encode_obs(obs, obs_bufs)\n                    obs = None\n                if reset_returns_info:\n                    p.send((obs, info))\n                else:\n                    p.send(obs)\n            elif cmd == \"close\":\n                p.send(env.close())\n                p.close()\n                break\n            elif cmd == \"render\":\n                p.send(env.render(**data) if hasattr(env, \"render\") else None)\n            elif cmd == \"seed\":\n                if hasattr(env, \"seed\"):\n                    p.send(env.seed(data))\n                else:\n                    env.reset(seed=data)\n                    p.send(None)\n            elif cmd == \"getattr\":\n                p.send(getattr(env, data) if hasattr(env, data) else None)\n            elif cmd == \"setattr\":\n                setattr(env.unwrapped, data[\"key\"], data[\"value\"])\n            elif cmd == \"check_success\":\n                p.send(env.check_success())\n            elif cmd == \"get_segmentation_of_interest\":\n                p.send(env.get_segmentation_of_interest(data))\n            elif cmd == \"get_sim_state\":\n                p.send(env.get_sim_state())\n            elif cmd == \"set_init_state\":\n                obs = env.set_init_state(data)\n                p.send(obs)\n            elif cmd == \"reconfigure\":\n                env.close()\n                seed = data.pop(\"seed\")\n                env = OffScreenRenderEnv(**data)\n                env.seed(seed)\n                p.send(None)\n            else:\n                p.close()\n                raise NotImplementedError\n    except KeyboardInterrupt:\n        p.close()\n\n\nclass ReconfigureSubprocEnvWorker(SubprocEnvWorker):\n    def __init__(self, env_fn: Callable[[], gym.Env], share_memory: bool = False):\n        self.parent_remote, self.child_remote = Pipe()\n        self.share_memory = share_memory\n        self.buffer: Optional[dict | tuple | ShArray] = None\n        if self.share_memory:\n            dummy = env_fn()\n            obs_space = dummy.observation_space\n            dummy.close()\n            del dummy\n            self.buffer = _setup_buf(obs_space)\n        args = (\n            self.parent_remote,\n            self.child_remote,\n            CloudpickleWrapper(env_fn),\n            self.buffer,\n        )\n        self.process = Process(target=_worker, args=args, daemon=True)\n        self.process.start()\n        self.child_remote.close()\n        EnvWorker.__init__(self, env_fn)\n\n    def reconfigure_env_fn(self, env_fn_param):\n        self.parent_remote.send([\"reconfigure\", env_fn_param])\n        return self.parent_remote.recv()\n\n\nclass ReconfigureSubprocEnv(SubprocVectorEnv):\n    def __init__(self, env_fns: list[Callable[[], gym.Env]], **kwargs: Any) -> None:\n        def worker_fn(fn: Callable[[], gym.Env]) -> ReconfigureSubprocEnvWorker:\n            return ReconfigureSubprocEnvWorker(fn, share_memory=False)\n\n        BaseVectorEnv.__init__(self, env_fns, worker_fn, **kwargs)\n\n    def reconfigure_env_fns(self, env_fns, id=None):\n        self._assert_is_not_closed()\n        id = self._wrap_id(id)\n        if self.is_async:\n            self._assert_id(id)\n\n        for j, i in enumerate(id):\n            self.workers[i].reconfigure_env_fn(env_fns[j])\n"
  },
  {
    "path": "verl/experimental/vla/fsdp_workers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe main entry point to run the PPO algorithm\n\"\"\"\n\nimport asyncio\nimport contextlib\nimport logging\nimport os\n\nimport torch\nimport torch.distributed\nfrom packaging import version\nfrom torch.distributed.device_mesh import init_device_mesh\nfrom torch.distributed.fsdp import FSDPModule\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp._unshard_param_utils import _get_module_fsdp_state, _unshard_params_for_summon\nfrom torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType\n\nfrom verl import DataProto\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import get_device_id, get_device_name, get_torch_device, set_expandable_segments\nfrom verl.utils.flops_counter import FlopsCounter\nfrom verl.utils.fsdp_utils import fsdp_version, set_reshard_after_forward\nfrom verl.utils.import_utils import import_external_libs\nfrom verl.utils.memory_utils import aggressive_empty_cache\nfrom verl.utils.profiler import DistProfiler, log_gpu_memory_usage, simple_timer\nfrom verl.utils.profiler.performance import reduce_timing, topk_reduce_ratio_min_max\nfrom verl.workers.config import HFModelConfig\nfrom verl.workers.fsdp_workers import ActorRolloutRefWorker\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\ndevice_name = get_device_name()\n\n\nclass RobActorRolloutRefWorker(ActorRolloutRefWorker):\n    \"\"\"\n    This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy\n    or a hybrid engine based on the config.rollout\n    \"\"\"\n\n    fsdp_unshard_exit_stack = contextlib.ExitStack()\n\n    def _build_rollout(self, trust_remote_code=False):\n        self.base_sync_done = False\n        world_size = torch.distributed.get_world_size()\n        dp = world_size\n        infer_tp = self.config.rollout.tensor_model_parallel_size\n        rollout_device_mesh = init_device_mesh(\n            device_name, mesh_shape=(dp, infer_tp), mesh_dim_names=[\"dp\", \"infer_tp\"]\n        )\n        # 3. init trainer and rollout random states\n        self.torch_random_states = get_torch_device().get_rng_state()\n        gen_dp_rank = rollout_device_mesh[\"dp\"].get_local_rank()\n        get_torch_device().manual_seed(gen_dp_rank + 1000)  # make sure all tp ranks have the same random states\n        self.gen_random_states = get_torch_device().get_rng_state()\n        get_torch_device().set_rng_state(self.torch_random_states)\n\n        fsdp_ver = fsdp_version(self.actor_module_fsdp)\n        if torch.distributed.get_world_size() == 1 and fsdp_ver == 1:\n            FSDP.set_state_dict_type(\n                self.actor_module_fsdp,\n                state_dict_type=StateDictType.FULL_STATE_DICT,\n                state_dict_config=FullStateDictConfig(),\n            )\n        elif fsdp_ver == 1:\n            FSDP.set_state_dict_type(\n                self.actor_module_fsdp,\n                state_dict_type=StateDictType.SHARDED_STATE_DICT,\n                state_dict_config=ShardedStateDictConfig(),\n            )\n        elif fsdp_ver == 2:\n            # FSDP2 already handles state dict logic via torch.distributed.checkpoint APIs.\n            pass\n        else:\n            raise NotImplementedError(f\"Unsupported fsdp version {fsdp_ver}\")\n\n        self._register_dispatch_collect_info(\"rollout\", dp_rank=self.rank, is_collect=True)\n\n        if self.config.get(\"algorithm\", \"grpo\") == \"sac\":\n            from verl.experimental.vla.sac.naive_rollout_pi05 import PI0RolloutRob\n\n            self.rollout = PI0RolloutRob(\n                module=self.actor_module_fsdp, model_config=self.config.model, tokenizer=self.tokenizer\n            )\n        else:\n            from verl.experimental.vla.naive_rollout_rob import NaiveRolloutRob\n\n            self.rollout = NaiveRolloutRob(module=self.actor_module_fsdp, model_config=self.config.model)\n\n        model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model, dataclass_type=HFModelConfig)\n        self.model_config = model_config\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def switch_to_rollout(self):\n        loop = asyncio.get_event_loop()\n        loop.run_until_complete(self.rollout_mode())\n        log_gpu_memory_usage(\"After switch to rollout mode\", logger=logger)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def switch_to_train(self):\n        loop = asyncio.get_event_loop()\n        loop.run_until_complete(self.trainer_mode())\n        log_gpu_memory_usage(\"After switch to trainer mode\", logger=logger)\n\n    async def rollout_mode(self):\n        \"\"\"Context switch hybridengine to rollout mode.\"\"\"\n\n        self.actor_module_fsdp.eval()\n\n        aggressive_empty_cache(force_sync=True)\n        self.base_sync_done = True\n\n        # important: need to manually set the random states of each tp to be identical.\n        self.torch_random_states = get_torch_device().get_rng_state()\n        get_torch_device().set_rng_state(self.gen_random_states)\n\n        if fsdp_version(self.actor_module_fsdp) == 1:\n            fsdp_unshard_exit_stack = contextlib.ExitStack()\n            optional_state = _get_module_fsdp_state(self.actor_module_fsdp)\n            if optional_state is None:\n                self.fsdp_unshard_exit_stack = fsdp_unshard_exit_stack\n            states_and_modules = ([optional_state], [self.actor_module_fsdp])\n\n            for state, fsdp_module in zip(*states_and_modules, strict=False):\n                fsdp_unshard_exit_stack.enter_context(\n                    _unshard_params_for_summon(\n                        module=fsdp_module,\n                        state=state,\n                        writeback=False,\n                        rank0_only=False,\n                        offload_to_cpu=False,\n                        with_grads=False,\n                    )\n                )\n\n            self.fsdp_unshard_exit_stack = fsdp_unshard_exit_stack\n        elif fsdp_version(self.actor_module_fsdp) == 2:\n            self.actor_module_fsdp.unshard()\n            for m in self.actor_module_fsdp.modules():\n                if isinstance(m, FSDPModule) or hasattr(m, \"unshard\"):\n                    m.unshard()\n            if version.parse(torch.__version__) < version.parse(\"2.8\"):\n                set_reshard_after_forward(self.actor_module_fsdp, False)\n            else:\n                self.actor_module_fsdp.set_reshard_after_forward(False)\n        else:\n            raise NotImplementedError(f\"Unsupported fsdp version {fsdp_version(self.actor_module_fsdp)}\")\n\n        logger.info(\"rollout mode\")\n\n    async def trainer_mode(self):\n        \"\"\"Context switch hybridengine to trainer mode.\"\"\"\n\n        self.actor_module_fsdp.train()\n\n        # add empty cache after each compute\n        aggressive_empty_cache(force_sync=True)\n        set_expandable_segments(True)\n\n        # restore random states\n        self.gen_random_states = get_torch_device().get_rng_state()\n        get_torch_device().set_rng_state(self.torch_random_states)\n\n        if fsdp_version(self.actor_module_fsdp) == 1:\n            if self.fsdp_unshard_exit_stack is not None:\n                self.fsdp_unshard_exit_stack.close()\n                self.fsdp_unshard_exit_stack = None\n        elif fsdp_version(self.actor_module_fsdp) == 2:\n            self.actor_module_fsdp.reshard()\n            for m in self.actor_module_fsdp.modules():\n                if isinstance(m, FSDPModule) or hasattr(m, \"reshard\"):\n                    m.reshard()\n            if version.parse(torch.__version__) < version.parse(\"2.8\"):\n                set_reshard_after_forward(self.actor_module_fsdp, True)\n            else:\n                self.actor_module_fsdp.set_reshard_after_forward(True)\n        else:\n            raise NotImplementedError(f\"Unsupported fsdp version {fsdp_version(self.actor_module_fsdp)}\")\n\n        logger.info(\"trainer mode\")\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"rollout\"), blocking=False)\n    @DistProfiler.annotate(color=\"red\", role=\"rollout_generate\")\n    def generate_sequences(self, prompts: DataProto):\n        # Support all hardwares\n        assert self._is_rollout\n        prompts = prompts.to(get_device_id())\n\n        meta_info = {\n            \"eos_token_id\": self.model_config.generation_config.eos_token_id\n            if self.model_config.generation_config is not None\n            else self.model_config.tokenizer.eos_token_id,\n            \"pad_token_id\": self.model_config.generation_config.pad_token_id\n            if self.model_config.generation_config is not None\n            else self.model_config.tokenizer.pad_token_id,\n        }\n        prompts.meta_info.update(meta_info)\n\n        timing_generate = {}\n\n        with simple_timer(\"generate_sequences\", timing_generate):\n            output = self.rollout.generate_sequences(prompts=prompts)\n\n        timing_generate_topk_ratio, timing_generate_min, timing_generate_max = topk_reduce_ratio_min_max(\n            timing_generate[\"generate_sequences\"]\n        )\n        timing_generate = reduce_timing(timing_generate)\n        timing_generate.update(\n            {\n                \"generation_timing/max\": timing_generate_max,\n                \"generation_timing/min\": timing_generate_min,\n                \"generation_timing/topk_ratio\": timing_generate_topk_ratio,\n            }\n        )\n        output.meta_info[\"metrics\"] = timing_generate\n        output = output.to(\"cpu\")\n\n        # clear kv cache\n        get_torch_device().empty_cache()\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        # This is used to import external_lib into the huggingface systems\n        import_external_libs(self.config.model.get(\"external_lib\", None))\n\n        # Initialize QAT config before _build_model_optimizer\n        self._init_qat_config()\n\n        from omegaconf import OmegaConf\n\n        override_model_config = OmegaConf.to_container(self.config.model.get(\"override_config\", OmegaConf.create()))\n\n        from verl.experimental.vla.models import register_vla_models\n\n        register_vla_models()\n\n        from transformers import AutoProcessor\n\n        self.processor = AutoProcessor.from_pretrained(self.config.model.path, trust_remote_code=True)\n\n        if self._is_actor or self._is_rollout:\n            # we need the model for actor and rollout\n            if self._is_actor:\n                optim_config = self.config.actor.optim\n                fsdp_config = self.config.actor.fsdp_config\n            else:\n                optim_config = None\n                fsdp_config = OmegaConf.create()\n            self.actor_module_fsdp, self.actor_optimizer, self.actor_lr_scheduler, self.actor_model_config = (\n                self._build_model_optimizer(\n                    model_path=self.config.model.path,\n                    fsdp_config=fsdp_config,\n                    optim_config=optim_config,\n                    override_model_config=override_model_config,\n                    enable_gradient_checkpointing=self.config.model.get(\"enable_gradient_checkpointing\", False),\n                    trust_remote_code=self.config.model.get(\"trust_remote_code\", False),\n                )\n            )\n\n            if fsdp_version(self.actor_module_fsdp) == 1:\n                # get the original unwrapped module\n                self.actor_module = self.actor_module_fsdp._fsdp_wrapped_module\n\n        if self._is_actor:\n            OmegaConf.set_struct(self.config.actor, True)\n            if self.config.get(\"algorithm\") == \"sac\":\n                from verl.experimental.vla.sac.sac_actor import RobDataParallelSACActor\n\n                self.actor = RobDataParallelSACActor(\n                    config=self.config.actor,\n                    actor_module=self.actor_module_fsdp,\n                    actor_optimizer=self.actor_optimizer,\n                    tokenizer=self.tokenizer,\n                )\n            else:\n                from verl.experimental.vla.dp_rob import RobDataParallelPPOActor\n\n                self.actor = RobDataParallelPPOActor(\n                    config=self.config.actor, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer\n                )\n\n        if self._is_rollout:\n            self._build_rollout(trust_remote_code=self.config.model.get(\"trust_remote_code\", False))\n\n        if self._is_actor:\n            self.flops_counter = FlopsCounter(self.actor_model_config)\n            self.checkpoint_manager = FSDPCheckpointManager(\n                model=self.actor_module_fsdp,\n                optimizer=self.actor.actor_optimizer,\n                lr_scheduler=self.actor_lr_scheduler,\n                processing_class=self.processor if self.processor is not None else self.tokenizer,\n                checkpoint_config=self.config.actor.checkpoint,\n                trust_remote_code=self.config.model.trust_remote_code,\n            )\n\n        torch.distributed.barrier()\n"
  },
  {
    "path": "verl/experimental/vla/main_ppo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\n\nimport datasets\nimport hydra\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom verl import DataProto\nfrom verl.trainer.constants_ppo import get_ppo_ray_runtime_env\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.trainer.ppo.utils import Role\nfrom verl.utils.device import is_cuda_available\n\nfrom .rob_ray_trainer import RobRayPPOTrainer\n\nlogger = logging.getLogger(__name__)\n\n\ndef calculate_reward(data: DataProto, return_dict: bool = False) -> torch.Tensor:\n    complete_tensor = data.batch[\"complete\"]\n    batch_size, num_steps = complete_tensor.shape[:2]\n    traj_has_complete = torch.any(complete_tensor, dim=(1, 2))  # shape: [batch_size]\n    reward_per_traj = traj_has_complete.float()\n    reward_per_step = reward_per_traj.unsqueeze(1).expand(batch_size, num_steps)\n    if return_dict:\n        return {\"reward_tensor\": reward_per_step}\n    else:\n        return reward_per_step\n\n\n@hydra.main(config_path=\"config\", config_name=\"rob_ppo_trainer\", version_base=None)\ndef main(config):\n    if not ray.is_initialized():\n        default_runtime_env = get_ppo_ray_runtime_env()\n        ray_init_kwargs = config.ray_kwargs.get(\"ray_init\", {})\n        runtime_env_kwargs = ray_init_kwargs.get(\"runtime_env\", {})\n        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)\n        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, \"runtime_env\": runtime_env})\n        logger.info(f\"ray init kwargs: {ray_init_kwargs}\")\n        ray.init(**OmegaConf.to_container(ray_init_kwargs))\n\n    # Apply controller nsight profiling if configured\n    if (\n        is_cuda_available\n        and config.global_profiler.tool == \"nsys\"\n        and config.global_profiler.get(\"steps\") is not None\n        and len(config.global_profiler.get(\"steps\", [])) > 0\n    ):\n        from verl.utils.import_utils import is_nvtx_available\n\n        assert is_nvtx_available(), \"nvtx is not available in CUDA platform. Please 'pip3 install nvtx'\"\n        nsight_options = OmegaConf.to_container(\n            config.global_profiler.global_tool_config.nsys.controller_nsight_options\n        )\n        main_task_with_options = main_task.options(runtime_env={\"nsight\": nsight_options})\n        ray.get(main_task_with_options.remote(config))\n    else:\n        ray.get(main_task.remote(config))\n\n    # [Optional] get the path of the timeline trace file from the configuration, default to None\n    # This file is used for performance analysis\n    timeline_json_file = config.ray_kwargs.get(\"timeline_json_file\", None)\n    if timeline_json_file:\n        ray.timeline(filename=timeline_json_file)\n\n\n@ray.remote\ndef main_task(config):\n    # print initial config\n    from pprint import pprint\n\n    from omegaconf import OmegaConf\n\n    from verl.utils.fs import copy_local_path_from_hdfs\n\n    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values\n    OmegaConf.resolve(config)\n\n    # download the checkpoint from hdfs\n    local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path)\n\n    # instantiate tokenizer\n    from verl.utils import hf_tokenizer\n\n    tokenizer = hf_tokenizer(local_path)\n\n    # define worker classes\n    if config.actor_rollout_ref.actor.strategy in [\"fsdp\", \"fsdp2\"]:\n        assert config.actor_rollout_ref.actor.strategy == config.critic.strategy\n        from verl.experimental.vla.workers.env.env_worker import EnvWorker\n        from verl.single_controller.ray import RayWorkerGroup\n\n        from .fsdp_workers import RobActorRolloutRefWorker\n\n        ray_worker_group_cls = RayWorkerGroup\n\n    else:\n        raise NotImplementedError\n\n    role_worker_mapping = {\n        # Role.Critic: ray.remote(RobActorRolloutRefWorker),\n        Role.ActorRollout: ray.remote(RobActorRolloutRefWorker),\n        # Role.RefPolicy: ray.remote(RobActorRolloutRefWorker),\n        Role.Env: ray.remote(EnvWorker),\n    }\n\n    train_rollout_pool_id = \"train_rollout_pool\"\n\n    num_nodes_actor_rollout = config.trainer.nnodes\n    train_rollout_gpu_num = config.trainer.n_rollout_gpus_per_node\n    env_gpu_num = config.trainer.n_env_gpus_per_node\n    if config.env.disagg_sim.enable:\n        # disaggregated sim and actor rollout\n        num_nodes_sim = config.env.disagg_sim.nnodes\n    else:\n        # colocated sim and actor rollout\n        num_nodes_sim = config.trainer.nnodes\n\n    resource_pool_spec = {\n        train_rollout_pool_id: [train_rollout_gpu_num] * num_nodes_actor_rollout,\n        \"env_gpu_pool\": [env_gpu_num] * num_nodes_sim,\n    }\n    mapping = {\n        Role.ActorRollout: train_rollout_pool_id,\n        # Role.Critic: global_pool_id,\n        # Role.RefPolicy: global_pool_id,\n        Role.Env: \"env_gpu_pool\",\n    }\n\n    reward_fn = calculate_reward\n    val_reward_fn = calculate_reward\n\n    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n\n    # Create training and validation datasets.\n    train_dataset = datasets.load_dataset(\"parquet\", data_files=config.data.train_files)[\"train\"]\n    val_dataset = datasets.load_dataset(\"parquet\", data_files=config.data.val_files)[\"train\"]\n\n    trainer = RobRayPPOTrainer(\n        config=config,\n        tokenizer=tokenizer,\n        role_worker_mapping=role_worker_mapping,\n        resource_pool_manager=resource_pool_manager,\n        ray_worker_group_cls=ray_worker_group_cls,\n        reward_fn=reward_fn,\n        val_reward_fn=val_reward_fn,\n        train_dataset=train_dataset,\n        val_dataset=val_dataset,\n    )\n    trainer.init_workers()\n    trainer.fit()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/experimental/vla/main_sac.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\nfrom pprint import pprint\n\nimport datasets\nimport hydra\nimport ray\nimport torch\nfrom omegaconf import OmegaConf\n\nfrom verl import DataProto\nfrom verl.experimental.vla.sac.sac_ray_trainer import RobRaySACTrainer\nfrom verl.trainer.constants_ppo import get_ppo_ray_runtime_env\nfrom verl.trainer.ppo.ray_trainer import ResourcePoolManager\nfrom verl.trainer.ppo.utils import Role\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.fs import copy_local_path_from_hdfs\n\nlogger = logging.getLogger(__name__)\n\n\ndef calculate_reward(data: DataProto, return_dict: bool = False) -> torch.Tensor:\n    complete_tensor = data.batch[\"complete\"]\n    reward_per_step = complete_tensor.float()\n    if return_dict:\n        return {\"reward_tensor\": reward_per_step}\n    else:\n        return reward_per_step\n\n\n@hydra.main(config_path=\"config\", config_name=\"rob_sac_trainer\", version_base=None)\ndef main(config):\n    if not ray.is_initialized():\n        default_runtime_env = get_ppo_ray_runtime_env()\n        ray_init_kwargs = config.ray_kwargs.get(\"ray_init\", {})\n        runtime_env_kwargs = ray_init_kwargs.get(\"runtime_env\", {})\n        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)\n        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, \"runtime_env\": runtime_env})\n        logger.info(f\"ray init kwargs: {ray_init_kwargs}\")\n        ray.init(**OmegaConf.to_container(ray_init_kwargs))\n    ray.get(main_task.remote(config))\n\n\n@ray.remote\ndef main_task(config):\n    # print initial config\n    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values\n    OmegaConf.resolve(config)\n\n    # download the checkpoint from hdfs\n    local_path = copy_local_path_from_hdfs(config.actor_rollout_ref.model.path)\n\n    # instantiate tokenizer\n    tokenizer = hf_tokenizer(local_path)\n\n    # define worker classes\n    if config.actor_rollout_ref.actor.strategy in [\"fsdp\", \"fsdp2\"]:\n        assert config.actor_rollout_ref.actor.strategy == config.critic.strategy\n        from verl.experimental.vla.workers.env.env_worker import EnvWorker\n        from verl.single_controller.ray import RayWorkerGroup\n\n        from .fsdp_workers import RobActorRolloutRefWorker\n\n        ray_worker_group_cls = RayWorkerGroup\n    else:\n        raise NotImplementedError\n\n    role_worker_mapping = {\n        Role.ActorRollout: ray.remote(RobActorRolloutRefWorker),\n        Role.Env: ray.remote(EnvWorker),\n    }\n\n    # setup resource pool manager\n    train_rollout_gpu_num = config.trainer.n_rollout_gpus_per_node\n    train_rollout_nodes_num = config.trainer.nnodes\n    env_gpu_num = config.trainer.n_env_gpus_per_node\n    env_nodes_num = config.env.disagg_sim.nnodes if config.env.disagg_sim.enable else config.trainer.nnodes\n\n    resource_pool_spec = {\n        \"train_rollout_pool\": [train_rollout_gpu_num] * train_rollout_nodes_num,\n        \"env_gpu_pool\": [env_gpu_num] * env_nodes_num,\n    }\n    mapping = {\n        Role.ActorRollout: \"train_rollout_pool\",\n        Role.Env: \"env_gpu_pool\",\n    }\n    resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=mapping)\n\n    # create datasets\n    train_dataset = datasets.load_dataset(\"parquet\", data_files=config.data.train_files)[\"train\"]\n    val_dataset = datasets.load_dataset(\"parquet\", data_files=config.data.val_files)[\"train\"]\n\n    # instantiate trainer and start training\n    trainer = RobRaySACTrainer(\n        config=config,\n        tokenizer=tokenizer,\n        role_worker_mapping=role_worker_mapping,\n        resource_pool_manager=resource_pool_manager,\n        ray_worker_group_cls=ray_worker_group_cls,\n        train_dataset=train_dataset,\n        val_dataset=val_dataset,\n    )\n\n    trainer.init_workers()\n    trainer.fit()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/experimental/vla/models/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .register_vla_models import register_vla_models\n\n__all__ = [\n    \"register_vla_models\",\n]\n"
  },
  {
    "path": "verl/experimental/vla/models/modules/mlp.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch.nn as nn\nimport torch.nn.init as init\n\n\nclass MLP(nn.Module):\n    \"\"\"\n    A configurable Multi-Layer Perceptron (MLP) module.\n    It supports dynamic layer construction, multiple activation functions,\n    and various weight initialization strategies.\n\n    Attributes:\n        input_dim (int): The number of input features.\n        hidden_dims (list of int): List containing the number of units in each hidden layer.\n        output_dim (int): The number of output units.\n        activation (str): The non-linear activation function to use.\n            Options: 'relu', 'tanh', 'sigmoid', 'leaky_relu', 'elu', 'selu', 'none'.\n        init_method (str): The weight initialization strategy.\n            Options: 'kaiming', 'xavier', 'normal', 'orthogonal'.\n        output_init_scale (float): Scale for uniform initialization of output layer weights.\n    \"\"\"\n\n    def __init__(\n        self,\n        input_dim: int,\n        hidden_dims: list[int],\n        output_dim: int,\n        activation: str = \"relu\",\n        init_method: str = \"kaiming\",\n        output_init_scale: float = 3e-3,\n    ):\n        super().__init__()\n\n        self.input_dim = input_dim\n        self.hidden_dims = hidden_dims\n        self.output_dim = output_dim\n        self.activation_name = activation.lower()\n        self.init_method = init_method.lower()\n        self.output_init_scale = float(output_init_scale)\n\n        layers = []\n        current_dim = input_dim\n\n        for h_dim in hidden_dims:\n            layers.append(nn.Linear(current_dim, h_dim))\n            act_layer = self._get_activation(self.activation_name)\n            if act_layer is not None:\n                layers.append(act_layer)\n            current_dim = h_dim\n\n        layers.append(nn.Linear(current_dim, output_dim))\n        self.network = nn.Sequential(*layers)\n        self.apply(self.init_weights)\n\n    def _get_activation(self, name: str):\n        \"\"\"\n        Factory method to return a *fresh* activation layer instance based on string name.\n        Available options: 'relu', 'tanh', 'sigmoid', 'leaky_relu', 'elu', 'selu', 'none'.\n        \"\"\"\n        name = name.lower()\n        if name == \"relu\":\n            return nn.ReLU()\n        if name == \"tanh\":\n            return nn.Tanh()\n        if name == \"sigmoid\":\n            return nn.Sigmoid()\n        if name == \"leaky_relu\":\n            return nn.LeakyReLU(0.2)\n        if name == \"elu\":\n            return nn.ELU()\n        if name == \"selu\":\n            return nn.SELU()\n        if name == \"none\":\n            return None\n        return nn.ReLU()\n\n    def init_weights(self, m: nn.Module):\n        \"\"\"\n        Initialize weights for Linear layers.\n\n        Hidden layers follow init_method.\n        Output layer uses small uniform init (±output_init_scale) to keep initial outputs near 0.\n        \"\"\"\n        if not isinstance(m, nn.Linear):\n            return\n\n        # Identify the output layer by matching out_features to the requested output_dim\n        # (works because only the last Linear has out_features == self.output_dim in this MLP)\n        is_output_layer = m.out_features == self.output_dim\n\n        if is_output_layer:\n            init.uniform_(m.weight, -self.output_init_scale, self.output_init_scale)\n        else:\n            if self.init_method == \"kaiming\":\n                if self.activation_name == \"leaky_relu\":\n                    init.kaiming_normal_(m.weight, a=0.2, nonlinearity=\"leaky_relu\")\n                else:\n                    init.kaiming_normal_(m.weight, nonlinearity=\"relu\")\n            elif self.init_method == \"xavier\":\n                init.xavier_normal_(m.weight)\n            elif self.init_method == \"normal\":\n                init.normal_(m.weight, mean=0.0, std=0.02)\n            elif self.init_method == \"orthogonal\":\n                init.orthogonal_(m.weight)\n\n        if m.bias is not None:\n            init.constant_(m.bias, 0.0)\n\n    def forward(self, x):\n        return self.network(x)\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/configuration_prismatic.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# from https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/\n# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/\n\"\"\"\nconfiguration_prismatic.py\n\nHuggingFace-style configuration definition for Prismatic VLMs, inheriting from `transformers.PretrainedConfig`.\nDefault configuration specifies `siglip-224px+7b`.\n\"\"\"\n\nfrom typing import Any, Optional\n\nfrom transformers import PretrainedConfig\nfrom transformers.models.auto import CONFIG_MAPPING\n\n# === Utilities for Mapping Prismatic names to HF names ===\n# fmt: off\nVISION_BACKBONE_TO_RESOLUTION: dict[str, list[int]] = {\n    \"clip-vit-l\": [224], \"siglip-vit-so400m\": [224], \"dinov2-vit-l\": [224], \"in1k-vit-l\": [224],\n\n    \"clip-vit-l-336px\": [336],\n    \"siglip-vit-so400m-384px\": [384],\n\n    \"dinoclip-vit-l-336px\": [336, 336],\n    \"dinosiglip-vit-so-224px\": [224, 224],\n    \"dinosiglip-vit-so-384px\": [384, 384],\n}\nVISION_BACKBONE_TO_TIMM_ID: dict[str, list[str]] = {\n    \"clip-vit-l\": [\"vit_large_patch14_clip_224.openai\"],\n    \"clip-vit-l-336px\": [\"vit_large_patch14_clip_336.openai\"],\n\n    \"dinov2-vit-l\": [\"vit_large_patch14_reg4_dinov2.lvd142m\"],\n    \"in1k-vit-l\": [\"vit_large_patch16_224.augreg_in21k_ft_in1k\"],\n\n    \"siglip-vit-so400m\": [\"vit_so400m_patch14_siglip_224\"],\n    \"siglip-vit-so400m-384px\": [\"vit_so400m_patch14_siglip_384\"],\n\n    \"dinoclip-vit-l-336px\": [\"vit_large_patch14_reg4_dinov2.lvd142m\", \"vit_large_patch14_clip_336.openai\"],\n    \"dinosiglip-vit-so-224px\": [\"vit_large_patch14_reg4_dinov2.lvd142m\", \"vit_so400m_patch14_siglip_224\"],\n    \"dinosiglip-vit-so-384px\": [\"vit_large_patch14_reg4_dinov2.lvd142m\", \"vit_so400m_patch14_siglip_384\"],\n}\nTIMM_OVERRIDE_ACT_LAYER: dict[str, list[Optional[str]]] = {\n    \"clip-vit-l\": [\"quick_gelu\"], \"clip-vit-l-336px\": [\"quick_gelu\"],\n    \"dinov2-vit-l\": [None], \"in1k-vit-l\": [None],\n    \"siglip-vit-so400m\": [None], \"siglip-vit-so400m-384px\": [None],\n    \"dinoclip-vit-l-336px\": [None, \"quick_gelu\"],\n    \"dinosiglip-vit-so-224px\": [None, None], \"dinosiglip-vit-so-384px\": [None, None]\n}\n\nLLM_BACKBONE_TO_HF_PATH = {\n    \"llama2-7b-pure\": \"meta-llama/Llama-2-7b-hf\", \"llama2-13b-pure\": \"meta-llama/Llama-2-13b-hf\",\n    \"llama2-7b-chat\": \"meta-llama/Llama-2-7b-chat-hf\", \"llama2-13b-chat\": \"meta-llama/Llama-2-13b-chat-hf\",\n\n    \"vicuna-v15-7b\": \"lmsys/vicuna-7b-v1.5\", \"vicuna-v15-13b\": \"lmsys/vicuna-13b-v1.5\",\n\n    \"mistral-v0.1-7b-pure\": \"mistralai/Mistral-7B-v0.1\",\n    \"mistral-v0.1-7b-instruct\": \"mistralai/Mistral-7B-Instruct-v0.1\",\n\n    \"phi-2-3b\": \"microsoft/phi-2\",\n}\nLLM_BACKBONE_TO_HF_METACLASS = {\n    \"llama2-7b-pure\": \"llama\", \"llama2-13b-pure\": \"llama\", \"llama2-7b-chat\": \"llama\", \"llama2-13b-chat\": \"llama\",\n    \"vicuna-v15-7b\": \"llama\", \"vicuna-v15-13b\": \"llama\",\n\n    \"mistral-v0.1-7b-pure\": \"mistral\", \"mistral-v0.1-7b-instruct\": \"mistral\",\n\n    \"phi-2-3b\": \"phi\",\n}\n\nVALID_VISION_BACKBONES = set(VISION_BACKBONE_TO_RESOLUTION.keys())\nVALID_LLM_BACKBONES = set(LLM_BACKBONE_TO_HF_PATH)\n# fmt: on\n\n\nclass PrismaticConfig(PretrainedConfig):\n    model_type: str = \"prismatic\"\n    is_composition: bool = False\n\n    def __init__(\n        self,\n        vision_backbone_id: str = \"siglip-vit-so400m\",\n        llm_backbone_id: str = \"vicuna-v15-7b\",\n        arch_specifier: str = \"no-align+gelu-mlp\",\n        use_fused_vision_backbone: Optional[bool] = None,\n        image_resize_strategy: str = \"letterbox\",\n        text_config: Optional[dict[str, Any]] = None,\n        llm_max_length: int = 2048,\n        pad_token_id: int = 32000,\n        pad_to_multiple_of: int = 64,\n        output_projector_states: bool = False,\n        **kwargs: str,\n    ) -> None:\n        if vision_backbone_id not in VALID_VISION_BACKBONES:\n            raise ValueError(f\"Vision backbone `{vision_backbone_id}` not in {VALID_VISION_BACKBONES = }\")\n\n        if llm_backbone_id not in VALID_LLM_BACKBONES:\n            raise ValueError(f\"LLM backbone `{llm_backbone_id}` not in {VALID_LLM_BACKBONES = }\")\n\n        # Set Prismatic Configuration Fields\n        self.vision_backbone_id = vision_backbone_id\n        self.llm_backbone_id = llm_backbone_id\n        self.arch_specifier = arch_specifier\n        self.output_projector_states = output_projector_states\n\n        # [Contract] All vision backbone parameters are lists =>> supports fused backbones with different preprocessing\n        self.use_fused_vision_backbone = (\n            use_fused_vision_backbone\n            if use_fused_vision_backbone is not None\n            else any(self.vision_backbone_id.startswith(v) for v in [\"dinoclip\", \"dinosiglip\"])\n        )\n\n        self.timm_model_ids = VISION_BACKBONE_TO_TIMM_ID[self.vision_backbone_id]\n        self.timm_override_act_layers = TIMM_OVERRIDE_ACT_LAYER[self.vision_backbone_id]\n        self.image_sizes = VISION_BACKBONE_TO_RESOLUTION[self.vision_backbone_id]\n        self.image_resize_strategy = image_resize_strategy\n\n        self.hf_llm_id = LLM_BACKBONE_TO_HF_PATH[self.llm_backbone_id]\n        self.llm_max_length = llm_max_length\n        self.pad_token_id, self.pad_to_multiple_of = pad_token_id, pad_to_multiple_of\n\n        # [IMPORTANT] HF Utilities actually look for a `text_config` field... we need to use that specific naming!\n        self.text_config = (\n            CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]](**text_config)\n            if text_config is not None\n            else CONFIG_MAPPING[LLM_BACKBONE_TO_HF_METACLASS[self.llm_backbone_id]]()\n        )\n\n        # Dispatch **kwargs to super() =>> note that `pad_token_id` collides, so we pass it in here as well...\n        super().__init__(pad_token_id=pad_token_id, **kwargs)\n\n\nclass OpenVLAConfig(PrismaticConfig):\n    model_type: str = \"openvla\"\n\n    def __init__(\n        self,\n        norm_stats: Optional[dict[str, dict[str, dict[str, dict[str, list[float]]]]]] = None,\n        n_action_bins: int = 256,\n        **kwargs: str,\n    ) -> None:\n        self.norm_stats, self.n_action_bins = norm_stats, n_action_bins\n\n        super().__init__(**kwargs)\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/constants.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# from https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/\n\n\n\"\"\"\nImportant constants for VLA training and evaluation.\n\nAttempts to automatically identify the correct constants to set based on the Python command used to launch\ntraining or evaluation. If it is unclear, defaults to using the LIBERO simulation benchmark constants.\n\"\"\"\n\nimport sys\nfrom enum import Enum\n\n# Llama 2 token constants\nIGNORE_INDEX = -100\nACTION_TOKEN_BEGIN_IDX = 31743\nSTOP_INDEX = 2  # '</s>'\n\n\n# Defines supported normalization schemes for action and proprioceptive state.\nclass NormalizationType(str, Enum):\n    # fmt: off\n    NORMAL = \"normal\"               # Normalize to Mean = 0, Stdev = 1\n    BOUNDS = \"bounds\"               # Normalize to Interval = [-1, 1]\n    BOUNDS_Q99 = \"bounds_q99\"       # Normalize [quantile_01, ..., quantile_99] --> [-1, ..., 1]\n    # fmt: on\n\n\n# Define constants for each robot platform\nLIBERO_CONSTANTS = {\n    \"NUM_ACTIONS_CHUNK\": 8,\n    \"ACTION_DIM\": 7,\n    \"PROPRIO_DIM\": 8,\n    \"ACTION_PROPRIO_NORMALIZATION_TYPE\": NormalizationType.BOUNDS_Q99,\n}\n\nALOHA_CONSTANTS = {\n    \"NUM_ACTIONS_CHUNK\": 25,\n    \"ACTION_DIM\": 14,\n    \"PROPRIO_DIM\": 14,\n    \"ACTION_PROPRIO_NORMALIZATION_TYPE\": NormalizationType.BOUNDS,\n}\n\nBRIDGE_CONSTANTS = {\n    \"NUM_ACTIONS_CHUNK\": 5,\n    \"ACTION_DIM\": 7,\n    \"PROPRIO_DIM\": 7,\n    \"ACTION_PROPRIO_NORMALIZATION_TYPE\": NormalizationType.BOUNDS_Q99,\n}\n\n\n# Function to detect robot platform from command line arguments\ndef detect_robot_platform():\n    cmd_args = \" \".join(sys.argv).lower()\n\n    if \"libero\" in cmd_args:\n        return \"LIBERO\"\n    elif \"aloha\" in cmd_args:\n        return \"ALOHA\"\n    elif \"bridge\" in cmd_args:\n        return \"BRIDGE\"\n    else:\n        # Default to LIBERO if unclear\n        return \"LIBERO\"\n\n\n# Determine which robot platform to use\nROBOT_PLATFORM = detect_robot_platform()\n\n# Set the appropriate constants based on the detected platform\nif ROBOT_PLATFORM == \"LIBERO\":\n    constants = LIBERO_CONSTANTS\nelif ROBOT_PLATFORM == \"ALOHA\":\n    constants = ALOHA_CONSTANTS\nelif ROBOT_PLATFORM == \"BRIDGE\":\n    constants = BRIDGE_CONSTANTS\n\n# Assign constants to global variables\nNUM_ACTIONS_CHUNK = constants[\"NUM_ACTIONS_CHUNK\"]\nACTION_DIM = constants[\"ACTION_DIM\"]\nPROPRIO_DIM = constants[\"PROPRIO_DIM\"]\nACTION_PROPRIO_NORMALIZATION_TYPE = constants[\"ACTION_PROPRIO_NORMALIZATION_TYPE\"]\n\n# Print which robot platform constants are being used (for debugging)\nprint(f\"Using {ROBOT_PLATFORM} constants:\")\nprint(f\"  NUM_ACTIONS_CHUNK = {NUM_ACTIONS_CHUNK}\")\nprint(f\"  ACTION_DIM = {ACTION_DIM}\")\nprint(f\"  PROPRIO_DIM = {PROPRIO_DIM}\")\nprint(f\"  ACTION_PROPRIO_NORMALIZATION_TYPE = {ACTION_PROPRIO_NORMALIZATION_TYPE}\")\nprint(\"If needed, manually set the correct constants in `prismatic/vla/constants.py`!\")\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/modeling_prismatic.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# from https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/\n# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/\n\n\"\"\"\nmodeling_prismatic.py\n\nCore HuggingFace-style PrismaticPreTrainedModel and PrismaticForConditionalGeneration class definitions.\nInherits from the default `transformers.PretrainedModel`. Meant to be standalone and self-contained,\nbut exactly replicate the logic in `prismatic.models.vlms.prismatic.py`.\n\"\"\"\n\nimport logging\nfrom dataclasses import dataclass\nfrom functools import partial\nfrom typing import Any, Callable, ClassVar, Optional\n\nimport numpy as np\nimport timm\nimport tokenizers\nimport torch\nimport torch.nn as nn\nimport transformers\nfrom timm.models.vision_transformer import LayerScale\nfrom transformers import AutoModelForCausalLM, PretrainedConfig, PreTrainedModel\nfrom transformers.modeling_outputs import ModelOutput\n\nfrom .configuration_prismatic import OpenVLAConfig, PrismaticConfig\nfrom .constants import (\n    ACTION_DIM,\n    ACTION_PROPRIO_NORMALIZATION_TYPE,\n    ACTION_TOKEN_BEGIN_IDX,\n    IGNORE_INDEX,\n    NUM_ACTIONS_CHUNK,\n    STOP_INDEX,\n    NormalizationType,\n)\nfrom .train_utils import (\n    get_current_action_mask,\n    get_next_actions_mask,\n)\n\n# Set up logger\nlogger = logging.getLogger(__name__)\n\n\n# === Utility Functions for Monkey-Patching ===\ndef unpack_tuple(fn: Callable[[Any], tuple[Any]]) -> Callable[[Any], Any]:\n    def wrapper(*args: Any, **kwargs: Any) -> Any:\n        result = fn(*args, **kwargs)\n        return result[0] if isinstance(result, tuple) else result\n\n    return wrapper\n\n\n# HF Transformers overwrites parameters with names containing `gamma`; we're going to patch VisionBackbone.LayerScale.\n#   =>> TIMM :: https://github.com/huggingface/pytorch-image-models/blob/main/timm/models/vision_transformer.py#L109\n#   =>> Transformers :: https://github.com/huggingface/transformers/blob/main/src/transformers/modeling_utils.py#L3960\ndef _ls_new_forward(self, x: torch.Tensor) -> torch.Tensor:\n    return x.mul_(self.scale_factor) if self.inplace else x * self.scale_factor\n\n\ndef ls_apply_patch(ls_module: LayerScale):\n    ls_module.scale_factor = nn.Parameter(ls_module.gamma.clone())\n    ls_module.forward = _ls_new_forward.__get__(ls_module, LayerScale)\n    del ls_module.gamma\n\n\n# === Prismatic Vision Backbone (nn.Module) Definitions (w/ Fused Backbone Support) ===\nclass PrismaticVisionBackbone(nn.Module):\n    \"\"\"\n    Vision backbone for Prismatic models that handles image feature extraction.\n\n    Supports both single backbone (e.g., SigLIP) and fused backbone (e.g., SigLIP + DINOv2) configurations.\n    For fused backbones, features from both models are concatenated along the feature dimension.\n    \"\"\"\n\n    def __init__(\n        self,\n        use_fused_vision_backbone: bool,\n        image_sizes: list[int],\n        timm_model_ids: list[str],\n        timm_override_act_layers: list[Optional[str]],\n    ) -> None:\n        \"\"\"\n        Initialize the vision backbone.\n\n        Args:\n            use_fused_vision_backbone: Whether to use two backbones and fuse their features\n            image_sizes: List of image sizes for each backbone\n            timm_model_ids: List of TIMM model IDs to use for each backbone\n            timm_override_act_layers: List of activation layer overrides for each backbone\n        \"\"\"\n        super().__init__()\n        self.use_fused_vision_backbone = use_fused_vision_backbone\n        self.num_images_in_input = 1  # Default value, can be overridden later\n\n        # Validate number of (fused) vision backbones\n        if len(timm_model_ids) > 2:\n            raise ValueError(\"Prismatic models only support up to 2 (fused) vision backbones!\")\n\n        # Create primary featurizer\n        self.featurizer = self._create_featurizer(\n            model_id=timm_model_ids[0], img_size=image_sizes[0], act_layer=timm_override_act_layers[0]\n        )\n        self.embed_dim = self.featurizer.embed_dim\n\n        # Create secondary featurizer if using fused backbone\n        if self.use_fused_vision_backbone:\n            self.fused_featurizer = self._create_featurizer(\n                model_id=timm_model_ids[1], img_size=image_sizes[1], act_layer=timm_override_act_layers[1]\n            )\n            self.embed_dim += self.fused_featurizer.embed_dim\n\n        # Patch LayerScale modules for HF compatibility\n        self._patch_layer_scales()\n\n    def _create_featurizer(self, model_id: str, img_size: int, act_layer: Optional[str]) -> nn.Module:\n        \"\"\"\n        Create a TIMM-based featurizer model with appropriate configurations.\n\n        Args:\n            model_id: The TIMM model ID to load\n            img_size: Input image size for the model\n            act_layer: Override for the activation layer type\n\n        Returns:\n            A configured featurizer model\n        \"\"\"\n        featurizer = timm.create_model(\n            model_id,\n            pretrained=False,\n            num_classes=0,\n            img_size=img_size,\n            act_layer=act_layer,\n        )\n\n        # Monkey-patch the forward function to extract the second-to-last layer features\n        num_blocks = len(featurizer.blocks)\n        featurizer.forward = unpack_tuple(partial(featurizer.get_intermediate_layers, n={num_blocks - 2}))\n\n        return featurizer\n\n    def _patch_layer_scales(self) -> None:\n        \"\"\"\n        Patch all LayerScale modules to be compatible with HF's parameter naming.\n\n        HF Transformers overwrites parameters with names containing 'gamma',\n        so we need to rename and modify the forward method.\n        \"\"\"\n        # Patch primary featurizer\n        for module in self.featurizer.modules():\n            if isinstance(module, LayerScale):\n                ls_apply_patch(module)\n\n        # Patch secondary featurizer if it exists\n        if self.use_fused_vision_backbone:\n            for module in self.fused_featurizer.modules():\n                if isinstance(module, LayerScale):\n                    ls_apply_patch(module)\n\n    def get_num_patches(self) -> int:\n        \"\"\"\n        Returns the number of vision patches output by the vision backbone.\n\n        Returns:\n            Number of patches per image\n        \"\"\"\n        return self.featurizer.patch_embed.num_patches\n\n    def get_num_images_in_input(self) -> int:\n        \"\"\"\n        Returns the number of input images for the vision backbone.\n\n        Returns:\n            Number of images expected in the input\n        \"\"\"\n        return self.num_images_in_input\n\n    def set_num_images_in_input(self, num_images_in_input: int) -> None:\n        \"\"\"\n        Sets the number of input images for the vision backbone.\n\n        Args:\n            num_images_in_input: Number of images to expect in the input\n        \"\"\"\n        self.num_images_in_input = num_images_in_input\n\n    def forward(self, pixel_values: torch.Tensor) -> torch.Tensor:\n        \"\"\"\n        Implements the forward pass for the vision backbone.\n\n        If `self.use_fused_vision_backbone == True`, uses both SigLIP and DINOv2 transformers to extract visual features\n        (otherwise uses SigLIP only). Allows multi-image inputs (but only for fused vision backbone).\n\n        Args:\n            pixel_values (torch.Tensor): Pixels for input image(s), (B, C, H, W).\n        \"\"\"\n        if self.num_images_in_input == 1:\n            if not self.use_fused_vision_backbone:\n                return self.featurizer(pixel_values)\n\n            # Split `pixel_values :: [bsz, 2 * 3, resolution, resolution]` =>> featurize =>> channel stack\n            img, img_fused = torch.split(pixel_values, [3, 3], dim=1)\n            patches, patches_fused = self.featurizer(img), self.fused_featurizer(img_fused)\n\n            return torch.cat([patches, patches_fused], dim=2)\n\n        else:\n            assert self.use_fused_vision_backbone, \"Multi-image inputs require using fused backbone!\"\n\n            # Split `pixel_values` into individual images (each with 6 channels: 3 for SigLIP + 3 for DINOv2)\n            images = torch.split(pixel_values, [6] * self.num_images_in_input, dim=1)\n\n            # Process each image and collect patches\n            all_patches = []\n            for img in images:\n                # Split each image further into two stacks of channels (each with 3 channels)\n                img_regular, img_fused = torch.split(img, [3, 3], dim=1)\n\n                # Get patches from both SigLIP and DINOv2 vision transformers\n                patches = self.featurizer(img_regular)\n                patches_fused = self.fused_featurizer(img_fused)\n\n                # Concatenate SigLIP and DINOv2 patches along the hidden dimension\n                combined_patches = torch.cat([patches, patches_fused], dim=2)\n                all_patches.append(combined_patches)\n\n            # Concatenate all patches along the patch dimension\n            return torch.cat(all_patches, dim=1)\n\n\n# === Prismatic Projector (nn.Module) Definitions ===\nclass PrismaticProjector(nn.Module):\n    def __init__(self, use_fused_vision_backbone: bool, vision_dim: int, llm_dim: int) -> None:\n        super().__init__()\n        self.use_fused_vision_backbone = use_fused_vision_backbone\n        self.vision_dim, self.llm_dim = vision_dim, llm_dim\n\n        # Switch on `use_fused_vision_backbone` =>> use slightly different MLPs and projection factors!\n        if not self.use_fused_vision_backbone:\n            self.fc1 = nn.Linear(self.vision_dim, self.llm_dim, bias=True)\n            self.fc2 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)\n            self.act_fn1 = nn.GELU()\n        else:\n            initial_projection_dim = 4 * vision_dim\n            self.fc1 = nn.Linear(self.vision_dim, initial_projection_dim, bias=True)\n            self.fc2 = nn.Linear(initial_projection_dim, self.llm_dim, bias=True)\n            self.fc3 = nn.Linear(self.llm_dim, self.llm_dim, bias=True)\n            self.act_fn1 = nn.GELU()\n            self.act_fn2 = nn.GELU()\n\n    def forward(self, img_patches: torch.Tensor) -> torch.Tensor:\n        if not self.use_fused_vision_backbone:\n            projected_features = self.fc1(img_patches)\n            projected_features = self.act_fn1(projected_features)\n            projected_features = self.fc2(projected_features)\n        else:\n            projected_features = self.fc1(img_patches)\n            projected_features = self.act_fn1(projected_features)\n            projected_features = self.fc2(projected_features)\n            projected_features = self.act_fn2(projected_features)\n            projected_features = self.fc3(projected_features)\n\n        return projected_features\n\n\n# === Main HF Class Definitions ===\n@dataclass\nclass PrismaticCausalLMOutputWithPast(ModelOutput):\n    \"\"\"Base class for Prismatic casual (visually-conditioned) language model outputs; also exposes visual features.\"\"\"\n\n    loss: Optional[torch.FloatTensor] = None\n    logits: torch.FloatTensor = None\n    past_key_values: Optional[tuple[tuple[torch.FloatTensor]]] = None\n    hidden_states: Optional[tuple[torch.FloatTensor, ...]] = None\n    attentions: Optional[tuple[torch.FloatTensor]] = None\n\n    # Additions for VLMs\n    projector_features: Optional[torch.FloatTensor] = None\n\n\nclass PrismaticPreTrainedModel(PreTrainedModel):\n    config_class: PretrainedConfig = PrismaticConfig\n    base_model_prefix: str = \"model\"\n    supports_gradient_checkpointing: bool = True\n\n    _no_split_modules: ClassVar[list[str]] = [\"PrismaticProjector\"]\n    _skip_keys_device_placement: str = \"past_key_values\"\n    _supports_flash_attn_2: bool = True\n\n    def _init_weights(self, module: nn.Module) -> None:\n        # Important :: this HF ported version is *not* meant for training from scratch; only inference and fine-tuning!\n        #   => As such, this init_weights code is not correct; if training VLMs from scratch, use the main codebase at\n        #      https://github.com/TRI-ML/prismatic-vlms\n        std = (\n            self.config.initializer_range\n            if hasattr(self.config, \"initializer_range\")\n            else self.config.text_config.initializer_range\n        )\n\n        if hasattr(module, \"class_embedding\"):\n            module.class_embedding.data.normal_(mean=0.0, std=std)\n\n        if isinstance(module, (nn.Linear | nn.Conv2d)):\n            module.weight.data.normal_(mean=0.0, std=std)\n            if module.bias is not None:\n                module.bias.data.zero_()\n        elif isinstance(module, nn.Embedding):\n            module.weight.data.normal_(mean=0.0, std=std)\n            if module.padding_idx is not None:\n                module.weight.data[module.padding_idx].zero_()\n\n    @property\n    def _supports_sdpa(self) -> bool:\n        \"\"\"Check LLM supports SDPA Attention\"\"\"\n        return self.language_model._supports_sdpa\n\n\nclass PrismaticForConditionalGeneration(PrismaticPreTrainedModel):\n    def __init__(self, config: PrismaticConfig) -> None:\n        super().__init__(config)\n\n        # [Validation] Lightweight Validate on `config` Fields + Dependency Versions\n        if config.use_fused_vision_backbone is None:\n            raise ValueError(\"Missing config field `use_fused_vision_backbone`\")\n\n        if timm.__version__ not in {\"0.9.10\", \"0.9.11\", \"0.9.12\", \"0.9.16\"}:\n            raise NotImplementedError(\n                \"TIMM Version must be >= 0.9.10 and < 1.0.0 (breaking); please raise a GitHub Issue \"\n                \"if you urgently need support for latest TIMM versions.\"\n            )\n\n        if (transformers.__version__ != \"4.40.1\") or (tokenizers.__version__ != \"0.19.1\"):\n            logger.warning(\n                f\"Expected `transformers==4.40.1` and `tokenizers==0.19.1` but got \"\n                f\"`transformers=={transformers.__version__}` and `tokenizers=={tokenizers.__version__}`; \"\n                f\"there might be inference-time regressions due to dependency changes. If in doubt, please\"\n                f\"use the above versions.\"\n            )\n\n        # Instantiate PrismaticVisionBackbone (w/ Potential Fused Backbone)\n        self.vision_backbone = PrismaticVisionBackbone(\n            config.use_fused_vision_backbone, config.image_sizes, config.timm_model_ids, config.timm_override_act_layers\n        )\n\n        # Create Multimodal Projector\n        self.projector = PrismaticProjector(\n            config.use_fused_vision_backbone,\n            vision_dim=self.vision_backbone.embed_dim,\n            llm_dim=config.text_config.hidden_size,\n        )\n\n        # Instantiate LLM Backbone\n        self.language_model = AutoModelForCausalLM.from_config(\n            config.text_config, attn_implementation=config._attn_implementation\n        )\n        self.vocab_size = config.text_config.vocab_size\n        self.pad_token_id = config.pad_token_id\n        self.llm_dim = config.text_config.hidden_size\n\n        # HF Boilerplate =>> initializes weights via `_init_weights()` and sets gradient checkpointing\n        self.post_init()\n\n    # === `PreTrainedModel` Boilerplate ===\n    def get_input_embeddings(self) -> nn.Module:\n        return self.language_model.get_input_embeddings()\n\n    def set_input_embeddings(self, value: nn.Module) -> None:\n        self.language_model.set_input_embeddings(value)\n\n    def get_output_embeddings(self) -> nn.Module:\n        return self.language_model.get_output_embeddings()\n\n    def set_output_embeddings(self, new_embeddings: nn.Module) -> None:\n        self.language_model.set_output_embeddings(new_embeddings)\n\n    def get_decoder(self) -> nn.Module:\n        return self.language_model.get_decoder()\n\n    def set_decoder(self, decoder: nn.Module) -> None:\n        self.language_model.set_decoder(decoder)\n\n    def tie_weights(self) -> None:\n        self.language_model.tie_weights()  # Note: `Llama-2` and `Mistral` don't tie weights (no-op)\n\n    def resize_token_embeddings(\n        self, new_num_tokens: Optional[int] = None, pad_to_multiple_of: Optional[int] = None\n    ) -> nn.Embedding:\n        updated_embeddings = self.language_model.resize_token_embeddings(new_num_tokens, pad_to_multiple_of)\n\n        # Update config/instance variables\n        self.config.text_config.vocab_size = updated_embeddings.num_embeddings\n        self.vocab_size = updated_embeddings.num_embeddings\n\n        return updated_embeddings\n\n    def _replace_input_embeddings(self, input_embeddings, all_actions_mask, noisy_action_features):\n        \"\"\"\n        Replace embeddings in input_embeddings at positions where all_actions_mask is True\n        with embeddings from noisy_action_features, using vectorized operations.\n\n        Args:\n            input_embeddings: Tensor of shape (B, S, D)\n            all_actions_mask: Boolean tensor of shape (B, S)\n            noisy_action_features: Tensor of shape (B, K, D) where K is the number of True values in mask per sample\n\n        Returns:\n            Modified input_embeddings tensor\n        \"\"\"\n        # Clone input to avoid modifying the original tensor\n        new_input_embeddings = input_embeddings.clone()\n\n        # Create a tensor with the same shape of input_embeddings to hold the noisy action features\n        repositioned_noisy_action_features = torch.zeros_like(input_embeddings)\n\n        # Create batch indices for splicing\n        batch_indices = torch.arange(input_embeddings.shape[0], device=input_embeddings.device)\n        batch_indices = batch_indices.unsqueeze(1).expand(-1, noisy_action_features.shape[1])\n\n        # Get indices where mask is True for each sample\n        masked_indices = torch.stack([torch.where(mask)[0] for mask in all_actions_mask])\n\n        # Move the noisy action features into their correct positions\n        repositioned_noisy_action_features[batch_indices, masked_indices] = noisy_action_features\n\n        # Combine original input embeddings and noisy action embeddings using the mask\n        new_input_embeddings = torch.where(\n            all_actions_mask.unsqueeze(-1), repositioned_noisy_action_features, new_input_embeddings\n        )\n\n        return new_input_embeddings\n\n    def _process_action_masks(self, labels):\n        \"\"\"Helper to get action masks from labels\"\"\"\n        current_action_mask = get_current_action_mask(labels)\n        next_actions_mask = get_next_actions_mask(labels)\n        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)\n        return all_actions_mask\n\n    def _process_vision_features(self, pixel_values, language_embeddings=None, use_film=False):\n        \"\"\"Process vision features with optional FiLM conditioning\"\"\"\n        if use_film:\n            # FiLM: Infuse language inputs into visual features\n            patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)\n        else:\n            patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)\n\n        # Project patch embeddings into language embedding space\n        return self.projector(patch_features)\n\n    def _process_proprio_features(self, projected_patch_embeddings, proprio, proprio_projector):\n        \"\"\"Process proprioceptive features and append to vision features\"\"\"\n        if proprio_projector is not None and proprio is not None:\n            # projected_patch_embeddings: (bsz, num_patches * num_images, llm_dim)\n            # proprio: (bsz, proprio_dim) or (propro_dim,)\n            proprio = proprio.reshape(projected_patch_embeddings.shape[0], -1)  # (bsz, proprio_dim)\n            proprio_features = proprio_projector(proprio)  # (bsz, llm_dim)\n            proprio_features = proprio_features.unsqueeze(dim=1)  # (bsz, 1, llm_dim)\n            # For simplicity, just append proprio token to the end of projected vision patch tokens\n            return torch.cat((projected_patch_embeddings, proprio_features), dim=1)\n        return projected_patch_embeddings\n\n    def _build_multimodal_attention(self, input_embeddings, projected_patch_embeddings, attention_mask):\n        \"\"\"Build multimodal embeddings and attention mask\"\"\"\n        # Update attention mask\n        projected_patch_attention_mask = None\n        if attention_mask is not None:\n            projected_patch_attention_mask = torch.full(\n                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),\n                fill_value=True,\n                dtype=attention_mask.dtype,\n                device=attention_mask.device,\n            )\n\n        # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)\n        multimodal_embeddings = torch.cat(\n            [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1\n        )\n\n        multimodal_attention_mask = None\n        if attention_mask is not None:\n            multimodal_attention_mask = torch.cat(\n                [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1\n            )\n\n        return multimodal_embeddings, multimodal_attention_mask\n\n    def _build_multimodal_labels(self, labels, projected_patch_embeddings):\n        \"\"\"Build multimodal labels with IGNORE_INDEX for patch embeddings\"\"\"\n        if labels is not None:\n            projected_patch_labels = torch.full(\n                (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),\n                fill_value=IGNORE_INDEX,\n                dtype=labels.dtype,\n                device=labels.device,\n            )\n            return torch.cat([labels[:, :1], projected_patch_labels, labels[:, 1:]], dim=1)\n        return None\n\n    # === Core Prismatic VLM `forward()` Logic ===\n    # def forward(\n    #     self,\n    #     input_ids: Optional[torch.LongTensor] = None,\n    #     attention_mask: Optional[torch.Tensor] = None,\n    #     pixel_values: Optional[torch.FloatTensor] = None,\n    #     labels: Optional[torch.LongTensor] = None,\n    #     inputs_embeds: Optional[torch.FloatTensor] = None,\n    #     past_key_values: Optional[List[torch.FloatTensor]] = None,\n    #     use_cache: Optional[bool] = None,\n    #     output_attentions: Optional[bool] = None,\n    #     output_hidden_states: Optional[bool] = None,\n    #     output_projector_features: Optional[bool] = None,\n    #     return_dict: Optional[bool] = None,\n    #     proprio=None,\n    #     proprio_projector=None,\n    #     noisy_actions=None,\n    #     noisy_action_projector=None,\n    #     diffusion_timestep_embeddings=None,\n    #     use_film: bool = False,\n    # ) -> Union[Tuple, PrismaticCausalLMOutputWithPast]:\n    #     \"\"\"Run a forward pass through the VLM, returning a PrismaticCausalLMOutputWithPast instance.\"\"\"\n    #     output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n    #     output_hidden_states = (\n    #         output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n    #     )\n    #     output_projector_features = output_projector_features if output_projector_features is not None else False\n    #     return_dict = return_dict if return_dict is not None else self.config.use_return_dict\n\n    #     # Respect `use_cache` only if not training (even if `gradient_checkpointing` is off)\n    #     use_cache = use_cache and not self.training\n\n    #     # Instantiate Placeholder for Projector Features\n    #     projected_patch_embeddings = None\n\n    #     # === Handle Generation with Cache (`input_ids.shape[1] == 1`) =>> requires `past_keys_values` ===\n    #     if input_ids.shape[1] == 1:\n    #         assert input_ids.shape[0] == 1, \"Generation is only currently supported for batch size of 1!\"\n    #         assert past_key_values is not None, \"You must provide `past_key_values` during cached generation!\"\n    #         assert labels is None, \"Unexpected key `labels` provided during cached generation!\"\n\n    #         language_model_output = self.language_model(\n    #             input_ids=input_ids,\n    #             attention_mask=None,\n    #             position_ids=None,\n    #             past_key_values=past_key_values,\n    #             inputs_embeds=None,\n    #             labels=None,\n    #             use_cache=use_cache,\n    #             output_attentions=output_attentions,\n    #             output_hidden_states=output_hidden_states,\n    #             return_dict=return_dict,\n    #         )\n\n    #     # === Handle Unimodal Forward ===\n    #     elif pixel_values is None:\n    #         assert (input_ids is not None) and (inputs_embeds is None), \\\n    #             \"Missing `input_ids` in language-only forward!\"\n    #         assert past_key_values is None, \\\n    #             \"Unexpected key `past_key_values` provided during language-only forward!\"\n\n    #         language_model_output = self.language_model(\n    #             input_ids=input_ids,\n    #             attention_mask=attention_mask,\n    #             position_ids=None,\n    #             past_key_values=None,\n    #             inputs_embeds=None,\n    #             labels=labels,\n    #             use_cache=use_cache,\n    #             output_attentions=output_attentions,\n    #             output_hidden_states=output_hidden_states,\n    #             return_dict=return_dict,\n    #         )\n\n    #     # === Handle Multimodal Forward ===\n    #     elif (input_ids.shape[0] == pixel_values.shape[0]) or (inputs_embeds.shape[0] == pixel_values.shape[0]):\n    #         assert past_key_values is None, \"Unexpected key `past_key_values` provided during multimodal forward!\"\n\n    #         #test\n    #\n    #         #test end\n\n    #         # Get input embeddings (from language model embeddings)\n    #         input_embeddings = self.get_input_embeddings()(input_ids)  # (B, seq_len, D)\n\n    #         # Extract action masks\n    #         all_actions_mask = self._process_action_masks(labels)\n\n    #         # Extract the language portion of the input embeddings (i.e. remove the action tokens portion)\n    #         language_embeddings = input_embeddings[~all_actions_mask].reshape(\n    #             input_embeddings.shape[0], -1, input_embeddings.shape[2]\n    #         )  # (B, lang_seq_len, llm_dim)\n\n    #         # Get visual features\n    #         projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)\n\n    #         # Add proprioceptive state if provided\n    #         projected_patch_embeddings = self._process_proprio_features(\n    #             projected_patch_embeddings, proprio, proprio_projector\n    #         )\n\n    #         # [Diffusion] Add diffusion timestep embedding if provided\n    #         if diffusion_timestep_embeddings is not None:\n    #             # For simplicity, just append diffusion timestep embedding to the end of projected vision patch tokens\n    #             projected_patch_embeddings = torch.cat(\n    #                 (projected_patch_embeddings, diffusion_timestep_embeddings), dim=1\n    #             )\n\n    #         # Process action embeddings\n    #         if noisy_actions is not None:\n    #             # Get mask corresponding to all action tokens\n    #             all_actions_mask = self._process_action_masks(labels)\n\n    #             # Reshape noisy actions into individual action tokens\n    #             # noisy_actions: (B, chunk_len, action_dim) -> (B, chunk_len * action_dim, 1)\n    #             B = noisy_actions.shape[0]\n    #             noisy_actions = noisy_actions.reshape(B, -1).unsqueeze(-1)\n\n    #             # Project noisy action tokens into language model embedding space\n    #             noisy_action_features = noisy_action_projector(noisy_actions)  # (B, chunk_len * action_dim, llm_dim)\n\n    #             # Replace embeddings of the action tokens with noisy action embeddings\n    #             input_embeddings = self._replace_input_embeddings(\n    #                 input_embeddings, all_actions_mask, noisy_action_features\n    #             )\n    #         else:\n    #             # Replace the embeddings of the action tokens with zeros\n    #             # (Later on, the positional embeddings will be added to them)\n    #             all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n    #             input_embeddings = input_embeddings * ~all_actions_mask\n\n    #         # Build multimodal embeddings & attention mask\n    #         multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n    #             input_embeddings, projected_patch_embeddings, attention_mask\n    #         )\n\n    #         # Build labels for multimodal sequence if needed\n    #         multimodal_labels = self._build_multimodal_labels(labels, projected_patch_embeddings)\n\n    #         # Dispatch to language model\n    #         language_model_output = self.language_model(\n    #             input_ids=None,\n    #             attention_mask=multimodal_attention_mask,\n    #             position_ids=None,\n    #             past_key_values=None,\n    #             inputs_embeds=multimodal_embeddings,\n    #             labels=multimodal_labels,\n    #             use_cache=use_cache,\n    #             output_attentions=output_attentions,\n    #             output_hidden_states=output_hidden_states,\n    #             return_dict=return_dict,\n    #         )\n\n    #     # === Otherwise =>> Assume Invalid! ===\n    #     elif (input_ids.shape[0] != pixel_values.shape[0]) or (inputs_embeds.shape[0] != pixel_values.shape[0]):\n    #         raise ValueError(\"Non-homogenous batch of (text, image) input \\\n    #                          -- forward() does not support mixed batches!\")\n\n    #     else:\n    #         raise ValueError(\n    #             \"Invalid PrismaticForConditionalGeneration `forward()` call with provided arguments:\\n\"\n    #             f\"=> `input_ids` = {input_ids is not None}\\n\"\n    #             f\"=> `attention_mask` = {attention_mask is not None}\\n\"\n    #             f\"=> `pixel_values` = {pixel_values is not None}\\n\"\n    #             f\"=> `labels` = {labels is not None}\\n\"\n    #             f\"=> `input_embeds` = {inputs_embeds is not None}\\n\"\n    #             f\"=> `past_key_values` = {past_key_values is not None}\\n\"\n    #             f\"=> `use_cache` = {use_cache}\"\n    #         )\n\n    #     # Unpack `language_model_output` and return PrismaticCausalLMOutputWithPast (or tuple if not `return_dict`)\n    #     if not return_dict:\n    #         if output_projector_features and (projected_patch_embeddings is not None):\n    #             return *language_model_output, projected_patch_embeddings\n\n    #         return language_model_output\n\n    #     return PrismaticCausalLMOutputWithPast(\n    #         loss=language_model_output.loss,\n    #         logits=language_model_output.logits,\n    #         past_key_values=language_model_output.past_key_values,\n    #         hidden_states=language_model_output.hidden_states,\n    #         attentions=language_model_output.attentions,\n    #         projector_features=projected_patch_embeddings,\n    #     )\n\n    # === GenerationMixin Methods ===\n    def prepare_inputs_for_generation(\n        self,\n        input_ids: Optional[torch.Tensor] = None,\n        past_key_values: Optional[list[torch.FloatTensor]] = None,\n        inputs_embeds: Optional[torch.FloatTensor] = None,\n        pixel_values: Optional[torch.FloatTensor] = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        **kwargs: str,\n    ) -> dict[str, torch.Tensor]:\n        \"\"\"Borrowed from `LlamaForCausalLM` and simplified for batch size = 1; mirrors original PrismaticVLM logic.\"\"\"\n        if ((input_ids is not None) and (input_ids.shape[0] > 1)) or (\n            (inputs_embeds is not None) and (inputs_embeds.shape[0] > 1)\n        ):\n            raise ValueError(\"Generation with batch size > 1 is not currently supported!\")\n\n        # Handle `past_key_values` (cache) =>> assume `input_ids` just has unprocessed tokens\n        if past_key_values is not None:\n            input_ids = input_ids[:, -1:]\n\n        # If `input_embeds` are passed, we only want to use them in the 1st generation step\n        if inputs_embeds is not None and past_key_values is None:\n            model_inputs = {\"input_embeds\": inputs_embeds}\n        else:\n            model_inputs = {\"input_ids\": input_ids}\n\n        # Make sure `pixel_values` are preserved in `model_inputs`\n        model_inputs.update(\n            {\n                \"attention_mask\": attention_mask,\n                \"pixel_values\": pixel_values,\n                \"past_key_values\": past_key_values,\n                \"use_cache\": kwargs.get(\"use_cache\"),\n            }\n        )\n\n        return model_inputs\n\n    # Defer to Language Model (all handle this differently, with different return types)\n    def _reorder_cache(self, *args, **kwargs) -> Any:\n        return self.language_model._reorder_cache(*args, **kwargs)\n\n    def _prepare_input_for_action_prediction_verl(self, input_ids, attention_mask):\n        \"\"\"Prepares input for action prediction by adding necessary tokens\"\"\"\n        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens\n        placeholder_action_token_ids = (\n            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)\n        )\n        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)\n\n        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)\n        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX\n        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)\n\n        # Extend the attention mask to fit the new shape of input\n        # Note: Only batch size == 1 supported right now\n        mask_extension = (\n            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))\n            .to(attention_mask.device)\n            .to(attention_mask.dtype)\n        )\n        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)\n\n        return input_ids, attention_mask\n\n    def _prepare_labels_for_action_prediction_verl(self, labels, input_ids):\n        \"\"\"Creates labels tensor for action prediction if not provided\"\"\"\n        # Extend labels tensor with fake action labels\n        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1\n        labels_extension = (\n            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)\n            * ARBITRARY_ACTION_TOKEN_IDX\n        )\n        labels = torch.cat([labels, labels_extension], dim=-1)\n\n        # Replace last label token with stop token\n        labels[:, -1] = STOP_INDEX\n\n        return labels\n\n    def _verl_discrete_compute_logits(\n        self,\n        input_embeddings,\n        all_actions_mask,\n        projected_patch_embeddings,\n        attention_mask,\n        labels,\n        NUM_PATCHES,\n        NUM_PROMPT_TOKENS,\n        action_head=None,\n    ):  # contintue!!!!!\n        \"\"\"Run L1 regression-based continuous action prediction or discrete action tokens prediction.\"\"\"\n        # Zero out action token embeddings\n        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n        input_embeddings = input_embeddings * ~all_actions_mask\n\n        # Build multimodal embeddings and attention mask\n        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n            input_embeddings, projected_patch_embeddings, attention_mask\n        )\n\n        # Forward pass through language model\n        language_model_output = self.language_model(\n            input_ids=None,\n            attention_mask=multimodal_attention_mask,\n            position_ids=None,\n            past_key_values=None,\n            inputs_embeds=multimodal_embeddings,\n            labels=None,\n            use_cache=None,\n            output_attentions=False,\n            output_hidden_states=False,\n            return_dict=True,\n        )\n\n        # Extract hidden states for action tokens\n        # last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)\n        # actions_hidden_states = last_hidden_states[\n        #     :,\n        #     NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n        #     :,\n        # ]  # (B, act_chunk_len, D)\n\n        # Handle different prediction methods\n        # if action_head is not None:\n        #     # L1 regression prediction\n        #     normalized_actions = action_head.predict_action(actions_hidden_states)\n        #     normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n        #     normalized_actions = normalized_actions.float().cpu().detach().numpy()\n        # else:\n        # Discrete token-based prediction\n\n        compute_logits = language_model_output.logits[\n            :,\n            NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n        ]\n\n        return compute_logits\n\n    # def forward(\n    #     self,\n    #     input_ids: Optional[torch.LongTensor] = None,\n    #     unnorm_key: Optional[str] = None,\n    #     proprio=None,\n    #     proprio_projector=None,\n    #     action_head=None,\n    #     noisy_action_projector=None,\n    #     use_film: bool = False,\n    #     **kwargs: str,\n    # ) :\n    #     \"\"\"Predict actions from input sequence, with options for different prediction methods.\n\n    #     Args:\n    #         input_ids: Input token ids\n    #         unnorm_key: Key for unnormalization statistics\n    #         proprio: Proprioceptive features\n    #         proprio_projector: Projector for proprioceptive features\n    #         action_head: Optional head for L1 regression or diffusion-based prediction\n    #         noisy_action_projector: Projector for noisy actions in diffusion-based prediction\n    #         use_film: Whether to use FiLM conditioning\n    #         **kwargs: Additional arguments including pixel_values and attention_mask\n\n    #     Returns:\n    #         Tuple of (unnormalized_actions, action_hidden_states)\n    #     \"\"\"\n    #     # If the special empty token ('') does not already appear after the colon (':') token in the prompt\n    #     # (after \"OUT:\" or \"ASSISTANT:\"), insert it to match the inputs seen at training time\n    #     # if not torch.all(input_ids[:, -1] == 29871):\n    #     #     input_ids = torch.cat(\n    #     #         (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1\n    #     #     )\n    #     #print(\"!!!!!!!!!!!!!!Entering forward!!!!!!!!!!\")\n    #     pixel_values = kwargs[\"pixel_values\"]\n    #     attention_mask = kwargs[\"attention_mask\"]\n\n    #     # Create fake labels tensor (needed for action mask)\n    #     labels = input_ids.clone()\n    #     labels[:] = IGNORE_INDEX\n\n    #     # Get number of tokens in prompt (excluding the start token)\n    #     NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token\n\n    #     # Prepare inputs by adding necessary tokens\n    #     #input_ids, attention_mask = self._prepare_input_for_action_prediction_verl(input_ids, attention_mask)\n\n    #     #test\n    #     placeholder_action_token_ids = (\n    #         torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)\n    #     )\n    #     input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)\n\n    #     # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)\n    #     stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX\n    #     input_ids = torch.cat([input_ids, stop_token_id], dim=-1)\n\n    #     # Extend the attention mask to fit the new shape of input\n    #     # Note: Only batch size == 1 supported right now\n    #     mask_extension = (\n    #         torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))\n    #         .to(attention_mask.device)\n    #         .to(attention_mask.dtype)\n    #     )\n    #     attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)\n\n    #     #return input_ids, attention_mask\n\n    #     #test end\n\n    #     # Update labels tensor for action mask computation later\n    #     #labels = self._prepare_labels_for_action_prediction_verl(labels, input_ids)\n    #     #test\n\n    #     ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1\n    #     labels_extension = (\n    #         torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)\n    #         * ARBITRARY_ACTION_TOKEN_IDX\n    #     )\n    #     labels = torch.cat([labels, labels_extension], dim=-1)\n\n    #     # Replace last label token with stop token\n    #     labels[:, -1] = STOP_INDEX\n\n    #     #return labels\n\n    #     #test ed\n\n    #     # Get input embeddings and action masks\n\n    #     input_embeddings = self.get_input_embeddings()(input_ids)\n\n    #     #all_actions_mask = self._process_action_masks(labels)\n    #     #test\n    #     #current_action_mask = get_current_action_mask(labels)\n    #     newline_positions = labels != IGNORE_INDEX\n\n    #     # Calculate cumulative sum to identify regions between newlines\n    #     cumsum = torch.cumsum(newline_positions, dim=1)\n\n    #     # Create the mask\n    #     mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)\n\n    #     # Extract the action part only\n    #     action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX\n    #     current_action_mask = action_tokens_only_mask * mask\n\n    #     #next_actions_mask = get_next_actions_mask(labels)\n    #     newline_positions = labels != IGNORE_INDEX\n\n    #     # Calculate cumulative sum to identify regions between newlines\n    #     cumsum = torch.cumsum(newline_positions, dim=1)\n\n    #     # Create the mask\n    #     mask = cumsum > ACTION_DIM\n\n    #     # Extract the action part only\n    #     action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX\n    #     next_actions_mask = action_tokens_only_mask * mask\n\n    #     all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)\n\n    #     #test end\n\n    #     # Extract language embeddings\n    #     language_embeddings = input_embeddings[~all_actions_mask].reshape(\n    #         input_embeddings.shape[0], -1, input_embeddings.shape[2]\n    #     )\n\n    #     # Process vision features\n    #     #projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)\n    #     #test\n    #     if use_film:\n    #         # FiLM: Infuse language inputs into visual features\n    #         raise ValueError\n    #         patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)\n    #     else:\n    #         patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)\n\n    #     projected_patch_embeddings = self.projector(patch_features)\n    #     #test end\n\n    #     # Add proprioceptive features if provided\n    #     use_proprio = proprio_projector is not None and proprio is not None\n    #     if use_proprio:\n    #         proprio = torch.Tensor(proprio).to(projected_patch_embeddings.device,\n    #                                            dtype=projected_patch_embeddings.dtype)\n    #         projected_patch_embeddings = self._process_proprio_features(\n    #             projected_patch_embeddings, proprio, proprio_projector\n    #         )\n\n    #     # Use diffusion if provided, otherwise use regression or discrete prediction\n    #     use_diffusion = noisy_action_projector is not None and hasattr(action_head, \"noise_scheduler\")\n\n    #     # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)\n    #     NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()\n    #     if use_proprio:\n    #         NUM_PATCHES += 1\n    #     if use_diffusion:\n    #         NUM_PATCHES += 1\n\n    #     if use_diffusion:\n    #         raise ValueError\n    #         # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion\n    #         noise = torch.randn(\n    #             size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype\n    #         )\n\n    #         # Run diffusion-based prediction\n    #         normalized_actions, actions_hidden_states = self._run_diffusion_prediction(\n    #             input_embeddings,\n    #             all_actions_mask,\n    #             noise,\n    #             action_head,\n    #             projected_patch_embeddings,\n    #             labels,\n    #             attention_mask,\n    #             NUM_PATCHES,\n    #             NUM_PROMPT_TOKENS,\n    #             noisy_action_projector,\n    #         )\n    #     else:\n    #         # Run regression or discrete token-based prediction\n    #         # compute_logits = self._verl_discrete_compute_logits(\n    #         #     input_embeddings,\n    #         #     all_actions_mask,\n    #         #     projected_patch_embeddings,\n    #         #     attention_mask,\n    #         #     labels,\n    #         #     NUM_PATCHES,\n    #         #     NUM_PROMPT_TOKENS,\n    #         #     action_head,\n    #         # )\n\n    #         #test\n\n    #         all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n    #         input_embeddings = input_embeddings * ~all_actions_mask\n\n    #         # Build multimodal embeddings and attention mask\n    #         # multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n    #         #     input_embeddings, projected_patch_embeddings, attention_mask\n    #         # )\n    #         #test\n\n    #         projected_patch_attention_mask = None\n    #         if attention_mask is not None:\n    #             projected_patch_attention_mask = torch.full(\n    #                 (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),\n    #                 fill_value=True,\n    #                 dtype=attention_mask.dtype,\n    #                 device=attention_mask.device,\n    #             )\n\n    #         # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)\n    #         multimodal_embeddings = torch.cat(\n    #             [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1\n    #         )\n\n    #         multimodal_attention_mask = None\n    #         if attention_mask is not None:\n    #             multimodal_attention_mask = torch.cat(\n    #                 [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1\n    #             )\n\n    #         #return multimodal_embeddings, multimodal_attention_mask\n\n    #         #test end\n\n    #         # Forward pass through language model\n    #         language_model_output = self.language_model(\n    #             input_ids=None,\n    #             attention_mask=multimodal_attention_mask,\n    #             position_ids=None,\n    #             past_key_values=None,\n    #             inputs_embeds=multimodal_embeddings,\n    #             labels=None,\n    #             use_cache=None,\n    #             output_attentions=False,\n    #             output_hidden_states=False,\n    #             return_dict=True,\n    #         )\n\n    #         compute_logits = language_model_output.logits[\n    #                     :,\n    #                     NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + \\\n    #                     NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n    #                 ]\n\n    #         #test end\n\n    #     return compute_logits\n\n    def forward(\n        self,\n        input_ids: Optional[torch.LongTensor] = None,\n        pixel_values=None,\n        attention_mask=None,\n        # labels=None,\n        proprio=None,\n        proprio_projector=None,\n        action_head=None,\n        noisy_action_projector=None,\n        use_film: bool = False,\n        **kwargs: str,\n    ):\n        \"\"\"Predict actions from input sequence, with options for different prediction methods.\n\n        Args:\n            input_ids: Input token ids\n            unnorm_key: Key for unnormalization statistics\n            proprio: Proprioceptive features\n            proprio_projector: Projector for proprioceptive features\n            action_head: Optional head for L1 regression or diffusion-based prediction\n            noisy_action_projector: Projector for noisy actions in diffusion-based prediction\n            use_film: Whether to use FiLM conditioning\n            **kwargs: Additional arguments including pixel_values and attention_mask\n\n        Returns:\n            Tuple of (unnormalized_actions, action_hidden_states)\n        \"\"\"\n        # If the special empty token ('') does not already appear after the colon (':') token in the prompt\n        # (after \"OUT:\" or \"ASSISTANT:\"), insert it to match the inputs seen at training time\n        # if not torch.all(input_ids[:, -1] == 29871):\n        #     input_ids = torch.cat(\n        #         (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1\n        #     )\n\n        # pixel_values = kwargs[\"pixel_values\"]\n        # attention_mask = kwargs[\"attention_mask\"]\n\n        # Create fake labels tensor (needed for action mask)\n        labels = input_ids.clone()\n        labels[:] = IGNORE_INDEX\n\n        # # Get number of tokens in prompt (excluding the start token)\n        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token\n\n        # # Prepare inputs by adding necessary tokens\n        # #input_ids, attention_mask = self._prepare_input_for_action_prediction_verl(input_ids, attention_mask)\n\n        # #test\n        placeholder_action_token_ids = (\n            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)\n        )\n        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)\n\n        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)\n        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX\n        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)\n\n        # Extend the attention mask to fit the new shape of input\n        # Note: Only batch size == 1 supported right now\n        mask_extension = (\n            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))\n            .to(attention_mask.device)\n            .to(attention_mask.dtype)\n        )\n        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)\n\n        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1\n        labels_extension = (\n            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)\n            * ARBITRARY_ACTION_TOKEN_IDX\n        )\n        labels = torch.cat([labels, labels_extension], dim=-1)\n\n        # # Replace last label token with stop token\n        labels[:, -1] = STOP_INDEX\n\n        # Get input embeddings and action masks\n\n        # NUM_PROMPT_TOKENS = kwargs[\"num_prompt_tokens\"]\n\n        input_embeddings = self.get_input_embeddings()(input_ids)\n\n        # all_actions_mask = self._process_action_masks(labels)\n        # test\n        # current_action_mask = get_current_action_mask(labels)\n        newline_positions = labels != IGNORE_INDEX\n\n        # Calculate cumulative sum to identify regions between newlines\n        cumsum = torch.cumsum(newline_positions, dim=1)\n\n        # Create the mask\n        mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)\n\n        # Extract the action part only\n        action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX\n        current_action_mask = action_tokens_only_mask * mask\n\n        # next_actions_mask = get_next_actions_mask(labels)\n        newline_positions = labels != IGNORE_INDEX\n\n        # Calculate cumulative sum to identify regions between newlines\n        cumsum = torch.cumsum(newline_positions, dim=1)\n\n        # Create the mask\n        mask = cumsum > ACTION_DIM\n\n        # Extract the action part only\n        action_tokens_only_mask = labels > ACTION_TOKEN_BEGIN_IDX\n        next_actions_mask = action_tokens_only_mask * mask\n\n        all_actions_mask = current_action_mask | next_actions_mask  # (B, seq_len)\n\n        # test end\n\n        # Extract language embeddings\n        language_embeddings = input_embeddings[~all_actions_mask].reshape(\n            input_embeddings.shape[0], -1, input_embeddings.shape[2]\n        )\n\n        # Process vision features\n        # projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)\n        # test\n        if use_film:\n            # FiLM: Infuse language inputs into visual features\n            raise ValueError\n            patch_features = self.vision_backbone(pixel_values, language_embeddings)  # (bsz, 256 * num_images, D)\n        else:\n            patch_features = self.vision_backbone(pixel_values)  # (bsz, 256 * num_images, D)\n\n        projected_patch_embeddings = self.projector(patch_features)\n        # test end\n\n        # Add proprioceptive features if provided\n        use_proprio = proprio_projector is not None and proprio is not None\n        if use_proprio:\n            proprio = torch.Tensor(proprio).to(\n                projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype\n            )\n            projected_patch_embeddings = self._process_proprio_features(\n                projected_patch_embeddings, proprio, proprio_projector\n            )\n\n        # Use diffusion if provided, otherwise use regression or discrete prediction\n        use_diffusion = noisy_action_projector is not None and hasattr(action_head, \"noise_scheduler\")\n\n        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)\n        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()\n        if use_proprio:\n            NUM_PATCHES += 1\n        if use_diffusion:\n            NUM_PATCHES += 1\n\n        if use_diffusion:\n            raise ValueError\n            # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion\n            noise = torch.randn(\n                size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype\n            )\n\n            # Run diffusion-based prediction\n            normalized_actions, actions_hidden_states = self._run_diffusion_prediction(\n                input_embeddings,\n                all_actions_mask,\n                noise,\n                action_head,\n                projected_patch_embeddings,\n                labels,\n                attention_mask,\n                NUM_PATCHES,\n                NUM_PROMPT_TOKENS,\n                noisy_action_projector,\n            )\n        else:\n            # Run regression or discrete token-based prediction\n            # compute_logits = self._verl_discrete_compute_logits(\n            #     input_embeddings,\n            #     all_actions_mask,\n            #     projected_patch_embeddings,\n            #     attention_mask,\n            #     labels,\n            #     NUM_PATCHES,\n            #     NUM_PROMPT_TOKENS,\n            #     action_head,\n            # )\n\n            # test\n\n            all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n            input_embeddings = input_embeddings * ~all_actions_mask\n\n            # Build multimodal embeddings and attention mask\n            # multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n            #     input_embeddings, projected_patch_embeddings, attention_mask\n            # )\n            # test\n\n            projected_patch_attention_mask = None\n            if attention_mask is not None:\n                projected_patch_attention_mask = torch.full(\n                    (projected_patch_embeddings.shape[0], projected_patch_embeddings.shape[1]),\n                    fill_value=True,\n                    dtype=attention_mask.dtype,\n                    device=attention_mask.device,\n                )\n\n            # Build multimodal embeddings & attention mask; insert embeddings after <BOS> token (1:)\n            multimodal_embeddings = torch.cat(\n                [input_embeddings[:, :1, :], projected_patch_embeddings, input_embeddings[:, 1:, :]], dim=1\n            )\n\n            multimodal_attention_mask = None\n            if attention_mask is not None:\n                multimodal_attention_mask = torch.cat(\n                    [attention_mask[:, :1], projected_patch_attention_mask, attention_mask[:, 1:]], dim=1\n                )\n\n            # return multimodal_embeddings, multimodal_attention_mask\n\n            # test end\n\n            # Forward pass through language model\n            language_model_output = self.language_model(\n                input_ids=None,\n                attention_mask=multimodal_attention_mask,\n                position_ids=None,\n                past_key_values=None,\n                inputs_embeds=multimodal_embeddings,\n                labels=None,\n                use_cache=None,\n                output_attentions=False,\n                output_hidden_states=False,\n                return_dict=True,\n            )\n\n            compute_logits = language_model_output.logits[\n                :,\n                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n            ]\n\n            # test end\n\n        return compute_logits\n\n\nclass OpenVLAForActionPrediction(PrismaticForConditionalGeneration):\n    config_class: PretrainedConfig = OpenVLAConfig\n    _supports_sdpa = True\n\n    def __init__(self, config: OpenVLAConfig) -> None:\n        super().__init__(config)\n        self.norm_stats = config.norm_stats\n\n        # Compute action bins\n        self.bins = np.linspace(-1, 1, config.n_action_bins)\n        self.bin_centers = (self.bins[:-1] + self.bins[1:]) / 2.0\n\n        # Compute vocab size for de-tokenization -- revert added \"multiple of\"\n        self.vocab_size = self.config.text_config.vocab_size - self.config.pad_to_multiple_of\n\n    def _prepare_input_for_action_prediction(self, input_ids, attention_mask):\n        \"\"\"Prepares input for action prediction by adding necessary tokens\"\"\"\n        # Add (ACTION_DIM * NUM_ACTIONS_CHUNK) placeholder tokens to input_ids to simulate action tokens\n        placeholder_action_token_ids = (\n            torch.ones((input_ids.shape[0], ACTION_DIM * NUM_ACTIONS_CHUNK)).to(input_ids.device).to(input_ids.dtype)\n        )\n        input_ids = torch.cat([input_ids, placeholder_action_token_ids], dim=-1)\n\n        # Add stop token to sequence (needed in non-causal bi-directional self-attention, as it appears at train time)\n        stop_token_id = torch.ones((input_ids.shape[0], 1)).to(input_ids.device).to(input_ids.dtype) * STOP_INDEX\n        input_ids = torch.cat([input_ids, stop_token_id], dim=-1)\n\n        # Extend the attention mask to fit the new shape of input\n        # Note: Only batch size == 1 supported right now\n        mask_extension = (\n            torch.ones((attention_mask.shape[0], input_ids.shape[-1] - attention_mask.shape[-1]))\n            .to(attention_mask.device)\n            .to(attention_mask.dtype)\n        )\n        attention_mask = torch.cat([attention_mask, mask_extension], dim=-1)\n\n        return input_ids, attention_mask\n\n    def _prepare_labels_for_action_prediction(self, labels, input_ids):\n        \"\"\"Creates labels tensor for action prediction if not provided\"\"\"\n        # Extend labels tensor with fake action labels\n        ARBITRARY_ACTION_TOKEN_IDX = ACTION_TOKEN_BEGIN_IDX + 1\n        labels_extension = (\n            torch.ones((labels.shape[0], input_ids.shape[-1] - labels.shape[-1])).to(labels.device).to(labels.dtype)\n            * ARBITRARY_ACTION_TOKEN_IDX\n        )\n        labels = torch.cat([labels, labels_extension], dim=-1)\n\n        # Replace last label token with stop token\n        labels[:, -1] = STOP_INDEX\n\n        return labels\n\n    def _unnormalize_actions(self, normalized_actions, unnorm_key=None):\n        \"\"\"Unnormalize actions using dataset statistics\"\"\"\n        action_norm_stats = self.get_action_stats(unnorm_key)\n\n        if ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS:\n            mask = action_norm_stats.get(\"mask\", np.ones_like(action_norm_stats[\"min\"], dtype=bool))\n            action_high, action_low = np.array(action_norm_stats[\"max\"]), np.array(action_norm_stats[\"min\"])\n        elif ACTION_PROPRIO_NORMALIZATION_TYPE == NormalizationType.BOUNDS_Q99:\n            mask = action_norm_stats.get(\"mask\", np.ones_like(action_norm_stats[\"q01\"], dtype=bool))\n            action_high, action_low = np.array(action_norm_stats[\"q99\"]), np.array(action_norm_stats[\"q01\"])\n        else:\n            raise ValueError(\"Unsupported action/proprio normalization type detected!\")\n\n        actions = np.where(\n            mask,\n            0.5 * (normalized_actions + 1) * (action_high - action_low + 1e-8) + action_low,\n            normalized_actions,\n        )\n\n        return actions\n\n    def _run_diffusion_prediction(\n        self,\n        input_embeddings,\n        all_actions_mask,\n        noise,\n        action_head,\n        projected_patch_embeddings,\n        labels,\n        attention_mask,\n        NUM_PATCHES,\n        NUM_PROMPT_TOKENS,\n        noisy_action_projector,\n    ):\n        \"\"\"Run diffusion-based action prediction\"\"\"\n        # Set diffusion timestep values\n        action_head.noise_scheduler.set_timesteps(action_head.num_diffusion_steps)\n        # Clone embedding for reuse in each timestep\n        orig_projected_patch_embeddings = projected_patch_embeddings.clone()\n        curr_noisy_actions = noise\n\n        # Reverse diffusion: Iteratively denoise to generate action prediction\n        for t in action_head.noise_scheduler.timesteps:\n            # Get diffusion model's noise prediction (conditioned on VLA latent embedding, current noisy action\n            # embedding, and diffusion timestep embedding)\n            timesteps = torch.Tensor([t]).to(labels.device)\n            diffusion_timestep_embeddings = (\n                action_head.time_encoder(timesteps).to(curr_noisy_actions.dtype).to(curr_noisy_actions.device)\n            )  # (B, llm_dim)\n            diffusion_timestep_embeddings = diffusion_timestep_embeddings.unsqueeze(1)  # (B, 1, llm_dim)\n\n            # [Diffusion] Replace the embeddings of the action tokens with noisy actions\n            # (Later on, the positional embeddings will be added to them)\n\n            # For simplicity, append diffusion timestep embedding to the end of projected vision tokens\n            projected_patch_embeddings = torch.cat(\n                (orig_projected_patch_embeddings, diffusion_timestep_embeddings), dim=1\n            )\n\n            # Reshape and project noisy actions into language embedding space\n            B = curr_noisy_actions.shape[0]\n            orig_curr_noisy_actions_shape = curr_noisy_actions.shape\n            curr_noisy_actions = curr_noisy_actions.reshape(B, -1).unsqueeze(-1)\n            noisy_action_features = noisy_action_projector(curr_noisy_actions)\n            curr_noisy_actions = curr_noisy_actions.reshape(orig_curr_noisy_actions_shape)\n\n            # Replace action token embeddings with noisy action embeddings\n            input_embeddings = self._replace_input_embeddings(\n                input_embeddings.clone(), all_actions_mask, noisy_action_features\n            )\n\n            # Build multimodal embeddings and attention mask\n            multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n                input_embeddings, projected_patch_embeddings, attention_mask\n            )\n\n            # Forward pass through language model\n            language_model_output = self.language_model(\n                input_ids=None,\n                attention_mask=multimodal_attention_mask,\n                position_ids=None,\n                past_key_values=None,\n                inputs_embeds=multimodal_embeddings,\n                labels=None,\n                use_cache=None,\n                output_attentions=False,\n                output_hidden_states=True,\n                return_dict=True,\n            )\n\n            # Extract hidden states for action portion of response\n            last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)\n            actions_hidden_states = last_hidden_states[\n                :,\n                NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n                :,\n            ]  # (B, act_chunk_len, D)\n\n            # Predict noise and update noisy actions: x_t -> x_{t-1}\n            noise_pred = action_head.predict_noise(actions_hidden_states)\n            curr_noisy_actions = action_head.noise_scheduler.step(noise_pred, t, curr_noisy_actions).prev_sample\n\n        curr_noisy_actions = curr_noisy_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n\n        # Return final actions\n        return curr_noisy_actions.float().cpu().detach().numpy(), actions_hidden_states\n\n    def _regression_or_discrete_prediction(\n        self,\n        input_embeddings,\n        all_actions_mask,\n        projected_patch_embeddings,\n        attention_mask,\n        labels,\n        NUM_PATCHES,\n        NUM_PROMPT_TOKENS,\n        action_head=None,\n    ):\n        \"\"\"Run L1 regression-based continuous action prediction or discrete action tokens prediction.\"\"\"\n        # Zero out action token embeddings\n        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n        input_embeddings = input_embeddings * ~all_actions_mask\n\n        # Build multimodal embeddings and attention mask\n        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n            input_embeddings, projected_patch_embeddings, attention_mask\n        )\n\n        # Forward pass through language model\n        language_model_output = self.language_model(\n            input_ids=None,\n            attention_mask=multimodal_attention_mask,\n            position_ids=None,\n            past_key_values=None,\n            inputs_embeds=multimodal_embeddings,\n            labels=None,\n            use_cache=None,\n            output_attentions=False,\n            output_hidden_states=True,\n            return_dict=True,\n        )\n\n        # Extract hidden states for action tokens\n        last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)\n        actions_hidden_states = last_hidden_states[\n            :,\n            NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n            :,\n        ]  # (B, act_chunk_len, D)\n\n        # Handle different prediction methods\n        if action_head is not None:\n            # L1 regression prediction\n            normalized_actions = action_head.predict_action(actions_hidden_states)\n            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n            normalized_actions = normalized_actions.float().cpu().detach().numpy()\n        else:\n            # Discrete token-based prediction\n            predicted_action_token_ids = (\n                language_model_output.logits[\n                    :,\n                    NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n                ]\n                .argmax(dim=2)\n                .cpu()\n                .numpy()\n            )\n            discretized_actions = self.vocab_size - predicted_action_token_ids\n            discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)\n            normalized_actions = self.bin_centers[discretized_actions]\n            normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n\n        return normalized_actions, actions_hidden_states\n\n    def _verl_discrete_prediction(\n        self,\n        input_embeddings,\n        all_actions_mask,\n        projected_patch_embeddings,\n        attention_mask,\n        labels,\n        NUM_PATCHES,\n        NUM_PROMPT_TOKENS,\n        action_head=None,\n        do_sample=True,\n        temperature=1,\n    ):\n        \"\"\"Run L1 regression-based continuous action prediction or discrete action tokens prediction.\"\"\"\n        # Zero out action token embeddings\n        all_actions_mask = all_actions_mask.unsqueeze(-1)  # (B, seq_len, 1)\n        input_embeddings = input_embeddings * ~all_actions_mask\n\n        # Build multimodal embeddings and attention mask\n        multimodal_embeddings, multimodal_attention_mask = self._build_multimodal_attention(\n            input_embeddings, projected_patch_embeddings, attention_mask\n        )\n\n        # Forward pass through language model\n        language_model_output = self.language_model(\n            input_ids=None,\n            attention_mask=multimodal_attention_mask,\n            position_ids=None,\n            past_key_values=None,\n            inputs_embeds=multimodal_embeddings,\n            labels=None,\n            use_cache=None,\n            output_attentions=False,\n            output_hidden_states=False,\n            return_dict=True,\n        )\n\n        # Extract hidden states for action tokens\n        # last_hidden_states = language_model_output.hidden_states[-1]  # (B, seq_len, D)\n        # actions_hidden_states = last_hidden_states[\n        #     :,\n        #     NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n        #     :,\n        # ]  # (B, act_chunk_len, D)\n\n        # Handle different prediction methods\n        # if action_head is not None:\n        #     # L1 regression prediction\n        #     normalized_actions = action_head.predict_action(actions_hidden_states)\n        #     normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n        #     normalized_actions = normalized_actions.float().cpu().detach().numpy()\n        # else:\n        # Discrete token-based prediction\n\n        # test\n        # NUM_PROMPT_TOKENS = NUM_PROMPT_TOKENS + NUM_PATCHES\n        # j = torch.arange(language_model_output.logits.shape[1], device=NUM_PROMPT_TOKENS.device)\n        # start = NUM_PROMPT_TOKENS.unsqueeze(1)\n        # end = start + ACTION_DIM * NUM_ACTIONS_CHUNK\n        # mask_2d = (j >= start) & (j < end)\n        # mask = mask_2d.unsqueeze(-1)\n        # actions_masks = mask.expand_as(language_model_output.logits)\n\n        NUM_PROMPT_TOKENS = NUM_PROMPT_TOKENS + NUM_PATCHES\n        batch_size = language_model_output.logits.shape[0]\n        device = language_model_output.logits.device\n\n        start_indices = NUM_PROMPT_TOKENS.unsqueeze(1)  # [batch_size, 1]\n        position_offsets = torch.arange(ACTION_DIM * NUM_ACTIONS_CHUNK, device=device).unsqueeze(0)  # [1, seq_length]\n        seq_indices = start_indices + position_offsets  # [batch_size, ACTION_DIM*NUM_ACTIONS_CHUNK]\n        # test end\n        # test add\n        # print(\"language_model_output\",language_model_output.logits.shape[-1])\n        # print(\"self.vocab_size\",self.vocab_size) 32000\n        # topk_values, topk_indices = torch.topk(language_model_output.logits, k=256, dim=-1)\n        # print(topk_indices)\n        # assert language_model_output.logits.shape[-1] == self.vocab_size\n        # test add\n        if not do_sample:\n            # org\n            # reponse_ids = language_model_output.logits[\n            #         :,\n            #         NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES +\\\n            #         NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n            #     ].argmax(dim=2)\n            # reponse_ids = language_model_output.logits[actions_masks].argmax(dim=2)\n            # org end\n\n            # padding\n            # reponse_ids = language_model_output.logits[\n            #     torch.arange(batch_size, device=device).unsqueeze(-1),\n            #     seq_indices,\n            #     :\n            # ].argmax(dim=2)\n            # padding end\n\n            # padding + only get last 256 token\n            reponse_ids_logits = language_model_output.logits[\n                torch.arange(batch_size, device=device).unsqueeze(-1), seq_indices, :\n            ]\n            start_index = self.vocab_size - 256\n            response_last256 = reponse_ids_logits[..., -256 - 64 : -64]  # Shape: [batch_size, seq_len, 256]\n            last256_argmax = response_last256.argmax(dim=-1)  # Shape: [batch_size, seq_len]\n            reponse_ids = last256_argmax + start_index  # Shape: [batch_size, seq_len]\n            # padding + only get last 256 token end\n\n            predicted_action_token_ids = reponse_ids.cpu().numpy()\n\n        else:\n            assert temperature > 0\n            # org\n            # action_logits  = language_model_output.logits[\n            #         :,\n            #         NUM_PATCHES + NUM_PROMPT_TOKENS : NUM_PATCHES + \\\n            #         NUM_PROMPT_TOKENS + ACTION_DIM * NUM_ACTIONS_CHUNK,\n            #     ]\n            # action_logits = language_model_output.logits[actions_masks]\n            # org end\n\n            action_logits = language_model_output.logits[\n                torch.arange(batch_size, device=device).unsqueeze(-1), seq_indices, :\n            ]\n            # padding\n            # scaled_logits = action_logits / temperature\n            # probs = torch.softmax(scaled_logits, dim=-1)\n            # probs_flat = probs.reshape(-1, probs.shape[-1])  # (B*act_chunk_len, vocab_size)\n            # sampled_indices_flat = torch.multinomial(probs_flat, num_samples=1)  # (B*act_chunk_len, 1)\n            # reponse_ids = sampled_indices_flat.view(action_logits.shape[0], -1)\n            # padding end\n\n            # padding + only get last 256 token\n            action_logits_last256 = action_logits[..., -256 - 64 : -64]\n            scaled_logits = action_logits_last256 / temperature\n            probs = torch.softmax(scaled_logits, dim=-1)\n            assert probs.shape[-1] == 256\n            probs_flat = probs.reshape(-1, probs.shape[-1])\n            sampled_indices_flat = torch.multinomial(probs_flat, num_samples=1)\n            original_ids_flat = sampled_indices_flat + (self.vocab_size - 256)\n            reponse_ids = original_ids_flat.view(action_logits.shape[0], -1)\n            # padding + only get last 256 token end\n\n            predicted_action_token_ids = reponse_ids.cpu().numpy()\n\n        discretized_actions = self.vocab_size - predicted_action_token_ids\n        discretized_actions = np.clip(discretized_actions - 1, a_min=0, a_max=self.bin_centers.shape[0] - 1)\n        normalized_actions = self.bin_centers[discretized_actions]\n        # normalized_actions = normalized_actions.reshape(NUM_ACTIONS_CHUNK, ACTION_DIM)\n        normalized_actions = normalized_actions.reshape(-1, ACTION_DIM)\n\n        return normalized_actions, reponse_ids\n        # return normalized_actions, actions_hidden_states\n\n    def predict_action(\n        self,\n        input_ids: Optional[torch.LongTensor] = None,\n        unnorm_key: Optional[str] = None,\n        proprio=None,\n        proprio_projector=None,\n        action_head=None,\n        noisy_action_projector=None,\n        use_film: bool = False,\n        **kwargs: str,\n    ) -> np.ndarray:\n        \"\"\"Predict actions from input sequence, with options for different prediction methods.\n\n        Args:\n            input_ids: Input token ids\n            unnorm_key: Key for unnormalization statistics\n            proprio: Proprioceptive features\n            proprio_projector: Projector for proprioceptive features\n            action_head: Optional head for L1 regression or diffusion-based prediction\n            noisy_action_projector: Projector for noisy actions in diffusion-based prediction\n            use_film: Whether to use FiLM conditioning\n            **kwargs: Additional arguments including pixel_values and attention_mask\n\n        Returns:\n            Tuple of (unnormalized_actions, action_hidden_states)\n        \"\"\"\n        # If the special empty token ('') does not already appear after the colon (':') token in the prompt\n        # (after \"OUT:\" or \"ASSISTANT:\"), insert it to match the inputs seen at training time\n        if not torch.all(input_ids[:, -1] == 29871):\n            input_ids = torch.cat(\n                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1\n            )\n\n        pixel_values = kwargs[\"pixel_values\"]\n        attention_mask = kwargs[\"attention_mask\"]\n\n        # Create fake labels tensor (needed for action mask)\n        labels = input_ids.clone()\n        labels[:] = IGNORE_INDEX\n\n        # Get number of tokens in prompt (excluding the start token)\n        NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token\n\n        # Prepare inputs by adding necessary tokens\n        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)\n\n        # Update labels tensor for action mask computation later\n        labels = self._prepare_labels_for_action_prediction(labels, input_ids)\n\n        # Get input embeddings and action masks\n        input_embeddings = self.get_input_embeddings()(input_ids)\n        all_actions_mask = self._process_action_masks(labels)\n\n        # Extract language embeddings\n        language_embeddings = input_embeddings[~all_actions_mask].reshape(\n            input_embeddings.shape[0], -1, input_embeddings.shape[2]\n        )\n\n        # Process vision features\n        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)\n\n        # Add proprioceptive features if provided\n        use_proprio = proprio_projector is not None and proprio is not None\n        if use_proprio:\n            proprio = torch.Tensor(proprio).to(\n                projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype\n            )\n            projected_patch_embeddings = self._process_proprio_features(\n                projected_patch_embeddings, proprio, proprio_projector\n            )\n\n        # Use diffusion if provided, otherwise use regression or discrete prediction\n        use_diffusion = noisy_action_projector is not None and hasattr(action_head, \"noise_scheduler\")\n\n        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)\n        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()\n        if use_proprio:\n            NUM_PATCHES += 1\n        if use_diffusion:\n            NUM_PATCHES += 1\n\n        if use_diffusion:\n            # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion\n            noise = torch.randn(\n                size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype\n            )\n\n            # Run diffusion-based prediction\n            normalized_actions, actions_hidden_states = self._run_diffusion_prediction(\n                input_embeddings,\n                all_actions_mask,\n                noise,\n                action_head,\n                projected_patch_embeddings,\n                labels,\n                attention_mask,\n                NUM_PATCHES,\n                NUM_PROMPT_TOKENS,\n                noisy_action_projector,\n            )\n        else:\n            # Run regression or discrete token-based prediction\n            normalized_actions, actions_hidden_states = self._regression_or_discrete_prediction(\n                input_embeddings,\n                all_actions_mask,\n                projected_patch_embeddings,\n                attention_mask,\n                labels,\n                NUM_PATCHES,\n                NUM_PROMPT_TOKENS,\n                action_head,\n            )\n\n        # Unnormalize predicted actions\n        actions = self._unnormalize_actions(normalized_actions, unnorm_key)\n\n        return actions, actions_hidden_states\n\n    def generate_action_verl(\n        self,\n        input_ids: Optional[torch.LongTensor] = None,\n        unnorm_key: Optional[str] = None,\n        proprio=None,\n        proprio_projector=None,\n        action_head=None,\n        noisy_action_projector=None,\n        use_film: bool = False,\n        **kwargs: str,\n    ) -> np.ndarray:\n        \"\"\"Predict actions from input sequence, with options for different prediction methods.\n\n        Args:\n            input_ids: Input token ids\n            unnorm_key: Key for unnormalization statistics\n            proprio: Proprioceptive features\n            proprio_projector: Projector for proprioceptive features\n            action_head: Optional head for L1 regression or diffusion-based prediction\n            noisy_action_projector: Projector for noisy actions in diffusion-based prediction\n            use_film: Whether to use FiLM conditioning\n            **kwargs: Additional arguments including pixel_values and attention_mask\n\n        Returns:\n            Tuple of (unnormalized_actions, action_hidden_states)\n        \"\"\"\n        # If the special empty token ('') does not already appear after the colon (':') token in the prompt\n        # (after \"OUT:\" or \"ASSISTANT:\"), insert it to match the inputs seen at training time\n        # if not torch.all(input_ids[:, -1] == 29871):\n        #     input_ids = torch.cat(\n        #         (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1\n        #     )\n\n        pixel_values = kwargs[\"pixel_values\"]\n        attention_mask = kwargs[\"attention_mask\"]\n        do_sample = kwargs[\"do_sample\"]\n        temperature = kwargs[\"temperature\"]\n\n        # Create fake labels tensor (needed for action mask)\n        labels = input_ids.clone()\n        labels[:] = IGNORE_INDEX\n\n        # Get number of tokens in prompt (excluding the start token)\n        # NUM_PROMPT_TOKENS = input_ids.shape[-1] - 1  # Subtract action tokens and stop token\n        # test\n        padding_idx = kwargs[\"padding_idx\"]\n        num_prompt_tokens = input_ids.ne(padding_idx).sum(dim=1) - 1\n        # test end\n\n        # Prepare inputs by adding necessary tokens\n        input_ids, attention_mask = self._prepare_input_for_action_prediction(input_ids, attention_mask)\n\n        # Update labels tensor for action mask computation later\n        labels = self._prepare_labels_for_action_prediction(labels, input_ids)\n\n        # here to convert padding from before to last\n        # test\n        padding_mask = input_ids.ne(padding_idx)\n        assert torch.all(padding_mask == attention_mask.ne(0))\n        # print(\"in predict_action padding_mask:\", padding_mask)\n        padding_mask = padding_mask.int()\n        sorted_indices = torch.argsort(padding_mask, dim=1, descending=True, stable=True)\n        input_ids = torch.gather(input_ids, 1, sorted_indices)\n        attention_mask = torch.gather(attention_mask, 1, sorted_indices)\n        labels = torch.gather(labels, 1, sorted_indices)\n        assert not use_film\n        # test end\n\n        # Get input embeddings and action masks\n        input_embeddings = self.get_input_embeddings()(input_ids)\n        all_actions_mask = self._process_action_masks(labels)\n\n        # Extract language embeddings\n        language_embeddings = input_embeddings[~all_actions_mask].reshape(\n            input_embeddings.shape[0], -1, input_embeddings.shape[2]\n        )\n\n        # Process vision features\n        projected_patch_embeddings = self._process_vision_features(pixel_values, language_embeddings, use_film)\n\n        # Add proprioceptive features if provided\n        use_proprio = proprio_projector is not None and proprio is not None\n        if use_proprio:\n            proprio = torch.Tensor(proprio).to(\n                projected_patch_embeddings.device, dtype=projected_patch_embeddings.dtype\n            )\n            projected_patch_embeddings = self._process_proprio_features(\n                projected_patch_embeddings, proprio, proprio_projector\n            )\n\n        # Use diffusion if provided, otherwise use regression or discrete prediction\n        use_diffusion = noisy_action_projector is not None and hasattr(action_head, \"noise_scheduler\")\n\n        # Calculate number of patches (including proprio token and/or diffusion timestep embedding if present)\n        NUM_PATCHES = self.vision_backbone.get_num_patches() * self.vision_backbone.get_num_images_in_input()\n        if use_proprio:\n            NUM_PATCHES += 1\n        if use_diffusion:\n            NUM_PATCHES += 1\n\n        if use_diffusion:\n            raise ValueError\n            # # Sample random noise with shape equal to output action, used as the starting state for reverse diffusion\n            # noise = torch.randn(\n            #     size=(1, NUM_ACTIONS_CHUNK, ACTION_DIM), device=input_embeddings.device, dtype=input_embeddings.dtype\n            # )\n\n            # # Run diffusion-based prediction\n            # normalized_actions, actions_hidden_states = self._run_diffusion_prediction(\n            #     input_embeddings,\n            #     all_actions_mask,\n            #     noise,\n            #     action_head,\n            #     projected_patch_embeddings,\n            #     labels,\n            #     attention_mask,\n            #     NUM_PATCHES,\n            #     NUM_PROMPT_TOKENS,\n            #     noisy_action_projector,\n            # )\n        else:\n            # Run regression or discrete token-based prediction\n            normalized_actions, reponse_ids = self._verl_discrete_prediction(\n                input_embeddings,\n                all_actions_mask,\n                projected_patch_embeddings,\n                attention_mask,\n                labels,\n                NUM_PATCHES,\n                num_prompt_tokens,\n                action_head,\n                do_sample=do_sample,\n                temperature=temperature,\n            )\n\n        # Unnormalize predicted actions\n        actions = self._unnormalize_actions(normalized_actions, unnorm_key)\n        # verl add!\n        actions = actions.reshape(-1, NUM_ACTIONS_CHUNK, ACTION_DIM)\n        #\n        return actions, reponse_ids\n\n    @staticmethod\n    def _check_unnorm_key(norm_stats: dict[str, dict[str, Any]], unnorm_key: Optional[str]) -> str:\n        \"\"\"Validate and resolve the unnormalization key for action statistics\"\"\"\n        if unnorm_key is None:\n            assert len(norm_stats) == 1, (\n                f\"Your model was trained on more than one dataset, \"\n                f\"please pass a `unnorm_key` from the following options to choose the statistics \"\n                f\"used for un-normalizing actions: {norm_stats.keys()}\"\n            )\n            unnorm_key = next(iter(norm_stats.keys()))\n\n        assert unnorm_key in norm_stats, (\n            f\"The `unnorm_key` you chose is not in the set of available dataset statistics, \"\n            f\"please choose from: {norm_stats.keys()}\"\n        )\n        return unnorm_key\n\n    def get_action_dim(self, unnorm_key: Optional[str] = None) -> int:\n        \"\"\"Get the dimensionality of the policy's action space.\"\"\"\n        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)\n        return len(self.norm_stats[unnorm_key][\"action\"][\"min\"])\n\n    def get_action_stats(self, unnorm_key: Optional[str] = None) -> dict[str, Any]:\n        \"\"\"Get all the logged statistics for the given dataset.\"\"\"\n        unnorm_key = self._check_unnorm_key(self.norm_stats, unnorm_key)\n        return self.norm_stats[unnorm_key][\"action\"]\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/processing_prismatic.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# from https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/\n# form https://huggingface.co/Haozhan72/Openvla-oft-SFT-libero10-trajall/blob/main/\n\n\"\"\"\nprocessing_prismatic.py\n\nHuggingFace-style preprocessor definitions for Prismatic VLMs, inheriting from `ProcessorMixin`. Default configuration\nspecifies `siglip-224px+7b`.\n\"\"\"\n\nfrom typing import Any, ClassVar, Optional\n\nimport timm.data\nimport torch\nimport torchvision.transforms.functional as TVF\nfrom PIL import Image\nfrom torchvision.transforms import CenterCrop, Compose, Normalize, Resize, ToTensor\nfrom transformers import PreTrainedTokenizerBase\nfrom transformers.image_processing_utils import BatchFeature, ImageProcessingMixin\nfrom transformers.processing_utils import ProcessorMixin\nfrom transformers.tokenization_utils import PaddingStrategy, PreTokenizedInput, TextInput, TruncationStrategy\nfrom transformers.utils import TensorType\n\n\n# === Image Processing ===\ndef letterbox_pad_transform(image: Image.Image, padding_fill_value: tuple[int, int, int]) -> Image.Image:\n    \"\"\"Given a PIL.Image, pad to square by adding a symmetric border around the height/width.\"\"\"\n    (w, h), max_wh = image.size, max(image.size)\n    horizontal_pad, vertical_pad = int((max_wh - w) / 2), int((max_wh - h) / 2)\n    padding = (horizontal_pad, vertical_pad, horizontal_pad, vertical_pad)\n\n    return TVF.pad(image, padding, fill=padding_fill_value, padding_mode=\"constant\")\n\n\nclass PrismaticImageProcessor(ImageProcessingMixin):\n    model_input_names: ClassVar[list[str]] = [\"pixel_values\"]\n\n    def __init__(\n        self,\n        use_fused_vision_backbone: bool = False,\n        image_resize_strategy: str = \"letterbox\",\n        input_sizes: Optional[list[tuple[int, int, int]]] = None,\n        interpolations: Optional[list[str]] = None,\n        means: Optional[list[tuple[float, float, float]]] = None,\n        stds: Optional[list[tuple[float, float, float]]] = None,\n        **kwargs: str,\n    ) -> None:\n        \"\"\"\n        Initialize a PrismaticImageProcessor as a wrapper around a torchvision transform; this transform will be\n        created by TIMM, and edited to follow our custom `image_resize_strategy` logic.\n        @param use_fused_vision_backbone: Boolean indicating single or fused (dual) vision backbone\n        @param image_resize_strategy: Prismatic image resize strategy in < resize-naive | resize-crop | letterbox >\n        @param input_size: [TIMM :: `data_cfg`] Input image size as tuple (channels, width, height)\n        @param interpolation: [TIMM :: `data_cfg`] Interpolation as string (default: \"bicubic\")\n        @param mean: [TIMM :: `data_cfg`] Normalization mean as float tuple (or two-tuple if `fused_backbone`)\n        @param std: [TIMM :: `data_cfg`] Normalization std as float tuple (or two-tuple if `fused_backbone`)\n        \"\"\"\n        self.use_fused_vision_backbone = use_fused_vision_backbone\n        self.image_resize_strategy = image_resize_strategy\n\n        # Handle `None` default values\n        input_sizes = [(3, 224, 224)] if input_sizes is None else input_sizes\n        means = [(0.5, 0.5, 0.5)] if means is None else means\n        stds = [(0.5, 0.5, 0.5)] if stds is None else stds\n\n        # TIMM `data_cfg` Parameters\n        self.input_sizes, self.interpolations, self.means, self.stds = input_sizes, interpolations, means, stds\n\n        # Grab torchvision transforms via TIMM =>> need to parse for specific \"functional\" transform values!\n        self.tvf_resize_params, self.tvf_crop_params, self.tvf_normalize_params = [], [], []\n        self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None\n\n        for idx in range(len(input_sizes)):\n            transform = timm.data.create_transform(\n                input_size=self.input_sizes[idx],\n                interpolation=self.interpolations[idx],\n                mean=self.means[idx],\n                std=self.stds[idx],\n                crop_pct=1.0,  # Set to 1.0 to ignore cropping (initial Resize sets `input_size`)\n                crop_mode=\"center\",  # Default crop mode -- no-op when `crop_pct == 1.0`\n                is_training=False,  # No image augmentations when loading the transform!\n            )\n\n            # [Validation] Ensure appropriate transform structure, expected sizes\n            if not (\n                isinstance(transform, Compose)\n                and (len(transform.transforms) == 4)\n                and isinstance(transform.transforms[0], Resize)\n                and isinstance(transform.transforms[1], CenterCrop)\n                and isinstance(transform.transforms[2], ToTensor)\n                and isinstance(transform.transforms[3], Normalize)\n                and (transform.transforms[0].size == self.input_sizes[idx][-1])\n                and (transform.transforms[1].size == self.input_sizes[idx][-2:])\n            ):\n                raise ValueError(f\"Unexpected TIMM image transformation structure/sizes: `{transform}`\")\n\n            # HF Image Processors *must* be JSON-serializable; as such, cannot have torchvision. as an attribute.\n            #   => Instead, we're going to parse the transform and call \"torchvision.transforms.functional\" (`tvf`)\n            resize_t, crop_t, norm_t = transform.transforms[0], transform.transforms[1], transform.transforms[3]\n            self.tvf_resize_params.append(\n                {\n                    \"size\": resize_t.size,\n                    \"interpolation\": TVF.pil_modes_mapping[resize_t.interpolation],\n                    \"max_size\": None,\n                    \"antialias\": True,\n                }\n            )\n            self.tvf_crop_params.append({\"output_size\": crop_t.size})\n            self.tvf_normalize_params.append(\n                {\n                    \"mean\": norm_t.mean.float().numpy().tolist(),\n                    \"std\": norm_t.std.float().numpy().tolist(),\n                    \"inplace\": False,\n                }\n            )\n            self.tvf_do_letterbox, self.tvf_letterbox_fill = False, None\n\n            # Handle Prismatic `image_resize_strategy`\n            if self.image_resize_strategy == \"resize-naive\":\n                self.tvf_resize_params[idx][\"size\"] = (resize_t.size, resize_t.size)\n            elif self.image_resize_strategy == \"letterbox\":\n                self.tvf_do_letterbox, self.tvf_letterbox_fill = True, tuple([int(x * 255) for x in self.means[idx]])\n            elif self.image_resize_strategy == \"resize-crop\":\n                pass\n            else:\n                raise ValueError(f\"Image resize strategy `{self.image_resize_strategy}` is not supported!\")\n\n        # Dispatch **kwargs to super()\n        super().__init__(**kwargs)\n\n    def apply_transform(self, img: Image.Image) -> torch.Tensor:\n        \"\"\"Apply `functional` variant of TIMM's Transform = Compose([Resize -> CenterCrop -> ToTensor -> Normalize])\"\"\"\n        if self.tvf_do_letterbox:\n            img = letterbox_pad_transform(img, self.tvf_letterbox_fill)\n\n        # [Contract] Fused Backbones expect \"channel-stacked\" inputs; we'll unpack on the model side!\n        imgs_t = []\n        for idx in range(len(self.input_sizes)):\n            img_idx = TVF.resize(img, **self.tvf_resize_params[idx])\n            img_idx = TVF.center_crop(img_idx, **self.tvf_crop_params[idx])\n            img_idx_t = TVF.to_tensor(img_idx)\n            img_idx_t = TVF.normalize(img_idx_t, **self.tvf_normalize_params[idx])\n            imgs_t.append(img_idx_t)\n\n        # [Contract] `imgs_t` is a list of Tensors of shape [3, input_size, input_size]; stack along dim = 0\n        img_t = torch.vstack(imgs_t)\n\n        return img_t\n\n    def preprocess(\n        self,\n        images: Image.Image | list[Image.Image],\n        return_tensors: Optional[str | TensorType] = None,\n        **_: str,\n    ) -> BatchFeature:\n        \"\"\"\n        Preprocess an image (or batch of images); note that unlike the `transformers :: BaseImageProcessor` we\n        explicitly only handle PIL.Image.Image instances for simplicity.\n        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.\n        @param return_tensors: BatchFeature default Tensor format (e.g., \"pt\" for torch); if None, returns np.ndarray\n        @return: Instance of `transformers :: BatchFeature` with a single key \"pixel_values\"\n        \"\"\"\n        if not isinstance(images, list):\n            images = [images]\n\n        # Apply `self.img_transform` to each image (will return list of torch.Tensors); stack into \"batched\" Tensor\n        pixel_values = torch.stack([self.apply_transform(img.convert(\"RGB\")) for img in images])\n\n        # Return BatchFeature =>> note that for compatibility, constructor expects Dict[str, np.ndarray], so we convert\n        return BatchFeature(data={\"pixel_values\": pixel_values.float().numpy()}, tensor_type=return_tensors)\n\n    def __call__(self, images: Image.Image | list[Image.Image], **kwargs) -> BatchFeature:\n        return self.preprocess(images, **kwargs)\n\n\n# === PrismaticProcessor =>> Wraps both ImageProcessor and Tokenizer ===\n#   =>> https://github.com/huggingface/transformers/blob/main/src/transformers/models/llava/processing_llava.py\nclass PrismaticProcessor(ProcessorMixin):\n    attributes: ClassVar[list[str]] = [\"image_processor\", \"tokenizer\"]\n    image_processor_class: str = \"AutoImageProcessor\"\n    tokenizer_class: str = \"AutoTokenizer\"\n\n    def __init__(\n        self,\n        image_processor: Optional[ImageProcessingMixin] = None,\n        tokenizer: Optional[PreTrainedTokenizerBase] = None,\n    ) -> None:\n        super().__init__(image_processor, tokenizer)\n\n    def __call__(\n        self,\n        text: TextInput | PreTokenizedInput | list[TextInput] | list[PreTokenizedInput],\n        images: Image.Image | list[Image.Image],\n        padding: bool | str | PaddingStrategy = False,\n        truncation: Optional[bool | str | TruncationStrategy] = None,\n        max_length: Optional[int] = None,\n        return_tensors: Optional[str | TensorType] = TensorType.PYTORCH,\n    ) -> BatchFeature:\n        \"\"\"\n        Preprocess a given (batch) of text/images for a Prismatic VLM; forwards text to the underlying LLM's tokenizer,\n        forwards images to PrismaticImageProcessor.\n        @param text: The (batch) of text to encode; must be a string or list of strings.\n        @param images: A (batch of) PIL.Image.Image instance(s) to preprocess.\n        @param padding: Sequence padding strategy (if multiple specified) in < True = \"longest\" | \"max_length\" | False >\n        @param truncation: Truncation strategy for the output sequences; requires `max_length` to be specified\n        @param max_length: Maximum length (in tokens) to truncate\n        @param return_tensors: Type of return tensors (usually \"pt\" or TensorType.PYTORCH)\n        @return: BatchFeature with keys for `input_ids`, `attention_mask` and `pixel_values`.\n        \"\"\"\n        pixel_values = self.image_processor(images, return_tensors=return_tensors)[\"pixel_values\"]\n        text_inputs = self.tokenizer(\n            text, return_tensors=return_tensors, padding=padding, truncation=truncation, max_length=max_length\n        )\n\n        # [Validate] Need same number of images and text inputs!\n        if pixel_values.shape[0] != text_inputs.input_ids.shape[0]:\n            raise ValueError(\"Batch is malformed; expected same number of images and text inputs!\")\n\n        return BatchFeature(data={**text_inputs, \"pixel_values\": pixel_values})\n\n    # === Tokenizer Dispatch Utilities =>> check `PreTrainedTokenizerBase` for documentation ===\n    def batch_decode(\n        self,\n        sequences: list[int] | list[list[int]] | torch.Tensor | Any,  # `Any` = np.ndarray | tf.Tensor\n        skip_special_tokens: bool = False,\n        clean_up_tokenization_spaces: Optional[bool] = None,\n        **kwargs: str,\n    ) -> list[str]:\n        return self.tokenizer.batch_decode(\n            sequences=sequences,\n            skip_special_tokens=skip_special_tokens,\n            clean_up_tokenization_spaces=clean_up_tokenization_spaces,\n            **kwargs,\n        )\n\n    def decode(\n        self,\n        token_ids: int | list[int] | torch.Tensor | Any,  # `Any` = np.ndarray | tf.Tensor\n        skip_special_tokens: bool = False,\n        clean_up_tokenization_spaces: Optional[bool] = None,\n        **kwargs: str,\n    ) -> str:\n        return self.tokenizer.decode(\n            token_ids=token_ids,\n            skip_special_tokens=skip_special_tokens,\n            clean_up_tokenization_spaces=clean_up_tokenization_spaces,\n            **kwargs,\n        )\n\n    @property\n    def model_input_names(self) -> list[str]:\n        tokenizer_input_names = self.tokenizer.model_input_names\n        image_processor_input_names = self.image_processor.model_input_names\n\n        return list(dict.fromkeys(tokenizer_input_names + image_processor_input_names))\n"
  },
  {
    "path": "verl/experimental/vla/models/openvla_oft/train_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# from https://github.com/PRIME-RL/SimpleVLA-RL/blob/main/verl/utils/vla_utils/openvla_oft/\n\n\"\"\"Utils for training/fine-tuning scripts.\"\"\"\n\nimport torch\n\nfrom .constants import ACTION_DIM, ACTION_TOKEN_BEGIN_IDX, IGNORE_INDEX\n\n\ndef get_current_action_mask(token_ids):\n    # Create a tensor marking positions of IGNORE_INDEX\n    newline_positions = token_ids != IGNORE_INDEX\n\n    # Calculate cumulative sum to identify regions between newlines\n    cumsum = torch.cumsum(newline_positions, dim=1)\n\n    # Create the mask\n    mask = (1 <= cumsum) & (cumsum <= ACTION_DIM)\n\n    # Extract the action part only\n    action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX\n    mask = action_tokens_only_mask * mask\n\n    return mask\n\n\ndef get_next_actions_mask(token_ids):\n    # Create a tensor marking positions of IGNORE_INDEX\n    newline_positions = token_ids != IGNORE_INDEX\n\n    # Calculate cumulative sum to identify regions between newlines\n    cumsum = torch.cumsum(newline_positions, dim=1)\n\n    # Create the mask\n    mask = cumsum > ACTION_DIM\n\n    # Extract the action part only\n    action_tokens_only_mask = token_ids > ACTION_TOKEN_BEGIN_IDX\n    mask = action_tokens_only_mask * mask\n\n    return mask\n\n\ndef compute_token_accuracy(predicted_token_ids, ground_truth_token_ids, mask):\n    correct_preds = (predicted_token_ids == ground_truth_token_ids) & mask\n    accuracy = correct_preds.sum().float() / mask.sum().float()\n    return accuracy\n\n\ndef compute_actions_l1_loss(action_tokenizer, predicted_token_ids, ground_truth_token_ids, mask):\n    pred_continuous_actions = torch.tensor(\n        action_tokenizer.decode_token_ids_to_actions(predicted_token_ids[mask].cpu().numpy())\n    )\n    true_continuous_actions = torch.tensor(\n        action_tokenizer.decode_token_ids_to_actions(ground_truth_token_ids[mask].cpu().numpy())\n    )\n    l1_loss = torch.nn.functional.l1_loss(pred_continuous_actions, true_continuous_actions)\n    return l1_loss\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .configuration_pi0_torch import PI0TorchConfig\nfrom .modeling_pi0_torch import PI0ForActionPrediction\n\n__all__ = [\"PI0TorchConfig\", \"PI0ForActionPrediction\"]\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/configuration_pi0_torch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom transformers import PretrainedConfig\n\n\nclass PI0TorchConfig(PretrainedConfig):\n    model_type = \"pi0_torch\"\n\n    def __init__(self, **kwargs):\n        super().__init__(**kwargs)\n\n        self.state_norm_stats = kwargs.get(\"state_norm_stats\", {})\n        self.action_norm_stats = kwargs.get(\"action_norm_stats\", {})\n        self.pi05_enabled = kwargs.get(\"pi05_enabled\", False)\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/model/modeling_pi0.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 Giga Team. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# from https://github.com/open-gigaai/giga-models\n\n\nimport math\n\nimport torch\nimport torch.nn.functional as F  # noqa: N812\nfrom diffusers.configuration_utils import ConfigMixin, register_to_config\nfrom diffusers.models.modeling_utils import ModelMixin\nfrom torch import Tensor, nn\n\nfrom .paligemma_with_expert import PaliGemmaWithExpertModel\n\n\ndef get_safe_dtype(dtype: torch.dtype, device: str | torch.device) -> torch.dtype:\n    \"\"\"Mps is currently not compatible with float64.\"\"\"\n    if isinstance(device, torch.device):\n        device = device.type\n    if device == \"mps\" and dtype == torch.float64:\n        return torch.float32\n    else:\n        return dtype\n\n\ndef create_sinusoidal_pos_embedding(\n    time: torch.Tensor, dimension: int, min_period: float, max_period: float, device: str | torch.device = \"cpu\"\n) -> Tensor:\n    \"\"\"Computes sine-cosine positional embedding vectors for scalar\n    positions.\"\"\"\n    if dimension % 2 != 0:\n        raise ValueError(f\"dimension ({dimension}) must be divisible by 2\")\n\n    if time.ndim != 1:\n        raise ValueError(\"The time tensor is expected to be of shape `(batch_size, )`.\")\n\n    dtype = get_safe_dtype(torch.float64, device)\n    fraction = torch.linspace(0.0, 1.0, dimension // 2, dtype=dtype, device=device)\n    period = min_period * (max_period / min_period) ** fraction\n\n    # Compute the outer product\n    scaling_factor = 1.0 / period * 2 * math.pi\n    sin_input = scaling_factor[None, :] * time[:, None]\n    pos_emb = torch.cat([torch.sin(sin_input), torch.cos(sin_input)], dim=1)\n    return pos_emb\n\n\ndef make_att_2d_masks(pad_masks: torch.Tensor, att_masks: torch.Tensor) -> torch.Tensor:\n    \"\"\"Copied from big_vision.\n\n    Tokens can attend to valid inputs tokens which have a cumulative mask_ar\n    smaller or equal to theirs. This way `mask_ar` int[B, N] can be used to\n    setup several types of attention, for example:\n\n      [[1 1 1 1 1 1]]: pure causal attention.\n\n      [[0 0 0 1 1 1]]: prefix-lm attention. The first 3 tokens can attend between\n          themselves and the last 3 tokens have a causal attention. The first\n          entry could also be a 1 without changing behaviour.\n\n      [[1 0 1 0 1 0 0 1 0 0]]: causal attention between 4 blocks. Tokens of a\n          block can attend all previous blocks and all tokens on the same block.\n\n    Args:\n      pad_masks: bool[B, N] indicating valid (true) vs. padding (false) tokens.\n      att_masks: int[B, N] defining attention type. A `1` at a position\n                 indicates the start of a new causal block.\n\n    Returns:\n        A 2D boolean attention mask of shape (B, N, N).\n    \"\"\"\n    if att_masks.ndim != 2:\n        raise ValueError(att_masks.ndim)\n    if pad_masks.ndim != 2:\n        raise ValueError(pad_masks.ndim)\n\n    cumsum = torch.cumsum(att_masks, dim=1)\n    att_2d_masks = cumsum[:, None, :] <= cumsum[:, :, None]\n    pad_2d_masks = pad_masks[:, None, :] * pad_masks[:, :, None]\n    att_2d_masks = att_2d_masks & pad_2d_masks\n    return att_2d_masks\n\n\nclass PI0Model(ModelMixin, ConfigMixin):\n    \"\"\"pi0: A Vision-Language-Action Flow Model for General Robot Control.\n\n    [Paper](https://www.physicalintelligence.company/download/pi0.pdf)\n    [Jax code](https://github.com/Physical-Intelligence/openpi)\n\n    ┌──────────────────────────────┐\n    │               actions        │\n    │               ▲              │\n    │              ┌┴─────┐        │\n    │  kv cache    │Gemma │        │\n    │  ┌──────────►│Expert│        │\n    │  │           │      │        │\n    │ ┌┴────────┐  │x 10  │        │\n    │ │         │  └▲──▲──┘        │\n    │ │PaliGemma│   │  │           │\n    │ │         │   │  robot state │\n    │ │         │   noise          │\n    │ └▲──▲─────┘                  │\n    │  │  │                        │\n    │  │  image(s)                 │\n    │  language tokens             │\n    └──────────────────────────────┘\n    \"\"\"\n\n    @register_to_config\n    def __init__(\n        self,\n        max_state_dim: int = 32,\n        max_action_dim: int = 32,\n        proj_width: int = 1024,\n        n_action_steps: int = 50,\n        num_steps: int = 10,\n        use_cache: bool = True,\n        pi05_enabled: bool = False,\n    ):\n        super().__init__()\n\n        # Store the parameters\n        self.max_state_dim = max_state_dim\n        self.max_action_dim = max_action_dim\n        self.proj_width = proj_width\n        self.n_action_steps = n_action_steps\n        self.num_steps = num_steps\n        self.use_cache = use_cache\n        self.pi05_enabled = pi05_enabled\n\n        self.paligemma_with_expert = PaliGemmaWithExpertModel(\n            pi05_enabled=pi05_enabled,\n        )\n\n        # Projections are float32\n        if self.pi05_enabled:\n            self.time_mlp_in = nn.Linear(self.proj_width, self.proj_width, dtype=torch.float32)\n            self.time_mlp_out = nn.Linear(self.proj_width, self.proj_width, dtype=torch.float32)\n        else:\n            self.state_proj = nn.Linear(self.max_state_dim, self.proj_width, dtype=torch.float32)\n            self.action_time_mlp_in = nn.Linear(self.proj_width * 2, self.proj_width, dtype=torch.float32)\n            self.action_time_mlp_out = nn.Linear(self.proj_width, self.proj_width, dtype=torch.float32)\n\n        self.action_in_proj = nn.Linear(self.max_action_dim, self.proj_width, dtype=torch.float32)\n        self.action_out_proj = nn.Linear(self.proj_width, self.max_action_dim, dtype=torch.float32)\n\n    def forward(\n        self,\n        images: list[torch.Tensor],\n        img_masks: list[torch.Tensor],\n        lang_tokens: torch.Tensor,\n        lang_masks: torch.Tensor,\n        state: torch.Tensor,\n        x_t: torch.Tensor,\n        timestep: torch.Tensor,\n    ) -> Tensor:\n        \"\"\"Full forward pass for one diffusion denoising step.\n\n        Args:\n            images: List of image tensors, each shaped (B, C, H, W) after batching.\n            img_masks: List of boolean masks corresponding to images, each (B,).\n            lang_tokens: Language token ids (B, L).\n            lang_masks: Language attention mask (B, L) with True for valid tokens.\n            state: State tensor (B, state_dim) if pi05 is disabled else ignored.\n            x_t: Noisy action tokens (B, n_action_steps, action_dim).\n            timestep: Diffusion timestep as float tensor (B,).\n\n        Returns:\n            Predicted v_t with shape (B, n_action_steps, action_dim).\n        \"\"\"\n        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)\n        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, timestep)\n\n        pad_masks = torch.cat([prefix_pad_masks, suffix_pad_masks], dim=1)\n        att_masks = torch.cat([prefix_att_masks, suffix_att_masks], dim=1)\n\n        att_2d_masks = make_att_2d_masks(pad_masks, att_masks)\n        position_ids = torch.cumsum(pad_masks, dim=1) - 1\n\n        (_, suffix_out), _ = self.paligemma_with_expert.forward(\n            attention_mask=att_2d_masks,\n            position_ids=position_ids,\n            past_key_values=None,\n            inputs_embeds=[prefix_embs, suffix_embs],\n            use_cache=False,\n            fill_kv_cache=False,\n            adarms_cond=[None, adarms_cond],\n        )\n        suffix_out = suffix_out[:, -self.n_action_steps :]\n        # Original openpi code, upcast attention output\n        suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype)\n        v_t = self.action_out_proj(suffix_out)\n        return v_t\n\n    def sample_noise(self, shape: tuple[int, ...], device: torch.device | str) -> torch.Tensor:\n        \"\"\"Generate Gaussian noise for the action trajectory.\n\n        Args:\n            shape: Desired output shape, typically (B, n_action_steps, action_dim).\n            device: Target device string or torch.device.\n\n        Returns:\n            A float32 tensor of standard normal samples with the given shape.\n        \"\"\"\n        noise = torch.normal(\n            mean=0.0,\n            std=1.0,\n            size=shape,\n            dtype=torch.float32,\n            device=device,\n        )\n        return noise\n\n    def embed_prefix(\n        self,\n        images: list[torch.Tensor],\n        img_masks: list[torch.Tensor],\n        lang_tokens: torch.Tensor,\n        lang_masks: torch.Tensor,\n    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        \"\"\"Embed visual and language inputs as the transformer prefix.\n\n        Args:\n            images: List of (B, C, H, W) tensors.\n            img_masks: List of (B,) boolean masks for image presence.\n            lang_tokens: (B, L) token ids.\n            lang_masks: (B, L) boolean mask; True indicates valid tokens.\n\n        Returns:\n            A tuple of (embs, pad_masks, att_masks):\n              - embs: (B, Np, D) concatenated image and language embeddings\n              - pad_masks: (B, Np) valid token mask\n              - att_masks: (B, Np) attention mask scheme selector\n        \"\"\"\n        # Optimize: batch process images and pre-allocate tensors\n        num_images = len(images)\n\n        # Stack images and masks for batch processing\n        images_stacked = torch.stack(images, dim=0)  # (num_images, bsize, ...)\n        img_masks_stacked = torch.stack(img_masks, dim=0)  # (num_images, bsize)\n\n        # Batch embed all images at once\n        # Reshape to (num_images * bsize, ...)\n        orig_shape = images_stacked.shape\n        images_flat = images_stacked.reshape(-1, *orig_shape[2:])\n        img_embs_flat = self.paligemma_with_expert.embed_image(images_flat)\n\n        # Reshape back to (num_images, bsize, num_img_embs, emb_dim)\n        bsize = orig_shape[1]\n        img_embs = img_embs_flat.reshape(num_images, bsize, *img_embs_flat.shape[1:])\n\n        # Normalize image embeddings\n        img_emb_dim = img_embs.shape[-1]\n        num_img_embs = img_embs.shape[2]\n\n        # Expand masks: (num_images, bsize) -> (num_images, bsize, num_img_embs)\n        img_masks_expanded = img_masks_stacked[:, :, None].expand(num_images, bsize, num_img_embs)\n\n        # Reshape to (bsize, num_images * num_img_embs, emb_dim)\n        img_embs_concat = img_embs.transpose(0, 1).reshape(bsize, num_images * num_img_embs, img_emb_dim)\n        img_masks_concat = img_masks_expanded.transpose(0, 1).reshape(bsize, num_images * num_img_embs)\n\n        # Process language embeddings\n        lang_emb = self.paligemma_with_expert.embed_language_tokens(lang_tokens)\n        lang_emb_dim = lang_emb.shape[-1]\n        lang_emb = lang_emb * math.sqrt(lang_emb_dim)\n        lang_emb = lang_emb.to(dtype=img_embs_concat.dtype)\n\n        num_lang_embs = lang_emb.shape[1]\n        total_seq_len = num_images * num_img_embs + num_lang_embs\n\n        # Pre-allocate final tensors\n        embs = torch.empty(\n            bsize, total_seq_len, img_emb_dim, dtype=img_embs_concat.dtype, device=img_embs_concat.device\n        )\n        pad_masks = torch.empty(bsize, total_seq_len, dtype=torch.bool, device=img_embs_concat.device)\n\n        # Fill pre-allocated tensors\n        embs[:, : num_images * num_img_embs] = img_embs_concat\n        embs[:, num_images * num_img_embs :] = lang_emb\n        pad_masks[:, : num_images * num_img_embs] = img_masks_concat\n        pad_masks[:, num_images * num_img_embs :] = lang_masks\n\n        # Create attention masks (all zeros for full attention between image and language)\n        att_masks = torch.zeros(total_seq_len, dtype=torch.bool, device=pad_masks.device)\n        att_masks = att_masks[None, :].expand(bsize, total_seq_len)\n\n        return embs, pad_masks, att_masks\n\n    def embed_suffix(\n        self, state: torch.Tensor, noisy_actions: torch.Tensor, timestep: torch.Tensor\n    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor | None]:\n        \"\"\"Embed state, action and time tokens as the transformer suffix.\n\n        Args:\n            state: (B, state_dim) robot state; ignored when pi05 is enabled.\n            noisy_actions: (B, n_action_steps, action_dim) current x_t.\n            timestep: (B,) diffusion time in [0, 1].\n\n        Returns:\n            (embs, pad_masks, att_masks, adarms_cond) where:\n              - embs: (B, Ns, D) suffix embeddings\n              - pad_masks: (B, Ns) valid mask\n              - att_masks: (B, Ns) causal scheme for suffix\n              - adarms_cond: (B, D) AdaRMS conditioning or None\n        \"\"\"\n        embs = []\n        pad_masks = []\n        att_masks = []\n\n        action_emb = self.action_in_proj(noisy_actions)\n        bsize = action_emb.shape[0]\n        dtype = action_emb.dtype\n        device = action_emb.device\n\n        # Embed state\n        if not self.pi05_enabled:\n            state_emb = self.state_proj(state)\n            embs.append(state_emb[:, None, :])\n\n            state_mask = torch.ones(bsize, 1, dtype=torch.bool, device=device)\n            pad_masks.append(state_mask)\n\n            # Set attention masks so that image and language inputs do not attend to state or actions\n            att_masks += [1]\n\n        # Embed timestep using sine-cosine positional encoding with sensitivity in the range [0, 1]\n        time_emb = create_sinusoidal_pos_embedding(\n            timestep, self.proj_width, min_period=4e-3, max_period=4.0, device=device\n        )\n        time_emb = time_emb.type(dtype=dtype)\n\n        if self.pi05_enabled:\n            # time MLP (for adaRMS)\n            time_emb = self.time_mlp_in(time_emb)\n            time_emb = F.silu(time_emb)\n            time_emb = self.time_mlp_out(time_emb)\n            time_emb = F.silu(time_emb)\n            action_expert_emb = action_emb\n            adarms_cond = time_emb\n        else:\n            # Fuse timestep + action information using an MLP\n            time_emb = time_emb[:, None, :].expand_as(action_emb)\n            action_time_emb = torch.cat([action_emb, time_emb], dim=2)\n\n            action_time_emb = self.action_time_mlp_in(action_time_emb)\n            action_time_emb = F.silu(action_time_emb)  # swish == silu\n            action_time_emb = self.action_time_mlp_out(action_time_emb)\n            action_expert_emb = action_time_emb\n            adarms_cond = None\n\n        # Add to input tokens\n        embs.append(action_expert_emb)\n\n        bsize, action_time_dim = action_expert_emb.shape[:2]\n        action_time_mask = torch.ones(bsize, action_time_dim, dtype=torch.bool, device=device)\n        pad_masks.append(action_time_mask)\n\n        # Set attention masks so that image, language and state inputs do not attend to action tokens\n        att_masks += [1] + ([0] * (self.n_action_steps - 1))\n\n        embs = torch.cat(embs, dim=1)\n        pad_masks = torch.cat(pad_masks, dim=1)\n        att_masks = torch.tensor(att_masks, dtype=embs.dtype, device=embs.device)\n        att_masks = att_masks[None, :].expand(bsize, len(att_masks))\n\n        return embs, pad_masks, att_masks, adarms_cond\n\n    @torch.no_grad()\n    def sample_actions(\n        self,\n        images: list[torch.Tensor],\n        img_masks: list[torch.Tensor],\n        lang_tokens: torch.Tensor,\n        lang_masks: torch.Tensor,\n        state: torch.Tensor,\n        noise: Tensor | None = None,\n    ) -> Tensor:\n        \"\"\"Run the full inference loop to predict an action trajectory.\n\n        Args:\n            images: List of (B, C, H, W) image tensors.\n            img_masks: List of (B,) boolean masks.\n            lang_tokens: (B, L) token ids.\n            lang_masks: (B, L) boolean mask for tokens.\n            state: (B, state_dim) robot state.\n            noise: Optional initial noise; if None, generated internally.\n\n        Returns:\n            Predicted actions with shape (B, n_action_steps, action_dim).\n        \"\"\"\n        bsize = lang_tokens.shape[0]\n        device = lang_tokens.device\n\n        if noise is None:\n            actions_shape = (bsize, self.n_action_steps, self.max_action_dim)\n            noise = self.sample_noise(actions_shape, device)\n\n        prefix_embs, prefix_pad_masks, prefix_att_masks = self.embed_prefix(images, img_masks, lang_tokens, lang_masks)\n        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)\n        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1\n\n        # Compute image and language key value cache\n        _, past_key_values = self.paligemma_with_expert.forward(\n            attention_mask=prefix_att_2d_masks,\n            position_ids=prefix_position_ids,\n            past_key_values=None,\n            inputs_embeds=[prefix_embs, None],\n            use_cache=self.use_cache,\n            fill_kv_cache=True,\n            adarms_cond=[None, None],\n        )\n\n        x_t = noise\n        dt = -1.0 / self.num_steps\n        timesteps = torch.arange(1.0, -dt / 2, dt, dtype=torch.float32, device=device)\n        for timestep in timesteps:\n            v_t = self.denoise_step(\n                state,\n                prefix_pad_masks,\n                past_key_values,\n                x_t,\n                timestep.expand(bsize),\n            )\n            x_t += dt * v_t\n\n        return x_t\n\n    def denoise_step(\n        self,\n        state: torch.Tensor,\n        prefix_pad_masks: torch.Tensor,\n        past_key_values: dict,\n        x_t: torch.Tensor,\n        timestep: torch.Tensor,\n    ) -> torch.Tensor:\n        \"\"\"Apply one denoising step of the noise x_t at a given timestep.\n\n        Args:\n            state: (B, state_dim) robot state.\n            prefix_pad_masks: (B, Np) prefix pad masks computed from embed_prefix.\n            past_key_values: KV cache dict for the prefix (images+language).\n            x_t: (B, n_action_steps, action_dim) current noisy actions.\n            timestep: (B,) current time in [0, 1].\n\n        Returns:\n            v_t prediction with shape (B, n_action_steps, action_dim).\n        \"\"\"\n        suffix_embs, suffix_pad_masks, suffix_att_masks, adarms_cond = self.embed_suffix(state, x_t, timestep)\n\n        suffix_len = suffix_pad_masks.shape[1]\n        batch_size = prefix_pad_masks.shape[0]\n        prefix_len = prefix_pad_masks.shape[1]\n        prefix_pad_2d_masks = prefix_pad_masks[:, None, :].expand(batch_size, suffix_len, prefix_len)\n\n        suffix_att_2d_masks = make_att_2d_masks(suffix_pad_masks, suffix_att_masks)\n\n        full_att_2d_masks = torch.cat([prefix_pad_2d_masks, suffix_att_2d_masks], dim=2)\n\n        prefix_offsets = torch.sum(prefix_pad_masks, dim=-1)[:, None]\n        position_ids = prefix_offsets + torch.cumsum(suffix_pad_masks, dim=1) - 1\n\n        outputs_embeds, _ = self.paligemma_with_expert.forward(\n            attention_mask=full_att_2d_masks,\n            position_ids=position_ids,\n            past_key_values=past_key_values,\n            inputs_embeds=[None, suffix_embs],\n            use_cache=self.use_cache,\n            fill_kv_cache=False,\n            adarms_cond=[None, adarms_cond],\n        )\n        suffix_out = outputs_embeds[1]\n        suffix_out = suffix_out[:, -self.n_action_steps :]\n        suffix_out = suffix_out.to(dtype=self.action_out_proj.weight.dtype)\n        v_t = self.action_out_proj(suffix_out)\n        return v_t\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/model/paligemma_with_expert.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 Giga Team. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# from https://github.com/open-gigaai/giga-models\n\n\nfrom typing import Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom torch import nn\nfrom transformers.activations import ACT2FN\nfrom transformers.modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling\nfrom transformers.models.auto import CONFIG_MAPPING\nfrom transformers.models.siglip.configuration_siglip import SiglipVisionConfig\nfrom transformers.models.siglip.modeling_siglip import (\n    SiglipEncoder,\n    SiglipMultiheadAttentionPoolingHead,\n    SiglipVisionEmbeddings,\n)\nfrom transformers.utils import can_return_tuple\n\nfrom verl.utils.device import get_device_name\n\n\ndef get_transformers_siglip_vision_config() -> SiglipVisionConfig:\n    return CONFIG_MAPPING[\"siglip_vision_model\"](\n        hidden_size=1152,\n        intermediate_size=4304,\n        num_channels=3,\n        num_attention_heads=16,\n        num_hidden_layers=27,\n        num_image_tokens=256,\n        patch_size=14,\n        projection_dim=2048,\n        projector_hidden_act=\"gelu_fast\",\n        torch_dtype=\"float32\",\n        vision_use_head=False,\n    )\n\n\nclass GemmaRMSNorm(nn.Module):\n    def __init__(self, dim: int, eps: float = 1e-6, use_ada_rms_norm: bool = False):\n        super().__init__()\n        self.eps = eps\n        self.use_ada_rms_norm = use_ada_rms_norm\n        if use_ada_rms_norm:\n            self.dense = nn.Linear(dim, dim * 3, bias=True)\n            nn.init.zeros_(self.dense.weight)\n        else:\n            self.weight = nn.Parameter(torch.zeros(dim))\n\n    def _norm(self, x):\n        return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)\n\n    def forward(self, x, cond: torch.Tensor | None = None):\n        normed_inputs = self._norm(x.float())\n\n        if self.use_ada_rms_norm:\n            modulation = self.dense(cond)\n            scale, shift, gate = torch.chunk(modulation.unsqueeze(1), 3, dim=-1)\n            normed_inputs = normed_inputs.float() * (1.0 + scale.float()) + shift.float()\n            return normed_inputs.type_as(x), gate.type_as(x)\n\n        # Llama does x.to(float16) * w whilst Gemma is (x * w).to(float16)\n        # See https://github.com/huggingface/transformers/pull/29402\n        output = normed_inputs * (1.0 + self.weight.float())\n        return output.type_as(x)\n\n    def extra_repr(self):\n        if self.use_ada_rms_norm:\n            return f\"{tuple(self.dense.weight.shape)}, eps={self.eps}, use_ada_rms_norm=True\"\n        else:\n            return f\"{tuple(self.weight.shape)}, eps={self.eps}\"\n\n\nclass SiglipVisionTransformer(nn.Module):\n    def __init__(self, config: SiglipVisionConfig):\n        super().__init__()\n        self.config = config\n        self.config._attn_implementation = \"sdpa\"\n        embed_dim = config.hidden_size\n\n        self.embeddings = SiglipVisionEmbeddings(config)\n        self.encoder = SiglipEncoder(config)\n        self.post_layernorm = nn.LayerNorm(embed_dim, eps=config.layer_norm_eps)\n        self.use_head = True if not hasattr(config, \"vision_use_head\") else config.vision_use_head\n        if self.use_head:\n            self.head = SiglipMultiheadAttentionPoolingHead(config)\n\n    @can_return_tuple\n    # @auto_docstring\n    def forward(\n        self,\n        pixel_values,\n        output_attentions: Optional[bool] = None,\n        output_hidden_states: Optional[bool] = None,\n        interpolate_pos_encoding: Optional[bool] = False,\n    ) -> BaseModelOutputWithPooling:\n        \"\"\"Forward pass of the SigLIP vision encoder.\n\n        Args:\n            pixel_values: Image tensor expected by SigLIP (B, C, H, W).\n            output_attentions: Whether to return attention maps.\n            output_hidden_states: Whether to return hidden states.\n            interpolate_pos_encoding: Enable pos-encoding interpolation for different sizes.\n\n        Returns:\n            BaseModelOutputWithPooling with last_hidden_state and optionally pooled output.\n        \"\"\"\n        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n        output_hidden_states = (\n            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n        )\n\n        hidden_states = self.embeddings(pixel_values, interpolate_pos_encoding=interpolate_pos_encoding)\n        hidden_states = hidden_states.to(dtype=torch.bfloat16)\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            encoder_outputs: BaseModelOutput = self.encoder(\n                inputs_embeds=hidden_states,\n                output_attentions=output_attentions,\n                output_hidden_states=output_hidden_states,\n            )\n            last_hidden_state = encoder_outputs.last_hidden_state\n            last_hidden_state = self.post_layernorm(last_hidden_state)\n\n            pooler_output = self.head(last_hidden_state) if self.use_head else None\n\n            return BaseModelOutputWithPooling(\n                last_hidden_state=last_hidden_state,\n                pooler_output=pooler_output,\n                hidden_states=encoder_outputs.hidden_states,\n                attentions=encoder_outputs.attentions,\n            )\n\n\n# Copied from transformers.models.paligemma.modeling_paligemma.PaliGemmaMultiModalProjector\nclass PaliGemmaMultiModalProjector(nn.Module):\n    def __init__(self, vision_hidden_size: int = 1152, projection_dim: int = 2048):\n        super().__init__()\n        self.linear = nn.Linear(vision_hidden_size, projection_dim, bias=True)\n\n    def forward(self, image_features: torch.Tensor) -> torch.Tensor:\n        \"\"\"Project vision features to the transformer hidden size.\"\"\"\n        hidden_states = self.linear(image_features)\n        return hidden_states\n\n\nclass RoPEEmbedding(nn.Module):\n    \"\"\"Precomputed RoPE embeddings for improved performance.\n\n    This implementation precomputes sin/cos values for a maximum sequence length, avoiding redundant trigonometric\n    calculations during forward passes.\n    \"\"\"\n\n    def __init__(self, dim: int, max_wavelength: int = 10_000, max_seq_len: int = 8192):\n        super().__init__()\n        self.dim = dim\n        self.max_wavelength = max_wavelength\n        self.max_seq_len = max_seq_len\n\n        # Precompute frequency exponents and inverse frequencies\n        d_half = dim // 2\n        freq_exponents = (2.0 / dim) * torch.arange(d_half, dtype=torch.float32)\n        inv_freq = 1.0 / (max_wavelength**freq_exponents)\n\n        # Precompute sin and cos for all positions up to max_seq_len\n        # Shape: [max_seq_len, d_half]\n        positions = torch.arange(max_seq_len, dtype=torch.float32)\n        freqs = torch.outer(positions, inv_freq)  # [max_seq_len, d_half]\n\n        # Precompute sin and cos values\n        # We expand to [max_seq_len, 1, d_half] for broadcasting in forward\n        cos_cached = torch.cos(freqs).unsqueeze(1)  # [max_seq_len, 1, d_half]\n        sin_cached = torch.sin(freqs).unsqueeze(1)  # [max_seq_len, 1, d_half]\n\n        # Register as buffers so they automatically move to the correct device with the model\n        self.register_buffer(\"cos_cached\", cos_cached, persistent=False)\n        self.register_buffer(\"sin_cached\", sin_cached, persistent=False)\n\n    def forward(self, x: torch.Tensor, positions: torch.LongTensor) -> torch.Tensor:\n        \"\"\"Applies RoPE positions [B, L] to x [B, L, H, D].\n\n        Args:\n            x: Input tensor of shape [B, L, H, D]\n            positions: Position indices of shape [B, L]\n\n        Returns:\n            Rotated tensor of shape [B, L, H, D]\n        \"\"\"\n        dtype = x.dtype\n        x = x.to(torch.float32)\n\n        # Index precomputed sin/cos values using positions\n        # positions: [B, L] -> cos/sin: [B, L, 1, d_half]\n        cos = self.cos_cached[positions]  # [B, L, 1, d_half]\n        sin = self.sin_cached[positions]  # [B, L, 1, d_half]\n\n        # Apply rotary embeddings\n        d_half = self.dim // 2\n        x1, x2 = x.split(d_half, dim=-1)  # Each: [B, L, H, d_half]\n\n        # Rotate: out1 = x1 * cos - x2 * sin, out2 = x2 * cos + x1 * sin\n        res = torch.empty_like(x)\n        res[..., :d_half] = x1 * cos - x2 * sin\n        res[..., d_half:] = x2 * cos + x1 * sin\n\n        return res.to(dtype)\n\n\nclass GemmaAttentionWithExpert(nn.Module):\n    def __init__(\n        self,\n        layer_idx: int,\n        # PaliGemma params\n        paligemma_hidden_size: int = 2048,\n        paligemma_num_attention_heads: int = 8,\n        paligemma_num_key_value_heads: int = 1,\n        paligemma_head_dim: int = 256,\n        paligemma_attention_bias: bool = False,\n        # Expert params\n        expert_hidden_size: int = 1024,\n        expert_num_attention_heads: int = 8,\n        expert_num_key_value_heads: int = 1,\n        expert_head_dim: int = 256,\n        expert_attention_bias: bool = False,\n        # RoPE params\n        rope_max_wavelength: int = 10_000,\n        rope_max_seq_len: int = 8192,\n    ):\n        super().__init__()\n        self.layer_idx = layer_idx\n        self.q_proj = nn.ModuleList(\n            [\n                nn.Linear(\n                    paligemma_hidden_size,\n                    paligemma_num_attention_heads * paligemma_head_dim,\n                    bias=paligemma_attention_bias,\n                ),\n                nn.Linear(expert_hidden_size, expert_num_attention_heads * expert_head_dim, bias=expert_attention_bias),\n            ]\n        )\n        self.k_proj = nn.ModuleList(\n            [\n                nn.Linear(\n                    paligemma_hidden_size,\n                    paligemma_num_key_value_heads * paligemma_head_dim,\n                    bias=paligemma_attention_bias,\n                ),\n                nn.Linear(expert_hidden_size, expert_num_key_value_heads * expert_head_dim, bias=expert_attention_bias),\n            ]\n        )\n        self.v_proj = nn.ModuleList(\n            [\n                nn.Linear(\n                    paligemma_hidden_size,\n                    paligemma_num_key_value_heads * paligemma_head_dim,\n                    bias=paligemma_attention_bias,\n                ),\n                nn.Linear(expert_hidden_size, expert_num_key_value_heads * expert_head_dim, bias=expert_attention_bias),\n            ]\n        )\n        self.o_proj = nn.ModuleList(\n            [\n                nn.Linear(\n                    paligemma_num_attention_heads * paligemma_head_dim,\n                    paligemma_hidden_size,\n                    bias=paligemma_attention_bias,\n                ),\n                nn.Linear(expert_num_attention_heads * expert_head_dim, expert_hidden_size, bias=expert_attention_bias),\n            ]\n        )\n\n        self.paligemma_num_attention_heads = paligemma_num_attention_heads\n        self.paligemma_num_key_value_heads = paligemma_num_key_value_heads\n        self.paligemma_head_dim = paligemma_head_dim\n        self.expert_num_attention_heads = expert_num_attention_heads\n        self.expert_num_key_value_heads = expert_num_key_value_heads\n        self.expert_head_dim = expert_head_dim\n\n        assert paligemma_head_dim == expert_head_dim\n        assert paligemma_num_attention_heads == expert_num_attention_heads\n        assert paligemma_num_key_value_heads == expert_num_key_value_heads\n        self.rope_embedding = RoPEEmbedding(\n            dim=paligemma_head_dim, max_wavelength=rope_max_wavelength, max_seq_len=rope_max_seq_len\n        )\n\n    def forward(\n        self,\n        inputs_embeds: list[Optional[torch.Tensor]],\n        position_ids: torch.LongTensor,\n        attention_mask: torch.Tensor,\n        use_cache: bool,\n        past_key_values: Optional[dict] = None,\n        fill_kv_cache: bool = False,\n    ) -> list[Optional[torch.Tensor]]:\n        \"\"\"Multi-source attention over PaliGemma and Expert streams.\n\n        Args:\n            inputs_embeds: [paligemma_embeds, expert_embeds]. Each is (B, L, D) or None.\n            position_ids: (B, L) rotary positions.\n            attention_mask: (B, L, L) attention mask.\n            use_cache: Whether to use KV cache.\n            past_key_values: Optional cache dict per layer.\n            fill_kv_cache: If True, fill cache; otherwise, append to it.\n\n        Returns:\n            List[Optional[Tensor]]: outputs per stream aligned to inputs order.\n        \"\"\"\n        query_states = []\n        key_states = []\n        value_states = []\n\n        if inputs_embeds[0] is not None:\n            # PaliGemma\n            hidden_states = inputs_embeds[0]\n            input_shape = hidden_states.shape[:-1]\n            hidden_shape = (*input_shape, -1, self.paligemma_head_dim)\n            query_states.append(self.q_proj[0](hidden_states).view(hidden_shape))\n            key_states.append(self.k_proj[0](hidden_states).view(hidden_shape))\n            value_states.append(self.v_proj[0](hidden_states).view(hidden_shape))\n\n        if inputs_embeds[1] is not None:\n            # Expert\n            hidden_states = inputs_embeds[1]\n            input_shape = hidden_states.shape[:-1]\n            hidden_shape = (*input_shape, -1, self.expert_head_dim)\n            query_states.append(self.q_proj[1](hidden_states).view(hidden_shape))\n            key_states.append(self.k_proj[1](hidden_states).view(hidden_shape))\n            value_states.append(self.v_proj[1](hidden_states).view(hidden_shape))\n\n        query_states = torch.cat(query_states, dim=1)\n        key_states = torch.cat(key_states, dim=1)\n        value_states = torch.cat(value_states, dim=1)\n\n        query_states = self.rope_embedding(query_states, position_ids)\n        key_states = self.rope_embedding(key_states, position_ids)\n\n        if use_cache:\n            if fill_kv_cache:\n                past_key_values[self.layer_idx] = {\n                    \"key_states\": key_states,\n                    \"value_states\": value_states,\n                }\n            else:\n                key_states = torch.cat([past_key_values[self.layer_idx][\"key_states\"], key_states], dim=1)\n                value_states = torch.cat([past_key_values[self.layer_idx][\"value_states\"], value_states], dim=1)\n\n        num_att_heads = self.paligemma_num_attention_heads  # Assume same for both\n        num_key_value_heads = self.paligemma_num_key_value_heads\n        head_dim = self.paligemma_head_dim\n        batch_size = query_states.shape[0]\n        query_states = query_states.transpose(1, 2)\n        key_states = key_states.transpose(1, 2)\n        value_states = value_states.transpose(1, 2)\n\n        if num_key_value_heads != num_att_heads:\n            # key_states: (B, num_kv_heads, L, D) -> (B, num_att_heads, L, D)\n            key_states = torch.repeat_interleave(key_states, num_att_heads // num_key_value_heads, dim=1)\n            value_states = torch.repeat_interleave(value_states, num_att_heads // num_key_value_heads, dim=1)\n\n        att_output = F.scaled_dot_product_attention(\n            query_states,\n            key_states,\n            value_states,\n            attn_mask=attention_mask[:, None, :, :],\n            is_causal=False,\n        )\n        att_output = att_output.permute(0, 2, 1, 3)\n        att_output = att_output.reshape(batch_size, -1, num_att_heads * head_dim)\n\n        outputs_embeds = []\n        start = 0\n        if inputs_embeds[0] is not None:\n            hidden_states = inputs_embeds[0]\n            end = start + hidden_states.shape[1]\n            if att_output.dtype != self.o_proj[0].weight.dtype:\n                att_output_i = att_output[:, start:end].to(self.o_proj[0].weight.dtype)\n            else:\n                att_output_i = att_output[:, start:end]\n            out_emb = self.o_proj[0](att_output_i)\n            outputs_embeds.append(out_emb)\n            start = end\n        else:\n            outputs_embeds.append(None)\n\n        if inputs_embeds[1] is not None:\n            hidden_states = inputs_embeds[1]\n            end = start + hidden_states.shape[1]\n            if att_output.dtype != self.o_proj[1].weight.dtype:\n                att_output_i = att_output[:, start:end].to(self.o_proj[1].weight.dtype)\n            else:\n                att_output_i = att_output[:, start:end]\n            out_emb = self.o_proj[1](att_output_i)\n            outputs_embeds.append(out_emb)\n        else:\n            outputs_embeds.append(None)\n\n        return outputs_embeds\n\n\nclass GemmaMLP(nn.Module):\n    def __init__(self, hidden_size: int = 1024, intermediate_size: int = 4096, hidden_act: str = \"gelu_pytorch_tanh\"):\n        super().__init__()\n        self.hidden_size = hidden_size\n        self.intermediate_size = intermediate_size\n        self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)\n        self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)\n        self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)\n        self.act_fn = ACT2FN[hidden_act]\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Gated MLP block used in both streams.\"\"\"\n        down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))\n        return down_proj\n\n\nclass GemmaDecoderLayerWithExpert(nn.Module):\n    def __init__(\n        self,\n        layer_idx: int,\n        pi05_enabled: bool,\n        # PaliGemma params\n        paligemma_hidden_size: int = 2048,\n        paligemma_num_attention_heads: int = 8,\n        paligemma_num_key_value_heads: int = 1,\n        paligemma_head_dim: int = 256,\n        paligemma_attention_bias: bool = False,\n        paligemma_intermediate_size: int = 16384,\n        paligemma_hidden_act: str = \"gelu_pytorch_tanh\",\n        paligemma_rms_norm_eps: float = 1e-6,\n        # Expert params\n        expert_hidden_size: int = 1024,\n        expert_num_attention_heads: int = 8,\n        expert_num_key_value_heads: int = 1,\n        expert_head_dim: int = 256,\n        expert_attention_bias: bool = False,\n        expert_intermediate_size: int = 4096,\n        expert_hidden_act: str = \"gelu_pytorch_tanh\",\n        expert_rms_norm_eps: float = 1e-6,\n        # RoPE params\n        rope_max_wavelength: int = 10_000,\n        rope_max_seq_len: int = 8192,\n    ):\n        super().__init__()\n        self.self_attn = GemmaAttentionWithExpert(\n            layer_idx,\n            paligemma_hidden_size,\n            paligemma_num_attention_heads,\n            paligemma_num_key_value_heads,\n            paligemma_head_dim,\n            paligemma_attention_bias,\n            expert_hidden_size,\n            expert_num_attention_heads,\n            expert_num_key_value_heads,\n            expert_head_dim,\n            expert_attention_bias,\n            rope_max_wavelength,\n            rope_max_seq_len,\n        )\n\n        self.mlps = nn.ModuleList(\n            [\n                GemmaMLP(paligemma_hidden_size, paligemma_intermediate_size, paligemma_hidden_act),\n                GemmaMLP(expert_hidden_size, expert_intermediate_size, expert_hidden_act),\n            ]\n        )\n\n        self.input_layernorms = nn.ModuleList(\n            [\n                GemmaRMSNorm(paligemma_hidden_size, eps=paligemma_rms_norm_eps),\n                GemmaRMSNorm(expert_hidden_size, eps=expert_rms_norm_eps, use_ada_rms_norm=pi05_enabled),\n            ]\n        )\n        self.post_attention_layernorms = nn.ModuleList(\n            [\n                GemmaRMSNorm(paligemma_hidden_size, eps=paligemma_rms_norm_eps),\n                GemmaRMSNorm(expert_hidden_size, eps=expert_rms_norm_eps, use_ada_rms_norm=pi05_enabled),\n            ]\n        )\n\n        self.pi05_enabled = pi05_enabled\n\n    def gated_residual(self, x, y, gate):\n        if x is None or y is None:\n            return None\n        if gate is None:\n            return x + y\n        return x + y * gate\n\n    def forward(\n        self,\n        inputs_embeds: list[Optional[torch.Tensor]],\n        adarms_cond: list[Optional[torch.Tensor]],\n        position_ids: torch.LongTensor,\n        attention_mask: torch.Tensor,\n        use_cache: bool,\n        past_key_values: Optional[dict] = None,\n        fill_kv_cache: bool = False,\n    ) -> list[Optional[torch.Tensor]]:\n        \"\"\"Decoder layer with dual-stream attention and optional AdaRMS\n        modulation.\n\n        Args:\n            inputs_embeds: [paligemma, expert] embeds.\n            adarms_cond: Optional conditioning vectors for AdaRMS.\n            position_ids: (B, L) positions for RoPE.\n            attention_mask: (B, L, L) attention mask.\n            use_cache: Whether to use KV cache.\n            past_key_values: Optional cache dict.\n            fill_kv_cache: Whether to fill or reuse KV cache.\n\n        Returns:\n            List[Optional[Tensor]]: Updated hidden states per stream.\n        \"\"\"\n        residuals = list(inputs_embeds)\n        normed_embeds = []\n        attn_gates = []\n\n        for i, hidden_states in enumerate(inputs_embeds):\n            if hidden_states is not None:\n                if self.pi05_enabled and adarms_cond[i] is not None:\n                    normed_h, attn_gate = self.input_layernorms[i](hidden_states, adarms_cond[i])\n                    normed_embeds.append(normed_h)\n                    attn_gates.append(attn_gate)\n                else:\n                    normed_embeds.append(self.input_layernorms[i](hidden_states))\n                    attn_gates.append(None)\n            else:\n                normed_embeds.append(None)\n                attn_gates.append(None)\n\n        attn_outputs = self.self_attn(\n            normed_embeds, position_ids, attention_mask, use_cache, past_key_values, fill_kv_cache\n        )\n\n        after_attn_embeds = []\n        for i, (residual, attn_output, attn_gate) in enumerate(zip(residuals, attn_outputs, attn_gates, strict=False)):\n            if residual is not None:\n                after_attn_embeds.append(self.gated_residual(residual, attn_output, attn_gate))\n            else:\n                after_attn_embeds.append(None)\n\n        outputs = []\n        for i, hidden_states in enumerate(after_attn_embeds):\n            if hidden_states is not None:\n                residual = hidden_states\n                if self.pi05_enabled and adarms_cond[i] is not None:\n                    normed_h, mlp_gate = self.post_attention_layernorms[i](hidden_states, adarms_cond[i])\n                else:\n                    normed_h = self.post_attention_layernorms[i](hidden_states)\n                    mlp_gate = None\n\n                mlp_out = self.mlps[i](normed_h)\n                outputs.append(self.gated_residual(residual, mlp_out, mlp_gate))\n            else:\n                outputs.append(None)\n\n        return outputs, past_key_values\n\n\nclass PaliGemmaWithExpertModel(nn.Module):\n    def __init__(\n        self,\n        pi05_enabled: bool = False,\n        # Paligemma params\n        paligemma_vocab_size: int = 257152,\n        paligemma_pad_token_id: int = 0,\n        paligemma_num_hidden_layers: int = 18,\n        paligemma_hidden_size: int = 2048,\n        paligemma_num_attention_heads: int = 8,\n        paligemma_num_key_value_heads: int = 1,\n        paligemma_attention_bias: bool = False,\n        paligemma_intermediate_size: int = 16384,\n        paligemma_hidden_act: str = \"gelu_pytorch_tanh\",\n        paligemma_rms_norm_eps: float = 1e-6,\n        # Expert params\n        expert_hidden_size: int = 1024,\n        expert_num_attention_heads: int = 8,\n        expert_num_key_value_heads: int = 1,\n        expert_head_dim: int = 256,\n        expert_attention_bias: bool = False,\n        expert_intermediate_size: int = 4096,\n        expert_hidden_act: str = \"gelu_pytorch_tanh\",\n        expert_rms_norm_eps: float = 1e-6,\n        # RoPE params\n        rope_max_wavelength: int = 10_000,\n        rope_max_seq_len: int = 8192,\n    ):\n        super().__init__()\n        self.pi05_enabled = pi05_enabled\n\n        siglip_vision_config = get_transformers_siglip_vision_config()\n\n        # Vision and projection\n        self.vision_tower = SiglipVisionTransformer(siglip_vision_config)\n        self.multi_modal_projector = PaliGemmaMultiModalProjector(\n            vision_hidden_size=siglip_vision_config.hidden_size, projection_dim=siglip_vision_config.projection_dim\n        )\n        self.paligemma_hidden_size = paligemma_hidden_size\n\n        # Language embed\n        self.embed_tokens = nn.Embedding(paligemma_vocab_size, paligemma_hidden_size, paligemma_pad_token_id)\n\n        # Decoder layers\n        self.layers = nn.ModuleList(\n            [\n                GemmaDecoderLayerWithExpert(\n                    layer_idx=i,\n                    pi05_enabled=pi05_enabled,\n                    paligemma_hidden_size=paligemma_hidden_size,\n                    paligemma_num_attention_heads=paligemma_num_attention_heads,\n                    paligemma_num_key_value_heads=paligemma_num_key_value_heads,\n                    paligemma_head_dim=paligemma_hidden_size // paligemma_num_attention_heads,\n                    paligemma_attention_bias=paligemma_attention_bias,  # gemma default\n                    paligemma_intermediate_size=paligemma_intermediate_size,\n                    paligemma_hidden_act=paligemma_hidden_act,\n                    paligemma_rms_norm_eps=paligemma_rms_norm_eps,  # gemma default\n                    expert_hidden_size=expert_hidden_size,\n                    expert_num_attention_heads=expert_num_attention_heads,\n                    expert_num_key_value_heads=expert_num_key_value_heads,\n                    expert_head_dim=expert_head_dim,\n                    expert_attention_bias=expert_attention_bias,\n                    expert_intermediate_size=expert_intermediate_size,\n                    expert_hidden_act=expert_hidden_act,\n                    expert_rms_norm_eps=expert_rms_norm_eps,\n                    rope_max_wavelength=rope_max_wavelength,\n                    rope_max_seq_len=rope_max_seq_len,\n                )\n                for i in range(paligemma_num_hidden_layers)\n            ]\n        )\n\n        # Final norms\n        self.norms = nn.ModuleList(\n            [\n                GemmaRMSNorm(paligemma_hidden_size, eps=1e-6),\n                GemmaRMSNorm(expert_hidden_size, eps=expert_rms_norm_eps, use_ada_rms_norm=pi05_enabled),\n            ]\n        )\n\n    def embed_image(self, image: torch.Tensor) -> torch.Tensor:\n        \"\"\"Encode images with SigLIP and project to hidden size.\"\"\"\n        image_outputs = self.vision_tower(image)\n        selected_image_feature = image_outputs.last_hidden_state\n        image_features = self.multi_modal_projector(selected_image_feature)\n        return image_features\n\n    def embed_language_tokens(self, tokens: torch.Tensor) -> torch.Tensor:\n        \"\"\"Embed token ids into continuous vectors.\"\"\"\n        return self.embed_tokens(tokens)\n\n    def forward(\n        self,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n        past_key_values: Optional[dict] = None,\n        inputs_embeds: list[torch.FloatTensor] = None,\n        use_cache: Optional[bool] = None,\n        fill_kv_cache: Optional[bool] = None,\n        adarms_cond: list[torch.FloatTensor] = None,\n    ) -> tuple[list[Optional[torch.Tensor]], dict]:\n        \"\"\"Run the stacked dual-stream decoder with optional caching and\n        AdaRMS.\n\n        Args:\n            attention_mask: (B, L, L) attention mask for both streams.\n            position_ids: (B, L) RoPE positions.\n            past_key_values: Optional KV cache dict to reuse.\n            inputs_embeds: [paligemma_embeds, expert_embeds].\n            use_cache: Whether to use KV cache.\n            fill_kv_cache: If True, populate cache from inputs.\n            adarms_cond: Optional per-stream modulation vectors for AdaRMS.\n\n        Returns:\n            (outputs_embeds, past_key_values): outputs per stream and the KV cache.\n        \"\"\"\n        inputs_embeds = [\n            input_embed.to(dtype=torch.bfloat16) if input_embed is not None else None for input_embed in inputs_embeds\n        ]\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            if use_cache and past_key_values is None:\n                past_key_values = {}\n\n            hidden_states_list = inputs_embeds\n            for layer in self.layers:\n                # FSDP will make a copy of the \"past_key_values\" dictionary, which needs to be reassigned.\n                hidden_states_list, past_key_values = layer(\n                    hidden_states_list,\n                    adarms_cond=adarms_cond,\n                    position_ids=position_ids,\n                    attention_mask=attention_mask,\n                    use_cache=use_cache,\n                    past_key_values=past_key_values,\n                    fill_kv_cache=fill_kv_cache,\n                )\n\n            outputs_embeds = []\n            for i, hidden_states in enumerate(hidden_states_list):\n                if hidden_states is not None:\n                    if self.pi05_enabled and adarms_cond[i] is not None:\n                        out_emb, _ = self.norms[i](hidden_states, adarms_cond[i])\n                    else:\n                        out_emb = self.norms[i](hidden_states)\n                    outputs_embeds.append(out_emb)\n                else:\n                    outputs_embeds.append(None)\n\n            return outputs_embeds, past_key_values\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/modeling_pi0_torch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport math\nfrom typing import Literal\n\nimport torch\nfrom onnx_ir import Tensor\nfrom torch import nn\nfrom torch.distributed.fsdp import register_fsdp_forward_method\nfrom transformers import PreTrainedModel\nfrom typing_extensions import override\n\nfrom verl.protocol import DataProto\nfrom verl.utils.device import get_device_name\n\nfrom ...sac.base import SupportSACTraining\nfrom ..modules.mlp import MLP\nfrom .configuration_pi0_torch import PI0TorchConfig\nfrom .model.modeling_pi0 import PI0Model, make_att_2d_masks\nfrom .pi0_utils import (\n    ImageTransform,\n    Normalize,\n    PromptTokenizerTransform,\n    Unnormalize,\n)\nfrom .policy.base import Pi0Output\n\n\ndef beta_schedule(step, beta0, beta_min, T):\n    progress = min(step / T, 1.0)\n    beta = beta_min + (beta0 - beta_min) * 0.5 * (1 + math.cos(math.pi * progress))\n    return beta\n\n\nclass PI0ForActionPrediction(PreTrainedModel, SupportSACTraining):\n    config_class = PI0TorchConfig\n    base_model_prefix = \"pi0_torch\"\n\n    def __init__(self, config: PI0TorchConfig):\n        super().__init__(config)\n        self.model: PI0Model = None\n        self.state_norm_stats = config.state_norm_stats\n        self.action_norm_stats = config.action_norm_stats\n        self.pi05_enabled = config.pi05_enabled\n\n        assert self.state_norm_stats, \"state_norm_stats must be provided in PI0TorchConfig\"\n        assert self.action_norm_stats, \"action_norm_stats must be provided in PI0TorchConfig\"\n        assert isinstance(self.pi05_enabled, bool), \"pi05_enabled must be provided in PI0TorchConfig\"\n\n        # Input transforms\n        self.state_normalize_transform = Normalize(self.state_norm_stats, use_quantiles=self.pi05_enabled)\n        self.action_normalize_transform = Normalize(self.action_norm_stats, use_quantiles=self.pi05_enabled)\n        self.image_transform = ImageTransform(resize_imgs_with_padding=(224, 224), enable_image_aug=False)\n        max_length = 200 if self.pi05_enabled else 48\n        self.prompt_tokenizer_transform = PromptTokenizerTransform(max_length=max_length, discrete_state_input=False)\n\n        # Output transforms\n        self.state_unnormalize_transform = Unnormalize(self.state_norm_stats, use_quantiles=self.pi05_enabled)\n        self.action_unnormalize_transform = Unnormalize(self.action_norm_stats, use_quantiles=self.pi05_enabled)\n\n        # Flow SDE parameters\n        self._to(get_device_name())\n        self.flow_sde_enable = bool(getattr(config, \"flow_sde_enable\", True))\n        self.flow_sde_noise_level = float(getattr(config, \"flow_sde_noise_level\", 0.5))\n        self.flow_sde_rollout_noise_scale = float(getattr(config, \"flow_sde_rollout_noise_scale\", 1.0))\n        self.flow_sde_train_noise_scale = float(getattr(config, \"flow_sde_train_noise_scale\", 1.0))\n        self.flow_sde_initial_beta = float(getattr(config, \"flow_sde_initial_beta\", 1.0))\n        self.flow_sde_beta_min = float(getattr(config, \"flow_sde_beta_min\", 0.02))\n        self.flow_sde_beta_schedule_T = int(getattr(config, \"flow_sde_beta_schedule_T\", 2000))\n        self.register_buffer(\"flow_sde_step\", torch.zeros((), dtype=torch.long))\n\n        ##### SAC Algorithm Support #####\n        if getattr(self.config, \"sac_enable\", False):\n            head_num = int(getattr(self.config, \"critic_head_num\", 10))\n            attn_heads = int(getattr(self.config, \"critic_prefix_attn_heads\", 8))\n\n            self.critic_state_token = nn.Parameter(torch.zeros(1, 1, 2048))\n            self.target_state_token = nn.Parameter(torch.zeros(1, 1, 2048))\n            nn.init.normal_(self.critic_state_token, mean=0.0, std=0.02)\n            self.target_state_token.data.copy_(self.critic_state_token.data)\n\n            self.critic_prefix_cross_attn = nn.MultiheadAttention(\n                embed_dim=2048,\n                num_heads=attn_heads,\n                batch_first=True,\n            )\n            self.target_prefix_cross_attn = nn.MultiheadAttention(\n                embed_dim=2048,\n                num_heads=attn_heads,\n                batch_first=True,\n            )\n\n            self.critic_heads = nn.ModuleList(\n                [\n                    MLP(\n                        input_dim=2150,\n                        hidden_dims=[2048, 1024, 256],\n                        output_dim=1,\n                        activation=\"relu\",\n                        init_method=\"kaiming\",\n                    )\n                    for _ in range(head_num)\n                ]\n            )\n\n            self.target_network_heads = nn.ModuleList(\n                [\n                    MLP(\n                        input_dim=2150,\n                        hidden_dims=[2048, 1024, 256],\n                        output_dim=1,\n                        activation=\"relu\",\n                        init_method=\"kaiming\",\n                    )\n                    for _ in range(head_num)\n                ]\n            )\n\n            self.target_network_heads.load_state_dict(self.critic_heads.state_dict())\n            self.target_prefix_cross_attn.load_state_dict(self.critic_prefix_cross_attn.state_dict())\n\n    def _to(self, device: torch.device | str):\n        self.state_normalize_transform.to(device)\n        self.state_unnormalize_transform.to(device)\n        self.action_normalize_transform.to(device)\n        self.action_unnormalize_transform.to(device)\n        return self\n\n    def forward(\n        self,\n        images: list[torch.Tensor],\n        img_masks: list[torch.Tensor],\n        lang_tokens: torch.Tensor,\n        lang_masks: torch.Tensor,\n        state: torch.Tensor,\n        x_t: torch.Tensor,\n        timestep: torch.Tensor,\n    ) -> Tensor:\n        \"\"\"Full forward pass for one diffusion denoising step.\n\n        Args:\n            images: List of image tensors, each shaped (B, C, H, W) after batching.\n            img_masks: List of boolean masks corresponding to images, each (B,).\n            lang_tokens: Language token ids (B, L).\n            lang_masks: Language attention mask (B, L) with True for valid tokens.\n            state: State tensor (B, state_dim) if pi05 is disabled else ignored.\n            x_t: Noisy action tokens (B, n_action_steps, action_dim).\n            timestep: Diffusion timestep as float tensor (B,).\n\n        Returns:\n            Predicted v_t with shape (B, n_action_steps, action_dim).\n        \"\"\"\n\n        if self.model is None:\n            raise RuntimeError(\"PI0ForActionPrediction.model is not initialized. Did from_pretrained() run?\")\n\n        return self.model(\n            images,\n            img_masks,\n            lang_tokens,\n            lang_masks,\n            state,\n            x_t,\n            timestep,\n        )\n\n    @torch.no_grad()\n    def sample_actions(\n        self,\n        env_obs: DataProto,\n        tokenizer,\n        validate: bool = False,\n    ) -> tuple[Pi0Output, dict, dict]:\n        \"\"\"Run one forward pass from raw inputs to final action sequence.\n\n        Args:\n            env_obs: The environment observations as DataProto.\n            tokenizer: The tokenizer used for prompt tokenization.\n\n        Returns:\n            A tuple of (pi0_output, s, a):\n                - pi0_output: The Pi0Output containing the predicted actions.\n                - s: Dictionary of tensors representing the states, with keys\n                    - \"images\": torch.Tensor of shape (B, n_images, C, H, W)\n                    - \"image_masks\": torch.Tensor of shape (B, n_images)\n                    - \"lang_tokens\": torch.Tensor of shape (B, L)\n                    - \"lang_masks\": torch.Tensor of shape (B, L)\n                    - \"states\": torch.Tensor of shape (B, state_dim)\n                - a: Dictionary of tensors representing actions, with key:\n                    - \"full_action\": torch.Tensor of shape (B, action_steps, action_dim)\n        \"\"\"\n\n        from .policy.libero_policy import LiberoPi0Input\n\n        pi0_input = LiberoPi0Input.from_env_obs(env_obs)\n\n        # Input transforms\n        state = self.state_normalize_transform(pi0_input.state)\n        images, _ = self.image_transform.call_batch(pi0_input.images)\n        lang_tokens, lang_masks = self.prompt_tokenizer_transform.call_batch(\n            {\"task\": pi0_input.task, \"observation.state\": state}, tokenizer\n        )\n\n        if self.flow_sde_enable and not validate:\n            prefix_features = self.model.embed_prefix(\n                images=images,\n                img_masks=pi0_input.img_masks,\n                lang_tokens=lang_tokens,\n                lang_masks=lang_masks,\n            )\n            pred_action, rollout_log_probs = self._sample_actions_flow_sde(\n                state_features=(prefix_features, state),\n                noise_scale=self.flow_sde_rollout_noise_scale,\n                requires_grad=False,\n                return_log_prob=True,\n            )\n        else:\n            pred_action = self.model.sample_actions(images, pi0_input.img_masks, lang_tokens, lang_masks, state=state)\n            rollout_log_probs = torch.zeros(pred_action.shape[0], device=pred_action.device, dtype=torch.float32)\n\n        # Output transforms\n        from .policy.libero_policy import LiberoPi0Output\n\n        pi0_output = LiberoPi0Output.from_model_output({\"full_action\": self.action_unnormalize_transform(pred_action)})\n        s = {\n            \"states\": state,\n            \"images\": torch.stack(images, dim=1),\n            \"image_masks\": torch.stack(pi0_input.img_masks, dim=1),\n            \"lang_tokens\": lang_tokens,\n            \"lang_masks\": lang_masks,\n        }\n        a = {\n            \"full_action\": pred_action,\n            \"log_probs\": rollout_log_probs,\n        }\n\n        return pi0_output, s, a\n\n    @classmethod\n    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):\n        config = kwargs.pop(\"config\", None)\n\n        if config is None:\n            config = PI0TorchConfig.from_pretrained(pretrained_model_name_or_path)\n\n        policy = cls(config)\n        policy.model = PI0Model.from_pretrained(pretrained_model_name_or_path)\n        return policy\n\n    # def load_state_dict(self, state_dict, strict: bool = True, assign: bool = False):\n    #     filtered_state_dict = {\n    #         key: value\n    #         for key, value in state_dict.items()\n    #         if key.startswith(\"model.\")\n    #     }\n    #     return super().load_state_dict(filtered_state_dict, strict=False, assign=assign)\n\n    def freeze_vision_tower(self) -> None:\n        \"\"\"Freeze the vision tower parameters.\"\"\"\n\n        if self.model is None:\n            raise RuntimeError(\"PI0ForActionPrediction.model is not initialized. Did from_pretrained() run?\")\n        vision_tower = self.model.paligemma_with_expert.vision_tower\n        vision_tower.requires_grad_(False)\n        vision_tower.eval()\n\n    def bc_loss(\n        self,\n        state_features: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],\n        actions: dict[str, torch.Tensor],\n        valids: torch.Tensor,\n    ) -> torch.Tensor:\n        \"\"\"Calculate the BC loss for the actor.\"\"\"\n\n        prefix_features, states = state_features\n        _, prefix_pad_masks, _ = prefix_features\n        action_tensor = actions[\"full_action\"]\n\n        batch_size = action_tensor.shape[0]\n        device = action_tensor.device\n\n        noise = self.model.sample_noise(action_tensor.shape, device=device)\n        gamma1 = torch.empty((batch_size,), device=device).uniform_(0, 1).pow(1 / 1.5)\n        gamma2 = torch.empty((batch_size,), device=device).uniform_(0, 1).pow(1 / 1.0)\n        time = (gamma1 / (gamma1 + gamma2)) * 0.999 + 0.001\n        time = time.to(dtype=torch.float32, device=device)\n\n        time_expanded = time[:, None, None]\n        x_t = time_expanded * noise + (1.0 - time_expanded) * action_tensor\n        u_t = noise - action_tensor\n\n        past_key_values = self._build_kv_cache_from_prefix(prefix_features)\n        model_pred = self.model.denoise_step(\n            states,\n            prefix_pad_masks,\n            past_key_values,\n            x_t,\n            time,\n        )\n\n        loss = torch.nn.functional.mse_loss(u_t, model_pred, reduction=\"none\").mean(dim=-1).mean(dim=-1)\n        valid_f = valids.float().to(loss.device)\n        return (loss * valid_f).sum() / valid_f.sum().clamp_min(1.0)\n\n    # --- SAC Algorithm Support ---\n\n    def _multi_heads_value(\n        self, value_heads: nn.ModuleList, input_tensor: torch.Tensor, method: Literal[\"cat\", \"min\"] = \"cat\"\n    ) -> torch.Tensor:\n        q_values = [head(input_tensor) for head in value_heads]\n        if method == \"cat\":\n            q_values = torch.cat(q_values, dim=-1)\n        elif method == \"min\":\n            q_values = torch.min(torch.cat(q_values, dim=-1), dim=-1).values\n        else:\n            raise ValueError(f\"Unknown method: {method}\")\n\n        return q_values\n\n    def _cross_attention_pool_prefix(\n        self,\n        prefix_embs: torch.Tensor,\n        prefix_pad_masks: torch.Tensor,\n        use_target_network: bool,\n    ) -> torch.Tensor:\n        cross_attn = self.target_prefix_cross_attn if use_target_network else self.critic_prefix_cross_attn\n        state_token = self.target_state_token if use_target_network else self.critic_state_token\n\n        batch_size = prefix_embs.shape[0]\n        query = state_token.expand(batch_size, -1, -1)\n        key_padding_mask = ~prefix_pad_masks.to(dtype=torch.bool)\n\n        pooled, _ = cross_attn(\n            query=query,\n            key=prefix_embs,\n            value=prefix_embs,\n            key_padding_mask=key_padding_mask,\n            need_weights=False,\n        )\n        return pooled.squeeze(1)\n\n    def _gaussian_log_prob(\n        self,\n        sample: torch.Tensor,\n        mean: torch.Tensor,\n        std: torch.Tensor,\n    ) -> torch.Tensor:\n        std_safe = std.clamp_min(1e-6)\n        log_prob = -0.5 * (((sample - mean) / std_safe) ** 2 + 2.0 * torch.log(std_safe) + math.log(2.0 * math.pi))\n        return log_prob.mean(dim=(-1, -2))\n\n    def flow_sde_beta(self) -> torch.Tensor:\n        beta = beta_schedule(\n            int(self.flow_sde_step.item()),\n            beta0=self.flow_sde_initial_beta,\n            beta_min=self.flow_sde_beta_min,\n            T=self.flow_sde_beta_schedule_T,\n        )\n        return torch.tensor(beta, device=self.flow_sde_step.device, dtype=torch.float32)\n\n    def _sample_actions_flow_sde(\n        self,\n        state_features: tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor],\n        noise_scale: float,\n        requires_grad: bool,\n        return_log_prob: bool,\n    ) -> tuple[torch.Tensor, torch.Tensor | None]:\n        \"\"\"\n        add noise to the action sampling process using Flow-SDE method.\n        see https://arxiv.org/abs/2510.25889\n        \"\"\"\n\n        prefix_features, states = state_features\n        prefix_embs, prefix_pad_masks, _ = prefix_features\n        batch_size = prefix_embs.shape[0]\n        device = prefix_embs.device\n        beta = self.flow_sde_beta().to(device=device, dtype=prefix_embs.dtype)\n\n        past_key_values = self._build_kv_cache_from_prefix(prefix_features)\n        actions_shape = (batch_size, self.model.n_action_steps, self.model.max_action_dim)\n        x_t = torch.randn(actions_shape, device=device, dtype=prefix_embs.dtype)\n\n        timesteps = torch.linspace(1.0, 0.0, self.model.num_steps + 1, dtype=torch.float32, device=device)\n        step_log_probs: list[torch.Tensor] = []\n\n        for idx in range(self.model.num_steps):\n            t_cur = timesteps[idx]\n            t_next = timesteps[idx + 1]\n            delta = (t_cur - t_next).clamp_min(1e-6)\n\n            if requires_grad:\n                v_t = self.model.denoise_step(\n                    states,\n                    prefix_pad_masks,\n                    past_key_values,\n                    x_t,\n                    t_cur.expand(batch_size),\n                )\n            else:\n                with torch.no_grad():\n                    v_t = self.model.denoise_step(\n                        states,\n                        prefix_pad_masks,\n                        past_key_values,\n                        x_t,\n                        t_cur.expand(batch_size),\n                    )\n\n            t_cur_safe = t_cur.clamp(min=1e-4, max=1.0 - 1e-4)\n            t_cur_exp = t_cur_safe.view(1, 1, 1)\n            t_next_exp = t_next.view(1, 1, 1)\n            delta_exp = delta.view(1, 1, 1)\n\n            x0_pred = x_t - v_t * t_cur_exp\n            x1_pred = x_t + v_t * (1.0 - t_cur_exp)\n\n            if noise_scale > 0:\n                sigma_schedule = self.flow_sde_noise_level * noise_scale * torch.sqrt(t_cur_safe / (1.0 - t_cur_safe))\n                sigma = beta * sigma_schedule\n                sigma_exp = sigma.view(1, 1, 1)\n                x0_weight = 1.0 - t_next_exp\n                x1_weight = t_next_exp - sigma_exp.pow(2) * delta_exp / (2.0 * t_cur_exp)\n                x_mean = x0_pred * x0_weight + x1_pred * x1_weight\n                sigma_t = torch.sqrt(delta_exp) * sigma_exp\n                eps = torch.randn_like(x_t)\n                x_prev = x_mean + sigma_t * eps\n            else:\n                x0_weight = 1.0 - t_next_exp\n                x1_weight = t_next_exp\n                x_mean = x0_pred * x0_weight + x1_pred * x1_weight\n                sigma_t = torch.zeros_like(x_mean)\n                x_prev = x_mean\n\n            if return_log_prob:\n                step_log_probs.append(self._gaussian_log_prob(x_prev, x_mean, sigma_t))\n\n            x_t = x_prev\n\n        if return_log_prob:\n            log_probs = torch.stack(step_log_probs, dim=1).sum(dim=1)\n        else:\n            log_probs = None\n\n        return x_t, log_probs\n\n    def _build_kv_cache_from_prefix(\n        self,\n        prefix_features: tuple[torch.Tensor, torch.Tensor, torch.Tensor],\n    ):\n        \"\"\"Build KV cache for prefix. No grad needed.\"\"\"\n        prefix_embs, prefix_pad_masks, prefix_att_masks = prefix_features\n        prefix_att_2d_masks = make_att_2d_masks(prefix_pad_masks, prefix_att_masks)\n        prefix_position_ids = torch.cumsum(prefix_pad_masks, dim=1) - 1\n\n        with torch.no_grad():\n            _, past_key_values = self.model.paligemma_with_expert.forward(\n                attention_mask=prefix_att_2d_masks,\n                position_ids=prefix_position_ids,\n                past_key_values=None,\n                inputs_embeds=[prefix_embs, None],\n                use_cache=self.model.use_cache,\n                fill_kv_cache=True,\n                adarms_cond=[None, None],\n            )\n        return past_key_values\n\n    @override\n    def sac_init(self):\n        \"\"\"Initialize SAC-related components.\"\"\"\n\n        self.freeze_vision_tower()\n\n        register_fsdp_forward_method(self, \"bc_loss\")\n        register_fsdp_forward_method(self, \"sac_forward_critic\")\n        register_fsdp_forward_method(self, \"sac_forward_actor\")\n        register_fsdp_forward_method(self, \"sac_update_target_network\")\n        register_fsdp_forward_method(self, \"sac_forward_state_features\")\n\n    @override\n    def sac_forward_actor(\n        self,\n        state_features: tuple[\n            tuple[torch.Tensor, torch.Tensor, torch.Tensor],\n            torch.Tensor,\n        ],\n        is_first_micro_batch: bool = False,\n    ) -> tuple[torch.Tensor, torch.Tensor | None, dict[str, float]]:\n        actions, log_probs = self._sample_actions_flow_sde(\n            state_features=state_features,\n            noise_scale=self.flow_sde_train_noise_scale,\n            requires_grad=True,\n            return_log_prob=True,\n        )\n        if is_first_micro_batch:\n            self.flow_sde_step.add_(1)\n        actor_metrics: dict[str, float] = {}\n        if self.flow_sde_enable:\n            actor_metrics = {\n                \"flow_sde_beta\": float(self.flow_sde_beta().item()),\n                \"flow_sde_step\": float(self.flow_sde_step.item()),\n            }\n        return actions, log_probs, actor_metrics\n\n    @override\n    def sac_forward_critic(\n        self,\n        a: dict[str, torch.Tensor],\n        state_features: tuple[\n            tuple[torch.Tensor, torch.Tensor, torch.Tensor],\n            torch.Tensor,\n        ],\n        *,\n        use_target_network: bool = False,\n        method: Literal[\"cat\", \"min\"] = \"cat\",\n        requires_grad: bool = False,\n    ):\n        critic_head = self.target_network_heads if use_target_network else self.critic_heads\n        for p in critic_head.parameters():\n            p.requires_grad_(requires_grad)\n        prefix_cross_attn = self.target_prefix_cross_attn if use_target_network else self.critic_prefix_cross_attn\n        for p in prefix_cross_attn.parameters():\n            p.requires_grad_(requires_grad)\n        if use_target_network:\n            self.target_state_token.requires_grad_(requires_grad)\n        else:\n            self.critic_state_token.requires_grad_(requires_grad)\n\n        prefix_features, states = state_features\n        prefix_embs, prefix_pad_masks, _ = prefix_features\n        pooled_prefix_embs = self._cross_attention_pool_prefix(\n            prefix_embs=prefix_embs,\n            prefix_pad_masks=prefix_pad_masks,\n            use_target_network=use_target_network,\n        )  # (B, 2048)\n        actions = a[\"full_action\"][:, :10, :7]  # (B, 10, 7)\n        flattened_actions = actions.reshape(actions.shape[0], -1)  # (B, 70)\n        critic_input = torch.cat([pooled_prefix_embs, states, flattened_actions], dim=-1)\n\n        q_values = self._multi_heads_value(critic_head, critic_input, method=method)\n\n        return q_values\n\n    @override\n    def sac_get_critic_parameters(self) -> list[torch.nn.Parameter]:\n        critic_head_params = [p for head in self.critic_heads for p in head.parameters()]\n        critic_prefix_cross_attn_params = list(self.critic_prefix_cross_attn.parameters())\n        return critic_head_params + critic_prefix_cross_attn_params + [self.critic_state_token]\n\n    @override\n    def sac_get_named_actor_parameters(self) -> list[tuple[str, torch.nn.Parameter]]:\n        named_parameters = [(name, param) for name, param in self.model.named_parameters() if param.requires_grad]\n        return named_parameters\n\n    @override\n    def sac_forward_state_features(\n        self, s: dict[str, torch.Tensor]\n    ) -> tuple[tuple[torch.Tensor, torch.Tensor, torch.Tensor], torch.Tensor]:\n        with torch.no_grad():\n            prefix_features = self.model.embed_prefix(\n                images=s[\"images\"].unbind(dim=1),\n                img_masks=s[\"image_masks\"].unbind(dim=1),\n                lang_tokens=s[\"lang_tokens\"],\n                lang_masks=s[\"lang_masks\"],\n            )\n        return (prefix_features, s[\"states\"])\n\n    @override\n    @torch.no_grad()\n    def sac_update_target_network(self, tau: float):\n        for t_head, head in zip(self.target_network_heads, self.critic_heads, strict=True):\n            t_sd = t_head.state_dict()\n            h_sd = head.state_dict()\n            for k in t_sd.keys():\n                t_sd[k].mul_(1.0 - tau).add_(h_sd[k], alpha=tau)\n            t_head.load_state_dict(t_sd, strict=True)\n\n        t_cross_attn_sd = self.target_prefix_cross_attn.state_dict()\n        cross_attn_sd = self.critic_prefix_cross_attn.state_dict()\n        for k in t_cross_attn_sd.keys():\n            t_cross_attn_sd[k].mul_(1.0 - tau).add_(cross_attn_sd[k], alpha=tau)\n        self.target_prefix_cross_attn.load_state_dict(t_cross_attn_sd, strict=True)\n\n        self.target_state_token.data.mul_(1.0 - tau).add_(self.critic_state_token.data, alpha=tau)\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/pi0_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 Giga Team. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# from https://github.com/open-gigaai/giga-models\n\n\nfrom typing import Any\n\nimport torch\nimport torch.nn.functional as F\nfrom torchvision import transforms\n\n\nclass Normalize:\n    \"\"\"Normalize robot state vectors using mean/std or quantiles.\n\n    Args:\n        stats: A dict containing either {'mean', 'std'} or {'q01', 'q99'}.\n        use_quantiles: If True, use quantile based normalization.\n    \"\"\"\n\n    def __init__(self, stats: dict[str, Any], *, use_quantiles: bool = False) -> None:\n        self.EPSILON = 1e-6\n        self.stats = stats\n        self.use_quantiles = use_quantiles\n\n        required_attrs = [\"mean\", \"std\"]\n        if self.use_quantiles:\n            required_attrs = [\"q01\", \"q99\"]\n\n        for attr in required_attrs:\n            if attr not in stats:\n                raise AttributeError(f\"stats object is missing the following attribute: {attr}\")\n\n        if self.use_quantiles:\n            self.q01 = torch.tensor(stats[\"q01\"], dtype=torch.float32)\n            self.q99 = torch.tensor(stats[\"q99\"], dtype=torch.float32)\n        else:\n            self.mean = torch.tensor(stats[\"mean\"], dtype=torch.float32)\n            self.std = torch.tensor(stats[\"std\"], dtype=torch.float32)\n\n    def to(self, device: torch.device | str) -> None:\n        if self.use_quantiles:\n            self.q01 = self.q01.to(device)\n            self.q99 = self.q99.to(device)\n        else:\n            self.mean = self.mean.to(device)\n            self.std = self.std.to(device)\n\n    def __call__(self, x: torch.Tensor) -> torch.Tensor:\n        x_dim = x.shape[-1]\n        if self.use_quantiles:\n            return (x - self.q01[..., :x_dim]) / (\n                self.q99[..., :x_dim] - self.q01[..., :x_dim] + self.EPSILON\n            ) * 2.0 - 1.0\n        else:\n            return (x - self.mean[..., :x_dim]) / (self.std[..., :x_dim] + self.EPSILON)\n\n\nclass Unnormalize:\n    def __init__(self, stats, *, use_quantiles: bool = False):\n        self.EPSILON = 1e-6\n        self.stats = stats\n        self.use_quantiles = use_quantiles\n\n        if self.use_quantiles:\n            self.q01 = torch.tensor(stats[\"q01\"], dtype=torch.float32)\n            self.q99 = torch.tensor(stats[\"q99\"], dtype=torch.float32)\n        else:\n            self.mean = torch.tensor(stats[\"mean\"], dtype=torch.float32)\n            self.std = torch.tensor(stats[\"std\"], dtype=torch.float32)\n\n    def to(self, device: torch.device | str) -> None:\n        if self.use_quantiles:\n            self.q01 = self.q01.to(device)\n            self.q99 = self.q99.to(device)\n        else:\n            self.mean = self.mean.to(device)\n            self.std = self.std.to(device)\n\n    def __call__(self, x: torch.Tensor) -> torch.Tensor:\n        x_dim = x.shape[-1]\n        if self.use_quantiles:\n            return (x + 1.0) / 2.0 * (self.q99[..., :x_dim] - self.q01[..., :x_dim] + self.EPSILON) + self.q01[\n                ..., :x_dim\n            ]\n        else:\n            return x * (self.std[..., :x_dim] + self.EPSILON) + self.mean[..., :x_dim]\n\n\nclass DeltaActions:\n    \"\"\"Repacks absolute actions into delta action space.\"\"\"\n\n    def __init__(self):\n        # If the robot has mobile base, masks of base action are False and it doesn't need to be specified explicitly.\n        self.mask = torch.tensor([True, True, True, True, True, True, False, True, True, True, True, True, True, False])\n\n    def to(self, device: torch.device | str) -> None:\n        self.mask = self.mask.to(device)\n\n    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:\n        if \"action\" not in data or \"observation.state\" not in data:\n            return data\n        state, action = data[\"observation.state\"], data[\"action\"]\n        dims = self.mask.shape[-1]\n        action[..., :dims] -= torch.where(self.mask, state[..., :dims], torch.zeros_like(state[..., :dims])).unsqueeze(\n            -2\n        )\n        data[\"action\"] = action\n        return data\n\n\nclass AbsoluteActions:\n    \"\"\"Repacks delta actions into absolute action space.\"\"\"\n\n    def __init__(self):\n        # If the robot has mobile base, masks of base action are False and it doesn't need to be specified explicitly.\n        self.mask = torch.tensor([True, True, True, True, True, True, False, True, True, True, True, True, True, False])\n\n    def to(self, device: torch.device | str) -> None:\n        self.mask = self.mask.to(device)\n\n    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:\n        if \"action\" not in data or \"observation.state\" not in data:\n            return data\n        state, action = data[\"observation.state\"], data[\"action\"]\n        dims = self.mask.shape[-1]\n        action[..., :dims] += torch.where(self.mask, state[..., :dims], torch.zeros_like(state[..., :dims])).unsqueeze(\n            -2\n        )\n        data[\"action\"] = action\n        return data\n\n\nclass AlohaInputs:\n    \"\"\"Inputs for the Aloha policy.\"\"\"\n\n    def __init__(self, adapt_to_pi: bool = True) -> None:\n        self.joint_flip_mask = torch.tensor([1, -1, -1, 1, 1, 1, 1, 1, -1, -1, 1, 1, 1, 1])\n        self.adapt_to_pi = adapt_to_pi\n\n    def to(self, device: torch.device | str) -> None:\n        self.joint_flip_mask = self.joint_flip_mask.to(device)\n\n    def _gripper_from_angular_inv(self, value: torch.Tensor) -> torch.Tensor:\n        # Directly inverts the gripper_from_angular function.\n        value = _unnormalize(value, min_val=-0.6213, max_val=1.4910)\n        return value - 0.5476\n\n    def _gripper_to_angular(self, value: torch.Tensor) -> torch.Tensor:\n        # Aloha transforms the gripper positions into a linear space. The following code\n        # reverses this transformation to be consistent with pi0 which is pretrained in\n        # angular space.\n        #\n        # These values are coming from the Aloha code:\n        # PUPPET_GRIPPER_POSITION_OPEN, PUPPET_GRIPPER_POSITION_CLOSED\n        value = _unnormalize(value, min_val=0.01844, max_val=0.05800)\n\n        # This is the inverse of the angular to linear transformation inside the Interbotix code.\n        def linear_to_radian(linear_position, arm_length, horn_radius):\n            value = (horn_radius**2 + linear_position**2 - arm_length**2) / (2 * horn_radius * linear_position)\n            return torch.arcsin(torch.clip(value, -1.0, 1.0))\n\n        # The constants are taken from the Interbotix code.\n        value = linear_to_radian(value, arm_length=0.036, horn_radius=0.022)\n\n        # pi0 gripper data is normalized (0, 1) between encoder counts (2405, 3110).\n        # There are 4096 total encoder counts and aloha uses a zero of 2048.\n        # Converting this to radians means that the normalized inputs are between (0.5476, 1.6296)\n        return _normalize(value, min_val=0.5476, max_val=1.6296)\n\n    def _encode_actions_inv(self, actions: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            actions[:, :14] = self.joint_flip_mask * actions[:, :14]\n            actions[:, [6, 13]] = self._gripper_from_angular_inv(actions[:, [6, 13]])\n        return actions\n\n    def _decode_state(self, state: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            # Flip the joints.\n            state[:14] = self.joint_flip_mask * state[:14]\n            # Reverse the gripper transformation that is being applied by the Aloha runtime.\n            state[[6, 13]] = self._gripper_to_angular(state[[6, 13]])\n        return state\n\n    def _decode_aloha(self, state: torch.Tensor) -> torch.Tensor:\n        # state is [left_arm_joint_angles, left_arm_gripper, right_arm_joint_angles, right_arm_gripper]\n        # dim sizes: [6, 1, 6, 1]\n        state = self._decode_state(state)\n        return state\n\n    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:\n        \"\"\"Decode Aloha-specific input formats into the pi0 training/runtime\n        format.\"\"\"\n        state = self._decode_aloha(data[\"observation.state\"])\n        data[\"observation.state\"] = state\n        # Actions are only available during training.\n        if \"action\" in data:\n            actions = data[\"action\"]\n            actions = self._encode_actions_inv(actions)\n            data[\"action\"] = actions\n        return data\n\n    # VeRL: Batch Inference\n\n    def _encode_actions_inv_batch(self, actions: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            actions[..., :14] = self.joint_flip_mask * actions[..., :14]\n            actions[..., [6, 13]] = self._gripper_from_angular_inv(actions[..., [6, 13]])\n        return actions\n\n    def _decode_state_batch(self, state: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            state[..., :14] = self.joint_flip_mask * state[..., :14]\n            state[..., [6, 13]] = self._gripper_to_angular(state[..., [6, 13]])\n        return state\n\n    def call_batch(self, data: dict[str, Any]) -> dict[str, Any]:\n        state = self._decode_state_batch(data[\"observation.state\"])\n        data[\"observation.state\"] = state\n        if \"action\" in data:\n            actions = data[\"action\"]\n            actions = self._encode_actions_inv_batch(actions)\n            data[\"action\"] = actions\n        return data\n\n\nclass AlohaOutputs:\n    \"\"\"Outputs for the Aloha policy.\"\"\"\n\n    def __init__(self, original_action_dim: int, adapt_to_pi: bool = True):\n        \"\"\"\n        Args:\n            original_action_dim: int. The original action dimension of the policy. dual-arm robot has 14 dims and mobile\n                                      dual-arm robot has 16 dims.\n            adapt_to_pi: bool. If true, this will convert the joint and gripper values from the standard Aloha space to\n            the space used by the pi internal runtime which was used to train the base model.\n        \"\"\"\n        self.joint_flip_mask = torch.tensor([1, -1, -1, 1, 1, 1, 1, 1, -1, -1, 1, 1, 1, 1])\n        self.original_action_dim = original_action_dim\n        self.adapt_to_pi = adapt_to_pi\n\n    def to(self, device: torch.device | str) -> None:\n        self.joint_flip_mask = self.joint_flip_mask.to(device)\n\n    def _gripper_from_angular(self, value: torch.Tensor) -> torch.Tensor:\n        # Convert from the gripper position used by pi0 to the gripper position that is used by Aloha.\n        # Note that the units are still angular but the range is different.\n\n        # We do not scale the output since the trossen model predictions are already in radians.\n        # See the comment in _gripper_to_angular for a derivation of the constant\n        value = value + 0.5476\n\n        # These values are coming from the Aloha code:\n        # PUPPET_GRIPPER_JOINT_OPEN, PUPPET_GRIPPER_JOINT_CLOSE\n        return _normalize(value, min_val=-0.6213, max_val=1.4910)\n\n    def _encode_actions(self, actions: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            # Flip the joints.\n            actions[:, :14] = self.joint_flip_mask * actions[:, :14]\n            actions[:, [6, 13]] = self._gripper_from_angular(actions[:, [6, 13]])\n        return actions\n\n    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:\n        actions = data[\"action\"][:, : self.original_action_dim]\n        return {\"action\": self._encode_actions(actions)}\n\n    # VeRL: Batch Inference\n\n    def _encode_actions_batch(self, actions: torch.Tensor) -> torch.Tensor:\n        if self.adapt_to_pi:\n            actions[..., :14] = self.joint_flip_mask * actions[..., :14]\n            actions[..., [6, 13]] = self._gripper_from_angular(actions[..., [6, 13]])\n        return actions\n\n    def call_batch(self, data: dict[str, Any]) -> dict[str, Any]:\n        actions = data[\"action\"][..., : self.original_action_dim]\n        return {\"action\": self._encode_actions_batch(actions)}\n\n\nclass PadStatesAndActions:\n    \"\"\"Zero-pads states and actions to the model action dimension.\"\"\"\n\n    def __init__(self, action_dim: int) -> None:\n        self.action_dim = action_dim\n\n    def _pad_to_dim(self, x: torch.Tensor, target_dim: int, axis: int = -1) -> torch.Tensor:\n        \"\"\"Pad an array to the target dimension with zeros along the specified\n        axis.\"\"\"\n        current_dim = x.shape[axis]\n        if current_dim < target_dim:\n            shape = list(x.shape)\n            shape[-1] = target_dim\n            new_vector = torch.zeros(*shape, dtype=x.dtype, device=x.device)\n            new_vector[..., :current_dim] = x\n            x = new_vector\n        return x\n\n    def __call__(self, data: dict[str, Any]) -> dict[str, Any]:\n        data[\"observation.state\"] = self._pad_to_dim(data[\"observation.state\"], self.action_dim, axis=-1)\n        if \"action\" in data:\n            data[\"action\"] = self._pad_to_dim(data[\"action\"], self.action_dim, axis=-1)\n        return data\n\n\ndef _normalize(x: torch.Tensor, min_val: float, max_val: float) -> torch.Tensor:\n    return (x - min_val) / (max_val - min_val)\n\n\ndef _unnormalize(x: torch.Tensor, min_val: float, max_val: float) -> torch.Tensor:\n    return x * (max_val - min_val) + min_val\n\n\ndef resize_with_pad(img: torch.Tensor, width: int, height: int, pad_value: float = -1.0) -> torch.Tensor:\n    \"\"\"Resize an image to fit inside the given (width, height) while preserving\n    aspect ratio, then pad with the specified value so that the final image\n    exactly matches the target size.\n\n    Args:\n        img: Input image, shape (C, H, W), with values typically in [0, 1].\n        width: Target width (W).\n        height: Target height (H).\n        pad_value: Value to use for padding, defaults to -1.\n\n    Returns:\n        A torch.Tensor of shape (C, height, width).\n    \"\"\"\n    # Validate input dimensions\n    if img.ndim != 3:\n        raise ValueError(f\"(C,H,W) expected, but got {img.shape}\")\n\n    cur_height, cur_width = img.shape[1:]\n\n    ratio = max(cur_width / width, cur_height / height)\n    resized_height = int(cur_height / ratio)\n    resized_width = int(cur_width / ratio)\n    resized_img = F.interpolate(\n        img.unsqueeze(0), size=(resized_height, resized_width), mode=\"bilinear\", align_corners=False\n    ).squeeze(0)\n\n    pad_height = max(0, int(height - resized_height))\n    pad_width = max(0, int(width - resized_width))\n\n    pad_top = pad_height // 2\n    pad_bottom = pad_height - pad_top\n    pad_left = pad_width // 2\n    pad_right = pad_width - pad_left\n\n    padded_img = F.pad(resized_img, (pad_left, pad_right, pad_top, pad_bottom), value=pad_value)\n    return padded_img.squeeze(0)\n\n\nclass ImageTransform:\n    def __init__(\n        self,\n        resize_imgs_with_padding: tuple[int, int],\n        present_img_keys: list[str] | None = None,\n        enable_image_aug: bool = False,\n    ) -> None:\n        self.resize_imgs_with_padding = resize_imgs_with_padding\n        self.present_img_keys = present_img_keys\n        if self.present_img_keys is None:\n            self.present_img_keys = [\n                \"observation.images.cam_high\",\n                \"observation.images.cam_left_wrist\",\n                \"observation.images.cam_right_wrist\",\n            ]\n        self.enable_image_aug = enable_image_aug\n        self.width, self.height = resize_imgs_with_padding\n        if self.enable_image_aug:\n            self.color_jitter_transform = transforms.ColorJitter(\n                brightness=0.3,\n                contrast=0.4,\n                saturation=0.5,\n            )\n            self.pose_transform = transforms.Compose(\n                [\n                    transforms.RandomCrop(int(self.width * 0.95), int(self.height * 0.95)),\n                    transforms.Resize((self.width, self.height)),\n                    transforms.RandomRotation((-5, 5)),\n                ]\n            )\n\n    def __call__(self, data: dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:\n        \"\"\"Preprocesses input images: optionally scales and pads to a fixed size,\n        then maps the pixel range from [0,1] to [-1,1].\n\n        Returns two lists:\n            images: The processed image arrays (C, H, W).\n            img_masks: A list of boolean masks of the same length as images, currently fixed to True.\n        \"\"\"\n        images = []\n        img_masks = []\n\n        for key in self.present_img_keys:\n            if key not in data:\n                raise ValueError(\n                    f\"{key} not found in data. Please check the present_img_keys in the config or the dataset.\"\n                )\n\n            img = data[key]\n            # [C, H, W] -> preprocess\n            if self.resize_imgs_with_padding is not None:\n                original_height, original_width = img.shape[1:]\n                target_height, target_width = self.resize_imgs_with_padding\n                if original_height != target_height or original_width != target_width:\n                    img = resize_with_pad(img, *self.resize_imgs_with_padding, pad_value=0)\n\n            if self.enable_image_aug:\n                if \"wrist\" not in key:\n                    img = self.pose_transform(img)\n                img = self.color_jitter_transform(img)\n\n            # Normalize pixel values to [-1, 1]\n            img = img * 2.0 - 1.0\n\n            images.append(img)\n            img_masks.append(torch.tensor(True, dtype=torch.bool, device=img.device))\n\n        return images, img_masks\n\n    # VeRL: Batch Inference\n\n    def call_batch(self, data: dict[str, torch.Tensor]) -> tuple[list[torch.Tensor], list[torch.Tensor]]:\n        images = []\n        img_masks = []\n\n        for key in self.present_img_keys:\n            if key not in data:\n                raise ValueError(\n                    f\"{key} not found in data. Please check the present_img_keys in the config or the dataset.\"\n                )\n\n            img = data[key]\n            if img.ndim != 4:\n                raise ValueError(f\"(B,C,H,W) expected, but got {img.shape}\")\n\n            if self.resize_imgs_with_padding is not None:\n                original_height, original_width = img.shape[2:]\n                target_height, target_width = self.resize_imgs_with_padding\n                if original_height != target_height or original_width != target_width:\n                    ratio = max(original_width / target_width, original_height / target_height)\n                    resized_height = int(original_height / ratio)\n                    resized_width = int(original_width / ratio)\n                    img = F.interpolate(img, size=(resized_height, resized_width), mode=\"bilinear\", align_corners=False)\n                    pad_height = max(0, int(target_height - resized_height))\n                    pad_width = max(0, int(target_width - resized_width))\n                    pad_top = pad_height // 2\n                    pad_bottom = pad_height - pad_top\n                    pad_left = pad_width // 2\n                    pad_right = pad_width - pad_left\n                    img = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom), value=0)\n\n            if self.enable_image_aug:\n                imgs = []\n                for sample in img:\n                    if \"wrist\" not in key:\n                        sample = self.pose_transform(sample)\n                    sample = self.color_jitter_transform(sample)\n                    imgs.append(sample)\n                img = torch.stack(imgs, dim=0)\n\n            img = img / 255.0 * 2.0 - 1.0  # pi05 libero\n            images.append(img)\n            img_masks.append(torch.ones((img.shape[0],), dtype=torch.bool, device=img.device))\n\n        return images, img_masks\n\n\nclass PromptTokenizerTransform:\n    def __init__(self, max_length: int, discrete_state_input: bool = False) -> None:\n        # self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_model_path)\n        self.tokenizer_max_length = max_length\n        self.discrete_state_input = discrete_state_input\n\n    def __call__(self, data: dict[str, Any], tokenizer) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"Tokenize the text input.\n\n        Args:\n            data: Dict containing 'task' string and optionally 'observation.state' tensor to infer device.\n\n        Returns:\n            A tuple of (lang_tokens, lang_masks), both as torch tensors on the inferred device.\n        \"\"\"\n        task = data[\"task\"].strip().replace(\"_\", \" \").replace(\"\\n\", \" \")\n\n        # Infer device from observation.state if available\n        device = data[\"observation.state\"].device if \"observation.state\" in data else torch.device(\"cpu\")\n\n        if self.discrete_state_input:\n            assert \"observation.state\" in data, \"discrete_state_input is True, but observation.state is not found.\"\n            discretized_state = (\n                torch.bucketize(data[\"observation.state\"], torch.linspace(-1, 1, 256 + 1, device=device)[:-1]) - 1\n            )\n            state_values = \" \".join([str(int(x)) for x in discretized_state.tolist()])\n            task = f\"Task: {task}, State: {state_values};\\nAction: \"\n        else:\n            # PaliGemma prompt has to end with a new line in Pi0\n            task = f\"{task}\\n\"\n\n        tokenized_prompt = tokenizer(\n            task,\n            padding=\"max_length\",\n            padding_side=\"right\",\n            max_length=self.tokenizer_max_length,\n            return_tensors=\"pt\",\n        )\n        lang_tokens = tokenized_prompt[\"input_ids\"][0].to(dtype=torch.int32, device=device)\n        lang_masks = tokenized_prompt[\"attention_mask\"][0].to(dtype=torch.bool, device=device)\n\n        return lang_tokens, lang_masks\n\n    # VeRL: Batch Inference\n\n    def call_batch(self, data: dict[str, Any], tokenizer) -> tuple[torch.Tensor, torch.Tensor]:\n        task = data[\"task\"]\n        if hasattr(task, \"tolist\") and not isinstance(task, str):\n            tasks = task.tolist()\n        else:\n            tasks = list(task)\n        tasks = [str(t).strip().replace(\"_\", \" \").replace(\"\\n\", \" \") for t in tasks]\n\n        device = data[\"observation.state\"].device if \"observation.state\" in data else torch.device(\"cpu\")\n\n        if self.discrete_state_input:\n            assert \"observation.state\" in data, \"discrete_state_input is True, but observation.state is not found.\"\n            state = data[\"observation.state\"]\n            discretized_state = torch.bucketize(state, torch.linspace(-1, 1, 256 + 1, device=device)[:-1]) - 1\n            state_values = [\" \".join([str(int(x)) for x in row.tolist()]) for row in discretized_state]\n            tasks = [\n                f\"Task: {task_item}, State: {state_value};\\nAction: \"\n                for task_item, state_value in zip(tasks, state_values, strict=False)\n            ]\n        else:\n            tasks = [f\"{task_item}\\n\" for task_item in tasks]\n\n        tokenized_prompt = tokenizer(\n            tasks,\n            padding=\"max_length\",\n            padding_side=\"right\",\n            max_length=self.tokenizer_max_length,\n            return_tensors=\"pt\",\n        )\n        lang_tokens = tokenized_prompt[\"input_ids\"].to(dtype=torch.int32, device=device)\n        lang_masks = tokenized_prompt[\"attention_mask\"].to(dtype=torch.bool, device=device)\n\n        return lang_tokens, lang_masks\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/policy/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/policy/base.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom abc import ABC, abstractmethod\n\nimport torch\n\n\nclass Pi0Input(ABC):\n    def __init__(self):\n        # three images for pi0 input with keys:\n        # [\n        #     'observation.images.cam_high',\n        #     'observation.images.cam_left_wrist',\n        #     'observation.images.cam_right_wrist',\n        # ],\n        # each with shape (B, C, H, W)\n        self.images: dict[str, torch.Tensor] = {}\n\n        # image masks corresponding to the images, each with shape (B,)\n        self.img_masks: list[torch.Tensor] = []\n\n        # task description as a list of strings\n        self.task: list[str] = []\n\n        # robot state with shape (B, state_dim)\n        self.state: torch.Tensor = None\n\n    @classmethod\n    @abstractmethod\n    def from_env_obs(cls, env_obs) -> \"Pi0Input\": ...\n\n\nclass Pi0Output:\n    def __init__(self):\n        self.action: torch.Tensor = None\n\n    @classmethod\n    @abstractmethod\n    def from_model_output(cls, model_output) -> \"Pi0Output\": ...\n"
  },
  {
    "path": "verl/experimental/vla/models/pi0_torch/policy/libero_policy.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom typing_extensions import override\n\nfrom verl.protocol import DataProto\n\nfrom .base import Pi0Input, Pi0Output\n\nPI0_MAX_STATE_DIM = 32\nPI0_ACTION_CHUNK_SIZE = 10\nLIBERO_ACTION_DIM = 7\n\n\nclass LiberoPi0Input(Pi0Input):\n    @override\n    @classmethod\n    def from_env_obs(cls, env_obs: DataProto) -> \"LiberoPi0Input\":\n        input = cls()\n\n        # Process images\n        images = env_obs.batch[\"full_image\"]\n        wrist_images = env_obs.batch[\"wrist_image\"]\n\n        batch_size = images.shape[0]\n        cam_high = images.permute(0, 3, 1, 2)\n        left_wrist = wrist_images.permute(0, 3, 1, 2)  # (B, H, W, C) -> (B, C, H, W)\n        empty_images = torch.zeros(\n            (batch_size, 3, cam_high.shape[2], cam_high.shape[3]),\n            device=env_obs.batch.device,\n            dtype=torch.bfloat16,\n        )\n\n        input.images = {\n            \"observation.images.cam_high\": cam_high.to(torch.bfloat16),\n            \"observation.images.cam_left_wrist\": left_wrist.to(torch.bfloat16),\n            \"observation.images.cam_right_wrist\": empty_images,\n        }\n        input.img_masks = [\n            torch.ones((batch_size,), device=env_obs.batch.device, dtype=torch.bool),\n            torch.ones((batch_size,), device=env_obs.batch.device, dtype=torch.bool),\n            torch.zeros((batch_size,), device=env_obs.batch.device, dtype=torch.bool),\n        ]\n\n        # Process other data\n        input.task = list(env_obs.non_tensor_batch[\"task_descriptions\"])\n\n        state = env_obs.batch[\"state\"]\n        input.state = torch.nn.functional.pad(\n            state, (0, max(0, PI0_MAX_STATE_DIM - state.shape[-1])), \"constant\", 0\n        ).to(env_obs.batch.device, dtype=torch.float32)\n\n        return input\n\n\nclass LiberoPi0Output(Pi0Output):\n    @override\n    @classmethod\n    def from_model_output(cls, model_output: dict) -> \"LiberoPi0Output\":\n        output = cls()\n        output.action = model_output[\"full_action\"][:, :PI0_ACTION_CHUNK_SIZE, :LIBERO_ACTION_DIM]\n        return output\n"
  },
  {
    "path": "verl/experimental/vla/models/register_vla_models.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n\"\"\"Utility helpers to register custom VLA models with Hugging Face Auto classes.\"\"\"\n\nfrom transformers import AutoConfig, AutoImageProcessor, AutoProcessor\n\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\n\nfrom .openvla_oft.configuration_prismatic import OpenVLAConfig\nfrom .openvla_oft.modeling_prismatic import OpenVLAForActionPrediction\nfrom .openvla_oft.processing_prismatic import PrismaticImageProcessor, PrismaticProcessor\nfrom .pi0_torch import PI0ForActionPrediction, PI0TorchConfig\n\n_REGISTERED_MODELS = {\n    \"openvla_oft\": False,\n    \"pi0_torch\": False,\n}\nAutoModelForVision2Seq = get_auto_model_for_vision2seq()\n\n\ndef register_openvla_oft() -> None:\n    \"\"\"Register the OpenVLA OFT model and processors.\"\"\"\n    if _REGISTERED_MODELS[\"openvla_oft\"]:\n        return\n\n    AutoConfig.register(\"openvla\", OpenVLAConfig)\n    AutoImageProcessor.register(OpenVLAConfig, PrismaticImageProcessor)\n    AutoProcessor.register(OpenVLAConfig, PrismaticProcessor)\n    AutoModelForVision2Seq.register(OpenVLAConfig, OpenVLAForActionPrediction)\n\n    _REGISTERED_MODELS[\"openvla_oft\"] = True\n\n\ndef register_pi0_torch_model() -> None:\n    \"\"\"Register the PI0 wrapper with the HF auto classes.\"\"\"\n    if _REGISTERED_MODELS[\"pi0_torch\"]:\n        return\n\n    AutoConfig.register(\"pi0_torch\", PI0TorchConfig)\n    AutoModelForVision2Seq.register(PI0TorchConfig, PI0ForActionPrediction)\n\n    _REGISTERED_MODELS[\"pi0_torch\"] = True\n\n\ndef register_vla_models() -> None:\n    \"\"\"Register all custom VLA models with Hugging Face.\"\"\"\n    register_openvla_oft()\n    register_pi0_torch_model()\n"
  },
  {
    "path": "verl/experimental/vla/naive_rollout_rob.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nIn single GPU rollout, the sequences are generated directly by sampling from the model.\nThe output will contain\n1. output_ids\n2. attention_masks (left padding)\n3. eos_masks\n4. log_probs\n\"\"\"\n\nimport json\nimport logging\nimport os\n\nimport torch\nfrom PIL import Image\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.nn.utils.rnn import pad_sequence\n\nfrom verl import DataProto\nfrom verl.experimental.vla.envs.action_utils import center_crop_image, resize_image\nfrom verl.experimental.vla.models.openvla_oft.modeling_prismatic import OpenVLAForActionPrediction\nfrom verl.experimental.vla.models.openvla_oft.processing_prismatic import PrismaticProcessor\nfrom verl.utils.device import get_device_id, get_device_name, get_torch_device\nfrom verl.workers.rollout.base import BaseRollout\n\nlogger = logging.getLogger(__name__)\n\n\n__all__ = [\"NaiveRolloutRob\"]\n\n\ndef pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False):\n    \"\"\"\n    pad a 2D tensors (e.g. responses, logprobs) in the last dim to max_seq_length.\n    input shape: [bs, seq_length]\n    output shape: [bs, max_seq_length]\n    (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad\n    \"\"\"\n    if tensors.shape[-1] >= max_seq_len:\n        return tensors\n    pad_tuple = (max_seq_len - tensors.shape[-1], 0) if left_pad else (0, max_seq_len - tensors.shape[-1])\n    return torch.nn.functional.pad(tensors, pad_tuple, \"constant\", pad_token_id)\n\n\ndef process_input(task_descriptions, images_and_states, processor):\n    batchdata = {\"input_ids\": [], \"attention_mask\": [], \"pixel_values\": []}\n\n    for i in range(len(task_descriptions)):\n        task_description = task_descriptions[i]\n        image = resize_image(images_and_states[\"full_image\"][i].cpu().numpy(), (224, 224))\n        image = Image.fromarray(image).convert(\"RGB\")\n        image = center_crop_image(image)\n        prompt = f\"In: What action should the robot take to {task_description.lower()}?\\nOut:\"\n        batch_feature = processor(prompt, image)\n\n        input_ids = batch_feature[\"input_ids\"]\n        attention_mask = batch_feature[\"attention_mask\"]\n        pixel_values = batch_feature[\"pixel_values\"]\n\n        if not torch.all(input_ids[:, -1] == 29871):\n            input_ids = torch.cat(\n                (input_ids, torch.unsqueeze(torch.Tensor([29871]).long(), dim=0).to(input_ids.device)), dim=1\n            )\n            attention_mask = torch.cat(\n                (attention_mask, torch.unsqueeze(torch.Tensor([True]).bool(), dim=0).to(attention_mask.device)), dim=1\n            )\n\n        batchdata[\"input_ids\"].append(input_ids)\n        batchdata[\"attention_mask\"].append(attention_mask)\n        batchdata[\"pixel_values\"].append(pixel_values)\n\n    device = get_device_id()\n\n    batchdata[\"input_ids\"] = [x.transpose(0, 1) for x in batchdata[\"input_ids\"]]\n    batchdata[\"attention_mask\"] = [x.transpose(0, 1) for x in batchdata[\"attention_mask\"]]\n    batchdata[\"input_ids\"] = (\n        pad_sequence(batchdata[\"input_ids\"], batch_first=True, padding_value=processor.tokenizer.pad_token_id)\n        .squeeze(-1)\n        .to(device)\n    )\n    batchdata[\"attention_mask\"] = (\n        pad_sequence(batchdata[\"attention_mask\"], batch_first=True, padding_value=0).squeeze(-1).to(device)\n    )\n\n    padding_mask = batchdata[\"input_ids\"].ne(processor.tokenizer.pad_token_id)\n    assert torch.all(padding_mask == batchdata[\"attention_mask\"].ne(0))\n    padding_mask = ~padding_mask\n    padding_mask = padding_mask.int()\n    sorted_indices = torch.argsort(padding_mask, dim=1, descending=True, stable=True)\n    batchdata[\"input_ids\"] = torch.gather(batchdata[\"input_ids\"], 1, sorted_indices)\n    batchdata[\"attention_mask\"] = torch.gather(batchdata[\"attention_mask\"], 1, sorted_indices)\n\n    batchdata[\"pixel_values\"] = torch.cat(batchdata[\"pixel_values\"], dim=0).to(device)\n    assert torch.all(batchdata[\"attention_mask\"].ne(0) == batchdata[\"input_ids\"].ne(processor.tokenizer.pad_token_id))\n\n    return batchdata\n\n\nclass NaiveRolloutRob(BaseRollout):\n    def __init__(\n        self,\n        model_config: dict,\n        module: torch.nn.Module = None,\n    ):\n        self.model_config = model_config\n        if module is not None:\n            self.module = module\n        else:\n            self.module = OpenVLAForActionPrediction.from_pretrained(model_config[\"path\"], trust_remote_code=True)\n        self.module.vision_backbone.set_num_images_in_input(1)\n        self.processor = PrismaticProcessor.from_pretrained(model_config[\"path\"], trust_remote_code=True)\n        dataset_statistics_path = os.path.join(model_config[\"path\"], \"dataset_statistics.json\")\n        if os.path.isfile(dataset_statistics_path):\n            with open(dataset_statistics_path) as f:\n                norm_stats = json.load(f)\n            if isinstance(self.module, FSDP):\n                self.module.module.norm_stats = norm_stats\n            else:\n                self.module.norm_stats = norm_stats\n        self.module.eval()\n\n    @torch.no_grad()\n    def _generate_one_step(self, prompts: dict, do_sample, temperature, max_prompt_length):\n        idx = prompts[\"input_ids\"]  # (bs, prompt_length)\n        attention_mask = prompts[\"attention_mask\"]  # left-padded attention_mask\n        pixel_values = prompts[\"pixel_values\"]\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            actions, response = self.module.generate_action_verl(\n                input_ids=idx,\n                pixel_values=pixel_values,\n                attention_mask=attention_mask,\n                padding_idx=self.processor.tokenizer.pad_token_id,\n                do_sample=do_sample,\n                unnorm_key=\"libero_10_no_noops\",\n                temperature=temperature,\n            )\n\n        assert self.processor.tokenizer.pad_token_id is not None\n\n        assert idx.ndim == 2\n        idx = pad_sequence_to_length(\n            idx, max_seq_len=max_prompt_length, pad_token_id=self.processor.tokenizer.pad_token_id, left_pad=True\n        )\n\n        assert attention_mask.ndim == 2\n        attention_mask = pad_sequence_to_length(\n            attention_mask, max_seq_len=max_prompt_length, pad_token_id=0, left_pad=True\n        )\n\n        device_type = get_device_name()\n        assert idx.device.type == device_type\n        assert response.device.type == device_type\n        assert attention_mask.device.type == device_type\n        assert pixel_values.device.type == device_type\n        batch = {\n            \"responses\": response,\n            \"input_ids\": idx,\n            \"attention_mask\": attention_mask,\n            \"pixel_values\": pixel_values,\n            \"action\": actions,\n        }\n\n        return batch\n\n    # @conditional_profiler(name=\"generate_sequences\", path=\"traces/rollout\", max_steps=5)\n    @torch.no_grad()\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Generate sequences\"\"\"\n        # make sampling args can be overriden by inputs\n        do_sample = prompts.meta_info[\"do_sample\"]\n        temperature = prompts.meta_info[\"temperature\"]\n        max_prompt_length = prompts.meta_info[\"prompt_length\"]\n        # TODO: split into micro-batches\n        task_descriptions = prompts.non_tensor_batch[\"task_descriptions\"]\n        images_and_states = {\"full_image\": prompts.batch[\"full_image\"]}\n        vla_input = process_input(task_descriptions, images_and_states, self.processor)\n\n        vla_output = self._generate_one_step(vla_input, do_sample, temperature, max_prompt_length)\n        # batch = TensorDict(vla_output)\n        batch = DataProto.from_dict(tensors=vla_output)\n        return batch\n\n    async def update_weights(self, weights_iterator, **kwargs):\n        prefix = \"_fsdp_wrapped_module.\"\n        target_state_dict = self.module.state_dict()\n        loaded_tensors_count = 0\n        for name, param in weights_iterator:\n            cleaned_name = name.replace(prefix, \"\")\n            if cleaned_name in target_state_dict:\n                target_tensor = target_state_dict[cleaned_name]\n                try:\n                    target_tensor.copy_(param, non_blocking=True)\n                    loaded_tensors_count += 1\n                except Exception as e:\n                    logger.warning(f\"Warning: Failed to copy tensor '{cleaned_name}'. Error: {e}\")\n            else:\n                logger.warning(f\"Warning: Failed to copy tensor '{cleaned_name}'. Model has no such key.\")\n        logger.info(f\"Rollout model weights updated. Loaded {loaded_tensors_count} tensors one by one.\")\n\n    async def release(self):\n        if self.module.device.type == get_device_name():\n            logger.info(\"Releasing rollout model to CPU.\")\n            self.module.cpu()\n            self.device = torch.device(\"cpu\")\n            get_torch_device().empty_cache()\n\n    async def resume(self, **kwargs):\n        if self.module.device.type == \"cpu\":\n            target_device = get_device_name()\n            logger.info(f\"Resuming rollout model to device: {target_device}.\")\n            self.module.to(target_device)\n            self.device = torch.device(target_device)\n"
  },
  {
    "path": "verl/experimental/vla/prepare_libero_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPreprocess the Geometry3k dataset to parquet format\n\"\"\"\n\nimport argparse\nimport os\nimport random\n\nimport numpy as np\nimport torch\nfrom datasets import Dataset\nfrom libero.libero import get_libero_path\nfrom libero.libero.benchmark import Benchmark, get_benchmark, get_benchmark_dict\n\n\ndef patched_get_task_init_states(self, i):\n    init_states_path = os.path.join(\n        get_libero_path(\"init_states\"),\n        self.tasks[i].problem_folder,\n        self.tasks[i].init_states_file,\n    )\n    init_states = torch.load(init_states_path, weights_only=False)\n    return init_states\n\n\nBenchmark.get_task_init_states = patched_get_task_init_states\n\n\ndef compute_total_num_group_envs(task_suite: Benchmark):\n    total_num_group_envs = 0\n    trial_id_bins = []\n    for task_id in range(task_suite.get_num_tasks()):\n        task_num_trials = len(task_suite.get_task_init_states(task_id))\n        trial_id_bins.append(task_num_trials)\n\n        total_num_group_envs += task_num_trials\n\n    cumsum_trial_id_bins = np.cumsum(trial_id_bins)\n    return total_num_group_envs, cumsum_trial_id_bins\n\n\ndef build_dataset_for_suite(task_suite_name: str, local_save_dir: str):\n    task_suite = get_benchmark(task_suite_name)()\n    total_num_group_envs, cumsum_trial_id_bins = compute_total_num_group_envs(task_suite)\n    print(f\"\\n[Suite: {task_suite_name}]\")\n    print(f\"Total number of group envs: {total_num_group_envs}\")\n    print(f\"Cumsum trial id bins: {cumsum_trial_id_bins}\")\n\n    def get_state_ids_for_task(task_id):\n        start_id = 0 if task_id == 0 else cumsum_trial_id_bins[task_id - 1]\n        end_id = cumsum_trial_id_bins[task_id]\n        return list(range(start_id, end_id))\n\n    all_task_ids = list(range(task_suite.get_num_tasks()))\n    if len(all_task_ids) > 1:\n        train_task_num = max(1, len(all_task_ids) - 1)\n        train_task_ids = sorted(random.sample(all_task_ids, train_task_num))\n        ood_test_task_ids = sorted(list(set(all_task_ids) - set(train_task_ids)))\n    else:\n        train_task_ids = all_task_ids\n        ood_test_task_ids = []\n\n    print(\"\\n[Data Split Plan]\")\n    print(f\"Training Task IDs: {train_task_ids}\")\n    print(f\"OOD Test Task IDs: {ood_test_task_ids}\")\n\n    train_metadata = []\n    test_metadata = []\n\n    # Train split + ID test split\n    for task_id in train_task_ids:\n        all_trials = get_state_ids_for_task(task_id)\n        random.shuffle(all_trials)\n\n        train_count = int(len(all_trials) * 0.8)\n        train_count = min(train_count, 40)\n        selected_train_trials = all_trials[:train_count]\n        selected_id_test_trials = all_trials[train_count:]\n\n        for state_id in selected_train_trials:\n            train_metadata.append({\"task_id\": task_id, \"state_id\": state_id, \"data_source\": \"train\"})\n\n        for state_id in selected_id_test_trials[:10]:\n            test_metadata.append({\"task_id\": task_id, \"state_id\": state_id, \"data_source\": \"test_in_distribution\"})\n\n    # OOD split\n    for ood_task_id in ood_test_task_ids:\n        ood_all_trials = get_state_ids_for_task(ood_task_id)\n        random.shuffle(ood_all_trials)\n        selected_ood_trials = ood_all_trials[:20]\n        for state_id in selected_ood_trials:\n            test_metadata.append(\n                {\"task_id\": ood_task_id, \"state_id\": state_id, \"data_source\": \"test_out_of_distribution\"}\n            )\n\n    print(f\"Generated {len(train_metadata)} training samples.\")\n    print(f\"Generated {len(test_metadata)} testing samples.\")\n    print(\"-\" * 20)\n\n    train_ds_meta = Dataset.from_list(train_metadata)\n    test_ds_meta = Dataset.from_list(test_metadata)\n\n    def map_and_process(example, idx):\n        task_id = example[\"task_id\"]\n        state_id = example[\"state_id\"]\n        data_source = example[\"data_source\"]\n        split = \"train\" if data_source == \"train\" else \"test\"\n        task = task_suite.get_task(task_id)\n\n        data = {\n            \"data_source\": data_source,\n            \"prompt\": task.language,\n            \"state_ids\": state_id,\n            \"task_ids\": task_id,\n            \"ability\": \"robot\",\n            \"extra_info\": {\n                \"split\": split,\n                \"state_ids\": state_id,\n                \"index\": idx,\n                \"task\": task,\n                \"task_ids\": task_id,\n                \"task_suite_name\": task_suite_name,\n            },\n        }\n        return data\n\n    print(\"Mapping and processing training dataset...\")\n    train_dataset = train_ds_meta.map(map_and_process, with_indices=True, num_proc=8)\n    print(\"Mapping and processing test dataset...\")\n    test_dataset = test_ds_meta.map(map_and_process, with_indices=True, num_proc=8)\n\n    suite_save_dir = os.path.join(local_save_dir, task_suite_name)\n    os.makedirs(suite_save_dir, exist_ok=True)\n    print(f\"Saving training dataset to {os.path.join(suite_save_dir, 'train.parquet')}\")\n    train_dataset.to_parquet(os.path.join(suite_save_dir, \"train.parquet\"))\n    print(f\"Saving test dataset to {os.path.join(suite_save_dir, 'test.parquet')}\")\n    test_dataset.to_parquet(os.path.join(suite_save_dir, \"test.parquet\"))\n    print(\"\\nDataset generation complete!\")\n\n    print(\"\\n--- Verification ---\")\n    print(\"Train dataset data sources:\", train_dataset.unique(\"data_source\"))\n    print(\"Test dataset data sources:\", test_dataset.unique(\"data_source\"))\n    print(\"Train dataset length:\", len(train_dataset))\n    print(\"Test dataset length:\", len(test_dataset))\n\n\ndef resolve_task_suites(task_suite_name: str) -> list[str]:\n    benchmark_dict = get_benchmark_dict()\n    available_suites = sorted(\n        [name for name in benchmark_dict.keys() if name.startswith(\"libero_\") and name != \"libero_100\"]\n    )\n\n    requested = task_suite_name.strip().lower()\n    if requested == \"all\":\n        return available_suites\n\n    suites = [suite.strip().lower() for suite in requested.split(\",\") if suite.strip()]\n    invalid = [suite for suite in suites if suite not in available_suites]\n    if invalid:\n        raise ValueError(\n            f\"Unknown task_suite_name: {invalid}. Available suites: {available_suites}. \"\n            \"You can also pass --task_suite_name all\"\n        )\n    return suites\n\n\nif __name__ == \"__main__\":\n    parser = argparse.ArgumentParser()\n    parser.add_argument(\n        \"--task_suite_name\",\n        default=\"libero_spatial\",\n        help=\"Task suite name. Support single suite (e.g. libero_spatial), multiple suites split by comma, or all.\",\n    )\n    parser.add_argument(\n        \"--local_save_dir\", default=\"~/data/libero_rl\", help=\"The save directory for the preprocessed dataset.\"\n    )\n    args = parser.parse_args()\n\n    random.seed(42)\n    np.random.seed(42)\n    local_save_dir = os.path.expanduser(args.local_save_dir)\n    os.makedirs(local_save_dir, exist_ok=True)\n    task_suites = resolve_task_suites(args.task_suite_name)\n    print(f\"Will process task suites: {task_suites}\")\n\n    for task_suite_name in task_suites:\n        build_dataset_for_suite(task_suite_name=task_suite_name, local_save_dir=local_save_dir)\n"
  },
  {
    "path": "verl/experimental/vla/requirements_vla.txt",
    "content": "# libero\ntimm<1.0.0 \nimageio\ndraccus==0.8.0\neinops\nhuggingface_hub\njson-numpy\njsonlines\nmatplotlib\nrich\nsentencepiece==0.1.99\n# dlimp @ git+https://github.com/moojink/dlimp_openvla\ndiffusers==0.30.3\nimageio\nuvicorn\nfastapi\njson-numpy\nwandb==0.19.11\nprotobuf==3.20.3"
  },
  {
    "path": "verl/experimental/vla/rob_ray_trainer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPPO Trainer with Ray-based single controller.\nThis trainer supports model-agonistic model initialization with huggingface\n\"\"\"\n\nimport asyncio\nimport uuid\nfrom collections import defaultdict\nfrom pprint import pprint\n\nimport numpy as np\nimport torch\nfrom omegaconf import OmegaConf\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.experimental.dataset.sampler import AbstractCurriculumSampler\nfrom verl.protocol import pad_dataproto_to_divisor, unpad_dataproto\nfrom verl.single_controller.ray import RayClassWithInitArgs\nfrom verl.single_controller.ray.base import create_colocated_worker_cls\nfrom verl.trainer.ppo.core_algos import agg_loss\nfrom verl.trainer.ppo.metric_utils import (\n    compute_data_metrics,\n    compute_throughout_metrics,\n    compute_timing_metrics,\n    process_validation_metrics,\n)\nfrom verl.trainer.ppo.ray_trainer import RayPPOTrainer, apply_kl_penalty, compute_advantage\nfrom verl.trainer.ppo.reward import compute_reward\nfrom verl.trainer.ppo.utils import Role\nfrom verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.metric import reduce_metrics\n\n\ndef compute_response_mask(config, data: DataProto) -> torch.Tensor:\n    \"\"\"Compute the attention mask for the response part of the sequence.\n\n    This function extracts the portion of the attention mask that corresponds to the model's response,\n    which is used for masking computations that should only apply to response tokens.\n\n    Args:\n        data (DataProto): The data containing batched model outputs and inputs.\n\n    Returns:\n        torch.Tensor: The attention mask for the response tokens.\n    \"\"\"\n    complete = data.batch[\"complete\"]  # shape: [batch_size, num_steps, chunk_size]\n\n    complete_traj = complete.view(complete.shape[0], -1)  # # shape: [batch_size, num_steps * chunk_size]\n    batch_size, action_steps = complete_traj.shape\n\n    step_indices = torch.arange(action_steps, device=complete.device).unsqueeze(0).expand(batch_size, -1)\n\n    first_true_idx_approx = torch.argmax(complete_traj.long(), dim=1)\n\n    has_any_true = complete_traj.any(dim=1)\n\n    final_first_true_idx = torch.where(\n        has_any_true, first_true_idx_approx, torch.tensor(action_steps - 1, device=complete.device)\n    )\n\n    mask_traj = step_indices <= final_first_true_idx.unsqueeze(1)\n\n    mask = mask_traj.view(complete.shape)  # shape: [batch_size, num_steps, chunk_size]\n    mask = mask.repeat_interleave(config.env.actor.model.action_dim, dim=-1)  # eapand to action dim\n    return mask\n\n\ndef flatten_trajectories(data: DataProto) -> DataProto:\n    batch_size, num_steps = data.batch[\"action\"].shape[:2]\n    new_batch_fields = {}\n    for key, tensor in data.batch.items():\n        if len(tensor.shape) >= 2 and tensor.shape[0] == batch_size and tensor.shape[1] == num_steps:\n            # (B, S, H, W) -> (B*S, H, W)\n            new_shape = (batch_size * num_steps, *tensor.shape[2:])\n            new_batch_fields[key] = tensor.reshape(new_shape)\n        elif len(tensor.shape) == 1 and tensor.shape[0] == batch_size:\n            # [e1, e2] -> [e1, e1, ..., e2, e2, ...] (S times each)\n            new_batch_fields[key] = tensor.repeat_interleave(num_steps)\n        else:\n            new_batch_fields[key] = tensor\n    new_data = DataProto.from_dict(tensors=new_batch_fields, meta_info=data.meta_info)\n    return new_data\n\n\n# def filter_by_acc(data: DataProto, accuracy_lower_bound, accuracy_upper_bound) -> torch.Tensor:\n\n\nclass RobRayPPOTrainer(RayPPOTrainer):\n    \"\"\"Distributed PPO trainer using Ray for scalable reinforcement learning.\n\n    This trainer orchestrates distributed PPO training across multiple nodes and GPUs,\n    managing actor rollouts, critic training, and reward computation with Ray backend.\n    Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration.\n    \"\"\"\n\n    def _start_profiling(self, do_profile: bool) -> None:\n        \"\"\"Start profiling for all worker groups including env workers.\"\"\"\n        super()._start_profiling(do_profile)\n        if do_profile and hasattr(self, \"env_wg\"):\n            self.env_wg.start_profile(role=\"env\", profile_step=self.global_steps)\n\n    def _stop_profiling(self, do_profile: bool) -> None:\n        \"\"\"Stop profiling for all worker groups including env workers.\"\"\"\n        super()._stop_profiling(do_profile)\n        if do_profile and hasattr(self, \"env_wg\"):\n            self.env_wg.stop_profile()\n\n    def init_workers(self):\n        self.resource_pool_manager.create_resource_pool()\n\n        if self.config.env.disagg_sim.enable:\n            # pin EnvWorker to Simulator GPU nodes\n            self.resource_pool_manager.get_resource_pool(Role.Env).accelerator_type = \"sim\"\n            self.resource_pool_manager.get_resource_pool(Role.ActorRollout).accelerator_type = \"train_rollout\"\n\n        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}\n        resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)\n        actor_rollout_cls = RayClassWithInitArgs(\n            cls=self.role_worker_mapping[Role.ActorRollout],\n            config=self.config.actor_rollout_ref,\n            role=\"actor_rollout\",\n        )\n        self.resource_pool_to_cls[resource_pool][\"actor_rollout\"] = actor_rollout_cls\n\n        assert Role.Env in self.role_worker_mapping\n        if Role.Env in self.role_worker_mapping:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Env)\n            env_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.Env], config=self.config.env)\n            self.resource_pool_to_cls[resource_pool][\"env\"] = env_cls\n\n        # initialize WorkerGroup\n        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,\n        # you should not use `create_colocated_worker_cls`.\n        # Instead, directly pass different resource pool to different worker groups.\n        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.\n        all_wg = {}\n        wg_kwargs = {}  # Setting up kwargs for RayWorkerGroup\n        if OmegaConf.select(self.config.trainer, \"ray_wait_register_center_timeout\") is not None:\n            wg_kwargs[\"ray_wait_register_center_timeout\"] = self.config.trainer.ray_wait_register_center_timeout\n        if OmegaConf.select(self.config.global_profiler, \"steps\") is not None:\n            wg_kwargs[\"profile_steps\"] = OmegaConf.select(self.config.global_profiler, \"steps\")\n            # Only require nsight worker options when tool is nsys\n            if OmegaConf.select(self.config.global_profiler, \"tool\") == \"nsys\":\n                assert (\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                    is not None\n                ), \"worker_nsight_options must be set when using nsys with profile_steps\"\n                wg_kwargs[\"worker_nsight_options\"] = OmegaConf.to_container(\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                )\n        wg_kwargs[\"device_name\"] = self.device_name\n\n        for resource_pool, class_dict in self.resource_pool_to_cls.items():\n            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n            wg_dict = self.ray_worker_group_cls(\n                resource_pool=resource_pool,\n                ray_cls_with_init=worker_dict_cls,\n                **wg_kwargs,\n            )\n            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n            all_wg.update(spawn_wg)\n\n        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory\n        self.actor_rollout_wg = all_wg[\"actor_rollout\"]\n        self.actor_rollout_wg.init_model()\n        self.env_wg = all_wg[\"env\"]\n\n        # create async rollout manager and request scheduler\n        self.async_rollout_mode = False\n        if self.config.actor_rollout_ref.rollout.mode == \"async_envloop\":\n            from verl.experimental.vla.env_loop import EnvLoop\n\n            self.async_rollout_mode = True\n            self.async_rollout_manager = EnvLoop(\n                config=self.config, rollout_wg=self.actor_rollout_wg, env_wg=self.env_wg\n            )\n\n    def _get_gen_batch(self, batch: DataProto) -> DataProto:\n        # pop those keys for generation\n        batch_keys_to_pop = []\n        non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys())\n        gen_batch = batch.pop(\n            batch_keys=batch_keys_to_pop,\n            non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),\n        )\n\n        return gen_batch\n\n    def _reset_envs(self, gen_batch: DataProto) -> asyncio.Future:\n        initial_state_ids = gen_batch.non_tensor_batch[\"state_ids\"]\n        task_ids = gen_batch.non_tensor_batch[\"task_ids\"]\n        reset_prompts = DataProto.from_dict(non_tensors={\"state_ids\": initial_state_ids, \"task_ids\": task_ids})\n        reset_future = self.env_wg.reset_envs_to_state_ids(reset_prompts)\n        return reset_future\n\n    def fit(self):\n        \"\"\"\n        The training loop of PPO.\n        The driver process only need to call the compute functions of the worker group through RPC\n        to construct the PPO dataflow.\n        The light-weight advantage computation is done on the driver process.\n        \"\"\"\n        from omegaconf import OmegaConf\n\n        from verl.utils.tracking import Tracking\n\n        logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        self.global_steps = 0\n\n        # load checkpoint before doing anything\n        self._load_checkpoint()\n\n        # perform validation before training\n        # currently, we only support validation using the reward_function.\n        if self.val_reward_fn is not None and self.config.trainer.get(\"val_before_train\", True):\n            val_metrics = self._validate()\n            assert val_metrics, f\"{val_metrics=}\"\n            pprint(f\"Initial validation metrics: {val_metrics}\")\n            logger.log(data=val_metrics, step=self.global_steps)\n            if self.config.trainer.get(\"val_only\", False):\n                return\n\n        # add tqdm\n        progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc=\"Training Progress\")\n\n        # we start from step 1\n        self.global_steps += 1\n        last_val_metrics = None\n        self.max_steps_duration = 0\n\n        prev_step_profile = False\n        curr_step_profile = (\n            self.global_steps in self.config.global_profiler.steps\n            if self.config.global_profiler.steps is not None\n            else False\n        )\n        next_step_profile = False\n\n        for epoch in range(self.config.trainer.total_epochs):\n            train_iter = iter(self.train_dataloader)\n            next_batch_dict = next(train_iter)\n            need_validate = False\n            dataloader_len = len(self.train_dataloader)\n            print(f\"Starting epoch {epoch}, dataloader length: {dataloader_len}\")\n            for step_idx in range(dataloader_len):\n                batch_dict = next_batch_dict\n                try:\n                    next_batch_dict = next(train_iter)\n                except StopIteration:\n                    next_batch_dict = None\n\n                metrics = {}\n                timing_raw = {}\n\n                with marked_timer(\"start_profile\", timing_raw):\n                    self._start_profiling(\n                        not prev_step_profile and curr_step_profile\n                        if self.config.global_profiler.profile_continuous_steps\n                        else curr_step_profile\n                    )\n\n                batch: DataProto = DataProto.from_single_dict(batch_dict)\n\n                # add uid to batch\n                batch.non_tensor_batch[\"uid\"] = np.array([str(uuid.uuid4()) for _ in range(len(batch))], dtype=object)\n\n                gen_batch = self._get_gen_batch(batch)\n\n                # pass global_steps to trace\n                gen_batch.meta_info[\"global_steps\"] = self.global_steps\n                # pass generation config to gen_batch\n                gen_batch.meta_info[\"do_sample\"] = True\n                gen_batch.meta_info[\"temperature\"] = self.config.actor_rollout_ref.rollout.temperature\n                gen_batch.meta_info[\"prompt_length\"] = self.config.actor_rollout_ref.rollout.prompt_length\n                gen_batch.meta_info[\"eos_token_id\"] = self.tokenizer.eos_token_id\n                gen_batch.meta_info[\"n_samples\"] = self.config.actor_rollout_ref.rollout.n\n                gen_batch.meta_info[\"pad_token_id\"] = self.tokenizer.pad_token_id\n\n                gen_batch = gen_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n\n                is_last_step = self.global_steps >= self.total_training_steps\n\n                if step_idx == 0 or need_validate:\n                    # reset env workers in first step\n                    # if validation on last step, the reset was not executed and need to be done here\n                    reset_future = self._reset_envs(gen_batch)\n\n                need_validate = (\n                    self.val_reward_fn is not None\n                    and self.config.trainer.test_freq > 0\n                    and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)\n                )\n\n                with marked_timer(\"step\", timing_raw):\n                    # generate a batch\n                    with marked_timer(\"gen\", timing_raw, color=\"red\"):\n                        gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch, reset_future)\n\n                    # prepare for next batch's env reset\n                    if step_idx != dataloader_len - 1 and not need_validate:\n                        next_batch: DataProto = DataProto.from_single_dict(next_batch_dict)\n                        next_gen_batch = self._get_gen_batch(next_batch)\n                        next_gen_batch = next_gen_batch.repeat(\n                            repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True\n                        )\n                        reset_future = self._reset_envs(next_gen_batch)\n\n                    # repeat to align with repeated responses in rollout\n                    batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n                    batch = gen_batch_output\n\n                    if \"response_mask\" not in batch.batch.keys():\n                        batch.batch[\"response_mask\"] = compute_response_mask(self.config, batch)\n\n                    with marked_timer(\"reward\", timing_raw, color=\"yellow\"):\n                        # compute reward model score\n                        reward_tensor, reward_extra_infos_dict = compute_reward(batch, self.reward_fn)\n\n                    batch.batch[\"reward_tensor\"] = reward_tensor\n                    batch = flatten_trajectories(batch)\n\n                    batch.meta_info[\"global_token_num\"] = torch.sum(batch.batch[\"attention_mask\"], dim=-1).tolist()\n\n                    # recompute old_log_probs\n                    with marked_timer(\"old_log_prob\", timing_raw, color=\"blue\"):\n                        old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)\n                        entropys = old_log_prob.batch[\"entropys\"]\n                        response_masks = batch.batch[\"response_mask\"]\n                        actor_config = self.config.actor_rollout_ref.actor\n                        entropy_agg = agg_loss(\n                            loss_mat=entropys,\n                            loss_mask=response_masks,\n                            loss_agg_mode=actor_config.loss_agg_mode,\n                            loss_scale_factor=actor_config.loss_scale_factor,\n                        )\n                        old_log_prob_metrics = {\"actor/entropy\": entropy_agg.detach().item()}\n                        metrics.update(old_log_prob_metrics)\n                        old_log_prob.batch.pop(\"entropys\")\n                        batch = batch.union(old_log_prob)\n\n                        if \"rollout_log_probs\" in batch.batch.keys():\n                            # TODO: we may want to add diff of probs too.\n                            from verl.utils.debug.metrics import calculate_debug_metrics\n\n                            metrics.update(calculate_debug_metrics(batch))\n\n                    if self.use_reference_policy:\n                        # compute reference log_prob\n                        with marked_timer(\"ref\", timing_raw, color=\"olive\"):\n                            if not self.ref_in_actor:\n                                ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)\n                            else:\n                                ref_log_prob = self.actor_rollout_wg.compute_ref_log_prob(batch)\n                            batch = batch.union(ref_log_prob)\n\n                    # compute values\n                    if self.use_critic:\n                        with marked_timer(\"values\", timing_raw, color=\"cyan\"):\n                            values = self.critic_wg.compute_values(batch)\n                            batch = batch.union(values)\n\n                    with marked_timer(\"adv\", timing_raw, color=\"brown\"):\n                        # we combine with rule-based rm\n                        reward_extra_infos_dict: dict[str, list] = None\n\n                        token_level_scores = torch.zeros_like(response_masks, dtype=torch.float32)\n                        flipped_mask = response_masks.flip(dims=[1])\n                        indices_in_flipped = torch.argmax(flipped_mask.long(), dim=1)\n\n                        last_true_indices = response_masks.shape[-1] - 1 - indices_in_flipped\n                        rows_with_response = response_masks.any(dim=1)\n                        effective_rewards = batch.batch[\"reward_tensor\"] * rows_with_response.to(\n                            batch.batch[\"reward_tensor\"].dtype\n                        )\n                        row_indices = torch.arange(response_masks.shape[0], device=token_level_scores.device)\n\n                        token_level_scores[row_indices, last_true_indices] = effective_rewards\n                        batch.batch[\"token_level_scores\"] = token_level_scores\n                        if reward_extra_infos_dict:\n                            batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})\n\n                        # compute rewards. apply_kl_penalty if available\n                        if self.config.algorithm.use_kl_in_reward:\n                            batch, kl_metrics = apply_kl_penalty(\n                                batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty\n                            )\n                            metrics.update(kl_metrics)\n                        else:\n                            batch.batch[\"token_level_rewards\"] = batch.batch[\"token_level_scores\"]\n\n                        # compute advantages, executed on the driver process\n\n                        norm_adv_by_std_in_grpo = self.config.algorithm.get(\n                            \"norm_adv_by_std_in_grpo\", True\n                        )  # GRPO adv normalization factor\n\n                        batch = compute_advantage(\n                            batch,\n                            adv_estimator=self.config.algorithm.adv_estimator,\n                            gamma=self.config.algorithm.gamma,\n                            lam=self.config.algorithm.lam,\n                            num_repeat=self.config.actor_rollout_ref.rollout.n,\n                            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n                            config=self.config.algorithm,\n                        )\n\n                    # update critic\n                    if self.use_critic:\n                        with marked_timer(\"update_critic\", timing_raw, color=\"pink\"):\n                            critic_output = self.critic_wg.update_critic(batch)\n                        critic_output_metrics = reduce_metrics(critic_output.meta_info[\"metrics\"])\n                        metrics.update(critic_output_metrics)\n\n                    # implement critic warmup\n                    if self.config.trainer.critic_warmup <= self.global_steps:\n                        # update actor\n                        with marked_timer(\"update_actor\", timing_raw, color=\"red\"):\n                            batch.meta_info[\"multi_turn\"] = self.config.actor_rollout_ref.rollout.multi_turn.enable\n                            actor_output = self.actor_rollout_wg.update_actor(batch)\n                        actor_output_metrics = reduce_metrics(actor_output.meta_info[\"metrics\"])\n                        metrics.update(actor_output_metrics)\n\n                    # Log rollout generations if enabled\n                    rollout_data_dir = self.config.trainer.get(\"rollout_data_dir\", None)\n                    if rollout_data_dir:\n                        with marked_timer(\"dump_rollout_generations\", timing_raw, color=\"green\"):\n                            inputs = self.tokenizer.batch_decode(batch.batch[\"prompts\"], skip_special_tokens=True)\n                            outputs = self.tokenizer.batch_decode(batch.batch[\"responses\"], skip_special_tokens=True)\n                            scores = batch.batch[\"token_level_scores\"].sum(-1).cpu().tolist()\n                            sample_gts = [\n                                item.non_tensor_batch.get(\"reward_model\", {}).get(\"ground_truth\", None)\n                                for item in batch\n                            ]\n\n                            if \"request_id\" in batch.non_tensor_batch:\n                                reward_extra_infos_dict.setdefault(\n                                    \"request_id\",\n                                    batch.non_tensor_batch[\"request_id\"].tolist(),\n                                )\n\n                            self._dump_generations(\n                                inputs=inputs,\n                                outputs=outputs,\n                                gts=sample_gts,\n                                scores=scores,\n                                reward_extra_infos_dict=reward_extra_infos_dict,\n                                dump_path=rollout_data_dir,\n                            )\n\n                # validate\n                if need_validate:\n                    with marked_timer(\"testing\", timing_raw, color=\"green\"):\n                        val_metrics: dict = self._validate()\n                        if is_last_step:\n                            last_val_metrics = val_metrics\n                    metrics.update(val_metrics)\n\n                # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.\n                esi_close_to_expiration = should_save_ckpt_esi(\n                    max_steps_duration=self.max_steps_duration,\n                    redundant_time=self.config.trainer.esi_redundant_time,\n                )\n                # Check if the conditions for saving a checkpoint are met.\n                # The conditions include a mandatory condition (1) and\n                # one of the following optional conditions (2/3/4):\n                # 1. The save frequency is set to a positive value.\n                # 2. It's the last training step.\n                # 3. The current step number is a multiple of the save frequency.\n                # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.\n                if self.config.trainer.save_freq > 0 and (\n                    is_last_step or self.global_steps % self.config.trainer.save_freq == 0 or esi_close_to_expiration\n                ):\n                    if esi_close_to_expiration:\n                        print(\"Force saving checkpoint: ESI instance expiration approaching.\")\n                    with marked_timer(\"save_checkpoint\", timing_raw, color=\"green\"):\n                        self._save_checkpoint()\n\n                with marked_timer(\"stop_profile\", timing_raw):\n                    next_step_profile = (\n                        self.global_steps + 1 in self.config.global_profiler.steps\n                        if self.config.global_profiler.steps is not None\n                        else False\n                    )\n                    self._stop_profiling(\n                        curr_step_profile and not next_step_profile\n                        if self.config.global_profiler.profile_continuous_steps\n                        else curr_step_profile\n                    )\n                    prev_step_profile = curr_step_profile\n                    curr_step_profile = next_step_profile\n\n                steps_duration = timing_raw[\"step\"]\n                self.max_steps_duration = max(self.max_steps_duration, steps_duration)\n\n                # training metrics\n                metrics.update(\n                    {\n                        \"training/global_step\": self.global_steps,\n                        \"training/epoch\": epoch,\n                    }\n                )\n                # collect metrics\n                metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))\n                metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))\n                # TODO: implement actual tflpo and theoretical tflpo\n                n_gpus = self.resource_pool_manager.get_n_gpus()\n                metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))\n\n                # this is experimental and may be changed/removed in the future in favor of a general-purpose one\n                if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):\n                    self.train_dataloader.sampler.update(batch=batch)\n\n                # TODO: make a canonical logger that supports various backend\n                logger.log(data=metrics, step=self.global_steps)\n\n                progress_bar.update(1)\n                self.global_steps += 1\n\n                if (\n                    hasattr(self.config.actor_rollout_ref.actor, \"profiler\")\n                    and self.config.actor_rollout_ref.actor.profiler.tool == \"torch_memory\"\n                ):\n                    self.actor_rollout_wg.dump_memory_snapshot(\n                        tag=f\"post_update_step{self.global_steps}\", sub_dir=f\"step{self.global_steps}\"\n                    )\n\n                if is_last_step:\n                    pprint(f\"Final validation metrics: {last_val_metrics}\")\n                    progress_bar.close()\n                    return\n\n                # this is experimental and may be changed/removed in the future\n                # in favor of a general-purpose data buffer pool\n                if hasattr(self.train_dataset, \"on_batch_end\"):\n                    # The dataset may be changed after each training batch\n                    self.train_dataset.on_batch_end(batch=batch)\n\n    def _validate(self):\n        data_source_lst = []\n        reward_extra_infos_dict: dict[str, list] = defaultdict(list)\n\n        # Lists to collect samples for the table\n        sample_scores = []\n        sample_turns = []\n        sample_uids = []\n\n        for test_data in self.val_dataloader:\n            test_batch = DataProto.from_single_dict(test_data)\n            if len(test_batch) < self.config.data.val_batch_size:\n                print(f\"drop last batch in val_dataloader, len {len(test_batch)}\")\n                break\n\n            if \"uid\" not in test_batch.non_tensor_batch:\n                test_batch.non_tensor_batch[\"uid\"] = np.array(\n                    [str(uuid.uuid4()) for _ in range(len(test_batch))], dtype=object\n                )\n\n            test_gen_batch = self._get_gen_batch(test_batch)\n            test_gen_batch.meta_info = {\n                \"eos_token_id\": self.tokenizer.eos_token_id,\n                \"pad_token_id\": self.tokenizer.pad_token_id,\n                \"prompt_length\": self.config.actor_rollout_ref.rollout.prompt_length,\n                \"recompute_log_prob\": False,\n                \"do_sample\": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,\n                \"temperature\": self.config.actor_rollout_ref.rollout.temperature,\n                \"n_samples\": self.config.actor_rollout_ref.rollout.n,\n                \"validate\": True,\n                \"global_steps\": self.global_steps,\n            }\n\n            test_gen_batch = test_gen_batch.repeat(\n                repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True\n            )\n\n            sample_uids.extend(test_gen_batch.non_tensor_batch[\"uid\"])\n\n            # pad to be divisible by dp_size\n            size_divisor = self.config.env.train.num_envs * self.config.env.rollout.pipeline_stage_num\n            test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor)\n            reset_future = self._reset_envs(test_gen_batch_padded)\n            test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(\n                test_gen_batch_padded, reset_future\n            )\n\n            # unpad\n            test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)\n\n            print(\"validation generation end\")\n\n            test_batch = test_output_gen_batch\n            test_batch.meta_info[\"validate\"] = True\n\n            # evaluate using reward_function\n            if self.val_reward_fn is None:\n                raise ValueError(\"val_reward_fn must be provided for validation.\")\n            result = self.val_reward_fn(test_batch, return_dict=True)\n            reward_tensor = result[\"reward_tensor\"]\n            scores = reward_tensor.sum(-1).cpu().tolist()\n            sample_scores.extend(scores)\n\n            reward_extra_infos_dict[\"reward\"].extend(scores)\n            print(f\"len reward_extra_infos_dict['reward']: {len(reward_extra_infos_dict['reward'])}\")\n            if \"reward_extra_info\" in result:\n                for key, lst in result[\"reward_extra_info\"].items():\n                    reward_extra_infos_dict[key].extend(lst)\n                    print(f\"len reward_extra_infos_dict['{key}']: {len(reward_extra_infos_dict[key])}\")\n\n            # collect num_turns of each prompt\n            if \"__num_turns__\" in test_batch.non_tensor_batch:\n                sample_turns.append(test_batch.non_tensor_batch[\"__num_turns__\"])\n\n            data_source_lst.append(test_batch.non_tensor_batch.get(\"data_source\", [\"unknown\"] * reward_tensor.shape[0]))\n\n        for key_info, lst in reward_extra_infos_dict.items():\n            assert len(lst) == 0 or len(lst) == len(sample_scores), f\"{key_info}: {len(lst)=}, {len(sample_scores)=}\"\n\n        data_sources = np.concatenate(data_source_lst, axis=0)\n\n        data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict)\n        metric_dict = {}\n        for data_source, var2metric2val in data_src2var2metric2val.items():\n            core_var = \"acc\" if \"acc\" in var2metric2val else \"reward\"\n            for var_name, metric2val in var2metric2val.items():\n                n_max = max([int(name.split(\"@\")[-1].split(\"/\")[0]) for name in metric2val.keys()])\n                for metric_name, metric_val in metric2val.items():\n                    if (\n                        (var_name == core_var)\n                        and any(metric_name.startswith(pfx) for pfx in [\"mean\", \"maj\", \"best\"])\n                        and (f\"@{n_max}\" in metric_name)\n                    ):\n                        metric_sec = \"val-core\"\n                    else:\n                        metric_sec = \"val-aux\"\n                    pfx = f\"{metric_sec}/{data_source}/{var_name}/{metric_name}\"\n                    metric_dict[pfx] = metric_val\n\n        if len(sample_turns) > 0:\n            sample_turns = np.concatenate(sample_turns)\n            metric_dict[\"val-aux/num_turns/min\"] = sample_turns.min()\n            metric_dict[\"val-aux/num_turns/max\"] = sample_turns.max()\n            metric_dict[\"val-aux/num_turns/mean\"] = sample_turns.mean()\n\n        return metric_dict\n"
  },
  {
    "path": "verl/experimental/vla/run_pi05_libero_sac.sh",
    "content": "set -x\nlibero_train_path=$HOME/data/libero_rl/train.parquet\nlibero_test_path=$HOME/data/libero_rl/test.parquet\n\ntrain_files=$libero_train_path\ntest_files=$libero_test_path\n\nOUTPUT_DIR=${MLP_MODEL_OUTPUT:-\"$HOME/models/vla_libero_grpo\"}\nVIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-\"$HOME\"}/video\nSFT_MODEL_PATH=${SFT_MODEL_PATH:-\"$HOME/data/pi05_libero_torch\"}\nTOKENIZER_PATH=\"$SFT_MODEL_PATH\"\n\n# Physical Node Config\nNUM_NODES=1                                    # number of nodes\nNUM_GPUS=4                                     # total number of gpus per node\n\n# Role Config\nNUM_ENV_GPUS=2                                 # number of gpus for env workers per node\nNUM_ROLLOUT_GPUS=$((NUM_GPUS - NUM_ENV_GPUS))  # number of gpus for rollout workers per node\n\n# Rollout Config\n# NOTE: TRAIN_BATCH_SIZE * ROLLOUT_N == NUM_ENV_GPUS * NUM_STAGE * NUM_ENV\nTRAIN_BATCH_SIZE=32                            # batch size for dataloaders per step\nROLLOUT_N=1                                    # response number for each prompt (for GRPO)\nNUM_STAGE=2                                    # number of pipeline stages\nNUM_ENV=8                                      # number of envs per env worker\n\nNUM_ACTION_CHUNKS=10                           # number of action chunks\nMAX_EPISODE_STEPS=512                          # max episode steps for each env\n                                               # max_interactions = MAX_EPISODE_STEPS / num_action_chunks\n\n# Training Config\nMINI_BATCH_SIZE=1024                           # mini batch size (batch size per GPU, automatically multiplied by ROLLOUT_N)\nMICRO_BATCH_SIZE=8                             # micro batch size (per GPU, for gradient accumulation, should divide MINI_BATCH_SIZE)\n\n\n\n# isaac or libero\n# libero means original libero benchmark with mujoco sim\n# isaac means libero benchmark using isaac sim\nSIM_TYPE=${SIM_TYPE:-\"libero\"}\nPROJECT_NAME=\"vla_libero_RL\"\nEXPERIMENT_NAME=\"${SIM_TYPE}_reinforce_plus_plus\"\n\nISSC_PYTHON=\"/workspace/isaaclab/_isaac_sim/python.sh\"\nPYTHON=python\nif [ -f \"$ISSC_PYTHON\" ]; then\n    PYTHON=$ISSC_PYTHON\nfi\n\n# avoiding warnings\nmkdir /root/LIBERO/libero/libero/../datasets\ngpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -n 1)\n\n# force osmesa in Hopper\nif echo \"$gpu_name\" | grep \"NVIDIA H\"; then\n    echo \"enable MUJOCO_GL=osmesa in Hopper\"\n    export MUJOCO_GL=osmesa\nfi\n\nexport VERL_LOGGING_LEVEL=INFO\n\n$PYTHON -m verl.experimental.vla.main_sac \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=$TRAIN_BATCH_SIZE \\\n    data.val_batch_size=4 \\\n    actor_rollout_ref.rollout.n=$ROLLOUT_N \\\n    env.train.num_envs=$NUM_ENV \\\n    data.max_prompt_length=256 \\\n    data.max_response_length=128 \\\n    env.rollout.pipeline_stage_num=$NUM_STAGE \\\n    env.train.simulator_type=$SIM_TYPE \\\n    env.actor.model.num_action_chunks=$NUM_ACTION_CHUNKS \\\n    env.actor.model.action_dim=7 \\\n    env.train.only_eval=False \\\n    env.train.max_episode_steps=$MAX_EPISODE_STEPS \\\n    env.train.video_cfg.save_video=True \\\n    env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \\\n    env.train.seed=42 \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.actor.fsdp_config.wrap_policy.transformer_layer_cls_to_wrap=[SiglipEncoderLayer,GemmaDecoderLayerWithExpert] \\\n    actor_rollout_ref.model.path=$SFT_MODEL_PATH \\\n    actor_rollout_ref.model.tokenizer_path=$TOKENIZER_PATH \\\n    actor_rollout_ref.rollout.mode=async_envloop \\\n    actor_rollout_ref.actor.optim.lr=5e-6 \\\n    actor_rollout_ref.actor.optim.warmup_style=constant \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$MINI_BATCH_SIZE \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$MICRO_BATCH_SIZE \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.grad_clip=1 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.num_images_in_input=1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.model.trust_remote_code=False \\\n    +actor_rollout_ref.model.override_config.attn_implementation=eager \\\n    actor_rollout_ref.actor.entropy_coeff=0. \\\n    actor_rollout_ref.rollout.temperature=1.6 \\\n    actor_rollout_ref.rollout.prompt_length=512 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=hf \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    +actor_rollout_ref.algorithm='sac' \\\n    algorithm.kl_ctrl.kl_coef=0.00 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=$PROJECT_NAME \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.default_local_dir=$OUTPUT_DIR \\\n    trainer.n_gpus_per_node=$NUM_GPUS \\\n    +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \\\n    +trainer.n_rollout_gpus_per_node=$NUM_ROLLOUT_GPUS \\\n    +trainer.rollout_interval=20 \\\n    trainer.nnodes=$NUM_NODES \\\n    trainer.save_freq=300 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=100 \\\n    trainer.val_only=False \\\n    trainer.val_before_train=False"
  },
  {
    "path": "verl/experimental/vla/run_pi05_libero_sac_disagg.sh",
    "content": "set -x\nlibero_train_path=$HOME/data/libero_rl/libero_spatial/train.parquet\nlibero_test_path=$HOME/data/libero_rl/libero_spatial/test.parquet\n\ntrain_files=$libero_train_path\ntest_files=$libero_test_path\n\nOUTPUT_DIR=${MLP_MODEL_OUTPUT:-\"$HOME/models/vla_libero_grpo\"}\nVIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-\"$HOME\"}/video\nSFT_MODEL_PATH=${SFT_MODEL_PATH:-\"$HOME/data/pi05_libero_torch\"}\nTOKENIZER_PATH=\"$SFT_MODEL_PATH\"\n\n# Physical Node Config\nNUM_GPUS=8                                     # total number of gpus per node\n\n# Role Config\nNUM_NODES=1                                    # number of nodes for rollout\nSIM_NODES=1                                    # number of nodes for sim                \nNUM_ENV_GPUS=4                                 # number of gpus for env workers per node\nNUM_ROLLOUT_GPUS=8                             # number of gpus for rollout workers per node\n\n# Rollout Config\n# NOTE: TRAIN_BATCH_SIZE * ROLLOUT_N == NUM_ENV_GPUS * NUM_STAGE * NUM_ENV\nTRAIN_BATCH_SIZE=64                            # batch size for dataloaders per step\nROLLOUT_N=1                                    # response number for each prompt (for GRPO)\nNUM_STAGE=2                                    # number of pipeline stages\nNUM_ENV=8                                      # number of envs per env worker\n\nNUM_ACTION_CHUNKS=10                           # number of action chunks\nMAX_EPISODE_STEPS=220                          # max episode steps for each env\n                                               # max_interactions = MAX_EPISODE_STEPS / num_action_chunks\n\n# Training Config\nMINI_BATCH_SIZE=1024                           # mini batch size (batch size per GPU, automatically multiplied by ROLLOUT_N)\nMICRO_BATCH_SIZE=16                            # micro batch size (per GPU, for gradient accumulation, should divide MINI_BATCH_SIZE)\n\n\n\n# isaac or libero\n# libero means original libero benchmark with mujoco sim\n# isaac means libero benchmark using isaac sim\nSIM_TYPE=${SIM_TYPE:-\"libero\"}\nPROJECT_NAME=\"pi05-libero-sac\"\nEXPERIMENT_NAME=\"${SIM_TYPE}_reinforce_plus_plus\"\n\nISSC_PYTHON=\"/workspace/isaaclab/_isaac_sim/python.sh\"\nPYTHON=python\nif [ -f \"$ISSC_PYTHON\" ]; then\n    PYTHON=$ISSC_PYTHON\nfi\n\n# avoiding warnings\nmkdir /root/LIBERO/libero/libero/../datasets\ngpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -n 1)\n\n# force osmesa in Hopper\nif echo \"$gpu_name\" | grep \"NVIDIA H\"; then\n    echo \"enable MUJOCO_GL=osmesa in Hopper\"\n    export MUJOCO_GL=osmesa\nfi\n\nexport VERL_LOGGING_LEVEL=INFO\n\n$PYTHON -m verl.experimental.vla.main_sac \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=$TRAIN_BATCH_SIZE \\\n    data.val_batch_size=$TRAIN_BATCH_SIZE \\\n    actor_rollout_ref.rollout.n=$ROLLOUT_N \\\n    env.train.num_envs=$NUM_ENV \\\n    data.max_prompt_length=256 \\\n    data.max_response_length=128 \\\n    env.disagg_sim.enable=True \\\n    env.disagg_sim.nnodes=$SIM_NODES \\\n    env.rollout.pipeline_stage_num=$NUM_STAGE \\\n    env.train.simulator_type=$SIM_TYPE \\\n    env.actor.model.num_action_chunks=$NUM_ACTION_CHUNKS \\\n    env.actor.model.action_dim=7 \\\n    env.train.only_eval=False \\\n    env.train.max_episode_steps=$MAX_EPISODE_STEPS \\\n    env.train.video_cfg.save_video=True \\\n    env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \\\n    env.train.seed=42 \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.actor.fsdp_config.wrap_policy.transformer_layer_cls_to_wrap=[SiglipEncoderLayer,GemmaDecoderLayerWithExpert] \\\n    actor_rollout_ref.model.path=$SFT_MODEL_PATH \\\n    actor_rollout_ref.model.tokenizer_path=$TOKENIZER_PATH \\\n    actor_rollout_ref.rollout.mode=async_envloop \\\n    actor_rollout_ref.actor.optim.lr=5e-6 \\\n    actor_rollout_ref.actor.optim.warmup_style=constant \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=$MINI_BATCH_SIZE \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=$MICRO_BATCH_SIZE \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.strategy=fsdp2 \\\n    critic.strategy=fsdp2 \\\n    actor_rollout_ref.actor.grad_clip=1 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.num_images_in_input=1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.model.trust_remote_code=False \\\n    +actor_rollout_ref.model.override_config.attn_implementation=eager \\\n    actor_rollout_ref.actor.entropy_coeff=0. \\\n    actor_rollout_ref.rollout.temperature=1.6 \\\n    actor_rollout_ref.rollout.prompt_length=512 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=hf \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    +actor_rollout_ref.algorithm='sac' \\\n    algorithm.kl_ctrl.kl_coef=0.00 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=$PROJECT_NAME \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.default_local_dir=$OUTPUT_DIR \\\n    trainer.n_gpus_per_node=$NUM_GPUS \\\n    +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \\\n    +trainer.n_rollout_gpus_per_node=$NUM_ROLLOUT_GPUS \\\n    +trainer.rollout_interval=20 \\\n    trainer.nnodes=$NUM_NODES \\\n    trainer.save_freq=500 \\\n    trainer.test_freq=10 \\\n    trainer.total_epochs=1000 \\\n    trainer.val_only=False \\\n    algorithm.adv_estimator=reinforce_plus_plus \\\n    trainer.val_before_train=False\n"
  },
  {
    "path": "verl/experimental/vla/run_simpleVLA_isaac_disagg.sh",
    "content": "#!/bin/bash\nset -x\n\necho \"remember to set ray param < --resources='{\\\"sim\\\"/\\\"actor_rollout\\\":1}' > if using disagg sim\"\n\nlibero_train_path=$HOME/data/libero_rl/train.parquet\nlibero_test_path=$HOME/data/libero_rl/test.parquet\n\ntrain_files=$libero_train_path\ntest_files=$libero_test_path\n\nOUTPUT_DIR=${MLP_MODEL_OUTPUT:-\"$HOME/models/vla_libero_grpo\"}\nVIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-\"$HOME\"}/video\nSFT_MODEL_PATH=${SFT_MODEL_PATH:-\"$HOME/data/Openvla-oft-SFT-libero10-trajall\"}\n\n# for rollout and train\nNUM_NODES=1\n# for simulator\nSIM_NODES=1\nNUM_ENV_GPUS=8\nNUM_ROLLOUT_GPUS=8\nSTAGE_NUM=2\nBATCH_SIZE=16\n# rollout.n should equal to num_envs for isaac env\nROLLOUT_N=8\n\n# 512 is required for libero benchmark, but you can reduce it in debugging to run faster\nMAX_EPISODE_STEPS=512\n\n# isaac or libero\n# libero means original libero benchmark with mujoco sim\n# isaac means libero benchmark using isaac sim\nSIM_TYPE=${SIM_TYPE:-\"isaac\"}\nPROJECT_NAME=\"vla-disagg-issac\"\nEXPERIMENT_NAME=\"${SIM_TYPE}_rl\"\n\nISSC_PYTHON=\"/workspace/isaaclab/_isaac_sim/python.sh\"\nPYTHON=python\nif [ -f \"$ISSC_PYTHON\" ]; then\n    PYTHON=$ISSC_PYTHON\nfi\n\n# avoiding warnings\nmkdir /root/LIBERO/libero/libero/../datasets\n\n\n$PYTHON -m verl.experimental.vla.main_ppo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=${BATCH_SIZE} \\\n    data.val_batch_size=${BATCH_SIZE} \\\n    actor_rollout_ref.rollout.n=$ROLLOUT_N \\\n    env.train.num_envs=$ROLLOUT_N \\\n    data.max_prompt_length=256 \\\n    data.max_response_length=128 \\\n    env.rollout.pipeline_stage_num=$STAGE_NUM \\\n    env.train.simulator_type=$SIM_TYPE \\\n    env.actor.model.num_action_chunks=8 \\\n    env.actor.model.action_dim=7 \\\n    env.train.only_eval=False \\\n    env.train.max_episode_steps=$MAX_EPISODE_STEPS \\\n    env.train.video_cfg.save_video=True \\\n    env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \\\n    env.train.seed=42 \\\n    env.disagg_sim.enable=True \\\n    env.disagg_sim.nnodes=$SIM_NODES \\\n    +actor_rollout_ref.algorithm='grpo' \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.model.path=$SFT_MODEL_PATH \\\n    actor_rollout_ref.rollout.mode=async_envloop \\\n    actor_rollout_ref.actor.optim.lr=5e-6 \\\n    actor_rollout_ref.actor.optim.warmup_style=constant \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.grad_clip=1 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.num_images_in_input=1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.model.trust_remote_code=False \\\n    actor_rollout_ref.actor.entropy_coeff=0. \\\n    actor_rollout_ref.rollout.temperature=1.6 \\\n    actor_rollout_ref.rollout.prompt_length=512 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=hf \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.00 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=$PROJECT_NAME \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.default_local_dir=$OUTPUT_DIR \\\n    trainer.n_gpus_per_node=$NUM_ROLLOUT_GPUS \\\n    +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \\\n    +trainer.n_rollout_gpus_per_node=$NUM_ROLLOUT_GPUS \\\n    trainer.nnodes=$NUM_NODES \\\n    trainer.save_freq=30 \\\n    trainer.test_freq=-1 \\\n    trainer.total_epochs=20 \\\n    trainer.val_only=False \\\n    trainer.total_training_steps=10000 \\\n    algorithm.adv_estimator=reinforce_plus_plus \\\n    trainer.val_before_train=False  $@\n\n\n"
  },
  {
    "path": "verl/experimental/vla/run_simpleVLA_libero_grpo.sh",
    "content": "set -x\nlibero_train_path=$HOME/data/libero_rl/train.parquet\nlibero_test_path=$HOME/data/libero_rl/test.parquet\n\n\ntrain_files=$libero_train_path\ntest_files=$libero_test_path\n\nOUTPUT_DIR=${MLP_MODEL_OUTPUT:-\"$HOME/models/vla_libero_grpo\"}\nVIDEO_OUTPUT=${MLP_MODEL_OUTPUT:-\"$HOME\"}/video\nSFT_MODEL_PATH=${SFT_MODEL_PATH:-\"$HOME/data/Openvla-oft-SFT-libero10-trajall\"}\n\nNUM_NODES=1\nNUM_GPUS=8\nNUM_ENV_GPUS=4\n# rollout.n should equal to num_envs for isaac env\nROLLOUT_N=8\n# isaac or libero\n# libero means original libero benchmark with mujoco sim\n# isaac means libero benchmark using isaac sim\nSIM_TYPE=${SIM_TYPE:-\"isaac\"}\nPROJECT_NAME=\"vla_libero_RL\"\nEXPERIMENT_NAME=\"${SIM_TYPE}_reinforce_plus_plus\"\n\nISSC_PYTHON=\"/workspace/isaaclab/_isaac_sim/python.sh\"\nPYTHON=python\nif [ -f \"$ISSC_PYTHON\" ]; then\n    PYTHON=$ISSC_PYTHON\nfi\n\n# avoiding warnings\nmkdir /root/LIBERO/libero/libero/../datasets\ngpu_name=$(nvidia-smi --query-gpu=name --format=csv,noheader,nounits | head -n 1)\n\n# force osmesa in Hopper\nif echo \"$gpu_name\" | grep \"NVIDIA H\"; then\n    echo \"enable MUJOCO_GL=osmesa in Hopper\"\n    export MUJOCO_GL=osmesa\nfi\n\n\n$PYTHON -m verl.experimental.vla.main_ppo \\\n    data.train_files=\"$train_files\" \\\n    data.val_files=\"$test_files\" \\\n    data.train_batch_size=8 \\\n    data.val_batch_size=8 \\\n    actor_rollout_ref.rollout.n=$ROLLOUT_N \\\n    env.train.num_envs=$ROLLOUT_N \\\n    data.max_prompt_length=256 \\\n    data.max_response_length=128 \\\n    env.rollout.pipeline_stage_num=2 \\\n    env.train.simulator_type=$SIM_TYPE \\\n    env.actor.model.num_action_chunks=8 \\\n    env.actor.model.action_dim=7 \\\n    env.train.only_eval=False \\\n    env.train.max_episode_steps=512 \\\n    env.train.video_cfg.save_video=True \\\n    env.train.video_cfg.video_base_dir=${VIDEO_OUTPUT} \\\n    env.train.seed=42 \\\n    +actor_rollout_ref.algorithm='grpo' \\\n    actor_rollout_ref.actor.fsdp_config.model_dtype=bfloat16 \\\n    actor_rollout_ref.model.path=$SFT_MODEL_PATH \\\n    actor_rollout_ref.rollout.mode=async_envloop \\\n    actor_rollout_ref.actor.optim.lr=5e-6 \\\n    actor_rollout_ref.actor.optim.warmup_style=constant \\\n    actor_rollout_ref.actor.ppo_mini_batch_size=128 \\\n    actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu=2 \\\n    actor_rollout_ref.actor.use_dynamic_bsz=False \\\n    actor_rollout_ref.actor.grad_clip=1 \\\n    actor_rollout_ref.actor.clip_ratio_high=0.28 \\\n    actor_rollout_ref.actor.clip_ratio_low=0.2 \\\n    actor_rollout_ref.actor.num_images_in_input=1 \\\n    actor_rollout_ref.model.enable_gradient_checkpointing=False \\\n    actor_rollout_ref.model.use_remove_padding=False \\\n    actor_rollout_ref.model.trust_remote_code=False \\\n    actor_rollout_ref.actor.entropy_coeff=0. \\\n    actor_rollout_ref.rollout.temperature=1.6 \\\n    actor_rollout_ref.rollout.prompt_length=512 \\\n    actor_rollout_ref.rollout.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.rollout.tensor_model_parallel_size=1 \\\n    actor_rollout_ref.rollout.name=hf \\\n    actor_rollout_ref.rollout.gpu_memory_utilization=0.9 \\\n    actor_rollout_ref.rollout.free_cache_engine=False \\\n    actor_rollout_ref.ref.log_prob_micro_batch_size=16 \\\n    actor_rollout_ref.ref.fsdp_config.param_offload=True \\\n    algorithm.kl_ctrl.kl_coef=0.00 \\\n    trainer.logger=['console'] \\\n    trainer.project_name=$PROJECT_NAME \\\n    trainer.experiment_name=$EXPERIMENT_NAME \\\n    trainer.default_local_dir=$OUTPUT_DIR \\\n    trainer.n_gpus_per_node=$NUM_GPUS \\\n    +trainer.n_env_gpus_per_node=$NUM_ENV_GPUS \\\n    +trainer.n_rollout_gpus_per_node=$((NUM_GPUS - NUM_ENV_GPUS)) \\\n    trainer.nnodes=$NUM_NODES \\\n    trainer.save_freq=30 \\\n    trainer.test_freq=30 \\\n    trainer.total_epochs=20 \\\n    trainer.val_only=False \\\n    trainer.total_training_steps=10000 \\\n    algorithm.adv_estimator=reinforce_plus_plus \\\n    trainer.val_before_train=False  $@\n\n\n"
  },
  {
    "path": "verl/experimental/vla/sac/base.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Literal, Optional\n\nimport torch\n\nfrom verl import DataProto\n\n\nclass SupportSACTraining:\n    \"\"\"\n    Base class for Soft Actor-Critic (SAC).\n\n    Subclasses implement a Policy that can be plugged directly into SAC training.\n    This implementation requires the actor and critic to be integrated within a\n    single model instance, e.g., sharing a backbone with an additional MLP head\n    that outputs critic values (Q/V) alongside the actor's action distribution.\n\n    Note:\n        This class intentionally does NOT inherit from `abc.ABC`.\n        The root model may be wrapped or transformed by FSDP (Fully Sharded\n        Data Parallel), which performs runtime class substitution; using\n        `ABCMeta` can break FSDP's class rewriting mechanism.\n    \"\"\"\n\n    def sac_init(self):\n        raise NotImplementedError(\"Subclasses must implement sac_init method.\")\n\n    def sac_get_critic_parameters(self) -> list[torch.nn.Parameter]:\n        \"\"\"Get the parameters of the critic head for optimization.\n\n        Returns:\n            A list of torch.nn.Parameter objects representing the critic head parameters.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_get_critic_parameters method.\")\n\n    def sac_get_named_actor_parameters(self) -> list[tuple[str, torch.nn.Parameter]]:\n        \"\"\"Get named actor parameters for optimization/EMA updates.\n\n        Returns:\n            A list of (name, parameter) tuples representing actor-side trainable parameters.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_get_named_actor_parameters method.\")\n\n    def sac_forward_critic(\n        self,\n        a: dict[str, torch.Tensor],\n        state_features: Any,\n        *,\n        use_target_network: bool = False,\n        method: Literal[\"cat\", \"min\"] = \"cat\",\n        requires_grad: bool = False,\n    ) -> torch.Tensor:\n        \"\"\"Compute Q-values for given state-action pairs.\n        Args:\n            a: Dictionary of tensors representing actions, with key:\n                - \"full_action\": torch.Tensor of shape (B, action_steps, action_dim)\n            state_features: Any data structure representing the processed state features.\n            use_target_network: Whether to use the target critic network heads.\n            method: Method to combine multiple heads' outputs (\"cat\" or \"min\").\n            requires_grad: Whether to enable gradients for the critic head parameters.\n\n        Returns:\n            q_values: torch.Tensor of shape (B, num_heads) if method is \"cat\",\n                      or (B, 1) if method is \"min\", representing the computed Q-values\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_forward_critic method.\")\n\n    def sac_forward_actor(\n        self,\n        state_features: Any,\n        is_first_micro_batch: bool = False,\n    ) -> tuple[torch.Tensor, Optional[torch.Tensor], dict[str, float]]:\n        \"\"\"Compute actions and their log probabilities from state features.\n\n        Args:\n            state_features: Any data structure representing the processed state features.\n            is_first_micro_batch: Whether the current forward corresponds to the first\n                micro batch of the actor update step.\n\n        Returns:\n            actions: torch.Tensor of shape (B, n_action_steps, action_dim), sampled actions.\n            log_probs: Optional torch.Tensor of shape (B,), log probabilities of sampled actions.\n                Can be None when SAC is configured to train without entropy/log-prob terms.\n            metrics: Scalar metrics produced by actor forward, used by outer trainer for logging.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_forward_actor method.\")\n\n    def sac_forward_state_features(self, s: dict[str, torch.Tensor]) -> Any:\n        \"\"\"Compute state features needed for SAC actor and critic.\n\n        Args:\n            s: Dictionary of tensors representing the states, with keys\n                - \"images\": torch.Tensor of shape (B, n_images, C, H, W)\n                - \"image_masks\": torch.Tensor of shape (B, n_images)\n                - \"lang_tokens\": torch.Tensor of shape (B, L)\n                - \"lang_masks\": torch.Tensor of shape (B, L)\n                - \"states\": torch.Tensor of shape (B, state_dim)\n        Returns:\n            state_features: Any data structure representing the processed state features.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_forward_state_features method.\")\n\n    def bc_loss(\n        self,\n        state_features: Any,\n        actions: dict[str, torch.Tensor],\n        valids: torch.Tensor,\n    ) -> torch.Tensor:\n        \"\"\"Compute behavior cloning loss for actor regularization.\"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement bc_loss method.\")\n\n    def sac_update_target_network(self, tau: float):\n        \"\"\"Update the target network heads using Polyak averaging.\n\n        Args:\n            tau: The interpolation parameter for Polyak averaging.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement sac_update_target_network method.\")\n\n\nclass BaseSACActor(ABC):\n    @abstractmethod\n    def update_policy(self, data: DataProto) -> dict:\n        \"\"\"\n        Update the policy using the provided data batch.\n\n        Args:\n            data: DataProto containing the following entries in `data.batch`:\n                - \"a0.full_action\": Tensor of shape (B, action_steps, action_dim),\n                    representing the current action chunk for each sample.\n                - \"a1.full_action\": Tensor of shape (B, action_steps, action_dim),\n                    representing the next action chunk for each sample.\n                - \"s0.states\": Tensor of shape (B, state_dim),\n                    representing the current environment or agent state.\n                - \"s1.states\": Tensor of shape (B, state_dim),\n                    representing the next environment or agent state.\n                - \"s0.images\": Tensor of shape (B, n_images, C, H, W),\n                    containing current visual observations.\n                - \"s1.images\": Tensor of shape (B, n_images, C, H, W),\n                    containing next-step visual observations.\n                - \"s0.image_masks\": Tensor of shape (B, n_images),\n                    indicating valid images per sample.\n                - \"s1.image_masks\": Tensor of shape (B, n_images),\n                    indicating valid images per sample.\n                - \"s0.lang_tokens\": Tensor of shape (B, max_seq_len),\n                    tokenized language instructions.\n                - \"s1.lang_tokens\": Tensor of shape (B, max_seq_len),\n                    tokenized language instructions for the next step.\n                - \"s0.lang_masks\": Tensor of shape (B, max_seq_len),\n                    attention masks for language tokens.\n                - \"s1.lang_masks\": Tensor of shape (B, max_seq_len),\n                    attention masks for language tokens for the next step.\n                - \"rewards\": Tensor of shape (B,),\n                    chunk-level scalar rewards aligned to the next step.\n                - \"response_mask\": Tensor of shape (B, action_steps),\n                    mask indicating whether each sample has a valid response.\n        \"\"\"\n\n        raise NotImplementedError(\"Subclasses must implement update_policy method.\")\n"
  },
  {
    "path": "verl/experimental/vla/sac/naive_rollout_pi05.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nIn single GPU rollout, the sequences are generated directly by sampling from the model.\nThe output will contain\n1. output_ids\n2. attention_masks (left padding)\n3. eos_masks\n4. log_probs\n\"\"\"\n\nimport logging\nfrom typing import Any\n\nimport torch\n\nfrom verl import DataProto\nfrom verl.experimental.vla.naive_rollout_rob import NaiveRolloutRob\nfrom verl.utils.device import get_device_id, get_device_name\n\nlogger = logging.getLogger(__name__)\n\n__all__ = [\"PI0RolloutRob\"]\n\n\nclass PI0RolloutRob(NaiveRolloutRob):\n    def __init__(\n        self,\n        model_config: dict,\n        module: torch.nn.Module,\n        tokenizer: Any,\n    ):\n        self.model_config = model_config\n        self.module = module\n        self.tokenizer = tokenizer\n\n        from torch.distributed.fsdp import register_fsdp_forward_method\n\n        register_fsdp_forward_method(self.module, \"sample_actions\")\n        register_fsdp_forward_method(self.module, \"sac_forward_state_features\")\n        register_fsdp_forward_method(self.module, \"sac_forward_critic\")\n\n    @torch.no_grad()\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Generate sequences\"\"\"\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            prompts.to(get_device_id())\n            validate = bool(prompts.meta_info.get(\"validate\", False))\n            output, s, a = self.module.sample_actions(\n                prompts,\n                tokenizer=self.tokenizer,\n                validate=validate,\n            )\n            state_features = self.module.sac_forward_state_features(s)\n            critic_value = (\n                self.module.sac_forward_critic(\n                    {\"full_action\": a[\"full_action\"]},\n                    state_features,\n                    use_target_network=False,\n                    method=\"min\",\n                    requires_grad=False,\n                )\n                .detach()\n                .float()\n                .reshape(-1)\n            )\n\n        tensor_batch = {\n            \"action\": output.action,\n            \"full_action\": a[\"full_action\"],\n            \"images\": s[\"images\"],\n            \"image_masks\": s[\"image_masks\"],\n            \"lang_tokens\": s[\"lang_tokens\"],\n            \"lang_masks\": s[\"lang_masks\"],\n            \"states\": s[\"states\"],\n            \"critic_value\": critic_value,\n        }\n\n        ret = DataProto.from_dict(tensor_batch)\n\n        return ret\n"
  },
  {
    "path": "verl/experimental/vla/sac/replay_pool.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom dataclasses import dataclass\nfrom typing import Any, Optional, Sequence\n\nimport torch\nfrom tensordict import TensorDict\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@dataclass\nclass _DualPoolState:\n    positive_pool: Optional[TensorDict] = None\n    negative_pool: Optional[TensorDict] = None\n    positive_size: int = 0\n    negative_size: int = 0\n    positive_position: int = 0\n    negative_position: int = 0\n\n\nclass SACReplayPool:\n    \"\"\"Task-aware SAC Replay Pool.\n\n    For each task_id we maintain two independent pools:\n    - positive pool\n    - negative pool\n\n    `single_pool_capacity` is the size of each single pool.\n    \"\"\"\n\n    def __init__(\n        self,\n        single_pool_capacity: int,\n        pool_device: str = \"cpu\",\n        sample_device: str = \"cpu\",\n    ):\n        self.single_pool_capacity = int(single_pool_capacity)\n        self.pool_device = pool_device\n        self.sample_device = sample_device\n\n        self.task_pools: dict[str, _DualPoolState] = {}\n\n        self.size = 0\n        self.positive_size = 0\n        self.negative_size = 0\n\n        self.rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0\n\n    def add_batch(self, batch: TensorDict, task_ids: Sequence[Any]):\n        \"\"\"Add a batch of samples into task-specific positive/negative pools.\"\"\"\n\n        if batch.batch_size[0] == 0:\n            return\n\n        if len(task_ids) != batch.batch_size[0]:\n            raise ValueError(f\"task_ids length ({len(task_ids)}) must match batch size ({batch.batch_size[0]}).\")\n\n        valid_mask = batch[\"valids\"].to(torch.bool)\n        valid_indices = torch.nonzero(valid_mask, as_tuple=False).squeeze(-1)\n        if valid_indices.numel() == 0:\n            return\n\n        batch = self._index_select_batch(batch, valid_indices.to(batch.device))\n        selected = valid_indices.cpu().tolist()\n        task_ids = [task_ids[i] for i in selected]\n\n        positive_mask = self._extract_positive_mask(batch)\n\n        grouped_indices: dict[str, dict[str, list[int]]] = {}\n        for idx in range(batch.batch_size[0]):\n            task_key = self._normalize_task_id(task_ids[idx])\n            if task_key not in grouped_indices:\n                grouped_indices[task_key] = {\"positive\": [], \"negative\": []}\n            if bool(positive_mask[idx].item()):\n                grouped_indices[task_key][\"positive\"].append(idx)\n            else:\n                grouped_indices[task_key][\"negative\"].append(idx)\n\n        for task_key, groups in grouped_indices.items():\n            pool_state = self._get_or_create_task_pool(task_key, batch)\n\n            if groups[\"positive\"]:\n                positive_idx = torch.tensor(groups[\"positive\"], device=batch.device, dtype=torch.long)\n                positive_batch = self._index_select_batch(batch, positive_idx)\n                self._insert_block_to_pool(pool_state, positive_batch, is_positive_pool=True)\n\n            if groups[\"negative\"]:\n                negative_idx = torch.tensor(groups[\"negative\"], device=batch.device, dtype=torch.long)\n                negative_batch = self._index_select_batch(batch, negative_idx)\n                self._insert_block_to_pool(pool_state, negative_batch, is_positive_pool=False)\n\n        self._refresh_global_stats()\n\n    def sample_batch(\n        self,\n        batch_size: int,\n        positive_sample_ratio: float = 0.5,\n        return_sample_info: bool = False,\n    ) -> TensorDict | tuple[TensorDict, dict]:\n        \"\"\"Sample a batch from all task-specific pools.\"\"\"\n\n        if self.size == 0:\n            raise ValueError(\"Replay pool is empty, unable to sample.\")\n\n        positive_sample_ratio = max(0.0, min(1.0, float(positive_sample_ratio)))\n        target_positive = int(round(batch_size * positive_sample_ratio))\n        target_negative = batch_size - target_positive\n\n        sampled_positive = min(target_positive, self.positive_size)\n        sampled_negative = min(target_negative, self.negative_size)\n\n        deficit = batch_size - sampled_positive - sampled_negative\n        if deficit > 0:\n            remaining_positive = self.positive_size - sampled_positive\n            remaining_negative = self.negative_size - sampled_negative\n\n            if remaining_positive >= remaining_negative:\n                extra_positive = min(deficit, remaining_positive)\n                sampled_positive += extra_positive\n                deficit -= extra_positive\n\n                extra_negative = min(deficit, remaining_negative)\n                sampled_negative += extra_negative\n                deficit -= extra_negative\n            else:\n                extra_negative = min(deficit, remaining_negative)\n                sampled_negative += extra_negative\n                deficit -= extra_negative\n\n                extra_positive = min(deficit, remaining_positive)\n                sampled_positive += extra_positive\n                deficit -= extra_positive\n\n        sampled_parts = []\n        if sampled_positive > 0:\n            sampled_parts.append(self._sample_from_task_pools(sampled_positive, is_positive_pool=True))\n        if sampled_negative > 0:\n            sampled_parts.append(self._sample_from_task_pools(sampled_negative, is_positive_pool=False))\n\n        sampled_count = sampled_positive + sampled_negative\n        if len(sampled_parts) == 1:\n            sampled_batch = sampled_parts[0]\n        else:\n            sampled_batch = TensorDict(\n                {key: torch.cat([part[key] for part in sampled_parts], dim=0) for key in sampled_parts[0].keys()},\n                batch_size=[sampled_count],\n                device=self.sample_device,\n            )\n\n        if sampled_count < batch_size:\n            sampled_batch = self._pad_sampled_batch(sampled_batch, target_batch_size=batch_size)\n        else:\n            sampled_batch = TensorDict(\n                {key: value for key, value in sampled_batch.items()},\n                batch_size=[batch_size],\n                device=self.sample_device,\n            )\n\n        shuffle_idx = torch.randperm(batch_size, device=self.sample_device)\n        sampled_batch = TensorDict(\n            {key: value.index_select(0, shuffle_idx) for key, value in sampled_batch.items()},\n            batch_size=[batch_size],\n            device=self.sample_device,\n        )\n\n        if not return_sample_info:\n            return sampled_batch\n\n        sample_info = {\n            \"actual_positive_sample_ratio\": sampled_positive / max(batch_size, 1),\n            \"positive_size\": self.positive_size,\n            \"negative_size\": self.negative_size,\n            \"task_count\": len(self.task_pools),\n        }\n        return sampled_batch, sample_info\n\n    def insert_and_resample(\n        self,\n        source: TensorDict,\n        task_ids: Sequence[Any],\n    ) -> TensorDict:\n        \"\"\"Insert source into replay pool and sample a batch with the same size.\"\"\"\n\n        self.add_batch(source, task_ids=task_ids)\n        return self.sample_batch(source.batch_size[0])\n\n    def save(self, directory: str):\n        \"\"\"Save the replay pool to a directory.\"\"\"\n\n        os.makedirs(directory, exist_ok=True)\n        filepath = f\"{directory}/sac_replay_pool_rank_{self.rank}.pt\"\n\n        tasks_payload: dict[str, dict[str, Any]] = {}\n        for task_id, pool_state in self.task_pools.items():\n            assert pool_state.positive_pool is not None\n            assert pool_state.negative_pool is not None\n            tasks_payload[task_id] = {\n                \"positive_pool\": pool_state.positive_pool.cpu(),\n                \"negative_pool\": pool_state.negative_pool.cpu(),\n                \"positive_size\": pool_state.positive_size,\n                \"negative_size\": pool_state.negative_size,\n                \"positive_position\": pool_state.positive_position,\n                \"negative_position\": pool_state.negative_position,\n            }\n\n        payload = {\n            \"meta_info\": {\n                \"version\": 3,\n                \"single_pool_capacity\": self.single_pool_capacity,\n                \"pool_device\": self.pool_device,\n                \"sample_device\": self.sample_device,\n                \"size\": self.size,\n                \"positive_size\": self.positive_size,\n                \"negative_size\": self.negative_size,\n                \"task_count\": len(self.task_pools),\n            },\n            \"tasks\": tasks_payload,\n        }\n\n        torch.save(payload, filepath)\n        logger.info(\n            f\"[Rank {self.rank}] Task replay pool saved to {filepath} with \\\n               size={self.size}, tasks={len(self.task_pools)}\"\n        )\n\n    def load(self, directory: str):\n        \"\"\"Load the replay pool from a directory.\"\"\"\n\n        filepath = f\"{directory}/sac_replay_pool_rank_{self.rank}.pt\"\n        if not os.path.exists(filepath):\n            return False\n\n        payload = torch.load(filepath, weights_only=False)\n        meta_info = payload[\"meta_info\"]\n        tasks_payload = payload[\"tasks\"]\n\n        self.single_pool_capacity = int(meta_info[\"single_pool_capacity\"])\n        self.task_pools = {}\n\n        for task_id, task_payload in tasks_payload.items():\n            pool_state = _DualPoolState(\n                positive_pool=task_payload[\"positive_pool\"].to(self.pool_device),\n                negative_pool=task_payload[\"negative_pool\"].to(self.pool_device),\n                positive_size=int(task_payload[\"positive_size\"]),\n                negative_size=int(task_payload[\"negative_size\"]),\n                positive_position=int(task_payload[\"positive_position\"]),\n                negative_position=int(task_payload[\"negative_position\"]),\n            )\n            self.task_pools[task_id] = pool_state\n\n        self._refresh_global_stats()\n        logger.info(\n            f\"[Rank {self.rank}] Task replay pool loaded from {filepath} with \\\n              size={self.size}, tasks={len(self.task_pools)}\"\n        )\n        return True\n\n    @classmethod\n    def from_path(\n        cls,\n        directory: str,\n    ) -> \"SACReplayPool\":\n        \"\"\"Load a replay pool from a file.\"\"\"\n\n        rank = torch.distributed.get_rank() if torch.distributed.is_initialized() else 0\n        filepath = f\"{directory}/sac_replay_pool_rank_{rank}.pt\"\n        payload = torch.load(filepath, weights_only=False)\n        meta_info = payload[\"meta_info\"]\n\n        replay_pool = cls(\n            single_pool_capacity=int(meta_info[\"single_pool_capacity\"]),\n            pool_device=meta_info[\"pool_device\"],\n            sample_device=meta_info[\"sample_device\"],\n        )\n        replay_pool.rank = rank\n\n        loaded = replay_pool.load(directory)\n        if not loaded:\n            raise RuntimeError(f\"Failed to load replay pool from {filepath}.\")\n\n        return replay_pool\n\n    def _insert_block_to_pool(\n        self,\n        pool_state: _DualPoolState,\n        source: TensorDict,\n        is_positive_pool: bool,\n    ):\n        \"\"\"Insert a block of data from source into one task pool.\"\"\"\n\n        source_size = source.batch_size[0]\n        if source_size == 0:\n            return\n\n        length = min(source_size, self.single_pool_capacity)\n        idx = torch.arange(length, device=self.pool_device)\n\n        if is_positive_pool:\n            assert pool_state.positive_pool is not None\n            idx = (pool_state.positive_position + idx) % self.single_pool_capacity\n            for key in source.keys():\n                pool_state.positive_pool[key].index_copy_(0, idx, source[key][:length].to(self.pool_device))\n\n            pool_state.positive_position = (pool_state.positive_position + length) % self.single_pool_capacity\n            pool_state.positive_size = min(pool_state.positive_size + length, self.single_pool_capacity)\n        else:\n            assert pool_state.negative_pool is not None\n            idx = (pool_state.negative_position + idx) % self.single_pool_capacity\n            for key in source.keys():\n                pool_state.negative_pool[key].index_copy_(0, idx, source[key][:length].to(self.pool_device))\n\n            pool_state.negative_position = (pool_state.negative_position + length) % self.single_pool_capacity\n            pool_state.negative_size = min(pool_state.negative_size + length, self.single_pool_capacity)\n\n    def _get_or_create_task_pool(self, task_id: str, sample: TensorDict) -> _DualPoolState:\n        if task_id in self.task_pools:\n            return self.task_pools[task_id]\n\n        logger.info(\n            f\"Initializing replay pools for task_id={task_id} with single_pool_capacity={self.single_pool_capacity}\"\n        )\n        pool_template = TensorDict(\n            {\n                key: torch.zeros(\n                    (self.single_pool_capacity, *value.shape[1:]),\n                    dtype=value.dtype,\n                    device=self.pool_device,\n                )\n                for key, value in sample.items()\n            },\n            batch_size=[self.single_pool_capacity],\n            device=self.pool_device,\n        )\n        pool_state = _DualPoolState(\n            positive_pool=pool_template.clone(),\n            negative_pool=pool_template.clone(),\n            positive_size=0,\n            negative_size=0,\n            positive_position=0,\n            negative_position=0,\n        )\n        self.task_pools[task_id] = pool_state\n        return pool_state\n\n    def _extract_positive_mask(self, batch: TensorDict) -> torch.Tensor:\n        positive_mask = batch[\"positive_sample_mask\"].to(torch.bool)\n        if positive_mask.ndim == 1:\n            return positive_mask\n        return positive_mask.reshape(positive_mask.shape[0], -1).any(dim=1)\n\n    def _pad_sampled_batch(self, sampled_batch: TensorDict, target_batch_size: int) -> TensorDict:\n        current_size = sampled_batch.batch_size[0]\n        if current_size >= target_batch_size:\n            return sampled_batch\n\n        pad_size = target_batch_size - current_size\n        pad_idx = torch.zeros(pad_size, dtype=torch.long, device=self.sample_device)\n        padded_batch = TensorDict(\n            {key: torch.cat([value, value.index_select(0, pad_idx)], dim=0) for key, value in sampled_batch.items()},\n            batch_size=[target_batch_size],\n            device=self.sample_device,\n        )\n\n        valid_tensor = padded_batch[\"valids\"].clone()\n        if valid_tensor.dtype == torch.bool:\n            valid_tensor[current_size:] = False\n        else:\n            valid_tensor[current_size:] = 0\n        padded_batch[\"valids\"] = valid_tensor\n\n        return padded_batch\n\n    def _index_select_batch(self, batch: TensorDict, idx: torch.Tensor) -> TensorDict:\n        length = int(idx.numel())\n        return TensorDict(\n            {key: value.index_select(0, idx) for key, value in batch.items()},\n            batch_size=[length],\n            device=batch.device,\n        )\n\n    def _sample_from_task_pools(self, batch_size: int, is_positive_pool: bool) -> TensorDict:\n        task_sizes = {\n            task_id: (pool_state.positive_size if is_positive_pool else pool_state.negative_size)\n            for task_id, pool_state in self.task_pools.items()\n            if (pool_state.positive_size if is_positive_pool else pool_state.negative_size) > 0\n        }\n\n        allocation = self._allocate_counts_across_tasks(task_sizes, batch_size)\n\n        sampled_parts = []\n        for task_id, count in allocation.items():\n            if count == 0:\n                continue\n            sampled_parts.append(self._sample_from_single_task_pool(self.task_pools[task_id], count, is_positive_pool))\n\n        if len(sampled_parts) == 1:\n            return sampled_parts[0]\n\n        return TensorDict(\n            {key: torch.cat([part[key] for part in sampled_parts], dim=0) for key in sampled_parts[0].keys()},\n            batch_size=[batch_size],\n            device=self.sample_device,\n        )\n\n    def _sample_from_single_task_pool(\n        self,\n        pool_state: _DualPoolState,\n        batch_size: int,\n        is_positive_pool: bool,\n    ) -> TensorDict:\n        pool = pool_state.positive_pool if is_positive_pool else pool_state.negative_pool\n        size = pool_state.positive_size if is_positive_pool else pool_state.negative_size\n        assert pool is not None\n\n        idx = torch.randperm(size, device=self.pool_device)[:batch_size]\n        return TensorDict(\n            {key: value.index_select(0, idx).to(self.sample_device) for key, value in pool.items()},\n            batch_size=[batch_size],\n            device=self.sample_device,\n        )\n\n    def _allocate_counts_across_tasks(self, task_sizes: dict[str, int], total_count: int) -> dict[str, int]:\n        total_available = sum(task_sizes.values())\n        if total_count > total_available:\n            raise ValueError(f\"Requested {total_count} samples but only {total_available} available across task pools.\")\n\n        allocation: dict[str, int] = {task_id: 0 for task_id in task_sizes}\n        task_order = list(task_sizes.keys())\n\n        remaining = total_count\n        while remaining > 0:\n            progressed = False\n            for task_id in task_order:\n                if allocation[task_id] < task_sizes[task_id]:\n                    allocation[task_id] += 1\n                    remaining -= 1\n                    progressed = True\n                    if remaining == 0:\n                        break\n\n            if not progressed:\n                raise RuntimeError(\"No eligible task pool left while allocation is still remaining.\")\n\n        return allocation\n\n    def _refresh_global_stats(self):\n        self.positive_size = sum(state.positive_size for state in self.task_pools.values())\n        self.negative_size = sum(state.negative_size for state in self.task_pools.values())\n        self.size = self.positive_size + self.negative_size\n\n    def _normalize_task_id(self, task_id: Any) -> str:\n        if isinstance(task_id, torch.Tensor):\n            task_id = task_id.item()\n        return str(task_id)\n\n    def __repr__(self):\n        return (\n            f\"SACReplayPool(single_pool_capacity={self.single_pool_capacity}, size={self.size}, \"\n            f\"positive_size={self.positive_size}, negative_size={self.negative_size}, \"\n            f\"task_count={len(self.task_pools)}, pool_device={self.pool_device}, sample_device={self.sample_device})\"\n        )\n\n    def __len__(self):\n        return self.size\n"
  },
  {
    "path": "verl/experimental/vla/sac/sac_actor.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nSingle Process Actor\n\"\"\"\n\nimport logging\nimport os\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nimport torch.nn.functional as F\nfrom tensordict import TensorDict\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom typing_extensions import override\n\nfrom verl.experimental.vla.sac.replay_pool import SACReplayPool\nfrom verl.protocol import DataProto\nfrom verl.utils.device import get_device_id, get_device_name\n\nfrom .base import BaseSACActor, SupportSACTraining\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef get_dict_from_prefix(tensordict: TensorDict, prefix: str) -> dict:\n    \"\"\"Extract a sub-dictionary from a TensorDict based on a given prefix.\n\n    Args:\n        tensordict: The input TensorDict containing various keys.\n        prefix: The prefix string to filter keys.\n    Returns:\n        A dictionary containing key-value pairs from the TensorDict\n        where the keys start with the specified prefix. The prefix is removed\n        from the keys in the resulting dictionary.\n    \"\"\"\n\n    result = {}\n    prefix_length = len(prefix)\n    for key in tensordict.keys():\n        if key.startswith(prefix):\n            new_key = key[prefix_length:]\n            result[new_key] = tensordict[key]\n    return result\n\n\ndef merge_nested_dicts_or_tuples(a: dict | tuple, b: dict | tuple) -> dict | tuple:\n    \"\"\"Merge two nested structures (dictionaries or tuples) by concatenating tensors\n    along the first dimension.\n    \"\"\"\n\n    if isinstance(a, dict) and isinstance(b, dict):\n        merged = {}\n        for key in a.keys():\n            merged[key] = merge_nested_dicts_or_tuples(a[key], b[key])\n        return merged\n    elif isinstance(a, tuple) and isinstance(b, tuple):\n        merged = []\n        for item_a, item_b in zip(a, b, strict=False):\n            merged.append(merge_nested_dicts_or_tuples(item_a, item_b))\n        return tuple(merged)\n    else:\n        return torch.cat([a, b], dim=0)\n\n\ndef split_nested_dicts_or_tuples(data: dict | tuple, split_num: int) -> list[dict | tuple]:\n    \"\"\"Split a nested structure (dictionary or tuple) into smaller chunks along the first dimension.\"\"\"\n\n    if isinstance(data, torch.Tensor):\n        split_tensors = torch.chunk(data, split_num, dim=0)\n        return list(split_tensors)\n    elif isinstance(data, dict):\n        split_dicts = [dict() for _ in range(split_num)]\n        for key, value in data.items():\n            split_values = split_nested_dicts_or_tuples(value, split_num)\n            for i in range(split_num):\n                split_dicts[i][key] = split_values[i]\n        return split_dicts\n    elif isinstance(data, tuple):\n        split_tuples = [list() for _ in range(split_num)]\n        for item in data:\n            split_items = split_nested_dicts_or_tuples(item, split_num)\n            for i in range(split_num):\n                split_tuples[i].append(split_items[i])\n        return [tuple(split_tuple) for split_tuple in split_tuples]\n    else:\n        raise TypeError(\"Input data must be a torch.Tensor, dict, or tuple.\")\n\n\ndef valid_mean(x: torch.Tensor, valid: torch.Tensor) -> torch.Tensor:\n    \"\"\"Compute the mean of tensor `x` over valid entries indicated by `valid` mask.\n\n    Args:\n        x: Tensor of shape (B, ...) containing values to average.\n        valid: Tensor of shape (B,) indicating valid entries (1 for valid, 0 for invalid).\n\n    Returns:\n        Scalar tensor (mean over valid samples only)\n    \"\"\"\n    x = x.squeeze(-1)\n    valid_f = valid.float().to(x.device)\n    denom = valid_f.sum().clamp_min(1.0)\n    return (x * valid_f).sum() / denom\n\n\nclass RobDataParallelSACActor(BaseSACActor):\n    def __init__(\n        self,\n        config,\n        actor_module: SupportSACTraining,\n        actor_optimizer: torch.optim.Optimizer,\n        tokenizer=None,\n    ):\n        super().__init__()\n        self.config = config\n        self.sac_config = config.sac\n        self.device = get_device_name()\n\n        self.actor_optimizer = actor_optimizer\n        self.actor_module = actor_module\n        self.actor_module.sac_init()\n        self.tokenizer = tokenizer\n\n        self.replay_pool = SACReplayPool(\n            single_pool_capacity=self.config.replay_pool_single_size,\n            sample_device=self.device,\n        )\n        self.replay_pool.load(self.config.replay_pool_save_dir)\n\n        self._init_alpha()\n        self._init_critic()\n\n        self.actor_ema_enabled = bool(self.config.get(\"actor_ema_enabled\", True))\n        self.actor_ema_decay = float(self.config.get(\"actor_ema_decay\", 0.995))\n        self.actor_ema_shadow: dict[str, torch.Tensor] = {}\n        self.actor_ema_initialized = False\n        self.bc_loss_coef = float(self.sac_config.get(\"bc_loss_coef\", 0.5))\n\n    def _init_critic(self):\n        \"\"\"Initialize the critic optimizer.\"\"\"\n\n        self.critic_optimizer = torch.optim.Adam(\n            self.actor_module.sac_get_critic_parameters(),\n            lr=self.config.critic_lr,\n            weight_decay=self.config.critic_weight_decay,\n        )\n        self.critic_scheduler = torch.optim.lr_scheduler.ConstantLR(self.critic_optimizer, factor=1.0)\n\n    def _init_alpha(self):\n        \"\"\"Initialize the alpha optimizer for automatic entropy tuning.\"\"\"\n\n        self.auto_entropy = self.sac_config.get(\"auto_entropy\", False)\n\n        if self.auto_entropy:\n            self.target_entropy = torch.tensor(float(self.sac_config.get(\"target_entropy\", -32.0)), device=self.device)\n\n            # Initialize raw_alpha parameter\n            self.alpha_type = self.sac_config.get(\"alpha_type\", \"softplus\")\n            if self.alpha_type == \"exp\":\n                self.raw_alpha = torch.nn.Parameter(\n                    np.log(np.exp(self.sac_config.get(\"initial_alpha\", 1))) * torch.ones(1, device=self.device),\n                    requires_grad=True,\n                )\n            elif self.alpha_type == \"softplus\":\n                self.raw_alpha = torch.nn.Parameter(\n                    np.log(np.exp(self.sac_config.get(\"initial_alpha\", 0.01)) - 1) * torch.ones(1, device=self.device),\n                    requires_grad=True,\n                )\n            else:\n                return NotImplementedError(f\"Unsupported alpha_type: {self.alpha_type}\")\n\n            # build alpha optimizer and scheduler\n            self.alpha_optimizer = torch.optim.Adam([self.raw_alpha], lr=self.sac_config.get(\"alpha_lr\", 3e-4))\n            self.alpha_scheduler = torch.optim.lr_scheduler.ConstantLR(self.alpha_optimizer, factor=1.0)\n\n    def _init_actor_ema(self):\n        if self.actor_ema_initialized:\n            return\n\n        self.actor_ema_shadow = {}\n\n        if not self.actor_ema_enabled:\n            self.actor_ema_initialized = True\n            return\n\n        for name, param in self.actor_module.sac_get_named_actor_parameters():\n            self.actor_ema_shadow[name] = param.detach().clone().to(dtype=torch.float32)\n\n        self.actor_ema_initialized = True\n\n    @torch.no_grad()\n    def _update_actor_ema(self):\n        if not self.actor_ema_enabled:\n            return\n\n        one_minus_decay = 1.0 - self.actor_ema_decay\n        for name, param in self.actor_module.sac_get_named_actor_parameters():\n            shadow = self.actor_ema_shadow[name]\n            shadow.mul_(self.actor_ema_decay).add_(param.detach().to(dtype=torch.float32), alpha=one_minus_decay)\n\n    @torch.no_grad()\n    def _apply_actor_ema_to_actor_module(self):\n        if not self.actor_ema_enabled:\n            return\n\n        for name, param in self.actor_module.sac_get_named_actor_parameters():\n            shadow = self.actor_ema_shadow[name]\n            param.copy_(shadow.to(device=param.device, dtype=param.dtype))\n\n    def _get_alpha(self) -> torch.Tensor:\n        if self.auto_entropy:\n            if self.alpha_type == \"exp\":\n                return self.raw_alpha.exp()\n            elif self.alpha_type == \"softplus\":\n                return torch.nn.functional.softplus(self.raw_alpha)\n            else:\n                return NotImplementedError(f\"Unsupported alpha_type: {self.alpha_type}\")\n        else:\n            return torch.tensor(float(self.sac_config.get(\"initial_alpha\", 0.2)), device=self.device)\n\n    def _calculate_actor_loss(\n        self,\n        log_probs: Optional[torch.Tensor],\n        q_values: torch.Tensor,\n        valids: torch.Tensor,\n    ) -> torch.Tensor:\n        \"\"\"Calculate actor loss using the SAC loss function.\n\n        Args:\n            log_probs: Tensor of shape (B,) representing the log probabilities of actions.\n            q_values: Tensor of shape (B,) representing the Q-values for the actions.\n            valids: Tensor of shape (B,) indicating valid samples (1 for valid, 0 for invalid).\n\n        Returns:\n            Tensor of shape (1,) representing the actor loss.\n        \"\"\"\n\n        alpha = self._get_alpha()\n        if log_probs is None:\n            loss = -q_values\n        else:\n            loss = alpha * log_probs - q_values\n        actor_loss = (loss * valids).sum() / (valids.sum().clamp_min(1.0))\n\n        return actor_loss\n\n    def _calculate_alpha_loss(self, log_probs: Optional[torch.Tensor], valids: torch.Tensor) -> torch.Tensor:\n        \"\"\"Calculate alpha loss for automatic entropy tuning.\n\n        Args:\n            log_probs: Tensor of shape (B,) representing the log probabilities of actions.\n            valids: Tensor of shape (B,) indicating valid samples (1 for valid, 0 for invalid).\n\n        Returns:\n            Tensor of shape (1,) representing the alpha loss.\n        \"\"\"\n\n        if log_probs is None:\n            return torch.tensor(0.0, device=valids.device)\n\n        alpha_loss = -self._get_alpha() * (log_probs.detach() + self.target_entropy)\n        alpha_loss = (alpha_loss * valids).sum() / (valids.sum().clamp_min(1.0))\n        return alpha_loss\n\n    def _calculate_critic_loss(\n        self,\n        q_predict: torch.Tensor,\n        q_target: torch.Tensor,\n        rewards: torch.Tensor,\n        dones: torch.Tensor,\n        next_log_prob: Optional[torch.Tensor],\n        valids: torch.Tensor,\n    ) -> torch.Tensor:\n        \"\"\"Calculate critic loss using the SAC loss function.\n\n        Args:\n            q_predict: Tensor of shape (B, critic_num) representing predicted Q-values.\n            q_target: Tensor of shape (B,) representing target Q-values.\n            rewards: Tensor of shape (B,) representing rewards.\n            dones: Tensor of shape (B,) representing done flags.\n            next_log_prob: Tensor of shape (B,) representing log probabilities of next actions.\n            valids: Tensor of shape (B,) indicating valid samples (1 for valid, 0 for invalid).\n\n        Returns:\n            Tensor of shape (1,) representing the critic loss.\n        \"\"\"\n\n        gamma = self.sac_config.gamma\n        alpha = self._get_alpha()\n\n        with torch.no_grad():\n            if next_log_prob is None:\n                y = rewards + gamma * (1.0 - dones) * q_target\n            else:\n                y = rewards + gamma * (1.0 - dones) * (q_target - alpha * next_log_prob)\n\n        y = y.unsqueeze(1).expand_as(q_predict)  # (B, critic_num)\n        valid_mask = valids.unsqueeze(1)\n        mse = F.mse_loss(q_predict, y, reduction=\"none\")\n        per_critic = (mse * valid_mask).sum(dim=0) / valid_mask.sum().clamp_min(1.0)\n        critic_loss = per_critic.sum()\n        return critic_loss\n\n    def _forward_critic(\n        self, micro_batch: TensorDict, resample=True\n    ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n        s0 = get_dict_from_prefix(micro_batch, \"s0.\")\n        s1 = get_dict_from_prefix(micro_batch, \"s1.\")\n        a0 = get_dict_from_prefix(micro_batch, \"a0.\")\n        a1 = get_dict_from_prefix(micro_batch, \"a1.\")\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            with torch.no_grad():\n                s = merge_nested_dicts_or_tuples(s0, s1)\n                state_features = self.actor_module.sac_forward_state_features(s)\n                s0_state_features, s1_state_features = split_nested_dicts_or_tuples(state_features, 2)\n                if resample:\n                    a1_actions, log_probs_1, _ = self.actor_module.sac_forward_actor(\n                        s1_state_features,\n                        is_first_micro_batch=False,\n                    )\n                    a1 = {\"full_action\": a1_actions}\n                else:\n                    log_probs_1 = None\n\n            q_values_0 = self.actor_module.sac_forward_critic(\n                a0,\n                s0_state_features,\n                use_target_network=False,\n                method=\"cat\",\n                requires_grad=True,\n            )\n            q_values_1 = self.actor_module.sac_forward_critic(\n                a1,\n                s1_state_features,\n                use_target_network=True,\n                method=\"min\",\n                requires_grad=False,\n            )\n\n            critic_loss = self._calculate_critic_loss(\n                q_predict=q_values_0,\n                q_target=q_values_1,\n                rewards=micro_batch[\"rewards\"],\n                dones=micro_batch[\"dones\"],\n                next_log_prob=log_probs_1,\n                valids=micro_batch[\"valids\"],\n            )\n        return critic_loss, q_values_0, q_values_1\n\n    def _forward_actor(\n        self,\n        micro_batch: TensorDict,\n        is_first_micro_batch: bool,\n    ) -> tuple[torch.Tensor, Optional[torch.Tensor], torch.Tensor, dict[str, float]]:\n        micro_batch = micro_batch.to(get_device_id())\n        s0 = get_dict_from_prefix(micro_batch, \"s0.\")\n\n        with torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            s0_state_features = self.actor_module.sac_forward_state_features(s0)\n            a0_actions, log_probs_0, actor_forward_metrics = self.actor_module.sac_forward_actor(\n                s0_state_features,\n                is_first_micro_batch=is_first_micro_batch,\n            )\n            q_values_0 = self.actor_module.sac_forward_critic(\n                {\"full_action\": a0_actions},\n                s0_state_features,\n                use_target_network=False,\n                method=\"min\",\n                requires_grad=False,\n            )\n\n            sac_loss = self._calculate_actor_loss(\n                log_probs=log_probs_0,\n                q_values=q_values_0,\n                valids=micro_batch[\"valids\"],\n            )\n            if self.bc_loss_coef > 0:\n                bc_loss = self.actor_module.bc_loss(\n                    state_features=s0_state_features,\n                    actions={\"full_action\": a0_actions},\n                    valids=micro_batch[\"valids\"],\n                )\n                actor_loss = sac_loss + self.bc_loss_coef * bc_loss\n            else:\n                actor_loss = sac_loss\n        return actor_loss, log_probs_0, q_values_0, actor_forward_metrics\n\n    def _force_set_lr(self, opt: torch.optim.Optimizer, lr: float):\n        for pg in opt.param_groups:\n            pg[\"lr\"] = lr\n\n    @override\n    def update_policy(self, data: DataProto):\n        if not self.actor_ema_initialized:\n            self._init_actor_ema()\n\n        # self._force_set_lr(self.actor_optimizer, 5e-6)\n        # self._force_set_lr(self.critic_optimizer, 1e-4)\n\n        if \"empty_batch\" not in data.meta_info:\n            task_ids = data.batch[\"task_ids\"]\n            self.replay_pool.add_batch(\n                data.select(\n                    [\n                        \"a0.full_action\",\n                        \"a1.full_action\",\n                        \"s0.states\",\n                        \"s1.states\",\n                        \"s0.images\",\n                        \"s1.images\",\n                        \"s0.image_masks\",\n                        \"s1.image_masks\",\n                        \"s0.lang_tokens\",\n                        \"s1.lang_tokens\",\n                        \"s0.lang_masks\",\n                        \"s1.lang_masks\",\n                        \"rewards\",\n                        \"dones\",\n                        \"valids\",\n                        \"positive_sample_mask\",\n                    ]\n                ).batch,\n                task_ids=task_ids,\n            )\n\n        replay_positive_sample_ratio = float(self.sac_config.get(\"critic_replay_positive_sample_ratio\", 0.5))\n        critic_batch, critic_replay_sample_info = self.replay_pool.sample_batch(\n            self.config.ppo_mini_batch_size,\n            positive_sample_ratio=replay_positive_sample_ratio,\n            return_sample_info=True,\n        )\n        micro_batches = critic_batch.split(self.config.ppo_micro_batch_size_per_gpu)\n        global_steps = data.meta_info[\"global_steps\"]\n        grad_accum_steps = len(micro_batches) * torch.distributed.get_world_size()\n\n        actor_logprobs_list, actor_qvalues_list = [], []\n        critic_qvalues_0_list, critic_qvalues_1_list = [], []\n        actor_loss_list, critic_loss_list, alpha_loss_list = [], [], []\n        actor_forward_metrics: dict[str, float] = {}\n\n        # Training critic\n        self.critic_optimizer.zero_grad()\n        for batch_idx, micro_batch in enumerate(micro_batches):\n            logger.info(f\"[{batch_idx + 1}/{len(micro_batches)}] critic micro batch \")\n\n            micro_batch = micro_batch.to(get_device_id())\n            raw_critic_loss, q_values_0, q_values_1 = self._forward_critic(micro_batch, resample=True)\n            (raw_critic_loss / grad_accum_steps).backward()\n            critic_loss_list.append(raw_critic_loss.detach().item())\n            critic_qvalues_0_list.append(q_values_0.mean(dim=-1).detach())\n            critic_qvalues_1_list.append(q_values_1.detach())\n        critic_grad_norm = torch.nn.utils.clip_grad_norm_(\n            self.actor_module.sac_get_critic_parameters(), max_norm=self.config.grad_clip\n        )\n        self.critic_optimizer.step()\n        self.critic_scheduler.step()\n\n        update_actor = (\n            global_steps >= self.config.critic_warmup_steps and global_steps % self.config.actor_update_interval == 0\n        )\n        if update_actor:\n            replay_positive_sample_ratio = float(self.sac_config.get(\"actor_replay_positive_sample_ratio\", 0.5))\n            actor_batch, actor_replay_sample_info = self.replay_pool.sample_batch(\n                self.config.ppo_mini_batch_size,\n                positive_sample_ratio=replay_positive_sample_ratio,\n                return_sample_info=True,\n            )\n            micro_batches = actor_batch.split(self.config.ppo_micro_batch_size_per_gpu)\n\n            # Training actor\n            self.actor_optimizer.zero_grad()\n            for batch_idx, micro_batch in enumerate(micro_batches):\n                logger.info(f\"[{batch_idx + 1}/{len(micro_batches)}] actor micro batch \")\n\n                micro_batch = micro_batch.to(get_device_id())\n                raw_actor_loss, log_probs, q_values, actor_forward_metrics_mb = self._forward_actor(\n                    micro_batch,\n                    is_first_micro_batch=(batch_idx == 0),\n                )\n                (raw_actor_loss / grad_accum_steps).backward()\n                actor_loss_list.append(raw_actor_loss.detach().item())\n                if log_probs is not None:\n                    actor_logprobs_list.append(log_probs.detach())\n                actor_qvalues_list.append(q_values.detach())\n                actor_forward_metrics.update(actor_forward_metrics_mb)\n            actor_grad_norm = self._optimizer_step()\n            self._update_actor_ema()\n            self._apply_actor_ema_to_actor_module()\n\n            # Training alpha\n            # NOTE: We reuse the log-probabilities computed during the actor forward pass\n            # to update the entropy temperature (alpha), instead of re-forwarding\n            # the actor after the policy update (saving compute).\n            if self.auto_entropy and actor_logprobs_list:\n                self.alpha_optimizer.zero_grad()\n                for micro_batch, log_probs in zip(micro_batches, actor_logprobs_list, strict=False):\n                    micro_batch = micro_batch.to(get_device_id())\n                    raw_alpha_loss = self._calculate_alpha_loss(log_probs, micro_batch[\"valids\"])\n                    (raw_alpha_loss / grad_accum_steps).backward()\n                    alpha_loss_list.append(raw_alpha_loss.detach().item())\n                torch.distributed.all_reduce(self.raw_alpha.grad, op=torch.distributed.ReduceOp.SUM)\n                alpha_grad_norm = torch.nn.utils.clip_grad_norm_(self.raw_alpha, max_norm=self.config.grad_clip)\n                self.alpha_optimizer.step()\n                self.alpha_scheduler.step()\n\n        # Update target networks\n        self.actor_module.sac_update_target_network(self.sac_config.tau)\n\n        # Save replay pool\n        if global_steps % self.config.replay_pool_save_interval == 0:\n            self.replay_pool.save(self.config.replay_pool_save_dir)\n\n        # Log metrics\n        positive_qvalue_mean = (\n            torch.cat(critic_qvalues_0_list)[\n                (critic_batch[\"positive_sample_mask\"].to(torch.bool) & critic_batch[\"valids\"].to(torch.bool)).to(\n                    torch.cat(critic_qvalues_0_list).device\n                )\n            ]\n            .mean()\n            .detach()\n            .item()\n            if critic_qvalues_0_list\n            and (critic_batch[\"positive_sample_mask\"].to(torch.bool) & critic_batch[\"valids\"].to(torch.bool)).any()\n            else 0.0\n        )\n        negative_qvalue_mean = (\n            torch.cat(critic_qvalues_0_list)[\n                (~critic_batch[\"positive_sample_mask\"].to(torch.bool) & critic_batch[\"valids\"].to(torch.bool)).to(\n                    torch.cat(critic_qvalues_0_list).device\n                )\n            ]\n            .mean()\n            .detach()\n            .item()\n            if critic_qvalues_0_list\n            and (~critic_batch[\"positive_sample_mask\"].to(torch.bool) & critic_batch[\"valids\"].to(torch.bool)).any()\n            else 0.0\n        )\n        metrics = {\n            \"data/reward_mean\": valid_mean(critic_batch[\"rewards\"], critic_batch[\"valids\"]).detach().item(),\n            \"data/valid_ratio\": critic_batch[\"valids\"].float().mean().item(),\n            \"sac/critic_replay_sampled_ratio\": critic_replay_sample_info[\"actual_positive_sample_ratio\"],\n            \"sac/actor_replay_sampled_ratio\": actor_replay_sample_info[\"actual_positive_sample_ratio\"]\n            if update_actor\n            else 0.0,\n            \"sac/replay_pool_positive_size\": critic_replay_sample_info[\"positive_size\"],\n            \"sac/replay_pool_negative_size\": critic_replay_sample_info[\"negative_size\"],\n            \"sac/replay_task_count\": critic_replay_sample_info[\"task_count\"],\n            \"sac/alpha\": self._get_alpha().detach().item(),\n            \"sac/actor_ema_enabled\": float(self.actor_ema_enabled),\n            \"sac/actor_ema_decay\": self.actor_ema_decay,\n            \"sac/replay_pool_size\": len(self.replay_pool),\n            \"critic/loss\": sum(critic_loss_list) / len(critic_loss_list) if critic_loss_list else 0.0,\n            \"critic/lr\": self.critic_optimizer.param_groups[0][\"lr\"],\n            \"critic/grad_norm\": critic_grad_norm.detach().item(),\n            \"critic/qvalue0_mean\": (\n                valid_mean(torch.cat(critic_qvalues_0_list), critic_batch[\"valids\"]).detach().item()\n                if critic_qvalues_0_list\n                else 0.0\n            ),\n            \"critic/qvalue1_mean\": (\n                valid_mean(torch.cat(critic_qvalues_1_list), critic_batch[\"valids\"]).detach().item()\n                if critic_qvalues_1_list\n                else 0.0\n            ),\n            \"critic/positive_qvalue_mean\": positive_qvalue_mean,\n            \"critic/negative_qvalue_mean\": negative_qvalue_mean,\n            \"critic/diff_pos_neg_qvalue_mean\": positive_qvalue_mean - negative_qvalue_mean,\n        }\n        if update_actor:\n            metrics.update(\n                {\n                    \"actor/loss\": sum(actor_loss_list) / len(actor_loss_list),\n                    \"actor/lr\": self.actor_optimizer.param_groups[0][\"lr\"],\n                    \"actor/grad_norm\": actor_grad_norm.detach().item(),\n                    \"actor/logprob_mean\": (\n                        valid_mean(torch.cat(actor_logprobs_list), actor_batch[\"valids\"]).detach().item()\n                        if actor_logprobs_list\n                        else 0.0\n                    ),\n                    \"actor/qvalue_mean\": valid_mean(torch.cat(actor_qvalues_list), actor_batch[\"valids\"])\n                    .detach()\n                    .item(),\n                    \"sac/alpha_lr\": self.alpha_optimizer.param_groups[0][\"lr\"]\n                    if self.auto_entropy and actor_logprobs_list\n                    else 0.0,\n                    \"sac/alpha_loss\": sum(alpha_loss_list) / len(alpha_loss_list)\n                    if self.auto_entropy and alpha_loss_list\n                    else 0.0,\n                    \"sac/alpha_grad_norm\": alpha_grad_norm.detach().item()\n                    if self.auto_entropy and actor_logprobs_list\n                    else 0.0,\n                }\n            )\n            metrics.update({f\"actor/{k}\": v for k, v in actor_forward_metrics.items()})\n\n        return metrics\n\n    def _optimizer_step(self) -> torch.Tensor:\n        assert self.config.grad_clip is not None\n\n        if isinstance(self.actor_module, FSDP):\n            grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip)\n        self.actor_optimizer.step()\n        return grad_norm\n"
  },
  {
    "path": "verl/experimental/vla/sac/sac_ray_trainer.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nfrom pprint import pprint\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nfrom omegaconf import OmegaConf\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.single_controller.ray import RayClassWithInitArgs\nfrom verl.single_controller.ray.base import create_colocated_worker_cls\nfrom verl.trainer.ppo.ray_trainer import RayPPOTrainer\nfrom verl.trainer.ppo.utils import Role\nfrom verl.utils.checkpoint.checkpoint_manager import should_save_ckpt_esi\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.metric import reduce_metrics\n\n\ndef compute_avg_positive_trajectory_length(batch: DataProto) -> float:\n    dones = batch.batch[\"dones\"].bool()  # (B, T)\n    positive_mask = batch.batch[\"positive_sample_mask\"]  # (B, T)\n    positive_traj = positive_mask.any(dim=1)  # (B,)\n\n    if positive_traj.sum() == 0:\n        return 0.0\n\n    B, T = dones.shape\n    done_idx = torch.argmax(dones.int(), dim=1)  # (B,)\n    traj_lens = done_idx + 1\n\n    return traj_lens[positive_traj].float().mean().item()\n\n\ndef flatten_trajectories(data: DataProto) -> DataProto:\n    batch_size, num_steps = data.batch[\"action\"].shape[:2]\n    new_batch_fields = {}\n    for key, tensor in data.batch.items():\n        if len(tensor.shape) >= 2 and tensor.shape[0] == batch_size and tensor.shape[1] == num_steps:\n            # (B, S, H, W) -> (B*S, H, W)\n            new_shape = (batch_size * num_steps, *tensor.shape[2:])\n            new_batch_fields[key] = tensor.reshape(new_shape)\n        elif len(tensor.shape) == 1 and tensor.shape[0] == batch_size:\n            # [e1, e2] -> [e1, e1, ..., e2, e2, ...] (S times each)\n            new_batch_fields[key] = tensor.repeat_interleave(num_steps)\n        else:\n            new_batch_fields[key] = tensor\n    new_data = DataProto.from_dict(tensors=new_batch_fields, meta_info=data.meta_info)\n    return new_data\n\n\ndef add_transition_prefixes(data: DataProto) -> DataProto:\n    batch = data.batch\n    step_key = \"action\" if \"action\" in batch else \"full_action\"\n    if step_key not in batch:\n        return data\n\n    num_steps = batch[step_key].shape[1]\n    if num_steps <= 1:\n        return data\n\n    def drop_last(tensor: torch.Tensor) -> torch.Tensor:\n        return tensor[:, :-1, ...]\n\n    def shift_next(tensor: torch.Tensor) -> torch.Tensor:\n        return tensor[:, 1:, ...]\n\n    state_keys = [\"states\", \"images\", \"image_masks\", \"lang_tokens\", \"lang_masks\"]\n    action_keys = [\"full_action\", \"action\"]\n\n    for key in state_keys:\n        if key in batch:\n            batch[f\"s0.{key}\"] = drop_last(batch[key])\n            batch[f\"s1.{key}\"] = shift_next(batch[key])\n\n    for key in action_keys:\n        if key in batch:\n            batch[f\"a0.{key}\"] = drop_last(batch[key])\n            batch[f\"a1.{key}\"] = shift_next(batch[key])\n\n    batch_size = batch[step_key].shape[0]\n    for key, tensor in list(batch.items()):\n        if tensor.ndim >= 2 and tensor.shape[0] == batch_size and tensor.shape[1] == num_steps:\n            batch[key] = drop_last(tensor)\n\n    return data\n\n\nclass RobRaySACTrainer(RayPPOTrainer):\n    def _start_profiling(self, do_profile: bool) -> None:\n        \"\"\"Start profiling for all worker groups including env workers.\"\"\"\n        super()._start_profiling(do_profile)\n        if do_profile and hasattr(self, \"env_wg\"):\n            self.env_wg.start_profile(role=\"env\", profile_step=self.global_steps)\n\n    def _stop_profiling(self, do_profile: bool) -> None:\n        \"\"\"Stop profiling for all worker groups including env workers.\"\"\"\n        super()._stop_profiling(do_profile)\n        if do_profile and hasattr(self, \"env_wg\"):\n            self.env_wg.stop_profile()\n\n    def init_workers(self):\n        self.resource_pool_manager.create_resource_pool()\n\n        if self.config.env.disagg_sim.enable:\n            # pin EnvWorker to Simulator GPU nodes\n            self.resource_pool_manager.get_resource_pool(Role.Env).accelerator_type = \"sim\"\n            self.resource_pool_manager.get_resource_pool(Role.ActorRollout).accelerator_type = \"train_rollout\"\n\n        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}\n        resource_pool = self.resource_pool_manager.get_resource_pool(Role.ActorRollout)\n        actor_rollout_cls = RayClassWithInitArgs(\n            cls=self.role_worker_mapping[Role.ActorRollout],\n            config=self.config.actor_rollout_ref,\n            role=\"actor_rollout\",\n        )\n        self.resource_pool_to_cls[resource_pool][\"actor_rollout\"] = actor_rollout_cls\n\n        assert Role.Env in self.role_worker_mapping\n        if Role.Env in self.role_worker_mapping:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Env)\n            env_cls = RayClassWithInitArgs(self.role_worker_mapping[Role.Env], config=self.config.env)\n            self.resource_pool_to_cls[resource_pool][\"env\"] = env_cls\n\n        # initialize WorkerGroup\n        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,\n        # you should not use `create_colocated_worker_cls`.\n        # Instead, directly pass different resource pool to different worker groups.\n        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.\n        all_wg = {}\n        wg_kwargs = {}  # Setting up kwargs for RayWorkerGroup\n        if OmegaConf.select(self.config.trainer, \"ray_wait_register_center_timeout\") is not None:\n            wg_kwargs[\"ray_wait_register_center_timeout\"] = self.config.trainer.ray_wait_register_center_timeout\n        if OmegaConf.select(self.config.global_profiler, \"steps\") is not None:\n            wg_kwargs[\"profile_steps\"] = OmegaConf.select(self.config.global_profiler, \"steps\")\n            # Only require nsight worker options when tool is nsys\n            if OmegaConf.select(self.config.global_profiler, \"tool\") == \"nsys\":\n                assert (\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                    is not None\n                ), \"worker_nsight_options must be set when using nsys with profile_steps\"\n                wg_kwargs[\"worker_nsight_options\"] = OmegaConf.to_container(\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                )\n        wg_kwargs[\"device_name\"] = self.device_name\n\n        for resource_pool, class_dict in self.resource_pool_to_cls.items():\n            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n            wg_dict = self.ray_worker_group_cls(\n                resource_pool=resource_pool,\n                ray_cls_with_init=worker_dict_cls,\n                **wg_kwargs,\n            )\n            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n            all_wg.update(spawn_wg)\n\n        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory\n        self.actor_rollout_wg = all_wg[\"actor_rollout\"]\n        self.actor_rollout_wg.init_model()\n        self.env_wg = all_wg[\"env\"]\n\n        # create async rollout manager and request scheduler\n        self.async_rollout_mode = False\n        if self.config.actor_rollout_ref.rollout.mode == \"async_envloop\":\n            from verl.experimental.vla.env_loop import EnvLoop\n\n            self.async_rollout_mode = True\n            self.async_rollout_manager = EnvLoop(\n                config=self.config, rollout_wg=self.actor_rollout_wg, env_wg=self.env_wg\n            )\n\n    def _get_gen_batch(self, batch: DataProto) -> DataProto:\n        # pop those keys for generation\n        batch_keys_to_pop = []\n        non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys())\n        gen_batch = batch.pop(\n            batch_keys=batch_keys_to_pop,\n            non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),\n        )\n\n        return gen_batch\n\n    def _reset_envs(self, gen_batch: DataProto) -> asyncio.Future:\n        initial_state_ids = gen_batch.non_tensor_batch[\"state_ids\"]\n        task_ids = gen_batch.non_tensor_batch[\"task_ids\"]\n        reset_prompts = DataProto.from_dict(non_tensors={\"state_ids\": initial_state_ids, \"task_ids\": task_ids})\n        reset_future = self.env_wg.reset_envs_to_state_ids(reset_prompts)\n        return reset_future\n\n    def _next_rollout_batch(self, train_iter) -> Optional[DataProto]:\n        try:\n            batch_dict = next(train_iter)\n        except StopIteration:\n            return None\n\n        rollout_batch = DataProto.from_single_dict(batch_dict)\n        rollout_batch = self._get_gen_batch(rollout_batch)\n        rollout_batch = rollout_batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n        rollout_batch.meta_info[\"task_ids\"] = np.asarray(rollout_batch.non_tensor_batch[\"task_ids\"], dtype=np.int64)\n        rollout_batch.meta_info[\"global_steps\"] = self.global_steps\n\n        return rollout_batch\n\n    def _prepare_actor_input(self, rollout_output: Optional[DataProto]) -> DataProto:\n        # dones\n        complete_any = rollout_output.batch[\"complete\"].any(dim=-1)  # (B, T)\n        dones_step = complete_any.clone()\n        dones_step[:, -2] = True\n        rollout_output.batch[\"dones\"] = dones_step.float()\n\n        # reward (sparse reward with step penalty)\n        sparse_rewards = complete_any.float()\n        rollout_output.batch[\"valids\"] = (~rollout_output.batch[\"complete\"]).any(dim=-1).float()\n        step_penalty = float(self.config.env.train.get(\"step_penalty\", 0.0))\n        rollout_output.batch[\"rewards\"] = sparse_rewards - step_penalty * rollout_output.batch[\"valids\"]\n        rollout_output.batch[\"rewards\"][:, -2] = -1.0\n\n        # mark samples in successful trajectories as positive samples\n        rollout_output.batch[\"positive_sample_mask\"] = (\n            sparse_rewards.any(dim=-1).unsqueeze(-1).repeat_interleave(rollout_output.batch[\"action\"].shape[1], dim=-1)\n        )\n\n        # task id\n        rollout_output.batch[\"task_ids\"] = torch.as_tensor(\n            rollout_output.meta_info[\"task_ids\"],\n            dtype=torch.long,\n            device=rollout_output.batch[\"action\"].device,\n        )\n\n        rollout_output.meta_info[\"global_token_num\"] = [0]\n        rollout_output.meta_info[\"data/trajectory_avg_reward\"] = (\n            sparse_rewards.any(dim=-1).mean(dtype=torch.float32).item()\n        )\n        rollout_output.meta_info[\"data/avg_positive_trajectory_length\"] = compute_avg_positive_trajectory_length(\n            rollout_output\n        )\n\n        rollout_output = add_transition_prefixes(rollout_output)\n        rollout_output = flatten_trajectories(rollout_output)\n\n        return rollout_output\n\n    def fit(self):\n        from omegaconf import OmegaConf\n\n        from verl.utils.tracking import Tracking\n\n        logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        self.global_steps = 0\n\n        # load checkpoint before doing anything\n        self._load_checkpoint()\n\n        # perform validation before training\n        # currently, we only support validation using the reward_function.\n        if self.config.trainer.get(\"val_before_train\", True):\n            val_metrics = self._validate()\n            assert val_metrics, f\"{val_metrics=}\"\n            pprint(f\"Initial validation metrics: {val_metrics}\")\n            logger.log(data=val_metrics, step=self.global_steps)\n            if self.config.trainer.get(\"val_only\", False):\n                return\n\n        # add tqdm\n        self.total_training_steps = (\n            self.config.trainer.total_epochs * len(self.train_dataloader) * self.config.trainer.rollout_interval\n        )\n        progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc=\"Training Progress\")\n\n        self.global_steps += 1\n        last_val_metrics = None\n        self.max_steps_duration = 0\n\n        prev_step_profile = False\n        curr_step_profile = (\n            self.global_steps in self.config.global_profiler.steps\n            if self.config.global_profiler.steps is not None\n            else False\n        )\n        next_step_profile = False\n\n        for epoch in range(self.config.trainer.total_epochs):\n            train_iter = iter(self.train_dataloader)\n            reset_future = None\n            next_rollout_batch = self._next_rollout_batch(train_iter)\n            if next_rollout_batch is None:\n                continue\n\n            print(f\"Starting epoch {epoch}, dataloader length: {len(self.train_dataloader)}\")\n            while next_rollout_batch is not None:\n                for training_step in range(self.config.trainer.rollout_interval):\n                    metrics = {}\n                    timing_raw = {}\n\n                    # === start profiling ===\n                    with marked_timer(\"start_profile\", timing_raw):\n                        self._start_profiling(\n                            not prev_step_profile and curr_step_profile\n                            if self.config.global_profiler.profile_continuous_steps\n                            else curr_step_profile\n                        )\n\n                    with marked_timer(\"step\", timing_raw):\n                        # === rollout ===\n                        # Determine whether to perform rollout:\n                        # enable at start and early warmup, disable during critic warmup phase\n                        warm_rollout_steps = int(getattr(self.config.actor_rollout_ref.actor, \"warm_rollout_steps\", 0))\n                        need_rollout = (training_step == 0) or self.global_steps < warm_rollout_steps\n                        if (\n                            warm_rollout_steps\n                            <= self.global_steps\n                            < self.config.actor_rollout_ref.actor.critic_warmup_steps\n                        ):\n                            need_rollout = False\n                        if need_rollout and next_rollout_batch is None:\n                            break\n\n                        actor_input = None\n                        if need_rollout:\n                            with marked_timer(\"rollout\", timing_raw):\n                                # execute rollout\n                                rollout_batch = next_rollout_batch\n                                assert rollout_batch is not None\n                                if reset_future is None:\n                                    reset_future = self._reset_envs(rollout_batch)\n                                with marked_timer(\"generate\", timing_raw, color=\"red\"):\n                                    rollout_output = self.async_rollout_manager.generate_sequences(\n                                        rollout_batch, reset_future\n                                    )\n\n                                # prepare for next batch's env reset\n                                next_rollout_batch = self._next_rollout_batch(train_iter)\n                                if next_rollout_batch is not None:\n                                    reset_future = self._reset_envs(next_rollout_batch)\n\n                                # compute rewards and other metrics, and prepare for actor update\n                                actor_input = self._prepare_actor_input(rollout_output)\n\n                        # === update policy ===\n                        with marked_timer(\"update_actor\", timing_raw, color=\"red\"):\n                            if actor_input is not None:\n                                actor_output = self.actor_rollout_wg.update_actor(actor_input)\n                            else:\n                                actor_output = self.actor_rollout_wg.update_actor(\n                                    DataProto(\n                                        meta_info={\n                                            \"empty_batch\": True,\n                                            \"global_steps\": self.global_steps,\n                                            \"global_token_num\": [0],\n                                        }\n                                    )\n                                )\n                        actor_output_metrics = reduce_metrics(actor_output.meta_info[\"metrics\"])\n                        metrics.update(actor_output_metrics)\n\n                    # === validate ===\n                    is_last_step = self.global_steps >= self.total_training_steps\n                    if (\n                        self.config.trainer.test_freq > 0\n                        and (is_last_step or self.global_steps % self.config.trainer.test_freq == 0)\n                        and self.global_steps >= self.config.actor_rollout_ref.actor.critic_warmup_steps\n                    ):\n                        with marked_timer(\"testing\", timing_raw, color=\"green\"):\n                            val_metrics: dict = self._validate()\n                            if is_last_step:\n                                last_val_metrics = val_metrics\n                        metrics.update(val_metrics)\n                        reset_future = None\n\n                    # === save checkpoint ===\n                    # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.\n                    esi_close_to_expiration = should_save_ckpt_esi(\n                        max_steps_duration=self.max_steps_duration,\n                        redundant_time=self.config.trainer.esi_redundant_time,\n                    )\n                    # Check if the conditions for saving a checkpoint are met.\n                    # The conditions include a mandatory condition (1) and\n                    # one of the following optional conditions (2/3/4):\n                    # 1. The save frequency is set to a positive value.\n                    # 2. It's the last training step.\n                    # 3. The current step number is a multiple of the save frequency.\n                    # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.\n                    if self.config.trainer.save_freq > 0 and (\n                        is_last_step\n                        or self.global_steps % self.config.trainer.save_freq == 0\n                        or esi_close_to_expiration\n                    ):\n                        if esi_close_to_expiration:\n                            print(\"Force saving checkpoint: ESI instance expiration approaching.\")\n                        with marked_timer(\"save_checkpoint\", timing_raw, color=\"green\"):\n                            self._save_checkpoint()\n\n                    # === stop profiling ===\n                    with marked_timer(\"stop_profile\", timing_raw):\n                        next_step_profile = (\n                            self.global_steps + 1 in self.config.global_profiler.steps\n                            if self.config.global_profiler.steps is not None\n                            else False\n                        )\n                        self._stop_profiling(\n                            curr_step_profile and not next_step_profile\n                            if self.config.global_profiler.profile_continuous_steps\n                            else curr_step_profile\n                        )\n                        prev_step_profile = curr_step_profile\n                        curr_step_profile = next_step_profile\n\n                    steps_duration = timing_raw[\"step\"]\n                    self.max_steps_duration = max(self.max_steps_duration, steps_duration)\n\n                    # === training metrics ===\n                    metrics.update(\n                        {\n                            \"training/global_step\": self.global_steps,\n                            \"training/epoch\": epoch,\n                        }\n                    )\n                    metrics.update({f\"timing_s/{name}\": value for name, value in timing_raw.items()})\n                    if actor_input is not None:\n                        metrics[\"data/trajectory_avg_reward\"] = actor_input.meta_info[\"data/trajectory_avg_reward\"]\n                        metrics[\"data/avg_positive_trajectory_length\"] = actor_input.meta_info[\n                            \"data/avg_positive_trajectory_length\"\n                        ]\n                    logger.log(data=metrics, step=self.global_steps)\n\n                    progress_bar.update(1)\n                    self.global_steps += 1\n\n                    if (\n                        hasattr(self.config.actor_rollout_ref.actor, \"profiler\")\n                        and self.config.actor_rollout_ref.actor.profiler.tool == \"torch_memory\"\n                    ):\n                        self.actor_rollout_wg.dump_memory_snapshot(\n                            tag=f\"post_update_step{self.global_steps}\", sub_dir=f\"step{self.global_steps}\"\n                        )\n\n                    if is_last_step:\n                        pprint(f\"Final validation metrics: {last_val_metrics}\")\n                        progress_bar.close()\n                        return\n\n    def _validate(self) -> dict:\n        metric_list = []\n        val_iter = iter(self.val_dataloader)\n        test_batch = self._next_rollout_batch(val_iter)\n        while test_batch is not None:\n            if len(test_batch) < self.config.data.val_batch_size:\n                print(f\"drop last batch in val_dataloader, len {len(test_batch)}\")\n                break\n\n            test_batch.meta_info[\"validate\"] = True\n            reset_future = self._reset_envs(test_batch)\n            rollout_output = self.async_rollout_manager.generate_sequences(test_batch, reset_future)\n            self._prepare_actor_input(rollout_output)\n            test_batch = self._next_rollout_batch(val_iter)\n            actor_input = self._prepare_actor_input(rollout_output)\n\n            metric_list.append(\n                {\n                    \"val/avg_reward\": actor_input.meta_info[\"data/trajectory_avg_reward\"],\n                    \"val/avg_positive_trajectory_length\": actor_input.meta_info[\"data/avg_positive_trajectory_length\"],\n                }\n            )\n\n        metrics = {}\n        if metric_list:\n            metrics[\"val/avg_reward\"] = np.mean([m[\"val/avg_reward\"] for m in metric_list])\n            metrics[\"val/avg_positive_trajectory_length\"] = np.mean(\n                [m[\"val/avg_positive_trajectory_length\"] for m in metric_list]\n            )\n\n        return metrics\n"
  },
  {
    "path": "verl/experimental/vla/workers/env/env_loop_wg_test.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\n\nimport ray\nfrom omegaconf import OmegaConf\n\nfrom verl import DataProto\nfrom verl.experimental.vla.naive_rollout_rob import NaiveRolloutRob\n\n# from verl.workers.env.env_worker import EnvWorker\nfrom verl.experimental.vla.workers.env.env_worker import EnvWorker\nfrom verl.single_controller.ray.base import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\nif not ray.is_initialized():\n    ray.init()\n\n    # for debugging\n    # ray.init(\n    #     runtime_env={\n    #         \"env_vars\": {\"RAY_DEBUG_POST_MORTEM\": \"1\"},\n    #     }\n    # )\n\nENV_WORKERS_NUM = 1\nSTAGE_NUM = 1\n# NUM_ENVS_PER_ITER = 32\n\n# NUM_ENVS_PER_STAGE = 8\n# NUM_ENVS_PER_ITER = STAGE_NUM * NUM_ENVS_PER_STAGE\n# NUM_ENVS_PER_ITER = 8\n# NUM_ENVS_PER_ITER = 32\nNUM_ENVS_PER_ITER = 2\nNUM_ENVS_PER_WORKER = NUM_ENVS_PER_ITER // ENV_WORKERS_NUM\n# NUM_ENVS_PER_WORKER_PER_STAGE = NUM_ENVS_PER_STAGE // ENV_WORKERS_NUM\nGROUP_SIZE = 2  # real group size = GROUP_SIZE * STAGE_NUM\nGROUP_NUM_PER_ITER = NUM_ENVS_PER_ITER * STAGE_NUM // GROUP_SIZE\nBATCH_SIZE_PER_GPU = 2\nNUM_ACTS_CHUNKS = 8\nMAX_EPISODE_STEPS = 32\nMAX_INFER_STEPS = MAX_EPISODE_STEPS // NUM_ACTS_CHUNKS\ncfg_dict = {\n    \"rollout\": {\"pipeline_stage_num\": STAGE_NUM},\n    \"train\": {\n        \"use_fixed_reset_state_ids\": False,\n        \"ignore_terminations\": False,\n        # \"auto_reset\": True,\n        \"auto_reset\": False,\n        \"max_episode_steps\": MAX_EPISODE_STEPS,\n        \"use_rel_reward\": False,\n        \"reward_coef\": 1.0,\n        \"only_eval\": False,\n        \"use_ordered_reset_state_ids\": False,\n        # \"num_images_in_input\": 1,\n        \"init_params\": {\n            \"camera_depths\": False,\n            \"camera_heights\": 256,\n            \"camera_widths\": 256,\n            \"camera_names\": [\"agentview\", \"robot0_eye_in_hand\"],\n        },\n        \"video_cfg\": {\n            \"save_video\": True,\n            \"video_base_dir\": \"/tmp/videos\",\n        },\n        \"task_suite_name\": \"libero_10\",\n        \"num_envs\": NUM_ENVS_PER_WORKER,\n        \"simulator_type\": \"isaac\",\n        \"seed\": 0,\n    },\n    \"enable_offload\": False,\n    \"actor\": {\"model\": {\"num_action_chunks\": NUM_ACTS_CHUNKS, \"action_dim\": 7}},\n    \"runner\": {\"only_eval\": False},\n}\nenv_cfg = OmegaConf.create(cfg_dict)\n\ngpu_pool = RayResourcePool([ENV_WORKERS_NUM], use_gpu=True)\n# RayEnvWorker = ray.remote(num_gpus=1)(EnvWorker)\nray_cls_with_init = RayClassWithInitArgs(cls=ray.remote(EnvWorker), config=env_cfg)\n\nenv_wg = RayWorkerGroup(gpu_pool, ray_cls_with_init)\n\n\ndef restructure_data_proto(data_proto: DataProto) -> list[DataProto]:\n    total_batch_size = len(data_proto)\n    tensors = data_proto.batch\n    non_tensors = data_proto.non_tensor_batch\n\n    full_image_tensor = tensors[\"full_image\"]\n    state_tensor = tensors[\"state\"]\n    task_descriptions_np = non_tensors[\"task_descriptions\"]\n    if total_batch_size != ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER:\n        raise ValueError(\n            f\"Total batch size {total_batch_size} does not match the expected size \"\n            f\"ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER = \"\n            f\"{ENV_WORKERS_NUM * STAGE_NUM * NUM_ENVS_PER_WORKER}\"\n        )\n\n    image_rest_shape = (ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER) + full_image_tensor.shape[1:]\n    state_rest_shape = (ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER) + state_tensor.shape[1:]\n    reshaped_full_image = full_image_tensor.view(image_rest_shape)\n    reshaped_state = state_tensor.view(state_rest_shape)\n\n    reshaped_task_descriptions = task_descriptions_np.reshape(ENV_WORKERS_NUM, STAGE_NUM, NUM_ENVS_PER_WORKER)\n    stages_data_list = []\n    for stage_idx in range(STAGE_NUM):\n        stage_images = reshaped_full_image[:, stage_idx, :]\n        stage_states = reshaped_state[:, stage_idx, :]\n        stage_tasks = reshaped_task_descriptions[:, stage_idx, :]\n        final_images = stage_images.reshape(ENV_WORKERS_NUM * NUM_ENVS_PER_WORKER, *full_image_tensor.shape[1:])\n        final_states = stage_states.reshape(ENV_WORKERS_NUM * NUM_ENVS_PER_WORKER, *state_tensor.shape[1:])\n        final_tasks = stage_tasks.flatten().tolist()\n\n        stage_dp = DataProto.from_dict(\n            tensors={\"full_image\": final_images, \"state\": final_states},\n            non_tensors={\"task_descriptions\": final_tasks},\n            meta_info={\"do_sample\": True, \"temperature\": 1.6, \"prompt_length\": 512},\n        )\n        stages_data_list.append(stage_dp)\n    return stages_data_list\n\n\nasync def run():\n    # breakpoint()\n    env_wg.init_worker()\n    env_wg.init_simulator()\n\n    reset_state_ids_tensordict = DataProto.from_dict(\n        non_tensors={\"state_ids\": [0] * NUM_ENVS_PER_ITER * STAGE_NUM, \"task_ids\": [0] * NUM_ENVS_PER_ITER * STAGE_NUM}\n    )\n\n    reset_result = env_wg.reset_envs_to_state_ids(reset_state_ids_tensordict)\n    print(f\"reset_envs_to_state_ids result: {reset_result}\")\n    stages_data_list = restructure_data_proto(reset_result)\n\n    RayNaiveRolloutRob = ray.remote(num_gpus=1)(NaiveRolloutRob)\n\n    model_config = {\"path\": \"Haozhan72/Openvla-oft-SFT-libero10-trajall\"}\n    rollout_workers = RayNaiveRolloutRob.remote(model_config)\n\n    env_obs_refs = {}\n    rollout_refs = {}\n    traj = [[], []]\n\n    for _ in range(MAX_INFER_STEPS):\n        for stage_id in range(STAGE_NUM):\n            if _ == 0:\n                rollout_refs[stage_id] = rollout_workers.generate_sequences.remote(stages_data_list[stage_id])\n            else:\n                # env_batch = env_obs_refs[stage_id]\n                env_batch: DataProto = env_obs_refs[stage_id].get()\n                env_batch_traj = env_batch.select(batch_keys=[\"rews\", \"terminations\", \"truncations\"])\n                traj[stage_id][-1].update({\"env\": env_batch_traj})\n                obs = env_batch\n                obs.meta_info.update({\"do_sample\": True, \"temperature\": 1.6, \"prompt_length\": 512})\n                rollout_refs[stage_id] = rollout_workers.generate_sequences.remote(obs)\n        for stage_id in range(STAGE_NUM):\n            batch: DataProto = ray.get(rollout_refs[stage_id])\n            traj[stage_id].append({\"model\": batch})\n            action = batch.batch[\"action\"]\n            action = action.cpu().numpy()\n            # already in env\n            data = DataProto.from_dict(non_tensors={\"actions\": action}, meta_info={\"stage_id\": stage_id})\n            env_obs_refs[stage_id] = env_wg.env_interact_step(data)\n\n    env_wg.finish_rollout()\n\n\nasyncio.run(run())\nray.timeline(filename=\"2stage_pipeline_timeline_wg.json\")\n"
  },
  {
    "path": "verl/experimental/vla/workers/env/env_manager.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport gc\nimport logging\nimport os\nimport subprocess\nfrom typing import Optional\n\nimport torch\nimport torch.multiprocessing as mp\n\nfrom verl.utils.device import get_torch_device\n\nlogger = logging.getLogger(__name__)\n\n\ndef cleanup_device_tensors():\n    gc.collect()\n    get_torch_device().empty_cache()\n\n\ndef get_gpu_numa_node(gpu_id: int) -> int:\n    try:\n        try:\n            import pynvml\n\n            pynvml.nvmlInit()\n            handle = pynvml.nvmlDeviceGetHandleByIndex(gpu_id)\n            # Get PCI bus info\n            pci_info = pynvml.nvmlDeviceGetPciInfo(handle)\n            pci_bus_id = pci_info.busId\n        except ImportError:\n            # Fallback to nvidia-smi\n            result = subprocess.run(\n                [\n                    \"nvidia-smi\",\n                    \"--query-gpu=pci.bus_id\",\n                    \"--format=csv,noheader,nounits\",\n                    f\"--id={gpu_id}\",\n                ],\n                capture_output=True,\n                text=True,\n                check=True,\n            )\n            pci_bus_id = result.stdout.strip()\n\n        # Extract bus number from PCI bus ID (format: 0000:XX:YY.Z)\n        bus_number = pci_bus_id.split(\":\")[1]\n\n        # Get NUMA node from sysfs\n        numa_node_path = f\"/sys/bus/pci/devices/0000:{bus_number}:00.0/numa_node\"\n        if os.path.exists(numa_node_path):\n            with open(numa_node_path) as f:\n                numa_node = int(f.read().strip())\n                if numa_node >= 0:\n                    return numa_node\n\n        # Fallback: try to get from lscpu\n        result = subprocess.run([\"lscpu\"], capture_output=True, text=True, check=True)\n        numa_nodes = 0\n        for line in result.stdout.split(\"\\n\"):\n            if \"NUMA node(s):\" in line:\n                numa_nodes = int(line.split(\":\")[1].strip())\n                break\n\n        # If we can't determine the exact NUMA node, distribute evenly\n        return gpu_id % numa_nodes if numa_nodes > 0 else 0\n\n    except Exception as e:\n        logger.error(f\"Warning: Could not determine NUMA node for GPU {gpu_id}: {e}\")\n        return 0\n\n\ndef get_numa_cpus(numa_node: int) -> list:\n    try:\n        # Read from sysfs\n        cpulist_path = f\"/sys/devices/system/node/node{numa_node}/cpulist\"\n        if os.path.exists(cpulist_path):\n            with open(cpulist_path) as f:\n                cpulist = f.read().strip()\n\n            # Parse CPU list (e.g., \"0-7,16-23\" or \"0,1,2,3\")\n            cpus = []\n            for part in cpulist.split(\",\"):\n                if \"-\" in part:\n                    start, end = map(int, part.split(\"-\"))\n                    cpus.extend(range(start, end + 1))\n                else:\n                    cpus.append(int(part))\n            return cpus\n    except Exception as e:\n        logger.error(f\"Warning: Could not get CPU list for NUMA node {numa_node}: {e}\")\n\n    # Fallback: return all available CPUs\n    return list(range(os.cpu_count() or 1))\n\n\ndef set_process_numa_affinity(gpu_id: int) -> None:\n    try:\n        numa_node = get_gpu_numa_node(gpu_id)\n        cpus = get_numa_cpus(numa_node)\n\n        if not cpus:\n            logger.error(f\"Warning: No CPUs found for NUMA node {numa_node}\")\n            return\n\n        os.sched_setaffinity(0, cpus)\n        try:\n            subprocess.run(\n                [\"numactl\", \"--membind\", str(numa_node), \"--\"],\n                check=False,\n                capture_output=True,\n            )\n        except FileNotFoundError:\n            pass  # numactl not available, that's ok\n\n    except Exception as e:\n        logger.error(f\"Warning: Could not set NUMA affinity for GPU {gpu_id}: {e}\")\n\n\ndef recursive_to_own(obj):\n    if isinstance(obj, torch.Tensor):\n        return obj.clone() if obj.is_shared() else obj\n    elif isinstance(obj, list):\n        return [recursive_to_own(elem) for elem in obj]\n    elif isinstance(obj, tuple):\n        return tuple(recursive_to_own(elem) for elem in obj)\n    elif isinstance(obj, dict):\n        return {k: recursive_to_own(v) for k, v in obj.items()}\n    else:\n        return obj\n\n\nclass EnvManager:\n    def __init__(self, cfg, rank, world_size, env_cls, stage_id: int = 0):\n        self.cfg = cfg\n        self.rank = rank\n        self.world_size = world_size\n        self.stage_id = stage_id\n        self.process: Optional[mp.Process] = None\n        self.command_queue: Optional[mp.Queue] = None\n        self.result_queue: Optional[mp.Queue] = None\n        self.state_buffer: Optional[bytes] = None\n\n        self.env_cls = env_cls\n\n    def start_simulator(self):\n        \"\"\"Start simulator process with shared memory queues\"\"\"\n        if self.process:\n            logger.info(f\"Simulator process already running for rank {self.rank}\")\n            return\n\n        self.context = mp.get_context(\"spawn\")\n        # Create shared memory queues\n        self.command_queue = self.context.Queue()\n        self.result_queue = self.context.Queue()\n\n        # Start simulator process\n        self.process = self.context.Process(\n            target=_simulator_worker,\n            args=(\n                self.cfg,\n                self.rank,\n                self.world_size,\n                self.stage_id,\n                self.env_cls,\n                self.command_queue,\n                self.result_queue,\n                self.state_buffer,\n                True,\n            ),\n        )\n        self.process.start()\n\n        # Wait for initialization\n        result = self.result_queue.get(timeout=180)\n        if result[\"status\"] != \"ready\":\n            raise RuntimeError(f\"Simulator initialization failed: {result}\")\n\n    def stop_simulator(self):\n        if not self.process:\n            return\n\n        # Request state save\n        self.command_queue.put({\"method\": \"get_state\", \"args\": [], \"kwargs\": {}})\n\n        # Get saved state\n        result = self.result_queue.get(timeout=180)\n        if result[\"status\"] == \"success\":\n            self.state_buffer = result[\"data\"]\n\n        self.command_queue.put({\"method\": \"shutdown\"})\n        self.command_queue.close()\n        self.result_queue.close()\n        self.command_queue = None\n        self.result_queue = None\n        self.process.join(timeout=5)\n\n        self.command_queue = None\n        self.result_queue = None\n        if self.process.is_alive():\n            self.process.terminate()\n            self.process.join()\n\n        self.process = None\n\n    def __getattr__(self, name):\n        if name in [\n            \"cfg\",\n            \"rank\",\n            \"world_size\",\n            \"stage_id\",\n            \"process\",\n            \"command_queue\",\n            \"result_queue\",\n            \"state_buffer\",\n            \"env_cls\",\n            \"context\",\n        ]:\n            return super().__getattr__(name)\n\n        def method_proxy(*args, **kwargs):\n            if self.process is None or not self.process.is_alive():\n                raise RuntimeError(\"Simulator not running\")\n\n            args = recursive_to_own(args)\n            kwargs = recursive_to_own(kwargs)\n            self.command_queue.put({\"method\": name, \"args\": args, \"kwargs\": kwargs})\n\n            result = self.result_queue.get()\n            result = recursive_to_own(result)\n            if result[\"status\"] == \"error\":\n                raise Exception(result[\"error\"])\n            return result[\"data\"]\n\n        return method_proxy\n\n    def get_all_state_ids(self):\n        \"\"\"Get all available state IDs from the environment.\"\"\"\n        if self.process is None or not self.process.is_alive():\n            raise RuntimeError(\"Simulator not running\")\n\n        self.command_queue.put({\"method\": \"get_all_state_ids\", \"args\": [], \"kwargs\": {}})\n        result = self.result_queue.get()\n        result = recursive_to_own(result)\n        if result[\"status\"] == \"error\":\n            raise Exception(result[\"error\"])\n        return result[\"data\"]\n\n    def reset_envs_to_state_ids(self, state_ids_list, task_ids_list):\n        \"\"\"Reset environments to specified state IDs.\"\"\"\n        if self.process is None or not self.process.is_alive():\n            raise RuntimeError(\"Simulator not running\")\n\n        state_ids_list = recursive_to_own(state_ids_list)\n        task_ids_list = recursive_to_own(task_ids_list)\n\n        self.command_queue.put(\n            {\n                \"method\": \"reset_envs_to_state_ids\",\n                \"args\": [state_ids_list, task_ids_list],\n                \"kwargs\": {},\n            }\n        )\n\n        result = self.result_queue.get()\n        result = recursive_to_own(result)\n        if result[\"status\"] == \"error\":\n            raise Exception(result[\"error\"])\n        return result[\"data\"]\n\n    def __setattr__(self, name, value):\n        # Handle special attributes that should be set on self\n        if name in [\n            \"cfg\",\n            \"rank\",\n            \"world_size\",\n            \"stage_id\",\n            \"process\",\n            \"command_queue\",\n            \"result_queue\",\n            \"state_buffer\",\n            \"env_cls\",\n            \"context\",\n        ]:\n            super().__setattr__(name, value)\n            return\n\n        if self.process is None or not self.process.is_alive():\n            raise RuntimeError(f\"Simulator not running to set attribute {name} to {value}\")\n\n        value = recursive_to_own(value)\n        self.command_queue.put(\n            {\n                \"method\": \"__setattr__\",\n                \"args\": [name, value],\n                \"kwargs\": {},\n            }\n        )\n\n        result = self.result_queue.get()\n        result = recursive_to_own(result)\n        if result[\"status\"] == \"error\":\n            raise Exception(result[\"error\"])\n\n\ndef _simulator_worker(\n    cfg,\n    rank,\n    world_size,\n    stage_id,\n    env_cls,\n    command_queue,\n    result_queue,\n    state_buffer,\n    bind_numa=True,\n):\n    \"\"\"Worker process for simulator\"\"\"\n    # Set NUMA affinity for the process to match the GPU rank\n    import logging\n    import os\n\n    pid = os.getpid()\n    logger = logging.getLogger(f\"simulator_worker_{rank}_{pid}\")\n\n    if bind_numa:\n        set_process_numa_affinity(rank)\n    try:\n        try:\n            env = env_cls(cfg, rank, world_size, stage_id=stage_id)\n        except TypeError:\n            env = env_cls(cfg, rank, world_size)\n\n        if state_buffer:\n            env.load_state(state_buffer)\n\n        # Signal ready\n        result_queue.put({\"status\": \"ready\"})\n\n        # Main command processing loop\n        while True:\n            try:\n                command = command_queue.get()\n                logger.debug(f\"Received command method: {command['method']}\")\n\n                if command[\"method\"] == \"shutdown\":\n                    env.close()\n                    break\n\n                method_name = command[\"method\"]\n                args = command.get(\"args\", [])\n                kwargs = command.get(\"kwargs\", {})\n                if method_name == \"__setattr__\":\n                    # Handle attribute setting\n                    attr_name, attr_value = args\n                    setattr(env, attr_name, attr_value)\n                    result_queue.put({\"status\": \"success\", \"data\": None})\n                elif hasattr(env, method_name):\n                    method = getattr(env, method_name)\n                    assert callable(method), f\"Method {method_name} is not callable\"\n                    result = method(*args, **kwargs)\n                    result_queue.put({\"status\": \"success\", \"data\": result})\n                else:\n                    logger.error(f\"Method '{method_name}' not found\")\n                    result_queue.put(\n                        {\n                            \"status\": \"error\",\n                            \"error\": f\"Method '{method_name}' not found\",\n                        }\n                    )\n\n            except Exception as e:\n                logger.exception(e)\n                result_queue.put({\"status\": \"error\", \"error\": str(e)})\n\n    except Exception as e:\n        logger.exception(e)\n        result_queue.put({\"status\": \"error\", \"error\": str(e)})\n\n    finally:\n        command_queue.close()\n        result_queue.close()\n"
  },
  {
    "path": "verl/experimental/vla/workers/env/env_worker.py",
    "content": "# Copyright 2025 The RLinf Authors.\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     https://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport itertools\n\nimport torch\nfrom omegaconf import DictConfig\nfrom torch.distributed.device_mesh import init_device_mesh\n\nfrom verl import DataProto\nfrom verl.experimental.vla.workers.env.env_manager import EnvManager\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import (\n    get_device_name,\n)\nfrom verl.utils.distributed import initialize_global_process_group_ray\nfrom verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig\n\n\ndef put_tensor_cpu(data_dict):\n    for key, value in data_dict.items():\n        if isinstance(value, dict):\n            data_dict[key] = put_tensor_cpu(value)\n        if isinstance(value, torch.Tensor):\n            data_dict[key] = value.cpu().contiguous()\n    return data_dict\n\n\ndef create_env_batch(obs, rews, dones, infos, meta=None):\n    ret_dict = {\"obs\": obs, \"rews\": rews, \"dones\": dones, \"infos\": infos}\n    if meta is not None:\n        ret_dict.update(meta=meta)\n\n    ret_dict = put_tensor_cpu(ret_dict)\n    return ret_dict\n\n\ndef create_env_batch_dataproto(obs, rews, terminations, truncations, infos, meta=None):\n    ret_dict = {\"obs\": obs, \"rews\": rews, \"terminations\": terminations, \"truncations\": truncations, \"infos\": infos}\n    if meta is not None:\n        ret_dict.update(meta=meta)\n\n    ret_dict = put_tensor_cpu(ret_dict)\n    tensor_batch = {\n        \"full_image\": ret_dict[\"obs\"][\"images_and_states\"][\"full_image\"],\n        \"wrist_image\": ret_dict[\"obs\"][\"images_and_states\"][\"wrist_image\"],\n        \"state\": ret_dict[\"obs\"][\"images_and_states\"][\"state\"],\n        \"rews\": ret_dict[\"rews\"],\n        \"terminations\": ret_dict[\"terminations\"],\n        \"truncations\": ret_dict[\"truncations\"],\n    }\n    non_tensor_batch = {\"task_descriptions\": obs[\"task_descriptions\"]}\n    output = DataProto.from_dict(tensors=tensor_batch, non_tensors=non_tensor_batch)\n\n    return output\n\n\nclass EnvWorker(Worker, DistProfilerExtension):\n    def __init__(self, config: DictConfig):\n        Worker.__init__(self)\n        self.cfg = config\n        self.train_video_cnt = 0\n        self.eval_video_cnt = 0\n\n        self.simulator_list = []\n        self.last_obs_list = []\n        self.last_dones_list = []\n        self.eval_simulator_list = []\n\n        self.stage_num = self.cfg.rollout.pipeline_stage_num\n        initialize_global_process_group_ray(timeout_second=None)\n        device_name = get_device_name()\n        env_device_mesh = init_device_mesh(device_name, mesh_shape=(self.world_size, 1), mesh_dim_names=[\"dp\", \"tp\"])\n        self._register_dispatch_collect_info(\"env\", dp_rank=env_device_mesh[\"dp\"].get_local_rank(), is_collect=True)\n\n        # Initialize profiler\n        omega_profiler_config = config.train.get(\"profiler\", {})\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    @DistProfiler.annotate(color=\"green\", role=\"env_init\")\n    def init_worker(self):\n        if self.cfg.train.simulator_type == \"libero\":\n            from verl.experimental.vla.envs.libero_env.libero_env import LiberoEnv\n\n            for stage_id in range(self.stage_num):\n                self.simulator_list.append(\n                    EnvManager(\n                        self.cfg.train,\n                        rank=self._rank,\n                        world_size=self._world_size,\n                        env_cls=LiberoEnv,\n                        stage_id=stage_id,\n                    )\n                )\n\n        elif self.cfg.train.simulator_type == \"isaac\":\n            from verl.experimental.vla.envs.isaac_env.isaac_env import IsaacEnv\n\n            for stage_id in range(self.stage_num):\n                self.simulator_list.append(\n                    EnvManager(\n                        self.cfg.train,\n                        rank=self._rank,\n                        world_size=self._world_size,\n                        env_cls=IsaacEnv,\n                        stage_id=stage_id,\n                    )\n                )\n        else:\n            raise NotImplementedError(f\"Simulator type {self.cfg.train.simulator_type} not implemented\")\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    @DistProfiler.annotate(color=\"green\", role=\"env_init_simulator\")\n    def init_simulator(self):\n        for i in range(self.stage_num):\n            self.simulator_list[i].start_simulator()\n        return\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"env\"), blocking=False)\n    @DistProfiler.annotate(color=\"red\", role=\"env_interact_step\")\n    def env_interact_step(self, data: DataProto) -> dict:\n        \"\"\"\n        This function is used to interact with the environment.\n        \"\"\"\n        chunk_actions: torch.Tensor = data.non_tensor_batch[\"actions\"]\n        chunk_values = data.non_tensor_batch[\"critic_values\"]\n        stage_id: int = data.meta_info[\"stage_id\"]\n\n        # Pi0.5 Libero is not required\n        # TODO: prepare actions according to simulator type\n        # chunk_actions = prepare_actions(\n        #     simulator_type=self.cfg.train.simulator_type,\n        #     raw_chunk_actions=chunk_actions,\n        #     num_action_chunks=self.cfg.actor.model.num_action_chunks,\n        #     action_dim=self.cfg.actor.model.action_dim,\n        # )\n\n        env_info_list = {}\n\n        extracted_obs, chunk_rewards, chunk_terminations, chunk_truncations, infos = self.simulator_list[\n            stage_id\n        ].chunk_step(chunk_actions, chunk_values=chunk_values)\n        chunk_dones = torch.logical_or(chunk_terminations, chunk_truncations)\n\n        if chunk_dones.any():\n            if \"final_info\" in infos:\n                final_info = infos[\"final_info\"]\n                for key in final_info[\"episode\"]:\n                    env_info_list[key] = final_info[\"episode\"][key][chunk_dones[:, -1]].cpu()\n\n        env_batch = create_env_batch_dataproto(\n            obs=extracted_obs,\n            rews=chunk_rewards,\n            terminations=chunk_terminations,\n            truncations=chunk_truncations,\n            infos=infos,\n            meta=env_info_list,\n        )\n        return env_batch\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def get_all_state_ids(self):\n        \"\"\"Get all available state IDs from the environment.\"\"\"\n        state_ids = self.simulator_list[0].get_all_state_ids()\n        return state_ids\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"env\"), blocking=False)\n    @DistProfiler.annotate(color=\"blue\", role=\"env_reset_envs_to_state_ids\")\n    def reset_envs_to_state_ids(self, data: DataProto):\n        \"\"\"Reset environments to specified state IDs.\n\n        Args:\n            state_ids: State IDs to reset environments to\n        \"\"\"\n        state_ids_list = list(data.non_tensor_batch[\"state_ids\"])\n        task_ids_list = list(data.non_tensor_batch[\"task_ids\"])\n\n        assert len(state_ids_list) == self.cfg.train.num_envs * self.stage_num, (\n            f\"state_ids_list length is {len(state_ids_list)}, but should be {self.cfg.train.num_envs * self.stage_num}\"\n        )\n        result_list = []\n        for stage_id in range(self.stage_num):\n            if self.cfg.train.simulator_type == \"isaac\":\n                assert (\n                    len(\n                        set(\n                            state_ids_list[\n                                stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs\n                            ]\n                        )\n                    )\n                    == 1\n                ), \"rollout.n should equal to num_envs for isaac\"\n\n            result = self.simulator_list[stage_id].reset_envs_to_state_ids(\n                state_ids_list[stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs],\n                task_ids_list[stage_id * self.cfg.train.num_envs : (stage_id + 1) * self.cfg.train.num_envs],\n            )\n            result_list.append(result)\n        output_tensor_dict = {}\n        output_non_tensor_dict = {}\n\n        # Handle nested 'images_and_states'\n        images_and_states_list = [d[0][\"images_and_states\"] for d in result_list]\n        if images_and_states_list:\n            # Assuming all dicts in the list have the same keys\n            for k in images_and_states_list[0].keys():\n                if isinstance(images_and_states_list[0][k], torch.Tensor):\n                    output_tensor_dict[k] = torch.cat([d[k] for d in images_and_states_list])\n\n        # Handle 'task_descriptions'\n        task_descriptions_list = [d[0][\"task_descriptions\"] for d in result_list]\n        output_non_tensor_dict[\"task_descriptions\"] = list(itertools.chain.from_iterable(task_descriptions_list))\n\n        output = DataProto.from_dict(tensors=output_tensor_dict, non_tensors=output_non_tensor_dict)\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    @DistProfiler.annotate(color=\"gray\", role=\"env_finish_rollout\")\n    def finish_rollout(self, mode=\"train\"):\n        # reset\n        if mode == \"train\":\n            if self.cfg.train.video_cfg.save_video:\n                for i in range(self.stage_num):\n                    self.simulator_list[i].flush_video(video_sub_dir=f\"stage_{i}\")\n"
  },
  {
    "path": "verl/interactions/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/interactions/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\n\nclass BaseInteraction:\n    def __init__(self, config: dict[str, Any]):\n        self.config = config\n        self.name: str = config.get(\"name\", \"interaction_agent\")  # More general agent default role name\n\n    async def start_interaction(self, instance_id: Optional[str] = None, **kwargs) -> str:\n        \"\"\"Create a tool instance.\n\n        Args:\n            instance_id: The instance id of the tool.\n\n        Returns:\n            The instance id of the tool.\n        \"\"\"\n        if instance_id is None:\n            return str(uuid4())\n        else:\n            return instance_id\n\n    async def generate_response(\n        self, instance_id: str, messages: list[dict[str, Any]], **kwargs\n    ) -> tuple[bool, str, float, dict[str, Any]]:  # More clear response generation method\n        \"\"\"\n        Generates a response for the current turn of interaction.\n        Returns a tuple containing:\n        - should_terminate_sequence (bool): True if the interaction sequence should end.\n        - response_content (str): The textual content of the response.\n        - current_turn_score (float): The score for this specific turn/response.\n        - additional_data (dict): Any extra information or metadata.\n        \"\"\"\n        should_terminate_sequence: bool = False  # if True, end rollout\n        response_content: str = \"Your current result seems acceptable.\"\n        current_turn_score: float = 0.8\n        additional_data: dict[str, Any] = {}\n        return should_terminate_sequence, response_content, current_turn_score, additional_data\n\n    async def calculate_score(self) -> float:  # More clear score calculation method\n        \"\"\"\n        Calculates a score for the interaction,\n        potentially considering aspects like partial exposure & in-context task switching.\n        should be invoke at turn-level\n        \"\"\"\n        # ...implement the logic to calculate turn-level score...\n        score = 0.0\n        return score\n\n    async def finalize_interaction(self) -> None:  # More clear interaction end and resource release method\n        \"\"\"\n        Finalizes the interaction session and releases any associated state or resources.\n        Simulates: release state\n        \"\"\"\n        # ...implement the logic to release state...\n        pass\n"
  },
  {
    "path": "verl/interactions/gsm8k_interaction.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom verl.utils.reward_score import gsm8k\n\nfrom .base import BaseInteraction\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass Gsm8kInteraction(BaseInteraction):\n    \"\"\"A demo interaction for calculating the reward of gsm8k.\n\n    - `start_interaction`: start a interaction instance for a trajectory.\n    - `generate_response`: generate the response of the assistant.\n    - `calculate_score`: calculate the score of the interaction.\n    - `finalize_interaction`: finalize the interaction instance.\n    \"\"\"\n\n    def __init__(self, config: dict):\n        super().__init__(config)\n        self._instance_dict = {}\n\n    async def start_interaction(\n        self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs\n    ) -> str:\n        if instance_id is None:\n            instance_id = str(uuid4())\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"ground_truth\": ground_truth,\n            \"reward\": 0.0,\n        }\n        return instance_id\n\n    async def generate_response(\n        self, instance_id: str, messages: list[dict[str, Any]], **kwargs\n    ) -> tuple[bool, str, float, dict]:\n        content = \"\"\n        for i in range(len(messages) - 1, -1, -1):\n            item = messages[i]\n            if item.get(\"role\") == \"assistant\":\n                content = item.get(\"content\")\n                break\n\n        self._instance_dict[instance_id][\"response\"] = content\n\n        reward = await self.calculate_score(instance_id)\n        if reward == 1.0:\n            response = \"Your response is correct!\"\n            should_terminate_sequence = True\n        else:\n            response = \"Your response is incorrect! You need to reflect on your answer and try again.\"\n            should_terminate_sequence = False\n\n        return should_terminate_sequence, response, reward, {}\n\n    async def calculate_score(self, instance_id: str, **kwargs) -> float:\n        return gsm8k.compute_score(\n            self._instance_dict[instance_id][\"response\"],\n            self._instance_dict[instance_id][\"ground_truth\"],\n            method=\"strict\",\n            format_score=0.0,\n            score=1.0,\n        )\n\n    async def finalize_interaction(self, instance_id: str, **kwargs) -> None:\n        del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/interactions/utils/__init__.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/interactions/utils/interaction_registry.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport importlib.util\nimport logging\nimport os\nimport sys\n\nfrom omegaconf import OmegaConf\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef get_interaction_class(cls_name):\n    \"\"\"Dynamically import and return the interaction class.\"\"\"\n    module_name, class_name = cls_name.rsplit(\".\", 1)\n    if module_name not in sys.modules:\n        spec = importlib.util.find_spec(module_name)\n        module = importlib.util.module_from_spec(spec)\n        sys.modules[module_name] = module\n        spec.loader.exec_module(module)\n    else:\n        module = sys.modules[module_name]\n\n    interaction_cls = getattr(module, class_name)\n    return interaction_cls\n\n\ndef initialize_interactions_from_config(interaction_config_file):\n    \"\"\"Initialize interactions from configuration file.\n\n    Args:\n        interaction_config_file: Path to the interaction configuration file.\n\n    Returns:\n        dict: A dictionary mapping interaction names to BaseInteraction instances.\n    \"\"\"\n    interaction_config = OmegaConf.load(interaction_config_file)\n    interaction_map = {}\n\n    for interaction_item in interaction_config.interaction:\n        cls_name = interaction_item.class_name\n        interaction_cls = get_interaction_class(cls_name)\n\n        # Extract config and name\n        config = OmegaConf.to_container(interaction_item.config, resolve=True)\n\n        # Get the interaction name - either from config or derive from class name\n        name = interaction_item.get(\"name\", None)\n        if name is None:\n            # If no name is specified, use the class name as default\n            class_simple_name = cls_name.split(\".\")[-1]\n            # Remove \"Interaction\" suffix if present, otherwise use full class name\n            if class_simple_name.endswith(\"Interaction\"):\n                name = class_simple_name[:-11].lower()  # Remove \"Interaction\" (11 chars)\n            else:\n                name = class_simple_name.lower()\n\n        # Check for duplicate names\n        if name in interaction_map:\n            raise ValueError(f\"Duplicate interaction name '{name}' found. Each interaction must have a unique name.\")\n\n        # Inject the name into the config\n        config[\"name\"] = name\n\n        # Create the interaction instance\n        interaction = interaction_cls(config=config)\n        interaction_map[name] = interaction\n\n        logger.info(f\"Initialized interaction '{name}' with class '{cls_name}'\")\n\n    return interaction_map\n"
  },
  {
    "path": "verl/interactions/weather_interaction.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom .base import BaseInteraction\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass WeatherInteraction(BaseInteraction):\n    \"\"\"A demo interaction for handling weather-related queries.\n\n    - `start_interaction`: start a interaction instance for a trajectory.\n    - `generate_response`: generate the response of the assistant.\n    - `calculate_score`: calculate the score of the interaction.\n    - `finalize_interaction`: finalize the interaction instance.\n    \"\"\"\n\n    def __init__(self, config: dict):\n        super().__init__(config)\n        self._instance_dict = {}\n\n    async def start_interaction(\n        self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs\n    ) -> str:\n        if instance_id is None:\n            instance_id = str(uuid4())\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"ground_truth\": ground_truth,\n            \"reward\": 0.0,\n        }\n        return instance_id\n\n    async def generate_response(\n        self, instance_id: str, messages: list[dict[str, Any]], **kwargs\n    ) -> tuple[bool, str, float, dict]:\n        content = \"no tool call\"\n        for i in range(len(messages) - 1, -1, -1):\n            item = messages[i]\n            if item.get(\"role\") == \"tool\":\n                content = item.get(\"content\")\n                break\n        self._instance_dict[instance_id][\"response\"] = content\n\n        reward = await self.calculate_score(instance_id)\n        if reward == 1.0:\n            response = \"Thank you for your weather query!\"\n            should_terminate_sequence = True\n        else:\n            response = \"Please use the weather tool to get the weather information.\"\n            should_terminate_sequence = True\n        return should_terminate_sequence, response, reward, {}\n\n    async def calculate_score(self, instance_id: str, **kwargs) -> float:\n        # For weather interaction, we can implement a more complex scoring logic\n        # For now, we'll just return a default score of 1.0\n        if self._instance_dict[instance_id][\"response\"] == \"no tool call\":\n            return 0.0\n        return 1.0\n\n    async def finalize_interaction(self, instance_id: str, **kwargs) -> None:\n        del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/model_merger/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/model_merger/__main__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nThis module is used to merge huggingface model and test verl checkpoints from FSDP and Megatron backends.\n\nTo merge FSDP checkpoints:\n```sh\npython -m verl.model_merger merge \\\n    --backend fsdp \\\n    --local_dir checkpoints/verl_fsdp_gsm8k_examples/qwen2_5_0b5_fsdp_saveload/global_step_1/actor \\\n    --target_dir /path/to/merged_hf_model\n```\n\nTo merge Megatron checkpoints:\n```sh\npython -m verl.model_merger merge \\\n    --backend megatron \\\n    --tie-word-embedding \\\n    --local_dir checkpoints/verl_megatron_gsm8k_examples/qwen2_5_0b5_megatron_saveload/global_step_1/actor \\\n    --target_dir /path/to/merged_hf_model\n```\n\nor use distribtued merge for large models like dpskv3 671B\n\n```sh\ntorchrun --nproc_per_node 1 --nnodes 8 --node_rank ${RANK} -m verl.model_merger merge\\\n    --backend megatron \\\n    --local_dir ./checkpoints/global_step_1/actor \\\n    --target_dir /path/to/merged_hf_model\n```\n\n\nFor more details, please refer to documentation:\nhttps://verl.readthedocs.io/en/latest/advance/checkpoint.html#convert-fsdp-and-megatron-checkpoints-to-huggingface-format-model\n\"\"\"\n\nfrom .base_model_merger import generate_config_from_args, parse_args\n\n\ndef main():\n    args = parse_args()\n    config = generate_config_from_args(args)\n    print(f\"config: {config}\")\n\n    if config.backend == \"fsdp\":\n        from .fsdp_model_merger import FSDPModelMerger\n\n        merger = FSDPModelMerger(config)\n    elif config.backend == \"megatron\":\n        from .megatron_model_merger import MegatronModelMerger\n\n        merger = MegatronModelMerger(config)\n    else:\n        raise NotImplementedError(f\"Unknown backend: {config.backend}\")\n\n    merger.merge_and_save()\n    merger.cleanup()\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/model_merger/base_model_merger.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport argparse\nimport os\nimport warnings\nfrom abc import ABC, abstractmethod\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nimport torch\nfrom accelerate import init_empty_weights\nfrom transformers import (\n    AutoConfig,\n    AutoModelForCausalLM,\n    AutoModelForTokenClassification,\n    GenerationConfig,\n)\n\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\n\nAutoModelForVision2Seq = get_auto_model_for_vision2seq()\n\n\ndef parse_args():\n    parser = argparse.ArgumentParser(description=\"verl model merger\")\n    subparsers = parser.add_subparsers(dest=\"operation\", required=True, help=\"Specify 'merge' or 'test' operation.\")\n\n    base_op_parser = argparse.ArgumentParser(add_help=False)\n    base_op_parser.add_argument(\n        \"--backend\", type=str, required=True, choices=[\"fsdp\", \"megatron\"], help=\"The backend of the model\"\n    )\n    base_op_parser.add_argument(\"--local_dir\", type=str, default=None, help=\"Path to the saved model checkpoints.\")\n    base_op_parser.add_argument(\n        \"--tie-word-embedding\",\n        action=\"store_true\",\n        help=\"Whether to tie word embedding weights (currently only Megatron supported)\",\n    )\n    base_op_parser.add_argument(\"--trust-remote-code\", action=\"store_true\", help=\"Whether to trust remote code\")\n    base_op_parser.add_argument(\n        \"--is-value-model\",\n        action=\"store_true\",\n        help=\"Whether the model is a value model (currently only Megatron supported)\",\n    )\n    base_op_parser.add_argument(\n        \"--use_cpu_initialization\",\n        action=\"store_true\",\n        help=\"Whether to use CPU initialization for the model. This is useful for large models that cannot \"\n        \"fit into GPU memory during initialization.\",\n    )\n\n    merge_parser = subparsers.add_parser(\"merge\", parents=[base_op_parser], help=\"Merge model checkpoints and save.\")\n    merge_parser.add_argument(\n        \"--target_dir\", default=\"tmp\", type=str, help=\"Directory to save the merged huggingface model\"\n    )\n    merge_parser.add_argument(\n        \"--hf_upload_path\", default=None, type=str, help=\"Hugging Face repository ID to upload the model\"\n    )\n    merge_parser.add_argument(\n        \"--private\", action=\"store_true\", help=\"Whether to upload the model to a private Hugging Face repository\"\n    )\n\n    test_parser = subparsers.add_parser(\n        \"test\", parents=[base_op_parser], help=\"Test merged model against a reference Hugging Face model\"\n    )\n    test_parser.add_argument(\n        \"--test_hf_dir\", type=str, required=True, help=\"Path to the reference Hugging Face model directory for testing\"\n    )\n\n    args = parser.parse_args()\n    return args\n\n\n@dataclass\nclass ModelMergerConfig:\n    \"\"\"Configuration for model merger operations.\n\n    Args:\n        operation (str): Operation type - 'merge' or 'test'.\n        backend (str): Backend type for the model ('fsdp' or 'megatron').\n        target_dir (Optional[str]): Directory to save the merged huggingface model. Defaults to \"tmp\".\n        hf_upload_path (Optional[str]): Hugging Face repository ID to upload the model. Defaults to None.\n        private (bool): Whether to upload the model to a private Hugging Face repository. Defaults to False.\n        test_hf_dir (Optional[str]): Path to the reference Hugging Face model directory for testing. Defaults to None.\n        tie_word_embedding (bool): Whether to tie word embedding weights (currently only Megatron\n            supported). Defaults to False.\n        trust_remote_code (bool): Whether to trust remote code. Defaults to False.\n        is_value_model (bool): Whether the model is a value model (currently only Megatron\n            supported). Defaults to False.\n        local_dir (Optional[str]): Path to the saved model checkpoints. Defaults to None.\n        hf_model_config_path (Optional[str]): Path to HuggingFace model configuration files. Defaults to None.\n        hf_upload (bool): Whether to upload to HuggingFace (computed automatically). Not for initialization.\n        use_cpu_initialization (bool): Whether to use CPU initialization for large models. Defaults to False.\n    \"\"\"\n\n    operation: str  # 'merge' or 'test'\n    backend: str\n    target_dir: Optional[str] = \"tmp\"\n    hf_upload_path: Optional[str] = None\n    private: bool = False\n    test_hf_dir: Optional[str] = None\n    tie_word_embedding: bool = False\n    trust_remote_code: bool = False\n    is_value_model: bool = False\n    local_dir: Optional[str] = None\n    hf_model_config_path: Optional[str] = None\n    hf_upload: bool = field(init=False)\n    use_cpu_initialization: bool = False\n\n    def __post_init__(self):\n        self.hf_upload = self.operation == \"merge\" and bool(self.hf_upload_path)\n        if self.operation == \"test\":\n            self.target_dir = None\n            self.hf_upload_path = None\n            self.private = False\n\n\ndef generate_config_from_args(args: argparse.Namespace) -> ModelMergerConfig:\n    common_config_args = {\n        \"operation\": args.operation,\n        \"backend\": args.backend,\n        \"tie_word_embedding\": args.tie_word_embedding,\n        \"trust_remote_code\": args.trust_remote_code,\n        \"is_value_model\": args.is_value_model,\n        \"local_dir\": args.local_dir,\n        \"hf_model_config_path\": os.path.join(args.local_dir, \"huggingface\"),\n        \"use_cpu_initialization\": args.use_cpu_initialization,\n    }\n\n    if args.operation == \"merge\":\n        config = ModelMergerConfig(\n            **common_config_args,\n            target_dir=args.target_dir,\n            hf_upload_path=args.hf_upload_path,\n            private=args.private,\n            test_hf_dir=None,\n        )\n        os.makedirs(config.target_dir, exist_ok=True)\n    elif args.operation == \"test\":\n        config = ModelMergerConfig(\n            **common_config_args,\n            test_hf_dir=args.test_hf_dir,\n            # the following args are not used by test operation\n            target_dir=None,\n            hf_upload_path=None,\n            private=False,\n        )\n    else:\n        raise NotImplementedError(f\"Unknown operation: {args.operation}\")\n    return config\n\n\nclass BaseModelMerger(ABC):\n    \"\"\"\n    Abstract base class for merging distributed model checkpoints into HuggingFace format.\n\n    This class provides common functionality for converting model checkpoints from different\n    distributed training backends (FSDP, Megatron) into standard HuggingFace format that\n    can be easily loaded and used for inference or further training.\n\n    The merger supports two main operations:\n    - merge: Convert and save checkpoints to HuggingFace format\n    - test: Validate merged checkpoints against a reference model\n\n    Args:\n        config (ModelMergerConfig): Configuration object containing paths, backend type,\n            and operation parameters.\n\n    Attributes:\n        config (ModelMergerConfig): The configuration object passed during initialization.\n        hf_model_config_path (str): Path to the HuggingFace model configuration files.\n        model_config (PretrainedConfig): Loaded HuggingFace model configuration.\n    \"\"\"\n\n    def __init__(self, config: ModelMergerConfig):\n        self.config = config\n        self.hf_model_config_path = config.hf_model_config_path\n        self.model_config = AutoConfig.from_pretrained(\n            self.hf_model_config_path, trust_remote_code=self.config.trust_remote_code\n        )\n\n    def get_transformers_auto_model_class(self):\n        has_remote_code = hasattr(self.model_config, \"auto_map\") and any(\n            self.model_config.architectures[0] in val for val in self.model_config.auto_map.values()\n        )\n        if has_remote_code:\n            auto_class = next(\n                k for k, v in self.model_config.auto_map.items() if self.model_config.architectures[0] in v\n            )\n            match auto_class:\n                case \"AutoModelForCausalLM\":\n                    return AutoModelForCausalLM\n                case \"AutoModelForTokenClassification\":\n                    return AutoModelForTokenClassification\n                case \"AutoModelForVision2Seq\":\n                    return AutoModelForVision2Seq\n                case \"AutoModelForImageTextToText\":\n                    return AutoModelForVision2Seq\n                case _:\n                    raise NotImplementedError(f\"Unknown auto class {auto_class}\")\n        else:\n            if \"ForTokenClassification\" in self.model_config.architectures[0]:\n                return AutoModelForTokenClassification\n            elif \"ForCausalLM\" in self.model_config.architectures[0]:\n                return AutoModelForCausalLM\n            elif \"ForConditionalGeneration\" in self.model_config.architectures[0]:\n                return AutoModelForVision2Seq\n\n            raise NotImplementedError(f\"Unknown architecture {self.model_config.architectures}\")\n\n    def patch_model_generation_config(self, model):\n        \"\"\"\n        The generation_config created from model config may be different to the pretrained model,\n        this may lead to error when generating: https://github.com/volcengine/verl/issues/1246\n\n        This function patch the generation_config created from model config to the pretrained model.\n        \"\"\"\n        if model.can_generate():\n            try:\n                model.generation_config = GenerationConfig.from_pretrained(self.hf_model_config_path)\n            except OSError:\n                print(\n                    f\"Warning: Generation config file not found in {self.hf_model_config_path}, using a \"\n                    f\"generation config created from the model config.\"\n                )\n        return model\n\n    def _load_lora_train_meta(self) -> Optional[dict[str, object]]:\n        if not self.config.local_dir:\n            return None\n\n        meta_path = os.path.join(self.config.local_dir, \"lora_train_meta.json\")\n        if not os.path.exists(meta_path):\n            return None\n\n        import json\n\n        try:\n            with open(meta_path, encoding=\"utf-8\") as f:\n                lora_meta = json.load(f)\n        except Exception as e:\n            warnings.warn(f\"Failed to read LoRA metadata from {meta_path}: {e}\", stacklevel=2)\n            return None\n\n        result = {}\n        if \"r\" in lora_meta:\n            try:\n                result[\"r\"] = int(lora_meta[\"r\"])\n            except (TypeError, ValueError):\n                warnings.warn(f\"Invalid LoRA rank in {meta_path}: {lora_meta['r']}\", stacklevel=2)\n\n        if \"lora_alpha\" in lora_meta:\n            try:\n                result[\"lora_alpha\"] = int(lora_meta[\"lora_alpha\"])\n            except (TypeError, ValueError):\n                warnings.warn(f\"Invalid lora_alpha in {meta_path}: {lora_meta['lora_alpha']}\", stacklevel=2)\n\n        if \"task_type\" in lora_meta:\n            task_type = lora_meta[\"task_type\"]\n            if task_type is None:\n                pass\n            elif isinstance(task_type, str):\n                result[\"task_type\"] = task_type\n            else:\n                warnings.warn(f\"Invalid task_type in {meta_path}: {task_type}\", stacklevel=2)\n\n        return result if len(result) > 0 else None\n\n    def save_lora_adapter(self, state_dict: dict[str, torch.Tensor]):\n        \"\"\"\n        Save lora adapter to safetensors.\n\n        Returns:\n            lora_path: str, the path to the lora adapter. None if no lora adapter found.\n\n        Note:\n            This function change the 'state_dict' in place.\n        \"\"\"\n        lora_params_names = [name for name in state_dict.keys() if \"lora_\" in name]\n\n        if len(lora_params_names) == 0:\n            return None\n\n        import json\n        from typing import OrderedDict\n\n        import peft\n        from safetensors.torch import save_file\n\n        lora_params = OrderedDict()\n        target_modules = set()\n        lora_key = None\n\n        for name in lora_params_names:\n            lora_key = name.replace(\".default.weight\", \".weight\")\n            target_modules.add(lora_key.split(\".\")[-3])\n            lora_params[lora_key] = state_dict.pop(name)\n\n        inferred_lora_rank = min(lora_params[lora_key].shape[0], lora_params[lora_key].shape[1])\n        lora_meta = self._load_lora_train_meta()\n\n        lora_rank = inferred_lora_rank\n        lora_alpha = 0\n        task_type = None\n\n        if lora_meta is not None:\n            meta_rank = lora_meta.get(\"r\")\n            if meta_rank is not None and meta_rank > 0:\n                if meta_rank != inferred_lora_rank:\n                    warnings.warn(\n                        f\"LoRA rank mismatch between metadata ({meta_rank}) and adapter weights \"\n                        f\"({inferred_lora_rank}); using metadata rank.\",\n                        stacklevel=2,\n                    )\n                lora_rank = meta_rank\n\n            meta_alpha = lora_meta.get(\"lora_alpha\")\n            if meta_alpha is not None:\n                lora_alpha = meta_alpha\n\n            meta_task_type = lora_meta.get(\"task_type\")\n            if meta_task_type is not None:\n                task_type = meta_task_type\n\n        if lora_alpha == 0:\n            warnings.warn(\n                \"LoRA alpha metadata is missing or equals 0; falling back to lora_alpha=0. \"\n                \"Please verify checkpoint LoRA metadata (lora_train_meta.json).\",\n                stacklevel=2,\n            )\n\n        peft_dict = {\n            \"r\": lora_rank,\n            \"lora_alpha\": lora_alpha,\n            \"target_modules\": list(target_modules),\n        }\n        if task_type is not None:\n            peft_dict[\"task_type\"] = task_type\n        peft_config = peft.LoraConfig(**peft_dict).to_dict()\n        peft_config[\"task_type\"] = peft_config[\"task_type\"].value if peft_config[\"task_type\"] else None\n        peft_config[\"peft_type\"] = peft_config[\"peft_type\"].value if peft_config[\"peft_type\"] else None\n        peft_config[\"target_modules\"] = list(peft_config[\"target_modules\"])\n\n        lora_path = os.path.join(self.config.target_dir, \"lora_adapter\")\n        os.makedirs(lora_path, exist_ok=True)\n        with open(os.path.join(lora_path, \"adapter_config.json\"), \"w\", encoding=\"utf-8\") as f:\n            json.dump(peft_config, f, ensure_ascii=False, indent=4)\n        save_file(lora_params, os.path.join(lora_path, \"adapter_model.safetensors\"))\n\n        for name in list(state_dict.keys()):\n            key = (\n                name.replace(\"base_model.model.\", \"\")\n                .replace(\".base_layer.weight\", \".weight\")\n                .replace(\".base_layer.bias\", \".bias\")\n            )\n            state_dict[key] = state_dict.pop(name)\n\n        return lora_path\n\n    def save_hf_model_and_tokenizer(self, state_dict: dict[str, torch.Tensor]):\n        auto_model_class = self.get_transformers_auto_model_class()\n        with init_empty_weights():\n            model = auto_model_class.from_config(\n                self.model_config, torch_dtype=torch.bfloat16, trust_remote_code=self.config.trust_remote_code\n            )\n        model.to_empty(device=\"cpu\")\n        model = self.patch_model_generation_config(model)\n\n        lora_path = self.save_lora_adapter(state_dict)\n        if lora_path:\n            print(f\"Saving lora adapter to {lora_path}\")\n\n        print(f\"Saving model to {self.config.target_dir}\")\n        model.save_pretrained(self.config.target_dir, state_dict=state_dict)\n        del state_dict\n        del model\n\n        processor = hf_processor(self.hf_model_config_path, trust_remote_code=self.config.trust_remote_code)\n        tokenizer = hf_tokenizer(self.hf_model_config_path, trust_remote_code=self.config.trust_remote_code)\n        if processor is not None:\n            print(f\"Saving processor to {self.config.target_dir}\")\n            processor.save_pretrained(self.config.target_dir)\n        if tokenizer is not None:\n            print(f\"Saving tokenizer to {self.config.target_dir}\")\n            tokenizer.save_pretrained(self.config.target_dir)\n\n    def upload_to_huggingface(self):\n        import requests\n        from huggingface_hub import HfApi\n        from huggingface_hub.utils import HfHubHTTPError, RepositoryNotFoundError\n\n        api = HfApi()\n        try:\n            # Attempt to create repository\n            api.create_repo(repo_id=self.config.hf_upload_path, private=self.config.private, exist_ok=True)\n        except HfHubHTTPError as e:\n            # Handle authentication/API errors\n            if e.response.status_code == 401:\n                raise PermissionError(\n                    \"Hugging Face authentication failed. Verify your token is valid and has write permissions.\"\n                ) from e\n            elif e.response.status_code == 404:\n                raise RepositoryNotFoundError(f\"Repository path not found: {self.config.hf_upload_path}\") from e\n            else:\n                raise ConnectionError(f\"Failed to create repository ({e.response.status_code}): {e}\") from e\n        except requests.exceptions.ConnectionError as e:\n            raise ConnectionError(\"Network connection failed. Check your internet connection.\") from e\n\n        try:\n            # Attempt folder upload\n            api.upload_folder(folder_path=self.config.target_dir, repo_id=self.config.hf_upload_path, repo_type=\"model\")\n        except HfHubHTTPError as e:\n            if e.response.status_code == 401:\n                raise PermissionError(\"Authentication failed during upload. Token may have expired.\") from e\n            else:\n                raise RuntimeError(f\"Upload failed ({e.response.status_code}): {e}\") from e\n        except requests.exceptions.ConnectionError as e:\n            raise ConnectionError(\"Network interruption during upload. Try again with stable connection.\") from e\n        except OSError as e:\n            raise FileNotFoundError(f\"Local folder error: {self.config.target_dir} - {str(e)}\") from e\n        except Exception as e:\n            raise RuntimeError(f\"Unexpected error during upload: {str(e)}\") from e\n\n    @abstractmethod\n    def merge_and_save(self):\n        raise NotImplementedError(\"Subclasses should implement this method\")\n\n    @abstractmethod\n    def cleanup(self):\n        raise NotImplementedError(\"Subclasses should implement this method to clean up resources if needed\")\n"
  },
  {
    "path": "verl/model_merger/fsdp_model_merger.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nfrom concurrent.futures import ThreadPoolExecutor\nfrom pathlib import Path\n\nimport numpy as np\nimport torch\nfrom torch.distributed._tensor import Placement, Shard\n\ntry:\n    # for torch 2.5+\n    from torch.distributed.tensor import DTensor\nexcept ImportError:\n    from torch.distributed._tensor import DTensor\n\nfrom tqdm import tqdm\n\nfrom .base_model_merger import BaseModelMerger\n\n\nclass FSDPModelMerger(BaseModelMerger):\n    \"\"\"\n    Model merger for FSDP (Fully Sharded Data Parallel) checkpoints.\n\n    This class handles the conversion of FSDP distributed checkpoints into HuggingFace format.\n    FSDP shards model parameters across multiple processes, and this merger reconstructs\n    the full model by loading and concatenating the sharded parameters from all ranks.\n\n    The merger supports various FSDP configurations including:\n    - Pure FSDP (single dimension sharding)\n    - FSDP + DDP (data parallel + fully sharded data parallel)\n    - DTensor-based sharding with custom device meshes\n\n    Key features:\n    - Automatic detection of world size from checkpoint filenames\n    - Support for DTensor and non-DTensor checkpoints\n    - Parallel loading of checkpoint shards for efficiency\n    - Validation against reference HuggingFace models\n\n    Example:\n        To merge FSDP checkpoints:\n        ```python\n        config = ModelMergerConfig(\n            operation=\"merge\",\n            backend=\"fsdp\",\n            local_dir=\"path/to/fsdp/checkpoints\",\n            target_dir=\"path/to/output\"\n        )\n        merger = FSDPModelMerger(config)\n        merger.merge_and_save()\n        ```\n    \"\"\"\n\n    def _get_world_size(self) -> int:\n        \"\"\"_summary_\n        From FSDP json config file, extract the world size.\n\n        Returns:\n            int: world size\n        \"\"\"\n        config_path = Path(self.config.local_dir) / \"fsdp_config.json\"\n        if not config_path.exists():\n            raise FileNotFoundError(f\"Config file {config_path} does not exist.\")\n\n        with open(config_path) as f:\n            config = json.load(f)\n\n        # Extract world size from the config\n        world_size = config.get(\"world_size\", None)\n        if world_size is None:\n            raise ValueError(\"World size not found in the config file.\")\n\n        return world_size\n\n    def _load_rank_zero_state_dict(self, world_size: int) -> dict:\n        return torch.load(\n            Path(self.config.local_dir) / f\"model_world_size_{world_size}_rank_0.pt\",\n            map_location=\"cpu\",\n            weights_only=False,\n        )\n\n    def _extract_device_mesh_info(self, state_dict: dict, world_size: int) -> tuple[np.ndarray, tuple[str, ...]]:\n        \"\"\"\n        Retrieves sharding information (device_mesh, mesh_dim_names) from a DTensor in the state_dict.\n        If no DTensor is found, infers a simple FSDP mesh based on world_size.\n        \"\"\"\n        pivot_key = sorted(list(state_dict.keys()))[0]\n        weight = state_dict[pivot_key]\n\n        if isinstance(weight, DTensor):\n            # get sharding info\n            device_mesh = weight.device_mesh\n            mesh = device_mesh.mesh\n            mesh_dim_names = device_mesh.mesh_dim_names\n        else:\n            # for non-DTensor\n            mesh = np.array([world_size], dtype=np.int64)\n            mesh_dim_names = (\"fsdp\",)\n\n        return mesh, mesh_dim_names\n\n    def _calculate_shard_configuration(\n        self, mesh: np.ndarray, mesh_dim_names: tuple[str, ...]\n    ) -> tuple[int, tuple[int, ...]]:\n        \"\"\"Calculates the total number of shards and the shape of the device mesh.\"\"\"\n        assert mesh_dim_names in ((\"fsdp\",), (\"ddp\", \"fsdp\")), f\"Unsupported mesh_dim_names {mesh_dim_names}\"\n\n        if \"tp\" in mesh_dim_names:\n            # TODO: \"tp\" is not supported yet due to the above assert\n            total_shards = mesh.shape[-1] * mesh.shape[-2]\n            mesh_shape = (mesh.shape[-2], mesh.shape[-1])\n        else:\n            total_shards = mesh.shape[-1]\n            mesh_shape = (mesh.shape[-1],)\n\n        return total_shards, mesh_shape\n\n    def _merge_by_placement(self, tensors: list[torch.Tensor], placement: Placement) -> torch.Tensor:\n        \"\"\"Merges a list of tensors based on their DTensor placement\"\"\"\n        if placement.is_replicate():\n            return tensors[0]\n        elif placement.is_partial():\n            raise NotImplementedError(\"Partial placement is not supported yet\")\n        elif placement.is_shard():\n            return torch.cat(tensors, dim=placement.dim).contiguous()\n\n        raise NotImplementedError(f\"Unsupported placement: {placement}\")\n\n    def _load_and_merge_state_dicts(\n        self, world_size: int, total_shards: int, mesh_shape: tuple[int, ...], mesh_dim_names: tuple[str, ...]\n    ) -> dict[str, torch.Tensor]:\n        model_state_dict_lst = [None] * total_shards\n\n        def process_one_shard(rank: int, model_state_dict_lst: list):\n            model_path = Path(self.config.local_dir) / f\"model_world_size_{world_size}_rank_{rank}.pt\"\n            state_dict = torch.load(model_path, map_location=\"cpu\", weights_only=False)\n            model_state_dict_lst[rank] = state_dict\n            return state_dict\n\n        with ThreadPoolExecutor(max_workers=min(32, os.cpu_count())) as executor:\n            futures = [executor.submit(process_one_shard, rank, model_state_dict_lst) for rank in range(total_shards)]\n            for future in tqdm(futures, desc=f\"Loading {total_shards} FSDP shards\", total=total_shards):\n                future.result()\n\n        # Merge state dicts from all shards\n        state_dict = {}\n        param_placements: dict[str, list] = {}\n\n        for key in set(model_state_dict_lst[0].keys()):\n            state_dict[key] = []\n            for model_state_shard in model_state_dict_lst:\n                # add tensor shard in order of rank to state_dict[key]\n                tensor = model_state_shard.pop(key)\n                if isinstance(tensor, DTensor):\n                    state_dict[key].append(tensor._local_tensor.bfloat16())\n\n                    placements = tuple(tensor.placements)\n                    # replicated placement at dp dimension can be discarded\n                    if mesh_dim_names[0] in (\"dp\", \"ddp\"):\n                        placements = placements[1:]\n\n                    if key not in param_placements:\n                        param_placements[key] = placements\n                    else:\n                        assert param_placements[key] == placements\n                else:\n                    state_dict[key].append(tensor.bfloat16())\n\n        del model_state_dict_lst\n\n        # Merge tensors\n        for key in sorted(state_dict):\n            if not isinstance(state_dict[key], list):\n                print(f\"No need to merge key {key}\")\n                continue\n            if key in param_placements:\n                # merge shards\n                placements: tuple[Shard] = param_placements[key]\n                if len(mesh_shape) == 1:\n                    # 1-D list, FSDP without TP\n                    assert len(placements) == 1\n                    shards = state_dict[key]\n                    state_dict[key] = self._merge_by_placement(shards, placements[0])\n                else:\n                    # 2-D list, FSDP + TP\n                    raise NotImplementedError(\"FSDP + TP is not supported yet\")\n            else:\n                state_dict[key] = torch.cat(state_dict[key], dim=0)\n\n        return state_dict\n\n    def merge_and_save(self):\n        world_size = self._get_world_size()\n        rank_zero_state_dict = self._load_rank_zero_state_dict(world_size)\n\n        mesh, mesh_dim_names = self._extract_device_mesh_info(rank_zero_state_dict, world_size)\n        print(f\"Got device mesh {mesh}, mesh_dim_names {mesh_dim_names}\")\n\n        total_shards, mesh_shape = self._calculate_shard_configuration(mesh, mesh_dim_names)\n        print(f\"Processing model shards with {total_shards} {mesh_shape} in total\")\n\n        merged_state_dict = self._load_and_merge_state_dicts(world_size, total_shards, mesh_shape, mesh_dim_names)\n\n        if self.config.operation == \"test\":\n            if not self.config.test_hf_dir:\n                raise ValueError(\"test_hf_dir must be provided for test operation\")\n            self._validate_state_dict(merged_state_dict)\n        elif self.config.operation == \"merge\":\n            self.save_hf_model_and_tokenizer(merged_state_dict)\n            if self.config.hf_upload:\n                self.upload_to_huggingface()\n        else:\n            raise ValueError(f\"Unknown operation: {self.config.operation}\")\n\n    def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]):\n        auto_model_class = self.get_transformers_auto_model_class()\n\n        hf_model = auto_model_class.from_pretrained(self.config.test_hf_dir, torch_dtype=torch.bfloat16)\n        hf_state_dict = hf_model.state_dict()\n        del hf_model\n\n        hf_model_keys = set(hf_state_dict.keys())\n        collected_keys = set(state_dict.keys())\n\n        missing_keys = hf_model_keys - collected_keys\n        assert len(missing_keys) == 0, f\"Missing keys in collected state dict: {list(sorted(missing_keys))}\"\n\n        extra_keys = collected_keys - hf_model_keys\n        assert len(extra_keys) == 0, f\"Extra keys in collected state dict: {list(sorted(extra_keys))}\"\n\n        for key in hf_model_keys:\n            hf_shape = hf_state_dict[key].shape\n            collected_shape = state_dict[key].shape\n            assert hf_shape == collected_shape, (\n                f\"Shape mismatch for key '{key}': original {hf_shape} vs collected {collected_shape}\"\n            )\n\n            hf_dtype = hf_state_dict[key].dtype\n            collected_dtype = state_dict[key].dtype\n            assert hf_dtype == collected_dtype, (\n                f\"Dtype mismatch for key '{key}': original {hf_dtype} vs collected {collected_dtype}\"\n            )\n\n            torch.testing.assert_close(hf_state_dict[key], state_dict[key], atol=1e-6, rtol=1e-6)\n\n        print(\"FSDP checks passed: The merged state_dict matches the hf model saved by FSDPCheckpointManager.\")\n\n    def cleanup(self):\n        \"\"\"Cleanup temporary files if needed.\"\"\"\n        # FSDP merger does not create temporary files, so no cleanup is needed.\n        pass\n"
  },
  {
    "path": "verl/model_merger/megatron_model_merger.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nimport warnings\nfrom contextlib import contextmanager\nfrom pathlib import Path\nfrom typing import Any, Callable, ContextManager\n\nimport numpy as np\nimport torch\nimport torch.distributed as dist\n\ntry:\n    # NPU patch\n    import mindspeed.megatron_adaptor  # noqa: F401\nexcept ImportError:\n    pass\n\nfrom accelerate import init_empty_weights\nfrom megatron.core import mpu\nfrom megatron.core.models.gpt.gpt_model import ModelType\nfrom megatron.core.tensor_parallel.random import model_parallel_cuda_manual_seed\nfrom safetensors.torch import load_file\nfrom transformers import (\n    AutoConfig,\n    PretrainedConfig,\n)\n\nfrom verl.models.mcore import hf_to_mcore_config\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device\nfrom verl.utils.distributed import set_numa_affinity\nfrom verl.utils.megatron.dist_checkpointing import load_dist_checkpointing\nfrom verl.utils.megatron_utils import get_model\nfrom verl.utils.tokenizer import hf_processor, hf_tokenizer\n\nfrom .base_model_merger import BaseModelMerger, ModelMergerConfig\n\n\n@contextmanager\ndef noop_context() -> Any:\n    yield\n\n\ndef get_dynamic_pipeline_shards(layer_num: int, pp_size: int) -> list[int]:\n    \"\"\"Calculate the pipeline sharding configuration for Megatron-LM.\n\n    Args:\n        layer_num: Total number of layers in the model.\n        pp_size: Number of pipeline parallel ranks.\n\n    Returns:\n        layer number of each pp rank. Make the sharding of the pipeline as uniform as possible.\n    \"\"\"\n    if layer_num < pp_size:\n        raise ValueError(f\"layer_num {layer_num} must be greater than pp_size {pp_size}.\")\n\n    if pp_size < 1:\n        raise ValueError(f\"pp_size must be at least 1, got {pp_size}.\")\n    if pp_size == 1:\n        return [layer_num]\n\n    if pp_size == 2:\n        return [\n            layer_num // 2,\n            layer_num - layer_num // 2,\n        ]\n\n    middle_size = pp_size - 2\n    shards_strategy = []\n    for middle_layer_num in range(layer_num):\n        first_last_layer_num = layer_num - middle_layer_num * middle_size\n        first_layer_num = first_last_layer_num // 2\n        last_layer_num = first_last_layer_num - first_last_layer_num // 2\n        if 0 < first_layer_num <= middle_layer_num and 0 < last_layer_num <= middle_layer_num:\n            shards_strategy.append(\n                (\n                    [first_layer_num] + [middle_layer_num] * middle_size + [last_layer_num],\n                    abs(first_layer_num - middle_layer_num),\n                )\n            )\n\n    # sort by diff of layer_num, to make it as uniform as possible\n    res = sorted(shards_strategy, key=lambda x: x[1])[0][0]\n    assert sum(res) == layer_num, f\"sum(res)={sum(res)} != layer_num={layer_num}, pp_size={pp_size}\"\n    return res\n\n\nclass MegatronModelMerger(BaseModelMerger):\n    \"\"\"\n    Model merger for Megatron-LM distributed checkpoints.\n\n    This class handles the conversion of Megatron-LM distributed checkpoints into HuggingFace format.\n    Megatron-LM uses tensor parallelism, pipeline parallelism, and data parallelism to distribute\n    large language models across multiple GPUs. This merger reconstructs the full model by\n    loading distributed checkpoints and applying the necessary transformations.\n\n    Key features:\n    - Support for tensor parallel, pipeline parallel, and data parallel configurations\n    - Automatic parameter name mapping from Megatron to HuggingFace conventions\n    - Handling of QKV and gate-up tensor splitting/merging\n    - Support for tied word embeddings and value models\n    - Integration with Megatron's distributed checkpointing system\n\n    The merger handles various model architectures and configurations:\n    - Standard transformer models (GPT-style)\n    - Models with tied word embeddings\n    - Value models for reinforcement learning\n    - Multi-layer attention (MLA) architectures\n    - Mixture of Experts (MoE) models\n\n    Args:\n        config (ModelMergerConfig): Configuration object with Megatron-specific settings\n            including tie_word_embedding and is_value_model flags.\n\n    Example:\n        To merge Megatron checkpoints:\n        ```python\n        config = ModelMergerConfig(\n            operation=\"merge\",\n            backend=\"megatron\",\n            local_dir=\"path/to/megatron/checkpoints\",\n            target_dir=\"path/to/output\",\n            tie_word_embedding=True\n        )\n        merger = MegatronModelMerger(config)\n        merger.merge_and_save()\n        ```\n    \"\"\"\n\n    def __init__(self, config: ModelMergerConfig):\n        super().__init__(config)\n        # Currently we use only 1 rank to merge the dist_ckpt, we will move to multi-process save shortly afterwards\n        if \"WORLD_SIZE\" not in os.environ:\n            os.environ[\"RANK\"] = \"0\"\n            os.environ[\"LOCAL_RANK\"] = \"0\"\n            os.environ[\"WORLD_SIZE\"] = \"1\"\n            os.environ[\"MASTER_ADDR\"] = \"localhost\"\n            os.environ[\"MASTER_PORT\"] = \"12355\"\n\n        set_numa_affinity()\n        torch.distributed.init_process_group(get_nccl_backend())\n\n        self.rank = torch.distributed.get_rank()\n        self.world_size = torch.distributed.get_world_size()\n        local_rank = os.environ.get(\"LOCAL_RANK\", 0)\n        get_torch_device().set_device(f\"{get_device_name()}:{local_rank}\")\n\n        mpu.initialize_model_parallel(\n            tensor_model_parallel_size=1,\n            pipeline_model_parallel_size=self.world_size,\n            virtual_pipeline_model_parallel_size=None,\n            context_parallel_size=1,\n            expert_model_parallel_size=1,\n        )\n        model_parallel_cuda_manual_seed(0)\n        self.hf_config = AutoConfig.from_pretrained(\n            self.config.hf_model_config_path, trust_remote_code=self.config.trust_remote_code\n        )\n        print(self.hf_config, flush=True)\n\n        self.params_mapping = {\n            # megatron core gpt model name, huggingface model name\n            # NOTICE: It's a little bit tricky, when 2 keys have the same prefix, we need to make sure the\n            # longer key within the containing relationship is processed first.\n            \"embedding.word_embeddings\": \"model.embed_tokens\",\n            # input layer norm for dpskv3\n            \"input_layernorm.weight\": \"input_layernorm.weight\",\n            \"input_layernorm.bias\": \"input_layernorm.bias\",\n            # attn\n            \"self_attention.linear_qkv.layer_norm_weight\": \"input_layernorm.weight\",\n            \"self_attention.linear_qkv.layer_norm_bias\": \"input_layernorm.bias\",\n            \"self_attention.linear_qkv\": \"self_attn.qkv_proj\",\n            \"self_attention.q_layernorm\": \"self_attn.q_norm\",\n            \"self_attention.k_layernorm\": \"self_attn.k_norm\",\n            \"self_attention.linear_proj\": \"self_attn.o_proj\",\n            # mla\n            \"self_attention.linear_q_proj\": \"self_attn.q_proj\",\n            \"self_attention.linear_q_down_proj\": \"self_attn.q_a_proj\",\n            \"self_attention.linear_q_up_proj.layer_norm_weight\": \"self_attn.q_a_layernorm.weight\",\n            \"self_attention.linear_q_up_proj\": \"self_attn.q_b_proj\",\n            \"self_attention.linear_kv_down_proj\": \"self_attn.kv_a_proj_with_mqa\",\n            \"self_attention.linear_kv_up_proj.layer_norm_weight\": \"self_attn.kv_a_layernorm.weight\",\n            \"self_attention.linear_kv_up_proj\": \"self_attn.kv_b_proj\",\n            # mlp\n            \"pre_mlp_layernorm\": \"post_attention_layernorm\",\n            \"mlp.linear_fc1.layer_norm_weight\": \"post_attention_layernorm.weight\",\n            \"mlp.linear_fc1.layer_norm_bias\": \"post_attention_layernorm.bias\",\n            \"mlp.linear_fc1\": \"mlp.gate_up_proj\",\n            \"mlp.linear_fc2\": \"mlp.down_proj\",\n            # moe\n            \"mlp.router.expert_bias\": \"mlp.gate.e_score_correction_bias\",\n            \"mlp.router\": \"mlp.gate\",\n            \"mlp.shared_experts.linear_fc1\": \"mlp.shared_experts.gate_up_proj\",\n            \"mlp.shared_experts.linear_fc2\": \"mlp.shared_experts.down_proj\",\n            \"linear_fc1\": \"gate_up_proj\",\n            \"linear_fc2\": \"down_proj\",\n            # output\n            \"final_layernorm\": \"norm\",\n            \"output_layer\": \"lm_head\",\n        }\n\n        if \"Qwen2MoeForCausalLM\" in self.hf_config.architectures:\n            self.params_mapping[\"mlp.shared_experts.linear_fc1\"] = \"mlp.shared_expert.gate_up_proj\"\n            self.params_mapping[\"mlp.shared_experts.linear_fc2\"] = \"mlp.shared_expert.down_proj\"\n            self.params_mapping[\"mlp.shared_experts.gate_weight\"] = \"mlp.shared_expert_gate.weight\"\n\n    def _load_state_dicts(self, model_ckpt_path: str) -> dict[str, Any]:\n        \"\"\"_summary_\n        Use Megatron dist_checkpointing to load the model state dicts from the checkpoint directory.\n\n        Args:\n            model_ckpt_path (str): Path to the model checkpoint directory.\n\n        Returns:\n            State dict containing the model parameters.\n        \"\"\"\n\n        # init hf config\n        self.pipeline_shards = get_dynamic_pipeline_shards(self.hf_config.num_hidden_layers, self.world_size)\n        print(f\"Pipeline shards: {self.pipeline_shards}, total layers: {sum(self.pipeline_shards)}\")\n\n        tf_config = hf_to_mcore_config(\n            self.hf_config,\n            torch.bfloat16,\n            num_layers_in_first_pipeline_stage=self.pipeline_shards[0] if len(self.pipeline_shards) > 1 else None,\n            num_layers_in_last_pipeline_stage=self.pipeline_shards[-1] if len(self.pipeline_shards) > 2 else None,\n        )\n        tf_config.use_cpu_initialization = self.config.use_cpu_initialization\n        tie_word_embeddings = getattr(self.hf_config, \"tie_word_embeddings\", False)\n\n        # init megatron model\n        def megatron_model_provider(pre_process, post_process):\n            from verl.models.mcore import init_mcore_model\n\n            parallel_model = init_mcore_model(\n                tf_config,\n                self.hf_config,\n                pre_process,\n                post_process,\n                share_embeddings_and_output_weights=tie_word_embeddings,\n                value=False,\n            )\n            return parallel_model\n\n        context: Callable[..., ContextManager] = (\n            init_empty_weights if self.config.use_cpu_initialization else noop_context\n        )\n        with context():\n            whole_model = get_model(\n                model_provider_func=megatron_model_provider,\n                model_type=ModelType.encoder_or_decoder,\n                wrap_with_ddp=False,\n                transformer_config=tf_config,\n            )\n\n        if self.config.use_cpu_initialization:\n            # convert meta device to empty tensor so it can use `copy_` function\n            whole_model[0].module = whole_model[0].module.to_empty(device=\"cpu\")\n\n        # load state dicts\n        sharded_state_dict = {}\n        for vpp_rank, model in enumerate(whole_model):\n            key = f\"model{vpp_rank}\" if len(whole_model) > 1 else \"model\"\n            mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)\n            sharded_state_dict[key] = model.sharded_state_dict()\n        model_state_dict = load_dist_checkpointing(sharded_state_dict, model_ckpt_path)\n        model_state_dict_list = []\n        for vpp_rank, model in enumerate(whole_model):\n            key = f\"model{vpp_rank}\" if len(whole_model) > 1 else \"model\"\n            mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)\n            model_state_dict_list.append(model_state_dict[key])\n\n        return model_state_dict_list\n\n    def _check_megatron_state_key(self, key: str) -> bool:\n        \"\"\"\n        Checks if the key is a valid Megatron state key.\n\n        Now the model merger only supports keys that start with \"decoder/embedding/output_layer\" in TransformerLayer.\n        Shall not use key starts with \"model.\"\n        \"\"\"\n        if key.startswith(\"model.\"):\n            raise ValueError(\n                f\"Invalid key {key} in Megatron state_dict. Expected keys to start with \"\n                f\"'decoder/embedding/output_layer' in TransformerLayer.\"\n            )\n\n        skip_checking_keys = [\"embedding.word_embeddings\", \"output_layer\"]\n        for skip_key in skip_checking_keys:\n            if skip_key in key:\n                print(f\"skip checking key {key}\")\n                return\n\n        # Exclude extra state keys\n        if not key.startswith(\"decoder\"):\n            raise ValueError(\n                f\"Invalid key {key} in Megatron state_dict. Expected keys to start with 'decoder' in TransformerLayer.\"\n            )\n\n    def _split_tensors(\n        self, key: str, tensor: torch.Tensor, config: PretrainedConfig, is_value_model: bool = False\n    ) -> list[torch.Tensor]:\n        \"\"\"\n        Splits a tensor into multiple tensors based on the name.\n        This is used to handle qkv and gate_up tensors.\n        \"\"\"\n        if \"linear_fc1.weight\" in key:\n            # if the tensor is gate and proj\n            gate_lst = []\n            up_lst = []\n            gate, up = tensor.chunk(2)\n            gate_lst.append(gate)\n            up_lst.append(up)\n            gate = torch.cat(gate_lst, dim=0)\n            up = torch.cat(up_lst, dim=0)\n            return [gate, up]\n        elif \"self_attention.linear_qkv.\" in key and \"layer_norm\" not in key:\n            # if the tensor is qkv, for each param on tp, split into q, k, v\n            # concat q, k, v separately.\n            q_lst, k_lst, v_lst = [], [], []\n            assert config.num_attention_heads % config.num_key_value_heads == 0\n            num_q_per_kv = config.num_attention_heads // config.num_key_value_heads\n            assert tensor.shape[0] % (num_q_per_kv + 2) == 0, (\n                f\"Tensor shape {tensor.shape} is not divisible by {num_q_per_kv + 2}\"\n            )\n            kv_size = tensor.shape[0] // (num_q_per_kv + 2)\n            split_size = [kv_size * num_q_per_kv, kv_size, kv_size]\n\n            num_query_groups_per_partition = config.num_key_value_heads\n            for chunk in tensor.chunk(num_query_groups_per_partition):\n                split_size = [\n                    kv_size * num_q_per_kv // num_query_groups_per_partition,\n                    kv_size // num_query_groups_per_partition,\n                    kv_size // num_query_groups_per_partition,\n                ]\n                q, k, v = chunk.split(split_size)\n                q_lst.append(q)\n                k_lst.append(k)\n                v_lst.append(v)\n\n            return [torch.cat(q_lst, dim=0), torch.cat(k_lst, dim=0), torch.cat(v_lst, dim=0)]\n        else:\n            return [tensor]\n\n    def _merge_state_dicts(self, model_state_dict_list: list[dict[str, Any]]) -> dict[str, torch.Tensor]:\n        state_dict = {}\n        layers_cum = 0\n        if self.world_size > 1:\n            pipeline_cumsum = np.cumsum(self.pipeline_shards)\n            layers_cum = 0 if self.rank == 0 else pipeline_cumsum[self.rank - 1]\n\n        print(f\"{layers_cum=}\")\n        for model_state_dict in model_state_dict_list:\n            layers_handled = 0\n            keys = model_state_dict.keys()\n            for key in keys:\n                if \"extra_state\" in key:\n                    continue\n                if self.config.tie_word_embedding and (\"output_layer\" in key):\n                    print(\"skip lm_head and reward_head loading because of tie_word_embeddings\")\n                    continue\n\n                self._check_megatron_state_key(key)\n                hf_name = self._replace_name(key, self.params_mapping)\n                assert hf_name is not None, f\"Failed to convert layer name [{key}] from megatron to huggingface.\"\n                if \"model.layers.\" in hf_name:\n                    local_layer_no = int(hf_name.split(\".\")[2])\n                    layers_handled = max(local_layer_no, layers_handled)\n                    global_layer_no = local_layer_no + layers_cum\n                    new_key_list = hf_name.split(\".\")\n                    new_key_list[2] = str(global_layer_no)\n                    hf_name = \".\".join(new_key_list)\n                else:\n                    warnings.warn(f\"hf_name {hf_name} will not be fixed with layer number\", stacklevel=2)\n\n                if \"mlp.experts.\" in hf_name and \".weight\" in hf_name:\n                    name_prefix, expert_id = hf_name.split(\".weight\")\n                    for proj in [\"gate_up\", \"down\"]:\n                        if f\"{proj}_proj\" in hf_name:\n                            hf_name = hf_name.replace(\n                                f\"mlp.experts.{proj}_proj.weight{expert_id}\",\n                                f\"mlp.experts.{expert_id}.{proj}_proj.weight\",\n                            )\n\n                tensor = model_state_dict[key]\n                split_tensor = self._split_tensors(\n                    key, tensor, self.hf_config, is_value_model=self.config.is_value_model\n                )\n\n                if len(split_tensor) == 1:\n                    state_dict[hf_name] = split_tensor[0]\n                elif len(split_tensor) == 3:\n                    # split qkv\n                    for n, d in zip([\"q\", \"k\", \"v\"], split_tensor, strict=True):\n                        state_dict[hf_name.replace(\"qkv\", n)] = d\n                elif len(split_tensor) == 2:\n                    # split gate up\n                    state_dict[hf_name.replace(\"gate_up\", \"gate\")] = split_tensor[0]\n                    state_dict[hf_name.replace(\"gate_up\", \"up\")] = split_tensor[1]\n                shape_info = (\n                    split_tensor.shape if isinstance(split_tensor, torch.Tensor) else [t.shape for t in split_tensor]\n                )\n                print(f\"converted {key} to {hf_name} with shape {shape_info}\")\n\n            layers_cum += layers_handled + 1  # zero based\n\n        return state_dict\n\n    def save_hf_model_and_tokenizer(self, merged_state_dict):\n        if self.world_size == 1:\n            return super().save_hf_model_and_tokenizer(merged_state_dict)\n\n        from safetensors.torch import save_file\n\n        layer_num = self.hf_config.num_hidden_layers\n\n        # FIXME: make configurable\n        saves_per_layer = 1 if layer_num < 30 else 2\n        saves_total = saves_per_layer * layer_num\n        saves_indexes = {}\n\n        # calculate the layer start index and key chunks\n        layer_this_rank = self.pipeline_shards[self.rank]\n        pipeline_cumsum = np.cumsum(self.pipeline_shards)\n        layer_start = 0 if self.rank == 0 else pipeline_cumsum[self.rank - 1]\n        keys = list(merged_state_dict.keys())\n        keys_chunk = np.array_split(np.array(keys), layer_this_rank * saves_per_layer)\n        numel = 0\n\n        assert len(keys_chunk) == layer_this_rank * saves_per_layer, (\n            f\"Expected {len(keys_chunk)} chunks, but got {layer_this_rank * saves_per_layer} for rank {self.rank}.\"\n        )\n\n        # save to model shards manually\n        target_dir = Path(self.config.target_dir)\n        for i, keys in enumerate(keys_chunk):\n            sd_to_save = {k: merged_state_dict[k] for k in keys}\n            numel += sum([sd_to_save[i].numel() for i in sd_to_save])\n            save_idx = layer_start * saves_per_layer + i\n            save_path = target_dir / f\"model-{save_idx + 1:05d}-of-{saves_total:05d}.safetensors\"\n\n            save_file(sd_to_save, save_path)\n            for k in keys:\n                saves_indexes[k] = str(save_path.name)\n\n        tensor = torch.tensor([numel]).to(get_device_name())\n        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)\n        numel = tensor.cpu().item()\n\n        all_save_indexes = [{} for _ in range(self.world_size)]\n        dist.all_gather_object(all_save_indexes, saves_indexes)\n        saves_indexes = {k: v for i in all_save_indexes for k, v in i.items()}\n        if self.rank == 0:\n            with open(target_dir / \"model.safetensors.index.json\", \"w\") as f:\n                json.dump(\n                    {\n                        \"metadata\": {\n                            \"total_size\": numel,\n                        },\n                        \"weight_map\": saves_indexes,\n                    },\n                    f,\n                    indent=4,\n                )\n            print(f\"model saved to {target_dir} with {numel=}\")\n\n            self.model_config.save_pretrained(self.config.target_dir)\n\n            processor = hf_processor(self.hf_model_config_path, trust_remote_code=self.config.trust_remote_code)\n            tokenizer = hf_tokenizer(self.hf_model_config_path, trust_remote_code=self.config.trust_remote_code)\n            if processor is not None:\n                print(f\"Saving processor to {self.config.target_dir}\")\n                processor.save_pretrained(self.config.target_dir)\n            if tokenizer is not None:\n                print(f\"Saving tokenizer to {self.config.target_dir}\")\n                tokenizer.save_pretrained(self.config.target_dir)\n\n    def merge_and_save(self):\n        from verl.utils.megatron_utils import get_dist_checkpoint_path\n\n        model_ckpt_path = get_dist_checkpoint_path(self.config.local_dir)\n\n        model_state_dict = self._load_state_dicts(model_ckpt_path)\n        merged_state_dict = self._merge_state_dicts(model_state_dict)\n        del model_state_dict\n\n        if self.config.operation == \"test\":\n            if not self.config.test_hf_dir:\n                raise ValueError(\"test_hf_dir must be provided for test operation\")\n            self._validate_state_dict(merged_state_dict)\n        elif self.config.operation == \"merge\":\n            self.save_hf_model_and_tokenizer(merged_state_dict)\n            if self.config.hf_upload:\n                self.upload_to_huggingface()\n        else:\n            raise ValueError(f\"Unknown operation: {self.config.operation}\")\n\n    def _validate_state_dict(self, state_dict: dict[str, torch.Tensor]):\n        \"\"\"\n        Compares the merged Megatron state_dict against a reference safetensors model.\n        Applies necessary name mappings from Megatron to Hugging Face conventions using _replace_name.\n        \"\"\"\n        ref_state_dict = load_file(Path(self.config.test_hf_dir) / \"model.safetensors\")\n\n        for name, loaded_weight in state_dict.items():\n            # name = self._replace_name(original_name, self.params_mapping)\n            if not name or name.endswith(\".bias\") and name not in ref_state_dict:\n                continue\n            if \"rotary_emb.inv_freq\" in name:\n                continue\n            if \"lm_head.weight\" in name:\n                if self.config.is_value_model or self.config.tie_word_embedding:\n                    continue\n            if name not in ref_state_dict:\n                raise RuntimeError(f\"key: {name} not exist in state_dict\")\n            param = ref_state_dict[name]\n            assert loaded_weight.dtype == param.dtype\n            torch.testing.assert_close(loaded_weight.to(\"cpu\"), param, atol=1e-2, rtol=5e-2)\n\n    def _replace_name(self, megatron_name: str, name_mapping: dict[str, str]) -> str:\n        for m_name, v_name in name_mapping.items():\n            if m_name not in megatron_name:\n                continue\n\n            megatron_name = megatron_name.replace(\"decoder\", \"model\")\n            param_name = megatron_name.replace(m_name, v_name)\n\n            return param_name\n\n        return None  # Return None if no mapping found\n\n    def cleanup(self):\n        torch.distributed.destroy_process_group()\n"
  },
  {
    "path": "verl/models/README.md",
    "content": "# Models\nCommon modelzoo such as huggingface/transformers stuggles when using Pytorch native model parallelism. Following the design principle of vLLM, we keep a simple, parallelizable, highly-optimized with packed inputs in verl. \n## Adding a New Huggingface Model\n### Step 1: Copy the model file from HF to verl\n- Add a new file under verl/models/hf\n- Copy ONLY the model file from huggingface/transformers/models to verl/models/hf\n\n### Step 2: Modify the model file to use packed inputs\n- Remove all the code related to inference (kv cache)\n- Modify the inputs to include only\n    - input_ids (total_nnz,)\n    - cu_seqlens (total_nnz + 1,)\n    - max_seqlen_in_batch: int\n- Note that this requires using flash attention with causal mask.\n\n### Step 2.5: Add tests\n- Add a test to compare this version and the huggingface version\n- Following the infrastructure and add tests to tests/models/hf\n\n### Step 3: Add a function to apply tensor parallelism\n- Please follow\n    - https://pytorch.org/docs/stable/distributed.tensor.parallel.html\n    - https://pytorch.org/tutorials/intermediate/TP_tutorial.html\n- General comments\n    - Tensor Parallelism in native Pytorch is NOT auto-parallelism. The way it works is to specify how model parameters and input/output reshards using configs. These configs are then registered as hooks to perform input/output resharding before/after model forward.\n\n### Step 4: Add a function to apply data parallelism\n- Please use FSDP2 APIs\n- See demo here https://github.com/pytorch/torchtitan/blob/main/torchtitan/parallelisms/parallelize_llama.py#L413\n\n### Step 5: Add a function to apply pipeline parallelism\n- Comes in Pytorch 2.4\n- Currently only in alpha in nightly version\n- Check torchtitan for more details\n\n"
  },
  {
    "path": "verl/models/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/models/llama/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/models/llama/megatron/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .modeling_llama_megatron import (\n    ParallelLlamaForCausalLM,\n    # rmpad with megatron\n    ParallelLlamaForCausalLMRmPad,\n    # rmpad with megatron and pipeline parallelism\n    ParallelLlamaForCausalLMRmPadPP,\n    ParallelLlamaForValueRmPad,\n    ParallelLlamaForValueRmPadPP,\n    # original model with megatron\n    ParallelLlamaModel,\n)\n\n__all__ = [\n    \"ParallelLlamaForCausalLM\",\n    \"ParallelLlamaForCausalLMRmPad\",\n    \"ParallelLlamaForCausalLMRmPadPP\",\n    \"ParallelLlamaForValueRmPad\",\n    \"ParallelLlamaForValueRmPadPP\",\n    \"ParallelLlamaModel\",\n]\n"
  },
  {
    "path": "verl/models/llama/megatron/checkpoint_utils/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/models/llama/megatron/checkpoint_utils/llama_loader.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    print(f\"get megatron data parallel size: {mpu.get_data_parallel_world_size()}\")\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef load_state_dict_to_megatron_llama(\n    state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False\n):\n    \"\"\"Load merged state_dict to sharded Megatron module in training.\"\"\"\n    from megatron.core import DistributedDataParallel as LocalDDP\n    from megatron.core import mpu\n    from megatron.core.transformer.module import Float16Module\n    from torch.nn.parallel import DistributedDataParallel as torchDDP\n\n    from verl.utils.logger import print_rank_0\n    from verl.utils.megatron_utils import unwrap_model\n\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    def fetch_params(module):\n        for param in module.parameters():\n            torch.distributed.fetch(\n                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()\n            )\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if torch.distributed.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers, (\n        f\"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size \"\n        f\"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}\"\n    )\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        gpt_model_module = _get_gpt_model(models[i])\n        assert len(gpt_model_module.model.layers) == num_layers_per_model\n\n    def _fetch_tensor(tensor, name) -> torch.Tensor:\n        \"\"\"fetch tensor\"\"\"\n        nonlocal state_dict\n        if tensor is not None:\n            tensor.data.copy_(state_dict[name])\n\n    def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards\"\"\"\n        nonlocal state_dict\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if name in state_dict:\n            full_weight = state_dict[name]\n\n            if mutate_func is not None:\n                full_weight = mutate_func(full_weight)\n            tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n            if tensor is not None:\n                tensor.data.copy_(tensor_chunk[tp_rank])\n        else:\n            print(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards\"\"\"\n        nonlocal state_dict\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if name in state_dict:\n            full_weight = state_dict[name]\n\n            if mutate_func is not None:\n                full_weight = mutate_func(full_weight)\n            tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n            if tensor is not None:\n                tensor.data.copy_(tensor_chunk[tp_rank])\n        else:\n            print(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor:\n        \"\"\"fetch gate_up tensor in tp shards\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if gate_name in state_dict and up_name in state_dict:\n            gate_weight = state_dict[gate_name]\n            up_weight = state_dict[up_name]\n            new_gate_up_weight = torch.empty(\n                config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                intermediate_size_tp = config.intermediate_size // tp_size\n                gate_weight_tp = gate_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                up_weight_tp = up_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                new_gate_up_weight[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)].copy_(\n                    torch.cat([gate_weight_tp, up_weight_tp], dim=0)\n                )\n\n            tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0)\n            if tensor is not None:\n                tensor.data.copy_(tensor_chunk[tp_rank])\n        else:\n            print(f\"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        assert q_name in state_dict and k_name in state_dict and v_name in state_dict\n        full_weight_q = state_dict[q_name]\n        full_weight_k = state_dict[k_name]\n        full_weight_v = state_dict[v_name]\n\n        hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n        if config.num_key_value_heads >= tp_size:\n            q_size_tp = config.hidden_size // tp_size\n            kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n            total_size = q_size_tp + 2 * kv_size_tp\n            new_weight_qkv = torch.empty(\n                total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                k_part = full_weight_k[i * kv_size_tp : (i + 1) * kv_size_tp]\n                v_part = full_weight_v[i * kv_size_tp : (i + 1) * kv_size_tp]\n                new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0))\n\n        else:\n            q_size_tp = config.hidden_size // tp_size\n            kv_size_tp = hidden_size_per_head\n            total_size = q_size_tp + 2 * kv_size_tp\n            new_weight_qkv = torch.empty(\n                total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head\n                end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head\n                k_part = full_weight_k[start_idx:end_idx]\n                v_part = full_weight_v[start_idx:end_idx]\n                new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0))\n\n        tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0)\n        if tensor is not None:\n            tensor.data.copy_(tensor_chunk[tp_rank])\n\n    # Embeddings\n    # -------------------\n    print_rank_0(\"loading embeddings...\")\n    gpt_model_module = _get_gpt_model(models[0])\n    embed_tokens_weight = None\n    if pp_rank == 0:\n        embed_tokens_weight = gpt_model_module.model.embed_tokens.weight\n    _fetch_tp_shard_tensor_vocab(embed_tokens_weight, \"model.embed_tokens.weight\")\n\n    # Transformer layers\n    # -------------------\n    layer_map = _megatron_calc_layer_map(config)\n\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    num_layer_per_pp = config.num_hidden_layers // pp_size\n    vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()\n\n    layer_list = []\n    if vpp_size is not None:\n        for vpp_rank in range(vpp_size):\n            num_layer_vpp_chunk = num_layer_per_pp // vpp_size\n            num_layer_this_model = num_layer_vpp_chunk\n            offset = vpp_rank * (config.num_hidden_layers // mpu.get_virtual_pipeline_model_parallel_world_size()) + (\n                mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk\n            )\n            layer_list.extend(list(range(offset, offset + num_layer_this_model)))\n    else:\n        num_layer_this_model = num_layer_per_pp\n        offset = pp_rank * num_layer_per_pp\n        layer_list.extend(list(range(offset, offset + num_layer_this_model)))\n\n    for layer in layer_list:\n        print_rank_0(f\"loading layer #{layer}...\")\n        layer_name = f\"model.layers.{layer}\"\n        dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer]\n\n        gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank])\n        sync_layer = gpt_model_module.model.layers[dst_layer_idx]\n\n        _fetch_tensor(\n            sync_layer.input_layernorm.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.input_layernorm.weight\",\n        )\n\n        _fetch_tp_shard_tensor_qkv(\n            sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.self_attn.q_proj.weight\",\n            f\"{layer_name}.self_attn.k_proj.weight\",\n            f\"{layer_name}.self_attn.v_proj.weight\",\n        )\n\n        _fetch_tp_shard_tensor(\n            sync_layer.self_attn.o_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.self_attn.o_proj.weight\",\n            chunk_dim=1,\n        )\n\n        _fetch_tensor(\n            sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.post_attention_layernorm.weight\",\n        )\n\n        _fetch_tp_shard_tensor_gate_up(\n            sync_layer.mlp.gate_up_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.mlp.gate_proj.weight\",\n            f\"{layer_name}.mlp.up_proj.weight\",\n        )\n\n        _fetch_tp_shard_tensor(\n            sync_layer.mlp.down_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.mlp.down_proj.weight\",\n            chunk_dim=1,\n        )\n    # Final Layernorm\n    # -------------------\n    print_rank_0(\"loading final layernorm...\")\n    gpt_model_module = _get_gpt_model(models[-1])\n    _fetch_tensor(\n        getattr(gpt_model_module.model.norm, \"weight\", None),\n        \"model.norm.weight\",\n    )\n\n    print_rank_0(\"loading lm_head...\")\n    if pp_rank + 1 == pp_size:\n        lm_head_weight = gpt_model_module.lm_head.weight\n\n        if is_value_model:\n            if \"lm_head.weight\" in state_dict and state_dict[\"lm_head.weight\"].shape[0] == 1:\n                _fetch_tensor(lm_head_weight, \"lm_head.weight\")\n                print_rank_0(\"load lm_head weight\")\n            elif \"reward_head.weight\" in state_dict and state_dict[\"reward_head.weight\"].shape[0] == 1:\n                _fetch_tensor(lm_head_weight, \"reward_head.weight\")\n                print_rank_0(\"load lm_head from value_head weight\")\n            else:\n                _fetch_tensor(None, \"lm_head.weight\")\n                print_rank_0(\"fail to match lm_head in value_model\")\n        else:\n            _fetch_tp_shard_tensor(lm_head_weight, \"lm_head.weight\")\n\n    dist.barrier()\n    get_torch_device().empty_cache()\n    print_rank_0(f\"loading megatron ckpt done, time elapsed {time.time() - start_time}s\")\n"
  },
  {
    "path": "verl/models/llama/megatron/checkpoint_utils/llama_loader_depracated.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    print(f\"get megatron data parallel size: {mpu.get_data_parallel_world_size()}\")\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef load_state_dict_to_megatron_llama(\n    state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False\n):\n    \"\"\"Load merged state_dict to sharded Megatron module in training.\"\"\"\n    from megatron.core import DistributedDataParallel as LocalDDP\n    from megatron.core import mpu\n    from megatron.core.transformer.module import Float16Module\n    from torch.nn.parallel import DistributedDataParallel as torchDDP\n\n    from verl.utils.logger import print_rank_0\n    from verl.utils.megatron_utils import unwrap_model\n\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    def broadcast_params(module):\n        for param in module.parameters():\n            torch.distributed.broadcast(\n                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()\n            )\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if torch.distributed.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers, (\n        f\"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size \"\n        f\"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}\"\n    )\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        gpt_model_module = _get_gpt_model(models[i])\n        assert len(gpt_model_module.model.layers) == num_layers_per_model\n\n    def _broadcast_tensor(tensor, name) -> torch.Tensor:\n        \"\"\"broadcast tensor from rank0 across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                weight = state_dict[name]\n                tensor_shape = weight.shape\n            else:\n                tensor_shape = None\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not in state_dict, skip load\")\n            return\n\n        if tensor is None:\n            tensor = torch.empty(\n                tensor_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        if torch.distributed.get_rank() == 0:\n            tensor.data.copy_(weight)\n        dist.broadcast(tensor, src=0, group=mp_group)\n\n    def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                full_weight = state_dict[name]\n\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                full_weight = state_dict[name]\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            gate_weight = state_dict[gate_name]\n            up_weight = state_dict[up_name]\n            new_gate_up_weight = torch.empty(\n                config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                intermediate_size_tp = config.intermediate_size // tp_size\n                gate_weight_tp = gate_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                up_weight_tp = up_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                new_gate_up_weight[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)].copy_(\n                    torch.cat([gate_weight_tp, up_weight_tp], dim=0)\n                )\n\n            tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank() == 0:} tensor {gate_name, up_name} shape \"\n                f\"{tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            assert q_name in state_dict and k_name in state_dict and v_name in state_dict\n            full_weight_q = state_dict[q_name]\n            full_weight_k = state_dict[k_name]\n            full_weight_v = state_dict[v_name]\n\n            hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                new_weight_qkv = torch.empty(\n                    total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                )\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    k_part = full_weight_k[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    v_part = full_weight_v[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(\n                        torch.cat([q_part, k_part, v_part], dim=0)\n                    )\n\n            else:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                new_weight_qkv = torch.empty(\n                    total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                )\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head\n                    end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head\n                    k_part = full_weight_k[start_idx:end_idx]\n                    v_part = full_weight_v[start_idx:end_idx]\n                    new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(\n                        torch.cat([q_part, k_part, v_part], dim=0)\n                    )\n\n            tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name, k_name, v_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    if dp_rank == 0:\n        # Embeddings\n        # -------------------\n        print_rank_0(\"loading embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        embed_tokens_weight = None\n        if pp_rank == 0:\n            embed_tokens_weight = gpt_model_module.model.embed_tokens.weight\n        _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, \"model.embed_tokens.weight\")\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n\n        for layer in range(config.num_hidden_layers):\n            print_rank_0(f\"loading layer #{layer}...\")\n            layer_name = f\"model.layers.{layer}\"\n            dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank])\n            sync_layer = gpt_model_module.model.layers[dst_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.input_layernorm.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.input_layernorm.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attn.o_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                chunk_dim=1,\n            )\n\n            _broadcast_tensor(\n                sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.gate_up_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.down_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                chunk_dim=1,\n            )\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"loading final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.model.norm, \"weight\", None),\n            \"model.norm.weight\",\n        )\n\n        print_rank_0(\"loading lm_head...\")\n        lm_head_weight = None\n        if pp_rank + 1 == pp_size:\n            lm_head_weight = gpt_model_module.lm_head.weight\n\n        if is_value_model:\n            if \"lm_head.weight\" in state_dict and state_dict[\"lm_head.weight\"].shape[0] == 1:\n                _broadcast_tensor(lm_head_weight, \"lm_head.weight\")\n                print_rank_0(\"load lm_head weight\")\n            elif \"reward_head.weight\" in state_dict and state_dict[\"reward_head.weight\"].shape[0] == 1:\n                _broadcast_tensor(lm_head_weight, \"reward_head.weight\")\n                print_rank_0(\"load lm_head from value_head weight\")\n            else:\n                _broadcast_tensor(None, \"lm_head.weight\")\n                print_rank_0(\"fail to match lm_head in value_model\")\n        else:\n            _broadcast_tp_shard_tensor(lm_head_weight, \"lm_head.weight\")\n    dist.barrier()\n    # Broadcast weights inside data parallel groups\n    for wrapped_model in wrapped_models:\n        broadcast_params(wrapped_model)\n\n    get_torch_device().empty_cache()\n    print_rank_0(f\"loading megatron ckpt done, time elapsed {time.time() - start_time}s\")\n"
  },
  {
    "path": "verl/models/llama/megatron/checkpoint_utils/llama_saver.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\nfrom megatron.core import mpu\nfrom megatron.core.distributed import DistributedDataParallel as LocalDDP\nfrom megatron.core.transformer.module import Float16Module\nfrom torch.nn.parallel import DistributedDataParallel as torchDDP\n\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.logger import print_rank_0\nfrom verl.utils.megatron_utils import unwrap_model\n\n\ndef _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0):\n    \"\"\"given TP,DP,PP rank to get the global rank.\"\"\"\n\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    dp_size = mpu.get_data_parallel_world_size()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    assert tp_size * dp_size * pp_size == torch.distributed.get_world_size(), (\n        f\"{tp_size} x {dp_size} x {pp_size} != {torch.distributed.get_world_size()}\"\n    )\n    # We only support TP-DP-PP grouping, for correctness when resharding\n    return (pp_rank * dp_size + dp_rank) * tp_size + tp_rank\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef merge_megatron_ckpt_llama(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False):\n    \"\"\"Merge sharded parameters of a Megatron module into a merged checkpoint.\n\n    Args:\n        wrapped_models (list of megatron.core.distributed.DistributedDataParallel):\n            The local DDP wrapped megatron modules.\n        config (str or None):\n            HF config for model\n        dtype: model params type\n        is_value_model: if model is value model\n        tie_word_embeddings: tie_word_embeddings, not used in llama, only to keep same interface with qwen2\n    Returns:\n        state_dict (dict):\n            The merged state_dict in rank 0, and an empty dictionary in other ranks.\n    \"\"\"\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if dist.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        assert len(models[i].model.layers) == num_layers_per_model, (\n            \"len model layers {} not equal to num_layers_per_model {}\".format(\n                len(models[i].model.layers), num_layers_per_model\n            )\n        )\n\n    state_dict = dict()\n\n    def _get_cpu_tensor(tensor: torch.Tensor):\n        if tensor is None:\n            return None\n        if tensor.device == torch.device(\"cpu\"):\n            return tensor.detach().clone()\n        return tensor.detach().cpu()\n\n    def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        if torch.distributed.get_rank() == src_rank:\n            if tensor is None:\n                weight = None\n                tensor_shape = None\n            else:\n                weight = tensor\n                tensor_shape = weight.shape\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not exist, skip collect\")\n            return\n\n        if weight is None:\n            weight = torch.empty(\n                tensor_shape,\n                dtype=dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n\n        dist.broadcast(weight, src=src_rank, group=mp_group)\n\n        if torch.distributed.get_rank() == 0:\n            state_dict[name] = _get_cpu_tensor(weight)\n\n    def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=concat_dim)\n            if mutate_func is not None:\n                full_tensor = mutate_func(full_tensor)\n            state_dict[name] = full_tensor\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            intermediate_size_tp = config.intermediate_size // tp_size\n            gate_weight_list = []\n            up_weight_list = []\n            for i in range(tp_size):\n                gate_up_weight_tp = full_tensor[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)]\n                gate_weight_tp = gate_up_weight_tp[:intermediate_size_tp]\n                up_weight_tp = gate_up_weight_tp[intermediate_size_tp:]\n                gate_weight_list.append(gate_weight_tp)\n                up_weight_list.append(up_weight_tp)\n\n            state_dict[gate_name] = torch.cat(gate_weight_list, dim=0)\n            state_dict[up_name] = torch.cat(up_weight_list, dim=0)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank):\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            q_weight_list = []\n            k_weight_list = []\n            v_weight_list = []\n            hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_part = qkv_part[:q_size_tp]\n                    k_part = qkv_part[q_size_tp : q_size_tp + kv_size_tp]\n                    v_part = qkv_part[q_size_tp + kv_size_tp : total_size]\n                    q_weight_list.append(q_part)\n                    k_weight_list.append(k_part)\n                    v_weight_list.append(v_part)\n            else:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_part = qkv_part[:q_size_tp]\n                    k_part = qkv_part[q_size_tp : q_size_tp + kv_size_tp]\n                    v_part = qkv_part[q_size_tp + kv_size_tp : total_size]\n                    q_weight_list.append(q_part)\n                    if i * config.num_key_value_heads % tp_size == 0:\n                        k_weight_list.append(k_part)\n                        v_weight_list.append(v_part)\n\n            state_dict[q_name] = torch.cat(q_weight_list, dim=0)\n            state_dict[k_name] = torch.cat(k_weight_list, dim=0)\n            state_dict[v_name] = torch.cat(v_weight_list, dim=0)\n\n    # empty cache before collecting weights\n    get_torch_device().empty_cache()\n    # Embeddings\n    # -------------------\n    if dp_rank == 0:\n        # Embeddings\n        # -------------------\n        print_rank_0(\"collecting embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        _broadcast_tp_shard_tensor(\n            gpt_model_module.model.embed_tokens.weight if pp_rank == 0 else None,\n            \"model.embed_tokens.weight\",\n            src_pp_rank=0,\n        )\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n        for layer in range(config.num_hidden_layers):\n            print_rank_0(f\"collecting layer #{layer}...\")\n            layer_name = f\"model.layers.{layer}\"\n            src_pp_rank, src_virtual_pp_rank, src_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[src_virtual_pp_rank])\n            sync_layer = gpt_model_module.model.layers[src_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.input_layernorm.weight,\n                f\"{layer_name}.input_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.weight,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attn.o_proj.weight,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tensor(\n                sync_layer.post_attention_layernorm.weight,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.gate_up_proj.weight,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.down_proj.weight,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"collecting final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.model.norm, \"weight\", None),\n            \"model.norm.weight\",\n            src_pp_rank=pp_size - 1,\n        )\n\n        print_rank_0(\"collecting lm_head...\")\n\n        if is_value_model:\n            if pp_rank == pp_size - 1:\n                print(f\"gpt_model_module.lm_head.weight: {gpt_model_module.lm_head.weight.shape}\")\n            _broadcast_tensor(\n                gpt_model_module.lm_head.weight if pp_rank == pp_size - 1 else None,\n                \"lm_head.weight\",\n                src_pp_rank=pp_size - 1,\n            )\n            _broadcast_tensor(\n                gpt_model_module.reward_head.weight\n                if pp_rank == pp_size - 1 and getattr(gpt_model_module, \"reward_weight\", None) is not None\n                else None,\n                \"reward_head.weight\",\n                src_pp_rank=pp_size - 1,\n            )\n\n        else:\n            _broadcast_tp_shard_tensor(\n                getattr(gpt_model_module.lm_head, \"weight\", None) if pp_rank == pp_size - 1 else None,\n                \"lm_head.weight\",\n                src_pp_rank=pp_size - 1,\n            )\n\n    dist.barrier()\n\n    get_torch_device().empty_cache()\n    if torch.distributed.get_rank() == 0:\n        if dtype not in [torch.float16, torch.bfloat16, torch.float32]:\n            print(f'Unknown/unsupported dtype to save: {dtype}\"')\n            exit(1)\n        for k, v in state_dict.items():\n            if dtype != v.dtype:\n                state_dict[k] = v.to(dtype)\n\n    print_rank_0(f\"merge megatron ckpt done, time elapsed {time.time() - start_time}s\")\n    return state_dict\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .parallel_attention import ParallelLlamaAttention\nfrom .parallel_decoder import ParallelLlamaDecoderLayer, ParallelLlamaDecoderLayerRmPad\nfrom .parallel_linear import (\n    LinearForLastLayer,\n    MergedColumnParallelLinear,\n    QKVParallelLinear,\n)\nfrom .parallel_mlp import ParallelLlamaMLP\nfrom .parallel_rmsnorm import ParallelLlamaRMSNorm\n\n__all__ = [\n    \"LinearForLastLayer\",\n    \"MergedColumnParallelLinear\",\n    \"QKVParallelLinear\",\n    \"ParallelLlamaAttention\",\n    \"ParallelLlamaDecoderLayer\",\n    \"ParallelLlamaDecoderLayerRmPad\",\n    \"ParallelLlamaMLP\",\n    \"ParallelLlamaRMSNorm\",\n]\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/parallel_attention.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport math\nfrom typing import Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom flash_attn.layers.rotary import apply_rotary_emb\nfrom megatron.core import ModelParallelConfig, tensor_parallel\nfrom megatron.core import parallel_state as mpu\nfrom torch import nn\nfrom transformers import LlamaConfig\nfrom transformers.utils import is_flash_attn_2_available\n\nfrom verl.models.llama.megatron.layers.parallel_linear import QKVParallelLinear\nfrom verl.utils.megatron import tensor_parallel as tp_utils\n\n\nclass LlamaRotaryEmbedding(nn.Module):\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):\n        super().__init__()\n\n        self.dim = dim\n        self.max_position_embeddings = max_position_embeddings\n        self.base = base\n        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        # Build here to make `torch.jit.trace` work.\n        self._set_cos_sin_cache(\n            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()\n        )\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n    def forward(self, x, seq_len=None):\n        # x: [bs, num_attention_heads, seq_len, head_size]\n        if seq_len > self.max_seq_len_cached:\n            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)\n\n        return (\n            self.cos_cached[:seq_len].to(dtype=x.dtype),\n            self.sin_cached[:seq_len].to(dtype=x.dtype),\n        )\n\n\nclass LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):\n    \"\"\"LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n        t = t / self.scaling_factor\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n\nclass LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):\n    \"\"\"LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n\n        if seq_len > self.max_position_embeddings:\n            base = self.base * (\n                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)\n            ) ** (self.dim / (self.dim - 2))\n            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n            self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n\nclass LlamaLlama3ScalingRotaryEmbedding(LlamaRotaryEmbedding):\n    def __init__(self, dim, config, max_position_embeddings=2048, base=10000, device=None):\n        super().__init__(dim, max_position_embeddings, base, device)\n\n        self.factor = config.rope_scaling[\"factor\"]  # `8` in the original implementation\n        self.high_freq_factor = config.rope_scaling[\"high_freq_factor\"]  # `1` in the original implementation\n        self.low_freq_factor = config.rope_scaling[\"low_freq_factor\"]  # `4` in the original implementation\n        self.old_context_len = config.rope_scaling[\n            \"original_max_position_embeddings\"\n        ]  # `8192` in the original implementation\n\n        low_freq_wavelen = self.old_context_len / self.low_freq_factor\n        high_freq_wavelen = self.old_context_len / self.high_freq_factor\n\n        wavelen = 2 * math.pi / self.inv_freq\n        # wavelen < high_freq_wavelen: do nothing; wavelen > low_freq_wavelen: divide by factor\n        inv_freq_llama = torch.where(wavelen > low_freq_wavelen, self.inv_freq / self.factor, self.inv_freq)\n        # otherwise: interpolate between the two, using a smooth factor\n        smooth_factor = (self.old_context_len / wavelen - self.low_freq_factor) / (\n            self.high_freq_factor - self.low_freq_factor\n        )\n        smoothed_inv_freq = (1 - smooth_factor) * inv_freq_llama / self.factor + smooth_factor * inv_freq_llama\n        is_medium_freq = ~(wavelen < high_freq_wavelen) * ~(wavelen > low_freq_wavelen)\n        inv_freq = torch.where(is_medium_freq, smoothed_inv_freq, inv_freq_llama)\n\n        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        # Build here to make `torch.jit.trace` work.\n        self._set_cos_sin_cache(\n            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()\n        )\n\n\ndef rotate_half(x):\n    \"\"\"Rotates half the hidden dims of the input.\"\"\"\n    x1 = x[..., : x.shape[-1] // 2]\n    x2 = x[..., x.shape[-1] // 2 :]\n    return torch.cat((-x2, x1), dim=-1)\n\n\ndef apply_rotary_pos_emb(q, k, cos, sin, position_ids):\n    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n    return q_embed, k_embed\n\n\ndef repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:\n    \"\"\"\n    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,\n    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)\n    \"\"\"\n    batch, num_key_value_heads, slen, head_dim = hidden_states.shape\n    if n_rep == 1:\n        return hidden_states\n    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)\n    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)\n\n\nclass ParallelLlamaAttention(nn.Module):\n    \"\"\"Multi-headed attention from 'Attention Is All You Need' paper\"\"\"\n\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config = config\n        self.megatron_config = megatron_config\n        self.hidden_size = config.hidden_size\n        self.num_heads = config.num_attention_heads\n        self.head_dim = self.hidden_size // self.num_heads\n        self.num_key_value_heads = config.num_key_value_heads\n        self.num_key_value_groups = self.num_heads // self.num_key_value_heads\n        self.max_position_embeddings = config.max_position_embeddings\n        self.rope_theta = config.rope_theta\n\n        # assign values after tp\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        assert self.num_heads % tp_size == 0, (\n            f\"num_head must be divisible by tp_size. Got num_head={self.num_heads}, tp_size={tp_size}\"\n        )\n        assert self.num_key_value_heads % tp_size == 0, (\n            f\"num_key_value_heads must be divisible by tp_size. Got num_key_value_heads=\"\n            f\"{self.num_key_value_heads}, tp_size={tp_size}\"\n        )\n\n        self.num_heads_per_tp = self.num_heads // tp_size\n        self.num_key_value_heads_per_tp = self.num_key_value_heads // tp_size\n        self.hidden_size_per_tp = self.hidden_size // tp_size\n\n        if (self.head_dim * self.num_heads) != self.hidden_size:\n            raise ValueError(\n                f\"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and \"\n                f\"`num_heads`: {self.num_heads}).\"\n            )\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear()\n\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            assert row_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, megatron_config)\n            tp_utils.update_kwargs_with_config(row_kwargs, megatron_config)\n\n        # [self.q_size, self.k_size, self.v_size]\n        self.qkv_proj = QKVParallelLinear(\n            input_size=self.hidden_size,\n            num_heads=self.num_heads,\n            num_key_value_heads=self.num_key_value_heads,\n            head_dim=self.head_dim,\n            bias=config.attention_bias,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n        self.q_size = self.num_heads_per_tp * self.head_dim\n        self.k_size = self.num_key_value_heads_per_tp * self.head_dim\n        self.v_size = self.num_key_value_heads_per_tp * self.head_dim\n\n        self.o_proj = tensor_parallel.RowParallelLinear(\n            input_size=self.num_heads * self.head_dim,\n            output_size=self.hidden_size,\n            bias=config.attention_bias,\n            input_is_parallel=True,\n            skip_bias_add=False,\n            **row_kwargs,\n        )\n\n        self._init_rope()\n\n    def _init_rope(self):\n        if self.config.rope_scaling is None:\n            self.rotary_emb = LlamaRotaryEmbedding(\n                self.head_dim,\n                max_position_embeddings=self.max_position_embeddings,\n                base=self.rope_theta,\n            )\n        else:\n            rope_type_key = \"type\" if \"type\" in self.config.rope_scaling else \"rope_type\"\n            scaling_type = self.config.rope_scaling[rope_type_key]\n            scaling_factor = self.config.rope_scaling[\"factor\"]\n            if scaling_type == \"linear\":\n                self.rotary_emb = LlamaLinearScalingRotaryEmbedding(\n                    self.head_dim,\n                    max_position_embeddings=self.max_position_embeddings,\n                    scaling_factor=scaling_factor,\n                    base=self.rope_theta,\n                )\n            elif scaling_type == \"dynamic\":\n                self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(\n                    self.head_dim,\n                    max_position_embeddings=self.max_position_embeddings,\n                    scaling_factor=scaling_factor,\n                    base=self.rope_theta,\n                )\n            elif scaling_type == \"llama3\":\n                self.rotary_emb = LlamaLlama3ScalingRotaryEmbedding(\n                    self.head_dim,\n                    self.config,\n                    max_position_embeddings=self.max_position_embeddings,\n                    base=self.rope_theta,\n                )\n            else:\n                raise ValueError(f\"Unknown RoPE scaling type {scaling_type}\")\n\n    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):\n        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n        bsz, q_len, _ = hidden_states.size()\n        qkv = self.qkv_proj(hidden_states)[0]\n        query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)\n\n        query_states = query_states.view(bsz, q_len, self.num_heads_per_tp, self.head_dim).transpose(1, 2)\n        key_states = key_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2)\n        value_states = value_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2)\n\n        kv_seq_len = key_states.shape[-2]\n        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)\n        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)\n\n        key_states = repeat_kv(key_states, self.num_key_value_groups)\n        value_states = repeat_kv(value_states, self.num_key_value_groups)\n\n        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)\n\n        if attn_weights.size() != (bsz, self.num_heads_per_tp, q_len, kv_seq_len):\n            raise ValueError(\n                f\"Attention weights should be of size {(bsz, self.num_heads_per_tp, q_len, kv_seq_len)}, \"\n                f\"but is {attn_weights.size()}\"\n            )\n\n        if attention_mask is not None:\n            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):\n                raise ValueError(\n                    f\"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}\"\n                )\n            attn_weights = attn_weights + attention_mask\n\n        # upcast attention to fp32\n        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)\n        attn_output = torch.matmul(attn_weights, value_states)\n\n        if attn_output.size() != (bsz, self.num_heads_per_tp, q_len, self.head_dim):\n            raise ValueError(\n                f\"`attn_output` should be of size {(bsz, self.num_heads_per_tp, q_len, self.head_dim)}, \"\n                f\"but is {attn_output.size()}\"\n            )\n\n        attn_output = attn_output.transpose(1, 2).contiguous()\n        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size_per_tp)\n        attn_output = self.o_proj(attn_output)[0]\n        return attn_output\n\n\n\"\"\"\nRemove padding Attention\n- Using Flash-attn 2\n- Compatible with sequence parallel\n\"\"\"\n\n\nif is_flash_attn_2_available():\n    from flash_attn import flash_attn_varlen_func\n    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa: F401\n\n\ndef apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length):\n    batch_size = position_ids.shape[0]\n\n    q = pad_input(q, indices, batch_size, sequence_length)  # (batch_size, seqlen, num_head, head_dim)\n    k = pad_input(k, indices, batch_size, sequence_length)\n    cos = cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]\n    sin = sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n\n    q_embed = index_first_axis(rearrange(q_embed, \"b s ... -> (b s) ...\"), indices)\n    k_embed = index_first_axis(rearrange(k_embed, \"b s ... -> (b s) ...\"), indices)\n\n    return q_embed, k_embed\n\n\n# use flash-attn rotary embeddings with rmpad\n# cos/sin shoudl be: (seq_length, rotary_dim / 2)\ndef apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seqlen):\n    q_embed = apply_rotary_emb(\n        q, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen\n    )\n    k_embed = apply_rotary_emb(\n        k, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen\n    )\n    return q_embed, k_embed\n\n\nclass ParallelLlamaAttentionRmPad(ParallelLlamaAttention):\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: torch.Tensor = None,\n        max_seqlen_in_batch: int = None,\n    ):\n        total_nnz, _, _ = hidden_states.size()  # This is the total_nnz padded after sequence parallel\n\n        if self.megatron_config.sequence_parallel:\n            total_nnz = total_nnz * mpu.get_tensor_model_parallel_world_size()\n\n        qkv = self.qkv_proj(hidden_states)[0]\n        query_states, key_states, value_states = qkv.split(\n            [self.q_size, self.k_size, self.v_size], dim=-1\n        )  # (total_nnz, 1, hidden_size)\n\n        if self.megatron_config.sequence_parallel:\n            sequence_parallel_pad = total_nnz - cu_seqlens[-1]\n            total_nnz = cu_seqlens[-1]  # total_nnz before sp padding\n            query_states = query_states[:total_nnz]\n            key_states = key_states[:total_nnz]\n            value_states = value_states[:total_nnz]\n\n        # Flash attention requires the input to have the shape\n        # batch_size x seq_length x head_dime x hidden_dim\n        # therefore we just need to keep the original shape\n        query_states = query_states.view(total_nnz, self.num_heads_per_tp, self.head_dim)\n        key_states = key_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim)\n        value_states = value_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim)\n\n        cos, sin = self.rotary_emb(value_states, seq_len=sequence_length)\n        cos, sin = cos[:, : cos.shape[1] // 2], sin[:, : sin.shape[1] // 2]  # flash attn only needs half\n        query_states, key_states = apply_rotary_pos_emb_rmpad_flash(\n            query_states, key_states, cos, sin, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen_in_batch\n        )\n        # query_states, key_states = apply_rotary_pos_emb_rmpad(query_states, key_states, cos, sin,\n        # position_ids, indices,\n\n        # TODO: llama does not have dropout in the config??\n        # It is recommended to use dropout with FA according to the docs\n        # when training.\n        dropout_rate = 0.0  # if not self.training else self.attn_dropout\n\n        # In PEFT, usually we cast the layer norms in float32 for training stability reasons\n        # therefore the input hidden states gets silently casted in float32. Hence, we need\n        # cast them back in float16 just to be sure everything works as expected.\n        # This might slowdown training & inference so it is recommended to not cast the LayerNorms\n        # in fp32. (LlamaRMSNorm handles it correctly)\n        input_dtype = query_states.dtype\n        if input_dtype == torch.float32:\n            query_states = query_states.to(torch.float16)\n            key_states = key_states.to(torch.float16)\n            value_states = value_states.to(torch.float16)\n\n        attn_output_unpad = flash_attn_varlen_func(\n            query_states,\n            key_states,\n            value_states,\n            cu_seqlens_q=cu_seqlens,\n            cu_seqlens_k=cu_seqlens,\n            max_seqlen_q=max_seqlen_in_batch,\n            max_seqlen_k=max_seqlen_in_batch,\n            dropout_p=dropout_rate,\n            softmax_scale=None,\n            causal=True,\n        )\n\n        attn_output_unpad = attn_output_unpad.to(input_dtype)\n        attn_output_unpad = attn_output_unpad.reshape(total_nnz, 1, self.hidden_size_per_tp).contiguous()\n\n        # sequence parallel reduce_scatter is performed inside RowColumnParallel if enabled\n        # Here we need to repad\n        if self.megatron_config.sequence_parallel:\n            attn_output_unpad = F.pad(attn_output_unpad, pad=(0, 0, 0, 0, 0, sequence_parallel_pad))\n\n        attn_output_unpad = self.o_proj(attn_output_unpad)[0]\n        return attn_output_unpad\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/parallel_decoder.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\nfrom megatron.core import ModelParallelConfig\nfrom torch import nn\nfrom transformers import LlamaConfig\n\nfrom verl.utils.megatron_utils import TransformerConfig, convert_config\n\nfrom .parallel_attention import ParallelLlamaAttention, ParallelLlamaAttentionRmPad\nfrom .parallel_mlp import ParallelLlamaMLP\nfrom .parallel_rmsnorm import ParallelLlamaRMSNorm\n\n\nclass ParallelLlamaDecoderLayer(nn.Module):\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, layer_idx: int):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.layer_idx = layer_idx\n        self.hidden_size = config.hidden_size\n        self.self_attn = ParallelLlamaAttention(config=config, megatron_config=megatron_config)\n\n        self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config)\n        self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config)\n        self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:\n        \"\"\"\n        Args:\n            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`\n            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size\n                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.\n            output_attentions (`bool`, *optional*):\n                Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n                returned tensors for more detail.\n            use_cache (`bool`, *optional*):\n                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding\n                (see `past_key_values`).\n            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states\n        \"\"\"\n\n        residual = hidden_states\n\n        hidden_states = self.input_layernorm(hidden_states)\n\n        # Note: sequence parallel is hidden inside ColumnParallelLinear\n        # reduce scatter is hidden inside RowParallelLinear\n\n        # Self Attention\n        hidden_states = self.self_attn(\n            hidden_states=hidden_states,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n        )\n\n        # TODO: add sequence parallel operator reduce_scatter here\n\n        hidden_states = residual + hidden_states\n\n        # Fully Connected\n        residual = hidden_states\n        hidden_states = self.post_attention_layernorm(hidden_states)\n\n        # TODO: add sequence parallel operator all_gather here\n\n        hidden_states = self.mlp(hidden_states)\n\n        # TODO: add sequence parallel operator reduce_scatter here\n\n        hidden_states = residual + hidden_states\n\n        outputs = hidden_states\n\n        return outputs\n\n\nclass ParallelLlamaDecoderLayerRmPad(nn.Module):\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, layer_idx: int):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.layer_idx = layer_idx\n        self.hidden_size = config.hidden_size\n        self.self_attn = ParallelLlamaAttentionRmPad(config=config, megatron_config=megatron_config)\n\n        self.mlp = ParallelLlamaMLP(config, megatron_config=megatron_config)\n        self.input_layernorm = ParallelLlamaRMSNorm(config, megatron_config)\n        self.post_attention_layernorm = ParallelLlamaRMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:\n        residual = hidden_states  # (total_nnz // sp, 1, hidden_size)\n\n        hidden_states = self.input_layernorm(hidden_states)\n\n        # Self Attention\n        # (total_nnz // sp, 1, hidden_size) -> all-gather (total_nnz, 1, hidden_size)\n        # -> col + row -> reduce-scatter -> (total_nnz // sp, 1, hidden_size)\n        hidden_states = self.self_attn(\n            hidden_states=hidden_states,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        hidden_states = residual + hidden_states\n\n        # Fully Connected\n        # shape changes same as attn\n        residual = hidden_states\n        hidden_states = self.post_attention_layernorm(hidden_states)\n        hidden_states = self.mlp(hidden_states)\n        hidden_states = residual + hidden_states\n\n        outputs = hidden_states\n\n        return outputs\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/parallel_linear.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023 The vLLM team.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py\n\nimport torch\nfrom megatron.core import tensor_parallel\n\n\nclass QKVParallelLinear(tensor_parallel.ColumnParallelLinear):\n    def __init__(\n        self,\n        input_size,\n        num_heads,\n        num_key_value_heads,\n        head_dim,\n        *,\n        bias=True,\n        gather_output=True,\n        skip_bias_add=False,\n        **kwargs,\n    ):\n        # Keep input parameters, and already restrict the head numbers\n        self.input_size = input_size\n        self.q_output_size = num_heads * head_dim\n        self.kv_output_size = num_key_value_heads * head_dim\n        self.head_dim = head_dim\n        self.gather_output = gather_output\n        self.skip_bias_add = skip_bias_add\n\n        input_size = self.input_size\n        output_size = (num_heads + 2 * num_key_value_heads) * self.head_dim\n\n        super().__init__(\n            input_size=input_size,\n            output_size=output_size,\n            bias=bias,\n            gather_output=gather_output,\n            skip_bias_add=skip_bias_add,\n            **kwargs,\n        )\n\n\nclass MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear):\n    def __init__(\n        self,\n        input_size,\n        gate_ouput_size,\n        up_output_size,\n        *,\n        bias=True,\n        gather_output=True,\n        skip_bias_add=False,\n        **kwargs,\n    ):\n        # Keep input parameters, and already restrict the head numbers\n        self.input_size = input_size\n        self.output_size = gate_ouput_size + up_output_size\n        self.gather_output = gather_output\n        self.skip_bias_add = skip_bias_add\n\n        super().__init__(\n            input_size=self.input_size,\n            output_size=self.output_size,\n            bias=bias,\n            gather_output=gather_output,\n            skip_bias_add=skip_bias_add,\n            **kwargs,\n        )\n\n\nclass LinearForLastLayer(torch.nn.Linear):\n    def __init__(\n        self,\n        input_size,\n        output_size,\n        *,\n        config,\n        bias=True,\n    ):\n        super().__init__(in_features=input_size, out_features=output_size, bias=bias)\n        self.sequence_parallel = config.sequence_parallel\n        if self.sequence_parallel:\n            self.weight.sequence_parallel = True\n\n    def forward(\n        self,\n        input_,\n        weight=None,\n        runtime_gather_output=None,\n    ):\n        logits = super().forward(input_)\n        logits = logits.float()\n        if self.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits, None\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/parallel_mlp.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom megatron.core import ModelParallelConfig, tensor_parallel\nfrom megatron.core import parallel_state as mpu\nfrom torch import nn\nfrom transformers.activations import ACT2FN\n\nfrom verl.models.llama.megatron.layers.parallel_linear import MergedColumnParallelLinear\nfrom verl.utils.megatron import tensor_parallel as tp_utils\n\n\nclass ParallelLlamaMLP(nn.Module):\n    def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None:\n        super().__init__()\n        self.config = config\n        self.hidden_size = config.hidden_size\n        self.intermediate_size = config.intermediate_size\n        # The weight is only [hidden_size, intermediate_size // model_parallel_world_size]\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear()\n\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            assert row_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(row_kwargs, megatron_config)\n            tp_utils.update_kwargs_with_config(column_kwargs, megatron_config)\n\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        self.gate_up_proj = MergedColumnParallelLinear(\n            input_size=self.hidden_size,\n            gate_ouput_size=self.intermediate_size,\n            up_output_size=self.intermediate_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n        self.gate_size = self.intermediate_size // tp_size\n\n        self.down_proj = tensor_parallel.RowParallelLinear(\n            input_size=self.intermediate_size,\n            output_size=self.hidden_size,\n            bias=False,\n            input_is_parallel=True,\n            skip_bias_add=False,\n            **row_kwargs,\n        )\n\n        self.act_fn = ACT2FN[config.hidden_act]\n\n    def forward(self, x):\n        gate_up = self.gate_up_proj(x)[0]\n        gate, up = gate_up.split(self.gate_size, dim=-1)\n        return self.down_proj(self.act_fn(gate) * up)[0]\n"
  },
  {
    "path": "verl/models/llama/megatron/layers/parallel_rmsnorm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport numbers\n\nimport torch\nfrom megatron.core import ModelParallelConfig\nfrom torch import nn\nfrom transformers import LlamaConfig\n\nfrom verl.utils.megatron import sequence_parallel as sp_utils\n\n\nclass ParallelLlamaRMSNorm(nn.Module):\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        \"\"\"\n        LlamaRMSNorm is equivalent to T5LayerNorm\n        \"\"\"\n        super().__init__()\n        if isinstance(config.hidden_size, numbers.Integral):\n            normalized_shape = (config.hidden_size,)\n        self.normalized_shape = torch.Size(normalized_shape)\n        self.weight = nn.Parameter(torch.ones(self.normalized_shape))\n        self.variance_epsilon = config.rms_norm_eps\n\n        if megatron_config.sequence_parallel:\n            sp_utils.mark_parameter_as_sequence_parallel(self.weight)\n\n    def forward(self, hidden_states):\n        from apex.normalization.fused_layer_norm import fused_rms_norm_affine\n\n        return fused_rms_norm_affine(\n            input=hidden_states,\n            weight=self.weight,\n            normalized_shape=self.normalized_shape,\n            eps=self.variance_epsilon,\n            memory_efficient=True,\n        )\n"
  },
  {
    "path": "verl/models/llama/megatron/modeling_llama_megatron.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"PyTorch LLaMA model with Megatron-style acceleration.\"\"\"\n\nfrom typing import Optional\n\nimport torch\nimport torch.utils.checkpoint\nfrom megatron.core import ModelParallelConfig, mpu, tensor_parallel\nfrom torch import nn\nfrom transformers.modeling_outputs import BaseModelOutputWithPast\nfrom transformers.models.llama.configuration_llama import LlamaConfig\nfrom transformers.models.llama.modeling_llama import CausalLMOutputWithPast\n\nfrom verl.utils.megatron import sequence_parallel as sp_utils\nfrom verl.utils.megatron import tensor_parallel as tp_utils\nfrom verl.utils.megatron_utils import TransformerConfig, convert_config\n\nfrom .layers import ParallelLlamaDecoderLayer, ParallelLlamaDecoderLayerRmPad, ParallelLlamaRMSNorm\n\n\"\"\"\nTODO: \n1. Add weight initialization. Here we need to be careful on TP weight init.\n2. Add sequence parallel\n3. Load checkpoint from meta LLama pretrained checkpoint\n\"\"\"\n\n\n# Copied from transformers.models.bart.modeling_bart._make_causal_mask\ndef _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device):\n    \"\"\"\n    Make causal mask used for bi-directional self-attention.\n    \"\"\"\n    bsz, tgt_len = input_ids_shape\n    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)\n    mask_cond = torch.arange(mask.size(-1), device=device)\n    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)\n    mask = mask.to(dtype)\n    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len)\n\n\n# Copied from transformers.models.bart.modeling_bart._expand_mask\ndef _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):\n    \"\"\"\n    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.\n    \"\"\"\n    bsz, src_len = mask.size()\n    tgt_len = tgt_len if tgt_len is not None else src_len\n\n    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)\n\n    inverted_mask = 1.0 - expanded_mask\n\n    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)\n\n\nclass ParallelLlamaModel(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]\n\n    Args:\n        config: LlamaConfig\n    \"\"\"\n\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config)\n        self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n            num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n        )\n\n        self.layers = nn.ModuleList(\n            [ParallelLlamaDecoderLayer(config, megatron_config) for _ in range(config.num_hidden_layers)]\n        )\n        self.norm = ParallelLlamaRMSNorm(config, megatron_config)\n\n    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask\n    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds):\n        # create causal mask\n        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n        combined_attention_mask = None\n        if input_shape[-1] > 1:\n            combined_attention_mask = _make_causal_mask(\n                input_shape,\n                inputs_embeds.dtype,\n                device=inputs_embeds.device,\n            )\n\n        if attention_mask is not None:\n            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(\n                inputs_embeds.device\n            )\n            combined_attention_mask = (\n                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask\n            )\n\n        return combined_attention_mask\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (batch_size, seq_length)\n            attention_mask: attention_mask. shape (batch_size, seq_length)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        batch_size, seq_length = input_ids.shape\n        inputs_embeds = self.embed_tokens(input_ids)\n        # embed positions\n\n        attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds)\n\n        hidden_states = inputs_embeds\n\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                attention_mask=attention_mask,\n                position_ids=position_ids,\n            )\n\n            hidden_states = layer_outputs\n\n        hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelLlamaForCausalLM(nn.Module):\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.model = ParallelLlamaModel(config, megatron_config=megatron_config)\n        self.vocab_size = config.vocab_size\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n\n        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\n        outputs = self.model(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n        )\n\n        hidden_states = outputs\n        logits = self.lm_head(hidden_states)[0]\n\n        logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)\n\n        logits = logits.float()\n        return CausalLMOutputWithPast(\n            loss=None,\n            logits=logits,\n            past_key_values=None,\n            hidden_states=None,\n            attentions=None,\n        )\n\n\nfrom flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa: F401, E402\n\n\nclass ParallelLlamaModelRmPad(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]\n\n    Args:\n        config: LlamaConfig\n    \"\"\"\n\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        self.megatron_config = megatron_config\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config)\n        self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n            num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n        )\n\n        self.layers = nn.ModuleList(\n            [ParallelLlamaDecoderLayerRmPad(config, megatron_config) for _ in range(config.num_hidden_layers)]\n        )\n        self.norm = ParallelLlamaRMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (1, totol_nnz)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        inputs_embeds = self.embed_tokens(input_ids)  # (1, total_nnz) -> (1, total_nnz, hidden_size)\n\n        # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size)\n        inputs_embeds = inputs_embeds.transpose(0, 1)\n        if self.megatron_config.sequence_parallel:\n            inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds)\n\n        hidden_states = inputs_embeds\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                position_ids=position_ids,\n                sequence_length=sequence_length,\n                indices=indices,\n                cu_seqlens=cu_seqlens,\n                max_seqlen_in_batch=max_seqlen_in_batch,\n            )\n\n            hidden_states = layer_outputs\n\n        hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelLlamaForCausalLMRmPad(nn.Module):\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.megatron_config = megatron_config\n        self.model = ParallelLlamaModelRmPad(config, megatron_config=megatron_config)\n        self.vocab_size = config.vocab_size\n        self._init_head(config)\n\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n    def _forward_head(self, hidden_states):\n        # all_gather from sequence parallel region is performed inside lm_head\n        logits = self.lm_head(hidden_states)[0]\n        logits = logits.float()  # (total_nnz_padded, 1, vocab_size // tp)\n        logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)  # (total_nnz_padded, 1, vocab_size)\n        return logits\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n        batch_size, sequence_length = input_ids.shape\n\n        # remove padding here\n        input_ids, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(\n            input_ids.unsqueeze(dim=-1), attention_mask\n        )  # (total_nnz, 1)\n\n        # pad input_ids to multiple of tp for all tp ranks\n        # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap\n        if self.megatron_config.sequence_parallel:\n            input_ids = sp_utils.pad_to_sequence_parallel(input_ids)\n\n        input_ids = input_ids.transpose(0, 1)  # (1, total_nnz+pad)\n\n        outputs = self.model(\n            input_ids=input_ids,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        hidden_states = outputs\n\n        logits = self._forward_head(hidden_states)\n\n        # remove padding from sequence parallel\n        if self.megatron_config.sequence_parallel:\n            totol_nnz = cu_seqlens[-1]\n            logits = logits[:totol_nnz]  # (total_nnz_padded)\n\n        logits = torch.squeeze(logits, dim=1)  # remove the artificial batch dimension\n        # add removed padding back\n        logits = pad_input(\n            logits, indices, batch_size, seqlen=sequence_length\n        )  # (batch_size, sequence_length, vocab_size)\n\n        return CausalLMOutputWithPast(\n            loss=None,\n            logits=logits,\n            past_key_values=None,\n            hidden_states=None,\n            attentions=None,\n        )\n\n\nclass ParallelLlamaForValueRmPad(ParallelLlamaForCausalLMRmPad):\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = nn.Linear(in_features=config.hidden_size, out_features=1, bias=False)\n        # lm_head is effectively the same as sequence parallel\n        sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)\n\n    def _forward_head(self, hidden_states):\n        logits = self.lm_head(hidden_states)  # (total_nnz_padded // tp, 1, 1)\n        logits = logits.float()\n        if self.megatron_config.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        output = super().forward(input_ids, attention_mask, position_ids)\n        output.logits = torch.squeeze(output.logits, dim=-1)\n        return output\n\n\n\"\"\"\nSupport pipeline parallelism\n\"\"\"\n\n\nclass ParallelLlamaModelRmPadPP(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]\n    This model definition supports pipeline parallelism. To support pp and vpp,\n    - This model only contains layer in this pp stage and vpp chunk\n    - When calling get_model in Megatron, this rank will instantiate all the vpp chunks in this pp.\n    Args:\n        config: LlamaConfig\n    \"\"\"\n\n    def __init__(self, config: LlamaConfig, megatron_config: ModelParallelConfig, pre_process, post_process):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        self.pre_process = pre_process\n        self.post_process = post_process\n        self.megatron_config = megatron_config\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config)\n        if pre_process:\n            self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n                num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n            )\n        else:\n            self.embed_tokens = None\n\n        pp_rank = mpu.get_pipeline_model_parallel_rank()\n        pp_size = megatron_config.pipeline_model_parallel_size\n        self.num_layer_per_pp = config.num_hidden_layers // pp_size\n        vpp_size = megatron_config.virtual_pipeline_model_parallel_size\n        vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank()\n\n        if vpp_size is not None:\n            self.layers = nn.ModuleList()\n            self.num_layer_vpp_chunk = self.num_layer_per_pp // vpp_size\n            self.num_layer_this_model = self.num_layer_vpp_chunk\n            offset = vpp_rank * (config.num_hidden_layers // vpp_size) + (pp_rank * self.num_layer_vpp_chunk)\n        else:\n            self.num_layer_this_model = self.num_layer_per_pp\n            offset = pp_rank * self.num_layer_per_pp\n\n        self.layers = nn.ModuleList()\n        for i in range(self.num_layer_this_model):\n            layer = ParallelLlamaDecoderLayerRmPad(config, megatron_config, layer_idx=offset + i)\n            self.layers.add_module(f\"{i}\", layer)\n\n        if post_process:\n            self.norm = ParallelLlamaRMSNorm(config, megatron_config)\n        else:\n            self.norm = None\n\n    def set_input_tensor(self, input_tensor):\n        \"\"\"Set input tensor to be used instead of forward()'s input.\n\n        When doing pipeline parallelism the input from the previous\n        stage comes from communication, not from the input, so the\n        model's forward_step_func won't have it. This function is thus\n        used by internal code to bypass the input provided by the\n        forward_step_func\"\"\"\n        self.input_tensor = input_tensor\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (1, totol_nnz)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        if self.pre_process:\n            inputs_embeds = self.embed_tokens(input_ids)  # (1, total_nnz) -> (1, total_nnz, hidden_size)\n\n            # vocab parallel embedding will not do sequence parallel reduce-scatter in open source megatron\n            # so need to deal with it by handle here:\n            # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size)\n            inputs_embeds = inputs_embeds.transpose(0, 1)\n            if self.megatron_config.sequence_parallel:\n                inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds)\n\n            hidden_states = inputs_embeds\n        else:\n            # self.hidden_states should be passed by Megatron\n            hidden_states = self.input_tensor\n\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                position_ids=position_ids,\n                sequence_length=sequence_length,\n                indices=indices,\n                cu_seqlens=cu_seqlens,\n                max_seqlen_in_batch=max_seqlen_in_batch,\n            )\n\n            hidden_states = layer_outputs\n\n        if self.post_process:\n            hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelLlamaForCausalLMRmPadPP(nn.Module):\n    def __init__(\n        self,\n        config: LlamaConfig,\n        megatron_config: ModelParallelConfig,\n        pre_process,\n        post_process,\n        share_embeddings_and_output_weights=False,\n    ):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.megatron_config = megatron_config\n        self.model = ParallelLlamaModelRmPadPP(\n            config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process\n        )\n        assert share_embeddings_and_output_weights is False, (\n            \"Llama Model not supports sharing embedding and output weights\"\n        )\n        self.share_embeddings_and_output_weights = share_embeddings_and_output_weights\n        self.vocab_size = config.vocab_size\n        self.pre_process = pre_process\n        self.post_process = post_process\n        if post_process:\n            self._init_head(config)\n\n    def set_input_tensor(self, input_tensor):\n        \"\"\"Set input tensor to be used instead of forward()'s input.\n\n        When doing pipeline parallelism the input from the previous\n        stage comes from communication, not from the input, so the\n        model's forward_step_func won't have it. This function is thus\n        used by internal code to bypass the input provided by the\n        forward_step_func\"\"\"\n        assert len(input_tensor) == 1\n        self.model.set_input_tensor(input_tensor[0])\n\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n    def _forward_head(self, hidden_states):\n        # all_gather from sequence parallel region is performed inside lm_head\n        # logits shape before forward_head hidden_states.shape: [4, 32, 4096]\n        logits = self.lm_head(hidden_states)[0]\n        # logits shape after forward_head logits.shape: [8, 32, 8]\n        logits = logits.float()  # (total_nnz_padded, 1, vocab_size // tp)\n        return logits\n\n    def forward(\n        self,\n        # original input\n        *,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n\n        # Note that input_ids, attention_mask and position_ids should be passed to every pp layer.\n        # In the first pp, input_ids will be used, in other pp layers hidden_states will be used inside self.model\n        batch_size, sequence_length = input_ids.shape\n        # remove padding here\n        input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(\n            input_ids.unsqueeze(dim=-1), attention_mask\n        )  # (total_nnz, 1)\n\n        # pad input_ids to multiple of tp for all tp ranks\n        # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap\n        if self.megatron_config.sequence_parallel:\n            input_ids_rmpad = sp_utils.pad_to_sequence_parallel(input_ids_rmpad)\n\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz+pad)\n\n        outputs = self.model(\n            input_ids=input_ids_rmpad,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        if self.post_process:\n            hidden_states = outputs\n            # print(f'hidden_states.shape = {hidden_states.shape}') # torch.Size([4, 32, 4096])\n            logits = self._forward_head(hidden_states)\n            logits = torch.squeeze(logits, dim=1)  # remove the artificial batch dimension # torch.Size([8, 32, 16])\n\n            # remove padding from sequence parallel\n            if self.megatron_config.sequence_parallel:\n                totol_nnz = cu_seqlens[-1]\n                logits = logits[:totol_nnz]  # (total_nnz_padded)\n            # add removed padding back. If input is already rmpad, we let the caller pad_input\n            logits = pad_input(\n                logits, indices, batch_size, seqlen=sequence_length\n            )  # (batch_size, sequence_length, vocab_size)\n\n            return CausalLMOutputWithPast(\n                loss=None,\n                logits=logits,\n                past_key_values=None,\n                hidden_states=None,\n                attentions=None,\n            )\n        else:\n            return outputs\n\n\nclass ParallelLlamaForValueRmPadPP(ParallelLlamaForCausalLMRmPadPP):\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = nn.Linear(in_features=config.hidden_size, out_features=1, bias=False)\n        # lm_head is effectively the same as sequence parallel\n        sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)\n\n    def _forward_head(self, hidden_states):\n        logits = self.lm_head(hidden_states)  # (total_nnz_padded // tp, 1, 1)\n        logits = logits.float()\n        if self.megatron_config.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits\n\n    def forward(\n        self,\n        *,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids)\n        if self.post_process:\n            output.logits = torch.squeeze(output.logits, dim=-1)\n            return output\n        else:\n            return output\n"
  },
  {
    "path": "verl/models/mcore/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom verl.models.mcore.patch import apply_patch_megatron_v012_with_torch_v28\n\nfrom .registry import (\n    get_mcore_forward_fn,\n    get_mcore_forward_fused_fn,\n    get_mcore_forward_fused_no_padding_fn,\n    get_mcore_forward_no_padding_fn,\n    get_mcore_weight_converter,\n    hf_to_mcore_config,\n    init_mcore_model,\n)\n\n__all__ = [\n    \"hf_to_mcore_config\",\n    \"init_mcore_model\",\n    \"get_mcore_forward_fn\",\n    \"get_mcore_weight_converter\",\n    \"get_mcore_forward_fused_fn\",\n    \"get_mcore_forward_fused_no_padding_fn\",\n    \"get_mcore_forward_no_padding_fn\",\n]\n\napply_patch_megatron_v012_with_torch_v28()\n"
  },
  {
    "path": "verl/models/mcore/bridge.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\ntry:\n    from megatron.bridge import AutoBridge\n    from megatron.bridge.models.conversion.param_mapping import AutoMapping\n    from megatron.bridge.peft.canonical_lora import CanonicalLoRA\n    from megatron.bridge.peft.dora import DoRA\n    from megatron.bridge.peft.lora import LoRA, VLMLoRA\nexcept ImportError:\n    # `pip install verl[mcore]` or\n    print(\"Megatron-Bridge package not found. Please install Megatron-Bridge with `pip install megatron-bridge`\")\n    raise\n\nimport torch\nfrom megatron.core import tensor_parallel\n\n\ndef _ensure_model_list(model):\n    return model if isinstance(model, list) else [model]\n\n\nclass LinearForLastLayer(torch.nn.Linear):\n    \"\"\"\n    A custom linear layer implementation for the last layer of a model.\n\n    This layer extends PyTorch's Linear module with functionality specifically designed\n    for handling the final layer in transformer models with sequence parallelism.\n\n    Attributes:\n        sequence_parallel: Boolean indicating whether sequence parallelism is enabled\n    \"\"\"\n\n    def __init__(\n        self,\n        input_size,\n        output_size,\n        *,\n        sequence_parallel: bool,\n    ):\n        \"\"\"\n        Initializes the LinearForLastLayer.\n\n        Args:\n            input_size: The size of the input features\n            output_size: The size of the output features\n            sequence_parallel (bool): Whether sequence parallelism is enabled\n        \"\"\"\n        super().__init__(in_features=input_size, out_features=output_size, bias=False)\n        self.sequence_parallel = sequence_parallel\n        if self.sequence_parallel:\n            self.weight.sequence_parallel = True\n\n    def forward(\n        self,\n        input_,\n        weight=None,\n        runtime_gather_output=None,\n    ):\n        \"\"\"\n        Forward pass for the linear layer.\n\n        This method computes the linear transformation and handles sequence parallelism\n        if enabled, gathering outputs from different sequence parallel regions.\n\n        Args:\n            input_: Input tensor\n            weight: Placeholder for compatibility\n            runtime_gather_output: Placeholder for compatibility\n\n        Returns:\n            tuple: (logits, None) where logits is the output of the linear transformation\n        \"\"\"\n        logits = super().forward(input_)\n        logits = logits.float()\n        if self.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits, None\n\n\n# Make Megatron-Bridge AutoMapping treats the custom last layer as replicated.\nAutoMapping.register_module_type(\"LinearForLastLayer\", \"replicated\")\n\n\ndef make_value_model(hidden_size, sequence_parallel):\n    \"\"\"Creates a pre-wrap hook that replace the output layer with a value head.\n\n    Args:\n        hidden_size (int): The hidden size of the model's transformer layers.\n        sequence_parallel (bool): Whether sequence parallelism is enabled.\n\n    Returns:\n        A hook function that can be used as a `pre_wrap_hook` in Megatron-Bridge.\n        The hook itself takes the model as input and prepares it for value head activation.\n    \"\"\"\n\n    from megatron.core import parallel_state\n\n    def hook(model):\n        model_post_process = []\n        if (\n            parallel_state.get_pipeline_model_parallel_world_size() > 1\n            and parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None\n        ):\n            for i in range(parallel_state.get_virtual_pipeline_model_parallel_world_size()):\n                model_post_process.append(parallel_state.is_pipeline_last_stage(ignore_virtual=False, vp_stage=i))\n        else:\n            model_post_process.append(parallel_state.is_pipeline_last_stage())\n\n        model_list = _ensure_model_list(model)\n        assert len(model_post_process) == len(model_list), \"Model list length and post process list length must match.\"\n\n        for index, model_chunk in enumerate(model_list):\n            if not model_post_process[index]:\n                continue\n\n            model_chunk.output_layer = LinearForLastLayer(\n                input_size=hidden_size,\n                output_size=1,\n                sequence_parallel=sequence_parallel,\n            )\n\n    return hook\n\n\ndef freeze_moe_router(model):\n    \"\"\"Pre-wrap hook to freeze MoE router parameters.\n\n    Args:\n        model: List of MegatronModule instances or single module\n\n    Returns:\n        The model with frozen router parameters\n    \"\"\"\n    for model_chunk in _ensure_model_list(model):\n        if hasattr(model_chunk, \"decoder\") and hasattr(model_chunk.decoder, \"layers\"):\n            for layer in model_chunk.decoder.layers:\n                if hasattr(layer.mlp, \"router\"):\n                    if hasattr(layer.mlp.router, \"weight\"):\n                        layer.mlp.router.weight.requires_grad = False\n                    if hasattr(layer.mlp.router, \"bias\"):\n                        layer.mlp.router.bias.requires_grad = False\n                if hasattr(layer.mlp, \"shared_experts\"):\n                    if (\n                        hasattr(layer.mlp.shared_experts, \"gate_weight\")\n                        and layer.mlp.shared_experts.gate_weight is not None\n                    ):\n                        layer.mlp.shared_experts.gate_weight.requires_grad = False\n                    if (\n                        hasattr(layer.mlp.shared_experts, \"gate_bias\")\n                        and layer.mlp.shared_experts.gate_bias is not None\n                    ):\n                        layer.mlp.shared_experts.gate_bias.requires_grad = False\n\n    return model\n\n\n__all__ = [\n    \"AutoBridge\",\n    \"make_value_model\",\n    \"freeze_moe_router\",\n    \"LoRA\",\n    \"VLMLoRA\",\n    \"DoRA\",\n    \"CanonicalLoRA\",\n]\n"
  },
  {
    "path": "verl/models/mcore/config_converter.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# convert huggingface config to mcore transformer config\n\n\nimport warnings\nfrom typing import TypeVar\n\nimport torch\nimport torch.nn.functional as F\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.transformer import MLATransformerConfig, TransformerConfig\nfrom transformers import PretrainedConfig\n\nT = TypeVar(\"T\", bound=TransformerConfig)\n\n\ndef _get_base_transformer_config(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> dict:\n    \"\"\"\n    Create a base TransformerConfig with common parameters across different model architectures.\n    TODO: (ycl) use dataclass or converter config?\n\n    Args:\n        hf_config: HuggingFace model configuration\n        dtype: Data type for the model\n        override_transformer_config_kwargs: Additional parameters to override defaults\n\n    Returns:\n        TransformerConfig with common parameters\n    \"\"\"\n\n    # Common parallel state parameters\n    overlap_p2p_comm = (\n        mpu.get_virtual_pipeline_model_parallel_world_size() is not None\n        and mpu.get_virtual_pipeline_model_parallel_world_size() > 1\n    )\n    batch_p2p_comm = False\n\n    # Base configuration with common parameters\n    base_config = {\n        # Model architecture parameters\n        \"num_layers\": hf_config.num_hidden_layers,\n        \"hidden_size\": hf_config.hidden_size,\n        \"num_attention_heads\": hf_config.num_attention_heads,\n        \"num_query_groups\": hf_config.num_key_value_heads,\n        \"ffn_hidden_size\": hf_config.intermediate_size,\n        \"attention_dropout\": hf_config.attention_dropout,\n        \"hidden_dropout\": getattr(hf_config, \"hidden_dropout\", 0.0),\n        \"kv_channels\": getattr(hf_config, \"head_dim\", None),\n        \"layernorm_epsilon\": hf_config.rms_norm_eps,\n        \"add_bias_linear\": True,\n        # Activation and normalization\n        \"activation_func\": F.silu,\n        \"normalization\": \"RMSNorm\",\n        \"gated_linear_unit\": True,\n        # Data types\n        \"pipeline_dtype\": dtype,\n        \"params_dtype\": dtype,\n        \"bf16\": dtype is torch.bfloat16,\n        # Parallel configuration\n        \"tensor_model_parallel_size\": mpu.get_tensor_model_parallel_world_size(),\n        \"pipeline_model_parallel_size\": mpu.get_pipeline_model_parallel_world_size(),\n        \"expert_model_parallel_size\": mpu.get_expert_model_parallel_world_size(),\n        \"expert_tensor_parallel_size\": mpu.get_expert_tensor_parallel_world_size(),\n        \"virtual_pipeline_model_parallel_size\": mpu.get_virtual_pipeline_model_parallel_world_size(),\n        \"context_parallel_size\": mpu.get_context_parallel_world_size(),\n        \"overlap_p2p_comm\": overlap_p2p_comm,\n        \"batch_p2p_comm\": batch_p2p_comm,\n        \"sequence_parallel\": mpu.get_tensor_model_parallel_world_size() > 1,\n        # Common settings\n        \"variable_seq_lengths\": True,\n        \"masked_softmax_fusion\": True,\n        \"moe_token_dispatcher_type\": \"alltoall\",\n    }\n\n    # Update with any provided overrides\n    # override_transformer_config_kwargs as kwargs shall never be none\n    base_config.update(override_transformer_config_kwargs)\n\n    return base_config\n\n\ndef _get_mla_transformer_config(\n    hf_config: PretrainedConfig, mla_rope_config: dict, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> dict:\n    \"\"\"\n    Create a MLATransformerConfig with common parameters across different model architectures.\n    This is specifically for MLA models like DeepseekV3.\n\n    Args:\n        hf_config: HuggingFace model configuration\n        mla_rope_config: MLA specific RoPE configuration\n        dtype: Data type for the model\n        override_transformer_config_kwargs: Additional parameters to override defaults\n\n    Returns:\n        MLATransformerConfig with common parameters\n    \"\"\"\n    base_config = _get_base_transformer_config(hf_config=hf_config, dtype=dtype, **override_transformer_config_kwargs)\n    mla_config = {\n        # MLA specific parameters\n        \"q_lora_rank\": hf_config.q_lora_rank,\n        \"kv_lora_rank\": hf_config.kv_lora_rank,\n        \"qk_head_dim\": hf_config.qk_nope_head_dim,\n        \"qk_pos_emb_head_dim\": hf_config.qk_rope_head_dim,\n        \"v_head_dim\": hf_config.v_head_dim,\n        \"rotary_base\": hf_config.rope_theta,\n        \"rotary_scaling_factor\": mla_rope_config[\"factor\"],\n        \"rope_type\": mla_rope_config[\"type\"],\n        \"max_position_embeddings\": mla_rope_config[\"original_max_position_embeddings\"],\n        \"beta_fast\": mla_rope_config[\"beta_fast\"],\n        \"beta_slow\": mla_rope_config[\"beta_slow\"],\n        \"mscale\": mla_rope_config[\"mscale\"],\n        \"mscale_all_dim\": mla_rope_config[\"mscale_all_dim\"],\n    }\n\n    base_config.update(mla_config)\n    return base_config\n\n\ndef check_and_construct_configs(original_config: dict, cls: type[T]) -> T:\n    \"\"\"\n    Check and disable incompatible configurations for older Megatron version.\n\n    Args:\n        original_config (dict): The original model configuration.\n\n    Returns:\n        dict: The updated model configuration with incompatible settings disabled.\n    \"\"\"\n    removed_keys = []\n    for key in original_config.keys():\n        if not hasattr(cls, key):\n            removed_keys.append(key)\n    if removed_keys:\n        warnings.warn(\n            f\"The following keys are not supported in the current Megatron version and will be removed: {removed_keys}\",\n            stacklevel=2,\n        )\n        for key in removed_keys:\n            original_config.pop(key)\n\n    original_config = mapping_string_to_attn_backend(original_config)\n    if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:\n        print(f\"Overridden {cls.__name__} init config: {original_config}\")\n    return cls(**original_config)\n\n\ndef hf_to_mcore_config_dense(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    # for LlamaForCausalLM or Qwen2ForCausalLM\n    qkv_bias = True if \"Qwen2\" in hf_config.architectures[0] else getattr(hf_config, \"attention_bias\", False)\n    qk_layernorm = True if \"Qwen3\" in hf_config.architectures[0] else False\n\n    args: dict = _get_base_transformer_config(\n        hf_config=hf_config,\n        dtype=dtype,\n        use_cpu_initialization=False,\n        add_bias_linear=False,\n        add_qkv_bias=qkv_bias,\n        qk_layernorm=qk_layernorm,\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    return check_and_construct_configs(args, TransformerConfig)\n\n\ndef hf_to_mcore_config_qwen2moe(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    args: dict = _get_base_transformer_config(\n        hf_config=hf_config,\n        dtype=dtype,\n        use_cpu_initialization=False,\n        add_bias_linear=False,\n        layernorm_epsilon=hf_config.rms_norm_eps,\n        # MoE specific\n        moe_ffn_hidden_size=hf_config.moe_intermediate_size,\n        moe_router_bias_update_rate=0.001,\n        moe_router_topk=hf_config.num_experts_per_tok,\n        num_moe_experts=hf_config.num_experts,\n        moe_shared_expert_intermediate_size=hf_config.shared_expert_intermediate_size,\n        moe_aux_loss_coeff=hf_config.router_aux_loss_coef,\n        # moe_aux_loss_coeff=0.0,\n        moe_router_load_balancing_type=\"none\",  # turn off aux_loss as it hurts perf in RL\n        moe_shared_expert_overlap=True,\n        moe_grouped_gemm=True,\n        moe_router_score_function=\"softmax\",\n        # Other optimizations\n        persist_layer_norm=True,\n        bias_activation_fusion=True,\n        bias_dropout_fusion=True,\n        # Qwen specific\n        moe_router_pre_softmax=True,\n        add_qkv_bias=True,\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    return check_and_construct_configs(args, TransformerConfig)\n\n\ndef hf_to_mcore_config_mixtral(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    args: dict = _get_base_transformer_config(\n        hf_config=hf_config,\n        dtype=dtype,\n        use_cpu_initialization=False,\n        add_bias_linear=False,\n        layernorm_epsilon=hf_config.rms_norm_eps,\n        # MoE specific\n        num_moe_experts=hf_config.num_local_experts,\n        moe_aux_loss_coeff=hf_config.router_aux_loss_coef,\n        moe_router_topk=hf_config.num_experts_per_tok,\n        moe_router_pre_softmax=True,\n        moe_router_load_balancing_type=\"none\",  # turn off aux_loss as it hurts perf in RL\n        moe_router_score_function=\"softmax\",\n        moe_shared_expert_intermediate_size=None,  # mixtral has no shared expert\n        moe_shared_expert_overlap=False,  # mixtral has no shared expert\n        moe_ffn_hidden_size=hf_config.intermediate_size,\n        moe_router_bias_update_rate=0.001,\n        # moe_permute_fusion=True, # need TE 2.1+\n        moe_grouped_gemm=True,\n        # Other optimizations\n        persist_layer_norm=True,\n        apply_rope_fusion=True,\n        bias_activation_fusion=True,\n        bias_dropout_fusion=True,\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    return check_and_construct_configs(args, TransformerConfig)\n\n\ndef hf_to_mcore_config_qwen3moe(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    args: dict = _get_base_transformer_config(\n        hf_config=hf_config,\n        dtype=dtype,\n        use_cpu_initialization=False,\n        add_bias_linear=False,\n        layernorm_epsilon=hf_config.rms_norm_eps,\n        # MoE specific\n        moe_ffn_hidden_size=hf_config.moe_intermediate_size,\n        moe_router_bias_update_rate=0.001,\n        moe_router_topk=hf_config.num_experts_per_tok,\n        num_moe_experts=hf_config.num_experts,\n        moe_aux_loss_coeff=hf_config.router_aux_loss_coef,\n        # moe_aux_loss_coeff=0.0,\n        moe_router_load_balancing_type=\"none\",  # turn off aux_loss as it hurts perf in RL\n        moe_grouped_gemm=True,\n        moe_router_score_function=\"softmax\",\n        # Other optimizations\n        persist_layer_norm=True,\n        bias_activation_fusion=True,\n        bias_dropout_fusion=True,\n        # Qwen specific\n        moe_router_pre_softmax=False,\n        qk_layernorm=True,\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    return check_and_construct_configs(args, TransformerConfig)\n\n\ndef hf_to_mcore_config_dpskv3(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> MLATransformerConfig:\n    # DeepseekV3ForCausalLM\n    from megatron.core.config import set_experimental_flag\n    from megatron.core.transformer.enums import AttnBackend\n\n    set_experimental_flag(True)\n\n    from .patch import apply_patch\n\n    apply_patch()\n\n    mla_rope_config = {\n        \"beta_fast\": 32,\n        \"beta_slow\": 1,\n        \"factor\": 1,\n        \"mscale\": 1.0,\n        \"mscale_all_dim\": 1.0,\n        \"original_max_position_embeddings\": 4096,\n        \"type\": \"rope\",\n    }\n    if \"rope_scaling\" in hf_config and hf_config.rope_scaling is not None:\n        mla_rope_config.update(hf_config.rope_scaling)\n    moe_layer_freq = [1] * hf_config.num_hidden_layers\n    for i in range(min(hf_config.first_k_dense_replace, hf_config.num_hidden_layers)):\n        moe_layer_freq[i] = 0\n\n    # disable MTP and quantization for now\n    if \"num_nextn_predict_layers\" in hf_config:\n        assert hf_config.num_nextn_predict_layers == 0, (\n            \"MTP is not supported for now, please modify the config.json to set num_nextn_predict_layers to 0\"\n        )\n    assert \"quantization_config\" not in hf_config or not hf_config.quantization_config, (\n        \"quantization is not supported for now, please modify the config.json to remove quantization_config\"\n    )\n\n    args: dict = _get_mla_transformer_config(\n        hf_config=hf_config,\n        mla_rope_config=mla_rope_config,\n        dtype=dtype,\n        # Additional parameters\n        use_cpu_initialization=False,\n        add_bias_linear=False,\n        attention_backend=AttnBackend.fused,\n        qk_layernorm=True,\n        # Standard MoE parameters\n        moe_ffn_hidden_size=hf_config.moe_intermediate_size,\n        moe_token_dispatcher_type=\"alltoall\",\n        moe_router_bias_update_rate=0.001,\n        moe_router_enable_expert_bias=True,\n        moe_router_topk=hf_config.num_experts_per_tok,\n        num_moe_experts=hf_config.n_routed_experts,\n        moe_shared_expert_intermediate_size=hf_config.moe_intermediate_size * hf_config.n_shared_experts,\n        moe_aux_loss_coeff=getattr(hf_config, \"aux_loss_alpha\", 0.001),\n        moe_router_load_balancing_type=\"seq_aux_loss\",\n        moe_shared_expert_overlap=True,\n        # moe_permute_fusion=True, # need TE 2.1+\n        moe_grouped_gemm=True,\n        moe_router_score_function=\"sigmoid\",\n        moe_router_pre_softmax=True,\n        moe_router_topk_scaling_factor=hf_config.routed_scaling_factor,\n        moe_layer_freq=moe_layer_freq,\n        # mcore 0.12 moe\n        moe_router_dtype=\"fp64\",\n        disable_bf16_reduced_precision_matmul=True,\n        # Other optimizations\n        # deallocate_pipeline_outputs=True,\n        # gradient_accumulation_fusion=True,\n        persist_layer_norm=True,\n        bias_activation_fusion=True,\n        bias_dropout_fusion=True,\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    transformer_config = check_and_construct_configs(args, MLATransformerConfig)\n    # MTP\n    if \"num_nextn_predict_layers\" in hf_config:\n        transformer_config.mtp_num_layers = hf_config.num_nextn_predict_layers\n        transformer_config.mtp_loss_scaling_factor = 0.1\n\n    return transformer_config\n\n\ndef hf_to_mcore_config_qwen2_5_vl(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    # Qwen2_5_VLForConditionalGeneration\n\n    args = _get_base_transformer_config(\n        hf_config=hf_config,\n        dtype=dtype,\n        add_bias_linear=False,\n        # qwen specific\n        add_qkv_bias=True,\n        mrope_section=hf_config.rope_scaling[\"mrope_section\"],\n    )\n    # override_transformer_config_kwargs as kwargs shall never be none\n    args.update(override_transformer_config_kwargs)\n    args = mapping_string_to_attn_backend(args)\n    return TransformerConfig(**args)\n\n\ndef hf_to_mcore_config_llama4(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    # Llama4ForConditionalGeneration\n    raise NotImplementedError(\"Llama4ForConditionalGeneration is not supported yet\")\n\n\ndef mapping_string_to_attn_backend(args: dict) -> dict:\n    if \"attention_backend\" in args and isinstance(args[\"attention_backend\"], str):\n        from megatron.core.transformer.enums import AttnBackend\n\n        args[\"attention_backend\"] = AttnBackend[args[\"attention_backend\"]]\n    return args\n"
  },
  {
    "path": "verl/models/mcore/loader.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\nfrom .saver import _megatron_calc_global_rank\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef load_state_dict_to_megatron_gptmodel(state_dict, wrapped_models, config, params_dtype, is_value_model=False):\n    \"\"\"Load merged state_dict to sharded Megatron module in training.\"\"\"\n    from megatron.core import DistributedDataParallel as LocalDDP\n    from megatron.core import mpu\n    from megatron.core.transformer.module import Float16Module\n    from torch.nn.parallel import DistributedDataParallel as torchDDP\n\n    from verl.utils.logger import print_rank_0\n    from verl.utils.megatron_utils import unwrap_model\n\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    def broadcast_params(module):\n        for param in module.parameters():\n            torch.distributed.broadcast(\n                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()\n            )\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    cp_rank = mpu.get_context_parallel_rank()\n    src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=cp_rank)\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if torch.distributed.get_rank() == src_rank:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        gpt_model_module = _get_gpt_model(models[i])\n        assert len(gpt_model_module.decoder.layers) == num_layers_per_model\n\n    def _broadcast_tensor(tensor, name) -> torch.Tensor:\n        \"\"\"broadcast tensor from rank0 across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        if torch.distributed.get_rank() == src_rank:\n            if name in state_dict:\n                weight = state_dict[name]\n                tensor_shape = weight.shape\n            else:\n                tensor_shape = None\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not in state_dict, skip load\")\n            return\n\n        if tensor is None:\n            tensor = torch.empty(\n                tensor_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        if torch.distributed.get_rank() == src_rank:\n            tensor.data.copy_(weight)\n        dist.broadcast(tensor, src=src_rank, group=mp_group)\n\n    def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == src_rank:\n            if name in state_dict:\n                full_weight = state_dict[name]\n\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == src_rank:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=src_rank, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == src_rank:\n            if name in state_dict:\n                full_weight = state_dict[name]\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == src_rank:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=src_rank, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == src_rank:\n            gate_weight = state_dict[gate_name]\n            up_weight = state_dict[up_name]\n            new_gate_up_weight = torch.empty(\n                config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                intermediate_size_tp = config.intermediate_size // tp_size\n                gate_weight_tp = gate_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                up_weight_tp = up_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                new_gate_up_weight[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)].copy_(\n                    torch.cat([gate_weight_tp, up_weight_tp], dim=0)\n                )\n\n            tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank() == src_rank:} tensor {gate_name, up_name} shape \"\n                f\"{tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == src_rank:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=src_rank, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == src_rank:\n            assert q_name in state_dict and k_name in state_dict and v_name in state_dict\n            full_weight_q = state_dict[q_name]\n            full_weight_k = state_dict[k_name]\n            full_weight_v = state_dict[v_name]\n\n            hidden_size_per_head = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                sizes = [total_size * tp_size]\n                if not bias:\n                    sizes.append(config.hidden_size)\n                new_weight_qkv = torch.empty(*sizes, dtype=params_dtype, device=get_device_id())\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    k_part = full_weight_k[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    v_part = full_weight_v[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    num_query_groups_per_partition = models[0].config.num_query_groups // tp_size\n                    new_weight_qkv_this_tp = new_weight_qkv[i * total_size : (i + 1) * total_size]\n                    q_part_per_head = torch.chunk(q_part, num_query_groups_per_partition, dim=0)\n                    k_part_per_head = torch.chunk(k_part, num_query_groups_per_partition, dim=0)\n                    v_part_per_head = torch.chunk(v_part, num_query_groups_per_partition, dim=0)\n                    total_size_per_head = total_size // num_query_groups_per_partition\n                    for j in range(num_query_groups_per_partition):\n                        new_weight_qkv_this_tp[j * total_size_per_head : (j + 1) * total_size_per_head].copy_(\n                            torch.cat([q_part_per_head[j], k_part_per_head[j], v_part_per_head[j]], dim=0)\n                        )\n\n            else:\n                q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                sizes = [total_size * tp_size]\n                if not bias:\n                    sizes.append(config.hidden_size)\n                new_weight_qkv = torch.empty(*sizes, dtype=params_dtype, device=get_device_id())\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head\n                    end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head\n                    k_part = full_weight_k[start_idx:end_idx]\n                    v_part = full_weight_v[start_idx:end_idx]\n                    new_weight_qkv_this_tp = new_weight_qkv[i * total_size : (i + 1) * total_size]\n                    q_part_per_head = torch.chunk(q_part, config.num_attention_heads, dim=0)\n                    k_part_per_head = torch.chunk(k_part, config.num_attention_heads, dim=0)\n                    v_part_per_head = torch.chunk(v_part, config.num_attention_heads, dim=0)\n                    total_size_per_head = total_size // config.num_attention_heads\n                    for j in range(config.num_attention_heads):\n                        new_weight_qkv_this_tp[j * total_size_per_head : (j + 1) * total_size_per_head].copy_(\n                            torch.cat([q_part_per_head[j], k_part_per_head[j], v_part_per_head[j]], dim=0)\n                        )\n\n            tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name, k_name, v_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == src_rank:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=src_rank, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    if dp_rank == 0:\n        # Embeddings\n        # -------------------\n        print_rank_0(\"loading embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        embed_tokens_weight = None\n        if pp_rank == 0:\n            embed_tokens_weight = gpt_model_module.embedding.word_embeddings.weight\n        _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, \"model.embed_tokens.weight\")\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n\n        for layer in range(config.num_hidden_layers):\n            layer_name = f\"model.layers.{layer}\"\n            print_rank_0(f\"loading layer #{layer}, with layer_name model.layers.{layer}...\")\n            dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank])\n            sync_layer = gpt_model_module.decoder.layers[dst_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.self_attention.linear_qkv.layer_norm_weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.input_layernorm.weight\",\n            )\n\n            if f\"{layer_name}.self_attn.q_norm.weight\" in state_dict:\n                _broadcast_tensor(\n                    sync_layer.self_attention.q_layernorm.weight if dst_pp_rank == pp_rank else None,\n                    f\"{layer_name}.self_attn.q_norm.weight\",\n                )\n                _broadcast_tensor(\n                    sync_layer.self_attention.k_layernorm.weight if dst_pp_rank == pp_rank else None,\n                    f\"{layer_name}.self_attn.k_norm.weight\",\n                )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attention.linear_qkv.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n            )\n            if f\"{layer_name}.self_attn.q_proj.bias\" in state_dict:\n                _broadcast_tp_shard_tensor_qkv(\n                    sync_layer.self_attention.linear_qkv.bias if dst_pp_rank == pp_rank else None,\n                    f\"{layer_name}.self_attn.q_proj.bias\",\n                    f\"{layer_name}.self_attn.k_proj.bias\",\n                    f\"{layer_name}.self_attn.v_proj.bias\",\n                    bias=True,\n                )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attention.linear_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                chunk_dim=1,\n            )\n            _broadcast_tensor(\n                sync_layer.mlp.linear_fc1.layer_norm_weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.linear_fc1.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.linear_fc2.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                chunk_dim=1,\n            )\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"loading final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.decoder.final_layernorm, \"weight\", None),\n            \"model.norm.weight\",\n        )\n\n        print_rank_0(\"loading lm_head...\")\n        lm_head_weight = None\n        if pp_rank + 1 == pp_size:\n            lm_head_weight = gpt_model_module.output_layer.weight\n\n        if is_value_model:\n            # if torch.distributed.get_rank() == src_rank:\n            if \"lm_head.weight\" in state_dict and state_dict[\"lm_head.weight\"].shape[0] == 1:\n                _broadcast_tensor(lm_head_weight, \"lm_head.weight\")\n            elif \"reward_head.weight\" in state_dict and state_dict[\"reward_head.weight\"].shape[0] == 1:\n                _broadcast_tensor(lm_head_weight, \"reward_head.weight\")\n                print_rank_0(\"load lm_head from value_head weight\")\n            elif \"score.weight\" in state_dict and state_dict[\"score.weight\"].shape[0] == 1:\n                _broadcast_tensor(lm_head_weight, \"score.weight\")\n                print_rank_0(\"load lm_head from score weight\")\n            else:\n                _broadcast_tensor(None, \"lm_head.weight\")\n                print_rank_0(\"fail to match lm_head in value_model\")\n            # else:\n\n            #     _broadcast_tensor(lm_head_weight, \"lm_head.weight\")\n\n        else:\n            _broadcast_tp_shard_tensor(lm_head_weight, \"lm_head.weight\")\n    dist.barrier()\n    # Broadcast weights inside data parallel groups\n    for wrapped_model in wrapped_models:\n        broadcast_params(wrapped_model)\n    pass\n    get_torch_device().empty_cache()\n    print_rank_0(f\"loading megatron ckpt done, time elapsed {time.time() - start_time}s\")\n"
  },
  {
    "path": "verl/models/mcore/mbridge.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# VANILLA_MBRIDGE\ntry:\n    from verl.models.mcore.patch import apply_patch_mbridge\n\n    apply_patch_mbridge()\n    from mbridge import AutoBridge\n    from mbridge.utils.post_creation_callbacks import freeze_moe_router, make_value_model\nexcept ImportError:\n    print(\"mbridge package not found. Please install mbridge with `pip install verl[mcore]` or `pip install mbridge`\")\n    raise\n\n__all__ = [\"AutoBridge\", \"make_value_model\", \"freeze_moe_router\"]\n"
  },
  {
    "path": "verl/models/mcore/model_forward.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom torch.nested._internal.nested_tensor import NestedTensor\n\nfrom verl.utils.megatron_utils import unwrap_model\nfrom verl.workers.config import MtpConfig\n\nfrom .util import (\n    postprocess_bshd,\n    postprocess_bshd_no_padding,\n    postprocess_packed_seqs,\n    postprocess_thd_no_padding,\n    preprocess_bshd,\n    preprocess_bshd_no_padding,\n    preprocess_packed_seqs,\n    preprocess_thd_no_padding,\n)\n\n\ndef model_forward_gen(vision_model: bool = False):\n    def model_forward(\n        model,\n        input_ids,\n        attention_mask,\n        position_ids,\n        multi_modal_inputs: dict,\n        logits_processor=None,\n        logits_processor_args: dict = None,\n        value_model=False,\n        data_format: str = \"thd\",\n        mtp_config: MtpConfig = None,\n    ):\n        \"\"\"Forward pass for models with sequence packing.\"\"\"\n        assert data_format in [\"thd\", \"bshd\"], \"data_format must be 'thd' or 'bshd'\"\n        pre_process = (\n            unwrap_model(model).pre_process if not vision_model else False\n        )  # vision model does not need pre_process, because we pack the input_ids to thd in the forward function\n        post_process = unwrap_model(model).post_process\n        sp = unwrap_model(model).config.sequence_parallel\n        fp8 = unwrap_model(model).config.fp8\n        use_fp8_padding = fp8 in [\"e4m3\", \"hybrid\"]\n\n        model_kwargs = {}\n        if \"pixel_values\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values\"] = multi_modal_inputs[\"pixel_values\"].to(input_ids.device)\n        if \"image_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"image_grid_thw\"] = multi_modal_inputs[\"image_grid_thw\"].to(input_ids.device)\n        if \"pixel_values_videos\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values_videos\"] = multi_modal_inputs[\"pixel_values_videos\"].to(input_ids.device)\n        if \"video_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"video_grid_thw\"] = multi_modal_inputs[\"video_grid_thw\"].to(input_ids.device)\n\n        batch_size, seq_len = attention_mask.shape[:2]\n        mtp_enable_train = mtp_config and mtp_config.enable_train\n\n        if data_format == \"thd\":\n            input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(\n                input_ids,\n                attention_mask,\n                pre_process=pre_process or (post_process and mtp_enable_train),\n                use_fp8_padding=use_fp8_padding,\n            )\n            input_ids_rmpad = input_ids_rmpad.contiguous()\n\n            # when pp > 1 and processor is not None, we need to pass the labels and loss_mask to the model\n            if mtp_enable_train and post_process:\n                args = {\n                    k: preprocess_packed_seqs(v, attention_mask, pre_process=True, use_fp8_padding=use_fp8_padding)[0]\n                    for k, v in logits_processor_args.items()\n                }\n                model_kwargs[\"labels\"] = args[\"label\"].contiguous()\n                model_kwargs[\"loss_mask\"] = args[\"label_mask\"].contiguous()\n\n            input_args = dict(\n                input_ids=input_ids_rmpad,\n                attention_mask=None,\n                position_ids=position_ids if not vision_model else None,  # vision models will calculate position_ids\n                packed_seq_params=packed_seq_params,\n                **model_kwargs,\n            )\n\n            if vision_model:\n                # workaround for supporting sequence packing with context parallelism\n                # cp split with sequence packing will make model lose vision token information, so we need to keep\n                # the original input_ids and pack them after vision embedding is calculated,\n                # cooporate with mbridge\n                input_args[\"input_ids\"] = input_ids\n                input_args[\"attention_mask\"] = attention_mask\n\n            output_orig = model(**input_args)\n\n            if post_process and logits_processor is not None:\n                args = {\n                    k: preprocess_packed_seqs(v, attention_mask, pre_process=True, use_fp8_padding=use_fp8_padding)[0]\n                    for k, v in logits_processor_args.items()\n                }\n                output_dict = logits_processor(output_orig, **args)\n                output = {\n                    k: postprocess_packed_seqs(\n                        v, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process\n                    )\n                    for k, v in output_dict.items()\n                }\n            else:\n                output = postprocess_packed_seqs(\n                    output_orig, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process\n                )\n        elif data_format == \"bshd\":\n            \"\"\"\n            data_format: \"thd\" or \"bshd\", default is \"thd\",\n            why we need this?\n                for some new models, GPT-OSS, the thd format is not supported, so we need to use the bshd format.\n            When using the bshd format, we have to add paddings to the input_ids to meet the longest sequence length, \n            so it is recommended to disable dynamic batch size and set batch size to 1\n            \"\"\"\n            assert fp8 is None, \"fp8 is not supported for bshd format yet\"\n\n            batch_size, sequence_length = attention_mask.shape[:2]\n            position_ids_for_preprocess = (\n                torch.arange(sequence_length, device=input_ids.device).unsqueeze(0).expand(batch_size, -1)\n                if vision_model\n                else position_ids\n            )\n            pre_process_for_bshd = True if vision_model else pre_process\n            new_input_ids, new_attention_mask, new_position_ids = preprocess_bshd(\n                input_ids,\n                attention_mask,\n                position_ids_for_preprocess,\n                sequence_parallel=sp,\n                pre_process=pre_process_for_bshd,\n            )\n            output_orig = model(\n                input_ids=new_input_ids,\n                position_ids=None if vision_model else new_position_ids,\n                attention_mask=new_attention_mask,\n                **model_kwargs,\n            )\n            if post_process and logits_processor is not None:\n                args = {\n                    k: preprocess_bshd(\n                        v, attention_mask, position_ids_for_preprocess, sequence_parallel=sp, pre_process=True\n                    )[0]\n                    for k, v in logits_processor_args.items()\n                }\n                output_dict = logits_processor(output_orig, **args)\n                output = {\n                    k: postprocess_bshd(\n                        v, new_attention_mask, attention_mask, sequence_length, post_process=post_process\n                    )\n                    for k, v in output_dict.items()\n                }\n            else:\n                output = postprocess_bshd(\n                    output_orig, new_attention_mask, attention_mask, sequence_length, post_process=post_process\n                )\n        if value_model and post_process:\n            output = output[..., 0]\n        return output\n\n    return model_forward\n\n\ndef _convert_to_nested_tensor(v, input_ids_lengths):\n    \"\"\"Convert regular tensor to NestedTensor, slicing according to input_ids_lengths.\n\n    Args:\n        v: Tensor to convert, shape [batch, seq_len]\n        input_ids_lengths: List of valid lengths for each sample\n\n    Returns:\n        Converted NestedTensor\n    \"\"\"\n    if isinstance(v, NestedTensor):\n        return v\n\n    batch_size = v.shape[0]\n    assert len(input_ids_lengths) == batch_size, (\n        f\"len(input_ids_lengths)={len(input_ids_lengths)} != batch_size={batch_size}\"\n    )\n\n    v_split_list = []\n    for i in range(batch_size):\n        vi = v[i]\n        target_len = input_ids_lengths[i]\n        if vi.shape[0] > target_len:\n            vi = vi[:target_len]\n        elif vi.shape[0] < target_len:\n            vi = torch.cat([vi, torch.ones(target_len - vi.shape[0], dtype=vi.dtype, device=vi.device)])\n        v_split_list.append(vi)\n\n    v = torch.nested.nested_tensor(v_split_list, layout=torch.jagged)\n    return v\n\n\ndef gptmodel_forward_no_padding(\n    model,\n    input_ids,\n    multi_modal_inputs: dict,\n    logits_processor=None,\n    logits_processor_args: dict = None,\n    value_model=False,\n    vision_model=False,\n    pad_token_id=None,\n    data_format: str = \"thd\",\n    mtp_enable_train: bool = False,\n):\n    \"\"\"Default forward pass for GPT models with optional sequence packing.\"\"\"\n\n    assert data_format in [\"thd\", \"bshd\"], \"data_format must be 'thd' or 'bshd'\"\n    pre_process = unwrap_model(model).pre_process\n    post_process = unwrap_model(model).post_process\n\n    fp8 = unwrap_model(model).config.fp8\n    use_fp8_padding = fp8 in [\"e4m3\", \"hybrid\"]\n\n    model_kwargs = {}\n    if \"pixel_values\" in multi_modal_inputs:\n        model_kwargs[\"pixel_values\"] = multi_modal_inputs[\"pixel_values\"].to(input_ids.device)\n    if \"image_grid_thw\" in multi_modal_inputs:\n        model_kwargs[\"image_grid_thw\"] = multi_modal_inputs[\"image_grid_thw\"].to(input_ids.device)\n    if \"pixel_values_videos\" in multi_modal_inputs:\n        model_kwargs[\"pixel_values_videos\"] = multi_modal_inputs[\"pixel_values_videos\"].to(input_ids.device)\n    if \"video_grid_thw\" in multi_modal_inputs:\n        model_kwargs[\"video_grid_thw\"] = multi_modal_inputs[\"video_grid_thw\"].to(input_ids.device)\n\n    batch_size = input_ids.shape[0]\n    if data_format == \"thd\":\n        input_ids_rmpad, packed_seq_params, position_ids_rmpad = preprocess_thd_no_padding(\n            input_ids, pre_process=pre_process or (post_process and mtp_enable_train), use_fp8_padding=use_fp8_padding\n        )\n        input_ids_rmpad = input_ids_rmpad.contiguous()\n\n        args = {}\n        if mtp_enable_train and post_process:\n            # Use input_ids sequence length to ensure label and loss_mask alignment\n            input_ids_offsets = input_ids.offsets()\n            input_ids_lengths = input_ids_offsets.diff().tolist()\n\n            for k in [\"label\", \"loss_mask\"]:\n                v = logits_processor_args[k]\n                v = _convert_to_nested_tensor(v, input_ids_lengths)\n                logits_processor_args[k] = v\n                args[k] = preprocess_thd_no_padding(\n                    v, pre_process=True, need_roll=True, use_fp8_padding=use_fp8_padding\n                )[0]\n\n            model_kwargs[\"labels\"] = args[\"label\"].contiguous()\n            model_kwargs[\"loss_mask\"] = args[\"loss_mask\"].contiguous()\n\n        if logits_processor_args and \"loss_mask\" in logits_processor_args:\n            logits_processor_args.pop(\"loss_mask\")\n\n        # For VLM model, need to pass bshd format `input_ids` and `attention_mask`.\n        attention_mask = None\n        if vision_model:\n            input_ids_rmpad = input_ids.to_padded_tensor(pad_token_id)\n            seqlens_in_batch = input_ids.offsets().diff()\n            attention_mask = torch.zeros_like(input_ids_rmpad, dtype=torch.bool)\n            for i, seqlen in enumerate(seqlens_in_batch):\n                attention_mask[i, :seqlen] = True\n\n        output_orig = model(\n            input_ids=input_ids_rmpad,\n            attention_mask=attention_mask,\n            position_ids=position_ids_rmpad if not vision_model else None,  # vision models will calculate position_ids\n            packed_seq_params=packed_seq_params,\n            **model_kwargs,\n        )\n\n        if post_process and logits_processor is not None:\n            args = {\n                k: preprocess_thd_no_padding(\n                    v, pre_process=True, need_roll=(k == \"label\"), use_fp8_padding=use_fp8_padding\n                )[0]\n                for k, v in logits_processor_args.items()\n            }\n            output_dict = logits_processor(output_orig, **args)\n            output = {\n                k: postprocess_thd_no_padding(v, packed_seq_params, input_ids, batch_size, post_process=post_process)\n                for k, v in output_dict.items()\n            }\n        else:\n            output = postprocess_thd_no_padding(\n                output_orig, packed_seq_params, input_ids, batch_size, post_process=post_process\n            )\n    else:\n        \"\"\"\n        data_format: \"thd\" or \"bshd\", default is \"thd\",\n        why we need this?\n            for some new models, GPT-OSS, the thd format is not supported, so we need to use the bshd format.\n        When using the bshd format, we have to add paddings to the input_ids to meet the longest sequence length, \n        so it is recommended to disable dynamic batch size and set batch size to 1\n        \"\"\"\n\n        input_ids_bshd, attention_mask_bshd, position_ids_bshd = preprocess_bshd_no_padding(\n            input_ids, pre_process=pre_process or (post_process and mtp_enable_train), use_fp8_padding=use_fp8_padding\n        )\n\n        if mtp_enable_train and post_process:\n            args = {}\n            # Use input_ids sequence length to ensure label and loss_mask alignment\n            input_ids_offsets = input_ids.offsets()\n            input_ids_lengths = input_ids_offsets.diff().tolist()\n\n            for k in [\"label\", \"loss_mask\"]:\n                v = logits_processor_args[k]\n                v = _convert_to_nested_tensor(v, input_ids_lengths)\n                logits_processor_args[k] = v\n                args[k] = preprocess_bshd_no_padding(\n                    v, pre_process=True, need_roll=True, use_fp8_padding=use_fp8_padding\n                )[0]\n            model_kwargs[\"labels\"] = args[\"label\"].contiguous()\n            model_kwargs[\"loss_mask\"] = args[\"loss_mask\"].contiguous()\n\n        if logits_processor_args and \"loss_mask\" in logits_processor_args:\n            logits_processor_args.pop(\"loss_mask\")\n\n        output_orig = model(\n            input_ids=input_ids_bshd,\n            attention_mask=attention_mask_bshd,\n            position_ids=None if vision_model else position_ids_bshd,\n            **model_kwargs,\n        )\n        if post_process and logits_processor is not None:\n            args = {\n                k: preprocess_bshd_no_padding(\n                    v, pre_process=True, need_roll=(k == \"label\"), use_fp8_padding=use_fp8_padding\n                )[0]\n                for k, v in logits_processor_args.items()\n            }\n            output_dict = logits_processor(output_orig, **args)\n            output = {\n                k: postprocess_bshd_no_padding(v, attention_mask_bshd, post_process=post_process)\n                for k, v in output_dict.items()\n            }\n        else:\n            output = postprocess_bshd_no_padding(output_orig, attention_mask_bshd, post_process=post_process)\n\n    if value_model and post_process:\n        # output = output[..., 0]\n        # while using nested tensor, the advanced indexing operation above will result in an error at backward, i.e.\n        # ValueError: NestedTensor _nested_select_backward_default(grad_output: t, self: jt_all, dim: any, index: any)\n        # so we use `squeeze` to remove the last dimension\n        output = output.squeeze(-1)\n\n    return output\n"
  },
  {
    "path": "verl/models/mcore/model_forward_1f1b_overlap.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable, Optional\n\nimport torch\nfrom megatron.core.models.common.model_chunk_schedule_plan import TransformerModelChunkSchedulePlan\nfrom megatron.core.models.gpt.gpt_model import GPTModel\nfrom megatron.core.utils import make_viewless_tensor\nfrom torch import Tensor\n\nfrom verl.models.mcore.util import preprocess_packed_seqs\nfrom verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\nfrom verl.utils.megatron_utils import unwrap_model\nfrom verl.utils.model import CausalLMOutputForPPO\n\nfrom .util import postprocess_packed_seqs, postprocess_packed_seqs_for_dict_output\n\n\ndef gptmodel_forward_1f1b_overlap(\n    model: GPTModel,\n    input_ids: Tensor,\n    position_ids: Tensor,\n    attention_mask: Tensor,\n    labels: Tensor = None,\n    labels_mask: Tensor = None,\n    multi_modal_inputs: Optional[dict] = None,\n    logits_processor: Optional[Callable] = None,\n    logits_processor_args: Optional[dict] = None,\n    temperature: float = 1.0,\n) -> TransformerModelChunkSchedulePlan:\n    pre_process: bool = unwrap_model(model).pre_process\n    post_process: bool = unwrap_model(model).post_process\n    assert logits_processor is None, \"only support fused kernel\"\n    batch_size, seq_len = attention_mask.shape[:2]\n    input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process)\n    input_ids_rmpad = input_ids_rmpad.contiguous()\n\n    schedule_plan = model.build_schedule_plan(\n        input_ids=input_ids_rmpad,\n        attention_mask=attention_mask,\n        labels=labels,\n        position_ids=position_ids,\n        packed_seq_params=packed_seq_params,\n    )\n    if post_process:\n        attention_mask_out = attention_mask\n\n        def _postprocess(\n            self,\n            hidden_states,\n            input_ids,\n            position_ids,\n            labels,\n            rotary_pos_emb,\n            rotary_pos_cos,\n            rotary_pos_sin,\n            mtp_in_postprocess=None,\n            loss_mask=None,\n            decoder_input=None,\n            attention_mask=None,\n            inference_params=None,\n            packed_seq_params=None,\n            sequence_len_offset=None,\n            runtime_gather_output=None,\n            extra_block_kwargs=None,\n            inference_context=None,\n        ):\n            \"\"\"patched from https://github.com/NVIDIA/Megatron-LM/blob/core_r0.14.0/megatron/core/models/gpt/gpt_model.py#L412\"\"\"\n            \"\"\"Postprocesses decoder hidden states to generate logits or compute loss.\n\n            Applies Multi-Token Prediction if enabled, generates output logits through\n            the output layer, and computes language model loss when labels are provided.\n            \"\"\"\n            from megatron.core import parallel_state\n            from megatron.core.tensor_parallel import gather_from_sequence_parallel_region\n\n            in_inference_mode = inference_context is not None and not self.training\n            if in_inference_mode:\n                assert runtime_gather_output, \"Inference must always gather TP logits\"\n\n            # logits and loss\n            output_weight = None\n            if self.share_embeddings_and_output_weights:\n                output_weight = self.shared_embedding_or_output_weight()\n\n            if mtp_in_postprocess:\n                hidden_states = self.mtp(\n                    input_ids=input_ids,\n                    position_ids=position_ids,\n                    hidden_states=hidden_states,\n                    attention_mask=attention_mask,\n                    inference_params=inference_params,\n                    rotary_pos_emb=rotary_pos_emb,\n                    rotary_pos_cos=rotary_pos_cos,\n                    rotary_pos_sin=rotary_pos_sin,\n                    packed_seq_params=packed_seq_params,\n                    sequence_len_offset=sequence_len_offset,\n                    embedding=self.embedding,\n                    **(extra_block_kwargs or {}),\n                )\n\n            if not self.post_process:\n                return hidden_states\n\n            if self.mtp_process:\n                from megatron.core.transformer.multi_token_prediction import (\n                    MTPLossAutoScaler,\n                    MTPLossLoggingHelper,\n                    roll_tensor,\n                )\n\n                mtp_labels = labels.clone()\n                hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)\n                hidden_states = hidden_states_list[0]\n                if loss_mask is None:\n                    # if loss_mask is not provided, use all ones as loss_mask\n                    loss_mask = torch.ones_like(mtp_labels)\n                for mtp_layer_number in range(self.config.mtp_num_layers):\n                    # output\n                    mtp_logits, _ = self.output_layer(\n                        hidden_states_list[mtp_layer_number + 1],\n                        weight=output_weight,\n                        runtime_gather_output=runtime_gather_output,\n                    )\n                    # Calc loss for the current Multi-Token Prediction (MTP) layers.\n                    mtp_labels, _ = roll_tensor(mtp_labels, shifts=-1, dims=-1, cp_group=self.cp_group)\n                    loss_mask, num_tokens = roll_tensor(loss_mask, shifts=-1, dims=-1, cp_group=self.cp_group)\n                    mtp_loss = self.compute_language_model_loss(mtp_labels, mtp_logits)\n                    mtp_loss = loss_mask * mtp_loss\n                    if self.training:\n                        # TODO(shifangx): remove the use of parallel_state here\n                        # after moving loss logging to loss_func in pretrain_gpt.py\n                        MTPLossLoggingHelper.save_loss_to_tracker(\n                            torch.sum(mtp_loss) / num_tokens,\n                            mtp_layer_number,\n                            self.config.mtp_num_layers,\n                            avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),\n                        )\n                    mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers\n                    if self.config.calculate_per_token_loss:\n                        hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss)\n                    else:\n                        hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss / num_tokens)\n\n            if logits_processor is not None:\n                logits, _ = self.output_layer(\n                    hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output\n                )\n                output_orig = logits.transpose(0, 1).contiguous()\n                args = {\n                    k: preprocess_packed_seqs(v, attention_mask_out, pre_process=True)[0]\n                    for k, v in logits_processor_args.items()\n                }\n                output_dict = logits_processor(output_orig, **args)\n                output = {\n                    k: postprocess_packed_seqs(\n                        v, packed_seq_params, attention_mask_out, batch_size, seq_len, post_process=post_process\n                    )\n                    for k, v in output_dict.items()\n                }\n            else:\n                # fused kernel\n\n                labels_rmpad, _ = preprocess_packed_seqs(labels, attention_mask, pre_process=True)\n                labels_mask_rmpad, _ = preprocess_packed_seqs(labels_mask, attention_mask, pre_process=True)\n                labels_rmpad = labels_rmpad.contiguous()\n                labels_mask_rmpad = labels_mask_rmpad.contiguous()\n\n                output = CausalLMOutputForPPO(\n                    loss=None,\n                    logits=None,\n                    past_key_values=None,\n                    hidden_states=hidden_states,\n                    attentions=None,\n                )\n                if self.config.sequence_parallel:\n                    hidden_states = gather_from_sequence_parallel_region(hidden_states)\n                logprobs, entropy = linear_cross_entropy(\n                    hidden_states,\n                    self.output_layer.weight,\n                    labels_rmpad,\n                    temperature,\n                    \"none\",\n                    parallel_state.get_tensor_model_parallel_group(),\n                )\n                output.entropy = entropy\n                output.log_probs = logprobs\n\n                output = postprocess_packed_seqs_for_dict_output(\n                    labels_mask_rmpad,\n                    output,\n                    packed_seq_params,\n                    attention_mask,\n                    batch_size,\n                    seq_len,\n                    post_process=post_process,\n                )\n            output_ = [output[\"log_probs\"]]\n            # TODO NOW 1f1b overlap only support one tensor output\n            # if \"entropy\" in output:\n            #     output_.append(output[\"entropy\"])\n            output_ = tuple(output_)\n            return output_\n\n        def _custom_post_process_node_forward_impl(self, hidden_states):\n            if self.gpt_model.decoder.final_layernorm and not self.gpt_model.mtp_process:\n                hidden_states = self.gpt_model.decoder.final_layernorm(hidden_states)\n                # TENorm produces a \"viewed\" tensor. This will result in schedule.py's\n                # deallocate_output_tensor() throwing an error, so a viewless tensor is\n                # created to prevent this.\n                hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)\n\n            # Run GPTModel._postprocess\n            output = self.gpt_model._postprocess(\n                hidden_states=hidden_states,\n                input_ids=self.chunk_state.input_ids,\n                position_ids=self.chunk_state.position_ids,\n                labels=self.chunk_state.labels,\n                decoder_input=self.chunk_state.decoder_input,\n                rotary_pos_emb=self.chunk_state.rotary_pos_emb,\n                rotary_pos_cos=self.chunk_state.rotary_pos_cos,\n                rotary_pos_sin=self.chunk_state.rotary_pos_sin,\n                mtp_in_postprocess=False,\n                loss_mask=self.chunk_state.loss_mask,\n                attention_mask=self.chunk_state.attention_mask,\n                packed_seq_params=self.chunk_state.packed_seq_params,\n                sequence_len_offset=self.chunk_state.sequence_len_offset,\n                runtime_gather_output=self.chunk_state.runtime_gather_output,\n                extra_block_kwargs=self.chunk_state.extra_block_kwargs,\n            )\n            return output\n\n        schedule_plan.post_process.forward_impl = _custom_post_process_node_forward_impl.__get__(\n            schedule_plan.post_process, schedule_plan.post_process.__class__\n        )\n        unwrap_model(model)._postprocess = _postprocess.__get__(unwrap_model(model), unwrap_model(model).__class__)\n\n    return schedule_plan\n"
  },
  {
    "path": "verl/models/mcore/model_forward_fused.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections import OrderedDict\nfrom typing import Optional\n\nimport megatron.core as mcore\nimport torch\nfrom megatron.core import parallel_state\nfrom megatron.core.config_logger import has_config_logger_enabled, log_config_to_disk\nfrom megatron.core.inference.contexts import BaseInferenceContext\nfrom megatron.core.models.gpt.gpt_model import GPTModel\nfrom megatron.core.packed_seq_params import PackedSeqParams\nfrom megatron.core.tensor_parallel.mappings import gather_from_sequence_parallel_region\nfrom megatron.core.utils import deprecate_inference_params\nfrom packaging import version\nfrom torch import Tensor\n\nfrom verl.models.mcore.util import preprocess_packed_seqs, preprocess_thd_no_padding\nfrom verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\nfrom verl.utils.megatron_utils import unwrap_model\nfrom verl.utils.model import CausalLMOutputForPPO\n\nfrom .util import postprocess_packed_seqs_for_dict_output, postprocess_thd_no_padding\n\n\ndef _get_patching_model(model: torch.nn.Module):\n    model = unwrap_model(model)\n    if isinstance(model, GPTModel):\n        return model\n\n    if not (hasattr(model, \"language_model\") and isinstance(model.language_model, GPTModel)):\n        print(f\"Model {model.__class__.__name__} is not a supported for fused forward\")\n        return None\n\n    return model.language_model\n\n\ndef patch_fused_forward(model: torch.nn.Module):\n    assert version.parse(mcore.__version__) >= version.parse(\"0.13.0\"), (\n        \"Fused forward patching requires mecore >= 0.13.0\"\n    )\n    model = _get_patching_model(model)\n    if model is not None:\n        model.forward_backup = model.forward\n        model.forward = _fused_GPTModel_forward.__get__(model, model.__class__)\n\n\ndef unpatch_fused_forward(model: torch.nn.Module):\n    model = _get_patching_model(model)\n    if model is not None:\n        model.forward = model.forward_backup\n\n\ndef fused_forward_model_gen(vision_model: bool = False):\n    def fused_forward_model(\n        model,\n        input_ids: Tensor,\n        position_ids: Tensor,\n        attention_mask: Tensor,\n        labels: Tensor,\n        labels_mask: Tensor,\n        temperature: float,\n        multi_modal_inputs: dict,\n    ):\n        pre_process: bool = (\n            unwrap_model(model).pre_process if not vision_model else False\n        )  # vision model does not need pre_process, because we pack the input_ids to thd in the forward function\n        post_process: bool = unwrap_model(model).post_process\n\n        model_kwargs = {}\n        if \"pixel_values\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values\"] = multi_modal_inputs[\"pixel_values\"].to(input_ids.device)\n        if \"image_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"image_grid_thw\"] = multi_modal_inputs[\"image_grid_thw\"].to(input_ids.device)\n        if \"pixel_values_videos\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values_videos\"] = multi_modal_inputs[\"pixel_values_videos\"].to(input_ids.device)\n        if \"video_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"video_grid_thw\"] = multi_modal_inputs[\"video_grid_thw\"].to(input_ids.device)\n\n        batch_size, seq_len = attention_mask.shape[:2]\n        input_ids_rmpad, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=pre_process)\n        input_ids_rmpad = input_ids_rmpad.contiguous()\n        labels_rmpad, _ = preprocess_packed_seqs(labels, attention_mask, pre_process=True)\n        labels_mask_rmpad, _ = preprocess_packed_seqs(labels_mask, attention_mask, pre_process=True)\n        labels_rmpad = labels_rmpad.contiguous()\n        labels_mask_rmpad = labels_mask_rmpad.contiguous()\n\n        input_args = dict(\n            input_ids=input_ids_rmpad,\n            attention_mask=None,\n            position_ids=position_ids if not vision_model else None,  # vision models will calculate position_ids\n            packed_seq_params=packed_seq_params,\n            labels=labels_rmpad,\n            temperature=temperature,\n            **model_kwargs,\n        )\n\n        if vision_model:\n            # workaround for supporting sequence packing with context parallelism\n            # cp split with sequence packing will make model lose vision token information, so we need to keep\n            # the original input_ids and pack them after vision embedding is calculated,\n            # cooporate with mbridge\n            input_args[\"input_ids\"] = input_ids\n            input_args[\"attention_mask\"] = attention_mask\n\n        output_orig: CausalLMOutputForPPO = model(**input_args)\n\n        if post_process:\n            # output_orig is in type of CausalLMOutputForPPO\n            output = postprocess_packed_seqs_for_dict_output(\n                labels_mask_rmpad,\n                output_orig,\n                packed_seq_params,\n                attention_mask,\n                batch_size,\n                seq_len,\n                post_process=post_process,\n            )\n        else:\n            output = output_orig\n        return output\n\n    return fused_forward_model\n\n\ndef fused_forward_no_padding_gen(vision_model: bool = False):\n    def fused_forward_no_padding(\n        model,\n        input_ids: Tensor,\n        labels: Tensor,\n        multi_modal_inputs: dict,\n        temperature: float,\n        calculate_entropy: bool,\n        pad_token_id: int,\n    ):\n        pre_process = unwrap_model(model).pre_process\n        post_process = unwrap_model(model).post_process\n\n        fp8 = unwrap_model(model).config.fp8\n        use_fp8_padding = fp8 in [\"e4m3\", \"hybrid\"]\n\n        input_ids_rmpad, packed_seq_params, _ = preprocess_thd_no_padding(\n            input_ids, pre_process=pre_process, use_fp8_padding=use_fp8_padding\n        )\n        input_ids_rmpad = input_ids_rmpad.contiguous()\n\n        model_kwargs = {}\n        if \"pixel_values\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values\"] = multi_modal_inputs[\"pixel_values\"].to(input_ids.device)\n        if \"image_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"image_grid_thw\"] = multi_modal_inputs[\"image_grid_thw\"].to(input_ids.device)\n        if \"pixel_values_videos\" in multi_modal_inputs:\n            model_kwargs[\"pixel_values_videos\"] = multi_modal_inputs[\"pixel_values_videos\"].to(input_ids.device)\n        if \"video_grid_thw\" in multi_modal_inputs:\n            model_kwargs[\"video_grid_thw\"] = multi_modal_inputs[\"video_grid_thw\"].to(input_ids.device)\n\n        attention_mask = None\n        if vision_model:\n            input_ids_rmpad = input_ids.to_padded_tensor(pad_token_id)\n            seqlens_in_batch = input_ids.offsets().diff().to(input_ids.device)\n            max_seq_len = input_ids_rmpad.shape[1]\n            attention_mask = torch.arange(max_seq_len, device=input_ids.device).unsqueeze(\n                0\n            ) < seqlens_in_batch.unsqueeze(1)\n\n        labels_rmpad, _, _ = preprocess_thd_no_padding(\n            labels, pre_process=True, need_roll=True, use_fp8_padding=use_fp8_padding\n        )\n        labels_rmpad = labels_rmpad.contiguous()\n        output_orig: CausalLMOutputForPPO = model(\n            input_ids=input_ids_rmpad,\n            attention_mask=attention_mask,\n            position_ids=None,\n            packed_seq_params=packed_seq_params,\n            labels=labels_rmpad,\n            temperature=temperature,\n            **model_kwargs,\n        )\n\n        if not post_process:\n            return output_orig\n\n        log_probs = output_orig.log_probs\n        if log_probs.dim() == 1:\n            log_probs = log_probs.unsqueeze(0)\n        log_probs = postprocess_thd_no_padding(\n            log_probs, packed_seq_params, input_ids, input_ids.shape[0], post_process=post_process\n        )\n\n        output = {\"log_probs\": log_probs}\n\n        if calculate_entropy:\n            entropy = output_orig.entropy\n            if entropy.dim() == 1:\n                entropy = entropy.unsqueeze(0)\n            entropy = postprocess_thd_no_padding(\n                entropy, packed_seq_params, input_ids, input_ids.shape[0], post_process=post_process\n            )\n            output[\"entropy\"] = entropy\n\n        return output\n\n    return fused_forward_no_padding\n\n\ndef _fused_GPTModel_forward(\n    model,\n    input_ids: Tensor,\n    position_ids: Tensor,\n    attention_mask: Tensor,\n    decoder_input: Tensor = None,\n    labels: Tensor = None,\n    inference_context: BaseInferenceContext = None,\n    packed_seq_params: PackedSeqParams = None,\n    extra_block_kwargs: dict = None,\n    runtime_gather_output: Optional[bool] = None,\n    *,\n    inference_params: Optional[BaseInferenceContext] = None,\n    loss_mask: Optional[Tensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> CausalLMOutputForPPO:\n    \"\"\"\n    Patch self._postprocess in forward for GPT models to enable fused kernel support.\n    https://github.com/NVIDIA/Megatron-LM/blob/core_v0.13.0/megatron/core/models/gpt/gpt_model.py\n\n    TODO: Currently we still need to patch `forward` because we need to pass `temperature`\n    explicitly to `self._postprocess` when calling, maybe there can be a better way to handle this?\n    \"\"\"\n\n    inference_context = deprecate_inference_params(inference_context, inference_params)\n\n    preproc_output = model._preprocess(\n        input_ids=input_ids,\n        position_ids=position_ids,\n        decoder_input=decoder_input,\n        inference_context=inference_context,\n        packed_seq_params=packed_seq_params,\n    )\n\n    (decoder_input, rotary_pos_emb, rotary_pos_cos, rotary_pos_sin, sequence_len_offset) = preproc_output[:5]\n\n    # Run decoder.\n    hidden_states = model.decoder(\n        hidden_states=decoder_input,\n        attention_mask=attention_mask,\n        inference_context=inference_context,\n        rotary_pos_emb=rotary_pos_emb,\n        rotary_pos_cos=rotary_pos_cos,\n        rotary_pos_sin=rotary_pos_sin,\n        packed_seq_params=packed_seq_params,\n        sequence_len_offset=sequence_len_offset,\n        **(extra_block_kwargs or {}),\n        **kwargs,\n    )\n\n    if not model.post_process:\n        return hidden_states\n\n    output = CausalLMOutputForPPO(\n        loss=None,\n        logits=None,\n        past_key_values=None,\n        hidden_states=hidden_states,\n        attentions=None,\n    )\n\n    if model.config.sequence_parallel:\n        hidden_states = gather_from_sequence_parallel_region(hidden_states)\n\n    # Get the output weight - use embedding weight if output_layer is None or weight is shared\n    if hasattr(model, \"output_layer\") and model.output_layer is not None and model.output_layer.weight is not None:\n        output_weight = model.output_layer.weight\n    else:\n        # When embeddings are tied, use the embedding weight\n        output_weight = model.embedding.word_embeddings.weight\n\n    logprobs, entropy = linear_cross_entropy(\n        hidden_states,\n        output_weight,\n        labels,\n        temperature,\n        \"none\",\n        parallel_state.get_tensor_model_parallel_group(),\n    )\n\n    if has_config_logger_enabled(model.config):\n        payload = OrderedDict(\n            {\n                \"input_ids\": input_ids,\n                \"position_ids\": position_ids,\n                \"attention_mask\": attention_mask,\n                \"decoder_input\": decoder_input,\n                \"logprobs\": logprobs,\n                \"entropy\": entropy,\n            }\n        )\n        log_config_to_disk(model.config, payload, prefix=\"input_and_logits\")\n\n    output.entropy = entropy\n    output.log_probs = logprobs\n\n    return output\n"
  },
  {
    "path": "verl/models/mcore/model_initializer.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# use mcore transformer config to initialize the model\nimport inspect\nfrom abc import ABC, abstractmethod\n\nfrom megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec, get_gpt_mtp_block_spec\nfrom megatron.core.models.gpt.gpt_model import GPTModel\n\nfrom .config_converter import PretrainedConfig, TransformerConfig\n\n\nclass BaseModelInitializer(ABC):\n    \"\"\"Base class for model initializers.\"\"\"\n\n    def __init__(self, tfconfig: TransformerConfig, hf_config: PretrainedConfig):\n        self.tfconfig = tfconfig\n        self.hf_config = hf_config\n        self.has_vp_stage = inspect.signature(get_gpt_decoder_block_spec).parameters.get(\"vp_stage\", None) is not None\n\n    @abstractmethod\n    def get_transformer_layer_spec(self, vp_stage=None):\n        \"\"\"Get the transformer layer specification.\n        https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/gpt/gpt_layer_specs.py\"\"\"\n        pass\n\n    def get_rope_scaling_args(self) -> dict:\n        \"\"\"Get rope scaling args.\"\"\"\n        rope_scaling_args = {}\n        if \"rope_scaling\" in self.hf_config:\n            if self.hf_config.rope_scaling is not None:\n                # assert self.hf_config.rope_scaling[\"type\"] == \"linear\", \"only linear scaling is supported for now\"\n                rope_scaling_args[\"seq_len_interpolation_factor\"] = self.hf_config.rope_scaling[\"factor\"]\n        return rope_scaling_args\n\n    def initialize(\n        self,\n        pre_process: bool = True,\n        post_process: bool = True,\n        share_embeddings_and_output_weights: bool = False,\n        value: bool = False,\n        **extra_kwargs,\n    ) -> GPTModel:\n        \"\"\"Initialize a GPT model with the given configuration.\n        https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/models/gpt/gpt_model.py\n\n        Args:\n            pre_process (bool): include embedding layer.\n            post_process (bool): including an output layer.\n            share_embeddings_and_output_weights (bool): input embeddings and output logit weights are shared.\n            value (bool): add an extra linear layer for classification or regression.\n\n        Returns:\n            GPTModel: An initialized GPT model instance\n        \"\"\"\n        vp_stage = extra_kwargs.get(\"vp_stage\", None)\n        transformer_layer_spec = self.get_transformer_layer_spec(vp_stage=vp_stage)\n        rope_scaling_args = self.get_rope_scaling_args()\n        mtp_block_spec = extra_kwargs.get(\"mtp_block_spec\", None)\n        model = GPTModel(\n            config=self.tfconfig,\n            transformer_layer_spec=transformer_layer_spec,\n            vocab_size=self.hf_config.vocab_size,\n            max_sequence_length=self.hf_config.max_position_embeddings,\n            pre_process=pre_process,\n            post_process=post_process,\n            share_embeddings_and_output_weights=share_embeddings_and_output_weights,\n            position_embedding_type=\"rope\",\n            rotary_base=self.hf_config.rope_theta,\n            **rope_scaling_args,\n            mtp_block_spec=mtp_block_spec,\n            **({} if not self.has_vp_stage else {\"vp_stage\": vp_stage}),\n        )\n\n        if post_process and value:\n            from verl.models.llama.megatron.layers.parallel_linear import LinearForLastLayer\n\n            model.output_layer = LinearForLastLayer(\n                input_size=self.tfconfig.hidden_size, output_size=1, config=self.tfconfig\n            )\n\n        return model\n\n\nclass DenseModel(BaseModelInitializer):\n    \"\"\"Initializer for dense models like Llama and Qwen2.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        assert self.tfconfig.normalization == \"RMSNorm\", \"only RMSNorm is supported for now\"\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        return get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n\n\nclass Qwen2MoEModel(BaseModelInitializer):\n    \"\"\"Initializer for Qwen2 MoE models.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        assert self.tfconfig.normalization == \"RMSNorm\", \"only RMSNorm is supported for now\"\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n\n        # Patch layer spec for shared experts\n        for i in range(len(transformer_layer_spec.layer_specs)):\n            transformer_layer_spec.layer_specs[i].submodules.mlp.submodules.shared_experts.params[\"gate\"] = True\n\n        return transformer_layer_spec\n\n    def initialize(self, **kwargs):\n        # Qwen default freeze_moe_router: true\n        model = super().initialize(**kwargs)\n        freeze_moe_router = kwargs.get(\"freeze_moe_router\", True)\n        if freeze_moe_router:\n            for layer in model.decoder.layers:\n                layer.mlp.router.weight.requires_grad = False\n        return model\n\n\nclass MixtralModel(BaseModelInitializer):\n    \"\"\"Initializer for Mixtral models.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        assert self.tfconfig.normalization == \"RMSNorm\", \"only RMSNorm is supported for now\"\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n        return transformer_layer_spec\n\n    def initialize(self, **kwargs):\n        model = super().initialize(**kwargs)\n        freeze_moe_router = kwargs.get(\"freeze_moe_router\", False)\n        if freeze_moe_router:\n            for layer in model.decoder.layers:\n                layer.mlp.router.weight.requires_grad = False\n        return model\n\n\nclass Qwen3MoEModel(BaseModelInitializer):\n    \"\"\"Initializer for Qwen3 MoE models.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        assert self.tfconfig.normalization == \"RMSNorm\", \"only RMSNorm is supported for now\"\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n        return transformer_layer_spec\n\n    def initialize(self, **kwargs):\n        # Qwen default freeze_moe_router: true\n        model = super().initialize(**kwargs)\n        freeze_moe_router = kwargs.get(\"freeze_moe_router\", True)\n        if freeze_moe_router:\n            for layer in model.decoder.layers:\n                layer.mlp.router.weight.requires_grad = False\n        return model\n\n\nclass DeepseekV3Model(BaseModelInitializer):\n    \"\"\"Initializer for DeepseekV3 models.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n        return transformer_layer_spec\n\n    def get_rope_scaling_args(self) -> dict:\n        \"\"\"Get rope scaling args.\"\"\"\n        rope_scaling_args = {}\n        return rope_scaling_args\n\n    def initialize(\n        self,\n        **kwargs,\n    ):\n        vp_stage = kwargs.get(\"vp_stage\", None)\n        freeze_moe_router = kwargs.get(\"freeze_moe_router\", True)\n        if freeze_moe_router:\n            self.tfconfig.moe_router_load_balancing_type = \"none\"\n        # MTP\n        if self.tfconfig.mtp_num_layers is not None and self.tfconfig.mtp_num_layers > 0:\n            transformer_layer_spec = self.get_transformer_layer_spec(vp_stage=vp_stage)\n            mtp_block_spec = get_gpt_mtp_block_spec(\n                self.tfconfig, transformer_layer_spec, use_transformer_engine=True, vp_stage=vp_stage\n            )\n            kwargs[\"mtp_block_spec\"] = mtp_block_spec\n\n        model = super().initialize(**kwargs)\n        if freeze_moe_router:\n            for layer in model.decoder.layers:\n                if hasattr(layer.mlp, \"router\"):\n                    layer.mlp.router.weight.requires_grad = False\n        return model\n\n\nclass Qwen25VLModel(BaseModelInitializer):\n    \"\"\"Initializer for Qwen2.5 VL models.\"\"\"\n\n    def get_transformer_layer_spec(self, vp_stage=None):\n        extra_kwargs = {} if not self.has_vp_stage else {\"vp_stage\": vp_stage}\n        transformer_layer_spec = get_gpt_decoder_block_spec(self.tfconfig, use_transformer_engine=True, **extra_kwargs)\n        return transformer_layer_spec\n\n    def initialize(\n        self,\n        pre_process=None,\n        post_process=None,\n        share_embeddings_and_output_weights=False,\n        value=False,\n        **extra_kwargs,\n    ):\n        tfconfig = self.tfconfig\n        hf_config = self.hf_config\n        # Qwen2_5_VLForConditionalGeneration\n        from copy import deepcopy\n\n        transformer_layer_spec = self.get_transformer_layer_spec()\n\n        from megatron.core.extensions.transformer_engine import TEColumnParallelLinear, TERowParallelLinear\n        from megatron.core.models.gpt.moe_module_specs import MLPSubmodules\n        from megatron.core.models.vision.vit_layer_specs import get_vit_layer_with_transformer_engine_spec\n\n        from .qwen2_5_vl import Qwen2_5VLModel, get_vision_model_config, get_vision_projection_config\n\n        vision_transformer_config = get_vision_model_config(deepcopy(tfconfig))\n        vision_transformer_config.pipeline_model_parallel_size = 1\n        vision_transformer_config.first_pipeline_num_layers = None\n\n        vision_projection_config = get_vision_projection_config(\n            deepcopy(tfconfig),\n            vision_transformer_config.hidden_size,\n            spatial_merge_size=hf_config.vision_config.spatial_merge_size,\n        )\n        vision_projection_layer_spec = MLPSubmodules(\n            linear_fc1=TEColumnParallelLinear,\n            linear_fc2=TERowParallelLinear,\n        )\n        vision_transformer_layer_spec = get_vit_layer_with_transformer_engine_spec()\n\n        qwen25_vl_model = Qwen2_5VLModel(\n            language_transformer_config=tfconfig,\n            language_transformer_layer_spec=transformer_layer_spec,\n            language_vocab_size=hf_config.vocab_size,\n            language_max_sequence_length=hf_config.max_position_embeddings,\n            vision_transformer_config=vision_transformer_config,\n            vision_transformer_layer_spec=vision_transformer_layer_spec,\n            vision_projection_config=vision_projection_config,\n            vision_projection_layer_spec=vision_projection_layer_spec,\n            vision_projection_type=\"mlp\",\n            language_rotary_base=hf_config.rope_theta,\n            pre_process=pre_process,\n            post_process=post_process,\n            add_decoder=True,\n            add_encoder=True,\n            parallel_output=True,\n            language_share_embeddings_and_output_weights=share_embeddings_and_output_weights,\n        )\n\n        if post_process and value:\n            from verl.models.llama.megatron.layers.parallel_linear import LinearForLastLayer\n\n            qwen25_vl_model.language_model.output_layer = LinearForLastLayer(\n                input_size=tfconfig.hidden_size, output_size=1, config=tfconfig\n            )\n\n        return qwen25_vl_model\n"
  },
  {
    "path": "verl/models/mcore/mtp_patch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n# Copyright 2025 Meituan Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable\n\nimport torch\nfrom megatron.core import parallel_state\nfrom megatron.core.models.gpt.gpt_model import GPTModel\nfrom megatron.core.transformer.multi_token_prediction import (\n    MTPLossAutoScaler,\n    MTPLossLoggingHelper,\n    roll_tensor,\n)\n\ntry:\n    from megatron.core.utils import unwrap_model\nexcept ImportError:\n    from verl.utils.megatron_utils import unwrap_model\n\n\ndef _get_patching_model(model: torch.nn.Module):\n    model = unwrap_model(model)\n    if isinstance(model, GPTModel):\n        return model\n\n    if not (hasattr(model, \"language_model\") and isinstance(model.language_model, GPTModel)):\n        print(f\"Model {model.__class__.__name__} is not a supported for fused forward\")\n        return None\n\n    return model.language_model\n\n\ndef patch_postprocess(model: torch.nn.Module):\n    model = _get_patching_model(model)\n    if model is not None:\n        model._postprocess_backup = model._postprocess\n        model._postprocess = _megatron_gptmodel_postprocess.__get__(model, model.__class__)\n\n\ndef unpatch_postprocess(model: torch.nn.Module):\n    model = _get_patching_model(model)\n    if model is not None:\n        model._postprocess = model._postprocess_backup\n\n\n# copy from https://github.com/NVIDIA/Megatron-LM/blob/23e092f41ec8bc659020e401ddac9576c1cfed7e/megatron/core/models/gpt/gpt_model.py\n# patch the postprocess method of GPTModel to support advanced features like MTP, 1f1b overlap, etc.\ndef _megatron_gptmodel_postprocess(\n    self,\n    hidden_states,\n    input_ids,\n    position_ids,\n    labels,\n    rotary_pos_emb,\n    rotary_pos_cos,\n    rotary_pos_sin,\n    mtp_in_postprocess=None,\n    loss_mask=None,\n    decoder_input=None,\n    attention_mask=None,\n    inference_params=None,\n    packed_seq_params=None,\n    sequence_len_offset=None,\n    runtime_gather_output=None,\n    extra_block_kwargs=None,\n    inference_context=None,\n):\n    \"\"\"Postprocesses decoder hidden states to generate logits or compute loss.\n\n    Applies Multi-Token Prediction if enabled, generates output logits through\n    the output layer, and computes language model loss when labels are provided.\n    \"\"\"\n\n    # logits and loss\n    output_weight = None\n    if self.share_embeddings_and_output_weights:\n        output_weight = self.shared_embedding_or_output_weight()\n\n    if mtp_in_postprocess and labels is not None:\n        hidden_states = self.mtp(\n            input_ids=input_ids,\n            position_ids=position_ids,\n            hidden_states=hidden_states,\n            attention_mask=attention_mask,\n            inference_params=inference_params,\n            rotary_pos_emb=rotary_pos_emb,\n            rotary_pos_cos=rotary_pos_cos,\n            rotary_pos_sin=rotary_pos_sin,\n            packed_seq_params=packed_seq_params,\n            sequence_len_offset=sequence_len_offset,\n            embedding=self.embedding,\n            **(extra_block_kwargs or {}),\n        )\n\n    if not self.post_process:\n        return hidden_states\n\n    # Skip when mtp_num_layers is None or 0\n    if self.config.mtp_num_layers and labels is not None:\n        mtp_labels = labels.clone()\n\n        hidden_states_list = torch.chunk(hidden_states, 1 + self.config.mtp_num_layers, dim=0)\n        hidden_states = hidden_states_list[0]\n        if loss_mask is None:\n            # if loss_mask is not provided, use all ones as loss_mask\n            loss_mask = torch.ones_like(mtp_labels)\n        for mtp_layer_number in range(self.config.mtp_num_layers):\n            # Calc loss for the current Multi-Token Prediction (MTP) layers.\n            mtp_labels, _ = roll_tensor(\n                mtp_labels,\n                shifts=-1,\n                dims=-1,\n                cp_group=self.cp_group,\n                packed_seq_params=packed_seq_params,\n            )\n            loss_mask, num_tokens = roll_tensor(\n                loss_mask,\n                shifts=-1,\n                dims=-1,\n                cp_group=self.cp_group,\n                packed_seq_params=packed_seq_params,\n            )\n\n            # Compute mtp loss without storing logits to save memory.\n            mtp_loss = self.compute_output_layer_and_language_model_loss(\n                hidden_states_list[mtp_layer_number + 1],\n                labels=mtp_labels,\n                weight=self.shared_embedding_or_output_weight(),\n                sequence_parallel_enabled=self.output_layer.sequence_parallel,\n                column_parallel_linear=self.output_layer,\n                col_linear_kwargs={\n                    \"weight\": output_weight,\n                    \"runtime_gather_output\": runtime_gather_output,\n                },\n            )\n\n            mtp_loss = loss_mask * mtp_loss\n            if self.training:\n                # TODO(shifangx): remove the use of parallel_state here\n                # after moving loss logging to loss_func in pretrain_gpt.py\n                MTPLossLoggingHelper.save_loss_to_tracker(\n                    torch.sum(mtp_loss) / num_tokens,\n                    mtp_layer_number,\n                    self.config.mtp_num_layers,\n                    avg_group=parallel_state.get_data_parallel_group(with_context_parallel=True),\n                )\n            mtp_loss_scale = self.config.mtp_loss_scaling_factor / self.config.mtp_num_layers\n            if self.config.calculate_per_token_loss:\n                hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss)\n            else:\n                hidden_states = MTPLossAutoScaler.apply(hidden_states, mtp_loss_scale * mtp_loss / num_tokens)\n\n    logits, _ = self.output_layer(hidden_states, weight=output_weight, runtime_gather_output=runtime_gather_output)\n    # [s b h] => [b s h]\n    return logits.transpose(0, 1).contiguous()\n\n\ndef patch_mtp_layer_get_embeddings(model: torch.nn.Module):\n    \"\"\"Patch the _get_embeddings method of MultiTokenPredictionLayer\"\"\"\n    from megatron.core.models.gpt.gpt_model import GPTModel\n    from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer\n\n    # Unwrap each model in the actor_module to get the actual GPTModel\n    model = _get_patching_model(model)\n    # Collect all MultiTokenPredictionLayer instances\n    target_layers = []\n\n    if isinstance(model, GPTModel):\n        # Check if GPTModel has MTP and find the layers\n        if hasattr(model, \"mtp\") and hasattr(model.mtp, \"layers\"):\n            for layer in model.mtp.layers:\n                if isinstance(layer, MultiTokenPredictionLayer):\n                    target_layers.append(layer)\n    elif hasattr(model, \"layers\"):\n        # Check if any layer in the model is MultiTokenPredictionLayer\n        for layer in model.layers:\n            if isinstance(layer, MultiTokenPredictionLayer):\n                target_layers.append(layer)\n\n    if target_layers:\n        for layer in target_layers:\n            layer._get_embeddings_backup = layer._get_embeddings\n            layer._get_embeddings = _patched_get_embeddings_for_detach.__get__(layer, layer.__class__)\n        print(f\"Found and patched {len(target_layers)} MTP layer(s) in any of the actor modules\")\n        return True\n    else:\n        print(\"No MTP layers found to patch in any of the actor modules\")\n        return False\n\n\ndef unpatch_mtp_layer_get_embeddings(model: torch.nn.Module):\n    \"\"\"Unpatch the _get_embeddings method of MultiTokenPredictionLayer\"\"\"\n    from megatron.core.models.gpt.gpt_model import GPTModel\n    from megatron.core.transformer.multi_token_prediction import MultiTokenPredictionLayer\n\n    # Unwrap each model in the actor_module to get the actual GPTModel\n    model = _get_patching_model(model)\n\n    # Collect all MultiTokenPredictionLayer instances\n    target_layers = []\n\n    if isinstance(model, GPTModel):\n        # Check if GPTModel has MTP and find the layers\n        if hasattr(model, \"mtp\") and hasattr(model.mtp, \"layers\"):\n            for layer in model.mtp.layers:\n                if isinstance(layer, MultiTokenPredictionLayer):\n                    target_layers.append(layer)\n    elif hasattr(model, \"layers\"):\n        # Check if any layer in the model is MultiTokenPredictionLayer\n        for layer in model.layers:\n            if isinstance(layer, MultiTokenPredictionLayer):\n                target_layers.append(layer)\n\n    unpatched_count = 0\n    for layer in target_layers:\n        if hasattr(layer, \"_get_embeddings_backup\"):\n            layer._get_embeddings = layer._get_embeddings_backup\n            delattr(layer, \"_get_embeddings_backup\")\n            unpatched_count += 1\n\n    if unpatched_count > 0:\n        print(f\"Unpatched {unpatched_count} MTP layer(s)\")\n        return True\n    return False\n\n\ndef _patched_get_embeddings_for_detach(\n    self,\n    input_ids: torch.Tensor,\n    position_ids: torch.Tensor,\n    embedding: Callable,\n    hidden_states: torch.Tensor,\n    packed_seq_params=None,\n):\n    \"\"\"\n    Patched version of _get_embeddings method for MultiTokenPredictionLayer.\n\n    This is a modified version that you can customize according to your needs.\n    The original implementation is preserved below with modifications.\n    \"\"\"\n\n    # You can modify the logic here as needed\n    # For example, you could:\n    # - Change the shift amount in roll_tensor\n    # - Apply custom transformations to input_ids or position_ids\n    # - Add debugging information\n    # - Modify the embedding computation\n\n    # Original logic with custom modifications\n    from megatron.core.transformer.multi_token_prediction import roll_tensor\n    from megatron.core.utils import make_viewless_tensor\n\n    # Calc logits for the current Multi-Token Prediction (MTP) layers.\n    input_ids, _ = roll_tensor(\n        input_ids,\n        shifts=-1,  # You can modify this shift value\n        dims=-1,\n        cp_group=self.cp_group,\n        packed_seq_params=packed_seq_params,\n    )\n    position_ids, _ = roll_tensor(\n        position_ids,\n        shifts=-1,  # You can modify this shift value\n        dims=-1,\n        cp_group=self.cp_group,\n        packed_seq_params=packed_seq_params,\n    )\n\n    # embedding computation - you can modify this part\n    decoder_input = embedding(input_ids=input_ids, position_ids=position_ids)\n\n    # Apply custom transformations if needed\n    # For example: decoder_input = some_custom_function(decoder_input)\n\n    hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)\n\n    # detach decoder_input and hidden_states\n    decoder_input = decoder_input.detach()\n    hidden_states = hidden_states.detach()\n\n    return input_ids, position_ids, decoder_input, hidden_states\n"
  },
  {
    "path": "verl/models/mcore/patch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# there is some bug in mcore 0.12, so we need to patch it\n# 1. `get_query_key_value_tensors` in `multi_latent_attention.py` works wrong when packed_seq_params is not None\n\n\ndef apply_patch():\n    import megatron.core\n    import torch\n    import torch.nn.functional as F\n    from megatron.core import parallel_state, tensor_parallel\n    from megatron.core.transformer.multi_latent_attention import (\n        MLASelfAttention,\n        MultiLatentAttention,\n        apply_rotary_pos_emb,\n        deprecate_inference_params,\n        gather_from_sequence_parallel_region,\n        gather_from_tensor_model_parallel_region,\n        scatter_to_sequence_parallel_region,\n    )\n    from packaging import version\n\n    mcore_ge_013 = version.parse(megatron.core.__version__) >= version.parse(\"0.13.0\")\n\n    def patch_get_query_key_value_tensors(\n        self,\n        hidden_states,\n        key_value_states=None,\n        position_ids=None,\n        packed_seq_params=None,\n        inference_context=None,\n        *,\n        inference_params=None,\n    ):\n        \"\"\"\n        Derives `query`, `key` and `value` tensors from `hidden_states`.\n        \"\"\"\n        # s = sequence length, b = batch size, h = hidden size, n = num attention heads\n        # Attention heads [s, b, n*h]\n        assert hidden_states.ndim == 3, f\"hidden_states should be 3D, [s, b, n*h], got {hidden_states.ndim}D\"\n\n        inference_context = deprecate_inference_params(inference_context, inference_params)\n\n        # =========================================\n        # Prepare RoPE and seqlen related params\n        # =========================================\n        rotary_seq_len = self.rotary_pos_emb.get_rotary_seq_len(\n            inference_context, None, hidden_states, self.config, packed_seq_params\n        )\n\n        # rotary_pos_emb:[s, b, 1, 64]\n        mscale = 1.0\n        if self.config.rope_type == \"rope\":\n            packed_seq = packed_seq_params is not None and packed_seq_params.qkv_format == \"thd\"\n            try:\n                # In case of TypeError: RotaryEmbedding.forward() got an unexpected keyword argument 'packed_seq'\n                rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len, packed_seq=packed_seq)\n            except TypeError:\n                rotary_pos_emb = self.rotary_pos_emb(rotary_seq_len)\n        else:\n            rotary_pos_emb, mscale = self.rotary_pos_emb(rotary_seq_len)\n\n        # =========================================\n        # QKV down projection and layernorm\n        # =========================================\n        if self.config.q_lora_rank is not None:\n            # if linear_q_down_proj is ColumnParallelLinear:\n            #     q_compressed: [s, b, q_lora_rank / TP]\n            # elif linear_q_down_proj is Linear:\n            #     q_compressed: [s / TP, b, q_lora_rank]\n            q_compressed, _ = self.linear_q_down_proj(hidden_states)\n\n            # When output is sharded (ColumnParallelLinear), two things are needed to be\n            # identical to a normal Linear.\n            #   1. Manually gather output to restore output dim q_lora_rank;\n            #   2. Scatter sequence back to s / TP if sequence-parallel since it was\n            #      gathered by ColumnParallelLinear.\n            if q_compressed.size(-1) != self.config.q_lora_rank:\n                q_compressed = gather_from_tensor_model_parallel_region(q_compressed)\n                if self.config.sequence_parallel:\n                    q_compressed = scatter_to_sequence_parallel_region(q_compressed)\n\n            q_compressed = self.q_layernorm(q_compressed)\n        else:\n            q_compressed = hidden_states\n\n        # if linear_kv_down_proj is ColumnParallelLinear:\n        #     kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim) / TP]\n        # elif linear_kv_down_proj is Linear:\n        #     kv_combined: [s / TP, b, (kv_lora_rank + qk_pos_emb_head_dim)]\n        kv_combined, _ = self.linear_kv_down_proj(hidden_states)\n        if kv_combined.size(-1) != self.config.kv_lora_rank + self.config.qk_pos_emb_head_dim:\n            # kv_combined: [s, b, (kv_lora_rank + qk_pos_emb_head_dim)]\n            kv_combined = gather_from_tensor_model_parallel_region(kv_combined)\n            # kv_compressed:[s, b, kv_lora_rank], k_pos_emb: [s, b, qk_pos_emb_head_dim]\n            kv_compressed, k_pos_emb = torch.split(\n                kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1\n            )\n            if self.config.sequence_parallel:\n                # kv_compressed:[s / TP, b, kv_lora_rank]\n                kv_compressed = scatter_to_sequence_parallel_region(kv_compressed)\n        else:\n            # kv_compressed:[s / TP, b, kv_lora_rank], k_pos_emb: [s / TP, b, qk_pos_emb_head_dim]\n            kv_compressed, k_pos_emb = torch.split(\n                kv_combined, [self.config.kv_lora_rank, self.config.qk_pos_emb_head_dim], dim=-1\n            )\n            if parallel_state.get_tensor_model_parallel_world_size() > 1:\n                # k_pos_emb: [s, b, qk_pos_emb_head_dim]\n                k_pos_emb = gather_from_sequence_parallel_region(k_pos_emb)\n\n        kv_compressed = self.kv_layernorm(kv_compressed)\n\n        # =========================================\n        # QKV up projection and RoPE apply\n        # =========================================\n        def qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb):\n            if self.config.q_lora_rank is not None:\n                q, _ = self.linear_q_up_proj(q_compressed)\n            else:\n                # hidden_states:[s, b, 2048], q: [s, b, n * 192]\n                q, _ = self.linear_q_proj(q_compressed)\n\n            q_len, bsz, _ = q.size()\n\n            # q: [s, b, n, 192]\n            q = q.view(q_len, bsz, self.num_attention_heads_per_partition, self.q_head_dim)\n\n            # kv: [s, b, 2048]\n            kv, _ = self.linear_kv_up_proj(kv_compressed)\n\n            # kv: [s, b, n, 256]\n            kv = kv.view(\n                q_len,\n                bsz,\n                self.num_attention_heads_per_partition,\n                self.config.qk_head_dim + self.config.v_head_dim,\n            )\n\n            cp_size = parallel_state.get_context_parallel_world_size()\n            if inference_context is not None:\n                # add offset to the sequence start for inference\n                sequence_start = inference_context.sequence_len_offset\n                sequence_end = sequence_start + q_len\n                rotary_pos_emb = rotary_pos_emb[sequence_start:sequence_end]\n            elif packed_seq_params is None or cp_size == 1:\n                # Shorten rotary_pos_emb to the sequence length when inference_params\n                # is not provided. This makes sure we can run forward directly with\n                # any sequence length. During training, the sequence length is always\n                # the full rotary_pos_emb length, except for sequence packing + CP.\n                # When sequence packing and context parallel are both enabled, the\n                # position embedding will not split rotary_pos_emb, so it may exceed\n                # the sequence length on this CP rank, but we need the full rotary_pos_emb\n                # to cover the full sequence, so we do not shorten it here.\n                rotary_pos_emb = rotary_pos_emb[0:q_len]\n\n            # [s, b, 64] -> [s, b, 1, 64]\n            k_pos_emb = torch.unsqueeze(k_pos_emb, 2)\n\n            # q: [s, b, n, 128], q_pos_emb: [s, b, n, 64]\n            q_no_pe, q_pos_emb = torch.split(q, [self.config.qk_head_dim, self.config.qk_pos_emb_head_dim], dim=-1)\n\n            # k_no_pe: [s, b, n, 128], value: [s, b, n, 128]\n            k_no_pe, value = torch.split(kv, [self.config.qk_head_dim, self.config.v_head_dim], dim=-1)\n\n            if packed_seq_params is not None:\n                cu_seqlens_q = packed_seq_params.cu_seqlens_q\n                cu_seqlens_kv = packed_seq_params.cu_seqlens_kv\n                q_pos_emb = q_pos_emb.squeeze(1)\n                k_pos_emb = k_pos_emb.squeeze(1)\n                q_no_pe = q_no_pe.squeeze(1)\n                k_no_pe = k_no_pe.squeeze(1)\n                value = value.squeeze(1)\n            else:\n                cu_seqlens_q = cu_seqlens_kv = None\n\n            # q_pos_emb: [s, b, n, 64], k_pos_emb:[s, b, 1, 64]\n            q_pos_emb = apply_rotary_pos_emb(\n                q_pos_emb,\n                rotary_pos_emb,\n                config=self.config,\n                cu_seqlens=cu_seqlens_q,\n                mscale=mscale,\n            )\n            k_pos_emb = apply_rotary_pos_emb(\n                k_pos_emb,\n                rotary_pos_emb,\n                config=self.config,\n                cu_seqlens=cu_seqlens_kv,\n                mscale=mscale,\n            )\n\n            # query: [s, b, n, 192]\n            query = torch.cat([q_no_pe, q_pos_emb], dim=-1)\n            if packed_seq_params is not None:\n                k_pos_emb = k_pos_emb.expand(-1, self.num_attention_heads_per_partition, -1)\n                key = torch.cat([k_no_pe, k_pos_emb], dim=-1)\n            else:\n                # key: [s, b, n, 192]\n                k_pos_emb = k_pos_emb.expand(-1, -1, self.num_attention_heads_per_partition, -1)\n                key = torch.cat([k_no_pe, k_pos_emb], dim=-1)\n\n            query = query.contiguous()\n            key = key.contiguous()\n            value = value.contiguous()\n            return query, key, value\n\n        if self.recompute_up_proj:\n            self.qkv_up_checkpoint = tensor_parallel.CheckpointWithoutOutput()\n            query, key, value = self.qkv_up_checkpoint.checkpoint(\n                qkv_up_proj_and_rope_apply, q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb\n            )\n        else:\n            query, key, value = qkv_up_proj_and_rope_apply(q_compressed, kv_compressed, k_pos_emb, rotary_pos_emb)\n\n        return query, key, value\n\n    def patch_forward(\n        self,\n        hidden_states,\n        attention_mask,\n        key_value_states=None,\n        inference_context=None,\n        rotary_pos_emb=None,\n        rotary_pos_cos=None,\n        rotary_pos_sin=None,\n        attention_bias=None,\n        packed_seq_params=None,\n        position_ids=None,\n        sequence_len_offset=None,\n        *,\n        inference_params=None,\n        **kwargs,\n    ):\n        \"\"\"Forward pass for multi-latent attention\"\"\"\n        assert attention_bias is None, \"Attention bias should not be passed into MLA.\"\n        assert rotary_pos_cos is None and rotary_pos_sin is None, \"MLA does not support Flash Decoding\"\n\n        # hidden_states: [sq, b, h]\n\n        inference_context = deprecate_inference_params(inference_context, inference_params)\n\n        # =====================\n        # Query, Key, and Value\n        # =====================\n        # Get the query, key and value tensors based on the type of attention -\n        # self or cross attn.\n        # query: [96, 1, 16, 128], key:[96, 1, 16, 128], value:[96, 1, 16, 128]\n        qkv = self.get_query_key_value_tensors(\n            hidden_states,\n            key_value_states,\n            position_ids,\n            packed_seq_params,\n            inference_context=inference_context,\n        )\n        query, key, value = qkv[:3]\n        q_compressed = None\n        # kv_compressed = None\n        if len(qkv) > 4:\n            q_compressed = qkv[3]\n            # kv_compressed = qkv[4]\n\n        # ===================================================\n        # Adjust key, value for inference\n        # ===================================================\n        # rotary_pos_emb = None\n        if mcore_ge_013:\n            query, key, value, _, attn_mask_type, _ = self._adjust_key_value_for_inference(\n                inference_context, query, key, value, rotary_pos_emb=None\n            )\n        else:\n            query, key, value, _, attn_mask_type = self._adjust_key_value_for_inference(\n                inference_context, query, key, value, rotary_pos_emb=None\n            )\n\n        # TODO: Currently, TE can only accept contiguous tensors for MLA\n        query = query.contiguous()\n        key = key.contiguous()\n        value = value.contiguous()\n\n        # ==================================\n        # core attention computation\n        # ==================================\n        # Need corresponding TE change\n        non_dsa_thd_qkv_format = (\n            packed_seq_params\n            and packed_seq_params.qkv_format == \"thd\"\n            and getattr(self.config, \"experimental_attention_variant\", None) is None\n        )\n        v_dim = value.shape[-1]\n        if non_dsa_thd_qkv_format and query.shape[-1] != v_dim:\n            value = F.pad(value, [0, query.shape[-1] - v_dim])\n            self.core_attention.hidden_size_per_attention_head_v = value.shape[-1]\n        if self.checkpoint_core_attention and self.training:\n            core_attn_out = self._checkpointed_attention_forward(\n                query, key, value, attention_mask, packed_seq_params=packed_seq_params\n            )\n        else:\n            extra_kwargs = {}\n            if getattr(self.config, \"experimental_attention_variant\", None) == \"dsa\":\n                # For dsa we need to pass in the original hidden states and the compressed\n                # query representation.\n                extra_kwargs[\"x\"] = hidden_states\n                extra_kwargs[\"qr\"] = q_compressed\n            core_attn_out = self.core_attention(\n                query,\n                key,\n                value,\n                attention_mask,\n                packed_seq_params=packed_seq_params,\n                attn_mask_type=attn_mask_type,\n                **extra_kwargs,\n            )\n        if non_dsa_thd_qkv_format:\n            if core_attn_out.ndim == 2:\n                core_attn_out = core_attn_out.reshape(*core_attn_out.shape[:-1], -1, value.shape[-1])\n            if query.shape[-1] != v_dim:\n                core_attn_out = core_attn_out[..., :v_dim]\n            # reshape to same output shape as unpacked case\n            # (t, np, hn) -> (t, b=1, h=np*hn)\n            # t is the pack size = sum (sq_i)\n            # note that batch is a dummy dimension in the packed case\n            core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)\n\n        if self.recompute_up_proj:\n            assert self.qkv_up_checkpoint is not None\n            self.qkv_up_checkpoint.discard_output_and_register_recompute(core_attn_out)\n            self.qkv_up_checkpoint = None\n\n        # =================\n        # Output. [sq, b, h]\n        # =================\n        output, bias = self.linear_proj(core_attn_out)\n\n        return output, bias\n\n    # This patch targets mcore 0.12 MLA behavior only.\n    # For newer mcore, upstream MLA already has packed-seq + CP handling and\n    # overriding it with the legacy implementation can break RoPE shapes.\n    if not mcore_ge_013:\n        MLASelfAttention.get_query_key_value_tensors = patch_get_query_key_value_tensors\n\n    MultiLatentAttention.forward = patch_forward\n\n\ndef apply_patch_mbridge():\n    try:\n        from megatron.core.utils import get_tensor_model_parallel_group_if_none\n    except ImportError:\n        import warnings\n\n        import megatron.core.utils\n        import torch\n        from megatron.core import parallel_state\n\n        def get_tensor_model_parallel_group_if_none(tp_group, is_expert=False, check_initialized=True):\n            \"\"\"Issue a deprecation warning if tp_group is None and return the default tp group.\"\"\"\n            if not torch.distributed.is_initialized():\n                return None\n            if tp_group is None:\n                if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:\n                    warnings.warn(\n                        \"Warning: tp_group is None, using default tp group. Passing tp_group will be mandatory soon\",\n                        DeprecationWarning,\n                        stacklevel=2,\n                    )\n                if is_expert:\n                    tp_group = parallel_state.get_expert_tensor_parallel_group(check_initialized=check_initialized)\n                else:\n                    tp_group = parallel_state.get_tensor_model_parallel_group(check_initialized=check_initialized)\n            return tp_group\n\n        megatron.core.utils.get_tensor_model_parallel_group_if_none = get_tensor_model_parallel_group_if_none\n\n\ndef apply_patch_megatron_v012_with_torch_v28():\n    # Error due to missing serialization_format in _write_item of megatron v012;\n    # resolved by using megatron v013's implementation.\n    import inspect\n    import logging\n    import os\n    from pathlib import Path\n\n    import megatron.core\n    import torch\n    from megatron.core.dist_checkpointing.strategies.async_utils import _disable_gc\n    from megatron.core.dist_checkpointing.strategies.filesystem_async import _process_memory\n    from packaging import version\n    from torch import multiprocessing as mp\n    from torch.distributed.checkpoint.filesystem import _write_item\n\n    if (\n        version.parse(torch.__version__).base_version != \"2.8.0\"\n        or version.parse(megatron.core.__version__).base_version != \"0.12.1\"\n    ):\n        return\n\n    WriteBucket = tuple[Path, str, tuple[list, list]]\n\n    @staticmethod\n    @_disable_gc()\n    def write_preloaded_data_patch(\n        transform_list,\n        local_proc_idx: int,\n        write_bucket: WriteBucket,\n        results_queue: mp.SimpleQueue,\n        count_queue: mp.JoinableQueue,\n        use_fsync: bool,\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Performs actual data saving to storage.\n\n        Args:\n            local_proc_idx (int): index of a local process that performs writing\n            write_bucket (WriteBucket): data to write to storage\n            results_queue (mp.Queue): queue to return the write results\n                to the proxy checkpoint process.\n            count_queue (mp.JoinableQueue): queue to marks worker task as completed\n            use_fsync (bool): if True, calls os.fsync at the end of saving\n\n        Returns: None, the write result are put into the `queue`\n        \"\"\"\n        logger = logging.getLogger(__name__)\n        logger.debug(f\"{local_proc_idx} started\")\n        mem_before = _process_memory()\n        use_msc = kwargs.get(\"use_msc\", False)\n        local_results = []\n        try:\n            file_name, storage_key, (bytes_data, tensor_data) = write_bucket\n            extra_kwargs = {}\n            if \"serialization_format\" in inspect.signature(_write_item).parameters:\n                from torch.distributed.checkpoint.filesystem import SerializationFormat\n\n                extra_kwargs[\"serialization_format\"] = SerializationFormat.TORCH_SAVE\n            if use_msc:\n                import multistorageclient as msc\n\n                open_file = msc.open\n            else:\n                open_file = open\n            with open_file(file_name, \"wb\") as stream:\n                for write_item, data in bytes_data:\n                    local_results.append(\n                        _write_item(*transform_list, stream, data, write_item, storage_key, **extra_kwargs)\n                    )\n\n                for write_item, tensor in tensor_data:\n                    assert tensor.is_cpu\n                    local_results.append(\n                        _write_item(*transform_list, stream, tensor, write_item, storage_key, **extra_kwargs)\n                    )\n\n                if use_fsync:\n                    if use_msc:\n                        stream.fsync()\n                    else:\n                        os.fsync(stream.fileno())\n            local_output = (local_proc_idx, local_results)\n        except Exception as e:\n            logger.debug(f\"{local_proc_idx} failed\")\n            local_output = (local_proc_idx, e)  # type: ignore[assignment]\n\n        results_queue.put(local_output)\n        # Signal this process is done.\n        count_queue.get()\n        count_queue.task_done()\n\n        mem_after = _process_memory()\n        logger.debug(f\"{local_proc_idx} consumed: {mem_after - mem_before}, before: {mem_before}, after: {mem_after}\")\n\n    from megatron.core.dist_checkpointing.strategies.filesystem_async import FileSystemWriterAsync\n\n    FileSystemWriterAsync.write_preloaded_data = write_preloaded_data_patch\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom .model import Qwen2_5VLModel\nfrom .vision_config import get_vision_model_config, get_vision_projection_config\n\n__all__ = [\"Qwen2_5VLModel\", \"get_vision_model_config\", \"get_vision_projection_config\"]\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/attention.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom megatron.core.transformer.attention import *\n\nfrom .rope_utils import apply_rotary_pos_emb_absolute\n\n\nclass Qwen2_5VLSelfAttention(SelfAttention):\n    \"\"\"\n    Overrides the SelfAttention class, the difference is that qwen2_5_vl uses apply_rotary_pos_emb_absolute\n    instead of apply_rotary_pos_emb\n    \"\"\"\n\n    def forward(\n        self,\n        hidden_states: Tensor,\n        attention_mask: Tensor,\n        key_value_states: Optional[Tensor] = None,\n        inference_context: Optional[BaseInferenceContext] = None,\n        rotary_pos_emb: Optional[Union[Tensor, Tuple[Tensor, Tensor]]] = None,\n        rotary_pos_cos: Optional[Tensor] = None,\n        rotary_pos_sin: Optional[Tensor] = None,\n        attention_bias: Optional[Tensor] = None,\n        packed_seq_params: Optional[PackedSeqParams] = None,\n        sequence_len_offset: Optional[int] = None,\n        *,\n        inference_params: Optional[BaseInferenceContext] = None,\n        rotary_pos_cos_sin: Optional[Tensor] = None,\n    ) -> Tuple[Tensor, Tensor]:\n        \"\"\"\n        Perform a forward pass through the attention module.\n\n        Args:\n            hidden_states (Tensor): Hidden states.\n            attention_mask (Tensor): Attention mask.\n            key_value_states (Optional[Tensor]): Key/value states (for cross attention).\n            inference_context (Optional[BaseInferenceContext]): Inference context that manages\n                KV cache.\n            rotary_pos_emb (Optional[Union[Tensor, Tuple[Tensor, Tensor]]]): Rotary\n                embedding tensor(s).\n            rotary_pos_cos (Optional[Tensor]): Rotary embedding cosine.\n            rotary_pos_sin (Optional[Tensor]): Rotary embedding sine.\n            attention_bias (Optional[Tensor]): Attention bias.\n            packed_seq_params (Optional[PackedSeqparams]): Parameters used for THD format.\n            sequence_len_offset (Optional[int]): Sequence length offset used for\n                inference CUDA graphs.\n\n        Return:\n            (Tuple[Tensor, Tensor]) Attention output and bias.\n\n        \"\"\"\n\n        inference_context = deprecate_inference_params(inference_context, inference_params)\n\n        if inference_context and inference_context.is_dynamic_batching():\n            assert flash_decode_and_prefill_kernel is not None, (\n                \"Internal use only: install package `nvidia_chunked_flash_attn`.\"\n            )\n\n        # hidden_states: [sq, b, h]\n        if self.config.flash_decode and not self.training and inference_context is not None:\n            rotary_pos_emb = None\n        else:\n            assert rotary_pos_cos is None and rotary_pos_sin is None\n\n        # For self attention we just duplicate the rotary_pos_emb if it isn't already\n        if rotary_pos_emb is not None and not isinstance(rotary_pos_emb, tuple):\n            rotary_pos_emb = (rotary_pos_emb,) * 2\n\n        # =====================\n        # Query, Key, and Value\n        # =====================\n        # Get the query, key and value tensors based on the type of attention -\n        # self or cross attn.\n        query, key, value = self.get_query_key_value_tensors(hidden_states, key_value_states)\n\n        # ===================================================\n        # Adjust key, value, and rotary_pos_emb for inference\n        # ===================================================\n\n        # This branch only runs in the decode phase of flash decoding and returns after the linear\n        # projection. This conditional is not used in the prefill phase or non-flash-decoding cases.\n        if (\n            self.config.flash_decode\n            and inference_context is not None\n            and inference_context.is_decode_only()\n            and not self.training\n            and rotary_pos_cos is not None\n        ):\n            assert self.layer_number in inference_context.key_value_memory_dict\n            assert inference_context.sequence_len_offset is not None\n            inference_key_memory, inference_value_memory = inference_context.key_value_memory_dict[self.layer_number]\n            output = self.flash_decode(\n                sequence_len_offset=sequence_len_offset,\n                query_layer=query,\n                key_layer=key,\n                value_layer=value,\n                inference_key_memory=inference_key_memory,\n                inference_value_memory=inference_value_memory,\n                rotary_cos=rotary_pos_cos,\n                rotary_sin=rotary_pos_sin,\n            )\n            out = output.transpose(0, 1).contiguous()\n            context_layer = out.view(out.size(0), out.size(1), -1)\n            output, bias = self.linear_proj(context_layer)\n            return output, bias\n\n        # Use latest mcore 0.13 API and forward-compatible with previous versions.\n        outputs = self._adjust_key_value_for_inference(\n            inference_context,\n            query,\n            key,\n            value,\n            rotary_pos_emb,\n            rotary_pos_cos,\n            rotary_pos_sin,\n            sequence_len_offset,\n        )\n\n        query, key, value, rotary_pos_emb, attn_mask_type = outputs[:5]\n\n        if packed_seq_params is not None:\n            query = query.squeeze(1)\n            key = key.squeeze(1)\n            value = value.squeeze(1)\n\n        # ================================================\n        # relative positional embedding (rotary embedding)\n        # ================================================\n        if rotary_pos_emb is not None and not self.config.flash_decode:\n            q_pos_emb, k_pos_emb = rotary_pos_emb\n\n            if packed_seq_params is not None:\n                if packed_seq_params.cu_seqlens_q_padded is not None:\n                    cu_seqlens_q = packed_seq_params.cu_seqlens_q_padded\n                else:\n                    cu_seqlens_q = packed_seq_params.cu_seqlens_q\n                if packed_seq_params.cu_seqlens_kv_padded is not None:\n                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv_padded\n                else:\n                    cu_seqlens_kv = packed_seq_params.cu_seqlens_kv\n            else:\n                cu_seqlens_q = cu_seqlens_kv = None\n\n            if q_pos_emb is not None:\n                # TODO VIJAY: simplify\n                if inference_context is None or inference_context.is_static_batching():\n                    query = apply_rotary_pos_emb_absolute(query, q_pos_emb, config=self.config, cu_seqlens=cu_seqlens_q)\n                else:\n                    query = inference_context.apply_rotary_emb_query(query, q_pos_emb, self.config, cu_seqlens_q)\n            if k_pos_emb is not None:\n                key = apply_rotary_pos_emb_absolute(key, k_pos_emb, config=self.config, cu_seqlens=cu_seqlens_kv)\n\n            # TODO, can apply positional embedding to value_layer so it has\n            # absolute positional embedding.\n            # otherwise, only relative positional embedding takes effect\n            # value_layer = apply_rotary_pos_emb(value_layer, k_pos_emb)\n\n        # ==================================\n        # core attention computation\n        # ==================================\n\n        if self.checkpoint_core_attention and self.training:\n            core_attn_out = self._checkpointed_attention_forward(\n                query,\n                key,\n                value,\n                attention_mask,\n                attn_mask_type=attn_mask_type,\n                attention_bias=attention_bias,\n                packed_seq_params=packed_seq_params,\n            )\n        else:\n            if inference_context is None or inference_context.is_static_batching():\n                # Static batching attention kernel.\n                core_attn_out = self.core_attention(\n                    query,\n                    key,\n                    value,\n                    attention_mask,\n                    attn_mask_type=attn_mask_type,\n                    attention_bias=attention_bias,\n                    packed_seq_params=packed_seq_params,\n                )\n\n            else:\n                # Dynamic batching attention kernel.\n                q, k, v = (query, key, value)\n                cu_query_lengths, max_seqlen_q = inference_context.cu_query_lengths()\n                cu_kv_lengths, max_seqlen_k = inference_context.cu_kv_lengths()\n\n                core_attn_out = self.flash_decode_and_prefill(\n                    q, k, v, max_seqlen_q, max_seqlen_k, cu_query_lengths, cu_kv_lengths\n                )\n                core_attn_out = core_attn_out.squeeze(0).unsqueeze(1)\n                core_attn_out = rearrange(core_attn_out, \"s b h d -> s b (h d)\")\n\n        if packed_seq_params is not None and packed_seq_params.qkv_format == \"thd\":\n            # reshape to same output shape as unpacked case\n            # (t, np, hn) -> (t, b=1, h=np*hn)\n            # t is the pack size = sum (sq_i)\n            # note that batch is a dummy dimension in the packed case\n            core_attn_out = core_attn_out.reshape(core_attn_out.size(0), 1, -1)\n\n        # =================\n        # Output. [sq, b, h]\n        # =================\n\n        output, bias = self.linear_proj(core_attn_out)\n\n        return output, bias\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/model.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\n\nimport torch\nfrom megatron.core import InferenceParams, mpu, tensor_parallel\nfrom megatron.core.models.gpt.gpt_model import GPTModel\n\n# from .transformer_config import Qwen2VLTransformerConfig\nfrom megatron.core.packed_seq_params import PackedSeqParams\nfrom megatron.core.transformer import MegatronModule\nfrom megatron.core.transformer.spec_utils import ModuleSpec\nfrom megatron.core.transformer.transformer_config import TransformerConfig\n\nfrom verl.models.mcore.util import preprocess_packed_seqs\n\nfrom .attention import Qwen2_5VLSelfAttention\nfrom .vision_model import Qwen2_5VisionModel\n\n\n# Note: This is under development and may be missing features.\nclass Qwen2_5VLModel(MegatronModule):\n    \"\"\"Qwen2.5VL multi-modal model.\n\n    Args:\n        language_transformer_config (TransformerConfig): Transformer config for the language model.\n        language_transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers of the\n            language model.\n        language_vocab_size (int): Language model vocabulary size.\n        language_max_sequence_length (int): Language model maximum sequence length. This is used for\n            positional embedding.\n        vision_transformer_config (TransformerConfig): Transformer config for the vision model.\n        vision_transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers of the\n            vision model.\n        vision_projection_config (TransformerConfig): Config for the projection from vision model outputs to\n            language model inputs.\n        vision_projection_layer_spec (ModuleSpec): Specifies the module to use for the vision\n            projection.\n        vision_projection_type (str): Type of the vision projection to use. Default is a 2-layer MLP.\n        parallel_output (bool): Do not gather the outputs, keep them split across tensor parallel ranks. This\n            is typically True for training and False for inference.\n        language_rotary_percent (float): Percent of rotary dimension to use for rotary position embeddings\n            in the language model. Defaults to 1.0.\n        pre_process (bool): Include the embedding layer in the gpt decoder (used with pipeline parallelism).\n            Defaults to True.\n        post_process (bool): Include an output layer and a layernorm in the gpt decoder (used with pipeline\n            parallelism). Defaults to True.\n        add_encoder (bool): Construct the encoder module (used with pipeline parallelism). Defaults to True.\n            When we use pipelining, the encoder\n            will live on only a subset of the pipeline stages (specifically, only the first stage).\n        add_decoder (bool): Construct the decoder module (used with pipeline parallelism). Defaults to True.\n            When we use pipelining, the decoder\n            will live on only a subset of the pipeline stages (specifically, every stage after the first one).\n        img_h (int): The height of each image that the ViT will see.\n        img_w (int): The width of each image that the ViT will see.\n        patch_dim (int): The size of each patch side.\n        img_embedding_idx (int): Index in the language_embeddings tensor where image_embeddings should be\n            inserted. Defaults to 0.\n    \"\"\"\n\n    def __init__(\n        self,\n        language_transformer_config: TransformerConfig,\n        language_transformer_layer_spec: ModuleSpec,\n        language_vocab_size: int,\n        language_max_sequence_length: int,\n        vision_transformer_config: TransformerConfig,\n        vision_transformer_layer_spec: ModuleSpec,\n        vision_projection_config: TransformerConfig,\n        vision_projection_layer_spec: ModuleSpec,\n        vision_projection_type: str = \"mlp\",\n        parallel_output: bool = True,\n        language_rotary_percent: float = 1.0,\n        pre_process: bool = True,\n        post_process: bool = True,\n        add_encoder: bool = True,\n        add_decoder: bool = True,\n        language_rotary_base: int = 10000,\n        fp16_lm_cross_entropy: bool = False,\n        language_share_embeddings_and_output_weights: bool = False,\n        image_token_id: int = 151655,\n        video_token_id: int = 151656,\n    ) -> None:\n        super().__init__(config=language_transformer_config)\n\n        # patch self_attention to use qwen2_5_vl attention\n        vision_transformer_layer_spec.submodules.self_attention.module = Qwen2_5VLSelfAttention\n        for layer_spec in language_transformer_layer_spec.layer_specs:\n            layer_spec.submodules.self_attention.module = Qwen2_5VLSelfAttention\n\n        logging.getLogger(__name__).warning(\"Qwen2VL model is under development and may be missing features.\")\n\n        self.pre_process = pre_process\n        self.post_process = post_process\n        self.add_encoder = add_encoder\n        self.add_decoder = add_decoder\n\n        self.encoder_hidden_state = None\n        self.vision_model = None\n        self.vision_projection = None\n        self.language_model = None\n        self.image_token_id = image_token_id\n        self.video_token_id = video_token_id\n\n        self.square_merge_size = vision_projection_config.ffn_hidden_size // vision_transformer_config.hidden_size\n\n        # This attribute is needed to check if an all-reduce is required\n        # on the word embeddings inside `finalize_model_grads._allreduce_word_embedding_grads`.\n        self.share_embeddings_and_output_weights = False\n        if self.pre_process:\n            self.vision_model = Qwen2_5VisionModel(\n                vision_transformer_config,\n                vision_transformer_layer_spec,\n                vision_projection_config,\n                vision_projection_layer_spec,\n                projection_type=vision_projection_type,\n                pre_process=True,\n                post_process=True,\n            )\n\n        self.language_model = GPTModel(\n            config=language_transformer_config,\n            transformer_layer_spec=language_transformer_layer_spec,\n            vocab_size=language_vocab_size,\n            max_sequence_length=language_max_sequence_length,\n            parallel_output=parallel_output,\n            position_embedding_type=\"mrope\",\n            rotary_percent=language_rotary_percent,\n            pre_process=self.pre_process,\n            post_process=self.post_process,\n            rotary_base=language_rotary_base,\n            fp16_lm_cross_entropy=fp16_lm_cross_entropy,\n            share_embeddings_and_output_weights=language_share_embeddings_and_output_weights,\n            scatter_embedding_sequence_parallel=False,\n        )\n        assert mpu.get_context_parallel_world_size() <= 1, \"please use mbridge for qwen2_5_vl with context parallelism\"\n        self.share_embeddings_and_output_weights = self.language_model.share_embeddings_and_output_weights\n\n    def shared_embedding_or_output_weight(self):\n        \"\"\"This is a convenience method to surface the language model's word embeddings, which is\n        necessary for `finalize_model_grads._allreduce_word_embedding_grads`.\"\"\"\n        if self.add_decoder:\n            return self.language_model.shared_embedding_or_output_weight()\n        return None\n\n    def set_input_tensor(self, input_tensor) -> None:\n        # This is usually handled in schedules.py but some inference code still\n        # gives us non-lists or None\n        if not isinstance(input_tensor, list):\n            input_tensor = [input_tensor]\n        assert len(input_tensor) == 1, \"input_tensor should only be length 1 for Qwen2VL\"\n\n        if self.pre_process:\n            self.encoder_hidden_state = input_tensor[0]\n        else:\n            self.language_model.set_input_tensor(input_tensor[0])\n\n    def freeze(self, freeze_language_model: bool, freeze_vision_model: bool, freeze_vision_projection: bool):\n        \"\"\"Freeze model modules.\n\n        Make specific modules non-trainable by setting requires_grad to False for the module's parameters.\n\n        Args:\n            freeze_language_model (bool): Freeze the language model module.\n            freeze_vision_model (bool): Freeze the vision model module.\n            freeze_vision_projection (bool): Freeze the vision projection module.\n        \"\"\"\n        modules = []\n        if freeze_language_model and self.language_model is not None:\n            modules.append(self.language_model)\n        if freeze_vision_model and self.vision_model is not None:\n            modules.append(self.vision_model)\n        if freeze_vision_projection and self.vision_projection is not None:\n            modules.append(self.vision_projection)\n\n        for module in modules:\n            for param in module.parameters():\n                param.requires_grad = False\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: torch.Tensor,\n        attention_mask: torch.Tensor = None,\n        labels: torch.Tensor = None,\n        inference_params: InferenceParams = None,\n        packed_seq_params: PackedSeqParams = None,\n        extra_block_kwargs: dict = None,\n        pixel_values: torch.Tensor = None,\n        pixel_values_videos: torch.Tensor = None,\n        image_grid_thw: torch.Tensor = None,\n        video_grid_thw: torch.Tensor = None,\n        **kwargs,\n    ) -> torch.Tensor:\n        \"\"\"Forward function of the Qwen2VL model.\n        ### there is a workaround for supporting sequence packing with context parallelism\n        # cp split with sequence packing will make model lose vision token information, so we need to keep\n        # the original input_ids and pack them after vision embedding is calculated,\n        # cooporate with verl's models/mcore/model_forward.py\n        # pack the combined_embeddings to thd here, we check if packed_seq_params is None to determine if\n        #  we need to pack the combined_embeddings to thd\n        # this function needs the position_ids and attention_mask in BSHD format, no matter use packed_seq or not\n\n        Args:\n            image_data (torch.Tensor): input image of shape [total_thw_size, n_features].\n            input_ids (torch.Tensor): input text ids [batch, text_seq_len].\n            position_ids (torch.Tensor): input text position ids [batch, text_seq_len].\n            attention_mask (torch.Tensor): attention mask for the language model [batch, 1, combined_seq_len,\n                combined_seq_len].\n            labels (torch.Tensor): Optional target text labels [batch, combined_seq_len].\n            inference_params (InferenceParams): Inference-time parameters including KV cache.\n\n            video_start_index:\n                0 -- all video\n                len(video_seq) -- all image\n                others -- mixture\n            *_input_mask: should not be None in the first PP stage\n        Returns:\n            output (torch.Tensor): Loss of shape [b, s] if labels are provided, otherwise logits of shape\n                [b, s, vocab_size].\n        \"\"\"\n        video_start_index = 0\n        vision_grid_thw = None\n        vision_data = None\n        if image_grid_thw is not None:\n            image_mask = input_ids == self.image_token_id\n            vision_grid_thw = image_grid_thw\n            vision_data = pixel_values\n            video_start_index = image_mask.sum().item()\n        if video_grid_thw is not None:\n            video_mask = input_ids == self.video_token_id\n            if vision_grid_thw is not None:\n                vision_grid_thw = torch.cat([vision_grid_thw, video_grid_thw], dim=0)\n                vision_data = torch.cat([vision_data, pixel_values_videos], dim=0)\n            else:\n                vision_grid_thw = video_grid_thw\n                vision_data = pixel_values_videos\n        use_inference_kv_cache = (\n            inference_params is not None and \"image_tokens_count\" in inference_params.key_value_memory_dict\n        )\n        if use_inference_kv_cache:\n            raise NotImplementedError()\n\n        if self.pre_process:\n            vision_embeds = None\n            if vision_grid_thw is not None and vision_grid_thw.shape[0] > 0:\n                vision_embeds = self.vision_model(\n                    vision_data=vision_data,  # If None, vision model should use intermediate outputs (EPP > 1)\n                    grid_thw=vision_grid_thw,  # should provided in each EPP stage\n                )\n\n            # If running inference, the language model KV cache will be updated for image token positions.\n            # Here we store the image tokens sequence length, which can be used as an offset to the KV cache later.\n            if inference_params is not None:\n                raise NotImplementedError()\n                # inference_params.key_value_memory_dict[\"image_tokens_count\"] = (\n                #     vision_embeddings.shape[0]\n                # )\n\n            # If running inference, we can skip image token computation if they were computed already earlier\n            # for this sample.\n            if use_inference_kv_cache:\n                language_embeddings: torch.Tensor = self.language_model.embedding(\n                    input_ids=input_ids,\n                    position_ids=None,  # NOTE: disable\n                )  # [text_seq_len, b, h_language]\n                # NOTE: why not cat here? is it the combined embeddings useless?\n                combined_embeddings = language_embeddings\n            elif vision_embeds is not None:\n                if video_start_index == 0:\n                    image_embeds = None\n                    video_embeds = vision_embeds\n                elif video_start_index == vision_embeds.shape[0]:\n                    image_embeds = vision_embeds\n                    video_embeds = None\n                elif 0 < video_start_index < vision_embeds.shape[0]:\n                    image_embeds = vision_embeds[:video_start_index]\n                    video_embeds = vision_embeds[video_start_index:]\n                else:\n                    raise ValueError(\n                        f\"Expect video token start index in range [0, {vision_embeds.shape[0]}], but got \"\n                        f\"{video_start_index}\"\n                    )\n\n                combined_embeddings = self.language_model.embedding(\n                    input_ids=input_ids,\n                    position_ids=None,  # NOTE: disable\n                )  # [text_seq_len, b, h_language]\n\n                if image_embeds is not None or video_embeds is not None:\n                    combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()\n                    if image_embeds is not None:\n                        image_mask = (input_ids == self.image_token_id).contiguous()\n                        if image_mask.sum() > 0:\n                            combined_embeddings = combined_embeddings.clone()\n                            combined_embeddings[image_mask] = image_embeds.to(\n                                dtype=combined_embeddings.dtype, device=combined_embeddings.device\n                            )\n                    if video_embeds is not None:\n                        video_mask = (input_ids == self.video_token_id).contiguous()\n                        if video_mask.sum() > 0:\n                            combined_embeddings = combined_embeddings.clone()\n                            combined_embeddings[video_mask] = video_embeds.to(\n                                dtype=combined_embeddings.dtype, device=combined_embeddings.device\n                            )\n                    combined_embeddings = combined_embeddings.transpose(0, 1).contiguous()\n\n            else:\n                combined_embeddings = self.language_model.embedding(\n                    input_ids=input_ids,\n                    position_ids=None,  # NOTE: disable\n                )  # [text_seq_len, b, h_language]\n\n            if packed_seq_params is not None:\n                combined_embeddings = (\n                    preprocess_packed_seqs(\n                        combined_embeddings.transpose(0, 1).contiguous(), attention_mask, pre_process=True\n                    )[0]\n                    .transpose(0, 1)\n                    .contiguous()\n                )\n            if self.config.sequence_parallel:\n                combined_embeddings = tensor_parallel.scatter_to_sequence_parallel_region(combined_embeddings)\n                combined_embeddings = combined_embeddings.contiguous()\n        else:\n            combined_embeddings = None\n        from .rope_utils import get_rope_index\n\n        # BSHD\n        position_ids, _ = get_rope_index(\n            input_ids,\n            image_grid_thw=image_grid_thw,\n            video_grid_thw=video_grid_thw,\n            attention_mask=attention_mask,\n        )\n        # THD\n        if packed_seq_params is not None:\n            position_ids = (\n                preprocess_packed_seqs(position_ids.permute(1, 2, 0), attention_mask, pre_process=True)[0]\n                .permute(2, 0, 1)\n                .contiguous()\n            )\n            attention_mask = None\n\n        output = self.language_model(\n            input_ids=None,\n            position_ids=position_ids,  # None in encoder\n            attention_mask=attention_mask,  # None in encoder\n            decoder_input=combined_embeddings,  # only not None in the first decoder PP stage\n            labels=labels,  # only not None in the last decoder PP stage\n            # inference_params=inference_params,  # currently always None\n            packed_seq_params=packed_seq_params,  # currently always None\n            **(extra_block_kwargs or {}),\n            **kwargs,\n        )\n\n        return output\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/rope_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom __future__ import annotations\n\nimport logging\nfrom typing import Optional\n\nimport torch\nfrom megatron.core.models.common.embeddings.rope_utils import *\nfrom megatron.core.models.common.embeddings.rope_utils import _apply_rotary_pos_emb_bshd\nfrom torch import Tensor\n\nlogger = logging.getLogger(__name__)\n\n\n# Slightly modified from Qwen2VLForConditionalGeneration.get_rope_index\ndef get_rope_index(\n    input_ids: Optional[torch.LongTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    second_per_grid_ts: Optional[torch.Tensor] = None,\n    attention_mask: Optional[torch.Tensor] = None,\n):\n    \"\"\"\n    Calculate the 3D rope index based on image and video's temporal, height and width in LLM.\n\n    Explanation:\n\n        Each embedding sequence contains vision embedding and text embedding or just contains text embedding.\n\n        For pure text embedding sequence, the rotary position embedding has no difference with modern LLMs.\n\n        Examples:\n\n            input_ids: [T T T T T], here T is for text.\n            temporal position_ids: [0, 1, 2, 3, 4]\n            height position_ids: [0, 1, 2, 3, 4]\n            width position_ids: [0, 1, 2, 3, 4]\n\n        For vision and text embedding sequence, we calculate 3D rotary position embedding for vision part\n        and 1D rotary position embedding for text part.\n\n        Examples:\n\n            Temporal (Time): 3 patches, representing different segments of the video in time.\n            Height: 2 patches, dividing each frame vertically.\n            Width: 2 patches, dividing each frame horizontally.\n            We also have some important parameters:\n            fps (Frames Per Second): The video's frame rate, set to 1. This means one frame is processed each\n            second.\n            tokens_per_second: This is a crucial parameter. It dictates how many \"time-steps\" or \"temporal\n                               tokens\" are conceptually packed into a one-second interval of the video.\n                               In this case, we have 25 tokens per second. So each second of the video will be\n                               represented with 25 separate time points. It essentially defines the temporal\n                               granularity.\n            temporal_patch_size: The number of frames that compose one temporal patch. Here, it's 2 frames.\n            interval: The step size for the temporal position IDs, calculated as tokens_per_second *\n            temporal_patch_size / fps. In this case, 25 * 2 / 1 = 50. This means that each temporal patch will be\n            have a difference of 50 in the temporal position IDs.\n            input_ids: [V V V V V V V V V V V V T T T T T], here V is for vision.\n            vision temporal position_ids: [0, 0, 0, 0, 50, 50, 50, 50, 100, 100, 100, 100]\n            vision height position_ids: [0, 0, 1, 1, 0, 0, 1, 1, 0, 0, 1, 1]\n            vision width position_ids: [0, 1, 0, 1, 0, 1, 0, 1, 0, 1, 0, 1]\n            text temporal position_ids: [101, 102, 103, 104, 105]\n            text height position_ids: [101, 102, 103, 104, 105]\n            text width position_ids: [101, 102, 103, 104, 105]\n            Here we calculate the text start position_ids as the max vision position_ids plus 1.\n\n    Args:\n        input_ids (`torch.LongTensor` of shape `(batch_size, sequence_length)`):\n            Indices of input sequence tokens in the vocabulary. Padding will be ignored by default should you provide\n            it.\n        image_grid_thw (`torch.LongTensor` of shape `(num_images, 3)`, *optional*):\n            The temporal, height and width of feature shape of each image in LLM.\n        video_grid_thw (`torch.LongTensor` of shape `(num_videos, 3)`, *optional*):\n            The temporal, height and width of feature shape of each video in LLM.\n        second_per_grid_ts (`torch.Tensor` of shape `(num_videos)`, *optional*):\n            The time interval (in seconds) for each grid along the temporal dimension in the 3D position IDs.\n        attention_mask (`torch.Tensor` of shape `(batch_size, sequence_length)`, *optional*):\n            Mask to avoid performing attention on padding token indices. Mask values selected in `[0, 1]`:\n\n            - 1 for tokens that are **not masked**,\n            - 0 for tokens that are **masked**.\n\n    Returns:\n        position_ids (`torch.LongTensor` of shape `(3, batch_size, sequence_length)`)\n        mrope_position_deltas (`torch.Tensor` of shape `(batch_size)`)\n    \"\"\"\n    spatial_merge_size = 2\n    tokens_per_second = 2\n    image_token_id = 151655\n    video_token_id = 151656\n    vision_start_token_id = 151652\n    mrope_position_deltas = []\n    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):\n        total_input_ids = input_ids\n        if attention_mask is None:\n            attention_mask = torch.ones_like(total_input_ids)\n        position_ids = torch.ones(\n            3,\n            input_ids.shape[0],\n            input_ids.shape[1],\n            dtype=input_ids.dtype,\n            device=input_ids.device,\n        )\n        image_index, video_index = 0, 0\n        attention_mask = attention_mask.to(total_input_ids.device)\n        for i, input_ids in enumerate(total_input_ids):\n            input_ids = input_ids[attention_mask[i] == 1]\n            image_nums, video_nums = 0, 0\n            vision_start_indices = torch.argwhere(input_ids == vision_start_token_id).squeeze(1)\n            vision_tokens = input_ids[vision_start_indices + 1]\n            image_nums = (vision_tokens == image_token_id).sum()\n            video_nums = (vision_tokens == video_token_id).sum()\n            input_tokens = input_ids.tolist()\n            llm_pos_ids_list: list = []\n            st = 0\n            remain_images, remain_videos = image_nums, video_nums\n            for _ in range(image_nums + video_nums):\n                if image_token_id in input_tokens and remain_images > 0:\n                    ed_image = input_tokens.index(image_token_id, st)\n                else:\n                    ed_image = len(input_tokens) + 1\n                if video_token_id in input_tokens and remain_videos > 0:\n                    ed_video = input_tokens.index(video_token_id, st)\n                else:\n                    ed_video = len(input_tokens) + 1\n                if ed_image < ed_video:\n                    t, h, w = (\n                        image_grid_thw[image_index][0],\n                        image_grid_thw[image_index][1],\n                        image_grid_thw[image_index][2],\n                    )\n                    second_per_grid_t = 0\n                    image_index += 1\n                    remain_images -= 1\n                    ed = ed_image\n\n                else:\n                    t, h, w = (\n                        video_grid_thw[video_index][0],\n                        video_grid_thw[video_index][1],\n                        video_grid_thw[video_index][2],\n                    )\n                    if second_per_grid_ts is not None:\n                        second_per_grid_t = second_per_grid_ts[video_index]\n                    else:\n                        second_per_grid_t = 1.0\n                    video_index += 1\n                    remain_videos -= 1\n                    ed = ed_video\n                llm_grid_t, llm_grid_h, llm_grid_w = (\n                    t.item(),\n                    h.item() // spatial_merge_size,\n                    w.item() // spatial_merge_size,\n                )\n                text_len = ed - st\n\n                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n                range_tensor = torch.arange(llm_grid_t).view(-1, 1)\n                expanded_range = range_tensor.expand(-1, llm_grid_h * llm_grid_w)\n\n                time_tensor = expanded_range * second_per_grid_t * tokens_per_second\n\n                time_tensor_long = time_tensor.long()\n                t_index = time_tensor_long.flatten()\n\n                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()\n                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()\n                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)\n                st = ed + llm_grid_t * llm_grid_h * llm_grid_w\n\n            if st < len(input_tokens):\n                st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n                text_len = len(input_tokens) - st\n                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n            llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)\n            position_ids[..., i, attention_mask[i] == 1] = llm_positions.to(position_ids.device)\n            mrope_position_deltas.append(llm_positions.max() + 1 - len(total_input_ids[i]))\n        mrope_position_deltas = torch.tensor(mrope_position_deltas, device=input_ids.device).unsqueeze(1)\n        return position_ids, mrope_position_deltas\n    else:\n        if attention_mask is not None:\n            position_ids = attention_mask.long().cumsum(-1) - 1\n            position_ids.masked_fill_(attention_mask == 0, 1)\n            position_ids = position_ids.unsqueeze(0).expand(3, -1, -1).to(attention_mask.device)\n            max_position_ids = position_ids.max(0, keepdim=False)[0].max(-1, keepdim=True)[0]\n            mrope_position_deltas = max_position_ids + 1 - attention_mask.shape[-1]\n        else:\n            position_ids = (\n                torch.arange(input_ids.shape[1], device=input_ids.device)\n                .view(1, 1, -1)\n                .expand(3, input_ids.shape[0], -1)\n            )\n            mrope_position_deltas = torch.zeros(\n                [input_ids.shape[0], 1],\n                device=input_ids.device,\n                dtype=input_ids.dtype,\n            )\n\n        return position_ids, mrope_position_deltas\n\n\ndef apply_rotary_pos_emb_thd_absolute(\n    t: Tensor, cu_seqlens: Tensor, freqs: Tensor, rotary_interleaved: bool = False\n) -> Tensor:\n    \"\"\"A baseline implementation of applying RoPE for `thd` format.\n\n    Args:\n        t (Tensor): Input tensor T is of shape [t, h, d]\n        cu_seqlens(Tensor):  Cumulative sum of sequence lengths in a batch for `t`,\n        with shape [b + 1] and dtype torch.int32.\n        freqs (Tensor): Rotary Positional embedding tensor freq is of shape [max_s, 1, 1, d]\n\n    Returns:\n        Tensor: Shape [t, h, d]. The input tensor after applying RoPE.\n    \"\"\"\n    return _apply_rotary_pos_emb_bshd(t[:, None], freqs, rotary_interleaved=rotary_interleaved).squeeze(1)\n\n\ndef apply_rotary_pos_emb_absolute(\n    t: Tensor,\n    freqs: Tensor,\n    config: TransformerConfig,\n    cu_seqlens: Optional[Tensor] = None,\n):\n    \"\"\"\n    Reroute to the appropriate apply_rotary_pos_emb function depending on\n    bshd (conventional) / thd (packed seq) format\n\n    In Qwen2-VL, the shape of freqs is (seq_length, bs, 1, 2 * dim) instead of [max_seqlen, 1, 1, 2 * dim]\n    \"\"\"\n\n    if config.apply_rope_fusion:\n        if cu_seqlens is None:\n            # NOTE: TE backends do not support mRoPE in bshd format when bs > 1\n            if freqs.shape[1] > 1:\n                return _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved)\n            else:\n                return fused_apply_rotary_pos_emb(t, freqs)\n        else:\n            # NOTE: as expected, thd format can use bshd\n            return fused_apply_rotary_pos_emb(t[:, None], freqs).squeeze(1)\n    else:\n        if cu_seqlens is None:\n            return _apply_rotary_pos_emb_bshd(t, freqs, rotary_interleaved=config.rotary_interleaved)\n        else:\n            return apply_rotary_pos_emb_thd_absolute(t, cu_seqlens, freqs, rotary_interleaved=config.rotary_interleaved)\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/vision_config.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom megatron.core import parallel_state\nfrom megatron.core.transformer import TransformerConfig\n\n\ndef get_vision_model_config(config: TransformerConfig) -> TransformerConfig:\n    # Given a Transformer Config from decoder, build vision encoder config\n    # diff: out_hidden_size & intermediate_size\n\n    # mlp: hidden_size -> intermediate_size -> embed_dim, silu\n    # NOTE: here we provide a workaround to solve the wrong layer amount when VPP of decoder is on\n    if config.num_layers in [28, 36]:\n        config.ffn_hidden_size = 3420\n    else:\n        config.ffn_hidden_size = 3456\n\n    if parallel_state.get_virtual_pipeline_model_parallel_world_size() is not None:\n        config.num_layers = 32 * parallel_state.get_virtual_pipeline_model_parallel_world_size()  # depth\n    else:\n        config.num_layers = 32  # depth\n    config.num_attention_heads = 16  # num_heads\n    config.add_bias_linear = True  # all nn.Linear has bias (MLP, attn)\n    config.add_qkv_bias = True  # qkv_proj in attn has bias\n    config.hidden_size = 1280  # hidden_size\n    config.hidden_dropout = 0.0\n    config.attention_dropout = 0.0\n\n    # config.gated_linear_unit = False # no gated\n    # config.activation_func = quick_gelu # hidden_act\n    config.kv_channels = config.hidden_size // config.num_attention_heads\n    config.num_query_groups = config.num_attention_heads  # no GQA\n    config.layernorm_zero_centered_gamma = False  # False\n    config.apply_query_key_layer_scaling = False  # factor=math.sqrt(head_dim)\n    config.bias_activation_fusion = False  # no swiglu, set false\n    config.bias_dropout_fusion = False  # no dropout, set false\n    config.attention_softmax_in_fp32 = True  # use True\n    # config.normalization = 'LayerNorm' # use RMSNorm\n    config.seq_length = 1\n\n    config.tp_comm_overlap = False\n    config.sequence_parallel = False\n    config.temporal_patch_size = 2\n    config.patch_size = 14\n    config.in_channels = 3\n    config.spatial_merge_size = 2\n\n    config.fullatt_block_indexes = [7, 15, 23, 31]\n    config._qwen2_5_vl_window_size = 112\n    return config\n\n\ndef get_vision_projection_config(\n    config: TransformerConfig, embed_dim: int, spatial_merge_size: int\n) -> TransformerConfig:\n    # merger:\n    # context_dim = hidden_size * merge_size**2\n    # out_hidden_size = hidden_size\n    # context_dim -> context_dim -> out_hidden_size\n    # MLP:\n    # input_size -> ffn_hidden_size -> hidden_size\n    # spec: LN -> Linear(bias=True) -> GELU -> Linear(bias=True)\n    config.gated_linear_unit = False\n    config.bias_activation_fusion = False\n    config.add_bias_linear = True\n    config.ffn_hidden_size = embed_dim * (spatial_merge_size**2)\n    config.activation_func = torch.nn.functional.gelu\n    config.tp_comm_overlap = False\n    config.sequence_parallel = False\n    return config\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/vision_model.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\nfrom megatron.core import InferenceParams\nfrom megatron.core.models.common.vision_module.vision_module import VisionModule\nfrom megatron.core.models.vision.multimodal_projector import MultimodalProjector\nfrom megatron.core.packed_seq_params import PackedSeqParams\nfrom megatron.core.transformer.enums import ModelType\nfrom megatron.core.transformer.spec_utils import ModuleSpec\nfrom megatron.core.transformer.transformer_config import TransformerConfig\nfrom torch import nn\nfrom torch.nn import functional as F\n\nfrom .vision_transformer_block import Qwen2_5VisionTransformerBlock as TransformerBlock\n\n\n# copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py\nclass PatchEmbed(nn.Module):\n    def __init__(\n        self,\n        patch_size: int = 14,\n        temporal_patch_size: int = 2,\n        in_channels: int = 3,\n        embed_dim: int = 1152,\n    ) -> None:\n        super().__init__()\n        self.patch_size = patch_size\n        self.temporal_patch_size = temporal_patch_size\n        self.in_channels = in_channels\n        self.embed_dim = embed_dim\n\n        kernel_size = [temporal_patch_size, patch_size, patch_size]\n        self.proj = nn.Conv3d(in_channels, embed_dim, kernel_size=kernel_size, stride=kernel_size, bias=False)\n\n    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n        target_dtype = self.proj.weight.dtype\n        hidden_states = hidden_states.view(\n            -1, self.in_channels, self.temporal_patch_size, self.patch_size, self.patch_size\n        )\n        hidden_states = self.proj(hidden_states.to(dtype=target_dtype)).view(-1, self.embed_dim)\n        return hidden_states\n\n\n# copied from https://github.com/huggingface/transformers/blob/main/src/transformers/models/qwen2_vl/modeling_qwen2_vl.py\nclass VisionRotaryEmbedding(nn.Module):\n    def __init__(self, dim: int, theta: float = 10000.0) -> None:\n        super().__init__()\n        inv_freq = 1.0 / (theta ** (torch.arange(0, dim, 2, dtype=torch.float) / dim))\n        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n    def forward(self, seqlen: int) -> torch.Tensor:\n        seq = torch.arange(seqlen, device=self.inv_freq.device, dtype=self.inv_freq.dtype)\n        freqs = torch.outer(seq, self.inv_freq)\n        return freqs.float()\n\n\nclass Qwen2_5VisionModel(VisionModule):\n    \"\"\"Qwen2.5 ViT vision model.\n\n    Args:\n        transformer_config (TransformerConfig): Transformer config.\n        transformer_layer_spec (ModuleSpec): Specifies module to use for transformer layers.\n        ln_pre_impl (ModuleSpec or type): Specifies the layer norm type to use for ln_pre.\n        add_class_token (bool, optional): Include a class token. Defaults to True.\n        class_token_len (int): Class token length. Defaults to 1 but 8 may be faster.\n        patch_dim (int): Image patch size.\n        img_h (int): Input image height.\n        img_w (int): Input image width.\n    \"\"\"\n\n    def __init__(\n        self,\n        transformer_config: TransformerConfig,\n        transformer_layer_spec: ModuleSpec,\n        projection_config: TransformerConfig,\n        projection_layer_spec: ModuleSpec,\n        projection_type: str = \"mlp\",\n        pre_process: bool = True,\n        post_process: bool = False,\n    ) -> None:\n        super().__init__(config=transformer_config)\n\n        self.spatial_merge_size = transformer_config.spatial_merge_size\n\n        embed_dim = transformer_config.hidden_size\n        num_heads = transformer_config.num_attention_heads\n        temporal_patch_size = transformer_config.temporal_patch_size\n        patch_size = transformer_config.patch_size\n        in_channels = transformer_config.in_channels\n\n        self.patch_size = transformer_config.patch_size\n        self.fullatt_block_indexes = transformer_config.fullatt_block_indexes\n        self.window_size = transformer_config._qwen2_5_vl_window_size\n        self.spatial_merge_unit = self.spatial_merge_size * self.spatial_merge_size\n\n        self.max_sequence_length = transformer_config.seq_length\n        self.patch_embed = PatchEmbed(\n            patch_size=patch_size,\n            temporal_patch_size=temporal_patch_size,\n            in_channels=in_channels,\n            embed_dim=embed_dim,\n        )\n\n        head_dim = embed_dim // num_heads\n        self.rotary_pos_emb = VisionRotaryEmbedding(head_dim // 2)\n\n        self.model_type = ModelType.encoder_or_decoder\n        self.pre_process = pre_process\n        self.post_process = post_process\n\n        # Transformer layers.\n        # TODO: Follow-up changes will make pre and post_process configurable. They are needed for supporting\n        # pipeline parallelism.\n        # NOTE: a final layer norm and/or linear layer present in some implementations are omitted here.\n        self.decoder = TransformerBlock(\n            config=transformer_config,\n            spec=transformer_layer_spec,\n            pre_process=self.pre_process,\n            post_process=self.post_process,\n            post_layer_norm=True,\n        )\n\n        self.merge_hidden_size = projection_config.ffn_hidden_size\n        self.square_merge_size = self.merge_hidden_size // embed_dim\n\n        if self.post_process:\n            self.projection = MultimodalProjector(\n                projection_config, projection_layer_spec, projection_type, projection_config.ffn_hidden_size\n            )\n        else:\n            self.projection = None\n\n        self.input_tensor = None\n\n    def set_input_tensor(self, input_tensor: torch.Tensor) -> None:\n        \"\"\"Sets input tensor to the model.\n\n        Args:\n            input_tensor (Tensor): Sets the input tensor for the model.\n        \"\"\"\n        if self.pre_process:  # always True\n            self.input_tensor = input_tensor\n        else:\n            raise NotImplementedError()\n\n    def rot_pos_emb(self, grid_thw):\n        pos_ids = []\n        for t, h, w in grid_thw:\n            hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)\n            hpos_ids = hpos_ids.reshape(\n                h // self.spatial_merge_size,\n                self.spatial_merge_size,\n                w // self.spatial_merge_size,\n                self.spatial_merge_size,\n            )\n            hpos_ids = hpos_ids.permute(0, 2, 1, 3)\n            hpos_ids = hpos_ids.flatten()\n\n            wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)\n            wpos_ids = wpos_ids.reshape(\n                h // self.spatial_merge_size,\n                self.spatial_merge_size,\n                w // self.spatial_merge_size,\n                self.spatial_merge_size,\n            )\n            wpos_ids = wpos_ids.permute(0, 2, 1, 3)\n            wpos_ids = wpos_ids.flatten()\n            pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))\n        pos_ids = torch.cat(pos_ids, dim=0).to(grid_thw.device)\n        max_grid_size = grid_thw[:, 1:].max()\n        rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size).to(grid_thw.device)\n        rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)\n        return rotary_pos_emb\n\n    def get_window_index(self, grid_thw):\n        window_index: list = []\n        cu_window_seqlens: list = [0]\n        window_index_id = 0\n        vit_merger_window_size = self.window_size // self.spatial_merge_size // self.patch_size\n\n        for grid_t, grid_h, grid_w in grid_thw:\n            llm_grid_h, llm_grid_w = (\n                grid_h // self.spatial_merge_size,\n                grid_w // self.spatial_merge_size,\n            )\n            index = torch.arange(grid_t * llm_grid_h * llm_grid_w).reshape(grid_t, llm_grid_h, llm_grid_w)\n            pad_h = vit_merger_window_size - llm_grid_h % vit_merger_window_size\n            pad_w = vit_merger_window_size - llm_grid_w % vit_merger_window_size\n            num_windows_h = (llm_grid_h + pad_h) // vit_merger_window_size\n            num_windows_w = (llm_grid_w + pad_w) // vit_merger_window_size\n            index_padded = F.pad(index, (0, pad_w, 0, pad_h), \"constant\", -100)\n            index_padded = index_padded.reshape(\n                grid_t,\n                num_windows_h,\n                vit_merger_window_size,\n                num_windows_w,\n                vit_merger_window_size,\n            )\n            index_padded = index_padded.permute(0, 1, 3, 2, 4).reshape(\n                grid_t,\n                num_windows_h * num_windows_w,\n                vit_merger_window_size,\n                vit_merger_window_size,\n            )\n            seqlens = (index_padded != -100).sum([2, 3]).reshape(-1)\n            index_padded = index_padded.reshape(-1)\n            index_new = index_padded[index_padded != -100]\n            window_index.append(index_new + window_index_id)\n            cu_seqlens_tmp = seqlens.cumsum(0) * self.spatial_merge_unit + cu_window_seqlens[-1]\n            cu_window_seqlens.extend(cu_seqlens_tmp.tolist())\n            window_index_id += (grid_t * llm_grid_h * llm_grid_w).item()\n        window_index = torch.cat(window_index, dim=0)\n\n        return window_index, cu_window_seqlens\n\n    def forward(\n        self,\n        vision_data: Optional[torch.Tensor],\n        grid_thw: torch.Tensor,\n        inference_params: Optional[InferenceParams] = None,\n        extra_block_kwargs: dict = None,\n    ) -> torch.Tensor:\n        \"\"\"Forward function of the Qwen2 Vision Model. This function passes the input tensors\n        through the embedding layer and then the transformer.\n\n        Args:\n            x (torch.Tensor): input image/video data of shape [n_tokens, n_dims]\n            grid_thw (torch.Tensor): the size tensor indicates grid size of each image/frame\n            packed_seq_params (PackedSeqParams): parameters to build attention mask in the backend\n\n        Returns:\n            x (torch.Tensor): output after final transformer block of shape [b, s, h].\n        \"\"\"\n        assert grid_thw is not None\n        assert self.input_tensor is None\n        assert inference_params is None\n\n        # Rotary positional embeddings (embedding is None for PP intermediate devices)\n        vision_data = self.patch_embed(vision_data)\n        window_index, cu_window_seqlens = self.get_window_index(grid_thw)\n        cu_window_seqlens = torch.tensor(\n            cu_window_seqlens,\n            device=vision_data.device,\n            dtype=torch.int32,\n        )\n        cu_window_seqlens = torch.unique_consecutive(cu_window_seqlens)\n\n        seq_len, _ = vision_data.size()\n        vision_data = vision_data.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)\n        vision_data = vision_data[window_index, :, :]\n        vision_data = vision_data.reshape(seq_len, 1, -1)\n\n        rotary_pos_emb = self.rot_pos_emb(grid_thw)\n        rotary_pos_emb = rotary_pos_emb.reshape(seq_len // self.spatial_merge_unit, self.spatial_merge_unit, -1)\n        rotary_pos_emb = rotary_pos_emb[window_index, :, :]\n        rotary_pos_emb = rotary_pos_emb.reshape(seq_len, 1, 1, -1).repeat(1, 1, 1, 2)\n\n        hidden_states = self.decoder(\n            hidden_states=vision_data,\n            attention_mask=None,\n            inference_params=inference_params,\n            rotary_pos_emb=rotary_pos_emb,\n            packed_seq_params=self.build_packed_seq_params(None, cu_window_seqlens),\n            packed_seq_params_full=self.build_packed_seq_params(grid_thw),\n            fullatt_block_indexes=self.fullatt_block_indexes,\n            **(extra_block_kwargs or {}),\n        )\n\n        hidden_states = self.projection(hidden_states.view(-1, self.merge_hidden_size))\n        reverse_indices = torch.argsort(window_index)\n        return hidden_states[reverse_indices, :]\n\n    def build_packed_seq_params(\n        self,\n        grid_thw: Optional[torch.Tensor],\n        cu_seqlens: Optional[torch.Tensor] = None,\n    ) -> PackedSeqParams:\n        # NOTE: each frame is a sequence (rather than each grid)\n        if grid_thw is not None:\n            seqlens = torch.repeat_interleave(grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0])\n            cu_seqlens = seqlens.cumsum(dim=0)\n            cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0).int()\n        else:\n            seqlens = cu_seqlens[1:] - cu_seqlens[:-1]\n\n        max_seqlen_q = seqlens.max()\n        return PackedSeqParams(\n            cu_seqlens_q=cu_seqlens,\n            cu_seqlens_kv=cu_seqlens,\n            qkv_format=\"thd\",\n            max_seqlen_q=max_seqlen_q,\n            max_seqlen_kv=max_seqlen_q,\n        )\n"
  },
  {
    "path": "verl/models/mcore/qwen2_5_vl/vision_transformer_block.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright (c) 2024 Alibaba PAI Team.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom megatron.core.transformer.transformer_block import *\n\n\nclass Qwen2_5VisionTransformerBlock(TransformerBlock):\n    def _checkpointed_forward(\n        self,\n        hidden_states: Tensor,\n        attention_mask: Tensor,\n        context: Tensor,\n        context_mask: Tensor,\n        rotary_pos_emb: Tensor,\n        attention_bias: Tensor,\n        packed_seq_params: PackedSeqParams,\n        packed_seq_params_full: PackedSeqParams,\n        fullatt_block_indexes,\n    ):\n        \"\"\"Forward method with activation checkpointing.\"\"\"\n\n        def custom(start: int, end: int):\n            def custom_forward(hidden_states, attention_mask, context, context_mask, rotary_pos_emb):\n                for index in range(start, end):\n                    if index in fullatt_block_indexes:\n                        packed_seq_params_now = packed_seq_params_full\n                    else:\n                        packed_seq_params_now = packed_seq_params\n                    layer = self._get_layer(index)\n                    hidden_states, context = layer(\n                        hidden_states=hidden_states,\n                        attention_mask=attention_mask,\n                        context=context,\n                        context_mask=context_mask,\n                        rotary_pos_emb=rotary_pos_emb,\n                        attention_bias=attention_bias,\n                        inference_context=None,\n                        packed_seq_params=packed_seq_params_now,\n                    )\n                return hidden_states, context\n\n            return custom_forward\n\n        def checkpoint_handler(forward_func):\n            \"\"\"Determines whether to use the `te_checkpoint` or `tensor_parallel.checkpoint`\"\"\"\n            if self.config.fp8:\n                return te_checkpoint(\n                    forward_func,\n                    self.config.distribute_saved_activations,\n                    tensor_parallel.random.get_cuda_rng_tracker,\n                    parallel_state.get_tensor_model_parallel_group(),\n                    hidden_states,\n                    attention_mask,\n                    context,\n                    context_mask,\n                    rotary_pos_emb,\n                )\n            else:\n                return tensor_parallel.checkpoint(\n                    forward_func,\n                    self.config.distribute_saved_activations,\n                    hidden_states,\n                    attention_mask,\n                    context,\n                    context_mask,\n                    rotary_pos_emb,\n                )\n\n        if self.config.recompute_method == \"uniform\":\n            # Uniformly divide the total number of Transformer layers and checkpoint\n            # the input activation of each divided chunk.\n            # A method to further reduce memory usage reducing checkpoints.\n            layer_idx = 0\n            while layer_idx < self.num_layers_per_pipeline_rank:\n                hidden_states, context = checkpoint_handler(\n                    custom(layer_idx, layer_idx + self.config.recompute_num_layers)\n                )\n\n                layer_idx += self.config.recompute_num_layers\n\n        elif self.config.recompute_method == \"block\":\n            # Checkpoint the input activation of only a set number of individual\n            # Transformer layers and skip the rest.\n            # A method fully use the device memory removing redundant re-computation.\n            recompute_skip_num_layers = 0\n            for layer_idx in range(self.num_layers_per_pipeline_rank):\n                # Skip recomputation when input grad computation is not needed.\n                # Need to have at least one input tensor with gradient computation\n                # for re-enterant autograd engine.\n                if self.config.fp8 and not hidden_states.requires_grad:\n                    recompute_skip_num_layers += 1\n                if (\n                    layer_idx >= recompute_skip_num_layers\n                    and layer_idx < self.config.recompute_num_layers + recompute_skip_num_layers\n                ):\n                    hidden_states, context = checkpoint_handler(custom(layer_idx, layer_idx + 1))\n                else:\n                    hidden_states, context = custom(layer_idx, layer_idx + 1)(\n                        hidden_states, attention_mask, context, context_mask, rotary_pos_emb\n                    )\n        else:\n            raise ValueError(\"Invalid activation recompute method.\")\n\n        return hidden_states\n\n    def forward(\n        self,\n        hidden_states: Union[Tensor, WrappedTensor],\n        attention_mask: Optional[Tensor],\n        context: Optional[Tensor] = None,\n        context_mask: Optional[Tensor] = None,\n        rotary_pos_emb: Optional[Tensor] = None,\n        rotary_pos_cos: Optional[Tensor] = None,\n        rotary_pos_sin: Optional[Tensor] = None,\n        attention_bias: Optional[Tensor] = None,\n        inference_context: Optional[BaseInferenceContext] = None,\n        packed_seq_params: Optional[PackedSeqParams] = None,\n        sequence_len_offset: Optional[Tensor] = None,\n        packed_seq_params_full: PackedSeqParams = None,\n        fullatt_block_indexes=None,\n        *,\n        inference_params: Optional[BaseInferenceContext] = None,\n    ):\n        \"\"\"\n        Perform the forward pass through the transformer block.\n\n        This method handles the core computation of the transformer, including\n        self-attention, optional cross-attention, and feed-forward operations.\n\n        Args:\n            hidden_states (Union[Tensor, WrappedTensor]): Input tensor of shape [s, b, h]\n                where s is the sequence length, b is the batch size, and h is the hidden size.\n                Can be passed as a WrappedTensor during inference to avoid an obsolete\n                reference in the calling function.\n            attention_mask (Tensor): Boolean tensor of shape [1, 1, s, s] for masking\n                self-attention.\n            context (Tensor, optional): Context tensor for cross-attention.\n            context_mask (Tensor, optional): Mask for cross-attention context\n            rotary_pos_emb (Tensor, optional): Rotary positional embeddings.\n            attention_bias (Tensor): Bias tensor for Q * K.T of shape in shape broadcastable\n                to [b, num_head, sq, skv], e.g. [1, 1, sq, skv].\n                Used as an alternative to apply attention mask for TE cuDNN attention.\n            inference_context (BaseInferenceContext, optional): Parameters for inference-time\n                optimizations.\n            packed_seq_params (PackedSeqParams, optional): Parameters for packed sequence\n                processing.\n\n        Returns:\n            Union[Tensor, Tuple[Tensor, Tensor]]: The output hidden states tensor of shape\n            [s, b, h], and optionally the updated context tensor if cross-attention is used.\n        \"\"\"\n\n        inference_context = deprecate_inference_params(inference_context, inference_params)\n\n        # Delete the obsolete reference to the initial input tensor if necessary\n        if isinstance(hidden_states, WrappedTensor):\n            hidden_states = hidden_states.unwrap()\n\n        if not self.pre_process:\n            # See set_input_tensor()\n            hidden_states = self.input_tensor\n\n        # Update the inference parameters with the current batch size in case it is variable\n        if inference_context and not self.training:\n            inference_context.current_batch_size = hidden_states.size(1)\n\n        # Viewless tensor.\n        # - We only need to create a viewless tensor in the case of micro batch\n        #   size (mbs) == 1, since in this case, 'hidden_states.transpose()'\n        #   above creates a view tensor, and '.contiguous()' is a pass-through.\n        #   For mbs >= 2, '.contiguous()' creates a new tensor, eliminating\n        #   the need to make it viewless.\n        #\n        #   However, we don't explicitly check mbs == 1 here because\n        #   make_viewless_tensor() has negligible overhead when its input\n        #   is already viewless.\n        #\n        # - For the 'else' case above, calling make_viewless_tensor() here is\n        #   likely redundant, since p2p_communication.py (likely originator)\n        #   already creates viewless tensors. That said, make_viewless_tensor()\n        #   is called here to be future-proof and corner-case-proof.\n        hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)\n\n        if self.config.sequence_parallel:\n            rng_context = tensor_parallel.get_cuda_rng_tracker().fork()\n        else:\n            rng_context = nullcontext()\n\n        # If fp8_recipe is delayed, wrap the entire pass with get_fp8_context(),\n        # otherwise do nothing extra at the outer level\n        # if we are using other fp8 recipes, then the context manager enter&exit are free\n        # we can wrap fp8_context within the for loop over layers, so that we can fine-grained\n        # control which layer will be fp8 or bf16\n        use_outer_fp8_context = self.config.fp8 and self.config.fp8_recipe == Fp8Recipe.delayed\n        use_inner_fp8_context = self.config.fp8 and self.config.fp8_recipe != Fp8Recipe.delayed\n        outer_fp8_context = get_fp8_context(self.config) if use_outer_fp8_context else nullcontext()\n\n        with rng_context, outer_fp8_context:\n            # Forward pass.\n            if self.config.recompute_granularity == \"full\" and self.training:\n                hidden_states = self._checkpointed_forward(\n                    hidden_states=hidden_states,\n                    attention_mask=attention_mask,\n                    context=context,\n                    context_mask=context_mask,\n                    rotary_pos_emb=rotary_pos_emb,\n                    attention_bias=attention_bias,\n                    packed_seq_params=packed_seq_params,\n                    packed_seq_params_full=packed_seq_params_full,\n                    fullatt_block_indexes=fullatt_block_indexes,\n                )\n            else:\n                for l_no, layer in enumerate(self.layers):\n                    inner_fp8_context = (\n                        get_fp8_context(self.config, layer.layer_number - 1) if use_inner_fp8_context else nullcontext()\n                    )\n                    if l_no in fullatt_block_indexes:\n                        packed_seq_params_now = packed_seq_params_full\n                    else:\n                        packed_seq_params_now = packed_seq_params\n                    with self.offload_context, inner_fp8_context:\n                        hidden_states, context = layer(\n                            hidden_states=hidden_states,\n                            attention_mask=attention_mask,\n                            context=context,\n                            context_mask=context_mask,\n                            rotary_pos_emb=rotary_pos_emb,\n                            rotary_pos_cos=rotary_pos_cos,\n                            rotary_pos_sin=rotary_pos_sin,\n                            attention_bias=attention_bias,\n                            inference_context=inference_context,\n                            packed_seq_params=packed_seq_params_now,\n                            sequence_len_offset=sequence_len_offset,\n                        )\n\n                    if (\n                        torch.is_grad_enabled()\n                        and self.config.cpu_offloading\n                        and self.group_prefetch_offload_commit_async is not None\n                    ):\n                        hidden_states = self.group_prefetch_offload_commit_async(hidden_states)\n\n        # Final layer norm.\n        if self.final_layernorm is not None:\n            hidden_states = self.final_layernorm(hidden_states)\n            # TENorm produces a \"viewed\" tensor. This will result in schedule.py's\n            # deallocate_output_tensor() throwing an error, so a viewless tensor is\n            # created to prevent this.\n            hidden_states = make_viewless_tensor(inp=hidden_states, requires_grad=True, keep_graph=True)\n\n        return hidden_states\n"
  },
  {
    "path": "verl/models/mcore/readme.md",
    "content": "updated 20251222\n\n# The ways verl integrates megatron-core\nThere has been 3 ways that verl integrates megatron-core as it training backend:\n1. the codes inside this directory, which defines the conversion for new models one by one. (deprecated now)\n2. through [mbridge](https://github.com/ISEEKYAN/mbridge) (will be deprecated at about v0.8)\n3. through [megatron-bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) (the official way for further development)\n\nThere is a configure option of `megatron.use_mbridge` to choose way#1 (false) or way#2 (true), and after the megatron-bridge is integrated we have a new option `megatron.vanilla_mbridge` to choose way#2 (true) or way#3 (false)\n\nNow since we deprecated the way#1, the option `use_mbridge` will be asserted to be true and will be removed after v0.7. The default `vanilla_mbridge` is true for now and will be false one the megatron-bridge backend turns default.\n\nWith the bridge way(#2 or #3), we can directly load and save the megatron model weight through HuggingFace format, and we can use any megatron version >= 0.13 to adopt new megatron optimization feature as handy as possible by directly add overrided megatron configs such as `+actor_rollout_ref.actor.megatron.override_transformer_config.recompute_method=uniform`.\n\n# How to support new models\n1. Make sure the model is supported by your inference engine (vLLM or SGLang or TensorRT-LLM) with correct version.\n2. Make sure the model is supported by the bridge\n   - If it is a model of new architecture, open an issue to `megatron-bridge` or contribute your implementation to `megatron-bridge`. Be cautious to have a matched version of `Megatron` and `TransformerEngine`\n   - If it is a private model, implement your private model with `mbridge` or `megatron-bridge`(prefered).\n\n3. Now the model is supported, just change the model path of your scripts and run the scritps.\n\n\n\n\n\n# #Below are deprecated since 2025.12#\n# verl Megatron-Core Models\nNow we use [mbridge](https://github.com/iseekyan/mbridge) to support megatron models. And we will migrate to [megatron-bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) in the future.\n\nWith the mbridge, we can use allmost all the Megatron-Core features to support new models with little effort. And no offline weights conversion is needed, all the weights conversion is done online. We can directly save the mcore model to huggingface format during training.\n\nAlso, we can easily upgrade the mcore version to the latest version. In most cases, the upgrade is seamless. (except when the mcore API changes and we need to update the verl code accordingly)\n\n## How to support new models\n1. make sure the model is supported by vLLM\n2. Support the model in [mbridge](https://github.com/iseekyan/mbridge), see its currently supported models for example.\n    - we will migrate to [megatron-bridge](https://github.com/NVIDIA-NeMo/Megatron-Bridge) in the future.\n3. Register the model forward function in verl, see the example in `verl/verl/models/mcore/registry.py`.\n\n\n\n# #Below are deprecated since 2025.10#\nThe earlier versions of verl use `Megatron-LM` 0.4 and workaround huggingface model classes. To better use the latest features and speedup of modern Megatron, we are migrating to `Megatron-Core`(mcore), and use the recommended `GPTModel` class for all language models. With mcore `GPTModel`, we can use the latest features like `context parallel`, `expert parallel`, `dist_checkpointing`, etc. and we can update mcore with little effort in the future for new features.\n\nThe migration has been successful with the help of the mcore team and the community. What we have done is:\n1. update `Megatron` version to `0.14.0`\n2. migrate `LlamaForCausalLM` and `Qwen2ForCausalLM` to mcore `GPTModel`\n3. support sequence packing/thd format.\n4. support `tensor parallel`, `pipeline parallel`, `sequence parallel`, `virtual pipeline parallel`, `context parallel`.\n5. support the mcore `dist_checkpointing` feature and a basic offline weighs conversion script from huggingface to mcore `dist_checkpointing` format.\n\nWe are working on the following features:\n- support `Qwen2MoeForCausalLM`\n- support `MixtralForCausalLM`\n- support `DeepseekV3ForCausalLM`\n- support `expert parallel`\n\nFeatures we invite the community to contribute:\n- better scripts for offline weights conversion from huggingface to mcore `dist_checkpointing` format.\n    - conversion of large models with multiple GPUs\n    - conversion of large models with single GPU\n- refactor the `megatron_checkpoint_manager.py` by `dist_checkpointing` format.\n- support llama4\n- support qwen2.5-vl\n\nTo track the progress of verl mcore integration, please refer to the [mcore integration issue](https://github.com/volcengine/verl/issues/1033).\n\n## How things work now\nTo engage the community in contributing, here are the key steps in our mcore integration process and features under development. \n\nThe huggingface `transformers` is the de facto standard of model zoo while mcore is good at computation efficiency. The main challenge is conversion between the two.\nmain steps:\n1. modelling the huggingface model with mcore `GPTModel`\n    - a. convert the huggingface config to mcore `TransformerConfig`\n    - b. init the mcore `GPTModel` with the converted config\n    - c. load the huggingface model weights to the `GPTModel`\n2. online weight conversion from mcore to huggingface (due to the rollout engine `vLLM` is using huggingface format)\n    - a. bridge the gap between mcore and huggingface weights format and name mapping\n    - b. online resharding the mcore weights to rollout engine\n        - this part is very complicated with multiple parallel strategies composition between mcore and rollout engine\n3. support the mcore features in verl\n    - a. support `tensor parallel`, `pipeline parallel`, `sequence parallel`, `virtual pipeline parallel`, `context parallel`\n    - b. support recompute and other mcore speed up features\n\n4. checkpointing\n    - a. support recovering the verl training.\n    - b. support exporting the mcore checkpoint to huggingface format, for downstream inference.\n\n### Modelling the huggingface model with mcore `GPTModel`\nThe first step is to convert huggingface config to mcore `TransformerConfig` and init the mcore `GPTModel` with the converted config. See code in `verl/models/mcore/config_converter.py` and `verl/verl/models/mcore/models/model_initializer.py`. The corresponding model forward code is in `verl/verl/models/mcore/models/model_forward.py`.\n\nThere are two ways of loading the huggingface model weights to the `GPTModel`\n1. Runtime loading\n    - every rank loads the entire huggingface model weights and then shard and convert to mcore weights.\n    - speed is slow and memory consumption is high.\n    - this way is deprecated and will not support new models.\n2. Offline loading\n    - use offline script to convert the huggingface model weights to mcore weights and save with mcore `dist_checkpointing` format.\n    - online loading and sharding is automatically done by mcore `dist_checkpointing` format. The speed is fast and memory consumption is low.\n    - the offline script is in `verl/scripts/converter_hf_to_mcore.py`.\n\n### online weight conversion from mcore to huggingface\nSee function `convert_megatron_model_to_transformers_model` in `verl/utils/megatron_utils.py` for the details.\n\nIt should be refatored for extensibility and better performance.\n\n### support the mcore features in verl\nMost of the features of `GPTModel` is out-of-the-box supported in verl through changing the `TransformerConfig`, except those about parallel strategies, such as `expert parallel`. \nFeatures about parallel strategies should be supported with changes about the online weights conversion(especially the resharding part) and verl work dispatching.\n\n### checkpointing\nThe existing checkpointing code is in `verl/utils/checkpoint/megatron_checkpoint_manager.py`. And the script to convert checkpoint to huggingface format is in `verl/scripts/model_merger`.\n\nThe existing checkpoint format simply saves every rank's weights and optimizer states. It should be refactored by `dist_checkpointing` format.\n\n\n## How to support new models\n1. make sure the model is supported by vLLM\n2. modelling the huggingface model with mcore `GPTModel` (The [Pai-Megatron-Path](https://github.com/alibaba/Pai-Megatron-Patch/tree/main) is a good reference)\n    - a. convert the huggingface config to mcore `TransformerConfig`\n    - b. init the mcore `GPTModel` with the converted config\n    - c. load the huggingface model weights to the `GPTModel`\n    - d. for VLM the interface might be different, it is ok to add a new model class with GPTModel as its module.\n3. offline weights conversion from huggingface to mcore `dist_checkpointing` format\n4. support online weights conversion from mcore to huggingface\n    - it is recommended to initialize a vLLM model with the converted mcore weights, and then test if the generating sequence is correct.\n\n\n## How to scale up to larger models like deepseek-v3 or other 100B+ models\nThe greatest challenge for scaling up to larger models is the memory consumption.\n\nThe necessary features under development for scaling up are\n1. Training engine part\n    - expert parallel\n2. Rollout engine part\n    - pipeline parallel\n    - expert parallel\n    - more efficient and general weight resharding and loading\n3. Offline weights conversion\n    - support weights larger than single GPU memory\n"
  },
  {
    "path": "verl/models/mcore/registry.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nRegistry module for model architecture components.\n\"\"\"\n\nfrom enum import Enum\nfrom typing import Callable\n\nimport torch\nimport torch.nn as nn\n\nfrom .model_forward import gptmodel_forward_no_padding, model_forward_gen\nfrom .model_forward_fused import fused_forward_model_gen, fused_forward_no_padding_gen\n\n\nclass SupportedVLM(Enum):\n    QWEN2_5_VL = \"Qwen2_5_VLForConditionalGeneration\"\n    QWEN3_MOE_VL = \"Qwen3VLMoeForConditionalGeneration\"\n    QWEN3_VL = \"Qwen3VLForConditionalGeneration\"\n    QWEN3_5_MOE_VL = \"Qwen3_5MoeForConditionalGeneration\"\n    QWEN3_5_VL = \"Qwen3_5ForConditionalGeneration\"\n\n\nsupported_vlm = [member.value for member in SupportedVLM]\n\n\ndef get_mcore_forward_fn(hf_config) -> Callable:\n    \"\"\"\n    Get the forward function for given model architecture.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    if hf_config.architectures[0] in supported_vlm:\n        return model_forward_gen(True)\n    else:\n        # default to language model\n        return model_forward_gen(False)\n\n\ndef get_mcore_forward_no_padding_fn(hf_config) -> Callable:\n    \"\"\"\n    Get the forward function for given model architecture.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    return gptmodel_forward_no_padding\n\n\ndef get_mcore_forward_fused_fn(hf_config) -> Callable:\n    \"\"\"\n    Get the forward function for given model architecture.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    if hf_config.architectures[0] in supported_vlm:\n        return fused_forward_model_gen(True)\n    else:\n        # default to language model\n        return fused_forward_model_gen(False)\n\n\ndef get_mcore_forward_fused_no_padding_fn(hf_config) -> Callable:\n    \"\"\"\n    Get the fused forward function for no-padding inputs.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    if hf_config.architectures[0] in supported_vlm:\n        return fused_forward_no_padding_gen(True)\n    else:\n        # default to language model\n        return fused_forward_no_padding_gen(False)\n\n\n# ruff: noqa\n\n########################################################\n# below is the deprecated code\n########################################################\n\nfrom .config_converter import (\n    PretrainedConfig,\n    TransformerConfig,\n    hf_to_mcore_config_dense,\n    hf_to_mcore_config_dpskv3,\n    hf_to_mcore_config_llama4,\n    hf_to_mcore_config_mixtral,\n    hf_to_mcore_config_qwen2_5_vl,\n    hf_to_mcore_config_qwen2moe,\n    hf_to_mcore_config_qwen3moe,\n)\nfrom .model_initializer import (\n    BaseModelInitializer,\n    DeepseekV3Model,\n    DenseModel,\n    MixtralModel,\n    Qwen2MoEModel,\n    Qwen3MoEModel,\n    Qwen25VLModel,\n)\nfrom .weight_converter import (\n    McoreToHFWeightConverterDense,\n    McoreToHFWeightConverterDpskv3,\n    McoreToHFWeightConverterMixtral,\n    McoreToHFWeightConverterQwen2_5_VL,\n    McoreToHFWeightConverterQwen2Moe,\n    McoreToHFWeightConverterQwen3Moe,\n)\n\n\nclass SupportedModel(Enum):\n    LLAMA = \"LlamaForCausalLM\"  # tested\n    QWEN2 = \"Qwen2ForCausalLM\"  # tested\n    QWEN2_MOE = \"Qwen2MoeForCausalLM\"  # pending\n    DEEPSEEK_V3 = \"DeepseekV3ForCausalLM\"  # not tested\n    MIXTRAL = \"MixtralForCausalLM\"  # tested\n    QWEN2_5_VL = \"Qwen2_5_VLForConditionalGeneration\"  # not supported\n    LLAMA4 = \"Llama4ForConditionalGeneration\"  # not tested\n    QWEN3 = \"Qwen3ForCausalLM\"  # tested\n    QWEN3_MOE = \"Qwen3MoeForCausalLM\"  # tested\n    GLM4_MOE = \"Glm4MoeForCausalLM\"\n    QWEN3_TOKEN_CLASSIFICATION = \"Qwen3ForTokenClassification\"\n    LLAMA_TOKEN_CLASSIFICATION = \"LlamaForTokenClassification\"\n    QWEN3_MOE_VL = \"Qwen3VLMoeForConditionalGeneration\"\n    QWEN3_VL = \"Qwen3VLForConditionalGeneration\"\n    GPT_OSS = \"GptOssForCausalLM\"\n    MIMO = \"MiMoForCausalLM\"\n\n\n# Registry for model configuration converters\nMODEL_CONFIG_CONVERTER_REGISTRY: dict[SupportedModel, Callable[[PretrainedConfig, torch.dtype], TransformerConfig]] = {\n    SupportedModel.LLAMA: hf_to_mcore_config_dense,\n    SupportedModel.QWEN2: hf_to_mcore_config_dense,\n    SupportedModel.QWEN2_MOE: hf_to_mcore_config_qwen2moe,\n    SupportedModel.DEEPSEEK_V3: hf_to_mcore_config_dpskv3,\n    SupportedModel.MIXTRAL: hf_to_mcore_config_mixtral,\n    SupportedModel.QWEN2_5_VL: hf_to_mcore_config_qwen2_5_vl,\n    SupportedModel.LLAMA4: hf_to_mcore_config_llama4,\n    SupportedModel.QWEN3: hf_to_mcore_config_dense,\n    SupportedModel.QWEN3_MOE: hf_to_mcore_config_qwen3moe,\n    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: hf_to_mcore_config_dense,\n    SupportedModel.LLAMA_TOKEN_CLASSIFICATION: hf_to_mcore_config_dense,\n}\n\n# Registry for model initializers\nMODEL_INITIALIZER_REGISTRY: dict[SupportedModel, type[BaseModelInitializer]] = {\n    SupportedModel.LLAMA: DenseModel,\n    SupportedModel.QWEN2: DenseModel,\n    SupportedModel.QWEN2_MOE: Qwen2MoEModel,\n    SupportedModel.MIXTRAL: MixtralModel,\n    SupportedModel.DEEPSEEK_V3: DeepseekV3Model,\n    SupportedModel.QWEN2_5_VL: Qwen25VLModel,\n    SupportedModel.LLAMA4: DenseModel,\n    SupportedModel.QWEN3: DenseModel,\n    SupportedModel.QWEN3_MOE: Qwen3MoEModel,\n    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: DenseModel,\n    SupportedModel.LLAMA_TOKEN_CLASSIFICATION: DenseModel,\n}\n\n# Registry for model forward functions\nMODEL_FORWARD_REGISTRY: dict[SupportedModel, Callable] = {\n    SupportedModel.LLAMA: model_forward_gen(),\n    SupportedModel.QWEN2: model_forward_gen(),\n    SupportedModel.QWEN2_MOE: model_forward_gen(),\n    SupportedModel.MIXTRAL: model_forward_gen(),\n    SupportedModel.DEEPSEEK_V3: model_forward_gen(),\n    SupportedModel.LLAMA4: model_forward_gen(),\n    SupportedModel.QWEN3: model_forward_gen(),\n    SupportedModel.QWEN3_MOE: model_forward_gen(),\n    SupportedModel.QWEN2_5_VL: model_forward_gen(True),\n    SupportedModel.QWEN3_MOE_VL: model_forward_gen(True),\n    SupportedModel.QWEN3_VL: model_forward_gen(True),\n    SupportedModel.GLM4_MOE: model_forward_gen(),\n    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: model_forward_gen(),\n    SupportedModel.LLAMA_TOKEN_CLASSIFICATION: model_forward_gen(),\n    SupportedModel.GPT_OSS: model_forward_gen(),\n    SupportedModel.MIMO: model_forward_gen(),\n}\n\n# Registry for model forward functions\nMODEL_FORWARD_NOPAD_REGISTRY: dict[SupportedModel, Callable] = {\n    SupportedModel.LLAMA: gptmodel_forward_no_padding,\n    SupportedModel.QWEN2: gptmodel_forward_no_padding,\n    SupportedModel.QWEN2_MOE: gptmodel_forward_no_padding,\n    SupportedModel.MIXTRAL: gptmodel_forward_no_padding,\n    SupportedModel.DEEPSEEK_V3: gptmodel_forward_no_padding,\n    SupportedModel.QWEN2_5_VL: gptmodel_forward_no_padding,\n    SupportedModel.QWEN3_MOE_VL: gptmodel_forward_no_padding,\n    SupportedModel.QWEN3_VL: gptmodel_forward_no_padding,\n    SupportedModel.LLAMA4: gptmodel_forward_no_padding,\n    SupportedModel.QWEN3: gptmodel_forward_no_padding,\n    SupportedModel.QWEN3_MOE: gptmodel_forward_no_padding,\n    SupportedModel.GLM4_MOE: gptmodel_forward_no_padding,\n    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: gptmodel_forward_no_padding,\n    SupportedModel.LLAMA_TOKEN_CLASSIFICATION: gptmodel_forward_no_padding,\n    SupportedModel.GPT_OSS: gptmodel_forward_no_padding,\n    SupportedModel.MIMO: gptmodel_forward_no_padding,\n}\n\n# Registry for model forward functions\nMODEL_FORWARD_FUSED_REGISTRY: dict[SupportedModel, Callable] = {\n    SupportedModel.LLAMA: fused_forward_model_gen(),\n    SupportedModel.QWEN2: fused_forward_model_gen(),\n    SupportedModel.QWEN2_MOE: fused_forward_model_gen(),\n    SupportedModel.MIXTRAL: fused_forward_model_gen(),\n    SupportedModel.QWEN2_5_VL: fused_forward_model_gen(True),\n    SupportedModel.QWEN3_MOE_VL: fused_forward_model_gen(True),\n    SupportedModel.QWEN3_VL: fused_forward_model_gen(True),\n    SupportedModel.LLAMA4: fused_forward_model_gen(),\n    SupportedModel.QWEN3: fused_forward_model_gen(),\n    SupportedModel.QWEN3_MOE: fused_forward_model_gen(),\n    SupportedModel.DEEPSEEK_V3: fused_forward_model_gen(),\n    SupportedModel.GLM4_MOE: fused_forward_model_gen(),\n    SupportedModel.GPT_OSS: fused_forward_model_gen(),\n    SupportedModel.MIMO: fused_forward_model_gen(),\n}\n\n# Registry for model weight converters\nMODEL_WEIGHT_CONVERTER_REGISTRY: dict[SupportedModel, type] = {\n    SupportedModel.LLAMA: McoreToHFWeightConverterDense,\n    SupportedModel.QWEN2: McoreToHFWeightConverterDense,\n    SupportedModel.QWEN2_MOE: McoreToHFWeightConverterQwen2Moe,\n    SupportedModel.MIXTRAL: McoreToHFWeightConverterMixtral,\n    SupportedModel.DEEPSEEK_V3: McoreToHFWeightConverterDpskv3,\n    SupportedModel.QWEN3: McoreToHFWeightConverterDense,\n    SupportedModel.QWEN3_MOE: McoreToHFWeightConverterQwen3Moe,\n    SupportedModel.QWEN2_5_VL: McoreToHFWeightConverterQwen2_5_VL,\n    SupportedModel.QWEN3_TOKEN_CLASSIFICATION: McoreToHFWeightConverterDense,\n    SupportedModel.LLAMA_TOKEN_CLASSIFICATION: McoreToHFWeightConverterDense,\n}\n\n\ndef get_supported_model(model_type: str) -> SupportedModel:\n    try:\n        return SupportedModel(model_type)\n    except ValueError as err:\n        supported_models = [e.value for e in SupportedModel]\n        raise NotImplementedError(\n            f\"Model Type: {model_type} not supported. Supported models: {supported_models}\"\n        ) from err\n\n\ndef hf_to_mcore_config(\n    hf_config: PretrainedConfig, dtype: torch.dtype, **override_transformer_config_kwargs\n) -> TransformerConfig:\n    \"\"\"Convert huggingface PretrainedConfig to mcore TransformerConfig.\n\n    Args:\n        hf_config: The huggingface PretrainedConfig.\n        dtype: The dtype of the model.\n        **override_transformer_config_kwargs: The kwargs to override the transformer config.\n\n    Returns:\n        The mcore TransformerConfig.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    model = get_supported_model(hf_config.architectures[0])\n    return MODEL_CONFIG_CONVERTER_REGISTRY[model](hf_config, dtype, **override_transformer_config_kwargs)\n\n\ndef init_mcore_model(\n    tfconfig: TransformerConfig,\n    hf_config: PretrainedConfig,\n    pre_process: bool = True,\n    post_process: bool = None,\n    *,\n    share_embeddings_and_output_weights: bool = False,\n    value: bool = False,\n    **extra_kwargs,  # may be used for vlm and moe\n) -> nn.Module:\n    \"\"\"\n    Initialize a Mcore model.\n\n    Args:\n        tfconfig: The transformer config.\n        hf_config: The HuggingFace config.\n        pre_process: Optional pre-processing function.\n        post_process: Optional post-processing function.\n        share_embeddings_and_output_weights: Whether to share embeddings and output weights.\n        value: Whether to use value.\n        **extra_kwargs: Additional keyword arguments.\n\n    Returns:\n        The initialized model.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    model = get_supported_model(hf_config.architectures[0])\n    initializer_cls = MODEL_INITIALIZER_REGISTRY[model]\n    initializer = initializer_cls(tfconfig, hf_config)\n    return initializer.initialize(\n        pre_process=pre_process,\n        post_process=post_process,\n        share_embeddings_and_output_weights=share_embeddings_and_output_weights,\n        value=value,\n        **extra_kwargs,\n    )\n\n\ndef get_mcore_weight_converter(hf_config: PretrainedConfig, dtype: torch.dtype) -> Callable:\n    \"\"\"\n    Get the weight converter for given model architecture.\n    \"\"\"\n    assert len(hf_config.architectures) == 1, \"Only one architecture is supported for now\"\n    model = get_supported_model(hf_config.architectures[0])\n    tfconfig = hf_to_mcore_config(hf_config, dtype)\n    return MODEL_WEIGHT_CONVERTER_REGISTRY[model](hf_config, tfconfig)\n"
  },
  {
    "path": "verl/models/mcore/saver.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\nfrom megatron.core import mpu\nfrom megatron.core.distributed import DistributedDataParallel as LocalDDP\nfrom megatron.core.transformer.module import Float16Module\nfrom torch.nn.parallel import DistributedDataParallel as torchDDP\n\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.logger import print_rank_0\nfrom verl.utils.megatron_utils import unwrap_model\n\n\ndef _megatron_calc_global_rank(\n    tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0, cp_rank: int = 0, ep_rank: int = 0\n):\n    \"\"\"Calculate global rank with support for CP/EP parallelism\"\"\"\n\n    # Get parallel sizes for each dimension\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    dp_size = mpu.get_data_parallel_world_size()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    cp_size = mpu.get_context_parallel_world_size()\n    # ep_size = mpu.get_expert_model_parallel_world_size()\n\n    # Verify total GPU count matches (must be consistent with parallel_state.py)\n    total_size = tp_size * dp_size * pp_size * cp_size\n    assert total_size == torch.distributed.get_world_size(), (\n        f\"{tp_size}x{dp_size}x{pp_size}x{cp_size} != {torch.distributed.get_world_size()}\"\n    )\n\n    # Core calculation logic (corresponds to RankGenerator order parameter)\n    # Assumes default order is \"tp-cp-ep-dp-pp\"\n    return ((pp_rank * dp_size + dp_rank) * cp_size + cp_rank) * tp_size + tp_rank\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef merge_megatron_ckpt_gptmodel(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False):\n    \"\"\"Merge sharded parameters of a Megatron module into a merged checkpoint.\n\n    Args:\n        wrapped_models (list of megatron.core.distributed.DistributedDataParallel):\n            The local DDP wrapped megatron modules.\n        config (str or None):\n            HF config for model\n        dtype: model params type\n        is_value_model: if model is value model\n        tie_word_embeddings: tie_word_embeddings\n    Returns:\n        state_dict (dict):\n            The merged state_dict in rank 0, and an empty dictionary in other ranks.\n    \"\"\"\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    cp_rank = mpu.get_context_parallel_rank()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if dist.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        assert len(models[i].decoder.layers) == num_layers_per_model, (\n            \"len model layers {} not equal to num_layers_per_model {}\".format(\n                len(models[i].decoder.layers), num_layers_per_model\n            )\n        )\n\n    state_dict = dict()\n\n    def _get_cpu_tensor(tensor: torch.Tensor):\n        if tensor is None:\n            return None\n        if tensor.device == torch.device(\"cpu\"):\n            return tensor.detach().clone()\n        return tensor.detach().cpu()\n\n    def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n\n        if torch.distributed.get_rank() == src_rank:\n            if tensor is None:\n                weight = None\n                tensor_shape = None\n            else:\n                weight = tensor\n                tensor_shape = weight.shape\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not exist, skip collect\")\n            return\n\n        if weight is None:\n            weight = torch.empty(\n                tensor_shape,\n                dtype=dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n\n        dist.broadcast(weight, src=src_rank, group=mp_group)\n\n        if torch.distributed.get_rank() == 0:\n            state_dict[name] = _get_cpu_tensor(weight)\n\n    def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        # tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=concat_dim)\n            if mutate_func is not None:\n                full_tensor = mutate_func(full_tensor)\n            state_dict[name] = full_tensor\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        # tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            intermediate_size_tp = config.intermediate_size // tp_size\n            gate_weight_list = []\n            up_weight_list = []\n            for i in range(tp_size):\n                gate_up_weight_tp = full_tensor[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)]\n                gate_weight_tp = gate_up_weight_tp[:intermediate_size_tp]\n                up_weight_tp = gate_up_weight_tp[intermediate_size_tp:]\n                gate_weight_list.append(gate_weight_tp)\n                up_weight_list.append(up_weight_tp)\n\n            state_dict[gate_name] = torch.cat(gate_weight_list, dim=0)\n            state_dict[up_name] = torch.cat(up_weight_list, dim=0)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank):\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        # tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank, cp_rank=cp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            q_weight_list = []\n            k_weight_list = []\n            v_weight_list = []\n            hidden_size_per_head = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    num_query_groups_per_partition = wrapped_models[0].config.num_query_groups // tp_size\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_size_chunk = q_size_tp // num_query_groups_per_partition\n                    kv_size_chunk = kv_size_tp // num_query_groups_per_partition\n                    for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition):\n                        q_part = qkv_part_chunk[:q_size_chunk]\n                        k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk]\n                        v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :]\n                        q_weight_list.append(q_part)\n                        k_weight_list.append(k_part)\n                        v_weight_list.append(v_part)\n            else:\n                q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    num_query_groups_per_partition = wrapped_models[0].config.num_query_groups // tp_size\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_size_chunk = q_size_tp // num_query_groups_per_partition\n                    kv_size_chunk = kv_size_tp // num_query_groups_per_partition\n                    for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition):\n                        q_part = qkv_part_chunk[:q_size_chunk]\n                        k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk]\n                        v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :]\n                        q_weight_list.append(q_part)\n                        if i * config.num_key_value_heads % tp_size == 0:\n                            k_weight_list.append(k_part)\n                            v_weight_list.append(v_part)\n\n            state_dict[q_name] = torch.cat(q_weight_list, dim=0)\n            state_dict[k_name] = torch.cat(k_weight_list, dim=0)\n            state_dict[v_name] = torch.cat(v_weight_list, dim=0)\n\n    # empty cache before collecting weights\n    get_torch_device().empty_cache()\n    # Embeddings\n    # -------------------\n    if dp_rank == 0 and cp_rank == 0:  # models are identical across cp ranks\n        # Embeddings\n        # -------------------\n        print_rank_0(\"collecting embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        _broadcast_tp_shard_tensor(\n            gpt_model_module.embedding.word_embeddings.weight if pp_rank == 0 else None,\n            \"model.embed_tokens.weight\",\n            src_pp_rank=0,\n        )\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n        for layer in range(config.num_hidden_layers):\n            print_rank_0(f\"collecting layer #{layer}...\")\n            layer_name = f\"model.layers.{layer}\"\n            src_pp_rank, src_virtual_pp_rank, src_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[src_virtual_pp_rank])\n            sync_layer = gpt_model_module.decoder.layers[src_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.self_attention.linear_qkv.layer_norm_weight,\n                f\"{layer_name}.input_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            if gpt_model_module.config.qk_layernorm:\n                _broadcast_tensor(\n                    sync_layer.self_attention.q_layernorm.weight,\n                    f\"{layer_name}.self_attn.q_norm.weight\",\n                    src_pp_rank=src_pp_rank,\n                )\n                _broadcast_tensor(\n                    sync_layer.self_attention.k_layernorm.weight,\n                    f\"{layer_name}.self_attn.k_norm.weight\",\n                    src_pp_rank=src_pp_rank,\n                )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attention.linear_qkv.weight,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            if gpt_model_module.config.add_qkv_bias:\n                _broadcast_tp_shard_tensor_qkv(\n                    sync_layer.self_attention.linear_qkv.bias,\n                    f\"{layer_name}.self_attn.q_proj.bias\",\n                    f\"{layer_name}.self_attn.k_proj.bias\",\n                    f\"{layer_name}.self_attn.v_proj.bias\",\n                    src_pp_rank=src_pp_rank,\n                )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attention.linear_proj.weight,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tensor(\n                sync_layer.mlp.linear_fc1.layer_norm_weight,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.linear_fc1.weight,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.linear_fc2.weight,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"collecting final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.decoder.final_layernorm, \"weight\", None),\n            \"model.norm.weight\",\n            src_pp_rank=pp_size - 1,\n        )\n\n        if tie_word_embeddings:\n            print_rank_0(\"tie word embedding skip load lm_head...\")\n        else:\n            print_rank_0(\"collecting lm_head...\")\n\n            if is_value_model:\n                lm_head_weight = None\n                if pp_rank == pp_size - 1:\n                    lm_head_weight = getattr(gpt_model_module.output_layer, \"weight\", None)\n                _broadcast_tensor(lm_head_weight, \"lm_head.weight\", src_pp_rank=pp_size - 1)\n\n            else:\n                _broadcast_tp_shard_tensor(\n                    getattr(gpt_model_module.output_layer, \"weight\", None) if pp_rank == pp_size - 1 else None,\n                    \"lm_head.weight\",\n                    src_pp_rank=pp_size - 1,\n                )\n\n    dist.barrier()\n    get_torch_device().empty_cache()\n    if torch.distributed.get_rank() == 0:\n        for k, v in state_dict.items():\n            if dtype != v.dtype:\n                state_dict[k] = v.to(dtype)\n\n    print_rank_0(f\"merge megatron ckpt done, time elapsed {time.time() - start_time}s\")\n    return state_dict\n\n\ndef merge_megatron_ckpt_gptmodel_qwen_moe(\n    wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False\n):\n    raise NotImplementedError(\"merge_megatron_ckpt_gptmodel_qwen_moe is not implemented\")\n\n\ndef merge_megatron_ckpt_gptmodel_qwen2_5_vl(\n    wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False\n):\n    raise NotImplementedError(\"merge_megatron_ckpt_gptmodel_qwen2_5_vl is not implemented\")\n\n\ndef merge_megatron_ckpt_gptmodel_dpskv3(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False):\n    raise NotImplementedError(\"merge_megatron_ckpt_gptmodel_dpskv3 is not implemented\")\n\n\ndef merge_megatron_ckpt_gptmodel_mixtral(\n    wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False\n):\n    raise NotImplementedError(\"merge_megatron_ckpt_gptmodel_mixtral is not implemented\")\n"
  },
  {
    "path": "verl/models/mcore/util.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport math\nimport os\nfrom typing import Optional\n\nimport torch\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.packed_seq_params import PackedSeqParams\n\nfrom verl.utils.model import CausalLMOutputForPPO\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef _compute_fp8_thd_align_size(align_size: int) -> tuple[int, int]:\n    \"\"\"Compute FP8 alignment sizes for thd-format sequences.\n\n    For FP8 block quantization, each sequence must be padded to a multiple of\n    lcm(16, align_size), and the total padded length must be divisible by\n    (align_size * 128) for TransformerEngine compatibility.\n\n    Returns (per_seq_align_size, total_align_size).\n    \"\"\"\n    return math.lcm(16, align_size), align_size * 128\n\n\ndef preprocess_packed_seqs(\n    input_ids: torch.Tensor, attention_mask: torch.Tensor, pre_process: bool = True, use_fp8_padding: bool = False\n) -> tuple[torch.Tensor, PackedSeqParams]:\n    \"\"\"\n    Preprocess packed sequences\n    CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1\n    gets second and second last chunks, and so on), this is for load balancing with causal masking.\n    See https://github.com/NVIDIA/TransformerEngine/issues/1368\n    \"\"\"\n    batch_size = input_ids.shape[0]\n\n    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    cp_size = mpu.get_context_parallel_world_size()\n    cp_rank = mpu.get_context_parallel_rank()\n    align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size\n    if use_fp8_padding:\n        per_seq_align, total_align = _compute_fp8_thd_align_size(align_size)\n        align_size = per_seq_align\n\n    pad_size = (align_size - seqlens_in_batch % align_size) % align_size\n    seqlens_in_batch_padded = seqlens_in_batch + pad_size\n\n    cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device)\n    cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0)\n    cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device)\n    cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0)\n\n    if use_fp8_padding:\n        pad_size_last = (total_align - cu_seqlens_padded[-1] % total_align) % total_align\n        cu_seqlens_padded[-1] += pad_size_last\n        seqlens_in_batch_padded[-1] += pad_size_last\n\n    # ----------------------------------------------------------------------------\n    # Move the index information needed in the subsequent loop to the CPU at once,\n    # to avoid frequent .item() calls in the loop that cause D2H synchronization\n    # ----------------------------------------------------------------------------\n    seqlens_in_batch_cpu: list[int] = seqlens_in_batch.tolist()  # original valid lengths\n    seqlens_in_batch_padded_cpu: list[int] = seqlens_in_batch_padded.tolist()  # lengths after padding\n    cu_seqlens_padded_cpu: list[int] = cu_seqlens_padded.tolist()  # start positions (after padding)\n\n    # Pure Python int calculation to avoid further synchronization\n    max_seqlen_in_batch = max(seqlens_in_batch_padded_cpu)\n\n    shape = list(input_ids.shape[1:])\n    shape[0] = sum(seqlens_in_batch_padded_cpu) // cp_size\n    if pre_process:\n        input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device)\n        for i in range(batch_size):\n            # Use Python int, so no GPU→CPU sync in the loop\n            if cp_size <= 1:\n                seqlen = seqlens_in_batch_cpu[i]\n                start_idx = cu_seqlens_padded_cpu[i]\n                input_ids_rmpad[start_idx : start_idx + seqlen] = input_ids[i, attention_mask[i]]\n                continue\n\n            seqlen_padded_i = seqlens_in_batch_padded_cpu[i]\n            seqlen = seqlen_padded_i // cp_size\n            half_seqlen = seqlen // 2\n            start_idx = cu_seqlens_padded_cpu[i] // cp_size\n            # split to 2 chunks\n            d = input_ids[i, attention_mask[i]]\n            input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[\n                half_seqlen * cp_rank : half_seqlen * (cp_rank + 1)\n            ]\n\n            remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1)\n            remain_end = seqlen_padded_i - half_seqlen * cp_rank\n            remain_end = min(remain_end, d.shape[0])\n            remain_len = remain_end - remain_start\n            if remain_len > 0:\n                input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[\n                    remain_start:remain_end\n                ]\n\n    packed_seq_params = PackedSeqParams(\n        qkv_format=\"thd\",\n        cu_seqlens_q=cu_seqlens_padded,\n        max_seqlen_q=max_seqlen_in_batch,\n        cu_seqlens_kv=cu_seqlens_padded,\n        max_seqlen_kv=max_seqlen_in_batch,\n        cu_seqlens_q_padded=cu_seqlens_padded,\n        cu_seqlens_kv_padded=cu_seqlens_padded,\n    )\n    if pre_process:\n        return input_ids_rmpad.unsqueeze(0), packed_seq_params\n    else:\n        return input_ids, packed_seq_params\n\n\ndef postprocess_packed_seqs(\n    output: torch.Tensor,\n    packed_seq_params: PackedSeqParams,\n    attention_mask: torch.Tensor,\n    batch_size: int,\n    seq_len: int,\n    post_process: bool = True,\n) -> torch.Tensor:\n    \"\"\"\n    Postprocess packed sequences\n    \"\"\"\n    if not post_process:\n        return output\n\n    # -------------------------------------------------------------------------\n    # Move the lengths and offsets needed for subsequent Python-level indexing to the CPU in advance,\n    # to avoid a large number of .item() calls in the loop\n    # -------------------------------------------------------------------------\n    cu_padded_cpu: list[int] = packed_seq_params.cu_seqlens_q_padded.tolist()\n    seq_lens_cpu: list[int] = attention_mask.sum(dim=1, dtype=torch.int32).cpu().tolist()\n\n    shape = [batch_size, seq_len] + list(output.shape[2:])  # 1,packed, dim -> batch_size, seq_len, dim\n    output_new = torch.zeros(shape, dtype=output.dtype, device=output.device)\n\n    cp_size = mpu.get_context_parallel_world_size()\n    # all gather output across context parallel group\n    if cp_size > 1:\n        # output shape: [1, packed_len, hidden_dim]\n        # need to gather across cp group and concatenate in sequence dimension\n        output_list = [torch.empty_like(output, dtype=output.dtype) for _ in range(cp_size)]\n        torch.distributed.all_gather(output_list, output.detach(), group=mpu.get_context_parallel_group())\n        output_list[mpu.get_context_parallel_rank()] = output\n    else:\n        output_list = [output]\n    for i in range(batch_size):\n        if cp_size <= 1:\n            s = seq_lens_cpu[i]\n            start_idx = cu_padded_cpu[i]\n            output_new[i, attention_mask[i]] = output[0][start_idx : start_idx + s]\n            continue\n        s_len_padded_chunk = (cu_padded_cpu[i + 1] - cu_padded_cpu[i]) // cp_size\n        half_seqlen = s_len_padded_chunk // 2\n        s_len = seq_lens_cpu[i]\n        s_len_padded = s_len_padded_chunk * cp_size\n        tmp = torch.empty(s_len_padded, *output.shape[2:], device=output.device, dtype=output.dtype)\n        for j in range(cp_size):\n            o = output_list[j][0]\n            # split to 2 chunks\n            packed_start_idx = cu_padded_cpu[i] // cp_size\n            o0, o1 = (\n                o[packed_start_idx : packed_start_idx + half_seqlen],\n                o[packed_start_idx + half_seqlen : packed_start_idx + s_len_padded_chunk],\n            )\n            tmp[j * half_seqlen : (j + 1) * half_seqlen] = o0\n            tmp[s_len_padded - (j + 1) * half_seqlen : s_len_padded - j * half_seqlen] = o1\n        output_new[i, attention_mask[i]] = tmp[:s_len]\n\n    return output_new\n\n\ndef preprocess_bshd(\n    input_ids: torch.Tensor,\n    attention_mask: torch.Tensor,\n    position_ids: torch.Tensor,\n    sequence_parallel: bool = False,\n    pre_process: bool = True,\n):\n    \"\"\"\n    Remove left padding from input_ids, attention_mask and position_ids\n    return new_input_ids, new_attention_mask, new_position_ids\n    \"\"\"\n    assert attention_mask.ndim == 2\n    assert position_ids.ndim == 2\n    cp_size = mpu.get_context_parallel_world_size()\n    assert cp_size == 1, \"Context parallel size without seq_pack is not supported\"\n    batch_size = input_ids.shape[0]\n    shape = list(input_ids.shape)  # batch_size, seq_len,...\n    seq_lens = attention_mask.sum(dim=1)\n    seq_len = seq_lens.max().item()\n    if sequence_parallel:\n        sp_world_size = mpu.get_tensor_model_parallel_world_size()\n        pad_size = (sp_world_size - seq_len % sp_world_size) % sp_world_size\n        seq_len = seq_len + pad_size\n    shape[1] = seq_len\n    if pre_process:\n        new_input_ids = torch.zeros(dtype=input_ids.dtype, device=input_ids.device, size=shape)\n    new_attention_mask = torch.zeros(\n        dtype=attention_mask.dtype, device=attention_mask.device, size=(batch_size, seq_len)\n    )\n    new_position_ids = torch.zeros(dtype=position_ids.dtype, device=position_ids.device, size=(batch_size, seq_len))\n    for i in range(batch_size):\n        if pre_process:\n            new_input_ids[i, : seq_lens[i]] = input_ids[i, attention_mask[i]]\n        new_attention_mask[i, : seq_lens[i]] = attention_mask[i, attention_mask[i]]\n        new_position_ids[i, : seq_lens[i]] = position_ids[i, attention_mask[i]]\n    if pre_process:\n        return new_input_ids, new_attention_mask, new_position_ids\n    else:\n        return input_ids, new_attention_mask, new_position_ids\n\n\ndef postprocess_bshd(\n    result,\n    attention_mask: torch.Tensor,\n    original_attention_mask: torch.Tensor,\n    origin_seqlen: int,\n    post_process: bool = True,\n):\n    \"\"\"\n    Recover left padding from result\n    return result\n    \"\"\"\n    if not post_process:\n        return result\n    shape = list(result.shape)\n    batch_size = shape[0]\n    shape[1] = origin_seqlen\n    new_result = torch.zeros(dtype=result.dtype, device=result.device, size=shape)\n    for i in range(batch_size):\n        new_result[i, original_attention_mask[i]] = result[i, attention_mask[i]]\n    return new_result\n\n\ndef postprocess_packed_seqs_for_dict_output(\n    labels_mask: torch.Tensor,\n    output: CausalLMOutputForPPO,\n    packed_seq_params: PackedSeqParams,\n    attention_mask: torch.Tensor,\n    batch_size: int,\n    seq_len: int,\n    post_process: bool = True,\n) -> dict[str, torch.Tensor]:\n    \"\"\"_summary_\n    For fused kernels, the output is a dictionary with keys like 'log_probs', 'entropy', etc.\n    This function post-processes each tensor in the output dictionary.\n    Args:\n        output (CausalLMOutputForPPO): _description_\n        packed_seq_params (PackedSeqParams): _description_\n        attention_mask (torch.Tensor): _description_\n        batch_size (int): _description_\n        seq_len (int): _description_\n        post_process (bool, optional): _description_. Defaults to True.\n    Returns:\n        CausalLMOutputForPPO: _description_\n    \"\"\"\n    ret = {}\n    output.entropy = output.entropy.view(1, -1)\n    output.log_probs = output.log_probs.view(1, -1)\n    output.log_probs = output.log_probs.masked_fill(~labels_mask, 0.0)\n    ret[\"entropy\"] = postprocess_packed_seqs(\n        output.entropy, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process\n    )\n    ret[\"log_probs\"] = postprocess_packed_seqs(\n        output.log_probs, packed_seq_params, attention_mask, batch_size, seq_len, post_process=post_process\n    )\n    return ret\n\n\n### No padding versions for model engine\n### inputs are nested tensors\ndef preprocess_thd_no_padding(\n    input_ids: torch.Tensor, pre_process: bool = True, need_roll: bool = False, use_fp8_padding: bool = False\n) -> tuple[torch.Tensor, PackedSeqParams, Optional[torch.Tensor]]:\n    \"\"\"\n    Preprocess packed sequences\n    CP splits sequence into CP*2 chunks, and each GPU gets 2 chunks (GPU0 gets first and last chunks, GPU1\n    gets second and second last chunks, and so on), this is for load balancing with causal masking.\n    See https://github.com/NVIDIA/TransformerEngine/issues/1368\n    \"\"\"\n    batch_size = input_ids.shape[0]\n\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    cp_size = mpu.get_context_parallel_world_size()\n    cp_rank = mpu.get_context_parallel_rank()\n    align_size = tp_size * cp_size * 2 if cp_size > 1 else tp_size\n    seqlens_in_batch = input_ids.offsets().diff()\n\n    if use_fp8_padding:\n        per_seq_align, total_align = _compute_fp8_thd_align_size(align_size)\n        align_size = per_seq_align\n\n    pad_size = (align_size - seqlens_in_batch % align_size) % align_size\n    seqlens_in_batch_padded = seqlens_in_batch + pad_size\n\n    cu_seqlens = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device)\n    cu_seqlens[1:] = torch.cumsum(seqlens_in_batch, dim=0)\n    cu_seqlens_padded = torch.zeros(batch_size + 1, dtype=torch.int32, device=input_ids.device)\n    cu_seqlens_padded[1:] = torch.cumsum(seqlens_in_batch_padded, dim=0)\n\n    if use_fp8_padding:\n        # Pad the last sequence so total length is divisible by total_align for TE\n        pad_size_last = (total_align - cu_seqlens_padded[-1] % total_align) % total_align\n        cu_seqlens_padded[-1] += pad_size_last\n        seqlens_in_batch_padded[-1] += pad_size_last\n\n    # ----------------------------------------------------------------------------\n    # Move the index information needed in the subsequent loop to the CPU at once,\n    # to avoid frequent .item() calls in the loop that cause D2H synchronization\n    # ----------------------------------------------------------------------------\n    seqlens_in_batch_cpu: list[int] = seqlens_in_batch.tolist()  # original valid lengths\n    seqlens_in_batch_padded_cpu: list[int] = seqlens_in_batch_padded.tolist()  # lengths after padding\n    cu_seqlens_padded_cpu: list[int] = cu_seqlens_padded.tolist()  # start positions (after padding)\n\n    # Pure Python int calculation to avoid further synchronization\n    max_seqlen_in_batch = max(seqlens_in_batch_padded_cpu)\n\n    shape = list(input_ids.shape[1:])\n    shape[0] = sum(seqlens_in_batch_padded_cpu) // cp_size\n    if pre_process:\n        input_ids_rmpad = torch.zeros(shape, dtype=input_ids.dtype, device=input_ids.device)\n        position_ids_rmpad = torch.zeros(shape, dtype=torch.long, device=input_ids.device)\n        if need_roll:\n            saved_roll_dict = {}\n            saved_position_roll_dict = {}\n        for i in range(batch_size):\n            # Use Python int, so no GPU→CPU sync in the loop\n            if cp_size <= 1:\n                seqlen = seqlens_in_batch_cpu[i]\n                start_idx = cu_seqlens_padded_cpu[i]\n                input_ids_rmpad[start_idx : start_idx + seqlen] = input_ids[i]\n                # Build position_ids: 0, 1, 2, ..., seqlen-1 for this sequence\n                position_ids_rmpad[start_idx : start_idx + seqlen] = torch.arange(\n                    seqlen, dtype=torch.long, device=input_ids.device\n                )\n                continue\n\n            seqlen_padded_i = seqlens_in_batch_padded_cpu[i]\n            seqlen = seqlen_padded_i // cp_size\n            half_seqlen = seqlen // 2\n            start_idx = cu_seqlens_padded_cpu[i] // cp_size\n            # split to 2 chunks\n            d = input_ids[i]\n            # If the number of elements in `d` is smaller than the required\n            # alignment size, pad the tensor with zeros so that its total\n            # length matches `align_size`. This ensures size alignment for\n            # downstream operations (e.g., communication or memory alignment).\n            if d.numel() < align_size:\n                original_size = d.numel()\n                pad = torch.zeros(align_size - d.numel(), dtype=d.dtype, device=d.device)\n                d = torch.cat([d, pad], dim=0)\n                logger.warning_once(\n                    f\"Padding tensor for context parallel alignment, original_size={original_size}, \"\n                    f\"align_size={align_size}\"\n                )\n\n            input_ids_rmpad[start_idx : start_idx + half_seqlen] = d[\n                half_seqlen * cp_rank : half_seqlen * (cp_rank + 1)\n            ]\n\n            # Build position_ids for the first chunk\n            position_ids_rmpad[start_idx : start_idx + half_seqlen] = torch.arange(\n                half_seqlen * cp_rank, half_seqlen * (cp_rank + 1), dtype=torch.long, device=input_ids.device\n            )\n\n            remain_start = seqlen_padded_i - half_seqlen * (cp_rank + 1)\n            remain_end = seqlen_padded_i - half_seqlen * cp_rank\n            remain_end = min(remain_end, d.shape[0])\n            remain_len = remain_end - remain_start\n            if remain_len > 0:\n                input_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = d[\n                    remain_start:remain_end\n                ]\n                # Build position_ids for the remaining chunk\n                position_ids_rmpad[start_idx + half_seqlen : start_idx + half_seqlen + remain_len] = torch.arange(\n                    seqlen_padded_i - remain_len, seqlen_padded_i, dtype=torch.long, device=input_ids.device\n                )\n\n            if need_roll:\n                # Handle roll for cp_size > 1 case\n                saved_roll_dict[start_idx + half_seqlen - 1] = d[(cp_rank + 1) * half_seqlen]\n                saved_position_roll_dict[start_idx + half_seqlen - 1] = position_ids_rmpad[start_idx + half_seqlen - 1]\n                if remain_len > 0:\n                    if remain_end == d.shape[0]:\n                        saved_roll_dict[start_idx + half_seqlen + remain_len - 1] = d[0]\n                        saved_position_roll_dict[start_idx + half_seqlen + remain_len - 1] = 0\n                    else:\n                        saved_roll_dict[start_idx + half_seqlen + remain_len - 1] = d[remain_end]\n                        saved_position_roll_dict[start_idx + half_seqlen + remain_len - 1] = position_ids_rmpad[\n                            start_idx + half_seqlen + remain_len - 1\n                        ]\n\n        if need_roll:\n            input_ids_rmpad = torch.roll(input_ids_rmpad, shifts=-1, dims=0)\n            position_ids_rmpad = torch.roll(position_ids_rmpad, shifts=-1, dims=0)\n            if len(saved_roll_dict) > 0:\n                for k, v in saved_roll_dict.items():\n                    input_ids_rmpad[k] = v\n                for k, v in saved_position_roll_dict.items():\n                    position_ids_rmpad[k] = v\n\n    packed_seq_params = PackedSeqParams(\n        qkv_format=\"thd\",\n        cu_seqlens_q=cu_seqlens_padded,\n        max_seqlen_q=max_seqlen_in_batch,\n        cu_seqlens_kv=cu_seqlens_padded,\n        max_seqlen_kv=max_seqlen_in_batch,\n        cu_seqlens_q_padded=cu_seqlens_padded,\n        cu_seqlens_kv_padded=cu_seqlens_padded,\n    )\n    if pre_process:\n        return input_ids_rmpad.unsqueeze(0), packed_seq_params, position_ids_rmpad.unsqueeze(0)\n    else:\n        return input_ids, packed_seq_params, None\n\n\ndef postprocess_thd_no_padding(\n    output: torch.Tensor,\n    packed_seq_params: PackedSeqParams,\n    input_ids: torch.Tensor,\n    batch_size: int,\n    post_process: bool = True,\n) -> torch.Tensor:\n    \"\"\"\n    Postprocess packed sequences\n    \"\"\"\n    if not post_process:\n        return output\n\n    # -------------------------------------------------------------------------\n    # Move the lengths and offsets needed for subsequent Python-level indexing to the CPU in advance,\n    # to avoid a large number of .item() calls in the loop\n    # -------------------------------------------------------------------------\n    cu_padded_cpu: list[int] = packed_seq_params.cu_seqlens_q_padded.tolist()\n    # The reason why we use input_ids.offsets() instead of packed_seq_params.cu_seqlens_q.diff()\n    # is that the latter one is the padded length, while the former one is the original length.\n    cu_seqlens = input_ids.offsets()\n    seq_lens_cpu: list[int] = cu_seqlens.diff().tolist()\n\n    output_new = []\n\n    cp_size = mpu.get_context_parallel_world_size()\n    # all gather output across context parallel group\n    if cp_size > 1:\n        # output shape: [1, packed_len, hidden_dim]\n        # need to gather across cp group and concatenate in sequence dimension\n        output_list = [torch.empty_like(output) for _ in range(cp_size)]\n        torch.distributed.all_gather(output_list, output.detach(), group=mpu.get_context_parallel_group())\n        output_list[mpu.get_context_parallel_rank()] = output\n    else:\n        output_list = [output]\n\n    for i in range(batch_size):\n        if cp_size <= 1:\n            s = seq_lens_cpu[i]\n            start_idx = cu_padded_cpu[i]\n            output_new.append(output[0][start_idx : start_idx + s])\n            continue\n        s_len_padded_chunk = (cu_padded_cpu[i + 1] - cu_padded_cpu[i]) // cp_size\n        half_seqlen = s_len_padded_chunk // 2\n        s_len = seq_lens_cpu[i]\n        s_len_padded = s_len_padded_chunk * cp_size\n        tmp = torch.empty(s_len_padded, *output.shape[2:], device=output.device)\n        for j in range(cp_size):\n            o = output_list[j][0]\n            # split to 2 chunks\n            packed_start_idx = cu_padded_cpu[i] // cp_size\n            o0, o1 = (\n                o[packed_start_idx : packed_start_idx + half_seqlen],\n                o[packed_start_idx + half_seqlen : packed_start_idx + s_len_padded_chunk],\n            )\n            tmp[j * half_seqlen : (j + 1) * half_seqlen] = o0\n            tmp[s_len_padded - (j + 1) * half_seqlen : s_len_padded - j * half_seqlen] = o1\n        output_new.append(tmp[:s_len])\n\n    output_new_tensor = torch.nested.as_nested_tensor(output_new, layout=torch.jagged)\n\n    return output_new_tensor\n\n\ndef preprocess_bshd_no_padding(\n    input_ids: torch.Tensor, pre_process: bool = True, need_roll: bool = False, use_fp8_padding: bool = False\n):\n    \"\"\"\n    Preprocess bshd sequences\n    return \"input_ids, attention_mask, position_ids\"\n    \"\"\"\n    cp_size = mpu.get_context_parallel_world_size()\n    # TODO: support context parallel size > 1\n    assert cp_size == 1, \"Context parallel size without bshd is not supported yet\"\n\n    batch_size = input_ids.shape[0]\n    seqlens_in_batch = input_ids.offsets().diff()\n    max_seqlen = seqlens_in_batch.max().item()\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    if tp_size > 1:\n        sp_world_size = tp_size\n        pad_size = (sp_world_size - max_seqlen % sp_world_size) % sp_world_size\n        max_seqlen = max_seqlen + pad_size\n    if use_fp8_padding:\n        # For FP8 block quantization, batch_size * max_seqlen / tp_size must be divisible by 128.\n        # We need: max_seqlen % tp_size == 0 (for SP) AND batch_size * max_seqlen % (128 * tp_size) == 0.\n        # Compute the required alignment for max_seqlen:\n        fp8_total_align = 128 * tp_size\n        fp8_seq_align = fp8_total_align // math.gcd(batch_size, fp8_total_align)\n        # Also ensure tp alignment for SP\n        fp8_seq_align = math.lcm(fp8_seq_align, tp_size)\n        max_seqlen = ((max_seqlen + fp8_seq_align - 1) // fp8_seq_align) * fp8_seq_align\n\n    attention_mask = torch.zeros(batch_size, max_seqlen, dtype=torch.bool, device=input_ids.device)\n    input_ids_bshd = torch.zeros(batch_size, max_seqlen, dtype=input_ids.dtype, device=input_ids.device)\n    for i in range(batch_size):\n        attention_mask[i, : seqlens_in_batch[i]] = True\n        input_ids_bshd[i, : seqlens_in_batch[i]] = input_ids[i]\n    position_ids = torch.arange(max_seqlen, dtype=torch.long, device=input_ids.device)\n    position_ids = position_ids.unsqueeze(0).expand_as(input_ids_bshd)\n    if need_roll:\n        input_ids_bshd = torch.roll(input_ids_bshd, shifts=-1, dims=1)\n\n    return input_ids_bshd, attention_mask, position_ids\n\n\ndef postprocess_bshd_no_padding(\n    output: torch.Tensor,\n    attention_mask: torch.Tensor,\n    post_process: bool = True,\n) -> torch.Tensor:\n    \"\"\"\n    Postprocess bshd sequences\n    \"\"\"\n    if not post_process:\n        return output\n\n    batch_size = output.shape[0]\n    output_new = []\n\n    for i in range(batch_size):\n        mask = attention_mask[i].bool()\n        output_new.append(output[i][mask])\n\n    output_new_tensor = torch.nested.as_nested_tensor(output_new, layout=torch.jagged)\n\n    return output_new_tensor\n"
  },
  {
    "path": "verl/models/mcore/weight_converter.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n# Copyright Amazon.com, Inc. or its affiliates. All Rights Reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# online convert mcore weight to pure huggingface weight, no any fusion\n# including format conversion and name mapping\n# not including resharding\nimport torch\nfrom megatron.core.transformer import TransformerConfig\nfrom transformers import PretrainedConfig\n\n\nclass McoreToHFWeightConverterBase:\n    def __init__(self, hf_config: PretrainedConfig, mcore_config: TransformerConfig):\n        self.hf_config = hf_config\n        self.mcore_config = mcore_config\n\n    def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> torch.Tensor:\n        raise NotImplementedError\n\n\nclass McoreToHFWeightConverterDense(McoreToHFWeightConverterBase):\n    def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # 'decoder.layers.0.self_attention.linear_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_qkv.layer_norm_weight'\n        # 'decoder.layers.0.self_attention.linear_qkv.weight'\n        # 'decoder.layers.0.self_attention.linear_qkv.bias'\n        layer_number = name.split(\".\")[2]\n        convert_names = []\n        if \"self_attention.linear_qkv.bias\" in name or \"self_attention.linear_qkv.weight\" in name:\n            param_type = name.split(\".\")[-1]\n            assert param_type == \"bias\" or param_type == \"weight\"\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.q_proj.{param_type}\")\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.k_proj.{param_type}\")\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.v_proj.{param_type}\")\n            assert len(params) == 3\n        elif \"self_attention.linear_proj.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.o_proj.weight\")\n            assert len(params) == 1\n        elif \"self_attention.linear_qkv.layer_norm_weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.input_layernorm.weight\")\n            assert len(params) == 1\n        elif \"self_attention.q_layernorm.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.q_norm.weight\")\n            assert len(params) == 1\n        elif \"self_attention.k_layernorm.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.self_attn.k_norm.weight\")\n            assert len(params) == 1\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # 'decoder.layers.0.mlp.linear_fc1.layer_norm_weight'\n        # 'decoder.layers.0.mlp.linear_fc1.weight'\n        # 'decoder.layers.0.mlp.linear_fc2.weight'\n        layer_number = name.split(\".\")[2]\n        convert_names = []\n        if \"mlp.linear_fc1.weight\" in name:\n            # split gate_proj and up_proj\n            convert_names.append(f\"model.layers.{layer_number}.mlp.gate_proj.weight\")\n            convert_names.append(f\"model.layers.{layer_number}.mlp.up_proj.weight\")\n            assert len(params) == 2\n        elif \"mlp.linear_fc1.layer_norm_weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.post_attention_layernorm.weight\")\n            assert len(params) == 1\n        elif \"mlp.linear_fc2.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.mlp.down_proj.weight\")\n            assert len(params) == 1\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n\n    def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        direct_name_mapping = {\n            \"embedding.word_embeddings.weight\": \"model.embed_tokens.weight\",\n            \"decoder.final_layernorm.weight\": \"model.norm.weight\",\n            \"output_layer.weight\": \"lm_head.weight\",\n        }\n        if name in direct_name_mapping:\n            return [direct_name_mapping[name]], [params_one_group[0]]\n\n        if \"self_attention\" in name:\n            return self._convert_attention_param(name, params_one_group)\n        elif \"mlp\" in name:\n            return self._convert_mlp_param(name, params_one_group)\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n\n\nclass McoreToHFWeightConverterQwen2Moe(McoreToHFWeightConverterDense):\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # 'decoder.layers.0.pre_mlp_layernorm.weight',\n        # 'decoder.layers.0.mlp.router.weight',\n        # 'decoder.layers.0.mlp.shared_experts.gate_weight',\n        # 'decoder.layers.0.mlp.shared_experts.linear_fc1.weight',\n        # 'decoder.layers.0.mlp.shared_experts.linear_fc2.weight'\n        # moe1\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight0',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight1',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight2',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight3',\n        # moe2\n        # 'decoder.layers.0.mlp.experts.linear_fc2.weight0',\n        # 'decoder.layers.0.mlp.experts.linear_fc2.weight1',\n        layer_number = name.split(\".\")[2]\n        convert_names = []\n        if \"pre_mlp_layernorm\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.post_attention_layernorm.weight\")\n            assert len(params) == 1\n        elif \"mlp.router.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.mlp.gate.weight\")\n            assert len(params) == 1\n        elif \"shared_experts.gate_weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.mlp.shared_expert_gate.weight\")\n            assert len(params) == 1\n        elif \"shared_experts.linear_fc1.weight\" in name:  # split gate_proj and up_proj\n            convert_names.append(f\"model.layers.{layer_number}.mlp.shared_expert.gate_proj.weight\")\n            convert_names.append(f\"model.layers.{layer_number}.mlp.shared_expert.up_proj.weight\")\n            assert len(params) == 2\n        elif \"shared_experts.linear_fc2.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.mlp.shared_expert.down_proj.weight\")\n            assert len(params) == 1\n        elif \"mlp.experts.linear_fc1\" in name:  # split gate_proj and up_proj\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight\")\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight\")\n            assert len(params) == 2\n        elif \"mlp.experts.linear_fc2\" in name:\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight\")\n            assert len(params) == 1\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n\n\nclass McoreToHFWeightConverterQwen2_5_VL(McoreToHFWeightConverterDense):\n    def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        direct_name_mapping = {\n            \"language_model.embedding.word_embeddings.weight\": \"model.embed_tokens.weight\",\n            \"language_model.decoder.final_layernorm.weight\": \"model.norm.weight\",\n            \"language_model.output_layer.weight\": \"lm_head.weight\",\n            \"vision_model.patch_embed.proj.weight\": \"visual.patch_embed.proj.weight\",\n            \"vision_model.decoder.final_layernorm.weight\": \"visual.merger.ln_q.weight\",\n            \"vision_model.projection.encoder.linear_fc1.weight\": \"visual.merger.mlp.0.weight\",\n            \"vision_model.projection.encoder.linear_fc1.bias\": \"visual.merger.mlp.0.bias\",\n            \"vision_model.projection.encoder.linear_fc2.weight\": \"visual.merger.mlp.2.weight\",\n            \"vision_model.projection.encoder.linear_fc2.bias\": \"visual.merger.mlp.2.bias\",\n        }\n        if name in direct_name_mapping:\n            return [direct_name_mapping[name]], [params_one_group[0]]\n\n        if \"self_attention\" in name:\n            return self._convert_attention_param(name, params_one_group)\n        elif \"mlp\" in name:\n            return self._convert_mlp_param(name, params_one_group)\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n\n    def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        model_type, _, _, layer_number = name.split(\".\")[:4]\n\n        convert_names = []\n        if model_type == \"language_model\":\n            name_map_after_layer = {\n                \"self_attention.linear_qkv.bias\": [\n                    \"self_attn.q_proj.bias\",\n                    \"self_attn.k_proj.bias\",\n                    \"self_attn.v_proj.bias\",\n                ],\n                \"self_attention.linear_qkv.weight\": [\n                    \"self_attn.q_proj.weight\",\n                    \"self_attn.k_proj.weight\",\n                    \"self_attn.v_proj.weight\",\n                ],\n                \"self_attention.linear_proj.weight\": \"self_attn.o_proj.weight\",\n                \"self_attention.linear_qkv.layer_norm_weight\": \"input_layernorm.weight\",\n            }\n            name_after_layer = \".\".join(name.split(\".\")[-3:])\n            mapped_name = name_map_after_layer.get(name_after_layer)\n            if isinstance(mapped_name, list):\n                assert len(params) == len(mapped_name)\n                for one in mapped_name:\n                    convert_names.append(f\"model.layers.{layer_number}.{one}\")\n            else:\n                assert len(params) == 1\n                convert_names.append(f\"model.layers.{layer_number}.{mapped_name}\")\n        elif model_type == \"vision_model\":\n            name_map_after_layer = {\n                \"self_attention.linear_proj.weight\": \"attn.proj.weight\",\n                \"self_attention.linear_proj.bias\": \"attn.proj.bias\",\n                \"self_attention.linear_qkv.layer_norm_weight\": \"norm1.weight\",\n            }\n            name_after_layer = \".\".join(name.split(\".\")[-3:])\n            mapped_name = name_map_after_layer.get(name_after_layer, None)\n            if mapped_name is None:\n                assert \"linear_qkv\" in name_after_layer\n                assert len(params) == 3\n                new_param = torch.cat(params, dim=0)\n                params = [new_param]\n                if \"bias\" in name_after_layer:\n                    convert_names.append(f\"visual.blocks.{layer_number}.attn.qkv.bias\")\n                else:\n                    convert_names.append(f\"visual.blocks.{layer_number}.attn.qkv.weight\")\n            else:\n                assert len(params) == 1\n                convert_names.append(f\"visual.blocks.{layer_number}.{mapped_name}\")\n        else:\n            raise NotImplementedError(f\"Unsupported model type: {model_type}\")\n        return convert_names, params\n\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        model_type, _, _, layer_number = name.split(\".\")[:4]\n\n        convert_names = []\n        if model_type == \"language_model\":\n            name_map_after_layer = {\n                \"mlp.linear_fc1.weight\": [\"mlp.gate_proj.weight\", \"mlp.up_proj.weight\"],\n                \"mlp.linear_fc1.bias\": [\"mlp.gate_proj.bias\", \"mlp.up_proj.bias\"],\n                \"mlp.linear_fc2.weight\": \"mlp.down_proj.weight\",\n                \"mlp.linear_fc2.bias\": \"mlp.down_proj.bias\",\n                \"mlp.linear_fc1.layer_norm_weight\": \"post_attention_layernorm.weight\",\n            }\n            name_after_layer = \".\".join(name.split(\".\")[-3:])\n            mapped_name = name_map_after_layer.get(name_after_layer)\n            if isinstance(mapped_name, list):\n                assert len(params) == len(mapped_name)\n                for one in mapped_name:\n                    convert_names.append(f\"model.layers.{layer_number}.{one}\")\n            else:\n                assert len(params) == 1\n                convert_names.append(f\"model.layers.{layer_number}.{mapped_name}\")\n\n        elif model_type == \"vision_model\":\n            name_map_after_layer = {\n                \"mlp.linear_fc1.weight\": [\"mlp.gate_proj.weight\", \"mlp.up_proj.weight\"],\n                \"mlp.linear_fc1.bias\": [\"mlp.gate_proj.bias\", \"mlp.up_proj.bias\"],\n                \"mlp.linear_fc2.weight\": \"mlp.down_proj.weight\",\n                \"mlp.linear_fc2.bias\": \"mlp.down_proj.bias\",\n                \"mlp.linear_fc1.layer_norm_weight\": \"norm2.weight\",\n            }\n            name_after_layer = \".\".join(name.split(\".\")[-3:])\n            mapped_name = name_map_after_layer.get(name_after_layer)\n            if isinstance(mapped_name, list):\n                assert len(params) == len(mapped_name)\n                for one in mapped_name:\n                    convert_names.append(f\"visual.blocks.{layer_number}.{one}\")\n            else:\n                assert len(params) == 1\n                convert_names.append(f\"visual.blocks.{layer_number}.{mapped_name}\")\n        else:\n            raise NotImplementedError(f\"Unsupported model type: {model_type}\")\n        return convert_names, params\n\n\nclass McoreToHFWeightConverterDpskv3(McoreToHFWeightConverterBase):\n    def _convert_attention_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # mcore\n        # 'decoder.layers.0.input_layernorm.weight'\n        # 'decoder.layers.0.self_attention.linear_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_q_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_kv_down_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_kv_up_proj.layer_norm_weight'\n        # 'decoder.layers.0.self_attention.linear_kv_up_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_q_down_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_q_up_proj.weight'\n        # 'decoder.layers.0.self_attention.linear_q_up_proj.layer_norm_weight'\n        # hf\n        # 'model.layers.0.input_layernorm.weight'\n        # 'model.layers.0.self_attn.o_proj.weight'\n        # 'model.layers.0.self_attn.q_proj.weight'\n        # 'model.layers.0.self_attn.kv_a_proj_with_mqa.weight'\n        # 'model.layers.0.self_attn.kv_a_layernorm.weight'\n        # 'model.layers.0.self_attn.kv_b_proj.weight'\n        # 'model.layers.0.self_attn.q_a_proj.weight'\n        # 'model.layers.0.self_attn.q_b_proj.weight'\n        # 'model.layers.0.self_attn.q_a_layernorm.weight'\n        name_map_after_layer = {\n            \"input_layernorm.weight\": \"input_layernorm.weight\",\n            \"self_attention.linear_proj.weight\": \"self_attn.o_proj.weight\",\n            \"self_attention.linear_q_proj.weight\": \"self_attn.q_proj.weight\",\n            \"self_attention.linear_kv_down_proj.weight\": \"self_attn.kv_a_proj_with_mqa.weight\",\n            \"self_attention.linear_kv_up_proj.layer_norm_weight\": \"self_attn.kv_a_layernorm.weight\",\n            \"self_attention.linear_kv_up_proj.weight\": \"self_attn.kv_b_proj.weight\",\n            \"self_attention.linear_q_down_proj.weight\": \"self_attn.q_a_proj.weight\",\n            \"self_attention.linear_q_up_proj.weight\": \"self_attn.q_b_proj.weight\",\n            \"self_attention.linear_q_up_proj.layer_norm_weight\": \"self_attn.q_a_layernorm.weight\",\n        }\n        assert len(params) == 1\n        convert_names = []\n        layer_number = name.split(\".\")[2]\n        name_after_layer = name.split(f\".{layer_number}.\")[1]\n        convert_names.append(f\"model.layers.{layer_number}.{name_map_after_layer[name_after_layer]}\")\n        return convert_names, params\n\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # mcore dense\n        # 'decoder.layers.0.mlp.linear_fc1.layer_norm_weight'\n        # 'decoder.layers.0.mlp.linear_fc2.weight'\n        # 'decoder.layers.0.mlp.linear_fc1.weight'\n        #       ---\n        # 'decoder.layers.1.mlp.shared_experts.linear_fc1.weight'\n        #       ---\n        # 'decoder.layers.1.mlp.shared_experts.linear_fc2.weight'\n        # hf dense\n        # 'model.layers.0.post_attention_layernorm.weight'\n        # 'model.layers.0.mlp.down_proj.weight'\n        # 'model.layers.0.mlp.gate_proj.weight'\n        # 'model.layers.0.mlp.up_proj.weight'\n        # 'model.layers.1.mlp.shared_experts.gate_proj.weight'\n        # 'model.layers.1.mlp.shared_experts.up_proj.weight'\n        # 'model.layers.1.mlp.shared_experts.down_proj.weight'\n\n        # mcore moe\n        # 'decoder.layers.1.pre_mlp_layernorm.weight'\n        # 'decoder.layers.1.mlp.router.weight'\n        # 'decoder.layers.1.mlp.router.expert_bias'\n        # 'decoder.layers.1.mlp.experts.linear_fc1.weight0'\n        #       ---\n        # 'decoder.layers.1.mlp.experts.linear_fc2.weight0'\n        # hf moe\n        # 'model.layers.1.post_attention_layernorm.weight'\n        # 'model.layers.1.mlp.gate.weight'\n        # 'model.layers.1.mlp.gate.e_score_correction_bias'\n        # 'model.layers.1.mlp.experts.0.gate_proj.weight'\n        # 'model.layers.1.mlp.experts.0.up_proj.weight'\n        # 'model.layers.1.mlp.experts.0.down_proj.weight'\n\n        name_map_after_layer = {\n            \"mlp.linear_fc1.layer_norm_weight\": \"post_attention_layernorm.weight\",\n            \"mlp.linear_fc2.weight\": \"mlp.down_proj.weight\",\n            \"mlp.shared_experts.linear_fc2.weight\": \"mlp.shared_experts.down_proj.weight\",\n            \"mlp.linear_fc1.weight\": [\"mlp.gate_proj.weight\", \"mlp.up_proj.weight\"],\n            \"mlp.shared_experts.linear_fc1.weight\": [\n                \"mlp.shared_experts.gate_proj.weight\",\n                \"mlp.shared_experts.up_proj.weight\",\n            ],\n            \"pre_mlp_layernorm.weight\": \"post_attention_layernorm.weight\",\n            \"mlp.router.weight\": \"mlp.gate.weight\",\n            \"mlp.router.expert_bias\": \"mlp.gate.e_score_correction_bias\",\n        }\n        convert_names = []\n        layer_number = name.split(\".\")[2]\n        name_after_layer = name.split(f\".{layer_number}.\")[1]\n        if name_after_layer in name_map_after_layer:\n            mapped_name = name_map_after_layer[name_after_layer]\n            if isinstance(mapped_name, list):\n                assert len(params) == len(mapped_name)\n                for one in mapped_name:\n                    convert_names.append(f\"model.layers.{layer_number}.{one}\")\n            else:\n                assert len(params) == 1\n                convert_names.append(f\"model.layers.{layer_number}.{mapped_name}\")\n        else:\n            if \"mlp.experts.linear_fc1.weight\" in name:\n                expert_id = name.split(\"weight\")[-1]\n                convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight\")\n                convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight\")\n                assert len(params) == 2\n            elif \"mlp.experts.linear_fc2.weight\" in name:\n                expert_id = name.split(\"weight\")[-1]\n                convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight\")\n                assert len(params) == 1\n            else:\n                raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n\n        return convert_names, params\n\n    def _convert_mtp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        assert self.mcore_config.mtp_num_layers == 1, \"only support one mtp layer for now\"\n        assert self.mcore_config.num_layers == 61, \"only support 61 layers for now\"\n        direct_name_mapping = {\n            \"mtp.layers.0.enorm.weight\": \"model.layers.61.enorm.weight\",\n            \"mtp.layers.0.hnorm.weight\": \"model.layers.61.hnorm.weight\",\n            \"mtp.layers.0.eh_proj.weight\": \"model.layers.61.eh_proj.weight\",\n            \"mtp.layers.0.final_layernorm.weight\": \"model.layers.61.shared_head.norm.weight\",\n        }\n        if name in direct_name_mapping:\n            return [direct_name_mapping[name]], [params[0]]\n        assert \"mtp.layers.0.transformer_layer\" in name, \"only support transformer layer for now\"\n        # use proxy name to convert\n        proxy_name = name.replace(\"mtp.layers.0.transformer_layer\", \"decoder.layers.61\")\n        if \"self_attention\" in proxy_name or \"input_layernorm.weight\" in proxy_name:\n            convert_names, params = self._convert_attention_param(proxy_name, params)\n        elif \"mlp\" in proxy_name:\n            convert_names, params = self._convert_mlp_param(proxy_name, params)\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n\n    def convert_param(self, name: str, params_one_group: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        direct_name_mapping = {\n            \"embedding.word_embeddings.weight\": \"model.embed_tokens.weight\",\n            \"decoder.final_layernorm.weight\": \"model.norm.weight\",\n            \"output_layer.weight\": \"lm_head.weight\",\n        }\n        if name in direct_name_mapping:\n            return [direct_name_mapping[name]], [params_one_group[0]]\n        if \"mtp\" in name:\n            return self._convert_mtp_param(name, params_one_group)\n        elif \"self_attention\" in name or \"input_layernorm.weight\" in name:\n            return self._convert_attention_param(name, params_one_group)\n        elif \"mlp\" in name:\n            return self._convert_mlp_param(name, params_one_group)\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n\n\nclass McoreToHFWeightConverterMixtral(McoreToHFWeightConverterDense):\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # decoder.layers.0.mlp.router.weight\n        # decoder.layers.0.mlp.experts.linear_fc1.weight0 - weight7\n        # decoder.layers.0.mlp.experts.linear_fc2.weight0 - weight7\n\n        layer_number = name.split(\".\")[2]\n        convert_names = []\n        if \"pre_mlp_layernorm\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.post_attention_layernorm.weight\")\n        elif \"mlp.router.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.block_sparse_moe.gate.weight\")\n        elif \"mlp.experts.linear_fc1.weight\" in name:\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w1.weight\")\n            convert_names.append(f\"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w3.weight\")\n        elif \"mlp.experts.linear_fc2.weight\" in name:\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.block_sparse_moe.experts.{expert_id}.w2.weight\")\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n\n\nclass McoreToHFWeightConverterQwen3Moe(McoreToHFWeightConverterDense):\n    def _convert_mlp_param(self, name: str, params: list[torch.Tensor]) -> tuple[list[str], list[torch.Tensor]]:\n        # qwen3 moe no share expert\n\n        # 'decoder.layers.0.pre_mlp_layernorm.weight',\n        # 'decoder.layers.0.mlp.router.weight',\n        # moe1\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight0',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight1',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight2',\n        # 'decoder.layers.0.mlp.experts.linear_fc1.weight3',\n        # moe2\n        # 'decoder.layers.0.mlp.experts.linear_fc2.weight0',\n        # 'decoder.layers.0.mlp.experts.linear_fc2.weight1',\n        layer_number = name.split(\".\")[2]\n        convert_names = []\n        if \"pre_mlp_layernorm\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.post_attention_layernorm.weight\")\n            assert len(params) == 1\n        elif \"mlp.router.weight\" in name:\n            convert_names.append(f\"model.layers.{layer_number}.mlp.gate.weight\")\n            assert len(params) == 1\n        elif \"mlp.experts.linear_fc1\" in name:  # split gate_proj and up_proj\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.gate_proj.weight\")\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.up_proj.weight\")\n            assert len(params) == 2\n        elif \"mlp.experts.linear_fc2\" in name:\n            expert_id = name.split(\"weight\")[-1]\n            convert_names.append(f\"model.layers.{layer_number}.mlp.experts.{expert_id}.down_proj.weight\")\n            assert len(params) == 1\n        else:\n            raise NotImplementedError(f\"Unsupported parameter name: {name}\")\n        return convert_names, params\n"
  },
  {
    "path": "verl/models/qwen2/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/models/qwen2/megatron/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .modeling_qwen2_megatron import (\n    ParallelQwen2ForCausalLM,\n    # rmpad with megatron\n    ParallelQwen2ForCausalLMRmPad,\n    # rmpad with megatron and pipeline parallelism\n    ParallelQwen2ForCausalLMRmPadPP,\n    ParallelQwen2ForValueRmPad,\n    ParallelQwen2ForValueRmPadPP,\n    # original model with megatron\n    ParallelQwen2Model,\n)\n\n__all__ = [\n    \"ParallelQwen2ForCausalLM\",\n    \"ParallelQwen2ForCausalLMRmPad\",\n    \"ParallelQwen2ForCausalLMRmPadPP\",\n    \"ParallelQwen2ForValueRmPad\",\n    \"ParallelQwen2ForValueRmPadPP\",\n    \"ParallelQwen2Model\",\n]\n"
  },
  {
    "path": "verl/models/qwen2/megatron/checkpoint_utils/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef load_state_dict_to_megatron_qwen2(\n    state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False\n):\n    \"\"\"Load merged state_dict to sharded Megatron module in training.\"\"\"\n    from megatron.core import DistributedDataParallel as LocalDDP\n    from megatron.core import mpu\n    from megatron.core.transformer.module import Float16Module\n    from torch.nn.parallel import DistributedDataParallel as torchDDP\n\n    from verl.utils.logger import print_rank_0\n    from verl.utils.megatron_utils import unwrap_model\n\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    def fetch_params(module):\n        for param in module.parameters():\n            torch.distributed.fetch(\n                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()\n            )\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if torch.distributed.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers, (\n        f\"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size: \"\n        f\"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}\"\n    )\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        gpt_model_module = _get_gpt_model(models[i])\n        assert len(gpt_model_module.model.layers) == num_layers_per_model\n\n    def _fetch_tensor(tensor, name) -> torch.Tensor:\n        \"\"\"fetch tensor\"\"\"\n        nonlocal state_dict\n        if tensor is not None:\n            tensor = tensor.data.copy_(state_dict[name], non_blocking=True)\n\n    def _fetch_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards\"\"\"\n        nonlocal state_dict\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if name in state_dict:\n            full_weight = state_dict[name]\n\n            if mutate_func is not None:\n                full_weight = mutate_func(full_weight)\n            tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n            if tensor is not None:\n                tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True)\n        else:\n            print(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards\"\"\"\n        nonlocal state_dict\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if name in state_dict:\n            full_weight = state_dict[name]\n\n            if mutate_func is not None:\n                full_weight = mutate_func(full_weight)\n            tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n            if tensor is not None:\n                tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True)\n        else:\n            print(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor:\n        \"\"\"fetch gate_up tensor in tp shards\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        if gate_name in state_dict and up_name in state_dict:\n            gate_weight = state_dict[gate_name]\n            up_weight = state_dict[up_name]\n            new_gate_up_weight = torch.empty(\n                config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                intermediate_size_tp = config.intermediate_size // tp_size\n                gate_weight_tp = gate_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                up_weight_tp = up_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                new_gate_up_weight[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)].copy_(\n                    torch.cat([gate_weight_tp, up_weight_tp], dim=0)\n                )\n\n            tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0)\n            if tensor is not None:\n                tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True)\n        else:\n            print(f\"tp_shard tensor:[{gate_name}, {up_name}] not in state_dict, skip loading\")\n\n    def _fetch_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor:\n        \"\"\"fetch tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        assert q_name in state_dict and k_name in state_dict and v_name in state_dict\n        full_weight_q = state_dict[q_name]\n        full_weight_k = state_dict[k_name]\n        full_weight_v = state_dict[v_name]\n\n        hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n        if config.num_key_value_heads >= tp_size:\n            q_size_tp = config.hidden_size // tp_size\n            kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n            total_size = q_size_tp + 2 * kv_size_tp\n            if not bias:\n                new_weight_qkv = torch.empty(\n                    total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                )\n            else:\n                new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id())\n            for i in range(tp_size):\n                q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                k_part = full_weight_k[i * kv_size_tp : (i + 1) * kv_size_tp]\n                v_part = full_weight_v[i * kv_size_tp : (i + 1) * kv_size_tp]\n                new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0))\n\n        else:\n            q_size_tp = config.hidden_size // tp_size\n            kv_size_tp = hidden_size_per_head\n            total_size = q_size_tp + 2 * kv_size_tp\n            if not bias:\n                new_weight_qkv = torch.empty(\n                    total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                )\n            else:\n                new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id())\n            for i in range(tp_size):\n                q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head\n                end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head\n                k_part = full_weight_k[start_idx:end_idx]\n                v_part = full_weight_v[start_idx:end_idx]\n                new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(torch.cat([q_part, k_part, v_part], dim=0))\n\n        tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0)\n        if tensor is not None:\n            tensor = tensor.data.copy_(tensor_chunk[tp_rank], non_blocking=True)\n\n    # Embeddings\n    # -------------------\n    print_rank_0(\"loading embeddings...\")\n    gpt_model_module = _get_gpt_model(models[0])\n    if pp_rank == 0:\n        embed_tokens_weight = gpt_model_module.model.embed_tokens.weight\n        _fetch_tp_shard_tensor_vocab(embed_tokens_weight, \"model.embed_tokens.weight\")\n\n    # Transformer layers\n    # -------------------\n    layer_map = _megatron_calc_layer_map(config)\n\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    num_layer_per_pp = config.num_hidden_layers // pp_size\n    vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()\n\n    layer_list = []\n    if vpp_size is not None:\n        for vpp_rank in range(vpp_size):\n            num_layer_vpp_chunk = num_layer_per_pp // vpp_size\n            num_layer_this_model = num_layer_vpp_chunk\n            offset = vpp_rank * (config.num_hidden_layers // mpu.get_virtual_pipeline_model_parallel_world_size()) + (\n                mpu.get_pipeline_model_parallel_rank() * num_layer_vpp_chunk\n            )\n            layer_list.extend(list(range(offset, offset + num_layer_this_model)))\n    else:\n        num_layer_this_model = num_layer_per_pp\n        offset = pp_rank * num_layer_per_pp\n        layer_list.extend(list(range(offset, offset + num_layer_this_model)))\n\n    for layer in layer_list:\n        print(f\"{torch.distributed.get_rank()} loading layer #{layer}...\")\n        layer_name = f\"model.layers.{layer}\"\n        dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer]\n\n        print(\n            f\"{torch.distributed.get_rank()} offset: {offset}, num_layer_this_model: {num_layer_this_model}, \"\n            f\"layer_name: {layer_name}, layer_map[layer]: {layer_map[layer]}\"\n        )\n\n        gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank])\n        sync_layer = gpt_model_module.model.layers[dst_layer_idx]\n\n        _fetch_tensor(\n            sync_layer.input_layernorm.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.input_layernorm.weight\",\n        )\n\n        _fetch_tp_shard_tensor_qkv(\n            sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.self_attn.q_proj.weight\",\n            f\"{layer_name}.self_attn.k_proj.weight\",\n            f\"{layer_name}.self_attn.v_proj.weight\",\n        )\n\n        _fetch_tp_shard_tensor_qkv(\n            sync_layer.self_attn.qkv_proj.bias if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.self_attn.q_proj.bias\",\n            f\"{layer_name}.self_attn.k_proj.bias\",\n            f\"{layer_name}.self_attn.v_proj.bias\",\n            bias=True,\n        )\n\n        _fetch_tp_shard_tensor(\n            sync_layer.self_attn.o_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.self_attn.o_proj.weight\",\n            chunk_dim=1,\n        )\n\n        _fetch_tensor(\n            sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.post_attention_layernorm.weight\",\n        )\n\n        _fetch_tp_shard_tensor_gate_up(\n            sync_layer.mlp.gate_up_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.mlp.gate_proj.weight\",\n            f\"{layer_name}.mlp.up_proj.weight\",\n        )\n\n        _fetch_tp_shard_tensor(\n            sync_layer.mlp.down_proj.weight if dst_pp_rank == pp_rank else None,\n            f\"{layer_name}.mlp.down_proj.weight\",\n            chunk_dim=1,\n        )\n    # Final Layernorm\n    # -------------------\n    print_rank_0(\"loading final layernorm...\")\n    gpt_model_module = _get_gpt_model(models[-1])\n    _fetch_tensor(\n        getattr(gpt_model_module.model.norm, \"weight\", None),\n        \"model.norm.weight\",\n    )\n\n    if tie_word_embeddings:\n        print_rank_0(\"tie_word_embeddings skip load lm_head\")\n    else:\n        print_rank_0(\"loading lm_head...\")\n        if pp_rank + 1 == pp_size:\n            lm_head_weight = gpt_model_module.lm_head.weight\n\n            if is_value_model:\n                if \"lm_head.weight\" in state_dict and state_dict[\"lm_head.weight\"].shape[0] == 1:\n                    _fetch_tensor(lm_head_weight, \"lm_head.weight\")\n                    print_rank_0(\"load lm_head from value_head weight\")\n                elif \"reward_head.weight\" in state_dict and state_dict[\"reward_head.weight\"].shape[0] == 1:\n                    _fetch_tensor(lm_head_weight, \"reward_head.weight\")\n                    print_rank_0(\"load lm_head from value_head weight\")\n                else:\n                    _fetch_tensor(None, \"lm_head.weight\")\n                    print_rank_0(\"fail to match lm_head in value_model\")\n\n            else:\n                _fetch_tp_shard_tensor(lm_head_weight, \"lm_head.weight\")\n\n    dist.barrier()\n    get_torch_device().empty_cache()\n    print_rank_0(f\"loading megatron ckpt done, time elapsed {time.time() - start_time}s\")\n"
  },
  {
    "path": "verl/models/qwen2/megatron/checkpoint_utils/qwen2_loader_depracated.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef load_state_dict_to_megatron_qwen2(\n    state_dict, wrapped_models, config, params_dtype, is_value_model=False, tie_word_embeddings=False\n):\n    \"\"\"Load merged state_dict to sharded Megatron module in training.\"\"\"\n    from megatron.core import DistributedDataParallel as LocalDDP\n    from megatron.core import mpu\n    from megatron.core.transformer.module import Float16Module\n    from torch.nn.parallel import DistributedDataParallel as torchDDP\n\n    from verl.utils.logger import print_rank_0\n    from verl.utils.megatron_utils import unwrap_model\n\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    def broadcast_params(module):\n        for param in module.parameters():\n            torch.distributed.broadcast(\n                param.data, src=mpu.get_data_parallel_src_rank(), group=mpu.get_data_parallel_group()\n            )\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if torch.distributed.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers, (\n        f\"num_layers_per_model: {num_layers_per_model} * pp_size: {pp_size} * virtual_pp_size: \"\n        f\"{virtual_pp_size} != config.num_hidden_layers: {config.num_hidden_layers}\"\n    )\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        gpt_model_module = _get_gpt_model(models[i])\n        assert len(gpt_model_module.model.layers) == num_layers_per_model\n\n    def _broadcast_tensor(tensor, name) -> torch.Tensor:\n        \"\"\"broadcast tensor from rank0 across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                weight = state_dict[name]\n                tensor_shape = weight.shape\n            else:\n                tensor_shape = None\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not in state_dict, skip load\")\n            return\n\n        if tensor is None:\n            tensor = torch.empty(\n                tensor_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        if torch.distributed.get_rank() == 0:\n            tensor.data.copy_(weight)\n        dist.broadcast(tensor, src=0, group=mp_group)\n\n    def _broadcast_tp_shard_tensor_vocab(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                full_weight = state_dict[name]\n\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor(tensor, name, chunk_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            if name in state_dict:\n                full_weight = state_dict[name]\n                if mutate_func is not None:\n                    full_weight = mutate_func(full_weight)\n                tensor_chunk = torch.chunk(full_weight, tp_size, dim=chunk_dim)\n                chunk_shape = tensor_chunk[0].shape\n            else:\n                chunk_shape = None\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            gate_weight = state_dict[gate_name]\n            up_weight = state_dict[up_name]\n            new_gate_up_weight = torch.empty(\n                config.intermediate_size * 2, config.hidden_size, dtype=params_dtype, device=get_device_id()\n            )\n            for i in range(tp_size):\n                intermediate_size_tp = config.intermediate_size // tp_size\n                gate_weight_tp = gate_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                up_weight_tp = up_weight[i * intermediate_size_tp : (i + 1) * intermediate_size_tp]\n                new_gate_up_weight[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)].copy_(\n                    torch.cat([gate_weight_tp, up_weight_tp], dim=0)\n                )\n\n            tensor_chunk = torch.chunk(new_gate_up_weight, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank() == 0:} tensor {gate_name, up_name} shape \"\n                f\"{tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, bias=False) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_rank = mpu.get_tensor_model_parallel_rank()\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        if torch.distributed.get_rank() == 0:\n            assert q_name in state_dict and k_name in state_dict and v_name in state_dict\n            full_weight_q = state_dict[q_name]\n            full_weight_k = state_dict[k_name]\n            full_weight_v = state_dict[v_name]\n\n            hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                if not bias:\n                    new_weight_qkv = torch.empty(\n                        total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                    )\n                else:\n                    new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id())\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    k_part = full_weight_k[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    v_part = full_weight_v[i * kv_size_tp : (i + 1) * kv_size_tp]\n                    new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(\n                        torch.cat([q_part, k_part, v_part], dim=0)\n                    )\n\n            else:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                if not bias:\n                    new_weight_qkv = torch.empty(\n                        total_size * tp_size, config.hidden_size, dtype=params_dtype, device=get_device_id()\n                    )\n                else:\n                    new_weight_qkv = torch.empty(total_size * tp_size, dtype=params_dtype, device=get_device_id())\n                for i in range(tp_size):\n                    q_part = full_weight_q[i * q_size_tp : (i + 1) * q_size_tp]\n                    start_idx = i * config.num_key_value_heads // tp_size * hidden_size_per_head\n                    end_idx = (i * config.num_key_value_heads // tp_size + 1) * hidden_size_per_head\n                    k_part = full_weight_k[start_idx:end_idx]\n                    v_part = full_weight_v[start_idx:end_idx]\n                    new_weight_qkv[i * total_size : (i + 1) * total_size].copy_(\n                        torch.cat([q_part, k_part, v_part], dim=0)\n                    )\n\n            tensor_chunk = torch.chunk(new_weight_qkv, tp_size, dim=0)\n            chunk_shape = tensor_chunk[0].shape\n        else:\n            chunk_shape = None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=0, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name, k_name, v_name}] not in state_dict, skip loading\")\n            return\n\n        if tensor is None:\n            sync_tensor = torch.empty(\n                chunk_shape,\n                dtype=params_dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n        else:\n            assert tensor.shape == chunk_shape, (\n                f\"rank #{torch.distributed.get_rank()} tensor {q_name} shape {tensor.shape} != {chunk_shape}\"\n            )\n            sync_tensor = torch.empty_like(tensor, device=get_device_id(), requires_grad=False)\n\n        for i in range(tp_size):\n            if torch.distributed.get_rank() == 0:\n                sync_tensor.data.copy_(tensor_chunk[i])\n            dist.broadcast(sync_tensor, src=0, group=mp_group)\n            if (i == tp_rank) and (tensor is not None):\n                tensor.data.copy_(sync_tensor)\n\n    if dp_rank == 0:\n        # Embeddings\n        # -------------------\n        print_rank_0(\"loading embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        embed_tokens_weight = None\n        if pp_rank == 0:\n            embed_tokens_weight = gpt_model_module.model.embed_tokens.weight\n        _broadcast_tp_shard_tensor_vocab(embed_tokens_weight, \"model.embed_tokens.weight\")\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n\n        for layer in range(config.num_hidden_layers):\n            print_rank_0(f\"loading layer #{layer}...\")\n            layer_name = f\"model.layers.{layer}\"\n            dst_pp_rank, dst_virtual_pp_rank, dst_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[dst_virtual_pp_rank])\n            sync_layer = gpt_model_module.model.layers[dst_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.input_layernorm.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.input_layernorm.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.bias if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.q_proj.bias\",\n                f\"{layer_name}.self_attn.k_proj.bias\",\n                f\"{layer_name}.self_attn.v_proj.bias\",\n                bias=True,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attn.o_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                chunk_dim=1,\n            )\n\n            _broadcast_tensor(\n                sync_layer.post_attention_layernorm.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.gate_up_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.down_proj.weight if dst_pp_rank == pp_rank else None,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                chunk_dim=1,\n            )\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"loading final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.model.norm, \"weight\", None),\n            \"model.norm.weight\",\n        )\n\n        if tie_word_embeddings:\n            print_rank_0(\"tie_word_embeddings skip load lm_head\")\n        else:\n            print_rank_0(\"loading lm_head...\")\n            lm_head_weight = None\n            if pp_rank + 1 == pp_size:\n                lm_head_weight = gpt_model_module.lm_head.weight\n\n            if is_value_model:\n                if \"lm_head.weight\" in state_dict and state_dict[\"lm_head.weight\"].shape[0] == 1:\n                    _broadcast_tensor(lm_head_weight, \"lm_head.weight\")\n                    print_rank_0(\"load lm_head from value_head weight\")\n                elif \"reward_head.weight\" in state_dict and state_dict[\"reward_head.weight\"].shape[0] == 1:\n                    _broadcast_tensor(lm_head_weight, \"reward_head.weight\")\n                    print_rank_0(\"load lm_head from value_head weight\")\n                else:\n                    _broadcast_tensor(None, \"lm_head.weight\")\n                    print_rank_0(\"fail to match lm_head in value_model\")\n\n            else:\n                _broadcast_tp_shard_tensor(lm_head_weight, \"lm_head.weight\")\n\n    dist.barrier()\n    # Broadcast weights inside data parallel groups\n    for wrapped_model in wrapped_models:\n        broadcast_params(wrapped_model)\n\n    get_torch_device().empty_cache()\n    print_rank_0(f\"loading megatron ckpt done, time elapsed {time.time() - start_time}s\")\n"
  },
  {
    "path": "verl/models/qwen2/megatron/checkpoint_utils/qwen2_saver.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport time\n\nimport torch\nimport torch.distributed as dist\nfrom megatron.core import mpu\nfrom megatron.core.distributed import DistributedDataParallel as LocalDDP\nfrom megatron.core.transformer.module import Float16Module\nfrom torch.nn.parallel import DistributedDataParallel as torchDDP\n\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.logger import print_rank_0\nfrom verl.utils.megatron_utils import unwrap_model\n\n\ndef _megatron_calc_global_rank(tp_rank: int = 0, dp_rank: int = 0, pp_rank: int = 0):\n    \"\"\"given TP,DP,PP rank to get the global rank.\"\"\"\n\n    tp_size = mpu.get_tensor_model_parallel_world_size()\n    dp_size = mpu.get_data_parallel_world_size()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    assert tp_size * dp_size * pp_size == torch.distributed.get_world_size(), (\n        f\"{tp_size} x {dp_size} x {pp_size} != {torch.distributed.get_world_size()}\"\n    )\n    # We only support TP-DP-PP grouping, for correctness when resharding\n    return (pp_rank * dp_size + dp_rank) * tp_size + tp_rank\n\n\ndef _megatron_calc_layer_map(config):\n    \"\"\"Calculate the mapping of global layer_idx to local layer_idx\n    Returns:\n        layer_map (Dict: int -> tuple(int, int, int)):\n            mapping from the global layer index to\n            a tuple of (pp_rank, virtual_pp_rank, layer_idx inside model)\n    \"\"\"\n    from megatron.core import mpu\n\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n\n    layer_map = dict()\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    for pp_rank_idx in range(pp_size):\n        for virtual_pp_rank_idx in range(virtual_pp_size):\n            layer_offset = (\n                virtual_pp_rank_idx * (config.num_hidden_layers // virtual_pp_size) + pp_rank_idx * num_layers_per_model\n            )\n            for layer_idx in range(num_layers_per_model):\n                layer_map[layer_offset + layer_idx] = (\n                    pp_rank_idx,\n                    virtual_pp_rank_idx,\n                    layer_idx,\n                )\n    return layer_map\n\n\ndef merge_megatron_ckpt_qwen2(wrapped_models, config, dtype, is_value_model=False, tie_word_embeddings=False):\n    \"\"\"Merge sharded parameters of a Megatron module into a merged checkpoint.\n\n    Args:\n        wrapped_models (list of megatron.core.distributed.DistributedDataParallel):\n            The local DDP wrapped megatron modules.\n        config (str or None):\n            HF config for model\n        dtype: model params type\n        is_value_model: if model is value model\n        tie_word_embeddings: tie_word_embeddings\n    Returns:\n        state_dict (dict):\n            The merged state_dict in rank 0, and an empty dictionary in other ranks.\n    \"\"\"\n    start_time = time.time()\n\n    def _get_gpt_model(model):\n        return model\n\n    dp_rank = mpu.get_data_parallel_rank()\n    pp_size = mpu.get_pipeline_model_parallel_world_size()\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    virtual_pp_size = mpu.get_virtual_pipeline_model_parallel_world_size() or 1\n    mp_group = mpu.get_model_parallel_group()\n\n    if dist.get_rank() == 0:\n        assert mp_group.rank() == 0, f\"mp_rank:[{mp_group.rank}] != 0 on rank #0\"\n        assert pp_rank == 0, f\"pp_rank:[{pp_rank}] != 0 on rank #0\"\n        assert dp_rank == 0, f\"dp_rank:[{dp_rank}] != 0 on rank #0\"\n\n    if not isinstance(wrapped_models, list | tuple):\n        wrapped_models = list(wrapped_models)\n\n    assert len(wrapped_models) == virtual_pp_size\n    num_layers_per_model = config.num_hidden_layers // pp_size // virtual_pp_size\n    assert num_layers_per_model * pp_size * virtual_pp_size == config.num_hidden_layers\n\n    models = [None] * len(wrapped_models)\n\n    for i, wrapped_model in enumerate(wrapped_models):\n        models[i] = unwrap_model(wrapped_model, (torchDDP, LocalDDP, Float16Module))\n        assert len(models[i].model.layers) == num_layers_per_model, (\n            \"len model layers {} not equal to num_layers_per_model {}\".format(\n                len(models[i].model.layers), num_layers_per_model\n            )\n        )\n\n    state_dict = dict()\n\n    def _get_cpu_tensor(tensor: torch.Tensor):\n        if tensor is None:\n            return None\n        if tensor.device == torch.device(\"cpu\"):\n            return tensor.detach().clone()\n        return tensor.detach().cpu()\n\n    def _broadcast_tensor(tensor, name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        if torch.distributed.get_rank() == src_rank:\n            if tensor is None:\n                weight = None\n                tensor_shape = None\n            else:\n                weight = tensor\n                tensor_shape = weight.shape\n        else:\n            weight = None\n            tensor_shape = None\n\n        obj_list = [tensor_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        tensor_shape = obj_list[0]\n\n        if tensor_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tensor:[{name}] not exist, skip collect\")\n            return\n\n        if weight is None:\n            weight = torch.empty(\n                tensor_shape,\n                dtype=dtype,\n                device=get_device_id(),\n                requires_grad=False,\n            )\n\n        dist.broadcast(weight, src=src_rank, group=mp_group)\n\n        if torch.distributed.get_rank() == 0:\n            state_dict[name] = _get_cpu_tensor(weight)\n\n    def _broadcast_tp_shard_tensor(tensor, name, src_pp_rank, concat_dim=0, mutate_func=None) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=concat_dim)\n            if mutate_func is not None:\n                full_tensor = mutate_func(full_tensor)\n            state_dict[name] = full_tensor\n\n    def _broadcast_tp_shard_tensor_gate_up(tensor, gate_name, up_name, src_pp_rank) -> torch.Tensor:\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{gate_name, up_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            intermediate_size_tp = config.intermediate_size // tp_size\n            gate_weight_list = []\n            up_weight_list = []\n            for i in range(tp_size):\n                gate_up_weight_tp = full_tensor[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)]\n                gate_weight_tp = gate_up_weight_tp[:intermediate_size_tp]\n                up_weight_tp = gate_up_weight_tp[intermediate_size_tp:]\n                gate_weight_list.append(gate_weight_tp)\n                up_weight_list.append(up_weight_tp)\n\n            state_dict[gate_name] = torch.cat(gate_weight_list, dim=0)\n            state_dict[up_name] = torch.cat(up_weight_list, dim=0)\n\n    def _broadcast_tp_shard_tensor_qkv(tensor, q_name, k_name, v_name, src_pp_rank):\n        \"\"\"broadcast tensor in tp shards across mp_group\"\"\"\n        nonlocal state_dict\n        nonlocal mp_group\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=src_pp_rank)\n\n        chunk_shape = tensor.shape if torch.distributed.get_rank() == src_rank else None\n\n        obj_list = [chunk_shape]\n        dist.broadcast_object_list(obj_list, src=src_rank, group=mp_group)\n        chunk_shape = obj_list[0]\n        if chunk_shape is None:\n            # all or none ranks in the mp_group should reach here\n            print_rank_0(f\"tp_shard tensor:[{q_name}] not exist, skip collecting\")\n            return\n\n        buffer_tensor = torch.empty(\n            chunk_shape,\n            dtype=dtype,\n            device=get_device_id(),\n            requires_grad=False,\n        )\n\n        chunk_tensors = [None] * tp_size\n\n        for i in range(tp_size):\n            cur_src_rank = _megatron_calc_global_rank(tp_rank=i, dp_rank=0, pp_rank=src_pp_rank)\n            sync_tensor = tensor if torch.distributed.get_rank() == cur_src_rank else buffer_tensor\n            dist.broadcast(sync_tensor, src=cur_src_rank, group=mp_group)\n\n            if torch.distributed.get_rank() == 0:\n                chunk_tensors[i] = _get_cpu_tensor(sync_tensor)\n\n        if torch.distributed.get_rank() == 0:\n            full_tensor = torch.concat(chunk_tensors, dim=0)\n            q_weight_list = []\n            k_weight_list = []\n            v_weight_list = []\n            hidden_size_per_head = config.hidden_size // config.num_attention_heads\n\n            if config.num_key_value_heads >= tp_size:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_part = qkv_part[:q_size_tp]\n                    k_part = qkv_part[q_size_tp : q_size_tp + kv_size_tp]\n                    v_part = qkv_part[q_size_tp + kv_size_tp : total_size]\n                    q_weight_list.append(q_part)\n                    k_weight_list.append(k_part)\n                    v_weight_list.append(v_part)\n            else:\n                q_size_tp = config.hidden_size // tp_size\n                kv_size_tp = hidden_size_per_head\n                total_size = q_size_tp + 2 * kv_size_tp\n                for i in range(tp_size):\n                    qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                    q_part = qkv_part[:q_size_tp]\n                    k_part = qkv_part[q_size_tp : q_size_tp + kv_size_tp]\n                    v_part = qkv_part[q_size_tp + kv_size_tp : total_size]\n                    q_weight_list.append(q_part)\n                    if i * config.num_key_value_heads % tp_size == 0:\n                        k_weight_list.append(k_part)\n                        v_weight_list.append(v_part)\n\n            state_dict[q_name] = torch.cat(q_weight_list, dim=0)\n            state_dict[k_name] = torch.cat(k_weight_list, dim=0)\n            state_dict[v_name] = torch.cat(v_weight_list, dim=0)\n\n    # empty cache before collecting weights\n    get_torch_device().empty_cache()\n    # Embeddings\n    # -------------------\n    if dp_rank == 0:\n        # Embeddings\n        # -------------------\n        print_rank_0(\"collecting embeddings...\")\n        gpt_model_module = _get_gpt_model(models[0])\n        _broadcast_tp_shard_tensor(\n            gpt_model_module.model.embed_tokens.weight if pp_rank == 0 else None,\n            \"model.embed_tokens.weight\",\n            src_pp_rank=0,\n        )\n\n        # Transformer layers\n        # -------------------\n        layer_map = _megatron_calc_layer_map(config)\n        for layer in range(config.num_hidden_layers):\n            print_rank_0(f\"collecting layer #{layer}...\")\n            layer_name = f\"model.layers.{layer}\"\n            src_pp_rank, src_virtual_pp_rank, src_layer_idx = layer_map[layer]\n\n            gpt_model_module = _get_gpt_model(models[src_virtual_pp_rank])\n            sync_layer = gpt_model_module.model.layers[src_layer_idx]\n\n            _broadcast_tensor(\n                sync_layer.input_layernorm.weight,\n                f\"{layer_name}.input_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.weight,\n                f\"{layer_name}.self_attn.q_proj.weight\",\n                f\"{layer_name}.self_attn.k_proj.weight\",\n                f\"{layer_name}.self_attn.v_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_qkv(\n                sync_layer.self_attn.qkv_proj.bias,\n                f\"{layer_name}.self_attn.q_proj.bias\",\n                f\"{layer_name}.self_attn.k_proj.bias\",\n                f\"{layer_name}.self_attn.v_proj.bias\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.self_attn.o_proj.weight,\n                f\"{layer_name}.self_attn.o_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tensor(\n                sync_layer.post_attention_layernorm.weight,\n                f\"{layer_name}.post_attention_layernorm.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor_gate_up(\n                sync_layer.mlp.gate_up_proj.weight,\n                f\"{layer_name}.mlp.gate_proj.weight\",\n                f\"{layer_name}.mlp.up_proj.weight\",\n                src_pp_rank=src_pp_rank,\n            )\n\n            _broadcast_tp_shard_tensor(\n                sync_layer.mlp.down_proj.weight,\n                f\"{layer_name}.mlp.down_proj.weight\",\n                concat_dim=1,\n                src_pp_rank=src_pp_rank,\n            )\n\n        # Final Layernorm\n        # -------------------\n        print_rank_0(\"collecting final layernorm...\")\n        gpt_model_module = _get_gpt_model(models[-1])\n        _broadcast_tensor(\n            getattr(gpt_model_module.model.norm, \"weight\", None),\n            \"model.norm.weight\",\n            src_pp_rank=pp_size - 1,\n        )\n\n        if tie_word_embeddings:\n            print_rank_0(\"tie word embedding skip load lm_head...\")\n        else:\n            print_rank_0(\"collecting lm_head...\")\n\n            if is_value_model:\n                _broadcast_tensor(\n                    gpt_model_module.lm_head.weight if pp_rank == pp_size - 1 else None,\n                    \"lm_head.weight\",\n                    src_pp_rank=pp_size - 1,\n                )\n                _broadcast_tensor(\n                    gpt_model_module.reward_head.weight\n                    if pp_rank == pp_size - 1 and getattr(gpt_model_module, \"reward_weight\", None) is not None\n                    else None,\n                    \"reward_head.weight\",\n                    src_pp_rank=pp_size - 1,\n                )\n\n            else:\n                _broadcast_tp_shard_tensor(\n                    getattr(gpt_model_module.lm_head, \"weight\", None) if pp_rank == pp_size - 1 else None,\n                    \"lm_head.weight\",\n                    src_pp_rank=pp_size - 1,\n                )\n\n    dist.barrier()\n\n    get_torch_device().empty_cache()\n    if torch.distributed.get_rank() == 0:\n        for k, v in state_dict.items():\n            if dtype != v.dtype:\n                state_dict[k] = v.to(dtype)\n\n    print_rank_0(f\"merge megatron ckpt done, time elapsed {time.time() - start_time}s\")\n    return state_dict\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .parallel_attention import ParallelQwen2Attention\nfrom .parallel_decoder import ParallelQwen2DecoderLayer, ParallelQwen2DecoderLayerRmPad\nfrom .parallel_mlp import ParallelQwen2MLP\nfrom .parallel_rmsnorm import ParallelQwen2RMSNorm\n\n__all__ = [\n    \"ParallelQwen2Attention\",\n    \"ParallelQwen2DecoderLayer\",\n    \"ParallelQwen2DecoderLayerRmPad\",\n    \"ParallelQwen2MLP\",\n    \"ParallelQwen2RMSNorm\",\n]\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/parallel_attention.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport math\nfrom typing import Optional\n\nimport torch.nn.functional as F\nfrom einops import rearrange\nfrom transformers.utils import is_flash_attn_2_available\n\nif is_flash_attn_2_available():\n    from flash_attn import flash_attn_varlen_func\n    from flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa: F401\n\nimport torch\nfrom flash_attn.layers.rotary import apply_rotary_emb\nfrom megatron.core import ModelParallelConfig, tensor_parallel\nfrom megatron.core import parallel_state as mpu\nfrom torch import nn\nfrom transformers import Qwen2Config\n\nfrom verl.models.qwen2.megatron.layers.parallel_linear import QKVParallelLinear\nfrom verl.utils.megatron import tensor_parallel as tp_utils\n\n\nclass Qwen2RotaryEmbedding(nn.Module):\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):\n        super().__init__()\n\n        self.dim = dim\n        self.max_position_embeddings = max_position_embeddings\n        self.base = base\n        inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n        self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        # Build here to make `torch.jit.trace` work.\n        self._set_cos_sin_cache(\n            seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()\n        )\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n    def forward(self, x, seq_len=None):\n        # x: [bs, num_attention_heads, seq_len, head_size]\n        if seq_len > self.max_seq_len_cached:\n            self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)\n\n        return (\n            self.cos_cached[:seq_len].to(dtype=x.dtype),\n            self.sin_cached[:seq_len].to(dtype=x.dtype),\n        )\n\n\nclass Qwen2LinearScalingRotaryEmbedding(Qwen2RotaryEmbedding):\n    \"\"\"Qwen2RotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n        t = t / self.scaling_factor\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n\nclass Qwen2DynamicNTKScalingRotaryEmbedding(Qwen2RotaryEmbedding):\n    \"\"\"Qwen2RotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla\"\"\"\n\n    def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):\n        self.scaling_factor = scaling_factor\n        super().__init__(dim, max_position_embeddings, base, device)\n\n    def _set_cos_sin_cache(self, seq_len, device, dtype):\n        self.max_seq_len_cached = seq_len\n\n        if seq_len > self.max_position_embeddings:\n            base = self.base * (\n                (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)\n            ) ** (self.dim / (self.dim - 2))\n            inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))\n            self.register_buffer(\"inv_freq\", inv_freq, persistent=False)\n\n        t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)\n\n        freqs = torch.einsum(\"i,j->ij\", t, self.inv_freq)\n        # Different from paper, but it uses a different permutation in order to obtain the same calculation\n        emb = torch.cat((freqs, freqs), dim=-1)\n        self.register_buffer(\"cos_cached\", emb.cos().to(dtype), persistent=False)\n        self.register_buffer(\"sin_cached\", emb.sin().to(dtype), persistent=False)\n\n\ndef rotate_half(x):\n    \"\"\"Rotates half the hidden dims of the input.\"\"\"\n    x1 = x[..., : x.shape[-1] // 2]\n    x2 = x[..., x.shape[-1] // 2 :]\n    return torch.cat((-x2, x1), dim=-1)\n\n\ndef apply_rotary_pos_emb(q, k, cos, sin, position_ids):\n    cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n    return q_embed, k_embed\n\n\ndef repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:\n    \"\"\"\n    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,\n    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)\n    \"\"\"\n    batch, num_key_value_heads, slen, head_dim = hidden_states.shape\n    if n_rep == 1:\n        return hidden_states\n    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)\n    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)\n\n\nclass ParallelQwen2Attention(nn.Module):\n    \"\"\"Multi-headed attention from 'Attention Is All You Need' paper\"\"\"\n\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config = config\n        self.megatron_config = megatron_config\n        self.hidden_size = config.hidden_size\n        self.num_heads = config.num_attention_heads\n        self.head_dim = self.hidden_size // self.num_heads\n        self.num_key_value_heads = config.num_key_value_heads\n        self.num_key_value_groups = self.num_heads // self.num_key_value_heads\n        self.max_position_embeddings = config.max_position_embeddings\n        self.rope_theta = config.rope_theta\n\n        # assign values after tp\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n        assert self.num_heads % tp_size == 0, (\n            f\"num_head must be divisible by tp_size. Got num_head={self.num_heads}, tp_size={tp_size}\"\n        )\n        assert self.num_key_value_heads % tp_size == 0, (\n            f\"num_key_value_heads must be divisible by tp_size. Got num_key_value_heads=\"\n            f\"{self.num_key_value_heads}, tp_size={tp_size}\"\n        )\n\n        self.num_heads_per_tp = self.num_heads // tp_size\n        self.num_key_value_heads_per_tp = self.num_key_value_heads // tp_size\n        self.hidden_size_per_tp = self.hidden_size // tp_size\n\n        if (self.head_dim * self.num_heads) != self.hidden_size:\n            raise ValueError(\n                f\"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size} and \"\n                f\"`num_heads`: {self.num_heads}).\"\n            )\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear()\n\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            assert row_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, megatron_config)\n            tp_utils.update_kwargs_with_config(row_kwargs, megatron_config)\n\n        # [self.q_size, self.k_size, self.v_size]\n        self.qkv_proj = QKVParallelLinear(\n            input_size=self.hidden_size,\n            num_heads=self.num_heads,\n            num_key_value_heads=self.num_key_value_heads,\n            head_dim=self.head_dim,\n            # bias=config.attention_bias,\n            bias=True,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n        self.q_size = self.num_heads_per_tp * self.head_dim\n        self.k_size = self.num_key_value_heads_per_tp * self.head_dim\n        self.v_size = self.num_key_value_heads_per_tp * self.head_dim\n\n        self.o_proj = tensor_parallel.RowParallelLinear(\n            input_size=self.num_heads * self.head_dim,\n            output_size=self.hidden_size,\n            # bias=config.attention_bias,\n            bias=False,\n            input_is_parallel=True,\n            skip_bias_add=False,\n            **row_kwargs,\n        )\n\n        self._init_rope()\n\n    def _init_rope(self):\n        self.rotary_emb = Qwen2RotaryEmbedding(\n            self.head_dim,\n            max_position_embeddings=self.max_position_embeddings,\n            base=self.rope_theta,\n        )\n\n    def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):\n        return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n        bsz, q_len, _ = hidden_states.size()\n        qkv = self.qkv_proj(hidden_states)[0]\n        query_states, key_states, value_states = qkv.split([self.q_size, self.k_size, self.v_size], dim=-1)\n\n        query_states = query_states.view(bsz, q_len, self.num_heads_per_tp, self.head_dim).transpose(1, 2)\n        key_states = key_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2)\n        value_states = value_states.view(bsz, q_len, self.num_key_value_heads_per_tp, self.head_dim).transpose(1, 2)\n\n        kv_seq_len = key_states.shape[-2]\n        cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)\n        query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)\n\n        key_states = repeat_kv(key_states, self.num_key_value_groups)\n        value_states = repeat_kv(value_states, self.num_key_value_groups)\n\n        attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)\n\n        if attn_weights.size() != (bsz, self.num_heads_per_tp, q_len, kv_seq_len):\n            raise ValueError(\n                f\"Attention weights should be of size {(bsz, self.num_heads_per_tp, q_len, kv_seq_len)}, \"\n                f\"but is {attn_weights.size()}\"\n            )\n\n        if attention_mask is not None:\n            if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):\n                raise ValueError(\n                    f\"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}\"\n                )\n            attn_weights = attn_weights + attention_mask\n\n        # upcast attention to fp32\n        attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)\n        attn_output = torch.matmul(attn_weights, value_states)\n\n        if attn_output.size() != (bsz, self.num_heads_per_tp, q_len, self.head_dim):\n            raise ValueError(\n                f\"`attn_output` should be of size {(bsz, self.num_heads_per_tp, q_len, self.head_dim)}, \"\n                f\"but is {attn_output.size()}\"\n            )\n\n        attn_output = attn_output.transpose(1, 2).contiguous()\n        attn_output = attn_output.reshape(bsz, q_len, self.hidden_size_per_tp)\n        attn_output = self.o_proj(attn_output)[0]\n        return attn_output\n\n\n\"\"\"\nRemove padding Attention\n- Using Flash-attn 2\n- Compatible with sequence parallel\n\"\"\"\n\n\ndef apply_rotary_pos_emb_rmpad(q, k, cos, sin, position_ids, indices, sequence_length):\n    batch_size = position_ids.shape[0]\n\n    q = pad_input(q, indices, batch_size, sequence_length)  # (batch_size, seqlen, num_head, head_dim)\n    k = pad_input(k, indices, batch_size, sequence_length)\n    cos = cos[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]\n    sin = sin[position_ids].unsqueeze(2)  # [bs, seq_len, 1, dim]\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n\n    q_embed = index_first_axis(rearrange(q_embed, \"b s ... -> (b s) ...\"), indices)\n    k_embed = index_first_axis(rearrange(k_embed, \"b s ... -> (b s) ...\"), indices)\n\n    return q_embed, k_embed\n\n\n# use flash-attn rotary embeddings with rmpad\n# cos/sin shoudl be: (seq_length, rotary_dim / 2)\ndef apply_rotary_pos_emb_rmpad_flash(q, k, cos, sin, cu_seqlens, max_seqlen):\n    q_embed = apply_rotary_emb(\n        q, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen\n    )\n    k_embed = apply_rotary_emb(\n        k, cos, sin, interleaved=False, inplace=False, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen\n    )\n    return q_embed, k_embed\n\n\nclass ParallelQwen2AttentionRmPad(ParallelQwen2Attention):\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: torch.Tensor = None,\n        max_seqlen_in_batch: int = None,\n    ):\n        total_nnz, _, _ = hidden_states.size()  # This is the total_nnz padded after sequence parallel\n\n        if self.megatron_config.sequence_parallel:\n            total_nnz = total_nnz * mpu.get_tensor_model_parallel_world_size()\n\n        qkv = self.qkv_proj(hidden_states)[0]\n        query_states, key_states, value_states = qkv.split(\n            [self.q_size, self.k_size, self.v_size], dim=-1\n        )  # (total_nnz, 1, hidden_size)\n\n        if self.megatron_config.sequence_parallel:\n            sequence_parallel_pad = total_nnz - cu_seqlens[-1]\n            total_nnz = cu_seqlens[-1]  # total_nnz before sp padding\n            query_states = query_states[:total_nnz]\n            key_states = key_states[:total_nnz]\n            value_states = value_states[:total_nnz]\n\n        # Flash attention requires the input to have the shape\n        # batch_size x seq_length x head_dime x hidden_dim\n        # therefore we just need to keep the original shape\n        query_states = query_states.view(total_nnz, self.num_heads_per_tp, self.head_dim)\n        key_states = key_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim)\n        value_states = value_states.view(total_nnz, self.num_key_value_heads_per_tp, self.head_dim)\n\n        cos, sin = self.rotary_emb(value_states, seq_len=sequence_length)\n        cos, sin = cos[:, : cos.shape[1] // 2], sin[:, : sin.shape[1] // 2]  # flash attn only needs half\n        query_states, key_states = apply_rotary_pos_emb_rmpad_flash(\n            query_states, key_states, cos, sin, cu_seqlens=cu_seqlens, max_seqlen=max_seqlen_in_batch\n        )\n        # query_states, key_states = apply_rotary_pos_emb_rmpad(query_states, key_states, cos, sin,\n        # position_ids, indices,\n\n        # It is recommended to use dropout with FA according to the docs\n        # when training.\n        dropout_rate = 0.0  # if not self.training else self.attn_dropout\n\n        # In PEFT, usually we cast the layer norms in float32 for training stability reasons\n        # therefore the input hidden states gets silently casted in float32. Hence, we need\n        # cast them back in float16 just to be sure everything works as expected.\n        # This might slowdown training & inference so it is recommended to not cast the LayerNorms\n        # in fp32. (Qwen2RMSNorm handles it correctly)\n        input_dtype = query_states.dtype\n        if input_dtype == torch.float32:\n            query_states = query_states.to(torch.float16)\n            key_states = key_states.to(torch.float16)\n            value_states = value_states.to(torch.float16)\n\n        attn_output_unpad = flash_attn_varlen_func(\n            query_states,\n            key_states,\n            value_states,\n            cu_seqlens_q=cu_seqlens,\n            cu_seqlens_k=cu_seqlens,\n            max_seqlen_q=max_seqlen_in_batch,\n            max_seqlen_k=max_seqlen_in_batch,\n            dropout_p=dropout_rate,\n            softmax_scale=None,\n            causal=True,\n        )\n\n        attn_output_unpad = attn_output_unpad.to(input_dtype)\n        attn_output_unpad = attn_output_unpad.reshape(total_nnz, 1, self.hidden_size_per_tp).contiguous()\n\n        # sequence parallel reduce_scatter is performed inside RowColumnParallel if enabled\n        # Here we need to repad\n        if self.megatron_config.sequence_parallel:\n            attn_output_unpad = F.pad(attn_output_unpad, pad=(0, 0, 0, 0, 0, sequence_parallel_pad))\n\n        attn_output_unpad = self.o_proj(attn_output_unpad)[0]\n        return attn_output_unpad\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/parallel_decoder.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\nfrom megatron.core import ModelParallelConfig\nfrom torch import nn\nfrom transformers import Qwen2Config\n\nfrom verl.utils.megatron_utils import TransformerConfig, convert_config\n\nfrom .parallel_attention import ParallelQwen2Attention, ParallelQwen2AttentionRmPad\nfrom .parallel_mlp import ParallelQwen2MLP\nfrom .parallel_rmsnorm import ParallelQwen2RMSNorm\n\n\nclass ParallelQwen2DecoderLayer(nn.Module):\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, layer_idx: int):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.layer_idx = layer_idx\n        self.hidden_size = config.hidden_size\n        self.self_attn = ParallelQwen2Attention(config=config, megatron_config=megatron_config)\n\n        self.mlp = ParallelQwen2MLP(config, megatron_config=megatron_config)\n        self.input_layernorm = ParallelQwen2RMSNorm(config, megatron_config)\n        self.post_attention_layernorm = ParallelQwen2RMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:\n        \"\"\"\n        Args:\n            hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`\n            attention_mask (`torch.FloatTensor`, *optional*): attention mask of size\n                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.\n            output_attentions (`bool`, *optional*):\n                Whether or not to return the attentions tensors of all attention layers. See `attentions` under\n                returned tensors for more detail.\n            use_cache (`bool`, *optional*):\n                If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding\n                (see `past_key_values`).\n            past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states\n        \"\"\"\n\n        residual = hidden_states\n\n        hidden_states = self.input_layernorm(hidden_states)\n\n        # Note: sequence parallel is hidden inside ColumnParallelLinear\n        # reduce scatter is hidden inside RowParallelLinear\n\n        # Self Attention\n        hidden_states = self.self_attn(\n            hidden_states=hidden_states,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n        )\n\n        # TODO: add sequence parallel operator reduce_scatter here\n\n        hidden_states = residual + hidden_states\n\n        # Fully Connected\n        residual = hidden_states\n        hidden_states = self.post_attention_layernorm(hidden_states)\n\n        # TODO: add sequence parallel operator all_gather here\n\n        hidden_states = self.mlp(hidden_states)\n\n        # TODO: add sequence parallel operator reduce_scatter here\n\n        hidden_states = residual + hidden_states\n\n        outputs = hidden_states\n\n        return outputs\n\n\nclass ParallelQwen2DecoderLayerRmPad(nn.Module):\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, layer_idx: int):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.hidden_size = config.hidden_size\n        self.layer_idx = layer_idx\n        self.self_attn = ParallelQwen2AttentionRmPad(config=config, megatron_config=megatron_config)\n\n        self.mlp = ParallelQwen2MLP(config, megatron_config=megatron_config)\n        self.input_layernorm = ParallelQwen2RMSNorm(config, megatron_config)\n        self.post_attention_layernorm = ParallelQwen2RMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        hidden_states: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple[torch.FloatTensor, Optional[tuple[torch.FloatTensor, torch.FloatTensor]]]:\n        residual = hidden_states  # (total_nnz // sp, 1, hidden_size)\n\n        hidden_states = self.input_layernorm(hidden_states)\n\n        # Self Attention\n        # (total_nnz // sp, 1, hidden_size) -> all-gather (total_nnz, 1, hidden_size)\n        # -> col + row -> reduce-scatter -> (total_nnz // sp, 1, hidden_size)\n        hidden_states = self.self_attn(\n            hidden_states=hidden_states,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        hidden_states = residual + hidden_states\n\n        # Fully Connected\n        # shape changes same as attn\n        residual = hidden_states\n        hidden_states = self.post_attention_layernorm(hidden_states)\n        hidden_states = self.mlp(hidden_states)\n        hidden_states = residual + hidden_states\n\n        outputs = hidden_states\n\n        return outputs\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/parallel_linear.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023 The vLLM team.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/vllm-project/vllm/blob/main/vllm/model_executor/layers/linear.py\n\n\nfrom megatron.core import tensor_parallel\n\n\nclass QKVParallelLinear(tensor_parallel.ColumnParallelLinear):\n    def __init__(\n        self,\n        input_size,\n        num_heads,\n        num_key_value_heads,\n        head_dim,\n        *,\n        bias=True,\n        gather_output=True,\n        skip_bias_add=False,\n        **kwargs,\n    ):\n        # Keep input parameters, and already restrict the head numbers\n        self.input_size = input_size\n        self.q_output_size = num_heads * head_dim\n        self.kv_output_size = num_key_value_heads * head_dim\n        self.head_dim = head_dim\n        self.gather_output = gather_output\n        self.skip_bias_add = skip_bias_add\n\n        input_size = self.input_size\n        output_size = (num_heads + 2 * num_key_value_heads) * self.head_dim\n\n        super().__init__(\n            input_size=input_size,\n            output_size=output_size,\n            bias=bias,\n            gather_output=gather_output,\n            skip_bias_add=skip_bias_add,\n            **kwargs,\n        )\n\n\nclass MergedColumnParallelLinear(tensor_parallel.ColumnParallelLinear):\n    def __init__(\n        self,\n        input_size,\n        gate_ouput_size,\n        up_output_size,\n        *,\n        bias=True,\n        gather_output=True,\n        skip_bias_add=False,\n        **kwargs,\n    ):\n        # Keep input parameters, and already restrict the head numbers\n        self.input_size = input_size\n        self.output_size = gate_ouput_size + up_output_size\n        self.gather_output = gather_output\n        self.skip_bias_add = skip_bias_add\n\n        super().__init__(\n            input_size=self.input_size,\n            output_size=self.output_size,\n            bias=bias,\n            gather_output=gather_output,\n            skip_bias_add=skip_bias_add,\n            **kwargs,\n        )\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/parallel_mlp.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom megatron.core import ModelParallelConfig, tensor_parallel\nfrom megatron.core import parallel_state as mpu\nfrom torch import nn\nfrom transformers.activations import ACT2FN\n\nfrom verl.models.qwen2.megatron.layers.parallel_linear import MergedColumnParallelLinear\nfrom verl.utils.megatron import tensor_parallel as tp_utils\n\n\nclass ParallelQwen2MLP(nn.Module):\n    def __init__(self, config, megatron_config: ModelParallelConfig = None) -> None:\n        super().__init__()\n        self.config = config\n        self.hidden_size = config.hidden_size\n        self.intermediate_size = config.intermediate_size\n        # The weight is only [hidden_size, intermediate_size // model_parallel_world_size]\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        row_kwargs = tp_utils.get_default_kwargs_for_row_parallel_linear()\n\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            assert row_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(row_kwargs, megatron_config)\n            tp_utils.update_kwargs_with_config(column_kwargs, megatron_config)\n\n        tp_size = mpu.get_tensor_model_parallel_world_size()\n\n        self.gate_up_proj = MergedColumnParallelLinear(\n            input_size=self.hidden_size,\n            gate_ouput_size=self.intermediate_size,\n            up_output_size=self.intermediate_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n        self.gate_size = self.intermediate_size // tp_size\n\n        self.down_proj = tensor_parallel.RowParallelLinear(\n            input_size=self.intermediate_size,\n            output_size=self.hidden_size,\n            bias=False,\n            input_is_parallel=True,\n            skip_bias_add=False,\n            **row_kwargs,\n        )\n\n        self.act_fn = ACT2FN[config.hidden_act]\n\n    def forward(self, x):\n        gate_up = self.gate_up_proj(x)[0]\n        gate, up = gate_up.split(self.gate_size, dim=-1)\n        return self.down_proj(self.act_fn(gate) * up)[0]\n"
  },
  {
    "path": "verl/models/qwen2/megatron/layers/parallel_rmsnorm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport numbers\n\nimport torch\nfrom apex.normalization.fused_layer_norm import fused_rms_norm_affine\nfrom megatron.core import ModelParallelConfig\nfrom torch import nn\nfrom transformers import Qwen2Config\n\nfrom verl.utils.megatron import sequence_parallel as sp_utils\n\n\nclass ParallelQwen2RMSNorm(nn.Module):\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        \"\"\"\n        Qwen2RMSNorm is equivalent to T5LayerNorm\n        \"\"\"\n        super().__init__()\n        if isinstance(config.hidden_size, numbers.Integral):\n            normalized_shape = (config.hidden_size,)\n        self.normalized_shape = torch.Size(normalized_shape)\n        self.weight = nn.Parameter(torch.ones(self.normalized_shape))\n        self.variance_epsilon = config.rms_norm_eps\n\n        if megatron_config.sequence_parallel:\n            sp_utils.mark_parameter_as_sequence_parallel(self.weight)\n\n    def forward(self, hidden_states):\n        return fused_rms_norm_affine(\n            input=hidden_states,\n            weight=self.weight,\n            normalized_shape=self.normalized_shape,\n            eps=self.variance_epsilon,\n            memory_efficient=True,\n        )\n"
  },
  {
    "path": "verl/models/qwen2/megatron/modeling_qwen2_megatron.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n#\n# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX\n# and OPT implementations in this library. It has been modified from its\n# original forms to accommodate minor architectural differences compared\n# to GPT-NeoX and OPT used by the Meta AI team that trained the model.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"PyTorch Qwen2 model.\"\"\"\n\nfrom typing import Optional\n\nimport torch\nimport torch.utils.checkpoint\nfrom megatron.core import ModelParallelConfig, mpu, parallel_state, tensor_parallel\nfrom torch import nn\nfrom transformers.modeling_outputs import BaseModelOutputWithPast\nfrom transformers.models.qwen2.configuration_qwen2 import Qwen2Config\nfrom transformers.models.qwen2.modeling_qwen2 import CausalLMOutputWithPast\n\nfrom verl.utils.device import get_device_name\nfrom verl.utils.megatron import sequence_parallel as sp_utils\nfrom verl.utils.megatron import tensor_parallel as tp_utils\nfrom verl.utils.megatron_utils import TransformerConfig, convert_config\n\nfrom .layers import ParallelQwen2DecoderLayer, ParallelQwen2DecoderLayerRmPad, ParallelQwen2RMSNorm\n\n\"\"\"\nTODO: \n1. Add weight initialization. Here we need to be careful on TP weight init.\n2. Add sequence parallel\n3. Load checkpoint from Qwen2 pretrained checkpoint\n\"\"\"\n\n\n# Copied from transformers.models.bart.modeling_bart._make_causal_mask\ndef _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device):\n    \"\"\"\n    Make causal mask used for bi-directional self-attention.\n    \"\"\"\n    bsz, tgt_len = input_ids_shape\n    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)\n    mask_cond = torch.arange(mask.size(-1), device=device)\n    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)\n    mask = mask.to(dtype)\n    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len)\n\n\n# Copied from transformers.models.bart.modeling_bart._expand_mask\ndef _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):\n    \"\"\"\n    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.\n    \"\"\"\n    bsz, src_len = mask.size()\n    tgt_len = tgt_len if tgt_len is not None else src_len\n\n    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)\n\n    inverted_mask = 1.0 - expanded_mask\n\n    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)\n\n\nclass ParallelQwen2Model(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]\n\n    Args:\n        config: Qwen2Config\n    \"\"\"\n\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, megatron_config)\n        self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n            num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n        )\n\n        self.layers = nn.ModuleList(\n            [ParallelQwen2DecoderLayer(config, megatron_config) for _ in range(config.num_hidden_layers)]\n        )\n        self.norm = ParallelQwen2RMSNorm(config, megatron_config)\n\n    # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask\n    def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds):\n        # create causal mask\n        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n        combined_attention_mask = None\n        if input_shape[-1] > 1:\n            combined_attention_mask = _make_causal_mask(\n                input_shape,\n                inputs_embeds.dtype,\n                device=inputs_embeds.device,\n            )\n\n        if attention_mask is not None:\n            # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n            expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(\n                inputs_embeds.device\n            )\n            combined_attention_mask = (\n                expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask\n            )\n\n        return combined_attention_mask\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (batch_size, seq_length)\n            attention_mask: attention_mask. shape (batch_size, seq_length)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        batch_size, seq_length = input_ids.shape\n        inputs_embeds = self.embed_tokens(input_ids)\n        # embed positions\n\n        attention_mask = self._prepare_decoder_attention_mask(attention_mask, (batch_size, seq_length), inputs_embeds)\n\n        hidden_states = inputs_embeds\n\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                attention_mask=attention_mask,\n                position_ids=position_ids,\n            )\n\n            hidden_states = layer_outputs\n\n        hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelQwen2ForCausalLM(nn.Module):\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.model = ParallelQwen2Model(config, megatron_config=megatron_config)\n        self.vocab_size = config.vocab_size\n\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n\n        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\n        outputs = self.model(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n        )\n\n        hidden_states = outputs\n        logits = self.lm_head(hidden_states)[0]\n\n        logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)\n\n        logits = logits.float()\n        return CausalLMOutputWithPast(\n            loss=None,\n            logits=logits,\n            past_key_values=None,\n            hidden_states=None,\n            attentions=None,\n        )\n\n\nfrom flash_attn.bert_padding import index_first_axis, pad_input, unpad_input  # noqa: F401, E402\n\n\nclass ParallelQwen2ModelRmPad(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]\n\n    Args:\n        config: Qwen2Config\n    \"\"\"\n\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        self.megatron_config = megatron_config\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config)\n        self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n            num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n        )\n\n        self.layers = nn.ModuleList(\n            [ParallelQwen2DecoderLayerRmPad(config, megatron_config) for _ in range(config.num_hidden_layers)]\n        )\n        self.norm = ParallelQwen2RMSNorm(config, megatron_config)\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (1, totol_nnz)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        inputs_embeds = self.embed_tokens(input_ids)  # (1, total_nnz) -> (1, total_nnz, hidden_size)\n\n        # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size)\n        inputs_embeds = inputs_embeds.transpose(0, 1)\n        if self.megatron_config.sequence_parallel:\n            inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds)\n\n        hidden_states = inputs_embeds\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                position_ids=position_ids,\n                sequence_length=sequence_length,\n                indices=indices,\n                cu_seqlens=cu_seqlens,\n                max_seqlen_in_batch=max_seqlen_in_batch,\n            )\n\n            hidden_states = layer_outputs\n\n        hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelQwen2ForCausalLMRmPad(nn.Module):\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.megatron_config = megatron_config\n        self.model = ParallelQwen2ModelRmPad(config, megatron_config=megatron_config)\n        self.vocab_size = config.vocab_size\n        self._init_head(config)\n\n    def _init_head(self, config: Qwen2Config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            **column_kwargs,\n        )\n\n    def _forward_head(self, hidden_states):\n        # all_gather from sequence parallel region is performed inside lm_head\n        logits = self.lm_head(hidden_states)[0]\n        logits = logits.float()  # (total_nnz_padded, 1, vocab_size // tp)\n        logits = tensor_parallel.gather_from_tensor_model_parallel_region(logits)  # (total_nnz_padded, 1, vocab_size)\n        return logits\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n        batch_size, sequence_length = input_ids.shape\n\n        # remove padding here\n        input_ids, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(\n            input_ids.unsqueeze(dim=-1), attention_mask\n        )  # (total_nnz, 1)\n\n        # pad input_ids to multiple of tp for all tp ranks\n        # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap\n        if self.megatron_config.sequence_parallel:\n            input_ids = sp_utils.pad_to_sequence_parallel(input_ids)\n\n        input_ids = input_ids.transpose(0, 1)  # (1, total_nnz+pad)\n\n        outputs = self.model(\n            input_ids=input_ids,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        hidden_states = outputs\n\n        logits = self._forward_head(hidden_states)\n\n        # remove padding from sequence parallel\n        if self.megatron_config.sequence_parallel:\n            totol_nnz = cu_seqlens[-1]\n            logits = logits[:totol_nnz]  # (total_nnz_padded)\n\n        logits = torch.squeeze(logits, dim=1)  # remove the artificial batch dimension\n        # add removed padding back\n        logits = pad_input(\n            logits, indices, batch_size, seqlen=sequence_length\n        )  # (batch_size, sequence_length, vocab_size)\n\n        return CausalLMOutputWithPast(\n            loss=None,\n            logits=logits,\n            past_key_values=None,\n            hidden_states=None,\n            attentions=None,\n        )\n\n\nclass ParallelQwen2ForValueRmPad(ParallelQwen2ForCausalLMRmPad):\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = nn.Linear(in_features=config.hidden_size, out_features=1, bias=False)\n        # lm_head is effectively the same as sequence parallel\n        sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)\n\n    def _forward_head(self, hidden_states):\n        logits = self.lm_head(hidden_states)  # (total_nnz_padded // tp, 1, 1)\n        logits = logits.float()\n        if self.megatron_config.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits\n\n    def forward(\n        self,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        output = super().forward(input_ids, attention_mask, position_ids)\n        output.logits = torch.squeeze(output.logits, dim=-1)\n        return output\n\n\n\"\"\"\nSupport pipeline parallelism\n\"\"\"\n\n\nclass ParallelQwen2ModelRmPadPP(nn.Module):\n    \"\"\"\n    Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`Qwen2DecoderLayer`]\n    This model definition supports pipeline parallelism. To support pp and vpp,\n    - This model only contains layer in this pp stage and vpp chunk\n    - When calling get_model in Megatron, this rank will instantiate all the vpp chunks in this pp.\n    Args:\n        config: Qwen2Config\n    \"\"\"\n\n    def __init__(self, config: Qwen2Config, megatron_config: ModelParallelConfig, pre_process, post_process):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.padding_idx = config.pad_token_id\n        self.vocab_size = config.vocab_size\n        self.pre_process = pre_process\n        self.post_process = post_process\n        self.megatron_config = megatron_config\n        embedding_kwargs = tp_utils.get_default_kwargs_for_parallel_embedding()\n        if megatron_config is not None:\n            assert embedding_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(embedding_kwargs, self.megatron_config)\n        if pre_process:\n            self.embed_tokens = tensor_parallel.VocabParallelEmbedding(\n                num_embeddings=config.vocab_size, embedding_dim=config.hidden_size, **embedding_kwargs\n            )\n        else:\n            self.embed_tokens = None\n\n        pp_rank = mpu.get_pipeline_model_parallel_rank()\n        pp_size = megatron_config.pipeline_model_parallel_size\n        self.num_layer_per_pp = config.num_hidden_layers // pp_size\n        vpp_size = megatron_config.virtual_pipeline_model_parallel_size\n        vpp_rank = mpu.get_virtual_pipeline_model_parallel_rank()\n\n        if vpp_size is not None:\n            self.num_layer_vpp_chunk = self.num_layer_per_pp // vpp_size\n            self.num_layer_this_model = self.num_layer_vpp_chunk\n            offset = vpp_rank * (config.num_hidden_layers // vpp_size) + (pp_rank * self.num_layer_vpp_chunk)\n        else:\n            self.num_layer_this_model = self.num_layer_per_pp\n            offset = pp_rank * self.num_layer_per_pp\n\n        self.layers = nn.ModuleList()\n        for i in range(self.num_layer_this_model):\n            layer = ParallelQwen2DecoderLayerRmPad(config, megatron_config, layer_idx=i + offset)\n            self.layers.add_module(f\"{i}\", layer)\n\n        if post_process:\n            self.norm = ParallelQwen2RMSNorm(config, megatron_config)\n        else:\n            self.norm = None\n\n    def set_input_tensor(self, input_tensor):\n        \"\"\"Set input tensor to be used instead of forward()'s input.\n\n        When doing pipeline parallelism the input from the previous\n        stage comes from communication, not from the input, so the\n        model's forward_step_func won't have it. This function is thus\n        used by internal code to bypass the input provided by the\n        forward_step_func\"\"\"\n        self.input_tensor = input_tensor\n\n    def forward(\n        self,\n        input_ids: torch.Tensor,\n        position_ids: Optional[torch.LongTensor] = None,\n        sequence_length: int = None,\n        indices: torch.Tensor = None,\n        cu_seqlens: int = None,\n        max_seqlen_in_batch: int = None,\n    ) -> tuple | BaseModelOutputWithPast:\n        \"\"\"\n\n        Args:\n            input_ids: input ids. shape (1, totol_nnz)\n            position_ids: position ids. shape (batch_size, seq_length)\n\n        Returns:\n\n        \"\"\"\n        if self.pre_process:\n            inputs_embeds = self.embed_tokens(input_ids)  # (1, total_nnz) -> (1, total_nnz, hidden_size)\n\n            # vocab parallel embedding will not do sequence parallel reduce-scatter in open source megatron\n            # so need to deal with it by handle here:\n            # (1, total_nnz, hidden_size) -> (total_nnz, 1, hidden_size) -> (total_nnz // sp, 1, hidden_size)\n            inputs_embeds = inputs_embeds.transpose(0, 1)\n            if self.megatron_config.sequence_parallel:\n                inputs_embeds = tensor_parallel.scatter_to_sequence_parallel_region(inputs_embeds)\n\n            hidden_states = inputs_embeds\n        else:\n            # self.hidden_states should be passed by Megatron\n            hidden_states = self.input_tensor\n\n        for idx, decoder_layer in enumerate(self.layers):\n            layer_outputs = decoder_layer(\n                hidden_states,\n                position_ids=position_ids,\n                sequence_length=sequence_length,\n                indices=indices,\n                cu_seqlens=cu_seqlens,\n                max_seqlen_in_batch=max_seqlen_in_batch,\n            )\n\n            hidden_states = layer_outputs\n\n        if self.post_process:\n            hidden_states = self.norm(hidden_states)\n\n        return hidden_states\n\n\nclass ParallelQwen2ForCausalLMRmPadPP(nn.Module):\n    def __init__(\n        self,\n        config: Qwen2Config,\n        megatron_config: ModelParallelConfig,\n        pre_process,\n        post_process,\n        share_embeddings_and_output_weights,\n    ):\n        super().__init__()\n        self.config: TransformerConfig = convert_config(config, megatron_config)\n        self.megatron_config = megatron_config\n        self.model = ParallelQwen2ModelRmPadPP(\n            config, megatron_config=megatron_config, pre_process=pre_process, post_process=post_process\n        )\n        self.share_embeddings_and_output_weights = share_embeddings_and_output_weights\n        self.vocab_size = config.vocab_size\n        self.pre_process = pre_process\n        self.post_process = post_process\n        if post_process:\n            self._init_head(config)\n        if pre_process or post_process:\n            self.setup_embeddings_and_output_layer()\n\n    def set_input_tensor(self, input_tensor):\n        \"\"\"Set input tensor to be used instead of forward()'s input.\n\n        When doing pipeline parallelism the input from the previous\n        stage comes from communication, not from the input, so the\n        model's forward_step_func won't have it. This function is thus\n        used by internal code to bypass the input provided by the\n        forward_step_func\"\"\"\n        assert len(input_tensor) == 1\n        self.model.set_input_tensor(input_tensor[0])\n\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = tensor_parallel.ColumnParallelLinear(\n            input_size=config.hidden_size,\n            output_size=config.vocab_size,\n            bias=False,\n            gather_output=False,\n            skip_bias_add=False,\n            skip_weight_param_allocation=self.pre_process and self.share_embeddings_and_output_weights,\n            **column_kwargs,\n        )\n\n    def setup_embeddings_and_output_layer(self) -> None:\n        \"\"\"Sets up embedding layer in first stage and output layer in last stage.\n\n        This function initializes word embeddings in the final stage when we are\n        using pipeline parallelism and sharing word embeddings, and sets up param\n        attributes on the embedding and output layers.\n        \"\"\"\n        # Set `is_embedding_or_output_parameter` attribute.\n        if self.pre_process:\n            self.model.embed_tokens.weight.is_embedding_or_output_parameter = True\n        if self.post_process and self.lm_head.weight is not None:\n            self.lm_head.weight.is_embedding_or_output_parameter = True\n\n        if not self.share_embeddings_and_output_weights:\n            return\n\n        if parallel_state.get_pipeline_model_parallel_world_size() == 1:\n            # Zero out wgrad if sharing embeddings between two layers on same\n            # pipeline stage to make sure grad accumulation into main_grad is\n            # correct and does not include garbage values (e.g., from torch.empty).\n            self.shared_embedding_or_output_weight().zero_out_wgrad = True\n            return\n\n        if parallel_state.is_pipeline_first_stage() and self.pre_process and not self.post_process:\n            self.shared_embedding_or_output_weight().shared_embedding = True\n\n        if self.post_process and not self.pre_process:\n            assert not parallel_state.is_pipeline_first_stage()\n            # set word_embeddings weights to 0 here, then copy first\n            # stage's weights using all_reduce below.\n            self.lm_head.weight.data.fill_(0)\n            self.lm_head.weight.shared = True\n            self.lm_head.weight.shared_embedding = True\n\n        if torch.distributed.is_initialized() and parallel_state.is_rank_in_embedding_group():\n            weight = self.shared_embedding_or_output_weight()\n            weight.data = weight.data.to(get_device_name())\n            torch.distributed.all_reduce(weight.data, group=parallel_state.get_embedding_group())\n\n    def shared_embedding_or_output_weight(self) -> torch.Tensor:\n        if self.pre_process:\n            return self.model.embed_tokens.weight\n        elif self.post_process:\n            return self.lm_head.weight\n        return None\n\n    def _forward_head(self, hidden_states):\n        # all_gather from sequence parallel region is performed inside lm_head\n        # print(f'logits shape before forward_head: {hidden_states.shape}, vocab_size = '\n        # f'{self.config.vocab_size}') # [4, 32, 4096]\n        output_weight = None\n        if self.share_embeddings_and_output_weights:\n            output_weight = self.shared_embedding_or_output_weight()\n        logits = self.lm_head(hidden_states, weight=output_weight)[0]\n        # print(f'logits shape after forward_head: {logits.shape}') # [8, 32, 8]\n        logits = logits.float()  # (total_nnz_padded, 1, vocab_size // tp)\n        return logits\n\n    def forward(\n        self,\n        # original input\n        *,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        r\"\"\"\n        Args:\n            labels (`torch.LongTensor` of shape `(batch_size, sequence_length)`, *optional*):\n                Labels for computing the masked language modeling loss. Indices should either be in `[0, ...,\n                config.vocab_size]` or -100 (see `input_ids` docstring). Tokens with indices set to `-100` are ignored\n                (masked), the loss is only computed for the tokens with labels in `[0, ..., config.vocab_size]`.\n\n        Returns:\n        ```\"\"\"\n\n        # Note that input_ids, attention_mask and position_ids should be passed to every pp layer.\n        # In the first pp, input_ids will be used, in other pp layers hidden_states will be used inside self.model\n        batch_size, sequence_length = input_ids.shape\n        # remove padding here\n        input_ids_rmpad, indices, cu_seqlens, max_seqlen_in_batch, *_ = unpad_input(\n            input_ids.unsqueeze(dim=-1), attention_mask\n        )  # (total_nnz, 1)\n\n        # pad input_ids to multiple of tp for all tp ranks\n        # TODO: for better performance, the sp padding should be removed at each layer. Not sure the performance gap\n        if self.megatron_config.sequence_parallel:\n            input_ids_rmpad = sp_utils.pad_to_sequence_parallel(input_ids_rmpad)\n\n        input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz+pad)\n\n        outputs = self.model(\n            input_ids=input_ids_rmpad,\n            position_ids=position_ids,\n            sequence_length=sequence_length,\n            indices=indices,\n            cu_seqlens=cu_seqlens,\n            max_seqlen_in_batch=max_seqlen_in_batch,\n        )\n\n        if self.post_process:\n            hidden_states = outputs\n            logits = self._forward_head(hidden_states)\n            logits = torch.squeeze(logits, dim=1)  # remove the artificial batch dimension # torch.Size([8, 32, 16])\n\n            # remove padding from sequence parallel\n            if self.megatron_config.sequence_parallel:\n                totol_nnz = cu_seqlens[-1]\n                logits = logits[:totol_nnz]  # (total_nnz_padded)\n            # add removed padding back. If input is already rmpad, we let the caller pad_input\n            logits = pad_input(\n                logits, indices, batch_size, seqlen=sequence_length\n            )  # (batch_size, sequence_length, vocab_size)\n\n            return CausalLMOutputWithPast(\n                loss=None,\n                logits=logits,\n                past_key_values=None,\n                hidden_states=None,\n                attentions=None,\n            )\n        else:\n            return outputs\n\n\nclass ParallelQwen2ForValueRmPadPP(ParallelQwen2ForCausalLMRmPadPP):\n    def _init_head(self, config):\n        column_kwargs = tp_utils.get_default_kwargs_for_column_parallel_linear()\n        if self.megatron_config is not None:\n            assert column_kwargs.get(\"config\", False), \"must have ModelParallelConfig\"\n            tp_utils.update_kwargs_with_config(column_kwargs, self.megatron_config)\n        self.lm_head = nn.Linear(in_features=config.hidden_size, out_features=1, bias=False)\n        # lm_head is effectively the same as sequence parallel\n        sp_utils.mark_parameter_as_sequence_parallel(self.lm_head.weight)\n\n    def _forward_head(self, hidden_states):\n        logits = self.lm_head(hidden_states)  # (total_nnz_padded // tp, 1, 1)\n        logits = logits.float()\n        if self.megatron_config.sequence_parallel:\n            logits = tensor_parallel.gather_from_sequence_parallel_region(logits, tensor_parallel_output_grad=False)\n        return logits\n\n    def forward(\n        self,\n        *,\n        input_ids: torch.LongTensor = None,\n        attention_mask: Optional[torch.Tensor] = None,\n        position_ids: Optional[torch.LongTensor] = None,\n    ) -> tuple | CausalLMOutputWithPast:\n        output = super().forward(input_ids=input_ids, attention_mask=attention_mask, position_ids=position_ids)\n        if self.post_process:\n            output.logits = torch.squeeze(output.logits, dim=-1)\n            return output\n        else:\n            return output\n"
  },
  {
    "path": "verl/models/registry.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport importlib\nfrom typing import Optional\n\nimport torch.nn as nn\n\n# Supported models in Megatron-LM\n# Architecture -> (module, class).\n_MODELS = {\n    \"LlamaForCausalLM\": (\n        \"llama\",\n        (\"ParallelLlamaForCausalLMRmPadPP\", \"ParallelLlamaForValueRmPadPP\", \"ParallelLlamaForCausalLMRmPad\"),\n    ),\n    \"Qwen2ForCausalLM\": (\n        \"qwen2\",\n        (\"ParallelQwen2ForCausalLMRmPadPP\", \"ParallelQwen2ForValueRmPadPP\", \"ParallelQwen2ForCausalLMRmPad\"),\n    ),\n    \"MistralForCausalLM\": (\n        \"mistral\",\n        (\"ParallelMistralForCausalLMRmPadPP\", \"ParallelMistralForValueRmPadPP\", \"ParallelMistralForCausalLMRmPad\"),\n    ),\n    \"ApertusForCausalLM\": (\n        \"apertus\",\n        (\"ParallelApertusForCausalLMRmPadPP\", \"ParallelApertusForValueRmPadPP\", \"ParallelApertusForCausalLMRmPad\"),\n    ),\n}\n\n\n# return model class\nclass ModelRegistry:\n    @staticmethod\n    def load_model_cls(model_arch: str, value=False) -> Optional[type[nn.Module]]:\n        if model_arch not in _MODELS:\n            return None\n\n        megatron = \"megatron\"\n\n        module_name, model_cls_name = _MODELS[model_arch]\n        if not value:  # actor/ref\n            model_cls_name = model_cls_name[0]\n        elif value:  # critic/rm\n            model_cls_name = model_cls_name[1]\n\n        module = importlib.import_module(f\"verl.models.{module_name}.{megatron}.modeling_{module_name}_megatron\")\n        return getattr(module, model_cls_name, None)\n\n    @staticmethod\n    def get_supported_archs() -> list[str]:\n        return list(_MODELS.keys())\n"
  },
  {
    "path": "verl/models/transformers/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom verl.models.transformers.monkey_patch import apply_monkey_patch\nfrom verl.models.transformers.tiled_mlp import apply_tiled_mlp_monkey_patch\n\n__all__ = [\n    \"apply_monkey_patch\",\n    \"apply_tiled_mlp_monkey_patch\",\n]\n"
  },
  {
    "path": "verl/models/transformers/apertus.py",
    "content": "# Copyright 2025 The SwissAI Initiative\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport sys\nfrom typing import Callable, Optional\n\nimport torch\n\nif sys.version_info >= (3, 11):\n    pass\nelse:\n    pass\n\nfrom transformers.cache_utils import Cache\nfrom transformers.models.apertus.modeling_apertus import apply_rotary_pos_emb\nfrom transformers.utils import logging\n\n# Import compatibility wrapper for flash_attn_supports_top_left_mask\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\nlogger = logging.get_logger(__name__)\n\n\ndef apertus_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    position_embeddings: tuple[torch.Tensor, torch.Tensor],\n    attention_mask: Optional[torch.Tensor],\n    past_key_value: Optional[Cache] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n    \"\"\"\n    Adapted from transformers 4.49.0 to support Ulysses sequence parallelism for transformers >= 4.48.0.\n\n    Key differences from Llama attention:\n    - QK normalization applied after Q/K projections\n\n        NOTE: This function has been tested only on transformers versions between 4.48.0 and 4.50.0.\n    \"\"\"\n    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS\n    from transformers.models.apertus.modeling_apertus import eager_attention_forward\n\n    bsz, q_len, _ = hidden_states.shape\n\n    query_states = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    key_states = self.k_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    value_states = self.v_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n\n    query_states = self.q_norm(query_states)\n    key_states = self.k_norm(key_states)\n\n    ########## AlltoAll for Ulysses ##########\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.config.num_attention_heads, ulysses_sp_size)\n\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n\n    full_q_len = query_states.size(2)\n\n    cos, sin = position_embeddings\n    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n\n    if past_key_value is not None:\n        # sin and cos are specific to RoPE models; cache_position needed for the static cache\n        cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position}\n        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)\n\n    attention_interface: Callable = eager_attention_forward\n    if self.config._attn_implementation != \"eager\":\n        if self.config._attn_implementation == \"sdpa\" and kwargs.get(\"output_attentions\", False):\n            logger.warning_once(\n                \"`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. \"\n                \"Falling back to eager attention. This warning can be removed using the argument \"\n                '`attn_implementation=\"eager\"` when loading the model.'\n            )\n        else:\n            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]\n\n    attn_output, attn_weights = attention_interface(\n        self,\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        dropout=0.0 if not self.training else self.attention_dropout,\n        scaling=self.scaling,\n        **kwargs,\n    )\n\n    attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1:\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n    attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()\n    attn_output = self.o_proj(attn_output)\n    return attn_output, attn_weights\n"
  },
  {
    "path": "verl/models/transformers/dense_common.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass\nfrom typing import Optional, Union\n\nimport torch\nfrom transformers.cache_utils import Cache\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\n\n@dataclass\nclass CausalLMOutputForPPO(CausalLMOutputWithPast):\n    log_probs: Optional[torch.FloatTensor] = None\n    entropy: Optional[torch.FloatTensor] = None\n\n\ndef forward_base_model(\n    self,\n    input_ids: Optional[torch.LongTensor] = None,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_values: Optional[Cache] = None,\n    inputs_embeds: Optional[torch.FloatTensor] = None,\n    use_cache: Optional[bool] = None,\n    output_attentions: Optional[bool] = None,\n    output_hidden_states: Optional[bool] = None,\n    return_dict: Optional[bool] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n) -> CausalLMOutputWithPast:\n    r\"\"\"\n    Copy paste LLaMa's forward\n    https://github.com/linkedin/Liger-Kernel/blob/main/src/liger_kernel/transformers/model/llama.py\n\n    This function should be generic enough for all pure text models.\n    ```\"\"\"\n\n    output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions\n    output_hidden_states = (\n        output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states\n    )\n\n    # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)\n    outputs = self.model(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        past_key_values=past_key_values,\n        inputs_embeds=inputs_embeds,\n        use_cache=use_cache,\n        output_attentions=output_attentions,\n        output_hidden_states=output_hidden_states,\n        return_dict=return_dict,\n        cache_position=cache_position,\n    )\n\n    return outputs\n\n\ndef forward_with_torch_backend(\n    self,\n    input_ids: torch.LongTensor = None,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_values: Optional[Union[\"Cache\", list[torch.FloatTensor]]] = None,\n    inputs_embeds: Optional[torch.FloatTensor] = None,\n    labels: Optional[torch.LongTensor] = None,\n    use_cache: Optional[bool] = None,\n    output_attentions: Optional[bool] = None,\n    output_hidden_states: Optional[bool] = None,\n    return_dict: Optional[bool] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n    logits_to_keep: int | torch.Tensor = 0,\n    temperature: float = 1.0,\n    **loss_kwargs,\n) -> tuple | CausalLMOutputForPPO:\n    from verl.utils.experimental.torch_functional import FusedLinearForPPO\n\n    outputs = forward_base_model(\n        self,\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        past_key_values=past_key_values,\n        inputs_embeds=inputs_embeds,\n        use_cache=use_cache,\n        output_attentions=output_attentions,\n        output_hidden_states=output_hidden_states,\n        cache_position=cache_position,\n    )\n\n    hidden_states = outputs[0]\n\n    if not return_dict:\n        raise NotImplementedError(\"forward_with_torch_backend has to return_dict\")\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_torch_backend, either labels or input_ids must be provided.\")\n\n    fused_linear_for_ppo = FusedLinearForPPO()\n    log_probs, entropy = fused_linear_for_ppo.forward(\n        hidden_states=hidden_states,\n        vocab_weights=self.lm_head.weight,\n        input_ids=rolled_labels,\n        temperature=temperature,\n    )\n\n    return CausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        past_key_values=outputs.past_key_values,\n        hidden_states=outputs.hidden_states,\n        attentions=outputs.attentions,\n    )\n\n\ndef forward_with_triton_backend(\n    self,\n    input_ids: torch.LongTensor = None,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_values: Optional[Union[\"Cache\", list[torch.FloatTensor]]] = None,\n    inputs_embeds: Optional[torch.FloatTensor] = None,\n    labels: Optional[torch.LongTensor] = None,\n    use_cache: Optional[bool] = None,\n    output_attentions: Optional[bool] = None,\n    output_hidden_states: Optional[bool] = None,\n    return_dict: Optional[bool] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n    logits_to_keep: int | torch.Tensor = 0,\n    temperature: float = 1.0,\n    **loss_kwargs,\n) -> tuple | CausalLMOutputForPPO:\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\n\n    outputs = forward_base_model(\n        self,\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        past_key_values=past_key_values,\n        inputs_embeds=inputs_embeds,\n        use_cache=use_cache,\n        output_attentions=output_attentions,\n        output_hidden_states=output_hidden_states,\n        return_dict=return_dict,\n        cache_position=cache_position,\n    )\n\n    hidden_states = outputs[0]\n\n    if not return_dict:\n        raise NotImplementedError(\"forward_with_triton_backend has to return_dict\")\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_triton_backend, either labels or input_ids must be provided.\")\n\n    log_probs, entropy = linear_cross_entropy(\n        hidden_states,\n        self.lm_head.weight,\n        rolled_labels,\n        temperature,\n        \"none\",\n    )\n\n    return CausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        past_key_values=outputs.past_key_values,\n        hidden_states=outputs.hidden_states,\n        attentions=outputs.attentions,\n    )\n"
  },
  {
    "path": "verl/models/transformers/glm4v.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport itertools\nimport logging\nimport os\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.distributed as dist\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward, fa_peft_integration_check\nfrom transformers.models.glm4v.modeling_glm4v import (\n    Glm4vCausalLMOutputWithPast,\n    Glm4vForConditionalGeneration,\n    Glm4vTextAttention,\n)\nfrom transformers.utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10\n\nfrom verl.utils.device import is_npu_available\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_group,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nif is_flash_attn_2_available():\n    from flash_attn import flash_attn_func, flash_attn_varlen_func\n\n    _flash_supports_window_size = \"window_size\" in inspect.signature(flash_attn_func).parameters\n    _flash_supports_deterministic = \"deterministic\" in inspect.signature(flash_attn_func).parameters\n    _flash_use_top_left_mask = not is_flash_attn_greater_or_equal_2_10()\n\nif is_npu_available:\n    from transformers.integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func\n    from transformers.integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func\n    from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask\n\n    _flash_supports_window_size = \"window_size\" in inspect.signature(flash_attn_func).parameters\n    _flash_supports_deterministic = \"deterministic\" in inspect.signature(flash_attn_func).parameters\n    _flash_use_top_left_mask = flash_attn_supports_top_left_mask()\n\n_flash_deterministic_enabled = os.getenv(\"FLASH_ATTENTION_DETERMINISTIC\", \"0\") == \"1\"\n\n\ndef get_rope_index(\n    processor,\n    input_ids: torch.Tensor,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    attention_mask: Optional[torch.Tensor] = None,\n) -> torch.Tensor:\n    \"\"\"\n    Gets the position ids for GLM4V in padding-free format.\n    The batch dim has been removed and the input_ids should be a 1D tensor representing a single example.\n    \"\"\"\n    spatial_merge_size = processor.image_processor.merge_size\n    image_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|image|>\")\n    video_start_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|begin_of_video|>\")\n    video_end_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|end_of_video|>\")\n\n    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):\n        if attention_mask is None:\n            attention_mask = torch.ones_like(input_ids)\n\n        position_ids = torch.ones(3, input_ids.size(0), dtype=input_ids.dtype, device=input_ids.device)  # (3, seqlen)\n        image_index, video_index = 0, 0\n        video_group_index = 0\n\n        input_ids_filtered = input_ids[attention_mask == 1]\n        input_tokens = input_ids_filtered.tolist()\n\n        input_token_type = []\n        video_check_flg = False\n        for token in input_tokens:\n            if token == video_start_token_id:\n                video_check_flg = True\n            elif token == video_end_token_id:\n                video_check_flg = False\n\n            if token == image_token_id and not video_check_flg:\n                input_token_type.append(\"image\")\n            elif token == image_token_id and video_check_flg:\n                input_token_type.append(\"video\")\n            else:\n                input_token_type.append(\"text\")\n\n        input_type_group = []\n        for key, group in itertools.groupby(enumerate(input_token_type), lambda x: x[1]):\n            group = list(group)\n            start_index = group[0][0]\n            end_index = group[-1][0] + 1\n            input_type_group.append((key, start_index, end_index))\n\n        llm_pos_ids_list = []\n        video_frame_num = 1\n\n        for modality_type, start_idx, end_idx in input_type_group:\n            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n\n            if modality_type == \"image\":\n                t, h, w = (\n                    image_grid_thw[image_index][0],\n                    image_grid_thw[image_index][1],\n                    image_grid_thw[image_index][2],\n                )\n                llm_grid_t, llm_grid_h, llm_grid_w = (\n                    t.item(),\n                    h.item() // spatial_merge_size,\n                    w.item() // spatial_merge_size,\n                )\n\n                t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()\n                h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()\n                w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()\n                llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)\n\n                image_index += 1\n                video_frame_num = 1\n\n            elif modality_type == \"video\":\n                t, h, w = (\n                    video_frame_num,\n                    video_grid_thw[video_index][1],\n                    video_grid_thw[video_index][2],\n                )\n\n                llm_grid_t, llm_grid_h, llm_grid_w = (\n                    t,\n                    h.item() // spatial_merge_size,\n                    w.item() // spatial_merge_size,\n                )\n\n                for t_idx in range(llm_grid_t):\n                    t_index = torch.tensor(t_idx).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()\n                    h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(1, -1, llm_grid_w).flatten()\n                    w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(1, llm_grid_h, -1).flatten()\n                    llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + st_idx)\n\n                video_group_index += 1\n\n                if video_group_index >= video_grid_thw[video_index][0]:\n                    video_index += 1\n                    video_group_index = 0\n\n                video_frame_num += 1\n\n            else:\n                text_len = end_idx - start_idx\n                llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n                video_frame_num = 1\n\n        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)\n        position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device)\n    else:\n        if attention_mask is not None:\n            position_ids = attention_mask.long().cumsum(-1) - 1\n            position_ids.masked_fill_(attention_mask == 0, 1)\n            position_ids = position_ids.unsqueeze(0).expand(3, -1).to(input_ids.device)\n        else:\n            position_ids = torch.arange(input_ids.shape[0], device=input_ids.device).view(1, -1).expand(3, -1)\n\n    return position_ids\n\n\ndef prepare_fa2_from_position_ids(\n    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor\n):\n    assert position_ids.ndim == 2  # (batch_size, seq_length)\n    query = query.contiguous().view(-1, query.size(-2), query.size(-1))\n    key = key.contiguous().view(-1, key.size(-2), key.size(-1))\n    value = value.contiguous().view(-1, value.size(-2), value.size(-1))\n    position_ids = position_ids.view(-1)\n    cu_seqlens = torch.cat(\n        (\n            (position_ids == 0).nonzero().view(-1).to(torch.int32),\n            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),\n        )\n    )\n    max_length = cu_seqlens.diff().max()  # use cu_seqlens to infer max_length for qwen2vl mrope\n    return (query, key, value, (cu_seqlens, cu_seqlens), (max_length, max_length))\n\n\ndef _custom_flash_attention_forward(\n    query_states: torch.Tensor,\n    key_states: torch.Tensor,\n    value_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor],\n    query_length: int,\n    is_causal: bool = True,\n    position_ids: Optional[torch.Tensor] = None,\n    use_top_left_mask: bool = False,\n    deterministic: Optional[bool] = None,\n    **kwargs,\n):\n    \"\"\"\n    Patches flash attention forward to handle 3D position ids in mrope. (3, batch_size, seq_length)\n    \"\"\"\n    # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).\n    flash_kwargs = {}\n\n    if _flash_supports_deterministic:\n        flash_kwargs[\"deterministic\"] = deterministic if deterministic is not None else _flash_deterministic_enabled\n\n    if kwargs.get(\"softcap\") is not None:\n        flash_kwargs[\"softcap\"] = kwargs.pop(\"softcap\")\n\n    query_states, key_states, value_states = fa_peft_integration_check(\n        query_states, key_states, value_states, target_dtype=torch.bfloat16\n    )\n\n    if position_ids is not None:\n        assert position_ids.ndim == 2  # (batch_size, seq_length / sp_size)\n\n    sp_size = get_ulysses_sequence_parallel_world_size()\n    if sp_size > 1:\n        # qkv: (batch_size, seq_length / sp_size, num_head, head_size)\n        validate_ulysses_config(query_states.size(2), sp_size)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)\n        position_ids_lst = [torch.empty_like(position_ids) for _ in range(sp_size)]\n        position_ids = dist.all_gather(position_ids_lst, position_ids, group=get_ulysses_sequence_parallel_group())\n        position_ids = torch.cat(position_ids_lst, dim=-1)  # (batch_size, seq_length)\n\n    if position_ids is not None and query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all():\n        batch_size = query_states.size(0)\n        q, k, v, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = prepare_fa2_from_position_ids(\n            query_states, key_states, value_states, position_ids\n        )\n        attn_output = flash_attn_varlen_func(\n            q=q,\n            k=k,\n            v=v,\n            cu_seqlens_q=cu_seqlens_q,\n            cu_seqlens_k=cu_seqlens_k,\n            max_seqlen_q=max_seqlen_q,\n            max_seqlen_k=max_seqlen_k,\n            dropout_p=kwargs.pop(\"dropout\", 0.0),\n            softmax_scale=kwargs.pop(\"softmax_scale\", None),\n            causal=is_causal,\n            **flash_kwargs,\n        )\n        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))\n    else:\n        attn_output = _flash_attention_forward(\n            query_states,\n            key_states,\n            value_states,\n            attention_mask,\n            query_length,\n            is_causal=is_causal,\n            use_top_left_mask=use_top_left_mask,\n            deterministic=deterministic,\n            **kwargs,\n        )  # do not pass position_ids to old flash_attention_forward\n\n    if sp_size > 1:\n        # (batch_size, seq_length, num_head, head_size)\n        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)\n\n    return attn_output\n\n\ndef glm4v_attn_forward(\n    self: \"Glm4vTextAttention\",\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46\n    **kwargs,\n) -> tuple[torch.Tensor, None, None]:\n    from transformers.models.glm4v.modeling_glm4v import apply_multimodal_rotary_pos_emb, repeat_kv\n\n    bsz, q_len, _ = hidden_states.size()  # q_len = seq_length / sp_size\n    query_states = self.q_proj(hidden_states)  # (batch_size, seq_length / sp_size, num_heads * head_size)\n    key_states = self.k_proj(hidden_states)\n    value_states = self.v_proj(hidden_states)\n\n    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n\n    # Because the input can be padded, the absolute sequence length depends on the max position id.\n    cos, sin = position_embeddings\n    query_states, key_states = apply_multimodal_rotary_pos_emb(\n        query_states, key_states, cos, sin, self.rope_scaling[\"mrope_section\"]\n    )\n    key_states = repeat_kv(key_states, self.num_key_value_groups)\n    value_states = repeat_kv(value_states, self.num_key_value_groups)\n    dropout_rate = 0.0 if not self.training else self.attention_dropout\n\n    # This is before the transpose\n    q_len = query_states.shape[2]\n\n    # FA2 uses non-transposed inputs\n    query_states = query_states.transpose(1, 2)\n    key_states = key_states.transpose(1, 2)\n    value_states = value_states.transpose(1, 2)\n\n    attn_output = _custom_flash_attention_forward(\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        query_length=q_len,\n        is_causal=getattr(self, \"is_causal\", True),\n        dropout=dropout_rate,\n        use_top_left_mask=_flash_use_top_left_mask,\n        position_ids=position_ids,  # important: pass position ids\n    )  # (batch_size, seq_length / sp_size, num_head, head_size)\n    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()\n    attn_output = self.o_proj(attn_output)\n    return attn_output, None\n\n\ndef _get_input_embeds(\n    model: \"Glm4vForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n):\n    inputs_embeds = model.get_input_embeddings()(input_ids)\n    if pixel_values is not None:\n        pixel_values = pixel_values.type(model.visual.dtype)\n        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        n_image_tokens = (input_ids == model.config.image_token_id).sum().item()\n        n_image_features = image_embeds.shape[0]\n        if n_image_tokens != n_image_features:\n            raise ValueError(\n                f\"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}\"\n            )\n\n        mask = input_ids == model.config.image_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        image_mask = mask_expanded.to(inputs_embeds.device)\n\n        image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)\n\n    if pixel_values_videos is not None:\n        pixel_values_videos = pixel_values_videos.type(model.visual.dtype)\n        video_embeds = model.visual(pixel_values_videos, grid_thw=video_grid_thw)\n        n_video_tokens = (input_ids == model.config.video_token_id).sum().item()\n        n_video_features = video_embeds.shape[0]\n        if n_video_tokens != n_video_features:\n            raise ValueError(\n                f\"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}\"\n            )\n\n        mask = input_ids == model.config.video_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        video_mask = mask_expanded.to(inputs_embeds.device)\n\n        video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)\n\n    if pixel_values is None and pixel_values_videos is None:  # handle mixed text-image data\n        pixel_values = torch.zeros((16, 1176), dtype=inputs_embeds.dtype, device=inputs_embeds.device)\n        image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long, device=inputs_embeds.device)\n        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        inputs_embeds += 0.0 * image_embeds.mean()\n\n    if attention_mask is not None:\n        attention_mask = attention_mask.to(inputs_embeds.device)\n\n    return inputs_embeds, attention_mask\n\n\ndef process_position_ids(position_ids: torch.Tensor) -> torch.Tensor:\n    if position_ids.ndim != 3 or position_ids.size(0) != 4:\n        # we concat the text position ids with the 3D vision position ids by default\n        # see https://github.com/huggingface/transformers/pull/39447\n        raise ValueError(\"position_ids should be a 3D tensor of shape (4, batch_size, seq_length).\")\n\n    return position_ids\n\n\n@dataclass\nclass Glm4vCausalLMOutputForPPO(Glm4vCausalLMOutputWithPast):\n    log_probs: Optional[torch.FloatTensor] = None\n    entropy: Optional[torch.FloatTensor] = None\n\n\ndef glm4v_base_forward(\n    self: \"Glm4vForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    labels: Optional[torch.LongTensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    **kwargs,\n):\n    kwargs[\"inputs_embeds\"], kwargs[\"attention_mask\"] = _get_input_embeds(\n        self, input_ids, attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw\n    )  # avoid lora module having multiple keyword arguments\n    return self.language_model(\n        input_ids=None,\n        **kwargs,\n    )\n\n\ndef glm4v_forward(\n    self: \"Glm4vForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    **kwargs,\n):\n    return self.model(\n        input_ids=input_ids,\n        attention_mask=attention_mask,\n        position_ids=process_position_ids(position_ids),\n        pixel_values=pixel_values,\n        pixel_values_videos=pixel_values_videos,\n        image_grid_thw=image_grid_thw,\n        video_grid_thw=video_grid_thw,\n        **kwargs,\n    )\n\n\ndef forward_with_normal_backend(\n    self: Glm4vForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> \"Glm4vCausalLMOutputWithPast\":\n    outputs = glm4v_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n    logits = self.lm_head(hidden_states)\n\n    return Glm4vCausalLMOutputWithPast(\n        logits=logits,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_torch_backend(\n    self: Glm4vForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> tuple | Glm4vCausalLMOutputForPPO:\n    from verl.utils.experimental.torch_functional import FusedLinearForPPO\n\n    outputs = glm4v_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_torch_backend, either labels or input_ids must be provided.\")\n\n    fused_linear_for_ppo = FusedLinearForPPO()\n    log_probs, entropy = fused_linear_for_ppo.forward(\n        hidden_states=hidden_states,\n        vocab_weights=self.lm_head.weight,\n        input_ids=rolled_labels,\n        temperature=temperature,\n    )\n    return Glm4vCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_triton_backend(\n    self: Glm4vForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> tuple | Glm4vCausalLMOutputForPPO:\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\n\n    outputs = glm4v_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_triton_backend, either labels or input_ids must be provided.\")\n\n    log_probs, entropy = linear_cross_entropy(\n        hidden_states,\n        self.lm_head.weight,\n        rolled_labels,\n        temperature,\n        \"none\",\n    )\n    return Glm4vCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n"
  },
  {
    "path": "verl/models/transformers/kimi_vl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\nimport torch.nn.functional as F\nfrom transformers.cache_utils import Cache\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward\n\nfrom verl.models.transformers.monkey_patch import is_transformers_version_in_range\n\n# Import compatibility wrapper for flash_attn_supports_top_left_mask\nfrom verl.utils.transformers_compat import flash_attn_supports_top_left_mask\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\n\n# Copied from transformers.models.llama.modeling_llama.rotate_half\ndef rotate_half(x):\n    \"\"\"Rotates half the hidden dims of the input.\"\"\"\n    x1 = x[..., : x.shape[-1] // 2]\n    x2 = x[..., x.shape[-1] // 2 :]\n    return torch.cat((-x2, x1), dim=-1)\n\n\n# Copied from transformers.models.llama.modeling_llama.apply_rotary_pos_emb\ndef apply_rotary_pos_emb(q, k, cos, sin, position_ids, unsqueeze_dim=1):\n    \"\"\"Applies Rotary Position Embedding to the query and key tensors.\n\n    Args:\n        q (`torch.Tensor`): The query tensor.\n        k (`torch.Tensor`): The key tensor.\n        cos (`torch.Tensor`): The cosine part of the rotary embedding.\n        sin (`torch.Tensor`): The sine part of the rotary embedding.\n        position_ids (`torch.Tensor`):\n            The position indices of the tokens corresponding to the query and key tensors. For example, this can be\n            used to pass offsetted position ids when working with a KV-cache.\n        unsqueeze_dim (`int`, *optional*, defaults to 1):\n            The 'unsqueeze_dim' argument specifies the dimension along which to unsqueeze cos[position_ids] and\n            sin[position_ids] so that they can be properly broadcasted to the dimensions of q and k. For example, note\n            that cos[position_ids] and sin[position_ids] have the shape [batch_size, seq_len, head_dim]. Then, if q and\n            k have the shape [batch_size, heads, seq_len, head_dim], then setting unsqueeze_dim=1 makes\n            cos[position_ids] and sin[position_ids] broadcastable to the shapes of q and k. Similarly, if q and k have\n            the shape [batch_size, seq_len, heads, head_dim], then set unsqueeze_dim=2.\n    Returns:\n        `tuple(torch.Tensor)` comprising of the query and key tensors rotated using the Rotary Position Embedding.\n    \"\"\"\n    cos = cos[position_ids].unsqueeze(unsqueeze_dim)\n    sin = sin[position_ids].unsqueeze(unsqueeze_dim)\n\n    b, h, s, d = q.shape\n    q = q.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)\n\n    b, h, s, d = k.shape\n    k = k.view(b, h, s, d // 2, 2).transpose(4, 3).reshape(b, h, s, d)\n\n    q_embed = (q * cos) + (rotate_half(q) * sin)\n    k_embed = (k * cos) + (rotate_half(k) * sin)\n    return q_embed, k_embed\n\n\n# Copied from transformers.models.llama.modeling_llama.repeat_kv\ndef repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:\n    \"\"\"\n    This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,\n    num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)\n    \"\"\"\n    batch, num_key_value_heads, slen, head_dim = hidden_states.shape\n    if n_rep == 1:\n        return hidden_states\n    hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)\n    return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)\n\n\ndef _ulysses_flash_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.LongTensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_value: Optional[Cache] = None,\n    output_attentions: bool = False,\n    use_cache: bool = False,\n    **kwargs,\n) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n    bsz, q_len, _ = hidden_states.size()\n\n    if self.q_lora_rank is None:\n        q = self.q_proj(hidden_states)\n    else:\n        q = self.q_b_proj(self.q_a_layernorm(self.q_a_proj(hidden_states)))\n    q = q.view(bsz, q_len, self.num_heads, self.q_head_dim).transpose(1, 2)\n\n    # Flash attention requires the input to have the shape\n    # batch_size x seq_length x head_dim x hidden_dim\n    # therefore we just need to keep the original shape\n    compressed_kv = self.kv_a_proj_with_mqa(hidden_states)\n    compressed_kv, k_pe = torch.split(compressed_kv, [self.kv_lora_rank, self.qk_rope_head_dim], dim=-1)\n    k_pe = k_pe.view(bsz, q_len, 1, self.qk_rope_head_dim).transpose(1, 2)\n    kv = (\n        self.kv_b_proj(self.kv_a_layernorm(compressed_kv))\n        .view(bsz, q_len, self.num_heads, self.qk_nope_head_dim + self.v_head_dim)\n        .transpose(1, 2)\n    )\n\n    k_nope, value_states = torch.split(kv, [self.qk_nope_head_dim, self.v_head_dim], dim=-1)\n\n    # patch\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.num_heads, ulysses_sp_size)\n\n        num_key_value_groups = self.config.num_attention_heads // self.config.num_key_value_heads\n        k_pe = repeat_kv(k_pe, ulysses_sp_size)  # to keep heads=1 after a2a\n        k_nope = repeat_kv(k_nope, num_key_value_groups)\n        value_states = repeat_kv(value_states, num_key_value_groups)\n        q = gather_seq_scatter_heads(q, seq_dim=2, head_dim=1)\n        k_pe = gather_seq_scatter_heads(k_pe, seq_dim=2, head_dim=1)\n        k_nope = gather_seq_scatter_heads(k_nope, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n        # (batch_size, num_head / sp_size, seq_length, head_size)\n        full_q_len = q.size(2)  # full_q_len = seq_length\n\n    else:\n        full_q_len = q_len\n\n    q_nope, q_pe = torch.split(q, [self.qk_nope_head_dim, self.qk_rope_head_dim], dim=-1)\n    cos, sin = self.rotary_emb(value_states, seq_len=full_q_len)\n    q_pe, k_pe = apply_rotary_pos_emb(q_pe, k_pe, cos, sin, position_ids)\n\n    query_states = k_pe.new_empty(bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim)\n    query_states[:, :, :, : self.qk_nope_head_dim] = q_nope\n    query_states[:, :, :, self.qk_nope_head_dim :] = q_pe\n\n    key_states = k_pe.new_empty(bsz, self.num_heads // ulysses_sp_size, full_q_len, self.q_head_dim)\n    key_states[:, :, :, : self.qk_nope_head_dim] = k_nope\n    key_states[:, :, :, self.qk_nope_head_dim :] = k_pe\n\n    if self.q_head_dim != self.v_head_dim:\n        value_states = F.pad(value_states, [0, self.q_head_dim - self.v_head_dim])\n\n    # TODO: These transpose are quite inefficient but Flash Attention requires the layout\n    # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache\n    # to be able to avoid many of these transpose/reshape/view.\n    query_states = query_states.transpose(1, 2)\n    key_states = key_states.transpose(1, 2)\n    value_states = value_states.transpose(1, 2)\n\n    dropout_rate = self.attention_dropout if self.training else 0.0\n\n    attn_output = _flash_attention_forward(\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        full_q_len,\n        dropout=dropout_rate,\n        sliding_window=None,\n        is_causal=self.is_causal,\n        use_top_left_mask=flash_attn_supports_top_left_mask(),\n        position_ids=position_ids,  # important: pass position ids\n        softmax_scale=self.softmax_scale,\n    )\n\n    if ulysses_sp_size > 1:\n        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)\n\n    if self.q_head_dim != self.v_head_dim:\n        attn_output = attn_output[:, :, :, : self.v_head_dim]\n\n    attn_output = attn_output.reshape(bsz, q_len, self.num_heads * self.v_head_dim).contiguous()\n    attn_output = self.o_proj(attn_output)\n\n    if is_transformers_version_in_range(min_version=\"4.53.0\"):\n        return attn_output, None\n    else:\n        return attn_output, None, None\n"
  },
  {
    "path": "verl/models/transformers/llama.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport sys\nfrom typing import Callable, Optional\n\nimport torch\n\nif sys.version_info >= (3, 11):\n    pass\nelse:\n    pass\n\nfrom transformers.cache_utils import Cache\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward\nfrom transformers.models.llama.modeling_llama import apply_rotary_pos_emb\nfrom transformers.utils import logging\n\n# Import compatibility wrapper for flash_attn_supports_top_left_mask\nfrom verl.utils.transformers_compat import flash_attn_supports_top_left_mask\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\nlogger = logging.get_logger(__name__)\n\n\ndef llama_flash_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.LongTensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_value: Optional[Cache] = None,\n    output_attentions: bool = False,\n    use_cache: bool = False,\n    cache_position: Optional[torch.LongTensor] = None,\n    position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46\n    **kwargs,\n) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n    \"\"\"\n    Adapted from transformers 4.47.1 to support Ulysses sequence parallelism.\n\n    NOTE: This function is used for transformers versions in the range [4.45.0, 4.47.1].\n    \"\"\"\n    output_attentions = False\n\n    bsz, q_len, _ = hidden_states.size()\n\n    query_states = self.q_proj(hidden_states)\n    key_states = self.k_proj(hidden_states)\n    value_states = self.v_proj(hidden_states)\n\n    # Flash attention requires the input to have the shape\n    # batch_size x seq_length x head_dim x hidden_dim\n    # therefore we just need to keep the original shape\n    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n\n    # trade off: repeat first and then all to all\n    # key_states = repeat_kv(key_states, self.num_key_value_groups)\n    # value_states = repeat_kv(value_states, self.num_key_value_groups)\n\n    ########## AlltoAll for Ulysses ##########\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.num_heads, ulysses_sp_size)\n\n        # (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n\n    full_q_len = query_states.size(2)  # full seq length\n\n    if position_embeddings is None:\n        logger.warning_once(\n            \"The attention layers in this model are transitioning from computing the RoPE embeddings internally \"\n            \"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed \"\n            \"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be \"\n            \"removed and `position_embeddings` will be mandatory.\"\n        )\n        cos, sin = self.rotary_emb(value_states, position_ids)\n    else:\n        cos, sin = position_embeddings\n    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n\n    if past_key_value is not None:\n        # sin and cos are specific to RoPE models; cache_position needed for the static cache\n        cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position}\n        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)\n\n    # TODO: These transpose are quite inefficient but Flash Attention requires the layout\n    # [batch_size, sequence_length, num_heads, head_dim]. We would need to refactor the KV cache\n    # to be able to avoid many of these transpose/reshape/view.\n    query_states = query_states.transpose(1, 2)\n    key_states = key_states.transpose(1, 2)\n    value_states = value_states.transpose(1, 2)\n\n    dropout_rate = self.attention_dropout if self.training else 0.0\n\n    # In PEFT, usually we cast the layer norms in float32 for training stability reasons\n    # therefore the input hidden states gets silently casted in float32. Hence, we need\n    # cast them back in the correct dtype just to be sure everything works as expected.\n    # This might slowdown training & inference so it is recommended to not cast the LayerNorms\n    # in fp32. (LlamaRMSNorm handles it correctly)\n\n    input_dtype = query_states.dtype\n    if input_dtype == torch.float32:\n        if torch.is_autocast_enabled():\n            target_dtype = torch.get_autocast_gpu_dtype()\n        # Handle the case where the model is quantized\n        elif hasattr(self.config, \"_pre_quantization_dtype\"):\n            target_dtype = self.config._pre_quantization_dtype\n        else:\n            target_dtype = self.q_proj.weight.dtype\n\n        logger.warning_once(\n            f\"The input hidden states seems to be silently casted in float32, this might be related to \"\n            f\"the fact you have upcasted embedding or layer norm layers in float32. We will cast back the \"\n            f\"input in {target_dtype}.\"\n        )\n\n        query_states = query_states.to(target_dtype)\n        key_states = key_states.to(target_dtype)\n        value_states = value_states.to(target_dtype)\n\n    attn_output = _flash_attention_forward(\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        full_q_len,\n        position_ids=position_ids,\n        dropout=dropout_rate,\n        sliding_window=getattr(self, \"sliding_window\", None),\n        use_top_left_mask=flash_attn_supports_top_left_mask(),\n        is_causal=self.is_causal,\n        **kwargs,\n    )\n\n    attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1:\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n    attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()\n    attn_output = self.o_proj(attn_output)\n\n    if not output_attentions:\n        attn_weights = None\n\n    return attn_output, attn_weights, past_key_value\n\n\ndef llama_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    position_embeddings: tuple[torch.Tensor, torch.Tensor],\n    attention_mask: Optional[torch.Tensor],\n    past_key_value: Optional[Cache] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n    \"\"\"\n    Adapted from transformers 4.49.0 to support Ulysses sequence parallelism for transformers >= 4.48.0.\n\n    NOTE: This function has been tested only on transformers versions between 4.48.0 and 4.50.0.\n    \"\"\"\n    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS\n    from transformers.models.llama.modeling_llama import eager_attention_forward\n\n    bsz, q_len, _ = hidden_states.shape\n\n    query_states = self.q_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    key_states = self.k_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    value_states = self.v_proj(hidden_states).view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n\n    ########## AlltoAll for Ulysses ##########\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.config.num_attention_heads, ulysses_sp_size)\n\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n\n    full_q_len = query_states.size(2)\n\n    cos, sin = position_embeddings\n    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n\n    if past_key_value is not None:\n        # sin and cos are specific to RoPE models; cache_position needed for the static cache\n        cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position}\n        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)\n\n    attention_interface: Callable = eager_attention_forward\n    if self.config._attn_implementation != \"eager\":\n        if self.config._attn_implementation == \"sdpa\" and kwargs.get(\"output_attentions\", False):\n            logger.warning_once(\n                \"`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. \"\n                \"Falling back to eager attention. This warning can be removed using the argument \"\n                '`attn_implementation=\"eager\"` when loading the model.'\n            )\n        else:\n            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]\n\n    attn_output, attn_weights = attention_interface(\n        self,\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        dropout=0.0 if not self.training else self.attention_dropout,\n        scaling=self.scaling,\n        **kwargs,\n    )\n\n    attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1:\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n    attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()\n    attn_output = self.o_proj(attn_output)\n    return attn_output, attn_weights\n"
  },
  {
    "path": "verl/models/transformers/monkey_patch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nApply monkey-patch function to models\n\"\"\"\n\nimport sys\nfrom types import SimpleNamespace\nfrom typing import Optional\n\nimport torch\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward\nfrom transformers.modeling_utils import PreTrainedModel\n\nfrom verl.utils.import_utils import is_trl_available\nfrom verl.utils.transformers_compat import is_transformers_version_in_range\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_group,\n    get_ulysses_sequence_parallel_world_size,\n    slice_input_tensor,\n)\n\n_PREFIX_GROUPER_PATCHED = False\n_PREFIX_GROUPER_SUPPORTED_ATTENTIONS = {\"flash_attention_2\", \"flash_attention_3\", \"sdpa\", \"flex_attention\", \"eager\"}\n\n\ndef _create_prefix_grouper_wrapper(original_fn):\n    \"\"\"Wrap attention function to support prefix_grouper in kwargs.\"\"\"\n\n    def wrapped(module, query, key, value, attention_mask, *args, **kwargs):\n        prefix_grouper = kwargs.pop(\"prefix_grouper\", None)\n        if prefix_grouper is None:\n            return original_fn(module, query, key, value, attention_mask, *args, **kwargs)\n\n        def attn_func(q, k, v, attn_mask, *inner_args, **inner_kwargs):\n            out, _ = original_fn(module, q, k, v, attn_mask, *inner_args, **inner_kwargs)\n            return out\n\n        return prefix_grouper.forward(attn_func, query, key, value, *args, **kwargs), None\n\n    return wrapped\n\n\ndef apply_prefix_grouper_patch():\n    \"\"\"Patch ALL_ATTENTION_FUNCTIONS to support prefix_grouper parameter.\"\"\"\n    global _PREFIX_GROUPER_PATCHED\n    if _PREFIX_GROUPER_PATCHED:\n        return\n\n    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS\n\n    patched = []\n    for name in list(ALL_ATTENTION_FUNCTIONS.keys()):\n        if name in _PREFIX_GROUPER_SUPPORTED_ATTENTIONS:\n            ALL_ATTENTION_FUNCTIONS[name] = _create_prefix_grouper_wrapper(ALL_ATTENTION_FUNCTIONS[name])\n            patched.append(name)\n\n    _PREFIX_GROUPER_PATCHED = True\n    print(f\"[PrefixGrouper] Patched: {patched}\")\n\n\ndef repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:\n    \"\"\"\n    This is the equivalent of torch.repeat_interleave(x, dim=2, repeats=n_rep). The hidden states go from (batch,\n    seqlen, num_key_value_heads, head_dim) to (batch, seqlen, num_attention_heads, head_dim)\n    \"\"\"\n    batch, slen, num_key_value_heads, head_dim = hidden_states.shape\n    if n_rep == 1:\n        return hidden_states\n    hidden_states = hidden_states[:, :, :, None, :].expand(batch, slen, num_key_value_heads, n_rep, head_dim)\n    return hidden_states.reshape(batch, slen, num_key_value_heads * n_rep, head_dim)\n\n\ndef _ulysses_flash_attention_forward(\n    query_states: torch.Tensor,\n    key_states: torch.Tensor,\n    value_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor],\n    query_length: int,\n    *args,\n    position_ids: Optional[torch.Tensor] = None,\n    **kwargs,\n):\n    \"\"\"Insert all-to-all before and after flash attention.\n    DeepSpeed-Ulysses: https://arxiv.org/pdf/2309.14509\n\n    For transformers>=4.55, the flash attention api has changed,\n    we need to pass the query_length after doing ulysses all2all.\n    See https://github.com/huggingface/transformers/issues/40399\n\n    Args:\n        query_states (torch.Tensor): (batch_size, seqlen/sp_size, nheads, head_dim)\n        key_states (torch.Tensor): (batch_size, seqlen/sp_size, nheads_k, head_dim)\n        value_states (torch.Tensor): (batch_size, seqlen/sp_size, nheads_k, head_dim)\n        position_ids (torch.Tensor, optional): (batch_size, seqlen/sp_size)\n\n    Returns:\n        torch.Tensor: (batch_size, seqlen/sp_size, nheads, head_dim)\n\n    \"\"\"\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    ########## AlltoAll for Ulysses ##########\n    # TODO: Disable sp for ViT, there's no elegent way to determine whether it's ViT or not.\n    # Use `position_ids` as condition since ViT doesn't pass it to flash attention.\n    if ulysses_sp_size > 1 and position_ids is not None:\n        # NOTE: repeat kv heads to be divided by sequence parallel. Instead of repeating nheads_q//nheads_k,\n        # we choose to repeat sp_size//nheads_k, since flash_attention supports MQA/GQA.\n        # For example:\n        # - nheads_k=4, sp=8, repeats=2\n        # - nheads_k=8, sp=8, repeats=1\n        # - nheads_k=16, sp=8, repeats=1\n        repeats = max(ulysses_sp_size // key_states.size(2), 1)\n        key_states = repeat_kv(key_states, repeats)\n        value_states = repeat_kv(value_states, repeats)\n\n        # (bsz, seq_len/n, n_head, head_dim) -> (bsz, seq_len, n_head/n, head_dim)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)\n\n        # TODO: all_gather position_ids because `prepare_fa2_from_position_ids` needs it, we can eliminate\n        # this all_gather by passing cu_seq_lens_q, cu_seq_lens_k, max_length_k, max_length_q explicitly.\n        # https://github.com/huggingface/transformers/pull/33932\n\n        # (bsz, seq_len/n) -> (bsz, seq_len)\n        position_ids_list = [torch.empty_like(position_ids) for _ in range(ulysses_sp_size)]\n        torch.distributed.all_gather(position_ids_list, position_ids, group=get_ulysses_sequence_parallel_group())\n        position_ids = torch.concat(position_ids_list, dim=-1)\n\n    # (bsz, seq_len, n_head/n, head_dim)\n    query_length = query_states.size(1)\n    attn_output = _flash_attention_forward(\n        query_states, key_states, value_states, attention_mask, query_length, *args, position_ids=position_ids, **kwargs\n    )\n\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1 and position_ids is not None:\n        # (bsz, seq_len, n_head/n, head_dim) -> (bsz, seq_len/n, n_head, head_dim)\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n\n    return attn_output\n\n\ndef patch_vlm_for_ulysses_input_slicing(model_class: type):\n    \"\"\"\n    Applies a monkey patch to the forward method of a given model class\n    to enable Ulysses sequence parallelism input slicing.\n    \"\"\"\n\n    def _create_ulysses_wrapped_decoder_forward(original_forward):\n        def ulysses_wrapped_decoder_forward(self, *args, **kwargs):\n            inputs_embeds = kwargs.get(\"inputs_embeds\")\n            position_ids = kwargs.get(\"position_ids\")\n            visual_pos_masks = kwargs.get(\"visual_pos_masks\")\n            deepstack_visual_embeds = kwargs.get(\"deepstack_visual_embeds\")\n            call_kwargs = kwargs.copy()\n\n            current_ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n            slice_now = (\n                inputs_embeds is not None\n                and current_ulysses_sp_size > 1\n                and getattr(self, \"_needs_initial_slice\", True)\n            )\n            if slice_now:\n                call_kwargs[\"inputs_embeds\"] = slice_input_tensor(inputs_embeds, dim=1, padding=False)\n                call_kwargs[\"position_ids\"] = slice_input_tensor(position_ids, dim=-1, padding=False)\n                # Also slice visual_pos_masks and deepstack_visual_embeds for Qwen3 VL models\n                if visual_pos_masks is not None:\n                    original_visual_mask = visual_pos_masks\n                    sliced_visual_mask = slice_input_tensor(visual_pos_masks, dim=1, padding=False)\n                    call_kwargs[\"visual_pos_masks\"] = sliced_visual_mask\n\n                    if deepstack_visual_embeds is not None:\n                        sliced_embeds = []\n\n                        num_visual_before = original_visual_mask.sum().item()\n                        num_visual_in_shard = sliced_visual_mask.sum().item()\n\n                        if num_visual_in_shard > 0 and num_visual_before > 0:\n                            # Calculate which visual embeddings belong to this shard\n                            # We need to find the offset of visual tokens in this shard\n                            from verl.utils.ulysses import get_ulysses_sequence_parallel_rank\n\n                            rank = get_ulysses_sequence_parallel_rank()\n                            seq_len = original_visual_mask.shape[1]\n                            local_seq_len = seq_len // current_ulysses_sp_size\n                            start_idx = rank * local_seq_len\n                            end_idx = start_idx + local_seq_len\n\n                            # Get total visual tokens before and up to the end of the shard's sequence slice\n                            # This correctly handles batches by summing across all samples\n                            visual_start = original_visual_mask[:, :start_idx].sum().item() if start_idx > 0 else 0\n                            visual_end = original_visual_mask[:, :end_idx].sum().item()\n\n                            # Slice each tensor in deepstack_visual_embeds\n                            for embed in deepstack_visual_embeds:\n                                sliced_embeds.append(embed[visual_start:visual_end])\n                        else:\n                            # No visual tokens in this shard, create empty tensors to maintain gradient flow\n                            for embed in deepstack_visual_embeds:\n                                sliced_embeds.append(embed[:0])\n                        call_kwargs[\"deepstack_visual_embeds\"] = sliced_embeds\n\n                self._needs_initial_slice = False\n            try:\n                return original_forward(self, *args, **call_kwargs)\n            finally:\n                if slice_now:\n                    self._needs_initial_slice = True\n\n        return ulysses_wrapped_decoder_forward\n\n    original_forward = model_class.forward\n    wrapped_forward = _create_ulysses_wrapped_decoder_forward(original_forward)\n    model_class.forward = wrapped_forward\n    print(f\"Monkey patch {model_class.__name__}.forward for Ulysses SP input slicing.\")\n\n\ndef patch_forward_with_backends(\n    model: PreTrainedModel,\n    use_fused_kernels: bool = False,\n    fused_kernels_backend: str = None,\n):\n    \"\"\"\n    Choose the forward function based on the model and backend.\n    Args:\n        model (PreTrainedModel): The model to apply the monkey patch.\n        use_fused_kernels (bool): Whether to use fused kernels.\n        fused_kernels_backend (str): The backend to use for fused kernels.\n    \"\"\"\n    if not use_fused_kernels or fused_kernels_backend not in [\"triton\", \"torch\"]:\n        print(\n            f\"Skipping monkey patch for {model.__class__.__name__} as use_fused_kernels is \"\n            f\"{use_fused_kernels} or fused_kernels_backend is {fused_kernels_backend}\"\n        )\n        return\n\n    forward_with_torch_backend_function = model.__class__.forward\n    forward_with_triton_backend_function = model.__class__.forward\n    if model.config.model_type in [\"qwen2_5_vl\", \"qwen2_vl\"]:\n        from verl.models.transformers.qwen2_vl import forward_with_torch_backend, forward_with_triton_backend\n\n        forward_with_torch_backend_function = forward_with_torch_backend\n        forward_with_triton_backend_function = forward_with_triton_backend\n    elif model.config.model_type in [\"qwen3_vl\", \"qwen3_vl_moe\"]:\n        from verl.models.transformers.qwen3_vl import forward_with_torch_backend, forward_with_triton_backend\n\n        forward_with_torch_backend_function = forward_with_torch_backend\n        forward_with_triton_backend_function = forward_with_triton_backend\n    elif model.config.model_type == \"glm4v\":\n        from verl.models.transformers.glm4v import forward_with_torch_backend, forward_with_triton_backend\n\n        forward_with_torch_backend_function = forward_with_torch_backend\n        forward_with_triton_backend_function = forward_with_triton_backend\n    else:\n        from verl.models.transformers.dense_common import forward_with_torch_backend, forward_with_triton_backend\n\n        forward_with_torch_backend_function = forward_with_torch_backend\n        forward_with_triton_backend_function = forward_with_triton_backend\n\n    if fused_kernels_backend == \"triton\":\n        model.__class__.forward = forward_with_triton_backend_function\n        print(f\"Using Triton backend for fused kernels in {model.__class__.__name__}\")\n    elif fused_kernels_backend == \"torch\":\n        model.__class__.forward = forward_with_torch_backend_function\n        print(f\"Using Torch backend for fused kernels in {model.__class__.__name__}\")\n    else:\n        raise ValueError(f\"Unsupported fused_kernels_backend: {fused_kernels_backend}. Choose 'triton' or 'torch'.\")\n\n\ndef apply_monkey_patch(\n    model: PreTrainedModel,\n    ulysses_sp_size: int = 1,\n    use_remove_padding: bool = True,\n    use_fused_kernels: bool = False,\n    fused_kernels_backend: str = None,\n    use_prefix_grouper: bool = False,\n    use_tiled_mlp: bool = False,\n    tiled_mlp_shards: int = 4,\n):\n    \"\"\"\n    Apply monkey patch to the models for ulysses sequence parallel, fused kernel, tiled MLP and prefix grouper.\n\n    In the end of this function forward function of the model is patched for fused kernel.\n    If the model is not supported with fused kernel, please return after patch.\n\n    Args:\n        model: The model to apply the monkey patch.\n        ulysses_sp_size: The size of ulysses sequence parallel.\n        use_remove_padding: Whether to use remove padding.\n        use_fused_kernels: Whether to use fused kernels.\n        fused_kernels_backend: The backend to use for fused kernels.\n        use_tiled_mlp: Whether to use TiledMLP for memory-efficient MLP computation.\n        tiled_mlp_shards: Number of shards for TiledMLP (higher = lower memory, slightly slower).\n    \"\"\"\n\n    # Apply TiledMLP monkey patch for memory-efficient MLP computation\n    if use_tiled_mlp:\n        from verl.models.transformers.tiled_mlp import apply_tiled_mlp_monkey_patch\n\n        model_type = getattr(model.config, \"model_type\", None)\n        apply_tiled_mlp_monkey_patch(num_shards=tiled_mlp_shards, model_type=model_type)\n    # Apply PrefixGrouper patch if enabled\n    if use_prefix_grouper:\n        apply_prefix_grouper_patch()\n\n    \"\"\"Replace _flash_attention_forward to _ulysses_flash_attention_forward\"\"\"\n    module = sys.modules[model.__module__]\n\n    try:\n        num_attention_heads, num_key_value_heads = model.config.num_attention_heads, model.config.num_key_value_heads\n    except AttributeError:\n        num_attention_heads, num_key_value_heads = (\n            model.config.text_config.num_attention_heads,\n            model.config.text_config.num_key_value_heads,\n        )\n\n    assert num_attention_heads % ulysses_sp_size == 0, (\n        f\"num_attention_heads {num_attention_heads} must be divisible by ulysses_sp_size {ulysses_sp_size}\"\n    )\n    assert num_key_value_heads % ulysses_sp_size == 0 or ulysses_sp_size % num_key_value_heads == 0, (\n        f\"num_key_value_heads {num_key_value_heads} must be divisible by ulysses_sp_size \"\n        f\"{ulysses_sp_size}or vise versa. Upon ulysses_sp_size % num_key_value_heads == 0,\"\n        f\"kv heads are repeated to ensure correctness.\"\n    )\n\n    if is_trl_available():\n        from trl import AutoModelForCausalLMWithValueHead  # type: ignore\n\n        def state_dict(self, *args, **kwargs):\n            return torch.nn.Module.state_dict(self, *args, **kwargs)\n\n        AutoModelForCausalLMWithValueHead.state_dict = state_dict\n        print(\"Monkey patch state_dict in AutoModelForCausalLMWithValueHead. \")\n\n    # TODO: VLM models only, unify monkey patch to LLM models.\n    if model.config.model_type in [\"qwen2_5_vl\", \"qwen2_vl\"]:\n        # Step 1: patch model to support image-text mixed data\n        if is_transformers_version_in_range(min_version=\"4.52.0\"):\n            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (\n                Qwen2_5_VLForConditionalGeneration,\n                Qwen2_5_VLModel,\n                Qwen2_5_VLTextModel,\n            )\n            from transformers.models.qwen2_vl.modeling_qwen2_vl import (\n                Qwen2VLForConditionalGeneration,\n                Qwen2VLModel,\n                Qwen2VLTextModel,\n            )\n        else:\n            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLForConditionalGeneration\n            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLModel as Qwen2_5_VLTextModel\n            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLForConditionalGeneration\n            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLModel as Qwen2VLTextModel\n\n            Qwen2_5_VLModel = SimpleNamespace(forward=None)\n            Qwen2VLModel = SimpleNamespace(forward=None)\n\n        from verl.models.transformers.qwen2_vl import forward_with_normal_backend, qwen2_vl_base_forward\n\n        Qwen2_5_VLModel.forward = qwen2_vl_base_forward\n        Qwen2VLModel.forward = qwen2_vl_base_forward\n        Qwen2_5_VLForConditionalGeneration.forward = forward_with_normal_backend\n        Qwen2VLForConditionalGeneration.forward = forward_with_normal_backend\n        print(f\"Monkey patch {model.__class__.__name__} model forward\")\n\n        # Step 2: patch attention to support ulysses parallelism\n        if is_transformers_version_in_range(min_version=\"4.54.0\"):\n            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import Qwen2_5_VLAttention\n            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLAttention\n        elif is_transformers_version_in_range(min_version=\"4.53.0\"):\n            raise RuntimeError(\"Transformers 4.53.* is bugged. Use transformers 4.54.0 or later.\")\n        else:\n            from transformers.models.qwen2_5_vl.modeling_qwen2_5_vl import (\n                Qwen2_5_VLFlashAttention2 as Qwen2_5_VLAttention,\n            )\n            from transformers.models.qwen2_vl.modeling_qwen2_vl import Qwen2VLFlashAttention2 as Qwen2VLAttention\n\n        if use_remove_padding or ulysses_sp_size > 1:\n            from verl.models.transformers.qwen2_vl import qwen2_vl_attn_forward\n\n            Qwen2_5_VLAttention.forward = qwen2_vl_attn_forward\n            Qwen2VLAttention.forward = qwen2_vl_attn_forward\n            print(f\"Monkey patch {model.__class__.__name__} attention layer\")\n\n        # Step 3: patch input for multimodal sequence parallelism\n        if ulysses_sp_size > 1:\n            patch_vlm_for_ulysses_input_slicing(Qwen2_5_VLTextModel)\n            patch_vlm_for_ulysses_input_slicing(Qwen2VLTextModel)\n\n    elif model.config.model_type in [\"qwen3_vl\", \"qwen3_vl_moe\"]:\n        # Step 1: patch model to support image-text mixed data\n        from transformers.models.qwen3_vl.modeling_qwen3_vl import (\n            Qwen3VLForConditionalGeneration,\n            Qwen3VLModel,\n            Qwen3VLTextModel,\n        )\n        from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import (\n            Qwen3VLMoeForConditionalGeneration,\n            Qwen3VLMoeModel,\n            Qwen3VLMoeTextModel,\n        )\n\n        from verl.models.transformers.qwen3_vl import (\n            forward_with_normal_backend,\n            patch_qwen3_vl_moe_sparse_moe_block_forward,\n            qwen3_vl_base_forward,\n        )\n\n        Qwen3VLModel.forward = qwen3_vl_base_forward\n        Qwen3VLMoeModel.forward = qwen3_vl_base_forward\n        Qwen3VLForConditionalGeneration.forward = forward_with_normal_backend\n        Qwen3VLMoeForConditionalGeneration.forward = forward_with_normal_backend\n        print(f\"Monkey patch {model.__class__.__name__} model forward\")\n\n        # Step 1.5: patch Qwen3VLMoeTextSparseMoeBlock to fix transformers 4.57.3 bug\n        if model.config.model_type == \"qwen3_vl_moe\" and is_transformers_version_in_range(max_version=\"4.57.3\"):\n            patch_qwen3_vl_moe_sparse_moe_block_forward()\n\n        # Step 2: patch input for multimodal sequence parallelism\n        if ulysses_sp_size > 1:\n            patch_vlm_for_ulysses_input_slicing(Qwen3VLTextModel)\n            patch_vlm_for_ulysses_input_slicing(Qwen3VLMoeTextModel)\n\n    elif model.config.model_type == \"glm4v\":\n        # Step 1: patch model to support image-text mixed data\n\n        from transformers.models.glm4v.modeling_glm4v import (\n            Glm4vForConditionalGeneration,\n            Glm4vModel,\n            Glm4vTextAttention,\n            Glm4vTextModel,\n        )\n\n        from verl.models.transformers.glm4v import forward_with_normal_backend, glm4v_base_forward\n\n        Glm4vModel.forward = glm4v_base_forward\n        Glm4vForConditionalGeneration.forward = forward_with_normal_backend\n        print(f\"Monkey patch {model.__class__.__name__} model forward\")\n\n        # Step 2: patch attention to support ulysses parallelism\n        if use_remove_padding or ulysses_sp_size > 1:\n            from verl.models.transformers.glm4v import glm4v_attn_forward\n\n            Glm4vTextAttention.forward = glm4v_attn_forward\n            print(f\"Monkey patch {model.__class__.__name__} attention layer\")\n\n        # Step 3: patch input for multimodal sequence parallelism\n        if ulysses_sp_size > 1:\n            patch_vlm_for_ulysses_input_slicing(Glm4vTextModel)\n\n    elif model.config.model_type == \"kimi_vl\":\n        if use_remove_padding or ulysses_sp_size > 1:\n            # TODO: Changes need to be made when transformers are adapted.\n            from verl.models.transformers.kimi_vl import _ulysses_flash_attn_forward\n\n            module.DeepseekV3FlashAttention2.forward = _ulysses_flash_attn_forward\n            print(\"Monkey patch FlashAttention2.forward in KimiVL\")\n\n        if ulysses_sp_size > 1:\n            patch_vlm_for_ulysses_input_slicing(module.DeepseekV3ForCausalLM)\n\n        if use_fused_kernels:\n            print(\"Not support fused kernels for KimiVL\")\n\n        return\n\n    if use_remove_padding or ulysses_sp_size > 1:\n        if hasattr(module, \"_flash_attention_forward\"):  # transformers <= 4.47.1 or legacy models\n            module._flash_attention_forward = _ulysses_flash_attention_forward\n            print(f\"Monkey patch _flash_attention_forward in {model.__module__}\")\n        else:\n            from transformers.integrations import flash_attention\n\n            flash_attention._flash_attention_forward = _ulysses_flash_attention_forward\n            print(f\"Monkey patch _flash_attention_forward in {flash_attention.__name__}\")\n\n    patch_forward_with_backends(model, use_fused_kernels=use_fused_kernels, fused_kernels_backend=fused_kernels_backend)\n"
  },
  {
    "path": "verl/models/transformers/npu_patch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\r\n#\r\n# Copyright 2025 The Qwen Team and The HuggingFace Inc. team\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\n\r\nimport torch\r\nimport torch.nn.functional as F\r\nimport torch_npu\r\nfrom torch import nn\r\nfrom transformers.activations import ACT2FN\r\nfrom transformers.models.qwen2 import modeling_qwen2\r\nfrom transformers.models.qwen2_5_vl import modeling_qwen2_5_vl\r\nfrom transformers.models.qwen3 import modeling_qwen3\r\nfrom transformers.models.qwen3_moe import modeling_qwen3_moe\r\nfrom transformers.models.qwen3_next import modeling_qwen3_next\r\nfrom transformers.models.qwen3_vl import modeling_qwen3_vl\r\nfrom transformers.models.qwen3_vl_moe import modeling_qwen3_vl_moe\r\nfrom transformers.utils import logging\r\n\r\nlogger = logging.get_logger(__name__)\r\n\r\n\r\ndef rms_norm_forward_npu(self, x):\r\n    \"\"\"NPU optimized implementation for RMSNorm.\"\"\"\r\n    if x.dtype != self.weight.dtype:\r\n        x = x.to(self.weight.dtype)\r\n    return torch_npu.npu_rms_norm(x, self.weight, epsilon=self.variance_epsilon)[0]\r\n\r\n\r\ndef silu_forward_npu(self, hidden_state):\r\n    \"\"\"NPU optimized implementation for SiLU in `forward` func in MLP layer.\"\"\"\r\n    gate_up = torch.cat((self.gate_proj(hidden_state), self.up_proj(hidden_state)), dim=-1)\r\n    return self.down_proj(torch_npu.npu_swiglu(gate_up, dim=-1))\r\n\r\n\r\ndef apply_rotary_pos_emb_npu(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):\r\n    \"\"\"NPU optimized implementation for RoPE.\"\"\"\r\n    cos = cos.unsqueeze(unsqueeze_dim)\r\n    sin = sin.unsqueeze(unsqueeze_dim)\r\n    q_embed = torch_npu.npu_rotary_mul(q, cos, sin)\r\n    k_embed = torch_npu.npu_rotary_mul(k, cos, sin)\r\n    return q_embed.to(q.dtype), k_embed.to(k.dtype)\r\n\r\n\r\ndef qwen3_next_rms_norm_forward_npu(self, x):\r\n    return torch_npu.npu_rms_norm(x.float(), 1.0 + self.weight.float(), epsilon=self.eps)[0].type_as(x)\r\n\r\n\r\ndef qwen3_next_rms_norm_forward_gated_npu(self, hidden_states, gate=None):\r\n    input_dtype = hidden_states.dtype\r\n    hidden_states = hidden_states.to(torch.float32)\r\n    hidden_states = torch_npu.npu_rms_norm(hidden_states, self.weight.float(), epsilon=self.variance_epsilon)[0]\r\n    hidden_states = hidden_states * F.silu(gate.to(torch.float32))\r\n    return hidden_states.to(input_dtype)\r\n\r\n\r\ndef qwen3_next_apply_rotary_pos_emb_npu(q, k, cos, sin, position_ids=None, unsqueeze_dim=1):\r\n    cos = cos.unsqueeze(unsqueeze_dim)\r\n    sin = sin.unsqueeze(unsqueeze_dim)\r\n\r\n    # Keep half or full tensor for later concatenation\r\n    rotary_dim = cos.shape[-1]\r\n    q_rot, q_pass = q[..., :rotary_dim], q[..., rotary_dim:]\r\n    k_rot, k_pass = k[..., :rotary_dim], k[..., rotary_dim:]\r\n\r\n    q_embed = torch_npu.npu_rotary_mul(q_rot, cos, sin).to(q.dtype)\r\n    k_embed = torch_npu.npu_rotary_mul(k_rot, cos, sin).to(k.dtype)\r\n    q_embed = torch.cat([q_embed, q_pass], dim=-1)\r\n    k_embed = torch.cat([k_embed, k_pass], dim=-1)\r\n    return q_embed, k_embed\r\n\r\n\r\nclass NPUGmmFunction(torch.autograd.Function):\r\n    @staticmethod\r\n    def forward(ctx, x, weight, group_list, group_list_type=1):\r\n        \"\"\"\r\n        Grouped Matmul(GMM) for Ascend NPU.\r\n\r\n        Args:\r\n            x (torch.Tensor): Input tensor, shape (tokens_num * top_k, hidden_size)\r\n            weight (torch.Tensor): Expert weights, shape (n_experts, hidden_size, intermediate_size)\r\n            group_list (torch.Tensor): Expert token counts, shape (n_experts,)\r\n                - type 0: cumsum of tokens per expert\r\n                - type 1: direct tokens per expert (default)\r\n        \"\"\"\r\n        ctx.save_for_backward(x, weight)\r\n        ctx.group_list = group_list\r\n        ctx.group_list_type = group_list_type\r\n\r\n        output = torch_npu.npu_grouped_matmul(\r\n            [x], [weight], bias=None, group_list=group_list, split_item=2, group_type=0, group_list_type=group_list_type\r\n        )[0]\r\n\r\n        return output\r\n\r\n    @staticmethod\r\n    def backward(ctx, grad_output):\r\n        x, weight = ctx.saved_tensors\r\n        group_list = ctx.group_list\r\n        group_list_type = ctx.group_list_type\r\n\r\n        dx = torch_npu.npu_grouped_matmul(\r\n            [grad_output],\r\n            [weight.transpose(1, 2)],\r\n            bias=None,\r\n            group_list=group_list,\r\n            split_item=2,\r\n            group_type=0,\r\n            group_list_type=group_list_type,\r\n        )[0]\r\n\r\n        dw = torch_npu.npu_grouped_matmul(\r\n            [x.transpose(0, 1)],\r\n            [grad_output],\r\n            bias=None,\r\n            group_list=group_list,\r\n            split_item=3,\r\n            group_type=2,\r\n            group_list_type=group_list_type,\r\n        )[0]\r\n\r\n        return dx, dw, None, None\r\n\r\n\r\ndef _qwen3_sparse_moe_routed_forward_npu(self, hidden_states: torch.Tensor):\r\n    \"\"\"\r\n    Shared NPU routed-expert path for Qwen3Moe/Qwen3Next sparse MoE blocks.\r\n\r\n    Returns:\r\n        tuple: (flattened_input, routed_hidden_states, router_logits)\r\n    \"\"\"\r\n    hidden_dim = hidden_states.shape[-1]\r\n    hidden_states = hidden_states.view(-1, hidden_dim)\r\n    # router_logits: (batch * sequence_length, n_experts)\r\n    router_logits = self.gate(hidden_states)\r\n\r\n    routing_weights = F.softmax(router_logits, dim=1, dtype=torch.float)\r\n    routing_weights, selected_experts = torch.topk(routing_weights, self.top_k, dim=-1)\r\n    if self.norm_topk_prob:  # only diff with mixtral sparse moe block!\r\n        routing_weights /= routing_weights.sum(dim=-1, keepdim=True)\r\n    # we cast back to the input dtype\r\n    routing_weights = routing_weights.to(hidden_states.dtype)\r\n\r\n    # Loop over all available experts in the model and perform the computation on each expert\r\n    # Concat all weights\r\n    input_dtype = hidden_states.dtype\r\n    up_weight_list = [e.up_proj.weight for e in self.experts]\r\n    gate_weight_list = [e.gate_proj.weight for e in self.experts]\r\n    down_weight_list = [e.down_proj.weight for e in self.experts]\r\n    w1 = torch.stack(up_weight_list).transpose(1, 2).to(input_dtype)\r\n    w2 = torch.stack(gate_weight_list).transpose(1, 2).to(input_dtype)\r\n    w3 = torch.stack(down_weight_list).transpose(1, 2).to(input_dtype)\r\n\r\n    permuted_tokens, row_ids_map = torch_npu.npu_moe_token_permute(hidden_states, selected_experts.to(torch.int32))\r\n    tokens_per_expert = torch.histc(selected_experts, bins=self.num_experts, min=0, max=self.num_experts)\r\n\r\n    up_res = NPUGmmFunction.apply(permuted_tokens, w1, tokens_per_expert)\r\n    gate_res = NPUGmmFunction.apply(permuted_tokens, w2, tokens_per_expert)\r\n    act_res = torch_npu.npu_swiglu(torch.cat([gate_res, up_res], dim=-1))\r\n    down_res = NPUGmmFunction.apply(act_res, w3, tokens_per_expert)\r\n\r\n    routed_hidden_states = torch_npu.npu_moe_token_unpermute(down_res, row_ids_map, probs=routing_weights)\r\n\r\n    return hidden_states, routed_hidden_states, router_logits\r\n\r\n\r\ndef qwen3_moe_sparse_moe_block_forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor:\r\n    \"\"\"NPU optimized implementation for `forward` in Qwen3MoeSparseMoeBlock.\"\"\"\r\n    output_shape = hidden_states.shape\r\n    _, routed_hidden_states, router_logits = _qwen3_sparse_moe_routed_forward_npu(self, hidden_states)\r\n    final_hidden_states = routed_hidden_states.reshape(output_shape)\r\n    return final_hidden_states, router_logits\r\n\r\n\r\ndef qwen3_next_sparse_moe_block_forward_npu(self, hidden_states: torch.Tensor) -> torch.Tensor:\r\n    \"\"\"NPU optimized implementation for `forward` in Qwen3NextSparseMoeBlock.\"\"\"\r\n    output_shape = hidden_states.shape\r\n    hidden_states, routed_hidden_states, router_logits = _qwen3_sparse_moe_routed_forward_npu(self, hidden_states)\r\n\r\n    shared_expert_output = self.shared_expert(hidden_states)\r\n    shared_expert_output = torch.sigmoid(self.shared_expert_gate(hidden_states)) * shared_expert_output\r\n\r\n    final_hidden_states = (routed_hidden_states + shared_expert_output).reshape(output_shape)\r\n    return final_hidden_states, router_logits\r\n\r\n\r\nclass NPUQwen3VLMoeTextExperts(nn.Module):\r\n    \"\"\"NPU optimized implementation for Qwen3VLMoeTextExperts.\"\"\"\r\n\r\n    def __init__(self, config):\r\n        super().__init__()\r\n        self.num_experts = config.num_experts\r\n        self.intermediate_size = config.moe_intermediate_size\r\n        self.hidden_size = config.hidden_size\r\n        self.expert_dim = self.intermediate_size\r\n        self.gate_up_proj = nn.Parameter(torch.empty(self.num_experts, self.hidden_size, 2 * self.expert_dim))\r\n        self.down_proj = nn.Parameter(torch.empty((self.num_experts, self.expert_dim, self.hidden_size)))\r\n        self.act_fn = ACT2FN[config.hidden_act]\r\n\r\n    def forward(\r\n        self, hidden_states: torch.Tensor, routing_weights: torch.Tensor, router_indices: torch.Tensor\r\n    ) -> torch.Tensor:\r\n        \"\"\"\r\n        When training it is more efficient to just loop over the experts and compute the output for each expert\r\n        as otherwise the memory would explode.\r\n\r\n        For inference we can sacrifice some memory and compute the output for all experts at once.\r\n        By repeating the inputs.\r\n\r\n        Args:\r\n            hidden_states (torch.Tensor): (batch_size * token_num, hidden_size)\r\n            routing_weights (torch.Tensor): (batch_size * token_num, num_experts)\r\n            router_indices (torch.Tensor): (batch_size * token_num, top_k)\r\n        Returns:\r\n            torch.Tensor\r\n        \"\"\"\r\n        batch_size = hidden_states.shape[0]\r\n        hidden_states = hidden_states.reshape(-1, self.hidden_size)  # (num_tokens, hidden_size)\r\n        if self.training:\r\n            permuted_hidden_states, row_ids_map = torch_npu.npu_moe_token_permute(\r\n                hidden_states, router_indices.to(torch.int32)\r\n            )\r\n            tokens_per_expert = torch.histc(router_indices, bins=self.num_experts, min=0, max=self.num_experts)\r\n            intermediate_hidden_states = NPUGmmFunction.apply(\r\n                permuted_hidden_states, self.gate_up_proj, tokens_per_expert\r\n            )\r\n            intermediate_activations = torch_npu.npu_swiglu(intermediate_hidden_states, dim=-1)\r\n            output = NPUGmmFunction.apply(intermediate_activations, self.down_proj, tokens_per_expert)\r\n            num_tokens = hidden_states.shape[0]\r\n            top_k = router_indices.shape[1]\r\n            batch_idx = torch.arange(num_tokens, device=routing_weights.device)\r\n            batch_idx = batch_idx.unsqueeze(1).expand(-1, top_k)\r\n            selected_probs = routing_weights[batch_idx, router_indices]\r\n            next_states = torch_npu.npu_moe_token_unpermute(output, row_ids_map, probs=selected_probs)\r\n            next_states = next_states.view(batch_size, -1, self.hidden_size)\r\n        else:\r\n            hidden_states = hidden_states.repeat(self.num_experts, 1)\r\n            hidden_states = hidden_states.view(self.num_experts, -1, self.hidden_size)\r\n            gate_up = torch.bmm(hidden_states, self.gate_up_proj)\r\n            gate, up = gate_up.chunk(2, dim=-1)  # not supported for DTensors\r\n            next_states = torch.bmm((up * self.act_fn(gate)), self.down_proj)\r\n            next_states = next_states.reshape(self.num_experts, batch_size, -1, self.hidden_size)\r\n            next_states = (\r\n                next_states * routing_weights.transpose(0, 1).view(self.num_experts, batch_size, -1)[..., None]\r\n            )\r\n            next_states = next_states.sum(dim=0)\r\n        return next_states\r\n\r\n\r\nclass NPUQwen3VLMoeTextSparseMoeBlock(nn.Module):\r\n    \"\"\"NPU optimized implementation for Qwen3VLMoeTextSparseMoeBlock.\"\"\"\r\n\r\n    def __init__(self, config):\r\n        super().__init__()\r\n        self.hidden_size = config.hidden_size\r\n        self.num_experts = config.num_experts\r\n        self.top_k = config.num_experts_per_tok\r\n        self.gate = nn.Linear(config.hidden_size, config.num_experts, bias=False)\r\n        self.experts = NPUQwen3VLMoeTextExperts(config)\r\n\r\n    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\r\n        batch_size = hidden_states.shape[0]\r\n        hidden_states = hidden_states.reshape(-1, self.hidden_size)\r\n        router_logits = self.gate(hidden_states)\r\n        routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)\r\n        routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)\r\n        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)\r\n        routing_weights = routing_weights.to(router_logits.dtype)\r\n        hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)\r\n        if not self.training:\r\n            routing_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)\r\n        routed_out = self.experts(hidden_states, routing_weights, router_indices)\r\n        return routed_out\r\n\r\n\r\n# Patches for Qwen2 Model\r\nmodeling_qwen2.Qwen2RMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen2.Qwen2MLP.forward = silu_forward_npu\r\nmodeling_qwen2.apply_rotary_pos_emb = apply_rotary_pos_emb_npu\r\n\r\n# Patches for Qwen2.5-VL Model\r\nmodeling_qwen2_5_vl.Qwen2RMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen2_5_vl.Qwen2_5_VLMLP.forward = silu_forward_npu\r\n\r\n# Patches for Qwen3 Model\r\nmodeling_qwen3.Qwen3RMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen3.Qwen3MLP.forward = silu_forward_npu\r\nmodeling_qwen3.apply_rotary_pos_emb = apply_rotary_pos_emb_npu\r\n\r\n# Patches for Qwen3 MoE Model\r\nmodeling_qwen3_moe.Qwen3MoeRMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen3_moe.Qwen3MoeSparseMoeBlock.forward = qwen3_moe_sparse_moe_block_forward_npu\r\nmodeling_qwen3_moe.apply_rotary_pos_emb = apply_rotary_pos_emb_npu\r\n\r\n# Patches for Qwen3 VL Model\r\nmodeling_qwen3_vl.Qwen3VLTextRMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen3_vl.Qwen3VLTextMLP.forward = silu_forward_npu\r\n\r\n# Patches for Qwen3-VL MoE Model\r\nmodeling_qwen3_vl_moe.Qwen3VLMoeTextSparseMoeBlock = NPUQwen3VLMoeTextSparseMoeBlock\r\nmodeling_qwen3_vl_moe.Qwen3VLMoeTextRMSNorm.forward = rms_norm_forward_npu\r\nmodeling_qwen3_vl_moe.apply_rotary_pos_emb = apply_rotary_pos_emb_npu\r\n\r\n# Patches for Qwen3 Next Model\r\nmodeling_qwen3_next.Qwen3NextSparseMoeBlock.forward = qwen3_next_sparse_moe_block_forward_npu\r\nmodeling_qwen3_next.Qwen3NextRMSNormGated.forward = qwen3_next_rms_norm_forward_gated_npu\r\nmodeling_qwen3_next.Qwen3NextRMSNorm.forward = qwen3_next_rms_norm_forward_npu\r\nmodeling_qwen3_next.apply_rotary_pos_emb = qwen3_next_apply_rotary_pos_emb_npu\r\n"
  },
  {
    "path": "verl/models/transformers/qwen2.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable, Optional\n\nimport torch\nfrom transformers.cache_utils import Cache\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward\nfrom transformers.models.llama.modeling_llama import apply_rotary_pos_emb, repeat_kv\nfrom transformers.utils import logging\n\n# Import compatibility wrapper for flash_attn_supports_top_left_mask\nfrom verl.utils.transformers_compat import flash_attn_supports_top_left_mask\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\nlogger = logging.get_logger(__name__)\n\n\ndef qwen2_flash_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    past_key_value: Optional[Cache] = None,\n    output_attentions: bool = False,\n    use_cache: bool = False,\n    cache_position: Optional[torch.LongTensor] = None,\n    position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46\n):\n    \"\"\"\n    Adapted from transformers 4.47.1 to support Ulysses sequence parallelism.\n\n    NOTE: This function is only tested on transformers versions between 4.45.0 and 4.47.1.\n    \"\"\"\n    bsz, q_len, _ = hidden_states.size()\n\n    query_states = self.q_proj(hidden_states)\n    key_states = self.k_proj(hidden_states)\n    value_states = self.v_proj(hidden_states)\n\n    query_states = query_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    key_states = key_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n    value_states = value_states.view(bsz, q_len, -1, self.head_dim).transpose(1, 2)\n\n    ########## AlltoAll for Ulysses ##########\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.num_heads, ulysses_sp_size)\n\n        # (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n\n    full_q_len = query_states.size(2)  # full seq length\n\n    if position_embeddings is None:\n        logger.warning_once(\n            \"The attention layers in this model are transitioning from computing the RoPE embeddings internally \"\n            \"through `position_ids` (2D tensor with the indexes of the tokens), to using externally computed \"\n            \"`position_embeddings` (Tuple of tensors, containing cos and sin). In v4.46 `position_ids` will be \"\n            \"removed and `position_embeddings` will be mandatory.\"\n        )\n        cos, sin = self.rotary_emb(value_states, position_ids)\n    else:\n        cos, sin = position_embeddings\n    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n\n    if past_key_value is not None:\n        cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position}  # Specific to RoPE models\n        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)\n\n    # repeat k/v heads if n_kv_heads < n_heads\n    key_states = repeat_kv(key_states, self.num_key_value_groups)\n    value_states = repeat_kv(value_states, self.num_key_value_groups)\n    dropout_rate = 0.0 if not self.training else self.attention_dropout\n\n    # In PEFT, usually we cast the layer norms in float32 for training stability reasons\n    # therefore the input hidden states gets silently casted in float32. Hence, we need\n    # cast them back in float16 just to be sure everything works as expected.\n    input_dtype = query_states.dtype\n    if input_dtype == torch.float32:\n        if torch.is_autocast_enabled():\n            target_dtype = torch.get_autocast_gpu_dtype()\n        # Handle the case where the model is quantized\n        elif hasattr(self.config, \"_pre_quantization_dtype\"):\n            target_dtype = self.config._pre_quantization_dtype\n        else:\n            target_dtype = self.q_proj.weight.dtype\n\n        logger.warning_once(\n            f\"The input hidden states seems to be silently casted in float32, this might be related to \"\n            f\"the fact you have upcasted embedding or layer norm layers in float32. We will cast back the \"\n            f\"input in {target_dtype}.\"\n        )\n\n        query_states = query_states.to(target_dtype)\n        key_states = key_states.to(target_dtype)\n        value_states = value_states.to(target_dtype)\n\n    # Reashape to the expected shape for Flash Attention\n    query_states = query_states.transpose(1, 2)\n    key_states = key_states.transpose(1, 2)\n    value_states = value_states.transpose(1, 2)\n\n    if (\n        self.config.use_sliding_window\n        and getattr(self.config, \"sliding_window\", None) is not None\n        and self.layer_idx >= self.config.max_window_layers\n    ):\n        sliding_window = self.config.sliding_window\n    else:\n        sliding_window = None\n\n    attn_output = _flash_attention_forward(\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        full_q_len,\n        position_ids=position_ids,\n        dropout=dropout_rate,\n        sliding_window=sliding_window,\n        is_causal=self.is_causal,\n        use_top_left_mask=flash_attn_supports_top_left_mask(),\n    )\n\n    # use full_q_len to reshape\n    attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1:\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n    attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()\n    attn_output = self.o_proj(attn_output)\n\n    if not output_attentions:\n        attn_weights = None\n\n    return attn_output, attn_weights, past_key_value\n\n\ndef qwen2_attn_forward(\n    self,\n    hidden_states: torch.Tensor,\n    position_embeddings: tuple[torch.Tensor, torch.Tensor],\n    attention_mask: Optional[torch.Tensor],\n    past_key_value: Optional[Cache] = None,\n    cache_position: Optional[torch.LongTensor] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, Optional[torch.Tensor], Optional[tuple[torch.Tensor]]]:\n    \"\"\"\n    Adapted from transformers 4.49.0 to support Ulysses sequence parallelism for transformers >= 4.48.0.\n\n    NOTE: This function has been tested only on transformers versions between 4.48.0 and 4.50.0.\n    \"\"\"\n    from transformers.modeling_utils import ALL_ATTENTION_FUNCTIONS\n\n    bsz, q_len, _ = hidden_states.shape\n    hidden_shape = (bsz, q_len, -1, self.head_dim)\n\n    query_states = self.q_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n    key_states = self.k_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n    value_states = self.v_proj(hidden_states).view(hidden_shape).transpose(1, 2)\n\n    ########## AlltoAll for Ulysses ##########\n    ulysses_sp_size = get_ulysses_sequence_parallel_world_size()\n\n    if ulysses_sp_size > 1:\n        validate_ulysses_config(self.config.num_attention_heads, ulysses_sp_size)\n\n        # (bsz, n_head, seq_len/n, head_dim) -> (bsz, n_head/n, seq_len, head_dim)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=2, head_dim=1)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=2, head_dim=1)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=2, head_dim=1)\n\n    full_q_len = query_states.size(2)\n\n    cos, sin = position_embeddings\n    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin)\n\n    if past_key_value is not None:\n        # sin and cos are specific to RoPE models; cache_position needed for the static cache\n        cache_kwargs = {\"sin\": sin, \"cos\": cos, \"cache_position\": cache_position}\n        key_states, value_states = past_key_value.update(key_states, value_states, self.layer_idx, cache_kwargs)\n\n    sliding_window = None\n    if (\n        self.config.use_sliding_window\n        and getattr(self.config, \"sliding_window\", None) is not None\n        and self.layer_idx >= self.config.max_window_layers\n    ):\n        sliding_window = self.config.sliding_window\n\n    from transformers.models.qwen2.modeling_qwen2 import eager_attention_forward\n\n    attention_interface: Callable = eager_attention_forward\n    if self.config._attn_implementation != \"eager\":\n        if self.config._attn_implementation == \"sdpa\" and kwargs.get(\"output_attentions\", False):\n            logger.warning_once(\n                \"`torch.nn.functional.scaled_dot_product_attention` does not support `output_attentions=True`. \"\n                \"Falling back to eager attention. This warning can be removed using the argument \"\n                '`attn_implementation=\"eager\"` when loading the model.'\n            )\n        else:\n            attention_interface = ALL_ATTENTION_FUNCTIONS[self.config._attn_implementation]\n\n    attn_output, attn_weights = attention_interface(\n        self,\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        dropout=0.0 if not self.training else self.attention_dropout,\n        scaling=self.scaling,\n        sliding_window=sliding_window,  # main diff with Llama\n        **kwargs,\n    )\n\n    attn_output = attn_output.reshape(bsz, full_q_len, -1, self.head_dim).contiguous()\n    ########## AlltoAll for Ulysses ##########\n    if ulysses_sp_size > 1:\n        # (bsz, seq_len, n_head/n, head_dim) -> (bsz, seq_len/n, n_head, head_dim)\n        attn_output = gather_heads_scatter_seq(attn_output, seq_dim=1, head_dim=2)\n    attn_output = attn_output.reshape(bsz, q_len, -1).contiguous()\n    attn_output = self.o_proj(attn_output)\n    return attn_output, attn_weights\n"
  },
  {
    "path": "verl/models/transformers/qwen2_vl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport logging\nimport os\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.distributed as dist\nfrom transformers.modeling_flash_attention_utils import _flash_attention_forward, fa_peft_integration_check\nfrom transformers.models.qwen2_vl.modeling_qwen2_vl import (\n    Qwen2VLAttention,\n    Qwen2VLCausalLMOutputWithPast,\n    Qwen2VLForConditionalGeneration,\n)\nfrom transformers.utils import is_flash_attn_2_available, is_flash_attn_greater_or_equal_2_10\n\nfrom verl.utils.device import is_npu_available\nfrom verl.utils.transformers_compat import is_transformers_version_in_range\nfrom verl.utils.ulysses import (\n    gather_heads_scatter_seq,\n    gather_seq_scatter_heads,\n    get_ulysses_sequence_parallel_group,\n    get_ulysses_sequence_parallel_world_size,\n    validate_ulysses_config,\n)\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nif is_flash_attn_2_available():\n    from flash_attn import flash_attn_func, flash_attn_varlen_func\n\n    _flash_supports_window_size = \"window_size\" in inspect.signature(flash_attn_func).parameters\n    _flash_supports_deterministic = \"deterministic\" in inspect.signature(flash_attn_func).parameters\n    _flash_use_top_left_mask = not is_flash_attn_greater_or_equal_2_10()\n\nif is_npu_available:\n    from transformers.integrations.npu_flash_attention import npu_flash_attn_func as flash_attn_func\n    from transformers.integrations.npu_flash_attention import npu_flash_attn_varlen_func as flash_attn_varlen_func\n    from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask\n\n    _flash_supports_window_size = \"window_size\" in inspect.signature(flash_attn_func).parameters\n    _flash_supports_deterministic = \"deterministic\" in inspect.signature(flash_attn_func).parameters\n    _flash_use_top_left_mask = flash_attn_supports_top_left_mask()\n\n_flash_deterministic_enabled = os.getenv(\"FLASH_ATTENTION_DETERMINISTIC\", \"0\") == \"1\"\n\n\ndef get_rope_index(\n    processor,\n    input_ids: torch.Tensor,\n    image_grid_thw: Optional[torch.Tensor] = None,\n    video_grid_thw: Optional[torch.Tensor] = None,\n    second_per_grid_ts: Optional[torch.Tensor] = None,\n    attention_mask: Optional[torch.Tensor] = None,\n) -> torch.Tensor:\n    \"\"\"\n    Gets the position ids for Qwen2-VL, it should be generated before sharding the sequence.\n    The batch dim has been removed and the input_ids should be a 1D tensor representing a single example.\n    https://github.com/huggingface/transformers/blob/v4.52.4/src/transformers/models/qwen2_5_vl/modeling_qwen2_5_vl.py#L1405\n    \"\"\"\n    spatial_merge_size = processor.image_processor.merge_size\n    tokens_per_second = 2\n    image_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|image_pad|>\")\n    video_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|video_pad|>\")\n    vision_start_token_id = processor.tokenizer.convert_tokens_to_ids(\"<|vision_start|>\")\n    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):\n        if attention_mask is None:\n            attention_mask = torch.ones_like(input_ids)\n\n        position_ids = torch.ones(3, input_ids.size(0), dtype=input_ids.dtype, device=input_ids.device)  # (3, seqlen)\n        image_index, video_index = 0, 0\n        input_ids = input_ids[attention_mask == 1]\n        image_nums, video_nums = 0, 0\n        vision_start_indices = torch.argwhere(input_ids == vision_start_token_id)\n        vision_tokens = input_ids[vision_start_indices + 1]\n        image_nums = (vision_tokens == image_token_id).sum()\n        video_nums = (vision_tokens == video_token_id).sum()\n        input_tokens = input_ids.tolist()\n        llm_pos_ids_list: list = []\n        st = 0\n        remain_images, remain_videos = image_nums, video_nums\n        for _ in range(image_nums + video_nums):\n            if image_token_id in input_tokens and remain_images > 0:\n                ed_image = input_tokens.index(image_token_id, st)\n            else:\n                ed_image = len(input_tokens) + 1\n            if video_token_id in input_tokens and remain_videos > 0:\n                ed_video = input_tokens.index(video_token_id, st)\n            else:\n                ed_video = len(input_tokens) + 1\n            if ed_image < ed_video:\n                t, h, w = (\n                    image_grid_thw[image_index][0],\n                    image_grid_thw[image_index][1],\n                    image_grid_thw[image_index][2],\n                )\n                second_per_grid_t = 0\n                image_index += 1\n                remain_images -= 1\n                ed = ed_image\n            else:\n                t, h, w = (\n                    video_grid_thw[video_index][0],\n                    video_grid_thw[video_index][1],\n                    video_grid_thw[video_index][2],\n                )\n                second_per_grid_t = second_per_grid_ts[video_index] if second_per_grid_ts is not None else 1.0\n\n                video_index += 1\n                remain_videos -= 1\n                ed = ed_video\n\n            llm_grid_t, llm_grid_h, llm_grid_w = (\n                t.item(),\n                h.item() // spatial_merge_size,\n                w.item() // spatial_merge_size,\n            )\n            text_len = ed - st\n\n            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n            t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)\n            t_index = (t_index * second_per_grid_t * tokens_per_second).long().flatten()\n            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()\n            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()\n            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)\n            st = ed + llm_grid_t * llm_grid_h * llm_grid_w\n\n        if st < len(input_tokens):\n            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n            text_len = len(input_tokens) - st\n            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)\n        position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device)\n    else:\n        if attention_mask is not None:\n            position_ids = attention_mask.long().cumsum(-1) - 1\n            position_ids.masked_fill_(attention_mask == 0, 1)\n            position_ids = position_ids.unsqueeze(0).expand(3, -1).to(input_ids.device)\n        else:\n            position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, -1).expand(3, -1)\n\n    return position_ids\n\n\ndef prepare_fa2_from_position_ids(\n    query: torch.Tensor, key: torch.Tensor, value: torch.Tensor, position_ids: torch.Tensor\n):\n    assert position_ids.ndim == 2  # (batch_size, seq_length)\n    query = query.contiguous().view(-1, query.size(-2), query.size(-1))\n    key = key.contiguous().view(-1, key.size(-2), key.size(-1))\n    value = value.contiguous().view(-1, value.size(-2), value.size(-1))\n    position_ids = position_ids.view(-1)\n    cu_seqlens = torch.cat(\n        (\n            (position_ids == 0).nonzero().view(-1).to(torch.int32),\n            torch.tensor(position_ids.size(), device=position_ids.device, dtype=torch.int32),\n        )\n    )\n    max_length = cu_seqlens.diff().max()  # use cu_seqlens to infer max_length for qwen2vl mrope\n    return (query, key, value, (cu_seqlens, cu_seqlens), (max_length, max_length))\n\n\ndef _custom_flash_attention_forward(\n    query_states: torch.Tensor,\n    key_states: torch.Tensor,\n    value_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor],\n    query_length: int,\n    is_causal: bool = True,\n    position_ids: Optional[torch.Tensor] = None,\n    sliding_window: Optional[int] = None,\n    use_top_left_mask: bool = False,\n    deterministic: Optional[bool] = None,\n    **kwargs,\n):\n    \"\"\"\n    Patches flash attention forward to handle 3D position ids in mrope. (3, batch_size, seq_length)\n    \"\"\"\n    # Assuming 4D tensors, key_states.shape[1] is the key/value sequence length (source length).\n    use_sliding_windows = (\n        _flash_supports_window_size and sliding_window is not None and key_states.shape[1] > sliding_window\n    )\n    flash_kwargs = {\"window_size\": (sliding_window, sliding_window)} if use_sliding_windows else {}\n\n    if _flash_supports_deterministic:\n        flash_kwargs[\"deterministic\"] = deterministic if deterministic is not None else _flash_deterministic_enabled\n\n    if kwargs.get(\"softcap\") is not None:\n        flash_kwargs[\"softcap\"] = kwargs.pop(\"softcap\")\n\n    query_states, key_states, value_states = fa_peft_integration_check(\n        query_states, key_states, value_states, target_dtype=torch.bfloat16\n    )\n\n    if position_ids is not None:\n        assert position_ids.ndim == 2  # (batch_size, seq_length / sp_size)\n\n    sp_size = get_ulysses_sequence_parallel_world_size()\n    if sp_size > 1:\n        # qkv: (batch_size, seq_length / sp_size, num_head, head_size)\n        validate_ulysses_config(query_states.size(2), sp_size)\n        query_states = gather_seq_scatter_heads(query_states, seq_dim=1, head_dim=2)\n        key_states = gather_seq_scatter_heads(key_states, seq_dim=1, head_dim=2)\n        value_states = gather_seq_scatter_heads(value_states, seq_dim=1, head_dim=2)\n        position_ids_lst = [torch.empty_like(position_ids) for _ in range(sp_size)]\n        position_ids = dist.all_gather(position_ids_lst, position_ids, group=get_ulysses_sequence_parallel_group())\n        position_ids = torch.cat(position_ids_lst, dim=-1)  # (batch_size, seq_length)\n\n    if position_ids is not None and query_length != 1 and not (torch.diff(position_ids, dim=-1) >= 0).all():\n        batch_size = query_states.size(0)\n        q, k, v, (cu_seqlens_q, cu_seqlens_k), (max_seqlen_q, max_seqlen_k) = prepare_fa2_from_position_ids(\n            query_states, key_states, value_states, position_ids\n        )\n        attn_output = flash_attn_varlen_func(\n            q=q,\n            k=k,\n            v=v,\n            cu_seqlens_q=cu_seqlens_q,\n            cu_seqlens_k=cu_seqlens_k,\n            max_seqlen_q=max_seqlen_q,\n            max_seqlen_k=max_seqlen_k,\n            dropout_p=kwargs.pop(\"dropout\", 0.0),\n            softmax_scale=kwargs.pop(\"softmax_scale\", None),\n            causal=is_causal,\n            **flash_kwargs,\n        )\n        attn_output = attn_output.view(batch_size, -1, attn_output.size(-2), attn_output.size(-1))\n    else:\n        attn_output = _flash_attention_forward(\n            query_states,\n            key_states,\n            value_states,\n            attention_mask,\n            query_length,\n            is_causal=is_causal,\n            sliding_window=sliding_window,\n            use_top_left_mask=use_top_left_mask,\n            deterministic=deterministic,\n            **kwargs,\n        )  # do not pass position_ids to old flash_attention_forward\n\n    if sp_size > 1:\n        # (batch_size, seq_length, num_head, head_size)\n        attn_output = gather_heads_scatter_seq(attn_output, head_dim=2, seq_dim=1)\n\n    return attn_output\n\n\ndef qwen2_vl_attn_forward(\n    self: \"Qwen2VLAttention\",\n    hidden_states: torch.Tensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    position_embeddings: Optional[tuple[torch.Tensor, torch.Tensor]] = None,  # will become mandatory in v4.46\n    **kwargs,\n) -> tuple[torch.Tensor, None, None]:\n    from transformers.models.qwen2_vl.modeling_qwen2_vl import apply_multimodal_rotary_pos_emb, repeat_kv\n\n    bsz, q_len, _ = hidden_states.size()  # q_len = seq_length / sp_size\n    query_states = self.q_proj(hidden_states)  # (batch_size, seq_length / sp_size, num_heads * head_size)\n    key_states = self.k_proj(hidden_states)\n    value_states = self.v_proj(hidden_states)\n\n    query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)\n    key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n    value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)\n\n    # Because the input can be padded, the absolute sequence length depends on the max position id.\n    cos, sin = position_embeddings\n    query_states, key_states = apply_multimodal_rotary_pos_emb(\n        query_states, key_states, cos, sin, self.rope_scaling[\"mrope_section\"]\n    )\n    key_states = repeat_kv(key_states, self.num_key_value_groups)\n    value_states = repeat_kv(value_states, self.num_key_value_groups)\n    dropout_rate = 0.0 if not self.training else self.attention_dropout\n\n    sliding_window = None\n    if (\n        self.config.use_sliding_window\n        and getattr(self.config, \"sliding_window\", None) is not None\n        and self.layer_idx >= self.config.max_window_layers\n    ):\n        sliding_window = self.config.sliding_window\n\n    # This is before the transpose\n    q_len = query_states.shape[2]\n\n    # FA2 uses non-transposed inputs\n    query_states = query_states.transpose(1, 2)\n    key_states = key_states.transpose(1, 2)\n    value_states = value_states.transpose(1, 2)\n\n    if position_ids.ndim == 3:\n        position_ids = position_ids[0]\n\n    attn_output = _custom_flash_attention_forward(\n        query_states,\n        key_states,\n        value_states,\n        attention_mask,\n        query_length=q_len,\n        is_causal=getattr(self, \"is_causal\", True),\n        dropout=dropout_rate,\n        sliding_window=sliding_window,\n        use_top_left_mask=_flash_use_top_left_mask,\n        position_ids=position_ids,  # important: pass position ids\n    )  # (batch_size, seq_length / sp_size, num_head, head_size)\n    attn_output = attn_output.reshape(bsz, q_len, self.hidden_size).contiguous()\n    attn_output = self.o_proj(attn_output)\n    if is_transformers_version_in_range(min_version=\"4.54.0\"):\n        return attn_output, None\n    else:\n        return attn_output, None, None\n\n\ndef _get_input_embeds(\n    model: \"Qwen2VLForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n):\n    inputs_embeds = model.get_input_embeddings()(input_ids)\n    if pixel_values is not None:\n        pixel_values = pixel_values.type(model.visual.dtype)\n        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        n_image_tokens = (input_ids == model.config.image_token_id).sum().item()\n        n_image_features = image_embeds.shape[0]\n        if n_image_tokens != n_image_features:\n            raise ValueError(\n                f\"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}\"\n            )\n\n        mask = input_ids == model.config.image_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        image_mask = mask_expanded.to(inputs_embeds.device)\n\n        image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)\n\n    if pixel_values_videos is not None:\n        pixel_values_videos = pixel_values_videos.type(model.visual.dtype)\n        video_embeds = model.visual(pixel_values_videos, grid_thw=video_grid_thw)\n        n_video_tokens = (input_ids == model.config.video_token_id).sum().item()\n        n_video_features = video_embeds.shape[0]\n        if n_video_tokens != n_video_features:\n            raise ValueError(\n                f\"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}\"\n            )\n\n        mask = input_ids == model.config.video_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        video_mask = mask_expanded.to(inputs_embeds.device)\n\n        video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)\n\n    if pixel_values is None and pixel_values_videos is None:  # handle mixed text-image data\n        config = model.config.vision_config\n        patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2\n        pixel_values = torch.zeros((16, patch_dim), dtype=inputs_embeds.dtype, device=inputs_embeds.device)\n        image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long, device=inputs_embeds.device)\n        image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        inputs_embeds += 0.0 * image_embeds.mean()\n\n    if attention_mask is not None:\n        attention_mask = attention_mask.to(inputs_embeds.device)\n\n    return inputs_embeds, attention_mask\n\n\ndef process_position_ids(position_ids: torch.Tensor) -> torch.Tensor:\n    if position_ids.ndim != 3 or position_ids.size(0) != 4:\n        # we concat the text position ids with the 3D vision position ids by default\n        # see https://github.com/huggingface/transformers/pull/39447\n        raise ValueError(\"position_ids should be a 3D tensor of shape (4, batch_size, seq_length).\")\n\n    if is_transformers_version_in_range(max_version=\"4.53.3\"):\n        # transformers < 4.54.0 only accepts vision position ids, so we discard the text position ids here\n        position_ids = position_ids[1:]\n\n    return position_ids\n\n\n@dataclass\nclass Qwen2VLCausalLMOutputForPPO(Qwen2VLCausalLMOutputWithPast):\n    log_probs: Optional[torch.FloatTensor] = None\n    entropy: Optional[torch.FloatTensor] = None\n\n\ndef qwen2_vl_base_forward(\n    self: \"Qwen2VLForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    labels: Optional[torch.LongTensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    **kwargs,\n):\n    kwargs[\"inputs_embeds\"], kwargs[\"attention_mask\"] = _get_input_embeds(\n        self, input_ids, attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw\n    )  # avoid lora module having multiple keyword arguments\n    return self.language_model(input_ids=None, **kwargs)\n\n\ndef qwen2_vl_forward(\n    self: \"Qwen2VLForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    position_ids: Optional[torch.LongTensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    **kwargs,\n):\n    if is_transformers_version_in_range(min_version=\"4.52.0\"):\n        return self.model(\n            input_ids=input_ids,\n            attention_mask=attention_mask,\n            position_ids=process_position_ids(position_ids),\n            pixel_values=pixel_values,\n            pixel_values_videos=pixel_values_videos,\n            image_grid_thw=image_grid_thw,\n            video_grid_thw=video_grid_thw,\n            **kwargs,\n        )\n    else:\n        inputs_embeds, attention_mask = _get_input_embeds(\n            self, input_ids, attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw\n        )\n        return self.model(\n            input_ids=None,\n            attention_mask=attention_mask,\n            position_ids=process_position_ids(position_ids),\n            inputs_embeds=inputs_embeds,\n            **kwargs,\n        )\n\n\ndef forward_with_normal_backend(\n    self: Qwen2VLForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> \"Qwen2VLCausalLMOutputWithPast\":\n    outputs = qwen2_vl_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n    logits = self.lm_head(hidden_states)\n\n    return Qwen2VLCausalLMOutputWithPast(\n        logits=logits,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_torch_backend(\n    self: Qwen2VLForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> tuple | Qwen2VLCausalLMOutputForPPO:\n    from verl.utils.experimental.torch_functional import FusedLinearForPPO\n\n    outputs = qwen2_vl_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_torch_backend, either labels or input_ids must be provided.\")\n\n    fused_linear_for_ppo = FusedLinearForPPO()\n    log_probs, entropy = fused_linear_for_ppo.forward(\n        hidden_states=hidden_states,\n        vocab_weights=self.lm_head.weight,\n        input_ids=rolled_labels,\n        temperature=temperature,\n    )\n    return Qwen2VLCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_triton_backend(\n    self: Qwen2VLForConditionalGeneration,\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> tuple | Qwen2VLCausalLMOutputForPPO:\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\n\n    outputs = qwen2_vl_forward(self, input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_triton_backend, either labels or input_ids must be provided.\")\n\n    log_probs, entropy = linear_cross_entropy(\n        hidden_states,\n        self.lm_head.weight,\n        rolled_labels,\n        temperature,\n        \"none\",\n    )\n    return Qwen2VLCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n"
  },
  {
    "path": "verl/models/transformers/qwen3_vl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nimport logging\nimport os\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport torch\nfrom transformers.models.qwen3_vl.modeling_qwen3_vl import (\n    Qwen3VLCausalLMOutputWithPast,\n    Qwen3VLForConditionalGeneration,\n)\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef get_rope_index(\n    processor,\n    input_ids: torch.Tensor,\n    image_grid_thw: Optional[torch.Tensor] = None,\n    video_grid_thw: Optional[torch.Tensor] = None,\n    attention_mask: Optional[torch.Tensor] = None,\n    **kwargs,\n) -> torch.Tensor:\n    \"\"\"\n    Gets the position ids for Qwen3-VL, it should be generated before sharding the sequence.\n    The batch dim has been removed and the input_ids should be a 1D tensor representing a single example.\n    https://github.com/huggingface/transformers/blob/v4.57.0/src/transformers/models/qwen3_vl/modeling_qwen3_vl.py#L916\n    \"\"\"\n    spatial_merge_size = processor.image_processor.merge_size\n    image_token_id = processor.image_token_id\n    video_token_id = processor.video_token_id\n    vision_start_token_id = processor.vision_start_token_id\n\n    # Since we use timestamps to separate videos,\n    # like <t1> <vision_start> <frame1> <vision_end> <t2> <vision_start> <frame2> <vision_end>,\n    # the video_grid_thw should also be split\n    if video_grid_thw is not None:\n        video_grid_thw = torch.repeat_interleave(video_grid_thw, video_grid_thw[:, 0], dim=0)\n        video_grid_thw[:, 0] = 1\n\n    if input_ids is not None and (image_grid_thw is not None or video_grid_thw is not None):\n        if attention_mask is None:\n            attention_mask = torch.ones_like(input_ids)\n\n        position_ids = torch.ones(3, input_ids.shape[0], dtype=input_ids.dtype, device=input_ids.device)\n        image_index, video_index = 0, 0\n        attention_mask = attention_mask.to(input_ids.device)\n        input_ids = input_ids[attention_mask == 1]\n        image_nums, video_nums = 0, 0\n        vision_start_indices = torch.argwhere(input_ids == vision_start_token_id)\n        vision_tokens = input_ids[vision_start_indices + 1]\n        image_nums = (vision_tokens == image_token_id).sum()\n        video_nums = (vision_tokens == video_token_id).sum()\n        input_tokens = input_ids.tolist()\n        llm_pos_ids_list: list = []\n        st = 0\n        remain_images, remain_videos = image_nums, video_nums\n        for _ in range(image_nums + video_nums):\n            if image_token_id in input_tokens and remain_images > 0:\n                ed_image = input_tokens.index(image_token_id, st)\n            else:\n                ed_image = len(input_tokens) + 1\n            if video_token_id in input_tokens and remain_videos > 0:\n                ed_video = input_tokens.index(video_token_id, st)\n            else:\n                ed_video = len(input_tokens) + 1\n            if ed_image < ed_video:\n                t, h, w = (\n                    image_grid_thw[image_index][0],\n                    image_grid_thw[image_index][1],\n                    image_grid_thw[image_index][2],\n                )\n                image_index += 1\n                remain_images -= 1\n                ed = ed_image\n            else:\n                t, h, w = (\n                    video_grid_thw[video_index][0],\n                    video_grid_thw[video_index][1],\n                    video_grid_thw[video_index][2],\n                )\n                video_index += 1\n                remain_videos -= 1\n                ed = ed_video\n\n            llm_grid_t, llm_grid_h, llm_grid_w = (\n                t.item(),\n                h.item() // spatial_merge_size,\n                w.item() // spatial_merge_size,\n            )\n            text_len = ed - st\n\n            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n            # t_index is always 0 because llm_grid_t is always 1\n            # (we use timestamps to encode the temporal information for videos)\n            t_index = torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w).flatten()\n            h_index = torch.arange(llm_grid_h).view(1, -1, 1).expand(llm_grid_t, -1, llm_grid_w).flatten()\n            w_index = torch.arange(llm_grid_w).view(1, 1, -1).expand(llm_grid_t, llm_grid_h, -1).flatten()\n            llm_pos_ids_list.append(torch.stack([t_index, h_index, w_index]) + text_len + st_idx)\n            st = ed + llm_grid_t * llm_grid_h * llm_grid_w\n\n        if st < len(input_tokens):\n            st_idx = llm_pos_ids_list[-1].max() + 1 if len(llm_pos_ids_list) > 0 else 0\n            text_len = len(input_tokens) - st\n            llm_pos_ids_list.append(torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx)\n\n        llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)\n        position_ids[..., attention_mask == 1] = llm_positions.to(position_ids.device)\n    else:\n        if attention_mask is not None:\n            position_ids = attention_mask.long().cumsum(-1) - 1\n            position_ids.masked_fill_(attention_mask == 0, 1)\n            position_ids = position_ids.unsqueeze(0).expand(3, -1).to(attention_mask.device)\n        else:\n            position_ids = torch.arange(input_ids.shape[1], device=input_ids.device).view(1, -1).expand(3, -1)\n\n    return position_ids\n\n\ndef _get_input_embeds(\n    model: \"Qwen3VLForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n):\n    inputs_embeds = model.get_input_embeddings()(input_ids)\n    image_mask, video_mask = None, None\n    if pixel_values is not None:\n        pixel_values = pixel_values.type(model.visual.dtype)\n        image_embeds, deepstack_image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        n_image_tokens = (input_ids == model.config.image_token_id).sum().item()\n        n_image_features = image_embeds.shape[0]\n        if n_image_tokens != n_image_features:\n            raise ValueError(\n                f\"Image features and image tokens do not match: tokens: {n_image_tokens}, features {n_image_features}\"\n            )\n\n        mask = input_ids == model.config.image_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        image_mask = mask_expanded.to(inputs_embeds.device)\n\n        image_embeds = image_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(image_mask, image_embeds)\n\n    if pixel_values_videos is not None:\n        pixel_values_videos = pixel_values_videos.type(model.visual.dtype)\n        video_embeds, deepstack_video_embeds = model.visual(pixel_values_videos, grid_thw=video_grid_thw)\n        n_video_tokens = (input_ids == model.config.video_token_id).sum().item()\n        n_video_features = video_embeds.shape[0]\n        if n_video_tokens != n_video_features:\n            raise ValueError(\n                f\"Video features and video tokens do not match: tokens: {n_video_tokens}, features {n_video_features}\"\n            )\n\n        mask = input_ids == model.config.video_token_id\n        mask_unsqueezed = mask.unsqueeze(-1)\n        mask_expanded = mask_unsqueezed.expand_as(inputs_embeds)\n        video_mask = mask_expanded.to(inputs_embeds.device)\n\n        video_embeds = video_embeds.to(inputs_embeds.device, inputs_embeds.dtype)\n        inputs_embeds = inputs_embeds.masked_scatter(video_mask, video_embeds)\n\n    visual_pos_masks = None\n    deepstack_visual_embeds = None\n    if image_mask is not None and video_mask is not None:\n        # aggregate visual_pos_masks and deepstack_visual_embeds\n        image_mask = image_mask[..., 0]\n        video_mask = video_mask[..., 0]\n        visual_pos_masks = image_mask | video_mask\n        deepstack_visual_embeds = []\n        image_mask_joint = image_mask[visual_pos_masks]\n        video_mask_joint = video_mask[visual_pos_masks]\n        for img_embed, vid_embed in zip(deepstack_image_embeds, deepstack_video_embeds, strict=False):\n            embed_joint = img_embed.new_zeros(visual_pos_masks.sum(), img_embed.shape[-1]).to(img_embed.device)\n            embed_joint[image_mask_joint, :] = img_embed\n            embed_joint[video_mask_joint, :] = vid_embed\n            deepstack_visual_embeds.append(embed_joint)\n    elif image_mask is not None:\n        image_mask = image_mask[..., 0]\n        visual_pos_masks = image_mask\n        deepstack_visual_embeds = deepstack_image_embeds\n    elif video_mask is not None:\n        video_mask = video_mask[..., 0]\n        visual_pos_masks = video_mask\n        deepstack_visual_embeds = deepstack_video_embeds\n\n    if pixel_values is None and pixel_values_videos is None:\n        config = model.config.vision_config\n        patch_dim = config.in_channels * config.temporal_patch_size * config.patch_size**2\n        pixel_values = torch.zeros((16, patch_dim), dtype=inputs_embeds.dtype, device=inputs_embeds.device)\n        image_grid_thw = torch.tensor([[1, 4, 4]], dtype=torch.long, device=inputs_embeds.device)\n        image_embeds, dummy_deepstack_image_embeds = model.visual(pixel_values, grid_thw=image_grid_thw)\n        inputs_embeds += 0.0 * image_embeds.mean()\n        for emb in dummy_deepstack_image_embeds or []:\n            inputs_embeds += 0.0 * emb.mean()\n\n    if attention_mask is not None:\n        attention_mask = attention_mask.to(inputs_embeds.device)\n\n    return {\n        \"inputs_embeds\": inputs_embeds,\n        \"attention_mask\": attention_mask,\n        \"visual_pos_masks\": visual_pos_masks,\n        \"deepstack_visual_embeds\": deepstack_visual_embeds,\n    }\n\n\n@dataclass\nclass Qwen3VLCausalLMOutputForPPO(Qwen3VLCausalLMOutputWithPast):\n    log_probs: Optional[torch.FloatTensor] = None\n    entropy: Optional[torch.FloatTensor] = None\n\n\ndef qwen3_vl_base_forward(\n    self: \"Qwen3VLForConditionalGeneration\",\n    input_ids: torch.LongTensor,\n    attention_mask: Optional[torch.Tensor] = None,\n    pixel_values: Optional[torch.FloatTensor] = None,\n    pixel_values_videos: Optional[torch.FloatTensor] = None,\n    image_grid_thw: Optional[torch.LongTensor] = None,\n    video_grid_thw: Optional[torch.LongTensor] = None,\n    **kwargs,\n):\n    input_kwargs = _get_input_embeds(\n        self, input_ids, attention_mask, pixel_values, pixel_values_videos, image_grid_thw, video_grid_thw\n    )  # avoid lora module having multiple keyword arguments\n    kwargs.update(input_kwargs)\n    return self.language_model(\n        input_ids=None,\n        **kwargs,\n    )\n\n\ndef forward_with_normal_backend(\n    self: \"Qwen3VLForConditionalGeneration\",\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> \"Qwen3VLCausalLMOutputForPPO\":\n    outputs = self.model(input_ids, **kwargs)\n    hidden_states = outputs[0]\n    logits = self.lm_head(hidden_states)\n\n    return Qwen3VLCausalLMOutputForPPO(\n        logits=logits,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_torch_backend(\n    self: \"Qwen3VLForConditionalGeneration\",\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> \"Qwen3VLCausalLMOutputForPPO\":\n    from verl.utils.experimental.torch_functional import FusedLinearForPPO\n\n    outputs = self.model(input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_torch_backend, either labels or input_ids must be provided.\")\n\n    fused_linear_for_ppo = FusedLinearForPPO()\n    log_probs, entropy = fused_linear_for_ppo.forward(\n        hidden_states=hidden_states,\n        vocab_weights=self.lm_head.weight,\n        input_ids=rolled_labels,\n        temperature=temperature,\n    )\n    return Qwen3VLCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef forward_with_triton_backend(\n    self: \"Qwen3VLForConditionalGeneration\",\n    input_ids: torch.LongTensor = None,\n    labels: Optional[torch.LongTensor] = None,\n    temperature: float = 1.0,\n    **kwargs,\n) -> \"Qwen3VLCausalLMOutputForPPO\":\n    from verl.utils.kernel.linear_cross_entropy import linear_cross_entropy\n\n    outputs = self.model(input_ids, **kwargs)\n    hidden_states = outputs[0]\n\n    # Loss calculations\n    if labels is not None:\n        rolled_labels = torch.roll(labels, shifts=-1, dims=-1)\n    elif input_ids is not None:\n        rolled_labels = torch.roll(input_ids, shifts=-1, dims=-1)\n    else:\n        raise RuntimeError(\"To use forward_with_triton_backend, either labels or input_ids must be provided.\")\n\n    log_probs, entropy = linear_cross_entropy(\n        hidden_states,\n        self.lm_head.weight,\n        rolled_labels,\n        temperature,\n        \"none\",\n    )\n    return Qwen3VLCausalLMOutputForPPO(\n        log_probs=log_probs,\n        entropy=entropy,\n        hidden_states=outputs.hidden_states,\n    )\n\n\ndef patch_qwen3_vl_moe_sparse_moe_block_forward():\n    \"\"\"\n    Monkey patch to fix a bug in transformers 4.57.3 where Qwen3VLMoeTextSparseMoeBlock.forward\n    incorrectly uses torch.zeros_like(hidden_states) instead of torch.zeros_like(router_logits)\n    when creating router_weights (line 148 in modeling_qwen3_vl_moe.py).\n\n    This is a minimal fix that only changes the problematic line while keeping the rest of the\n    original implementation intact.\n    \"\"\"\n    try:\n        from transformers.models.qwen3_vl_moe.modeling_qwen3_vl_moe import Qwen3VLMoeTextSparseMoeBlock\n    except ImportError:\n        # Model not available, skip patching\n        return\n\n    # Store the original forward method for reference\n    original_forward = Qwen3VLMoeTextSparseMoeBlock.forward\n\n    @functools.wraps(original_forward)\n    def patched_forward(self, hidden_states: torch.Tensor) -> torch.Tensor:\n        batch_size = hidden_states.shape[0]\n        hidden_states = hidden_states.reshape(-1, self.hidden_size)\n        router_logits = self.gate(hidden_states)\n        routing_weights = torch.nn.functional.softmax(router_logits, dim=-1, dtype=torch.float)\n        routing_weights, router_indices = torch.topk(routing_weights, self.top_k, dim=-1)\n        routing_weights = routing_weights / routing_weights.sum(dim=-1, keepdim=True)\n        # BUG FIX: Original code incorrectly uses hidden_states here, should use router_logits\n        routing_weights = routing_weights.to(router_logits.dtype)\n        router_weights = torch.zeros_like(router_logits).scatter_(1, router_indices, routing_weights)\n        hidden_states = hidden_states.reshape(batch_size, -1, self.hidden_size)\n        routed_out = self.experts(hidden_states, router_weights, router_indices)\n        return routed_out\n\n    # Apply the patch\n    Qwen3VLMoeTextSparseMoeBlock.forward = patched_forward\n    logger.info(\"Monkey patched Qwen3VLMoeTextSparseMoeBlock.forward to fix router_weights bug\")\n"
  },
  {
    "path": "verl/models/transformers/tiled_mlp.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nFSDP2-compatible TiledMLP implementation for memory-efficient MLP computation.\n\nThis module provides a tiled MLP implementation that reduces peak memory usage\nby processing the MLP forward/backward pass in chunks (tiles). This is particularly\nuseful for large models with FSDP2 training.\n\"\"\"\n\nimport threading\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\n\n\nclass GradientAccumulator:\n    \"\"\"Gradient accumulator for TiledMLP (FSDP compatible).\n\n    This class manages gradient accumulation across multiple shards during\n    the backward pass of TiledMLP. It ensures correct gradient computation\n    when processing input in chunks.\n    \"\"\"\n\n    def __init__(self, params: list[torch.nn.Parameter], total_shards: int, dtype: torch.dtype = None):\n        self.params = params\n        self.total_shards = total_shards\n        self.grad_accumulation_dtype = dtype or torch.float32\n        self.accumulated_grads = {}\n        self.hooks = []\n        self.lock = threading.Lock()\n\n        for param in self.params:\n            if param.grad is not None:\n                self.accumulated_grads[param] = param.grad.to(self.grad_accumulation_dtype)\n                param.grad = None\n            else:\n                self.accumulated_grads[param] = torch.zeros_like(param, dtype=self.grad_accumulation_dtype)\n\n    def install_hooks(self, is_last_shard: bool):\n        \"\"\"Install gradient hooks for the current shard.\"\"\"\n        self._remove_hooks()\n\n        def create_hook(param):\n            def hook(grad):\n                with self.lock:\n                    grad_to_accum_dtype = grad.to(self.grad_accumulation_dtype)\n                    self.accumulated_grads[param] += grad_to_accum_dtype\n\n                    if is_last_shard:\n                        param.grad = None  # Critical: prevent double accumulation\n                        final_grad = self.accumulated_grads[param].to(param.dtype)\n                        return final_grad\n                    return None\n\n            return hook\n\n        for param in self.params:\n            if param.requires_grad:\n                hook = param.register_hook(create_hook(param))\n                self.hooks.append(hook)\n\n    def _remove_hooks(self):\n        \"\"\"Remove all registered hooks.\"\"\"\n        for hook in self.hooks:\n            hook.remove()\n        self.hooks.clear()\n\n    def cleanup(self):\n        \"\"\"Cleanup hooks and resources.\"\"\"\n        self._remove_hooks()\n\n\nclass TiledMLP(torch.autograd.Function):\n    \"\"\"TiledMLP implementation for memory-efficient MLP computation.\n\n    This autograd function processes MLP forward/backward in tiles (chunks)\n    to reduce peak memory usage. Compatible with FSDP2.\n    \"\"\"\n\n    @staticmethod\n    def forward(ctx, fn, module, x, shards, compute_params):\n        ctx.fn = fn\n        ctx.module = module\n        ctx.shards = shards\n        ctx.compute_params = [p for p in compute_params if p.requires_grad]\n        ctx.save_for_backward(x)\n\n        # Split on dim=-2 (seqlen dimension) following Liger Kernel style\n        x_shards = list(torch.chunk(x, chunks=shards, dim=-2))\n        with torch.no_grad():\n            output_shards = [fn(module, x_shard) for x_shard in x_shards]\n        output_unsharded = torch.cat(output_shards, dim=-2)\n        return output_unsharded\n\n    @staticmethod\n    def backward(ctx, *grads):\n        fn = ctx.fn\n        (x,) = ctx.saved_tensors\n        module = ctx.module\n        shards = ctx.shards\n        compute_params = ctx.compute_params\n\n        x_requires_grad = x.requires_grad\n        x = x.detach()\n        x.requires_grad_(x_requires_grad)\n\n        # Flatten to [bs*seqlen, hidden_size]\n        hidden_size = x.shape[-1]\n        x_shape_orig = x.shape\n        x = x.view(-1, hidden_size)\n        incoming_grad = grads[0].view(-1, hidden_size)\n\n        # Pre-allocate input gradient\n        x_grad = torch.zeros_like(x)\n\n        # Split on dim=0\n        x_shards = list(torch.chunk(x, chunks=shards, dim=0))\n\n        grad_accumulator = GradientAccumulator(compute_params, shards, dtype=x.dtype)\n\n        for i, x_shard in enumerate(x_shards):\n            x_shard.requires_grad_(x_requires_grad)\n\n            shard_step = x_shards[i].shape[0]\n            shard_offset = i * x_shards[0].shape[0]\n\n            # narrow(0, ...) creates a contiguous view that can receive gradients\n            x_shard.grad = x_grad.narrow(0, shard_offset, shard_step)\n            incoming_grad_shard = incoming_grad.narrow(0, shard_offset, shard_step)\n\n            is_last_shard = i + 1 == shards\n            grad_accumulator.install_hooks(is_last_shard)\n\n            with torch.enable_grad():\n                output = fn(module, x_shard)\n            torch.autograd.backward(output, incoming_grad_shard)\n\n        grad_accumulator.cleanup()\n        del grad_accumulator\n\n        # Restore original shape\n        x_grad = x_grad.view(x_shape_orig) if x_requires_grad else None\n        return (None, None, x_grad, None, None)\n\n\ndef _mlp_forward_fn(module, x):\n    \"\"\"Forward function for LlamaMLP / Qwen2MLP / Qwen3MLP style.\"\"\"\n    return module.down_proj(module.act_fn(module.gate_proj(x)) * module.up_proj(x))\n\n\n# ============================================================================\n# Monkey Patch Functions\n# ============================================================================\n\n# Model type to MLP class mapping\n_MODEL_TYPE_TO_MLP_CLASS = {\n    \"llama\": (\"transformers.models.llama.modeling_llama\", \"LlamaMLP\"),\n    \"qwen2\": (\"transformers.models.qwen2.modeling_qwen2\", \"Qwen2MLP\"),\n    \"qwen2_5\": (\"transformers.models.qwen2.modeling_qwen2\", \"Qwen2MLP\"),  # Qwen2.5 uses Qwen2 MLP\n    \"qwen3\": (\"transformers.models.qwen3.modeling_qwen3\", \"Qwen3MLP\"),\n}\n\n\ndef apply_tiled_mlp_monkey_patch(\n    num_shards: int = 4,\n    model_type: Optional[str] = None,\n):\n    \"\"\"Apply TiledMLP monkey patch based on model_type.\n\n    This function MUST be called BEFORE model instantiation to take effect.\n    It patches the MLP classes in transformers library to use TiledMLP for\n    memory-efficient computation during training.\n\n    Args:\n        num_shards: Number of shards to split the input into. Higher values\n                   reduce peak memory but may slightly impact performance.\n        model_type: The model type string (e.g., \"llama\", \"qwen2\", \"qwen3\").\n                   If None, patches all supported model types.\n\n    Returns:\n        List of patched class names.\n    \"\"\"\n    if model_type is None:\n        types_to_patch = list(_MODEL_TYPE_TO_MLP_CLASS.keys())\n    elif model_type in _MODEL_TYPE_TO_MLP_CLASS:\n        types_to_patch = [model_type]\n    else:\n        raise ValueError(\n            f\"TiledMLP does not support model_type='{model_type}'. \"\n            f\"Supported types: {list(_MODEL_TYPE_TO_MLP_CLASS.keys())}. \"\n            f\"For SwiGLU-style MLPs, you can add support by extending _MODEL_TYPE_TO_MLP_CLASS \"\n            f\"in verl/models/transformers/tiled_mlp.py\"\n        )\n\n    patched_classes = []\n\n    for mtype in types_to_patch:\n        module_path, class_name = _MODEL_TYPE_TO_MLP_CLASS[mtype]\n        try:\n            import importlib\n\n            module = importlib.import_module(module_path)\n            mlp_class = getattr(module, class_name)\n            _patch_mlp_class(mlp_class, _mlp_forward_fn, num_shards)\n            if class_name not in patched_classes:\n                patched_classes.append(class_name)\n        except (ImportError, AttributeError) as e:\n            print(f\"Warning: Could not patch {mtype} MLP: {e}\")\n\n    if patched_classes:\n        print(f\"TiledMLP monkey patch applied to: {', '.join(patched_classes)} (shards={num_shards})\")\n\n    return patched_classes\n\n\ndef _patch_mlp_class(mlp_class: type[nn.Module], forward_fn, num_shards: int):\n    \"\"\"Patch a single MLP class to use TiledMLP.\"\"\"\n\n    def tiled_forward(self, x):\n        compute_params = [p for p in self.parameters() if p.requires_grad]\n        return TiledMLP.apply(forward_fn, self, x, num_shards, compute_params)\n\n    mlp_class.forward = tiled_forward\n"
  },
  {
    "path": "verl/models/weight_loader_registry.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\ndef get_weight_loader(arch: str):\n    from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel\n\n    _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY = {\n        \"LlamaForCausalLM\": load_state_dict_to_megatron_gptmodel,\n        \"Qwen2ForCausalLM\": load_state_dict_to_megatron_gptmodel,\n    }\n\n    if arch in _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY:\n        return _MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY[arch]\n    raise ValueError(\n        f\"Model architectures {arch} loader are not supported for now. Supported architectures: \"\n        f\"{_MODEL_WEIGHT_MEGATRON_LOADER_REGISTRY.keys()}\"\n    )\n\n\ndef get_weight_saver(arch: str):\n    from verl.models.mcore.saver import (\n        merge_megatron_ckpt_gptmodel,\n        merge_megatron_ckpt_gptmodel_dpskv3,\n        merge_megatron_ckpt_gptmodel_mixtral,\n        merge_megatron_ckpt_gptmodel_qwen2_5_vl,\n        merge_megatron_ckpt_gptmodel_qwen_moe,\n    )\n\n    _MODEL_WEIGHT_MEGATRON_SAVER_REGISTRY = {\n        \"LlamaForCausalLM\": merge_megatron_ckpt_gptmodel,\n        \"Qwen2ForCausalLM\": merge_megatron_ckpt_gptmodel,\n        \"MixtralForCausalLM\": merge_megatron_ckpt_gptmodel_mixtral,\n        \"Qwen2MoeForCausalLM\": merge_megatron_ckpt_gptmodel_qwen_moe,\n        \"Qwen2_5_VLForConditionalGeneration\": merge_megatron_ckpt_gptmodel_qwen2_5_vl,\n        \"DeepseekV3ForCausalLM\": merge_megatron_ckpt_gptmodel_dpskv3,\n        \"Qwen3ForCausalLM\": merge_megatron_ckpt_gptmodel,\n        \"Qwen3ForTokenClassification\": merge_megatron_ckpt_gptmodel,\n        \"Qwen3MoeForCausalLM\": merge_megatron_ckpt_gptmodel_qwen_moe,\n        \"LlamaForTokenClassification\": merge_megatron_ckpt_gptmodel,\n    }\n    if arch in _MODEL_WEIGHT_MEGATRON_SAVER_REGISTRY:\n        return _MODEL_WEIGHT_MEGATRON_SAVER_REGISTRY[arch]\n    raise ValueError(\n        f\"Model architectures {arch} saver are not supported for now. Supported architectures: \"\n        f\"{_MODEL_WEIGHT_MEGATRON_SAVER_REGISTRY.keys()}\"\n    )\n"
  },
  {
    "path": "verl/protocol.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nImplement base data transfer protocol between any two functions, modules.\nWe can subclass Protocol to define more detailed batch info with specific keys\n\"\"\"\n\nimport contextlib\nimport copy\nimport logging\nimport math\nimport os\nimport pickle\nfrom dataclasses import dataclass, field\nfrom typing import Any, Callable, Optional\n\nimport numpy as np\nimport ray\nimport tensordict\nimport torch\nimport torch.distributed\nfrom packaging import version\nfrom packaging.version import parse as parse_version\nfrom tensordict import TensorDict\nfrom torch.utils.data import DataLoader\n\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.py_functional import list_of_dict_to_dict_of_list, union_two_dict\nfrom verl.utils.torch_functional import allgather_dict_tensors\n\n__all__ = [\"DataProto\", \"union_tensor_dict\"]\n\nwith contextlib.suppress(Exception):\n    tensordict.set_lazy_legacy(False).set()\n    if parse_version(tensordict.__version__) < parse_version(\"0.10.0\"):\n        tensordict.set_list_to_stack(True).set()\n\n\nclass _DataProtoConfigMeta(type):\n    _config = {}\n\n    auto_padding_key = \"_verl_auto_padding\"\n\n    @property\n    def auto_padding(cls):\n        enabled_by_env = os.getenv(\"VERL_AUTO_PADDING\", \"FALSE\").upper() in [\"TRUE\", \"1\"]\n        return enabled_by_env or cls._config.get(cls.auto_padding_key, False)\n\n    @auto_padding.setter\n    def auto_padding(cls, enabled: bool):\n        assert isinstance(enabled, bool), f\"enabled must be a boolean, got {enabled} as {type(enabled)}\"\n        cls._config[cls.auto_padding_key] = enabled\n\n\nclass DataProtoConfig(metaclass=_DataProtoConfigMeta):\n    pass\n\n\n_padding_size_key = \"_padding_size_key_x123d\"\n\n\ndef pad_dataproto_to_divisor(data: \"DataProto\", size_divisor: int):\n    \"\"\"Pad a DataProto to size divisible by size_divisor\n\n    Args:\n        size_divisor (int): size divisor\n\n    Returns:\n        data: (DataProto): the padded DataProto\n        pad_size (int)\n    \"\"\"\n    assert isinstance(data, DataProto), \"data must be a DataProto\"\n    if len(data) % size_divisor != 0:\n        pad_size = size_divisor - len(data) % size_divisor\n        padding_protos = []\n        remaining_pad = pad_size\n        while remaining_pad > 0:\n            take_size = min(remaining_pad, len(data))\n            padding_protos.append(data[:take_size])\n            remaining_pad -= take_size\n        data_padded = DataProto.concat([data] + padding_protos)\n    else:\n        if len(data) == 0:\n            logging.warning(\"padding a DataProto with no item, no changed made\")\n        pad_size = 0\n        data_padded = data\n    return data_padded, pad_size\n\n\ndef unpad_dataproto(data: \"DataProto\", pad_size):\n    \"\"\"Unpad the data proto with pad_size. i.e. `data[:-pad_size]`\"\"\"\n    if pad_size != 0:\n        data = data[:-pad_size]\n    return data\n\n\ndef union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:\n    \"\"\"Union two tensordicts.\"\"\"\n    assert tensor_dict1.batch_size == tensor_dict2.batch_size, (\n        f\"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}\"\n    )\n    for key in tensor_dict2.keys():\n        if key not in tensor_dict1.keys():\n            tensor_dict1[key] = tensor_dict2[key]\n        else:\n            assert tensor_dict1[key].equal(tensor_dict2[key]), (\n                f\"{key} in tensor_dict1 and tensor_dict2 are not the same object\"\n            )\n\n    return tensor_dict1\n\n\ndef _array_equal(array1: np.ndarray, array2: np.ndarray, visited: set[int]) -> bool:\n    \"\"\"\n    Recursively compares two NumPy arrays for strict equality, with special\n    handling for object-dtype arrays, NaN values, and circular references.\n    This function assumes that the two arguments provided are NumPy arrays.\n\n    Args:\n        array1: The first NumPy array.\n        array2: The second NumPy array.\n\n    Returns:\n        True if the arrays' dtypes, shapes, and all elements are equal.\n    \"\"\"\n    # Check dtype and shape first, as this is the fastest failure path.\n    if array1.dtype != array2.dtype or array1.shape != array2.shape:\n        return False\n\n    # For non-object dtypes, use NumPy's implementation with equal_nan=True.\n    if array1.dtype != \"object\":\n        return np.array_equal(array1, array2, equal_nan=True)\n\n    # For object-dtype arrays, we must recursively compare each element.\n    # We delegate to _deep_equal to handle elements, as they could be any\n    # type, including other nested arrays or NaNs.\n    return all(_deep_equal(x, y, visited) for x, y in zip(array1.flat, array2.flat, strict=False))\n\n\ndef _deep_equal(a: Any, b: Any, visited: set[int]) -> bool:\n    \"\"\"\n    Recursively performs a deep comparison between two Python objects.\n    - Handles NaN values correctly (NaN == NaN evaluates to True).\n    - Handling circular references.\n    - Dispatches to _array_equal if both objects are NumPy arrays.\n    - Otherwise, uses standard '==' comparison.\n    \"\"\"\n    if type(a) is not type(b):\n        return False\n\n    # If we have seen this object ID before on this path, it's a cycle.\n    # Since we already know the types match, we can safely assume this part\n    # of the structure is equal.\n    obj_id = id(a)\n    if obj_id in visited:\n        return True\n\n    visited.add(obj_id)\n\n    # Perform the specific comparison based on type\n    result = False\n    if isinstance(a, float) and math.isnan(a) and math.isnan(b):\n        result = True\n    elif isinstance(a, np.ndarray):\n        # We know b is also an ndarray due to the initial type check\n        result = _array_equal(a, b, visited)\n    else:\n        # Standard equality for all other types\n        result = a == b\n\n    # Clean up the visited set on the way out of the recursion\n    visited.remove(obj_id)\n    return result\n\n\ndef union_numpy_dict(tensor_dict1: dict[str, np.ndarray], tensor_dict2: dict[str, np.ndarray]) -> dict[str, np.ndarray]:\n    for key, val in tensor_dict2.items():\n        if key in tensor_dict1:\n            assert isinstance(tensor_dict2[key], np.ndarray)\n            assert isinstance(tensor_dict1[key], np.ndarray)\n            # to properly deal with nan and object type\n            assert _deep_equal(tensor_dict1[key], tensor_dict2[key], visited=set()), (\n                f\"`{key}` in tensor_dict1 and tensor_dict2 are not the same object.\"\n            )\n        tensor_dict1[key] = val\n\n    return tensor_dict1\n\n\ndef fold_batch_dim(data: \"DataProto\", new_batch_size):\n    \"\"\"\n    Fold a batch dim from [bsz, xxx] into [new_bsz, bsz // new_bsz, xxx]\n    \"\"\"\n    batch_size = data.batch.batch_size[0]\n\n    assert batch_size % new_batch_size == 0\n\n    tensor: TensorDict = data.batch\n    non_tensor = data.non_tensor_batch\n\n    tensor = tensor.view(new_batch_size, -1)\n    tensor.auto_batch_size_(batch_dims=1)\n\n    for key, val in non_tensor.items():\n        non_tensor[key] = np.reshape(val, newshape=(new_batch_size, -1, *val.shape[1:]))\n\n    return type(data)(batch=tensor, non_tensor_batch=non_tensor, meta_info=data.meta_info)\n\n\ndef unfold_batch_dim(data: \"DataProto\", batch_dims=2):\n    \"\"\"\n    Unfold the first n dims as new batch dim\n    \"\"\"\n    tensor: TensorDict = data.batch\n    non_tensor = data.non_tensor_batch\n    tensor.auto_batch_size_(batch_dims=batch_dims)\n    tensor = tensor.view(-1)\n\n    batch_size = tensor.batch_size[0]\n\n    non_tensor_new = {}\n\n    for key, val in non_tensor.items():\n        non_tensor_new[key] = np.reshape(val, newshape=(batch_size, *val.shape[batch_dims:]))\n\n    return type(data)(batch=tensor, non_tensor_batch=non_tensor_new, meta_info=data.meta_info)\n\n\ndef serialize_single_tensor(obj: torch.Tensor) -> tuple[str, tuple[int, ...], int | memoryview]:\n    data = obj.flatten().contiguous().view(torch.uint8).numpy()\n    dtype = str(obj.dtype).removeprefix(\"torch.\")\n    return dtype, obj.shape, data\n\n\ndef serialize_tensordict(batch: TensorDict) -> tuple[tuple[int, ...], Optional[str], dict[str, tuple[str, Any]]]:\n    encoded_items: dict[str, tuple[Any]] = {}\n    for k, v in batch.items():\n        if not v.is_nested:\n            encoded_items[k] = serialize_single_tensor(v)\n        else:\n            layout = str(v.layout).removeprefix(\"torch.\")\n            data = [serialize_single_tensor(tensor) for tensor in v.unbind()]\n            encoded_items[k] = (layout, data)\n\n    batch_size = tuple(batch.batch_size)\n    device = str(batch.device) if batch.device is not None else None\n    return batch_size, device, encoded_items\n\n\ndef deserialize_single_tensor(arr: Any) -> torch.Tensor:\n    dtype, shape, data = arr\n\n    torch_dtype = getattr(torch, dtype)\n    assert isinstance(torch_dtype, torch.dtype)\n\n    buffer = bytearray(data)\n    # Create uint8 array\n    arr = torch.frombuffer(buffer, dtype=torch.uint8)\n    # Convert back to proper shape & type\n    return arr.view(torch_dtype).view(shape)\n\n\ndef deserialize_tensordict(arr: Any) -> TensorDict:\n    batch_size, device, encoded_items = arr\n    decoded_items: dict[str, Any] = {}\n\n    for k, v in encoded_items.items():\n        if len(v) == 3:\n            # decode single tensor\n            decoded_items[k] = deserialize_single_tensor(v)\n        elif len(v) == 2:\n            # decode nested tensor\n            layout, data = v\n            torch_layout = getattr(torch, layout)\n            decoded_items[k] = torch.nested.as_nested_tensor(\n                [deserialize_single_tensor(tensor) for tensor in data], layout=torch_layout\n            )\n        else:\n            raise ValueError(f\"Invalid tensor encoding format, expected length 2 or 3, got {len(v)}\")\n\n    return TensorDict(source=decoded_items, batch_size=batch_size, device=device)\n\n\ndef collate_fn(x: list[\"DataProtoItem\"]):\n    batch = []\n    non_tensor_batch = []\n    for data in x:\n        batch.append(data.batch)\n        non_tensor_batch.append(data.non_tensor_batch)\n    batch = torch.stack(batch).contiguous()\n    non_tensor_batch = list_of_dict_to_dict_of_list(non_tensor_batch)\n    for key, val in non_tensor_batch.items():\n        non_tensor_batch[key] = np.array(val, dtype=object)\n    return DataProto(batch=batch, non_tensor_batch=non_tensor_batch)\n\n\n@dataclass\nclass DataProtoItem:\n    # TODO(zhangchi.usc1992) add consistency check\n    batch: TensorDict = None\n    non_tensor_batch: dict = field(default_factory=dict)\n    meta_info: dict = field(default_factory=dict)\n\n\n@dataclass\nclass DataProto:\n    \"\"\"\n    A DataProto is a data structure that aims to provide a standard protocol for data exchange between functions.\n    It contains a batch (TensorDict) and a meta_info (Dict). The batch is a TensorDict https://pytorch.org/tensordict/.\n    TensorDict allows you to manipulate a dictionary of Tensors like a single Tensor. Ideally, the tensors with the\n    same batch size should be put inside batch.\n    \"\"\"\n\n    batch: TensorDict = None\n    non_tensor_batch: dict = field(default_factory=dict)\n    meta_info: dict = field(default_factory=dict)\n\n    def __post_init__(self):\n        # perform necessary checking\n        self.check_consistency()\n\n    def __len__(self):\n        if self.batch is not None:\n            return self.batch.batch_size[0]\n        elif self.non_tensor_batch is not None and len(self.non_tensor_batch) > 0:\n            random_key = list(self.non_tensor_batch.keys())[0]\n            return self.non_tensor_batch[random_key].shape[0]\n        else:\n            return 0\n\n    def __getitem__(self, item):\n        \"\"\"\n        Enhanced indexing for DataProto objects.\n\n        Args:\n            item: Can be one of:\n                - int: A single index\n                - slice: A slice object (start:stop:step)\n                - list: A list of indices\n                - numpy.ndarray: An array of indices\n                - torch.Tensor: A tensor of indices\n\n        Returns:\n            DataProto: For all indexing types except single integers\n            DataProtoItem: Only for single integer indices\n        \"\"\"\n        # Case 1: Slice object - use the slice method\n        if isinstance(item, slice):\n            return self.slice(item.start, item.stop, item.step)\n\n        # Case 2: List, numpy array, or torch tensor - use sel_idxs\n        elif isinstance(item, list | np.ndarray | torch.Tensor):\n            return self.select_idxs(item)\n\n        # Case 3: Single integer - return DataProtoItem for backward compatibility\n        elif isinstance(item, int | np.integer):\n            tensor_data = self.batch[item] if self.batch is not None else None\n            non_tensor_data = {key: val[item] for key, val in self.non_tensor_batch.items()}\n            return DataProtoItem(batch=tensor_data, non_tensor_batch=non_tensor_data, meta_info=self.meta_info)\n\n        # # Case 4: Unsupported type\n        else:\n            raise TypeError(f\"Indexing with {type(item)} is not supported\")\n\n    def __getstate__(self):\n        if version.parse(tensordict.__version__) >= version.parse(\"0.5.0\") and self.batch is not None:\n            # Check if batch is empty to avoid torch.cat error in consolidate\n            if len(self.batch.keys()) > 0:\n                batch = self.batch.contiguous().consolidate()\n            else:\n                batch = self.batch\n        else:\n            batch = self.batch\n\n        if os.getenv(\"VERL_DATAPROTO_SERIALIZATION_METHOD\") == \"numpy\":\n            if batch is not None:\n                batch = serialize_tensordict(self.batch)\n\n            return (\n                batch,\n                self.non_tensor_batch,\n                self.meta_info,\n            )\n        else:\n            import io\n\n            buffer = io.BytesIO()\n            torch.save(batch, buffer)\n            buffer_bytes = buffer.getvalue()\n            return buffer_bytes, self.non_tensor_batch, self.meta_info\n\n    def __setstate__(self, data):\n        batch_deserialized_bytes, non_tensor_batch, meta_info = data\n\n        if os.getenv(\"VERL_DATAPROTO_SERIALIZATION_METHOD\") == \"numpy\":\n            if batch_deserialized_bytes is not None:\n                self.batch = deserialize_tensordict(batch_deserialized_bytes)\n            else:\n                self.batch = None\n        else:\n            import io\n\n            batch_deserialized = io.BytesIO(initial_bytes=batch_deserialized_bytes)\n            batch = torch.load(\n                batch_deserialized,\n                weights_only=False,\n                map_location=\"cpu\" if not get_torch_device().is_available() else None,\n            )\n            self.batch = batch\n\n        self.non_tensor_batch = non_tensor_batch\n        self.meta_info = meta_info\n\n    def save_to_disk(self, filepath):\n        with open(filepath, \"wb\") as f:\n            pickle.dump(self, f)\n\n    @staticmethod\n    def load_from_disk(filepath) -> \"DataProto\":\n        with open(filepath, \"rb\") as f:\n            data = pickle.load(f)\n            return data\n\n    def print_size(self, prefix=\"\"):\n        size_of_tensordict = 0\n        if self.batch is not None:\n            for _, tensor in self.batch.items():\n                size_of_tensordict += tensor.element_size() * tensor.numel()\n        size_of_numpy_array = 0\n        for _, numpy_array in self.non_tensor_batch.items():\n            size_of_numpy_array += numpy_array.nbytes\n\n        size_of_numpy_array /= 1024**3\n        size_of_tensordict /= 1024**3\n\n        message = f\"Size of tensordict: {size_of_tensordict} GB, size of non_tensor_batch: {size_of_numpy_array} GB\"\n\n        if prefix:\n            message = f\"{prefix}, \" + message\n        print(message)\n\n    def check_consistency(self):\n        \"\"\"Check the consistency of the DataProto. Mainly for batch and non_tensor_batch\n        We expose this function as a public one so that user can call themselves directly\n        \"\"\"\n        if self.batch is not None:\n            assert len(self.batch.batch_size) == 1, \"only support num_batch_dims=1\"\n\n        if self.non_tensor_batch is not None:\n            for key, val in self.non_tensor_batch.items():\n                assert isinstance(val, np.ndarray)\n\n        if self.batch is not None and self.non_tensor_batch is not None and len(self.non_tensor_batch) != 0:\n            # TODO: we can actually lift this restriction if needed\n            assert len(self.batch.batch_size) == 1, \"only support num_batch_dims=1 when non_tensor_batch is not empty.\"\n\n            batch_size = self.batch.batch_size[0]\n            for key, val in self.non_tensor_batch.items():\n                assert isinstance(val, np.ndarray), (\n                    f\"data in the non_tensor_batch must be a numpy.array with dtype=object, but for \"\n                    f\"{key=}, got {type(val)=}\"\n                )\n                assert val.shape[0] == batch_size, (\n                    f\"key {key} length {len(val)} is not equal to batch size {batch_size}\"\n                )\n\n    @classmethod\n    def from_single_dict(cls, data: dict[str, torch.Tensor | np.ndarray], meta_info=None, auto_padding=False):\n        \"\"\"Create a DataProto from a dict of tensors and non_tensors\"\"\"\n        tensors = {}\n        non_tensors = {}\n\n        for key, val in data.items():\n            if isinstance(val, torch.Tensor):\n                tensors[key] = val\n            elif isinstance(val, np.ndarray):\n                non_tensors[key] = val\n            else:\n                raise ValueError(f\"Unsupported type in data {type(val)}\")\n\n        return cls.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info, auto_padding=auto_padding)\n\n    @classmethod\n    def from_dict(\n        cls,\n        tensors: Optional[dict[str, torch.Tensor]] = None,\n        non_tensors=None,\n        meta_info=None,\n        num_batch_dims=1,\n        auto_padding=False,\n    ):\n        \"\"\"Create a DataProto from a dict of tensors. This assumes that\n        1. All the tensor in tensors have the same dim0\n        2. Only dim0 is the batch dim\n        \"\"\"\n\n        assert num_batch_dims > 0, \"num_batch_dims must be greater than zero\"\n        if non_tensors is not None:\n            assert num_batch_dims == 1, \"only support num_batch_dims=1 when non_tensors is not None.\"\n\n        if tensors is None:\n            tensors = {}\n        if meta_info is None:\n            meta_info = {}\n        if non_tensors is None:\n            non_tensors = {}\n\n        assert isinstance(non_tensors, dict)\n\n        # get and check batch size\n        batch_size = None\n        pivot_key = None\n        for key, tensor in tensors.items():\n            if batch_size is None:\n                batch_size = tensor.shape[:num_batch_dims]\n                pivot_key = key\n            else:\n                current_batch = tensor.shape[:num_batch_dims]\n                assert batch_size == current_batch, (\n                    f\"Not all the tensor in tensors have the same batch size with batch_dims={num_batch_dims}. \"\n                    f\"Got {pivot_key} has {batch_size}, {key} has {current_batch}\"\n                )\n\n        for key, val in non_tensors.items():\n            if not isinstance(val, np.ndarray):\n                non_tensors[key] = np.array(val, dtype=object)\n\n        tensor_dict = TensorDict(source=tensors, batch_size=batch_size) if tensors else None\n        if auto_padding:\n            meta_info[DataProtoConfig.auto_padding_key] = True\n        return cls(batch=tensor_dict, non_tensor_batch=non_tensors, meta_info=meta_info)\n\n    @classmethod\n    def from_tensordict(\n        cls,\n        tensor_dict: TensorDict = None,\n        meta_info=None,\n        num_batch_dims=1,\n    ):\n        \"\"\"Create a DataProto from a TensorDict. This assumes that\n        1. All the tensor in tensor_dict have the same dim0\n        2. Only dim0 is the batch dim\n        \"\"\"\n        assert version.parse(tensordict.__version__) >= version.parse(\"0.10.0\"), (\n            \"Build DataProto from TensorDict at least requires tensordict version 0.10.0\"\n        )\n        from tensordict import NonTensorData, NonTensorStack\n\n        assert num_batch_dims > 0, \"num_batch_dims must be greater than zero\"\n        if not all(isinstance(val, torch.Tensor) for val in tensor_dict.values()):\n            assert num_batch_dims == 1, \"only support num_batch_dims=1 when tensor_dict contains non tensor data.\"\n\n        if meta_info is None:\n            meta_info = {}\n        batch = {}\n        non_tensor_batch = {}\n        batch_size = None\n        for key, val in tensor_dict.items():\n            if isinstance(val, torch.Tensor):\n                batch[key] = val\n                if batch_size is None:\n                    batch_size = val.shape[:num_batch_dims]\n            elif isinstance(val, NonTensorStack):\n                non_tensor_batch[key] = np.array([elem.data for elem in val], dtype=object)\n            elif isinstance(val, NonTensorData):\n                meta_info[key] = val.data\n\n        return cls(\n            batch=TensorDict(batch, batch_size=batch_size),\n            non_tensor_batch=non_tensor_batch,\n            meta_info=meta_info,\n        )\n\n    def to(self, device) -> \"DataProto\":\n        \"\"\"move the batch to device\n\n        Args:\n            device (torch.device, str): torch device\n\n        Returns:\n            DataProto: the current DataProto\n\n        \"\"\"\n        if self.batch is not None:\n            self.batch = self.batch.to(device)\n        return self\n\n    def select(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None, deepcopy=False) -> \"DataProto\":\n        \"\"\"Select a subset of the DataProto via batch_keys and meta_info_keys\n\n        Args:\n            batch_keys (list, optional): a list of strings indicating the keys in batch to select\n            meta_info_keys (list, optional): a list of keys indicating the meta info to select\n\n        Returns:\n            DataProto: the DataProto with the selected batch_keys and meta_info_keys\n        \"\"\"\n        # TODO (zhangchi.usc1992) whether to copy\n        if batch_keys is not None:\n            batch_keys = tuple(batch_keys)\n            sub_batch = self.batch.select(*batch_keys)\n        else:\n            sub_batch = self.batch\n\n        if non_tensor_batch_keys is not None:\n            non_tensor_batch = {key: val for key, val in self.non_tensor_batch.items() if key in non_tensor_batch_keys}\n        else:\n            non_tensor_batch = self.non_tensor_batch\n\n        if deepcopy:\n            non_tensor_batch = copy.deepcopy(non_tensor_batch)\n\n        if meta_info_keys is not None:\n            sub_meta_info = {key: val for key, val in self.meta_info.items() if key in meta_info_keys}\n        else:\n            sub_meta_info = self.meta_info\n\n        if deepcopy:\n            sub_meta_info = copy.deepcopy(sub_meta_info)\n\n        return type(self)(batch=sub_batch, non_tensor_batch=non_tensor_batch, meta_info=sub_meta_info)\n\n    def select_idxs(self, idxs):\n        \"\"\"\n        Select specific indices from the DataProto.\n\n        Args:\n            idxs (torch.Tensor or numpy.ndarray or list): Indices to select\n\n        Returns:\n            DataProto: A new DataProto containing only the selected indices\n        \"\"\"\n        if isinstance(idxs, list):\n            idxs = torch.tensor(idxs)\n            if idxs.dtype != torch.bool:\n                idxs = idxs.type(torch.int32)\n\n        if isinstance(idxs, np.ndarray):\n            idxs_np = idxs\n            idxs_torch = torch.from_numpy(idxs)\n        else:  # torch.Tensor\n            idxs_torch = idxs\n            idxs_np = idxs.detach().cpu().numpy()\n\n        batch_size = int(idxs_np.sum()) if idxs_np.dtype == bool else idxs_np.shape[0]\n\n        if self.batch is not None:\n            # Use TensorDict's built-in indexing capabilities\n            selected_batch = TensorDict(\n                source={key: tensor[idxs_torch] for key, tensor in self.batch.items()},\n                batch_size=(batch_size,),\n                device=self.batch.device,\n            )\n        else:\n            selected_batch = None\n\n        selected_non_tensor = {}\n        for key, val in self.non_tensor_batch.items():\n            selected_non_tensor[key] = val[idxs_np]\n\n        return type(self)(batch=selected_batch, non_tensor_batch=selected_non_tensor, meta_info=self.meta_info)\n\n    def slice(self, start=None, end=None, step=None):\n        \"\"\"\n        Slice the DataProto and return a new DataProto object.\n        This is an improved version of direct slicing which returns a DataProtoItem.\n\n        Args:\n            start (int, optional): Start index. Defaults to None (start from beginning).\n            end (int, optional): End index (exclusive). Defaults to None (go to end).\n            step (int, optional): Step size. Defaults to None (step=1).\n\n        Returns:\n            DataProto: A new DataProto containing the sliced data\n\n        Examples:\n            # Using the slice method directly\n            sliced_data = data_proto.slice(10, 20)\n\n            # Using enhanced indexing (returns DataProto)\n            sliced_data = data_proto[10:20]\n            sliced_data = data_proto[::2]  # Every other element\n\n            # Using list indexing (returns DataProto)\n            indices = [1, 5, 10]\n            selected_data = data_proto[indices]\n\n            # Single index still returns DataProtoItem\n            single_item = data_proto[5]\n        \"\"\"\n        # Create a slice object\n        slice_obj = slice(start, end, step)\n\n        # Handle the batch data\n        if self.batch is not None:\n            # Use TensorDict's built-in slicing capabilities\n            sliced_batch = self.batch[slice_obj]\n        else:\n            sliced_batch = None\n\n        # Handle the non-tensor batch data\n        sliced_non_tensor = {}\n        for key, val in self.non_tensor_batch.items():\n            sliced_non_tensor[key] = val[slice_obj]\n\n        # Return a new DataProto object\n        return type(self)(batch=sliced_batch, non_tensor_batch=sliced_non_tensor, meta_info=self.meta_info)\n\n    def pop(self, batch_keys=None, non_tensor_batch_keys=None, meta_info_keys=None) -> \"DataProto\":\n        \"\"\"Pop a subset of the DataProto via `batch_keys` and `meta_info_keys`\n\n        Args:\n            batch_keys (list, optional): a list of strings indicating the keys in batch to pop\n            meta_info_keys (list, optional): a list of keys indicating the meta info to pop\n\n        Returns:\n            DataProto: the DataProto with the poped batch_keys and meta_info_keys\n        \"\"\"\n        if batch_keys is None:\n            batch_keys = []\n        if meta_info_keys is None:\n            meta_info_keys = []\n        if non_tensor_batch_keys is None:\n            non_tensor_batch_keys = []\n\n        tensors = {}\n        # tensor batch\n        for key in batch_keys:\n            assert key in self.batch.keys()\n            tensors[key] = self.batch.pop(key)\n        non_tensors = {}\n        # non tensor batch\n        for key in non_tensor_batch_keys:\n            assert key in self.non_tensor_batch.keys()\n            non_tensors[key] = self.non_tensor_batch.pop(key)\n        meta_info = {}\n        for key in meta_info_keys:\n            assert key in self.meta_info.keys()\n            meta_info[key] = self.meta_info.pop(key)\n        return DataProto.from_dict(tensors=tensors, non_tensors=non_tensors, meta_info=meta_info)\n\n    def rename(self, old_keys=None, new_keys=None) -> \"DataProto\":\n        \"\"\"\n        Note that this function only rename the key in the batch\n        \"\"\"\n\n        def validate_input(keys):\n            if keys is not None:\n                if isinstance(keys, str):\n                    keys = [keys]\n                elif isinstance(keys, list):\n                    pass\n                else:\n                    raise TypeError(f\"keys must be a list or a string, but got {type(keys)}\")\n            return keys\n\n        old_keys = validate_input(old_keys)\n        new_keys = validate_input(new_keys)\n\n        if len(new_keys) != len(old_keys):\n            raise ValueError(\n                f\"new_keys and old_keys must have the same length, but got {len(new_keys)} and {len(old_keys)}\"\n            )\n\n        self.batch.rename_key_(tuple(old_keys), tuple(new_keys))\n\n        return self\n\n    def union(self, other: \"DataProto\") -> \"DataProto\":\n        \"\"\"Union with another DataProto. Union batch and meta_info separately.\n        Throw an error if\n\n        - there are conflict keys in batch and they are not equal\n        - the batch size of two data batch is not the same\n        - there are conflict keys in meta_info and they are not the same.\n\n        Args:\n            other (DataProto): another DataProto to union\n\n        Returns:\n            DataProto: the DataProto after union\n        \"\"\"\n        self.batch = union_tensor_dict(self.batch, other.batch)\n        self.non_tensor_batch = union_numpy_dict(self.non_tensor_batch, other.non_tensor_batch)\n        self.meta_info = union_two_dict(self.meta_info, other.meta_info)\n        return self\n\n    def make_iterator(self, mini_batch_size, epochs, seed=None, dataloader_kwargs=None):\n        r\"\"\"Make an iterator from the DataProto. This is built upon that TensorDict can be used as a normal Pytorch\n        dataset. See https://pytorch.org/tensordict/stable/tutorials/data_fashion for more details.\n\n\n        Args:\n            mini_batch_size (int): mini-batch size when iterating the dataset. We require that\n                ``batch.batch_size[0] % mini_batch_size == 0``.\n            epochs (int): number of epochs when iterating the dataset.\n            dataloader_kwargs (Any): internally, it returns a DataLoader over the batch. The\n                dataloader_kwargs is the kwargs passed to the DataLoader.\n\n        Returns:\n            Iterator: an iterator that yields a mini-batch data at a time. The total number of iteration\n                steps is ``self.batch.batch_size * epochs // mini_batch_size``\n        \"\"\"\n        assert self.batch.batch_size[0] % mini_batch_size == 0, f\"{self.batch.batch_size[0]} % {mini_batch_size} != 0\"\n        # we can directly create a dataloader from TensorDict\n        if dataloader_kwargs is None:\n            dataloader_kwargs = {}\n\n        if seed is not None:\n            generator = torch.Generator()\n            generator.manual_seed(seed)\n        else:\n            generator = None\n\n        assert isinstance(dataloader_kwargs, dict)\n        train_dataloader = DataLoader(\n            dataset=self, batch_size=mini_batch_size, collate_fn=collate_fn, generator=generator, **dataloader_kwargs\n        )\n\n        def get_data():\n            for _ in range(epochs):\n                for d in train_dataloader:\n                    d.meta_info = self.meta_info\n                    yield d\n\n        return iter(get_data())\n\n    def is_padding_enabled(self):\n        \"\"\"\n        Check if padding is enabled for the DataProto.\n        Returns:\n            bool: True if padding is enabled, False otherwise.\n        \"\"\"\n        dataproto_specific_padding = self.meta_info.get(DataProtoConfig.auto_padding_key, False)\n        return dataproto_specific_padding or DataProtoConfig.auto_padding\n\n    def padding(self, padding_size, padding_candidate=\"\"):\n        \"\"\"Pad the DataProto by concating with padding_candidate.repeat(padding_size)\n\n        Args:\n            padding_size (int): the number of repeated padding_candidate\n            padding_candidate: the item to be repeated and appended to the DataProto, only supporting [\"first\", \"last\"]\n        \"\"\"\n        if padding_size == 0:\n            return\n        padding_candidate = self.select_idxs([0 if padding_candidate == \"first\" else len(self) - 1])\n        padding_part = padding_candidate.repeat(padding_size)\n        padded_dp = DataProto.concat([self, padding_part])\n        self.batch = padded_dp.batch\n        self.non_tensor_batch = padded_dp.non_tensor_batch\n\n    def chunk(self, chunks: int) -> list[\"DataProto\"]:\n        \"\"\"Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split.\n\n        Args:\n            chunks (int): the number of chunks to split on dim=0\n\n        Returns:\n            List[DataProto]: a list of DataProto after splitting\n        \"\"\"\n        if not self.is_padding_enabled():\n            assert len(self) % chunks == 0, (\n                f\"only support equal chunk. Got size of DataProto {len(self)} and chunk {chunks}.\"\n            )\n\n        bsz_in_batch = None\n        if self.batch is not None:\n            batch_lst = self.batch.chunk(chunks=chunks, dim=0)\n            bsz_in_batch = np.array([batch.batch_size[0] for batch in batch_lst])\n            chunk_indices = np.cumsum(bsz_in_batch)[:-1]\n        else:\n            batch_lst = [None for _ in range(chunks)]\n\n        non_tensor_batch_lst = [{} for _ in range(chunks)]\n        for key, val in self.non_tensor_batch.items():\n            assert isinstance(val, np.ndarray)\n            if bsz_in_batch is not None:\n                non_tensor_lst = np.array_split(val, chunk_indices.tolist())\n            else:\n                non_tensor_lst = np.array_split(val, chunks)\n            assert len(non_tensor_lst) == chunks\n            for i in range(chunks):\n                non_tensor_batch_lst[i][key] = non_tensor_lst[i]\n\n        output = []\n        for i in range(chunks):\n            output.append(\n                type(self)(batch=batch_lst[i], non_tensor_batch=non_tensor_batch_lst[i], meta_info=self.meta_info)\n            )\n\n        return output\n\n    def split(self, split_size: int) -> list[\"DataProto\"]:\n        \"\"\"Split the batch among dim=0 into chunks. The meta_info is passed to each DataProto after split.\n\n        Args:\n            split_size (int): the size of each split\n\n        Returns:\n            List[DataProto]: a list of DataProto after splitting\n        \"\"\"\n        return [self[i : i + split_size] for i in range(0, len(self), split_size)]\n\n    @staticmethod\n    def concat(data: list[\"DataProto\"]) -> \"DataProto\":\n        \"\"\"Concat a list of DataProto. The batch is concatenated among dim=0.\n        The meta_info is merged, with special handling for metrics from different workers.\n\n        Args:\n            data (List[DataProto]): list of DataProto\n\n        Returns:\n            DataProto: concatenated DataProto\n        \"\"\"\n        batch_lst = []\n        for batch in data:\n            batch_lst.append(batch.batch)\n        new_batch = torch.cat(batch_lst, dim=0) if batch_lst[0] is not None else None\n\n        non_tensor_batch = list_of_dict_to_dict_of_list(list_of_dict=[d.non_tensor_batch for d in data])\n        for key, val in non_tensor_batch.items():\n            non_tensor_batch[key] = np.concatenate(val, axis=0)\n\n        # Merge meta_info with special handling for metrics\n        merged_meta_info = {}\n        if data:\n            # Merge non-metric meta_info and aggregate metrics from all workers.\n            all_metrics = []\n            for d in data:\n                for k, v in d.meta_info.items():\n                    if k == \"metrics\":\n                        if v is not None:\n                            if isinstance(v, list):\n                                all_metrics.extend(v)\n                            else:\n                                all_metrics.append(v)\n                    else:\n                        if k in merged_meta_info:\n                            # Ensure consistency for overlapping non-metric keys\n                            assert merged_meta_info[k] == v, f\"Conflicting values for meta_info key '{k}'\"\n                        else:\n                            merged_meta_info[k] = v\n\n            # Flatten list of dicts to dict of lists for consistent metrics structure\n            if all_metrics:\n                merged_meta_info[\"metrics\"] = list_of_dict_to_dict_of_list(all_metrics)\n\n        cls = type(data[0]) if len(data) > 0 else DataProto\n        return cls(batch=new_batch, non_tensor_batch=non_tensor_batch, meta_info=merged_meta_info)\n\n    def reorder(self, indices):\n        \"\"\"\n        Note that this operation is in-place\n        \"\"\"\n        indices_np = indices.detach().numpy()\n        self.batch = self.batch[indices]\n        self.non_tensor_batch = {key: val[indices_np] for key, val in self.non_tensor_batch.items()}\n\n    def repeat(self, repeat_times=2, interleave=True):\n        \"\"\"\n        Repeat the batch data a specified number of times.\n\n        Args:\n            repeat_times (int): Number of times to repeat the data.\n            interleave (bool): Whether to interleave the repeated data.\n\n        Returns:\n            DataProto: A new DataProto with repeated data.\n        \"\"\"\n        if self.batch is not None:\n            if interleave:\n                # Interleave the data\n                repeated_tensors = {\n                    key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items()\n                }\n            else:\n                # Stack the data\n                repeated_tensors = {\n                    key: tensor.unsqueeze(0).expand(repeat_times, *tensor.shape).reshape(-1, *tensor.shape[1:])\n                    for key, tensor in self.batch.items()\n                }\n\n            repeated_batch = TensorDict(\n                source=repeated_tensors,\n                batch_size=(self.batch.batch_size[0] * repeat_times,),\n            )\n        else:\n            repeated_batch = None\n\n        repeated_non_tensor_batch = {}\n        for key, val in self.non_tensor_batch.items():\n            if interleave:\n                repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0)\n            else:\n                repeated_non_tensor_batch[key] = np.tile(val, (repeat_times,) + (1,) * (val.ndim - 1))\n\n        return type(self)(\n            batch=repeated_batch,\n            non_tensor_batch=repeated_non_tensor_batch,\n            meta_info=self.meta_info,\n        )\n\n    def unfold_column_chunks(self, n_split: int, split_keys: Optional[list[str]] = None):\n        \"\"\"Split along the second dim into `n_split`, unfold it to the first dim (batch dim)\n        Useful in passing grouped tensors that doesn't want to be shuffled in dataset.\n        keys not in split_keys are repeated to match the shape\n        Note that if the `split_keys` is not provided, it will repeat all the keys in the second dim.\n        \"\"\"\n        if self.batch is not None:\n            unfolded_batch = {}\n            for key in self.batch.keys():\n                if key in split_keys if split_keys is not None else False:\n                    shape = list(self.batch[key].shape)\n                    shape[0] = self.batch[key].shape[0] * n_split\n                    shape[1] = self.batch[key].shape[1] // n_split\n                    unfolded_batch[key] = self.batch[key].reshape(*shape)\n                else:\n                    unfolded_batch[key] = torch.repeat_interleave(self.batch[key], n_split, dim=0)\n            # locate the `unfolded_batch` as a TensorDict on the same device as the original batch\n            unfolded_batch = TensorDict(\n                source=unfolded_batch, batch_size=(self.batch.batch_size[0] * n_split,), device=self.batch.device\n            )\n        else:\n            unfolded_batch = None\n\n        repeated_non_tensor_batch = {}\n        for key, val in self.non_tensor_batch.items():\n            if key in split_keys:\n                shape = list(val.shape)\n                shape[0] = val.shape[0] * n_split\n                shape[1] = val.shape[1] // n_split\n                repeated_non_tensor_batch[key] = val.reshape(*shape)\n            else:\n                repeated_non_tensor_batch[key] = np.repeat(val, n_split, axis=0)\n\n        return type(self)(\n            batch=unfolded_batch,\n            non_tensor_batch=repeated_non_tensor_batch,\n            meta_info=self.meta_info,\n        )\n\n    def sample_level_repeat(self, repeat_times):\n        \"\"\"\n        Repeat each row of the batch data a specified number of times.\n\n        Args:\n            repeat_times (torch.tensor, list, tuple, ndarray):  Number of times to repeat the data.\n\n        Returns:\n            DataProto: A new DataProto with repeated data.\n        \"\"\"\n        if isinstance(repeat_times, tuple):\n            repeat_times = list(repeat_times)\n        elif isinstance(repeat_times, torch.Tensor):\n            assert len(repeat_times.shape) == 1\n            repeat_times = repeat_times.tolist()\n        elif isinstance(repeat_times, np.ndarray):\n            assert len(repeat_times.shape) == 1\n            repeat_times = repeat_times.tolist()\n        else:\n            assert isinstance(repeat_times, list), (\n                f\"repeat_times type must be in [list, torch.Tensor, np.ndarray, tuple], got {type(repeat_times)}\"\n            )\n        repeat_times = torch.tensor(repeat_times)\n\n        if self.batch is not None:\n            # Interleave the data\n            repeated_tensors = {\n                key: tensor.repeat_interleave(repeat_times, dim=0) for key, tensor in self.batch.items()\n            }\n\n            repeated_batch = TensorDict(\n                source=repeated_tensors,\n                batch_size=(repeat_times.sum().item(),),\n                device=self.batch.device,\n            )\n        else:\n            repeated_batch = None\n\n        repeated_non_tensor_batch = {}\n        for key, val in self.non_tensor_batch.items():\n            repeated_non_tensor_batch[key] = np.repeat(val, repeat_times, axis=0)\n\n        return type(self)(\n            batch=repeated_batch,\n            non_tensor_batch=repeated_non_tensor_batch,\n            meta_info=self.meta_info,\n        )\n\n    def to_tensordict(self) -> TensorDict:\n        \"\"\"Convert this DataProto to TensorDict. Note that this requires tensordict version at least 0.10\n\n        Returns:\n\n        \"\"\"\n        assert parse_version(tensordict.__version__) >= parse_version(\"0.10\"), (\n            \"Convert DataProto to TensorDict at least requires tensordict version 0.10\"\n        )\n        tensor_batch = self.batch.to_dict()\n        non_tensor_batch = self.non_tensor_batch\n\n        from tensordict.tensorclass import NonTensorData, NonTensorStack\n\n        from verl.utils import tensordict_utils as tu\n\n        common_keys = set(tensor_batch.keys()) & set(non_tensor_batch.keys())\n        assert len(common_keys) == 0, f\"tensor_batch and non_tensor_batch have common keys {common_keys}\"\n\n        for key, val in non_tensor_batch.items():\n            assert isinstance(val, np.ndarray)\n            # Convert to NonTensorStack instead of plain list to handle nested structures\n            tensor_batch[key] = NonTensorStack.from_list([NonTensorData(item) for item in val])\n        output = tu.get_tensordict(tensor_dict=tensor_batch, non_tensor_dict=self.meta_info)\n        return output\n\n    def get_data_info(self) -> str:\n        \"\"\"Return formatted information about stored data with nested type details.\n\n        Returns:\n            str: Formatted string showing tensor details and recursive metadata types\n        \"\"\"\n        info = [\"batch\"]\n\n        for key, tensor in self.batch.items():\n            if hasattr(tensor, \"shape\") and hasattr(tensor, \"dtype\") and hasattr(tensor, \"device\"):\n                info.append(f\"  {key}: {tuple(tensor.shape)} ({tensor.dtype}) {tensor.device}\")\n            elif hasattr(tensor, \"shape\") and hasattr(tensor, \"dtype\"):\n                info.append(f\"  {key}: {tuple(tensor.shape)} ({tensor.dtype})\")\n            else:\n                info.append(f\"  {key}: {type(tensor).__name__}\")\n\n        info.append(\"non_tensor_batch\")\n        for key, array in self.non_tensor_batch.items():\n            info.append(f\"  {key}: ndarray{array.shape} ({array.dtype})\")\n\n        info.append(\"meta_info\")\n        for k, v in self.meta_info.items():\n            type_info = self._get_type_info(v)\n            info.append(f\"  {k}: {type_info}\")\n\n        return \"\\n\".join(info)\n\n    def _get_type_info(self, value):\n        \"\"\"Recursively get type information for nested structures\"\"\"\n        if isinstance(value, list):\n            elem_types = {self._get_type_info(v) for v in value[:3]}\n            return f\"list[{'|'.join(elem_types) if elem_types else '...'}]\"\n        if isinstance(value, tuple):\n            elem_types = [self._get_type_info(v) for v in value]\n            return f\"tuple({', '.join(elem_types)})\"\n        if isinstance(value, dict):\n            if not value:\n                return \"dict\"\n            k, v = next(iter(value.items()))\n            return f\"dict[{self._get_type_info(k)}: {self._get_type_info(v)}]\"\n        if isinstance(value, np.ndarray):\n            return f\"ndarray{value.shape} ({value.dtype})\"\n        return type(value).__name__\n\n\n@dataclass\nclass DataProtoFuture:\n    \"\"\"\n    DataProtoFuture aims to eliminate actual data fetching on driver. By doing so, the driver doesn't have to wait\n    for data so that asynchronous execution becomes possible.\n    DataProtoFuture contains a list of futures from another WorkerGroup of size world_size.\n    - collect_fn is a Callable that reduces the list of futures to a DataProto\n    - dispatch_fn is a Callable that partitions the DataProto into a list of DataProto of size world_size\n        and then select\n\n    Potential issue: we can optimize dispatch_fn(collect_fn) such that only needed data is fetched on destination\n    - DataProtoFuture only supports directly passing from the output of a method to another input. You can't perform any\n    operation on the DataProtoFuture in driver.\n    \"\"\"\n\n    collect_fn: Callable\n    futures: list[ray.ObjectRef]\n    dispatch_fn: Callable = None\n\n    @staticmethod\n    def concat(data: list[ray.ObjectRef]) -> \"DataProtoFuture\":\n        output = DataProtoFuture(collect_fn=DataProto.concat, futures=data)\n        return output\n\n    def chunk(self, chunks: int) -> list[\"DataProtoFuture\"]:\n        from functools import partial\n\n        arg_future_lst = []\n        for i in range(chunks):\n            # note that we can't directly pass i and chunks\n            def dispatch_fn(x, i, chunks):\n                return x.chunk(chunks=chunks)[i]\n\n            arg_future = DataProtoFuture(\n                collect_fn=self.collect_fn, dispatch_fn=partial(dispatch_fn, i=i, chunks=chunks), futures=self.futures\n            )\n            arg_future_lst.append(arg_future)\n        return arg_future_lst\n\n    def get(self):\n        output = ray.get(self.futures)  # dp_size.\n        for o in output:\n            assert isinstance(o, DataProto | TensorDict)\n\n        if isinstance(output[0], DataProto):\n            output = DataProto.concat(output)  # select dp, concat\n        elif isinstance(output[0], TensorDict):\n            from verl.utils.tensordict_utils import concat_tensordict\n\n            output = concat_tensordict(output)\n        else:\n            raise TypeError(f\"Unknown type {type(o[0])} in DataProtoFuture\")\n\n        if self.dispatch_fn is not None:\n            output = self.dispatch_fn(output)  # split in batch dim, select using dp\n        return output\n\n\nclass BatchData:\n    \"\"\"Uniform dispatch wrapper for batch data operations.\n\n    All type-specific logic (isinstance checks) is centralized here so that\n    callers (e.g. decorator.py) never need to branch on the concrete data type.\n\n    Usage::\n\n        # chunk a single data item into N pieces\n        chunks = BatchData(arg).chunk(chunks=N)\n\n        # concat a list of data items into one\n        merged = BatchData(output_list).concat()\n\n        # validate before dispatching\n        assert BatchData(arg).is_chunkable()\n        assert BatchData(output_list).is_concatable()\n    \"\"\"\n\n    _CHUNKABLE_TYPES = (TensorDict,)  # lazily extended with DataProto etc.\n    _CONCATABLE_TYPES = (TensorDict,)\n\n    def __init__(self, data):\n        self._data = data\n\n    # ---- validation ----------------------------------------------------------\n\n    def is_chunkable(self) -> bool:\n        \"\"\"Return True if the wrapped data supports chunk dispatch.\"\"\"\n        return isinstance(self._data, self._chunkable_types())\n\n    def is_concatable(self) -> bool:\n        \"\"\"Return True if the wrapped list of data supports concat collect.\"\"\"\n        data = self._data\n        if not isinstance(data, list | tuple) or len(data) == 0:\n            return False\n        return isinstance(data[0], self._concatable_types())\n\n    # ---- operations ----------------------------------------------------------\n\n    def chunk(self, chunks: int):\n        \"\"\"Split the wrapped data into *chunks* pieces along the batch dim.\n\n        Returns a tuple/list of the **original data types** (not BatchData).\n        \"\"\"\n        data = self._data\n        if isinstance(data, TensorDict):\n            from verl.utils.tensordict_utils import chunk_tensordict, contiguous\n\n            raw_chunks = chunk_tensordict(data, chunks)\n            return tuple(contiguous(val).consolidate() for val in raw_chunks)\n        # DataProto, DataProtoFuture, BatchMeta all expose .chunk()\n        return data.chunk(chunks=chunks)\n\n    def concat(self):\n        \"\"\"Concat the wrapped list of data items into a single result.\n\n        Returns the **original data type** (not BatchData).\n        \"\"\"\n        data = self._data\n        if not data:\n            raise ValueError(\"Cannot concatenate an empty list of data items.\")\n        sample = data[0]\n        if isinstance(sample, ray.ObjectRef):\n            return DataProtoFuture.concat(data)\n        if isinstance(sample, TensorDict):\n            from verl.utils.tensordict_utils import concat_tensordict\n\n            return concat_tensordict(data)\n        # DataProto, BatchMeta expose .concat() as classmethod / staticmethod\n        return type(sample).concat(data)\n\n    # ---- helpers (lazy type tuples to avoid import-order issues) -------------\n\n    @classmethod\n    def _chunkable_types(cls):\n        return (DataProto, DataProtoFuture, TensorDict)\n\n    @classmethod\n    def _concatable_types(cls):\n        return (DataProto, ray.ObjectRef, TensorDict)\n\n\ndef all_gather_data_proto(data: DataProto, process_group):\n    # Note that this is an inplace operator just like torch.distributed.all_gather\n    group_size = torch.distributed.get_world_size(group=process_group)\n    assert isinstance(data, DataProto)\n    prev_device = data.batch.device\n    data = data.to(get_device_id())\n    data.batch = allgather_dict_tensors(data.batch.contiguous(), size=group_size, group=process_group, dim=0)\n    data = data.to(prev_device)\n    # all gather non_tensor_batch\n    all_non_tensor_batch = [None for _ in range(group_size)]\n    torch.distributed.all_gather_object(all_non_tensor_batch, data.non_tensor_batch, group=process_group)\n    data.non_tensor_batch = {k: np.concatenate([d[k] for d in all_non_tensor_batch]) for k in data.non_tensor_batch}\n"
  },
  {
    "path": "verl/py.typed",
    "content": ""
  },
  {
    "path": "verl/single_controller/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\n\nfrom . import base\nfrom .base import *\n\nversion_folder = os.path.dirname(os.path.join(os.path.abspath(__file__)))\n\n# Note(haibin.lin): single_controller.__version__ is deprecated\nwith open(os.path.join(os.path.join(version_folder, os.pardir), \"version/version\")) as f:\n    __version__ = f.read().strip()\n\n\n__all__ = base.__all__\n"
  },
  {
    "path": "verl/single_controller/base/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .worker import Worker\nfrom .worker_group import ClassWithInitArgs, ResourcePool, WorkerGroup\n\n__all__ = [\"Worker\", \"WorkerGroup\", \"ClassWithInitArgs\", \"ResourcePool\"]\n"
  },
  {
    "path": "verl/single_controller/base/decorator.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nfrom functools import partial, wraps\nfrom types import FunctionType\n\nfrom verl.protocol import DataProtoFuture, _padding_size_key\nfrom verl.utils.py_functional import DynamicEnum\n\n# here we add a magic number of avoid user-defined function already have this attribute\nMAGIC_ATTR = \"attrs_3141562937\"\n\n\nclass Dispatch(DynamicEnum):\n    \"\"\"Enum class defining different dispatch modes for distributed computation.\n\n    Each mode represents a specific strategy for distributing data across\n    different ranks in a distributed system. The modes are used to control\n    how data is partitioned and processed across different worker groups.\n    \"\"\"\n\n    _registry = {}\n    _next_value = 0\n\n\ndef init_predefined_dispatch_mode():\n    Dispatch.register(\"RANK_ZERO\")\n    Dispatch.register(\"ONE_TO_ALL\")\n    Dispatch.register(\"ALL_TO_ALL\")\n    Dispatch.register(\"DP_COMPUTE\")\n    Dispatch.register(\"DP_COMPUTE_PROTO\")\n    Dispatch.register(\"DP_COMPUTE_PROTO_WITH_FUNC\")\n    Dispatch.register(\"DP_COMPUTE_METRIC\")\n    # This is a special dispatch mode for vllm ExternalRayDistributedExecutor\n    Dispatch.register(\"DIRECT_ROLLOUT_METHOD\")\n\n\nclass Execute(DynamicEnum):\n    \"\"\"Enum class defining different execution modes for distributed computation.\n\n    These modes control how a function should be executed across different ranks\n    in a distributed system.\n    \"\"\"\n\n    _registry = {}\n    _next_value = 0\n\n\ndef init_predefined_execute_mode():\n    Execute.register(\"ALL\")\n    Execute.register(\"RANK_ZERO\")\n\n\n# Initialize the two Dynamic Enum Classes\ninit_predefined_dispatch_mode()\ninit_predefined_execute_mode()\n\n\ndef _split_args_kwargs_data_proto(chunks, *args, **kwargs):\n    from verl.protocol import BatchData\n\n    splitted_args = []\n    for arg in args:\n        assert BatchData(arg).is_chunkable(), f\"arg of type {type(arg)} is not chunkable\"\n        chunked_arg = BatchData(arg).chunk(chunks=chunks)\n        assert len(chunked_arg) == chunks\n        splitted_args.append(chunked_arg)\n\n    splitted_kwargs = {}\n    for key, val in kwargs.items():\n        assert BatchData(val).is_chunkable(), f\"kwarg '{key}' of type {type(val)} is not chunkable\"\n        chunked_kwarg = BatchData(val).chunk(chunks=chunks)\n        assert len(chunked_kwarg) == chunks\n        splitted_kwargs[key] = chunked_kwarg\n\n    return splitted_args, splitted_kwargs\n\n\ndef _split_args_kwargs_data_proto_with_auto_padding(chunks, *args, **kwargs):\n    from verl.protocol import DataProto, DataProtoFuture\n\n    data_proto_len = None\n    padding_size = None\n\n    def _padding_and_split_data(obj, chunks):\n        nonlocal data_proto_len, padding_size\n        assert isinstance(obj, DataProto | DataProtoFuture)\n        if isinstance(obj, DataProto) and obj.is_padding_enabled():\n            # for padding, we only support DataProto with same length\n            if data_proto_len is None:\n                data_proto_len = len(obj)\n                padding_size = (chunks - (data_proto_len % chunks)) if (data_proto_len % chunks > 0) else 0\n            else:\n                assert data_proto_len == len(obj), (\n                    f\"expecting all arg share same length of {data_proto_len}, but got {len(obj)}\"\n                )\n            obj.padding(padding_size=padding_size)\n        return obj.chunk(chunks=chunks)\n\n    splitted_args = [_padding_and_split_data(arg, chunks) for arg in args]\n    splitted_kwargs = {key: _padding_and_split_data(val, chunks) for key, val in kwargs.items()}\n    if padding_size is not None:\n        splitted_kwargs[_padding_size_key] = padding_size\n\n    return splitted_args, splitted_kwargs\n\n\ndef dispatch_one_to_all(worker_group, *args, **kwargs):\n    args = tuple([arg] * worker_group.world_size for arg in args)\n    kwargs = {k: [v] * worker_group.world_size for k, v in kwargs.items()}\n    return args, kwargs\n\n\ndef dummy_direct_rollout_call(worker_group, *args, **kwargs):\n    raise NotImplementedError(\"Direct rollout call is forbidden.\")\n\n\ndef dispatch_all_to_all(worker_group, *args, **kwargs):\n    return args, kwargs\n\n\ndef collect_all_to_all(worker_group, output):\n    return output\n\n\ndef _concat_data_proto_or_future(output: list):\n    from verl.protocol import BatchData\n\n    # make sure all the elements in output has the same type\n    for o in output:\n        assert type(o) is type(output[0])\n\n    return BatchData(output).concat()\n\n\ndef dispatch_dp_compute(worker_group, *args, **kwargs):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n    for arg in args:\n        assert isinstance(arg, tuple | list) and len(arg) == worker_group.world_size\n    for k, v in kwargs.items():\n        assert isinstance(v, tuple | list) and len(v) == worker_group.world_size\n    return args, kwargs\n\n\ndef collect_dp_compute(worker_group, output):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n    assert len(output) == worker_group.world_size\n    return output\n\n\ndef dispatch_dp_compute_data_proto(worker_group, *args, **kwargs):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n    # Note: enable auto padding for dp compute DatapProto\n    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto_with_auto_padding(\n        worker_group.world_size,\n        *args,\n        **kwargs,\n    )\n    return splitted_args, splitted_kwargs\n\n\ndef dispatch_dp_compute_data_proto_with_func(worker_group, *args, **kwargs):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n    assert isinstance(args[0], FunctionType)  # NOTE: The first one args is a function!\n\n    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(worker_group.world_size, *args[1:], **kwargs)\n    splitted_args_with_func = [[args[0]] * worker_group.world_size] + splitted_args\n    return splitted_args_with_func, splitted_kwargs\n\n\ndef collect_dp_compute_data_proto(worker_group, output):\n    from verl.protocol import BatchData\n\n    assert BatchData(output).is_concatable(), (\n        f\"expecting concatable output, but got element type {type(output[0]) if output else 'empty'}\"\n    )\n\n    output = collect_dp_compute(worker_group, output)\n    return _concat_data_proto_or_future(output)\n\n\ndef dispatch_nd_compute(dp_rank_mapping: list[int], dp_size, worker_group, *args, **kwargs):\n    import os\n\n    from verl.single_controller.base.worker_group import WorkerGroup\n    from verl.utils.ray_utils import parallel_put\n\n    assert isinstance(worker_group, WorkerGroup)\n\n    max_workers = max(1, min(len(args[0]), os.cpu_count()))\n\n    args = [parallel_put(arg, max_workers=max_workers) for arg in args]\n    kwargs = {k: parallel_put(v, max_workers=max_workers) for k, v in kwargs.items()}\n\n    all_args = []\n    for arg in args:\n        assert isinstance(arg, tuple | list) and len(arg) == dp_size\n        transformed_args = []\n        for i in range(worker_group.world_size):\n            local_dp_rank = dp_rank_mapping[i]\n            transformed_args.append(arg[local_dp_rank])\n        all_args.append(transformed_args)\n    all_args = tuple(all_args)\n\n    all_kwargs = {}\n    for k, v in kwargs.items():\n        assert isinstance(v, tuple | list) and len(v) == dp_size\n        transformed_v = []\n        for i in range(worker_group.world_size):\n            local_dp_rank = dp_rank_mapping[i]\n            transformed_v.append(v[local_dp_rank])\n        all_kwargs[k] = transformed_v\n    return all_args, all_kwargs\n\n\ndef collect_nd_compute(collect_mask: list[bool], worker_group, output):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n    assert len(output) == worker_group.world_size\n\n    output_in_dp = []\n    for global_rank in range(worker_group.world_size):\n        collect_dp_rank = collect_mask[global_rank]\n        if collect_dp_rank:\n            output_in_dp.append(output[global_rank])\n    return output_in_dp\n\n\ndef dispatch_nd_compute_dataproto(dp_rank_mapping: list[int], dp_size, worker_group, *args, **kwargs):\n    splitted_args, splitted_kwargs = _split_args_kwargs_data_proto(dp_size, *args, **kwargs)\n    return dispatch_nd_compute(dp_rank_mapping, dp_size, worker_group, *splitted_args, **splitted_kwargs)\n\n\ndef collect_nd_compute_dataproto(collect_mask: list[bool], worker_group, output):\n    output = collect_nd_compute(collect_mask, worker_group, output)\n\n    from verl.protocol import BatchData\n\n    assert BatchData(output).is_concatable(), (\n        f\"expecting concatable output, but got element type {type(output[0]) if output else 'empty'}\"\n    )\n    return _concat_data_proto_or_future(output)\n\n\ndef dispatch_lazy_compute_data_proto(mesh_name, worker_group, *args, **kwargs):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n\n    # query dispatch info of the worker group\n    if mesh_name not in worker_group._dispatch_info:\n        worker_group._dispatch_info[mesh_name] = worker_group._query_dispatch_info(mesh_name)\n        assert len(worker_group._dispatch_info[mesh_name]) == worker_group.world_size\n\n    dp_rank_mapping = worker_group._dispatch_info[mesh_name]\n    # perform dispatch\n    dp_size = max(dp_rank_mapping) + 1\n    return dispatch_nd_compute_dataproto(dp_rank_mapping, dp_size, worker_group, *args, **kwargs)\n\n\ndef collect_lazy_compute_data_proto(mesh_name, worker_group, *args, **kwargs):\n    from verl.single_controller.base.worker_group import WorkerGroup\n\n    assert isinstance(worker_group, WorkerGroup)\n\n    # the dispatch info is stored in the worker group\n    assert mesh_name in worker_group._dispatch_info\n\n    if mesh_name not in worker_group._collect_info:\n        worker_group._collect_info[mesh_name] = worker_group._query_collect_info(mesh_name)\n        assert len(worker_group._collect_info[mesh_name]) == worker_group.world_size\n\n    # a boolean of whether the dp_rank is used for collect\n    collect_mask = worker_group._collect_info[mesh_name]\n    # perform dispatch\n    return collect_nd_compute_dataproto(collect_mask, worker_group, *args, **kwargs)\n\n\ndef make_nd_compute_dataproto_dispatch_fn(mesh_name):\n    return {\n        \"dispatch_fn\": partial(dispatch_lazy_compute_data_proto, mesh_name),\n        \"collect_fn\": partial(collect_lazy_compute_data_proto, mesh_name),\n    }\n\n\n# Global registry for dispatch mode.\nDISPATCH_MODE_FN_REGISTRY = {\n    Dispatch.ONE_TO_ALL: {\n        \"dispatch_fn\": dispatch_one_to_all,\n        \"collect_fn\": collect_all_to_all,\n    },\n    Dispatch.ALL_TO_ALL: {\n        \"dispatch_fn\": dispatch_all_to_all,\n        \"collect_fn\": collect_all_to_all,\n    },\n    Dispatch.DP_COMPUTE: {\"dispatch_fn\": dispatch_dp_compute, \"collect_fn\": collect_dp_compute},\n    Dispatch.DP_COMPUTE_PROTO: {\n        \"dispatch_fn\": dispatch_dp_compute_data_proto,\n        \"collect_fn\": collect_dp_compute_data_proto,\n    },\n    Dispatch.DP_COMPUTE_PROTO_WITH_FUNC: {\n        \"dispatch_fn\": dispatch_dp_compute_data_proto_with_func,\n        \"collect_fn\": collect_dp_compute_data_proto,\n    },\n    Dispatch.DP_COMPUTE_METRIC: {\"dispatch_fn\": dispatch_dp_compute_data_proto, \"collect_fn\": collect_dp_compute},\n    Dispatch.DIRECT_ROLLOUT_METHOD: {\n        \"dispatch_fn\": dummy_direct_rollout_call,\n        \"collect_fn\": dummy_direct_rollout_call,\n    },\n}\n\n\ndef get_predefined_dispatch_fn(dispatch_mode):\n    return DISPATCH_MODE_FN_REGISTRY[dispatch_mode]\n\n\ndef register_dispatch_mode(dispatch_mode_name, dispatch_fn, collect_fn):\n    \"\"\"\n    Register a new dispatch mode.\n    \"\"\"\n    dispatch_mode = Dispatch.register(dispatch_mode_name)\n    _check_dispatch_mode(dispatch_mode)\n    assert dispatch_mode not in DISPATCH_MODE_FN_REGISTRY, f\"dispatch_mode_name {dispatch_mode_name} already exists\"\n    DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {\"dispatch_fn\": dispatch_fn, \"collect_fn\": collect_fn}\n\n\ndef update_dispatch_mode(dispatch_mode, dispatch_fn, collect_fn):\n    \"\"\"\n    Update the dispatch mode.\n    \"\"\"\n    _check_dispatch_mode(dispatch_mode)\n    assert dispatch_mode in DISPATCH_MODE_FN_REGISTRY, f\"dispatch_mode {dispatch_mode} not found\"\n    DISPATCH_MODE_FN_REGISTRY[dispatch_mode] = {\"dispatch_fn\": dispatch_fn, \"collect_fn\": collect_fn}\n\n\ndef get_predefined_execute_fn(execute_mode):\n    \"\"\"\n    Note that here we only asks execute_all and execute_rank_zero to be implemented\n    Leave the choice of how these two functions handle argument 'blocking' to users\n    \"\"\"\n    predefined_execute_mode_fn = {\n        Execute.ALL: {\"execute_fn_name\": \"execute_all\"},\n        Execute.RANK_ZERO: {\"execute_fn_name\": \"execute_rank_zero\"},\n    }\n    return predefined_execute_mode_fn[execute_mode]\n\n\ndef _check_dispatch_mode(dispatch_mode):\n    assert isinstance(dispatch_mode, Dispatch | dict), (\n        f\"dispatch_mode must be a Dispatch or a Dict. Got {dispatch_mode}\"\n    )\n    if isinstance(dispatch_mode, dict):\n        necessary_keys = [\"dispatch_fn\", \"collect_fn\"]\n        for key in necessary_keys:\n            assert key in dispatch_mode, f\"key {key} should be in dispatch_mode if it is a dictionary\"\n\n\ndef _check_execute_mode(execute_mode):\n    assert isinstance(execute_mode, Execute), f\"execute_mode must be a Execute. Got {execute_mode}\"\n\n\ndef _materialize_futures(*args, **kwargs):\n    new_args = []\n    for arg in args:\n        if isinstance(arg, DataProtoFuture):\n            arg = arg.get()\n        # add more type to materialize\n        new_args.append(arg)\n    for k, v in kwargs.items():\n        if isinstance(v, DataProtoFuture):\n            kwargs[k] = v.get()\n\n    new_args = tuple(new_args)\n    return new_args, kwargs\n\n\ndef register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.ALL, blocking=True, materialize_futures=True):\n    \"\"\"Register a function with distributed execution configuration.\n\n    This decorator registers a function with specific dispatch and execution modes\n    for distributed computation. It handles both synchronous and asynchronous\n    functions, and optionally materializes futures before execution.\n\n    Args:\n        dispatch_mode:\n            Dispatch mode for computation distribution. Default: Dispatch.ALL_TO_ALL.\n        execute_mode:\n            Execute mode for computation distribution. Default: Execute.ALL.\n        blocking:\n            Whether the execution should be blocking. Defaults to True.\n        materialize_futures:\n            Whether to materialize the data before dispatching. Defaults to True.\n\n    Returns:\n        A decorator that wraps the original function with distributed execution\n        configuration.\n    \"\"\"\n\n    _check_dispatch_mode(dispatch_mode=dispatch_mode)\n    _check_execute_mode(execute_mode=execute_mode)\n\n    def decorator(func):\n        @wraps(func)\n        def inner(*args, **kwargs):\n            if materialize_futures:\n                args, kwargs = _materialize_futures(*args, **kwargs)\n            return func(*args, **kwargs)\n\n        @wraps(func)\n        async def async_inner(*args, **kwargs):\n            if materialize_futures:\n                args, kwargs = _materialize_futures(*args, **kwargs)\n            return await func(*args, **kwargs)\n\n        wrapper = async_inner if inspect.iscoroutinefunction(func) else inner\n        attrs = {\"dispatch_mode\": dispatch_mode, \"execute_mode\": execute_mode, \"blocking\": blocking}\n        setattr(wrapper, MAGIC_ATTR, attrs)\n        return wrapper\n\n    return decorator\n"
  },
  {
    "path": "verl/single_controller/base/worker.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nthe class for Worker\n\"\"\"\n\nimport os\nimport socket\nimport warnings\nfrom dataclasses import dataclass\n\nimport ray\n\nfrom verl.utils.device import (\n    get_torch_device,\n    get_visible_devices_keyword,\n    is_npu_available,\n)\n\nfrom .decorator import Dispatch, Execute, register\n\n\n@dataclass\nclass DistRankInfo:\n    tp_rank: int\n    dp_rank: int\n    pp_rank: int\n    cp_rank: int\n\n\n@dataclass\nclass DistGlobalInfo:\n    tp_size: int\n    dp_size: int\n    pp_size: int\n    cp_size: int\n\n\nclass WorkerHelper:\n    @staticmethod\n    def _get_node_ip():\n        if os.getenv(\"WG_BACKEND\", None) == \"ray\":\n            return ray.util.get_node_ip_address()\n        else:\n            raise NotImplementedError(\"WG_BACKEND now just support ray mode.\")\n\n    @staticmethod\n    def _get_free_port():\n        with socket.socket() as sock:\n            sock.bind((\"\", 0))\n            return sock.getsockname()[1]\n\n    def get_availale_master_addr_port(self):\n        warnings.warn(\n            \"This function is deprecated due to typo in name; Please use `get_available_master_addr_port` instead\",\n            stacklevel=2,\n        )\n        return self.get_available_master_addr_port()\n\n    def get_available_master_addr_port(self):\n        return self._get_node_ip().strip(\"[]\"), str(self._get_free_port())\n\n\n# we assume that in each WorkerGroup, there is a Master Worker\nclass Worker(WorkerHelper):\n    \"\"\"A distributed worker that handles initialization and configuration for distributed training.\n\n    This class manages worker initialization, configuration, and provides methods for executing\n    distributed operations. It handles communication settings, device configuration, and worker\n    metadata management.\n    \"\"\"\n\n    fused_worker_attr_name = \"fused_worker_dict\"\n\n    def _register_dispatch_collect_info(self, mesh_name: str, dp_rank: int, is_collect: bool):\n        \"\"\"Register the dp_rank for a given mesh name. This function is meant to be called by the worker\n\n        Args:\n            mesh_name (str):\n                Name of the mesh to register dp_rank for.\n            dp_rank (int):\n                dp_rank to register for the given mesh name.\n            is_collect (bool):\n                Whether the dp_rank is used for collect.\n        \"\"\"\n        if mesh_name in self.__dispatch_dp_rank or mesh_name in self.__collect_dp_rank:\n            raise ValueError(f\"mesh_name {mesh_name} has been registered\")\n        self.__dispatch_dp_rank[mesh_name] = dp_rank\n        self.__collect_dp_rank[mesh_name] = is_collect\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def _query_dispatch_info(self, mesh_name: str):\n        \"\"\"Query the dispatch info for a given mesh name.\n\n        Args:\n            mesh_name (str):\n                Name of the mesh to query dispatch info for.\n\n        Returns:\n            int:\n                The dp_rank for the given mesh name.\n        \"\"\"\n        assert mesh_name in self.__dispatch_dp_rank, f\"{mesh_name} is not registered in {self.__class__.__name__}\"\n        # note that each rank store its own dp_rank\n        return self.__dispatch_dp_rank[mesh_name]\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def _query_collect_info(self, mesh_name: str):\n        return self.query_collect_info(mesh_name)\n\n    def query_collect_info(self, mesh_name: str):\n        \"\"\"Query the collect info for a given mesh name.\n\n        Args:\n            mesh_name (str):\n                Name of the mesh to query collect info for.\n\n        Returns:\n            bool:\n                Whether the dp_rank is used for collect.\n        \"\"\"\n        assert mesh_name in self.__collect_dp_rank, f\"{mesh_name} is not registered in {self.__class__.__name__}\"\n        return self.__collect_dp_rank[mesh_name]\n\n    def get_dispatch_collect(self):\n        \"\"\"Get all registered dispatch and collect dp_ranks.\n\n        Returns:\n            dict[str, int]:\n                A dictionary mapping mesh names to their dispatch dp_ranks.\n            dict[str, bool]:\n                A dictionary mapping mesh names to whether they are used for collect.\n        \"\"\"\n        return {\"dispatch_dp_rank\": self.__dispatch_dp_rank, \"collect_dp_rank\": self.__collect_dp_rank}\n\n    def set_dispatch_collect(self, mesh_name: str, dispatch_dp_rank: dict[str, int], collect_dp_rank: dict[str, bool]):\n        \"\"\"Set the dispatch and collect dp_ranks for all registered meshes.\n\n        Args:\n            mesh_name (str): Mesh name to set dispatch and collect dp_ranks for.\n            dispatch_dp_rank (dict[str, int]):\n                A dictionary mapping mesh names to their dispatch dp_ranks.\n            collect_dp_rank (dict[str, bool]):\n                A dictionary mapping mesh names to whether they are used for collect.\n        \"\"\"\n        assert mesh_name not in self.__dispatch_dp_rank, (\n            f\"{mesh_name} is already registered, {self.__dispatch_dp_rank.keys()}\"\n        )\n        assert mesh_name not in self.__collect_dp_rank, (\n            f\"{mesh_name} is already registered, {self.__collect_dp_rank.keys()}\"\n        )\n        for dp_rank in dispatch_dp_rank.values():\n            self.__dispatch_dp_rank[mesh_name] = dp_rank\n        for is_collect in collect_dp_rank.values():\n            self.__collect_dp_rank[mesh_name] = is_collect\n\n    @classmethod\n    def env_keys(cls):\n        \"\"\"The keys of the environment variables that are used to configure the Worker.\"\"\"\n        return [\n            \"WORLD_SIZE\",\n            \"RANK\",\n            \"LOCAL_WORLD_SIZE\",\n            \"LOCAL_RANK\",\n            \"MASTER_ADDR\",\n            \"MASTER_PORT\",\n            get_visible_devices_keyword().upper(),\n        ]\n\n    def __init__(self, cuda_visible_devices=None) -> None:\n        \"\"\"Initialize the worker with environment settings and device configuration.\n\n        Args:\n            cuda_visible_devices (str, optional):\n                CUDA visible devices configuration. Defaults to None.\n        \"\"\"\n        # construct a meta from environment variable. Note that the import must be inside the class because\n        # it is executed remotely\n        import os\n\n        self._setup_env_cuda_visible_devices()\n\n        world_size = int(os.environ[\"WORLD_SIZE\"])\n        rank = int(os.environ[\"RANK\"])\n        self._rank = rank\n        self._world_size = world_size\n\n        master_addr = os.environ[\"MASTER_ADDR\"]\n        master_port = os.environ[\"MASTER_PORT\"]\n\n        local_world_size = int(os.getenv(\"LOCAL_WORLD_SIZE\", \"1\"))\n        local_rank = int(os.getenv(\"LOCAL_RANK\", \"0\"))\n\n        store = {\n            \"_world_size\": world_size,\n            \"_rank\": rank,\n            \"_local_world_size\": local_world_size,\n            \"_local_rank\": local_rank,\n            \"_master_addr\": master_addr,\n            \"_master_port\": master_port,\n        }\n        if cuda_visible_devices is not None:\n            store[f\"_{get_visible_devices_keyword()}\".lower()] = cuda_visible_devices\n\n        self._configure_with_store(store=store)\n\n        self.fused_worker_dict = {}\n        self.__dispatch_dp_rank = {}\n        self.__collect_dp_rank = {}\n\n    def get_fused_worker_by_name(self, worker_name: str):\n        \"\"\"Get a fused worker by its name.\n\n        Args:\n            worker_name (str):\n                Name of the worker to retrieve\n        \"\"\"\n        return self.fused_worker_dict.get(worker_name, None)\n\n    def _setup_env_cuda_visible_devices(self):\n        from verl.utils.ray_utils import ray_noset_visible_devices\n\n        is_ray_noset_visible_devices = ray_noset_visible_devices()\n\n        # Prevent use of clashing `{CUDA/HIP/ROCR}_VISIBLE_DEVICES``\n        rocr_val = os.environ.get(\"ROCR_VISIBLE_DEVICES\", None)\n        hip_val = os.environ.get(\"HIP_VISIBLE_DEVICES\", None)\n        cuda_val = os.environ.get(\"CUDA_VISIBLE_DEVICES\", None)\n        if hip_val:\n            # Switch the use of HIP_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES for consistency.\n            # Make sure that the HIP_VISIBLE_DEVICES is set to the same value as CUDA_VISIBLE_DEVICES\n            # at this point.\n            val = os.environ.pop(\"HIP_VISIBLE_DEVICES\")\n            hip_val = None\n            if cuda_val:\n                assert val == cuda_val, (\n                    f\"Please use the same HIP_VISIBLE_DEVICES or CUDA_VISIBLE_DEVICES, inconsistant values \"\n                    f\"found: {val} and {cuda_val}.\"\n                )\n            else:\n                cuda_val = val\n                os.environ[\"CUDA_VISIBLE_DEVICES\"] = val\n                # os.environ[\"HIP_VISIBLE_DEVICES\"] = val\n\n        if rocr_val:\n            # You must take care if both HIP/CUDA and ROCR env vars are set as they have\n            # different meanings. Both env vars accept either a list of ints or a\n            # list of UUIDs. The ROCR env var is processed first which then reduces\n            # the number of GPUs that HIP can select from.\n            # https://github.com/pytorch/pytorch/pull/144026\n            # To avoid the complexity of this, we simply gives out error if both are set\n            # (Also to keep consistency with ray's practice with 2.45.0).\n            # Otherwise, we will set ROCR_VISIBLE_DEVICES to CUDA_VISIBLE_DEVICES\n            # and remove ROCR_VISIBLE_DEVICES.\n            if cuda_val:\n                raise ValueError(\"Please don't set ROCR_VISIBLE_DEVICES when HIP/CUDA_VISIBLE_DEVICES is set.\")\n\n            cuda_val = os.environ.pop(\"ROCR_VISIBLE_DEVICES\")\n            os.environ[\"CUDA_VISIBLE_DEVICES\"] = cuda_val\n            rocr_val = None\n\n        if is_ray_noset_visible_devices:\n            # NOTE: Ray will automatically set the *_VISIBLE_DEVICES\n            # environment variable for each actor, unless\n            # RAY_EXPERIMENTAL_NOSET_*_VISIBLE_DEVICES is set,\n            # so we need to set local rank when the flag is set.\n            device_name = \"NPU\" if is_npu_available else \"GPU\"\n            local_rank = ray.get_runtime_context().get_accelerator_ids()[device_name][0]\n            os.environ[\"LOCAL_RANK\"] = local_rank\n            get_torch_device().set_device(int(local_rank))\n\n    def _configure_with_store(self, store: dict):\n        \"\"\"\n        This function should only be called inside by WorkerGroup\n        \"\"\"\n        store_env_dict = {f\"_{key.lower()}\": store.get(f\"_{key.lower()}\", None) for key in type(self).env_keys()}\n        self.__dict__.update(store_env_dict)  # this is hacky\n        # print(f\"__dict__: {self.__dict__}\")\n        for key in type(self).env_keys():\n            val = self.__dict__.get(f\"_{key.lower()}\", None)\n            if val is not None:\n                # print(f\"set {key} to {val}\")\n                os.environ[key] = str(val)\n        os.environ[\"REDIS_STORE_SERVER_HOST\"] = (\n            str(self._master_addr).replace(\"[\", \"\").replace(\"]\", \"\") if self._master_addr else \"\"\n        )\n\n    def get_master_addr_port(self):\n        \"\"\"Get the master address and port for distributed communication.\"\"\"\n        return self._master_addr, self._master_port\n\n    def get_cuda_visible_devices(self):\n        \"\"\"Get the CUDA visible devices configuration.\"\"\"\n        import os\n\n        visible_devices = os.environ.get(get_visible_devices_keyword().upper(), \"not set\")\n        return visible_devices\n\n    @property\n    def world_size(self):\n        \"\"\"Get the total number of workers in the distributed setup.\"\"\"\n        return self._world_size\n\n    @property\n    def rank(self):\n        \"\"\"Get the rank of this worker in the distributed setup.\"\"\"\n        return self._rank\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE_PROTO_WITH_FUNC)\n    def execute_with_func_generator(self, func, *args, **kwargs):\n        \"\"\"Execute a function with function generator dispatch mode.\n\n        Args:\n            func:\n                Function to execute\n            *args:\n                Positional arguments for the function\n            **kwargs:\n                Keyword arguments for the function\n        \"\"\"\n        ret_proto = func(self, *args, **kwargs)\n        return ret_proto\n\n    @register(dispatch_mode=Dispatch.ALL_TO_ALL, execute_mode=Execute.RANK_ZERO)\n    def execute_func_rank_zero(self, func, *args, **kwargs):\n        \"\"\"Execute a function in rank zero execution mode.\n\n        Args:\n            func:\n                Function to execute\n            *args:\n                Positional arguments for the function\n            **kwargs:\n                Keyword arguments for the function\n        \"\"\"\n        result = func(*args, **kwargs)\n        return result\n"
  },
  {
    "path": "verl/single_controller/base/worker_group.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nthe class of WorkerGroup\n\"\"\"\n\nimport logging\nimport signal\nimport threading\nimport time\nfrom typing import Any, Callable\n\nfrom .decorator import MAGIC_ATTR, Dispatch, get_predefined_dispatch_fn, get_predefined_execute_fn\n\n\nclass ResourcePool:\n    \"\"\"\n    Manages a pool of resources across multiple nodes, tracking process counts and GPU allocations.\n    The class provides methods to calculate world size, local world sizes, and local ranks\n    across all nodes in the pool.\n    \"\"\"\n\n    def __init__(self, process_on_nodes=None, max_colocate_count: int = 10, n_gpus_per_node=8) -> None:\n        \"\"\"Initialize the ResourcePool with node processes and GPU configuration.\n\n        Args:\n            process_on_nodes (List[int], optional): List of process counts per node. Defaults to empty list.\n            max_colocate_count (int, optional): Maximum number of processes that can be colocated. Defaults to 10.\n            n_gpus_per_node (int, optional): Number of GPUs available per node. Defaults to 8.\n        \"\"\"\n        if process_on_nodes is None:\n            process_on_nodes = []\n        self._store = process_on_nodes\n        self.max_colocate_count = max_colocate_count\n        self.n_gpus_per_node = n_gpus_per_node  # this is left for future huawei GPU that contains 16 GPUs per node\n\n    def add_node(self, process_count):\n        self._store.append(process_count)\n\n    @property\n    def world_size(self):\n        \"\"\"Total number of processes across all nodes in the pool.\"\"\"\n        return sum(self._store)\n\n    def __call__(self) -> Any:\n        return self._store\n\n    @property\n    def store(self):\n        return self._store\n\n    def local_world_size_list(self) -> list[int]:\n        \"\"\"Returns a flat list where each process has its local world size.\"\"\"\n        nested_local_world_size_list = [\n            [local_world_size for _ in range(local_world_size)] for local_world_size in self._store\n        ]\n        return [item for row in nested_local_world_size_list for item in row]\n\n    def local_rank_list(self) -> list[int]:\n        \"\"\"Returns a flat list of local ranks for all processes across all nodes.\"\"\"\n        nested_local_rank_list = [[i for i in range(local_world_size)] for local_world_size in self._store]\n        return [item for row in nested_local_rank_list for item in row]\n\n\nclass ClassWithInitArgs:\n    \"\"\"\n    Wrapper class that stores constructor arguments for deferred instantiation.\n    This class is particularly useful for remote class instantiation where\n    the actual construction needs to happen at a different time or location.\n    \"\"\"\n\n    def __init__(self, cls, *args, **kwargs) -> None:\n        \"\"\"Initialize the ClassWithInitArgs instance.\n\n        Args:\n            cls: The class to be instantiated later\n            *args: Positional arguments for the class constructor\n            **kwargs: Keyword arguments for the class constructor\n        \"\"\"\n        self.cls = cls\n        self.args = args\n        self.kwargs = kwargs\n\n        self.fused_worker_used = False\n\n    def __call__(self) -> Any:\n        \"\"\"Instantiate the stored class with the stored arguments.\"\"\"\n        return self.cls(*self.args, **self.kwargs)\n\n\ndef check_workers_alive(workers: list, is_alive: Callable, gap_time: float = 1) -> None:\n    \"\"\"Continuously monitors worker processes and raises SIGABRT if any worker dies.\n\n    Args:\n        workers (List):\n            List of worker objects to monitor\n        is_alive (Callable):\n            Function to check if a worker is alive\n        gap_time (float):\n            Time interval between checks\n    \"\"\"\n    import time\n\n    while True:\n        for worker in workers:\n            if not is_alive(worker):\n                logging.warning(f\"worker {worker} is not alive sending signal to main thread\")\n                signal.raise_signal(signal.SIGABRT)\n        time.sleep(gap_time)\n\n\nclass WorkerGroup:\n    \"\"\"\n    Base class for managing a group of workers in a distributed system.\n    The class provides methods for worker management, aliveness checking, and method binding.\n    \"\"\"\n\n    fused_worker_execute_fn_name = \"_fuw_execute\"\n\n    def __init__(self, resource_pool: ResourcePool, **kwargs) -> None:\n        self._is_init_with_detached_workers = resource_pool is None\n\n        self.fused_worker_used = False\n\n        if resource_pool is not None:\n            # handle the case when WorkGroup is attached to an existing one\n            self._process_dispatch_config = resource_pool()\n        else:\n            self._process_dispatch_config = None\n\n        self._workers = []\n        self._worker_names = []\n\n        self._dispatch_info = {}\n        self._collect_info = {}\n\n        self._master_addr = None\n        self._master_port = None\n\n        self._checker_thread: threading.Thread = None\n\n    def _is_worker_alive(self, worker):\n        \"\"\"Check if a worker is alive. Must be implemented by derived classes.\"\"\"\n        raise NotImplementedError(\"WorkerGroup._is_worker_alive called, should be implemented in derived class.\")\n\n    def _block_until_all_workers_alive(self) -> None:\n        \"\"\"Blocks until all workers in the group are alive.\"\"\"\n        while True:\n            all_state = [self._is_worker_alive(worker) for worker in self._workers]\n            if False in all_state:\n                time.sleep(1)\n            else:\n                break\n\n    def start_worker_aliveness_check(self, every_n_seconds=1) -> None:\n        \"\"\"Starts a background thread to monitor worker aliveness.\n\n        Args:\n            every_n_seconds (int): Interval between aliveness checks\n        \"\"\"\n        # before starting checking worker aliveness, make sure all workers are already alive\n        self._block_until_all_workers_alive()\n\n        self._checker_thread = threading.Thread(\n            target=check_workers_alive, args=(self._workers, self._is_worker_alive, every_n_seconds)\n        )\n        self._checker_thread.start()\n\n    @property\n    def world_size(self):\n        \"\"\"Number of workers in the group.\"\"\"\n        return len(self._workers)\n\n    def _bind_worker_method(self, user_defined_cls, func_generator):\n        \"\"\"Binds worker methods to the WorkerGroup based on registered attributes.\n\n        Args:\n            user_defined_cls (type): The class containing methods to bind\n            func_generator (Callable): Function that generates the bound method\n\n        Returns:\n            List[str]: List of method names that were successfully bound\n        \"\"\"\n        method_names = []\n        for method_name in dir(user_defined_cls):\n            try:\n                method = getattr(user_defined_cls, method_name)\n                assert callable(method), f\"{method_name} in {user_defined_cls} is not callable\"\n            except Exception:\n                # if it is a property, it will fail because Class doesn't have instance property\n                continue\n\n            if hasattr(method, MAGIC_ATTR):\n                # this method is decorated by register\n                attribute = getattr(method, MAGIC_ATTR)\n                assert isinstance(attribute, dict), f\"attribute must be a dictionary. Got {type(attribute)}\"\n                assert \"dispatch_mode\" in attribute, \"attribute must contain dispatch_mode in its key\"\n\n                dispatch_mode = attribute[\"dispatch_mode\"]\n                execute_mode = attribute[\"execute_mode\"]\n                blocking = attribute[\"blocking\"]\n\n                # get dispatch fn\n                if isinstance(dispatch_mode, Dispatch):\n                    # get default dispatch fn\n                    fn = get_predefined_dispatch_fn(dispatch_mode=dispatch_mode)\n                    dispatch_fn = fn[\"dispatch_fn\"]\n                    collect_fn = fn[\"collect_fn\"]\n                else:\n                    assert isinstance(dispatch_mode, dict)\n                    assert \"dispatch_fn\" in dispatch_mode\n                    assert \"collect_fn\" in dispatch_mode\n                    dispatch_fn = dispatch_mode[\"dispatch_fn\"]\n                    collect_fn = dispatch_mode[\"collect_fn\"]\n\n                # get execute_fn_name\n                execute_mode = get_predefined_execute_fn(execute_mode=execute_mode)\n                wg_execute_fn_name = execute_mode[\"execute_fn_name\"]\n\n                # get execute_fn from string\n                try:\n                    execute_fn = getattr(self, wg_execute_fn_name)\n                    assert callable(execute_fn), \"execute_fn must be callable\"\n                except Exception:\n                    print(f\"execute_fn {wg_execute_fn_name} is invalid\")\n                    raise\n\n                # bind a new method to the RayWorkerGroup\n                func = func_generator(\n                    self,\n                    method_name,\n                    dispatch_fn=dispatch_fn,\n                    collect_fn=collect_fn,\n                    execute_fn=execute_fn,\n                    blocking=blocking,\n                )\n\n                try:\n                    setattr(self, method_name, func)\n                    method_names.append(method_name)\n                except Exception as e:\n                    raise ValueError(f\"Fail to set method_name {method_name}\") from e\n\n        return method_names\n"
  },
  {
    "path": "verl/single_controller/ray/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .base import (\n    RayClassWithInitArgs,\n    RayResourcePool,\n    RayWorkerGroup,\n    ResourcePoolManager,\n    SubRayResourcePool,\n    create_colocated_worker_cls,\n    create_colocated_worker_cls_fused,\n)\n\n__all__ = [\n    \"RayClassWithInitArgs\",\n    \"RayResourcePool\",\n    \"SubRayResourcePool\",\n    \"RayWorkerGroup\",\n    \"ResourcePoolManager\",\n    \"create_colocated_worker_cls\",\n    \"create_colocated_worker_cls_fused\",\n]\n"
  },
  {
    "path": "verl/single_controller/ray/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport logging\nimport os\nimport socket\nfrom copy import deepcopy\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nimport numpy as np\nimport ray\nfrom ray.experimental.state.api import get_actor\nfrom ray.util.placement_group import PlacementGroup, placement_group\nfrom ray.util.scheduling_strategies import NodeAffinitySchedulingStrategy, PlacementGroupSchedulingStrategy\n\nfrom verl.protocol import DataProto, _padding_size_key\nfrom verl.single_controller.base import ClassWithInitArgs, ResourcePool, Worker, WorkerGroup\nfrom verl.single_controller.base.decorator import MAGIC_ATTR, Dispatch\nfrom verl.utils.device import get_device_name\nfrom verl.utils.py_functional import temp_env_var\n\n__all__ = [\"Worker\"]\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef get_random_string(length: int) -> str:\n    import random\n    import string\n\n    letters_digits = string.ascii_letters + string.digits\n    return \"\".join(random.choice(letters_digits) for _ in range(length))\n\n\ndef func_generator(self, method_name, dispatch_fn, collect_fn, execute_fn, blocking):\n    class Functor:\n        def __call__(this, *args, **kwargs):\n            args, kwargs = dispatch_fn(self, *args, **kwargs)\n            padding_count = kwargs.pop(_padding_size_key, 0)\n            output = execute_fn(method_name, *args, **kwargs)\n            if blocking:\n                output = ray.get(output)\n            output = collect_fn(self, output)\n            if padding_count > 0:\n                if isinstance(output, DataProto):\n                    indices = [i for i in range(len(output))][:-padding_count]\n                    output = output.select_idxs(indices)\n                elif isinstance(output, list):\n                    output = output[:-padding_count]\n            return output\n\n    # use class type to pass the method_name to get a better observability\n    return type(method_name, (Functor,), {})()\n\n\ndef sort_placement_group_by_node_ip(pgs: list[PlacementGroup]) -> list[PlacementGroup]:\n    \"\"\"\n    Sort the placement groups by node ip, all bundles in a single placement group should be on the same node.\n\n    FSDPCheckpointManager saves sharded model states and optimizer states in local storage, which requires RANK\n    to be consistent across nodes when resume from checkpoint.\n\n    With this function, if there's only one resource pool and there's no node change, RANK should be consistent\n    across nodes in multiple ray jobs, even if the whole ray cluster is restarted.\n    \"\"\"\n    node_ip = {node[\"NodeID\"]: node[\"NodeManagerAddress\"] for node in ray.nodes()}\n    pg_ip = {}\n    for pg in pgs:\n        specs = ray._private.state.state.placement_group_table(pg.id)\n        # all bunles should be on the same node\n        node_id = specs[\"bundles_to_node_id\"][0]\n        pg_ip[pg.id] = node_ip[node_id]\n    return sorted(pgs, key=lambda pg: pg_ip[pg.id])\n\n\n@ray.remote\ndef get_master_addr_port(master_port_range: Optional[list[int]] = None) -> tuple[str, str]:\n    addr = ray.util.get_node_ip_address().strip(\"[]\")\n\n    if master_port_range is None:\n        with socket.socket() as s:\n            s.bind((\"\", 0))\n            port = s.getsockname()[1]\n    else:\n        port = master_port_range[0]\n        while port < master_port_range[1]:\n            try:\n                with socket.socket() as s:\n                    s.bind((\"\", port))\n                    break\n            except OSError:\n                port += 1  # Increment port number if already in use\n                logger.info(\"Port %d is already in use, trying port %d\", port - 1, port)\n        else:\n            raise RuntimeError(f\"Could not find a free port in range {master_port_range}\")\n    return addr, str(port)\n\n\nclass RayResourcePool(ResourcePool):\n    def __init__(\n        self,\n        process_on_nodes: Optional[list[int]] = None,\n        use_gpu: bool = True,\n        name_prefix: str = None,\n        max_colocate_count: int = 10,\n        detached=False,\n        accelerator_type: Optional[str] = None,\n    ) -> None:\n        super().__init__(process_on_nodes, max_colocate_count)\n        self.use_gpu = use_gpu\n        # print(f\"in RayProcessDispatchConfiguration: name_prefix = {name_prefix}\")\n        self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix\n        self.pgs = None\n        self.detached = detached\n        self.accelerator_type = accelerator_type\n\n    def get_placement_groups(self, strategy=\"STRICT_PACK\", name=None, device_name=\"cuda\"):\n        if self.pgs is not None:\n            return self.pgs\n\n        pg_name_prefix = (\n            name if name else f\"{self.name_prefix}verl_group_{'_'.join([str(count) for count in self._store])}:\"\n        )\n        # print(f\"pg_name_prefix = {pg_name_prefix}\")\n        if device_name == \"npu\":\n            device_name = \"NPU\"\n        elif device_name == \"cuda\":\n            device_name = \"GPU\"\n\n        bundle = {\"CPU\": self.max_colocate_count}\n        if self.use_gpu:\n            bundle[device_name] = 1\n            if self.accelerator_type is not None:\n                bundle[self.accelerator_type] = 1e-4\n        pg_scheme = [[bundle.copy() for _ in range(process_count)] for process_count in self._store]\n\n        lifetime = \"detached\" if self.detached else None\n\n        pgs = [\n            placement_group(bundles=bundles, strategy=strategy, name=pg_name_prefix + str(idx), lifetime=lifetime)\n            for idx, bundles in enumerate(pg_scheme)\n        ]\n\n        ray.get([pg.ready() for pg in pgs])\n\n        self.pgs = sort_placement_group_by_node_ip(pgs)\n        return pgs\n\n\nclass SubRayResourcePool(RayResourcePool):\n    def __init__(\n        self,\n        placement_groups: list[PlacementGroup],\n        start_bundle_index: int,\n        subgroup_world_size: int,\n        **kwargs,\n    ) -> None:\n        super().__init__(**kwargs)\n        self.pgs = placement_groups\n        self.start_bundle_index = start_bundle_index\n        self.subgroup_world_size = subgroup_world_size\n\n    @property\n    def world_size(self):\n        return self.subgroup_world_size\n\n\n@dataclass\nclass ResourcePoolManager:\n    \"\"\"\n    Define a resource pool specification. Resource pool will be initialized first.\n    \"\"\"\n\n    resource_pool_spec: dict[str, list[int]]\n    mapping: dict[int, str]\n    resource_pool_dict: dict[str, RayResourcePool] = field(default_factory=dict)\n\n    def create_resource_pool(self):\n        \"\"\"Create Ray resource pools for distributed training.\n\n        Initializes resource pools based on the resource pool specification,\n        with each pool managing GPU resources across multiple nodes.\n        For FSDP backend, uses max_colocate_count=1 to merge WorkerGroups.\n        For Megatron backend, uses max_colocate_count>1 for different models.\n        \"\"\"\n        for resource_pool_name, process_on_nodes in self.resource_pool_spec.items():\n            # max_colocate_count means the number of WorkerGroups (i.e. processes) in each RayResourcePool\n            # For FSDP backend, using max_colocate_count=3: actor_critic_ref, rollout, reward model (optional)\n            # For Megatron backend, we recommend using max_colocate_count>1\n            # that can utilize different WorkerGroup for differnt models\n            resource_pool = RayResourcePool(\n                process_on_nodes=process_on_nodes, use_gpu=True, max_colocate_count=3, name_prefix=resource_pool_name\n            )\n            self.resource_pool_dict[resource_pool_name] = resource_pool\n\n        self._check_resource_available()\n\n    def get_resource_pool(self, role) -> RayResourcePool:\n        \"\"\"Get the resource pool of the worker_cls\"\"\"\n        return self.resource_pool_dict[self.mapping[role]]\n\n    def get_n_gpus(self) -> int:\n        \"\"\"Get the number of gpus in this cluster.\"\"\"\n        return sum([n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes])\n\n    def _check_resource_available(self):\n        \"\"\"Check if the resource pool can be satisfied in this ray cluster.\"\"\"\n        node_available_resources = ray._private.state.available_resources_per_node()\n        node_available_gpus = {\n            node: node_info.get(\"GPU\", 0) if \"GPU\" in node_info else node_info.get(\"NPU\", 0)\n            for node, node_info in node_available_resources.items()\n        }\n\n        # check total required gpus can be satisfied\n        total_available_gpus = sum(node_available_gpus.values())\n        total_required_gpus = sum(\n            [n_gpus for process_on_nodes in self.resource_pool_spec.values() for n_gpus in process_on_nodes]\n        )\n        if total_available_gpus < total_required_gpus:\n            raise ValueError(\n                f\"Total available GPUs {total_available_gpus} is less than total desired GPUs {total_required_gpus}\"\n            )\n\n\ndef extract_pg_from_exist(\n    resource_pools: dict[str, RayResourcePool], src_role_names: list[str], resource_pool: RayResourcePool\n) -> list:\n    src_pgs = [\n        pg\n        for role_name, resource_pool in resource_pools.items()\n        for pg in resource_pool.get_placement_groups()\n        if role_name in src_role_names\n    ]\n\n    sorted_src_pgs = sorted(src_pgs, key=lambda pg: pg.bundle_count, reverse=True)\n    sorted_process_on_nodes = sorted([(val, idx) for idx, val in enumerate(resource_pool.store)], reverse=True)\n\n    unsorted_pgs: list[tuple[int, PlacementGroup]] = []\n    searching_idx = 0\n    for request_process, original_idx in sorted_process_on_nodes:\n        assert searching_idx < len(sorted_src_pgs), f\"no enough nodes for request: searching {searching_idx} th node\"\n        assert request_process <= sorted_src_pgs[searching_idx].bundle_count, (\n            f\"requesting {request_process} processes, bundle count cannot satisfy\"\n        )\n        unsorted_pgs.append((original_idx, sorted_src_pgs[searching_idx]))\n        searching_idx += 1\n\n    return [pg for _, pg in sorted(unsorted_pgs)]\n\n\n# split a RayResourcePool or SubRayResourcePool into multiple SubRayResourcePool\ndef split_resource_pool(\n    resource_pool: RayResourcePool | SubRayResourcePool, split_size: int | list[int]\n) -> list[SubRayResourcePool]:\n    \"\"\"\n    Split a RayResourcePool into multiple SubRayResourcePool.\n    resouce_pool can also be a SubRayResourcePool (have been splited) for multiple-time spliting.\n\n    Args:\n        resource_pool (RayResourcePool | SubRayResourcePool): The resource pool to split.\n        split_size (int | list[int]): The size of each split. If int, all splits will have the same size.\n            If list[int], each element in the list represents the size of a split.\n\n    Returns:\n        list[SubRayResourcePool]: A list of SubRayResourcePool after splitting.\n    \"\"\"\n    # convert split_size to list[int]\n    if isinstance(split_size, int):\n        assert resource_pool.world_size % split_size == 0, \"split_size must be a divisor of world_size\"\n        num_replica = resource_pool.world_size // split_size\n        split_size_list = [split_size] * num_replica\n    else:\n        split_size_list = split_size\n\n    assert sum(split_size_list) == resource_pool.world_size, \"split_size must sum up to world_size\"\n\n    # judge if this resource pool has been splited\n    if isinstance(resource_pool, SubRayResourcePool):\n        start_bundle_idx_list = np.cumsum([resource_pool.start_bundle_index] + split_size_list[:-1])\n    else:\n        start_bundle_idx_list = np.cumsum([0] + split_size_list[:-1])\n\n    # ensure resource_pool.pgs has been initialized\n    placement_groups = resource_pool.get_placement_groups()\n    split_resource_pools = [\n        SubRayResourcePool(\n            process_on_nodes=resource_pool.store,\n            use_gpu=resource_pool.use_gpu,\n            name_prefix=f\"{resource_pool.name_prefix}_split_{split_idx}\",\n            max_colocate_count=resource_pool.max_colocate_count,\n            placement_groups=placement_groups,\n            start_bundle_index=start_bundle_idx_list[split_idx],\n            subgroup_world_size=split_size_list[split_idx],\n        )\n        for split_idx in range(len(split_size_list))\n    ]\n    return split_resource_pools\n\n\ndef merge_resource_pool(rp1: RayResourcePool, rp2: RayResourcePool) -> RayResourcePool:\n    assert rp1.use_gpu == rp2.use_gpu, \"Both RayResourcePool must either use_gpu or not\"\n    assert rp1.max_colocate_count == rp2.max_colocate_count, \"Both RayResourcePool must has the same max_colocate_count\"\n    assert rp1.n_gpus_per_node == rp2.n_gpus_per_node, \"Both RayResourcePool must has the same n_gpus_per_node\"\n    assert rp1.detached == rp2.detached, \"Detached ResourcePool cannot be merged with non-detached ResourcePool\"\n\n    new_store = rp1.store + rp2.store\n\n    merged = type(rp1)(\n        new_store, rp1.use_gpu, f\"{rp1.name_prefix}_{rp2.name_prefix}\", rp1.max_colocate_count, rp1.detached\n    )\n    merged.pgs = rp1.get_placement_groups(device_name=get_device_name()) + rp2.get_placement_groups(\n        device_name=get_device_name()\n    )\n\n    return merged\n\n\nclass RayClassWithInitArgs(ClassWithInitArgs):\n    \"\"\"A wrapper class for Ray actors with initialization arguments.\n\n    This class extends ClassWithInitArgs to provide additional functionality for\n    configuring and creating Ray actors with specific resource requirements and\n    scheduling strategies.\n    \"\"\"\n\n    def __init__(self, cls, *args, **kwargs) -> None:\n        # self._options = kwargs.pop('options', dict())\n        super().__init__(cls, *args, **kwargs)\n        self._options = {}\n        self._additional_resource = {}\n\n    def set_additional_resource(self, additional_resource):\n        \"\"\"Set additional resource requirements for the actor.\n\n        Args:\n            additional_resource: Dictionary specifying additional resource requirements\n        \"\"\"\n        self._additional_resource = additional_resource\n\n    def update_options(self, options: dict):\n        \"\"\"Update the Ray actor creation options.\n\n        Args:\n            options: Dictionary of options to update\n        \"\"\"\n        self._options.update(options)\n\n    def __call__(\n        self,\n        placement_group,\n        placement_group_bundle_idx,\n        use_gpu: bool = True,\n        num_gpus=1,\n        sharing_with=None,\n        device_name=\"cuda\",\n    ) -> Any:\n        \"\"\"Create and return a Ray actor with the configured options.\n\n        Args:\n            placement_group: Ray placement group for scheduling\n            placement_group_bundle_idx: Index of the bundle in the placement group\n            use_gpu: Whether to use GPU resources\n            num_gpus: Number of GPUs to allocate\n            sharing_with: Actor to share resources with\n            device_name: Device for training\n\n        Returns:\n            A Ray actor handle with the configured options\n        \"\"\"\n        if sharing_with is not None:\n            target_node_id = ray.get(sharing_with.get_node_id.remote())\n            visible_devices = ray.get(sharing_with.get_cuda_visible_devices.remote())\n            options = {\"scheduling_strategy\": NodeAffinitySchedulingStrategy(node_id=target_node_id, soft=False)}\n            return self.cls.options(**options).remote(*self.args, cuda_visible_devices=visible_devices, **self.kwargs)\n\n        options = {\n            \"scheduling_strategy\": PlacementGroupSchedulingStrategy(\n                placement_group=placement_group, placement_group_bundle_index=placement_group_bundle_idx\n            )\n        }\n        options.update(self._options)\n\n        if use_gpu and device_name == \"cuda\":\n            options[\"num_gpus\"] = num_gpus\n        if use_gpu and device_name == \"npu\":\n            options[\"resources\"] = {\"NPU\": num_gpus}\n\n        if len(self._additional_resource) > 1:\n            for k, v in self._additional_resource.items():\n                options[k] = v\n\n        # print(\"cls:\", self.cls)\n        # print(\"args: \", self.args)\n        # print(\"kwargs: \", self.kwargs)\n        return self.cls.options(**options).remote(*self.args, **self.kwargs)\n\n\nclass RayWorkerGroup(WorkerGroup):\n    \"\"\"A group of Ray workers that can be managed collectively.\n\n    This class extends WorkerGroup to provide Ray-specific functionality for\n    creating and managing groups of Ray actors with specific resource requirements\n    and scheduling strategies.\n    \"\"\"\n\n    def __init__(\n        self,\n        resource_pool: RayResourcePool = None,\n        ray_cls_with_init: RayClassWithInitArgs = None,\n        bin_pack: bool = True,\n        name_prefix: str = None,\n        detached=False,\n        worker_names=None,\n        worker_handles: list[ray.actor.ActorHandle] = None,\n        ray_wait_register_center_timeout: int = 300,\n        **kwargs,\n    ) -> None:\n        \"\"\"Initialize a RayWorkerGroup.\n\n        Args:\n            resource_pool: Resource pool for worker allocation\n            ray_cls_with_init: Class with initialization arguments for workers\n            bin_pack: Whether to use strict bin packing for resource allocation\n            name_prefix: Prefix for worker names\n            detached: Whether workers should be detached\n            worker_names: Names of existing workers to attach to\n            ray_wait_register_center_timeout: Timeout for waiting on register center\n            **kwargs: Additional keyword arguments\n        \"\"\"\n        self._master_addr = kwargs.pop(\"master_addr\", None)\n        self._master_port = kwargs.pop(\"master_port\", None)\n        self.use_gpu = kwargs.pop(\"use_gpu\", resource_pool.use_gpu if resource_pool is not None else True)\n        self._ray_master_port_range = kwargs.pop(\"master_port_range\", None)\n        super().__init__(resource_pool=resource_pool, **kwargs)\n        self.ray_cls_with_init = ray_cls_with_init\n        self.name_prefix = get_random_string(length=6) if name_prefix is None else name_prefix\n        self._ray_wait_register_center_timeout = ray_wait_register_center_timeout\n        # Whether the WorkerGroup is a Colocate WorkerGroup created by FusedWorker.\n        self.fused_worker_used = False if ray_cls_with_init is None else ray_cls_with_init.fused_worker_used\n        # if a WorkerGroup is spawned from Colocate WorkerGroup, this indicates which sub-class is binded to\n        # this WorkerGroup.\n        self.sub_cls_name = \"\"\n        self.device_name = kwargs.get(\"device_name\", \"cuda\")\n        self.profile_steps = kwargs.get(\"profile_steps\", None)\n        self.worker_nsight_options = kwargs.get(\"worker_nsight_options\", None)\n        self.customized_worker_env = kwargs.get(\"worker_env\", {})\n        if self.worker_nsight_options is not None and self.worker_nsight_options[\"capture-range-end\"] is None:\n            self.worker_nsight_options[\"capture-range-end\"] = f\"repeat-shutdown:{6 * len(self.profile_steps)}\"\n\n        if worker_names is not None and (not self.fused_worker_used):\n            assert self._is_init_with_detached_workers\n            self._worker_names = worker_names\n\n        if self._is_init_with_detached_workers:\n            self._init_with_detached_workers(worker_names=worker_names, worker_handles=worker_handles)\n        elif isinstance(resource_pool, SubRayResourcePool):\n            self._init_with_subresource_pool(\n                resource_pool=resource_pool,\n                ray_cls_with_init=ray_cls_with_init,\n                bin_pack=bin_pack,\n                detached=detached,\n                worker_env=self.customized_worker_env,\n            )\n        else:\n            self._init_with_resource_pool(\n                resource_pool=resource_pool,\n                ray_cls_with_init=ray_cls_with_init,\n                bin_pack=bin_pack,\n                detached=detached,\n                worker_env=self.customized_worker_env,\n            )\n\n        if ray_cls_with_init is not None:\n            self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)\n\n        self.wg_dict = None\n        self.method_names = []\n\n    def _is_worker_alive(self, worker: ray.actor.ActorHandle):\n        \"\"\"Check if a worker actor is still alive.\n\n        Args:\n            worker: Ray actor handle to check\n\n        Returns:\n            bool: True if the worker is alive, False otherwise\n        \"\"\"\n        worker_state_dict = get_actor(worker._actor_id.hex())\n        return worker_state_dict.get(\"state\", \"undefined\") == \"ALIVE\" if worker_state_dict is not None else False\n\n    def _init_with_detached_workers(self, worker_names, worker_handles):\n        # ray.get_actor holds a weak reference to the actor, which causes actors garbage collected unexpectedly\n        # if we only hold spawn RayWorkerGroup. By passing actor handle explicitly, spawn RayWorkerGroup have\n        # strong reference to these actors.\n        # https://github.com/ray-project/ray/pull/45699\n        workers = worker_handles if worker_handles else [ray.get_actor(name=name) for name in worker_names]\n        self._workers = workers\n        self._world_size = len(workers)\n\n    def _get_master_addr_port(self, pg, bundle_index=0, master_port_range=None):\n        \"\"\"Get master addr and port for this worker group\"\"\"\n        if self._master_addr is None and self._master_port is None:\n            self._master_addr, self._master_port = ray.get(\n                get_master_addr_port.options(\n                    scheduling_strategy=PlacementGroupSchedulingStrategy(\n                        placement_group=pg, placement_group_bundle_index=bundle_index\n                    ),\n                ).remote(master_port_range=master_port_range)\n            )\n        elif self._master_addr is not None and self._master_port is not None:\n            logger.debug(f\"{self._master_addr=} {self._master_port=}\")\n        else:\n            raise ValueError(\n                \"Both 'master_addr' and 'master_port' must be provided if you intend to manually specify them, \"\n                \"or neither should be provided to use Ray's default assignment.\"\n            )\n\n    def _init_with_resource_pool(\n        self,\n        resource_pool,\n        ray_cls_with_init,\n        bin_pack,\n        detached,\n        worker_env=None,\n    ):\n        \"\"\"Initialize the worker group by creating new workers from a resource pool.\n\n        Args:\n            resource_pool: Resource pool for worker allocation\n            ray_cls_with_init: Class with initialization arguments for workers\n            bin_pack: Whether to use strict bin packing for resource allocation\n            detached: Whether workers should be detached\n        \"\"\"\n        self.resource_pool = resource_pool\n        strategy = \"PACK\"\n        if bin_pack:\n            strategy = \"STRICT_PACK\"\n        pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)\n        world_size = resource_pool.world_size\n        self._world_size = world_size\n        # cia.add_kwarg(\"_world_size\", world_size)\n\n        rank = -1\n        local_world_size = resource_pool.store[0]\n        for pg_idx, pg in enumerate(sort_placement_group_by_node_ip(pgs)):\n            assert local_world_size <= pg.bundle_count, f\"when generating for {self.name_prefix}, for the \"\n            if pg_idx == 0:\n                self._get_master_addr_port(pg, bundle_index=0, master_port_range=self._ray_master_port_range)\n\n            for local_rank in range(local_world_size):\n                rank += 1\n                self._create_worker(\n                    rank=rank,\n                    pg_idx=pg_idx,\n                    pg=pg,\n                    local_rank=local_rank,\n                    resource_pool=resource_pool,\n                    ray_cls_with_init=ray_cls_with_init,\n                    worker_env=worker_env,\n                    detached=detached,\n                )\n\n    def _init_with_subresource_pool(self, resource_pool, ray_cls_with_init, bin_pack, detached, worker_env=None):\n        \"\"\"Initialize the worker group by creating new workers from a resource pool or sub resource pool.\n        Args:\n            resource_pool: Resource pool for worker allocation\n            ray_cls_with_init: Class with initialization arguments for workers\n            bin_pack: Whether to use strict bin packing for resource allocation\n            detached: Whether workers should be detached\n        \"\"\"\n        strategy = \"PACK\"\n        if bin_pack:\n            strategy = \"STRICT_PACK\"\n        pgs = resource_pool.get_placement_groups(strategy=strategy, device_name=self.device_name)\n        world_size = resource_pool.world_size\n        self._world_size = world_size\n\n        rank = -1\n        local_world_size = resource_pool.store[0]\n        self._get_master_addr_port(\n            pgs[resource_pool.start_bundle_index // local_world_size],\n            bundle_index=resource_pool.start_bundle_index % local_world_size,\n            master_port_range=self._ray_master_port_range,\n        )\n        for curr_rank in range(resource_pool.start_bundle_index, resource_pool.start_bundle_index + world_size):\n            pg_idx = curr_rank // local_world_size\n            pg = pgs[pg_idx]\n            local_rank = curr_rank % local_world_size\n            assert local_world_size <= pg.bundle_count, f\"when generating for {self.name_prefix}, for the \"\n\n            rank += 1\n            self._create_worker(\n                rank=rank,\n                pg_idx=pg_idx,\n                pg=pg,\n                local_rank=local_rank,\n                resource_pool=resource_pool,\n                ray_cls_with_init=ray_cls_with_init,\n                worker_env=worker_env,\n                detached=detached,\n            )\n\n    def _create_worker(self, rank, pg_idx, pg, local_rank, resource_pool, ray_cls_with_init, worker_env, detached):\n        world_size = resource_pool.world_size\n        use_gpu = resource_pool.use_gpu\n        if self.use_gpu and not use_gpu:\n            raise ValueError(\"use_gpu is True but resource_pool.use_gpu is False\")\n        local_world_size = resource_pool.store[0]\n        num_gpus = 1 / resource_pool.max_colocate_count\n\n        # we pass in environment variable at option so that Worker can use environment variable to set\n        env_vars = {\n            \"WORLD_SIZE\": str(world_size),\n            \"RANK\": str(rank),\n            \"WG_PREFIX\": self.name_prefix,\n            \"WG_BACKEND\": \"ray\",\n            \"RAY_LOCAL_WORLD_SIZE\": str(local_world_size),\n            \"MASTER_ADDR\": self._master_addr,\n            \"MASTER_PORT\": self._master_port,\n        }\n        if worker_env is not None:\n            logging.debug(f\"Appending ray class env, origin: {env_vars}, customized env: {worker_env}\")\n            conflict_env_vars = set(env_vars.keys()) & set(worker_env.keys())\n            if len(conflict_env_vars) > 0:\n                logging.error(\n                    f\"User customized env vars conflict with system env: {conflict_env_vars} \"\n                    f\"Overriding may cause unexpected behavior.\"\n                )\n                raise ValueError(f\"Cannot override protected system env: {conflict_env_vars}\")\n            env_vars.update(worker_env)\n        import re\n\n        cia_name = type(ray_cls_with_init.cls).__name__\n        match = re.search(r\"ActorClass\\(([^)]+)\\)\", cia_name)  # ray.remote(Obj) -> \"ActorClass(Obj)\"\n        cia_name = match.group(1) if match else cia_name  # \"ActorClass(Obj)\" -> \"Obj\"\n        name = f\"{self.name_prefix}{cia_name}_{pg_idx}:{local_rank}\"  # e.g. Worker_2:5\n\n        if self.profile_steps and self.device_name == \"cuda\":\n            ray_cls_with_init.update_options(\n                {\n                    \"runtime_env\": {\n                        \"env_vars\": env_vars,\n                        \"nsight\": self.worker_nsight_options,\n                    },\n                    \"name\": name,\n                }\n            )\n        else:\n            ray_cls_with_init.update_options({\"runtime_env\": {\"env_vars\": env_vars}, \"name\": name})\n\n        if detached:\n            ray_cls_with_init.update_options({\"lifetime\": \"detached\"})\n\n        # create a worker\n        worker = ray_cls_with_init(\n            placement_group=pg,\n            placement_group_bundle_idx=local_rank,\n            use_gpu=self.use_gpu,\n            num_gpus=num_gpus,\n            device_name=self.device_name,\n        )\n        self._workers.append(worker)\n        self._worker_names.append(name)\n\n    @property\n    def worker_names(self):\n        return self._worker_names\n\n    @classmethod\n    def from_detached(\n        cls,\n        name_prefix=None,\n        worker_names=None,\n        worker_handles=None,\n        ray_cls_with_init=None,\n        **kwargs,\n    ):\n        \"\"\"Create a worker group from existing detached workers.\n\n        Args:\n            name_prefix: Prefix for worker names\n            worker_names: Names of existing workers to attach to\n            ray_cls_with_init: Class with initialization arguments for workers\n\n        Returns:\n            A new RayWorkerGroup instance\n        \"\"\"\n        worker_group = cls(\n            resource_pool=None,\n            ray_cls_with_init=ray_cls_with_init,\n            name_prefix=name_prefix,\n            worker_names=worker_names,\n            worker_handles=worker_handles,\n            **kwargs,\n        )\n        return worker_group\n\n    def spawn(self, prefix_set):\n        \"\"\"Spawn to a dictionary of worker groups, each with a subset of method with prefix.\n\n        Args:\n            prefix_set: Set of prefixes to create worker groups for\n\n        Returns:\n            Dictionary of worker groups keyed by prefix\n        \"\"\"\n        if self.fused_worker_used:\n            return self.spawn_fused(prefix_set)\n\n        def _rebind_actor_methods(worker_group, actor_name):\n            prefix: str = actor_name + \"_\"\n            for method_name in dir(worker_group):\n                if method_name.startswith(prefix):\n                    original_method_name = method_name.removeprefix(prefix)\n                    method = getattr(worker_group, method_name)\n                    setattr(worker_group, original_method_name, method)\n\n        new_worker_group_dict = {}\n        for prefix in prefix_set:\n            new_worker_group = self.from_detached(\n                name_prefix=self.name_prefix,\n                worker_names=self._worker_names,\n                worker_handles=self._workers,\n                ray_cls_with_init=self.ray_cls_with_init,\n                profile_steps=self.profile_steps,\n                worker_nsight_options=self.worker_nsight_options,\n            )\n\n            _rebind_actor_methods(new_worker_group, prefix)\n            new_worker_group_dict[prefix] = new_worker_group\n        return new_worker_group_dict\n\n    def spawn_fused(self, prefix_set):\n        \"\"\"Create a dictionary of worker groups for fused workers.\n\n        Args:\n            prefix_set: Set of prefixes to create worker groups for\n\n        Returns:\n            Dictionary of worker groups keyed by prefix\n        \"\"\"\n        wg_dict = dict()\n        for key in prefix_set:\n            new_wg = deepcopy(self)\n            new_wg._bind_worker_method(self.ray_cls_with_init.cls.raw_cls_dict[key], func_generator)\n            new_wg.sub_cls_name = key\n            wg_dict[key] = new_wg\n        return wg_dict\n\n    def fuse(self, prefix_set):\n        \"\"\"Fuse multiple worker groups into the current worker group.\n\n        Args:\n            prefix_set: Set of prefixes to fuse into the worker group\n        \"\"\"\n        if self.wg_dict is None:\n            self.wg_dict = self.spawn(prefix_set)\n        for role_name, role_wg in self.wg_dict.items():\n            setattr(self, role_name, role_wg)\n        self.method_names = self._bind_worker_method(self.ray_cls_with_init.cls, func_generator)\n\n    def _execute_remote_single_worker(self, worker, method_name: str, *args, **kwargs):\n        \"\"\"Execute a method on a single worker remotely.\n\n        Args:\n            worker: The worker actor handle\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            Remote object reference to the method execution\n        \"\"\"\n        if self.fused_worker_used and method_name not in self.method_names:\n            remote_call = getattr(worker, self.fused_worker_execute_fn_name)\n            return remote_call.remote(f\"{self.sub_cls_name}_fwmn_{method_name}\", *args, **kwargs)\n        # fused worker not used\n        remote_call = getattr(worker, method_name)\n        return remote_call.remote(*args, **kwargs)\n\n    def execute_rank_zero_sync(self, method_name: str, *args, **kwargs):\n        \"\"\"Execute a method on rank zero worker synchronously.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            Result of the method execution\n        \"\"\"\n        return ray.get(self.execute_rank_zero_async(method_name, *args, **kwargs))\n\n    def execute_rank_zero_async(self, method_name: str, *args, **kwargs):\n        \"\"\"Execute a method on rank zero worker asynchronously.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            Remote object reference to the method execution\n        \"\"\"\n        return self._execute_remote_single_worker(self._workers[0], method_name, *args, **kwargs)\n\n    def execute_rank_zero(self, method_name: str, *args, **kwargs):\n        \"\"\"Alias for execute_rank_zero_async.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            Remote object reference to the method execution\n        \"\"\"\n        return self.execute_rank_zero_async(method_name, *args, **kwargs)\n\n    def execute_all(self, method_name: str, *args, **kwargs):\n        \"\"\"Alias for execute_all_async.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            List of remote object references to the method executions\n        \"\"\"\n        return self.execute_all_async(method_name, *args, **kwargs)\n\n    def execute_all_sync(self, method_name: str, *args, **kwargs):\n        \"\"\"Execute a method on all workers synchronously.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            List of results from all workers\n        \"\"\"\n        return ray.get(self.execute_all_async(method_name, *args, **kwargs))\n\n    def execute_all_async(self, method_name: str, *args, **kwargs):\n        \"\"\"Execute a method on all workers asynchronously.\n\n        Args:\n            method_name: Name of the method to execute\n            *args: Positional arguments for the method\n            **kwargs: Keyword arguments for the method\n\n        Returns:\n            List of remote object references to the method executions\n        \"\"\"\n        # Here, we assume that if all arguments in args and kwargs are lists,\n        # and their lengths match len(self._workers), we'll distribute each\n        # element in these lists to the corresponding worker\n        # print(f\"execute_all_async: method {method_name}({args}, {kwargs})\")\n        length = len(self._workers)\n        if all(isinstance(arg, list) for arg in args) and all(isinstance(kwarg, list) for kwarg in kwargs.values()):\n            if all(len(arg) == length for arg in args) and all(len(kwarg) == length for kwarg in kwargs.values()):\n                # print(f\"splitting args and kwargs into {length} shards\")\n                result = []\n                for i in range(length):\n                    sliced_args = tuple(arg[i] for arg in args)\n                    sliced_kwargs = {k: v[i] for k, v in kwargs.items()}\n                    result.append(\n                        self._execute_remote_single_worker(self._workers[i], method_name, *sliced_args, **sliced_kwargs)\n                    )\n                return result\n\n        return [self._execute_remote_single_worker(worker, method_name, *args, **kwargs) for worker in self._workers]\n\n    @property\n    def master_address(self):\n        return self._master_addr\n\n    @property\n    def master_port(self):\n        return self._master_port\n\n    @property\n    def workers(self):\n        return self._workers\n\n    @property\n    def world_size(self):\n        return self._world_size\n\n\n\"\"\"\nUtilities that enables creating workers inside the same ray.Actor,\nwith code written in separate ray.Actors.\n\"\"\"\n\n\n# deprecated, switching to FusedWorker\ndef _bind_workers_method_to_parent(cls, key, user_defined_cls):\n    \"\"\"\n    Binds the methods of each worker to the WorkerDict.\n    Note that we only bind public methods that are decorated by register\n    \"\"\"\n\n    for method_name in dir(user_defined_cls):\n        try:\n            method = getattr(user_defined_cls, method_name)\n            assert callable(method), f\"{method_name} in {user_defined_cls} is not callable\"\n        except Exception:\n            # if it is a property, it will fail because Class doesn't have instance property\n            continue\n\n        if hasattr(method, MAGIC_ATTR):\n\n            def generate_function(name, key=key):\n                def func(self, *args, **kwargs):\n                    # dispatch to the actual worker\n                    return getattr(self.worker_dict[key], name)(*args, **kwargs)\n\n                async def async_func(self, *args, **kwargs):\n                    # dispatch to the actual worker\n                    return await getattr(self.worker_dict[key], name)(*args, **kwargs)\n\n                wrapper = async_func if inspect.iscoroutinefunction(method) else func  # noqa: B023\n\n                return wrapper\n\n            func = generate_function(method_name)\n            # pass MAGIC_ATTR for outer worker group\n            attrs = getattr(method, MAGIC_ATTR)\n            setattr(func, MAGIC_ATTR, attrs)\n            try:\n                # bind direct rollout method to class without prefix\n                if attrs[\"dispatch_mode\"] == Dispatch.DIRECT_ROLLOUT_METHOD and \"rollout\" in key:\n                    assert not hasattr(cls, method_name), (\n                        f\"conflict direct rollout method {method_name} with role {key}\"\n                    )\n                    setattr(cls, method_name, func)\n                    print(f\"bind role {key} method {method_name} to class {cls}\")\n                else:\n                    method_name_with_prefix = key + \"_\" + method_name\n                    setattr(cls, method_name_with_prefix, func)\n            except Exception as e:\n                raise ValueError(f\"Fail to set method_name {method_name}\") from e\n\n\ndef _unwrap_ray_remote(cls):\n    if hasattr(cls, \"__ray_actor_class__\"):\n        cls = cls.__ray_actor_class__\n    return cls\n\n\ndef _determine_fsdp_megatron_base_class(mros: list):\n    \"\"\"\n    - megatron: base class should be MegatronWorker\n    - fsdp: base class should be Worker\n    \"\"\"\n    for cls in mros[0]:\n        if cls.__name__ == \"MegatronWorker\":\n            return cls\n        if cls.__name__ == \"Worker\":\n            return cls\n    raise ValueError(f\"Cannot determine base class for {mros}\")\n\n\n# deprecated, switching to FusedWorker\ndef create_colocated_worker_cls(class_dict: dict[str, RayClassWithInitArgs]):\n    \"\"\"\n    This function should return a class instance that delegates the calls to every\n    cls in cls_dict\n    \"\"\"\n    cls_dict = {}\n    init_args_dict = {}\n    worker_cls = _determine_fsdp_megatron_base_class(\n        [cls.cls.__ray_actor_class__.__mro__ for cls in class_dict.values()]\n    )\n    assert issubclass(worker_cls, Worker), f\"worker_cls {worker_cls} should be a subclass of Worker\"\n    print(f\"colocated worker base class {worker_cls}\")\n\n    for key, cls in class_dict.items():\n        cls_dict[key] = cls.cls\n        init_args_dict[key] = {\"args\": cls.args, \"kwargs\": cls.kwargs}\n\n    assert cls_dict.keys() == init_args_dict.keys()\n\n    # TODO: create a class with customizable name\n    class WorkerDict(worker_cls):\n        def __init__(self):\n            super().__init__()\n            self.worker_dict = {}\n            for key, user_defined_cls in cls_dict.items():\n                user_defined_cls = _unwrap_ray_remote(user_defined_cls)\n                # directly instantiate the class without remote\n                # in worker class, e.g. <verl.single_controller.base.worker.Worker>\n                # when DISABLE_WORKER_INIT == 1 it will return immediately\n                with temp_env_var(\"DISABLE_WORKER_INIT\", \"1\"):\n                    self.worker_dict[key] = user_defined_cls(\n                        *init_args_dict[key].get(\"args\", ()), **init_args_dict[key].get(\"kwargs\", {})\n                    )\n\n    # now monkey-patch the methods from inner class to WorkerDict\n    for key, user_defined_cls in cls_dict.items():\n        user_defined_cls = _unwrap_ray_remote(user_defined_cls)\n        _bind_workers_method_to_parent(WorkerDict, key, user_defined_cls)\n\n    remote_cls = ray.remote(WorkerDict)\n    remote_cls = RayClassWithInitArgs(cls=remote_cls)\n    return remote_cls\n\n\nFusedWorkerCLSName = \"FusedWorker\"\n\n\ndef create_colocated_worker_raw_cls(class_dict: dict[str, RayClassWithInitArgs]):\n    \"\"\"\n    This function returns a FusedWorker class.\n\n    `FusedWorker.{class_name}` -> FusedClass\n        Use `class_name` as a param to directly access the underlying class.\n\n    `FusedWorker._fuw_execute(\"{class_name}_fwmn_{method_name}\", *args, **kwargs)`\n        First param must be \"{class_name}_fwmn_{method_name}\" in order to access `method_name`\n        of underlying class `{class_name}`.\n\n    `FusedWorker.fused_worker_dict` -> {\"class_name\": FusedClass}\n        Stores all underlying classes.\n\n    `FusedClass.fused_worker_dict` -> {\"class_name\": FusedClass}\n        The same as `FusedWorker.fused_worker_dict`, enables underlying class to access other\n        underlying classes.\n    \"\"\"\n    raw_cls_dict = {cls_name: _unwrap_ray_remote(cia.cls) for cls_name, cia in class_dict.items()}\n    init_args_dict = {cls_name: cia.args for cls_name, cia in class_dict.items()}\n    init_kwargs_dict = {cls_name: cia.kwargs for cls_name, cia in class_dict.items()}\n    cls_names = list(class_dict.keys())\n\n    # FusedWorker_Actor_Critic\n    class_name_renamed = \"_\".join([FusedWorkerCLSName] + cls_names)\n\n    class FusedWorker(Worker):\n        def __init__(self, *args, **kwargs):\n            super().__init__(*args, **kwargs)\n            self.cls_names = cls_names\n            self.raw_cls_dict = raw_cls_dict\n            self.init_args_dict = init_args_dict\n            self.init_kwargs_dict = init_kwargs_dict\n\n            for cls_name, udc, ud_args, ud_kwargs in zip(\n                self.cls_names,\n                self.raw_cls_dict.values(),\n                self.init_args_dict.values(),\n                self.init_kwargs_dict.values(),\n                strict=True,\n            ):\n                with temp_env_var(\"DISABLE_WORKER_INIT\", \"1\"):\n                    udc._get_ray_actor_cls_name = lambda x, name_renamed=class_name_renamed: name_renamed\n                    udc._get_ray_method_prefix = lambda x, name_prefixed=cls_name: f\"{name_prefixed}_\"\n                    # cls_name = \"actor\", \"critic\", udc = ActorWorker, CriticWorker\n                    self.fused_worker_dict[cls_name] = udc(*ud_args, **ud_kwargs)\n                    setattr(self, cls_name, self.fused_worker_dict[cls_name])\n\n            # injecting fused_worker to each sub worker so they can be aware of existence of each other\n            for _, worker in self.fused_worker_dict.items():\n                setattr(worker, Worker.fused_worker_attr_name, self.fused_worker_dict)\n\n        def _fuw_execute(self, method_name: str, *args, **kwargs):\n            # for fused_worker, method_name is in a form of \"{cls_name}_fwmn_{method_name}\"\n            # where fwmn stands \"fused worker method name\"\n            names = method_name.split(\"_fwmn_\")\n            cls_name = names[0]\n            method_name = names[1]\n\n            assert cls_name in self.fused_worker_dict, (\n                f\"calling {cls_name}'s {method_name}, but {cls_name} not in fused_worker_dict\"\n            )\n            udc_method = getattr(self.fused_worker_dict[cls_name], method_name)\n            return udc_method(*args, **kwargs)\n\n    renamed_fused_worker_cls = type(class_name_renamed, (FusedWorker,), {})\n    renamed_fused_worker_cls.is_fused_worker = True\n    renamed_fused_worker_cls.raw_cls_dict = raw_cls_dict\n\n    return renamed_fused_worker_cls\n\n\ndef create_colocated_worker_cls_fused(class_dict: dict[str, RayClassWithInitArgs]):\n    \"\"\"\n    This function returns a RayClassWithInitArgs instance of FusedWorker, which is an replacement\n    of `create_colocated_worker_cls`. WorkerGroup constructed using this class will be a colocated\n    WorkerGroup, which will be referenced as `ColocateWorkerGroup` below.\n\n    `ColocateWorkerGroup.spawn(prefix_set)`\n        returns a dict of WorkerGroup {\"class_name\": WorkerGroup}, WorkerGroup in this dict will\n        have methods of underlying class `class_name` attached.\n\n    `ColocateWorkerGroup.fuse(prefix_set)`\n        After executing this function, `ColocateWorkerGroup.{class_name}` will return WorkerGroup\n        with methods of underlying class `class_name` attached.\n    \"\"\"\n    raw_colocated_worker_cls = create_colocated_worker_raw_cls(class_dict)\n\n    remote_cls = ray.remote(raw_colocated_worker_cls)\n    cia = RayClassWithInitArgs(cls=remote_cls)\n    cia.fused_worker_used = True\n\n    return cia\n"
  },
  {
    "path": "verl/third_party/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/third_party/torch/__init__.py",
    "content": "# official torch 2.6.0 set_model_state_dict API leads to OOM\n# this is a copy of torch/distributed/checkpoint from torch 2.7.0\n\n# From PyTorch:\n\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n\n# From Caffe2:\n\n# Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n\n# All contributions by Facebook:\n# Copyright (c) 2016 Facebook Inc.\n\n# All contributions by Google:\n# Copyright (c) 2015 Google Inc.\n# All rights reserved.\n\n# All contributions by Yangqing Jia:\n# Copyright (c) 2015 Yangqing Jia\n# All rights reserved.\n\n# All contributions by Kakao Brain:\n# Copyright 2019-2020 Kakao Brain\n\n# All contributions by Cruise LLC:\n# Copyright (c) 2022 Cruise LLC.\n# All rights reserved.\n\n# All contributions by Tri Dao:\n# Copyright (c) 2024 Tri Dao.\n# All rights reserved.\n\n# All contributions by Arm:\n# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates\n\n# All contributions from Caffe:\n# Copyright(c) 2013, 2014, 2015, the respective contributors\n# All rights reserved.\n\n# All other contributions:\n# Copyright(c) 2015, 2016 the respective contributors\n# All rights reserved.\n\n# Caffe2 uses a copyright model similar to Caffe: each contributor holds\n# copyright over their contributions to Caffe2. The project versioning records\n# all such contribution and copyright details. If a contributor wants to further\n# mark their specific copyright on a particular contribution, they should\n# indicate their copyright solely in the commit message of the change when it is\n# committed.\n\n# All rights reserved.\n\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n\n# 1. Redistributions of source code must retain the above copyright\n#    notice, this list of conditions and the following disclaimer.\n\n# 2. Redistributions in binary form must reproduce the above copyright\n#    notice, this list of conditions and the following disclaimer in the\n#    documentation and/or other materials provided with the distribution.\n\n# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n#    and IDIAP Research Institute nor the names of its contributors may be\n#    used to endorse or promote products derived from this software without\n#    specific prior written permission.\n\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n"
  },
  {
    "path": "verl/third_party/torch/distributed/__init__.py",
    "content": "# official torch 2.6.0 set_model_state_dict API leads to OOM\n# this is a copy of torch/distributed/checkpoint from torch 2.7.0\n\n# From PyTorch:\n\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n\n# From Caffe2:\n\n# Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n\n# All contributions by Facebook:\n# Copyright (c) 2016 Facebook Inc.\n\n# All contributions by Google:\n# Copyright (c) 2015 Google Inc.\n# All rights reserved.\n\n# All contributions by Yangqing Jia:\n# Copyright (c) 2015 Yangqing Jia\n# All rights reserved.\n\n# All contributions by Kakao Brain:\n# Copyright 2019-2020 Kakao Brain\n\n# All contributions by Cruise LLC:\n# Copyright (c) 2022 Cruise LLC.\n# All rights reserved.\n\n# All contributions by Tri Dao:\n# Copyright (c) 2024 Tri Dao.\n# All rights reserved.\n\n# All contributions by Arm:\n# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates\n\n# All contributions from Caffe:\n# Copyright(c) 2013, 2014, 2015, the respective contributors\n# All rights reserved.\n\n# All other contributions:\n# Copyright(c) 2015, 2016 the respective contributors\n# All rights reserved.\n\n# Caffe2 uses a copyright model similar to Caffe: each contributor holds\n# copyright over their contributions to Caffe2. The project versioning records\n# all such contribution and copyright details. If a contributor wants to further\n# mark their specific copyright on a particular contribution, they should\n# indicate their copyright solely in the commit message of the change when it is\n# committed.\n\n# All rights reserved.\n\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n\n# 1. Redistributions of source code must retain the above copyright\n#    notice, this list of conditions and the following disclaimer.\n\n# 2. Redistributions in binary form must reproduce the above copyright\n#    notice, this list of conditions and the following disclaimer in the\n#    documentation and/or other materials provided with the distribution.\n\n# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n#    and IDIAP Research Institute nor the names of its contributors may be\n#    used to endorse or promote products derived from this software without\n#    specific prior written permission.\n\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n"
  },
  {
    "path": "verl/third_party/torch/distributed/_state_dict_utils.py",
    "content": "# official torch 2.6.0 set_model_state_dict API leads to OOM\n# this is a copy of torch/distributed/checkpoint from torch 2.7.0\n\n# From PyTorch:\n\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n\n# From Caffe2:\n\n# Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n\n# All contributions by Facebook:\n# Copyright (c) 2016 Facebook Inc.\n\n# All contributions by Google:\n# Copyright (c) 2015 Google Inc.\n# All rights reserved.\n\n# All contributions by Yangqing Jia:\n# Copyright (c) 2015 Yangqing Jia\n# All rights reserved.\n\n# All contributions by Kakao Brain:\n# Copyright 2019-2020 Kakao Brain\n\n# All contributions by Cruise LLC:\n# Copyright (c) 2022 Cruise LLC.\n# All rights reserved.\n\n# All contributions by Tri Dao:\n# Copyright (c) 2024 Tri Dao.\n# All rights reserved.\n\n# All contributions by Arm:\n# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates\n\n# All contributions from Caffe:\n# Copyright(c) 2013, 2014, 2015, the respective contributors\n# All rights reserved.\n\n# All other contributions:\n# Copyright(c) 2015, 2016 the respective contributors\n# All rights reserved.\n\n# Caffe2 uses a copyright model similar to Caffe: each contributor holds\n# copyright over their contributions to Caffe2. The project versioning records\n# all such contribution and copyright details. If a contributor wants to further\n# mark their specific copyright on a particular contribution, they should\n# indicate their copyright solely in the commit message of the change when it is\n# committed.\n\n# All rights reserved.\n\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n\n# 1. Redistributions of source code must retain the above copyright\n#    notice, this list of conditions and the following disclaimer.\n\n# 2. Redistributions in binary form must reproduce the above copyright\n#    notice, this list of conditions and the following disclaimer in the\n#    documentation and/or other materials provided with the distribution.\n\n# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n#    and IDIAP Research Institute nor the names of its contributors may be\n#    used to endorse or promote products derived from this software without\n#    specific prior written permission.\n\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\n\n# ruff: noqa: B028, UP038, UP007, E721, E501\n# mypy: allow-untyped-defs\nimport copy\nimport io\nimport math\nimport weakref\nfrom collections.abc import Mapping, MutableMapping\nfrom typing import TYPE_CHECKING, Any, Callable, NamedTuple, Optional, Union, cast\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn.functional as F\nfrom torch.distributed._functional_collectives import AsyncCollectiveTensor\n\nif dist.is_available() or TYPE_CHECKING:\n    from torch.distributed import distributed_c10d\n    from torch.distributed._shard.sharded_tensor import ShardedTensor\n    from torch.distributed.tensor import DTensor, Replicate, distribute_tensor\n    from torch.distributed.tensor._utils import compute_local_shape_and_global_offset\n\n\ndef _identity_func(\n    obj: torch.Tensor,\n    pg: Optional[dist.ProcessGroup],\n    device: Optional[torch.device],\n    companion_obj: Any,\n) -> torch.Tensor:\n    return obj\n\n\ndef _all_gather_sharded_tensor(\n    sharded_tensor: \"ShardedTensor\",\n    pg: Optional[dist.ProcessGroup] = None,\n    device: Optional[torch.device] = None,\n) -> torch.Tensor:\n    if pg is None:\n        pg = distributed_c10d._get_default_group()\n    world_size = dist.get_world_size(pg)\n    shards = sharded_tensor.local_shards()\n    dim_0_size = sharded_tensor.size()[0]  # type: ignore[index]\n    tensor_numel = sharded_tensor.size().numel()  # type: ignore[union-attr]\n    chunk_size = math.ceil(dim_0_size / world_size) * tensor_numel // dim_0_size\n    pg_device = distributed_c10d._get_pg_default_device(pg) if device is None else device\n    if shards:\n        local_tensor = shards[0].tensor.flatten()\n        if local_tensor.device.type != pg_device.type:\n            local_tensor = local_tensor.to(pg_device)\n        num_padding = chunk_size - local_tensor.numel()\n        if num_padding > 0:\n            local_tensor = F.pad(local_tensor, [0, num_padding])\n    else:\n        local_tensor = torch.zeros(chunk_size, dtype=sharded_tensor.dtype, device=pg_device)\n\n    tensor = torch.empty(\n        chunk_size * world_size,\n        dtype=local_tensor.dtype,\n        device=pg_device,\n    )\n    dist.all_gather_into_tensor(tensor, local_tensor, group=pg)\n\n    tensor = tensor.narrow(0, 0, tensor_numel).reshape(sharded_tensor.size())\n    return tensor\n\n\nclass CompanionMismatch(Exception):\n    pass\n\n\ndef _iterate_state_dict(\n    iter_object: Any,\n    sharded_tensor_func: Callable,\n    dtensor_func: Callable,\n    tensor_func: Callable,\n    *,\n    pg: Optional[dist.ProcessGroup] = None,\n    device: Optional[torch.device] = None,\n    cpu_offload: bool = False,\n    companion_obj: Any = None,\n    ranks_only: tuple[int, ...] = (),\n    type_check: bool = True,\n    non_blocking: bool = True,\n) -> dict[str, Any]:\n    \"\"\"Iterate through the state dict, applying the given functions to each tensor type.\n\n    Args:\n        iter_object (Any): the target state_dict.\n        sharded_tensor_func (Callable): the function to apply to ShardedTensor\n        dtensor_func (Callable): the function to apply to DTensor\n        tensor_func (Callable): the function to apply to Tensor\n        pg (Optional[dist.ProcessGroup]): process group passed to tensor functions\n        device (Optional[torch.device]): device passed to tensor functions\n        cpu_offload (bool): whether to offload the tensors to CPU memory. This option is ignored\n            if a companion_obj is supplied.\n        companion_obj (Any): A companion object to the state dict. If this object\n            is supplied, we attempt to copy the tensor to the companion object.\n        ranks_only (Tuple[int, ...]): if this tuple is empty, all ranks will\n            have the same state_dicts. Otherwise only ranks that in ``ranks_only``\n            have the same state_dicts. Other ranks will get empty state_dicts.\n        type_check (bool): check if the instance data type is a supported type\n            that can be saved by DCP.  The current supported data types are\n            torch.Tensor, DTensor, int, float, str, list, dict, None.\n        non_blocking (bool): whether to use non-blocking copy when copying to the companion object.\n    \"\"\"\n    # TODO: should we use pytree?\n    cpu_device = torch.device(\"cpu\")\n    if isinstance(iter_object, ShardedTensor):\n        ret = sharded_tensor_func(iter_object, pg, device, companion_obj)\n    elif isinstance(iter_object, DTensor):\n        ret = dtensor_func(iter_object, pg, device, companion_obj)\n    elif isinstance(iter_object, torch.Tensor):\n        ret = tensor_func(iter_object, pg, device, companion_obj)\n    elif isinstance(iter_object, (int, float, str, bytes, io.BytesIO)) or iter_object is None:\n        ret = iter_object\n    elif isinstance(iter_object, dict):\n        if companion_obj is not None and (\n            not isinstance(companion_obj, dict) or set(companion_obj.keys()) != set(iter_object.keys())\n        ):\n            msg = \"\" if isinstance(companion_obj, dict) else f\"{set(companion_obj.keys())=} {set(iter_object.keys())=}\"\n            raise CompanionMismatch(msg)\n\n        ret = {\n            key: _iterate_state_dict(\n                value,\n                sharded_tensor_func,\n                dtensor_func,\n                tensor_func,\n                pg=pg,\n                device=device,\n                cpu_offload=cpu_offload,\n                companion_obj=companion_obj[key] if companion_obj is not None else None,\n                ranks_only=ranks_only,\n                type_check=type_check,\n                non_blocking=non_blocking,\n            )\n            for key, value in iter_object.items()\n        }\n    elif isinstance(iter_object, (list, tuple)):\n        if companion_obj is not None and (\n            not isinstance(companion_obj, (list, tuple)) or len(companion_obj) != len(iter_object)\n        ):\n            raise CompanionMismatch\n\n        ret = [\n            _iterate_state_dict(\n                v,\n                sharded_tensor_func,\n                dtensor_func,\n                tensor_func,\n                pg=pg,\n                device=device,\n                cpu_offload=cpu_offload,\n                companion_obj=companion_obj[idx] if companion_obj is not None else None,\n                ranks_only=ranks_only,\n                type_check=type_check,\n                non_blocking=non_blocking,\n            )\n            for idx, v in enumerate(iter_object)\n        ]\n        if isinstance(iter_object, tuple):\n            ret = tuple(ret)\n    elif not type_check:\n        ret = copy.deepcopy(iter_object)\n    else:\n        raise ValueError(f\"Unexpected value type {type(iter_object)}\")\n\n    if not ranks_only or dist.get_rank(pg) in ranks_only:\n        if isinstance(ret, torch.Tensor):\n            if cpu_offload and companion_obj is None:\n                ret = ret.to(cpu_device)\n\n            if companion_obj is not None:\n                if isinstance(companion_obj, DTensor):\n                    assert isinstance(ret, DTensor)\n                    companion_obj._local_tensor.copy_(ret._local_tensor, non_blocking=non_blocking)\n                else:\n                    companion_obj.copy_(ret, non_blocking=non_blocking)\n                ret = companion_obj\n    else:\n        ret = {} if isinstance(ret, dict) else None\n\n    return ret\n\n\ndef _gather_state_dict(\n    state_dict: dict[str, Any],\n    *,\n    pg: Optional[dist.ProcessGroup] = None,\n    device: Optional[torch.device] = None,\n    cpu_offload: bool = False,\n    ranks_only: tuple[int, ...] = (),\n    type_check: bool = True,\n) -> dict[str, Any]:\n    \"\"\"\n    Given a state_dict, this API gathers all the ShardedTensors or DTensors in\n    the state_dict.\n\n\n    Args:\n        state_dict (Dict[str, Any]): the target sharded state_dict.\n        pg (Optional[dist.ProcessGroup]): the process group that is used to\n            gather ShardedTensor. Note that gathering a DTensor will use\n            the DeviceMesh. So this argument will be ignored when gathering a\n            DTensor.\n        device: (Optional[torch.device]): the device that is used to\n            perform allgather for ShardedTensor. Note that gathering a DTensor\n            will use the DeviceMesh. So this argument will be ignored when\n            gathering a DTensor.\n        cpu_offload (bool): whether to offload the tensors to CPU memory. The\n            default value is False.\n        ranks_only: (Tuple[int, ...]): if this tuple is empty, all ranks will\n            have the same state_dicts. Otherwise only ranks that in ``ranks_only``\n            have the same state_dicts. Other ranks will get empty state_dicts.\n        type_check: (bool): check if the instance data type is a supported type\n            that can be saved by DCP.  The current supported data types are\n            torch.Tensor, DTensor, int, float, str, list, dict, None.\n\n    Returns:\n        The gathered state dictionary.\n    \"\"\"\n\n    def sharded_tensor_func(value, pg, device, companion_obj):\n        # ShardedTensor does not seem to record the original device type.\n        # So if the tensor is moved to CPU, we won't know the original type.\n        # As a result, we have to rely on the user to tell us the correct one.\n        cpu_device = torch.device(\"cpu\")\n        output_tensor = _all_gather_sharded_tensor(value, pg, device)\n        local_shard_device = value.local_shards()[0].tensor.device if value.local_shards() else cpu_device\n        if output_tensor.device != local_shard_device:\n            value = output_tensor.to(local_shard_device)\n        else:\n            value = output_tensor\n        return value\n\n    def dtensor_func(value, pg, device, companion_obj):\n        if value.device != value.device_mesh.device_type:\n            value = value.to(value.device_mesh.device_type)\n        # FSDP all_gather: [Shard(0)] -> [Replicate()]\n        # HSDP all_gather: [Replicate(), Shard(0)] -> [Replicate(), Replicate()]\n        # 2D FSDP + TP all_gather:\n        # - [Shard(0), Shard(n)] -> [Replicate(), Replicate()]\n        # - [Shard(0), Replicate()] -> [Replicate(), Replicate()]\n        placements = [Replicate() for _ in value.placements]\n        value = value.redistribute(\n            device_mesh=value.device_mesh,\n            placements=placements,\n        )\n        # Call `wait()` to force the tensor to be synchronous with respect\n        # to the main stream.\n        # See the discussion in https://github.com/pytorch/pytorch/pull/117799.\n        value = value.to_local()\n        if isinstance(value, AsyncCollectiveTensor):\n            value = value.wait()\n        return value\n\n    return _iterate_state_dict(\n        state_dict,\n        sharded_tensor_func,\n        dtensor_func,\n        _identity_func,\n        pg=pg,\n        device=device,\n        cpu_offload=cpu_offload,\n        ranks_only=ranks_only,\n        type_check=type_check,\n    )\n\n\ndef _offload_state_dict_to_cpu(\n    state_dict: dict[str, Any],\n    *,\n    ranks_only: tuple[int, ...] = (),\n    type_check: bool = True,\n) -> dict[str, Any]:\n    \"\"\"\n    Given a state_dict, this API offload all the tensors to CPU memory.\n\n    Args:\n        state_dict (Dict[str, Any]): the target state_dict.\n        pg (Optional[dist.ProcessGroup]): the process group that is used to\n            gather ShardedTensor. Note that gathering a DTensor will use\n            the DeviceMesh. So this argument will be ignored when gathering a\n            DTensor.\n        ranks_only: (Tuple[int, ...]): if this tuple is empty, all ranks will\n            have the same state_dicts. Otherwise only ranks that in ``ranks_only``\n            have the same state_dicts. Other ranks will get empty state_dicts.\n        type_check: (bool): check if the instance data type is a supported type\n            that can be saved by DCP.  The current supported data types are\n            torch.Tensor, DTensor, int, float, str, list, dict, None.\n\n    Returns:\n        The gathered state dictionary.\n    \"\"\"\n\n    ret = _iterate_state_dict(\n        state_dict,\n        _identity_func,\n        _identity_func,\n        _identity_func,\n        pg=None,\n        device=None,\n        cpu_offload=True,\n        ranks_only=ranks_only,\n        type_check=type_check,\n    )\n    return ret\n\n\n@torch.no_grad()\ndef _copy_state_dict(\n    state_dict: dict[str, Any],\n    copy_state_dict: dict[str, Any],\n    non_blocking: bool = False,\n    type_check: bool = True,\n) -> dict[str, Any]:\n    \"\"\"\n    Copies all tensors in a given state dict into a different state_dict with the\n    same structure. Additionally, a copied state dict with the same value references\n    is returned. Editing the keys on this state dict will not affect the\n    passed in copy_state_dict (but the value references are the same).\n\n    .. warning::\n        It is expected by this function that state_dict and copy_state_dict share\n        the same structure and data types.\n\n    .. warning::\n        The current supported data types are\n            torch.Tensor, DTensor, int, float, str, list, dict, None.\n\n    Args:\n        state_dict (Dict[str, Any]): the target state_dict.\n        copy_state_dict (Dict[str, Any]):\n            The state dict we are copying into. This state_dict must have exactly\n             the same structure as the source `state_dict`.\n        non_blocking: (bool): Whether copy ops should be performed asynchronously\n        type_check (bool): check if the instance data type is a supported type\n            that can be saved by DCP. The current supported data types are\n            torch.Tensor, DTensor, int, float, str, list, dict, None.\n\n    Returns:\n        State Dict copy\n    \"\"\"\n\n    return _iterate_state_dict(\n        state_dict,\n        _identity_func,\n        _identity_func,\n        _identity_func,\n        pg=None,\n        device=None,\n        cpu_offload=False,\n        ranks_only=(),\n        companion_obj=copy_state_dict,\n        type_check=type_check,\n        non_blocking=non_blocking,\n    )\n\n\n@torch.no_grad()\ndef _create_cpu_state_dict(\n    state_dict: dict[str, Any], pin_memory: bool = False, share_memory: bool = False\n) -> dict[str, Any]:\n    \"\"\"\n    Given a state_dict, create another state_dict with the same structure and elements.\n    However, all tensors in the returned state_dict are new tensors on CPU. These\n    tensors can be placed on pin_memory or share_memory based on the provided arguments.\n\n    .. warning::\n        Setting both `pin_memory` and `share_memory` to True significantly increases the\n        latency of this method because of the nuances which require us to register memory\n        as pinned directly as opposed to relying on the pin_memory cache allocator. This\n        option should only be used for long lived tensors which are required to be shared.\n        This is not the case as long as at least one of `pin_memory` or `share_memory` is\n         set to False.\n\n    \"\"\"\n\n    def tensor_func(\n        obj: torch.Tensor,\n        pg: Optional[dist.ProcessGroup],\n        device: Optional[torch.device],\n        _: Any,\n    ) -> torch.Tensor:\n        if len(obj.size()) == 0:\n            return torch.tensor(0, dtype=obj.dtype)\n\n        if share_memory:\n            t = torch.empty(*tuple(obj.size()), dtype=obj.dtype)\n            t = t.share_memory_()\n            if pin_memory:\n\n                def unpin_memory(t):\n                    succ = int(torch.cuda.cudart().cudaHostUnregister(t.data_ptr()))\n                    assert succ == 0, f\"Unpinning shared memory failed with error-code: {succ}\"\n\n                weakref.finalize(t, unpin_memory, t)\n                succ = int(\n                    torch.cuda.cudart().cudaHostRegister(\n                        t.data_ptr(),\n                        t.numel() * t.element_size(),\n                        1,  # lines up with 'cudaHostRegisterPortable'\n                    )\n                )\n                assert succ == 0, f\"Pinning shared memory failed with error-code: {succ}\"\n            return t\n        elif pin_memory:\n            return torch.empty(*tuple(obj.size()), dtype=obj.dtype).pin_memory()\n        else:\n            return torch.empty(*tuple(obj.size()), dtype=obj.dtype)\n\n    def dtensor_func(\n        obj: DTensor,\n        pg: Optional[dist.ProcessGroup],\n        device: Optional[torch.device],\n        _: Any,\n    ) -> DTensor:\n        if len(obj.size()) == 0:\n            return obj\n\n        if obj.device != torch.device(\"cpu\"):\n            ret = cast(DTensor, obj.to(device=\"cpu\"))\n        else:\n            ret = copy.deepcopy(obj)\n        ret._local_tensor = tensor_func(ret._local_tensor, pg, device, None)\n        return ret\n\n    ret = _iterate_state_dict(\n        state_dict,\n        _identity_func,\n        dtensor_func,\n        tensor_func,\n        pg=None,\n        device=None,\n        cpu_offload=False,\n        ranks_only=(),\n        type_check=False,\n    )\n    return ret\n\n\ndef _check_state_dict_similarity(\n    state_dict: dict[str, Any],\n    compared_state_dict: dict[str, Any],\n) -> bool:\n    \"\"\"\n    Given two state_dicts, check if the structures are the same. And\n    if a [key, tensor] pair exist in one state_dict there must be\n    the a corresponding pait, [key, other_tensor], in the other state_dict,\n    where tensor and other_tensor have the same size and dtype.\n\n    Return the check result.\n    \"\"\"\n\n    def tensor_func(\n        obj: torch.Tensor,\n        pg: Optional[dist.ProcessGroup],\n        device: Optional[torch.device],\n        companion_obj: Any,\n    ) -> torch.Tensor:\n        if companion_obj.dtype != obj.dtype or companion_obj.size() != obj.size():\n            raise CompanionMismatch\n        return obj\n\n    try:\n        _iterate_state_dict(\n            state_dict,\n            _identity_func,\n            _identity_func,\n            tensor_func,\n            pg=None,\n            device=None,\n            cpu_offload=False,\n            ranks_only=(),\n            companion_obj=compared_state_dict,\n            type_check=False,\n        )\n    except CompanionMismatch:\n        return False\n\n    return True\n\n\nclass _TensorInfo(NamedTuple):\n    size: torch.Size\n    dtype: torch.dtype\n\n\ndef _broadcast_tensors(\n    full_state_dict: dict[str, Any],\n    local_state_dict: dict[str, Any],\n    keys: list[str],\n    device: torch.device,\n    pg: Optional[dist.ProcessGroup] = None,\n) -> None:\n    tensors = []\n    for key in keys:\n        if dist.get_rank() == 0:\n            full_state = full_state_dict[key]\n            assert isinstance(full_state, torch.Tensor)\n            full_tensor = full_state.detach().to(device)\n        else:\n            tensor_info = full_state_dict[key]\n            full_tensor = torch.empty(\n                size=tensor_info.size,\n                device=device,\n                dtype=tensor_info.dtype,\n            )\n        tensors.append(full_tensor)\n        local_state = local_state_dict.get(key, None)\n        if local_state is None:\n            continue\n        elif isinstance(local_state, DTensor):\n            local_state_dict[key] = (local_state, full_tensor)\n        else:\n            local_state_dict[key] = full_tensor\n\n    if pg is None:\n        pg = dist.distributed_c10d._get_default_group()\n\n    if len(tensors) > 1:\n        dist._broadcast_coalesced(pg, tensors, 500, 0)\n    else:\n        dist.broadcast(tensors[0], src=0, group=pg)\n\n    _distribute_tensors(local_state_dict, keys, device, pg)\n\n\ndef _distribute_tensors(\n    local_state_dict: dict[str, Any],\n    keys: list[str],\n    device: torch.device,\n    pg: Optional[dist.ProcessGroup] = None,\n) -> None:\n    if pg is None:\n        pg = dist.distributed_c10d._get_default_group()\n    for key in keys:\n        _local_state = local_state_dict.get(key, None)\n        if _local_state is None or torch.is_tensor(_local_state):\n            continue\n\n        local_state = _local_state[0]\n        full_tensor = _local_state[1]\n\n        shape, offset = compute_local_shape_and_global_offset(\n            full_tensor.shape, local_state.device_mesh, local_state.placements\n        )\n        slices = [\n            slice(cur_offset, cur_offset + cur_shape) for cur_shape, cur_offset in zip(shape, offset, strict=False)\n        ]\n        if local_state.is_meta:\n            # Use .clone() here rather than view to clone and return only the sliced portion, minimizing memory access and cost.\n            local_tensor = full_tensor[slices].detach().clone()\n            # TODO: currently, we cannot handle strided sharding if the dp dimension is not even. For example,\n            # one of the case that is not yet supported is when placements = (Shard(0), _StridedShard(0, sf=2)).\n            ret = DTensor.from_local(\n                local_tensor,\n                local_state.device_mesh,\n                local_state.placements,\n                shape=local_state.shape,\n                stride=local_state.stride(),\n            )\n        else:\n            ret = local_state\n            # Copy full_tensor[slices] into local_state.to_local() to reduce memory footprint.\n            ret.to_local().copy_(full_tensor[slices])\n        local_state_dict[key] = ret\n\n\ndef _broadcast_state_dict(\n    full_state_dict: dict[str, Any],\n    local_state_dict: dict[str, Any],\n    device: torch.device,\n    pg: Optional[dist.ProcessGroup] = None,\n    strict: bool = False,\n    cpu_offload: bool = False,\n) -> None:\n    # Broadcast from rank0's `full_state_dict` to all ranks' `local_state_dict`.\n    # If strict is True, any keys in `local_state_dict` but not in `full_state_dict`\n    # will be removed from `local_state_dict`.\n    ret = {}\n    if dist.get_rank() == 0:\n        for key, value in full_state_dict.items():\n            if not torch.is_tensor(value):\n                ret[key] = value\n            elif value.dim() == 0:\n                ret[key] = value.cpu()\n            else:\n                ret[key] = _TensorInfo(value.size(), value.dtype)\n\n    broadcast_list = [ret]\n    dist.broadcast_object_list(broadcast_list, src=0, group=pg)\n    ret = broadcast_list[0]\n    # Gather values\n    keys = []\n    local_state_dict_keys = set(local_state_dict.keys())\n    global_keys = set()\n    for key, value in ret.items():\n        global_keys.add(key)\n        if not isinstance(value, _TensorInfo):\n            if key in local_state_dict:\n                local_state_dict[key] = value\n            continue\n\n        if dist.get_rank() == 0:\n            ret[key] = full_state_dict[key]\n\n        keys.append(key)\n        # Broadcast every tensor to avoid OOM for now.\n        if len(keys) >= 1:\n            _broadcast_tensors(ret, local_state_dict, keys, device, pg)\n            if cpu_offload:\n                for key in keys:\n                    local_state_dict[key] = local_state_dict[key].cpu()\n            keys.clear()\n\n    if strict:\n        if missing_keys := (local_state_dict_keys - global_keys):\n            for key in missing_keys:\n                local_state_dict.pop(key)\n\n    if keys:\n        _broadcast_tensors(ret, local_state_dict, keys, device, pg)\n        if cpu_offload:\n            for key in keys:\n                local_state_dict[key] = local_state_dict[key].cpu()\n\n\ndef _distribute_state_dict(\n    full_state_dict: dict[str, Any],\n    local_state_dict: dict[str, Any],\n    device: torch.device,\n    pg: Optional[dist.ProcessGroup] = None,\n) -> None:\n    # Full_state_dict = True, broadcast_from_rank0 = False here. Each rank has\n    # full_state_dict. Skip the broadcast in ``_broadcast_state_dict`` and\n    # distribute tensors in each rank\n    for key, value in full_state_dict.items():\n        if key not in full_state_dict:\n            continue\n        if not torch.is_tensor(value):\n            local_state_dict[key] = value\n        elif value.dim() == 0:\n            local_state_dict[key] = value.cpu()\n        else:\n            assert isinstance(value, torch.Tensor)\n            local_state = local_state_dict.get(key, None)\n            if local_state is None:\n                continue\n            elif isinstance(local_state, DTensor):\n                local_state_dict[key] = distribute_tensor(\n                    value.detach().to(device),\n                    local_state.device_mesh,\n                    local_state.placements,\n                )\n            else:\n                local_state_dict[key] = value.detach().to(device)\n\n\n# These APIs are from torch.distributed.checkpoint.\n# TODO: We should consolidate the code here as some not all modules can depend on\n# DCP.\nPATH_ITEM = Union[str, int]\nOBJ_PATH = tuple[PATH_ITEM, ...]\nFLATTEN_MAPPING = dict[str, OBJ_PATH]\nSTATE_DICT_TYPE = dict[str, Any]\nCONTAINER_TYPE = MutableMapping[PATH_ITEM, Any]\n\n\ndef _traverse_state_dict(\n    state_dict: STATE_DICT_TYPE,\n    visitor: Callable[[OBJ_PATH, Any], None],\n) -> None:\n    \"\"\"\n    Invoke ``visitor`` for each value recursively in ``state_dict``.\n    Mapping, list, and tuple will be flattened and other value types are treated\n    as the terminal values and will invoke ``visitor``.\n    \"\"\"\n\n    def _traverse_obj(path: OBJ_PATH, value: Any) -> None:\n        if isinstance(value, Mapping):\n            for k, v in value.items():\n                _traverse_obj(path + (str(k),), v)\n        elif isinstance(value, (list, tuple)):\n            for i, v in enumerate(value):\n                _traverse_obj(path + (i,), v)\n        else:\n            visitor(path, value)\n\n    for key, value in state_dict.items():\n        _traverse_obj((str(key),), value)\n\n\ndef _flatten_state_dict(\n    state_dict: STATE_DICT_TYPE,\n) -> tuple[STATE_DICT_TYPE, FLATTEN_MAPPING]:\n    \"\"\"\n    Flatten ``state_dict`` made of nested dicts and lists into a top level dictionary.\n\n    Use ``unflatten_state_dict`` to revert this process.\n    Returns:\n        A tuple with the flatten state_dict and a mapping from original to new state_dict.\n    N.B. The new keys are derived from the object paths, joined by dot.\n        For example: ``{ 'a': {'b':...}}`` results in the key `a.b`.\n    \"\"\"\n    flattened: STATE_DICT_TYPE = {}\n    mappings: FLATTEN_MAPPING = {}\n\n    def flat_copy(path: OBJ_PATH, value: Any) -> None:\n        new_fqn = \".\".join(map(str, path))\n        if new_fqn in flattened:\n            raise ValueError(f\"duplicated flatten key {new_fqn}\")\n        flattened[new_fqn] = value\n        mappings[new_fqn] = path\n\n    _traverse_state_dict(state_dict, flat_copy)\n    return flattened, mappings\n\n\ndef _set_element(root_dict: STATE_DICT_TYPE, path: OBJ_PATH, value: Any) -> None:\n    \"\"\"Set ``value`` in ``root_dict`` along the ``path`` object path.\"\"\"\n    cur_container = cast(CONTAINER_TYPE, root_dict)\n\n    def extend_list(lst: list[Any], idx: int) -> None:\n        while len(lst) <= idx:\n            lst.append(None)\n\n    for i in range(1, len(path)):\n        prev_key = path[i - 1]\n        key = path[i]\n        def_val: CONTAINER_TYPE | list[Any] = {} if type(key) == str else []\n\n        if isinstance(cur_container, Mapping):\n            cur_container = cast(CONTAINER_TYPE, cur_container.setdefault(prev_key, def_val))\n        else:\n            extend_list(cur_container, prev_key)\n            if cur_container[prev_key] is None:\n                cur_container[prev_key] = def_val\n            cur_container = cur_container[prev_key]\n\n    key = path[-1]\n    if type(key) == int:\n        extend_list(cast(list[Any], cur_container), key)\n\n    cur_container[key] = value\n\n\ndef _unflatten_state_dict(state_dict: STATE_DICT_TYPE, mapping: FLATTEN_MAPPING) -> STATE_DICT_TYPE:\n    \"\"\"Restore the original nested state_dict according to ``mapping`` and the flattened ``state_dict``.\"\"\"\n    nested: STATE_DICT_TYPE = {}\n    for key, value in state_dict.items():\n        _set_element(nested, mapping[key], value)\n    return nested\n"
  },
  {
    "path": "verl/third_party/torch/distributed/checkpoint/__init__.py",
    "content": "# official torch 2.6.0 set_model_state_dict API leads to OOM\n# this is a copy of torch/distributed/checkpoint from torch 2.7.0\n\n# From PyTorch:\n\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n\n# From Caffe2:\n\n# Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n\n# All contributions by Facebook:\n# Copyright (c) 2016 Facebook Inc.\n\n# All contributions by Google:\n# Copyright (c) 2015 Google Inc.\n# All rights reserved.\n\n# All contributions by Yangqing Jia:\n# Copyright (c) 2015 Yangqing Jia\n# All rights reserved.\n\n# All contributions by Kakao Brain:\n# Copyright 2019-2020 Kakao Brain\n\n# All contributions by Cruise LLC:\n# Copyright (c) 2022 Cruise LLC.\n# All rights reserved.\n\n# All contributions by Tri Dao:\n# Copyright (c) 2024 Tri Dao.\n# All rights reserved.\n\n# All contributions by Arm:\n# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates\n\n# All contributions from Caffe:\n# Copyright(c) 2013, 2014, 2015, the respective contributors\n# All rights reserved.\n\n# All other contributions:\n# Copyright(c) 2015, 2016 the respective contributors\n# All rights reserved.\n\n# Caffe2 uses a copyright model similar to Caffe: each contributor holds\n# copyright over their contributions to Caffe2. The project versioning records\n# all such contribution and copyright details. If a contributor wants to further\n# mark their specific copyright on a particular contribution, they should\n# indicate their copyright solely in the commit message of the change when it is\n# committed.\n\n# All rights reserved.\n\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n\n# 1. Redistributions of source code must retain the above copyright\n#    notice, this list of conditions and the following disclaimer.\n\n# 2. Redistributions in binary form must reproduce the above copyright\n#    notice, this list of conditions and the following disclaimer in the\n#    documentation and/or other materials provided with the distribution.\n\n# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n#    and IDIAP Research Institute nor the names of its contributors may be\n#    used to endorse or promote products derived from this software without\n#    specific prior written permission.\n\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n"
  },
  {
    "path": "verl/third_party/torch/distributed/checkpoint/state_dict.py",
    "content": "# official torch 2.6.0 set_model_state_dict API leads to OOM\n# this is a copy of torch/distributed/checkpoint from torch 2.7.0\n\n# From PyTorch:\n\n# Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)\n# Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)\n# Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)\n# Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)\n# Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)\n# Copyright (c) 2011-2013 NYU                      (Clement Farabet)\n# Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)\n# Copyright (c) 2006      Idiap Research Institute (Samy Bengio)\n# Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)\n\n# From Caffe2:\n\n# Copyright (c) 2016-present, Facebook Inc. All rights reserved.\n\n# All contributions by Facebook:\n# Copyright (c) 2016 Facebook Inc.\n\n# All contributions by Google:\n# Copyright (c) 2015 Google Inc.\n# All rights reserved.\n\n# All contributions by Yangqing Jia:\n# Copyright (c) 2015 Yangqing Jia\n# All rights reserved.\n\n# All contributions by Kakao Brain:\n# Copyright 2019-2020 Kakao Brain\n\n# All contributions by Cruise LLC:\n# Copyright (c) 2022 Cruise LLC.\n# All rights reserved.\n\n# All contributions by Tri Dao:\n# Copyright (c) 2024 Tri Dao.\n# All rights reserved.\n\n# All contributions by Arm:\n# Copyright (c) 2021, 2023-2024 Arm Limited and/or its affiliates\n\n# All contributions from Caffe:\n# Copyright(c) 2013, 2014, 2015, the respective contributors\n# All rights reserved.\n\n# All other contributions:\n# Copyright(c) 2015, 2016 the respective contributors\n# All rights reserved.\n\n# Caffe2 uses a copyright model similar to Caffe: each contributor holds\n# copyright over their contributions to Caffe2. The project versioning records\n# all such contribution and copyright details. If a contributor wants to further\n# mark their specific copyright on a particular contribution, they should\n# indicate their copyright solely in the commit message of the change when it is\n# committed.\n\n# All rights reserved.\n\n# Redistribution and use in source and binary forms, with or without\n# modification, are permitted provided that the following conditions are met:\n\n# 1. Redistributions of source code must retain the above copyright\n#    notice, this list of conditions and the following disclaimer.\n\n# 2. Redistributions in binary form must reproduce the above copyright\n#    notice, this list of conditions and the following disclaimer in the\n#    documentation and/or other materials provided with the distribution.\n\n# 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America\n#    and IDIAP Research Institute nor the names of its contributors may be\n#    used to endorse or promote products derived from this software without\n#    specific prior written permission.\n\n# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS \"AS IS\"\n# AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE\n# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE\n# ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE\n# LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR\n# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF\n# SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS\n# INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN\n# CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)\n# ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE\n# POSSIBILITY OF SUCH DAMAGE.\n\n# ruff: noqa: B028, UP038, UP007, E721\n# mypy: allow-untyped-defs\nimport contextlib\nimport functools\nimport gc\nimport warnings\nfrom collections.abc import Generator, Iterable\nfrom dataclasses import asdict, dataclass, field\nfrom itertools import chain\nfrom typing import Any, Callable, Optional, Union, cast, no_type_check\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nfrom torch.distributed._shard.sharded_tensor import ShardedTensor\nfrom torch.distributed.algorithms._checkpoint.checkpoint_wrapper import (\n    _CHECKPOINT_PREFIX,\n)\nfrom torch.distributed.fsdp import (\n    FullOptimStateDictConfig,\n    FullStateDictConfig,\n    OptimStateDictConfig,\n    ShardedOptimStateDictConfig,\n    ShardedStateDictConfig,\n    StateDictConfig,\n    StateDictType,\n)\nfrom torch.distributed.fsdp import (\n    FullyShardedDataParallel as FSDP,\n)\nfrom torch.distributed.fsdp._common_utils import (\n    FSDP_WRAPPED_MODULE,\n    _get_module_fsdp_state_if_fully_sharded_module,\n)\nfrom torch.distributed.tensor import DTensor\nfrom torch.nn.modules.module import _IncompatibleKeys\nfrom torch.nn.parallel import DistributedDataParallel as DDP\nfrom torch.utils._pytree import tree_map_only\n\nfrom verl.third_party.torch.distributed._state_dict_utils import (\n    _broadcast_state_dict,\n    _distribute_state_dict,\n    _flatten_state_dict,\n    _gather_state_dict,\n    _offload_state_dict_to_cpu,\n    _unflatten_state_dict,\n)\n\n__all__ = [\n    \"FQNS_T\",\n    \"PrimitiveType\",\n    \"ValueType\",\n    \"DictValueType\",\n    \"ListDictValueType\",\n    \"OptimizerStateType\",\n    \"StateDictOptions\",\n    \"get_model_state_dict\",\n    \"get_optimizer_state_dict\",\n    \"get_state_dict\",\n    \"set_model_state_dict\",\n    \"set_optimizer_state_dict\",\n    \"set_state_dict\",\n]\n\n\n_FLAT_PARAM = \"_flat_param\"\n_PG = \"param_groups\"\n_PARAMS = \"params\"\n_STATE = \"state\"\n\nFQNS_T = set[str]\nPrimitiveType = Union[DTensor, ShardedTensor, torch.Tensor, int, float, str]\nValueType = Union[PrimitiveType, list[PrimitiveType], tuple[PrimitiveType], dict[str, \"ValueType\"]]\nDictValueType = dict[str, ValueType]\nListDictValueType = list[DictValueType]\nOptimizerStateType = dict[str, DictValueType | ListDictValueType]\n\n\n_patched_state_dict: set[Callable] = set()\n\n\n@contextlib.contextmanager\ndef _gc_context():\n    is_enabled = gc.isenabled()\n    gc.disable()\n    try:\n        yield\n    finally:\n        if is_enabled:\n            gc.enable()\n\n\n@dataclass\nclass StateDictOptions:\n    \"\"\"\n    This dataclass specifies how get_state_dict/set_state_dict will work.\n\n    - ``full_state_dict``: if this is set to True, all the tensors in the\n      returned state_dict will be gathered. No ShardedTensor and DTensor\n      will be in the returned state_dict.\n\n    - ``cpu_offload``: offload all the tensors to cpu. To prevent CPU OOM, if\n      ``full_state_dict`` is also true, then only the rank0 will get the\n      state_dict and all other ranks will get empty state_dict.\n\n    - ``ignore_frozen_params``: if the value is True, the returned state_dict\n      won't contain any frozen parameters -- the ``requires_grad`` is False.\n      The default value is False.\n\n    - ``keep_submodule_prefixes`` (deprecated): when ``submodules`` is not None, this option\n      indicates whether to keep the submodule prefixes from the state_dict keys.\n      or example, if the submodule is ``module.pretrain`` and the full FQN of\n      the parameter is ``pretrain.layer1.weight`` of the param. When this option\n      is True, the parameter's key in the returned state_dict will be\n      ``pretrain.layer1.weight``. If the options is False, the key will be\n      ``layer1.weight``.\n      Note that if ``keep_submodule_prefixes`` is False, there may be conflicted\n      FQNs, hence there should be only one submodule in ``submodules``.\n\n    - ``strict``: the ``strict`` option when ``set_state_dict`` calls\n      model.load_state_dict().\n\n    - ``broadcast_from_rank0``: when the option is True, rank0 should receive a\n       full state_dict and will broadcast the tensors in the state_dict/\n       optim_state_dict one by one to other ranks. Other ranks will receive\n       the tensors and shard according to the local shards in the model and\n       optimizer. ``full_state_dict`` must be set to True when using this option.\n       This option currently only supports DTensor, not the legacy ShardedTensor.\n    \"\"\"\n\n    full_state_dict: bool = False\n    cpu_offload: bool = False\n    ignore_frozen_params: bool = False\n    keep_submodule_prefixes: bool = True\n    strict: bool = True\n    broadcast_from_rank0: bool = False\n    flatten_optimizer_state_dict: bool = False\n    dsd_fqn_modifiers: str = \"_fqn_modifiers\"\n\n\n@dataclass\nclass _StateDictInfo(StateDictOptions):\n    fqn_param_mapping: dict[\n        str | torch.Tensor,\n        FQNS_T | torch.Tensor,\n    ] = field(default_factory=dict)\n    shared_params_mapping: dict[\n        str | torch.Tensor,\n        FQNS_T | torch.Tensor,\n    ] = field(default_factory=dict)\n    submodule_prefixes: set[str] = field(default_factory=set)\n    handle_model: bool = True\n    handle_optim: bool = True\n    fsdp_context: Callable = contextlib.nullcontext\n    fsdp_modules: list[nn.Module] = field(default_factory=list)\n\n\n@functools.cache\ndef _get_fqns(\n    model: nn.Module,\n    name: str,\n    dsd_fqn_modifiers: str = \"_fqn_modifiers\",\n    skip_ddp_prefix: bool = True,\n    skip_compiler_prefix: bool = True,\n) -> FQNS_T:\n    \"\"\"\n    This API is used to convert the name of a parameter to the FQNs. For FSDP\n    without `use_orig_params`, the name of FlatParameter can be mapped to\n    multiple original parameters. As a result, the return type of this function\n    is `set[str]`.\n\n    Args:\n        module (nn.Module): the root model.\n        name (str): the name\n        skip_ddp_prefix (bool): whether to skip DDP's `module` prefix\n\n    Returns:\n        The canonical FQNs based on the model traversal.\n    \"\"\"\n\n    # Remove the checkpoint prefix, if it exists.\n    name = name.replace(_CHECKPOINT_PREFIX, \"\")\n    if \".\" not in name:\n        return {name}\n\n    obj_names = name.split(\".\")\n    fqn_obj_names = []\n    curr_obj = model\n    for i, curr_obj_name in enumerate(obj_names):\n        if isinstance(curr_obj, DDP):\n            assert curr_obj_name == \"module\"\n            curr_obj = curr_obj.module\n            if not skip_ddp_prefix:\n                fqn_obj_names.append(curr_obj_name)\n        elif isinstance(curr_obj, FSDP):\n            if i < len(obj_names) - 1 and obj_names[i + 1] == _FLAT_PARAM:\n                prefix = \".\".join(fqn_obj_names)\n                flat_param = getattr(curr_obj, _FLAT_PARAM)\n                if prefix:\n                    prefix = f\"{prefix}.\"\n                return {f\"{prefix}{fqn}\" for fqn in flat_param._fqns}\n            curr_obj = getattr(curr_obj, FSDP_WRAPPED_MODULE)\n            if curr_obj_name != FSDP_WRAPPED_MODULE:\n                fqn_obj_names.append(curr_obj_name)\n                curr_obj = getattr(curr_obj, curr_obj_name)\n        elif isinstance(curr_obj, torch._dynamo.eval_frame.OptimizedModule):\n            assert curr_obj_name == \"_orig_mod\"\n            curr_obj = curr_obj._orig_mod\n            if not skip_compiler_prefix:\n                fqn_obj_names.append(curr_obj_name)\n        else:\n            # In some modeuls, _fqn_modifiers would not shown in the state_dict keys,\n            # skip them in the fqn to ensure load stat dict successfully for them.\n            if hasattr(curr_obj, dsd_fqn_modifiers):\n                if removed_fqn := getattr(curr_obj, dsd_fqn_modifiers)().get(curr_obj_name):\n                    if hasattr(curr_obj, removed_fqn):\n                        curr_obj = getattr(curr_obj, removed_fqn)\n            fqn_obj_names.append(curr_obj_name)\n            if curr_obj_name == nn.modules.module._EXTRA_STATE_KEY_SUFFIX:\n                if i != len(obj_names) - 1:\n                    raise RuntimeError(\"Expect `_extra_state` to be the last obj name\")\n            else:\n                curr_obj = getattr(curr_obj, curr_obj_name)\n\n    return {\".\".join(fqn_obj_names).replace(_CHECKPOINT_PREFIX, \"\")}\n\n\nclass _EXTRA_STATE:\n    pass\n\n\ndef _iterate_valid_model_state(model, dsd_fqn_modifiers=\"_fqn_modifiers\"):\n    visited_modules: set[nn.Module] = set()\n\n    def recurse(module: nn.Module, curr_fqn: str) -> Generator:\n        visited_modules.add(module)\n\n        curr_fqn = f\"{curr_fqn}.\" if curr_fqn else \"\"\n        for name, submodule in module.named_children():\n            if submodule in visited_modules:\n                continue\n            # if user have state_dict_hooks in their model, they can add the state_dict key changes\n            # at dsd_fqn_modifiers in input to align with the function of state_dict_hook\n            if hasattr(module, dsd_fqn_modifiers) and name in getattr(module, dsd_fqn_modifiers)().values():\n                # skip _fqn_modifiers here thus remove the last `.` added\n                new_fqn = curr_fqn[:-1]\n            else:\n                new_fqn = f\"{curr_fqn}{name}\"\n            yield from recurse(submodule, new_fqn)\n\n        for name, obj in chain(module.named_buffers(recurse=False), module.named_parameters(recurse=False)):\n            if name in module._non_persistent_buffers_set:\n                continue\n            new_fqn = f\"{curr_fqn}{name}\"\n            yield new_fqn, obj\n\n        if getattr(module.__class__, \"get_extra_state\", nn.Module.get_extra_state) != nn.Module.get_extra_state:\n            new_fqn = f\"{curr_fqn}{nn.modules.module._EXTRA_STATE_KEY_SUFFIX}\"\n            yield new_fqn, _EXTRA_STATE()\n\n    yield from recurse(model, \"\")\n\n\ndef _verify_options(\n    model: nn.Module,\n    optims: tuple[torch.optim.Optimizer, ...],\n    optim_only: bool,\n    *,\n    submodules: Optional[set[nn.Module]] = None,\n    options: Optional[StateDictOptions] = None,\n) -> _StateDictInfo:\n    \"\"\"\n    Verify the model and options passed by the user and generates _StateDictInfo.\n    \"\"\"\n    if submodules:\n        warnings.warn(\n            \"Getting submodules only model/optim state_dict is deprecated and \"\n            \"will be removed in 2.5. This feature can be achieved by manually \"\n            \"filtering out the state_dict returned from get_state_dict.\",\n            FutureWarning,\n        )\n    if optim_only and not optims:\n        raise RuntimeError(\"Optimizers are not passed in but optim_only is set to True.\")\n\n    options = options or StateDictOptions()\n\n    fqn_param_mapping: dict[str | torch.Tensor, set[str] | torch.Tensor] = {}\n    shared_params_mapping: dict[str | torch.Tensor, set[str] | torch.Tensor] = {}\n    for name, param in _iterate_valid_model_state(model):\n        if isinstance(param, _EXTRA_STATE):\n            continue\n\n        fqns = _get_fqns(model, name)\n        fqn = fqn_param_mapping.get(param, None)\n        if fqn is not None:\n            cast(set[str], fqn_param_mapping[param]).update(fqns)\n            shared_params_mapping[param] = fqn_param_mapping[param]\n        else:\n            # We need to do copy as _get_fqns is lru_cached\n            fqn_param_mapping[param] = fqns.copy()\n        for fqn in fqns:\n            if not isinstance(param, _EXTRA_STATE):\n                fqn_param_mapping[fqn] = param\n\n    for param_, fqns_ in list(shared_params_mapping.items()):\n        for fqn in fqns_:\n            shared_params_mapping[fqn] = cast(torch.Tensor, param_)\n\n    submodule_prefixes: set[str] = set()\n    if submodules:\n        submodules = set(submodules)\n        for name, module in model.named_modules():\n            if module not in submodules:\n                continue\n            fqns = _get_fqns(model, name)\n            assert len(fqns) == 1, \"Submodule FQN should only have 1 instance\"\n            submodule_prefixes.update(f\"{fqn}.\" for fqn in fqns)\n\n    if options.broadcast_from_rank0 and not options.full_state_dict:\n        raise ValueError(\"full_state_dict must be True when broadcast_from_rank0 is True.\")\n    fsdp_modules = FSDP.fsdp_modules(model)\n    state_dict_config: StateDictConfig\n    optim_state_dict_config: OptimStateDictConfig\n    fsdp_context: Callable\n    if fsdp_modules:\n        # FSDP API only work if at least one FSDP instance exists.\n        if options.full_state_dict:\n            state_dict_config = FullStateDictConfig(offload_to_cpu=options.cpu_offload, rank0_only=options.cpu_offload)\n            optim_state_dict_config = FullOptimStateDictConfig(\n                offload_to_cpu=options.cpu_offload,\n                rank0_only=(options.cpu_offload or options.broadcast_from_rank0),\n            )\n            state_dict_type = StateDictType.FULL_STATE_DICT\n        else:\n            state_dict_config = ShardedStateDictConfig(\n                offload_to_cpu=options.cpu_offload,\n            )\n            optim_state_dict_config = ShardedOptimStateDictConfig(\n                offload_to_cpu=options.cpu_offload,\n            )\n            state_dict_type = StateDictType.SHARDED_STATE_DICT\n\n        @contextlib.contextmanager\n        def fsdp_state_dict_type_without_warning(\n            module,\n            state_dict_type,\n            state_dict_config,\n            optim_state_dict_config,\n        ):\n            with warnings.catch_warnings():\n                warnings.filterwarnings(\"ignore\", message=\"FSDP.state_dict_type\", category=FutureWarning)\n                with FSDP.state_dict_type(\n                    module=module,\n                    state_dict_type=state_dict_type,\n                    state_dict_config=state_dict_config,\n                    optim_state_dict_config=optim_state_dict_config,\n                ):\n                    yield\n\n        fsdp_context = functools.partial(\n            fsdp_state_dict_type_without_warning,\n            module=model,\n            state_dict_type=state_dict_type,\n            state_dict_config=state_dict_config,\n            optim_state_dict_config=optim_state_dict_config,\n        )\n    else:\n        fsdp_context = contextlib.nullcontext\n\n    return _StateDictInfo(\n        **asdict(options),\n        fqn_param_mapping=fqn_param_mapping,\n        shared_params_mapping=shared_params_mapping,\n        submodule_prefixes=submodule_prefixes,\n        fsdp_context=fsdp_context,\n        fsdp_modules=cast(list[nn.Module], fsdp_modules),\n        handle_model=not optim_only,\n        handle_optim=(len(optims) > 0),\n    )\n\n\ndef _verify_state_dict(\n    model_state_dict: dict[str, ValueType],\n    optim_state_dict: OptimizerStateType,\n    info: _StateDictInfo,\n) -> None:\n    for module in info.fsdp_modules:\n        fsdp_state = _get_module_fsdp_state_if_fully_sharded_module(module)\n        assert fsdp_state is not None, \"Expected a fsdp_state with a fsdp module.\"\n\n    # Verify if the model_state_dict and optim_state_dict are valid. This API\n    # should give the users an explicit error message to debug or report.\n    if (\n        info.handle_model\n        and not model_state_dict\n        and not info.submodule_prefixes\n        and not info.ignore_frozen_params\n        and not (info.cpu_offload and info.full_state_dict)\n        and info.strict\n        and not info.broadcast_from_rank0\n    ):\n        raise RuntimeError(\n            \"The option indicates that model state_dict is required to save \"\n            \"or load, but model state_dict is empty.\"\n            f\"rank = {dist.get_rank()=}.\"\n        )\n\n    if info.handle_optim:\n        if not optim_state_dict and not (info.cpu_offload and info.full_state_dict) and (not info.broadcast_from_rank0):\n            raise RuntimeError(\n                \"The option indicates that model state_dict is required to save, \"\n                f\"or load but optim state_dict is empty. {optim_state_dict}\"\n            )\n\n    for key in model_state_dict.keys():\n        if _FLAT_PARAM in key:\n            raise RuntimeError(f\"{key} contains {_FLAT_PARAM}. This can happen if the model is not the root module.\")\n\n\ndef _state_dict_fn(obj: nn.Module | torch.optim.Optimizer, api: str) -> Callable:\n    call = getattr(obj, api)\n    if call in _patched_state_dict:\n        call = functools.partial(getattr(obj.__class__, api), self=obj)\n    return call\n\n\ndef _maybe_full_or_cpu_state_dict(state_dict: dict[str, Any], info: _StateDictInfo) -> dict[str, Any]:\n    if info.full_state_dict:\n        ranks_only = () if (not info.cpu_offload or not torch.distributed.is_initialized()) else (0,)\n        return _gather_state_dict(state_dict, cpu_offload=info.cpu_offload, ranks_only=ranks_only)\n    elif info.cpu_offload:\n        return _offload_state_dict_to_cpu(state_dict)\n    else:\n        return state_dict\n\n\n@torch.no_grad()\ndef _get_model_state_dict(model: nn.Module, info: _StateDictInfo) -> dict[str, ValueType]:\n    if not info.handle_model:\n        return {}\n\n    with info.fsdp_context():\n        state_dict = _state_dict_fn(model, \"state_dict\")()\n\n    for key in list(state_dict.keys()):\n        fqns = _get_fqns(model, key)\n        assert len(fqns) == 1, (key, fqns)\n        fqn = next(iter(fqns))\n        if fqn != key:\n            # As we only support FSDP, DDP, and TP, the only cases are\n            # wrapper-based DDP and compiler. Verify if the assumption\n            # is correct.\n            def verify(key, fqn) -> bool:\n                if len(fqn) >= len(key):\n                    return False\n                fqn_split = fqn.split(\".\")\n                key_split = key.split(\".\")\n                fqn_idx = 0\n                for key_idx, key_name in enumerate(key_split):\n                    if key_name == fqn_split[fqn_idx]:\n                        fqn_idx += 1\n                        if fqn_idx == len(fqn_split):\n                            return key_idx == len(key_split) - 1\n                    elif key_name in (\"module\", \"_orig_mod\"):\n                        continue\n                    else:\n                        return False\n                return True\n\n            if not verify(key, fqn):\n                raise RuntimeError(f\"An unexpected key, {key}, exists. FQN is {fqn}\")\n            state_dict[fqn] = state_dict.pop(key)\n\n    if info.submodule_prefixes:\n        new_state_dict: dict[str, ValueType] = {}\n        # TODO: make this faster.\n        for fqn in state_dict.keys():\n            for prefix in info.submodule_prefixes:\n                if not fqn.startswith(prefix):\n                    continue\n                if info.keep_submodule_prefixes:\n                    new_state_dict[fqn] = state_dict[fqn]\n                else:\n                    new_fqn = fqn[len(prefix) :]\n                    new_state_dict[new_fqn] = state_dict[fqn]\n        state_dict = new_state_dict\n\n    if info.ignore_frozen_params:\n        for key, param in model.named_parameters():\n            if param.requires_grad:\n                continue\n            fqns = _get_fqns(model, key)\n            for fqn in fqns:\n                state_dict.pop(fqn)\n\n    for key, p in list(state_dict.items()):\n        if torch.is_tensor(p) and p.is_meta:\n            state_dict.pop(key)\n\n    return _maybe_full_or_cpu_state_dict(state_dict, info)\n\n\n@torch.no_grad()\ndef _load_model_state_dict(\n    model: nn.Module,\n    state_dict: dict[str, ValueType],\n    info: _StateDictInfo,\n) -> _IncompatibleKeys:\n    if not info.handle_model or (not state_dict and not info.broadcast_from_rank0):\n        return _IncompatibleKeys({}, {})\n\n    local_state_dict = {}\n    for key, value in _iterate_valid_model_state(model, info.dsd_fqn_modifiers):\n        fqns = _get_fqns(model, key, info.dsd_fqn_modifiers)\n        fqns_with_prefix = _get_fqns(\n            model,\n            key,\n            info.dsd_fqn_modifiers,\n            skip_ddp_prefix=False,\n            skip_compiler_prefix=False,\n        )\n\n        for fqn, fqn_with_prefix in zip(fqns, fqns_with_prefix, strict=False):\n            if (not info.broadcast_from_rank0 or dist.get_rank() == 0) and fqn != fqn_with_prefix:\n                load_value = state_dict.pop(fqn, None)\n                if load_value is None:\n                    if info.strict:\n                        raise RuntimeError(f\"Missing key: {fqn}.\")\n                else:\n                    state_dict[fqn_with_prefix] = load_value\n            local_state_dict[fqn_with_prefix] = value\n\n    assign = False\n    if info.broadcast_from_rank0 or info.full_state_dict:\n        devices = set()\n        for key, value in local_state_dict.items():\n            if torch.is_tensor(value) and value.dim() > 0:\n                devices.add(value.device)\n        # In lora state_dict, there could be multiple devices, with meta device inside.\n        # Take the other device in the broadcast/distribtue, and set assign to True\n        if torch.device(\"meta\") in devices:\n            devices.remove(torch.device(\"meta\"))\n            assign = True\n        if len(devices) == 0:\n            devices.add(dist.distributed_c10d._get_pg_default_device())\n        elif len(devices) > 1:\n            raise ValueError(\"Multiple devices found\")\n\n        if info.broadcast_from_rank0:\n            _broadcast_state_dict(\n                state_dict,\n                local_state_dict,\n                device=devices.pop(),\n                strict=info.strict,\n                cpu_offload=info.cpu_offload,\n            )\n        elif info.full_state_dict:\n            _distribute_state_dict(state_dict, local_state_dict, device=devices.pop())\n        for fqn, local_state in local_state_dict.items():\n            state_dict[fqn] = local_state\n\n    with info.fsdp_context():\n        return cast(\n            _IncompatibleKeys,\n            _state_dict_fn(model, \"load_state_dict\")(state_dict=state_dict, strict=info.strict, assign=assign),\n        )\n\n\ndef _init_optim_state(optim: torch.optim.Optimizer) -> None:\n    \"\"\"\n    Initialize optim states by calling the step() with zero grads.\n    \"\"\"\n    if optim.state:\n        # The optimizer state is initialized.\n        return\n\n    # There are some stateless optimizers like SGD. These optimizer will\n    # not return in the above condition. So if gradients exist, we should also\n    # return. If gradients do not exist, the following initialization should\n    # not disturb SGD because the gradients and lr are both zero.\n    for param_group in optim.param_groups:\n        for param in param_group[_PARAMS]:\n            if param.grad is not None:\n                return\n\n    for param_group in optim.param_groups:\n        for param in param_group[_PARAMS]:\n            if param.requires_grad:\n                param.grad = torch.zeros_like(param)\n\n    # Some optimizers will update parameters regardless of grads due to lr, so\n    # make lr to zero when calling `step()`.\n    lrs = []\n    for param_group in optim.param_groups:\n        if \"lr\" in param_group:\n            lrs.append(param_group[\"lr\"])\n            param_group[\"lr\"] = torch.tensor(0.0) if isinstance(param_group[\"lr\"], torch.Tensor) else 0.0\n    optim.step(closure=None)\n    # Whether to recover the \"lr\" should not matter too much as we will\n    # restore checkpointing later.\n    for param_group in optim.param_groups:\n        if \"lr\" in param_group:\n            param_group[\"lr\"] = lrs.pop(0)\n    optim.zero_grad(set_to_none=True)\n\n\ndef _flatten_optim_state_dict(state_dict: OptimizerStateType) -> dict[str, ValueType]:\n    \"\"\"\n    This API flattens the optimizer state_dict to support optimizer resharding for\n    MPMD, e.g., pipeline parallelism.\n\n    Without the API, the original optimizer state_dict looks like:\n    {\n        \"state\": {\n            \"layer1.weight\": {\n                \"step\": 10, \"exp_avg\": SomeTensor, \"exp_avg_sq\": SomeTensor\n            },\n            \"layer2.weight\": {\n                \"step\": 10, \"exp_avg\": SomeTensor, \"exp_avg_sq\": SomeTensor\n            },\n        },\n        \"param_group\": [\n            {\n                \"lr\": 0.0,\n                \"betas\": (0.9, 0.95), ...,\n                \"params\": [\"layer1.weight\", \"layer2.weight\"]\n            }\n        ]\n    }\n\n    With this API, the optimizer state_dict looks like:\n    {\n        \"state.layer1.weight.step\": 10,\n        \"state.layer2.weight.step\": 10,\n        \"state.layer1.weight.exp_avg\": SomeTensor,\n        \"state.layer2.weight.exp_avg\": SomeTensor,\n        \"state.layer1.weight.exp_avg_sq\": SomeTensor,\n        \"state.layer2.weight.exp_avg_sq\": SomeTensor,\n        \"param_group.layer1.weight.lr\" : 0.1,\n        \"param_group.layer2.weight.lr\" : 0.1,\n        \"param_group.layer1.weight.betas\" : (0.9, 0.95),\n        \"param_group.layer2.weight.betas\" : (0.9, 0.95),\n    }\n\n    Note that if any of the value is a container, like the betas in the example,\n    this API won't flattent it.\n    \"\"\"\n\n    def _raise_if_type_not_supported(v):\n        if not isinstance(v, (torch.Tensor, int, float)):\n            raise NotImplementedError(\n                f\"Flattening optimizer state_dict only supports tensor, int, float states now. Type is {type(v)}.\"\n            )\n\n    ret: dict[str, ValueType] = {}\n    for fqn, state in cast(DictValueType, state_dict[_STATE]).items():\n        for k, v in cast(DictValueType, state).items():\n            _raise_if_type_not_supported(v)\n            ret[f\"{_STATE}.{fqn}.{k}\"] = v\n\n    for param_group in cast(ListDictValueType, state_dict[_PG]):\n        fqns = param_group.pop(_PARAMS)\n        for fqn in cast(list[str], fqns):\n            for k, v in param_group.items():\n                ret[f\"{_PG}.{fqn}.{k}\"] = v\n    return ret\n\n\ndef _unflatten_optim_state_dict(\n    optim: torch.optim.Optimizer,\n    state_dict: dict[str, ValueType],\n    info: _StateDictInfo,\n) -> OptimizerStateType:\n    \"\"\"\n    This API unflattens the state_dict generated by _flatten_optim_state_dict().\n    See the docstring of _flatten_optim_state_dict() for more detail.\n    \"\"\"\n    state: DictValueType = {}\n    pg_state: ListDictValueType = []\n    return_osd: OptimizerStateType = {_STATE: state, _PG: pg_state}\n\n    for param_group in optim.param_groups:\n        pg_state.append({_PARAMS: []})\n        for param in param_group[_PARAMS]:\n            for fqn in info.fqn_param_mapping[param]:\n                # If a parameter is shared, only one of the FQN will be used.\n                # So we need to verify which if this fqn is actually used in\n                # the state_dict.\n                if fqn in info.shared_params_mapping:\n                    in_params = False\n                    for k in param_group.keys():\n                        if k == _PARAMS:\n                            continue\n                        flatten_key = f\"{_PG}.{fqn}.{k}\"\n                        if flatten_key in state_dict:\n                            in_params = True\n                        break\n                else:\n                    in_params = True\n\n                if not in_params:\n                    continue\n\n                params = pg_state[-1][_PARAMS]\n                assert isinstance(params, list)  # typing\n                params.append(fqn)\n                if not param.requires_grad:\n                    continue\n                state[fqn] = {}\n                for state_name in optim.state[param].keys():\n                    cast(DictValueType, state[fqn])[state_name] = state_dict[f\"{_STATE}.{fqn}.{state_name}\"]\n\n        first_param_fqn = cast(list[str], pg_state[-1][_PARAMS])[0]\n        for k in param_group.keys():\n            if k == _PARAMS:\n                continue\n            value = state_dict[f\"{_PG}.{first_param_fqn}.{k}\"]\n            if k not in pg_state[-1]:\n                pg_state[-1][k] = value\n            elif pg_state[-1][k] != value:\n                raise RuntimeError(\n                    \"All the parameters in the same parameter group should have \"\n                    f\"the same saved param_group value. But {first_param_fqn}.{k} \"\n                    f\"is {value} while other(s) is {pg_state[-1][k]}.\"\n                )\n\n    return return_osd\n\n\n@torch.no_grad()\ndef _get_optim_state_dict(\n    model: nn.Module,\n    optimizers: tuple[torch.optim.Optimizer, ...],\n    info: _StateDictInfo,\n) -> OptimizerStateType:\n    if not info.handle_optim:\n        return {}\n\n    optim_state_dict: OptimizerStateType = {_STATE: {}, _PG: []}\n    for optim in optimizers:\n        _init_optim_state(optim)\n        osd = _state_dict_fn(optim, \"state_dict\")()\n        if info.fsdp_modules:\n            with info.fsdp_context():\n                osd = FSDP.optim_state_dict(model, optim, osd)\n\n            # We need to specially handle FlatParameter FSDP as\n            # FlatParameter FSDP converts the FQNs.\n            # There are no easy ways to do this conversion systematically.\n            # We can only use a string replacment without correctness check.\n            if not osd:\n                continue\n            for k in list(osd[_STATE].keys()):\n                if \"_orig_mod\" in k:\n                    osd[_STATE][k.replace(\"_orig_mod.\", \"\")] = osd[_STATE].pop(k)\n            for g in osd[_PG]:\n                params = [k.replace(\"_orig_mod.\", \"\") for k in g[_PARAMS]]\n                g[_PARAMS] = params\n        else:\n            params = list(chain.from_iterable(g[_PARAMS] for g in optim.param_groups))\n            param_pid_mapping = dict(zip(params, range(len(params)), strict=False))\n            fqn_pid_mapping = {}\n            for key, param in model.named_parameters():\n                fqns = _get_fqns(model, key)\n                assert len(fqns) == 1\n                fqn = next(iter(fqns))\n                if param not in param_pid_mapping:\n                    continue\n                pid = param_pid_mapping[param]\n                fqn_pid_mapping[fqn] = pid\n                fqn_pid_mapping[pid] = fqn\n\n            for key in list(osd[_STATE].keys()):\n                fqn = fqn_pid_mapping[key]\n                osd[_STATE][fqn] = osd[_STATE].pop(key)\n\n            for group in osd[_PG]:\n                group[_PARAMS] = [fqn_pid_mapping[pid] for pid in group[_PARAMS]]\n\n        if not osd:\n            continue\n\n        cast(DictValueType, optim_state_dict[_STATE]).update(osd[_STATE])\n        cast(ListDictValueType, optim_state_dict[_PG]).extend(osd[_PG])\n\n    if info.flatten_optimizer_state_dict:\n        optim_state_dict = cast(OptimizerStateType, _flatten_optim_state_dict(optim_state_dict))\n\n    return _maybe_full_or_cpu_state_dict(optim_state_dict, info)\n\n\ndef _split_optim_state_dict(\n    model: nn.Module,\n    optim: torch.optim.Optimizer,\n    optim_state_dict: OptimizerStateType,\n    info: _StateDictInfo,\n) -> OptimizerStateType:\n    \"\"\"\n    Extract the corresponding optim state_dict from ``optim_state_dict`` for\n    ``optim`` and return the result optim state_dict.\n\n    Args:\n        model (nn.Module): the root model.\n        optim (torch.optim.Optimizer): the optimizer.\n        optim_state_dict (Dict[str, ValueType]): the superset optim state_dict that\n            contains the optim state_dict of ``optim``.\n        info (_StateDictInfo): state dict information.\n\n    Returns:\n        The optim state_dict of ``optim``.\n    \"\"\"\n\n    state: DictValueType = {}\n    pg_state: ListDictValueType = []\n    return_osd: OptimizerStateType = {_STATE: state, _PG: pg_state}\n    pg_mapping: dict[int, int] = {}\n\n    if all(isinstance(k, int) for k in cast(DictValueType, optim_state_dict[_STATE]).keys()):\n        return optim_state_dict\n\n    for param_group in optim.param_groups:\n        pg_state.append({_PARAMS: []})\n        for param in param_group[_PARAMS]:\n            for fqn in info.fqn_param_mapping[param]:\n                if fqn in info.shared_params_mapping:\n                    in_params = False\n                    for loaded_param_group in cast(ListDictValueType, optim_state_dict[_PG]):\n                        if fqn in cast(list[str], loaded_param_group[_PARAMS]):\n                            in_params = True\n                            break\n                else:\n                    in_params = True\n                if not in_params:\n                    continue\n\n                params = pg_state[-1][_PARAMS]\n                assert isinstance(params, list)\n                params.append(fqn)\n                if param.requires_grad:\n                    state[fqn] = cast(DictValueType, optim_state_dict[_STATE])[fqn]\n                for loaded_param_group in cast(ListDictValueType, optim_state_dict[_PG]):\n                    if fqn in cast(list[str], loaded_param_group[_PARAMS]):\n                        pg_mapping[id(loaded_param_group)] = len(return_osd[_PG]) - 1\n\n        if len(param_group[_PARAMS]) == 0:\n            # Param_group with empty params.\n            ret = []\n            for loaded_param_group in cast(ListDictValueType, optim_state_dict[_PG]):\n                if len(cast(list[str], loaded_param_group[_PARAMS])) == 0:\n                    ret.append(loaded_param_group)\n            if len(ret) != 1:\n                raise ValueError(\n                    \"There are param groups that have zero parameters. \"\n                    \"In such a case, DSD only support exactly one param group \"\n                    \"with zero parameters.\"\n                    \"But the loaded state_dict has zero or more than one param groups \"\n                    \"that have zero parameters.\"\n                )\n            if len(optim_state_dict[_PG]) != len(optim.param_groups):\n                raise ValueError(\n                    \"When there is a parameter group that has zero parameters, multiple optimizers are not supported.\"\n                )\n            pg_mapping[id(loaded_param_group)] = len(return_osd[_PG]) - 1\n\n    for param_group in cast(ListDictValueType, optim_state_dict[_PG]):\n        pg_idx = pg_mapping.get(id(param_group), -1)\n        if pg_idx == -1:\n            continue\n\n        for key, value in param_group.items():\n            if key == _PARAMS:\n                continue\n            # TODO: check if value is the same if exists.\n            pg_state[pg_idx][key] = value\n\n    return return_osd\n\n\n@torch.no_grad()\ndef _load_optim_state_dict(\n    model: nn.Module,\n    optimizers: tuple[torch.optim.Optimizer, ...],\n    state_dict: OptimizerStateType,\n    info: _StateDictInfo,\n) -> None:\n    if not info.handle_optim:\n        return\n\n    for optim in optimizers:\n        _init_optim_state(optim)\n        if state_dict:\n            if _STATE in state_dict:\n                optim_state_dict = _split_optim_state_dict(model, optim, state_dict, info)\n            else:\n                optim_state_dict = _unflatten_optim_state_dict(optim, cast(dict[str, ValueType], state_dict), info)\n        else:\n            optim_state_dict = {}\n        if info.fsdp_modules:\n            # We need to specially handle FlatParameter FSDP as\n            # FlatParameter FSDP converts the FQNs.\n            for original_fqn, _ in model.named_parameters():\n                fqns = _get_fqns(model, original_fqn)\n                fqns_with_compiler = _get_fqns(model, original_fqn, skip_compiler_prefix=False)\n                if fqns == fqns_with_compiler:\n                    continue\n\n                assert len(fqns) == 1\n                fqn = fqns.pop()\n                fqn_with_compiler = fqns_with_compiler.pop()\n                for g in optim_state_dict[_PG]:\n                    val = cast(dict[str, Any], g)\n                    params = [key.replace(fqn, fqn_with_compiler) for key in val[_PARAMS]]\n                    val[_PARAMS] = params\n                osd_state = cast(DictValueType, optim_state_dict[_STATE])\n                for k in list(osd_state.keys()):\n                    if fqn in k:\n                        osd_state[k.replace(fqn, fqn_with_compiler)] = osd_state.pop(k)\n\n            with info.fsdp_context():\n                optim_state_dict = FSDP.optim_state_dict_to_load(model, optim, optim_state_dict)\n        elif info.full_state_dict:\n            info.full_state_dict = False\n            local_state_dict = _get_optim_state_dict(model, (optim,), info)\n            info.full_state_dict = True\n            device = None\n\n            def _device(t):\n                if t.dim() > 0:\n                    nonlocal device\n                    if device is None:\n                        device = t.device\n                    elif device != t.device:\n                        raise ValueError(\"Device mismatch\")\n                return t\n\n            _ = tree_map_only(torch.Tensor, _device, local_state_dict)\n            assert device is not None\n            flatten_osd, osd_mapping = _flatten_state_dict(optim_state_dict)\n            flatten_local_osd, local_osd_mapping = _flatten_state_dict(local_state_dict)\n            if info.broadcast_from_rank0:\n                _broadcast_state_dict(flatten_osd, flatten_local_osd, device=device)\n            else:\n                _distribute_state_dict(flatten_osd, flatten_local_osd, device=device)\n            # The modifications listed seek to address the problem where optim might possess\n            # dissimilar parameters in comparison to optim_state_dict. This is achieved by\n            # incorporating differential parameters within local, which may result in optim\n            # having additional parameters ultimately.\n            for optim_key in flatten_osd.keys():\n                if optim_key not in flatten_local_osd:\n                    assert optim_key in osd_mapping\n                    flatten_local_osd[optim_key] = flatten_osd[optim_key]\n                    local_osd_mapping[optim_key] = osd_mapping[optim_key]\n            optim_state_dict = _unflatten_state_dict(flatten_local_osd, local_osd_mapping)\n            for pg in optim_state_dict[_PG]:\n                if _PARAMS not in pg:\n                    cast(dict[str, ValueType], pg)[_PARAMS] = []\n\n        # Note that we do not have to convert the FQN back to param id here if\n        # order in optim.param_groups[idx][_PARAMS] is the same as the one in\n        # optim_state_dict[_PG][idx][_PARAMS].\n        _state_dict_fn(optim, \"load_state_dict\")(state_dict=optim_state_dict)\n\n\ndef get_model_state_dict(\n    model: nn.Module,\n    *,\n    submodules: Optional[set[nn.Module]] = None,\n    options: Optional[StateDictOptions] = None,\n) -> dict[str, ValueType]:\n    \"\"\"\n    Return the model state_dict of ``model``.\n\n    See ``get_state_dict`` for the detail usage.\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        submodules (deprecated): Optional[set[nn.Module]]: only return the model parameters\n            that belong to the submodules.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be returned. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        The state_dict for ``model``.\n\n    :rtype: typing.Dict[str, ValueType]\n    \"\"\"\n    with _gc_context():\n        info = _verify_options(\n            model,\n            (),\n            optim_only=False,\n            submodules=submodules,\n            options=options,\n        )\n        model_state_dict = _get_model_state_dict(model, info)\n        _verify_state_dict(model_state_dict, {}, info)\n        return model_state_dict\n\n\ndef get_optimizer_state_dict(\n    model: nn.Module,\n    optimizers: torch.optim.Optimizer | Iterable[torch.optim.Optimizer],\n    *,\n    submodules: Optional[set[nn.Module]] = None,\n    options: Optional[StateDictOptions] = None,\n) -> OptimizerStateType:\n    \"\"\"\n    Return the combined state_dict for optimizers.\n\n    See ``get_state_dict`` for the detail usage.\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        optimizers (Union[None, Optimizer, Iterable[Optimizer]]):\n            The optimizers that are used to optimize ``model``.\n        submodules (deprecated): Optional[set[nn.Module]]: only return the model parameters\n            that belong to the submodules.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be returned. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        The state_dict for ``optimizers``.\n\n    :rtype: OptimizerStateType\n    \"\"\"\n    with _gc_context():\n        optimizers = (optimizers,) if isinstance(optimizers, torch.optim.Optimizer) else tuple(optimizers)\n        info = _verify_options(\n            model,\n            optimizers,\n            optim_only=True,\n            submodules=submodules,\n            options=options,\n        )\n        optim_state_dict = _get_optim_state_dict(model, optimizers, info)\n        _verify_state_dict({}, optim_state_dict, info)\n        return optim_state_dict\n\n\ndef get_state_dict(\n    model: nn.Module,\n    optimizers: torch.optim.Optimizer | Iterable[torch.optim.Optimizer],\n    *,\n    submodules: Optional[set[nn.Module]] = None,\n    options: Optional[StateDictOptions] = None,\n) -> tuple[dict[str, ValueType], OptimizerStateType]:\n    \"\"\"\n    Return the model state_dict and optimizers state_dict.\n\n    ``get_state_dict`` can process any module that is parallelized by PyTorch\n    FSDP/fully_shard, DDP/replicate, tensor_parallel/parallelize_module, and any\n    combination of these parallelisms. The main functions of ``get_state_dict``\n    are: 1.) returning a model and optimizer state_dict that can be resharded\n    with a different number of trainers and/or different parallelisms.\n    2.) hiding the parallelism-specific state_dict APIs. Users don't have to call\n    these APIs.\n    3.) sanity checking the result state_dict.\n\n    The keys of the result state dictionary are the canonical FQNs (Fully\n    Qualified Names).  A canonical FQN refers to the FQN based on a parameter's\n    position in an nn.Module hierarchy. More specifically, a canonical FQN to a\n    parameter is the FQN returned by ``module.named_parameters()`` or\n    ``module.named_buffers()`` when the module is not distributed by any\n    parallelisms. Since the optimizer internally uses parameter IDs to represent\n    a parameter, there will be a conversion from the parameter IDs to the\n    canonical FQNs when calling this API.\n\n    ``get_state_dict`` can also process a module that is not parallelized. In\n    such a case, ``get_state_dict`` only performs one function -- converting the\n    optimizer parameter IDs to the canonical FQNs.\n\n    Example:\n        >>> # xdoctest: +SKIP\n        >>> import torch\n        >>> from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n        >>> from torch.nn.parallel import DistributedDataParallel as DDP\n        >>> from torch.distributed.checkpoint.state_dict import get_state_dict\n\n        >>> fsdp_model = FSDP(copy.deepcopy(model))\n        >>> fsdp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)\n        >>> ddp_model = DDP(copy.deepcopy(model))\n        >>> ddp_optim = torch.optim.Adam(model.parameters(), lr=1e-3)\n\n\n        >>> ddp_state_dict, ddp_optim_state_dict = get_state_dict(ddp_model, ddp_optim)\n        >>> fsdp_state_dict, fsdp_optim_state_dict = get_state_dict(\n        ...     fsdp_model, fsdp_optim\n        ... )\n\n        >>> # if we simply call ddp_model.state_dict() and fsdp_model.state_dict(),\n        >>> # the asserts will fail.\n        >>> assert ddp_state_dict == fsdp_state_dict\n        >>> assert ddp_optim_state == fsdp_optim_state_dict\n\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        optimizers (Union[None, Optimizer, Iterable[Optimizer]]):\n            The optimizers that are used to optimize ``model``.\n        submodules (deprecated): Optional[set[nn.Module]]: only return the model parameters\n            that belong to the submodules.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be returned. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        ``Tuple`` that contain model state_dict and optimizer state_dict.\n\n    :rtype: typing.Tuple[typing.Dict[str, ValueType], OptimizerStateType]\n    \"\"\"\n\n    with _gc_context():\n        optimizers = (optimizers,) if isinstance(optimizers, torch.optim.Optimizer) else tuple(optimizers)\n        info = _verify_options(\n            model,\n            optimizers,\n            optim_only=False,\n            submodules=submodules,\n            options=options,\n        )\n        model_state_dict = _get_model_state_dict(model, info)\n        optim_state_dict = _get_optim_state_dict(model, optimizers, info)\n        _verify_state_dict(model_state_dict, optim_state_dict, info)\n        return model_state_dict, optim_state_dict\n\n\ndef _unflatten_model_state_dict(\n    model: nn.Module,\n    state_dict: dict[nn.Module, dict[str, ValueType]] | dict[str, ValueType],\n) -> dict[str, ValueType]:\n    if not state_dict:\n        return {}\n\n    if isinstance(next(iter(state_dict.keys())), nn.Module):\n        warnings.warn(\n            \"Passing model_state_dict as a ``Dict[nn.Module, Dict[str, Any]]``\"\n            \"is deprecated and will be removed in 2.5. If you need this \"\n            \"feature, please preprocessing the model_state_dict to achieve the \"\n            \"same functionality.\",\n            FutureWarning,\n        )\n        cast_state_dict = cast(dict[nn.Module, dict[str, ValueType]], state_dict)\n        new_state_dict: dict[str, ValueType] = {}\n        for submodule, sub_state_dict in cast_state_dict.items():\n            for name, m in model.named_modules():\n                if m != submodule:\n                    continue\n\n                fqns = _get_fqns(model, name)\n                assert len(fqns) == 1, \"FQNs for a submodule should only have 1 element\"\n                prefix = f\"{next(iter(fqns))}.\"\n                new_state_dict.update({prefix + subfqn: value for subfqn, value in sub_state_dict.items()})\n        return new_state_dict\n    else:\n        return cast(dict[str, ValueType], state_dict)\n\n\ndef set_model_state_dict(\n    model: nn.Module,\n    model_state_dict: dict[str, ValueType],\n    *,\n    options: Optional[StateDictOptions] = None,\n) -> _IncompatibleKeys:\n    \"\"\"Load the model state_dict.\n\n    The counterpart of ``get_model_state_dict`` to set the state_dict to the\n    model. See ``set_state_dict`` for the detail usage.\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        model_state_dict: (Dict[str, ValueType]):\n           the model state_dict to load. If the key of the ``model_state_dict``\n           is nn.Module, the key is a submodule of ``model`` and the value should\n           be the state_dict of the submodule. When loading the state_dict,\n           the prefix of the submodule will be append to the state_dict.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be loaded. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n            * **missing_keys** is a list of str containing the missing keys\n            * **unexpected_keys** is a list of str containing the unexpected keys\n\n    :type model_state_dict: typing.Dict[str, ValueType]\n    \"\"\"\n    model_state_dict: dict[str, ValueType] = _unflatten_model_state_dict(model, model_state_dict)\n    with _gc_context():\n        info = _verify_options(model, (), optim_only=False, options=options)\n\n        _verify_state_dict(model_state_dict, {}, info)\n        return _load_model_state_dict(model, model_state_dict, info)\n\n\ndef set_optimizer_state_dict(\n    model: nn.Module,\n    optimizers: torch.optim.Optimizer | Iterable[torch.optim.Optimizer],\n    optim_state_dict: OptimizerStateType,\n    *,\n    options: Optional[StateDictOptions] = None,\n) -> None:\n    \"\"\"Load the optimizers state_dict.\n\n    The counterpart of ``get_optimizer_state_dict`` to set the state_dict to the\n    optimizers. See ``set_state_dict`` for the detail usage.\n\n    WARN: ``set_optimizer_state_dict`` can only be called before ``backward()`` or after\n        ``step()`` is called on the optimizers. Otherwise, the optimizer states won't be\n        initialized correctly.\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        optimizers (Union[Optimizer, Iterable[Optimizer]]):\n            The optimizers that are used to optimize ``model``.\n        optim_state_dict: OptimizerStateType:\n            the optimizer state_dict to load.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be loaded. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        None\n\n    :type optim_state_dict: typing.OptimizerStateType\n    \"\"\"\n    with _gc_context():\n        optimizers = (optimizers,) if isinstance(optimizers, torch.optim.Optimizer) else tuple(optimizers)\n        info = _verify_options(model, optimizers, optim_only=True, options=options)\n\n        _verify_state_dict({}, optim_state_dict, info)\n        _load_optim_state_dict(model, optimizers, optim_state_dict, info)\n\n\ndef set_state_dict(\n    model: nn.Module,\n    optimizers: torch.optim.Optimizer | Iterable[torch.optim.Optimizer],\n    *,\n    model_state_dict: dict[str, ValueType],\n    optim_state_dict: OptimizerStateType,\n    options: Optional[StateDictOptions] = None,\n) -> _IncompatibleKeys:\n    \"\"\"Load the model state_dict and optimizers state_dict.\n\n    The counterpart of ``get_state_dict`` to set the state_dict to the model and\n    optimizers.  The given ``model_state_dict`` and ``optim_state_dict`` do not\n    have to be returned by ``get_state_dict`` but must meet the following\n    requirements: 1) all FQNs are canonical FQNs as defined in ``get_state_dict``,\n    2) if a tensor is sharded, it must be either a ShardedTensor or DTensor,\n    3) optimizer state_dict cannot contain the parameter IDs; the keys should be\n    the canonical FQNs.\n\n    WARN: ``set_state_dict`` can only be called before ``backward()`` or after ``step()``\n        is called on the optimizers. Otherwise, the optimizer states won't be initialized\n        correctly.\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        optimizers (Union[Optimizer, Iterable[Optimizer]]):\n            The optimizers that are used to optimize ``model``.\n        model_state_dict: (Union[Dict[nn.Module, Dict[str, ValueType]], Dict[str, ValueType]]):\n           the model state_dict to load. If the key of the ``model_state_dict``\n           is nn.Module, the key is a submodule of ``model`` and the value should\n           be the state_dict of the submodule. When loading the state_dict,\n           the prefix of the submodule will be append to the state_dict.\n        optim_state_dict: OptimizerStateType:\n            the optimizer state_dict to load.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be loaded. See\n            `StateDictOptions` for the details.\n\n    Returns:\n        ``NamedTuple`` with ``missing_keys`` and ``unexpected_keys`` fields:\n            * **missing_keys** is a list of str containing the missing keys of the model state_dict.\n            * **unexpected_keys** is a list of str containing the unexpected keys of the model state_dict.\n\n    :type model_state_dict: typing.Dict[str, ValueType]\n    :type optim_state_dict: typing.OptimizerStateType\n    \"\"\"\n\n    model_state_dict: dict[str, ValueType] = _unflatten_model_state_dict(model, model_state_dict)\n    with _gc_context():\n        optimizers = (optimizers,) if isinstance(optimizers, torch.optim.Optimizer) else tuple(optimizers)\n        info = _verify_options(model, optimizers, optim_only=not model_state_dict, options=options)\n\n        _verify_state_dict(model_state_dict, optim_state_dict, info)\n        _load_optim_state_dict(model, optimizers, optim_state_dict, info)\n        return _load_model_state_dict(model, model_state_dict, info)\n\n\n# TODO: correct the state_dict function signature.\n# TODO: this API is not yet fully tested. Make it private\n@no_type_check\ndef _patch_model_state_dict(\n    model: nn.Module,\n    *,\n    options: Optional[StateDictOptions] = None,\n) -> None:\n    \"\"\"Patch the ``state_dict`` and ``load_state_dict`` attributes of ``model``.\n\n    Patch the ``state_dict`` and ``load_state_dict`` attributes of ``model`` to\n    be a partial function to call ``get_state_dict`` and ``set_state_dict``.\n\n    Example:\n        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n        from torch.distributed.checkpoint.state_dict import patch_model_state_dict\n\n        model = fsdp(model)\n        patch_model_state_dict(model)\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be loaded. See\n            `StateDictOptions` for the details.\n    Returns:\n        None\n    \"\"\"\n\n    _state_dict_call = functools.partial(\n        get_model_state_dict,\n        model=model,\n        options=options,\n    )\n\n    def state_dict_call():\n        return _state_dict_call()\n\n    model.state_dict = state_dict_call\n\n    _load_state_dict_call = functools.partial(\n        set_model_state_dict,\n        model=model,\n        options=options,\n    )\n\n    def load_state_dict_call(state_dict: dict[str, Any]):\n        _load_state_dict_call(model_state_dict=state_dict)\n\n    model.load_state_dict = load_state_dict_call\n\n    _patched_state_dict.add(state_dict_call)\n    _patched_state_dict.add(load_state_dict_call)\n\n\n# TODO: correct the load_state_dict function signature.\n# TODO: this API is not yet fully tested. Make it private\n@no_type_check\ndef _patch_optimizer_state_dict(\n    model: nn.Module,\n    *,\n    optimizers: tuple[torch.optim.Optimizer, ...],\n    options: Optional[StateDictOptions] = None,\n) -> None:\n    \"\"\"Patch the ``state_dict`` and ``load_state_dict`` attributes of ``optimizers``.\n\n    Patch the ``state_dict`` and ``load_state_dict`` attributes of ``optimizers`` to\n    be a partial function to call ``get_state_dict`` and ``set_state_dict``.\n\n    Note that if there are multiple optimizers, all of the optimizers will be patched.\n    So users only need to call one of the state_dict() to get the full result.\n\n    Example:\n        from torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n        from torch.distributed.checkpoint.state_dict import patch_model_state_dict\n\n        model = fsdp(model)\n        patch_model_state_dict(model)\n\n    Args:\n        model (nn.Module): the nn.Module to the model.\n        options (StateDictOptions): the options to control how\n            model state_dict and optimizer state_dict should be loaded. See\n            `StateDictOptions` for the details.\n    Returns:\n        None\n    \"\"\"\n\n    _state_dict_call = functools.partial(\n        get_optimizer_state_dict,\n        model=model,\n        optimizers=optimizers,\n        options=options,\n    )\n\n    def state_dict_call():\n        return _state_dict_call()\n\n    _load_state_dict_call = functools.partial(\n        set_optimizer_state_dict,\n        model=model,\n        optimizers=optimizers,\n        options=options,\n    )\n\n    def load_state_dict_call(state_dict: dict[str, Any]):\n        _load_state_dict_call(optim_state_dict=state_dict)\n\n    _patched_state_dict.add(state_dict_call)\n    _patched_state_dict.add(load_state_dict_call)\n    optimizers = (optimizers,) if isinstance(optimizers, torch.optim.Optimizer) else tuple(optimizers)\n    for optim in optimizers:\n        optim.state_dict = state_dict_call\n        optim.load_state_dict = load_state_dict_call\n"
  },
  {
    "path": "verl/third_party/vllm/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom importlib.metadata import PackageNotFoundError, version\n\nfrom packaging import version as vs\n\nfrom verl.utils.device import is_npu_available\nfrom verl.utils.import_utils import is_sglang_available\n\n\ndef get_version(pkg):\n    try:\n        return version(pkg)\n    except PackageNotFoundError:\n        return None\n\n\npackage_name = \"vllm\"\npackage_version = get_version(package_name)\nvllm_version = None\nVLLM_SLEEP_LEVEL = 1\n\nif package_version is None:\n    if not is_sglang_available():\n        raise ValueError(\n            f\"vllm version {package_version} not supported and SGLang also not Found. Currently supported \"\n            f\"vllm versions are 0.7.0+\"\n        )\nelif is_npu_available:\n    # sleep_mode=2 is not supported on vllm-ascend for now, will remove this restriction when this ability is ready.\n    VLLM_SLEEP_LEVEL = 1\n    from vllm import LLM\n    from vllm.distributed import parallel_state\nelif vs.parse(package_version) >= vs.parse(\"0.7.0\"):\n    vllm_version = package_version\n    if vs.parse(package_version) >= vs.parse(\"0.8.5\"):\n        VLLM_SLEEP_LEVEL = 2\n    from vllm import LLM\n    from vllm.distributed import parallel_state\nelse:\n    if vs.parse(package_version) in [vs.parse(\"0.5.4\"), vs.parse(\"0.6.3\")]:\n        raise ValueError(\n            f\"vLLM version {package_version} support has been removed. vLLM 0.5.4 and 0.6.3 are no longer \"\n            f\"supported. Please use vLLM 0.7.0 or later.\"\n        )\n    if not is_sglang_available():\n        raise ValueError(\n            f\"vllm version {package_version} not supported and SGLang also not Found. Currently supported \"\n            f\"vllm versions are 0.7.0+\"\n        )\n\n__all__ = [\"LLM\", \"parallel_state\"]\n"
  },
  {
    "path": "verl/tools/__init__.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/tools/base_tool.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\n\nclass BaseTool:\n    \"\"\"Base class for tools.\n\n    A tool should support the following methods:\n\n    - `get_openai_tool_schema`: return the tool schema in OpenAI format.\n    - `create`: create a tool instance for a trajectory.\n    - `execute`: execute the tool.\n    - `calc_reward`: calculate the reward respect to tool state.\n    - `release`: release the tool instance.\n    \"\"\"\n\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        self.config = config\n        self.tool_schema = tool_schema or self.get_openai_tool_schema()\n        assert self.tool_schema is not None, \"Tool schema is not set!\"\n        self.name = self.tool_schema.function.name\n        print(json.dumps(self.tool_schema.model_dump(exclude_unset=True, exclude_none=True), indent=2))\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        return self.tool_schema\n\n    async def create(self, instance_id: Optional[str] = None, **kwargs) -> tuple[str, ToolResponse]:\n        \"\"\"Create a tool instance.\n\n        Args:\n            instance_id: The instance id of the tool.\n\n        Returns:\n            The instance id of the tool.\n            tool_creation_response: The response of the tool when creating the instance.\n        \"\"\"\n        if instance_id is None:\n            return str(uuid4()), ToolResponse()\n        else:\n            return instance_id, ToolResponse()\n\n    @rollout_trace_op\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        \"\"\"Execute the tool.\n\n        Args:\n            instance_id: The instance id of the tool.\n            parameters: The json string of the parameters of the tool.\n\n        Returns: tool_response, tool_reward_score, tool_metrics\n            tool_response: The ToolResponse object containing text, image, and/or video content.\n            tool_reward_score: The step reward score of the tool.\n            tool_metrics: The metrics of the tool.\n        \"\"\"\n        return ToolResponse(text=\"Updated the tool state.\"), 0.0, {}\n\n    async def calc_reward(self, instance_id: str, **kwargs) -> float:\n        \"\"\"Calculate the reward of the tool.\n\n        Args:\n            instance_id: The instance id of the tool.\n\n        Returns:\n            The reward of the tool.\n        \"\"\"\n        return 0.0\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        \"\"\"Release the tool instance.\n\n        Args:\n            instance_id: The instance id of the tool.\n        \"\"\"\n        pass\n"
  },
  {
    "path": "verl/tools/geo3k_tool.py",
    "content": "# Copyright 2023-2025 SGLang Team\n# Copyright Amazon.com, Inc. or its affiliates.\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom verl.utils.reward_score import geo3k\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nfrom .base_tool import BaseTool\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass Geo3kTool(BaseTool):\n    \"\"\"A demo tool for calculating the reward of geo3k.\n    - `get_openai_tool_schema`: return the tool schema in OpenAI format.\n    - `create`: create a tool instance for a trajectory.\n    - `execute`: execute the tool.\n    - `calc_reward`: calculate the reward respect to tool state.\n    - `release`: release the tool instance.\n    \"\"\"\n\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        \"\"\"\n        _tool_schema = OpenAIFunctionToolSchema.model_validate({\n            \"type\": \"function\",\n            \"function\": {\n                \"name\": \"calc_geo3k_reward\",\n                \"description\": \"A tool for calculating the reward of geo3k\",\n                \"parameters\": {\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"answer\": {\n                            \"type\": \"string\",\n                            \"description\": \"The answer to the question, enclosed in \\\\boxed{}\",\n                        },\n                    },\n                    \"required\": [\"answer\"],\n                },\n            }\n        })\n        \"\"\"\n        super().__init__(config, tool_schema)\n        self._instance_dict = {}\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        return self.tool_schema\n\n    async def create(\n        self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs\n    ) -> tuple[str, ToolResponse]:\n        if instance_id is None:\n            instance_id = str(uuid4())\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"ground_truth\": ground_truth,\n            \"reward\": 0.0,\n        }\n        return instance_id, ToolResponse()\n\n    @rollout_trace_op\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        answer = parameters.get(\"answer\", \"\")\n        if not isinstance(answer, str):\n            answer = str(answer)\n        self._instance_dict[instance_id][\"response\"] = answer\n        reward = await self.calc_reward(instance_id)\n        # penalty for non improved answer submission\n        tool_reward = 0.0 if reward > self._instance_dict[instance_id][\"reward\"] else -0.05\n        # update the reward\n        self._instance_dict[instance_id][\"reward\"] = reward\n        return ToolResponse(text=f\"Current parsed {answer=} {reward=}\"), tool_reward, {}\n\n    async def calc_reward(self, instance_id: str, **kwargs) -> float:\n        return geo3k.compute_score(\n            self._instance_dict[instance_id][\"response\"],\n            self._instance_dict[instance_id][\"ground_truth\"],\n            use_boxed=False,\n            format_score=0.0,\n        )\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/tools/gsm8k_tool.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom verl.utils.reward_score import gsm8k\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nfrom .base_tool import BaseTool\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass Gsm8kTool(BaseTool):\n    \"\"\"A demo tool for calculating the reward of gsm8k.\n\n    - `get_openai_tool_schema`: return the tool schema in OpenAI format.\n    - `create`: create a tool instance for a trajectory.\n    - `execute`: execute the tool.\n    - `calc_reward`: calculate the reward respect to tool state.\n    - `release`: release the tool instance.\n    \"\"\"\n\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        \"\"\"\n        _tool_schema = OpenAIFunctionToolSchema.model_validate({\n            \"type\": \"function\",\n            \"function\": {\n                \"name\": \"calc_gsm8k_reward\",\n                \"description\": \"A tool for calculating the reward of gsm8k\",\n                \"parameters\": {\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"answer\": {\n                            \"type\": \"string\",\n                            \"description\": \"The answer to the question\",\n                        },\n                    },\n                    \"required\": [\"answer\"],\n                },\n            }\n        })\n        \"\"\"\n        super().__init__(config, tool_schema)\n        self._instance_dict = {}\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        return self.tool_schema\n\n    async def create(\n        self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs\n    ) -> tuple[str, ToolResponse]:\n        if instance_id is None:\n            instance_id = str(uuid4())\n        if ground_truth is None:\n            ground_truth = kwargs.get(\"create_kwargs\", {}).get(\"ground_truth\", None)\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"ground_truth\": ground_truth,\n            \"reward\": 0.0,\n        }\n        return instance_id, ToolResponse()\n\n    @rollout_trace_op\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        answer = parameters.get(\"answer\", \"\")\n        if not isinstance(answer, str):\n            answer = str(answer)\n\n        if answer.startswith(\"#### \"):\n            self._instance_dict[instance_id][\"response\"] = answer\n        else:\n            self._instance_dict[instance_id][\"response\"] = \"#### \" + answer\n\n        reward = await self.calc_reward(instance_id)\n        # penalty for non improved answer submission\n        tool_reward = 0.0 if reward > self._instance_dict[instance_id][\"reward\"] else -0.05\n        # update the reward\n        self._instance_dict[instance_id][\"reward\"] = reward\n\n        return ToolResponse(text=f\"Current parsed {answer=} {reward=}\"), tool_reward, {}\n\n    async def calc_reward(self, instance_id: str, **kwargs) -> float:\n        return gsm8k.compute_score(\n            self._instance_dict[instance_id][\"response\"],\n            self._instance_dict[instance_id][\"ground_truth\"],\n            method=\"flexible\",\n            format_score=0.0,\n            score=1.0,\n        )\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/tools/image_zoom_in_tool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nimport threading\nfrom contextlib import ExitStack\nfrom enum import Enum\nfrom math import ceil, floor\nfrom typing import Any, Callable, Optional, TypeVar\nfrom uuid import uuid4\n\nimport ray\nimport ray.actor\nfrom qwen_vl_utils import fetch_image\n\nfrom .base_tool import BaseTool\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\nT = TypeVar(\"T\")\n\n\n# Adapted from verl/tools/sandbox_fusion_tools.py\nclass PoolMode(Enum):\n    \"\"\"Execution pool mode enumeration.\"\"\"\n\n    ThreadMode = 1\n    ProcessMode = 2\n\n\n@ray.remote(concurrency_groups={\"acquire\": 1, \"release\": 10})\nclass TokenBucketWorker:\n    \"\"\"Ray actor for rate limiting using token bucket algorithm.\"\"\"\n\n    def __init__(self, rate_limit: int):\n        self.rate_limit = rate_limit\n        self.current_count = 0  # For observability\n        self._semaphore = threading.Semaphore(rate_limit)\n\n    @ray.method(concurrency_group=\"acquire\")\n    def acquire(self):\n        \"\"\"Acquire a token from the bucket.\"\"\"\n        self._semaphore.acquire()\n        self.current_count += 1\n\n    @ray.method(concurrency_group=\"release\")\n    def release(self):\n        \"\"\"Release a token back to the bucket.\"\"\"\n        self._semaphore.release()\n        self.current_count -= 1\n\n    def get_current_count(self):\n        \"\"\"Get current number of acquired tokens.\"\"\"\n        return self.current_count\n\n\nclass VisualExecutionWorker:\n    \"\"\"Worker for executing visual processing operations with optional rate limiting.\"\"\"\n\n    def __init__(self, enable_global_rate_limit=True, rate_limit=10):\n        self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None\n\n    def _init_rate_limit(self, rate_limit):\n        \"\"\"Initialize singleton rate limiter.\"\"\"\n        return TokenBucketWorker.options(name=\"rate-limiter\", get_if_exists=True).remote(rate_limit)\n\n    def ping(self):\n        \"\"\"Health check method.\"\"\"\n        return True\n\n    def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:\n        \"\"\"Execute function with optional rate limiting.\"\"\"\n        if self.rate_limit_worker:\n            with ExitStack() as stack:\n                stack.callback(self.rate_limit_worker.release.remote)\n                ray.get(self.rate_limit_worker.acquire.remote())\n                try:\n                    return fn(*fn_args, **fn_kwargs)\n                except Exception as e:\n                    # TODO we should make this available to the tool caller\n                    logger.warning(f\"Error when executing visual processing: {e}\")\n        else:\n            return fn(*fn_args, **fn_kwargs)\n\n\ndef init_visual_execution_pool(\n    num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode = PoolMode.ThreadMode\n):\n    \"\"\"Initialize visual execution pool.\"\"\"\n    if mode == PoolMode.ThreadMode:\n        return (\n            ray.remote(VisualExecutionWorker)\n            .options(max_concurrency=num_workers)\n            .remote(enable_global_rate_limit=enable_global_rate_limit, rate_limit=rate_limit)\n        )\n    else:\n        raise NotImplementedError(\"Process mode is not implemented yet\")\n\n\nclass ImageZoomInTool(BaseTool):\n    \"\"\"A tool for zooming in on an image by cropping it based on a bounding box.\n\n    This tool provides a zoom-in functionality by cropping a region from an image,\n    with rate limiting and concurrent execution support through Ray.\n\n    Methods:\n        get_openai_tool_schema: Return the tool schema in OpenAI format\n        create: Create a tool instance for a trajectory\n        execute: Execute the zoom-in operation\n        calc_reward: Calculate the reward with respect to tool state\n        release: Release the tool instance\n    \"\"\"\n\n    MIN_DIMENSION = 28\n\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        \"\"\"\n        _tool_schema = OpenAIFunctionToolSchema.model_validate({\n            \"type\": \"function\",\n            \"function\": {\n                \"name\": \"image_zoom_in_tool\",\n                \"description\": (\n                    \"Zoom in on a specific region of an image by cropping it based on a bounding box (bbox) and an \"\n                    \"optional object label.\"\n                ),\n                \"parameters\": {\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"bbox_2d\": {\n                            \"type\": \"array\",\n                            \"items\":{\"type\":\"number\"},\n                            \"minItems\":4,\n                            \"maxItems\":4,\n                            \"description\": (\n                                \"The bounding box of the region to zoom in, as [x1, y1, x2, y2], where (x1, y1) is \"\n                                \"the top-left corner and (x2, y2) is the bottom-right corner.\"\n                            ),\n                        },\n                        \"label\": {\n                            \"type\": \"string\",\n                            \"description\": \"The name or label of the object in the specified bounding box (optional).\",\n                        },\n                    },\n                    \"required\": [\"bbox_2d\"],\n                },\n            }\n        })\n        \"\"\"\n        super().__init__(config, tool_schema)\n        self._instance_dict = {}\n\n        # Worker and rate limiting configuration\n        self.num_workers = config.get(\"num_workers\", 20)\n        self.rate_limit = config.get(\"rate_limit\", 50)\n        self.timeout = config.get(\"timeout\", 30)\n\n        self.enable_global_rate_limit = config.get(\"enable_global_rate_limit\", True)\n        self.execution_pool = init_visual_execution_pool(\n            num_workers=self.num_workers,\n            enable_global_rate_limit=self.enable_global_rate_limit,\n            rate_limit=self.rate_limit,\n            mode=PoolMode.ThreadMode,\n        )\n        logger.info(f\"Initialized ImageZoomInTool with config: {config}\")\n\n    def _validate_bbox(self, left: float, top: float, right: float, bottom: float) -> bool:\n        \"\"\"Validate the bounding box dimensions and aspect ratio.\"\"\"\n        try:\n            if not (left < right and top < bottom):\n                logger.warning(f\"Invalid bbox shape: left={left}, top={top}, right={right}, bottom={bottom}\")\n                return False\n\n            height = bottom - top\n            width = right - left\n\n            # Prevent division by zero for zero-sized boxes\n            if min(height, width) == 0:\n                logger.warning(f\"Bbox has zero width or height: left={left}, top={top}, right={right}, bottom={bottom}\")\n                return False\n\n            if max(height, width) / min(height, width) > 100:\n                logger.warning(f\"Bbox aspect ratio > 100: left={left}, top={top}, right={right}, bottom={bottom}\")\n                return False\n\n            return True\n        except Exception as e:\n            logger.warning(f\"Bbox validation error: {e}\")\n            return False\n\n    def _maybe_resize_bbox(self, bbox_2d: list[float], image_width: int, image_height: int) -> Optional[list[float]]:\n        \"\"\"\n        Clamp, validate, and potentially resize a bounding box.\n\n        This function ensures the final bounding box is within image bounds and meets the minimum\n        dimension requirements. If the initial box is too small, it attempts to expand it\n        from its center. It performs a final check to guarantee the output dimensions are valid.\n\n        Returns:\n            A valid bounding box as a list of coordinates, or None if validation fails.\n        \"\"\"\n        left, top, right, bottom = bbox_2d\n\n        # 1. Clamp the initial bounding box to the image dimensions.\n        left = max(0.0, float(left))\n        top = max(0.0, float(top))\n        right = min(float(image_width), float(right))\n        bottom = min(float(image_height), float(bottom))\n\n        # 2. If clamped bbox is invalid, return immediately.\n        if not self._validate_bbox(left, top, right, bottom):\n            return None\n\n        current_bbox = [left, top, right, bottom]\n        height = bottom - top\n        width = right - left\n\n        # 3. If the box is too small, attempt to resize it.\n        if height < self.MIN_DIMENSION or width < self.MIN_DIMENSION:\n            logger.info(f\"Bbox {width}x{height} is smaller than {self.MIN_DIMENSION}, attempting resize.\")\n            center_x = (left + right) / 2.0\n            center_y = (top + bottom) / 2.0\n\n            min_dim = min(height, width)\n            if min_dim == 0:  # Safeguard for zero-area boxes\n                return None\n\n            # 1. Calculate the target dimensions to make the smallest side MIN_DIMENSION.\n            ratio = self.MIN_DIMENSION / min_dim\n            target_width = width * ratio\n            target_height = height * ratio\n\n            # 2. If the target size is larger than the image, scale it down to fit.\n            #    This preserves the aspect ratio while respecting image boundaries.\n            if target_width > image_width:\n                scale_down = image_width / target_width\n                target_width = image_width\n                target_height *= scale_down\n\n            if target_height > image_height:\n                scale_down = image_height / target_height\n                target_height = image_height\n                target_width *= scale_down\n\n            # 3. Determine the coordinates for the box centered on the original center.\n            new_half_width = target_width / 2.0\n            new_half_height = target_height / 2.0\n            new_left = center_x - new_half_width\n            new_top = center_y - new_half_height\n\n            # 4. Shift the box if it extends beyond the image boundaries to keep its size.\n            if new_left < 0:\n                new_left = 0\n            if new_top < 0:\n                new_top = 0\n            if new_left + target_width > image_width:\n                new_left = image_width - target_width\n            if new_top + target_height > image_height:\n                new_top = image_height - target_height\n\n            new_right = new_left + target_width\n            new_bottom = new_top + target_height\n\n            # Use floor and ceil for final integer coordinates.\n            current_bbox = [floor(new_left), floor(new_top), ceil(new_right), ceil(new_bottom)]\n\n        # 4. Final validation on the resulting bounding box (either original or resized).\n        final_left, final_top, final_right, final_bottom = current_bbox\n        if not self._validate_bbox(final_left, final_top, final_right, final_bottom):\n            logger.warning(f\"Final bbox is invalid after processing: {current_bbox}\")\n            return None\n\n        final_height = floor(final_bottom) - floor(final_top)\n        final_width = floor(final_right) - floor(final_left)\n\n        if final_height < self.MIN_DIMENSION or final_width < self.MIN_DIMENSION:\n            logger.warning(\n                f\"Final bbox size ({final_width}x{final_height}) are still smaller than minimum ({self.MIN_DIMENSION}).\"\n                f\"Original bbox: {bbox_2d}, original image size: {image_width}x{image_height}\"\n            )\n            return None\n\n        return current_bbox\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        return self.tool_schema\n\n    async def create(self, instance_id: Optional[str] = None, **kwargs) -> tuple[str, ToolResponse]:\n        \"\"\"\n        Creates a new instance for image zoom-in tool.\n\n        This method initializes a new session for an image, which can then be used\n        for operations like zooming. It fetches the image from various sources\n        and stores it internally.\n\n        Args:\n            instance_id: An optional unique identifier for the instance. If not\n                provided, a new UUID will be generated.\n            **kwargs: Should contain 'image' key with image data, or 'create_kwargs'\n                containing {'image': image_data}. Image can be one of the following:\n                - A PIL.Image.Image object.\n                - A string containing an HTTP or HTTPS URL.\n                - A string containing a local file path.\n                - A string containing a file URI (e.g., \"file:///path/to/image.jpg\").\n                - A string containing a base64-encoded image in the format of \"data:image/jpeg;base64,...\"\n\n        Returns:\n            Tuple of (instance_id, ToolResponse)\n        \"\"\"\n        if instance_id is None:\n            instance_id = str(uuid4())\n\n        # Handle create_kwargs parameter if passed\n        create_kwargs = kwargs.get(\"create_kwargs\", {})\n        if create_kwargs:\n            kwargs.update(create_kwargs)\n\n        # Get image from kwargs\n        image = kwargs.get(\"image\")\n        if image is None:\n            raise ValueError(\"Missing required 'image' parameter in kwargs\")\n\n        img = fetch_image({\"image\": image})\n        self._instance_dict[instance_id] = {\n            \"image\": img,\n            \"response\": \"\",\n            \"reward\": 0.0,\n        }\n        return instance_id, ToolResponse()\n\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        bbox_2d = parameters.get(\"bbox_2d\")\n        label = parameters.get(\"label\", \"\")\n\n        if not bbox_2d or len(bbox_2d) != 4:\n            return (\n                ToolResponse(text=\"Error: bbox_2d parameter is missing or not a list of 4 numbers.\"),\n                -0.05,\n                {\"success\": False},\n            )\n\n        instance_data = self._instance_dict[instance_id]\n        image = instance_data[\"image\"]\n        image_width, image_height = image.size\n\n        try:\n            resized_bbox = self._maybe_resize_bbox(bbox_2d, image_width=image_width, image_height=image_height)\n\n            if resized_bbox is None:\n                error_msg = (\n                    f\"Error: The specified bounding box {bbox_2d} is invalid or results in a crop smaller than \"\n                    f\"the minimum size of {self.MIN_DIMENSION}x{self.MIN_DIMENSION}.\"\n                )\n                logger.warning(f\"Tool execution failed: {error_msg}\")\n                return ToolResponse(text=error_msg), -0.05, {\"success\": False}\n\n            cropped_image = image.crop(resized_bbox)\n            logger.info(f\"Cropped image size: {cropped_image.size}\")\n        except Exception as e:\n            logger.error(f\"Error processing image zoom-in: {e}\")\n            return ToolResponse(text=f\"Error processing image zoom-in: {e}\"), -0.05, {\"success\": False}\n\n        response_text = f\"Zoomed in on the image to the region {bbox_2d}.\"\n        if label:\n            response_text = f\"Zoomed in on the image to the region {bbox_2d} with label {label}.\"\n\n        return (\n            ToolResponse(\n                image=[cropped_image],\n                text=response_text,\n            ),\n            0.0,\n            {\"success\": True},\n        )\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        if instance_id in self._instance_dict:\n            del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/tools/mcp_base_tool.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport logging\nimport os\nfrom typing import Any, Optional\nfrom uuid import uuid4\n\nfrom fastmcp.exceptions import ClientError\n\nfrom verl.tools.utils.mcp_clients.McpClientManager import ClientManager\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nfrom .base_tool import BaseTool\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass MCPBaseTool(BaseTool):\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        super().__init__(config, tool_schema)\n        self._instance_dict = {}\n        self.timeout = config.get(\"timeout\", 30)\n\n        # TODO(hechanghao): create a global client manager to manage the rate limit, client and pool\n        logger.info(f\"Initialized MCPBaseTool with config: {config}\")\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        \"\"\"Return the OpenAI tool schema.\"\"\"\n        return self.tool_schema\n\n    async def create(self, instance_id: Optional[str] = None, **kwargs) -> tuple[str, ToolResponse]:\n        \"\"\"Create a tool instance.\n\n        Args:\n            instance_id: The instance id of the tool.\n\n        Returns:\n            The instance id of the tool.\n            tool_crtool_creation_response: The response of the tool when creating the instance.\n        \"\"\"\n        if instance_id is None:\n            instance_id = str(uuid4())\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"reward\": [],\n        }\n        return instance_id, ToolResponse()\n\n    async def _call_tool(self, instance_id, parameters) -> tuple[str, dict]:\n        err_msg = \"\"\n        metadata = {}\n        try:\n            call_tool_result = await ClientManager.call_tool(self.name, parameters, self.timeout)\n            logger.debug(f\"Tool result for instance {instance_id} with tool {self.name}: {call_tool_result.content}\")\n            result, metadata = self._parse_tool_result(call_tool_result.content)\n        except ClientError as e:\n            err_msg = f\"\\n Tool call failed: {e}\"\n        except ConnectionError as e:\n            err_msg = f\"\\n Connection failed: {e}\"\n        except Exception as e:\n            err_msg = f\"\\n An unexpected error occurred: {e}\"\n        finally:\n            if err_msg:\n                result = err_msg\n                metadata[\"api_request_error\"] = err_msg\n            else:\n                metadata[\"api_request_error\"] = None\n        return result, metadata\n\n    @rollout_trace_op\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        if self.name == \"\" or self.name is None or parameters is None:\n            error_msg = \"Error: 'parameters' is missing or empty.\"\n            logger.error(f\"[MCPTool] {error_msg} Received tool name: {self.name}, parameters: {parameters}\")\n            return ToolResponse(text=json.dumps({\"result\": error_msg})), 0.0, {}\n\n        try:\n            result_text, metadata = await self._call_tool(instance_id, parameters)\n\n            # Store results in instance dictionary\n            self._instance_dict[instance_id][\"reward\"].append(result_text.strip())\n\n            # Convert metadata to metrics\n            metrics = {\n                \"query_count\": metadata.get(\"query_count\", 0),\n                \"status\": metadata.get(\"status\", \"unknown\"),\n                \"total_results\": metadata.get(\"total_results\", 0),\n                \"api_request_error\": metadata.get(\"api_request_error\"),\n            }\n\n            return ToolResponse(text=result_text), 0.0, metrics\n\n        except Exception as e:\n            error_result = json.dumps({\"result\": f\"Tool execution failed: {e}\"})\n            logger.error(f\"[MCPBaseTool] Execution failed: {e}\")\n            return ToolResponse(text=error_result), 0.0, {\"error\": str(e)}\n\n    async def calc_reward(self, instance_id: str, **kwargs) -> str:\n        return self._instance_dict[instance_id][\"reward\"]\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        if instance_id in self._instance_dict:\n            del self._instance_dict[instance_id]\n\n    def _parse_tool_result(self, content: list) -> tuple[str, dict]:\n        tools_content = [part.text for part in filter(lambda x: x.type == \"text\", content)]\n        return \" \".join(tools_content), {}\n"
  },
  {
    "path": "verl/tools/mcp_search_tool.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport logging\nimport os\nimport re\n\nfrom verl.tools.mcp_base_tool import MCPBaseTool\n\nfrom .schemas import OpenAIFunctionToolSchema\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass MCPSearchTool(MCPBaseTool):\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        super().__init__(config, tool_schema)\n\n    def _parse_tool_result(self, content: list) -> tuple[str, dict]:\n        res = \"\"\n        res_cnt = 0\n        query_list = []\n        metadata = {\n            \"api_request_error\": \"\",\n            \"status\": \"unknown\",\n            \"total_results\": 0,\n        }\n        try:\n            for part in content:\n                if part.type != \"text\":\n                    continue\n                text = part.text.replace(\"'\", '\"')\n                query_match = re.search(r'query\"\\s*:\\s*\"([^\"]+)\"', text)\n                query = query_match.group(1) if query_match else \"\"\n                query_list.append(query)\n\n                title_matches = re.findall(r'\"title\"\\s*:', text)\n                title_count = len(title_matches)\n\n                results_match = re.search(r'\"results\"\\s*:\\s*(\\[.*?\\])', text, re.DOTALL)\n                results_content = results_match.group(1) if results_match else \"\"\n\n                res += results_content\n                res_cnt += title_count\n        except json.JSONDecodeError:\n            err_msg = \"json parse error.\"\n            logger.error(err_msg)\n            metadata[\"api_request_error\"] = err_msg\n            metadata[\"status\"] = \"error\"\n\n        # update metadata\n        metadata[\"status\"] = \"success\"\n        metadata[\"queries\"] = query_list\n        metadata[\"query_count\"] = len(query_list)\n        metadata[\"total_results\"] = res_cnt\n        return res, metadata\n"
  },
  {
    "path": "verl/tools/sandbox_fusion_tools.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nimport threading\nfrom contextlib import ExitStack\nfrom enum import Enum\nfrom typing import Any, Callable, Optional, TypeVar\nfrom uuid import uuid4\n\nimport ray\n\nfrom verl.tools.base_tool import BaseTool\nfrom verl.utils.reward_score.sandbox_fusion.utils import _process_single_case\nfrom verl.utils.rollout_trace import rollout_trace_op\n\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\nT = TypeVar(\"T\")\n\n\nclass PoolMode(Enum):\n    ThreadMode = 1\n    ProcessMode = 2\n\n\n@ray.remote(concurrency_groups={\"acquire\": 1, \"release\": 10})\nclass TokenBucketWorker:\n    def __init__(self, rate_limit: int):\n        self.rate_limit = rate_limit\n        # this only used for observalability\n        self.current_count = 0\n        self._semaphore = threading.Semaphore(rate_limit)\n\n    @ray.method(concurrency_group=\"acquire\")\n    def acquire(self):\n        self._semaphore.acquire()\n        self.current_count += 1\n\n    @ray.method(concurrency_group=\"release\")\n    def release(self):\n        self._semaphore.release()\n        self.current_count -= 1\n\n    def get_current_count(self):\n        return self.current_count\n\n\nclass ExecutionWorker:\n    def __init__(self, enable_global_rate_limit=True, rate_limit=10):\n        self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None\n\n    def _init_rate_limit(self, rate_limit):\n        # TODO validation for rate_limit\n        # A Singleton Rate Limitor\n        return TokenBucketWorker.options(name=\"rate-limiter\", get_if_exists=True).remote(rate_limit)\n\n    def ping(self):\n        return True\n\n    def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:\n        with ExitStack() as stack:\n            stack.callback(self.rate_limit_worker.release.remote)\n            ray.get(self.rate_limit_worker.acquire.remote())\n            try:\n                return fn(*fn_args, **fn_kwargs)\n            except Exception as e:\n                # TODO we should make this available to the tool caller\n                logger.warning(f\"Error when executing code: {e}\")\n\n\ndef init_execution_pool(\n    num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode = PoolMode.ThreadMode\n):\n    if mode == PoolMode.ThreadMode:\n        return (\n            ray.remote(ExecutionWorker)\n            .options(max_concurrency=num_workers)\n            .remote(enable_global_rate_limit=enable_global_rate_limit, rate_limit=rate_limit)\n        )\n    else:\n        raise NotImplementedError(\"Process mode is not implemented yet\")\n        # return ray.util.multiprocessing.Pool(processes=num_workers)\n\n\nclass SandboxFusionTool(BaseTool):\n    \"\"\"A tool for executing the code using sanbox fusion image.\n\n    - `get_openai_tool_schema`: return the tool schema in OpenAI format.\n    - `create`: create a tool instance for a trajectory.\n    - `execute`: execute the tool.\n    - `calc_reward`: calculate the reward respect to tool state.\n    - `release`: release the tool instance.\n    \"\"\"\n\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\n        \"\"\"\n        _tool_schema = OpenAIFunctionToolSchema.model_validate({\n            \"type\": \"function\",\n            \"function\": {\n                \"name\": \"code_interpreter\",\n                \"description\": \"A tool for execute code\",\n                \"parameters\": {\n                    \"type\": \"object\",\n                    \"properties\": {\n                        \"code\": {\n                            \"type\": \"string\",\n                            \"description\": \"code needs to be execute and grad\",\n                        },\n                    },\n                    \"required\": [\"code\"],\n                },\n            }\n        })\n        \"\"\"\n        super().__init__(config, tool_schema)\n        self._instance_dict = {}\n        # TODO: better documentation for the config\n        self.num_workers = config.get(\"num_workers\", 10)\n        self.rate_limit = config.get(\"rate_limit\", 10)\n        self.default_timeout = config.get(\"default_timeout\", 30)\n        self.default_language = config.get(\"default_language\", \"python\")\n        self.enable_global_rate_limit = config.get(\"enable_global_rate_limit\", True)\n        self.execution_pool = init_execution_pool(\n            num_workers=self.num_workers,\n            enable_global_rate_limit=self.enable_global_rate_limit,\n            rate_limit=self.rate_limit,\n            mode=PoolMode.ThreadMode,\n        )\n        self.sandbox_fusion_url = config.get(\"sandbox_fusion_url\", \"\")\n        self.memory_limit_mb = config.get(\"memory_limit_mb\", 1024)\n        if self.sandbox_fusion_url == \"\":\n            raise ValueError(\"sandbox_fusion_url is not set\")\n        log_msg = f\"Init SandboxFusionTool with config: {config}\"\n        logger.info(log_msg)\n\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\n        return self.tool_schema\n\n    async def create(\n        self, instance_id: Optional[str] = None, ground_truth: Optional[str] = None, **kwargs\n    ) -> tuple[str, ToolResponse]:\n        if instance_id is None:\n            instance_id = str(uuid4())\n        self._instance_dict[instance_id] = {\n            \"response\": \"\",\n            \"ground_truth\": ground_truth,\n            \"reward\": [],\n        }\n        return instance_id, ToolResponse()\n\n    @rollout_trace_op\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\n        code = parameters.get(\"code\", \"\")\n        timeout = parameters.get(\"timeout\", self.default_timeout)\n        language = parameters.get(\"language\", self.default_language)\n        if not isinstance(code, str):\n            code = str(code)\n\n        result = await self.execution_pool.execute.remote(self.execute_code, instance_id, code, timeout, language)\n        # sandbox has no score or metrics, use Nones\n        if isinstance(result, ToolResponse):\n            return result, None, None\n        return ToolResponse(text=None if result is None else str(result)), None, None\n\n    def execute_code(self, instance_id, code, timeout=30, language=\"python\"):\n        result_status, metadata = _process_single_case(\n            0, None, None, self.sandbox_fusion_url, code, timeout, self.memory_limit_mb, language\n        )\n        # we should always expect this since we don't have correct answer\n        if metadata[\"run_status\"] == \"Finished\":\n            actual_output = metadata[\"stdout\"] + metadata[\"stderr\"]\n            logger.debug(f\"actual_output from sandbox fusion: {actual_output},{instance_id}\")\n            return ToolResponse(text=actual_output)\n        else:\n            return ToolResponse(text=\"no stdout here\")\n\n    async def calc_reward(self, instance_id: str, **kwargs) -> str:\n        return self._instance_dict[instance_id][\"reward\"]\n\n    async def release(self, instance_id: str, **kwargs) -> None:\n        del self._instance_dict[instance_id]\n"
  },
  {
    "path": "verl/tools/schemas.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nfrom typing import Any, Literal\n\nfrom pydantic import BaseModel, Field, model_validator\n\n\nclass OpenAIFunctionPropertySchema(BaseModel):\n    \"\"\"The schema of a parameter in OpenAI format.\"\"\"\n\n    type: str\n    description: str | None = None\n    enum: list[str] | None = None\n\n\nclass OpenAIFunctionParametersSchema(BaseModel):\n    \"\"\"The schema of parameters in OpenAI format.\"\"\"\n\n    type: str\n    properties: dict[str, OpenAIFunctionPropertySchema]\n    required: list[str]\n\n\nclass OpenAIFunctionSchema(BaseModel):\n    \"\"\"The schema of a function in OpenAI format.\"\"\"\n\n    name: str\n    description: str\n    parameters: OpenAIFunctionParametersSchema = Field(\n        default_factory=lambda: OpenAIFunctionParametersSchema(type=\"object\", properties={}, required=[])\n    )\n    strict: bool = False\n\n\nclass OpenAIFunctionToolSchema(BaseModel):\n    \"\"\"The schema of a tool in OpenAI format.\"\"\"\n\n    type: str\n    function: OpenAIFunctionSchema\n\n\nclass OpenAIFunctionParsedSchema(BaseModel):\n    \"\"\"The parsed schema of a tool in OpenAI format.\"\"\"\n\n    name: str\n    arguments: str  # JSON string\n\n\nclass OpenAIFunctionCallSchema(BaseModel):\n    \"\"\"The parsed schema of a tool in OpenAI format.\"\"\"\n\n    name: str\n    arguments: dict[str, Any]\n\n    @staticmethod\n    def from_openai_function_parsed_schema(\n        parsed_schema: OpenAIFunctionParsedSchema,\n    ) -> tuple[\"OpenAIFunctionCallSchema\", bool]:\n        has_decode_error = False\n        try:\n            arguments = json.loads(parsed_schema.arguments)\n        except json.JSONDecodeError:\n            arguments = {}\n            has_decode_error = True\n        # If the arguments is not a dict, it means the arguments is not a valid JSON string\n        if not isinstance(arguments, dict):\n            arguments = {}\n            has_decode_error = True\n\n        return OpenAIFunctionCallSchema(name=parsed_schema.name, arguments=arguments), has_decode_error\n\n\nclass OpenAIFunctionToolCall(BaseModel):\n    \"\"\"The tool call in OpenAI format.\"\"\"\n\n    id: str\n    type: Literal[\"function\"] = \"function\"\n    function: OpenAIFunctionCallSchema\n\n\nclass ToolResponse(BaseModel):\n    \"\"\"The response from a tool execution.\"\"\"\n\n    text: str | None = None\n    image: list[Any] | None = None\n    video: list[Any] | None = None\n\n    @model_validator(mode=\"before\")\n    @classmethod\n    def initialize_request(cls, values):\n        if \"image\" in values and not isinstance(values[\"image\"], list):\n            raise ValueError(\n                f\"Image must be a list, but got {type(values['image'])}. Please check the tool.execute(). \"\n                f\"For single images, wrap in a list: [image]. \"\n                f\"Example: {{'image': [img1]}} or {{'image': [img1, img2, ...]}}.\"\n            )\n        if \"video\" in values and not isinstance(values[\"video\"], list):\n            raise ValueError(\n                f\"Video must be a list, but got {type(values['video'])}. Please check the tool.execute(). \"\n                f\"For single videos, wrap in a list: [video]. \"\n                f\"Example: {{'video': [video1]}} or {{'video': [video1, video2, ...]}}.\"\n            )\n\n        return values\n\n    def is_empty(self) -> bool:\n        return not self.text and not self.image and not self.video\n\n    def is_text_only(self) -> bool:\n        return self.text and not self.image and not self.video\n"
  },
  {
    "path": "verl/tools/search_tool.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\r\n# Copyright 2023-2024 SGLang Team\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\nimport json\r\nimport logging\r\nimport os\r\nimport threading\r\nfrom contextlib import ExitStack\r\nfrom enum import Enum\r\nfrom typing import Any, Callable, Optional, TypeVar\r\nfrom uuid import uuid4\r\n\r\nimport ray\r\nimport ray.actor\r\n\r\nfrom verl.tools.utils.search_r1_like_utils import perform_single_search_batch\r\nfrom verl.utils.rollout_trace import rollout_trace_op\r\n\r\nfrom .base_tool import BaseTool\r\nfrom .schemas import OpenAIFunctionToolSchema, ToolResponse\r\n\r\nlogger = logging.getLogger(__name__)\r\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\r\n\r\nT = TypeVar(\"T\")\r\n\r\n\r\n# Adapted from verl/tools/sandbox_fusion_tools.py\r\nclass PoolMode(Enum):\r\n    \"\"\"Execution pool mode enumeration.\"\"\"\r\n\r\n    ThreadMode = 1\r\n    ProcessMode = 2\r\n\r\n\r\n@ray.remote(concurrency_groups={\"acquire\": 1, \"release\": 10})\r\nclass TokenBucketWorker:\r\n    \"\"\"Ray actor for rate limiting using token bucket algorithm.\"\"\"\r\n\r\n    def __init__(self, rate_limit: int):\r\n        self.rate_limit = rate_limit\r\n        self.current_count = 0  # For observability\r\n        self._semaphore = threading.Semaphore(rate_limit)\r\n\r\n    @ray.method(concurrency_group=\"acquire\")\r\n    def acquire(self):\r\n        \"\"\"Acquire a token from the bucket.\"\"\"\r\n        self._semaphore.acquire()\r\n        self.current_count += 1\r\n\r\n    @ray.method(concurrency_group=\"release\")\r\n    def release(self):\r\n        \"\"\"Release a token back to the bucket.\"\"\"\r\n        self._semaphore.release()\r\n        self.current_count -= 1\r\n\r\n    def get_current_count(self):\r\n        \"\"\"Get current number of acquired tokens.\"\"\"\r\n        return self.current_count\r\n\r\n\r\nclass SearchExecutionWorker:\r\n    \"\"\"Worker for executing search operations with optional rate limiting.\"\"\"\r\n\r\n    def __init__(self, enable_global_rate_limit=True, rate_limit=10):\r\n        self.rate_limit_worker = self._init_rate_limit(rate_limit) if enable_global_rate_limit else None\r\n\r\n    def _init_rate_limit(self, rate_limit):\r\n        \"\"\"Initialize singleton rate limiter.\"\"\"\r\n        return TokenBucketWorker.options(name=\"rate-limiter\", get_if_exists=True).remote(rate_limit)\r\n\r\n    def ping(self):\r\n        \"\"\"Health check method.\"\"\"\r\n        return True\r\n\r\n    def execute(self, fn: Callable[..., T], *fn_args, **fn_kwargs) -> T:\r\n        \"\"\"Execute function with optional rate limiting.\"\"\"\r\n        if self.rate_limit_worker:\r\n            with ExitStack() as stack:\r\n                stack.callback(self.rate_limit_worker.release.remote)\r\n                ray.get(self.rate_limit_worker.acquire.remote())\r\n                try:\r\n                    return fn(*fn_args, **fn_kwargs)\r\n                except Exception as e:\r\n                    # TODO we should make this available to the tool caller\r\n                    logger.warning(f\"Error when executing search: {e}\")\r\n        else:\r\n            return fn(*fn_args, **fn_kwargs)\r\n\r\n\r\ndef init_search_execution_pool(\r\n    num_workers: int, enable_global_rate_limit=True, rate_limit=10, mode: PoolMode = PoolMode.ThreadMode\r\n):\r\n    \"\"\"Initialize search execution pool.\"\"\"\r\n    if mode == PoolMode.ThreadMode:\r\n        return (\r\n            ray.remote(SearchExecutionWorker)\r\n            .options(max_concurrency=num_workers)\r\n            .remote(enable_global_rate_limit=enable_global_rate_limit, rate_limit=rate_limit)\r\n        )\r\n    else:\r\n        raise NotImplementedError(\"Process mode is not implemented yet\")\r\n\r\n\r\nclass SearchTool(BaseTool):\r\n    \"\"\"Search tool for retrieving information using external retrieval services.\r\n\r\n    This tool provides search functionality with rate limiting and concurrent execution\r\n    support through Ray. It integrates with external retrieval services to perform\r\n    semantic search operations.\r\n\r\n    Methods:\r\n        get_openai_tool_schema: Return the tool schema in OpenAI format\r\n        create: Create a tool instance for a trajectory\r\n        execute: Execute the search tool\r\n        calc_reward: Calculate the reward with respect to tool state\r\n        release: Release the tool instance\r\n    \"\"\"\r\n\r\n    def __init__(self, config: dict, tool_schema: OpenAIFunctionToolSchema):\r\n        \"\"\"Initialize SearchTool with configuration and schema.\r\n\r\n        Args:\r\n            config: Configuration dictionary containing tool settings\r\n            tool_schema: OpenAI function tool schema definition\r\n\r\n        Example tool_schema:\r\n            {\r\n                \"type\": \"function\",\r\n                \"function\": {\r\n                    \"name\": \"search\",\r\n                    \"description\": \"Searches for relevant information based on queries.\",\r\n                    \"parameters\": {\r\n                        \"type\": \"object\",\r\n                        \"properties\": {\r\n                            \"query_list\": {\r\n                                \"type\": \"array\",\r\n                                \"items\": {\"type\": \"string\"},\r\n                                \"description\": \"List of search queries\"\r\n                            }\r\n                        },\r\n                        \"required\": [\"query_list\"]\r\n                    }\r\n                }\r\n            }\r\n        \"\"\"\r\n        super().__init__(config, tool_schema)\r\n        self._instance_dict = {}\r\n\r\n        # Worker and rate limiting configuration\r\n        self.num_workers = config.get(\"num_workers\", 120)\r\n        self.rate_limit = config.get(\"rate_limit\", 120)\r\n        self.timeout = config.get(\"timeout\", 30)\r\n\r\n        self.enable_global_rate_limit = config.get(\"enable_global_rate_limit\", True)\r\n        self.execution_pool = init_search_execution_pool(\r\n            num_workers=self.num_workers,\r\n            enable_global_rate_limit=self.enable_global_rate_limit,\r\n            rate_limit=self.rate_limit,\r\n            mode=PoolMode.ThreadMode,\r\n        )\r\n\r\n        # Retrieval service configuration\r\n        self.retrieval_service_url = config.get(\"retrieval_service_url\")\r\n        assert self.retrieval_service_url, \"Configuration must include 'retrieval_service_url'\"\r\n        self.topk = config.get(\"topk\", 3)\r\n        if self.retrieval_service_url == \"\":\r\n            raise ValueError(\"retrieval_service_url is not set\")\r\n\r\n        logger.info(f\"Initialized SearchTool with config: {config}\")\r\n\r\n    def get_openai_tool_schema(self) -> OpenAIFunctionToolSchema:\r\n        \"\"\"Return the OpenAI tool schema.\"\"\"\r\n        return self.tool_schema\r\n\r\n    async def create(self, instance_id: Optional[str] = None, **kwargs) -> tuple[str, ToolResponse]:\r\n        \"\"\"Create a tool instance.\r\n\r\n        Args:\r\n            instance_id: The instance id of the tool.\r\n\r\n        Returns:\r\n            The instance id of the tool.\r\n            tool_creation_response: The response of the tool when creating the instance.\r\n        \"\"\"\r\n        if instance_id is None:\r\n            instance_id = str(uuid4())\r\n        self._instance_dict[instance_id] = {\r\n            \"response\": \"\",\r\n            \"reward\": [],\r\n        }\r\n        return instance_id, ToolResponse()\r\n\r\n    def execute_search(self, instance_id: str, query_list: list, retrieval_service_url: str, topk: int, timeout: int):\r\n        \"\"\"Execute search operation using retrieval service.\r\n\r\n        Args:\r\n            instance_id: Tool instance ID\r\n            query_list: List of search queries\r\n            retrieval_service_url: URL of the retrieval service\r\n            topk: Number of top results to return\r\n            timeout: Request timeout in seconds\r\n\r\n        Returns:\r\n            Tuple of (result_text, metadata)\r\n        \"\"\"\r\n        result_text, metadata = perform_single_search_batch(\r\n            retrieval_service_url=retrieval_service_url,\r\n            query_list=query_list,\r\n            topk=topk,\r\n            concurrent_semaphore=None,  # Ray handles concurrency control\r\n            timeout=timeout,\r\n        )\r\n        logger.debug(f\"Search result for instance {instance_id}: {result_text}\")\r\n        return result_text, metadata\r\n\r\n    @rollout_trace_op\r\n    async def execute(self, instance_id: str, parameters: dict[str, Any], **kwargs) -> tuple[ToolResponse, float, dict]:\r\n        \"\"\"Execute the search tool.\r\n\r\n        Args:\r\n            instance_id: The instance ID of the tool\r\n            parameters: Tool parameters containing query_list and optional timeout\r\n\r\n        Returns: tool_response, tool_reward_score, tool_metrics\r\n            tool_response: The response str of the tool.\r\n            tool_reward_score: The step reward score of the tool.\r\n            tool_metrics: The metrics of the tool.\r\n        \"\"\"\r\n        timeout = self.timeout\r\n        query_list_from_params = parameters.get(\"query_list\")\r\n\r\n        if not query_list_from_params or not isinstance(query_list_from_params, list):\r\n            error_msg = \"Error: 'query_list' is missing, empty, or not a list in parameters.\"\r\n            logger.error(f\"[SearchTool] {error_msg} Received parameters: {parameters}\")\r\n            return ToolResponse(text=json.dumps({\"result\": error_msg})), 0.0, {}\r\n\r\n        # Execute search using Ray execution pool\r\n        try:\r\n            result_text, metadata = await self.execution_pool.execute.remote(\r\n                self.execute_search, instance_id, query_list_from_params, self.retrieval_service_url, self.topk, timeout\r\n            )\r\n\r\n            # Store results in instance dictionary\r\n            self._instance_dict[instance_id][\"reward\"].append(result_text.strip())\r\n\r\n            # Convert metadata to metrics\r\n            metrics = {\r\n                \"query_count\": metadata.get(\"query_count\", 0),\r\n                \"status\": metadata.get(\"status\", \"unknown\"),\r\n                \"total_results\": metadata.get(\"total_results\", 0),\r\n                \"api_request_error\": metadata.get(\"api_request_error\"),\r\n            }\r\n\r\n            return ToolResponse(text=result_text), 0.0, metrics\r\n\r\n        except Exception as e:\r\n            error_result = json.dumps({\"result\": f\"Search execution failed: {e}\"})\r\n            logger.error(f\"[SearchTool] Execution failed: {e}\")\r\n            return ToolResponse(text=error_result), 0.0, {\"error\": str(e)}\r\n\r\n    async def calc_reward(self, instance_id: str, **kwargs) -> str:\r\n        return self._instance_dict[instance_id][\"reward\"]\r\n\r\n    async def release(self, instance_id: str, **kwargs) -> None:\r\n        if instance_id in self._instance_dict:\r\n            del self._instance_dict[instance_id]\r\n"
  },
  {
    "path": "verl/tools/utils/__init__.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/tools/utils/mcp_clients/McpClientManager.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\nimport asyncio\r\nimport json\r\nimport logging\r\nfrom typing import Any\r\n\r\nfrom fastmcp import Client\r\nfrom fastmcp.client.transports import SSETransport\r\n\r\nfrom verl.tools.utils.mcp_clients.utils import TokenBucket, mcp2openai\r\n\r\nlogger = logging.getLogger(__name__)\r\n\r\n\r\nclass MCPClientManager:\r\n    rootServerName = \"mcpServers\"\r\n    initialized = False\r\n    clients = []\r\n    tool_client_mapping = {}\r\n    rate_limiter = None\r\n\r\n    async def initialize(self, config_path, rate_limit: float = 10.0):\r\n        if self.initialized:\r\n            return\r\n        \"\"\"Initialize the MCP Client Manager and start all clients\"\"\"\r\n        result = self._load_config(config_path)\r\n        servers = result[self.rootServerName]\r\n        exclude_sse_servers = {self.rootServerName: {}}\r\n        for server_name in servers.keys():\r\n            server = servers[server_name]\r\n            if \"auth_token\" in server:\r\n                transport = SSETransport(url=server[\"url\"], headers={\"Authorization\": f\"Bearer {server['auth_token']}\"})\r\n                client = Client(transport)\r\n                self.clients.append(client)\r\n            else:\r\n                exclude_sse_servers[self.rootServerName][server_name] = server\r\n\r\n        if exclude_sse_servers[self.rootServerName]:\r\n            self.clients.append(Client(exclude_sse_servers))\r\n\r\n        # Initialize rate limiter\r\n        self.rate_limiter = TokenBucket(rate_limit)\r\n        self.initialized = True\r\n\r\n    async def call_tool(self, tool_name, parameters, timeout):\r\n        # Apply rate limiting\r\n        while not self.rate_limiter.acquire():\r\n            await asyncio.sleep(0.1)\r\n\r\n        client = self.get_client_with_tool_name(tool_name)\r\n        async with client:\r\n            return await client.call_tool_mcp(tool_name, parameters)\r\n\r\n    async def fetch_tool_schemas(self, tool_selected_list: list[str]) -> list[dict]:\r\n        tool_schemas = []\r\n        for client in self.clients:\r\n            async with client:\r\n                tools = await client.list_tools_mcp()\r\n                for tool in tools.tools:\r\n                    if not tool_selected_list:\r\n                        self.tool_client_mapping[tool.name] = client\r\n                        tool_schemas.append(mcp2openai(tool))\r\n                    elif tool.name in tool_selected_list:\r\n                        self.tool_client_mapping[tool.name] = client\r\n                        tool_schemas.append(mcp2openai(tool))\r\n\r\n        return tool_schemas\r\n\r\n    def get_client_with_tool_name(self, tool_name: str):\r\n        return self.tool_client_mapping[tool_name]\r\n\r\n    def _load_config(self, file: str) -> dict[str, Any]:\r\n        try:\r\n            with open(file) as f:\r\n                return json.load(f)\r\n        except FileNotFoundError:\r\n            logger.warning(f'the \"{file}\" file was not found')\r\n        except Exception:\r\n            logger.error(f'there was an error reading the \"{file}\" file')\r\n\r\n        return {}\r\n\r\n\r\nClientManager = MCPClientManager()\r\n"
  },
  {
    "path": "verl/tools/utils/mcp_clients/utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport threading\nimport time\n\nfrom mcp import Tool\n\nlogger = logging.getLogger(__file__)\n\n\nclass TokenBucket:\n    def __init__(self, rate_limit: float):\n        self.rate_limit = rate_limit  # tokens per second\n        self.tokens = rate_limit\n        self.last_update = time.time()\n        self.lock = threading.Lock()\n\n    def acquire(self) -> bool:\n        with self.lock:\n            now = time.time()\n            # Add new tokens based on time elapsed\n            new_tokens = (now - self.last_update) * self.rate_limit\n            self.tokens = min(self.rate_limit, self.tokens + new_tokens)\n            self.last_update = now\n\n            if self.tokens >= 1:\n                self.tokens -= 1\n                return True\n            return False\n\n\ndef mcp2openai(mcp_tool: Tool) -> dict:\n    \"\"\"Convert a MCP Tool to an OpenAI ChatCompletionTool.\"\"\"\n    openai_format = {\n        \"type\": \"function\",\n        \"function\": {\n            \"name\": mcp_tool.name,\n            \"description\": mcp_tool.description,\n            \"parameters\": mcp_tool.inputSchema,\n            \"strict\": False,\n        },\n    }\n    if not openai_format[\"function\"][\"parameters\"].get(\"required\", None):\n        openai_format[\"function\"][\"parameters\"][\"required\"] = []\n    return openai_format\n"
  },
  {
    "path": "verl/tools/utils/search_r1_like_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\r\n# Copyright 2023-2024 SGLang Team\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n\r\nimport json\r\nimport logging\r\nimport threading\r\nimport time\r\nimport traceback\r\nimport uuid\r\nfrom typing import Any, Optional\r\n\r\nimport requests\r\n\r\nDEFAULT_TIMEOUT = 30  # Default search request timeout\r\nMAX_RETRIES = 10\r\nINITIAL_RETRY_DELAY = 1\r\nAPI_TIMEOUT = 10\r\n\r\nlogger = logging.getLogger(__name__)\r\n\r\n\r\ndef call_search_api(\r\n    retrieval_service_url: str,\r\n    query_list: list[str],\r\n    topk: int = 3,\r\n    return_scores: bool = True,\r\n    timeout: int = DEFAULT_TIMEOUT,\r\n) -> tuple[Optional[dict[str, Any]], Optional[str]]:\r\n    \"\"\"\r\n    Calls the remote search API to perform retrieval with retry logic for various errors,\r\n    using increasing delay between retries. Logs internal calls with a unique ID.\r\n\r\n    Args:\r\n        retrieval_service_url: The URL of the retrieval service API.\r\n        query_list: List of search queries.\r\n        topk: Number of top results to return.\r\n        return_scores: Whether to return scores.\r\n        timeout: Request timeout in seconds.\r\n\r\n    Returns:\r\n        A tuple (response_json, error_message).\r\n        If successful, response_json is the API's returned JSON object, error_message is None.\r\n        If failed after retries, response_json is None, error_message contains the error information.\r\n    \"\"\"\r\n    request_id = str(uuid.uuid4())\r\n    log_prefix = f\"[Search Request ID: {request_id}] \"\r\n\r\n    payload = {\"queries\": query_list, \"topk\": topk, \"return_scores\": return_scores}\r\n\r\n    headers = {\"Content-Type\": \"application/json\", \"Accept\": \"application/json\"}\r\n\r\n    last_error = None\r\n\r\n    for attempt in range(MAX_RETRIES):\r\n        try:\r\n            logger.info(\r\n                f\"{log_prefix}Attempt {attempt + 1}/{MAX_RETRIES}: Calling search API at {retrieval_service_url}\"\r\n            )\r\n            response = requests.post(\r\n                retrieval_service_url,\r\n                headers=headers,\r\n                json=payload,\r\n                timeout=timeout,\r\n            )\r\n\r\n            # Check for Gateway Timeout (504) and other server errors for retrying\r\n            if response.status_code in [500, 502, 503, 504]:\r\n                last_error = (\r\n                    f\"{log_prefix}API Request Error: Server Error ({response.status_code}) on attempt \"\r\n                    f\"{attempt + 1}/{MAX_RETRIES}\"\r\n                )\r\n                logger.warning(last_error)\r\n                if attempt < MAX_RETRIES - 1:\r\n                    delay = INITIAL_RETRY_DELAY * (attempt + 1)\r\n                    logger.info(f\"{log_prefix}Retrying after {delay} seconds...\")\r\n                    time.sleep(delay)\r\n                continue\r\n\r\n            # Check for other HTTP errors (e.g., 4xx)\r\n            response.raise_for_status()\r\n\r\n            # If successful (status code 2xx)\r\n            logger.info(f\"{log_prefix}Search API call successful on attempt {attempt + 1}\")\r\n            return response.json(), None\r\n\r\n        except requests.exceptions.ConnectionError as e:\r\n            last_error = f\"{log_prefix}Connection Error: {e}\"\r\n            logger.warning(last_error)\r\n            if attempt < MAX_RETRIES - 1:\r\n                delay = INITIAL_RETRY_DELAY * (attempt + 1)\r\n                logger.info(f\"{log_prefix}Retrying after {delay} seconds...\")\r\n                time.sleep(delay)\r\n            continue\r\n        except requests.exceptions.Timeout as e:\r\n            last_error = f\"{log_prefix}Timeout Error: {e}\"\r\n            logger.warning(last_error)\r\n            if attempt < MAX_RETRIES - 1:\r\n                delay = INITIAL_RETRY_DELAY * (attempt + 1)\r\n                logger.info(f\"{log_prefix}Retrying after {delay} seconds...\")\r\n                time.sleep(delay)\r\n            continue\r\n        except requests.exceptions.RequestException as e:\r\n            last_error = f\"{log_prefix}API Request Error: {e}\"\r\n            break  # Exit retry loop on other request errors\r\n        except json.JSONDecodeError as e:\r\n            raw_response_text = response.text if \"response\" in locals() else \"N/A\"\r\n            last_error = f\"{log_prefix}API Response JSON Decode Error: {e}, Response: {raw_response_text[:200]}\"\r\n            break  # Exit retry loop on JSON decode errors\r\n        except Exception as e:\r\n            last_error = f\"{log_prefix}Unexpected Error: {e}\"\r\n            break  # Exit retry loop on other unexpected errors\r\n\r\n    # If loop finishes without returning success, return the last recorded error\r\n    logger.error(f\"{log_prefix}Search API call failed. Last error: {last_error}\")\r\n    return None, last_error.replace(log_prefix, \"API Call Failed: \") if last_error else \"API Call Failed after retries\"\r\n\r\n\r\ndef _passages2string(retrieval_result):\r\n    \"\"\"Convert retrieval results to formatted string.\"\"\"\r\n    format_reference = \"\"\r\n    for idx, doc_item in enumerate(retrieval_result):\r\n        content = doc_item[\"document\"][\"contents\"]\r\n        title = content.split(\"\\n\")[0]\r\n        text = \"\\n\".join(content.split(\"\\n\")[1:])\r\n        format_reference += f\"Doc {idx + 1} (Title: {title})\\n{text}\\n\\n\"\r\n    return format_reference.strip()\r\n\r\n\r\ndef perform_single_search_batch(\r\n    retrieval_service_url: str,\r\n    query_list: list[str],\r\n    topk: int = 3,\r\n    concurrent_semaphore: Optional[threading.Semaphore] = None,\r\n    timeout: int = DEFAULT_TIMEOUT,\r\n) -> tuple[str, dict[str, Any]]:\r\n    \"\"\"\r\n    Performs a single batch search for multiple queries (original search tool behavior).\r\n\r\n    Args:\r\n        retrieval_service_url: The URL of the retrieval service API.\r\n        query_list: List of search queries.\r\n        topk: Number of top results to return.\r\n        concurrent_semaphore: Optional semaphore for concurrency control.\r\n        timeout: Request timeout in seconds.\r\n\r\n    Returns:\r\n        A tuple (result_text, metadata).\r\n        result_text: The search result JSON string.\r\n        metadata: Metadata dictionary for the batch search.\r\n    \"\"\"\r\n    logger.info(f\"Starting batch search for {len(query_list)} queries.\")\r\n\r\n    api_response = None\r\n    error_msg = None\r\n\r\n    try:\r\n        if concurrent_semaphore:\r\n            with concurrent_semaphore:\r\n                api_response, error_msg = call_search_api(\r\n                    retrieval_service_url=retrieval_service_url,\r\n                    query_list=query_list,\r\n                    topk=topk,\r\n                    return_scores=True,\r\n                    timeout=timeout,\r\n                )\r\n        else:\r\n            api_response, error_msg = call_search_api(\r\n                retrieval_service_url=retrieval_service_url,\r\n                query_list=query_list,\r\n                topk=topk,\r\n                return_scores=True,\r\n                timeout=timeout,\r\n            )\r\n    except Exception as e:\r\n        error_msg = f\"API Request Exception during batch search: {e}\"\r\n        logger.error(f\"Batch search: {error_msg}\")\r\n        traceback.print_exc()\r\n\r\n    metadata = {\r\n        \"query_count\": len(query_list),\r\n        \"queries\": query_list,\r\n        \"api_request_error\": error_msg,\r\n        \"api_response\": None,\r\n        \"status\": \"unknown\",\r\n        \"total_results\": 0,\r\n        \"formatted_result\": None,\r\n    }\r\n\r\n    result_text = json.dumps({\"result\": \"Search request failed or timed out after retries.\"}, ensure_ascii=False)\r\n\r\n    if error_msg:\r\n        metadata[\"status\"] = \"api_error\"\r\n        result_text = json.dumps({\"result\": f\"Search error: {error_msg}\"}, ensure_ascii=False)\r\n        logger.error(f\"Batch search: API error occurred: {error_msg}\")\r\n    elif api_response:\r\n        logger.debug(f\"Batch search: API Response: {api_response}\")\r\n        metadata[\"api_response\"] = api_response\r\n\r\n        try:\r\n            raw_results = api_response.get(\"result\", [])\r\n            if raw_results:\r\n                pretty_results = []\r\n                total_results = 0\r\n\r\n                for retrieval in raw_results:\r\n                    formatted = _passages2string(retrieval)\r\n                    pretty_results.append(formatted)\r\n                    total_results += len(retrieval) if isinstance(retrieval, list) else 1\r\n\r\n                final_result = \"\\n---\\n\".join(pretty_results)\r\n                result_text = json.dumps({\"result\": final_result}, ensure_ascii=False)\r\n                metadata[\"status\"] = \"success\"\r\n                metadata[\"total_results\"] = total_results\r\n                metadata[\"formatted_result\"] = final_result\r\n                logger.info(f\"Batch search: Successful, got {total_results} total results\")\r\n            else:\r\n                result_text = json.dumps({\"result\": \"No search results found.\"}, ensure_ascii=False)\r\n                metadata[\"status\"] = \"no_results\"\r\n                metadata[\"total_results\"] = 0\r\n                logger.info(\"Batch search: No results found\")\r\n        except Exception as e:\r\n            error_msg = f\"Error processing search results: {e}\"\r\n            result_text = json.dumps({\"result\": error_msg}, ensure_ascii=False)\r\n            metadata[\"status\"] = \"processing_error\"\r\n            logger.error(f\"Batch search: {error_msg}\")\r\n    else:\r\n        metadata[\"status\"] = \"unknown_api_state\"\r\n        result_text = json.dumps(\r\n            {\"result\": \"Unknown API state (no response and no error message).\"}, ensure_ascii=False\r\n        )\r\n        logger.error(\"Batch search: Unknown API state.\")\r\n\r\n    return result_text, metadata\r\n"
  },
  {
    "path": "verl/tools/utils/tool_registry.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nimport importlib\nimport logging\nimport os\nimport sys\nimport threading\nfrom enum import Enum\n\nfrom omegaconf import OmegaConf\n\nfrom verl.tools.schemas import OpenAIFunctionToolSchema\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass ToolType(Enum):\n    NATIVE = \"native\"\n    MCP = \"mcp\"\n\n\nasync def initialize_mcp_tool(tool_cls, tool_config) -> list:\n    from verl.tools.utils.mcp_clients.McpClientManager import ClientManager\n\n    tool_list = []\n    mcp_servers_config_path = tool_config.mcp.mcp_servers_config_path\n    tool_selected_list = tool_config.mcp.tool_selected_list if \"tool_selected_list\" in tool_config.mcp else None\n    await ClientManager.initialize(mcp_servers_config_path, tool_config.config.rate_limit)\n    # Wait for MCP client to be ready\n    max_retries = 10\n    retry_interval = 2  # seconds\n    for i in range(max_retries):\n        tool_schemas = await ClientManager.fetch_tool_schemas(tool_selected_list)\n        if tool_schemas:\n            break\n        if i < max_retries - 1:\n            logger.debug(f\"Waiting for MCP client to be ready, attempt {i + 1}/{max_retries}\")\n            await asyncio.sleep(retry_interval)\n    else:\n        raise RuntimeError(\"Failed to initialize MCP tools after maximum retries\")\n    # mcp registry\n    assert len(tool_schemas), \"mcp tool is empty\"\n    for tool_schema_dict in tool_schemas:\n        logger.debug(f\"tool_schema_dict: {tool_schema_dict}\")\n        tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)\n        tool = tool_cls(\n            config=OmegaConf.to_container(tool_config.config, resolve=True),\n            tool_schema=tool_schema,\n        )\n        tool_list.append(tool)\n    return tool_list\n\n\ndef get_tool_class(cls_name):\n    module_name, class_name = cls_name.rsplit(\".\", 1)\n    if module_name not in sys.modules:\n        spec = importlib.util.find_spec(module_name)\n        module = importlib.util.module_from_spec(spec)\n        sys.modules[module_name] = module\n        spec.loader.exec_module(module)\n    else:\n        module = sys.modules[module_name]\n\n    tool_cls = getattr(module, class_name)\n    return tool_cls\n\n\ndef initialize_tools_from_config(tools_config_file):\n    \"\"\"Initialize tools from config file.\n\n    Supports both NATIVE and MCP tool types. For MCP tools, a temporary event loop\n    is created only when needed and properly closed after use to prevent memory leaks.\n    \"\"\"\n    tools_config = OmegaConf.load(tools_config_file)\n    tool_list = []\n\n    # Lazy initialization for MCP support - only create event loop when needed\n    tmp_event_loop = None\n    thread = None\n\n    def get_mcp_event_loop():\n        \"\"\"Lazily create event loop and thread for MCP tools.\"\"\"\n        nonlocal tmp_event_loop, thread\n        if tmp_event_loop is None:\n            tmp_event_loop = asyncio.new_event_loop()\n            thread = threading.Thread(target=tmp_event_loop.run_forever, name=\"mcp tool list fetcher\", daemon=True)\n            thread.start()\n        return tmp_event_loop\n\n    def run_coroutine(coroutine):\n        \"\"\"Run coroutine in the MCP event loop.\"\"\"\n        loop = get_mcp_event_loop()\n        future = asyncio.run_coroutine_threadsafe(coroutine, loop)\n        return future.result()\n\n    try:\n        for tool_config in tools_config.tools:\n            cls_name = tool_config.class_name\n            tool_type = ToolType(tool_config.config.type)\n            tool_cls = get_tool_class(cls_name)\n\n            match tool_type:\n                case ToolType.NATIVE:\n                    if tool_config.get(\"tool_schema\", None) is None:\n                        tool_schema = None\n                    else:\n                        tool_schema_dict = OmegaConf.to_container(tool_config.tool_schema, resolve=True)\n                        tool_schema = OpenAIFunctionToolSchema.model_validate(tool_schema_dict)\n                    tool = tool_cls(\n                        config=OmegaConf.to_container(tool_config.config, resolve=True),\n                        tool_schema=tool_schema,\n                    )\n                    tool_list.append(tool)\n                case ToolType.MCP:\n                    mcp_tools = run_coroutine(initialize_mcp_tool(tool_cls, tool_config))\n                    tool_list.extend(mcp_tools)\n                case _:\n                    raise NotImplementedError\n    finally:\n        # Properly cleanup event loop if it was created\n        if tmp_event_loop is not None:\n            # stop first and then close\n            tmp_event_loop.call_soon_threadsafe(tmp_event_loop.stop)\n            if thread is not None and thread.is_alive():\n                thread.join(timeout=5.0)\n            tmp_event_loop.close()\n\n    return tool_list\n"
  },
  {
    "path": "verl/trainer/README.md",
    "content": "# verl Main Entrypoints\n\n## SFT Trainer\n- sft_trainer.py: SFT trainer based on model engine, support various backends: fsdp, megatron, veomni, torchtitan. Launched by `torchrun` and run in multi-controller mode.\n- **[EXPERIMENTAL]** sft_trainer_ray.py: SFT trainer based on model engine with single-controller mode. Launched by ray with a driver process coordinating multiple worker processes.\n\n## RL Trainer\n|trainer|description|sync/async|trainer/rollout|partial rollout|\n|----|----|----|----|----|\n|main_ppo.py|rollout until a batch is completed, then train|synchronous|colocated|No|\n|TBD|[kimi-1.5](https://arxiv.org/pdf/2501.12599) style trainer: streaming rollout with capped length partial rollout|asynchronous|colocated|Yes|\n|TBD|[Areal](https://arxiv.org/pdf/2505.24298) style trainer: fully decoupled trainer and rollout with staleness control|asynchronous|disaggregated|Yes|\n\n## Inference and Evaluation\n- main_generation_server.py: Launch standalone servers and generate responses for a specified prompt dataset.\n- main_eval.py: Evaluate the performance of generated responses with reward function on a specified prompt dataset.\n"
  },
  {
    "path": "verl/trainer/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/trainer/config/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom . import algorithm, config\nfrom .algorithm import *  # noqa: F401\nfrom .config import *  # noqa: F401\n\n__all__ = config.__all__ + algorithm.__all__\n"
  },
  {
    "path": "verl/trainer/config/_generated_ppo_megatron_trainer.yaml",
    "content": "# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'\n# in which it invokes 'python3 scripts/print_cfg.py --cfg job --config-name=ppo_megatron_trainer.yaml' to flatten the 'verl/trainer/config/ppo_megatron_trainer.yaml' config fields into a single file.\n# Do not modify this file directly.\n# The file is usually only for reference and never used.\n\nactor_rollout_ref:\n  actor:\n    optim:\n      _target_: verl.workers.config.McoreOptimizerConfig\n      lr: 1.0e-06\n      lr_warmup_steps_ratio: 0.0\n      total_training_steps: -1\n      weight_decay: 0.01\n      lr_warmup_steps: -1\n      betas:\n      - 0.9\n      - 0.999\n      clip_grad: 1.0\n      optimizer: adam\n      lr_warmup_init: 0.0\n      lr_decay_steps: null\n      lr_decay_style: constant\n      min_lr: 0.0\n      weight_decay_incr_style: constant\n      lr_wsd_decay_style: exponential\n      lr_wsd_decay_steps: null\n      use_checkpoint_opt_param_scheduler: false\n      override_optimizer_config: {}\n    megatron:\n      _target_: verl.workers.config.McoreEngineConfig\n      param_offload: false\n      grad_offload: false\n      optimizer_offload: false\n      tensor_model_parallel_size: 1\n      expert_model_parallel_size: 1\n      expert_tensor_parallel_size: null\n      pipeline_model_parallel_size: 1\n      virtual_pipeline_model_parallel_size: null\n      context_parallel_size: 1\n      sequence_parallel: true\n      use_distributed_optimizer: true\n      use_dist_checkpointing: false\n      dist_checkpointing_path: null\n      dist_checkpointing_prefix: ''\n      dist_ckpt_optim_fully_reshardable: true\n      distrib_optim_fully_reshardable_mem_efficient: false\n      seed: 42\n      override_ddp_config: {}\n      override_transformer_config:\n        recompute_granularity: null\n        recompute_modules:\n        - core_attn\n        recompute_method: null\n        recompute_num_layers: null\n        attention_backend: flash\n      override_mcore_model_config: {}\n      use_mbridge: true\n      vanilla_mbridge: true\n      use_remove_padding: true\n      forward_only: false\n      dtype: bfloat16\n      router_replay:\n        _target_: verl.workers.config.EngineRouterReplayConfig\n        mode: disabled\n        record_file: null\n        replay_file: null\n    _target_: verl.workers.config.McoreActorConfig\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: megatron\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: false\n    ppo_max_token_len_per_gpu: 16384\n    clip_ratio: 0.2\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.2\n    tau_pos: 1.0\n    tau_neg: 1.05\n    freeze_vision_tower: false\n    policy_loss:\n      _target_: verl.workers.config.PolicyLossConfig\n      loss_mode: vanilla\n      clip_cov_ratio: 0.0002\n      clip_cov_lb: 1.0\n      clip_cov_ub: 5.0\n      kl_cov_ratio: 0.0002\n      ppo_kl_coef: 0.1\n    clip_ratio_c: 3.0\n    loss_agg_mode: token-mean\n    loss_scale_factor: null\n    entropy_coeff: 0\n    calculate_entropy: false\n    use_kl_loss: false\n    use_prefix_grouper: false\n    use_torch_compile: true\n    kl_loss_coef: 0.001\n    kl_loss_type: low_var_kl\n    ppo_epochs: 1\n    shuffle: false\n    data_loader_seed: 42\n    checkpoint:\n      _target_: verl.trainer.config.CheckpointConfig\n      save_contents:\n      - model\n      - optimizer\n      - extra\n      load_contents: ${.save_contents}\n      async_save: false\n      mbridge_config: {}\n    use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    load_weight: true\n  ref:\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: megatron\n    use_torch_compile: ${oc.select:actor_rollout_ref.actor.use_torch_compile,true}\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    megatron:\n      _target_: verl.workers.config.McoreEngineConfig\n      param_offload: ${oc.select:actor_rollout_ref.actor.megatron.param_offload,False}\n      grad_offload: false\n      optimizer_offload: false\n      tensor_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.tensor_model_parallel_size,1}\n      expert_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.expert_model_parallel_size,1}\n      expert_tensor_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.expert_tensor_parallel_size,null}\n      pipeline_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.pipeline_model_parallel_size,1}\n      virtual_pipeline_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size,null}\n      context_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.context_parallel_size,1}\n      sequence_parallel: true\n      use_distributed_optimizer: true\n      use_dist_checkpointing: false\n      dist_checkpointing_path: null\n      dist_checkpointing_prefix: ''\n      dist_ckpt_optim_fully_reshardable: true\n      distrib_optim_fully_reshardable_mem_efficient: false\n      seed: ${oc.select:actor_rollout_ref.actor.megatron.seed,42}\n      override_ddp_config: {}\n      override_transformer_config: ${oc.select:actor_rollout_ref.actor.megatron.override_transformer_config,{}}\n      override_mcore_model_config: {}\n      use_mbridge: ${oc.select:actor_rollout_ref.actor.megatron.use_mbridge,False}\n      vanilla_mbridge: ${oc.select:actor_rollout_ref.actor.megatron.vanilla_mbridge,True}\n      use_remove_padding: ${oc.select:actor_rollout_ref.actor.megatron.use_remove_padding,True}\n      forward_only: true\n      dtype: bfloat16\n      router_replay:\n        _target_: verl.workers.config.EngineRouterReplayConfig\n        mode: disabled\n        record_file: null\n        replay_file: null\n    _target_: verl.workers.config.McoreActorConfig\n    load_weight: true\n  rollout:\n    _target_: verl.workers.config.RolloutConfig\n    name: ???\n    mode: async\n    nnodes: 0\n    n_gpus_per_node: ${oc.select:trainer.n_gpus_per_node,8}\n    temperature: 1.0\n    top_k: -1\n    top_p: 1\n    prompt_length: ${oc.select:data.max_prompt_length,512}\n    response_length: ${oc.select:data.max_response_length,512}\n    dtype: bfloat16\n    gpu_memory_utilization: 0.5\n    ignore_eos: false\n    enforce_eager: false\n    cudagraph_capture_sizes: null\n    free_cache_engine: true\n    tensor_model_parallel_size: 2\n    data_parallel_size: 1\n    expert_parallel_size: 1\n    pipeline_model_parallel_size: 1\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    enable_chunked_prefill: true\n    enable_prefix_caching: true\n    logprobs_mode: processed_logprobs\n    scheduling_policy: fcfs\n    load_format: dummy\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    disable_log_stats: true\n    do_sample: true\n    'n': 1\n    over_sample_rate: 0\n    multi_stage_wake_up: false\n    engine_kwargs:\n      vllm: {}\n      sglang: {}\n      trtllm: {}\n    val_kwargs:\n      _target_: verl.workers.config.SamplingConfig\n      top_k: -1\n      top_p: 1.0\n      temperature: 0\n      'n': 1\n      do_sample: false\n    multi_turn:\n      _target_: verl.workers.config.MultiTurnConfig\n      enable: false\n      max_assistant_turns: null\n      tool_config_path: null\n      max_user_turns: null\n      max_parallel_calls: 1\n      max_tool_response_length: 256\n      tool_response_truncate_side: middle\n      interaction_config_path: null\n      use_inference_chat_template: false\n      tokenization_sanity_check_mode: strict\n      format: hermes\n      num_repeat_rollouts: null\n    calculate_log_probs: false\n    agent:\n      _target_: verl.workers.config.AgentLoopConfig\n      num_workers: 8\n      default_agent_loop: single_turn_agent\n      agent_loop_config_path: null\n      custom_async_server:\n        _target_: verl.workers.config.CustomAsyncServerConfig\n        path: null\n        name: null\n    checkpoint_engine:\n      _target_: verl.workers.config.CheckpointEngineConfig\n      backend: naive\n      update_weights_bucket_megabytes: 2048\n      engine_kwargs: {}\n    trace:\n      _target_: verl.workers.config.TraceConfig\n      project_name: ${oc.select:trainer.project_name,null}\n      experiment_name: ${oc.select:trainer.experiment_name,null}\n      backend: null\n      token2text: false\n      max_samples_per_step_per_worker: null\n    skip_rollout: false\n    skip_dump_dir: /tmp/rollout_dump\n    skip_tokenizer_init: true\n    enable_rollout_routing_replay: false\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: ${oc.select:actor_rollout_ref.actor.profiler.enable,false}\n      all_ranks: ${oc.select:actor_rollout_ref.actor.profiler.all_ranks,false}\n      ranks: ${oc.select:actor_rollout_ref.actor.profiler.ranks,[]}\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.contents,[]}\n          level: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.level,level0}\n          analysis: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.analysis,false}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.discrete,false}\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.contents,[]}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.discrete,false}\n    prometheus:\n      _target_: verl.workers.config.PrometheusConfig\n      enable: false\n      port: 9090\n      file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml\n      served_model_name: ${oc.select:actor_rollout_ref.model.path,null}\n    quantization: null\n    quantization_config_file: null\n    mtp: ${oc.select:actor_rollout_ref.model.mtp, null}\n    qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,null}\n    layer_name_map:\n      qkv_layer_name: qkv\n      gate_proj_layer_name: gate_up\n  model:\n    _target_: verl.workers.config.HFModelConfig\n    path: ~/models/deepseek-llm-7b-chat\n    hf_config_path: null\n    tokenizer_path: null\n    use_shm: false\n    trust_remote_code: false\n    custom_chat_template: null\n    external_lib: null\n    override_config:\n      model_config: {}\n      moe_config:\n        freeze_moe_router: false\n    enable_gradient_checkpointing: true\n    enable_activation_offload: false\n    use_remove_padding: false\n    lora_rank: 0\n    lora_alpha: 16\n    target_modules: all-linear\n    exclude_modules: null\n    lora_adapter_path: null\n    use_liger: false\n    use_fused_kernels: false\n    fused_kernel_options:\n      impl_backend: torch\n    tiled_mlp:\n      enabled: false\n      num_shards: 4\n    mtp:\n      _target_: verl.workers.config.MtpConfig\n      enable: false\n      enable_train: false\n      enable_rollout: false\n      detach_encoder: false\n      mtp_loss_scaling_factor: 0.1\n      speculative_algorithm: EAGLE\n      speculative_num_steps: 3\n      speculative_eagle_topk: 1\n      speculative_num_draft_tokens: 4\n      method: mtp\n      num_speculative_tokens: 1\n    lora:\n      type: lora\n      merge: false\n      rank: 0\n      alpha: 32\n      dropout: 0.0\n      target_modules:\n      - linear_qkv\n      - linear_proj\n      - linear_fc1\n      - linear_fc2\n      exclude_modules: []\n      dropout_position: pre\n      lora_A_init_method: xavier\n      lora_B_init_method: zero\n      a2a_experimental: false\n      dtype: null\n      adapter_path: null\n      freeze_vision_model: true\n      freeze_vision_projection: true\n      freeze_language_model: true\n  hybrid_engine: true\n  nccl_timeout: 600\ndata:\n  tokenizer: null\n  use_shm: false\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1\n  val_max_samples: -1\n  prompt_key: prompt\n  reward_fn_key: data_source\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null\n  tool_config_path: ${oc.select:actor_rollout_ref.rollout.multi_turn.tool_config_path,\n    null}\n  return_raw_input_ids: false\n  return_raw_chat: true\n  return_full_prompt: false\n  shuffle: true\n  seed: null\n  dataloader_num_workers: 8\n  image_patch_size: 14\n  validation_shuffle: false\n  filter_overlong_prompts: false\n  filter_overlong_prompts_workers: 1\n  truncation: error\n  image_key: images\n  video_key: videos\n  trust_remote_code: false\n  custom_cls:\n    path: null\n    name: null\n  return_multi_modal_inputs: true\n  sampler:\n    class_path: null\n    class_name: null\n  datagen:\n    path: null\n    name: null\n  apply_chat_template_kwargs: {}\ncritic:\n  optim:\n    _target_: verl.workers.config.McoreOptimizerConfig\n    lr: 1.0e-05\n    lr_warmup_steps_ratio: 0.0\n    total_training_steps: -1\n    weight_decay: 0.01\n    lr_warmup_steps: -1\n    betas:\n    - 0.9\n    - 0.999\n    clip_grad: 1.0\n    optimizer: adam\n    lr_warmup_init: 0.0\n    lr_decay_steps: null\n    lr_decay_style: constant\n    min_lr: 0.0\n    weight_decay_incr_style: constant\n    lr_wsd_decay_style: exponential\n    lr_wsd_decay_steps: null\n    use_checkpoint_opt_param_scheduler: false\n    override_optimizer_config: {}\n  megatron:\n    _target_: verl.workers.config.McoreEngineConfig\n    param_offload: false\n    grad_offload: false\n    optimizer_offload: false\n    tensor_model_parallel_size: 1\n    expert_model_parallel_size: 1\n    expert_tensor_parallel_size: null\n    pipeline_model_parallel_size: 1\n    virtual_pipeline_model_parallel_size: null\n    context_parallel_size: 1\n    sequence_parallel: true\n    use_distributed_optimizer: true\n    use_dist_checkpointing: false\n    dist_checkpointing_path: null\n    dist_checkpointing_prefix: ''\n    dist_ckpt_optim_fully_reshardable: true\n    distrib_optim_fully_reshardable_mem_efficient: false\n    seed: 42\n    override_ddp_config: {}\n    override_transformer_config:\n      recompute_granularity: null\n      recompute_modules:\n      - core_attn\n      recompute_method: null\n      recompute_num_layers: null\n      attention_backend: flash\n    override_mcore_model_config: {}\n    use_mbridge: true\n    vanilla_mbridge: true\n    use_remove_padding: true\n    forward_only: false\n    dtype: bfloat16\n    router_replay:\n      _target_: verl.workers.config.EngineRouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n  _target_: verl.workers.config.McoreCriticConfig\n  rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n  strategy: megatron\n  enable: null\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${oc.select:actor_rollout_ref.model.path,\"~/models/deepseek-llm-7b-chat\"}\n    override_config:\n      model_config: {}\n      moe_config:\n        freeze_moe_router: false\n    external_lib: ${oc.select:actor_rollout_ref.model.external_lib,null}\n    trust_remote_code: ${oc.select:actor_rollout_ref.model.trust_remote_code,false}\n    _target_: verl.trainer.config.BaseModelConfig\n    lora:\n      type: lora\n      rank: 0\n      alpha: 32\n      dropout: 0.0\n      target_modules:\n      - linear_qkv\n      - linear_proj\n      - linear_fc1\n      - linear_fc2\n      exclude_modules: []\n      dropout_position: pre\n      lora_A_init_method: xavier\n      lora_B_init_method: zero\n      a2a_experimental: false\n      dtype: null\n      adapter_path: null\n      freeze_vision_model: true\n      freeze_vision_projection: true\n      freeze_language_model: true\n  ppo_mini_batch_size: ${oc.select:actor_rollout_ref.actor.ppo_mini_batch_size,256}\n  ppo_micro_batch_size: null\n  ppo_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size,null}\n  use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n  ppo_max_token_len_per_gpu: 32768\n  forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu}\n  ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1}\n  shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false}\n  data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n  cliprange_value: 0.5\n  loss_agg_mode: ${oc.select:actor_rollout_ref.actor.loss_agg_mode,token-mean}\n  checkpoint:\n    _target_: verl.trainer.config.CheckpointConfig\n    save_contents:\n    - model\n    - optimizer\n    - extra\n    load_contents: ${.save_contents}\n    async_save: false\n    mbridge_config: {}\n  profiler:\n    _target_: verl.utils.profiler.ProfilerConfig\n    tool: ${oc.select:global_profiler.tool,null}\n    enable: false\n    all_ranks: false\n    ranks: []\n    save_path: ${oc.select:global_profiler.save_path,null}\n    tool_config:\n      nsys:\n        _target_: verl.utils.profiler.config.NsightToolConfig\n        discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n      npu:\n        _target_: verl.utils.profiler.config.NPUToolConfig\n        contents: []\n        level: level0\n        analysis: true\n        discrete: false\n      torch:\n        _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n        contents: []\n        discrete: false\n      torch_memory:\n        _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n        trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n        stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n  nccl_timeout: 600\n  load_weight: true\ncustom_reward_function:\n  path: null\n  name: null\nreward_model:\n  num_workers: null\n  reward_manager: null\n  enable: null\n  enable_resource_pool: null\n  n_gpus_per_node: null\n  nnodes: null\n  reward_loop_source: null\n  reward_loop_module_path: null\n  reward_loop_class_name: null\n  model:\n    path: null\n    external_lib: null\n    trust_remote_code: null\n  rollout:\n    name: null\n    dtype: null\n    gpu_memory_utilization: null\n    enforce_eager: null\n    cudagraph_capture_sizes: null\n    free_cache_engine: null\n    data_parallel_size: null\n    expert_parallel_size: null\n    tensor_model_parallel_size: null\n    max_num_batched_tokens: null\n    max_model_len: null\n    max_num_seqs: null\n    load_format: null\n    engine_kwargs: null\n    limit_images: null\n    enable_chunked_prefill: null\n    enable_prefix_caching: null\n    disable_log_stats: null\n    skip_tokenizer_init: null\n    prompt_length: null\n    response_length: null\nsandbox_fusion:\n  url: null\n  max_concurrent: null\n  memory_limit_mb: null\nreward:\n  num_workers: 8\n  custom_reward_function:\n    path: null\n    name: compute_score\n  reward_manager:\n    _target_: verl.workers.config.reward_model.RewardManagerConfig\n    source: register\n    name: naive\n    module:\n      _target_: verl.trainer.config.config.ModuleConfig\n      path: null\n      name: custom_reward_manager\n  reward_model:\n    enable: false\n    enable_resource_pool: false\n    n_gpus_per_node: 8\n    nnodes: 0\n    model_path: null\n    rollout:\n      _target_: verl.workers.config.RolloutConfig\n      name: ???\n      dtype: bfloat16\n      gpu_memory_utilization: 0.5\n      enforce_eager: true\n      cudagraph_capture_sizes: null\n      free_cache_engine: true\n      data_parallel_size: 1\n      expert_parallel_size: 1\n      tensor_model_parallel_size: 2\n      max_num_batched_tokens: 8192\n      max_model_len: null\n      max_num_seqs: 1024\n      load_format: auto\n      engine_kwargs: {}\n      limit_images: null\n      enable_chunked_prefill: true\n      enable_prefix_caching: true\n      disable_log_stats: true\n      skip_tokenizer_init: false\n      prompt_length: 2048\n      response_length: 2048\n  sandbox_fusion:\n    url: null\n    max_concurrent: 64\n    memory_limit_mb: 1024\nalgorithm:\n  rollout_correction:\n    rollout_is: null\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    rollout_rs_threshold: null\n    bypass_mode: false\n    loss_type: ppo_clip\n    rollout_is_batch_normalize: false\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: true\n  use_kl_in_reward: false\n  kl_penalty: kl\n  kl_ctrl:\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: false\n  pf_ppo:\n    reweight_method: pow\n    weight_pow: 2.0\ntrainer:\n  balance_batch: true\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger:\n  - console\n  - wandb\n  log_val_generations: 0\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n  resume_mode: auto\n  resume_from_path: null\n  del_local_ckpt_after_load: false\n  val_before_train: true\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  rollout_data_dir: null\n  use_legacy_worker_impl: auto\nglobal_profiler:\n  _target_: verl.utils.profiler.ProfilerConfig\n  tool: null\n  steps: null\n  profile_continuous_steps: false\n  save_path: outputs/profile\n  global_tool_config:\n    nsys:\n      discrete: false\n      controller_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n      worker_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n        capture-range: cudaProfilerApi\n        capture-range-end: null\n        kill: none\n    torch_memory:\n      trace_alloc_max_entries: 100000\n      stack_depth: 32\n      context: all\n      stacks: all\n      kw_args: {}\ntransfer_queue:\n  enable: false\nray_kwargs:\n  ray_init:\n    num_cpus: null\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/_generated_ppo_torchtitan_trainer.yaml",
    "content": "# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'\n# in which it invokes 'python3 scripts/print_cfg.py --cfg job model_engine=torchtitan' to flatten the 'verl/trainer/config/ppo_trainer.yaml' config fields into a single file.\n# Do not modify this file directly.\n# The file is usually only for reference and never used.\n\nactor_rollout_ref:\n  actor:\n    optim:\n      _target_: verl.workers.config.TorchtitanOptimizerConfig\n      name: AdamW\n      lr: 1.0e-06\n      lr_warmup_steps_ratio: 0.0\n      total_training_steps: -1\n      weight_decay: 0.01\n      lr_warmup_steps: -1\n      betas:\n      - 0.9\n      - 0.999\n      clip_grad: 1.0\n      eps: 1.0e-08\n      decay_type: linear\n      min_lr_factor: 0.0\n    torchtitan:\n      _target_: verl.workers.config.TorchtitanEngineConfig\n      param_offload: false\n      optimizer_offload: false\n      wrap_policy:\n        min_num_params: 0\n      reshard_after_forward: default\n      forward_prefetch: false\n      use_orig_params: false\n      mixed_precision: false\n      use_torch_compile: true\n      entropy_from_logits_with_chunking: false\n      entropy_checkpointing: false\n      data_parallel_size: 1\n      data_parallel_replicate_size: 1\n      data_parallel_shard_size: 1\n      tensor_parallel_size: 1\n      expert_parallel_size: 1\n      pipeline_parallel_size: 1\n      context_parallel_size: 1\n      attn_type: flex\n      max_seq_len: null\n      strategy: torchtitan\n      seed: 42\n      full_determinism: false\n      forward_only: false\n      dtype: bfloat16\n    _target_: verl.workers.config.TorchTitanActorConfig\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: torchtitan\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: false\n    ppo_max_token_len_per_gpu: 16384\n    clip_ratio: 0.2\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.2\n    tau_pos: 1.0\n    tau_neg: 1.05\n    freeze_vision_tower: false\n    policy_loss:\n      _target_: verl.workers.config.PolicyLossConfig\n      loss_mode: vanilla\n      clip_cov_ratio: 0.0002\n      clip_cov_lb: 1.0\n      clip_cov_ub: 5.0\n      kl_cov_ratio: 0.0002\n      ppo_kl_coef: 0.1\n    clip_ratio_c: 3.0\n    loss_agg_mode: token-mean\n    loss_scale_factor: null\n    entropy_coeff: 0\n    calculate_entropy: false\n    use_kl_loss: false\n    use_prefix_grouper: false\n    use_torch_compile: true\n    kl_loss_coef: 0.001\n    kl_loss_type: low_var_kl\n    ppo_epochs: 1\n    shuffle: false\n    data_loader_seed: 42\n    checkpoint:\n      _target_: verl.trainer.config.CheckpointConfig\n      save_contents:\n      - model\n      - optimizer\n      - extra\n      load_contents: ${.save_contents}\n      async_save: false\n      mbridge_config: {}\n    use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n  ref:\n    optim:\n      _target_: verl.workers.config.TorchtitanOptimizerConfig\n      name: AdamW\n      lr: 0.001\n      lr_warmup_steps_ratio: 0.0\n      total_training_steps: -1\n      weight_decay: 0.01\n      lr_warmup_steps: -1\n      betas:\n      - 0.9\n      - 0.999\n      clip_grad: 1.0\n      eps: 1.0e-08\n      decay_type: linear\n      min_lr_factor: 0.0\n    torchtitan:\n      _target_: verl.workers.config.TorchtitanEngineConfig\n      param_offload: false\n      optimizer_offload: false\n      wrap_policy:\n        min_num_params: 0\n      reshard_after_forward: default\n      forward_prefetch: false\n      use_orig_params: false\n      mixed_precision: false\n      use_torch_compile: true\n      entropy_from_logits_with_chunking: false\n      entropy_checkpointing: false\n      data_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_size,1}\n      data_parallel_replicate_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_replicate_size,1}\n      data_parallel_shard_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_shard_size,1}\n      tensor_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.tensor_parallel_size,1}\n      expert_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.expert_parallel_size,1}\n      pipeline_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.pipeline_parallel_size,1}\n      context_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.context_parallel_size,1}\n      attn_type: ${oc.select:actor_rollout_ref.actor.torchtitan.attn_type,flex}\n      max_seq_len: null\n      strategy: torchtitan\n      seed: ${oc.select:actor_rollout_ref.actor.torchtitan.seed,42}\n      full_determinism: false\n      forward_only: true\n      dtype: bfloat16\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: torchtitan\n    use_torch_compile: ${oc.select:actor_rollout_ref.actor.use_torch_compile,true}\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    _target_: verl.workers.config.TorchTitanActorConfig\n  rollout:\n    _target_: verl.workers.config.RolloutConfig\n    name: ???\n    mode: async\n    nnodes: 0\n    n_gpus_per_node: ${oc.select:trainer.n_gpus_per_node,8}\n    temperature: 1.0\n    top_k: -1\n    top_p: 1\n    prompt_length: ${oc.select:data.max_prompt_length,512}\n    response_length: ${oc.select:data.max_response_length,512}\n    dtype: bfloat16\n    gpu_memory_utilization: 0.5\n    ignore_eos: false\n    enforce_eager: false\n    cudagraph_capture_sizes: null\n    free_cache_engine: true\n    tensor_model_parallel_size: 2\n    data_parallel_size: 1\n    expert_parallel_size: 1\n    pipeline_model_parallel_size: 1\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    enable_chunked_prefill: true\n    enable_prefix_caching: true\n    logprobs_mode: processed_logprobs\n    scheduling_policy: fcfs\n    load_format: dummy\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    disable_log_stats: true\n    do_sample: true\n    'n': 1\n    over_sample_rate: 0\n    multi_stage_wake_up: false\n    engine_kwargs:\n      vllm: {}\n      sglang: {}\n      trtllm: {}\n    val_kwargs:\n      _target_: verl.workers.config.SamplingConfig\n      top_k: -1\n      top_p: 1.0\n      temperature: 0\n      'n': 1\n      do_sample: false\n    multi_turn:\n      _target_: verl.workers.config.MultiTurnConfig\n      enable: false\n      max_assistant_turns: null\n      tool_config_path: null\n      max_user_turns: null\n      max_parallel_calls: 1\n      max_tool_response_length: 256\n      tool_response_truncate_side: middle\n      interaction_config_path: null\n      use_inference_chat_template: false\n      tokenization_sanity_check_mode: strict\n      format: hermes\n      num_repeat_rollouts: null\n    calculate_log_probs: false\n    agent:\n      _target_: verl.workers.config.AgentLoopConfig\n      num_workers: 8\n      default_agent_loop: single_turn_agent\n      agent_loop_config_path: null\n      custom_async_server:\n        _target_: verl.workers.config.CustomAsyncServerConfig\n        path: null\n        name: null\n    checkpoint_engine:\n      _target_: verl.workers.config.CheckpointEngineConfig\n      backend: naive\n      update_weights_bucket_megabytes: 2048\n      engine_kwargs: {}\n    trace:\n      _target_: verl.workers.config.TraceConfig\n      project_name: ${oc.select:trainer.project_name,null}\n      experiment_name: ${oc.select:trainer.experiment_name,null}\n      backend: null\n      token2text: false\n      max_samples_per_step_per_worker: null\n    skip_rollout: false\n    skip_dump_dir: /tmp/rollout_dump\n    skip_tokenizer_init: true\n    enable_rollout_routing_replay: false\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: ${oc.select:actor_rollout_ref.actor.profiler.enable,false}\n      all_ranks: ${oc.select:actor_rollout_ref.actor.profiler.all_ranks,false}\n      ranks: ${oc.select:actor_rollout_ref.actor.profiler.ranks,[]}\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.contents,[]}\n          level: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.level,level0}\n          analysis: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.analysis,false}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.discrete,false}\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.contents,[]}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.discrete,false}\n    prometheus:\n      _target_: verl.workers.config.PrometheusConfig\n      enable: false\n      port: 9090\n      file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml\n      served_model_name: ${oc.select:actor_rollout_ref.model.path,null}\n    quantization: null\n    quantization_config_file: null\n    mtp: ${oc.select:actor_rollout_ref.model.mtp, null}\n    qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,null}\n    layered_summon: false\n  model:\n    _target_: verl.workers.config.HFModelConfig\n    path: ~/models/deepseek-llm-7b-chat\n    hf_config_path: null\n    tokenizer_path: null\n    use_shm: false\n    trust_remote_code: false\n    custom_chat_template: null\n    external_lib: null\n    override_config: {}\n    enable_gradient_checkpointing: true\n    enable_activation_offload: false\n    use_remove_padding: true\n    lora_rank: 0\n    lora_alpha: 16\n    target_modules: all-linear\n    exclude_modules: null\n    lora_adapter_path: null\n    use_liger: false\n    use_fused_kernels: false\n    fused_kernel_options:\n      impl_backend: torch\n    tiled_mlp:\n      enabled: false\n      num_shards: 4\n    mtp:\n      _target_: verl.workers.config.MtpConfig\n      enable: false\n      enable_train: false\n      enable_rollout: false\n      detach_encoder: false\n      mtp_loss_scaling_factor: 0.1\n      speculative_algorithm: EAGLE\n      speculative_num_steps: 3\n      speculative_eagle_topk: 1\n      speculative_num_draft_tokens: 4\n      method: mtp\n      num_speculative_tokens: 1\n  hybrid_engine: true\n  nccl_timeout: 600\ndata:\n  tokenizer: null\n  use_shm: false\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1\n  val_max_samples: -1\n  prompt_key: prompt\n  reward_fn_key: data_source\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null\n  tool_config_path: ${oc.select:actor_rollout_ref.rollout.multi_turn.tool_config_path,\n    null}\n  return_raw_input_ids: false\n  return_raw_chat: true\n  return_full_prompt: false\n  shuffle: true\n  seed: null\n  dataloader_num_workers: 8\n  image_patch_size: 14\n  validation_shuffle: false\n  filter_overlong_prompts: false\n  filter_overlong_prompts_workers: 1\n  truncation: error\n  image_key: images\n  video_key: videos\n  trust_remote_code: false\n  custom_cls:\n    path: null\n    name: null\n  return_multi_modal_inputs: true\n  sampler:\n    class_path: null\n    class_name: null\n  datagen:\n    path: null\n    name: null\n  apply_chat_template_kwargs: {}\ncritic:\n  optim:\n    _target_: verl.workers.config.TorchtitanOptimizerConfig\n    name: AdamW\n    lr: 1.0e-05\n    lr_warmup_steps_ratio: 0.0\n    total_training_steps: -1\n    weight_decay: 0.01\n    lr_warmup_steps: -1\n    betas:\n    - 0.9\n    - 0.999\n    clip_grad: 1.0\n    eps: 1.0e-08\n    decay_type: linear\n    min_lr_factor: 0.0\n  torchtitan:\n    _target_: verl.workers.config.TorchtitanEngineConfig\n    param_offload: false\n    optimizer_offload: false\n    wrap_policy:\n      min_num_params: 0\n    reshard_after_forward: default\n    forward_prefetch: false\n    use_orig_params: false\n    mixed_precision: false\n    use_torch_compile: true\n    entropy_from_logits_with_chunking: false\n    entropy_checkpointing: false\n    data_parallel_size: 1\n    data_parallel_replicate_size: 1\n    data_parallel_shard_size: 1\n    tensor_parallel_size: 1\n    expert_parallel_size: 1\n    pipeline_parallel_size: 1\n    context_parallel_size: 1\n    attn_type: flex\n    max_seq_len: null\n    strategy: torchtitan\n    seed: 42\n    full_determinism: false\n    forward_only: false\n    dtype: bfloat16\n  _target_: verl.workers.config.TorchTitanCriticConfig\n  rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n  strategy: torchtitan\n  enable: null\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${oc.select:actor_rollout_ref.model.path,\"~/models/deepseek-llm-7b-chat\"}\n    override_config: {}\n    external_lib: ${oc.select:actor_rollout_ref.model.external_lib,null}\n    trust_remote_code: ${oc.select:actor_rollout_ref.model.trust_remote_code,false}\n    _target_: verl.trainer.config.BaseModelConfig\n  ppo_mini_batch_size: ${oc.select:actor_rollout_ref.actor.ppo_mini_batch_size,256}\n  ppo_micro_batch_size: null\n  ppo_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size,null}\n  use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n  ppo_max_token_len_per_gpu: 32768\n  forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu}\n  ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1}\n  shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false}\n  data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n  cliprange_value: 0.5\n  loss_agg_mode: ${oc.select:actor_rollout_ref.actor.loss_agg_mode,token-mean}\n  checkpoint:\n    _target_: verl.trainer.config.CheckpointConfig\n    save_contents:\n    - model\n    - optimizer\n    - extra\n    load_contents: ${.save_contents}\n    async_save: false\n    mbridge_config: {}\n  profiler:\n    _target_: verl.utils.profiler.ProfilerConfig\n    tool: ${oc.select:global_profiler.tool,null}\n    enable: false\n    all_ranks: false\n    ranks: []\n    save_path: ${oc.select:global_profiler.save_path,null}\n    tool_config:\n      nsys:\n        _target_: verl.utils.profiler.config.NsightToolConfig\n        discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n      npu:\n        _target_: verl.utils.profiler.config.NPUToolConfig\n        contents: []\n        level: level0\n        analysis: true\n        discrete: false\n      torch:\n        _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n        contents: []\n        discrete: false\n      torch_memory:\n        _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n        trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n        stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\ncustom_reward_function:\n  path: null\n  name: null\nreward_model:\n  num_workers: null\n  reward_manager: null\n  enable: null\n  enable_resource_pool: null\n  n_gpus_per_node: null\n  nnodes: null\n  reward_loop_source: null\n  reward_loop_module_path: null\n  reward_loop_class_name: null\n  model:\n    path: null\n    external_lib: null\n    trust_remote_code: null\n  rollout:\n    name: null\n    dtype: null\n    gpu_memory_utilization: null\n    enforce_eager: null\n    cudagraph_capture_sizes: null\n    free_cache_engine: null\n    data_parallel_size: null\n    expert_parallel_size: null\n    tensor_model_parallel_size: null\n    max_num_batched_tokens: null\n    max_model_len: null\n    max_num_seqs: null\n    load_format: null\n    engine_kwargs: null\n    limit_images: null\n    enable_chunked_prefill: null\n    enable_prefix_caching: null\n    disable_log_stats: null\n    skip_tokenizer_init: null\n    prompt_length: null\n    response_length: null\nsandbox_fusion:\n  url: null\n  max_concurrent: null\n  memory_limit_mb: null\nreward:\n  num_workers: 8\n  custom_reward_function:\n    path: null\n    name: compute_score\n  reward_manager:\n    _target_: verl.workers.config.reward_model.RewardManagerConfig\n    source: register\n    name: naive\n    module:\n      _target_: verl.trainer.config.config.ModuleConfig\n      path: null\n      name: custom_reward_manager\n  reward_model:\n    enable: false\n    enable_resource_pool: false\n    n_gpus_per_node: 8\n    nnodes: 0\n    model_path: null\n    rollout:\n      _target_: verl.workers.config.RolloutConfig\n      name: ???\n      dtype: bfloat16\n      gpu_memory_utilization: 0.5\n      enforce_eager: true\n      cudagraph_capture_sizes: null\n      free_cache_engine: true\n      data_parallel_size: 1\n      expert_parallel_size: 1\n      tensor_model_parallel_size: 2\n      max_num_batched_tokens: 8192\n      max_model_len: null\n      max_num_seqs: 1024\n      load_format: auto\n      engine_kwargs: {}\n      limit_images: null\n      enable_chunked_prefill: true\n      enable_prefix_caching: true\n      disable_log_stats: true\n      skip_tokenizer_init: false\n      prompt_length: 2048\n      response_length: 2048\n  sandbox_fusion:\n    url: null\n    max_concurrent: 64\n    memory_limit_mb: 1024\nalgorithm:\n  rollout_correction:\n    rollout_is: null\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    rollout_rs_threshold: null\n    bypass_mode: false\n    loss_type: ppo_clip\n    rollout_is_batch_normalize: false\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: true\n  use_kl_in_reward: false\n  kl_penalty: kl\n  kl_ctrl:\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: false\n  pf_ppo:\n    reweight_method: pow\n    weight_pow: 2.0\ntrainer:\n  balance_batch: true\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger:\n  - console\n  - wandb\n  log_val_generations: 0\n  rollout_data_dir: null\n  validation_data_dir: null\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n  resume_mode: auto\n  resume_from_path: null\n  val_before_train: true\n  val_only: false\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  del_local_ckpt_after_load: false\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  use_legacy_worker_impl: auto\nglobal_profiler:\n  _target_: verl.utils.profiler.ProfilerConfig\n  tool: null\n  steps: null\n  profile_continuous_steps: false\n  save_path: outputs/profile\n  global_tool_config:\n    nsys:\n      _target_: verl.utils.profiler.config.NsightToolConfig\n      discrete: false\n      controller_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n      worker_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n        capture-range: cudaProfilerApi\n        capture-range-end: null\n        kill: none\n    torch_memory:\n      trace_alloc_max_entries: 100000\n      stack_depth: 32\n      context: all\n      stacks: all\n      kw_args: {}\ntransfer_queue:\n  enable: false\nray_kwargs:\n  ray_init:\n    num_cpus: null\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/_generated_ppo_trainer.yaml",
    "content": "# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'\n# in which it invokes 'python3 scripts/print_cfg.py --cfg job ' to flatten the 'verl/trainer/config/ppo_trainer.yaml' config fields into a single file.\n# Do not modify this file directly.\n# The file is usually only for reference and never used.\n\nactor_rollout_ref:\n  actor:\n    optim:\n      _target_: verl.workers.config.FSDPOptimizerConfig\n      optimizer: AdamW\n      optimizer_impl: torch.optim\n      lr: 1.0e-06\n      lr_warmup_steps_ratio: 0.0\n      total_training_steps: -1\n      weight_decay: 0.01\n      lr_warmup_steps: -1\n      betas:\n      - 0.9\n      - 0.999\n      clip_grad: 1.0\n      min_lr_ratio: 0.0\n      num_cycles: 0.5\n      lr_scheduler_type: constant\n      zero_indexed_step: true\n      warmup_style: null\n      override_optimizer_config: null\n    fsdp_config:\n      _target_: verl.workers.config.FSDPEngineConfig\n      wrap_policy:\n        min_num_params: 0\n      param_offload: false\n      optimizer_offload: false\n      offload_policy: false\n      reshard_after_forward: true\n      fsdp_size: -1\n      forward_prefetch: false\n      model_dtype: fp32\n      use_orig_params: false\n      seed: 42\n      full_determinism: false\n      ulysses_sequence_parallel_size: 1\n      entropy_from_logits_with_chunking: false\n      use_torch_compile: true\n      entropy_checkpointing: false\n      forward_only: false\n      strategy: fsdp\n      dtype: bfloat16\n      qat:\n        _target_: verl.workers.config.QATEngineConfig\n        enable: false\n        mode: w4a16\n        group_size: 16\n        ignore_patterns:\n        - lm_head\n        - embed_tokens\n        - re:.*mlp.gate$\n        activation_observer: static_minmax\n        quantization_config_path: null\n    _target_: verl.workers.config.FSDPActorConfig\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: fsdp\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: false\n    ppo_max_token_len_per_gpu: 16384\n    clip_ratio: 0.2\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.2\n    tau_pos: 1.0\n    tau_neg: 1.05\n    freeze_vision_tower: false\n    policy_loss:\n      _target_: verl.workers.config.PolicyLossConfig\n      loss_mode: vanilla\n      clip_cov_ratio: 0.0002\n      clip_cov_lb: 1.0\n      clip_cov_ub: 5.0\n      kl_cov_ratio: 0.0002\n      ppo_kl_coef: 0.1\n    clip_ratio_c: 3.0\n    loss_agg_mode: token-mean\n    loss_scale_factor: null\n    entropy_coeff: 0\n    calculate_entropy: false\n    use_kl_loss: false\n    use_prefix_grouper: false\n    use_torch_compile: true\n    kl_loss_coef: 0.001\n    kl_loss_type: low_var_kl\n    ppo_epochs: 1\n    shuffle: false\n    data_loader_seed: 42\n    checkpoint:\n      _target_: verl.trainer.config.CheckpointConfig\n      save_contents:\n      - model\n      - optimizer\n      - extra\n      load_contents: ${.save_contents}\n      async_save: false\n      mbridge_config: {}\n    use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    grad_clip: 1.0\n    ulysses_sequence_parallel_size: 1\n    entropy_from_logits_with_chunking: false\n    entropy_checkpointing: false\n    use_remove_padding: ${oc.select:actor_rollout_ref.model.use_remove_padding,false}\n    calculate_sum_pi_squared: false\n    sum_pi_squared_checkpointing: false\n    qat:\n      enable: false\n      mode: w4a16\n      group_size: 16\n      ignore_patterns:\n      - lm_head\n      - embed_tokens\n      - re:.*mlp.gate$\n      activation_observer: static_minmax\n      quantization_config_path: null\n  ref:\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: ${actor_rollout_ref.actor.strategy}\n    use_torch_compile: ${oc.select:actor_rollout_ref.actor.use_torch_compile,true}\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    fsdp_config:\n      _target_: verl.workers.config.FSDPEngineConfig\n      wrap_policy:\n        min_num_params: 0\n      param_offload: false\n      optimizer_offload: false\n      offload_policy: false\n      reshard_after_forward: true\n      fsdp_size: -1\n      forward_prefetch: false\n      model_dtype: fp32\n      use_orig_params: false\n      seed: 42\n      full_determinism: false\n      ulysses_sequence_parallel_size: 1\n      entropy_from_logits_with_chunking: false\n      use_torch_compile: true\n      entropy_checkpointing: false\n      forward_only: true\n      strategy: fsdp\n      dtype: bfloat16\n      qat:\n        _target_: verl.workers.config.QATEngineConfig\n        enable: false\n        mode: w4a16\n        group_size: 16\n        ignore_patterns:\n        - lm_head\n        - embed_tokens\n        - re:.*mlp.gate$\n        activation_observer: static_minmax\n        quantization_config_path: null\n    _target_: verl.workers.config.FSDPActorConfig\n    ulysses_sequence_parallel_size: ${oc.select:actor_rollout_ref.actor.ulysses_sequence_parallel_size,1}\n    entropy_from_logits_with_chunking: false\n    entropy_checkpointing: false\n  rollout:\n    _target_: verl.workers.config.RolloutConfig\n    name: ???\n    mode: async\n    nnodes: 0\n    n_gpus_per_node: ${oc.select:trainer.n_gpus_per_node,8}\n    temperature: 1.0\n    top_k: -1\n    top_p: 1\n    prompt_length: ${oc.select:data.max_prompt_length,512}\n    response_length: ${oc.select:data.max_response_length,512}\n    dtype: bfloat16\n    gpu_memory_utilization: 0.5\n    ignore_eos: false\n    enforce_eager: false\n    cudagraph_capture_sizes: null\n    free_cache_engine: true\n    tensor_model_parallel_size: 2\n    data_parallel_size: 1\n    expert_parallel_size: 1\n    pipeline_model_parallel_size: 1\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    enable_chunked_prefill: true\n    enable_prefix_caching: true\n    logprobs_mode: processed_logprobs\n    scheduling_policy: fcfs\n    load_format: dummy\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    disable_log_stats: true\n    do_sample: true\n    'n': 1\n    over_sample_rate: 0\n    multi_stage_wake_up: false\n    engine_kwargs:\n      vllm: {}\n      sglang: {}\n      trtllm: {}\n    val_kwargs:\n      _target_: verl.workers.config.SamplingConfig\n      top_k: -1\n      top_p: 1.0\n      temperature: 0\n      'n': 1\n      do_sample: false\n    multi_turn:\n      _target_: verl.workers.config.MultiTurnConfig\n      enable: false\n      max_assistant_turns: null\n      tool_config_path: null\n      max_user_turns: null\n      max_parallel_calls: 1\n      max_tool_response_length: 256\n      tool_response_truncate_side: middle\n      interaction_config_path: null\n      use_inference_chat_template: false\n      tokenization_sanity_check_mode: strict\n      format: hermes\n      num_repeat_rollouts: null\n    calculate_log_probs: false\n    agent:\n      _target_: verl.workers.config.AgentLoopConfig\n      num_workers: 8\n      default_agent_loop: single_turn_agent\n      agent_loop_config_path: null\n      custom_async_server:\n        _target_: verl.workers.config.CustomAsyncServerConfig\n        path: null\n        name: null\n    checkpoint_engine:\n      _target_: verl.workers.config.CheckpointEngineConfig\n      backend: naive\n      update_weights_bucket_megabytes: 2048\n      engine_kwargs: {}\n    trace:\n      _target_: verl.workers.config.TraceConfig\n      project_name: ${oc.select:trainer.project_name,null}\n      experiment_name: ${oc.select:trainer.experiment_name,null}\n      backend: null\n      token2text: false\n      max_samples_per_step_per_worker: null\n    skip_rollout: false\n    skip_dump_dir: /tmp/rollout_dump\n    skip_tokenizer_init: true\n    enable_rollout_routing_replay: false\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: ${oc.select:actor_rollout_ref.actor.profiler.enable,false}\n      all_ranks: ${oc.select:actor_rollout_ref.actor.profiler.all_ranks,false}\n      ranks: ${oc.select:actor_rollout_ref.actor.profiler.ranks,[]}\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.contents,[]}\n          level: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.level,level0}\n          analysis: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.analysis,false}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.discrete,false}\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.contents,[]}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.discrete,false}\n    prometheus:\n      _target_: verl.workers.config.PrometheusConfig\n      enable: false\n      port: 9090\n      file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml\n      served_model_name: ${oc.select:actor_rollout_ref.model.path,null}\n    quantization: null\n    quantization_config_file: null\n    mtp: ${oc.select:actor_rollout_ref.model.mtp, null}\n    qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,null}\n    layered_summon: false\n  model:\n    _target_: verl.workers.config.HFModelConfig\n    path: ~/models/deepseek-llm-7b-chat\n    hf_config_path: null\n    tokenizer_path: null\n    use_shm: false\n    trust_remote_code: false\n    custom_chat_template: null\n    external_lib: null\n    override_config: {}\n    enable_gradient_checkpointing: true\n    enable_activation_offload: false\n    use_remove_padding: true\n    lora_rank: 0\n    lora_alpha: 16\n    target_modules: all-linear\n    exclude_modules: null\n    lora_adapter_path: null\n    use_liger: false\n    use_fused_kernels: false\n    fused_kernel_options:\n      impl_backend: torch\n    tiled_mlp:\n      enabled: false\n      num_shards: 4\n    mtp:\n      _target_: verl.workers.config.MtpConfig\n      enable: false\n      enable_train: false\n      enable_rollout: false\n      detach_encoder: false\n      mtp_loss_scaling_factor: 0.1\n      speculative_algorithm: EAGLE\n      speculative_num_steps: 3\n      speculative_eagle_topk: 1\n      speculative_num_draft_tokens: 4\n      method: mtp\n      num_speculative_tokens: 1\n  hybrid_engine: true\n  nccl_timeout: 600\ndata:\n  tokenizer: null\n  use_shm: false\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1\n  val_max_samples: -1\n  prompt_key: prompt\n  reward_fn_key: data_source\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null\n  tool_config_path: ${oc.select:actor_rollout_ref.rollout.multi_turn.tool_config_path,\n    null}\n  return_raw_input_ids: false\n  return_raw_chat: true\n  return_full_prompt: false\n  shuffle: true\n  seed: null\n  dataloader_num_workers: 8\n  image_patch_size: 14\n  validation_shuffle: false\n  filter_overlong_prompts: false\n  filter_overlong_prompts_workers: 1\n  truncation: error\n  image_key: images\n  video_key: videos\n  trust_remote_code: false\n  custom_cls:\n    path: null\n    name: null\n  return_multi_modal_inputs: true\n  sampler:\n    class_path: null\n    class_name: null\n  datagen:\n    path: null\n    name: null\n  apply_chat_template_kwargs: {}\ncritic:\n  optim:\n    _target_: verl.workers.config.FSDPOptimizerConfig\n    optimizer: AdamW\n    optimizer_impl: torch.optim\n    lr: 1.0e-05\n    lr_warmup_steps_ratio: 0.0\n    total_training_steps: -1\n    weight_decay: 0.01\n    lr_warmup_steps: -1\n    betas:\n    - 0.9\n    - 0.999\n    clip_grad: 1.0\n    min_lr_ratio: 0.0\n    num_cycles: 0.5\n    lr_scheduler_type: constant\n    zero_indexed_step: true\n    warmup_style: null\n    override_optimizer_config: null\n  model:\n    fsdp_config:\n      _target_: verl.workers.config.FSDPEngineConfig\n      wrap_policy:\n        min_num_params: 0\n      param_offload: false\n      optimizer_offload: false\n      offload_policy: false\n      reshard_after_forward: true\n      fsdp_size: -1\n      forward_prefetch: false\n      model_dtype: fp32\n      use_orig_params: false\n      seed: 42\n      full_determinism: false\n      ulysses_sequence_parallel_size: 1\n      entropy_from_logits_with_chunking: false\n      use_torch_compile: true\n      entropy_checkpointing: false\n      forward_only: false\n      strategy: fsdp\n      dtype: bfloat16\n      qat:\n        _target_: verl.workers.config.QATEngineConfig\n        enable: false\n        mode: w4a16\n        group_size: 16\n        ignore_patterns:\n        - lm_head\n        - embed_tokens\n        - re:.*mlp.gate$\n        activation_observer: static_minmax\n        quantization_config_path: null\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${oc.select:actor_rollout_ref.model.path,\"~/models/deepseek-llm-7b-chat\"}\n    override_config: {}\n    external_lib: ${oc.select:actor_rollout_ref.model.external_lib,null}\n    trust_remote_code: ${oc.select:actor_rollout_ref.model.trust_remote_code,false}\n    _target_: verl.workers.config.FSDPCriticModelCfg\n    use_shm: false\n    enable_gradient_checkpointing: true\n    enable_activation_offload: false\n    use_remove_padding: false\n    lora_rank: 0\n    lora_alpha: 16\n    target_modules: all-linear\n    tiled_mlp:\n      enabled: false\n      num_shards: 4\n  _target_: verl.workers.config.FSDPCriticConfig\n  rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n  strategy: fsdp\n  enable: null\n  ppo_mini_batch_size: ${oc.select:actor_rollout_ref.actor.ppo_mini_batch_size,256}\n  ppo_micro_batch_size: null\n  ppo_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size,null}\n  use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n  ppo_max_token_len_per_gpu: 32768\n  forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu}\n  ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1}\n  shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false}\n  data_loader_seed: 42\n  cliprange_value: 0.5\n  loss_agg_mode: ${oc.select:actor_rollout_ref.actor.loss_agg_mode,token-mean}\n  checkpoint:\n    _target_: verl.trainer.config.CheckpointConfig\n    save_contents:\n    - model\n    - optimizer\n    - extra\n    load_contents: ${.save_contents}\n    async_save: false\n    mbridge_config: {}\n  profiler:\n    _target_: verl.utils.profiler.ProfilerConfig\n    tool: ${oc.select:global_profiler.tool,null}\n    enable: false\n    all_ranks: false\n    ranks: []\n    save_path: ${oc.select:global_profiler.save_path,null}\n    tool_config:\n      nsys:\n        _target_: verl.utils.profiler.config.NsightToolConfig\n        discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n      npu:\n        _target_: verl.utils.profiler.config.NPUToolConfig\n        contents: []\n        level: level0\n        analysis: true\n        discrete: false\n      torch:\n        _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n        contents: []\n        discrete: false\n      torch_memory:\n        _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n        trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n        stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n  forward_micro_batch_size: ${oc.select:.ppo_micro_batch_size,null}\n  forward_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size_per_gpu,null}\n  ulysses_sequence_parallel_size: 1\n  grad_clip: 1.0\ncustom_reward_function:\n  path: null\n  name: null\nreward_model:\n  num_workers: null\n  reward_manager: null\n  enable: null\n  enable_resource_pool: null\n  n_gpus_per_node: null\n  nnodes: null\n  reward_loop_source: null\n  reward_loop_module_path: null\n  reward_loop_class_name: null\n  model:\n    path: null\n    external_lib: null\n    trust_remote_code: null\n  rollout:\n    name: null\n    dtype: null\n    gpu_memory_utilization: null\n    enforce_eager: null\n    cudagraph_capture_sizes: null\n    free_cache_engine: null\n    data_parallel_size: null\n    expert_parallel_size: null\n    tensor_model_parallel_size: null\n    max_num_batched_tokens: null\n    max_model_len: null\n    max_num_seqs: null\n    load_format: null\n    engine_kwargs: null\n    limit_images: null\n    enable_chunked_prefill: null\n    enable_prefix_caching: null\n    disable_log_stats: null\n    skip_tokenizer_init: null\n    prompt_length: null\n    response_length: null\nsandbox_fusion:\n  url: null\n  max_concurrent: null\n  memory_limit_mb: null\nreward:\n  num_workers: 8\n  custom_reward_function:\n    path: null\n    name: compute_score\n  reward_manager:\n    _target_: verl.workers.config.reward_model.RewardManagerConfig\n    source: register\n    name: naive\n    module:\n      _target_: verl.trainer.config.config.ModuleConfig\n      path: null\n      name: custom_reward_manager\n  reward_model:\n    enable: false\n    enable_resource_pool: false\n    n_gpus_per_node: 8\n    nnodes: 0\n    model_path: null\n    rollout:\n      _target_: verl.workers.config.RolloutConfig\n      name: ???\n      dtype: bfloat16\n      gpu_memory_utilization: 0.5\n      enforce_eager: true\n      cudagraph_capture_sizes: null\n      free_cache_engine: true\n      data_parallel_size: 1\n      expert_parallel_size: 1\n      tensor_model_parallel_size: 2\n      max_num_batched_tokens: 8192\n      max_model_len: null\n      max_num_seqs: 1024\n      load_format: auto\n      engine_kwargs: {}\n      limit_images: null\n      enable_chunked_prefill: true\n      enable_prefix_caching: true\n      disable_log_stats: true\n      skip_tokenizer_init: false\n      prompt_length: 2048\n      response_length: 2048\n  sandbox_fusion:\n    url: null\n    max_concurrent: 64\n    memory_limit_mb: 1024\nalgorithm:\n  rollout_correction:\n    rollout_is: null\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    rollout_rs_threshold: null\n    bypass_mode: false\n    loss_type: ppo_clip\n    rollout_is_batch_normalize: false\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: true\n  use_kl_in_reward: false\n  kl_penalty: kl\n  kl_ctrl:\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: false\n  pf_ppo:\n    reweight_method: pow\n    weight_pow: 2.0\ntrainer:\n  balance_batch: true\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger:\n  - console\n  - wandb\n  log_val_generations: 0\n  rollout_data_dir: null\n  validation_data_dir: null\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n  resume_mode: auto\n  resume_from_path: null\n  val_before_train: true\n  val_only: false\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  del_local_ckpt_after_load: false\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  use_legacy_worker_impl: auto\nglobal_profiler:\n  _target_: verl.utils.profiler.ProfilerConfig\n  tool: null\n  steps: null\n  profile_continuous_steps: false\n  save_path: outputs/profile\n  global_tool_config:\n    nsys:\n      _target_: verl.utils.profiler.config.NsightToolConfig\n      discrete: false\n      controller_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n      worker_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n        capture-range: cudaProfilerApi\n        capture-range-end: null\n        kill: none\n    torch_memory:\n      trace_alloc_max_entries: 100000\n      stack_depth: 32\n      context: all\n      stacks: all\n      kw_args: {}\ntransfer_queue:\n  enable: false\nray_kwargs:\n  ray_init:\n    num_cpus: null\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/_generated_ppo_veomni_trainer.yaml",
    "content": "# This reference configration yaml is automatically generated via 'scripts/generate_trainer_config.sh'\n# in which it invokes 'python3 scripts/print_cfg.py --cfg job model_engine=veomni' to flatten the 'verl/trainer/config/ppo_trainer.yaml' config fields into a single file.\n# Do not modify this file directly.\n# The file is usually only for reference and never used.\n\nactor_rollout_ref:\n  actor:\n    optim:\n      _target_: verl.workers.config.VeOmniOptimizerConfig\n      optimizer: adamw\n      lr: 1.0e-06\n      lr_min: 0.0\n      lr_start: 0.0\n      lr_warmup_steps_ratio: 0.0\n      lr_decay_ratio: 1.0\n      total_training_steps: -1\n      weight_decay: 0.01\n      lr_warmup_steps: -1\n      betas:\n      - 0.9\n      - 0.999\n      clip_grad: 1.0\n      lr_scheduler_type: cosine\n      override_optimizer_config: {}\n    veomni:\n      _target_: verl.workers.config.VeOmniEngineConfig\n      param_offload: false\n      optimizer_offload: false\n      fsdp_size: -1\n      ulysses_parallel_size: 1\n      expert_parallel_size: 1\n      mixed_precision: true\n      seed: 42\n      full_determinism: false\n      init_device: meta\n      enable_full_shard: true\n      ckpt_manager: dcp\n      forward_prefetch: true\n      strategy: veomni\n      use_torch_compile: false\n      forward_only: false\n      enable_fsdp_offload: false\n      enable_reentrant: false\n      attn_implementation: flash_attention_2\n      moe_implementation: fused\n      force_use_huggingface: false\n      activation_gpu_limit: 0.0\n    _target_: verl.workers.config.VeOmniActorConfig\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: veomni\n    ppo_mini_batch_size: 256\n    ppo_micro_batch_size: null\n    ppo_micro_batch_size_per_gpu: null\n    use_dynamic_bsz: false\n    ppo_max_token_len_per_gpu: 16384\n    clip_ratio: 0.2\n    clip_ratio_low: 0.2\n    clip_ratio_high: 0.2\n    tau_pos: 1.0\n    tau_neg: 1.05\n    freeze_vision_tower: false\n    policy_loss:\n      _target_: verl.workers.config.PolicyLossConfig\n      loss_mode: vanilla\n      clip_cov_ratio: 0.0002\n      clip_cov_lb: 1.0\n      clip_cov_ub: 5.0\n      kl_cov_ratio: 0.0002\n      ppo_kl_coef: 0.1\n    clip_ratio_c: 3.0\n    loss_agg_mode: token-mean\n    loss_scale_factor: null\n    entropy_coeff: 0\n    calculate_entropy: false\n    use_kl_loss: false\n    use_prefix_grouper: false\n    use_torch_compile: true\n    kl_loss_coef: 0.001\n    kl_loss_type: low_var_kl\n    ppo_epochs: 1\n    shuffle: false\n    data_loader_seed: 42\n    checkpoint:\n      _target_: verl.trainer.config.CheckpointConfig\n      save_contents:\n      - model\n      - optimizer\n      - extra\n      load_contents: ${.save_contents}\n      async_save: false\n      mbridge_config: {}\n    use_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n  ref:\n    rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n    strategy: veomni\n    use_torch_compile: ${oc.select:actor_rollout_ref.actor.use_torch_compile,true}\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: false\n      all_ranks: false\n      ranks: []\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        nsys:\n          _target_: verl.utils.profiler.config.NsightToolConfig\n          discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: []\n          level: level0\n          analysis: true\n          discrete: false\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: []\n          discrete: false\n        torch_memory:\n          _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n          trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n          stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n    router_replay:\n      _target_: verl.workers.config.RouterReplayConfig\n      mode: disabled\n      record_file: null\n      replay_file: null\n    veomni:\n      _target_: verl.workers.config.VeOmniEngineConfig\n      param_offload: ${oc.select:actor_rollout_ref.actor.veomni.param_offload,False}\n      optimizer_offload: false\n      fsdp_size: ${oc.select:actor_rollout_ref.actor.veomni.fsdp_size,-1}\n      ulysses_parallel_size: ${oc.select:actor_rollout_ref.actor.veomni.ulysses_parallel_size,1}\n      expert_parallel_size: ${oc.select:actor_rollout_ref.actor.veomni.expert_parallel_size,1}\n      mixed_precision: true\n      seed: ${oc.select:actor_rollout_ref.actor.veomni.seed,42}\n      full_determinism: false\n      init_device: meta\n      enable_full_shard: true\n      ckpt_manager: dcp\n      forward_prefetch: true\n      strategy: veomni\n      use_torch_compile: false\n      forward_only: true\n      enable_fsdp_offload: false\n      enable_reentrant: false\n      attn_implementation: ${oc.select:actor_rollout_ref.actor.veomni.attn_implementation,flash_attention_2}\n      moe_implementation: ${oc.select:actor_rollout_ref.actor.veomni.moe_implementation,fused}\n      force_use_huggingface: false\n      activation_gpu_limit: 0.0\n    _target_: verl.workers.config.VeOmniActorConfig\n  rollout:\n    _target_: verl.workers.config.RolloutConfig\n    name: ???\n    mode: async\n    nnodes: 0\n    n_gpus_per_node: ${oc.select:trainer.n_gpus_per_node,8}\n    temperature: 1.0\n    top_k: -1\n    top_p: 1\n    prompt_length: ${oc.select:data.max_prompt_length,512}\n    response_length: ${oc.select:data.max_response_length,512}\n    dtype: bfloat16\n    gpu_memory_utilization: 0.5\n    ignore_eos: false\n    enforce_eager: false\n    cudagraph_capture_sizes: null\n    free_cache_engine: true\n    tensor_model_parallel_size: 2\n    data_parallel_size: 1\n    expert_parallel_size: 1\n    pipeline_model_parallel_size: 1\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    enable_chunked_prefill: true\n    enable_prefix_caching: true\n    logprobs_mode: processed_logprobs\n    scheduling_policy: fcfs\n    load_format: dummy\n    log_prob_micro_batch_size: null\n    log_prob_micro_batch_size_per_gpu: null\n    log_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n    log_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n    disable_log_stats: true\n    do_sample: true\n    'n': 1\n    over_sample_rate: 0\n    multi_stage_wake_up: false\n    engine_kwargs:\n      vllm: {}\n      sglang: {}\n      trtllm: {}\n    val_kwargs:\n      _target_: verl.workers.config.SamplingConfig\n      top_k: -1\n      top_p: 1.0\n      temperature: 0\n      'n': 1\n      do_sample: false\n    multi_turn:\n      _target_: verl.workers.config.MultiTurnConfig\n      enable: false\n      max_assistant_turns: null\n      tool_config_path: null\n      max_user_turns: null\n      max_parallel_calls: 1\n      max_tool_response_length: 256\n      tool_response_truncate_side: middle\n      interaction_config_path: null\n      use_inference_chat_template: false\n      tokenization_sanity_check_mode: strict\n      format: hermes\n      num_repeat_rollouts: null\n    calculate_log_probs: false\n    agent:\n      _target_: verl.workers.config.AgentLoopConfig\n      num_workers: 8\n      default_agent_loop: single_turn_agent\n      agent_loop_config_path: null\n      custom_async_server:\n        _target_: verl.workers.config.CustomAsyncServerConfig\n        path: null\n        name: null\n    checkpoint_engine:\n      _target_: verl.workers.config.CheckpointEngineConfig\n      backend: naive\n      update_weights_bucket_megabytes: 2048\n      engine_kwargs: {}\n    trace:\n      _target_: verl.workers.config.TraceConfig\n      project_name: ${oc.select:trainer.project_name,null}\n      experiment_name: ${oc.select:trainer.experiment_name,null}\n      backend: null\n      token2text: false\n      max_samples_per_step_per_worker: null\n    skip_rollout: false\n    skip_dump_dir: /tmp/rollout_dump\n    skip_tokenizer_init: true\n    enable_rollout_routing_replay: false\n    profiler:\n      _target_: verl.utils.profiler.ProfilerConfig\n      tool: ${oc.select:global_profiler.tool,null}\n      enable: ${oc.select:actor_rollout_ref.actor.profiler.enable,false}\n      all_ranks: ${oc.select:actor_rollout_ref.actor.profiler.all_ranks,false}\n      ranks: ${oc.select:actor_rollout_ref.actor.profiler.ranks,[]}\n      save_path: ${oc.select:global_profiler.save_path,null}\n      tool_config:\n        npu:\n          _target_: verl.utils.profiler.config.NPUToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.contents,[]}\n          level: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.level,level0}\n          analysis: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.analysis,false}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.discrete,false}\n        torch:\n          _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n          contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.contents,[]}\n          discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.discrete,false}\n    prometheus:\n      _target_: verl.workers.config.PrometheusConfig\n      enable: false\n      port: 9090\n      file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml\n      served_model_name: ${oc.select:actor_rollout_ref.model.path,null}\n    quantization: null\n    quantization_config_file: null\n    mtp: ${oc.select:actor_rollout_ref.model.mtp, null}\n    qat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,null}\n    layered_summon: false\n  model:\n    _target_: verl.workers.config.HFModelConfig\n    path: ~/models/deepseek-llm-7b-chat\n    hf_config_path: null\n    tokenizer_path: null\n    use_shm: false\n    trust_remote_code: false\n    custom_chat_template: null\n    external_lib: null\n    override_config: {}\n    enable_gradient_checkpointing: true\n    enable_activation_offload: false\n    use_remove_padding: true\n    lora_rank: 0\n    lora_alpha: 16\n    target_modules: all-linear\n    exclude_modules: null\n    lora_adapter_path: null\n    use_liger: false\n    use_fused_kernels: false\n    fused_kernel_options:\n      impl_backend: torch\n    tiled_mlp:\n      enabled: false\n      num_shards: 4\n    mtp:\n      _target_: verl.workers.config.MtpConfig\n      enable: false\n      enable_train: false\n      enable_rollout: false\n      detach_encoder: false\n      mtp_loss_scaling_factor: 0.1\n      speculative_algorithm: EAGLE\n      speculative_num_steps: 3\n      speculative_eagle_topk: 1\n      speculative_num_draft_tokens: 4\n      method: mtp\n      num_speculative_tokens: 1\n  hybrid_engine: true\n  nccl_timeout: 600\ndata:\n  tokenizer: null\n  use_shm: false\n  train_files: ~/data/rlhf/gsm8k/train.parquet\n  val_files: ~/data/rlhf/gsm8k/test.parquet\n  train_max_samples: -1\n  val_max_samples: -1\n  prompt_key: prompt\n  reward_fn_key: data_source\n  max_prompt_length: 512\n  max_response_length: 512\n  train_batch_size: 1024\n  val_batch_size: null\n  tool_config_path: ${oc.select:actor_rollout_ref.rollout.multi_turn.tool_config_path,\n    null}\n  return_raw_input_ids: false\n  return_raw_chat: true\n  return_full_prompt: false\n  shuffle: true\n  seed: null\n  dataloader_num_workers: 8\n  image_patch_size: 14\n  validation_shuffle: false\n  filter_overlong_prompts: false\n  filter_overlong_prompts_workers: 1\n  truncation: error\n  image_key: images\n  video_key: videos\n  trust_remote_code: false\n  custom_cls:\n    path: null\n    name: null\n  return_multi_modal_inputs: true\n  sampler:\n    class_path: null\n    class_name: null\n  datagen:\n    path: null\n    name: null\n  apply_chat_template_kwargs: {}\ncritic:\n  optim:\n    _target_: verl.workers.config.VeOmniOptimizerConfig\n    optimizer: adamw\n    lr: 1.0e-05\n    lr_min: 0.0\n    lr_start: 0.0\n    lr_warmup_steps_ratio: 0.0\n    lr_decay_ratio: 1.0\n    total_training_steps: -1\n    weight_decay: 0.01\n    lr_warmup_steps: -1\n    betas:\n    - 0.9\n    - 0.999\n    clip_grad: 1.0\n    lr_scheduler_type: cosine\n    override_optimizer_config: {}\n  veomni:\n    _target_: verl.workers.config.VeOmniEngineConfig\n    param_offload: false\n    optimizer_offload: false\n    fsdp_size: -1\n    ulysses_parallel_size: 1\n    expert_parallel_size: 1\n    mixed_precision: true\n    seed: 42\n    full_determinism: false\n    init_device: meta\n    enable_full_shard: true\n    ckpt_manager: dcp\n    forward_prefetch: true\n    strategy: veomni\n    use_torch_compile: false\n    forward_only: false\n    enable_fsdp_offload: false\n    enable_reentrant: false\n    attn_implementation: flash_attention_2\n    moe_implementation: fused\n    force_use_huggingface: false\n    activation_gpu_limit: 0.0\n  _target_: verl.workers.config.VeOmniCriticConfig\n  rollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n  strategy: veomni\n  enable: null\n  model:\n    path: ~/models/deepseek-llm-7b-chat\n    tokenizer_path: ${oc.select:actor_rollout_ref.model.path,\"~/models/deepseek-llm-7b-chat\"}\n    override_config: {}\n    external_lib: ${oc.select:actor_rollout_ref.model.external_lib,null}\n    trust_remote_code: ${oc.select:actor_rollout_ref.model.trust_remote_code,false}\n    _target_: verl.trainer.config.BaseModelConfig\n  ppo_mini_batch_size: ${oc.select:actor_rollout_ref.actor.ppo_mini_batch_size,256}\n  ppo_micro_batch_size: null\n  ppo_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size,null}\n  use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n  ppo_max_token_len_per_gpu: 32768\n  forward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu}\n  ppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1}\n  shuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false}\n  data_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n  cliprange_value: 0.5\n  loss_agg_mode: ${oc.select:actor_rollout_ref.actor.loss_agg_mode,token-mean}\n  checkpoint:\n    _target_: verl.trainer.config.CheckpointConfig\n    save_contents:\n    - model\n    - optimizer\n    - extra\n    load_contents: ${.save_contents}\n    async_save: false\n    mbridge_config: {}\n  profiler:\n    _target_: verl.utils.profiler.ProfilerConfig\n    tool: ${oc.select:global_profiler.tool,null}\n    enable: false\n    all_ranks: false\n    ranks: []\n    save_path: ${oc.select:global_profiler.save_path,null}\n    tool_config:\n      nsys:\n        _target_: verl.utils.profiler.config.NsightToolConfig\n        discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n      npu:\n        _target_: verl.utils.profiler.config.NPUToolConfig\n        contents: []\n        level: level0\n        analysis: true\n        discrete: false\n      torch:\n        _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n        contents: []\n        discrete: false\n      torch_memory:\n        _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n        trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n        stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\ncustom_reward_function:\n  path: null\n  name: null\nreward_model:\n  num_workers: null\n  reward_manager: null\n  enable: null\n  enable_resource_pool: null\n  n_gpus_per_node: null\n  nnodes: null\n  reward_loop_source: null\n  reward_loop_module_path: null\n  reward_loop_class_name: null\n  model:\n    path: null\n    external_lib: null\n    trust_remote_code: null\n  rollout:\n    name: null\n    dtype: null\n    gpu_memory_utilization: null\n    enforce_eager: null\n    cudagraph_capture_sizes: null\n    free_cache_engine: null\n    data_parallel_size: null\n    expert_parallel_size: null\n    tensor_model_parallel_size: null\n    max_num_batched_tokens: null\n    max_model_len: null\n    max_num_seqs: null\n    load_format: null\n    engine_kwargs: null\n    limit_images: null\n    enable_chunked_prefill: null\n    enable_prefix_caching: null\n    disable_log_stats: null\n    skip_tokenizer_init: null\n    prompt_length: null\n    response_length: null\nsandbox_fusion:\n  url: null\n  max_concurrent: null\n  memory_limit_mb: null\nreward:\n  num_workers: 8\n  custom_reward_function:\n    path: null\n    name: compute_score\n  reward_manager:\n    _target_: verl.workers.config.reward_model.RewardManagerConfig\n    source: register\n    name: naive\n    module:\n      _target_: verl.trainer.config.config.ModuleConfig\n      path: null\n      name: custom_reward_manager\n  reward_model:\n    enable: false\n    enable_resource_pool: false\n    n_gpus_per_node: 8\n    nnodes: 0\n    model_path: null\n    rollout:\n      _target_: verl.workers.config.RolloutConfig\n      name: ???\n      dtype: bfloat16\n      gpu_memory_utilization: 0.5\n      enforce_eager: true\n      cudagraph_capture_sizes: null\n      free_cache_engine: true\n      data_parallel_size: 1\n      expert_parallel_size: 1\n      tensor_model_parallel_size: 2\n      max_num_batched_tokens: 8192\n      max_model_len: null\n      max_num_seqs: 1024\n      load_format: auto\n      engine_kwargs: {}\n      limit_images: null\n      enable_chunked_prefill: true\n      enable_prefix_caching: true\n      disable_log_stats: true\n      skip_tokenizer_init: false\n      prompt_length: 2048\n      response_length: 2048\n  sandbox_fusion:\n    url: null\n    max_concurrent: 64\n    memory_limit_mb: 1024\nalgorithm:\n  rollout_correction:\n    rollout_is: null\n    rollout_is_threshold: 2.0\n    rollout_rs: null\n    rollout_rs_threshold: null\n    bypass_mode: false\n    loss_type: ppo_clip\n    rollout_is_batch_normalize: false\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: true\n  use_kl_in_reward: false\n  kl_penalty: kl\n  kl_ctrl:\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: false\n  pf_ppo:\n    reweight_method: pow\n    weight_pow: 2.0\ntrainer:\n  balance_batch: true\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger:\n  - console\n  - wandb\n  log_val_generations: 0\n  rollout_data_dir: null\n  validation_data_dir: null\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n  resume_mode: auto\n  resume_from_path: null\n  val_before_train: true\n  val_only: false\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  del_local_ckpt_after_load: false\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  use_legacy_worker_impl: auto\nglobal_profiler:\n  _target_: verl.utils.profiler.ProfilerConfig\n  tool: null\n  steps: null\n  profile_continuous_steps: false\n  save_path: outputs/profile\n  global_tool_config:\n    nsys:\n      _target_: verl.utils.profiler.config.NsightToolConfig\n      discrete: false\n      controller_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n      worker_nsight_options:\n        trace: cuda,nvtx,cublas,ucx\n        cuda-memory-usage: 'true'\n        cuda-graph-trace: graph\n        capture-range: cudaProfilerApi\n        capture-range-end: null\n        kill: none\n    torch_memory:\n      trace_alloc_max_entries: 100000\n      stack_depth: 32\n      context: all\n      stacks: all\n      kw_args: {}\ntransfer_queue:\n  enable: false\nray_kwargs:\n  ray_init:\n    num_cpus: null\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/actor/actor.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# Target class for this configuration\n_target_: verl.workers.config.ActorConfig\n\n# Number of rollouts per update (mirrors actor rollout_n)\nrollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n\n# the abstract actor configs\n# fsdp, fsdp2 or megatron. must be set.\nstrategy: ???\n\n# Split each sample into sub-batches of this size for PPO\nppo_mini_batch_size: 256\n\n# [Deprecated] Global micro batch size\nppo_micro_batch_size: null\n\n# Local per-GPU micro batch size\nppo_micro_batch_size_per_gpu: null\n\n# Whether to automatically adjust batch size at runtime\n# oc.select: the default val for ref.log_prob_use_dynamic_bsz\nuse_dynamic_bsz: false\n\n# Max tokens per GPU in one PPO batch; affects gradient accumulation\n# Typically it should be: n * ${data.max_prompt_length} + ${data.max_response_length}\n# oc.select: the default val for ref.log_prob_max_token_len_per_gpu\nppo_max_token_len_per_gpu: 16384\n\n# PPO clip ratio\nclip_ratio: 0.2\n\n# Lower bound for asymmetric clipping (used in dual-clip PPO)\nclip_ratio_low: 0.2\n\n# Upper bound for asymmetric clipping (used in dual-clip PPO)\nclip_ratio_high: 0.2\n\n# Positive and negative tau for smoothing function in SAPO (https://arxiv.org/pdf/2511.20347)\n# default values used in the paper with Qwen3-30B-A3B-Base\ntau_pos: 1.0\n\n# negative tau for smoothing function in SAPO\ntau_neg: 1.05\n\n# Whether to freeze vision model, if set true, it will be freeze vision model\nfreeze_vision_tower: false\n\n# policy loss config\npolicy_loss:\n\n  # # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.PolicyLossConfig\n\n  # Loss function mode: vanilla / clip-cov / kl-cov /gpg from https://arxiv.org/abs/2505.22617\n  loss_mode: \"vanilla\"\n\n  # Ratio of tokens to be clipped for clip-cov loss\n  clip_cov_ratio: 0.0002\n\n  # Lower bound for clip-cov loss\n  clip_cov_lb: 1.0\n\n  # Upper bound for clip-cov loss\n  clip_cov_ub: 5.0\n\n  # Ratio of tokens to be applied kl penalty for kl-cov loss\n  kl_cov_ratio: 0.0002\n\n  # KL divergence penalty coefficient\n  ppo_kl_coef: 0.1\n\n# Constant C in Dual-clip PPO; clips when advantage < 0 and ratio > C\nclip_ratio_c: 3.0\n\n# Loss aggregation mode: \"token-mean\", \"seq-mean-token-sum\", \"seq-mean-token-mean\", or \"seq-mean-token-sum-norm\"\nloss_agg_mode: token-mean\n\n# Scale factor for \"seq-mean-token-sum-norm\" loss aggregation mode.\n# If null, uses response_length. Set to a constant to ensure consistent normalization.\nloss_scale_factor: null\n\n# Entropy regularization coefficient in PPO loss\nentropy_coeff: 0\n\n# When true, the actor forward will request entropy from the model\ncalculate_entropy: false\n\n# Whether to use KL loss instead of KL reward penalty. True for GRPO\nuse_kl_loss: false\n\n# Whether to enable PrefixGrouper shared-prefix forward\nuse_prefix_grouper: false\n\n# Whether to use torch.compile()\n# oc.select: the default val for ref.use_torch_compile\nuse_torch_compile: true\n\n# KL loss coefficient when use_kl_loss is enabled. For GRPO\nkl_loss_coef: 0.001\n\n# Type of KL divergence loss. Options: \"kl\"(k1), \"abs\", \"mse\"(k2), \"low_var_kl\"(k3), \"full\"\nkl_loss_type: low_var_kl\n\n# Number of PPO epochs per batch\nppo_epochs: 1\n\n# Shuffle training data across PPO epochs\nshuffle: false\n\n# The seed used to construct mini-batch\ndata_loader_seed: 42\n\n# checkpoint configs\ncheckpoint:\n\n  # Target dataclass for this configuration\n  _target_: verl.trainer.config.CheckpointConfig\n\n  # What to include in saved checkpoints\n  # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n  save_contents: ['model', 'optimizer', 'extra']\n\n  # For more flexibility, you can specify the contents to load from the checkpoint.\n  # .xxx refers to the local variable xxx from the same level of hierarchy similar to python pkg\n  load_contents: ${.save_contents}\n\n  # Whether to save checkpoints asynchronously. Only effective for Megatron as of now.\n  async_save: False\n  \n  # Mbridge config extension.\n  # when vanilla_mbridge=True, and your filesystem is a distributed filesystem,(which means you write a file in node A\n  # and you can read the file in node B immediately)\n  # set `mbridge_config.distributed_filesystem=True` and `mbridge_config.memory_efficient=True` to \n  # speed up the checkpoint saving by 10x speed.\n  mbridge_config: {}\n\n# optimizer configs\noptim:\n\n  # Learning rate\n  lr: 1e-6\n\n  # Warmup steps ratio (used if lr_warmup_steps is 0 or negative)\n  lr_warmup_steps_ratio: 0.0\n\n  # Total training steps (must be overridden at runtime)\n  total_training_steps: -1\n\n  # Weight decay\n  weight_decay: 0.01\n\n  # Prioritized. None, 0 or Negative values mean delegating to lr_warmup_steps_ratio.\n  lr_warmup_steps: -1\n\n\n# Whether to use custom fused kernels (e.g., FlashAttention, fused MLP)\nuse_fused_kernels: ${oc.select:actor_rollout_ref.model.use_fused_kernels,false}\n\n# profile the actor model in `update_policy` \nprofiler:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.utils.profiler.ProfilerConfig\n\n  # profiler tool, default same as profiler.tool in global config\n  # choices: nsys, npu, torch\n  tool: ${oc.select:global_profiler.tool,null}\n\n  # whether enable profile on Actor\n  enable: False\n  \n  # Whether to profile all ranks.\n  all_ranks: False\n\n  # The ranks that will be profiled. [] or [0,1,...]\n  ranks: []\n\n  # profile results saving path\n  save_path: ${oc.select:global_profiler.save_path,null}\n\n  # specific tool config which only related to the role\n  tool_config:\n\n    # nsys tool config\n    nsys:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NsightToolConfig\n    \n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n    \n    # npu config\n    npu:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NPUToolConfig\n\n      # Contents to profile, can be empty\n      # options: npu, cpu, memory, shapes, module, stack\n      contents: []\n\n      # Collection level, optional values: level_none, level0, level1, level2.\n      level: \"level0\"\n\n      # Whether to automatically parse the data.\n      analysis: True\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: False\n    \n    # torch profiler config\n    torch:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n      # Contents to profile, can be empty\n      # options: cuda, cpu, memory, shapes, stack\n      contents: []\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: false\n\n    # torch memory profiler config\n    torch_memory:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n\n      # Maximum number of memory allocation entries to track\n      trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n\n      # Stack trace depth for memory allocations\n      stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n\n# Router replay configuration for MoE models\nrouter_replay:\n\n  # Target dataclass for this configuration\n  _target_: verl.workers.config.RouterReplayConfig\n\n  # Router replay mode: disabled, R2, R3\n  # - R2: Use R2 routing strategy (record mode)\n  # - R3: Use R3 routing strategy (record mode)\n  mode: disabled\n\n  # File path to save recorded routing decisions\n  # Required when mode is 'record', 'R2', or 'R3'\n  record_file: null\n\n  # File path to load recorded routing decisions for replay\n  # Required when mode is 'replay'\n  replay_file: null\n\n"
  },
  {
    "path": "verl/trainer/config/actor/dp_actor.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# defaults specify the default config from each component\ndefaults:\n\n  # fsdp optimizer config\n  - ../optim@optim: fsdp\n\n  # fsdp engine config\n  - ../engine@fsdp_config: fsdp\n\n  # dp actor config, inheriting from trainer/config/actor/actor.yaml\n  - actor\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Target class for this configuration\n_target_: verl.workers.config.FSDPActorConfig\n\n# TODO(haibin.lin): switch to fsdp2\nstrategy: fsdp\n\n# Gradient clipping for actor updates, specific to the strategy.\ngrad_clip: 1.0\n\n# Sequence parallelism size for Ulysses-style model parallelism\n# oc.select: the default val for ref.ulysses_sequence_parallel_size\n# [DEPRECATED] use fsdp_config.ulysses_sequence_parallel_size instead\nulysses_sequence_parallel_size: 1\n\n# calculate entropy with chunking to reduce memory peak\nentropy_from_logits_with_chunking: False\n\n# recompute entropy\nentropy_checkpointing: False\n\n# Whether to remove padding tokens in inputs during training\nuse_remove_padding: ${oc.select:actor_rollout_ref.model.use_remove_padding,false}\n\n# This computes Σπ² needed for the Logit-Gradient Norm proxy W(τ) = Σ_t[1 - 2π_t + Σπ²]\n# c.f. https://yingru.notion.site/The-Optimal-Token-Baseline-399211a558b782cfa936014c0d42dfb8\ncalculate_sum_pi_squared: False\n\n# Enable gradient checkpointing for sum_pi_squared computation (saves memory)\nsum_pi_squared_checkpointing: False\n\n# QAT (Quantization-Aware Training) configuration\n# When enabled:\n#   - QAT is automatically applied to actor model during training\n#   - Fused scales (QKV/GateUp) are automatically enabled for training-inference consistency\n#   - Fast quantization is used when syncing weights to vLLM rollout\n# Supported modes: \"w4a16\" (NVFP4 weight-only)\n# Note: \"w4a4\" mode is included in the code but currently has KL divergence issues and is NOT recommended for use.\n# For usage examples, see: https://github.com/verl-project/verl-recipe/blob/main/qat/README.md\nqat:\n\n  # Whether to enable QAT\n  enable: false\n\n  # Quantization mode: \"w4a16\" (weight-only). \"w4a4\" is experimental and not recommended.\n  mode: \"w4a16\"\n\n  # Quantization group size (NVFP4 requires 16)\n  group_size: 16\n\n  # Patterns to ignore (e.g., lm_head, embed_tokens)\n  ignore_patterns:\n\n    - \"lm_head\"\n    - \"embed_tokens\"\n    - \"re:.*mlp.gate$\"\n\n  # Activation observer for W4A4 mode: \"static_minmax\", \"memoryless_minmax\", or \"minmax\"\n  activation_observer: \"static_minmax\"\n\n  # Path to vLLM quantization config JSON file\n  quantization_config_path: null\n"
  },
  {
    "path": "verl/trainer/config/actor/megatron_actor.yaml",
    "content": "# megatron actor config, inheriting from trainer/config/actor/actor.yaml\ndefaults:\n  # megatron optimizer config\n  - ../optim@optim: megatron\n\n  # megatron engine config\n  - ../engine@megatron: megatron\n\n  - actor\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.McoreActorConfig\n\nstrategy: megatron\n\nload_weight: True\n"
  },
  {
    "path": "verl/trainer/config/actor/torchtitan_actor.yaml",
    "content": "# torchtitan actor config, inheriting from trainer/config/actor/actor.yaml\ndefaults:\n  # torchtitan optimizer config\n  - ../optim@optim: torchtitan\n\n  # torchtitan engine config\n  - ../engine@torchtitan: torchtitan\n\n  - actor\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.TorchTitanActorConfig\n\nstrategy: torchtitan\n"
  },
  {
    "path": "verl/trainer/config/actor/veomni_actor.yaml",
    "content": "# veomni actor config, inheriting from trainer/config/actor/actor.yaml\ndefaults:\n  # veomni optimizer config\n  - ../optim@optim: veomni\n\n  # veomni engine config\n  - ../engine@veomni: veomni\n\n  - actor\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.VeOmniActorConfig\n\nstrategy: veomni\n"
  },
  {
    "path": "verl/trainer/config/algorithm/rollout_correction.yaml",
    "content": "# Rollout Correction: corrects off-policy distribution shifts\n# See documentation: docs/algo/rollout_corr.md\n# Use presets: RolloutCorrectionConfig.decoupled_seq_is(), .bypass_pg_is(), etc.\n\n# IS aggregation level: null (disabled), \"token\" (per-token), \"sequence\" (per-sequence)\nrollout_is: null\n\n# Upper threshold for IS weight truncation (typical: 2.0-5.0)\nrollout_is_threshold: 2.0\n\n# RS aggregation level: null (disabled), e.g. \"token_k1\", \"seq_sum_k1\", \"seq_mean_k3\"\nrollout_rs: null\n\n# Threshold for rejection sampling (string or float; see code docs)\nrollout_rs_threshold: null\n\n# Operating mode: false = Decoupled (3 policies), true = Bypass (2 policies)\nbypass_mode: false\n\n# Loss type in bypass mode (bypass_mode=true):\n# - \"ppo_clip\": PPO clipped objective (IS handled by ratio, default)\n# - \"reinforce\": REINFORCE with explicit IS weights (no PPO clipping)\nloss_type: ppo_clip\n\n# Batch normalize IS weights: false = raw weights, true = normalize to mean=1.0\nrollout_is_batch_normalize: false\n"
  },
  {
    "path": "verl/trainer/config/algorithm.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nfrom verl.base_config import BaseConfig\n\n__all__ = [\"AlgoConfig\", \"FilterGroupsConfig\", \"KLControlConfig\", \"RolloutCorrectionConfig\"]\n\n\n@dataclass\nclass KLControlConfig(BaseConfig):\n    \"\"\"Configuration for KL control.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        type (str): Type of KL control. Can be \"fixed\" or \"adaptive\".\n        kl_coef (float): Initial coefficient for KL penalty.\n        horizon (int): Horizon value for adaptive controller.\n        target_kl (float): Target KL divergence for adaptive controller.\n    \"\"\"\n\n    type: str = \"fixed\"\n    kl_coef: float = 0.001\n    horizon: int = 10000\n    target_kl: float = 0.1\n\n\n@dataclass\nclass FilterGroupsConfig(BaseConfig):\n    \"\"\"Configuration for filter groups (used in DAPO and Entropy).\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        enable (bool): Whether to enable filter groups.\n        metric (Optional[str]): Metric to use for filtering: \"acc\", \"score\", \"seq_reward\", \"seq_final_reward\", etc.\n        max_num_gen_batches (int): Non-positive values mean no upper limit.\n    \"\"\"\n\n    enable: bool = False\n    metric: Optional[str] = None\n    max_num_gen_batches: int = 0\n\n\n@dataclass\nclass RolloutCorrectionConfig(BaseConfig):\n    \"\"\"Configuration for Rollout Correction (addresses off-policy issues in RL training).\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Rollout Correction handles off-policiness from multiple sources:\n    1. Policy mismatch: Rollout policy (e.g., vLLM BF16) vs Training policy (e.g., FSDP FP32)\n    2. Model update staleness: Rollout data collected from older policy checkpoints\n    3. General off-policy scenarios: Any distribution shift between data collection and training\n\n    For more details, see:\n    \"When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch\"\n    https://richardli.xyz/rl-collapse\n\n    This typed config replaces the old dict-based approach and provides:\n    - Type safety and validation\n    - Clear documentation of all parameters\n    - Named factory methods for common presets (TIS, MIS, etc.)\n    - Sensible defaults\n\n    Args:\n        rollout_is (Optional[str]): IS weight aggregation level.\n            - None: No IS weights (metrics only)\n            - \"token\": Per-token IS weights (low variance, biased)\n            - \"sequence\": Per-sequence IS weights (unbiased, high variance)\n            Default: \"sequence\"\n\n        rollout_is_threshold (float): Upper threshold for IS weight truncation/rejection.\n            Typical range: 1.5-5.0 for token level, 2.0-10.0 for sequence level.\n            Default: 2.0\n\n        rollout_is_batch_normalize (bool): Apply batch normalization to IS weights.\n            - True: Normalize IS weights to have mean=1.0 within each batch\n            - False: Use raw (truncated) IS weights (standard)\n            - Reduces variance by ensuring average weight is 1.0 per batch\n            - Only affects IS weight values, not rejection sampling\n            Default: False (no batch normalization)\n\n        rollout_rs (Optional[str]): Rejection sampling aggregation modes.\n            Accepts a comma-delimited list (duplicates removed) of canonical options implemented in\n            ``rollout_corr_helper``:\n            - \"token_k1\": Token-level rejection with ``-log r`` (ratio thresholds supplied via\n              ``rollout_rs_threshold`` as ``lower_upper``)\n            - \"token_k2\": Token-level rejection with ``0.5 * (log r)^2`` (upper bound only)\n            - \"token_k3\": Token-level rejection with ``exp(log r) - 1 - log r`` (upper bound only)\n            - \"seq_sum_k1\": Sequence sum of ``-log r`` (ratio bounds)\n            - \"seq_sum_k2\": Sequence sum of rejection with ``0.5 * (log r)^2`` (upper bound only)\n            - \"seq_sum_k3\": Sequence sum of rejection with ``exp(log r) - 1 - log r`` (upper bound only)\n            - \"seq_mean_k1\": Sequence mean of ``-log r`` (ratio bounds)\n            - \"seq_mean_k2\": Sequence mean of rejection with ``0.5 * (log r)^2`` (upper bound only)\n            - \"seq_mean_k3\": Sequence mean of rejection with ``exp(log r) - 1 - log r`` (upper bound only)\n            - \"seq_max_k2\": Sequence max of rejection with ``0.5 * (log r)^2`` (upper bound only)\n            - \"seq_max_k3\": Sequence max of rejection with ``exp(log r) - 1 - log r`` (upper bound only)\n            names automatically. Default: None\n\n        rollout_rs_threshold (Optional[Union[str, float]]): Threshold specification for rejection sampling.\n            Provide one value per option (single entry is broadcast when multiple options are supplied).\n            Ratio-based modes (``*k1``) expect ``lower_upper`` strings; supplying a single float implies\n            only the upper ratio bound, with the lower bound inferred as its reciprocal. Divergence modes\n            (k2/k3) expect positive upper bounds (float or string). Default: None\n\n        bypass_mode (bool): Operating mode - bypass or decoupled.\n            - True: Bypass mode - reuse rollout_log_prob as old_log_prob (2 policies)\n              Uses compute_policy_loss_bypass_mode() with loss_type selection\n            - False: Decoupled mode - compute old_log_prob separately (3 policies)\n              Uses standard PPO loss with IS weight correction\n            Default: False (decoupled mode)\n\n        loss_type (str): Loss function type in bypass mode (bypass_mode=True).\n            - \"reinforce\": REINFORCE-style policy gradient with explicit IS weights\n              L = -E[w * log π(a|s) * A] where w = π_current / π_rollout\n            - \"ppo_clip\": PPO clipped objective (IS handled by ratio, no explicit weights)\n              L = -E[min(r*A, clip(r)*A)] where r = π_current / π_rollout\n            Default: \"ppo_clip\"\n\n    Example:\n        # Create with defaults\n        config = RolloutCorrectionConfig()\n\n        # Decoupled PPO mode presets (3 policies: π_rollout, π_old, π_θ)\n        # IS weights correct for gap between π_old and π_rollout\n        config = RolloutCorrectionConfig.decoupled_token_is()  # Token-TIS\n        config = RolloutCorrectionConfig.decoupled_seq_is()    # Seq-TIS\n        config = RolloutCorrectionConfig.decoupled_seq_is_rs() # Seq-MIS\n        config = RolloutCorrectionConfig.decoupled_geo_rs()    # Geo-RS (ratio mode)\n\n        # Bypass mode presets (2 policies: π_rollout = π_old, π_θ)\n        # loss_type controls the loss function\n        # PPO-clip presets (ratio handles IS, so no separate IS weights needed):\n        config = RolloutCorrectionConfig.bypass_ppo_clip()              # PPO-clip only\n        config = RolloutCorrectionConfig.bypass_ppo_clip_geo_rs()       # PPO-clip + Geo-RS\n        config = RolloutCorrectionConfig.bypass_ppo_clip_k3_rs()        # PPO-clip + K3-RS\n        # REINFORCE presets (explicit IS weights):\n        config = RolloutCorrectionConfig.bypass_pg_is()                 # REINFORCE + Seq-TIS\n        config = RolloutCorrectionConfig.bypass_pg_geo_rs()             # REINFORCE + Geo-RS\n        config = RolloutCorrectionConfig.bypass_pg_geo_rs_seq_tis()     # REINFORCE + Geo-RS + Seq-TIS\n        config = RolloutCorrectionConfig.bypass_pg_geo_rs_token_tis()   # REINFORCE + Geo-RS + Token-TIS\n\n        # Decoupled Geometric ratio presets (length-normalized IS ratio)\n        config = RolloutCorrectionConfig.decoupled_geo_rs_seq_tis()           # Decoupled Geo-RS + Seq-TIS\n        config = RolloutCorrectionConfig.decoupled_geo_rs_token_tis()         # Decoupled Geo-RS + Token-TIS\n\n        # Decoupled K3 KL Estimator presets (more stable for small KL values)\n        config = RolloutCorrectionConfig.decoupled_k3_rs()                    # Decoupled K3-RS\n        config = RolloutCorrectionConfig.decoupled_k3_rs_seq_tis()            # Decoupled K3-RS + Seq-TIS\n        config = RolloutCorrectionConfig.decoupled_k3_rs_token_tis()          # Decoupled K3-RS + Token-TIS\n\n    Reference:\n        Liu, Li, Fu, Wang, Liu, Shen (2025)\n        \"When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch\"\n        https://richardli.xyz/rl-collapse\n    \"\"\"\n\n    rollout_is: Optional[str] = \"sequence\"\n    rollout_is_threshold: float = 2.0\n    rollout_is_batch_normalize: bool = False\n    rollout_rs: Optional[str] = None\n    rollout_rs_threshold: Optional[str | float] = None\n    bypass_mode: bool = False\n    loss_type: str = \"ppo_clip\"\n\n    @classmethod\n    def decoupled_token_is(cls, threshold: float = 2.0) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled Mode with Token-level Importance Sampling.\n\n        IS weight correction at token level in decoupled mode (three policies).\n\n        Args:\n            threshold (float): Upper threshold for IS weights. Default: 2.0\n\n        Returns:\n            RolloutCorrectionConfig configured for decoupled mode with token-level IS\n        \"\"\"\n        return cls(rollout_is=\"token\", rollout_is_threshold=threshold, rollout_rs=None)\n\n    @classmethod\n    def decoupled_seq_is(cls, threshold: float = 2.0) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled Mode with Sequence-level Importance Sampling.\n\n        IS weight correction at sequence level in decoupled mode (three policies).\n\n        Args:\n            threshold (float): Upper threshold for IS weights. Default: 2.0\n\n        Returns:\n            RolloutCorrectionConfig configured for decoupled mode with sequence-level IS\n        \"\"\"\n        return cls(rollout_is=\"sequence\", rollout_is_threshold=threshold, rollout_rs=None)\n\n    @classmethod\n    def decoupled_seq_is_rs(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: Optional[str | float] = \"0.5_2.0\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled Mode with Sequence-level IS + Rejection Sampling.\n\n        Sequence-level IS with sequence-level rejection sampling in decoupled mode.\n        Rejects entire sequences based on sequence-level IS weight.\n\n        Args:\n            is_threshold (float): Upper threshold for IS weights. Default: 2.0\n            rs_threshold (Optional[Union[str, float]]): Upper threshold for rejection sampling. Default: 0.5_2.0\n\n        Returns:\n            RolloutCorrectionConfig configured for decoupled mode with sequence IS + RS\n        \"\"\"\n        return cls(\n            rollout_is=\"sequence\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_sum_k1\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def decoupled_geo_rs(\n        cls,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled Mode with Geometric Mean Rejection Sampling (ratio-based).\n\n        Uses geometric mean IS ratio E[log(r)] for rejection sampling at sequence level.\n        This is a ratio-based mode (ideal = 0.0) with [lower, upper] threshold bounds.\n        Length-normalized but still uses IS ratio semantics.\n\n        Args:\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for decoupled mode with Geo-RS\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def bypass_ppo_clip(cls) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with PPO-clip loss.\n\n        PPO clipped objective in bypass mode. The PPO ratio = π_θ/π_rollout\n        already handles IS correction, so no explicit IS weights are applied.\n\n        Skips old_log_prob computation for faster execution (2 policies instead of 3).\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with PPO-clip\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=None,\n            bypass_mode=True,\n            loss_type=\"ppo_clip\",\n        )\n\n    @classmethod\n    def bypass_ppo_clip_geo_rs(\n        cls,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with PPO-clip loss and Geometric Mean RS (ratio-based).\n\n        PPO clipped objective in bypass mode with geometric mean IS ratio RS.\n        Uses E[log(r)] (ideal = 0.0) with [lower, upper] threshold bounds.\n\n        Args:\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with PPO-clip + Geo-RS\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n            bypass_mode=True,\n            loss_type=\"ppo_clip\",\n        )\n\n    @classmethod\n    def bypass_ppo_clip_k3_rs(\n        cls,\n        rs_threshold: float = 0.01,\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with PPO-clip loss and K3 Rejection Sampling.\n\n        PPO clipped objective in bypass mode with K3 KL estimator RS to mask outliers.\n        K3 is more stable than K1 for small KL values.\n        The PPO ratio = π_θ/π_rollout already handles IS correction.\n\n        Args:\n            rs_threshold (float): Max allowed K3 divergence. Default: 0.01\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with PPO-clip + K3-RS\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k3\",\n            rollout_rs_threshold=rs_threshold,\n            bypass_mode=True,\n            loss_type=\"ppo_clip\",\n        )\n\n    @classmethod\n    def bypass_pg_is(cls, threshold: float = 2.0) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with REINFORCE loss and IS Correction.\n\n        Uses REINFORCE loss with explicit IS correction in bypass mode.\n        No PPO clipping.\n\n        Args:\n            threshold (float): Upper threshold for IS weights. Default: 2.0\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with REINFORCE + IS\n        \"\"\"\n        return cls(\n            rollout_is=\"sequence\",\n            rollout_is_threshold=threshold,\n            rollout_rs=None,\n            bypass_mode=True,\n            loss_type=\"reinforce\",\n        )\n\n    @classmethod\n    def bypass_pg_geo_rs(\n        cls,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with REINFORCE loss and Geometric Mean RS (ratio-based).\n\n        REINFORCE with geometric mean IS ratio rejection sampling in bypass mode.\n        Uses E[log(r)] (ideal = 0.0) with [lower, upper] threshold bounds.\n\n        Args:\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with REINFORCE + Geo-RS\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n            bypass_mode=True,\n            loss_type=\"reinforce\",\n        )\n\n    @classmethod\n    def decoupled_geo_rs_seq_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled mode with Geometric Mean RS and Sequence-level Truncated IS (ratio-based).\n\n        Combines the Geometric Mean Filter (ratio-based validity check) with\n        Clipped Sequence Weight (debiasing). Uses E[log(r)] (ideal = 0.0).\n\n        Args:\n            is_threshold (float): Upper threshold for sequence IS weights. Default: 2.0\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for Geo-RS-Seq-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"sequence\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def decoupled_geo_rs_token_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled mode with Geometric Mean RS and Token-level Truncated IS (ratio-based).\n\n        Combines the Geometric Mean Filter (ratio-based validity check) with\n        Token-level IS weights. Uses E[log(r)] (ideal = 0.0).\n\n        Args:\n            is_threshold (float): Upper threshold for token IS weights. Default: 2.0\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for Geo-RS-Token-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"token\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def bypass_pg_geo_rs_seq_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with REINFORCE loss, Geo-RS, and Sequence-level IS.\n\n        Combines geometric mean IS ratio rejection with sequence-level IS\n        in bypass mode with REINFORCE loss (no PPO clipping).\n        Uses E[log(r)] (ideal = 0.0) with [lower, upper] threshold bounds.\n\n        Args:\n            is_threshold (float): Upper threshold for sequence IS weights. Default: 2.0\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with REINFORCE + Geo-RS + Seq-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"sequence\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n            bypass_mode=True,\n            loss_type=\"reinforce\",\n        )\n\n    @classmethod\n    def bypass_pg_geo_rs_token_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: Optional[str | float] = \"0.999_1.001\",\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Bypass mode with REINFORCE loss, Geo-RS, and Token-level IS.\n\n        Combines geometric mean IS ratio rejection with token-level IS weights\n        in bypass mode with REINFORCE loss (no PPO clipping).\n        Uses E[log(r)] (ideal = 0.0) with [lower, upper] threshold bounds.\n\n        Token-level IS has lower variance but introduces bias.\n\n        Args:\n            is_threshold (float): Upper threshold for token IS weights. Default: 2.0\n            rs_threshold (Optional[Union[str, float]]): Geometric RS threshold (upper). Default: 0.999_1.001 (±0.1%)\n\n        Returns:\n            RolloutCorrectionConfig configured for bypass mode with REINFORCE + Geo-RS + Token-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"token\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k1\",\n            rollout_rs_threshold=rs_threshold,\n            bypass_mode=True,\n            loss_type=\"reinforce\",\n        )\n\n    @classmethod\n    def decoupled_k3_rs(\n        cls,\n        rs_threshold: float = 0.01,\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled mode with K3 KL Estimator Rejection Sampling.\n\n        Uses K3 KL estimator at sequence level for rejection sampling.\n        K3 = E[r - log(r) - 1] where r = π_train/π_rollout.\n        More stable than geometric mean for small KL values.\n\n        K3 >= 0 always (equals 0 when policies match exactly).\n\n        Args:\n            rs_threshold (float): Max allowed K3 divergence. Default: 0.01\n                Typical range: 0.001-0.1\n\n        Returns:\n            RolloutCorrectionConfig configured for K3 RS\n        \"\"\"\n        return cls(\n            rollout_is=None,\n            rollout_rs=\"seq_mean_k3\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def decoupled_k3_rs_seq_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: float = 0.01,\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled mode with K3 RS and Sequence-level Truncated IS.\n\n        Combines K3 KL estimator rejection with sequence-level IS weights.\n        K3 provides more stable outlier detection than geometric mean.\n\n        Args:\n            is_threshold (float): Upper threshold for sequence IS weights. Default: 2.0\n            rs_threshold (float): Max allowed K3 divergence. Default: 0.01\n\n        Returns:\n            RolloutCorrectionConfig configured for K3-RS-Seq-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"sequence\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k3\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def decoupled_k3_rs_token_tis(\n        cls,\n        is_threshold: float = 2.0,\n        rs_threshold: float = 0.01,\n    ) -> \"RolloutCorrectionConfig\":\n        \"\"\"Decoupled mode with K3 RS and Token-level Truncated IS.\n\n        Combines K3 KL estimator rejection with token-level IS weights.\n        K3 provides more stable outlier detection than geometric mean.\n        Token-level IS has lower variance but introduces bias.\n\n        Args:\n            is_threshold (float): Upper threshold for token IS weights. Default: 2.0\n            rs_threshold (float): Max allowed K3 divergence. Default: 0.01\n\n        Returns:\n            RolloutCorrectionConfig configured for K3-RS-Token-TIS\n        \"\"\"\n        return cls(\n            rollout_is=\"token\",\n            rollout_is_threshold=is_threshold,\n            rollout_rs=\"seq_mean_k3\",\n            rollout_rs_threshold=rs_threshold,\n        )\n\n    @classmethod\n    def disabled(cls) -> \"RolloutCorrectionConfig\":\n        \"\"\"Disabled - Metrics Only Mode.\n\n        Computes and logs off-policy metrics without applying correction.\n\n        Returns:\n            RolloutCorrectionConfig with all correction disabled\n        \"\"\"\n        return cls(rollout_is=None, rollout_rs=None)\n\n\n@dataclass\nclass AlgoConfig(BaseConfig):\n    \"\"\"Configuration for the algorithm.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        gamma (float): Discount factor for future rewards.\n        lam (float): Trade-off between bias and variance in the GAE estimator.\n        adv_estimator (str): Advantage estimator type: \"gae\", \"grpo\", \"reinforce_plus_plus\", etc.\n        norm_adv_by_std_in_grpo (bool): Whether to normalize advantages by std (specific to GRPO).\n        use_kl_in_reward (bool): Whether to enable in-reward KL penalty.\n        kl_penalty (str): How to estimate KL divergence: \"kl\", \"abs\", \"mse\", \"low_var_kl\", or \"full\".\n        kl_ctrl (KLControlConfig): KL control configuration.\n        use_pf_ppo (bool): Whether to enable preference feedback PPO.\n        pf_ppo (dict[str, Any]): Preference feedback PPO settings.\n        filter_groups (Optional[FilterGroupsConfig]): Filter groups configuration, used in DAPO and Entropy\n        rollout_correction (Optional[RolloutCorrectionConfig]): Rollout Correction configuration.\n            Addresses off-policy issues from policy mismatch, model staleness, and general distribution shifts.\n\n            Set to None to disable entirely. Use factory methods for common presets:\n            - RolloutCorrectionConfig.decoupled_token_is() - Decoupled mode with token-level IS\n            - RolloutCorrectionConfig.decoupled_seq_is() - Decoupled mode with sequence-level IS\n            - RolloutCorrectionConfig.decoupled_seq_is_rs() - Decoupled mode with sequence IS + RS\n            - RolloutCorrectionConfig.decoupled_k1_rs() - Decoupled mode with K1-RS (divergence)\n            - RolloutCorrectionConfig.decoupled_geo_rs() - Decoupled mode with Geo-RS (ratio)\n            - RolloutCorrectionConfig.bypass_ppo_clip() - Bypass mode with PPO-clip\n            - RolloutCorrectionConfig.bypass_ppo_clip_k1_rs() - Bypass mode with PPO-clip + K1-RS\n            - RolloutCorrectionConfig.bypass_pg_is() - Bypass mode with REINFORCE + IS\n            - RolloutCorrectionConfig.bypass_pg_k1_rs() - Bypass mode with REINFORCE + K1-RS\n\n            For backward compatibility, you can still pass a dict, which will be converted to\n            RolloutCorrectionConfig automatically.\n    \"\"\"\n\n    gamma: float = 1.0\n    lam: float = 1.0\n    adv_estimator: str = \"gae\"\n    norm_adv_by_std_in_grpo: bool = True\n    use_kl_in_reward: bool = False\n    kl_penalty: str = \"kl\"\n    kl_ctrl: KLControlConfig = field(default_factory=KLControlConfig)\n    use_pf_ppo: bool = False\n    pf_ppo: dict[str, Any] = field(default_factory=dict)\n    filter_groups: Optional[FilterGroupsConfig] = None\n    # Rollout Correction: corrects off-policy issues (policy mismatch, model staleness, distribution shifts)\n    # Set to None to disable, use RolloutCorrectionConfig presets (e.g., .tis(), .mis()), or pass dict\n    rollout_correction: Optional[RolloutCorrectionConfig] = None\n    # GDPO (Group reward-Decoupled Normalization Policy Optimization) settings.\n    # gdpo_reward_keys: keys in non_tensor_batch (from compute_score's return dict) that\n    #   correspond to individual reward dimensions, e.g. [\"format_reward\", \"accuracy_reward\"].\n    # gdpo_reward_weights: per-dimension weights for aggregation (default: equal weights).\n    gdpo_reward_keys: Optional[list[str]] = None\n    gdpo_reward_weights: Optional[list[float]] = None\n"
  },
  {
    "path": "verl/trainer/config/config.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nfrom verl.base_config import BaseConfig\n\n__all__ = [\"CheckpointConfig\", \"ProfileConfig\", \"BaseModelConfig\"]\n\n\n@dataclass\nclass CheckpointConfig(BaseConfig):\n    \"\"\"Configuration for model checkpointing.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        save_contents (list[str]): What to include in saved checkpoints.\n            Options: 'model', 'optimizer', 'extra', 'hf_model'.\n        load_contents (list[str]): Contents to load from checkpoint. Defaults to same as save_contents.\n        async_save (bool): Whether to save checkpoints asynchronously. Only implemented for Megatron as of now.\n    \"\"\"\n\n    save_contents: list[str] = field(default_factory=lambda: [\"model\", \"optimizer\", \"extra\"])\n    load_contents: list[str] = field(default_factory=lambda: [\"model\", \"optimizer\", \"extra\"])\n    async_save: bool = False\n    mbridge_config: dict[str, Any] = field(default_factory=dict)\n\n\n@dataclass\nclass ProfileConfig(BaseConfig):\n    \"\"\"Configuration for profiling.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        profile_ranks (Optional[list[int]]): List of ranks to profile. None means all ranks.\n        step_start (int): Starting step for profiling.\n        step_end (int): Ending step for profiling.\n        save_path (Optional[str]): Path to save profiling results.\n    \"\"\"\n\n    profile_ranks: Optional[list[int]] = None\n    step_start: int = -1\n    step_end: int = -1\n    save_path: Optional[str] = None\n\n\n@dataclass\nclass BaseModelConfig(BaseConfig):\n    \"\"\"Base configuration for a model.\n    Contains core settings for loading and initializing a pretrained model checkpoint.\n\n    Args:\n        path (str): Path to pretrained model weights.\n        tokenizer_path (Optional[str]): Tokenizer path (defaults to actor's model path if not set).\n        override_config (dict): Hugging Face config override.\n        external_lib (Optional[str]): External model implementation (optional).\n        trust_remote_code (bool): Whether to trust remote code from Hugging Face models.\n        lora (dict[str, Any]): LoRA configuration dictionary.\n    \"\"\"\n\n    path: str = \"~/models/deepseek-llm-7b-chat\"\n    tokenizer_path: Optional[str] = None\n    override_config: dict[str, Any] = field(default_factory=dict)\n    external_lib: Optional[str] = None\n    trust_remote_code: bool = False\n    lora: dict[str, Any] = field(default_factory=dict)\n\n\n@dataclass\nclass ModuleConfig(BaseConfig):\n    \"\"\"Configuration for external Python module, which can be loaded, executed (and optionally, ``import``ed).\n\n    Args:\n        path (str, optional): Path to the module file to load and execute.\n        name (str, optional): Name of the module to ``import``. Format: ``\"import.path.to.module\"``.\n            If ``None``, the module will be loaded with a hased name and\n                will not be added to ``sys.modules``, thus can not be ``import``ed as ``name``.\n    \"\"\"\n\n    path: Optional[str] = None\n    name: Optional[str] = None\n"
  },
  {
    "path": "verl/trainer/config/critic/critic.yaml",
    "content": "# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.CriticConfig\n\n# Number of rollouts per update (mirrors actor rollout_n)\nrollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n\n# fsdp or fsdp2 strategy used for critic model training\nstrategy: ???\n\n# whether to enable the critic worker.\n# by default it is only enabled if advantage estimator is gae\n# set it to True manually if you always want to enable critic worker\nenable: null\n\n# optimizer configs\noptim:\n\n  # Learning rate\n  lr: 1e-5\n\n  # Warmup steps ratio; total steps will be injected at runtime\n  lr_warmup_steps_ratio: 0.0\n\n  # Total training steps (must be overridden at runtime)\n  total_training_steps: -1\n\n  # Weight decay\n  weight_decay: 0.01\n\n  # Prioritized. None, 0 or Negative values mean delegating to lr_warmup_steps_ratio.\n  lr_warmup_steps: -1\n\n\n# model config for the critic\nmodel:\n\n  # Path to pretrained model weights\n  path: ~/models/deepseek-llm-7b-chat\n\n  # Tokenizer path (defaults to actor's model path)\n  tokenizer_path: ${oc.select:actor_rollout_ref.model.path,\"~/models/deepseek-llm-7b-chat\"}\n\n  # Hugging Face config override\n  override_config: {}\n\n  # External model implementation (optional)\n  external_lib: ${oc.select:actor_rollout_ref.model.external_lib,null}\n\n  # Whether to trust remote code from Hugging Face models\n  trust_remote_code: ${oc.select:actor_rollout_ref.model.trust_remote_code,false}\n\n# PPO mini-batch size per update\nppo_mini_batch_size: ${oc.select:actor_rollout_ref.actor.ppo_mini_batch_size,256}\n\n# [Deprecated] Global micro batch size\nppo_micro_batch_size: null\n\n# Local per-GPU micro batch size\nppo_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size,null}\n\n# Whether to automatically adjust batch size at runtime\nuse_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n\n# Max tokens per GPU in one PPO batch (doubled for critic)\nppo_max_token_len_per_gpu: 32768\n\n# Max token length per GPU in forward pass\nforward_max_token_len_per_gpu: ${.ppo_max_token_len_per_gpu}\n\n# Number of PPO epochs per batch\nppo_epochs: ${oc.select:actor_rollout_ref.actor.ppo_epochs,1}\n\n# Shuffle training data across PPO epochs\nshuffle: ${oc.select:actor_rollout_ref.actor.shuffle,false}\n\n# The seed used to construct mini-batch\ndata_loader_seed: 42\n\n# PPO value function clipping range\ncliprange_value: 0.5\n\n# Loss aggregation mode: \"token-mean\", \"seq-mean-token-sum\", or \"seq-mean-token-mean\"\nloss_agg_mode: ${oc.select:actor_rollout_ref.actor.loss_agg_mode,token-mean}\n\n# checkpoint configs\ncheckpoint:\n\n  # Target dataclass for this configuration\n  _target_: verl.trainer.config.CheckpointConfig\n\n  # What to include in saved checkpoints\n  # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n  save_contents: ['model', 'optimizer', 'extra']\n\n  # What to include when loading checkpoints\n  load_contents: ${.save_contents}\n\n  # Whether to save checkpoints asynchronously. Only effective for Megatron as of now.\n  async_save: False\n\n  # Mbridge config extension.\n  mbridge_config: {}\n\n# profile the critic model in `update_critic`\nprofiler:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.utils.profiler.ProfilerConfig\n\n  # profiler tool, default same as profiler.tool in global config\n  # choices: nsys, npu, torch, torch_memory\n  tool: ${oc.select:global_profiler.tool,null}\n\n  # whether enable profile on Critic\n  enable: False\n\n  # Whether to profile all ranks.\n  all_ranks: False\n\n  # The ranks that will be profiled. [] or [0,1,...]\n  ranks: []\n\n  # profile results saving path\n  save_path: ${oc.select:global_profiler.save_path,null}\n\n  # specific tool config which only related to the role\n  tool_config:\n\n    # nsys tool config\n    nsys:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NsightToolConfig\n    \n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n    \n    # npu config\n    npu:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NPUToolConfig\n\n      # Contents to profile, can be empty\n      # options: npu, cpu, memory, shapes, module, stack\n      contents: []\n\n      # Collection level, optional values: level_none, level0, level1, level2.\n      level: \"level0\"\n\n      # Whether to automatically parse the data.\n      analysis: True\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: False\n    \n    # torch profiler config\n    torch:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n      # Contents to profile, can be empty\n      # options: cuda, cpu, memory, shapes, stack\n      contents: []\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: false\n\n    # torch memory profiler config\n    torch_memory:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n\n      # Maximum number of memory allocation entries to track\n      trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n\n      # Stack trace depth for memory allocations\n      stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n      "
  },
  {
    "path": "verl/trainer/config/critic/dp_critic.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# defaults specify the default config from each component\ndefaults:\n\n  # fsdp optimizer config\n  - ../optim@optim: fsdp\n\n  # fsdp engine config\n  - ../engine@model.fsdp_config: fsdp\n\n  # dp actor config, inheriting from trainer/config/critic/critic.yaml\n  - critic\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.FSDPCriticConfig\n\n# distribution strategy. Options: fsdp (deprecating), fsdp2\nstrategy: fsdp\n\n# model config for the critic\nmodel:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.FSDPCriticModelCfg\n\n  # Whether to use shared memory for loading the model\n  use_shm: False\n\n  # Enable gradient checkpointing to save memory\n  enable_gradient_checkpointing: True\n\n  # Offload activations to CPU to reduce GPU memory usage\n  enable_activation_offload: False\n\n  # Use remove padding optimization (saves compute)\n  use_remove_padding: False\n\n  # Set to positive value to enable LoRA (e.g., 32)\n  lora_rank: 0\n\n  # LoRA scaling factor\n  lora_alpha: 16\n\n  # LoRA target modules: \"all-linear\" or list of linear projection layers\n  target_modules: all-linear\n\n  # TiledMLP configuration for memory-efficient MLP computation.\n  tiled_mlp:\n\n    # whether to enable TiledMLP\n    enabled: False\n\n    # number of shards to split the input\n    num_shards: 4\n\n# Forward-only batch size during inference (global)\nforward_micro_batch_size: ${oc.select:.ppo_micro_batch_size,null}\n\n# Forward-only batch size during inference (per GPU)\nforward_micro_batch_size_per_gpu: ${oc.select:.ppo_micro_batch_size_per_gpu,null}\n\n# Sequence parallelism size for Ulysses-style model parallelism\n# [DEPRECATED] use fsdp_config.ulysses_sequence_parallel_size instead\nulysses_sequence_parallel_size: 1\n\n# Gradient clipping for critic updates\ngrad_clip: 1.0\n"
  },
  {
    "path": "verl/trainer/config/critic/megatron_critic.yaml",
    "content": "# defaults specify the default config from each component\ndefaults:\n\n  # megatron optimizer config\n  - ../optim@optim: megatron\n\n  # megatron engine config\n  - ../engine@megatron: megatron\n\n  # dp actor config, inheriting from trainer/config/critic/critic.yaml\n  - critic\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.McoreCriticConfig\n\nstrategy: megatron\n\n# seconds, default is 10 minutes for torch, you can set it to a larger value if you have long-running operations like 32B or 72B model using megatron\nnccl_timeout: 600\n\n# model config for the critic\nmodel:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.BaseModelConfig\n\n  # override default empty mapping\n  override_config:\n\n    model_config: {}\n\n    moe_config:\n\n      freeze_moe_router: False\n\n  # LoRA (Low-Rank Adaptation) configuration for parameter-efficient fine-tuning\n  lora:\n      # LoRA type: \"lora\", \"vlm_lora\", \"canonical_lora\", or \"dora\"\n      type: lora\n\n      # LoRA rank (Dimension of the low-rank projection space.). Set to 0 to disable LoRA\n      rank: 0  # typical values: 8, 16, 32, 64\n      \n      #  Weighting factor for the low-rank projection. Defaults to 32\n      alpha: 32\n      \n      # Dropout rate for the low-rank projection. Defaults to 0.0\n      dropout: 0.0\n      \n      # A list of module names to apply LoRA to.\n      # For fused LoRA, Defaults to all linear layers ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2'].\n      # For canonical LoRA: [\"linear_q\", \"linear_k\", \"linear_v\", \"linear_proj\", \"linear_fc1_up\", \"linear_fc1_gate\", \"linear_fc2\"]\n      # - 'linear_qkv': Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention\n      # - 'linear_proj': Apply LoRA to the linear layer used for projecting the output of self-attention\n      # - 'linear_fc1': Apply LoRA to the first fully-connected layer in MLP\n      # - 'linear_fc2': Apply LoRA to the second fully-connected layer in MLP\n      # Target modules can also contain wildcards. For example, you can specify\n      # target_modules=['*.layers.0.*.linear_qkv', '*.layers.1.*.linear_qkv'] to add LoRA to only linear_qkv on the first two layers\n      # \n      # Note:\n      # For MLA (e.g., DeepSeek), you should use [\"linear_kv_down_proj\",\"linear_kv_up_proj\",\"linear_q_down_proj\",\"linear_q_up_proj\",\"linear_q_proj\"]\n      # Instead of \"linear_qkv\" or [\"linear_q\",\"linear_k\",\"linear_v\"]\n      # By default, MoE routers are excluded from LoRA adaptation, and you will need to specify \"router\" in target_modules to include them.\n      target_modules:\n        - linear_qkv\n        - linear_proj\n        - linear_fc1\n        - linear_fc2\n      \n      # A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name\n      # does not match any string in exclude_modules. If used, will require target_modules to be empty list or null\n      exclude_modules: []\n\n      # Position for applying dropout, can be 'pre' (before the low-rank projection) or 'post' (after). Defaults to 'pre'\n      dropout_position: pre\n\n      # Initialization method for the low-rank matrix A. Defaults to \"xavier\".\n      lora_A_init_method: xavier\n\n      # Initialization method for the low-rank matrix B. Defaults to \"zero\".\n      lora_B_init_method: zero\n\n      # Enables the experimental All-to-All (A2A) communication strategy. Defaults to False\n      a2a_experimental: False\n\n      # Parameter data type for LoRA weights. Default to null, which will use model's dtype.\n      dtype: null\n\n      # Path to pre-trained LoRA adapter weights (null to train from scratch)\n      adapter_path: null\n\n      # VLMLoRA additionally allows the user to specify whether the language or vision models should be frozen.\n      # For example, a common finetuning workload for multimodal models is to apply adapters to language model and fully\n      # finetune the vision model.\n      freeze_vision_model: True\n      freeze_vision_projection: True\n      freeze_language_model: True\n\n# Whether to load initial weights\nload_weight: True\n\n# seed for data loader\ndata_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n"
  },
  {
    "path": "verl/trainer/config/critic/torchtitan_critic.yaml",
    "content": "# defaults specify the default config from each component\ndefaults:\n\n  # torchtitan optimizer config\n  - ../optim@optim: torchtitan\n\n  # torchtitan engine config\n  - ../engine@torchtitan: torchtitan\n\n  # critic config, inheriting from trainer/config/critic/critic.yaml\n  - critic\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.TorchTitanCriticConfig\n\nstrategy: torchtitan\n\n# model config for the critic\nmodel:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.BaseModelConfig\n\n# seed for data loader\ndata_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n"
  },
  {
    "path": "verl/trainer/config/critic/veomni_critic.yaml",
    "content": "# defaults specify the default config from each component\ndefaults:\n\n  # veomni optimizer config\n  - ../optim@optim: veomni\n\n  # veomni engine config\n  - ../engine@veomni: veomni\n\n  # critic config, inheriting from trainer/config/critic/critic.yaml\n  - critic\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.VeOmniCriticConfig\n\nstrategy: veomni\n\n# model config for the critic\nmodel:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.BaseModelConfig\n\n# seed for data loader\ndata_loader_seed: ${oc.select:actor_rollout_ref.actor.data_loader_seed,null}\n\n"
  },
  {
    "path": "verl/trainer/config/data/legacy_data.yaml",
    "content": "# Tokenizer class or path. If null, it will be inferred from the model.\ntokenizer: null\n\n# Whether to use shared memory for data loading.\nuse_shm: False\n\n# Training set parquet. Can be a list or a single file.\n# The program will read all files into memory, so it can't be too large (< 100GB).\n# The path can be either a local path or an HDFS path.\n# For HDFS path, we provide utils to download it to DRAM and convert it to a local path.\ntrain_files: ~/data/rlhf/gsm8k/train.parquet\n\n# Validation parquet. Can be a list or a single file.\nval_files: ~/data/rlhf/gsm8k/test.parquet\n\n# Maximum sample length to be used.\n# Set to -1 to use full dataset, otherwise, randomly\n# select the specified number of samples from train dataset\ntrain_max_samples: -1\n\n# Maximum sample length to be used.\n# Set to -1 to use full dataset, otherwise, randomly\n# select the specified number of samples from val dataset\nval_max_samples: -1\n\n# The field in the dataset where the prompt is located. Default is 'prompt'.\nprompt_key: prompt\n\n# The field used to select the reward function (if using different ones per example).\nreward_fn_key: data_source\n\n# Maximum prompt length. All prompts will be left-padded to this length.\n# An error will be reported if the length is too long.\n# oc.select: default val for rollout.prompt_length\nmax_prompt_length: 512\n\n# Maximum response length. Rollout in RL algorithms (e.g. PPO) generates up to this length.\n# oc.select: default val for rollout.response_length\nmax_response_length: 512\n\n# Batch size sampled for one training iteration of different RL algorithms.\ntrain_batch_size: 1024\n\n# Batch size used during validation. Can be null.\nval_batch_size: null\n\n# use tool config to calculate true prompt length\ntool_config_path: ${oc.select:actor_rollout_ref.rollout.multi_turn.tool_config_path, null}\n\n# Whether to return the original input_ids without adding chat template.\n# This is used when the reward model's chat template differs from the policy.\n# If using a model-based RM with different templates, this should be True.\nreturn_raw_input_ids: False\n\n# Whether to return the original chat (prompt) without applying chat template.\nreturn_raw_chat: True\n\n# Whether to return the full prompt with chat template.\nreturn_full_prompt: False\n\n# Whether to shuffle the data in the dataloader.\nshuffle: True\n\n# Seed to use when shuffling the data\nseed: null\n\n# num dataloader workers\ndataloader_num_workers: 8\n\n# image patch size\nimage_patch_size: 14\n\n# Whether to shuffle the validation set.\nvalidation_shuffle: False\n\n# Whether to filter overlong prompts.\nfilter_overlong_prompts: False\n\n# Number of workers for filtering overlong prompts.\n# For large-scale datasets, filtering can be time-consuming.\n# Use multiprocessing to speed up. Default is 1.\nfilter_overlong_prompts_workers: 1\n\n# Truncate the input_ids or prompt if they exceed max_prompt_length.\n# Options: 'error', 'left', 'right', 'middle'. Default is 'error'.\ntruncation: error\n\n# The field in the multi-modal dataset where the image is located. Default is 'images'.\nimage_key: images\n\n# The field in the multi-modal dataset where the video is located.\nvideo_key: videos\n\n# If the remote tokenizer has a Python file, this flag determines whether to allow using it.\ntrust_remote_code: False\n\n# Optional: specify a custom dataset class path and name if overriding default loading behavior.\ncustom_cls:\n\n  # The path to the file containing your customized dataset class. If not specified, pre-implemented dataset will be used.\n  path: null\n\n  # The name of the dataset class within the specified file.\n  name: null\n\n# Whether to return multi-modal inputs in the dataset. Set to False if rollout generates new multi-modal inputs.\nreturn_multi_modal_inputs: True\n\n# settings related to data sampler\nsampler:\n\n  # the path to the module containing a curriculum class which implements the\n  # AbstractSampler interface\n  class_path: null\n\n  # the name of the curriculum class like `MySampler`\n  class_name: null\n\n# Data generation configuration for augmenting the dataset.\ndatagen:\n\n  # The path to the file containing your customized data generation class.\n  # E.g. 'pkg://verl.experimental.dynamic_dataset.dynamicgen_dataset'\n  path: null\n\n  # The class name of the data generation class within the specified file.\n  # E.g. 'MockDataGenerator'\n  name: null\n\n# Additional kwargs when calling tokenizer.apply_chat_template\napply_chat_template_kwargs: {}\n"
  },
  {
    "path": "verl/trainer/config/engine/automodel.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.AutomodelEngineConfig\n\n# Backend strategy identifier\nstrategy: automodel\n\n# Distributed training strategy: \"fsdp2\", \"megatron_fsdp\", or \"ddp\"\ndistributed_strategy: fsdp2\n\n# Parallelism sizes\ntp_size: 1\npp_size: 1\ncp_size: 1\nep_size: 1\ndp_replicate_size: 1\nsequence_parallel: false\ndefer_fsdp_grad_sync: true\n\n# Whether to offload model parameters to CPU\nparam_offload: false\n\n# Whether to offload optimizer state to CPU\noptimizer_offload: false\n\n# Whether to enable activation checkpointing\nactivation_checkpointing: false\n\n# Whether to enable FP8 training\nenable_fp8: false\n\n# Whether to enable torch.compile for the model\nenable_compile: false\n\n# Model data type for loading weights (\"fp32\", \"bf16\", \"fp16\")\nmodel_dtype: fp32\n\n# Attention implementation (\"sdpa\", \"flash_attention_2\", \"eager\", \"te\")\nattn_implementation: flash_attention_2\n\n# Backend settings\nbackend_config:\n  attn: sdpa                      # \"te\", \"sdpa\"\n  linear: te                      # \"torch\", \"te\"\n  rms_norm: torch_fp32            # \"torch\", \"torch_fp32\", \"te\"\n  rope_fusion: true\n  dispatcher: torch                # \"torch\", \"deepep\"\n  experts: gmm                    # \"gmm\", \"torch_mm\", \"torch\", \"te\"\n  gate_precision: null\n  enable_hf_state_dict_adapter: true\n  enable_fsdp_optimizations: false\n  fake_balanced_gate: false\n  fake_gate_noise: 0.0\n\n# MoE settings (MoEParallelizerConfig)\nmoe_config:\n  ignore_router_for_ac: false\n  reshard_after_forward: false\n  lm_head_precision: null\n  wrap_outer_model: true\n\n# Mixed precision policy (FSDP2 MixedPrecisionPolicy)\nmp_param_dtype: bf16\nmp_reduce_dtype: fp32\nmp_output_dtype: bf16\n\n# Random seed for reproducibility\nseed: 42\n\n# Whether to enable full determinism for distributed training, only for debugging\nfull_determinism: false\n\n# Whether to use forward only mode\nforward_only: false\n\n# Whether to use torch compile for entropy computation\nuse_torch_compile: false\n\n# Whether to use chunked entropy computation\nentropy_from_logits_with_chunking: false\n\n# Whether to use checkpointing for entropy computation\nentropy_checkpointing: false\n"
  },
  {
    "path": "verl/trainer/config/engine/fsdp.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.FSDPEngineConfig\n\n# policy for wrapping the model\nwrap_policy:\n\n  # Minimum number of parameters to trigger wrapping a layer with FSDP\n  min_num_params: 0\n\n# Whether to offload model parameters to CPU (trades speed for memory)\n# Note that this differs from the offload_policy in FSDP\nparam_offload: false\n\n# Whether to offload optimizer state to CPU\n# Note that this differs from the offload_policy in FSDP\noptimizer_offload: false\n\n# Only for FSDP2: offload param/grad/optimizer during train\noffload_policy: false\n\n# Reshard after forward pass to reduce memory footprint\n# For FSDP1, `false` enables `ShardingStrategy.SHARD_GRAD_OP`\nreshard_after_forward: true\n\n# Number of GPUs in each FSDP shard group; -1 means auto\nfsdp_size: -1\n\n# Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n# before the current forward computation.\nforward_prefetch: False\n\n# model dtype of fsdp\nmodel_dtype: fp32\n\n# Whether to use original parameters in fsdp. Only avaiable in fsdp1\nuse_orig_params: false\n\n# Random seed for reproducibility.\nseed: 42\n\n# Whether to enable full determinism for distributed training, only for debugging.\nfull_determinism: false\n\n# ulysses sequence parallel size\nulysses_sequence_parallel_size: 1\n\n# Whether to use entropy_from_logits_with_chunking in fsdp.\nentropy_from_logits_with_chunking: false\n\n# Whether to use torch compile in fsdp.\nuse_torch_compile: true\n\n# Whether to use entropy checkpointing in fsdp.\nentropy_checkpointing: false\n\n# Whether to use forward only in fsdp.\nforward_only: false\n\n# fsdp or fsdp2\nstrategy: fsdp\n\n# Mixed precision training param dtype\ndtype: bfloat16 # [\"bfloat16\", \"float16\"]\n\n# QAT (Quantization-Aware Training) configuration\nqat:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.QATEngineConfig\n\n  # Whether to enable QAT\n  enable: false\n\n  # Quantization mode: \"w4a16\" (weight-only). \"w4a4\" is experimental and not recommended.\n  mode: \"w4a16\"\n\n  # Quantization group size (NVFP4 requires 16)\n  group_size: 16\n\n  # Patterns to ignore (e.g., lm_head, embed_tokens)\n  ignore_patterns:\n\n    - \"lm_head\"\n    - \"embed_tokens\"\n    - \"re:.*mlp.gate$\"\n\n  # Activation observer for W4A4 mode\n  activation_observer: \"static_minmax\"\n\n  # Path to vLLM quantization config JSON file\n  quantization_config_path: null\n"
  },
  {
    "path": "verl/trainer/config/engine/megatron.yaml",
    "content": "# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.workers.config.McoreEngineConfig\n\n# Whether to offload model parameters to CPU\nparam_offload: False\n\n# Whether to offload gradients to CPU\ngrad_offload: False\n\n# Whether to offload optimizer state to CPU\noptimizer_offload: False\n\n# tensor model parallel size\ntensor_model_parallel_size: 1\n\n# expert model parallel size\nexpert_model_parallel_size: 1\n\n# expert tensor parallel size (null to be same as TP)\nexpert_tensor_parallel_size: null\n\n# pipeline model parallel size\npipeline_model_parallel_size: 1\n\n# virtual pipeline model parallel size\nvirtual_pipeline_model_parallel_size: null\n\n# context parallel size\ncontext_parallel_size: 1\n\n# sequence parallel\nsequence_parallel: True\n\n# Whether to use distributed optimizer\nuse_distributed_optimizer: True\n\n# Whether to use distributed checkpointing\nuse_dist_checkpointing: False\n\n# distributed checkpointing path\ndist_checkpointing_path: null\n\n# distributed checkpointing prefix, e.g. Nemo2 will append prefix 'module.' to the state dict keys\ndist_checkpointing_prefix: ''\n\n# Make optimizer distributed checkpoint fully reshardable (TP/PP/EP/DP) as opposed to plain DP reshardability\ndist_ckpt_optim_fully_reshardable: True\n\n# Use as little memory as possible during save and load by using Gloo.\n# Has effect only when `dist_ckpt_optim_fully_reshardable` is enabled\ndistrib_optim_fully_reshardable_mem_efficient: False\n\n# oc.select: default val for ref.megatron.seed\nseed: 42\n\n# Allow to override Distributed Data Parallel (DDP) config\noverride_ddp_config: {}\n\n# additional transformer config like: num_layers_in_first(/last)_pipeline_stage\n# oc.select: default val for ref.megatron.override_transformer_config\noverride_transformer_config:\n  # Recompute configuration, same as in megatron.training.arguments\n  # default use minimal performance-interference recompute methods\n  # Recompute granualarity, choices: [\"full\", \"selective\"]\n  recompute_granularity: null\n\n  # Recompute modules, multiple choices: [\"core_attn\", \"moe_act\", \"layernorm\", \"mla_up_proj\", \"mlp\", \"moe\"]\n  # Please use correct module in matched model\n  recompute_modules: [\"core_attn\"]\n\n  # 'uniform', 'block'\n  # 'uniform' divides the total number of transformer layers and checkpoints the input activation of each chunk\n  # 'block' checkpoints the specified number of layers per pipeline stage at the specified granularity\n  recompute_method: null\n\n  # 'full' will checkpoint the entire transformer layer and 'selective' only checkpoints memory intensive part of attention\n  recompute_num_layers: null\n\n  # Attention backend to use (flash,fused,unfused,local,auto). Defaults to auto in mcore, flash in verl\n  attention_backend: flash\n\noverride_mcore_model_config: {}\n\n# oc.select: default val for ref.megatron.use_mbridge\nuse_mbridge: True\n\n# oc.select: default val for ref.megatron.vanilla_mbridge\nvanilla_mbridge: True\n\n# whether to use thd format (sequence packing), if not, use bshd format, padding the input_ids to the longest sequence length\nuse_remove_padding: True\n\n# whether to use forward only\nforward_only: False\n\n# Mixed precision training param dtype\ndtype: bfloat16 # [\"bfloat16\", \"float16\"]\n\n# Router replay configuration for MoE models\nrouter_replay:\n\n  # Target dataclass for this configuration\n  _target_: verl.workers.config.EngineRouterReplayConfig\n\n  # Router replay mode: disabled, R2, R3\n  # - R2: Use R2 routing strategy (record mode)\n  # - R3: Use R3 routing strategy (record mode)\n  mode: disabled\n\n  # File path to save recorded routing decisions\n  # Required when mode is 'record', 'R2', or 'R3'\n  record_file: null\n\n  # File path to load recorded routing decisions for replay\n  # Required when mode is 'replay'\n  replay_file: null\n"
  },
  {
    "path": "verl/trainer/config/engine/torchtitan.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.TorchtitanEngineConfig\n\n# Whether to offload model parameters to CPU\nparam_offload: False\n\n# Whether to offload optimizer state to CPU\noptimizer_offload: False\n\n# policy for wrapping the model\nwrap_policy:\n  # Minimum number of parameters to trigger wrapping a layer with FSDP\n  min_num_params: 0\n\n# The policy for applying `reshard_after_forward` within an FSDP setup\n# Options: \"default\", \"always\", \"never\"\nreshard_after_forward: default\n\n# Prefetch the next forward-pass all-gather before the current forward computation.\nforward_prefetch: false\n\n# Whether to use original parameters\nuse_orig_params: false\n\n# Mixed precision configuration for FSDP\nmixed_precision: false\n\n# Whether to use torch compile\nuse_torch_compile: true\n\n# Whether to use entropy_from_logits_with_chunking\nentropy_from_logits_with_chunking: false\n\n# Whether to use entropy checkpointing\nentropy_checkpointing: false\n\n# Data parallel size (FSDP group size)\ndata_parallel_size: 1\n\n# Data parallel replicate size\ndata_parallel_replicate_size: 1\n\n# Data parallel shard size\ndata_parallel_shard_size: 1\n\n# Tensor parallel size\ntensor_parallel_size: 1\n\n# Expert parallel size\nexpert_parallel_size: 1\n\n# Pipeline parallel size\npipeline_parallel_size: 1\n\n# Context parallel size\ncontext_parallel_size: 1\n\n# Attention type for torchtitan's model (e.g., \"sdpa\", \"flex\", \"varlen\")\nattn_type: flex\n\n# Maximum sequence length for RoPE cache. If null, defaults to torchtitan's TrainingConfig.seq_len (2048).\nmax_seq_len: null\n\n# Strategy\nstrategy: torchtitan\n\n# Random seed for reproducibility\nseed: 42\n\n# Whether to enable full determinism for distributed training, only for debugging\nfull_determinism: false\n\n# Whether to use forward only\nforward_only: false\n\n# Mixed precision training param dtype\ndtype: bfloat16\n"
  },
  {
    "path": "verl/trainer/config/engine/veomni.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.VeOmniEngineConfig\n\n# Whether to offload model parameters to CPU\nparam_offload: False\n\n# Whether to offload optimizer state to CPU\noptimizer_offload: False\n\n# FSDP group size. -1 means use all available GPUs.\nfsdp_size: -1\n\nulysses_parallel_size: 1\n\nexpert_parallel_size: 1\n\nmixed_precision: true\n\n# Random seed for reproducibility.\nseed: 42\n\n# Whether to enable full determinism for distributed training, only for debugging.\nfull_determinism: false\n\ninit_device: meta\n\nenable_full_shard: true\n\nckpt_manager: dcp\n\n# Only for FSDP1: FSDP1 configuration, prefetch the next forward-pass all-gather\n# before the current forward computation.\nforward_prefetch: true\n\nstrategy: veomni\n\n# Whether to use torch compile in fsdp.\nuse_torch_compile: false\n\n# Whether to use forward only in fsdp.\nforward_only: false\n\nenable_fsdp_offload: false\n\nenable_reentrant: false\n\n# support eager, sdpa, flash_attention_2, flash_attention_3, veomni_flash_attention_2_with_sp,\n# veomni_flash_attention_3_with_sp and native-sparse\nattn_implementation: flash_attention_2\n\n# eager or fused\nmoe_implementation: fused\n\nforce_use_huggingface: false\n\nactivation_gpu_limit: 0.0\n"
  },
  {
    "path": "verl/trainer/config/evaluation.yaml",
    "content": "data:\n  path: /tmp/math_Qwen2-7B-Instruct.parquet\n  prompt_key: prompt\n  response_key: responses\n  data_source_key: data_source\n  reward_model_key: reward_model\n\ncustom_reward_function:\n  path: null\n  name: compute_score\n\nray_kwargs:\n  ray_init:\n    num_cpus: null # `None` means using all CPUs, which might cause hang if limited in systems like SLURM. Please set to a number allowed then.\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/legacy_reward_impl.yaml",
    "content": "custom_reward_function:\n  path: null\n  name: null\n\nreward_model:\n  num_workers: null\n  reward_manager: null\n  enable: null\n  enable_resource_pool: null\n  n_gpus_per_node: null\n  nnodes: null\n  reward_loop_source: null\n  reward_loop_module_path: null\n  reward_loop_class_name: null\n  model:\n    path: null\n    external_lib: null\n    trust_remote_code: null\n  rollout:\n    name: null\n    dtype: null\n    gpu_memory_utilization: null\n    enforce_eager: null\n    cudagraph_capture_sizes: null\n    free_cache_engine: null\n    data_parallel_size: null\n    expert_parallel_size: null\n    tensor_model_parallel_size: null\n    max_num_batched_tokens: null\n    max_model_len: null\n    max_num_seqs: null\n    load_format: null\n    engine_kwargs: null\n    limit_images: null\n    enable_chunked_prefill: null\n    enable_prefix_caching: null\n    disable_log_stats: null\n    skip_tokenizer_init: null\n\n    prompt_length: null\n    response_length: null\n\nsandbox_fusion:\n  url: null\n  max_concurrent: null\n  memory_limit_mb: null\n"
  },
  {
    "path": "verl/trainer/config/model/hf_model.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n_target_: verl.workers.config.HFModelConfig\n\n# path to the huggingface model\npath: ~/models/deepseek-llm-7b-chat\n\n# config to the huggingface config. In case it is not the same as path\nhf_config_path: null\n\n# path to the huggingface tokenizer. In case it is not the same as path\ntokenizer_path: null\n\n# whether to use shared memory for model loading\nuse_shm: False\n\n# whether to trust remote code.\ntrust_remote_code: False\n\n# custom chat template for the model\ncustom_chat_template: null\n\n# whether to use external libs for the model\nexternal_lib: null\n\n# override hf config\noverride_config: {}\n\n# whether to enable gradient checkpointing. Only valid when we use hf model definition\nenable_gradient_checkpointing: True\n\n# whether to enable activation offload. Only valid when we use hf model definition\nenable_activation_offload: False\n\n# whether to use remove padding. Only valid when we use hf model definition\nuse_remove_padding: True\n\n# Set to positive value to enable LoRA (e.g., 32)\nlora_rank: 0\n\n# LoRA scaling factor\nlora_alpha: 16\n\n# Target modules for LoRA adaptation\ntarget_modules: all-linear\n\n# Exclude modules from LoRA adaptation\nexclude_modules: null\n\n# Path to pre-trained LoRA adapter to load for continued training\nlora_adapter_path: null\n\n# whether to use liger. Only valid when we use hf model definition\nuse_liger: False\n\n# whether to use fused kernels.\nuse_fused_kernels: False\n\n# fused kernel options.\nfused_kernel_options:\n\n  # the implementation backend for fused kernels.\n  impl_backend: torch\n\n# TiledMLP configuration for memory-efficient MLP computation.\n# Reduces peak memory by processing MLP forward/backward in tiles.\ntiled_mlp:\n\n  # whether to enable TiledMLP\n  enabled: False\n\n  # number of shards to split the input. Higher values reduce peak memory but may slightly impact performance.\n  num_shards: 4\n\n# MTP\nmtp:\n\n  _target_: verl.workers.config.MtpConfig\n\n  enable: False\n  enable_train: False\n  enable_rollout: False\n\n  detach_encoder: False\n  mtp_loss_scaling_factor: 0.1\n\n  speculative_algorithm: EAGLE\n  speculative_num_steps: 3\n  speculative_eagle_topk: 1\n  speculative_num_draft_tokens: 4\n\n  method: mtp\n  num_speculative_tokens: 1\n"
  },
  {
    "path": "verl/trainer/config/model_engine/dp.yaml",
    "content": "# @package _global_\nmodel_engine: dp\n"
  },
  {
    "path": "verl/trainer/config/model_engine/torchtitan.yaml",
    "content": "# @package _global_\nmodel_engine: torchtitan\n"
  },
  {
    "path": "verl/trainer/config/model_engine/veomni.yaml",
    "content": "# @package _global_\nmodel_engine: veomni\n"
  },
  {
    "path": "verl/trainer/config/npu_profile/npu_profile.yaml",
    "content": "# Options for the npu profiler\noptions:\n\n  # Storage path of collected data.\n  save_path: ./profiler_data\n\n  # The roles that will be profiled. Only takes effect in discrete mode.\n  # optional values: all, rollout_generate, actor_compute_log_prob, actor_update and ref_compute_log_prob.\n  # \"all\" means all roles will be profiled.\n  roles: [\"all\"]\n\n  # Collection level, optional values: level_none, level0, level1, level2.\n  level: level0\n\n  # Whether to enable memory analysis.\n  with_memory: False\n\n  # Whether to record tensor shape.\n  record_shapes: False\n\n  # Whether to record Device-side performance data.\n  with_npu: True\n\n  # Whether to record Host-side performance data.\n  with_cpu: True\n\n  # Whether to record Python call stack information.\n  with_module: False\n\n  # Whether to record operator call stack information.\n  with_stack: False\n\n  # Whether to automatically parse the data.\n  analysis: True"
  },
  {
    "path": "verl/trainer/config/optim/automodel.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.AutomodelOptimizerConfig\n\noptimizer: AdamW\n\n# Module path to import optimizer from\noptimizer_impl: torch.optim\n\n# Learning rate (maps to max_lr in Automodel's OptimizerParamScheduler)\nlr: 1e-5\n\n# LR warmup steps ratio (used when lr_warmup_steps <= 0)\nlr_warmup_steps_ratio: 0.0\n\n# Total training steps (injected at runtime)\ntotal_training_steps: -1\n\n# Weight decay\nweight_decay: 0.01\n\n# LR warmup steps (set > 0 to override lr_warmup_steps_ratio)\nlr_warmup_steps: -1\n\n# Betas for Adam optimizer\nbetas: [0.9, 0.999]\n\n# Clip gradient norm\nclip_grad: 1.0\n\n# Initial LR ratio for warmup start (init_lr = lr * init_lr_ratio)\ninit_lr_ratio: 0.1\n\n# Minimum LR ratio after decay (min_lr = lr * min_lr_ratio)\nmin_lr_ratio: 0.01\n\n# LR scheduler type (Automodel OptimizerParamScheduler decay style)\n# Options: \"constant\", \"cosine\", \"linear\", \"inverse-square-root\"\nlr_scheduler_type: cosine\n\n# Weight decay increment style: \"constant\", \"linear\", or \"cosine\"\nwd_incr_style: constant\n\n# Kept for backward compatibility (unused by Automodel scheduler)\nnum_cycles: 0.5\nzero_indexed_step: true\n\n# Common optimizer kwargs\neps: 1e-8\nmaster_weights: false\nstore_param_remainders: false\nexp_avg_dtype: null       # \"fp32\", \"bf16\"\nexp_avg_sq_dtype: null    # \"fp32\", \"bf16\"\nmaster_weight_dtype: null # \"fp32\", \"bf16\"\n\n# Additional optimizer kwargs (passed directly to constructor)\noverride_optimizer_config: {}\n"
  },
  {
    "path": "verl/trainer/config/optim/fsdp.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.FSDPOptimizerConfig\n\n# Optimizer class name (e.g., \"AdamW\", \"AdamW8bit\", \"_AdamW\", \"Adam\")\noptimizer: AdamW\n\n# Module path to import optimizer\n# Examples: \"torch.optim\", \"torchao.optim\", \"bitsandbytes.optim\"\noptimizer_impl: torch.optim\n\n# Learning rate\nlr: 1e-3\n\n# LR warmup steps ratio\nlr_warmup_steps_ratio: 0.0\n\n# Total training steps\ntotal_training_steps: -1\n\n# Weight decay\nweight_decay: 0.01\n\n# LR warmup steps\nlr_warmup_steps: -1\n\n# Betas for Adam optimizer\nbetas: [0.9, 0.999]\n\n# Clip gradient\nclip_grad: 1.0\n\n# Minimum LR ratio for cosine schedule\nmin_lr_ratio: 0.0\n\n# Number of cosine cycles in LR schedule\nnum_cycles: 0.5\n\n# LR scheduler type: \"constant\" or \"cosine\"\nlr_scheduler_type: constant\n\n# Whether the LR schedule uses 0-indexed steps\nzero_indexed_step: true\n\n# deprecated\nwarmup_style: null\n\n# Additional optimizer-specific keyword arguments\n# Example for torchao with bf16 stochastic rounding:\n# optimizer_impl: torchao.optim\n# optimizer: _AdamW\n# override_optimizer_config:\n#   bf16_stochastic_round: true\noverride_optimizer_config: null\n"
  },
  {
    "path": "verl/trainer/config/optim/megatron.yaml",
    "content": "_target_: verl.workers.config.McoreOptimizerConfig\n\n# Learning rate\nlr: 1e-3\n\n# LR warmup steps ratio\nlr_warmup_steps_ratio: 0.0\n\n# Total training steps\ntotal_training_steps: -1\n\n# Weight decay\nweight_decay: 0.01\n\n# LR warmup steps\nlr_warmup_steps: -1\n\n# Betas for Adam optimizer\nbetas: [0.9, 0.999]\n\n# Clip gradient\nclip_grad: 1.0\n\n# optimizer type\noptimizer: adam\n\n# initial learning rate for warmup, default to 0.0\nlr_warmup_init: 0.0\n\nlr_decay_steps: null\n\n# select from constant/linear/cosine/inverse_square_root\nlr_decay_style: constant\n\n# minimum learning rate, default to 0.0\nmin_lr: 0.0\n\n# select from constant/linear/cosine\nweight_decay_incr_style: constant\n\n# select from constant/exponential/cosine\nlr_wsd_decay_style: exponential\n\nlr_wsd_decay_steps: null\n\n# use checkpoint optimizer parameter scheduler\nuse_checkpoint_opt_param_scheduler: False\n\noverride_optimizer_config: {}\n"
  },
  {
    "path": "verl/trainer/config/optim/torchtitan.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.TorchtitanOptimizerConfig\n\n# Optimizer name\nname: AdamW\n\n# Learning rate\nlr: 1e-3\n\n# LR warmup steps ratio\nlr_warmup_steps_ratio: 0.0\n\n# Total training steps\ntotal_training_steps: -1\n\n# Weight decay\nweight_decay: 0.01\n\n# LR warmup steps\nlr_warmup_steps: -1\n\n# Betas for Adam optimizer\nbetas: [0.9, 0.999]\n\n# Clip gradient\nclip_grad: 1.0\n\n# Epsilon for Adam optimizer\neps: 1e-8\n\n# Decay type: \"linear\", \"sqrt\", or \"cosine\"\ndecay_type: linear\n\n# Minimum LR factor for cosine schedule\nmin_lr_factor: 0.0\n"
  },
  {
    "path": "verl/trainer/config/optim/veomni.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.VeOmniOptimizerConfig\n\noptimizer: adamw\n\n# Learning rate\nlr: 1e-3\n\n# Minimum learning rate\nlr_min: 0.0\n\n# Starting learning rate for warmup\nlr_start: 0.0\n\n# LR warmup steps ratio\nlr_warmup_steps_ratio: 0.0\n\n# LR decay steps ratio\nlr_decay_ratio: 1.0\n\n# Total training steps\ntotal_training_steps: -1\n\n# Weight decay\nweight_decay: 0.01\n\n# LR warmup steps\nlr_warmup_steps: -1\n\n# Betas for Adam optimizer\nbetas: [0.9, 0.999]\n\n# Clip gradient\nclip_grad: 1.0\n\n# LR scheduler type: \"constant\" or \"cosine\"\nlr_scheduler_type: cosine\n\noverride_optimizer_config: {}\n"
  },
  {
    "path": "verl/trainer/config/ppo_megatron_trainer.yaml",
    "content": "# specify the default per-component configs\ndefaults:\n  # <folder_name>@<field_name>.<field_name>: <yaml_file_name>\n  # actor_rollout_ref.actor: trainer/config/actor/megatron_actor.yaml\n  - actor@actor_rollout_ref.actor: megatron_actor\n  # data: trainer/config/data/legacy_data.yaml\n  - data@data: legacy_data\n  # load the reference default config, then apply the fields in the current yaml\n  # Reference model config.\n  # Reference model will be enabled when actor.use_kl_loss or/and algorithm.use_kl_in_reward is/are True.\n  - ref@actor_rollout_ref.ref: megatron_ref\n  # Rollout model config.\n  - rollout@actor_rollout_ref.rollout: rollout\n  # Model config.\n  - model@actor_rollout_ref.model: hf_model\n  # Critic model config.\n  - critic@critic: megatron_critic\n  # legacy reward impl config, for backward compatibility\n  - legacy_reward_impl\n  # Reward model config.\n  - reward@reward: reward\n  # Rollout correction config.\n  - algorithm@algorithm.rollout_correction: rollout_correction\n  - _self_\n\nactor_rollout_ref:\n  hybrid_engine: True\n\n  nccl_timeout: 600 # seconds, default is 10 minutes for torch, you can set it to a larger value if you have long-running operations like 32B or 72B model using megatron\n\n  model:\n    override_config:\n      model_config: {}\n      moe_config:\n        freeze_moe_router: False\n\n    use_fused_kernels: False # Whether to use custom fused kernels (PostProcessing, for memory efficiency)\n\n    trust_remote_code: False\n\n    # Whether to remove padding tokens in inputs during training\n    use_remove_padding: false\n\n    # LoRA (Low-Rank Adaptation) configuration for parameter-efficient fine-tuning\n    lora:\n      # LoRA type: \"lora\", \"vlm_lora\", \"canonical_lora\", or \"dora\"\n      type: lora\n\n      # whether to sync weights / refit by either merging LoRA adapters into the base model weights before transferring to vLLM (for better inference speed but more refit time and potential precision loss). If this is False, it will load separate adapters.\n      merge: False\n\n      # LoRA rank (Dimension of the low-rank projection space.). Set to 0 to disable LoRA\n      rank: 0  # typical values: 8, 16, 32, 64\n      \n      #  Weighting factor for the low-rank projection. Defaults to 32\n      alpha: 32\n      \n      # Dropout rate for the low-rank projection. Defaults to 0.0\n      dropout: 0.0\n      \n      # A list of module names to apply LoRA to.\n      # For fused LoRA, Defaults to all linear layers ['linear_qkv', 'linear_proj', 'linear_fc1', 'linear_fc2'].\n      # For canonical LoRA: [\"linear_q\", \"linear_k\", \"linear_v\", \"linear_proj\", \"linear_fc1_up\", \"linear_fc1_gate\", \"linear_fc2\"]\n      # - 'linear_qkv': Apply LoRA to the fused linear layer used for query, key, and value projections in self-attention\n      # - 'linear_proj': Apply LoRA to the linear layer used for projecting the output of self-attention\n      # - 'linear_fc1': Apply LoRA to the first fully-connected layer in MLP\n      # - 'linear_fc2': Apply LoRA to the second fully-connected layer in MLP\n      # Target modules can also contain wildcards. For example, you can specify\n      # target_modules=['*.layers.0.*.linear_qkv', '*.layers.1.*.linear_qkv'] to add LoRA to only linear_qkv on the first two layers\n      # \n      # Note:\n      # For MLA (e.g., DeepSeek), you should use [\"linear_kv_down_proj\",\"linear_kv_up_proj\",\"linear_q_down_proj\",\"linear_q_up_proj\",\"linear_q_proj\"]\n      # Instead of \"linear_qkv\" or [\"linear_q\",\"linear_k\",\"linear_v\"]\n      # By default, MoE routers are excluded from LoRA adaptation, and you will need to specify \"router\" in target_modules to include them.\n      target_modules:\n        - linear_qkv\n        - linear_proj\n        - linear_fc1\n        - linear_fc2\n      \n      # A list of module names not to apply LoRa to. It will match all nn.Linear & nn.Linear-adjacent modules whose name\n      # does not match any string in exclude_modules. If used, will require target_modules to be empty list or None\n      exclude_modules: []\n\n      # Position for applying dropout, can be 'pre' (before the low-rank projection) or 'post' (after). Defaults to 'pre'\n      dropout_position: pre\n\n      # Initialization method for the low-rank matrix A. Defaults to \"xavier\".\n      lora_A_init_method: xavier\n\n      # Initialization method for the low-rank matrix B. Defaults to \"zero\".\n      lora_B_init_method: zero\n\n      # Enables the experimental All-to-All (A2A) communication strategy. Defaults to False\n      a2a_experimental: False\n\n      # Parameter data type for LoRA weights. Default to null, which will use model's dtype.\n      dtype: null\n\n      # Path to pre-trained LoRA adapter weights (null to train from scratch)\n      adapter_path: null\n\n      # VLMLoRA additionally allows the user to specify whether the language or vision models should be frozen.\n      # For example, a common finetuning workload for multimodal models is to apply adapters to language model and fully\n      # finetune the vision model.\n      freeze_vision_model: True\n      freeze_vision_projection: True\n      freeze_language_model: True\n\n  rollout:\n    quantization: null\n\n    layer_name_map:\n      qkv_layer_name: qkv\n      gate_proj_layer_name: gate_up\n\nalgorithm:\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.AlgoConfig\n  gamma: 1.0\n  lam: 1.0\n  adv_estimator: gae\n  norm_adv_by_std_in_grpo: True\n  use_kl_in_reward: False\n  kl_penalty: kl # how to estimate kl divergence\n  kl_ctrl:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.trainer.config.KLControlConfig\n    type: fixed\n    kl_coef: 0.001\n    horizon: 10000\n    target_kl: 0.1\n  use_pf_ppo: False\n  pf_ppo:\n    reweight_method: pow # [\"pow\", \"max_min\", \"max_random\"]\n    weight_pow: 2.0\n\ntrainer:\n  balance_batch: True\n  total_epochs: 30\n  total_training_steps: null\n  project_name: verl_examples\n  experiment_name: gsm8k\n  logger: [\"console\", \"wandb\"]\n  log_val_generations: 0\n  nnodes: 1\n  n_gpus_per_node: 8\n  save_freq: -1\n  esi_redundant_time: 0\n\n  # auto: find the last ckpt to resume. If can't find, start from scratch\n  resume_mode: auto # or disable or resume_path if resume_from_path is set\n  resume_from_path: null\n  del_local_ckpt_after_load: False\n  val_before_train: True\n  test_freq: -1\n  critic_warmup: 0\n  default_hdfs_dir: null\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  max_actor_ckpt_to_keep: null\n  max_critic_ckpt_to_keep: null\n  # The timeout for ray worker group to wait for the register center to be ready\n  ray_wait_register_center_timeout: 300\n  device: cuda\n  # Directory for logging rollout data; no dump if null\n  rollout_data_dir: null\n\n  # whether to use legacy worker implementation\n  #  mode: \"auto\", \"enable\", or \"disable\"\n  use_legacy_worker_impl: auto\n\nglobal_profiler:\n  _target_: verl.utils.profiler.ProfilerConfig\n  tool: null # choose between nsys, npu, torch, torch_memory\n  steps: null # profile steps\n  profile_continuous_steps: False\n  save_path: \"outputs/profile\" # profiler saving path\n  # Specific tool configs, can use +profiler.tool_config.[tool].xxx to config\n  global_tool_config:\n    # nsys config\n    nsys:\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: False\n\n      # controller Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n      ## reference https://docs.nvidia.com/nsight-systems/UserGuide/index.html\n      ## reference https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html\n      controller_nsight_options:\n        # Select the API(s) to be traced.\n        trace: \"cuda,nvtx,cublas,ucx\"\n\n        # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n        cuda-memory-usage: \"true\"\n\n        # CUDA graphs will be traced as a whole\n        cuda-graph-trace: \"graph\"\n\n      # worker Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n      worker_nsight_options:\n        # Select the API(s) to be traced.\n        trace: \"cuda,nvtx,cublas,ucx\"\n\n        # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n        cuda-memory-usage: \"true\"\n\n        # CUDA graphs will be traced as a whole\n        cuda-graph-trace: \"graph\"\n\n        # Profiling only in a range of torch.cuda.profiler.start and stop. Do not change this config.\n        capture-range: \"cudaProfilerApi\"\n\n        # Specify the desired behavior when a capture range ends.\n        # In verl we need the torch.cuda.profiler.start/stop pair to repeats n times.\n        # valid values are \"repeat-shutdown:n\" or null.\n        # For normal whole step profiling, n = len(profile_steps);\n        # but for discrete profiling, n = len(profile_steps) * Number(subtasks).\n        # Or you can just leave it null and the program will use n = len(profile_steps) * 6;\n        capture-range-end: null\n\n        # Send signal to the target application's process group. We let the program to exit by itself.\n        kill: none\n\n    # enable memory visualization for debugging memory usage\n    torch_memory:\n      #  Maximum number of allocation entries to record\n      trace_alloc_max_entries: 100_000\n      # The depth of the call stack to capture for each allocation\n      stack_depth: 32\n      # 'alloc': records only allocation events || 'state': records memory state changes || 'all': records both.\n      context: \"all\"\n      # 'python': records Python stacks || 'cpp': records C++ stacks (available in some versions) || 'all': records both.\n      stacks: \"all\"\n      # devices, record_context etc.\n      kw_args: {}\n\n# configs for TransferQueue\ntransfer_queue:\n  # Whether to enable transfer queue\n  enable: False\n\nray_kwargs:\n  ray_init:\n    num_cpus: null # `None` means using all CPUs, which might cause hang if limited in systems like SLURM. Please set to a number allowed then.\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/ppo_trainer.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# specify the default per-component configs\ndefaults:\n\n  - model_engine: dp\n\n  # <folder_name>@<field_name>.<field_name>: <yaml_file_name>\n  # actor_rollout_ref.actor: trainer/config/actor/dp_actor.yaml\n  - actor@actor_rollout_ref.actor: ${model_engine}_actor\n\n  # data: trainer/config/data/legacy_data.yaml\n  - data@data: legacy_data\n\n  # Reference model config.\n  # Reference model will be enabled when actor.use_kl_loss or/and algorithm.use_kl_in_reward is/are True.\n  - ref@actor_rollout_ref.ref: ${model_engine}_ref\n\n  # Rollout model config.\n  - rollout@actor_rollout_ref.rollout: rollout\n\n  # Model config.\n  - model@actor_rollout_ref.model: hf_model\n\n  # Critic model config.\n  - critic@critic: ${model_engine}_critic\n\n  # legacy reward impl config, for backward compatibility\n  - legacy_reward_impl\n\n  # Reward config.\n  - reward@reward: reward\n\n  # Rollout correction config.\n  - algorithm@algorithm.rollout_correction: rollout_correction\n\n  # load the reference default config, then apply the fields in the current yaml\n  # self config override anything above\n  - _self_\n\n# config for actor, rollout and reference model\nactor_rollout_ref:\n\n  # Whether it's a hybrid engine, currently only supports hybrid engine\n  hybrid_engine: true\n\n  # Timeout for operations executed against the process group\n  nccl_timeout: 600\n\n  # Rollout model config.\n  rollout:\n\n    # for huge model, layered summon can save memory (prevent OOM) but make it slower\n    layered_summon: False\n\n# config for the algorithm\nalgorithm:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.trainer.config.AlgoConfig\n\n  # Discount factor for future rewards\n  gamma: 1.0\n\n  # Trade-off between bias and variance in the GAE estimator\n  lam: 1.0\n\n  # Advantage estimator type: \"gae\", \"grpo\", \"reinforce_plus_plus\", etc.\n  adv_estimator: gae\n\n  # Whether to normalize advantages by std (specific to GRPO)\n  norm_adv_by_std_in_grpo: True\n\n  # Whether to enable in-reward KL penalty\n  use_kl_in_reward: False\n\n  # How to estimate KL divergence: \"kl\", \"abs\", \"mse\", \"low_var_kl\", or \"full\"\n  kl_penalty: kl\n\n  # KL control configuration\n  kl_ctrl:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.trainer.config.KLControlConfig\n\n    # KL control type: \"fixed\" or \"adaptive\"\n    type: fixed\n\n    # Initial coefficient for KL penalty\n    kl_coef: 0.001\n\n    # Horizon value for adaptive controller (if enabled)\n    horizon: 10000\n\n    # Target KL divergence (used for adaptive controller)\n    target_kl: 0.1\n\n  # Whether to enable preference feedback PPO\n  use_pf_ppo: False\n\n  # Preference feedback PPO settings\n  pf_ppo:\n\n    # Method for reweighting samples: \"pow\", \"max_min\", or \"max_random\"\n    reweight_method: pow\n\n    # Power used for weight scaling in \"pow\" method\n    weight_pow: 2.0\n\n# config for the trainer\ntrainer:\n\n  # Whether to balance batch sizes across distributed workers\n  balance_batch: True\n\n  # Number of epochs in training\n  total_epochs: 30\n\n  # Total training steps (can be set explicitly or derived from epochs)\n  total_training_steps: null\n\n  # Project name for experiment tracking (e.g., wandb)\n  project_name: verl_examples\n\n  # Experiment name for run identification in tracking tools\n  experiment_name: gsm8k\n\n  # Logging backends to use: \"console\", \"wandb\", etc.\n  logger: [\"console\", \"wandb\"]\n\n  # Number of generations to log during validation\n  log_val_generations: 0\n\n  # Directory for logging rollout data; no dump if null\n  rollout_data_dir: null\n\n  # Directory for logging validation data; no dump if null\n  validation_data_dir: null\n\n  # Number of nodes used in the training\n  nnodes: 1\n\n  # Number of GPUs per node\n  n_gpus_per_node: 8\n\n  # Save frequency (by iteration) for model checkpoints\n  save_freq: -1\n\n  # ESI refers to the elastic server instance used during training, similar to the training plan. For example,\n  # if you purchase 10 hours of computing power, the ESI will automatically shut down after 10 hours of training.\n  # To ensure a checkpoint is saved before ESI shuts down, the system will start saving a checkpoint in advance.\n  # The advance time is calculated as: Advance Time = Longest historical step duration + Checkpoint save duration + esi_redundant_time.\n  # Here, esi_redundant_time is a user-defined value that further extends the advance time for added safety.\n  esi_redundant_time: 0\n\n  # Resume mode: \"auto\", \"disable\", or \"resume_path\"\n  # \"auto\": resume from last checkpoint if available\n  # \"disable\": start from scratch\n  # \"resume_path\": resume from a user-defined path\n  resume_mode: auto\n\n  # Path to resume training from (only used when resume_mode is \"resume_path\")\n  resume_from_path: null\n\n  # Whether to run validation before training begins\n  val_before_train: True\n\n  # Whether to run validation only\n  val_only: False\n\n  # Validation frequency (in training iterations)\n  test_freq: -1\n\n  # Number of iterations to warm up the critic before updating policy\n  critic_warmup: 0\n\n  # Default path to distributed filesystem for saving checkpoints\n  default_hdfs_dir: null\n\n  # Whether to delete local checkpoints after loading\n  del_local_ckpt_after_load: False\n\n  # Default local directory for saving checkpoints\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n\n  # Maximum number of actor checkpoints to keep\n  max_actor_ckpt_to_keep: null\n\n  # Maximum number of critic checkpoints to keep\n  max_critic_ckpt_to_keep: null\n\n  # Timeout (in seconds) for Ray worker to wait for registration\n  ray_wait_register_center_timeout: 300\n\n  # Device to run training on (e.g., \"cuda\", \"cpu\")\n  device: cuda\n\n  # whether to use legacy worker implementation\n  #  mode: \"auto\", \"enable\", or \"disable\"\n  use_legacy_worker_impl: auto\n\n# profiler configs\nglobal_profiler:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.utils.profiler.ProfilerConfig\n\n  # Profiling tool: choose between nsys, npu, torch, torch_memory\n  tool: null\n\n  # profile steps\n  steps: null\n\n  # Whether to combine continuous steps into one database.\n  ## If True, worker.profiler.discrete must be False, [1,2] in one, [5] in another.\n  ## If False, [1] in one, [2] in another, [5] in another.\n  profile_continuous_steps: False\n\n  # Path to save profiling contents\n  save_path: \"outputs/profile\"\n\n  # Specific tool configs, can use +profiler.tool_config.[tool].xxx to config\n  global_tool_config:\n\n    # nsys config\n    nsys:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NsightToolConfig\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: False\n\n      # controller Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n      ## reference https://docs.nvidia.com/nsight-systems/UserGuide/index.html\n      ## reference https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html\n      controller_nsight_options:\n\n        # Select the API(s) to be traced.\n        trace: \"cuda,nvtx,cublas,ucx\"\n\n        # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n        cuda-memory-usage: \"true\"\n\n        # CUDA graphs will be traced as a whole\n        cuda-graph-trace: \"graph\"\n\n      # worker Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n      worker_nsight_options:\n\n        # Select the API(s) to be traced.\n        trace: \"cuda,nvtx,cublas,ucx\"\n\n        # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n        cuda-memory-usage: \"true\"\n\n        # CUDA graphs will be traced as a whole\n        cuda-graph-trace: \"graph\"\n\n        # Profiling only in a range of torch.cuda.profiler.start and stop. Do not change this config.\n        capture-range: \"cudaProfilerApi\"\n\n        # Specify the desired behavior when a capture range ends.\n        # In verl we need the torch.cuda.profiler.start/stop pair to repeats n times.\n        # valid values are \"repeat-shutdown:n\" or null.\n        # For normal whole step profiling, n = len(profile_steps);\n        # but for discrete profiling, n = len(profile_steps) * Number(subtasks).\n        # Or you can just leave it null and the program will use n = len(profile_steps) * 6;\n        capture-range-end: null\n\n        # Send signal to the target application's process group. We let the program to exit by itself.\n        kill: none\n\n    # enable memory visualization for debugging memory usage\n    torch_memory:\n\n      #  Maximum number of allocation entries to record\n      trace_alloc_max_entries: 100_000\n\n      # The depth of the call stack to capture for each allocation\n      stack_depth: 32\n\n      # 'alloc': records only allocation events || 'state': records memory state changes || 'all': records both.\n      context: \"all\"\n\n      # 'python': records Python stacks || 'cpp': records C++ stacks (available in some versions) || 'all': records both.\n      stacks: \"all\"\n\n      # devices, record_context etc.\n      kw_args: {}\n\n# configs for TransferQueue\ntransfer_queue:\n\n  # Whether to enable transfer queue\n  enable: False\n\n# configs related to ray\nray_kwargs:\n\n  # configs related to ray initialization\n  ray_init:\n\n    # Number of CPUs for Ray. Use a fixed number instead of null when using SLURM.\n    num_cpus: null\n\n  # Path to save Ray timeline JSON for performance profiling\n  timeline_json_file: null\n"
  },
  {
    "path": "verl/trainer/config/profiler/profiler.yaml",
    "content": "# Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n_target_: verl.utils.profiler.ProfilerConfig\n\n# profiler tool, default same as profiler.tool in global config\n# choices: nsys, npu, torch\ntool: torch\n\n# whether enable profile on Actor\nenable: False\n\n# Whether to profile all ranks.\nall_ranks: False\n\n# The ranks that will be profiled. [] or [0,1,...]\nranks: []\n\n# profile results saving path\nsave_path: \"outputs/profile\"\n\ntool_config:\n  npu:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.config.NPUToolConfig\n\n    # Contents to profile, can be empty\n    # options: npu, cpu, memory, shapes, module, stack\n    contents: [ ]\n\n    # Collection level, optional values: level_none, level0, level1, level2.\n    level: \"level0\"\n\n    # Whether to automatically parse the data.\n    analysis: True\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: False\n\n    name: npu\n\n\n  nsys:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.config.NsightToolConfig\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n\n    name: nsight\n\n  torch:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n    # Contents to profile, can be empty\n    # options: cuda, cpu, memory, shapes, stack\n    contents: []\n\n    # True for each task has its own database, False for all tasks in one training step share one database.\n    discrete: false\n\n    name: torch\n\n  torch_memory:\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n\n    # Maximum number of memory allocation entries to track\n    trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n\n    # Stack trace depth for memory allocations\n    stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n\n    name: torch_memory"
  },
  {
    "path": "verl/trainer/config/ref/dp_ref.yaml",
    "content": "# defaults specify the default config from each component\ndefaults:\n\n  # dp ref config, inheriting from trainer/config/ref/ref.yaml\n  - ref\n  \n  # fsdp engine config\n  - ../engine@fsdp_config: fsdp\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n# Target class for this configuration\n_target_: verl.workers.config.FSDPActorConfig\n\n# fsdp config\nfsdp_config:\n\n  # ref model is forward only\n  forward_only: True\n\n# sequence parallel size\n# same as actor_rollout_ref.actor.ulysses_sequence_parallel_size if it exists, otherwise 1\nulysses_sequence_parallel_size: ${oc.select:actor_rollout_ref.actor.ulysses_sequence_parallel_size,1}\n\n# calculate entropy with chunking to reduce memory peak\nentropy_from_logits_with_chunking: False\n\n# recompute entropy\nentropy_checkpointing: False\n"
  },
  {
    "path": "verl/trainer/config/ref/megatron_ref.yaml",
    "content": "# megatron ref config, inheriting from trainer/config/ref/ref.yaml\ndefaults:\n  - ref\n\n  # megatron engine config\n  - ../engine@megatron: megatron\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.McoreActorConfig\n\nstrategy: megatron\n\nmegatron:\n  seed: ${oc.select:actor_rollout_ref.actor.megatron.seed,42}\n  override_transformer_config: ${oc.select:actor_rollout_ref.actor.megatron.override_transformer_config,{}}\n  use_mbridge: ${oc.select:actor_rollout_ref.actor.megatron.use_mbridge,False}\n  vanilla_mbridge: ${oc.select:actor_rollout_ref.actor.megatron.vanilla_mbridge,True}\n  use_remove_padding: ${oc.select:actor_rollout_ref.actor.megatron.use_remove_padding,True}\n  tensor_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.tensor_model_parallel_size,1}\n  pipeline_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.pipeline_model_parallel_size,1}\n  virtual_pipeline_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.virtual_pipeline_model_parallel_size,null}\n  context_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.context_parallel_size,1}\n  expert_model_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.expert_model_parallel_size,1}\n  expert_tensor_parallel_size: ${oc.select:actor_rollout_ref.actor.megatron.expert_tensor_parallel_size,null}\n  param_offload: ${oc.select:actor_rollout_ref.actor.megatron.param_offload,False}\n  forward_only: True\n\nload_weight: True\n"
  },
  {
    "path": "verl/trainer/config/ref/ref.yaml",
    "content": "# Number of rollouts per update (mirrors actor rollout_n)\nrollout_n: ${oc.select:actor_rollout_ref.rollout.n,1}\n\n# actor_rollout_ref.ref: FSDP config same as actor. For models larger than 7B, it’s recommended to turn on offload for ref by default\nstrategy: ${actor_rollout_ref.actor.strategy}\n\n# whether to enable torch.compile\n# same as actor_rollout_ref.actor.use_torch_compile if it exists, otherwise 1\nuse_torch_compile: ${oc.select:actor_rollout_ref.actor.use_torch_compile,true}\n\n# [Will be deprecated, use log_prob_micro_batch_size_per_gpu]\n# The batch size for one forward pass in the computation of log_prob. Global batch size.\nlog_prob_micro_batch_size: null\n\n# The batch size for one forward pass in the computation of log_prob. Local batch size per GPU.\nlog_prob_micro_batch_size_per_gpu: null\n\n# enable dynamic batch size (sequence packing) for log_prob computation\n# same as actor_rollout_ref.actor.use_dynamic_bsz if it exists, otherwise false\nlog_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n\n# the max token length per GPU\n# same as actor_rollout_ref.actor.ppo_max_token_len_per_gpu if it exists, otherwise 16384\nlog_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n\n# profile the ref model in `compute_log_prob`\nprofiler:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.utils.profiler.ProfilerConfig\n\n  # choices: nsys, npu, torch, torch_memory\n  tool: ${oc.select:global_profiler.tool,null}\n\n  # whether enable profile on Ref\n  enable: False\n\n  # Whether to profile all ranks.\n  all_ranks: False\n\n  # The ranks that will be profiled. [] or [0,1,...]\n  ranks: []\n\n  # profile results saving path\n  save_path: ${oc.select:global_profiler.save_path,null}\n\n  # specific tool config which only related to the role\n  tool_config:\n\n    # nsys tool config\n    nsys:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NsightToolConfig\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: ${oc.select:global_profiler.global_tool_config.nsys.discrete}\n\n    # npu config\n    npu:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NPUToolConfig\n\n      # Contents to profile, can be empty\n      # options: npu, cpu, memory, shapes, module, stack\n      contents: []\n\n      # Collection level, optional values: level_none, level0, level1, level2.\n      level: \"level0\"\n\n      # Whether to automatically parse the data.\n      analysis: True\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: False\n\n    # torch profiler config\n    torch:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n      # Contents to profile, can be empty\n      # options: cuda, cpu, memory, shapes, stack\n      contents: []\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: false\n\n    # torch memory profiler config\n    torch_memory:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchMemoryToolConfig\n\n      # Maximum number of memory allocation entries to track\n      trace_alloc_max_entries: ${oc.select:global_profiler.global_tool_config.torch_memory.trace_alloc_max_entries,100000}\n\n      # Stack trace depth for memory allocations\n      stack_depth: ${oc.select:global_profiler.global_tool_config.torch_memory.stack_depth,32}\n\n# Router replay configuration for MoE models\nrouter_replay:\n\n  # Target dataclass for this configuration\n  _target_: verl.workers.config.RouterReplayConfig\n\n  # Router replay mode: disabled, R2, R3\n  # - R2: Use R2 routing strategy (record mode)\n  # - R3: Use R3 routing strategy (record mode)\n  mode: disabled\n\n  # File path to save recorded routing decisions\n  # Required when mode is 'record', 'R2', or 'R3'\n  record_file: null\n\n  # File path to load recorded routing decisions for replay\n  # Required when mode is 'replay'\n  replay_file: null"
  },
  {
    "path": "verl/trainer/config/ref/torchtitan_ref.yaml",
    "content": "# torchtitan ref config, inheriting from trainer/config/ref/ref.yaml\ndefaults:\n  # torchtitan optimizer config\n  - ../optim@optim: torchtitan\n\n  # torchtitan engine config\n  - ../engine@torchtitan: torchtitan\n\n  - ref\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.TorchTitanActorConfig\n\nstrategy: torchtitan\n\ntorchtitan:\n  seed: ${oc.select:actor_rollout_ref.actor.torchtitan.seed,42}\n  data_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_size,1}\n  data_parallel_replicate_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_replicate_size,1}\n  data_parallel_shard_size: ${oc.select:actor_rollout_ref.actor.torchtitan.data_parallel_shard_size,1}\n  tensor_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.tensor_parallel_size,1}\n  expert_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.expert_parallel_size,1}\n  pipeline_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.pipeline_parallel_size,1}\n  context_parallel_size: ${oc.select:actor_rollout_ref.actor.torchtitan.context_parallel_size,1}\n  attn_type: ${oc.select:actor_rollout_ref.actor.torchtitan.attn_type,flex}\n  forward_only: True\n"
  },
  {
    "path": "verl/trainer/config/ref/veomni_ref.yaml",
    "content": "# veomni ref config, inheriting from trainer/config/ref/ref.yaml\ndefaults:\n  - ref\n\n  # veomni engine config\n  - ../engine@veomni: veomni\n\n  # load the reference default config, then apply the fields in the current yaml\n  - _self_\n\n_target_: verl.workers.config.VeOmniActorConfig\n\nstrategy: veomni\n\nveomni:\n  seed: ${oc.select:actor_rollout_ref.actor.veomni.seed,42}\n  fsdp_size: ${oc.select:actor_rollout_ref.actor.veomni.fsdp_size,-1}\n  ulysses_parallel_size: ${oc.select:actor_rollout_ref.actor.veomni.ulysses_parallel_size,1}\n  expert_parallel_size: ${oc.select:actor_rollout_ref.actor.veomni.expert_parallel_size,1}\n  param_offload: ${oc.select:actor_rollout_ref.actor.veomni.param_offload,False}\n  attn_implementation: ${oc.select:actor_rollout_ref.actor.veomni.attn_implementation,flash_attention_2}\n  moe_implementation: ${oc.select:actor_rollout_ref.actor.veomni.moe_implementation,fused}\n  forward_only: True\n\n"
  },
  {
    "path": "verl/trainer/config/reward/reward.yaml",
    "content": "# configs for the reward computation\n\n# we launch num_workers reward managers to parallelize the reward computation\nnum_workers: 8\n\n# custom reward function definition\ncustom_reward_function:\n  # The path to the file containing your customized reward function.\n  # If not specified, pre-implemented reward functions will be used.\n  path: null\n  # The name of the reward function within the specified file. Default is 'compute_score'.\n  name: compute_score\n\n# reward manager, see `verl/trainer/config/reward_manager.yaml` for details.\nreward_manager:\n  _target_: verl.workers.config.reward_model.RewardManagerConfig\n  source: register\n  name: naive\n  module:\n    _target_: verl.trainer.config.config.ModuleConfig\n    path: null\n    name: custom_reward_manager\n\n# Inference config for reward models,\n# support both discriminative and generative models\nreward_model:\n  enable: False\n\n  # Whether to deploy the model to a separate resource pool.\n  # If true, n_gpus_per_node & nnodes will be used to determine the resource node.\n  enable_resource_pool: False\n  n_gpus_per_node: 8\n  nnodes: 0\n\n  model_path: null\n  rollout:\n    _target_: verl.workers.config.RolloutConfig\n    name: ???\n    dtype: bfloat16\n    gpu_memory_utilization: 0.5\n    enforce_eager: true\n    cudagraph_capture_sizes: null\n    free_cache_engine: true\n    data_parallel_size: 1\n    expert_parallel_size: 1\n    tensor_model_parallel_size: 2\n    max_num_batched_tokens: 8192\n    max_model_len: null\n    max_num_seqs: 1024\n    load_format: auto\n    engine_kwargs: {}\n    limit_images: null\n    enable_chunked_prefill: true\n    enable_prefix_caching: true\n    disable_log_stats: true\n    skip_tokenizer_init: false\n\n    prompt_length: 2048\n    response_length: 2048\n\n# Cloud/local sandbox fusion configuration for custom reward logic\nsandbox_fusion:\n  # Cloud /local function URL for sandbox execution\n  url: null\n  # Max concurrent requests allowed to sandbox\n  max_concurrent: 64\n  # Max memory limit for each sandbox process in MB\n  memory_limit_mb: 1024\n"
  },
  {
    "path": "verl/trainer/config/rollout/rollout.yaml",
    "content": "# Target class for this configuration\n_target_: verl.workers.config.RolloutConfig\n\n# actor_rollout_ref.rollout.name: hf/vllm/sglang/trtllm. The default value will be removed in the future\nname: ???\n\n# sync: LLM, async: AsyncLLM\nmode: async\n\n# Number of nodes for standalone rollout server, must be > 0 in one-step-off/fully async training.\nnnodes: 0\n\n# Number of GPUs per node for rollout server.\nn_gpus_per_node: ${oc.select:trainer.n_gpus_per_node,8}\n\n# Sampling temperature for rollout.\ntemperature: 1.0\n\n# Top-k sampling parameter. -1 for vLLM rollout, 0 for HF rollout.\ntop_k: -1\n\n# Top-p sampling parameter. Default 1.0.\ntop_p: 1\n\n# typically the same as data max prompt length\n# same as data.max_prompt_length if it exists\nprompt_length: ${oc.select:data.max_prompt_length,512}\n\n# typically the same as data max response length\n# same as data.max_response_length if it exists\nresponse_length: ${oc.select:data.max_response_length,512}\n\n# for vllm rollout\n# Rollout model parameters type. Align with actor model's FSDP/Megatron type.\ndtype: bfloat16\n\n# Fraction of GPU memory used by vLLM/SGLang/TRTLLM for KV cache.\ngpu_memory_utilization: 0.5\n\n# Whether to ignore EOS and continue generating after EOS is hit.\nignore_eos: False\n\n# Whether to disable CUDA graph. Default False to best performance.\nenforce_eager: False\n\n# batch size of cudagraph to capture. Require enforce_eager: False to use this option\n# Since cudagraph in inference engine can not be offloaded during update policy,\n# you can use smaller batch size to save memory used in cuda graph, eg: [1 ,2, 4, 8, 16, 32]\n# supported engines: vllm\ncudagraph_capture_sizes: null\n\n# Whether to free engine KVCache after generation.\nfree_cache_engine: True\n\n# TP size for rollout. Not effective for hf\ntensor_model_parallel_size: 2\n\n# DP size for rollout\ndata_parallel_size: 1\n\n# EP size for rollout\n# For MoE models in vllm, EP=1 refers to ETP parallel in fused_moe with TP*DP weight splits,\n# EP>1 (should satisfy EP=TP*DP) refers to EP parallel in fused_moe\nexpert_parallel_size: 1\n\n# PP size for rollout.\npipeline_model_parallel_size: 1\n\n# max number of tokens in a batch\nmax_num_batched_tokens: 8192\n\n# max length for rollout\nmax_model_len: null\n\n# max length of sequences\nmax_num_seqs: 1024\n\n# may get higher throughput when set to True. When activated, Please increase max_num_batched_tokens or decrease max_model_len.\nenable_chunked_prefill: True\n\n# Prefix caching kv-cache blocks is a popular optimization in LLM inference to avoid redundant prompt computations.\nenable_prefix_caching: True\n\n# logprobs mode for rollout logprobs\nlogprobs_mode: processed_logprobs\n\n# scheduling policy for vllm rollout\nscheduling_policy: fcfs\n\n# Which loader to use for rollout model weights: dummy, hf, megatron, etc.\n# safetensors (for huge model, and set use_shm=True); dummy: randomly init model weight\nload_format: dummy\n\n# [Will be deprecated, use log_prob_micro_batch_size_per_gpu] The batch size for one forward pass in the computation of log_prob. Global batch size.\nlog_prob_micro_batch_size: null\n\n# The batch size for one forward pass in the computation of log_prob. Local batch size per GPU.\nlog_prob_micro_batch_size_per_gpu: null\n\n# enable dynamic batch size (sequence packing) for log_prob computation\n# same as actor_rollout_ref.actor.use_dynamic_bsz if it exists, otherwise false\nlog_prob_use_dynamic_bsz: ${oc.select:actor_rollout_ref.actor.use_dynamic_bsz,false}\n\n# max token length for log_prob computation\n# same as actor_rollout_ref.actor.ppo_max_token_len_per_gpu if it exists, otherwise 16384\nlog_prob_max_token_len_per_gpu: ${oc.select:actor_rollout_ref.actor.ppo_max_token_len_per_gpu,16384}\n\n# disable logging statistics\ndisable_log_stats: True\n\n# for hf rollout\n# Whether to sample during training rollout. False uses greedy sampling.\ndo_sample: True\n\n# number of responses (i.e. num sample times). > 1 for grpo\nn: 1\n\n# The over_sample_rate parameter controls the early termination threshold for training rollouts,\n# where the system will abort remaining requests when (1 - over_sample_rate) * total_requests completions are reached.\nover_sample_rate: 0\n\n# Whether to wake up inference engine in multi-stage for SGLang\n# to reduce peak memory during training-rollout transition.\n# This is only effective for SGLang rollout.\nmulti_stage_wake_up: false\n\n# Extra inference engine arguments (vllm, sglang, trtllm), please refer vllm/sglang/trtllm official doc for detail\nengine_kwargs:\n\n  # vllm engine config\n  vllm: {}\n\n  # sglang engine config\n  sglang: {}\n\n  # trtllm engine config\n  trtllm: {}\n\n# Sampling parameters used during validation.\nval_kwargs:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.SamplingConfig\n\n  # sampling parameters for validation\n  # Top-k sampling parameter. -1 for vLLM rollout, 0 for HF rollout.\n  top_k: -1\n\n  # Top-p sampling parameter. Default 1.0.\n  top_p: 1.0\n\n  # Sampling temperature for rollout.\n  temperature: 0\n\n  # whether to repeat n times for validation\n  n: 1\n\n  # Whether to sample during training rollout. False uses greedy sampling.\n  do_sample: False\n\n# Multi-turn interaction config for tools or chat.\nmulti_turn:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.MultiTurnConfig\n\n  # set to True for multi-turn tool interaction tasks; should set rollout.name to sglang as well\n  enable: False\n\n  # null for no limit (default max_length // 3)\n  max_assistant_turns: null\n\n  # null for no tool\n  tool_config_path: null\n\n  # null for no limit (default max_length // 3)\n  max_user_turns: null\n\n  # max parallel call for tools in single turn\n  max_parallel_calls: 1\n\n  # max length of tool response\n  max_tool_response_length: 256\n\n  # truncate side of tool response: left, middle, right\n  tool_response_truncate_side: middle\n\n  # null for no interaction\n  interaction_config_path: null\n\n  # - When set to True, the model's default chat template is used for multi-turn rollout, which typically matches production behavior.\n  # - When set to False, the token ids recorded for training are used instead; unlike the default chat template, these always include the model's full output,\n  #   which may contain additional content such as reasoning content. This maintains the consistency between training and rollout, but it will lead to longer prompts.\n  use_inference_chat_template: False\n\n  # Tokenization is performed turn by turn and the resulting token ids are concatenated to form the full conversation.\n  # To ensure this matches the result of tokenizing the entire conversation at once, a sanity check is run at the end of each multi-turn rollout to compare the two sets of token ids.\n  # Some models are known to produce different tokenization results when tokenizing turn by turn vs. all at once. aThis behavior has already been validated for them.\n  # To reduce excessive warnings, you can turn off the sanity check for these models if you are using their default chat template:\n  # Qwen/QwQ-32B, Qwen/Qwen3-xxB\n  # - disable: disable tokenization sanity check\n  # - strict: enable strict tokenization sanity check (default)\n  # - ignore_strippable: ignore strippable tokens when checking tokenization sanity\n  tokenization_sanity_check_mode: strict\n\n  # Format of the multi-turn interaction. Options: hermes, llama3_json, ...\n  format: hermes\n\n  # Number of repeat rollouts for each interaction\n  num_repeat_rollouts: null\n\n# support logging rollout prob for debugging purpose\n# \"Truncated importance sampling\" requires rollout log probs, set to True when turning on Truncated importance sampling\ncalculate_log_probs: False\n\n# [Experimental] agent loop based rollout configs\nagent:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.AgentLoopConfig\n\n  # Number of agent loop workers\n  num_workers: 8\n\n  # default agent loop to use if `agent_name` not set in RL dataset\n  default_agent_loop: single_turn_agent\n\n  # custom agent loop config path, which should contain list of configs to initialize AgentLoop instances.\n  # https://hydra.cc/docs/advanced/instantiate_objects/overview/\n  #\n  # - name: react_agent\n  #   _target_: recipe.langgraph_agent.react_agent_loop.ReactAgentLoop\n  #   tools: [\"get_current_temperature\"]\n  # - name: math_expression\n  #   _target_: recipe.langgraph_agent.example.math_expression.MathExpressionReactAgentLoop\n  #   min_terms: 2\n  #   max_terms: 6\n  agent_loop_config_path: null\n\n  # custom async server configs\n  custom_async_server:\n\n    # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n    _target_: verl.workers.config.CustomAsyncServerConfig\n\n    # Path to the custom async server implementation\n    path: null\n\n    # Class name of the custom async server class (e.g. AsyncvLLMServer)\n    name: null\n\n# Checkpoint Engine config for update weights from trainer to rollout\ncheckpoint_engine:\n\n  # Target class for checkpoint engine config\n  _target_: verl.workers.config.CheckpointEngineConfig\n\n  # Backend for checkpoint engine: naive, nccl, nixl, hccl\n  backend: naive\n\n  # Specifies the tensor bucket size (in megabytes) for batch weight updates during rollout operations.\n  # This parameter controls the maximum payload size for a single weight update request.\n  # Reference: https://github.com/volcengine/verl/pull/2418\n  # Currently only supported in SGLang rollout implementations\n  # Larger values may improve throughput but increase memory overhead\n  # Detailed performance comparison:\n  # https://github.com/zhaochenyang20/Awesome-ML-SYS-Tutorial/issues/169#issuecomment-3070686720\n  # Default value (512MB) is optimized for typical GPU memory configurations\n  # For the best performance of `rebuild_cuda_tensor`, it is recommended to:\n  # 1. Enable `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES`\n  # 2. Manually set `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`\n  # when using Tensor Parallelism (TP) >= 8.\n  update_weights_bucket_megabytes: 2048\n\n  # Additional keyword arguments to pass to the checkpoint engine constructor\n  engine_kwargs: {}\n\n# trace rollout data\ntrace:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.TraceConfig\n\n  # Project name for experiment tracking (e.g., wandb)\n  project_name: ${oc.select:trainer.project_name,null}\n\n  # Experiment name for run identification in tracking tools\n  experiment_name: ${oc.select:trainer.experiment_name,null}\n\n  # trace backend, support mlflow, weave\n  backend: null\n\n  # whether translate token id to text in output\n  token2text: False\n\n  # Maximum number of unique samples to trace per agent worker per training step.\n  # If null, all samples are traced. If set to N, each agent loop worker will randomly\n  # select N unique samples to trace (including all their rollouts for GRPO).\n  # Total traces per step = max_samples_per_step_per_worker * num_workers * n_rollouts_per_sample\n  max_samples_per_step_per_worker: null\n\n# When enabled (True), the trainer will attempt to load previously generated rollout data from the specified directory instead of computing new rollouts.\n# If no cached data is found or loading fails, new rollouts will be generated and automatically saved.\n# This feature is useful for debugging or when you want to reuse computation results across multiple runs.\nskip_rollout: False\n\n# Specifies the filesystem path where rollout data should be cached when skip_rollout is enabled.\n# Note: Giving path under /tmp/ray/session* is not recommended as these are temporary Ray cluster directories.\nskip_dump_dir: /tmp/rollout_dump\n\n# Whether to skip tokenizer initialization for rollout engine\n# When enabled (True), the rollout assume token in token out for generation\nskip_tokenizer_init: True\n\n# Whether to enable rollout routing replay for MoE models\n# When enabled (True), the rollout will record the routing decisions.\nenable_rollout_routing_replay: False\n\n\n# profile the rollout model in `generate_sequence`\nprofiler:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.utils.profiler.ProfilerConfig\n\n  # profiler tool, default same as profiler.tool in global config\n  # choices: npu, torch\n  tool: ${oc.select:global_profiler.tool,null}\n\n  # whether enable profile on rollout\n  enable: ${oc.select:actor_rollout_ref.actor.profiler.enable,false}\n\n  # Whether to profile all ranks.\n  all_ranks: ${oc.select:actor_rollout_ref.actor.profiler.all_ranks,false}\n\n  # The ranks that will be profiled. [] or [0,1,...]\n  ranks: ${oc.select:actor_rollout_ref.actor.profiler.ranks,[]}\n\n  # profile results saving path\n  save_path: ${oc.select:global_profiler.save_path,null}\n\n  # specific tool config\n  tool_config:\n\n    # npu config\n    npu:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.NPUToolConfig\n\n      # Contents to profile, can be empty\n      # options: npu, cpu, memory, shapes, module, stack\n      contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.contents,[]}\n\n      # Collection level, optional values: level_none, level0, level1, level2.\n      level: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.level,level0}\n\n      # Whether to automatically parse the data.\n      analysis: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.analysis,false}\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.npu.discrete,false}\n\n    # torch profiler config\n    torch:\n\n      # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n      _target_: verl.utils.profiler.config.TorchProfilerToolConfig\n\n      # Contents to profile, can be empty\n      # options: cuda, cpu, memory, shapes, stack\n      contents: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.contents,[]}\n\n      # True for each task has its own database, False for all tasks in one training step share one database.\n      discrete: ${oc.select:actor_rollout_ref.actor.profiler.tool_config.torch.discrete,false}\n\n# prometheus configuration for vllm/sglang server mode\nprometheus:\n\n  # Required when using verl.utils.omega_conf_to_dataclass to instantiate dataclass configs\n  _target_: verl.workers.config.PrometheusConfig\n\n  # whether enable prometheus on server mode rollout\n  enable: false\n\n  # Port number that Prometheus listens on, default is 9090\n  port: 9090\n\n  # Path to Prometheus configuration file\n  file: /tmp/ray/session_latest/metrics/prometheus/prometheus.yml\n\n  # Specify served_model_name to avoid displaying overly long model paths in Grafana\n  served_model_name: ${oc.select:actor_rollout_ref.model.path,null}\n\n# type of quantization in vllm, currently support fp8 and torchao\nquantization: null\n\n# extra quantization information serialized in a config file, e.g. torchao_config.json\nquantization_config_file: null\n\n# MTP configuration, reuse model configuration\nmtp: ${oc.select:actor_rollout_ref.model.mtp, null}\n\n# QAT configuration (inherited from actor's engine config)\nqat: ${oc.select:actor_rollout_ref.actor.fsdp_config.qat,null}\n"
  },
  {
    "path": "verl/trainer/config/sft_trainer_engine.yaml",
    "content": "# Format checks enforced on CI:\n# 1. Comments must appear above each field.\n# 2. There must be a blank line between each field.\n# 3. Inline comments (after a field on the same line) are not allowed.\n# 4. Indentation level is respected for nested fields.\n\n# <folder_name>@<field_name>.<field_name>: <yaml_file_name>\n\ndefaults:\n  - model@model: hf_model\n  - engine@engine: fsdp\n  - optim@optim: fsdp\n  - profiler@profiler: profiler\n  - _self_\n\ndata:\n  train_batch_size: 256 # global batch size\n  micro_batch_size_per_gpu: 4  # this is also val batch size\n  max_token_len_per_gpu: 8192\n  use_dynamic_bsz: True\n  train_files: ~/data/gsm8k/train.parquet\n  val_files: null\n  train_max_samples: -1  # set to -1 to use full dataset\n  val_max_samples: -1  # set to -1 to use full dataset\n  # Multi-turn settings\n  messages_key: messages  # Key for messages list in multi-turn mode\n  tools_key: tools  # Key for tools list in multi-turn mode\n  enable_thinking_key: enable_thinking  # Whether to enable thinking in multi-turn mode\n  enable_thinking_default: none         # The default value when enable_thinking_key is not present in the dataset\n  pad_mode: no_padding\n  # for right padding\n  max_length: 1024\n  truncation: error\n  balance_dp_token: False # to be implement\n  custom_cls:\n    path: null\n    name: null\n  use_shm: False\n  apply_chat_template_kwargs: {}\n  num_workers: 8\n\n  # MultiTurnSFTDataset apply_chat_template to each turn separately and concat `input_ids`\n  # as a whole sequence, which may not equal to apply_chat_template to whole messages at once. \n  # For example, Qwen Thinking series models add <think></think> tags to last turn, please check\n  # your tokenizer chat template settings.\n  # Set to True to ignore input_ids mismatch and use the concatenated input_ids as the final input_ids.\n  ignore_input_ids_mismatch: False\n\n# Checkpoint configuration\ncheckpoint:\n  _target_: verl.trainer.config.CheckpointConfig\n  # What to include in saved checkpoints\n  # with 'hf_model' you can save whole model as hf format, now only use sharded model checkpoint to save space\n  save_contents: [\"model\", \"optimizer\", \"extra\"]\n\n  # For more flexibility, you can specify the contents to load from the checkpoint.\n  load_contents: ${checkpoint.save_contents}\n  # Mbridge config extension.\n  mbridge_config: {}\n\ntrainer:\n  default_local_dir: checkpoints/${trainer.project_name}/${trainer.experiment_name}\n  default_hdfs_dir: null\n  project_name: gsm8k-sft\n  experiment_name: test\n  total_epochs: 4\n  total_training_steps: null\n  logger: [ 'console', 'wandb' ]\n  seed: 1\n  save_freq: -1\n  test_freq: -1\n  max_ckpt_to_keep: null  # Maximum number of checkpoints to keep, set to null to keep all\n\n  # Resume mode: \"auto\", \"disable\", or \"resume_path\"\n  # \"auto\": resume from last checkpoint if available\n  # \"disable\": start from scratch\n  # \"resume_path\": resume from a user-defined path\n  resume_mode: auto\n\n  # Path to resume training from (used when resume_mode is \"resume_path\" or \"auto\")\n  resume_from_path: null  \n  device: cuda\n\n  nnodes: 1\n  n_gpus_per_node: 1\n\n  profile_interval: [-1, -1]\n\nglobal_profiler:\n  global_tool_config:\n    # nsys config\n    nsys:\n      # controller Nvidia Nsight Systems Options. Must set when profile_steps is not None.\n      ## reference https://docs.nvidia.com/nsight-systems/UserGuide/index.html\n      ## reference https://docs.ray.io/en/latest/ray-observability/user-guides/profiling.html\n      worker_nsight_options:\n        # Select the API(s) to be traced.\n        trace: \"cuda,nvtx,cublas,ucx\"\n\n        # Track the GPU memory usage by CUDA kernels. Must be string type \"true\" or \"false\".\n        cuda-memory-usage: \"true\"\n\n        # CUDA graphs will be traced as a whole\n        cuda-graph-trace: \"graph\"\n\n        # Profiling only in a range of torch.cuda.profiler.start and stop. Do not change this config.\n        capture-range: \"cudaProfilerApi\"\n\n        # Specify the desired behavior when a capture range ends.\n        # In verl we need the torch.cuda.profiler.start/stop pair to repeats n times.\n        # valid values are \"repeat-shutdown:n\" or null.\n        # For normal whole step profiling, n = len(profile_steps);\n        # but for discrete profiling, n = len(profile_steps) * Number(subtasks).\n        # Or you can just leave it null and the program will use n = len(profile_steps) * 6;\n        capture-range-end: null\n\n        # Send signal to the target application's process group. We let the program to exit by itself.\n        kill: none\n"
  },
  {
    "path": "verl/trainer/constants_ppo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\n\nfrom ray._private.runtime_env.constants import RAY_JOB_CONFIG_JSON_ENV_VAR\n\nPPO_RAY_RUNTIME_ENV = {\n    \"env_vars\": {\n        \"TOKENIZERS_PARALLELISM\": \"true\",\n        \"NCCL_DEBUG\": \"WARN\",\n        \"VLLM_LOGGING_LEVEL\": \"WARN\",\n        \"VLLM_ALLOW_RUNTIME_LORA_UPDATING\": \"true\",\n        \"CUDA_DEVICE_MAX_CONNECTIONS\": \"1\",\n        # TODO: disable compile cache due to cache corruption issue\n        # https://github.com/vllm-project/vllm/issues/31199\n        \"VLLM_DISABLE_COMPILE_CACHE\": \"1\",\n        # Needed for multi-processes colocated on same NPU device\n        # https://www.hiascend.com/document/detail/zh/canncommercial/83RC1/maintenref/envvar/envref_07_0143.html\n        \"HCCL_HOST_SOCKET_PORT_RANGE\": \"auto\",\n        \"HCCL_NPU_SOCKET_PORT_RANGE\": \"auto\",\n    },\n}\n\n\ndef get_ppo_ray_runtime_env():\n    \"\"\"\n    A filter function to return the PPO Ray runtime environment.\n    To avoid repeat of some environment variables that are already set.\n    \"\"\"\n    working_dir = (\n        json.loads(os.environ.get(RAY_JOB_CONFIG_JSON_ENV_VAR, \"{}\")).get(\"runtime_env\", {}).get(\"working_dir\", None)\n    )\n\n    runtime_env = {\n        \"env_vars\": PPO_RAY_RUNTIME_ENV[\"env_vars\"].copy(),\n        **({\"working_dir\": None} if working_dir is None else {}),\n    }\n    for key in list(runtime_env[\"env_vars\"].keys()):\n        if os.environ.get(key) is not None:\n            runtime_env[\"env_vars\"].pop(key, None)\n    return runtime_env\n"
  },
  {
    "path": "verl/trainer/main_eval.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nOffline evaluate the performance of a generated file using reward model and ground truth verifier.\nThe input is a parquet file that contains N generated sequences and (optional) the ground truth.\n\n\"\"\"\n\nfrom collections import defaultdict\n\nimport hydra\nimport numpy as np\nimport pandas as pd\nimport ray\nfrom omegaconf import OmegaConf\nfrom tqdm import tqdm\n\nfrom verl.trainer.ppo.reward import get_custom_reward_fn\nfrom verl.utils.fs import copy_to_local\n\n\n@ray.remote\ndef process_item(config, data_source, response_lst, reward_data):\n    reward_fn = get_custom_reward_fn(config)\n    ground_truth = reward_data[\"ground_truth\"]\n    score_lst = [reward_fn(data_source, r, ground_truth) for r in response_lst]\n    return data_source, np.mean(score_lst)\n\n\n@hydra.main(config_path=\"config\", config_name=\"evaluation\", version_base=None)\ndef main(config):\n    local_path = copy_to_local(config.data.path, use_shm=config.data.get(\"use_shm\", False))\n    dataset = pd.read_parquet(local_path)\n    responses = dataset[config.data.response_key]\n    data_sources = dataset[config.data.data_source_key]\n    reward_model_data = dataset[config.data.reward_model_key]\n\n    total = len(dataset)\n\n    # Initialize Ray\n    if not ray.is_initialized():\n        ray.init(**OmegaConf.to_container(config.ray_kwargs.get(\"ray_init\", {})))\n\n    # evaluate test_score based on data source\n    data_source_reward = defaultdict(list)\n    # Create remote tasks\n    remote_tasks = [\n        process_item.remote(config, data_sources[i], responses[i], reward_model_data[i]) for i in range(total)\n    ]\n\n    # Process results as they come in\n    with tqdm(total=total) as pbar:\n        while len(remote_tasks) > 0:\n            # Use ray.wait to get completed tasks\n            done_ids, remote_tasks = ray.wait(remote_tasks)\n            for result_id in done_ids:\n                data_source, score = ray.get(result_id)\n                data_source_reward[data_source].append(score)\n                pbar.update(1)\n\n    metric_dict = {}\n    for data_source, rewards in data_source_reward.items():\n        metric_dict[f\"test_score/{data_source}\"] = np.mean(rewards)\n\n    print(metric_dict)\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/trainer/main_generation_server.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nGenerate responses given a dataset of prompts\n\"\"\"\n\nimport os\n\nimport aiohttp\nimport hydra\nimport numpy as np\nimport ray\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n# os.environ['TORCH_COMPILE_DISABLE'] = '1'\n\nimport asyncio\nfrom pprint import pprint\n\nimport pandas as pd\nfrom omegaconf import OmegaConf\nfrom openai.types.chat import ChatCompletion\n\nfrom verl.utils.hdfs_io import makedirs\nfrom verl.workers.rollout.replica import get_rollout_replica_class\n\n\nasync def start_server(config):\n    tp_size = config.actor_rollout_ref.rollout.tensor_model_parallel_size\n    num_replicas = (config.trainer.n_gpus_per_node * config.trainer.nnodes) // tp_size\n    rollout_config = config.actor_rollout_ref.rollout\n    model_config = config.actor_rollout_ref.model\n    # create standalone rollout server\n    rollout_server_class = get_rollout_replica_class(config.actor_rollout_ref.rollout.name)\n    rollout_servers = [\n        rollout_server_class(\n            replica_rank=replica_rank,\n            config=rollout_config,\n            model_config=model_config,\n            gpus_per_node=config.trainer.n_gpus_per_node,\n        )\n        for replica_rank in range(num_replicas)\n    ]\n    await asyncio.gather(*[server.init_standalone() for server in rollout_servers])\n\n    server_handles = [server._server_handle for server in rollout_servers]\n    server_addresses = [server._server_address for server in rollout_servers]\n    assert len(server_handles) == num_replicas\n    assert len(server_addresses) == num_replicas\n\n    return server_handles, server_addresses\n\n\nasync def submit_request(server_address, **chat_complete_request):\n    try:\n        extra_headers = chat_complete_request.pop(\"extra_headers\", {})\n        timeout = aiohttp.ClientTimeout(total=None)\n        session = aiohttp.ClientSession(timeout=timeout)\n        async with session.post(\n            url=f\"http://{server_address}/v1/chat/completions\",\n            headers={\"Authorization\": \"Bearer token-abc123\", **extra_headers},\n            json=chat_complete_request,\n        ) as resp:\n            data = await resp.json()\n            return ChatCompletion(**data)\n    finally:\n        await session.close()\n\n\nasync def generate_per_replica(server_address, model_path: str, n_samples: int, sampling_params: dict, chat_lst: list):\n    # here we should sample n_samples for each chat_lst.\n    # we use aiohttp to avoid hang in AsyncOpenAI when the number of requests is large.\n\n    # client = AsyncOpenAI(\n    #     api_key=\"123-abc\",\n    #     base_url=f\"http://{server_address}/v1\",\n    # )\n\n    chat_complete_request = [\n        {\n            \"model\": model_path,\n            \"messages\": messages,\n            **sampling_params,\n        }\n        for messages in chat_lst\n        for _ in range(n_samples)\n    ]\n\n    tasks = [submit_request(server_address, **req) for req in chat_complete_request]\n    results = await asyncio.gather(*tasks)\n    return results\n\n\nasync def generate(\n    server_addresses: list, model_path: str, n_samples: int, sampling_params: dict, chat_numpy: np.ndarray\n):\n    num_replicas = len(server_addresses)\n    chat_sub_array = np.array_split(chat_numpy, num_replicas)\n    chat_sub_array = [chat.tolist() for chat in chat_sub_array]\n    assert len(server_addresses) == len(chat_sub_array)\n    results = await asyncio.gather(\n        *[\n            generate_per_replica(server_addresses[i], model_path, n_samples, sampling_params, chat_sub_array[i])\n            for i in range(num_replicas)\n        ]\n    )\n    return results\n\n\n@hydra.main(config_path=\"config\", config_name=\"ppo_trainer\", version_base=None)\ndef main(config):\n    ray.init(runtime_env={\"env_vars\": {\"TOKENIZERS_PARALLELISM\": \"true\", \"NCCL_DEBUG\": \"WARN\", \"VLLM_USE_V1\": \"1\"}})\n\n    pprint(OmegaConf.to_container(config, resolve=True))  # resolve=True will eval symbol values\n    OmegaConf.resolve(config)\n\n    n_samples = config.actor_rollout_ref.rollout.n\n\n    if config.actor_rollout_ref.rollout.temperature == 0.0:\n        assert n_samples == 1, \"When temperature=0, n_samples must be 1.\"\n    assert n_samples >= 1, \"n_samples should always >= 1\"\n\n    sampling_params = {\n        \"temperature\": config.actor_rollout_ref.rollout.temperature,\n        \"top_p\": config.actor_rollout_ref.rollout.top_p,\n        # \"top_k\": config.actor_rollout_ref.rollout.top_k,\n        \"max_tokens\": config.actor_rollout_ref.rollout.response_length,\n    }\n\n    from omegaconf import ListConfig\n\n    train_files = config.data.train_files\n    if not isinstance(train_files, list | ListConfig):\n        train_files = [train_files]\n\n    # read dataset. Note that the dataset should directly contain chat template format (e.g., a list of dictionary)\n\n    datasets = []\n    for train_file in train_files:\n        dataset = pd.read_parquet(train_file)\n        datasets.append(dataset)\n\n    # concat dataset\n    dataset = pd.concat(datasets, axis=0, ignore_index=True)\n    chat_lst = dataset[config.data.prompt_key].tolist()\n    chat_lst = [chat.tolist() for chat in chat_lst]\n    chat_numpy = np.array(chat_lst)\n\n    # start native server\n    server_handles, server_addresses = asyncio.run(start_server(config))\n\n    # run generate\n    gen_results = asyncio.run(\n        generate(server_addresses, config.actor_rollout_ref.model.path, n_samples, sampling_params, chat_numpy)\n    )\n\n    # reshape results into a numpy array\n    import itertools\n\n    results = list(itertools.chain.from_iterable(gen_results))\n\n    # extract content from results\n    results = np.array([result.choices[0].message.content for result in results])\n    results = np.reshape(results, (-1, n_samples))\n\n    assert results.shape == (len(chat_lst), n_samples)\n\n    results = results.tolist()\n\n    # add to the data frame\n    dataset[\"responses\"] = results\n\n    # write to a new parquet\n    output_dir = os.path.dirname(config.data.output_path)\n    makedirs(output_dir, exist_ok=True)\n    print(f\"Saving results to {config.data.output_path}\")\n    dataset.to_parquet(config.data.output_path)\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/trainer/main_ppo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nNote that we don't combine the main with ray_trainer as ray_trainer is used by other mpain.\n\"\"\"\n\nimport os\nimport socket\n\nimport hydra\nimport ray\nfrom omegaconf import OmegaConf\n\nfrom verl.experimental.dataset.sampler import AbstractSampler\nfrom verl.experimental.reward_loop import migrate_legacy_reward_impl\nfrom verl.trainer.constants_ppo import get_ppo_ray_runtime_env\nfrom verl.trainer.ppo.ray_trainer import RayPPOTrainer\nfrom verl.trainer.ppo.utils import need_critic, need_reference_policy\nfrom verl.utils.config import validate_config\nfrom verl.utils.device import auto_set_device, is_cuda_available\nfrom verl.utils.import_utils import load_extern_object\n\n\n@hydra.main(config_path=\"config\", config_name=\"ppo_trainer\", version_base=None)\ndef main(config):\n    \"\"\"Main entry point for PPO training with Hydra configuration management.\n\n    Args:\n        config: Hydra configuration dictionary containing training parameters.\n    \"\"\"\n    # Automatically set `config.trainer.device = npu` when running on Ascend NPU.\n    auto_set_device(config)\n    config = migrate_legacy_reward_impl(config)\n    run_ppo(config)\n\n\n# Define a function to run the PPO-like training process\ndef run_ppo(config, task_runner_class=None) -> None:\n    \"\"\"Initialize Ray cluster and run distributed PPO training process.\n\n    Args:\n        config: Training configuration object containing all necessary parameters\n                for distributed PPO training including Ray initialization settings,\n                model paths, and training hyperparameters.\n        task_runner_class: For recipe to change TaskRunner.\n    \"\"\"\n    # Check if Ray is not initialized\n    if not ray.is_initialized():\n        # Initialize Ray with a local cluster configuration\n        # Set environment variables in the runtime environment to control tokenizer parallelism,\n        # NCCL debug level, VLLM logging level, and allow runtime LoRA updating\n        # `num_cpus` specifies the number of CPU cores Ray can use, obtained from the configuration\n        default_runtime_env = get_ppo_ray_runtime_env()\n        ray_init_kwargs = config.ray_kwargs.get(\"ray_init\", {})\n        runtime_env_kwargs = ray_init_kwargs.get(\"runtime_env\", {})\n\n        if config.transfer_queue.enable:\n            # Add runtime environment variables for transfer queue\n            runtime_env_vars = runtime_env_kwargs.get(\"env_vars\", {})\n            runtime_env_vars[\"TRANSFER_QUEUE_ENABLE\"] = \"1\"\n            runtime_env_kwargs[\"env_vars\"] = runtime_env_vars\n\n        runtime_env = OmegaConf.merge(default_runtime_env, runtime_env_kwargs)\n        ray_init_kwargs = OmegaConf.create({**ray_init_kwargs, \"runtime_env\": runtime_env})\n        print(f\"ray init kwargs: {ray_init_kwargs}\")\n        ray.init(**OmegaConf.to_container(ray_init_kwargs))\n\n    if task_runner_class is None:\n        task_runner_class = ray.remote(num_cpus=1)(TaskRunner)  # please make sure main_task is not scheduled on head\n\n    # Create a remote instance of the TaskRunner class, and\n    # Execute the `run` method of the TaskRunner instance remotely and wait for it to complete\n    if (\n        is_cuda_available\n        and config.global_profiler.tool == \"nsys\"\n        and config.global_profiler.get(\"steps\") is not None\n        and len(config.global_profiler.get(\"steps\", [])) > 0\n    ):\n        from verl.utils.import_utils import is_nvtx_available\n\n        assert is_nvtx_available(), \"nvtx is not available in CUDA platform. Please 'pip3 install nvtx'\"\n        nsight_options = OmegaConf.to_container(\n            config.global_profiler.global_tool_config.nsys.controller_nsight_options\n        )\n        runner = task_runner_class.options(runtime_env={\"nsight\": nsight_options}).remote()\n    else:\n        runner = task_runner_class.remote()\n    ray.get(runner.run.remote(config))\n\n    # [Optional] get the path of the timeline trace file from the configuration, default to None\n    # This file is used for performance analysis\n    timeline_json_file = config.ray_kwargs.get(\"timeline_json_file\", None)\n    if timeline_json_file:\n        ray.timeline(filename=timeline_json_file)\n\n\nclass TaskRunner:\n    \"\"\"Ray remote class for executing distributed PPO training tasks.\n\n    This class encapsulates the main training logic and runs as a Ray remote actor\n    to enable distributed execution across multiple nodes and GPUs.\n\n    Attributes:\n        role_worker_mapping: Dictionary mapping Role enums to Ray remote worker classes\n        mapping: Dictionary mapping Role enums to resource pool IDs for GPU allocation\n    \"\"\"\n\n    def __init__(self):\n        self.role_worker_mapping = {}\n        self.mapping = {}\n\n    def add_actor_rollout_worker(self, config):\n        \"\"\"Add actor rollout worker based on the actor strategy.\"\"\"\n        from verl.single_controller.ray import RayWorkerGroup\n        from verl.trainer.ppo.ray_trainer import Role\n\n        use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n\n        # use new model engine implementation\n        if use_legacy_worker_impl == \"disable\":\n            from verl.workers.engine_workers import ActorRolloutRefWorker\n\n            actor_rollout_cls = ActorRolloutRefWorker\n            ray_worker_group_cls = RayWorkerGroup\n\n            lora_rank = config.actor_rollout_ref.model.get(\"lora\", {}).get(\"rank\", 0)\n            if lora_rank <= 0:\n                lora_rank = config.actor_rollout_ref.model.get(\"lora_rank\", 0)\n            ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get(\"lora_adapter_path\") is not None\n            # NOTE: In new model engine, ref policy and actor rollout are in same ActorRolloutRefWorker,\n            # while in legacy model engine, ref policy is in a separate ActorRolloutRefWorker.\n            if need_reference_policy(config) and not ref_in_actor:\n                role = Role.ActorRolloutRef\n            else:\n                role = Role.ActorRollout\n            self.role_worker_mapping[role] = ray.remote(actor_rollout_cls)\n            self.mapping[role] = \"global_pool\"\n            return actor_rollout_cls, ray_worker_group_cls\n\n        # Note: sync mode validation is now handled in RolloutConfig.__post_init__\n        # Always use async worker since sync mode is deprecated and rejected\n        if config.actor_rollout_ref.actor.strategy in {\"fsdp\", \"fsdp2\"}:\n            from verl.workers.fsdp_workers import AsyncActorRolloutRefWorker\n\n            actor_rollout_cls = AsyncActorRolloutRefWorker\n            ray_worker_group_cls = RayWorkerGroup\n\n        elif config.actor_rollout_ref.actor.strategy == \"megatron\":\n            from verl.workers.megatron_workers import AsyncActorRolloutRefWorker\n\n            actor_rollout_cls = AsyncActorRolloutRefWorker\n            ray_worker_group_cls = RayWorkerGroup\n\n        elif (\n            config.actor_rollout_ref.actor.strategy == \"veomni\"\n            or config.actor_rollout_ref.actor.strategy == \"torchtitan\"\n        ):\n            raise NotImplementedError(\n                f\"{config.actor_rollout_ref.actor.strategy} does not support legacy worker implementation\"\n            )\n\n        else:\n            raise NotImplementedError\n\n        self.role_worker_mapping[Role.ActorRollout] = ray.remote(actor_rollout_cls)\n        self.mapping[Role.ActorRollout] = \"global_pool\"\n        return actor_rollout_cls, ray_worker_group_cls\n\n    def add_critic_worker(self, config):\n        \"\"\"Add critic worker to role mapping.\"\"\"\n        use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n        if config.critic.strategy in {\"fsdp\", \"fsdp2\"}:\n            if use_legacy_worker_impl in [\"auto\", \"enable\"]:\n                from verl.workers.fsdp_workers import CriticWorker\n            elif use_legacy_worker_impl == \"disable\":\n                # we don't need to specialize critic worker. Just use TrainingWorker\n                from verl.workers.engine_workers import TrainingWorker\n\n                CriticWorker = TrainingWorker\n                print(\"Using new worker implementation\")\n            else:\n                raise ValueError(f\"Invalid use_legacy_worker_impl: {use_legacy_worker_impl}\")\n\n        elif config.critic.strategy == \"megatron\":\n            # TODO: switch this to TrainingWorker as well\n            from verl.workers.megatron_workers import CriticWorker\n\n        elif config.critic.strategy == \"veomni\" or config.critic.strategy == \"torchtitan\":\n            if use_legacy_worker_impl == \"disable\":\n                from verl.workers.engine_workers import TrainingWorker\n\n                CriticWorker = TrainingWorker\n                print(f\"Using new worker implementation for {config.critic.strategy}\")\n            else:\n                raise ValueError(\n                    f\"Invalid use_legacy_worker_impl for {config.critic.strategy}: {use_legacy_worker_impl}\"\n                )\n\n        else:\n            raise NotImplementedError\n\n        from verl.trainer.ppo.ray_trainer import Role\n\n        self.role_worker_mapping[Role.Critic] = ray.remote(CriticWorker)\n        self.mapping[Role.Critic] = \"global_pool\"\n\n    def init_resource_pool_mgr(self, config):\n        \"\"\"Initialize resource pool manager.\"\"\"\n\n        global_pool_id = \"global_pool\"\n        resource_pool_spec = {\n            global_pool_id: [config.trainer.n_gpus_per_node] * config.trainer.nnodes,\n        }\n\n        if config.reward.reward_model.enable_resource_pool:\n            if config.reward.reward_model.n_gpus_per_node <= 0:\n                raise ValueError(\"config.reward.reward_model.n_gpus_per_node must be greater than 0\")\n            if config.reward.reward_model.nnodes <= 0:\n                raise ValueError(\"config.reward.reward_model.nnodes must be greater than 0\")\n\n            reward_pool = [config.reward.reward_model.n_gpus_per_node] * config.reward.reward_model.nnodes\n            resource_pool_spec[\"reward_pool\"] = reward_pool\n        else:\n            config.reward.reward_model.nnodes = config.trainer.nnodes\n            config.reward.reward_model.n_gpus_per_node = config.trainer.n_gpus_per_node\n\n        from verl.trainer.ppo.ray_trainer import ResourcePoolManager\n\n        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=self.mapping)\n        return resource_pool_manager\n\n    def add_reward_model_resource_pool(self, config):\n        \"\"\"Add reward model worker if enabled.\"\"\"\n        from verl.trainer.ppo.ray_trainer import Role\n\n        if config.reward.reward_model.enable:\n            # we do not use reward model workers, so we only register reward model in resource pool\n            # without continue to register reward model worker in role mapping\n            if config.reward.reward_model.enable_resource_pool:\n                self.mapping[Role.RewardModel] = \"reward_pool\"\n            else:\n                self.mapping[Role.RewardModel] = \"global_pool\"\n\n    def add_ref_policy_worker(self, config, ref_policy_cls):\n        \"\"\"Add reference policy worker if KL loss or KL reward is used.\"\"\"\n        from verl.trainer.ppo.ray_trainer import Role\n\n        # Ref policy has been fused into ActorRolloutRefWorker in new model engine,\n        # we don't need to add a separate ref policy worker group.\n        use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n        if use_legacy_worker_impl == \"disable\":\n            return\n\n        if need_reference_policy(config):\n            self.role_worker_mapping[Role.RefPolicy] = ray.remote(ref_policy_cls)\n            self.mapping[Role.RefPolicy] = \"global_pool\"\n\n    def run(self, config):\n        \"\"\"Execute the main PPO training workflow.\n\n        This method sets up the distributed training environment, initializes\n        workers, datasets, and reward functions, then starts the training process.\n\n        Args:\n            config: Training configuration object containing all parameters needed\n                   for setting up and running the PPO training process.\n        \"\"\"\n        # Print the initial configuration. `resolve=True` will evaluate symbolic values.\n        from pprint import pprint\n\n        from omegaconf import OmegaConf\n\n        from verl.utils.fs import copy_to_local\n\n        print(f\"TaskRunner hostname: {socket.gethostname()}, PID: {os.getpid()}\")\n        pprint(OmegaConf.to_container(config, resolve=True))\n        OmegaConf.resolve(config)\n\n        actor_rollout_cls, ray_worker_group_cls = self.add_actor_rollout_worker(config)\n        self.add_critic_worker(config)\n\n        self.add_reward_model_resource_pool(config)\n\n        # Add a reference policy worker if KL loss or KL reward is used.\n        self.add_ref_policy_worker(config, actor_rollout_cls)\n\n        # validate config\n        validate_config(\n            config=config,\n            use_reference_policy=need_reference_policy(config),\n            use_critic=need_critic(config),\n        )\n\n        # Download the checkpoint from HDFS to the local machine.\n        # `use_shm` determines whether to use shared memory, which could lead to faster model loading if turned on\n        local_path = copy_to_local(\n            config.actor_rollout_ref.model.path, use_shm=config.actor_rollout_ref.model.get(\"use_shm\", False)\n        )\n\n        # Instantiate the tokenizer and processor.\n        from verl.utils import hf_processor, hf_tokenizer\n\n        trust_remote_code = config.data.get(\"trust_remote_code\", False)\n        tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)\n        # Used for multimodal LLM, could be None\n        processor = hf_processor(local_path, trust_remote_code=trust_remote_code, use_fast=True)\n\n        resource_pool_manager = self.init_resource_pool_mgr(config)\n\n        from verl.utils.dataset.rl_dataset import collate_fn\n\n        # Create training and validation datasets.\n        train_dataset = create_rl_dataset(\n            config.data.train_files,\n            config.data,\n            tokenizer,\n            processor,\n            is_train=True,\n            max_samples=config.data.get(\"train_max_samples\", -1),\n        )\n        val_dataset = create_rl_dataset(\n            config.data.val_files,\n            config.data,\n            tokenizer,\n            processor,\n            is_train=False,\n            max_samples=config.data.get(\"val_max_samples\", -1),\n        )\n        train_sampler = create_rl_sampler(config.data, train_dataset)\n\n        # Initialize the PPO trainer.\n        trainer = RayPPOTrainer(\n            config=config,\n            tokenizer=tokenizer,\n            processor=processor,\n            role_worker_mapping=self.role_worker_mapping,\n            resource_pool_manager=resource_pool_manager,\n            ray_worker_group_cls=ray_worker_group_cls,\n            train_dataset=train_dataset,\n            val_dataset=val_dataset,\n            collate_fn=collate_fn,\n            train_sampler=train_sampler,\n        )\n        # Initialize the workers of the trainer.\n        trainer.init_workers()\n\n        # Start the training process.\n        trainer.fit()\n\n\ndef create_rl_dataset(data_paths, data_config, tokenizer, processor, is_train=True, max_samples: int = -1):\n    \"\"\"Create a dataset.\n\n    Arguments:\n        data_paths: List of paths to data files.\n        data_config: The data config.\n        tokenizer (Tokenizer): The tokenizer.\n        processor (Processor): The processor.\n\n    Returns:\n        dataset (Dataset): The dataset.\n    \"\"\"\n\n    from verl.utils.dataset.rl_dataset import get_dataset_class\n\n    # Get the dataset class\n    dataset_cls = get_dataset_class(data_config)\n\n    # Instantiate the dataset using the determined dataset class\n    dataset = dataset_cls(\n        data_files=data_paths,\n        tokenizer=tokenizer,\n        processor=processor,\n        config=data_config,\n        max_samples=max_samples,\n    )\n\n    return dataset\n\n\ndef create_rl_sampler(data_config, dataset):\n    \"\"\"Create a sampler for the dataset.\n\n    Arguments:\n        data_config: The data config.\n        dataset (Dataset): The dataset.\n\n    Returns:\n        sampler (Sampler): The sampler.\n    \"\"\"\n    import torch\n    from torch.utils.data import SequentialSampler\n\n    # torch.utils.data.RandomSampler could not recover properly\n    from torchdata.stateful_dataloader.sampler import RandomSampler\n\n    if data_config.sampler is not None and data_config.sampler.get(\"class_path\", None) is not None:\n        curriculum_class = load_extern_object(\n            data_config.sampler.class_path,\n            data_config.sampler.class_name,\n        )\n        sampler = curriculum_class(\n            data_source=dataset,\n            data_config=data_config,\n        )\n        assert isinstance(sampler, AbstractSampler)\n        assert data_config.get(\"dataloader_num_workers\", 8) == 0, (\n            \"If using curriculum, num_workers must be 0 to prevent data caching. \"\n            \"If the dataloader caches data before the batch is done the \"\n            \"curriculum sampler won't have the opportunity to reorder it. \"\n        )\n\n    # Use a sampler to facilitate checkpoint resumption.\n    # If shuffling is enabled in the data configuration, create a random sampler.\n    elif data_config.shuffle:\n        train_dataloader_generator = torch.Generator()\n        seed = data_config.get(\"seed\")\n        if seed is not None:\n            train_dataloader_generator.manual_seed(seed)\n        sampler = RandomSampler(data_source=dataset, generator=train_dataloader_generator)\n    else:\n        # If shuffling is disabled, use a sequential sampler to iterate through the dataset in order.\n        sampler = SequentialSampler(data_source=dataset)\n\n    return sampler\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/trainer/ppo/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/trainer/ppo/core_algos.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 The HuggingFace Team. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nCore functions to implement PPO algorithms.\nThe function implemented in this file should be used by trainer with different distributed strategies to\nimplement PPO-like algorithms.\n\"\"\"\n\n__all__ = [\"register_adv_est\", \"get_adv_estimator_fn\", \"AdvantageEstimator\"]\n\nfrom collections import defaultdict\nfrom enum import Enum\nfrom typing import Any, Callable, Optional\n\nimport numpy as np\nimport torch\nfrom omegaconf import DictConfig\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.trainer.config import AlgoConfig\nfrom verl.utils import as_torch_index, group_mean_std\nfrom verl.utils.import_utils import deprecated\nfrom verl.workers.config import ActorConfig\n\nPolicyLossFn = Callable[\n    [\n        torch.Tensor,  # old_log_prob\n        torch.Tensor,  # log_prob\n        torch.Tensor,  # advantages\n        torch.Tensor,  # response_mask\n        str,  # loss_agg_mode\n        Optional[DictConfig | ActorConfig],  # config\n        torch.Tensor | None,  # rollout_log_probs\n    ],\n    tuple[torch.Tensor, dict[str, Any]],\n]\n\nPOLICY_LOSS_REGISTRY: dict[str, PolicyLossFn] = {}\n\n\ndef register_policy_loss(name: str) -> Callable[[PolicyLossFn], PolicyLossFn]:\n    \"\"\"Register a policy loss function with the given name.\n\n    Args:\n        name (str): The name to register the policy loss function under.\n\n    Returns:\n        function: Decorator function that registers the policy loss function.\n    \"\"\"\n\n    def decorator(func: PolicyLossFn) -> PolicyLossFn:\n        POLICY_LOSS_REGISTRY[name] = func\n        return func\n\n    return decorator\n\n\ndef get_policy_loss_fn(name):\n    \"\"\"Get the policy loss with a given name.\n\n    Args:\n        name: `(str)`\n            The name of the policy loss.\n\n    Returns:\n        `(callable)`: The policy loss function.\n    \"\"\"\n    loss_name = name\n    if loss_name not in POLICY_LOSS_REGISTRY:\n        raise ValueError(\n            f\"Unsupported loss mode: {loss_name}. Supported modes are: {list(POLICY_LOSS_REGISTRY.keys())}\"\n        )\n    return POLICY_LOSS_REGISTRY[loss_name]\n\n\nclass AdvantageEstimator(str, Enum):\n    \"\"\"Using an enumeration class to avoid spelling errors in adv_estimator.\n\n    Note(haibin.lin): this enum class is immutable after creation. Extending this\n    enum for new estimators may not be necessary since users can always just call\n    `verl.trainer.ppo.core_algos.register` with string name for a custom advantage\n    estimator instead.\n    \"\"\"\n\n    GAE = \"gae\"\n    GRPO = \"grpo\"\n    REINFORCE_PLUS_PLUS = \"reinforce_plus_plus\"\n    REINFORCE_PLUS_PLUS_BASELINE = \"reinforce_plus_plus_baseline\"\n    REMAX = \"remax\"\n    RLOO = \"rloo\"\n    OPO = \"opo\"\n    GRPO_PASSK = \"grpo_passk\"\n    GPG = \"gpg\"\n    RLOO_VECTORIZED = \"rloo_vectorized\"\n    GRPO_VECTORIZED = \"grpo_vectorized\"\n    OPTIMAL_TOKEN_BASELINE = \"optimal_token_baseline\"\n    TIR_OPTIMAL_TOKEN_BASELINE = \"tir_optimal_token_baseline\"\n    GDPO = \"gdpo\"\n\n\nADV_ESTIMATOR_REGISTRY: dict[str, Any] = {}\n\n\ndef register_adv_est(name_or_enum: str | AdvantageEstimator) -> Any:\n    \"\"\"Decorator to register a advantage estimator function with a given name.\n\n    Args:\n        name_or_enum: `(str)` or `(AdvantageEstimator)`\n            The name or enum of the advantage estimator.\n\n    \"\"\"\n\n    def decorator(fn):\n        name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum\n        if name in ADV_ESTIMATOR_REGISTRY and ADV_ESTIMATOR_REGISTRY[name] != fn:\n            raise ValueError(\n                f\"Adv estimator {name} has already been registered: {ADV_ESTIMATOR_REGISTRY[name]} vs {fn}\"\n            )\n        ADV_ESTIMATOR_REGISTRY[name] = fn\n        return fn\n\n    return decorator\n\n\ndef get_adv_estimator_fn(name_or_enum):\n    \"\"\"Get the advantage estimator function with a given name.\n\n    Args:\n        name_or_enum: `(str)` or `(AdvantageEstimator)`\n            The name or enum of the advantage estimator.\n\n    Returns:\n        `(callable)`: The advantage estimator function.\n    \"\"\"\n    name = name_or_enum.value if isinstance(name_or_enum, Enum) else name_or_enum\n    if name not in ADV_ESTIMATOR_REGISTRY:\n        raise ValueError(f\"Unknown advantage estimator simply: {name}\")\n    return ADV_ESTIMATOR_REGISTRY[name]\n\n\nclass AdaptiveKLController:\n    \"\"\"\n    Adaptive KL controller described in the paper:\n    https://arxiv.org/pdf/1909.08593.pdf\n    \"\"\"\n\n    def __init__(self, init_kl_coef, target_kl, horizon):\n        self.value = init_kl_coef\n        self.target = target_kl\n        self.horizon = horizon\n\n    def update(self, current_kl, n_steps):\n        \"\"\"Update the KL coefficient based on current KL divergence.\n\n        Args:\n            current_kl (float): Current KL divergence value.\n            n_steps (int): Number of steps taken.\n        \"\"\"\n        target = self.target\n        proportional_error = np.clip(current_kl / target - 1, -0.2, 0.2)\n        mult = 1 + proportional_error * n_steps / self.horizon\n        self.value *= mult\n\n\nclass FixedKLController:\n    \"\"\"Fixed KL controller.\"\"\"\n\n    def __init__(self, kl_coef):\n        self.value = kl_coef\n\n    def update(self, current_kl, n_steps):\n        \"\"\"Update method for fixed KL controller (no-op).\n\n        Args:\n            current_kl (float): Current KL divergence value (unused).\n            n_steps (int): Number of steps taken (unused).\n        \"\"\"\n        pass\n\n\ndef get_kl_controller(kl_ctrl):\n    \"\"\"Factory function to create appropriate KL controller based on configuration.\n\n    Args:\n        kl_ctrl: Configuration object containing KL controller settings.\n\n    Returns:\n        KL controller instance (FixedKLController or AdaptiveKLController).\n\n    Raises:\n        NotImplementedError: If controller type is not supported.\n        AssertionError: If adaptive controller horizon is not positive.\n    \"\"\"\n    if kl_ctrl.type == \"fixed\":\n        return FixedKLController(kl_coef=kl_ctrl.kl_coef)\n    elif kl_ctrl.type == \"adaptive\":\n        assert kl_ctrl.horizon > 0, f\"horizon must be larger than 0. Got {kl_ctrl.horizon}\"\n        return AdaptiveKLController(init_kl_coef=kl_ctrl.kl_coef, target_kl=kl_ctrl.target_kl, horizon=kl_ctrl.horizon)\n    else:\n        raise NotImplementedError\n\n\n@register_adv_est(AdvantageEstimator.GAE)  # or simply: @register_adv_est(\"gae\")\ndef compute_gae_advantage_return(\n    token_level_rewards: torch.Tensor,\n    values: torch.Tensor,\n    response_mask: torch.Tensor,\n    gamma: torch.Tensor,\n    lam: torch.Tensor,\n):\n    \"\"\"Adapted from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape is (bs, response_length)\n        values: `(torch.Tensor)`\n            shape is (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape is (bs, response_length). [EOS] mask. The token after [EOS] have mask zero.\n        gamma is `(float)`\n            discounted factor used in RL\n        lam: `(float)`\n            lambda value when computing Generalized Advantage Estimation (https://arxiv.org/abs/1506.02438)\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n\n    \"\"\"\n    with torch.no_grad():\n        nextvalues = 0\n        lastgaelam = 0\n        advantages_reversed = []\n        gen_len = token_level_rewards.shape[-1]\n\n        for t in reversed(range(gen_len)):\n            delta = token_level_rewards[:, t] + gamma * nextvalues - values[:, t]\n            lastgaelam_ = delta + gamma * lam * lastgaelam\n\n            # skip values and TD-error on observation tokens\n            nextvalues = values[:, t] * response_mask[:, t] + (1 - response_mask[:, t]) * nextvalues\n            lastgaelam = lastgaelam_ * response_mask[:, t] + (1 - response_mask[:, t]) * lastgaelam\n\n            advantages_reversed.append(lastgaelam)\n        advantages = torch.stack(advantages_reversed[::-1], dim=1)\n\n        returns = advantages + values\n        advantages = verl_F.masked_whiten(advantages, response_mask)\n    return advantages, returns\n\n\n# NOTE(sgm): this implementation only consider outcome supervision, where the reward is a scalar.\n@register_adv_est(AdvantageEstimator.GRPO)  # or simply: @register_adv_est(\"grpo\")\ndef compute_grpo_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    norm_adv_by_std_in_grpo: bool = True,\n    config: Optional[AlgoConfig] = None,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for GRPO, operating only on Outcome reward\n    (with only one scalar reward for each response).\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape is (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape is (bs, response_length)\n        index: `(np.ndarray)`\n            index array for grouping\n        epsilon: `(float)`\n            small value to avoid division by zero\n        norm_adv_by_std_in_grpo: `(bool)`\n            whether to scale the GRPO advantage\n        config: `(Optional[AlgoConfig])`\n            algorithm configuration object\n\n    Note:\n        If norm_adv_by_std_in_grpo is True, the advantage is scaled by the std, as in the original GRPO.\n        If False, the advantage is not scaled, as in Dr.GRPO (https://arxiv.org/abs/2503.20783).\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape is (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape is (bs, response_length)\n    \"\"\"\n    scores = token_level_rewards.sum(dim=-1)\n\n    id2score = defaultdict(list)\n    id2mean = {}\n    id2std = {}\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        for i in range(bsz):\n            id2score[index[i]].append(scores[i])\n        for idx in id2score:\n            if len(id2score[idx]) == 1:\n                id2mean[idx] = torch.tensor(0.0)\n                id2std[idx] = torch.tensor(1.0)\n            elif len(id2score[idx]) > 1:\n                scores_tensor = torch.stack(id2score[idx])\n                id2mean[idx] = torch.mean(scores_tensor)\n                id2std[idx] = torch.std(scores_tensor)\n            else:\n                raise ValueError(f\"no score in prompt index: {idx}\")\n        for i in range(bsz):\n            if norm_adv_by_std_in_grpo:\n                scores[i] = (scores[i] - id2mean[index[i]]) / (id2std[index[i]] + epsilon)\n            else:\n                scores[i] = scores[i] - id2mean[index[i]]\n        scores = scores.unsqueeze(-1) * response_mask\n\n    return scores, scores\n\n\n@register_adv_est(AdvantageEstimator.GRPO_VECTORIZED)\ndef compute_grpo_vectorized_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    norm_adv_by_std_in_grpo: bool = True,\n    config: Optional[AlgoConfig] = None,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Vectorized GRPO（outcome-only）:\n      For each group g:\n      a_i = \\\\frac{r_i - \\\\mu_g}{\\\\sigma_g} (or without dividing by \\\\sigma_g),\n      then broadcast the scalar across the token dimension (multiplied by response_mask).。\n    \"\"\"\n    with torch.no_grad():\n        scores = token_level_rewards.sum(dim=-1)\n        g = as_torch_index(index, device=scores.device)\n        mean_g, std_g, _ = group_mean_std(scores, g, eps=epsilon, device=scores.device)\n        if norm_adv_by_std_in_grpo:\n            scalars = (scores - mean_g[g]) / (std_g[g] + epsilon)\n        else:\n            scalars = scores - mean_g[g]\n        advantages = scalars.unsqueeze(-1) * response_mask\n        return advantages, advantages\n\n\n@register_adv_est(AdvantageEstimator.GDPO)  # or simply: @register_adv_est(\"gdpo\")\ndef compute_gdpo_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    norm_adv_by_std_in_grpo: bool = True,\n    config: Optional[AlgoConfig] = None,\n    non_tensor_batch: Optional[dict] = None,\n    batch: Optional[dict] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    GDPO: Group reward-Decoupled Normalization Policy Optimization.\n\n    Instead of summing all reward dimensions first (like GRPO), GDPO normalizes\n    each reward dimension independently within each group before aggregation.\n    This prevents a dominant reward signal from drowning out weaker ones.\n\n    Mathematical formulation:\n        Step 1 – Group-wise decoupled normalization (via GRPO per dimension):\n            For each reward dimension k, within each group g:\n            A_k = (r_k - μ_group(r_k)) / (σ_group(r_k) + ε)\n\n        Step 2 – Weighted aggregation:\n            A_sum = Σ_k w_k · A_k\n\n        Step 3 – Batch-level normalization (via masked_whiten):\n            A_final = whiten(A_sum, response_mask)\n\n    Args:\n        token_level_rewards: (bs, response_length) – standard token-level rewards.\n            Used as fallback when per-dimension rewards are not provided.\n        response_mask: (bs, response_length)\n        index: (bs,) – group id per sample (from ``uid``).\n        epsilon: Numerical stability constant.\n        norm_adv_by_std_in_grpo: Whether to normalize by std in GRPO.\n        config: Algorithm configuration (optional).\n        non_tensor_batch: Non-tensor batch data containing per-dimension reward scores.\n        batch: Batch data containing prompts, attention_mask, etc.\n\n    Note:\n        Ref GDPO (https://arxiv.org/abs/2601.05242).\n\n    Returns:\n        advantages: (bs, response_length)\n        returns: (bs, response_length) – same as advantages (outcome-only).\n    \"\"\"\n    score_list = None\n    reward_weights = None\n\n    if config is not None and non_tensor_batch is not None and batch is not None:\n        gdpo_reward_keys = config.get(\"gdpo_reward_keys\", None)\n        assert gdpo_reward_keys, (\n            \"GDPO requires 'algorithm.gdpo_reward_keys' listing the individual reward \"\n            \"component keys returned by compute_score (e.g. ['format_reward', 'accuracy_reward']).\"\n        )\n        device = token_level_rewards.device\n        prompt_length = batch[\"prompts\"].size(1)\n        valid_response_length = batch[\"attention_mask\"][:, prompt_length:].sum(dim=1) - 1\n\n        score_list = []\n        for key in gdpo_reward_keys:\n            assert key in non_tensor_batch, (\n                f\"GDPO reward key '{key}' not found in non_tensor_batch. \"\n                f\"Available keys: {list(non_tensor_batch.keys())}. \"\n                f\"Make sure your compute_score returns a dict containing '{key}'.\"\n            )\n            comp = non_tensor_batch[key]\n            rm_score = torch.tensor(np.asarray(comp, dtype=np.float32), device=device)\n            rm_scores = torch.zeros_like(response_mask, dtype=torch.float32)\n            rm_scores[torch.arange(rm_scores.size(0), device=device), valid_response_length] = rm_score\n            score_list.append(rm_scores)\n\n        gdpo_weights = config.get(\"gdpo_reward_weights\", None)\n        if gdpo_weights is not None:\n            reward_weights = list(gdpo_weights)\n\n    if score_list is None:\n        score_list = [token_level_rewards]\n\n    num_scores = len(score_list)\n\n    if reward_weights is not None:\n        weights = torch.tensor(reward_weights, dtype=torch.float32, device=token_level_rewards.device)\n    else:\n        weights = torch.ones(num_scores, dtype=torch.float32, device=token_level_rewards.device)\n\n    new_advantage = None\n\n    for i in range(num_scores):\n        normalized_score, _ = compute_grpo_outcome_advantage(\n            token_level_rewards=score_list[i],\n            response_mask=response_mask,\n            index=index,\n            epsilon=epsilon,\n            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n            config=config,\n        )\n\n        if new_advantage is None:\n            new_advantage = weights[i] * normalized_score\n        else:\n            new_advantage += weights[i] * normalized_score\n\n    advantages = verl_F.masked_whiten(new_advantage, response_mask) * response_mask\n\n    return advantages, advantages\n\n\n@register_adv_est(AdvantageEstimator.GRPO_PASSK)  # or simply: @register_adv_est(\"grpo_passk\")\ndef compute_grpo_passk_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    norm_adv_by_std_in_grpo: bool = True,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for Pass@k using a GRPO-style outcome reward formulation.\n    Only the best response per group gets a non-zero advantage: r_max - r_second_max.\n\n    Implemented as described in https://arxiv.org/abs/2503.19595.\n\n    Args:\n        token_level_rewards: (bs, response_length)\n        response_mask: (bs, response_length)\n        index: (bs,) → group ID per sample\n        epsilon: float for numerical stability\n        config: (AlgoConfig) algorithm settings, which contains \"norm_adv_by_std_in_grpo\"\n\n    Returns:\n        advantages: (bs, response_length)\n        returns: (bs, response_length)\n    \"\"\"\n    assert config is not None\n    # if True, normalize advantage by std within group\n    norm_adv_by_std_in_grpo = config.get(\"norm_adv_by_std_in_grpo\", True)\n    scores = token_level_rewards.sum(dim=-1)  # (bs,)\n    advantages = torch.zeros_like(scores)\n\n    id2scores = defaultdict(list)\n    id2indices = defaultdict(list)\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        for i in range(bsz):\n            idx = index[i]\n            id2scores[idx].append(scores[i])\n            id2indices[idx].append(i)\n\n        for idx in id2scores:\n            rewards = torch.stack(id2scores[idx])  # (k,)\n            if rewards.numel() < 2:\n                raise ValueError(\n                    f\"Pass@k requires at least 2 samples per group. Got {rewards.numel()} for group {idx}.\"\n                )\n            topk, topk_idx = torch.topk(rewards, 2)\n            r_max, r_second_max = topk[0], topk[1]\n            i_max = id2indices[idx][topk_idx[0].item()]\n            advantage = r_max - r_second_max\n            if norm_adv_by_std_in_grpo:\n                std = torch.std(rewards)\n                advantage = advantage / (std + epsilon)\n            advantages[i_max] = advantage\n\n    advantages = advantages.unsqueeze(-1) * response_mask\n    return advantages, advantages\n\n\n@register_adv_est(\n    AdvantageEstimator.REINFORCE_PLUS_PLUS_BASELINE\n)  # or simply: @register_adv_est(\"reinforce_plus_plus_baseline\")\ndef compute_reinforce_plus_plus_baseline_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: torch.Tensor,\n    epsilon: float = 1e-6,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for RF++-baseline (https://arxiv.org/abs/2501.03262), operating only on Outcome reward\n    (with only one scalar reward for each response).\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    response_length = token_level_rewards.shape[-1]\n    scores = token_level_rewards.sum(dim=-1)\n\n    id2score = defaultdict(list)\n    id2mean = {}\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        for i in range(bsz):\n            id2score[index[i]].append(scores[i])\n        for idx in id2score:\n            if len(id2score[idx]) == 1:\n                id2mean[idx] = torch.tensor(0.0)\n            elif len(id2score[idx]) > 1:\n                id2mean[idx] = torch.mean(torch.stack(id2score[idx]))\n            else:\n                raise ValueError(f\"no score in prompt index: {idx}\")\n        for i in range(bsz):\n            scores[i] = scores[i] - id2mean[index[i]]\n\n        scores = scores.unsqueeze(-1).tile([1, response_length]) * response_mask\n        scores = verl_F.masked_whiten(scores, response_mask) * response_mask\n\n    return scores, scores\n\n\n@register_adv_est(AdvantageEstimator.RLOO)  # or simply: @register_adv_est(\"rloo\")\ndef compute_rloo_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    scores = token_level_rewards.sum(dim=-1)\n\n    id2score = defaultdict(list)\n    id2mean = {}\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        for i in range(bsz):\n            id2score[index[i]].append(scores[i])\n        for idx in id2score:\n            if len(id2score[idx]) == 1:\n                id2mean[idx] = torch.tensor(0.0)\n            elif len(id2score[idx]) > 1:\n                id2mean[idx] = torch.mean(torch.stack(id2score[idx]))\n            else:\n                raise ValueError(f\"no score in prompt index: {idx}\")\n        for i in range(bsz):\n            response_num = len(id2score[index[i]])\n            if response_num > 1:\n                scores[i] = scores[i] * response_num / (response_num - 1) - id2mean[index[i]] * response_num / (\n                    response_num - 1\n                )\n        scores = scores.unsqueeze(-1) * response_mask\n\n    return scores, scores\n\n\n@register_adv_est(AdvantageEstimator.OPO)  # or simply: @register_adv_est(\"opo\")\ndef compute_opo_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for OPO based on https://arxiv.org/pdf/2505.23585\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    response_length = response_mask.sum(dim=-1)\n    scores = token_level_rewards.sum(dim=-1)\n\n    id2score = defaultdict(list)\n    id2len = defaultdict(list)\n    id2bsl = {}\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        for i in range(bsz):\n            id2score[index[i]].append(scores[i])\n            id2len[index[i]].append(response_length[i])\n\n        for idx in id2score:\n            if len(id2score[idx]) == 1:\n                id2bsl[idx] = torch.tensor(0.0)\n            elif len(id2score[idx]) > 1:\n                score_tensor = torch.stack(id2score[idx])\n                len_tensor = torch.stack(id2len[idx])\n                id2bsl[idx] = (len_tensor * score_tensor).sum() / len_tensor.sum()\n            else:\n                raise ValueError(f\"no score in prompt index: {idx}\")\n        for i in range(bsz):\n            scores[i] = scores[i] - id2bsl[index[i]]\n        scores = scores.unsqueeze(-1) * response_mask\n\n    return scores, scores\n\n\n@register_adv_est(AdvantageEstimator.REINFORCE_PLUS_PLUS)  # or simply: @register_adv_est(\"reinforce_plus_plus\")\ndef compute_reinforce_plus_plus_outcome_advantage(\n    token_level_rewards: torch.Tensor, response_mask: torch.Tensor, config: Optional[AlgoConfig] = None, **kwargs\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for REINFORCE++.\n    This implementation is based on the paper: https://arxiv.org/abs/2501.03262\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    assert config is not None\n    gamma = config.gamma\n    with torch.no_grad():\n        returns = torch.zeros_like(token_level_rewards)\n        running_return = 0\n\n        for t in reversed(range(token_level_rewards.shape[1])):\n            running_return = token_level_rewards[:, t] + gamma * running_return\n            returns[:, t] = running_return\n            # Reset after EOS\n            running_return = running_return * response_mask[:, t]\n\n        advantages = verl_F.masked_whiten(returns, response_mask)\n        advantages = advantages * response_mask\n\n    return advantages, returns\n\n\n@register_adv_est(AdvantageEstimator.REMAX)  # or simply: @register_adv_est(\"remax\")\ndef compute_remax_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    reward_baselines: torch.Tensor,\n    response_mask: torch.Tensor,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for ReMax, operating only on Outcome reward\n    This implementation is based on the paper: https://arxiv.org/abs/2310.10505\n    (with only one scalar reward for each response).\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        reward_baselines: `(torch.Tensor)`\n            shape: (bs,)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n\n    with torch.no_grad():\n        returns = (token_level_rewards * response_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1])\n        advantages = returns - reward_baselines.unsqueeze(-1) * response_mask\n\n    return advantages, returns\n\n\n@register_adv_est(AdvantageEstimator.GPG)  # or simply: @register_adv_est(\"gpg\")\ndef compute_gpg_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    f_norm: float = 1.0,\n    alpha: float = 1.0,\n    config=None,\n    **kwargs,\n):\n    \"\"\"\n    Compute advantage for GPG, operating only on Outcome reward\n    (with only one scalar reward for each response).\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        index: `(np.ndarray)`\n            shape: (bs,)\n        epsilon: (float)\n        f_norm: (float)\n        alpha: (float)\n        config: (dict) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    scores = token_level_rewards.sum(dim=-1)\n\n    id2score = defaultdict(list)\n    id2mean = {}\n    id2std = {}\n\n    with torch.no_grad():\n        bsz = scores.shape[0]\n        m = torch.count_nonzero(scores)\n        alpha = bsz / m.clamp(min=1)\n\n        for i in range(bsz):\n            id2score[index[i]].append(scores[i])\n\n        for idx in id2score:\n            if len(id2score[idx]) == 1:\n                id2mean[idx] = torch.tensor(0.0)\n                id2std[idx] = torch.tensor(1.0)\n            elif len(id2score[idx]) > 1:\n                scores_tensor = torch.stack(id2score[idx])\n                id2mean[idx] = torch.mean(scores_tensor)\n                id2std[idx] = torch.std(scores_tensor)\n            else:\n                raise ValueError(f\"no score in prompt index: {idx}\")\n        for i in range(bsz):\n            scores[i] = alpha * (scores[i] - id2mean[index[i]]) / (f_norm)\n        scores = scores.unsqueeze(-1) * response_mask\n\n    return scores, scores\n\n\n@register_adv_est(AdvantageEstimator.RLOO_VECTORIZED)  # or simply: @register_adv_est(\"rloo_vectorized\")\ndef compute_rloo_vectorized_outcome_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    epsilon: float = 1e-6,\n    config: Optional[AlgoConfig] = None,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantage for RLOO based on https://arxiv.org/abs/2402.14740\n\n    Args:\n        token_level_rewards: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n        config: (AlgoConfig) algorithm config\n\n    Returns:\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        Returns: `(torch.Tensor)`\n            shape: (bs, response_length)\n    \"\"\"\n    scores = token_level_rewards.sum(dim=-1)\n\n    with torch.no_grad():\n        inv = torch.from_numpy(np.unique(index, return_inverse=True)[1]).to(scores.device)\n\n        c = torch.bincount(inv)[inv].to(scores.dtype)\n        adv = ((c * scores - torch.bincount(inv, weights=scores)[inv]) / (c - 1).clamp_min(1)) * (c > 1)\n\n        adv = adv.unsqueeze(-1) * response_mask\n\n    return adv, adv\n\n\n@register_adv_est(AdvantageEstimator.OPTIMAL_TOKEN_BASELINE)\ndef compute_optimal_token_baseline_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    old_log_probs: torch.Tensor,\n    sum_pi_squared: torch.Tensor,\n    rollout_is_weights: torch.Tensor = None,\n    handle_zero_tail: bool = True,\n    epsilon: float = 1e-8,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantages using Optimal Token Baseline (OTB).\n\n    Unlike the group mean based baseline which uses a single baseline per trajectory,\n    this computes a unique baseline for each timestep using cumulative path variance.\n\n    Theory:\n        For each timestep t in each prompt group:\n            B_t* = E[G_t × W_t] / E[W_t]\n        where W_t = Σ_{j=1}^t ||s_j||² (cumulative path-variance proxy)\n        and ||s_j||² = 1 - 2π_j + Σπ²\n\n    The cumulative sum W_t captures the \"realized energy\" of trajectory has been up to timestep t,\n    giving higher weight to predicting rewards on high-variance paths.\n\n    Args:\n        token_level_rewards: Rewards at each token position [shape: (bs, response_length)]\n        response_mask: Binary mask for valid tokens (1) vs padding (0) [shape: (bs, response_length)]\n        index: Prompt indices for grouping trajectories from same prompt [shape: (bs,)]\n        old_log_probs: Log probabilities from training policy during generation [shape: (bs, response_length)]\n        sum_pi_squared: Sum of squared probabilities over vocabulary Σπ² [shape: (bs, response_length)]\n        rollout_is_weights: Pre-computed IS weights for W correction [shape: (bs, response_length)],\n            None if not using IS\n        handle_zero_tail: If True, zero baselines will be set in the portion of the longest trajectory\n            that extends beyond the second-longest trajectory in the prompt group.\n            Default: True\n        epsilon: Small constant for numerical stability (default: 1e-8)\n\n    Returns:\n        advantages: OTB advantage estimates [shape: (bs, response_length)]\n        returns: Cumulative rewards (returns) from each position [shape: (bs, response_length)]\n\n    Note on Rollout Importance Sampling:\n        When rollout_is_weights is provided, W_t is scaled by ρ̄²(t) to minimize MSE under truncated IS:\n            B_t* = Σ[G_t × ρ̄²(t) × W_t] / Σ[ρ̄²(t) × W_t]\n    \"\"\"\n    with torch.no_grad():\n        batch_size, seq_len = token_level_rewards.shape\n        device = token_level_rewards.device\n\n        # Compute returns (reward-to-go) for each timestep\n        returns = (token_level_rewards * response_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1])\n\n        # Step 1: Compute w_per_timestep = 1 - 2π_t + Σπ²)\n        pi_t = torch.exp(old_log_probs)\n        w_per_timestep = 1 - 2 * pi_t + sum_pi_squared\n\n        # Step 2: Apply rollout importance sampling correction (if enabled)\n        if rollout_is_weights is not None:\n            # Scale W by ρ̄² to minimize MSE under truncated IS\n            w_per_timestep = w_per_timestep * (rollout_is_weights**2)\n\n        # Step 3: Compute cumulative path-variance proxy: W_t = Σ_{j=1}^t w_j\n        # This measures accumulated variance from the start of the trajectory up to timestep t\n        w_cumulative = (w_per_timestep * response_mask).cumsum(dim=-1)\n\n        # Group trajectories by prompt\n        prompt_groups = defaultdict(list)\n        for i in range(batch_size):\n            prompt_groups[index[i]].append(i)\n\n        # Initialize baselines tensor [batch_size, seq_len]\n        baselines = torch.zeros_like(returns)\n\n        # Compute per-step baseline for each prompt group\n        for _, trajectory_indices in prompt_groups.items():\n            N = len(trajectory_indices)\n            if N == 1:\n                # Single trajectory - no baseline (advantage = return)\n                continue\n\n            traj_idx = torch.tensor(trajectory_indices, device=device)\n\n            # Extract group data [N, seq_len]\n            returns_group = returns[traj_idx]\n            w_cumulative_group = w_cumulative[traj_idx]\n            mask_group = response_mask[traj_idx]\n\n            # Compute per-timestep baseline: B_t = Σ[G_t × W_t] / Σ[W_t]\n            # where W_t = Σ_{j=1}^t ||s_j||² (cumulative path variance)\n            # Shape: [seq_len]\n            numerator = (returns_group * w_cumulative_group * mask_group).sum(dim=0)  # Sum over trajectories\n            denominator = (w_cumulative_group * mask_group).sum(dim=0) + epsilon\n\n            baseline_per_step = numerator / denominator  # [seq_len]\n\n            # Assign to all trajectories in this group\n            baselines[traj_idx] = baseline_per_step.unsqueeze(0).expand(N, -1)\n\n            if handle_zero_tail:\n                # Optionally zero out the portion of the longest trajectory that extends\n                # beyond the second-longest trajectory in the prompt group.\n                response_lengths = mask_group.sum(dim=-1)\n                sorted_lengths, _ = torch.sort(response_lengths)\n                max_length = int(sorted_lengths[-1].item())\n                second_max_length = int(sorted_lengths[-2].item())\n                max_length_idx = (response_lengths == max_length).nonzero(as_tuple=True)[0]\n                if max_length_idx.numel() == 1 and max_length > second_max_length:\n                    max_length_traj_idx = trajectory_indices[int(max_length_idx[0])]\n                    baselines[max_length_traj_idx, second_max_length:] = 0.0\n\n        # Compute advantages: A_t = G_t - B_t\n        advantages = (returns - baselines) * response_mask\n\n    return advantages, returns\n\n\n@register_adv_est(AdvantageEstimator.TIR_OPTIMAL_TOKEN_BASELINE)\ndef compute_multi_turn_optimal_token_baseline_advantage(\n    token_level_rewards: torch.Tensor,\n    response_mask: torch.Tensor,\n    index: np.ndarray,\n    old_log_probs: torch.Tensor,\n    sum_pi_squared: torch.Tensor,\n    rollout_is_weights: torch.Tensor = None,\n    handle_zero_tail: bool = True,\n    epsilon: float = 1e-8,\n    **kwargs,\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute advantages using Optimal Token Baseline (OTB).\n\n    Unlike the group mean based baseline which uses a single baseline per trajectory,\n    this computes a unique baseline for each timestep using cumulative path variance.\n\n    Theory:\n        For each timestep t in each prompt group:\n            B_t* = E[G_t × W_t] / E[W_t]\n        where W_t = Σ_{j=1}^t ||s_j||² (cumulative path-variance proxy)\n        and ||s_j||² = 1 - 2π_j + Σπ²\n\n    The cumulative sum W_t captures the \"realized energy\" of trajectory has been up to timestep t,\n    giving higher weight to predicting rewards on high-variance paths.\n\n    Args:\n        token_level_rewards: Rewards at each token position [shape: (bs, response_length)]\n        response_mask: Binary mask for valid tokens (1) vs padding (0) [shape: (bs, response_length)]\n        index: Prompt indices for grouping trajectories from same prompt [shape: (bs,)]\n        old_log_probs: Log probabilities from training policy during generation [shape: (bs, response_length)]\n        sum_pi_squared: Sum of squared probabilities over vocabulary Σπ² [shape: (bs, response_length)]\n        rollout_is_weights: Pre-computed IS weights for W correction [shape: (bs, response_length)],\n            None if not using IS\n        handle_zero_tail: If True, zero baselines will be set in the portion of the longest trajectory\n            that extends beyond the second-longest trajectory in the prompt group.\n            Default: False\n        epsilon: Small constant for numerical stability (default: 1e-8)\n\n    Returns:\n        advantages: OTB advantage estimates [shape: (bs, response_length)]\n        returns: Cumulative rewards (returns) from each position [shape: (bs, response_length)]\n\n    Note on Rollout Importance Sampling:\n        When rollout_is_weights is provided, W_t is scaled by ρ̄²(t) to minimize MSE under truncated IS:\n            B_t* = Σ[G_t × ρ̄²(t) × W_t] / Σ[ρ̄²(t) × W_t]\n    \"\"\"\n    with torch.no_grad():\n        # Compute returns (reward-to-go) for each timestep\n        token_returns = (token_level_rewards * response_mask).flip(dims=[-1]).cumsum(dim=-1).flip(dims=[-1])\n\n        # Step 1: Compute w_per_timestep = 1 - 2π_t + Σπ²)\n        pi_t = torch.exp(old_log_probs)\n        w_per_timestep = 1 - 2 * pi_t + sum_pi_squared\n\n        # Step 2: Apply rollout importance sampling correction (if enabled)\n        if rollout_is_weights is not None:\n            # Scale W by ρ̄² to minimize MSE under truncated IS\n            w_per_timestep = w_per_timestep * (rollout_is_weights**2)\n\n        # Step 3: Compute cumulative path-variance proxy: W_t = Σ_{j=1}^t w_j\n        # This measures accumulated variance from the start of the trajectory up to timestep t\n        w_cumulative = (w_per_timestep * response_mask).cumsum(dim=-1)\n\n        # Step 4: Concatenate returns and w_cumulative for each trajectory\n        # This allows us to compute baseline per timestep for each trajectory\n        response_lengths = response_mask.sum(dim=-1).to(dtype=torch.long)  # [shape: (bs * n, )]\n        max_response_length = int(response_lengths.max().item()) if response_lengths.numel() > 0 else 0\n        all_w_values = w_cumulative.new_zeros(\n            (len(response_lengths), max_response_length)\n        )  # [shape: (bs * n, max_response_length)]\n        all_returns = torch.zeros_like(all_w_values)\n        for i in range(len(response_lengths)):\n            length = int(response_lengths[i].item())\n            if length == 0:\n                continue\n            mask = response_mask[i].bool()\n            all_w_values[i, :length] = w_cumulative[i, mask]\n            all_returns[i, :length] = token_returns[i, mask]\n\n        # Group trajectories by prompt\n        prompt_groups = defaultdict(list)\n        for i in range(len(response_lengths)):\n            if response_lengths[i] == 0:\n                continue\n            prompt_groups[index[i]].append(i)\n\n        # Compute optimal baseline for each prompt group\n        baselines = torch.zeros_like(all_returns)\n\n        for _, trajectory_indices in prompt_groups.items():\n            N = len(trajectory_indices)\n            traj_idx = torch.tensor(trajectory_indices, device=all_returns.device)\n\n            if N == 1:\n                # Single trajectory - no baseline (keep original reward as advantage)\n                baselines[traj_idx[0]] = 0.0\n                continue\n\n            # Extract group data\n            w_group = all_w_values[traj_idx]  # [shape: (N, max_response_length)]\n            R_group = all_returns[traj_idx]  # [shape: (N, max_response_length)]\n            # Direct optimal baseline - single value for all in group\n            b_star = (R_group * w_group).sum(dim=0) / (w_group.sum(dim=0) + epsilon)\n            # Convert to match baselines dtype (epsilon can cause float64 promotion)\n            baselines[traj_idx] = b_star.to(baselines.dtype)\n\n            if handle_zero_tail:\n                # Optionally zero out the portion of the longest trajectory that extends\n                # beyond the second-longest trajectory in the prompt group.\n                response_lengths_group = response_lengths[traj_idx]\n                sorted_lengths, _ = torch.sort(response_lengths_group)\n                max_length = int(sorted_lengths[-1].item())\n                second_max_length = int(sorted_lengths[-2].item())\n                max_length_idx = (response_lengths_group == max_length).nonzero(as_tuple=True)[0]\n                if max_length_idx.numel() == 1 and max_length > second_max_length:\n                    max_length_traj_idx = trajectory_indices[int(max_length_idx[0])]\n                    baselines[max_length_traj_idx, second_max_length:] = 0.0\n\n        # Compute advantages\n        all_advantages = all_returns - baselines  # [shape: (bs * n, max_response_length)]\n\n        advantages = torch.zeros_like(token_returns)  # [shape: (bs * n, turn * response_length)]\n        for i in range(len(response_lengths)):\n            if response_lengths[i] == 0:\n                continue\n            advantages[i, response_mask[i].bool()] = all_advantages[i, : response_lengths[i]]\n\n        advantages = advantages * response_mask  # [shape: (bs * n * turn, response_length)]\n\n    return advantages, token_returns\n\n\ndef compute_rewards(token_level_scores, old_log_prob, ref_log_prob, kl_ratio):\n    \"\"\"Compute token-level rewards with KL penalty.\n\n    Args:\n        token_level_scores (torch.Tensor): Token-level reward scores.\n        old_log_prob (torch.Tensor): Log probabilities from current policy.\n        ref_log_prob (torch.Tensor): Log probabilities from reference policy.\n        kl_ratio (float): KL penalty coefficient.\n\n    Returns:\n        torch.Tensor: Token-level rewards with KL penalty applied.\n    \"\"\"\n    kl = old_log_prob - ref_log_prob\n    return token_level_scores - kl * kl_ratio\n\n\ndef agg_loss(\n    loss_mat: torch.Tensor,\n    loss_mask: torch.Tensor,\n    loss_agg_mode: str,\n    dp_size: int = 1,\n    batch_num_tokens: Optional[int] = None,\n    global_batch_size: Optional[int] = None,\n    loss_scale_factor: Optional[int] = None,\n):\n    \"\"\"\n    Aggregate the loss across global batch to ensure the loss is invariant to fsdp/megatron parallelism.\n\n    NOTE: The returned loss has different behaviors for different backend:\n    - FSDP: the loss is directly used for backward.\n    - Megatron: the loss should be scaled by `num_microbatches` and `cp_size` for pp schedule.\n\n    Args:\n        loss_mat: micro batch loss matrix, (bs, response_length)\n        loss_mask: micro batch loss mask, (bs, response_length)\n        loss_agg_mode: method to aggregate the loss matrix into a scalar\n        dp_size: data parallel size\n        batch_num_tokens: number of valid tokens in global batch\n        global_batch_size: global batch size\n        loss_scale_factor: scale factor for \"seq-mean-token-sum-norm\" mode. If None, uses loss_mask.shape[-1].\n            Set this to a constant value to ensure consistent normalization throughout training.\n\n    Returns:\n        loss: `a scalar torch.Tensor`\n            aggregated loss\n    \"\"\"\n    if loss_agg_mode == \"token-mean\":\n        if batch_num_tokens is None:\n            if dp_size > 1:\n                raise ValueError(\"(global) batch_num_tokens is required when dp_size > 1\")\n            batch_num_tokens = loss_mask.sum()\n        loss = verl_F.masked_sum(loss_mat, loss_mask) / batch_num_tokens * dp_size\n    elif loss_agg_mode in [\"seq-mean-token-sum\", \"seq-mean-token-sum-norm\"]:\n        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1)  # token-sum\n        seq_mask = (torch.sum(loss_mask, dim=-1) > 0).float()  # exclude fully masked sequences\n        if global_batch_size is None:\n            if dp_size > 1:\n                raise ValueError(\"global_batch_size is required when dp_size > 1\")\n            global_batch_size = seq_mask.sum()\n        loss = verl_F.masked_sum(seq_losses, seq_mask) / global_batch_size * dp_size  # seq-mean\n        if loss_agg_mode == \"seq-mean-token-sum-norm\":\n            if loss_scale_factor is None:\n                horizon = loss_mask.shape[-1]\n                loss_scale_factor = horizon\n            loss /= loss_scale_factor\n    elif loss_agg_mode == \"seq-mean-token-mean\":\n        seq_mask = torch.sum(loss_mask, dim=-1)  # per-sequence token count\n        seq_losses = torch.sum(loss_mat * loss_mask, dim=-1) / (seq_mask + 1e-8)  # token-mean\n        seq_mask = (seq_mask > 0).float()  # exclude fully masked sequences\n        if global_batch_size is None:\n            if dp_size > 1:\n                raise ValueError(\"global_batch_size is required when dp_size > 1\")\n            global_batch_size = seq_mask.sum()\n        loss = verl_F.masked_sum(seq_losses, seq_mask) / global_batch_size * dp_size  # seq-mean\n    else:\n        raise ValueError(f\"Invalid loss_agg_mode: {loss_agg_mode}\")\n\n    return loss\n\n\n@deprecated(\"verl.trainer.ppo.core_algos.compute_policy_loss_vanilla\")\ndef compute_policy_loss(\n    old_log_prob,\n    log_prob,\n    advantages,\n    response_mask,\n    cliprange=None,\n    cliprange_low=None,\n    cliprange_high=None,\n    clip_ratio_c=3.0,\n    loss_agg_mode: str = \"token-mean\",\n):\n    \"\"\"\n    Compute the clipped policy objective and related metrics for PPO.\n\n    Adapted from\n    https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        cliprange (float, optional):\n            Clipping parameter ε for standard PPO. See https://arxiv.org/abs/1707.06347.\n            Defaults to None (must be provided).\n        cliprange_low (float, optional):\n            Lower clip range for dual-clip PPO. Defaults to same as `cliprange`.\n        cliprange_high (float, optional):\n            Upper clip range for dual-clip PPO. Defaults to same as `cliprange`.\n        clip_ratio_c (float, optional):\n            Lower bound of the ratio for dual-clip PPO. See https://arxiv.org/pdf/1912.09729.\n            Defaults to 3.0.\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n    \"\"\"\n    assert clip_ratio_c > 1.0, (\n        \"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0,\"\n        + f\" but get the value: {clip_ratio_c}.\"\n    )\n\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    pg_losses1 = -advantages * ratio\n    if cliprange_low is None:\n        cliprange_low = cliprange\n    if cliprange_high is None:\n        cliprange_high = cliprange\n    pg_losses2 = -advantages * torch.clamp(\n        ratio, 1 - cliprange_low, 1 + cliprange_high\n    )  # - clip(ratio, 1-cliprange, 1+cliprange) * A\n    clip_pg_losses1 = torch.maximum(\n        pg_losses1, pg_losses2\n    )  # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A)\n    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)\n\n    pg_losses3 = -advantages * clip_ratio_c\n    clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)\n    pg_clipfrac_lower = verl_F.masked_mean(\n        torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), response_mask\n    )\n\n    pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)\n    pg_loss = agg_loss(loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n\n    return pg_loss, pg_clipfrac, ppo_kl, pg_clipfrac_lower\n\n\n@register_policy_loss(\"vanilla\")  # type: ignore[arg-type]\ndef compute_policy_loss_vanilla(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for PPO.\n\n    Adapted from\n    https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1122\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n        config: `(verl.trainer.config.ActorConfig)`:\n            config for the actor.\n        rollout_log_probs: `(torch.Tensor)`:\n            log probabilities of actions under the rollout policy, shape (batch_size, response_length).\n    \"\"\"\n\n    assert config is not None\n    assert not isinstance(config, AlgoConfig)\n    clip_ratio = config.clip_ratio  # Clipping parameter ε for standard PPO. See https://arxiv.org/abs/1707.06347.\n    clip_ratio_low = config.clip_ratio_low if config.clip_ratio_low is not None else clip_ratio\n    clip_ratio_high = config.clip_ratio_high if config.clip_ratio_high is not None else clip_ratio\n    clip_ratio_c = config.get(  # Lower bound of the ratio for dual-clip PPO. See https://arxiv.org/pdf/1912.09729.\n        \"clip_ratio_c\", 3.0\n    )\n\n    cliprange = clip_ratio\n    cliprange_low = clip_ratio_low\n    cliprange_high = clip_ratio_high\n\n    assert clip_ratio_c > 1.0, (\n        \"The lower bound of the clip_ratio_c for dual-clip PPO should be greater than 1.0,\"\n        + f\" but get the value: {clip_ratio_c}.\"\n    )\n\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    pg_losses1 = -advantages * ratio\n    if cliprange_low is None:\n        cliprange_low = cliprange\n    if cliprange_high is None:\n        cliprange_high = cliprange\n    pg_losses2 = -advantages * torch.clamp(\n        ratio, 1 - cliprange_low, 1 + cliprange_high\n    )  # - clip(ratio, 1-cliprange, 1+cliprange) * A\n    clip_pg_losses1 = torch.maximum(\n        pg_losses1, pg_losses2\n    )  # max(-ratio * A, -clip(ratio, 1-cliprange, 1+cliprange) * A)\n    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)\n\n    pg_losses3 = -advantages * clip_ratio_c\n    clip_pg_losses2 = torch.min(pg_losses3, clip_pg_losses1)\n    pg_clipfrac_lower = verl_F.masked_mean(\n        torch.gt(clip_pg_losses1, pg_losses3) * (advantages < 0).float(), response_mask\n    )\n\n    pg_losses = torch.where(advantages < 0, clip_pg_losses2, clip_pg_losses1)\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"dppo_tv\")\ndef compute_policy_loss_dppo_tv(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for DPPO-Binary-TV.\n\n    See https://arxiv.org/pdf/2602.04879 for more details.\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n        config: `(verl.trainer.config.ActorConfig)`:\n            config for the actor.\n        rollout_log_probs: `(torch.Tensor)`:\n            log probabilities of actions under the rollout policy, shape (batch_size, response_length).\n    \"\"\"\n\n    assert config is not None\n    assert not isinstance(config, AlgoConfig)\n    # Note: the clip_ratio is different from the standard PPO, it is the TV divergence threshold for DPPO.\n    clip_divergence = config.clip_ratio\n    clip_divergence_low = config.clip_ratio_low if config.clip_ratio_low is not None else clip_divergence\n    clip_divergence_high = config.clip_ratio_high if config.clip_ratio_high is not None else clip_divergence\n\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    # Instead of dual-clip PPO, we use truncated importance sampling (TIS) to clip the policy loss.\n    # However, a large threshold is recommended to avoid performance degradation due to the truncation bias.\n    # See Section 5.4 in https://arxiv.org/pdf/2602.04879 for more details.\n    clip_ratio_c = config.get(\"clip_ratio_c\", 20.0)\n    truncated_ratio = torch.clamp(ratio, max=clip_ratio_c)\n    truncated_ratio = truncated_ratio.detach()\n\n    # Compute valid mask for DPPO-Binary-TV\n    prob = torch.exp(log_prob)\n    old_prob = torch.exp(old_log_prob)\n    valid_positive_mask = (prob - old_prob) <= clip_divergence_high\n    valid_negative_mask = (prob - old_prob) >= -clip_divergence_low\n    valid_mask = torch.where(advantages > 0, valid_positive_mask, valid_negative_mask)\n    valid_mask = valid_mask.detach().float()\n\n    pg_losses = -advantages * truncated_ratio * log_prob * valid_mask\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n\n    pg_clipfrac = verl_F.masked_mean((1.0 - valid_mask).float(), response_mask)\n    pg_clipfrac_lower = verl_F.masked_mean((ratio > clip_ratio_c).float() * valid_mask, response_mask)\n\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"dppo_kl\")\ndef compute_policy_loss_dppo_kl(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for DPPO-Binary-KL.\n\n    See https://arxiv.org/pdf/2602.04879 for more details.\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n        config: `(verl.trainer.config.ActorConfig)`:\n            config for the actor.\n        rollout_log_probs: `(torch.Tensor)`:\n            log probabilities of actions under the rollout policy, shape (batch_size, response_length).\n    \"\"\"\n\n    assert config is not None\n    assert not isinstance(config, AlgoConfig)\n    # Note: the clip_ratio is different from the standard PPO, it is the KL divergence threshold for DPPO.\n    clip_divergence = config.clip_ratio\n    clip_divergence_low = config.clip_ratio_low if config.clip_ratio_low is not None else clip_divergence\n    clip_divergence_high = config.clip_ratio_high if config.clip_ratio_high is not None else clip_divergence\n\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    # Instead of dual-clip PPO, we use truncated importance sampling (TIS) to clip the policy loss.\n    # However, a large threshold is recommended to avoid performance degradation due to the truncation bias.\n    # See Section 5.4 in https://arxiv.org/pdf/2602.04879 for more details.\n    clip_ratio_c = config.get(\"clip_ratio_c\", 20.0)\n    truncated_ratio = torch.clamp(ratio, max=clip_ratio_c)\n    truncated_ratio = truncated_ratio.detach()\n\n    # Compute valid mask for DPPO-Binary-KL\n    prob = torch.exp(log_prob)\n    old_prob = torch.exp(old_log_prob)\n    binary_kl = old_prob * (old_log_prob - log_prob) + (1 - old_prob) * torch.log(\n        (1.0 - old_prob + 1e-8) / (1.0 - prob + 1e-8)\n    )\n    valid_positive_mask = (binary_kl <= clip_divergence_high) | (prob <= old_prob)\n    valid_negative_mask = (binary_kl <= clip_divergence_low) | (prob >= old_prob)\n    valid_mask = torch.where(advantages > 0, valid_positive_mask, valid_negative_mask)\n    valid_mask = valid_mask.detach().float()\n\n    pg_losses = -advantages * truncated_ratio * log_prob * valid_mask\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n\n    # For compatibility, return zero for pg_clipfrac_lower (not used in standard DPPO)\n    pg_clipfrac = verl_F.masked_mean((1.0 - valid_mask).float(), response_mask)\n    pg_clipfrac_lower = verl_F.masked_mean((ratio > clip_ratio_c).float() * valid_mask, response_mask)\n\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"gspo\")\ndef compute_policy_loss_gspo(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"seq-mean-token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for GSPO.\n\n    See https://arxiv.org/pdf/2507.18071 for more details.\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. For GSPO, it is recommended to use \"seq-mean-token-mean\".\n    \"\"\"\n\n    assert config is not None\n    assert isinstance(config, ActorConfig)\n    clip_ratio_low = config.clip_ratio_low if config.clip_ratio_low is not None else config.clip_ratio\n    clip_ratio_high = config.clip_ratio_high if config.clip_ratio_high is not None else config.clip_ratio\n\n    negative_approx_kl = log_prob - old_log_prob\n\n    # compute sequence-level importance ratio:\n    # si(θ) = (π_θ(yi|x)/π_θold(yi|x))^(1/|yi|) =\n    # exp [(1/|y_i|) * Σ_t log(π_θ(y_i,t|x,y_i,<t)/π_θold(y_i,t|x,y_i,<t))]\n    seq_lengths = torch.sum(response_mask, dim=-1).clamp(min=1)\n    negative_approx_kl_seq = torch.sum(negative_approx_kl * response_mask, dim=-1) / seq_lengths\n\n    # Combined ratio at token level:\n    # s_i,t(θ) = sg[s_i(θ)] · π_θ(y_i,t|x, y_i,<t) / sg[π_θ(y_i,t|x, y_i,<t)]\n    # In log space: log(s_i,t(θ)) = sg[log(s_i(θ))] + log_prob - sg[log_prob]\n    log_seq_importance_ratio = log_prob - log_prob.detach() + negative_approx_kl_seq.detach().unsqueeze(-1)\n    log_seq_importance_ratio = torch.clamp(log_seq_importance_ratio, max=10.0)  # clamp for numerical stability\n\n    # finaly exp() to remove log\n    seq_importance_ratio = torch.exp(log_seq_importance_ratio)\n\n    pg_losses1 = -advantages * seq_importance_ratio\n    pg_losses2 = -advantages * torch.clamp(seq_importance_ratio, 1 - clip_ratio_low, 1 + clip_ratio_high)\n    pg_losses = torch.maximum(pg_losses1, pg_losses2)\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    # for GSPO, we need to aggregate the loss at the sequence level (seq-mean-token-mean)\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=\"seq-mean-token-mean\", **config.global_batch_info\n    )\n\n    # For compatibility, return zero for pg_clipfrac_lower (not used in standard GSPO)\n    pg_clipfrac = verl_F.masked_mean(torch.gt(pg_losses2, pg_losses1).float(), response_mask)\n    pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)\n\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"sapo\")\ndef compute_policy_loss_sapo(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"seq-mean-token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the smoothed policy objective and related metrics for SAPO.\n\n    See https://arxiv.org/pdf/2511.20347 for more details.\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. For SAPO, it is recommended to use \"seq-mean-token-mean\".\n    \"\"\"\n\n    assert config is not None\n    assert isinstance(config, ActorConfig)\n\n    # temperature for positive and negative token updates\n    tau_pos = torch.as_tensor(config.tau_pos, dtype=advantages.dtype, device=advantages.device)\n    tau_neg = torch.as_tensor(config.tau_neg, dtype=advantages.dtype, device=advantages.device)\n\n    def gate_function(x, tau):\n        \"\"\"The gating function used in SAPO\"\"\"\n        return torch.sigmoid(tau * (x - 1.0)) * (4.0 / tau)\n\n    # compute IS at token level:\n    # r_{i,t}(θ) = π_θ(y_{i,t}|x, y_{i,<t}) / π_θold(y_{i,t}|x, y_{i,<t})]\n    # In log space: log(r_{i,t}(θ)) = log_prob - ol_log_prob\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    # finally exp() to remove log and get r_{i,t}(θ)\n    ratio = torch.exp(negative_approx_kl)\n\n    # tau_{i,t} is tau_pos if adv > 0 else tau_neg\n    taus = torch.where(\n        condition=advantages > 0,\n        input=tau_pos,  # if A_{i,t} > 0 we set to tau_pos\n        other=tau_neg,  # if A_{i,t} <= 0 we set to tau_neg\n    )\n\n    # compute the gates f_{i,t}(r_{i,t}(θ)) at token level\n    gates = gate_function(ratio, taus)\n\n    # compute policy gradient loss\n    pg_losses = -gates * advantages\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    # for SAPO, we need to aggregate the loss at the sequence level (seq-mean-token-mean)\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=\"seq-mean-token-mean\", **config.global_batch_info\n    )\n\n    # For compatibility, return zero for both pg_clipfrac and pg_clipfrac_lower (not used in SAPO)\n    pg_clipfrac = torch.tensor(0.0, device=pg_loss.device)\n    pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)\n    # compute KL for metrics tracking\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n    # return metrics dict\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"gpg\")\ndef compute_policy_loss_gpg(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"Adapted from\n    https://github.com/AMAP-ML/GPG/blob/main/VisualThinker-R1-Zero/src/open-r1-multimodal/src/open_r1/trainer/grpo_trainer.py#L495\n    Args:\n        log_prob: `(torch.Tensor)`\n            shape: (bs, response_length)\n        advantages: `(torch.Tensor)`\n            shape: (bs, response_length)\n        response_mask: `(torch.Tensor)`\n            shape: (bs, response_length)\n    return:\n        pg_loss: `a scalar torch.Tensor`\n            policy gradient loss computed via GPG\n    \"\"\"\n    assert config is not None\n    pg_losses = -log_prob * advantages\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n    return pg_loss, {}\n\n\n@register_policy_loss(\"clip_cov\")\ndef compute_policy_loss_clip_cov(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for Clip-Cov.\n\n    Adapted from\n    https://github.com/PRIME-RL/Entropy-Mechanism-of-RL/blob/main/verl/trainer/ppo/core_algos.py\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        cliprange (float, optional):\n            Clipping parameter ε for standard PPO. See https://arxiv.org/abs/1707.06347.\n            Defaults to None (must be provided).\n        cliprange_low (float, optional):\n            Lower clip range for dual-clip PPO. Defaults to same as `cliprange`.\n        cliprange_high (float, optional):\n            Upper clip range for dual-clip PPO. Defaults to same as `cliprange`.\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n        clip_cvo_ratio (float, optional):\n            Ratio for clipping the covariance. Defaults to 0.0002.\n        clip_cov_lb (float, optional):\n            Lower bound for clipping covariance. Defaults to 1.0.\n        clip_cov_ub (float, optional):\n            Upper bound for clipping covariance. Defaults to 5.0.\n    \"\"\"\n    assert config is not None\n    assert not isinstance(config, AlgoConfig), \"passing AlgoConfig not supported yet\"\n    assert config.policy_loss is not None\n\n    clip_cov_ratio = config.policy_loss.clip_cov_ratio if config.policy_loss.clip_cov_ratio is not None else 0.0002\n    cliprange = config.clip_ratio\n    cliprange_low = config.clip_ratio_low if config.clip_ratio_low is not None else cliprange\n    cliprange_high = config.clip_ratio_high if config.clip_ratio_high is not None else cliprange\n    clip_cov_ub = config.policy_loss.clip_cov_ub if config.policy_loss.clip_cov_ub is not None else 5.0\n    clip_cov_lb = config.policy_loss.clip_cov_lb if config.policy_loss.clip_cov_lb is not None else 1.0\n\n    assert clip_cov_ratio > 0, \"clip_ratio should be larger than 0.\"\n\n    negative_approx_kl = log_prob - old_log_prob\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    pg_losses1 = -advantages * ratio\n\n    if cliprange_low is None:\n        cliprange_low = cliprange\n    if cliprange_high is None:\n        cliprange_high = cliprange\n\n    corr = torch.ones_like(advantages)\n    pg_losses2 = -advantages * torch.clamp(ratio, 1 - cliprange_low, 1 + cliprange_high)\n    clip_by_origin = (pg_losses2 > pg_losses1) & (response_mask > 0)\n\n    cov_all = (advantages - verl_F.masked_mean(advantages, response_mask)) * (\n        log_prob - verl_F.masked_mean(log_prob.detach(), response_mask)\n    )\n    cov_all[response_mask == 0] = -torch.inf\n    cov_all[clip_by_origin] = -torch.inf\n\n    clip_num = max(int(clip_cov_ratio * response_mask.sum().item()), 1)\n    top_k_idx = (cov_all < clip_cov_ub) & (cov_all > clip_cov_lb) & (response_mask > 0)\n    top_k_idx = torch.nonzero(top_k_idx)\n\n    if len(top_k_idx) > 0:\n        perm = torch.randperm(len(top_k_idx))\n        top_k_idx = top_k_idx[perm[: min(clip_num, len(top_k_idx))]]\n    else:\n        top_k_idx = torch.empty((0, 2), device=cov_all.device, dtype=torch.long)\n\n    corr[top_k_idx[:, 0], top_k_idx[:, 1]] = 0\n\n    pg_clipfrac = verl_F.masked_mean((corr == 0).float(), response_mask)\n\n    pg_losses = torch.maximum(pg_losses1, pg_losses2) * corr\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"kl_cov\")\ndef compute_policy_loss_kl_cov(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for Clip-Cov.\n\n    Adapted from\n    https://github.com/PRIME-RL/Entropy-Mechanism-of-RL/blob/main/verl/trainer/ppo/core_algos.py\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n        kl_cov_ratio (float, optional):\n            Ratio for selecting the top-k covariance values. Defaults to 0.0002.\n        ppo_kl_coef (float, optional):\n            Coefficient for the KL penalty term in the loss. Defaults to 1.\n    \"\"\"\n    assert config is not None\n    assert not isinstance(config, AlgoConfig), \"passing AlgoConfig not supported yet\"\n    assert config.policy_loss is not None\n\n    kl_cov_ratio = config.policy_loss.kl_cov_ratio if config.policy_loss.kl_cov_ratio is not None else 0.0002\n    ppo_kl_coef = config.policy_loss.ppo_kl_coef if config.policy_loss.ppo_kl_coef is not None else 1.0\n\n    assert kl_cov_ratio > 0, \"kl_cov_ratio should be larger than 0.\"\n\n    negative_approx_kl = log_prob - old_log_prob\n    abs_kl = negative_approx_kl.abs()\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl_abs = verl_F.masked_mean(negative_approx_kl.abs(), response_mask)\n    pg_losses1 = -advantages * ratio\n    pg_losses_kl = -advantages * ratio + ppo_kl_coef * abs_kl\n    pg_losses = pg_losses1\n\n    all_valid = response_mask > 0\n    all_valid_idx = torch.nonzero(all_valid.reshape(-1), as_tuple=True)[0]\n    all_valid_adv = advantages[all_valid].detach().reshape(-1).cpu()\n    all_valid_logp = log_prob[all_valid].detach().reshape(-1).cpu()\n\n    k = min(kl_cov_ratio, len(all_valid_adv))\n\n    if k != 0:\n        cov_lst_all = (all_valid_adv - all_valid_adv.mean()) * (all_valid_logp - all_valid_logp.mean())\n        k_percent_nums = max(1, int(len(cov_lst_all) * kl_cov_ratio))\n        large_cov_idxs = torch.topk(cov_lst_all, k_percent_nums, largest=True).indices\n\n        if len(large_cov_idxs) != 0:\n            large_cov_idxs = all_valid_idx[large_cov_idxs]\n            pg_losses[large_cov_idxs // advantages.shape[1], large_cov_idxs % advantages.shape[1]] = pg_losses_kl[\n                large_cov_idxs // advantages.shape[1], large_cov_idxs % advantages.shape[1]\n            ]\n\n    # Apply rollout correction weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n    pg_metrics = {\n        \"actor/ppo_kl\": ppo_kl_abs.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"geo_mean\")\ndef compute_policy_loss_geo_mean(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for GMPO.\n\n    Adapted from paper https://arxiv.org/abs/2507.20673\n    https://github.com/callsys/GMPO/blob/main/train_zero_math_gmpo.py\n\n    Args:\n        old_log_prob (torch.Tensor):\n            Log-probabilities of actions under the old policy, shape (batch_size, response_length).\n        log_prob (torch.Tensor):\n            Log-probabilities of actions under the current policy, shape (batch_size, response_length).\n        advantages (torch.Tensor):\n            Advantage estimates for each action, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the loss, shape (batch_size, response_length).\n        loss_agg_mode (str, optional):\n            not used\n    \"\"\"\n\n    assert config is not None\n    assert not isinstance(config, AlgoConfig)\n    clip_ratio = config.clip_ratio  # Clipping parameter. See https://arxiv.org/abs/1707.06347.\n    clip_ratio_low = config.clip_ratio_low if config.clip_ratio_low is not None else clip_ratio\n    clip_ratio_high = config.clip_ratio_high if config.clip_ratio_high is not None else clip_ratio\n\n    cliprange = clip_ratio\n    cliprange_low = clip_ratio_low\n    cliprange_high = clip_ratio_high\n    if cliprange_low is None:\n        cliprange_low = cliprange\n    if cliprange_high is None:\n        cliprange_high = cliprange\n\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp negative_approx_kl for stability (uncomment it if you like)\n    # negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    # Clipping at token-level & Clipping wider\n    sgn_advantage = torch.sign(advantages)\n    negative_approx_kl_clamp = torch.clamp(negative_approx_kl, -cliprange_low, cliprange_high)\n    negative_approx_kl_min = torch.min(sgn_advantage * negative_approx_kl, sgn_advantage * negative_approx_kl_clamp)\n    negative_approx_kl_min = sgn_advantage * negative_approx_kl_min\n\n    # Geometric-Mean Policy Optimization\n    response_mask_sum = response_mask.sum(dim=-1)\n    ratio = torch.exp((negative_approx_kl_min * response_mask).sum(dim=-1) / (response_mask_sum + 1e-8))\n    # we only support sequence level advantage for now,\n    # otherwise, below would be not consistent with the paper\n    advantage = (advantages * response_mask).sum(dim=-1) / (response_mask_sum + 1e-8)\n    pg_losses = -advantage * ratio\n\n    # Apply rollout correction weights if provided\n    # For geo_mean, IS weights are 2D (batch_size, seq_length) and need to be aggregated to sequence level\n    if rollout_is_weights is not None:\n        # Aggregate token-level weights to sequence level using geometric mean for consistency\n        # Note: rollout_is_weights is always 2D regardless of aggregation mode\n        seq_is_weights = torch.exp(\n            (torch.log(rollout_is_weights + 1e-10) * response_mask).sum(dim=-1) / (response_mask_sum + 1e-8)\n        )\n        pg_losses = pg_losses * seq_is_weights\n\n    pg_loss = torch.mean(pg_losses)\n\n    # higher: ratio is too large that need clamp to clip_high (when adv > 0)\n    clipped = torch.ne(negative_approx_kl, negative_approx_kl_clamp)\n    pg_clipfrac = verl_F.masked_mean((clipped * (advantages > 0)).float(), response_mask)\n    pg_clipfrac_lower = verl_F.masked_mean((clipped * (advantages < 0)).float(), response_mask)\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"cispo\")\ndef compute_policy_loss_cispo(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[DictConfig | ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"\n    Compute the clipped policy objective and related metrics for CISPO.\n\n    See https://arxiv.org/pdf/2506.13585 for more details.\n    \"\"\"\n\n    assert config is not None\n    assert isinstance(config, ActorConfig)\n    clip_ratio_low = config.clip_ratio_low if config.clip_ratio_low is not None else config.clip_ratio\n    clip_ratio_high = config.clip_ratio_high if config.clip_ratio_high is not None else config.clip_ratio\n\n    # Compute importance sampling ratio: π_θ / π_θ_old\n    negative_approx_kl = log_prob - old_log_prob\n    # Clamp for numerical stability\n    negative_approx_kl = torch.clamp(negative_approx_kl, min=-20.0, max=20.0)\n    ratio = torch.exp(negative_approx_kl)\n    ppo_kl = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    # CISPO: Clip the importance sampling weights\n    # KEY: Apply stop gradient to the clipped ratio\n    # This prevents gradients from flowing through the ratio computation and clipping\n    # Gradients only flow through log_prob in the final loss term\n    clipped_ratio = torch.clamp(ratio, 1 - clip_ratio_low, 1 + clip_ratio_high)\n    clipped_ratio_sg = clipped_ratio.detach()\n\n    # CISPO objective function (to maximize): J = sg(clip(ratio)) * A * log π_θ\n    # Loss function (to minimize): L = -J = -sg(clip(ratio)) * A * log_prob\n    pg_losses = -clipped_ratio_sg * advantages * log_prob\n\n    # Track clipping statistics\n    pg_clipfrac = verl_F.masked_mean((ratio != clipped_ratio).float(), response_mask)\n\n    # Apply rollout importance sampling weights if provided\n    if rollout_is_weights is not None:\n        pg_losses = pg_losses * rollout_is_weights\n\n    pg_loss = agg_loss(\n        loss_mat=pg_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n    )\n\n    # For compatibility, return zero for pg_clipfrac_lower (not used in CISPO)\n    pg_clipfrac_lower = torch.tensor(0.0, device=pg_loss.device)\n\n    pg_metrics = {\n        \"actor/pg_clipfrac\": pg_clipfrac.detach().item(),\n        \"actor/ppo_kl\": ppo_kl.detach().item(),\n        \"actor/pg_clipfrac_lower\": pg_clipfrac_lower.detach().item(),\n    }\n    return pg_loss, pg_metrics\n\n\ndef compute_entropy_loss(logits, response_mask, loss_agg_mode: str = \"token-mean\"):\n    \"\"\"Compute categorical entropy loss (For backward compatibility)\n\n    Args:\n        logits (torch.Tensor): shape is (bs, response_length, vocab_size)\n        response_mask (torch.Tensor): shape is (bs, response_length)\n\n    Returns:\n        entropy: a scalar torch.Tensor\n\n    \"\"\"\n    # compute entropy\n    token_entropy = verl_F.entropy_from_logits(logits)  # (bs, response_len)\n    entropy_loss = agg_loss(loss_mat=token_entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n    return entropy_loss\n\n\ndef compute_value_loss(\n    vpreds: torch.Tensor,\n    returns: torch.Tensor,\n    values: torch.Tensor,\n    response_mask: torch.Tensor,\n    cliprange_value: float,\n    loss_agg_mode: str = \"token-mean\",\n):\n    \"\"\"\n    Compute the clipped value-function loss for PPO.\n\n    Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1151\n\n    Args:\n        vpreds (torch.FloatTensor):\n            Predicted values from the value head, shape (batch_size, response_length).\n        values (torch.FloatTensor):\n            Old (baseline) values from the value head, shape (batch_size, response_length).\n        returns (torch.FloatTensor):\n            Ground-truth returns, shape (batch_size, response_length).\n        response_mask (torch.Tensor):\n            Mask indicating which tokens to include in the value loss calculation.\n        cliprange_value (float):\n            Clip range for value prediction updates.\n        loss_agg_mode (str, optional):\n            Aggregation mode for `agg_loss`. Defaults to \"token-mean\".\n\n    Returns:\n        vf_loss (torch.FloatTensor):\n            A scalar tensor containing the aggregated value-function loss.\n        vf_clipfrac (float):\n            Fraction of elements where the clipped loss was used.\n    \"\"\"\n    vpredclipped = verl_F.clip_by_value(vpreds, values - cliprange_value, values + cliprange_value)\n    vf_losses1 = (vpreds - returns) ** 2\n    vf_losses2 = (vpredclipped - returns) ** 2\n    clipped_vf_losses = torch.max(vf_losses1, vf_losses2)\n    vf_loss = 0.5 * agg_loss(loss_mat=clipped_vf_losses, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n    vf_clipfrac = verl_F.masked_mean(torch.gt(vf_losses2, vf_losses1).float(), response_mask)\n    return vf_loss, vf_clipfrac\n\n\ndef kl_penalty(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:\n    \"\"\"Compute KL divergence given logprob and ref_logprob. Optionally using straight through to bind k2 on other\n    kl penalty compute method for unbiased KL gradient estimation.\n    See more description in http://joschu.net/blog/kl-approx.html\n\n    Args:\n        logprob:\n        ref_logprob:\n\n    Returns:\n        kl_estimate\n    \"\"\"\n    forward_score = kl_penalty_forward(logprob, ref_logprob, kl_penalty)\n    if not kl_penalty.endswith(\"+\") or kl_penalty in (\"mse\", \"k2\"):\n        return forward_score\n\n    \"\"\"\n    The expectation of k1 and k3 estimator is the expected value of KL, but the expected gradient of k1 and k3\n    estimator is not the expected gradient of KL. On the other hand k2 estimator gives right gradient estimator, \n    so we use a straight through trick here if the kl_penalty method ends with '+', e.g., k3+. \n    \"\"\"\n    backward_score = 0.5 * (logprob - ref_logprob).square()\n\n    return backward_score - backward_score.detach() + forward_score.detach()\n\n\ndef kl_penalty_forward(logprob: torch.FloatTensor, ref_logprob: torch.FloatTensor, kl_penalty) -> torch.FloatTensor:\n    \"\"\"Compute KL divergence given logprob and ref_logprob.\n    Copied from https://github.com/huggingface/trl/blob/main/trl/trainer/ppo_trainer.py#L1104\n    See more description in http://joschu.net/blog/kl-approx.html\n\n    Args:\n        logprob:\n        ref_logprob:\n\n    Returns:\n        kl_estimate\n    \"\"\"\n    if kl_penalty in (\"kl\", \"k1\"):\n        return logprob - ref_logprob\n\n    if kl_penalty == \"abs\":\n        return (logprob - ref_logprob).abs()\n\n    if kl_penalty in (\"mse\", \"k2\"):\n        return 0.5 * (logprob - ref_logprob).square()\n\n    # J. Schulman. Approximating kl divergence, 2020.\n    # # URL http://joschu.net/blog/kl-approx.html.\n    if kl_penalty in (\"low_var_kl\", \"k3\"):\n        kl = ref_logprob - logprob\n        # For numerical stability\n        kl = torch.clamp(kl, min=-20, max=20)\n        ratio = torch.exp(kl)\n        kld = (ratio - kl - 1).contiguous()\n        return torch.clamp(kld, min=-10, max=10)\n\n    if kl_penalty == \"full\":\n        # so, here logprob and ref_logprob should contain the logits for every token in vocabulary\n        raise NotImplementedError\n\n    raise NotImplementedError\n\n\ndef compute_pf_ppo_reweight_data(\n    data,\n    reweight_method: str = \"pow\",\n    weight_pow: float = 2.0,\n):\n    \"\"\"Reweight the data based on the token_level_scores.\n\n    Args:\n        data: DataProto object, containing batch, non_tensor_batch and meta_info\n        reweight_method: str, choices: \"pow\", \"max_min\", \"max_random\"\n        weight_pow: float, the power of the weight\n\n    Returns:\n\n    \"\"\"\n\n    @torch.no_grad()\n    def compute_weights(scores: torch.Tensor, reweight_method: str, weight_pow: float) -> torch.Tensor:\n        \"\"\"Compute importance weights for resampling based on scores.\n\n        Args:\n            scores (torch.Tensor): Tensor of scores to compute weights from.\n            reweight_method (str): Method for computing weights ('pow', 'max_min', 'max_random').\n            weight_pow (float): Power exponent for 'pow' method.\n\n        Returns:\n            torch.Tensor: Computed importance weights.\n\n        Raises:\n            ValueError: If reweight_method is not supported.\n        \"\"\"\n        if reweight_method == \"pow\":\n            weights = torch.pow(torch.abs(scores), weight_pow)\n        elif reweight_method == \"max_min\":\n            max_score = torch.max(scores)\n            min_score = torch.min(scores)\n            weights = torch.where((scores == max_score) | (scores == min_score), 1.0, 0.0)\n        elif reweight_method == \"max_random\":\n            max_score = torch.max(scores)\n            weights = torch.where(scores == max_score, 0.4, 0.1)\n        else:\n            raise ValueError(f\"Unsupported reweight_method: {reweight_method}\")\n        return weights\n\n    scores = data.batch[\"token_level_scores\"].sum(dim=-1)\n    weights = compute_weights(scores, reweight_method, weight_pow)\n    weights = torch.clamp(weights + 1e-8, min=1e-8)\n\n    batch_size = scores.shape[0]\n    sample_indices = torch.multinomial(weights, batch_size, replacement=True)\n\n    resampled_batch = {key: tensor[sample_indices] for key, tensor in data.batch.items()}\n\n    sample_indices_np = sample_indices.numpy()\n    resampled_non_tensor_batch = {}\n    for key, array in data.non_tensor_batch.items():\n        if isinstance(array, np.ndarray):\n            resampled_non_tensor_batch[key] = array[sample_indices_np]\n        else:\n            resampled_non_tensor_batch[key] = [array[i] for i in sample_indices_np]\n\n    resampled_meta_info = {}\n    for key, value in data.meta_info.items():\n        if isinstance(value, list) and len(value) == batch_size:\n            resampled_meta_info[key] = [value[i] for i in sample_indices_np]\n        else:\n            resampled_meta_info[key] = value\n\n    from copy import deepcopy\n\n    resampled_data = deepcopy(data)\n    resampled_data.batch = type(data.batch)(resampled_batch)\n    resampled_data.batch.batch_size = data.batch.batch_size\n    resampled_data.non_tensor_batch = resampled_non_tensor_batch\n    resampled_data.meta_info = resampled_meta_info\n\n    return resampled_data\n\n\ndef compute_policy_loss_reinforce(\n    rollout_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"seq-mean-token-sum\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: Optional[torch.Tensor] = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"Compute REINFORCE-style policy gradient loss with optional IS correction.\n\n    This function implements policy gradient (REINFORCE) with optional importance\n    sampling correction for rollout-training policy mismatch.\n\n    Mathematical formulation:\n        Without IS (rollout_is_weights=None):\n            L = -E[log π(a|s) * A(s,a)]\n            Gradient: ∇_θ L = -E[∇log π(a|s) * A] (standard REINFORCE)\n\n        With IS (rollout_is_weights provided):\n            L = -E_π_rollout[w * log π(a|s) * A(s,a)]\n            where w = π_current / π_rollout (truncated IS weight)\n            Gradient: ∇_θ L = -E[w * ∇log π(a|s) * A] (IS-corrected policy gradient)\n\n    Args:\n        rollout_log_prob: Log probabilities from rollout policy (e.g., vLLM BF16).\n            Shape: (batch_size, seq_length). Used for KL computation.\n        log_prob: Log probabilities from current training policy.\n            Shape: (batch_size, seq_length)\n        advantages: Advantage estimates for each token.\n            Shape: (batch_size, seq_length)\n        response_mask: Mask indicating valid tokens (1 for valid, 0 for padding).\n            Shape: (batch_size, seq_length). Should already include rejection sampling.\n        loss_agg_mode: Loss aggregation strategy (see agg_loss for details).\n        config: Actor config (required for global_batch_info).\n        rollout_is_weights: Pre-computed IS weights (π_current / π_rollout).\n            Shape: (batch_size, seq_length). None to disable IS correction.\n\n    Returns:\n        Tuple of (loss, metrics):\n            loss: Scalar policy gradient loss\n            metrics: Dictionary with \"actor/ppo_kl\"\n\n    Note:\n        Unlike PPO (compute_policy_loss_vanilla), this function:\n        - Does NOT use PPO clipping\n        - Uses log π(a|s) directly (not ratio)\n        - IS weights are applied as multiplicative factor\n    \"\"\"\n    assert config is not None, \"ActorConfig must be provided for REINFORCE loss\"\n\n    # Compute pure policy gradient loss with optional IS correction\n    # Standard REINFORCE: L = -E[log π(a|s) * A]\n    # With IS: L = -E[w * log π(a|s) * A] where w = π_current / π_rollout\n    if rollout_is_weights is not None:\n        # IS-corrected policy gradient: L = -E[stopgrad(w) · log π · A]\n        pg_losses = -advantages * log_prob * rollout_is_weights\n    else:\n        # Standard REINFORCE: L = -E[log π · A]\n        pg_losses = -advantages * log_prob\n\n    # Aggregate loss\n    pg_loss = agg_loss(\n        loss_mat=pg_losses,\n        loss_mask=response_mask,\n        loss_agg_mode=loss_agg_mode,\n        **config.global_batch_info,\n    )\n\n    # Compute KL divergence between current and rollout policy\n    negative_approx_kl = log_prob - rollout_log_prob\n    kl_divergence = verl_F.masked_mean(-negative_approx_kl, response_mask)\n\n    pg_metrics = {\n        \"actor/ppo_kl\": kl_divergence.detach().item(),\n    }\n\n    return pg_loss, pg_metrics\n\n\n@register_policy_loss(\"bypass_mode\")\ndef compute_policy_loss_bypass_mode(\n    old_log_prob: torch.Tensor,\n    log_prob: torch.Tensor,\n    advantages: torch.Tensor,\n    response_mask: torch.Tensor,\n    loss_agg_mode: str = \"token-mean\",\n    config: Optional[ActorConfig] = None,\n    rollout_is_weights: torch.Tensor | None = None,\n) -> tuple[torch.Tensor, dict[str, Any]]:\n    \"\"\"Bypass mode policy loss supporting both REINFORCE and PPO-clip.\n\n    This function is the entry point for bypass mode, where old_log_prob = rollout_log_prob.\n    It computes IS weights and rejection masks, then dispatches to either REINFORCE or\n    PPO-clip loss based on the loss_type configuration.\n\n    IMPORTANT - Bypass mode semantics:\n        In bypass mode, the trainer sets old_log_prob = rollout_log_prob.\n        This means:\n        - For REINFORCE: We use IS weights w = π_current / π_rollout explicitly\n        - For PPO-clip: The PPO ratio π_current / π_old = π_current / π_rollout\n          already incorporates the IS correction through clipping, so we do NOT\n          apply additional IS weights (would be double-counting)\n\n    Loss types:\n        - \"ppo_clip\" (default): PPO clipped objective (compute_policy_loss_vanilla)\n            L = -E[min(r*A, clip(r)*A)] where r = π_current / π_rollout\n            Note: IS weights are NOT applied (clipping handles the ratio)\n        - \"reinforce\": REINFORCE-style policy gradient with IS correction\n            L = -E[w * log π(a|s) * A] where w = π_current / π_rollout\n\n    Args:\n        old_log_prob: In bypass mode, this is actually rollout_log_prob.\n            Shape: (batch_size, seq_length)\n        log_prob: Current policy log probabilities.\n            Shape: (batch_size, seq_length)\n        advantages: Advantage estimates.\n            Shape: (batch_size, seq_length)\n        response_mask: Valid token mask (1=valid, 0=padding).\n            Shape: (batch_size, seq_length)\n        loss_agg_mode: Loss aggregation mode (passed to underlying loss function).\n        config: Actor config containing rollout_correction settings in policy_loss.\n        rollout_is_weights: Pre-computed IS weights (ignored, computed internally).\n\n    Config options (in config.policy_loss.rollout_correction):\n        loss_type: \"ppo_clip\" (default) or \"reinforce\"\n        rollout_is: IS aggregation level (\"token\", \"sequence\", or None)\n        rollout_is_threshold: Upper threshold for truncating IS weights (default: 2.0)\n        rollout_rs: Rejection sampling level (see rollout_corr_helper for supported modes)\n        rollout_rs_threshold: Threshold specification for rejection sampling\n        rollout_is_batch_normalize: Whether to normalize IS weights to mean=1.0\n\n    Returns:\n        Tuple of (loss, metrics):\n            loss: Scalar policy loss\n            metrics: Dictionary with rollout correction metrics and actor/ppo_kl\n    \"\"\"\n    from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_rejection_mask\n\n    assert config is not None, \"config is required for bypass_mode loss\"\n\n    # Extract rollout_correction config from policy_loss\n    rollout_corr_config = config.policy_loss.get(\"rollout_correction\", None) if hasattr(config, \"policy_loss\") else None\n\n    if rollout_corr_config is None:\n        raise ValueError(\n            \"rollout_correction config not found in policy_loss. \"\n            \"When using loss_mode='bypass_mode', ensure rollout_correction config is passed.\"\n        )\n\n    # Extract parameters\n    loss_type = rollout_corr_config.get(\"loss_type\", \"ppo_clip\")\n    rollout_is = rollout_corr_config.get(\"rollout_is\", None)\n    rollout_is_threshold = rollout_corr_config.get(\"rollout_is_threshold\", 2.0)\n    rollout_is_batch_normalize = rollout_corr_config.get(\"rollout_is_batch_normalize\", False)\n    rollout_rs = rollout_corr_config.get(\"rollout_rs\", None)\n    rollout_rs_threshold = rollout_corr_config.get(\"rollout_rs_threshold\", None)\n\n    # In bypass mode: old_log_prob IS rollout_log_prob\n    rollout_log_prob = old_log_prob\n\n    # Compute IS weights and rejection mask\n    # Note: For PPO-clip, we still compute IS weights for metrics, but don't apply them\n    with torch.no_grad():\n        rollout_is_weights_proto, modified_response_mask, rollout_metrics = (\n            compute_rollout_correction_and_rejection_mask(\n                old_log_prob=log_prob,  # Current policy (for IS ratio: π_current / π_rollout)\n                rollout_log_prob=rollout_log_prob,  # Rollout policy\n                response_mask=response_mask,\n                rollout_is=rollout_is,\n                rollout_is_threshold=rollout_is_threshold,\n                rollout_is_batch_normalize=rollout_is_batch_normalize,\n                rollout_rs=rollout_rs,\n                rollout_rs_threshold=rollout_rs_threshold,\n            )\n        )\n\n    # Extract IS weights tensor (or None if disabled)\n    computed_is_weights = rollout_is_weights_proto.batch[\"rollout_is_weights\"] if rollout_is_weights_proto else None\n\n    # Apply rejection mask (RS + veto)\n    effective_mask = modified_response_mask\n\n    # Dispatch to appropriate loss function based on loss_type\n    if loss_type == \"reinforce\":\n        # REINFORCE: Apply IS weights explicitly\n        pg_loss, pg_metrics = compute_policy_loss_reinforce(\n            rollout_log_prob=rollout_log_prob,\n            log_prob=log_prob,\n            advantages=advantages,\n            response_mask=effective_mask,\n            loss_agg_mode=loss_agg_mode,\n            config=config,\n            rollout_is_weights=computed_is_weights,\n        )\n\n    elif loss_type == \"ppo_clip\":\n        # PPO-clip: The ratio π_current/π_old = π_current/π_rollout already handles IS\n        # DO NOT apply IS weights - would be double-counting!\n        # The clipping mechanism constrains the effective IS ratio\n        pg_loss, pg_metrics = compute_policy_loss_vanilla(  # type: ignore[call-arg]\n            old_log_prob=rollout_log_prob,  # = old_log_prob in bypass mode\n            log_prob=log_prob,\n            advantages=advantages,\n            response_mask=effective_mask,\n            loss_agg_mode=loss_agg_mode,\n            config=config,\n            rollout_is_weights=None,  # Explicitly None - no IS weights for PPO-clip\n        )\n\n    else:\n        raise ValueError(f\"Invalid loss_type: {loss_type}. Must be 'reinforce' or 'ppo_clip'.\")\n\n    # Merge rollout correction metrics\n    pg_metrics.update(rollout_metrics)\n\n    return pg_loss, pg_metrics\n"
  },
  {
    "path": "verl/trainer/ppo/metric_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nMetrics related to the PPO trainer.\n\"\"\"\n\nfrom collections import defaultdict\nfrom functools import partial\nfrom typing import Any, Callable\n\nimport numpy as np\nimport torch\n\nimport verl.utils.torch_functional as verl_F\nfrom verl import DataProto\nfrom verl.utils.import_utils import deprecated\n\n\n@deprecated(\"verl.utils.metric.reduce_metrics\")\ndef reduce_metrics(metrics: dict[str, list[Any]]) -> dict[str, Any]:\n    \"\"\"\n    Reduces a dictionary of metric lists by computing the mean of each list.\n\n    Args:\n        metrics: A dictionary mapping metric names to lists of metric values.\n\n    Returns:\n        A dictionary with the same keys but with each list replaced by its mean value.\n\n    Example:\n        >>> metrics = {\"loss\": [1.0, 2.0, 3.0], \"accuracy\": [0.8, 0.9, 0.7]}\n        >>> reduce_metrics(metrics)\n        {\"loss\": 2.0, \"accuracy\": 0.8}\n    \"\"\"\n    from verl.utils.metric import reduce_metrics\n\n    return reduce_metrics(metrics)\n\n\ndef _compute_response_info(batch: DataProto) -> dict[str, Any]:\n    \"\"\"\n    Computes information about prompts and responses from a batch.\n\n    This is an internal helper function that extracts masks and lengths for prompts and responses.\n\n    Args:\n        batch: A DataProto object containing batch data with responses and attention masks.\n\n    Returns:\n        A dictionary containing:\n            - response_mask: Attention mask for the response tokens\n            - prompt_length: Tensor of prompt lengths for each item in the batch\n            - response_length: Tensor of response lengths for each item in the batch\n    \"\"\"\n    response_length = batch.batch[\"responses\"].shape[-1]\n\n    prompt_mask = batch.batch[\"attention_mask\"][:, :-response_length]\n    response_mask = batch.batch[\"attention_mask\"][:, -response_length:]\n\n    prompt_length = prompt_mask.sum(-1).float()\n    response_length = response_mask.sum(-1).float()  # (batch_size,)\n\n    return dict(\n        response_mask=response_mask,\n        prompt_length=prompt_length,\n        response_length=response_length,\n    )\n\n\ndef compute_data_metrics(batch: DataProto, use_critic: bool = True) -> dict[str, Any]:\n    \"\"\"\n    Computes various metrics from a batch of data for PPO training.\n\n    This function calculates metrics related to scores, rewards, advantages, returns, values,\n    and sequence lengths from a batch of data. It provides statistical information (mean, max, min)\n    for each metric category.\n\n    Args:\n        batch: A DataProto object containing batch data with token-level scores, rewards, advantages, etc.\n        use_critic: Whether to include critic-specific metrics. Defaults to True.\n\n    Returns:\n        A dictionary of metrics including:\n            - critic/score/mean, max, min: Statistics about sequence scores\n            - critic/rewards/mean, max, min: Statistics about sequence rewards\n            - critic/advantages/mean, max, min: Statistics about advantages\n            - critic/returns/mean, max, min: Statistics about returns\n            - critic/values/mean, max, min: Statistics about critic values (if use_critic=True)\n            - critic/vf_explained_var: Explained variance of the value function (if use_critic=True)\n            - response_length/mean, max, min, clip_ratio: Statistics about response lengths\n            - prompt_length/mean, max, min, clip_ratio: Statistics about prompt lengths\n            - num_turns/mean, max, min: Statistics about the number of multi-turn conversations\n    \"\"\"\n    sequence_score = batch.batch[\"token_level_scores\"].sum(-1)\n    sequence_reward = batch.batch[\"token_level_rewards\"].sum(-1)\n\n    advantages = batch.batch[\"advantages\"]\n    returns = batch.batch[\"returns\"]\n\n    max_response_length = batch.batch[\"responses\"].shape[-1]\n\n    prompt_mask = batch.batch[\"attention_mask\"][:, :-max_response_length].bool()\n    response_mask = batch.batch[\"response_mask\"].bool()\n\n    max_prompt_length = prompt_mask.size(-1)\n\n    response_info = _compute_response_info(batch)\n    prompt_length = response_info[\"prompt_length\"]\n    response_length = response_info[\"response_length\"]\n\n    aborted_mask = (response_length == 0).bool()\n    non_aborted_mask = ~aborted_mask\n\n    non_aborted_sequence_score = sequence_score[non_aborted_mask]\n    non_aborted_sequence_reward = sequence_reward[non_aborted_mask]\n\n    score_mean = torch.mean(non_aborted_sequence_score).detach().item()\n    score_max = torch.max(non_aborted_sequence_score).detach().item()\n    score_min = torch.min(non_aborted_sequence_score).detach().item()\n\n    reward_mean = torch.mean(non_aborted_sequence_reward).detach().item()\n    reward_max = torch.max(non_aborted_sequence_reward).detach().item()\n    reward_min = torch.min(non_aborted_sequence_reward).detach().item()\n\n    valid_adv = torch.masked_select(advantages, response_mask)\n    valid_returns = torch.masked_select(returns, response_mask)\n\n    if use_critic:\n        values = batch.batch[\"values\"]\n        valid_values = torch.masked_select(values, response_mask)\n        return_diff_var = torch.var(valid_returns - valid_values)\n        return_var = torch.var(valid_returns)\n\n    # Aborted samples and non-aborted response length statistics\n    # response_length_non_aborted/*: statistics computed on non-aborted samples only\n    aborted_ratio = torch.mean(aborted_mask.float()).detach().item()\n\n    non_aborted_response_length = response_length[non_aborted_mask]\n    if non_aborted_response_length.numel() > 0:\n        non_aborted_response_length_mean = torch.mean(non_aborted_response_length).detach().item()\n        non_aborted_response_length_max = torch.max(non_aborted_response_length).detach().item()\n        non_aborted_response_length_min = torch.min(non_aborted_response_length).detach().item()\n        non_aborted_response_length_clip_ratio = (\n            torch.mean(torch.eq(non_aborted_response_length, max_response_length).float()).detach().item()\n        )\n    else:\n        raise ValueError(\"All samples are aborted, this should not happen.\")\n\n    metrics = {\n        # score\n        \"critic/score/mean\": score_mean,\n        \"critic/score/max\": score_max,\n        \"critic/score/min\": score_min,\n        # reward\n        \"critic/rewards/mean\": reward_mean,\n        \"critic/rewards/max\": reward_max,\n        \"critic/rewards/min\": reward_min,\n        # adv\n        \"critic/advantages/mean\": torch.mean(valid_adv).detach().item(),\n        \"critic/advantages/max\": torch.max(valid_adv).detach().item(),\n        \"critic/advantages/min\": torch.min(valid_adv).detach().item(),\n        # returns\n        \"critic/returns/mean\": torch.mean(valid_returns).detach().item(),\n        \"critic/returns/max\": torch.max(valid_returns).detach().item(),\n        \"critic/returns/min\": torch.min(valid_returns).detach().item(),\n        **(\n            {\n                # values\n                \"critic/values/mean\": torch.mean(valid_values).detach().item(),\n                \"critic/values/max\": torch.max(valid_values).detach().item(),\n                \"critic/values/min\": torch.min(valid_values).detach().item(),\n                # vf explained var\n                \"critic/vf_explained_var\": (1.0 - return_diff_var / (return_var + 1e-5)).detach().item(),\n            }\n            if use_critic\n            else {}\n        ),\n        # response length\n        \"response_length/mean\": torch.mean(response_length).detach().item(),\n        \"response_length/max\": torch.max(response_length).detach().item(),\n        \"response_length/min\": torch.min(response_length).detach().item(),\n        \"response_length/clip_ratio\": torch.mean(torch.eq(response_length, max_response_length).float())\n        .detach()\n        .item(),\n        # response length (non-aborted only)\n        # These statistics exclude aborted samples to avoid skew from zeros\n        \"response_length_non_aborted/mean\": non_aborted_response_length_mean,\n        \"response_length_non_aborted/max\": non_aborted_response_length_max,\n        \"response_length_non_aborted/min\": non_aborted_response_length_min,\n        \"response_length_non_aborted/clip_ratio\": non_aborted_response_length_clip_ratio,\n        # aborted ratio\n        # Fraction of samples whose response length is zero\n        \"response/aborted_ratio\": aborted_ratio,\n        # prompt length\n        \"prompt_length/mean\": torch.mean(prompt_length).detach().item(),\n        \"prompt_length/max\": torch.max(prompt_length).detach().item(),\n        \"prompt_length/min\": torch.min(prompt_length).detach().item(),\n        \"prompt_length/clip_ratio\": torch.mean(torch.eq(prompt_length, max_prompt_length).float()).detach().item(),\n    }\n\n    # multi-turn conversation\n    if \"__num_turns__\" in batch.non_tensor_batch:\n        num_turns = batch.non_tensor_batch[\"__num_turns__\"]\n        metrics[\"num_turns/min\"] = num_turns.min()\n        metrics[\"num_turns/max\"] = num_turns.max()\n        metrics[\"num_turns/mean\"] = num_turns.mean()\n\n    if \"tool_call_counts\" in batch.non_tensor_batch:\n        tool_call_counts = batch.non_tensor_batch[\"tool_call_counts\"]\n        metrics[\"tool_call_counts/min\"] = tool_call_counts.min()\n        metrics[\"tool_call_counts/max\"] = tool_call_counts.max()\n        metrics[\"tool_call_counts/mean\"] = tool_call_counts.mean()\n\n    return metrics\n\n\ndef compute_timing_metrics(batch: DataProto, timing_raw: dict[str, float]) -> dict[str, Any]:\n    \"\"\"\n    Computes timing metrics for different processing stages in PPO training.\n\n    This function calculates both raw timing metrics (in seconds) and per-token timing metrics\n    (in milliseconds) for various processing stages like generation, reference computation,\n    value computation, advantage computation, and model updates.\n\n    Args:\n        batch: A DataProto object containing batch data with responses and attention masks.\n        timing_raw: A dictionary mapping stage names to their execution times in seconds.\n\n    Returns:\n        A dictionary containing:\n            - timing_s/{name}: Raw timing in seconds for each stage\n            - timing_per_token_ms/{name}: Per-token timing in milliseconds for each stage\n\n    Note:\n        Different stages use different token counts for normalization:\n        - \"gen\" uses only response tokens\n        - Other stages (\"ref\", \"values\", \"adv\", \"update_critic\", \"update_actor\") use all tokens\n          (prompt + response)\n    \"\"\"\n    response_info = _compute_response_info(batch)\n    num_prompt_tokens = torch.sum(response_info[\"prompt_length\"]).item()\n    num_response_tokens = torch.sum(response_info[\"response_length\"]).item()\n    num_overall_tokens = num_prompt_tokens + num_response_tokens\n\n    num_tokens_of_section = {\n        \"gen\": num_response_tokens,\n        **{name: num_overall_tokens for name in [\"ref\", \"values\", \"adv\", \"update_critic\", \"update_actor\"]},\n    }\n\n    return {\n        **{f\"timing_s/{name}\": value for name, value in timing_raw.items()},\n        **{\n            f\"timing_per_token_ms/{name}\": timing_raw[name] * 1000 / num_tokens_of_section[name]\n            for name in set(num_tokens_of_section.keys()) & set(timing_raw.keys())\n        },\n    }\n\n\ndef compute_throughout_metrics(batch: DataProto, timing_raw: dict[str, float], n_gpus: int) -> dict[str, Any]:\n    \"\"\"\n    Computes throughput metrics for PPO training.\n\n    This function calculates performance metrics related to token processing speed,\n    including the total number of tokens processed, time per step, and throughput\n    (tokens per second per GPU).\n\n    Args:\n        batch: A DataProto object containing batch data with meta information about token counts.\n        timing_raw: A dictionary mapping stage names to their execution times in seconds.\n                   Must contain a \"step\" key with the total step time.\n        n_gpus: Number of GPUs used for training.\n\n    Returns:\n        A dictionary containing:\n            - perf/total_num_tokens: Total number of tokens processed in the batch\n            - perf/time_per_step: Time taken for the step in seconds\n            - perf/throughput: Tokens processed per second per GPU\n\n    Note:\n        The throughput is calculated as total_tokens / (time * n_gpus) to normalize\n        across different GPU counts.\n    \"\"\"\n    total_num_tokens = sum(batch.meta_info[\"global_token_num\"])\n    time = timing_raw[\"step\"]\n    # estimated_flops, promised_flops = flops_function.estimate_flops(num_tokens, time)\n    # f'Actual TFLOPs/s/GPU​': estimated_flops/(n_gpus),\n    # f'Theoretical TFLOPs/s/GPU​': promised_flops,\n    return {\n        \"perf/total_num_tokens\": total_num_tokens,\n        \"perf/time_per_step\": time,\n        \"perf/throughput\": total_num_tokens / (time * n_gpus),\n    }\n\n\ndef compute_variance_proxy_metrics(batch: DataProto, gradient_norm: float = None) -> dict[str, float]:\n    \"\"\"\n    Compute variance proxy metrics using the simplified expected squared norm approach.\n\n    This metric provides a computationally efficient way to monitor gradient variance\n    during training. It works for any advantage estimator as long as sum_pi_squared\n    is available from the actor.\n\n    Theory:\n    - Full variance: Var(g̃) = E[||g̃||²] - ||g_true||²\n    - Simplified proxy (when ||g_true||² ≈ 0): Var(g̃) ≈ E[||g̃||²]\n    - Using W-score approximation: E[||g̃||²] ≈ E[A² × W(τ)]\n\n    Where W(τ) = Σ_t[1 - 2π_t(y_t) + Σπ²] is the score-norm proxy.\n    \"\"\"\n    metrics = {}\n\n    # Check if we have the necessary data (sum_pi_squared is required for W-score)\n    if \"sum_pi_squared\" not in batch.batch or \"old_log_probs\" not in batch.batch or \"advantages\" not in batch.batch:\n        return metrics\n\n    # Compute W(τ) = Σ_t[1 - 2π_t(y_t) + Σπ²]\n    pi_t = torch.exp(batch.batch[\"old_log_probs\"])\n    w_per_timestep = 1 - 2 * pi_t + batch.batch[\"sum_pi_squared\"]\n\n    # Get response mask to only consider valid tokens\n    response_mask = batch.batch[\"response_mask\"]\n\n    # Use pre-computed rollout IS weights from batch (for variance proxy consistency with training loss)\n    # IS weights are computed centrally in ray_trainer.py to avoid duplication\n    rollout_is_weights = None\n    if \"rollout_is_weights\" in batch.batch:\n        # Extract pre-computed IS weights from batch (already computed in trainer)\n        rollout_is_weights = batch.batch[\"rollout_is_weights\"]\n\n        # Scale W by (rollout IS weight)² for optimal baseline under biased estimation\n        w_per_timestep = w_per_timestep * (rollout_is_weights**2).detach()\n\n        # Note: IS weight statistics and mismatch metrics are logged in ray_trainer.py\n\n    # Get scalar advantages (mean over timesteps)\n    advantages = batch.batch[\"advantages\"]\n    # Compute mean advantage per trajectory using masked_mean\n    advantages_scalar = verl_F.masked_mean(advantages, response_mask, axis=-1)\n\n    # Compute W values (sum over timesteps)\n    w_values = verl_F.masked_sum(w_per_timestep, response_mask, axis=-1)\n\n    # ====== COMPUTE VARIANCE PROXIES ======\n    # Variance proxy should match the actual gradient computation:\n    # - If IS weights were computed/applied: use them in variance proxy calculation\n    # - Otherwise: compute on-policy variance proxy\n\n    # ====== PROXY 1: Signal Strength ||ḡ||² ======\n    # The squared norm of the mean gradient (provided from training loop)\n    proxy1_signal_strength = gradient_norm**2 if gradient_norm is not None else None\n\n    # ====== PROXY 2: Total Power E[||ĝ_τ||²] ======\n    # Measures the average of squared gradient norms (Signal + Noise)\n    if rollout_is_weights is not None:\n        # Off-policy with IS correction applied: use clamped weights consistently with actual gradient computation\n        rollout_is_weights_scalar = verl_F.masked_mean(rollout_is_weights, response_mask, axis=-1)\n        # Recover original W (before IS correction was applied in line 657)\n        # Clamp to avoid division by zero when IS weights are zero\n        w_original = verl_F.masked_sum(\n            w_per_timestep / torch.clamp((rollout_is_weights**2).detach(), min=1e-10), response_mask, axis=-1\n        )\n        # Clamp W to avoid negative values (which would cause NaN in sqrt)\n        w_original = torch.clamp(w_original, min=0.0)\n        # Proxy 2 for off-policy: E[ρ̄² × A² × W]\n        proxy2_total_power = ((rollout_is_weights_scalar**2) * (advantages_scalar**2) * w_original).mean()\n\n    else:\n        # On-policy Proxy 2: E[A² × W]\n        # Clamp W to avoid negative values (which would cause NaN in sqrt)\n        w_values_clamped = torch.clamp(w_values, min=0.0)\n        proxy2_total_power = (advantages_scalar**2 * w_values_clamped).mean()\n\n    # ====== PROXY 3: Pure Noise - Variance of Mean Vector ======\n    # Requires ||ḡ||² from actual batch gradient\n    # Formula: (1/(N-1)) × (Proxy2 - Proxy1)\n    proxy3_pure_noise = None\n    if proxy1_signal_strength is not None:\n        batch_size = advantages_scalar.shape[0]\n        if batch_size > 1:\n            proxy3_pure_noise = (1.0 / (batch_size - 1)) * (proxy2_total_power - proxy1_signal_strength)\n            # Ensure non-negative (can be negative due to numerical errors)\n            proxy3_pure_noise = max(\n                0.0, proxy3_pure_noise.item() if torch.is_tensor(proxy3_pure_noise) else proxy3_pure_noise\n            )\n\n    # Decompose into components for analysis\n    expected_a_squared = (advantages_scalar**2).mean()\n    expected_w = w_values.mean()\n\n    metrics.update(\n        {\n            # Proxy 1: Signal Strength ||ḡ||²\n            \"variance_proxy/proxy1_signal_strength\": (\n                proxy1_signal_strength if proxy1_signal_strength is not None else 0.0\n            ),\n            # Proxy 2: Total Power E[||ĝ_τ||²]\n            \"variance_proxy/proxy2_total_power\": proxy2_total_power.detach().item(),\n            # Proxy 3: Pure Noise - Variance of Mean Vector\n            \"variance_proxy/proxy3_pure_noise\": proxy3_pure_noise if proxy3_pure_noise is not None else 0.0,\n            # Component metrics for debugging\n            \"variance_proxy/expected_a_squared\": expected_a_squared.detach().item(),\n            \"variance_proxy/expected_w\": expected_w.detach().item(),\n        }\n    )\n\n    return metrics\n\n\ndef bootstrap_metric(\n    data: list[Any],\n    subset_size: int,\n    reduce_fns: list[Callable[[np.ndarray], float]],\n    n_bootstrap: int = 1000,\n    seed: int = 42,\n) -> list[tuple[float, float]]:\n    \"\"\"\n    Performs bootstrap resampling to estimate statistics of metrics.\n\n    This function uses bootstrap resampling to estimate the mean and standard deviation\n    of metrics computed by the provided reduction functions on random subsets of the data.\n\n    Args:\n        data: List of data points to bootstrap from.\n        subset_size: Size of each bootstrap sample.\n        reduce_fns: List of functions that compute a metric from a subset of data.\n        n_bootstrap: Number of bootstrap iterations. Defaults to 1000.\n        seed: Random seed for reproducibility. Defaults to 42.\n\n    Returns:\n        A list of tuples, where each tuple contains (mean, std) for a metric\n        corresponding to each reduction function in reduce_fns.\n\n    Example:\n        >>> data = [1, 2, 3, 4, 5]\n        >>> reduce_fns = [np.mean, np.max]\n        >>> bootstrap_metric(data, 3, reduce_fns)\n        [(3.0, 0.5), (4.5, 0.3)]  # Example values\n    \"\"\"\n    np.random.seed(seed)\n    data_np = np.array(data, dtype=object)\n    n_data = len(data_np)\n\n    # generate bootstrap indices, shape: (n_bootstrap, subset_size)\n    bootstrap_idxs = np.random.choice(n_data, size=(n_bootstrap, subset_size), replace=True)\n\n    # pre-allocate result array, shape: (n_fns, n_bootstrap)\n    n_fns = len(reduce_fns)\n    metric_results = np.empty((n_fns, n_bootstrap), dtype=np.float64)\n\n    # compute metric results for each bootstrap sample\n    for fn_idx, reduce_fn in enumerate(reduce_fns):\n        # bootstrap sample and compute metric\n        for boot_idx in range(n_bootstrap):\n            sample = data_np[bootstrap_idxs[boot_idx]]\n            metric_results[fn_idx, boot_idx] = reduce_fn(sample)\n\n    # compute mean and std for each metric function\n    result = [\n        (float(np.mean(metric_results[fn_idx])), float(np.std(metric_results[fn_idx]))) for fn_idx in range(n_fns)\n    ]\n    return result\n\n\ndef calc_maj_val(data: list[dict[str, Any]], vote_key: str, val_key: str) -> float:\n    \"\"\"\n    Calculate a value based on majority voting.\n\n    This function identifies the most common value for a specified vote key\n    in the data, then returns the corresponding value for that majority vote.\n\n    Args:\n        data: List of dictionaries, where each dictionary contains both vote_key and val_key.\n        vote_key: The key in each dictionary used for voting/counting.\n        val_key: The key in each dictionary whose value will be returned for the majority vote.\n\n    Returns:\n        The value associated with the most common vote.\n\n    Example:\n        >>> data = [\n        ...     {\"pred\": \"A\", \"val\": 0.9},\n        ...     {\"pred\": \"B\", \"val\": 0.8},\n        ...     {\"pred\": \"A\", \"val\": 0.7}\n        ... ]\n        >>> calc_maj_val(data, vote_key=\"pred\", val_key=\"val\")\n        0.9  # Returns the first \"val\" for the majority vote \"A\"\n    \"\"\"\n    vote2vals = defaultdict(list)\n    for d in data:\n        vote2vals[d[vote_key]].append(d[val_key])\n\n    vote2cnt = {k: len(v) for k, v in vote2vals.items()}\n    maj_vote = max(vote2cnt, key=vote2cnt.get)\n\n    maj_val = vote2vals[maj_vote][0]\n\n    return maj_val\n\n\ndef process_validation_metrics(\n    data_sources: list[str], sample_uids: list[str], infos_dict: dict[str, list[Any]], seed: int = 42\n) -> dict[str, dict[str, dict[str, float]]]:\n    \"\"\"\n    Process validation metrics into a structured format with statistical analysis.\n\n    This function organizes validation metrics by data source and prompt, then computes\n    various statistical measures including means, standard deviations, best/worst values,\n    and majority voting results. It also performs bootstrap sampling to estimate statistics\n    for different sample sizes.\n\n    Args:\n        data_sources: List of data source identifiers for each sample.\n        sample_uids: List of sample uids corresponding to each sample.\n        infos_dict: Dictionary mapping variable names to lists of values for each sample.\n        seed: Random seed for bootstrap sampling. Defaults to 42.\n\n    Returns:\n        A nested dictionary with the structure:\n        {\n            data_source: {\n                variable_name: {\n                    metric_name: value\n                }\n            }\n        }\n\n        Where metric_name includes:\n        - \"mean@N\": Mean value across N samples\n        - \"std@N\": Standard deviation across N samples\n        - \"best@N/mean\": Mean of the best values in bootstrap samples of size N\n        - \"best@N/std\": Standard deviation of the best values in bootstrap samples\n        - \"worst@N/mean\": Mean of the worst values in bootstrap samples\n        - \"worst@N/std\": Standard deviation of the worst values in bootstrap samples\n        - \"maj@N/mean\": Mean of majority voting results in bootstrap samples (if \"pred\" exists)\n        - \"maj@N/std\": Standard deviation of majority voting results (if \"pred\" exists)\n\n    Example:\n        >>> data_sources = [\"source1\", \"source1\", \"source2\"]\n        >>> sample_uids = [\"uid1\", \"uid1\", \"uid2\"]\n        >>> infos_dict = {\"score\": [0.8, 0.9, 0.7], \"pred\": [\"A\", \"A\", \"B\"]}\n        >>> result = process_validation_metrics(data_sources, sample_uids, infos_dict)\n        >>> # result will contain statistics for each data source and variable\n    \"\"\"\n    # Group metrics by data source, prompt and variable\n    data_src2uid2var2vals = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))\n    for sample_idx, data_source in enumerate(data_sources):\n        uid = sample_uids[sample_idx]\n        var2vals = data_src2uid2var2vals[data_source][uid]\n        for var_name, var_vals in infos_dict.items():\n            var2vals[var_name].append(var_vals[sample_idx])\n\n    np_mean = np.mean\n    np_std = np.std\n    reduce_fns_best_worst = [np.max, np.min]\n    n_bootstrap = 1000\n\n    # 2. cache ns list\n    def gen_ns(n_resps: int) -> list[int]:\n        if n_resps <= 1:\n            return []\n        ns = []\n        n = 2\n        while n < n_resps:\n            ns.append(n)\n            n *= 2\n        ns.append(n_resps)\n        return ns\n\n    ns_cache = {}\n\n    # 3. cache metric results\n    data_src2uid2var2metric = {}\n\n    # 4. flatten loop\n    for data_source, uid2var2vals in data_src2uid2var2vals.items():\n        # create uid dict\n        uid_dict = data_src2uid2var2metric.setdefault(data_source, {})\n\n        for uid, var2vals in uid2var2vals.items():\n            pred_vals = var2vals.get(\"pred\")\n            has_pred = pred_vals is not None\n            var_dict = uid_dict.setdefault(uid, {})\n\n            for var_name, var_vals in var2vals.items():\n                # skip empty or string values\n                if not var_vals or isinstance(var_vals[0], str):\n                    continue\n\n                # compute mean and std\n                n_resps = len(var_vals)\n                metric = {f\"mean@{n_resps}\": float(np_mean(var_vals))}\n\n                if n_resps > 1:\n                    metric[f\"std@{n_resps}\"] = float(np_std(var_vals))\n\n                    # cache ns list\n                    if n_resps not in ns_cache:\n                        ns_cache[n_resps] = gen_ns(n_resps)\n                    ns = ns_cache[n_resps]\n\n                    # compute best/worst metrics\n                    for n in ns:\n                        # compute best/worst metrics\n                        (bon_mean, bon_std), (won_mean, won_std) = bootstrap_metric(\n                            data=var_vals,\n                            subset_size=n,\n                            reduce_fns=reduce_fns_best_worst,\n                            n_bootstrap=n_bootstrap,\n                            seed=seed,\n                        )\n                        metric[f\"best@{n}/mean\"] = bon_mean\n                        metric[f\"best@{n}/std\"] = bon_std\n                        metric[f\"worst@{n}/mean\"] = won_mean\n                        metric[f\"worst@{n}/std\"] = won_std\n\n                        # compute maj metrics\n                        if has_pred:\n                            # create vote_data\n                            vote_data = [\n                                {\"val\": val, \"pred\": pred} for val, pred in zip(var_vals, pred_vals, strict=True)\n                            ]\n                            # compute maj metrics\n                            [(maj_n_mean, maj_n_std)] = bootstrap_metric(\n                                data=vote_data,\n                                subset_size=n,\n                                reduce_fns=[partial(calc_maj_val, vote_key=\"pred\", val_key=\"val\")],\n                                n_bootstrap=n_bootstrap,\n                                seed=seed,\n                            )\n                            metric[f\"maj@{n}/mean\"] = maj_n_mean\n                            metric[f\"maj@{n}/std\"] = maj_n_std\n\n                var_dict[var_name] = metric\n\n    # Aggregate metrics across uids\n    data_src2var2metric2uid_vals = defaultdict(lambda: defaultdict(lambda: defaultdict(list)))\n    for data_source, uid2var2metric in data_src2uid2var2metric.items():\n        for uid, var2metric in uid2var2metric.items():\n            for var_name, metric in var2metric.items():\n                for metric_name, metric_val in metric.items():\n                    data_src2var2metric2uid_vals[data_source][var_name][metric_name].append(metric_val)\n\n    data_src2var2metric2val = defaultdict(lambda: defaultdict(lambda: defaultdict(float)))\n    for data_source, var2metric2uid_vals in data_src2var2metric2uid_vals.items():\n        for var_name, metric2uid_vals in var2metric2uid_vals.items():\n            for metric_name, uid_vals in metric2uid_vals.items():\n                data_src2var2metric2val[data_source][var_name][metric_name] = np.mean(uid_vals)\n    return data_src2var2metric2val\n"
  },
  {
    "path": "verl/trainer/ppo/prefix_grouper_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom __future__ import annotations\n\nimport torch\nfrom prefix_grouper import PrefixGrouper\n\nfrom verl.utils.torch_functional import logprobs_from_logits\n\n\ndef build_position_ids_for_prefix_grouper(prefix_grouper: PrefixGrouper) -> torch.Tensor:\n    \"\"\"Build position_ids for PrefixGrouper where each response restarts from prefix_len.\"\"\"\n    num_samples = len(prefix_grouper.group_info)\n    max_len = prefix_grouper.padding_mask.size(1)\n    device = prefix_grouper.padding_mask.device\n\n    position_ids = torch.zeros(num_samples, max_len, dtype=torch.long, device=device)\n\n    for i, group in enumerate(prefix_grouper.group_info):\n        prefix_len = group.prefix_len\n\n        position_ids[i, :prefix_len] = torch.arange(prefix_len, device=device)\n        cur_pos = prefix_len\n        for suffix_len in group.suffix_lens:\n            if suffix_len > 0:\n                position_ids[i, cur_pos : cur_pos + suffix_len] = torch.arange(\n                    prefix_len, prefix_len + suffix_len, device=device\n                )\n                cur_pos += suffix_len\n\n    return position_ids\n\n\ndef build_pg_from_micro_batch(\n    micro_batch: dict,\n    pad_token_id: int,\n    padding_mode: str = \"right\",\n):\n    \"\"\"Build PrefixGrouper from micro_batch dict containing prompts, responses, response_mask, uid.\"\"\"\n    prompts = micro_batch[\"prompts\"]\n    responses = micro_batch[\"responses\"]\n    response_mask = micro_batch[\"response_mask\"]\n    uids = micro_batch[\"uid\"]\n\n    bs = responses.size(0)\n\n    group_sizes = []\n    cur = 1\n    for i in range(1, bs):\n        if uids[i] == uids[i - 1]:\n            cur += 1\n        else:\n            group_sizes.append(cur)\n            cur = 1\n    group_sizes.append(cur)\n\n    prefix_indices = []\n    cursor = 0\n    for gs in group_sizes:\n        prefix_indices.append(cursor)\n        cursor += gs\n    prefix_indices = torch.tensor(prefix_indices, device=prompts.device)\n\n    prefix_ids = prompts.index_select(0, prefix_indices)\n    prefix_mask = prefix_ids.ne(pad_token_id)\n\n    prefix_grouper = PrefixGrouper.from_ungrouped_masks(\n        prefix_mask=prefix_mask,\n        suffix_mask=response_mask,\n        group_sizes=group_sizes,\n        padding_mode=padding_mode,\n        device=prompts.device,\n    )\n\n    concat_input_ids = prefix_grouper.concat_input(prefix_ids, prefix_mask, responses, response_mask)\n\n    attention_mask = prefix_grouper.padding_mask\n\n    position_ids = build_position_ids_for_prefix_grouper(prefix_grouper)\n\n    return (\n        prefix_grouper,\n        concat_input_ids,\n        attention_mask,\n        position_ids,\n        responses,\n        response_mask,\n    )\n\n\ndef pg_forward(\n    model,\n    prefix_grouper,\n    concat_input_ids,\n    attention_mask,\n    position_ids,\n    completion_ids,\n    completion_mask,\n    *,\n    temperature=1.0,\n    padding_mode=\"right\",\n    include_prefix_last=1,\n    calculate_entropy=False,\n    entropy_fn=None,\n):\n    logits = model(\n        input_ids=concat_input_ids,\n        attention_mask=attention_mask,\n        position_ids=position_ids,\n        use_cache=False,\n        prefix_grouper=prefix_grouper,\n    ).logits\n\n    prefix_out, prefix_mask, suffix_out_raw, suffix_mask_raw = prefix_grouper.split_output(\n        logits, include_prefix_last=include_prefix_last\n    )\n\n    completion_ids_right = prefix_grouper.convert_padding(\n        completion_ids,\n        completion_mask,\n        padding_mode=padding_mode,\n    )\n\n    suffix_out = suffix_out_raw[:, :-1].float()\n    suffix_mask = suffix_mask_raw[:, 1:]\n\n    suffix_out /= temperature\n\n    log_probs = logprobs_from_logits(suffix_out, completion_ids_right)\n\n    entropy = None\n    if calculate_entropy and entropy_fn is not None:\n        entropy = entropy_fn(suffix_out)\n\n    return log_probs, entropy, suffix_mask\n\n\ndef forward_micro_batch_with_prefix_grouper(\n    micro_batch: dict,\n    model,\n    temperature: float,\n    calculate_entropy: bool,\n    device_name: str,\n    param_dtype,\n    use_chunking_entropy: bool = False,\n):\n    \"\"\"\n    Forward pass using PrefixGrouper for shared-prefix optimization.\n\n    Args:\n        micro_batch: Dict containing prompts, responses, response_mask, uid, etc.\n        model: The actor module.\n        temperature: Temperature for logits scaling.\n        calculate_entropy: Whether to compute entropy.\n        device_name: Device name for autocast.\n        param_dtype: Parameter dtype for autocast.\n        use_chunking_entropy: Whether to use chunking entropy function.\n\n    Returns:\n        tuple: (entropy, log_probs) where entropy may be None if not calculated.\n    \"\"\"\n    import verl.utils.torch_functional as verl_F\n\n    entropy_fn = None\n    if calculate_entropy:\n        if use_chunking_entropy:\n            entropy_fn = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_fn = verl_F.entropy_from_logits\n\n    pad_token_id = micro_batch.get(\"pad_token_id\", 0)\n\n    (\n        prefix_grouper,\n        concat_input_ids,\n        attention_mask,\n        position_ids,\n        responses,\n        response_mask,\n    ) = build_pg_from_micro_batch(\n        micro_batch,\n        pad_token_id=pad_token_id,\n        padding_mode=\"right\",\n    )\n\n    with torch.autocast(device_type=device_name, dtype=param_dtype):\n        log_probs, entropy, suffix_mask_from_pg = pg_forward(\n            model=model,\n            prefix_grouper=prefix_grouper,\n            concat_input_ids=concat_input_ids,\n            attention_mask=attention_mask,\n            position_ids=position_ids,\n            completion_ids=responses,\n            completion_mask=response_mask,\n            temperature=temperature,\n            padding_mode=\"right\",\n            include_prefix_last=1,\n            calculate_entropy=calculate_entropy,\n            entropy_fn=entropy_fn,\n        )\n\n    # Zero out padding positions\n    padding_mask = suffix_mask_from_pg == 0\n    log_probs = log_probs.masked_fill(padding_mask, 0.0)\n    if entropy is not None:\n        entropy = entropy.masked_fill(padding_mask, 0.0)\n\n    # Pad to target response length if needed\n    target_response_length = responses.size(1)\n    if log_probs.size(1) != target_response_length:\n        batch_size = log_probs.size(0)\n        current_len = log_probs.size(1)\n\n        full_log_probs = log_probs.new_zeros(batch_size, target_response_length)\n        full_log_probs[:, :current_len] = log_probs\n        log_probs = full_log_probs\n\n        if entropy is not None:\n            full_entropy = entropy.new_zeros(batch_size, target_response_length)\n            full_entropy[:, :current_len] = entropy\n            entropy = full_entropy\n\n    return entropy, log_probs\n"
  },
  {
    "path": "verl/trainer/ppo/ray_trainer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nPPO Trainer with Ray-based single controller.\nThis trainer supports model-agonistic model initialization with huggingface\n\"\"\"\n\nimport json\nimport os\nimport uuid\nfrom collections import defaultdict\nfrom copy import deepcopy\nfrom pprint import pprint\nfrom typing import Any, Optional\n\nimport numpy as np\nimport torch\nfrom omegaconf import OmegaConf, open_dict\nfrom torch.utils.data import Dataset, Sampler\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom tqdm import tqdm\n\nfrom verl import DataProto\nfrom verl.checkpoint_engine import CheckpointEngineManager\nfrom verl.experimental.dataset.sampler import AbstractCurriculumSampler\nfrom verl.protocol import pad_dataproto_to_divisor, unpad_dataproto\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayWorkerGroup, ResourcePoolManager\nfrom verl.single_controller.ray.base import create_colocated_worker_cls\nfrom verl.trainer.config import AlgoConfig\nfrom verl.trainer.ppo import core_algos\nfrom verl.trainer.ppo.core_algos import AdvantageEstimator, agg_loss\nfrom verl.trainer.ppo.metric_utils import (\n    compute_data_metrics,\n    compute_throughout_metrics,\n    compute_timing_metrics,\n    compute_variance_proxy_metrics,\n    process_validation_metrics,\n)\nfrom verl.trainer.ppo.reward import extract_reward\nfrom verl.trainer.ppo.utils import Role, WorkerType, need_critic, need_reference_policy, need_reward_model\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, should_save_ckpt_esi\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.debug import marked_timer\nfrom verl.utils.import_utils import load_class_from_fqn\nfrom verl.utils.metric import reduce_metrics\nfrom verl.utils.py_functional import rename_dict\nfrom verl.utils.rollout_skip import RolloutSkip\nfrom verl.utils.seqlen_balancing import calculate_workload, get_seqlen_balanced_partitions, log_seqlen_unbalance\nfrom verl.utils.torch_functional import masked_mean\nfrom verl.utils.tracking import ValidationGenerationsLogger\nfrom verl.workers.config import FSDPEngineConfig\nfrom verl.workers.utils.padding import left_right_2_no_padding, no_padding_2_padding\n\n\ndef apply_kl_penalty(data: DataProto, kl_ctrl: core_algos.AdaptiveKLController, kl_penalty=\"kl\"):\n    \"\"\"Apply KL penalty to the token-level rewards.\n\n    This function computes the KL divergence between the reference policy and current policy,\n    then applies a penalty to the token-level rewards based on this divergence.\n\n    Args:\n        data (DataProto): The data containing batched model outputs and inputs.\n        kl_ctrl (core_algos.AdaptiveKLController): Controller for adaptive KL penalty.\n        kl_penalty (str, optional): Type of KL penalty to apply. Defaults to \"kl\".\n\n    Returns:\n        tuple: A tuple containing:\n            - The updated data with token-level rewards adjusted by KL penalty\n            - A dictionary of metrics related to the KL penalty\n    \"\"\"\n    response_mask = data.batch[\"response_mask\"]\n    token_level_scores = data.batch[\"token_level_scores\"]\n    batch_size = data.batch.batch_size[0]\n\n    # compute kl between ref_policy and current policy\n    # When apply_kl_penalty, algorithm.use_kl_in_reward=True, so the reference model has been enabled.\n    kld = core_algos.kl_penalty(\n        data.batch[\"old_log_probs\"], data.batch[\"ref_log_prob\"], kl_penalty=kl_penalty\n    )  # (batch_size, response_length)\n    kld = kld * response_mask\n    beta = kl_ctrl.value\n\n    token_level_rewards = token_level_scores - beta * kld\n\n    current_kl = masked_mean(kld, mask=response_mask, axis=-1)  # average over sequence\n    current_kl = torch.mean(current_kl, dim=0).item()\n\n    # according to https://github.com/huggingface/trl/blob/951ca1841f29114b969b57b26c7d3e80a39f75a0/trl/trainer/ppo_trainer.py#L837\n    kl_ctrl.update(current_kl=current_kl, n_steps=batch_size)\n    data.batch[\"token_level_rewards\"] = token_level_rewards\n\n    metrics = {\"actor/reward_kl_penalty\": current_kl, \"actor/reward_kl_penalty_coeff\": beta}\n\n    return data, metrics\n\n\ndef compute_response_mask(data: DataProto):\n    \"\"\"Compute the attention mask for the response part of the sequence.\n\n    This function extracts the portion of the attention mask that corresponds to the model's response,\n    which is used for masking computations that should only apply to response tokens.\n\n    Args:\n        data (DataProto): The data containing batched model outputs and inputs.\n\n    Returns:\n        torch.Tensor: The attention mask for the response tokens.\n    \"\"\"\n    responses = data.batch[\"responses\"]\n    response_length = responses.size(1)\n    attention_mask = data.batch[\"attention_mask\"]\n    return attention_mask[:, -response_length:]\n\n\ndef compute_advantage(\n    data: DataProto,\n    adv_estimator: AdvantageEstimator,\n    gamma: float = 1.0,\n    lam: float = 1.0,\n    num_repeat: int = 1,\n    norm_adv_by_std_in_grpo: bool = True,\n    config: Optional[AlgoConfig] = None,\n) -> DataProto:\n    \"\"\"Compute advantage estimates for policy optimization.\n\n    This function computes advantage estimates using various estimators like GAE, GRPO, REINFORCE++, etc.\n    The advantage estimates are used to guide policy optimization in RL algorithms.\n\n    Args:\n        data (DataProto): The data containing batched model outputs and inputs.\n        adv_estimator (AdvantageEstimator): The advantage estimator to use (e.g., GAE, GRPO, REINFORCE++).\n        gamma (float, optional): Discount factor for future rewards. Defaults to 1.0.\n        lam (float, optional): Lambda parameter for GAE. Defaults to 1.0.\n        num_repeat (int, optional): Number of times to repeat the computation. Defaults to 1.\n        norm_adv_by_std_in_grpo (bool, optional): Whether to normalize advantages by standard deviation in\n            GRPO. Defaults to True.\n        config (dict, optional): Configuration dictionary for algorithm settings. Defaults to None.\n\n    Returns:\n        DataProto: The updated data with computed advantages and returns.\n    \"\"\"\n    # Back-compatible with trainers that do not compute response mask in fit\n    if \"response_mask\" not in data.batch.keys():\n        data.batch[\"response_mask\"] = compute_response_mask(data)\n    # prepare response group\n    if adv_estimator == AdvantageEstimator.GAE:\n        # Compute advantages and returns using Generalized Advantage Estimation (GAE)\n        advantages, returns = core_algos.compute_gae_advantage_return(\n            token_level_rewards=data.batch[\"token_level_rewards\"],\n            values=data.batch[\"values\"],\n            response_mask=data.batch[\"response_mask\"],\n            gamma=gamma,\n            lam=lam,\n        )\n        data.batch[\"advantages\"] = advantages\n        data.batch[\"returns\"] = returns\n        if config.get(\"use_pf_ppo\", False):\n            data = core_algos.compute_pf_ppo_reweight_data(\n                data,\n                config.pf_ppo.get(\"reweight_method\"),\n                config.pf_ppo.get(\"weight_pow\"),\n            )\n    elif adv_estimator == AdvantageEstimator.GRPO:\n        # Initialize the mask for GRPO calculation\n        grpo_calculation_mask = data.batch[\"response_mask\"]\n\n        # Call compute_grpo_outcome_advantage with parameters matching its definition\n        advantages, returns = core_algos.compute_grpo_outcome_advantage(\n            token_level_rewards=data.batch[\"token_level_rewards\"],\n            response_mask=grpo_calculation_mask,\n            index=data.non_tensor_batch[\"uid\"],\n            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n        )\n        data.batch[\"advantages\"] = advantages\n        data.batch[\"returns\"] = returns\n    else:\n        # handle all other adv estimator type other than GAE and GRPO\n        adv_estimator_fn = core_algos.get_adv_estimator_fn(adv_estimator)\n        adv_kwargs = {\n            \"token_level_rewards\": data.batch[\"token_level_rewards\"],\n            \"response_mask\": data.batch[\"response_mask\"],\n            \"config\": config,\n        }\n        if \"uid\" in data.non_tensor_batch:  # optional\n            adv_kwargs[\"index\"] = data.non_tensor_batch[\"uid\"]\n        if \"reward_baselines\" in data.batch:  # optional\n            adv_kwargs[\"reward_baselines\"] = data.batch[\"reward_baselines\"]\n        # GDPO: pass raw data for per-dimension reward extraction\n        if adv_estimator in (AdvantageEstimator.GDPO, \"gdpo\"):\n            adv_kwargs[\"non_tensor_batch\"] = data.non_tensor_batch\n            adv_kwargs[\"batch\"] = data.batch\n        # Add sum_pi_squared for Optimal Token Baseline\n        if adv_estimator in (AdvantageEstimator.OPTIMAL_TOKEN_BASELINE, AdvantageEstimator.TIR_OPTIMAL_TOKEN_BASELINE):\n            # Check if sum_pi_squared is available\n            assert \"sum_pi_squared\" in data.batch, (\n                \"Step-dependent optimal baseline requires sum_pi_squared from actor. \"\n                \"Please set actor.calculate_sum_pi_squared=True in config.\"\n            )\n            adv_kwargs[\"sum_pi_squared\"] = data.batch[\"sum_pi_squared\"]\n            # old_log_probs needed for path-variance proxy: w_t = 1 - 2*exp(old_log_probs) + sum_pi_squared\n            adv_kwargs[\"old_log_probs\"] = data.batch[\"old_log_probs\"]\n            # Get pre-computed rollout IS weights if available\n            rollout_is_weights = data.batch.get(\"rollout_is_weights\", None)\n            adv_kwargs[\"rollout_is_weights\"] = rollout_is_weights\n\n        # calculate advantage estimator\n        advantages, returns = adv_estimator_fn(**adv_kwargs)\n        data.batch[\"advantages\"] = advantages\n        data.batch[\"returns\"] = returns\n    return data\n\n\nclass RayPPOTrainer:\n    \"\"\"Distributed PPO trainer using Ray for scalable reinforcement learning.\n\n    This trainer orchestrates distributed PPO training across multiple nodes and GPUs,\n    managing actor rollouts, critic training, and reward computation with Ray backend.\n    Supports various model architectures including FSDP, Megatron, vLLM, and SGLang integration.\n    \"\"\"\n\n    # TODO: support each role have individual ray_worker_group_cls,\n    # i.e., support different backend of different role\n    def __init__(\n        self,\n        config,\n        tokenizer,\n        role_worker_mapping: dict[Role, WorkerType],\n        resource_pool_manager: ResourcePoolManager,\n        ray_worker_group_cls: type[RayWorkerGroup] = RayWorkerGroup,\n        processor=None,\n        train_dataset: Optional[Dataset] = None,\n        val_dataset: Optional[Dataset] = None,\n        collate_fn=None,\n        train_sampler: Optional[Sampler] = None,\n        device_name=None,\n    ):\n        \"\"\"\n        Initialize distributed PPO trainer with Ray backend.\n        Note that this trainer runs on the driver process on a single CPU/GPU node.\n\n        Args:\n            config: Configuration object containing training parameters.\n            tokenizer: Tokenizer used for encoding and decoding text.\n            role_worker_mapping (dict[Role, WorkerType]): Mapping from roles to worker classes.\n            resource_pool_manager (ResourcePoolManager): Manager for Ray resource pools.\n            ray_worker_group_cls (RayWorkerGroup, optional): Class for Ray worker groups. Defaults to RayWorkerGroup.\n            processor: Optional data processor, used for multimodal data\n            train_dataset (Optional[Dataset], optional): Training dataset. Defaults to None.\n            val_dataset (Optional[Dataset], optional): Validation dataset. Defaults to None.\n            collate_fn: Function to collate data samples into batches.\n            train_sampler (Optional[Sampler], optional): Sampler for the training dataset. Defaults to None.\n            device_name (str, optional): Device name for training (e.g., \"cuda\", \"cpu\"). Defaults to None.\n        \"\"\"\n\n        # Store the tokenizer for text processing\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.config = config\n\n        self.hybrid_engine = config.actor_rollout_ref.hybrid_engine\n        assert self.hybrid_engine, \"Currently, only support hybrid engine\"\n\n        if self.hybrid_engine:\n            assert Role.ActorRollout in role_worker_mapping or Role.ActorRolloutRef in role_worker_mapping, (\n                f\"{role_worker_mapping.keys()=}\"\n            )\n\n        self.role_worker_mapping = role_worker_mapping\n        self.resource_pool_manager = resource_pool_manager\n        self.use_reference_policy = need_reference_policy(self.config)\n\n        self.use_rm = need_reward_model(self.config)\n\n        self.use_critic = need_critic(self.config)\n        self.ray_worker_group_cls = ray_worker_group_cls\n        self.device_name = device_name if device_name else self.config.trainer.device\n        self.validation_generations_logger = ValidationGenerationsLogger(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n        )\n\n        # if ref_in_actor is True, the reference policy will be actor without lora applied\n        lora_rank = config.actor_rollout_ref.model.get(\"lora\", {}).get(\"rank\", 0)\n        if lora_rank <= 0:\n            lora_rank = config.actor_rollout_ref.model.get(\"lora_rank\", 0)\n        self.ref_in_actor = lora_rank > 0 or config.actor_rollout_ref.model.get(\"lora_adapter_path\") is not None\n\n        # define in-reward KL control\n        # kl loss control currently not suppoorted\n        if self.config.algorithm.use_kl_in_reward:\n            self.kl_ctrl_in_reward = core_algos.get_kl_controller(self.config.algorithm.kl_ctrl)\n\n        self.use_prefix_grouper = self.config.actor_rollout_ref.actor.get(\"use_prefix_grouper\", False)\n        self.use_legacy_worker_impl = config.trainer.get(\"use_legacy_worker_impl\", \"auto\")\n\n        self._create_dataloader(train_dataset, val_dataset, collate_fn, train_sampler)\n\n        self.checkpoint_manager = None\n\n    def _create_dataloader(self, train_dataset, val_dataset, collate_fn, train_sampler: Optional[Sampler]):\n        \"\"\"\n        Creates the train and validation dataloaders.\n        \"\"\"\n        # TODO: we have to make sure the batch size is divisible by the dp size\n        from verl.trainer.main_ppo import create_rl_dataset, create_rl_sampler\n\n        if train_dataset is None:\n            train_dataset = create_rl_dataset(\n                self.config.data.train_files,\n                self.config.data,\n                self.tokenizer,\n                self.processor,\n                max_samples=self.config.data.get(\"train_max_samples\", -1),\n            )\n        if val_dataset is None:\n            val_dataset = create_rl_dataset(\n                self.config.data.val_files,\n                self.config.data,\n                self.tokenizer,\n                self.processor,\n                max_samples=self.config.data.get(\"val_max_samples\", -1),\n            )\n        self.train_dataset, self.val_dataset = train_dataset, val_dataset\n\n        if train_sampler is None:\n            train_sampler = create_rl_sampler(self.config.data, self.train_dataset)\n        if collate_fn is None:\n            from verl.utils.dataset.rl_dataset import collate_fn as default_collate_fn\n\n            collate_fn = default_collate_fn\n\n        num_workers = self.config.data[\"dataloader_num_workers\"]\n\n        self.train_dataloader = StatefulDataLoader(\n            dataset=self.train_dataset,\n            batch_size=self.config.data.get(\"gen_batch_size\", self.config.data.train_batch_size),\n            num_workers=num_workers,\n            drop_last=True,\n            collate_fn=collate_fn,\n            sampler=train_sampler,\n        )\n\n        val_batch_size = self.config.data.val_batch_size  # Prefer config value if set\n        if val_batch_size is None:\n            val_batch_size = len(self.val_dataset)\n\n        self.val_dataloader = StatefulDataLoader(\n            dataset=self.val_dataset,\n            batch_size=val_batch_size,\n            num_workers=num_workers,\n            shuffle=self.config.data.get(\"validation_shuffle\", True),\n            drop_last=False,\n            collate_fn=collate_fn,\n        )\n\n        assert len(self.train_dataloader) >= 1, \"Train dataloader is empty!\"\n        assert len(self.val_dataloader) >= 1, \"Validation dataloader is empty!\"\n\n        print(\n            f\"Size of train dataloader: {len(self.train_dataloader)}, Size of val dataloader: \"\n            f\"{len(self.val_dataloader)}\"\n        )\n\n        total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs\n\n        if self.config.trainer.total_training_steps is not None:\n            total_training_steps = self.config.trainer.total_training_steps\n\n        self.total_training_steps = total_training_steps\n        print(f\"Total training steps: {self.total_training_steps}\")\n\n        try:\n            OmegaConf.set_struct(self.config, True)\n            with open_dict(self.config):\n                if OmegaConf.select(self.config, \"actor_rollout_ref.actor.optim\"):\n                    self.config.actor_rollout_ref.actor.optim.total_training_steps = total_training_steps\n                if OmegaConf.select(self.config, \"critic.optim\"):\n                    self.config.critic.optim.total_training_steps = total_training_steps\n        except Exception as e:\n            print(f\"Warning: Could not set total_training_steps in config. Structure missing? Error: {e}\")\n\n    def _dump_generations(self, inputs, outputs, gts, scores, reward_extra_infos_dict, dump_path):\n        \"\"\"Dump rollout/validation samples as JSONL.\"\"\"\n        os.makedirs(dump_path, exist_ok=True)\n        filename = os.path.join(dump_path, f\"{self.global_steps}.jsonl\")\n\n        n = len(inputs)\n        base_data = {\n            \"input\": inputs,\n            \"output\": outputs,\n            \"gts\": gts,\n            \"score\": scores,\n            \"step\": [self.global_steps] * n,\n        }\n\n        for k, v in reward_extra_infos_dict.items():\n            if len(v) == n:\n                base_data[k] = v\n\n        lines = []\n        for i in range(n):\n            entry = {k: v[i] for k, v in base_data.items()}\n            lines.append(json.dumps(entry, ensure_ascii=False))\n\n        with open(filename, \"w\") as f:\n            f.write(\"\\n\".join(lines) + \"\\n\")\n\n        print(f\"Dumped generations to {filename}\")\n\n    def _log_rollout_data(\n        self, batch: DataProto, reward_extra_infos_dict: dict, timing_raw: dict, rollout_data_dir: str\n    ):\n        \"\"\"Log rollout data to disk.\n        Args:\n            batch (DataProto): The batch containing rollout data\n            reward_extra_infos_dict (dict): Additional reward information to log\n            timing_raw (dict): Timing information for profiling\n            rollout_data_dir (str): Directory path to save the rollout data\n        \"\"\"\n        with marked_timer(\"dump_rollout_generations\", timing_raw, color=\"green\"):\n            inputs = self.tokenizer.batch_decode(batch.batch[\"prompts\"], skip_special_tokens=True)\n            outputs = self.tokenizer.batch_decode(batch.batch[\"responses\"], skip_special_tokens=True)\n            scores = batch.batch[\"token_level_scores\"].sum(-1).cpu().tolist()\n            sample_gts = [item.non_tensor_batch.get(\"reward_model\", {}).get(\"ground_truth\", None) for item in batch]\n\n            reward_extra_infos_to_dump = reward_extra_infos_dict.copy()\n            if \"request_id\" in batch.non_tensor_batch:\n                reward_extra_infos_dict.setdefault(\n                    \"request_id\",\n                    batch.non_tensor_batch[\"request_id\"].tolist(),\n                )\n\n            self._dump_generations(\n                inputs=inputs,\n                outputs=outputs,\n                gts=sample_gts,\n                scores=scores,\n                reward_extra_infos_dict=reward_extra_infos_to_dump,\n                dump_path=rollout_data_dir,\n            )\n\n    def _maybe_log_val_generations(self, inputs, outputs, scores):\n        \"\"\"Log a table of validation samples to the configured logger (wandb or swanlab)\"\"\"\n\n        generations_to_log = self.config.trainer.log_val_generations\n\n        if generations_to_log == 0:\n            return\n\n        import numpy as np\n\n        # Create tuples of (input, output, score) and sort by input text\n        samples = list(zip(inputs, outputs, scores, strict=True))\n        samples.sort(key=lambda x: x[0])  # Sort by input text\n\n        # Use fixed random seed for deterministic shuffling\n        rng = np.random.RandomState(42)\n        rng.shuffle(samples)\n\n        # Take first N samples after shuffling\n        samples = samples[:generations_to_log]\n\n        # Log to each configured logger\n        self.validation_generations_logger.log(self.config.trainer.logger, samples, self.global_steps)\n\n    def _get_gen_batch(self, batch: DataProto) -> DataProto:\n        reward_keys = set({\"data_source\", \"reward_model\", \"extra_info\", \"uid\"}) & batch.non_tensor_batch.keys()\n\n        # pop those keys for generation\n        batch_keys_to_pop = []\n        non_tensor_batch_keys_to_pop = set(batch.non_tensor_batch.keys()) - reward_keys\n        gen_batch = batch.pop(\n            batch_keys=batch_keys_to_pop,\n            non_tensor_batch_keys=list(non_tensor_batch_keys_to_pop),\n        )\n\n        # For agent loop, we need reward model keys to compute score.\n        gen_batch.non_tensor_batch.update(batch.non_tensor_batch)\n\n        return gen_batch\n\n    def _compute_reward_colocate(self, batch: DataProto) -> tuple[torch.Tensor, dict[str, Any]] | torch.Tensor:\n        \"\"\"\n        compute reward use colocate reward model\n        \"\"\"\n        assert self.reward_loop_manager is not None, \"RewardLoopManager is None\"\n        batch_reward = self.reward_loop_manager.compute_rm_score(batch)\n        return batch_reward\n\n    def _validate(self, merged: bool = False):\n        data_source_lst = []\n        reward_extra_infos_dict: dict[str, list] = defaultdict(list)\n\n        # Lists to collect samples for the table\n        sample_inputs = []\n        sample_outputs = []\n        sample_gts = []\n        sample_scores = []\n        sample_turns = []\n        sample_uids = []\n\n        for test_data in self.val_dataloader:\n            test_batch = DataProto.from_single_dict(test_data)\n\n            if \"uid\" not in test_batch.non_tensor_batch:\n                test_batch.non_tensor_batch[\"uid\"] = np.array(\n                    [str(uuid.uuid4()) for _ in range(len(test_batch.batch))], dtype=object\n                )\n\n            # repeat test batch\n            test_batch = test_batch.repeat(\n                repeat_times=self.config.actor_rollout_ref.rollout.val_kwargs.n, interleave=True\n            )\n\n            ground_truths = [\n                item.non_tensor_batch.get(\"reward_model\", {}).get(\"ground_truth\", None) for item in test_batch\n            ]\n            sample_gts.extend(ground_truths)\n\n            test_gen_batch = self._get_gen_batch(test_batch)\n            test_gen_batch.meta_info = {\n                \"eos_token_id\": self.tokenizer.eos_token_id,\n                \"pad_token_id\": self.tokenizer.pad_token_id,\n                \"recompute_log_prob\": False,\n                \"do_sample\": self.config.actor_rollout_ref.rollout.val_kwargs.do_sample,\n                \"validate\": True,\n                \"global_steps\": self.global_steps,\n            }\n            print(f\"test_gen_batch meta info: {test_gen_batch.meta_info}\")\n\n            # pad to be divisible by dp_size\n            size_divisor = self.config.actor_rollout_ref.rollout.agent.num_workers\n            test_gen_batch_padded, pad_size = pad_dataproto_to_divisor(test_gen_batch, size_divisor)\n            test_output_gen_batch_padded = self.async_rollout_manager.generate_sequences(test_gen_batch_padded)\n\n            if self.use_rm and \"rm_scores\" not in test_output_gen_batch_padded.batch.keys():\n                # for colocate reward models, we need to sleep rollout model\n                # to spare GPU memory for reward model\n                self.checkpoint_manager.sleep_replicas()\n                batch_reward = self._compute_reward_colocate(test_output_gen_batch_padded)\n                test_output_gen_batch_padded = test_output_gen_batch_padded.union(batch_reward)\n                # wake up rollout model\n                # replace with wake_up method once supported\n                self.checkpoint_manager.update_weights(self.global_steps)\n\n            # unpad\n            test_output_gen_batch = unpad_dataproto(test_output_gen_batch_padded, pad_size=pad_size)\n\n            print(\"validation generation end\")\n\n            # Store generated outputs\n            output_ids = test_output_gen_batch.batch[\"responses\"]\n            output_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in output_ids]\n            sample_outputs.extend(output_texts)\n\n            test_batch = test_batch.union(test_output_gen_batch)\n            test_batch.meta_info[\"validate\"] = True\n\n            # Store original inputs\n            input_ids = test_batch.batch[\"prompts\"]\n            # TODO: Can we keep special tokens except for padding tokens?\n            input_texts = [self.tokenizer.decode(ids, skip_special_tokens=True) for ids in input_ids]\n            sample_inputs.extend(input_texts)\n            sample_uids.extend(test_batch.non_tensor_batch[\"uid\"])\n\n            # evaluate using reward_function\n            reward_tensor, reward_extra_info = extract_reward(test_batch)\n\n            scores = reward_tensor.sum(-1).cpu().tolist()\n            sample_scores.extend(scores)\n\n            reward_extra_infos_dict[\"reward\"].extend(scores)\n            for key, values in reward_extra_info.items():\n                if key not in reward_extra_infos_dict:\n                    reward_extra_infos_dict[key] = []\n                if isinstance(values, np.ndarray):\n                    reward_extra_infos_dict[key].extend(values.tolist())\n                else:\n                    reward_extra_infos_dict[key].extend(values if isinstance(values, list) else [values])\n\n            # collect num_turns of each prompt\n            if \"__num_turns__\" in test_batch.non_tensor_batch:\n                sample_turns.append(test_batch.non_tensor_batch[\"__num_turns__\"])\n\n            data_source_lst.append(test_batch.non_tensor_batch.get(\"data_source\", [\"unknown\"] * reward_tensor.shape[0]))\n\n        self._maybe_log_val_generations(inputs=sample_inputs, outputs=sample_outputs, scores=sample_scores)\n\n        # dump generations\n        val_data_dir = self.config.trainer.get(\"validation_data_dir\", None)\n        if val_data_dir:\n            self._dump_generations(\n                inputs=sample_inputs,\n                outputs=sample_outputs,\n                gts=sample_gts,\n                scores=sample_scores,\n                reward_extra_infos_dict=reward_extra_infos_dict,\n                dump_path=val_data_dir,\n            )\n\n        for key_info, lst in reward_extra_infos_dict.items():\n            assert len(lst) == 0 or len(lst) == len(sample_scores), f\"{key_info}: {len(lst)=}, {len(sample_scores)=}\"\n\n        if merged:\n            print(\"_merge_validation_results validate result will be merged\")\n            return {\n                \"data_sources\": data_source_lst,\n                \"sample_uids\": sample_uids,\n                \"sample_turns\": sample_turns,\n                \"reward_extra_infos_dict\": reward_extra_infos_dict,\n            }\n        data_sources = np.concatenate(data_source_lst, axis=0)\n        return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns)\n\n    def _val_metrics_update(self, data_sources, sample_uids, reward_extra_infos_dict, sample_turns):\n        data_src2var2metric2val = process_validation_metrics(data_sources, sample_uids, reward_extra_infos_dict)\n        metric_dict = {}\n        for data_source, var2metric2val in data_src2var2metric2val.items():\n            core_var = \"acc\" if \"acc\" in var2metric2val else \"reward\"\n            for var_name, metric2val in var2metric2val.items():\n                n_max = max([int(name.split(\"@\")[-1].split(\"/\")[0]) for name in metric2val.keys()])\n                for metric_name, metric_val in metric2val.items():\n                    if (\n                        (var_name == core_var)\n                        and any(metric_name.startswith(pfx) for pfx in [\"mean\", \"maj\", \"best\"])\n                        and (f\"@{n_max}\" in metric_name)\n                    ):\n                        metric_sec = \"val-core\"\n                    else:\n                        metric_sec = \"val-aux\"\n                    pfx = f\"{metric_sec}/{data_source}/{var_name}/{metric_name}\"\n                    metric_dict[pfx] = metric_val\n\n        if len(sample_turns) > 0:\n            sample_turns = np.concatenate(sample_turns)\n            metric_dict[\"val-aux/num_turns/min\"] = sample_turns.min()\n            metric_dict[\"val-aux/num_turns/max\"] = sample_turns.max()\n            metric_dict[\"val-aux/num_turns/mean\"] = sample_turns.mean()\n\n        return metric_dict\n\n    def _merge_validation_results(self, result_a, result_b):\n        if result_a is None and result_b is None:\n            return {}\n        if result_a is None:\n            result_a = {\"data_sources\": [], \"sample_uids\": [], \"sample_turns\": [], \"reward_extra_infos_dict\": {}}\n        if result_b is None:\n            result_b = {\"data_sources\": [], \"sample_uids\": [], \"sample_turns\": [], \"reward_extra_infos_dict\": {}}\n\n        if not result_a.get(\"data_sources\") and not result_b.get(\"data_sources\"):\n            return {}\n\n        data_sources = np.concatenate(result_a[\"data_sources\"] + result_b[\"data_sources\"], axis=0)\n        sample_uids = result_a[\"sample_uids\"] + result_b[\"sample_uids\"]\n        sample_turns = result_a[\"sample_turns\"] + result_b[\"sample_turns\"]\n\n        reward_extra_infos_dict = {}\n        all_keys = set(result_a[\"reward_extra_infos_dict\"].keys()) | set(result_b[\"reward_extra_infos_dict\"].keys())\n        for key in all_keys:\n            list_a = result_a[\"reward_extra_infos_dict\"].get(key, [])\n            list_b = result_b[\"reward_extra_infos_dict\"].get(key, [])\n            reward_extra_infos_dict[key] = list_a + list_b\n\n        return self._val_metrics_update(data_sources, sample_uids, reward_extra_infos_dict, sample_turns)\n\n    def init_workers(self):\n        \"\"\"Initialize distributed training workers using Ray backend.\n\n        Creates:\n        1. Ray resource pools from configuration\n        2. Worker groups for each role (actor, critic, etc.)\n        \"\"\"\n        self.resource_pool_manager.create_resource_pool()\n\n        self.resource_pool_to_cls = {pool: {} for pool in self.resource_pool_manager.resource_pool_dict.values()}\n\n        # create actor and rollout\n        actor_role = Role.ActorRolloutRef if Role.ActorRolloutRef in self.role_worker_mapping else Role.ActorRollout\n        if self.hybrid_engine:\n            actor_rollout_resource_pool = self.resource_pool_manager.get_resource_pool(actor_role)\n            actor_rollout_cls = RayClassWithInitArgs(\n                cls=self.role_worker_mapping[actor_role],\n                config=self.config.actor_rollout_ref,\n                role=str(actor_role),\n            )\n            self.resource_pool_to_cls[actor_rollout_resource_pool][str(actor_role)] = actor_rollout_cls\n        else:\n            raise NotImplementedError\n\n        # create critic\n        if self.use_critic:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.Critic)\n\n            from verl.workers.config import CriticConfig\n\n            critic_cfg: CriticConfig = omega_conf_to_dataclass(self.config.critic)\n\n            if self.use_legacy_worker_impl == \"disable\":\n                # convert critic_cfg into TrainingWorkerConfig\n                from verl.workers.engine_workers import TrainingWorkerConfig\n\n                orig_critic_cfg = critic_cfg\n                if orig_critic_cfg.strategy == \"fsdp\":\n                    engine_config: FSDPEngineConfig = orig_critic_cfg.model.fsdp_config\n                    engine_config.infer_max_token_len_per_gpu = critic_cfg.ppo_infer_max_token_len_per_gpu\n                    engine_config.max_token_len_per_gpu = critic_cfg.ppo_max_token_len_per_gpu\n                else:\n                    raise NotImplementedError(f\"Unknown strategy {orig_critic_cfg.strategy=}\")\n\n                critic_cfg = TrainingWorkerConfig(\n                    model_type=\"value_model\",\n                    model_config=orig_critic_cfg.model_config,\n                    engine_config=engine_config,\n                    optimizer_config=orig_critic_cfg.optim,\n                    checkpoint_config=orig_critic_cfg.checkpoint,\n                )\n\n            critic_cls = RayClassWithInitArgs(cls=self.role_worker_mapping[Role.Critic], config=critic_cfg)\n            self.resource_pool_to_cls[resource_pool][str(Role.Critic)] = critic_cls\n\n        # create reference policy if needed\n        if self.use_reference_policy and Role.RefPolicy in self.role_worker_mapping:\n            resource_pool = self.resource_pool_manager.get_resource_pool(Role.RefPolicy)\n            ref_policy_cls = RayClassWithInitArgs(\n                self.role_worker_mapping[Role.RefPolicy],\n                config=self.config.actor_rollout_ref,\n                role=str(Role.RefPolicy),\n            )\n            self.resource_pool_to_cls[resource_pool][str(Role.RefPolicy)] = ref_policy_cls\n\n        # initialize WorkerGroup\n        # NOTE: if you want to use a different resource pool for each role, which can support different parallel size,\n        # you should not use `create_colocated_worker_cls`.\n        # Instead, directly pass different resource pool to different worker groups.\n        # See https://github.com/volcengine/verl/blob/master/examples/ray/tutorial.ipynb for more information.\n        all_wg = {}\n        wg_kwargs = {}  # Setting up kwargs for RayWorkerGroup\n        if OmegaConf.select(self.config.trainer, \"ray_wait_register_center_timeout\") is not None:\n            wg_kwargs[\"ray_wait_register_center_timeout\"] = self.config.trainer.ray_wait_register_center_timeout\n        if OmegaConf.select(self.config.global_profiler, \"steps\") is not None:\n            wg_kwargs[\"profile_steps\"] = OmegaConf.select(self.config.global_profiler, \"steps\")\n            # Only require nsight worker options when tool is nsys\n            if OmegaConf.select(self.config.global_profiler, \"tool\") == \"nsys\":\n                assert (\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                    is not None\n                ), \"worker_nsight_options must be set when using nsys with profile_steps\"\n                wg_kwargs[\"worker_nsight_options\"] = OmegaConf.to_container(\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                )\n        wg_kwargs[\"device_name\"] = self.device_name\n\n        for resource_pool, class_dict in self.resource_pool_to_cls.items():\n            if not class_dict:\n                continue\n            worker_dict_cls = create_colocated_worker_cls(class_dict=class_dict)\n            wg_dict = self.ray_worker_group_cls(\n                resource_pool=resource_pool,\n                ray_cls_with_init=worker_dict_cls,\n                **wg_kwargs,\n            )\n            spawn_wg = wg_dict.spawn(prefix_set=class_dict.keys())\n            all_wg.update(spawn_wg)\n\n        if self.use_critic:\n            self.critic_wg = all_wg[str(Role.Critic)]\n            if self.use_legacy_worker_impl == \"disable\":\n                self.critic_wg.reset()\n                # assign critic loss\n                from functools import partial\n\n                from verl.workers.utils.losses import value_loss\n\n                value_loss_ = partial(value_loss, config=orig_critic_cfg)\n                self.critic_wg.set_loss_fn(value_loss_)\n            else:\n                self.critic_wg.init_model()\n\n        if self.use_reference_policy and not self.ref_in_actor:\n            if str(Role.RefPolicy) in all_wg:\n                self.ref_policy_wg = all_wg[str(Role.RefPolicy)]\n                self.ref_policy_wg.init_model()\n            else:\n                # Model engine: ActorRolloutRefWorker\n                assert str(Role.ActorRolloutRef) in all_wg, f\"{all_wg.keys()=}\"\n                self.ref_policy_wg = all_wg[str(Role.ActorRolloutRef)]\n\n        # we should create rollout at the end so that vllm can have a better estimation of kv cache memory\n        self.actor_rollout_wg = all_wg[str(actor_role)]\n        self.actor_rollout_wg.init_model()\n\n        if self.ref_in_actor:\n            self.ref_policy_wg = self.actor_rollout_wg\n\n        # create reward loop manager\n        from verl.experimental.reward_loop import RewardLoopManager\n\n        # initalize reward loop manager\n        # reward model (colocate or standalone): get resource_pool\n        # no reward model: resource_pool = None\n        resource_pool = self.resource_pool_manager.get_resource_pool(Role.RewardModel) if self.use_rm else None\n        self.reward_loop_manager = RewardLoopManager(\n            config=self.config,\n            rm_resource_pool=resource_pool,\n        )\n\n        # create async rollout manager and request scheduler\n        # Note: mode is always \"async\" since sync mode is deprecated\n        self.async_rollout_mode = True\n\n        # Support custom AgentLoopManager via config\n        manager_class_fqn = self.config.actor_rollout_ref.rollout.get(\"agent\", {}).get(\"agent_loop_manager_class\")\n        if manager_class_fqn:\n            AgentLoopManager = load_class_from_fqn(manager_class_fqn, \"AgentLoopManager\")\n        else:\n            from verl.experimental.agent_loop import AgentLoopManager\n\n        # infrastructure overview: https://verl.readthedocs.io/en/latest/advance/reward_loop.html#architecture-design\n        # agent_reward_loop: streaming reward computation with actor rollout\n        # two conditions satisfied: (1) no reward model, or (2) reward model with extra resource pool\n        enable_agent_reward_loop = not self.use_rm or self.config.reward.reward_model.enable_resource_pool\n\n        # if enable_agent_reward_loop, we directly pass reward_loop_workers to agent loop manager\n        # to stream reward computation with actor rollout\n        reward_loop_worker_handles = self.reward_loop_manager.reward_loop_workers if enable_agent_reward_loop else None\n        self.async_rollout_manager = AgentLoopManager.create(\n            config=self.config,\n            worker_group=self.actor_rollout_wg,\n            rollout_resource_pool=actor_rollout_resource_pool,\n            reward_loop_worker_handles=reward_loop_worker_handles,\n        )\n        checkpoint_engine_config = omega_conf_to_dataclass(self.config.actor_rollout_ref.rollout.checkpoint_engine)\n        self.checkpoint_manager = CheckpointEngineManager(\n            config=checkpoint_engine_config,\n            trainer=self.actor_rollout_wg,\n            replicas=self.async_rollout_manager.rollout_replicas,\n        )\n\n        # sleep all replicas to load checkpoint\n        self.checkpoint_manager.sleep_replicas()\n\n    def _save_checkpoint(self):\n        from verl.utils.fs import local_mkdir_safe\n\n        # path: given_path + `/global_step_{global_steps}` + `/actor`\n        local_global_step_folder = os.path.join(\n            self.config.trainer.default_local_dir, f\"global_step_{self.global_steps}\"\n        )\n\n        print(f\"local_global_step_folder: {local_global_step_folder}\")\n        actor_local_path = os.path.join(local_global_step_folder, \"actor\")\n\n        actor_remote_path = (\n            None\n            if self.config.trainer.default_hdfs_dir is None\n            else os.path.join(self.config.trainer.default_hdfs_dir, f\"global_step_{self.global_steps}\", \"actor\")\n        )\n\n        remove_previous_ckpt_in_save = self.config.trainer.get(\"remove_previous_ckpt_in_save\", False)\n        if remove_previous_ckpt_in_save:\n            print(\n                \"Warning: remove_previous_ckpt_in_save is deprecated,\"\n                + \" set max_actor_ckpt_to_keep=1 and max_critic_ckpt_to_keep=1 instead\"\n            )\n        max_actor_ckpt_to_keep = (\n            self.config.trainer.get(\"max_actor_ckpt_to_keep\", None) if not remove_previous_ckpt_in_save else 1\n        )\n        max_critic_ckpt_to_keep = (\n            self.config.trainer.get(\"max_critic_ckpt_to_keep\", None) if not remove_previous_ckpt_in_save else 1\n        )\n\n        self.actor_rollout_wg.save_checkpoint(\n            actor_local_path, actor_remote_path, self.global_steps, max_ckpt_to_keep=max_actor_ckpt_to_keep\n        )\n\n        if self.use_critic:\n            critic_local_path = os.path.join(local_global_step_folder, str(Role.Critic))\n            critic_remote_path = (\n                None\n                if self.config.trainer.default_hdfs_dir is None\n                else os.path.join(\n                    self.config.trainer.default_hdfs_dir, f\"global_step_{self.global_steps}\", str(Role.Critic)\n                )\n            )\n            self.critic_wg.save_checkpoint(\n                critic_local_path, critic_remote_path, self.global_steps, max_ckpt_to_keep=max_critic_ckpt_to_keep\n            )\n\n        # save dataloader\n        local_mkdir_safe(local_global_step_folder)\n        dataloader_local_path = os.path.join(local_global_step_folder, \"data.pt\")\n        dataloader_state_dict = self.train_dataloader.state_dict()\n        torch.save(dataloader_state_dict, dataloader_local_path)\n\n        # latest checkpointed iteration tracker (for atomic usage)\n        if (\n            hasattr(self.config.actor_rollout_ref.actor.checkpoint, \"async_save\")\n            and self.config.actor_rollout_ref.actor.checkpoint.async_save\n        ) or (\n            \"async_save\" in self.config.actor_rollout_ref.actor.checkpoint\n            and self.config.actor_rollout_ref.actor.checkpoint[\"async_save\"]\n        ):\n            print(\"skip write latest_checkpointed_iteration.txt when async_save is True\")\n            return\n        local_latest_checkpointed_iteration = os.path.join(\n            self.config.trainer.default_local_dir, \"latest_checkpointed_iteration.txt\"\n        )\n        with open(local_latest_checkpointed_iteration, \"w\") as f:\n            f.write(str(self.global_steps))\n\n    def _load_checkpoint(self):\n        if self.config.trainer.resume_mode == \"disable\":\n            return 0\n\n        # load from hdfs\n        if self.config.trainer.default_hdfs_dir is not None:\n            raise NotImplementedError(\"load from hdfs is not implemented yet\")\n        else:\n            checkpoint_folder = self.config.trainer.default_local_dir  # TODO: check path\n            if not os.path.isabs(checkpoint_folder):\n                working_dir = os.getcwd()\n                checkpoint_folder = os.path.join(working_dir, checkpoint_folder)\n            global_step_folder = find_latest_ckpt_path(checkpoint_folder)  # None if no latest\n\n        # find global_step_folder\n        if self.config.trainer.resume_mode == \"auto\":\n            if global_step_folder is None:\n                print(\"Training from scratch\")\n                return 0\n        else:\n            if self.config.trainer.resume_mode == \"resume_path\":\n                assert isinstance(self.config.trainer.resume_from_path, str), \"resume ckpt must be str type\"\n                assert \"global_step_\" in self.config.trainer.resume_from_path, (\n                    \"resume ckpt must specify the global_steps\"\n                )\n                global_step_folder = self.config.trainer.resume_from_path\n                if not os.path.isabs(global_step_folder):\n                    working_dir = os.getcwd()\n                    global_step_folder = os.path.join(working_dir, global_step_folder)\n        print(f\"Load from checkpoint folder: {global_step_folder}\")\n        # set global step\n        self.global_steps = int(global_step_folder.split(\"global_step_\")[-1])\n\n        print(f\"Setting global step to {self.global_steps}\")\n        print(f\"Resuming from {global_step_folder}\")\n\n        actor_path = os.path.join(global_step_folder, \"actor\")\n        critic_path = os.path.join(global_step_folder, str(Role.Critic))\n        # load actor\n        self.actor_rollout_wg.load_checkpoint(\n            actor_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load\n        )\n        # load critic\n        if self.use_critic:\n            self.critic_wg.load_checkpoint(\n                critic_path, del_local_after_load=self.config.trainer.del_local_ckpt_after_load\n            )\n\n        # load dataloader,\n        # TODO: from remote not implemented yet\n        dataloader_local_path = os.path.join(global_step_folder, \"data.pt\")\n        if os.path.exists(dataloader_local_path):\n            dataloader_state_dict = torch.load(dataloader_local_path, weights_only=False)\n            self.train_dataloader.load_state_dict(dataloader_state_dict)\n        else:\n            print(f\"Warning: No dataloader state found at {dataloader_local_path}, will start from scratch\")\n\n    def _start_profiling(self, do_profile: bool) -> None:\n        \"\"\"Start profiling for all worker groups if profiling is enabled.\"\"\"\n        if do_profile:\n            self.actor_rollout_wg.start_profile(role=\"e2e\", profile_step=self.global_steps)\n            if self.use_reference_policy:\n                self.ref_policy_wg.start_profile(profile_step=self.global_steps)\n            if self.use_critic:\n                self.critic_wg.start_profile(profile_step=self.global_steps)\n\n    def _stop_profiling(self, do_profile: bool) -> None:\n        \"\"\"Stop profiling for all worker groups if profiling is enabled.\"\"\"\n        if do_profile:\n            self.actor_rollout_wg.stop_profile()\n            if self.use_reference_policy:\n                self.ref_policy_wg.stop_profile()\n            if self.use_critic:\n                self.critic_wg.stop_profile()\n\n    def _get_dp_size(self, worker_group, role: str) -> int:\n        \"\"\"Get data parallel size from worker group dispatch info.\n\n        This method retrieves the data parallel size by querying the dispatch info\n        for the specified role. The dispatch info is cached for subsequent calls.\n\n        Args:\n            worker_group: The worker group to query dispatch info from.\n            role: The role name (e.g., \"actor\", \"critic\") to get DP size for.\n\n        Returns:\n            The data parallel size (number of DP ranks).\n        \"\"\"\n        if role not in worker_group._dispatch_info:\n            dp_rank_mapping = worker_group._query_dispatch_info(role)\n            worker_group._dispatch_info[role] = dp_rank_mapping\n        else:\n            dp_rank_mapping = worker_group._dispatch_info[role]\n        return max(dp_rank_mapping) + 1\n\n    def _balance_batch(self, batch: DataProto, metrics, logging_prefix=\"global_seqlen\", keep_minibatch=False):\n        \"\"\"Reorder the data on single controller such that each dp rank gets similar total tokens.\n\n        When use_prefix_grouper is enabled, uses group-level balancing to keep samples with\n        the same uid together on the same rank for prefix sharing optimization.\n        \"\"\"\n        attention_mask = batch.batch[\"attention_mask\"]\n        batch_size = attention_mask.shape[0]\n        global_seqlen_lst = batch.batch[\"attention_mask\"].view(batch_size, -1).sum(-1)  # (train_batch_size,)\n        workload_lst = calculate_workload(global_seqlen_lst)\n        # Get dp_size from dispatch info to correctly balance across data parallel ranks\n        # Note: world_size may include tensor/pipeline parallel dimensions, but we only want DP\n        dp_size = self._get_dp_size(self.actor_rollout_wg, \"actor\")\n\n        # Use group-level balancing for PrefixGrouper to keep same-uid samples together\n        if getattr(self, \"use_prefix_grouper\", False) and \"uid\" in batch.non_tensor_batch:\n            from verl.utils.seqlen_balancing import get_group_balanced_partitions\n\n            uid_list = list(batch.non_tensor_batch[\"uid\"])\n            seqlen_list = global_seqlen_lst.tolist()\n\n            # Count number of uid groups\n            num_groups = len(set(uid_list))\n\n            if num_groups % dp_size != 0:\n                raise ValueError(\n                    f\"PrefixGrouper with balance_batch requires num_uid_groups ({num_groups}) \"\n                    f\"% dp_size ({dp_size}) == 0. \"\n                    f\"This ensures each rank gets equal number of groups. \"\n                    f\"Current batch_size={batch_size}, adjust batch_size to be a multiple of \"\n                    f\"dp_size * rollout.n.\"\n                )\n\n            global_partition_lst = get_group_balanced_partitions(\n                seqlen_list=seqlen_list,\n                uid_list=uid_list,\n                k_partitions=dp_size,\n            )\n\n        elif keep_minibatch:\n            # Decouple the DP balancing and mini-batching.\n            minibatch_size = self.config.actor_rollout_ref.actor.get(\"ppo_mini_batch_size\")\n            minibatch_num = len(workload_lst) // minibatch_size\n            global_partition_lst = [[] for _ in range(dp_size)]\n            for i in range(minibatch_num):\n                rearrange_minibatch_lst = get_seqlen_balanced_partitions(\n                    workload_lst[i * minibatch_size : (i + 1) * minibatch_size],\n                    k_partitions=dp_size,\n                    equal_size=True,\n                )\n                for j, part in enumerate(rearrange_minibatch_lst):\n                    global_partition_lst[j].extend([x + minibatch_size * i for x in part])\n        else:\n            global_partition_lst = get_seqlen_balanced_partitions(workload_lst, k_partitions=dp_size, equal_size=True)\n        # Place smaller micro-batches at both ends to reduce the bubbles in pipeline parallel.\n        # Skip reordering within partitions for PrefixGrouper to maintain uid grouping\n        if not getattr(self, \"use_prefix_grouper\", False):\n            for idx, partition in enumerate(global_partition_lst):\n                partition.sort(key=lambda x: (workload_lst[x], x))\n                ordered_partition = partition[::2] + partition[1::2][::-1]\n                global_partition_lst[idx] = ordered_partition\n\n        # reorder based on index. The data will be automatically equally partitioned by dispatch function\n        global_idx = torch.tensor([j for partition in global_partition_lst for j in partition])\n        batch.reorder(global_idx)\n        global_balance_stats = log_seqlen_unbalance(\n            seqlen_list=global_seqlen_lst.tolist(), partitions=global_partition_lst, prefix=logging_prefix\n        )\n        metrics.update(global_balance_stats)\n\n    def _compute_values(self, batch: DataProto) -> DataProto:\n        if self.use_legacy_worker_impl == \"disable\":\n            batch_td = batch.to_tensordict()\n            # step 2: convert from padding to nopadding\n            batch_td = left_right_2_no_padding(batch_td)\n            # step 3: add meta info\n            tu.assign_non_tensor(batch_td, compute_loss=False)\n            output = self.critic_wg.infer_batch(batch_td)\n            output = output.get()\n            values = tu.get(output, \"values\")\n            values = no_padding_2_padding(values, batch_td)\n            values = tu.get_tensordict({\"values\": values.float()})\n            values = DataProto.from_tensordict(values)\n        else:\n            values = self.critic_wg.compute_values(batch)\n        return values\n\n    def _compute_ref_log_prob(self, batch: DataProto) -> DataProto:\n        if self.use_legacy_worker_impl == \"disable\":\n            # step 1: convert dataproto to tensordict.\n            batch_td = batch.to_tensordict()\n            # step 2: convert from padding to nopadding\n            batch_td = left_right_2_no_padding(batch_td)\n            # step 3: add meta info\n            metadata = {\"calculate_entropy\": False, \"compute_loss\": False}\n            if self.ref_in_actor:\n                metadata[\"no_lora_adapter\"] = True\n            tu.assign_non_tensor(batch_td, **metadata)\n            if self.ref_in_actor:\n                output = self.actor_rollout_wg.compute_log_prob(batch_td)\n            else:\n                output = self.ref_policy_wg.compute_ref_log_prob(batch_td)\n            # gather output\n            log_probs = tu.get(output, \"log_probs\")\n            # step 4. No padding to padding\n            log_probs = no_padding_2_padding(log_probs, batch_td)\n            # step 5: rebuild a tensordict and convert to dataproto\n            ref_log_prob = tu.get_tensordict({\"ref_log_prob\": log_probs.float()})\n            ref_log_prob = DataProto.from_tensordict(ref_log_prob)\n        else:\n            ref_log_prob = self.ref_policy_wg.compute_ref_log_prob(batch)\n\n        return ref_log_prob\n\n    def _compute_old_log_prob(self, batch: DataProto):\n        if self.use_legacy_worker_impl == \"disable\":\n            # TODO: remove step 1, 2, 4 after we make the whole training tensordict and padding free\n            # step 1: convert dataproto to tensordict.\n            batch_td = batch.to_tensordict()\n            # step 2: convert from padding to nopadding\n            batch_td = left_right_2_no_padding(batch_td)\n            # step 3: add meta info\n            tu.assign_non_tensor(batch_td, calculate_entropy=True, compute_loss=False)\n            output = self.actor_rollout_wg.compute_log_prob(batch_td)\n            # gather output\n            entropy = tu.get(output, \"entropy\")\n            log_probs = tu.get(output, \"log_probs\")\n            routed_experts = tu.get(output, \"routed_experts\")\n            old_log_prob_mfu = tu.get(output, \"metrics\")[\"mfu\"]\n            # step 4. No padding to padding\n            entropy = no_padding_2_padding(entropy, batch_td)\n            log_probs = no_padding_2_padding(log_probs, batch_td)\n            # step 5: rebuild a tensordict and convert to dataproto\n            if routed_experts is not None:\n                old_log_prob = tu.get_tensordict(\n                    {\"old_log_probs\": log_probs.float(), \"entropys\": entropy.float(), \"routed_experts\": routed_experts}\n                )\n            else:\n                old_log_prob = tu.get_tensordict({\"old_log_probs\": log_probs.float(), \"entropys\": entropy.float()})\n            old_log_prob = DataProto.from_tensordict(old_log_prob)\n        else:\n            old_log_prob = self.actor_rollout_wg.compute_log_prob(batch)\n            old_log_prob_mfu = 0\n        return old_log_prob, old_log_prob_mfu\n\n    def _update_actor(self, batch: DataProto) -> DataProto:\n        rollout_config = self.config.actor_rollout_ref.rollout\n        batch.meta_info[\"multi_turn\"] = rollout_config.multi_turn.enable\n        # TODO: Make \"temperature\" single source of truth from generation.\n        batch.meta_info[\"temperature\"] = rollout_config.temperature\n        # update actor\n        if self.use_legacy_worker_impl == \"disable\":\n            batch_td = batch.to_tensordict()\n            # step 2: convert from padding to no-padding\n            batch_td = left_right_2_no_padding(batch_td)\n            calculate_entropy = self.config.actor_rollout_ref.actor.entropy_coeff != 0.0\n            ppo_mini_batch_size = self.config.actor_rollout_ref.actor.ppo_mini_batch_size\n            ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n\n            ppo_epochs = self.config.actor_rollout_ref.actor.ppo_epochs\n            seed = self.config.actor_rollout_ref.actor.data_loader_seed\n            shuffle = self.config.actor_rollout_ref.actor.shuffle\n            tu.assign_non_tensor(\n                batch_td,\n                calculate_entropy=calculate_entropy,\n                global_batch_size=ppo_mini_batch_size,\n                mini_batch_size=ppo_mini_batch_size,\n                epochs=ppo_epochs,\n                seed=seed,\n                dataloader_kwargs={\"shuffle\": shuffle},\n            )\n\n            actor_output = self.actor_rollout_wg.update_actor(batch_td)\n            actor_output = tu.get(actor_output, \"metrics\")\n            actor_output = rename_dict(actor_output, \"actor/\")\n            # modify key name\n            actor_output[\"perf/mfu/actor\"] = actor_output.pop(\"actor/mfu\")\n            actor_output = DataProto.from_single_dict(data={}, meta_info={\"metrics\": actor_output})\n        else:\n            actor_output = self.actor_rollout_wg.update_actor(batch)\n\n        return actor_output\n\n    def _update_critic(self, batch: DataProto) -> DataProto:\n        if self.use_legacy_worker_impl == \"disable\":\n            batch_td = batch.to_tensordict()\n            # step 2: convert from padding to no-padding\n            batch_td = left_right_2_no_padding(batch_td)\n            ppo_mini_batch_size = self.config.critic.ppo_mini_batch_size\n            ppo_mini_batch_size = ppo_mini_batch_size * self.config.actor_rollout_ref.rollout.n\n            ppo_epochs = self.config.critic.ppo_epochs\n            seed = self.config.critic.data_loader_seed\n            shuffle = self.config.critic.shuffle\n            tu.assign_non_tensor(\n                batch_td,\n                global_batch_size=ppo_mini_batch_size,\n                mini_batch_size=ppo_mini_batch_size,\n                epochs=ppo_epochs,\n                seed=seed,\n                dataloader_kwargs={\"shuffle\": shuffle},\n            )\n\n            output = self.critic_wg.train_mini_batch(batch_td)\n            output = output.get()\n            output = tu.get(output, \"metrics\")\n            output = rename_dict(output, \"critic/\")\n            # modify key name\n            output[\"perf/mfu/critic\"] = output.pop(\"critic/mfu\")\n            critic_output = DataProto.from_single_dict(data={}, meta_info={\"metrics\": output})\n        else:\n            critic_output = self.critic_wg.update_critic(batch)\n        return critic_output\n\n    def fit(self):\n        \"\"\"\n        The training loop of PPO.\n        The driver process only need to call the compute functions of the worker group through RPC\n        to construct the PPO dataflow.\n        The light-weight advantage computation is done on the driver process.\n        \"\"\"\n        from omegaconf import OmegaConf\n\n        from verl.utils.tracking import Tracking\n\n        logger = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        self.global_steps = 0\n\n        # load checkpoint and update weights before doing anything\n        self._load_checkpoint()\n        self.checkpoint_manager.update_weights(self.global_steps)\n\n        current_epoch = self.global_steps // len(self.train_dataloader)\n\n        # perform validation before training\n        # currently, we only support validation using the reward_function.\n        if self.config.trainer.get(\"val_before_train\", True):\n            val_metrics = self._validate()\n            assert val_metrics, f\"{val_metrics=}\"\n            pprint(f\"Initial validation metrics: {val_metrics}\")\n            logger.log(data=val_metrics, step=self.global_steps)\n            if self.config.trainer.get(\"val_only\", False):\n                return\n\n        if self.config.actor_rollout_ref.rollout.get(\"skip_rollout\", False):\n            rollout_skip = RolloutSkip(self.config, self.async_rollout_manager)\n            rollout_skip.wrap_generate_sequences()\n\n        # add tqdm\n        progress_bar = tqdm(total=self.total_training_steps, initial=self.global_steps, desc=\"Training Progress\")\n\n        # we start from step 1\n        self.global_steps += 1\n        last_val_metrics = None\n        self.max_steps_duration = 0\n\n        prev_step_profile = False\n        curr_step_profile = (\n            self.global_steps in self.config.global_profiler.steps\n            if self.config.global_profiler.steps is not None\n            else False\n        )\n        next_step_profile = False\n\n        for epoch in range(current_epoch, self.config.trainer.total_epochs):\n            for batch_dict in self.train_dataloader:\n                if hasattr(self.actor_rollout_wg, \"async_calls_finalize_fn_exec\"):\n                    self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=False)\n                metrics = {}\n                timing_raw = {}\n\n                with marked_timer(\"start_profile\", timing_raw):\n                    self._start_profiling(\n                        not prev_step_profile and curr_step_profile\n                        if self.config.global_profiler.profile_continuous_steps\n                        else curr_step_profile\n                    )\n                batch: DataProto = DataProto.from_single_dict(batch_dict)\n                batch.meta_info[\"temperature\"] = self.config.actor_rollout_ref.rollout.temperature\n\n                # add uid to batch\n                batch.non_tensor_batch[\"uid\"] = np.array(\n                    [str(uuid.uuid4()) for _ in range(len(batch.batch))], dtype=object\n                )\n\n                gen_batch = self._get_gen_batch(batch)\n\n                # pass global_steps to trace\n                gen_batch.meta_info[\"global_steps\"] = self.global_steps\n                gen_batch_output = gen_batch.repeat(\n                    repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True\n                )\n\n                is_last_step = self.global_steps >= self.total_training_steps\n                with marked_timer(\"step\", timing_raw):\n                    # generate a batch\n                    with marked_timer(\"gen\", timing_raw, color=\"red\"):\n                        if curr_step_profile:\n                            self.async_rollout_manager.start_profile()\n                        gen_batch_output = self.async_rollout_manager.generate_sequences(gen_batch_output)\n                        self.checkpoint_manager.sleep_replicas()\n                        if curr_step_profile:\n                            self.async_rollout_manager.stop_profile()\n\n                        timing_raw.update(gen_batch_output.meta_info[\"timing\"])\n                        gen_batch_output.meta_info.pop(\"timing\", None)\n\n                    if self.config.algorithm.adv_estimator == AdvantageEstimator.REMAX:\n                        with marked_timer(\"gen_max\", timing_raw, color=\"purple\"):\n                            gen_baseline_batch = deepcopy(gen_batch)\n                            gen_baseline_batch.meta_info[\"do_sample\"] = False\n                            if curr_step_profile:\n                                self.async_rollout_manager.start_profile()\n                            gen_baseline_output = self.async_rollout_manager.generate_sequences(gen_baseline_batch)\n                            self.checkpoint_manager.sleep_replicas()\n                            if curr_step_profile:\n                                self.async_rollout_manager.stop_profile()\n                            batch = batch.union(gen_baseline_output)\n                            # compute reward model score on batch\n                            rm_scores = None\n                            if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                                batch_reward = self._compute_reward_colocate(batch)\n                                batch = batch.union(batch_reward)\n\n                            # Compute or extract reward for REMAX baseline\n                            reward_baseline_tensor = batch.batch[\"rm_scores\"].sum(dim=-1)\n\n                            keys_to_pop = set(gen_baseline_output.batch.keys())\n                            if rm_scores is not None:\n                                keys_to_pop.update(rm_scores.batch.keys())\n                            batch.pop(batch_keys=list(keys_to_pop))\n\n                            batch.batch[\"reward_baselines\"] = reward_baseline_tensor\n\n                            del rm_scores, gen_baseline_batch, gen_baseline_output\n                    # repeat to align with repeated responses in rollout\n                    batch = batch.repeat(repeat_times=self.config.actor_rollout_ref.rollout.n, interleave=True)\n                    batch = batch.union(gen_batch_output)\n\n                    if \"response_mask\" not in batch.batch.keys():\n                        batch.batch[\"response_mask\"] = compute_response_mask(batch)\n                    # Balance the number of valid tokens across DP ranks.\n                    # NOTE: This usually changes the order of data in the `batch`,\n                    # which won't affect the advantage calculation (since it's based on uid),\n                    # but might affect the loss calculation (due to the change of mini-batching).\n                    if self.config.trainer.balance_batch:\n                        self._balance_batch(batch, metrics=metrics)\n\n                    # compute global_valid tokens\n                    batch.meta_info[\"global_token_num\"] = torch.sum(batch.batch[\"attention_mask\"], dim=-1).tolist()\n                    # get images_seqlens\n                    images_seqlens_all = []\n                    for multi_modal_input in batch.non_tensor_batch[\"multi_modal_inputs\"]:\n                        if \"image_grid_thw\" not in multi_modal_input.keys():\n                            continue\n                        images_seqlens_all.extend(multi_modal_input[\"images_seqlens\"].tolist())\n                    batch.meta_info[\"images_seqlens\"] = images_seqlens_all\n                    with marked_timer(\"reward\", timing_raw, color=\"yellow\"):\n                        # compute reward model score\n                        if self.use_rm and \"rm_scores\" not in batch.batch.keys():\n                            batch_reward = self._compute_reward_colocate(batch)\n                            batch = batch.union(batch_reward)\n\n                        # extract reward_tensor and reward_extra_infos_dict for training\n                        reward_tensor, reward_extra_infos_dict = extract_reward(batch)\n\n                    # Operating Mode Selection:\n                    # - Bypass mode: Sets old_log_probs = rollout_log_probs (2 policies: π_rollout, π_θ)\n                    # - Decoupled mode: Recomputes old_log_probs as proximal anchor (3 policies: π_rollout, π_old, π_θ)\n                    #   Note: π_old computed once per data batch, serves as stable reference during mini-batch updates\n                    rollout_corr_config = self.config.algorithm.get(\"rollout_correction\", None)\n                    bypass_recomputing_logprobs = rollout_corr_config and rollout_corr_config.get(\"bypass_mode\", False)\n                    if bypass_recomputing_logprobs:  # Use `rollout_log_probs`\n                        from verl.trainer.ppo.rollout_corr_helper import apply_bypass_mode\n\n                        apply_bypass_mode(\n                            batch=batch,\n                            rollout_corr_config=rollout_corr_config,\n                            policy_loss_config=self.config.actor_rollout_ref.actor.policy_loss,\n                        )\n                    else:  # Recompute old_log_probs\n                        with marked_timer(\"old_log_prob\", timing_raw, color=\"blue\"):\n                            old_log_prob, old_log_prob_mfu = self._compute_old_log_prob(batch)\n                            entropys = old_log_prob.batch[\"entropys\"]\n                            response_masks = batch.batch[\"response_mask\"]\n                            actor_config = self.config.actor_rollout_ref.actor\n                            entropy_agg = agg_loss(\n                                loss_mat=entropys,\n                                loss_mask=response_masks,\n                                loss_agg_mode=actor_config.loss_agg_mode,\n                                loss_scale_factor=actor_config.loss_scale_factor,\n                            )\n                            old_log_prob_metrics = {\n                                \"actor/entropy\": entropy_agg.detach().item(),\n                                \"perf/mfu/actor_infer\": old_log_prob_mfu,\n                            }\n                            metrics.update(old_log_prob_metrics)\n                            old_log_prob.batch.pop(\"entropys\")\n                            if \"routed_experts\" in batch.batch and \"routed_experts\" in old_log_prob.batch:\n                                raise ValueError(\n                                    \"Detected conflicting router replay configuration: \"\n                                    \"router_replay.mode='R2' and enable_rollout_routing_replay=True \"\n                                    \"cannot be enabled simultaneously. \"\n                                    \"The enable_rollout_routing_replay option is only used in R3 mode; \"\n                                    \"it should not be set when using R2 mode.\"\n                                )\n                            batch = batch.union(old_log_prob)\n                            if \"rollout_log_probs\" in batch.batch.keys():\n                                # TODO: we may want to add diff of probs too.\n                                from verl.utils.debug.metrics import calculate_debug_metrics\n\n                                metrics.update(calculate_debug_metrics(batch))\n\n                    assert \"old_log_probs\" in batch.batch, f'\"old_log_prob\" not in {batch.batch.keys()=}'\n\n                    if self.use_reference_policy:\n                        # compute reference log_prob\n                        with marked_timer(str(Role.RefPolicy), timing_raw, color=\"olive\"):\n                            ref_log_prob = self._compute_ref_log_prob(batch)\n                            batch = batch.union(ref_log_prob)\n\n                    # compute values\n                    if self.use_critic:\n                        with marked_timer(\"values\", timing_raw, color=\"cyan\"):\n                            values = self._compute_values(batch)\n                            batch = batch.union(values)\n\n                    with marked_timer(\"adv\", timing_raw, color=\"brown\"):\n                        # we combine with rule-based rm\n                        reward_extra_infos_dict: dict[str, list]\n                        batch.batch[\"token_level_scores\"] = reward_tensor\n\n                        if reward_extra_infos_dict:\n                            batch.non_tensor_batch.update({k: np.array(v) for k, v in reward_extra_infos_dict.items()})\n\n                        # compute rewards. apply_kl_penalty if available\n                        if self.config.algorithm.use_kl_in_reward:\n                            batch, kl_metrics = apply_kl_penalty(\n                                batch, kl_ctrl=self.kl_ctrl_in_reward, kl_penalty=self.config.algorithm.kl_penalty\n                            )\n                            metrics.update(kl_metrics)\n                        else:\n                            batch.batch[\"token_level_rewards\"] = batch.batch[\"token_level_scores\"]\n\n                        # Compute rollout correction: IS weights, rejection sampling, and metrics\n                        # Only runs in decoupled mode (computes once per batch using stable π_old)\n                        # In bypass mode, this is skipped - actor computes metrics from evolving π_θ vs π_rollout\n                        if (\n                            rollout_corr_config is not None\n                            and \"rollout_log_probs\" in batch.batch\n                            and not bypass_recomputing_logprobs  # Only in decoupled mode\n                        ):\n                            from verl.trainer.ppo.rollout_corr_helper import compute_rollout_correction_and_add_to_batch\n\n                            # Compute IS weights, apply rejection sampling, compute metrics\n                            batch, is_metrics = compute_rollout_correction_and_add_to_batch(batch, rollout_corr_config)\n                            # IS and off-policy metrics already have rollout_corr/ prefix\n                            metrics.update(is_metrics)\n\n                        # compute advantages, executed on the driver process\n                        norm_adv_by_std_in_grpo = self.config.algorithm.get(\n                            \"norm_adv_by_std_in_grpo\", True\n                        )  # GRPO adv normalization factor\n\n                        batch = compute_advantage(\n                            batch,\n                            adv_estimator=self.config.algorithm.adv_estimator,\n                            gamma=self.config.algorithm.gamma,\n                            lam=self.config.algorithm.lam,\n                            num_repeat=self.config.actor_rollout_ref.rollout.n,\n                            norm_adv_by_std_in_grpo=norm_adv_by_std_in_grpo,\n                            config=self.config.algorithm,\n                        )\n\n                    # update critic\n                    if self.use_critic:\n                        with marked_timer(\"update_critic\", timing_raw, color=\"pink\"):\n                            critic_output = self._update_critic(batch)\n                        critic_output_metrics = reduce_metrics(critic_output.meta_info[\"metrics\"])\n                        metrics.update(critic_output_metrics)\n\n                    # implement critic warmup\n                    if self.config.trainer.critic_warmup <= self.global_steps:\n                        # update actor\n                        with marked_timer(\"update_actor\", timing_raw, color=\"red\"):\n                            actor_output = self._update_actor(batch)\n\n                        # Check if the ESI (Elastic Server Instance)/training plan is close to expiration.\n                        esi_close_to_expiration = should_save_ckpt_esi(\n                            max_steps_duration=self.max_steps_duration,\n                            redundant_time=self.config.trainer.esi_redundant_time,\n                        )\n                        # Check if the conditions for saving a checkpoint are met.\n                        # The conditions include a mandatory condition (1) and\n                        # one of the following optional conditions (2/3/4):\n                        # 1. The save frequency is set to a positive value.\n                        # 2. It's the last training step.\n                        # 3. The current step number is a multiple of the save frequency.\n                        # 4. The ESI(Elastic Server Instance)/training plan is close to expiration.\n                        if self.config.trainer.save_freq > 0 and (\n                            is_last_step\n                            or self.global_steps % self.config.trainer.save_freq == 0\n                            or esi_close_to_expiration\n                        ):\n                            if esi_close_to_expiration:\n                                print(\"Force saving checkpoint: ESI instance expiration approaching.\")\n                            with marked_timer(\"save_checkpoint\", timing_raw, color=\"green\"):\n                                self._save_checkpoint()\n\n                        # update weights from trainer to rollout\n                        with marked_timer(\"update_weights\", timing_raw, color=\"red\"):\n                            self.checkpoint_manager.update_weights(self.global_steps)\n\n                        actor_output_metrics = reduce_metrics(actor_output.meta_info[\"metrics\"])\n                        metrics.update(actor_output_metrics)\n\n                    # Log rollout generations if enabled\n                    rollout_data_dir = self.config.trainer.get(\"rollout_data_dir\", None)\n                    if rollout_data_dir:\n                        self._log_rollout_data(batch, reward_extra_infos_dict, timing_raw, rollout_data_dir)\n\n                # validate\n                if self.config.trainer.test_freq > 0 and (\n                    is_last_step or self.global_steps % self.config.trainer.test_freq == 0\n                ):\n                    with marked_timer(\"testing\", timing_raw, color=\"green\"):\n                        val_metrics: dict = self._validate()\n                        if is_last_step:\n                            last_val_metrics = val_metrics\n                    metrics.update(val_metrics)\n\n                with marked_timer(\"stop_profile\", timing_raw):\n                    next_step_profile = (\n                        self.global_steps + 1 in self.config.global_profiler.steps\n                        if self.config.global_profiler.steps is not None\n                        else False\n                    )\n                    self._stop_profiling(\n                        curr_step_profile and not next_step_profile\n                        if self.config.global_profiler.profile_continuous_steps\n                        else curr_step_profile\n                    )\n                    prev_step_profile = curr_step_profile\n                    curr_step_profile = next_step_profile\n\n                steps_duration = timing_raw[\"step\"]\n                self.max_steps_duration = max(self.max_steps_duration, steps_duration)\n\n                # training metrics\n                metrics.update(\n                    {\n                        \"training/global_step\": self.global_steps,\n                        \"training/epoch\": epoch,\n                    }\n                )\n                # collect metrics\n                metrics.update(compute_data_metrics(batch=batch, use_critic=self.use_critic))\n                # GDPO per-component reward metrics\n                gdpo_reward_keys = self.config.algorithm.get(\"gdpo_reward_keys\", None)\n                if gdpo_reward_keys and self.config.algorithm.adv_estimator in (\"gdpo\", AdvantageEstimator.GDPO):\n                    for key in gdpo_reward_keys:\n                        if key in batch.non_tensor_batch:\n                            vals = np.asarray(batch.non_tensor_batch[key], dtype=np.float32)\n                            metrics[f\"gdpo/{key}/mean\"] = float(np.mean(vals))\n                            metrics[f\"gdpo/{key}/std\"] = float(np.std(vals))\n                            metrics[f\"gdpo/{key}/max\"] = float(np.max(vals))\n                            metrics[f\"gdpo/{key}/min\"] = float(np.min(vals))\n                metrics.update(compute_timing_metrics(batch=batch, timing_raw=timing_raw))\n                # TODO: implement actual tflpo and theoretical tflpo\n                n_gpus = self.resource_pool_manager.get_n_gpus()\n                metrics.update(compute_throughout_metrics(batch=batch, timing_raw=timing_raw, n_gpus=n_gpus))\n                # compute variance proxy metrics\n                gradient_norm = metrics.get(\"actor/grad_norm\", None)\n                metrics.update(compute_variance_proxy_metrics(batch=batch, gradient_norm=gradient_norm))\n                # Note: mismatch metrics (KL, PPL, etc.) are collected at line 1179 after advantage computation\n\n                # this is experimental and may be changed/removed in the future in favor of a general-purpose one\n                if isinstance(self.train_dataloader.sampler, AbstractCurriculumSampler):\n                    self.train_dataloader.sampler.update(batch=batch)\n\n                # TODO: make a canonical logger that supports various backend\n                logger.log(data=metrics, step=self.global_steps)\n\n                progress_bar.update(1)\n                self.global_steps += 1\n\n                if (\n                    hasattr(self.config.actor_rollout_ref.actor, \"profiler\")\n                    and self.config.actor_rollout_ref.actor.profiler.tool == \"torch_memory\"\n                ):\n                    self.actor_rollout_wg.dump_memory_snapshot(\n                        tag=f\"post_update_step{self.global_steps}\", sub_dir=f\"step{self.global_steps}\"\n                    )\n\n                if is_last_step:\n                    if hasattr(self.actor_rollout_wg, \"async_calls_finalize_fn_exec\"):\n                        self.actor_rollout_wg.async_calls_finalize_fn_exec(blocking=True)\n                    pprint(f\"Final validation metrics: {last_val_metrics}\")\n                    progress_bar.close()\n                    return\n\n                # this is experimental and may be changed/removed in the future\n                # in favor of a general-purpose data buffer pool\n                if hasattr(self.train_dataset, \"on_batch_end\"):\n                    # The dataset may be changed after each training batch\n                    self.train_dataset.on_batch_end(batch=batch)\n"
  },
  {
    "path": "verl/trainer/ppo/reward.py",
    "content": "# Copyright 2025 Individual Contributor: Thibaut Barroyer\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport inspect\nimport multiprocessing\nfrom functools import partial\nfrom typing import TYPE_CHECKING, Any, Optional, cast\n\nfrom verl import DataProto\nfrom verl.utils.reward_score import default_compute_score\n\nif TYPE_CHECKING:\n    from omegaconf import DictConfig\n\n    from verl.experimental.reward_loop.reward_manager.base import RawRewardFn, RewardManagerBase\n    from verl.trainer.config.config import ModuleConfig\n    from verl.workers.config.reward import RewardManagerConfig\n\n\ndef _call_with_kwargs(raw_fn, extra_kwargs, *args, **kwargs):\n    \"\"\"Calls `raw_fn` by merging `extra_kwargs` into call-time `kwargs`, with `extra_kwargs` taking precedence.\n\n    This function is used to merge additional keyword arguments with the original function's arguments.\n    \"\"\"\n    merged_kwargs = {**kwargs, **extra_kwargs}\n    return raw_fn(*args, **merged_kwargs)\n\n\nasync def _call_with_kwargs_async(raw_fn, extra_kwargs, *args, **kwargs):\n    \"\"\"Calls `raw_fn` by merging `extra_kwargs` into call-time `kwargs`, with `extra_kwargs` taking precedence.\n\n    This function is used to merge additional keyword arguments with the original function's arguments.\n    \"\"\"\n    merged_kwargs = {**kwargs, **extra_kwargs}\n    return await raw_fn(*args, **merged_kwargs)\n\n\ndef get_custom_reward_fn(config: DictConfig) -> Optional[RawRewardFn]:\n    \"\"\"Load and return a custom reward function from external file.\n\n    Dynamically imports a reward function from a specified file path and wraps\n    it with additional keyword arguments from the configuration.\n\n    Args:\n        config (dict): Configuration dictionary containing custom_reward_function\n                      settings with 'path', 'name', and 'reward_kwargs' fields.\n\n    Returns:\n        callable or None: Wrapped reward function with merged kwargs, or None\n                         if no custom reward function is configured.\n\n    Raises:\n        FileNotFoundError: If the specified reward function file doesn't exist.\n        RuntimeError: If there's an error loading the module from file.\n        AttributeError: If the specified function name isn't found in the module.\n    \"\"\"\n\n    reward_fn_config = config.reward.get(\"custom_reward_function\") or {}\n    module_path = reward_fn_config.get(\"path\")\n    if not module_path:\n        return None\n\n    fn_name = reward_fn_config.get(\"name\")\n    assert fn_name is not None\n\n    from verl.utils.import_utils import load_extern_object\n\n    raw_fn = load_extern_object(module_path=module_path, object_name=fn_name)\n\n    reward_kwargs = dict(reward_fn_config.get(\"reward_kwargs\", {}))\n    if not inspect.iscoroutinefunction(raw_fn):\n        return partial(_call_with_kwargs, raw_fn, reward_kwargs)\n    else:\n        return partial(_call_with_kwargs_async, raw_fn, reward_kwargs)\n\n\ndef load_reward_manager(config: DictConfig, tokenizer: Any, **reward_kwargs: Any) -> RewardManagerBase:\n    \"\"\"\n    Load and initialize a reward manager based on the configuration.\n\n    Args:\n        config: PPO trainer configuration object containing reward_model fields.\n        tokenizer: Tokenizer object used for processing text.\n        **reward_kwargs: Additional keyword arguments for the reward manager.\n\n    Returns:\n        An instance of the specified reward manager class.\n    \"\"\"\n\n    # Try to get a custom reward function based on the configuration\n    # user defined reward manager can be registered in custom_reward_fn\n    compute_score = get_custom_reward_fn(config)\n    final_compute_score = compute_score\n\n    reward_manager_cfg: RewardManagerConfig = config.reward.reward_manager\n    reward_manager_cls: type[RewardManagerBase]\n    if reward_manager_cfg.source == \"register\":\n        from verl.experimental.reward_loop.reward_manager import get_reward_manager_cls\n\n        reward_manager_cls = get_reward_manager_cls(reward_manager_cfg.name)\n    elif reward_manager_cfg.source == \"importlib\":\n        from verl.utils.import_utils import load_extern_object\n\n        module_cfg: ModuleConfig | None = reward_manager_cfg.module\n        assert module_cfg is not None and module_cfg.path is not None, (\n            f\"Module path is required when {reward_manager_cfg.source=}, but got {module_cfg=}\"\n        )\n        reward_manager_cls_name = reward_manager_cfg.name\n        reward_manager_cls = cast(\n            \"type[RewardManagerBase]\",\n            load_extern_object(module_path=module_cfg.path, object_name=reward_manager_cls_name),\n        )\n\n    if compute_score is None:\n        sandbox_config = config.reward.get(\"sandbox_fusion\")\n        sandbox_url = sandbox_config.get(\"url\") if sandbox_config else None\n        memory_limit_mb = sandbox_config.get(\"memory_limit_mb\", 1024) if sandbox_config else 1024\n        if sandbox_url:\n            sandbox_manager = multiprocessing.Manager()\n            # Create a semaphore to control concurrent access to the sandbox\n            _concurrent_semaphore = sandbox_manager.Semaphore(sandbox_config.get(\"max_concurrent\", 64))\n            final_compute_score = partial(\n                default_compute_score,\n                sandbox_fusion_url=sandbox_url,\n                concurrent_semaphore=_concurrent_semaphore,\n                memory_limit_mb=memory_limit_mb,\n            )\n        else:\n            final_compute_score = default_compute_score\n\n    # Instantiate and return the reward manager with the specified parameters\n    return reward_manager_cls(\n        config=config,\n        tokenizer=tokenizer,\n        compute_score=final_compute_score,\n        **reward_kwargs,\n    )\n\n\ndef extract_reward(batch: DataProto):\n    \"\"\"\n    Extract reward tensor and extra info from batch data.\n    \"\"\"\n    reward_tensor = batch.batch[\"rm_scores\"]\n    reward_extra_keys = batch.meta_info.get(\"reward_extra_keys\", [])\n    reward_extra_infos_dict = {key: batch.non_tensor_batch[key] for key in reward_extra_keys}\n    return reward_tensor, reward_extra_infos_dict\n"
  },
  {
    "path": "verl/trainer/ppo/rollout_corr_helper.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nRollout Correction Helper Module\n\nThis module provides a complete pipeline to address **off-policy issues** in RL training,\nincluding:\n1. Policy mismatch between rollout and training implementations (e.g., vLLM BFloat16 vs FSDP FP32)\n2. Model update staleness (training on trajectories from older checkpoints)\n3. General distribution shifts between data collection and training\n\nIts core capabilities include computing importance sampling (IS) weights,\nfiltering outlier samples via rejection sampling (RS), and\ntracking metrics to diagnose and correct off-policy issues.\n\n## Core Capabilities\n1. **Multi-Granularity Aggregation**:\n   - Importance Sampling (IS):\n        Token-level\n        Sequence-level\n   - Rejection Sampling (RS):\n        Divergence-based filters (token_k*, seq_sum_k*, seq_mean_k*, seq_max_k*)\n2. **Memory-Efficient Design**:\n   - Log-space computations to avoid numerical overflow/underflow.\n   - Fixed safety bounds (exp(±20)) for stable exponentiation.\n   - Metrics calculated without large intermediate tensors (prevents CUDA OOM).\n3. **Comprehensive Metrics Tracking**:\n   - IS/RS statistics (mean/max/min, effective sample size ESS, rejection rate).\n   - Off-policy diagnostics (KL divergence, perplexity PPL, log PPL difference, χ² divergence).\n   - Sequence-level breakdowns (deviation from ideal weights, outlier fraction).\n\n## Key Interfaces & Usage\n- compute_rollout_correction_and_rejection_mask(): compute IS weights + rejection mask.\n- compute_rollout_correction_weights(): only compute truncated IS weights (for variance\n  reduction, no outlier rejection).\n- compute_rollout_rejection_mask(): only filter outliers (for sample cleaning, no IS weight\n  computation).\n- compute_offpolicy_metrics(): called by core functions to calculate off-policy diagnostics\n  (KL/PPL/χ²) — no direct external calls needed.\n\n### Integration Notes\n- Used in `ray_trainer.py` via `compute_rollout_correction_and_add_to_batch()` (batch training pipeline).\n- Used in `dp_actor.py` for distributed worker computations (distributed training scenarios).\n- All functions support batch inputs and valid token masking (via `response_mask`).\n\n\n## References\n- \"When Speed Kills Stability: Demystifying RL Collapse from the Training-Inference Mismatch\": https://richardli.xyz/rl-collapse\n- Off-policy RL (theoretical basis for IS): https://fengyao.notion.site/off-policy-rl\n\"\"\"\n\nimport math\nfrom typing import Any, Optional\n\nimport torch\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.protocol import DataProto\nfrom verl.trainer.config.algorithm import RolloutCorrectionConfig\nfrom verl.workers.config.actor import PolicyLossConfig\n\n# Safety bound to prevent numerical overflow/underflow when exponentiating\n# exp(20) ≈ 485 million (upper limit for stable weights), exp(-20) ≈ 2e-9 (lower limit)\nSAFETY_BOUND = 20.0\n\nSUPPORTED_ROLLOUT_RS_OPTIONS: set[str] = {\n    \"token_k1\",\n    \"token_k2\",\n    \"token_k3\",\n    \"seq_sum_k1\",\n    \"seq_sum_k2\",\n    \"seq_sum_k3\",\n    \"seq_mean_k1\",\n    \"seq_mean_k2\",\n    \"seq_mean_k3\",\n    \"seq_max_k2\",\n    \"seq_max_k3\",\n}\nTOKEN_LEVEL_ROLLOUT_RS_OPTIONS: set[str] = {\"token_k1\", \"token_k2\", \"token_k3\"}\n\n\ndef _parse_rollout_rs_thresholds(\n    options: list[str], threshold_spec: Optional[str | float]\n) -> dict[str, dict[str, Optional[float]]]:\n    if threshold_spec is None:\n        raise ValueError(\"rollout_rs_threshold must be provided for rejection sampling.\")\n\n    if isinstance(threshold_spec, int | float):\n        raw_specs: list[str] = [str(threshold_spec)]\n    elif isinstance(threshold_spec, str):\n        raw_specs = [part.strip() for part in threshold_spec.split(\",\") if part.strip()]\n    else:\n        raise TypeError(\"rollout_rs_threshold must be a string or numeric value specifying per-option thresholds.\")\n\n    if not raw_specs:\n        raise ValueError(\"rollout_rs_threshold must contain at least one threshold value.\")\n\n    if len(raw_specs) not in (1, len(options)):\n        raise ValueError(\n            f\"rollout_rs_threshold expects either one threshold shared by all options or exactly \"\n            f\"{len(options)} thresholds to match the provided rollout_rs options.\"\n        )\n\n    if len(raw_specs) == 1 and len(options) > 1:\n        raw_specs = raw_specs * len(options)\n\n    thresholds: dict[str, dict[str, Optional[float]]] = {}\n    for option, spec in zip(options, raw_specs, strict=False):\n        if option.endswith(\"k1\"):\n            if \"_\" in spec:\n                lower_str, upper_str = spec.split(\"_\", 1)\n            else:\n                upper_str = spec\n                lower_str = str(1.0 / float(upper_str))\n            try:\n                lower = float(lower_str)\n                upper = float(upper_str)\n            except ValueError as exc:\n                raise ValueError(f\"Invalid numeric threshold '{spec}' for option '{option}'.\") from exc\n            if lower <= 0 or upper <= 0:\n                raise ValueError(f\"Thresholds for option '{option}' must be positive, got {spec}.\")\n            thresholds[option] = {\n                \"lower\": lower,\n                \"upper\": upper,\n            }\n        else:\n            if \"_\" in spec:\n                raise ValueError(\n                    f\"rollout_rs_threshold for option '{option}' must provide a single upper bound \"\n                    f\"without '_'. Received '{spec}'.\"\n                )\n            try:\n                upper = float(spec)\n            except ValueError as exc:\n                raise ValueError(f\"Invalid numeric threshold '{spec}' for option '{option}'.\") from exc\n            if upper <= 0:\n                raise ValueError(f\"Threshold for option '{option}' must be positive, got {spec}.\")\n            thresholds[option] = {\n                \"lower\": None,\n                \"upper\": upper,\n            }\n    return thresholds\n\n\ndef compute_rollout_rejection_mask(\n    log_ratio: torch.Tensor,\n    response_mask: torch.Tensor,\n    rollout_rs: str = \"token_k1\",\n    rollout_rs_threshold: Optional[str | float] = None,\n) -> tuple[torch.Tensor, dict[str, float]]:\n    \"\"\"Compute hard trust region mask using divergence estimators.\n\n    This function enforces a hard trust region constraint by masking tokens/sequences\n    where the estimated divergence (between training and rollout policies) exceeds\n    a threshold. Unlike PPO's soft clipping, this provides a hard boundary.\n\n    Multiple rejection criteria can be supplied via a comma separated `rollout_rs` string.\n    All requested options must pass for a token/sequence to remain valid.\n\n    Supported KL divergence-based modes (ideal = 0.0 unless noted):\n    - \"token_k{1,2,3}\": Token-level divergences.\n    - \"seq_sum_k{1,2,3}\": Sum of token divergences per sequence.\n    - \"seq_mean_k{1,2,3}\": Mean of token divergences per sequence.\n    - \"seq_max_k{2,3}\": Maximum token divergence per sequence.\n\n    Args:\n        log_ratio: Log ratio of training policy probability to rollout policy probability,\n            shape (batch_size, seq_length).\n        response_mask: Binary mask for valid tokens (1=valid, 0=padding),\n            shape (batch_size, seq_length).\n        rollout_rs: Comma separated rejection sampling options (e.g. \"token_k1,seq_sum_k3\").\n        rollout_rs_threshold: Threshold specification string (required). Provide one entry per\n            rollout_rs option separated by commas. Each entry must be a positive number.\n            For K1-style options (``*k1``), specify ``lower_upper`` (e.g. ``\"0.1_1.2\"``)\n            to denote lower/upper ratio bounds; other options accept a single upper bound.\n\n    Returns:\n        Tuple containing:\n            modified_response_mask: Response mask with trust region violations masked (0=rejected),\n                shape (batch_size, seq_length).\n            metrics: Dictionary of trust region metrics (all scalars).\n    \"\"\"\n    if rollout_rs is None or not isinstance(rollout_rs, str):\n        raise ValueError(\"rollout_rs must be a non-empty string (comma separated for multiple options).\")\n    if rollout_rs_threshold is None:\n        raise ValueError(\"rollout_rs_threshold must be provided for rejection sampling.\")\n\n    if log_ratio.shape[0] == 0:\n        return response_mask, {}\n\n    # rollout_rs supports chained criteria via comma separation (e.g. \"token_k1,seq_mean_k3\").\n    #            Every listed option must pass; combined_mask aggregates them via logical AND.\n    option_modes = [opt.strip() for opt in rollout_rs.split(\",\") if opt.strip()]\n    if not option_modes:\n        raise ValueError(\"rollout_rs must contain at least one valid option.\")\n\n    normalized_options: list[str] = []\n    seen: set[str] = set()\n    for opt in option_modes:\n        if opt not in SUPPORTED_ROLLOUT_RS_OPTIONS:\n            raise ValueError(\n                f\"Invalid rollout_rs option: {opt}. Must be one of {sorted(SUPPORTED_ROLLOUT_RS_OPTIONS)}.\"\n            )\n        if opt not in seen:\n            normalized_options.append(opt)\n            seen.add(opt)\n\n    threshold_specs = _parse_rollout_rs_thresholds(normalized_options, rollout_rs_threshold)\n\n    log_ratio_safe: torch.Tensor = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)\n    token_k1: torch.Tensor = -log_ratio_safe\n    token_k2: torch.Tensor = 0.5 * log_ratio_safe**2\n    token_k3: torch.Tensor = torch.exp(log_ratio_safe) - 1.0 - log_ratio_safe\n\n    response_mask_bool: torch.Tensor = response_mask.bool()\n    seq_valid_mask: torch.Tensor = response_mask.sum(dim=-1) > 0\n    # combined_mask accumulates per-option passes; any failure flips tokens to 0.\n    combined_mask: torch.Tensor = torch.ones_like(response_mask, dtype=log_ratio.dtype)\n    metrics: dict[str, float] = {}\n\n    def _sequence_sum(values: torch.Tensor) -> torch.Tensor:\n        return verl_F.masked_sum(values, response_mask, axis=-1)\n\n    def _sequence_mean(values: torch.Tensor) -> torch.Tensor:\n        return verl_F.masked_mean(values, response_mask, axis=-1)\n\n    def _sequence_max(values: torch.Tensor) -> torch.Tensor:\n        mask_bool = response_mask.bool()\n        neg_inf = torch.tensor(float(\"-inf\"), device=values.device, dtype=values.dtype)\n        masked_values = values.masked_fill(~mask_bool, neg_inf)\n        max_values = masked_values.max(dim=-1).values\n        return torch.where(max_values == neg_inf, torch.zeros_like(max_values), max_values)\n\n    for option_name in normalized_options:\n        thresholds_info = threshold_specs[option_name]\n        is_k1_option = option_name.endswith(\"k1\")\n        upper_value = thresholds_info[\"upper\"]\n        lower_value = thresholds_info[\"lower\"]\n        apply_lower_threshold = is_k1_option\n        lower_log: Optional[float] = None\n        upper_log: Optional[float] = None\n\n        if is_k1_option:\n            if lower_value is None or upper_value is None:\n                raise ValueError(\n                    f\"rollout_rs_threshold for option '{option_name}' must specify both lower and upper bounds.\"\n                )\n            lower_log = math.log(lower_value)\n            upper_log = math.log(upper_value)\n        else:\n            if upper_value is None:\n                raise ValueError(f\"rollout_rs_threshold for option '{option_name}' must specify an upper bound.\")\n\n        level = \"sequence\" if option_name not in TOKEN_LEVEL_ROLLOUT_RS_OPTIONS else \"token\"\n\n        per_token_stat: torch.Tensor\n        per_sequence_stat: Optional[torch.Tensor] = None\n        token_keep_bool: torch.Tensor\n\n        if option_name == \"token_k1\":\n            if lower_log is None:\n                raise ValueError(\"Threshold specification for token_k1 must include lower and upper bounds.\")\n            per_token_stat = token_k1\n            token_keep_bool = (per_token_stat >= lower_log) & (per_token_stat <= upper_log)\n        elif option_name == \"token_k2\":\n            per_token_stat = token_k2\n            token_keep_bool = per_token_stat <= upper_value\n        elif option_name == \"token_k3\":\n            per_token_stat = token_k3\n            token_keep_bool = per_token_stat <= upper_value\n        elif option_name.startswith(\"seq_sum\"):\n            if option_name.endswith(\"k1\"):\n                if lower_log is None:\n                    raise ValueError(\n                        f\"Threshold specification for option '{option_name}' must include lower and upper bounds.\"\n                    )\n                seq_stat = _sequence_sum(token_k1)\n                seq_keep_bool_direct = (seq_stat >= lower_log) & (seq_stat <= upper_log)\n            elif option_name.endswith(\"k2\"):\n                seq_stat = _sequence_sum(token_k2)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            elif option_name.endswith(\"k3\"):\n                seq_stat = _sequence_sum(token_k3)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            else:\n                raise ValueError(f\"Unsupported rollout_rs option: {option_name}.\")\n            per_sequence_stat = seq_stat\n            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)\n            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)\n        elif option_name.startswith(\"seq_mean\"):\n            if option_name.endswith(\"k1\"):\n                if lower_log is None:\n                    raise ValueError(\n                        f\"Threshold specification for option '{option_name}' must include lower and upper bounds.\"\n                    )\n                seq_stat = _sequence_mean(token_k1)\n                seq_keep_bool_direct = (seq_stat >= lower_log) & (seq_stat <= upper_log)\n            elif option_name.endswith(\"k2\"):\n                seq_stat = _sequence_mean(token_k2)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            elif option_name.endswith(\"k3\"):\n                seq_stat = _sequence_mean(token_k3)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            else:\n                raise ValueError(f\"Unsupported rollout_rs option: {option_name}.\")\n            per_sequence_stat = seq_stat\n            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)\n            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)\n        elif option_name.startswith(\"seq_max\"):\n            if option_name.endswith(\"k2\"):\n                seq_stat = _sequence_max(token_k2)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            elif option_name.endswith(\"k3\"):\n                seq_stat = _sequence_max(token_k3)\n                seq_keep_bool_direct = seq_stat <= upper_value\n            else:\n                raise ValueError(f\"Unsupported rollout_rs option: {option_name}.\")\n            per_sequence_stat = seq_stat\n            token_keep_bool = seq_keep_bool_direct.unsqueeze(-1).expand_as(response_mask_bool)\n            per_token_stat = seq_stat.unsqueeze(-1).expand_as(response_mask)\n        else:\n            raise ValueError(f\"Unsupported rollout_rs option: {option_name}.\")\n\n        metrics_upper_threshold = upper_log if is_k1_option else upper_value\n        metrics_lower_threshold = lower_log if (is_k1_option and lower_log is not None) else 0.0\n\n        token_keep_mask = token_keep_bool.to(dtype=log_ratio.dtype)\n        combined_mask = combined_mask * token_keep_mask\n        seq_keep_bool_tensor = (~((~token_keep_bool) & response_mask_bool)).all(dim=-1)\n\n        option_metrics = compute_rs_metrics(\n            option_name=option_name,\n            rs_statistic=per_token_stat,\n            response_mask=response_mask,\n            seq_valid_mask=seq_valid_mask,\n            level=level,\n            per_sequence_values=per_sequence_stat,\n            rollout_rs_threshold=metrics_upper_threshold,\n            rollout_rs_threshold_lower=metrics_lower_threshold,\n            apply_lower_threshold=apply_lower_threshold,\n        )\n        metrics.update(option_metrics)\n\n        token_masked_fraction = verl_F.masked_mean(1 - token_keep_mask, response_mask).item()\n        seq_valid_float = seq_valid_mask.float()\n        if seq_valid_float.sum() > 0:\n            seq_keep_float = seq_keep_bool_tensor.to(dtype=log_ratio.dtype)\n            seq_masked_fraction = (((1.0 - seq_keep_float) * seq_valid_float).sum() / seq_valid_float.sum()).item()\n        else:\n            seq_masked_fraction = 0.0\n        metrics[f\"rollout_rs_{option_name}_masked_fraction\"] = token_masked_fraction\n        metrics[f\"rollout_rs_{option_name}_seq_masked_fraction\"] = seq_masked_fraction\n\n    final_mask = combined_mask\n    metrics[\"rollout_rs_masked_fraction\"] = verl_F.masked_mean(1 - final_mask, response_mask).item()\n    final_keep_bool = (final_mask > 0.5) & response_mask_bool\n    seq_has_masked: torch.Tensor = (~final_keep_bool & response_mask_bool).any(dim=-1)\n    metrics[\"rollout_rs_seq_masked_fraction\"] = seq_has_masked.float().mean().item()\n\n    modified_response_mask: torch.Tensor = (response_mask * final_mask).to(dtype=response_mask.dtype)\n    return modified_response_mask, metrics\n\n\ndef compute_rs_metrics(\n    option_name: str,\n    rs_statistic: torch.Tensor,\n    response_mask: torch.Tensor,\n    seq_valid_mask: torch.Tensor,\n    *,\n    level: str,\n    per_sequence_values: Optional[torch.Tensor],\n    rollout_rs_threshold: float,\n    rollout_rs_threshold_lower: float,\n    apply_lower_threshold: bool,\n) -> dict[str, float]:\n    \"\"\"Compute metrics for hard trust region enforcement (per-option).\n\n    Args:\n        option_name: Original option string supplied by the user.\n        rs_statistic: Trust region statistic (per token) used for thresholding.\n        response_mask: Binary mask for valid tokens (1=valid, 0=padding).\n        seq_valid_mask: Boolean mask indicating sequences with at least one valid token.\n        level: \"token\" or \"sequence\" describing aggregation level.\n        per_sequence_values: Optional per-sequence statistic (same semantics as rs_statistic).\n        rollout_rs_threshold: Upper threshold.\n        rollout_rs_threshold_lower: Lower threshold (ignored if ``apply_lower_threshold`` is False).\n        apply_lower_threshold: Whether to mask/log metrics for values below the lower threshold.\n    \"\"\"\n    if not response_mask.any():\n        raise ValueError(\"response_mask must contain at least one valid token (1).\")\n\n    metrics: dict[str, float] = {}\n    prefix = f\"rollout_rs_{option_name}\"\n    mask_bool: torch.Tensor = response_mask.bool()\n\n    # Compute sequence statistics (used by several metrics).\n    if per_sequence_values is not None:\n        seq_values = per_sequence_values\n    else:\n        seq_values = verl_F.masked_mean(rs_statistic, response_mask, axis=-1)\n    if seq_values.dim() > 1:\n        seq_values = seq_values.squeeze(-1)\n    seq_values_valid = seq_values[seq_valid_mask]\n\n    # Mean of the statistic (always reported).\n    metrics[f\"{prefix}_mean\"] = verl_F.masked_mean(rs_statistic, response_mask).item()\n\n    # Max/min values.\n    if level == \"sequence\" and seq_values_valid.numel() > 0:\n        metrics[f\"{prefix}_max\"] = seq_values_valid.max().item()\n        metrics[f\"{prefix}_min\"] = seq_values_valid.min().item()\n    else:\n        metrics[f\"{prefix}_max\"] = rs_statistic.masked_fill(~mask_bool, float(\"-inf\")).max().item()\n        metrics[f\"{prefix}_min\"] = rs_statistic.masked_fill(~mask_bool, float(\"inf\")).min().item()\n\n    # Fractions above/below the thresholds.\n    if level == \"sequence\" and seq_values_valid.numel() > 0:\n        fraction_high = (seq_values_valid > rollout_rs_threshold).float().mean().item()\n        fraction_low = (\n            (seq_values_valid < rollout_rs_threshold_lower).float().mean().item() if apply_lower_threshold else 0.0\n        )\n    else:\n        fraction_high = verl_F.masked_mean((rs_statistic > rollout_rs_threshold).float(), response_mask).item()\n        fraction_low = (\n            verl_F.masked_mean((rs_statistic < rollout_rs_threshold_lower).float(), response_mask).item()\n            if apply_lower_threshold\n            else 0.0\n        )\n    metrics[f\"{prefix}_fraction_high\"] = fraction_high\n    metrics[f\"{prefix}_fraction_low\"] = fraction_low\n\n    # Standard deviation (clamped for stability).\n    mask_count: torch.Tensor = response_mask.sum()\n    if mask_count > 1:\n        if apply_lower_threshold:\n            clamp_min = rollout_rs_threshold_lower\n        else:\n            clamp_min = 0.0\n        stat_for_std: torch.Tensor = rs_statistic.clamp(min=clamp_min, max=rollout_rs_threshold)\n        mean_clamped: torch.Tensor = verl_F.masked_mean(stat_for_std, response_mask)\n        stat_var: torch.Tensor = verl_F.masked_mean(stat_for_std.square(), response_mask) - mean_clamped.square()\n        metrics[f\"{prefix}_std\"] = torch.sqrt(torch.clamp(stat_var, min=0.0)).item()\n    else:\n        metrics[f\"{prefix}_std\"] = 0.0\n\n    # Sequence-level summary metrics.\n    if seq_values_valid.numel() > 0:\n        metrics[f\"{prefix}_seq_mean\"] = seq_values_valid.mean().item()\n        metrics[f\"{prefix}_seq_std\"] = seq_values_valid.std().item() if seq_values_valid.numel() > 1 else 0.0\n        metrics[f\"{prefix}_seq_max\"] = seq_values_valid.max().item()\n        metrics[f\"{prefix}_seq_min\"] = seq_values_valid.min().item()\n        metrics[f\"{prefix}_seq_max_deviation\"] = (seq_values_valid - 0.0).abs().max().item()\n        metrics[f\"{prefix}_seq_fraction_high\"] = (seq_values_valid > rollout_rs_threshold).float().mean().item()\n        if apply_lower_threshold:\n            metrics[f\"{prefix}_seq_fraction_low\"] = (\n                (seq_values_valid < rollout_rs_threshold_lower).float().mean().item()\n            )\n    else:\n        metrics[f\"{prefix}_seq_mean\"] = 0.0\n        metrics[f\"{prefix}_seq_std\"] = 0.0\n        metrics[f\"{prefix}_seq_max\"] = 0.0\n        metrics[f\"{prefix}_seq_min\"] = 0.0\n        metrics[f\"{prefix}_seq_max_deviation\"] = 0.0\n        metrics[f\"{prefix}_seq_fraction_high\"] = 0.0\n        metrics[f\"{prefix}_seq_fraction_low\"] = 0.0\n\n    return metrics\n\n\ndef compute_rollout_correction_weights(\n    log_ratio: torch.Tensor,\n    response_mask: torch.Tensor,\n    rollout_is: str = \"token\",\n    rollout_is_threshold: float = 2.0,\n    rollout_is_batch_normalize: bool = False,\n) -> tuple[torch.Tensor, dict[str, float]]:\n    \"\"\"Compute importance sampling weights to correct for off-policy distribution shifts.\n\n    This function calculates IS weights (π_train / π_rollout) using log ratios for numerical stability.\n    It supports multiple aggregation levels and truncates extreme weights to prevent training instability.\n\n    Key design:\n    - Log-space computations to avoid overflow\n    - Truncation of extreme weights (TIS: Truncated Importance Sampling)\n    - Optional batch normalization (normalize to mean=1.0)\n    - Metrics tracking for weight distribution analysis\n\n    Args:\n        log_ratio: Log ratio of training policy probability to rollout policy probability,\n            shape (batch_size, seq_length).\n        response_mask: Binary mask for valid tokens (1=valid, 0=padding),\n            shape (batch_size, seq_length).\n        rollout_is: IS weight aggregation level, must be one of:\n            - \"token\": Per-token weights (biased, low variance)\n            - \"sequence\": Per-sequence weight (product of tokens; unbiased, high variance)\n        rollout_is_threshold: Upper threshold for truncating extreme weights (e.g., 2.0),\n            default 2.0.\n        rollout_is_batch_normalize: Whether to normalize IS weights to have mean=1.0 per batch,\n            default False.\n\n    Returns:\n        Tuple containing:\n            rollout_is_weights: Truncated IS weights (masked to zero for padding tokens),\n                shape (batch_size, seq_length). If batch_normalize=True, normalized to mean=1.0.\n            metrics: Dictionary of IS weight metrics (all scalars), including:\n                - rollout_is_mean/max/min: Statistic of weights (before batch normalization)\n                - rollout_is_eff_sample_size: Effective sample size (ESS)\n                - rollout_is_seq_*: Sequence-level weight statistics\n                - rollout_is_batch_norm_factor: Normalization factor (only if batch_normalize=True)\n    \"\"\"\n    # Validate input parameters\n    valid_is_levels = {\"token\", \"sequence\"}\n    if rollout_is not in valid_is_levels:\n        raise ValueError(f\"Invalid rollout_is: {rollout_is}. Must be one of {valid_is_levels}.\")\n    if rollout_is_threshold <= 0:\n        raise ValueError(f\"rollout_is_threshold must be positive, got {rollout_is_threshold}.\")\n\n    # Compute IS weights from log ratio (handles different aggregation levels)\n    if rollout_is == \"token\":\n        # Per-token IS weight: exp(log(π_train/π_rollout)) with safety clamp\n        log_ratio_for_metrics: torch.Tensor = log_ratio\n        log_ratio_safe: torch.Tensor = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)\n        rollout_is_weights: torch.Tensor = torch.exp(log_ratio_safe)\n\n    elif rollout_is == \"sequence\":\n        # Sequence-level IS weight: product of token ratios (exp(sum(log ratios)))\n        log_ratio_sum: torch.Tensor = verl_F.masked_sum(log_ratio, response_mask, axis=-1).unsqueeze(\n            -1\n        )  # Shape: (batch_size, 1)\n        log_ratio_for_metrics = log_ratio_sum\n\n        log_ratio_sum_safe: torch.Tensor = torch.clamp(log_ratio_sum, min=-SAFETY_BOUND, max=SAFETY_BOUND)\n        rollout_is_weights = torch.exp(log_ratio_sum_safe).expand_as(log_ratio)  # Broadcast to sequence length\n\n    else:\n        raise ValueError(f\"Unsupported rollout_is: {rollout_is}\")\n\n    # Zero out weights for padding tokens using response mask\n    rollout_is_weights = rollout_is_weights * response_mask\n\n    # Compute IS weight metrics (BEFORE truncation to get accurate fraction_high/low)\n    metrics: dict[str, float] = compute_is_metrics(\n        rollout_is_weights=rollout_is_weights,\n        log_ratio_for_metrics=log_ratio_for_metrics,\n        response_mask=response_mask,\n        rollout_is=rollout_is,\n        rollout_is_threshold=rollout_is_threshold,\n    )\n\n    # Truncate extreme weights (TIS: Truncated Importance Sampling)\n    rollout_is_weights = rollout_is_weights.clamp(max=rollout_is_threshold)\n\n    # Detach weights to prevent gradient flow (mathematically required by IS theory)\n    # IS weights change the measure, not the objective. See §3.2.2 in docs/algo/rollout_corr_math.md\n    rollout_is_weights = rollout_is_weights.detach()\n\n    # Apply batch normalization if requested\n    if rollout_is_batch_normalize:\n        # Compute mean based on aggregation level\n        mask_float = response_mask.to(dtype=rollout_is_weights.dtype)\n        if rollout_is == \"token\":\n            # Token-level: normalize over all token weights\n            if torch.distributed.is_available() and torch.distributed.is_initialized():\n                weights_mean = verl_F.distributed_masked_mean(rollout_is_weights, mask_float)\n            else:\n                weights_mean = verl_F.masked_mean(rollout_is_weights, response_mask)\n        elif rollout_is == \"sequence\":\n            # Sequence-level: normalize over sequence weights (one weight per sequence)\n            # For each sequence, compute mean over valid tokens (they all have the same weight)\n            # then average across sequences\n            seq_weights = verl_F.masked_mean(rollout_is_weights, response_mask, axis=-1)  # (batch_size,)\n            seq_mask = (response_mask.sum(dim=-1) > 0).to(dtype=rollout_is_weights.dtype)\n            if torch.distributed.is_available() and torch.distributed.is_initialized():\n                weights_mean = verl_F.distributed_masked_mean(seq_weights, seq_mask)\n            else:\n                weights_mean = (seq_weights * seq_mask).sum() / seq_mask.sum().clamp_min(1e-8)\n        else:\n            raise ValueError(f\"Unsupported rollout_is: {rollout_is}\")\n\n        # Normalize to mean=1.0 (avoid division by zero)\n        if weights_mean > 1e-8:\n            rollout_is_weights = rollout_is_weights / weights_mean\n            metrics[\"rollout_is_batch_norm_factor\"] = weights_mean.item()\n        else:\n            metrics[\"rollout_is_batch_norm_factor\"] = 1.0\n\n    return rollout_is_weights, metrics\n\n\ndef compute_is_metrics(\n    rollout_is_weights: torch.Tensor,\n    log_ratio_for_metrics: torch.Tensor,\n    response_mask: torch.Tensor,\n    rollout_is: str,\n    rollout_is_threshold: float,\n) -> dict[str, float]:\n    \"\"\"Compute comprehensive metrics for truncated importance sampling weights.\n\n    This function calculates statistics for truncated IS weights (TIS), using log-space\n    for accurate threshold checks and clamped weights for stable mean/std calculations.\n\n    Args:\n        rollout_is_weights: Truncated IS weights (π_train / π_rollout),\n            shape (batch_size, seq_length).\n        log_ratio_for_metrics: Log ratio of training to rollout probabilities (unclamped),\n            shape varies by aggregation level.\n        response_mask: Binary mask for valid tokens (1=valid, 0=padding),\n            shape (batch_size, seq_length).\n        rollout_is: IS weight aggregation level (matches compute_rollout_correction_weights).\n        rollout_is_threshold: Upper threshold for truncated IS weights.\n\n    Returns:\n        Dictionary of IS weight metrics (all scalars).\n    \"\"\"\n    if not response_mask.any():\n        raise ValueError(\"response_mask must contain at least one valid token (1).\")\n\n    metrics: dict[str, float] = {}\n    device: torch.device = rollout_is_weights.device\n    # Default lower threshold (reciprocal of upper threshold)\n    rollout_is_threshold_lower: float = 1.0 / rollout_is_threshold\n\n    # Precompute log thresholds for accurate checks\n    log_threshold_upper: torch.Tensor = torch.log(torch.tensor(rollout_is_threshold, device=device))\n    log_threshold_lower: torch.Tensor = torch.log(torch.tensor(rollout_is_threshold_lower, device=device))\n\n    # Compute metrics based on aggregation level\n    if rollout_is == \"sequence\":\n        # Sequence-level aggregation: use log-space for unclamped stats\n        log_max: torch.Tensor = log_ratio_for_metrics.max()\n        log_min: torch.Tensor = log_ratio_for_metrics.min()\n        metrics[\"rollout_is_max\"] = torch.exp(torch.clamp(log_max, max=SAFETY_BOUND)).item()\n        metrics[\"rollout_is_min\"] = torch.exp(log_min).item()\n\n        # Mean uses truncated weights to avoid overflow\n        metrics[\"rollout_is_mean\"] = verl_F.masked_mean(rollout_is_weights, response_mask).item()\n\n        # Fraction of weights exceeding thresholds (log-space for accuracy)\n        exceeds_upper: torch.Tensor = log_ratio_for_metrics > log_threshold_upper\n        below_lower: torch.Tensor = log_ratio_for_metrics < log_threshold_lower\n        metrics[\"rollout_is_ratio_fraction_high\"] = exceeds_upper.float().mean().item()\n        metrics[\"rollout_is_ratio_fraction_low\"] = below_lower.float().mean().item()\n\n    else:  # token-level\n        # Token-level aggregation: compute directly from truncated weights\n        metrics[\"rollout_is_mean\"] = verl_F.masked_mean(rollout_is_weights, response_mask).item()\n\n        # Fraction of tokens exceeding thresholds\n        rollout_is_above_threshold: torch.Tensor = rollout_is_weights > rollout_is_threshold\n        rollout_is_below_threshold: torch.Tensor = rollout_is_weights < rollout_is_threshold_lower\n        metrics[\"rollout_is_ratio_fraction_high\"] = verl_F.masked_mean(\n            rollout_is_above_threshold.float(), response_mask\n        ).item()\n        metrics[\"rollout_is_ratio_fraction_low\"] = verl_F.masked_mean(\n            rollout_is_below_threshold.float(), response_mask\n        ).item()\n\n        # Max/min (mask out padding tokens)\n        mask_bool: torch.Tensor = response_mask.bool()\n        metrics[\"rollout_is_max\"] = rollout_is_weights.masked_fill(~mask_bool, float(\"-inf\")).max().item()\n        metrics[\"rollout_is_min\"] = rollout_is_weights.masked_fill(~mask_bool, float(\"inf\")).min().item()\n\n    # Compute standard deviation (using clamped weights for stability)\n    mask_count: torch.Tensor = response_mask.sum()\n    if mask_count > 1:\n        weights_for_std: torch.Tensor = rollout_is_weights.clamp(\n            min=rollout_is_threshold_lower, max=rollout_is_threshold\n        )\n        mean_clamped: torch.Tensor = verl_F.masked_mean(weights_for_std, response_mask)\n        rollout_is_var: torch.Tensor = (\n            verl_F.masked_mean(weights_for_std.square(), response_mask) - mean_clamped.square()\n        )\n        metrics[\"rollout_is_std\"] = torch.sqrt(torch.clamp(rollout_is_var, min=0.0)).item()\n    else:\n        metrics[\"rollout_is_std\"] = 0.0\n\n    # Compute Effective Sample Size (ESS) for truncated weights\n    weights_for_ess: torch.Tensor = rollout_is_weights.clamp(min=rollout_is_threshold_lower, max=rollout_is_threshold)\n    mean_for_ess: torch.Tensor = verl_F.masked_mean(weights_for_ess, response_mask)\n    is_weights_normalized: torch.Tensor = weights_for_ess / (mean_for_ess + 1e-8)  # Avoid division by zero\n    metrics[\"rollout_is_eff_sample_size\"] = (\n        1.0 / verl_F.masked_mean(is_weights_normalized.square(), response_mask).item()\n    )\n\n    # Add sequence-level metrics if weights have batch dimension\n    if rollout_is_weights.dim() > 1:\n        seq_mean_weights: torch.Tensor = verl_F.masked_mean(rollout_is_weights, response_mask, axis=-1)\n\n        metrics[\"rollout_is_seq_mean\"] = seq_mean_weights.mean().item()\n        metrics[\"rollout_is_seq_std\"] = seq_mean_weights.std().item() if seq_mean_weights.numel() > 1 else 0.0\n        metrics[\"rollout_is_seq_max\"] = seq_mean_weights.max().item()\n        metrics[\"rollout_is_seq_min\"] = seq_mean_weights.min().item()\n\n        # Sequence deviation from ideal weight (1.0)\n        seq_deviation: torch.Tensor = (seq_mean_weights - 1.0).abs()\n        metrics[\"rollout_is_seq_max_deviation\"] = seq_deviation.max().item()\n\n        # Fraction of sequences with extreme weights\n        metrics[\"rollout_is_seq_fraction_high\"] = (seq_mean_weights > rollout_is_threshold).float().mean().item()\n        metrics[\"rollout_is_seq_fraction_low\"] = (seq_mean_weights < rollout_is_threshold_lower).float().mean().item()\n\n    return metrics\n\n\ndef compute_rollout_correction_and_rejection_mask(\n    old_log_prob: torch.Tensor,\n    rollout_log_prob: torch.Tensor,\n    response_mask: torch.Tensor,\n    rollout_is: Optional[str] = None,\n    rollout_is_threshold: Optional[float] = 2.0,\n    rollout_is_batch_normalize: bool = False,\n    rollout_rs: Optional[str] = None,\n    rollout_rs_threshold: Optional[str | float] = None,\n) -> tuple[Optional[DataProto], torch.Tensor, dict[str, float]]:\n    \"\"\"Unified interface for computing IS weights and rejection masks.\n\n    This function combines IS weight calculation (truncated) and rejection sampling (masked)\n    into a single pipeline.\n\n    Key design:\n    - Separation of IS weights (for variance reduction) and rejection masks (for sample filtering)\n    - Comprehensive metrics tracking for mismatch diagnosis\n\n    Args:\n        old_log_prob: Log probabilities from the training policy (e.g., FSDP FP32),\n            shape (batch_size, seq_length).\n        rollout_log_prob: Log probabilities from the rollout policy (e.g., vLLM BF16),\n            shape (batch_size, seq_length).\n        response_mask: Binary mask for valid tokens (1=valid, 0=padding),\n            shape (batch_size, seq_length).\n        rollout_is: IS weight aggregation level (see compute_rollout_correction_weights for options).\n            Set to None to disable IS weight computation.\n        rollout_is_threshold: Upper threshold for truncated IS weights (used if rollout_is is set),\n            default 2.0.\n        rollout_rs: Rejection sampling aggregation modes as a comma separated string\n            (see compute_rollout_rejection_mask for the full list). Set to None to disable\n            rejection sampling.\n        rollout_rs_threshold: Threshold specification string (see compute_rollout_rejection_mask for details).\n            Provide one threshold per option (comma separated). For K1-style options, specify\n            ``lower_upper`` to denote the lower/upper ratio bounds.\n        rollout_is_batch_normalize: Whether to normalize IS weights to have mean=1.0 per batch.\n            Default: False.\n\n    Returns:\n        Tuple containing:\n            rollout_is_weights_proto: DataProto with IS weights (None if rollout_is is None),\n                key \"rollout_is_weights\", shape (batch_size, seq_length).\n            modified_response_mask: Response mask with rejection sampling applied,\n                shape (batch_size, seq_length).\n            metrics: Dictionary of all metrics (prefixed with \"rollout_corr/\"), including:\n                - IS weight statistics\n                - Rejection sampling rates\n                - Policy mismatch metrics (KL, PPL, etc.)\n    \"\"\"\n    # Validate input masks\n    if not response_mask.any():\n        raise ValueError(\"response_mask must contain at least one valid token (1).\")\n    if old_log_prob.shape != rollout_log_prob.shape:\n        raise ValueError(\n            f\"old_log_prob shape {old_log_prob.shape} does not match rollout_log_prob shape {rollout_log_prob.shape}.\"\n        )\n    if old_log_prob.shape != response_mask.shape:\n        raise ValueError(\n            f\"log_prob shape {old_log_prob.shape} does not match response_mask shape {response_mask.shape}.\"\n        )\n\n    # Step 1: Compute log ratio (log(π_train / π_rollout))\n    log_ratio: torch.Tensor = old_log_prob - rollout_log_prob\n    metrics: dict[str, float] = {}\n\n    # Step 2: Compute IS weights (if enabled)\n    rollout_is_weights: Optional[torch.Tensor] = None\n    if rollout_is is not None and rollout_is_threshold is not None:\n        rollout_is_weights, is_metrics = compute_rollout_correction_weights(\n            log_ratio=log_ratio,\n            response_mask=response_mask,\n            rollout_is=rollout_is,\n            rollout_is_threshold=rollout_is_threshold,\n            rollout_is_batch_normalize=rollout_is_batch_normalize,\n        )\n        metrics.update(is_metrics)\n\n    # Step 3: Compute rejection mask (if enabled)\n    modified_response_mask: torch.Tensor = response_mask.clone()\n    if rollout_rs is not None:\n        if rollout_rs_threshold is None:\n            raise ValueError(\n                \"rollout_rs_threshold must be explicitly provided when rollout_rs is enabled. \"\n                \"Set rollout_rs_threshold to the desired threshold value.\"\n            )\n        modified_response_mask, rs_metrics = compute_rollout_rejection_mask(\n            log_ratio=log_ratio,\n            response_mask=response_mask,\n            rollout_rs=rollout_rs,\n            rollout_rs_threshold=rollout_rs_threshold,\n        )\n        metrics.update(rs_metrics)\n\n    # Step 4: Compute off-policy metrics (KL, PPL, χ², etc.)\n    offpolicy_metrics: dict[str, float] = compute_offpolicy_metrics(\n        old_log_prob=old_log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n    )\n    metrics.update(offpolicy_metrics)\n\n    # Step 6: Add \"rollout_corr/\" prefix to all metrics for logging consistency\n    metrics_scalar: dict[str, float] = {}\n    for key, value in metrics.items():\n        if isinstance(value, torch.Tensor):\n            metrics_scalar[f\"rollout_corr/{key}\"] = value.item()\n        else:\n            metrics_scalar[f\"rollout_corr/{key}\"] = value\n\n    # Step 7: Wrap IS weights in DataProto for consistency with API\n    rollout_is_weights_proto: Optional[DataProto] = None\n    if rollout_is_weights is not None:\n        rollout_is_weights_proto = DataProto.from_dict(tensors={\"rollout_is_weights\": rollout_is_weights})\n\n    return rollout_is_weights_proto, modified_response_mask, metrics_scalar\n\n\ndef compute_offpolicy_metrics(\n    old_log_prob: torch.Tensor,\n    rollout_log_prob: Optional[torch.Tensor],\n    response_mask: torch.Tensor,\n) -> dict[str, Any]:\n    \"\"\"Compute off-policy diagnostic metrics (helper function).\n\n    This helper function operates on raw tensors and is used internally by:\n    - compute_rollout_correction_and_rejection_mask() in this module (automatically included)\n    - Tests (test_rollout_corr.py, test_rollout_corr_integration.py)\n\n    These metrics help diagnose the off-policy gap between rollout and training policies,\n    which can arise from:\n    - Policy mismatch (e.g., vLLM BF16 vs FSDP FP32)\n    - Model staleness (training on trajectories from older checkpoints)\n    - General distribution shifts\n\n    Key metrics:\n    - kl: Direct KL divergence estimator KL(π_rollout || π_training)\n    - k3_kl: K3 KL estimator for stability (more stable for small KL)\n    - training_ppl: Perplexity of training policy\n    - rollout_ppl: Perplexity of rollout policy\n    - log_ppl_diff: Difference in log perplexities\n    - ppl_ratio: Ratio of training PPL to rollout PPL\n    - chi2_token: Token-level χ² divergence E[ρ²] - 1\n    - chi2_seq: Sequence-level χ² divergence E[(∏ρ_t)²] - 1\n\n    Args:\n        old_log_prob: Log probabilities from training policy, shape (batch_size, seq_length)\n        rollout_log_prob: Log probabilities from rollout policy, shape (batch_size, seq_length)\n        response_mask: Mask for valid tokens, shape (batch_size, seq_length)\n\n    Returns:\n        Dictionary of off-policy metrics (without prefix)\n    \"\"\"\n    # Validate that we have at least one valid token\n    assert response_mask.any(), \"Expected at least one valid token in response_mask\"\n\n    metrics = {}\n\n    # 1. Training policy perplexity (always available)\n    # Formula: exp(-1/|T| * Σ log π_training(y_t|y_<t))\n    # where |T| is the number of tokens generated by the model\n    mean_log_prob_training = verl_F.masked_mean(old_log_prob, response_mask, axis=-1)  # (batch_size,)\n    training_ppl = torch.exp(-mean_log_prob_training).mean()  # Batch mean of per-sequence PPL\n    metrics[\"training_ppl\"] = training_ppl.detach().item()\n\n    # Also log log-ppl for easier analysis (avoids exponential scale)\n    metrics[\"training_log_ppl\"] = (-mean_log_prob_training).mean().detach().item()\n\n    # 2. Compute rollout off-policy metrics (only if rollout_log_probs available)\n    if rollout_log_prob is not None:\n        # 2a. kl: Direct estimator for KL(π_rollout || π_training)\n        # This is the standard KL divergence: E[log(π_rollout) - log(π_training)]\n        # Positive value means rollout policy is more confident than training policy\n        metrics[\"kl\"] = verl_F.masked_mean(rollout_log_prob - old_log_prob, response_mask).detach().item()\n\n        # 2b. k3_kl: K3 estimator for KL(π_rollout || π_training)\n        # More stable for small KL values using: E[exp(log_ratio) - log_ratio - 1]\n        # Formula: KL ≈ E[r - log(r) - 1] where r = π_training/π_rollout\n        log_ratio = old_log_prob - rollout_log_prob\n        k3_kl_matrix = torch.exp(log_ratio) - log_ratio - 1\n        metrics[\"k3_kl\"] = verl_F.masked_mean(k3_kl_matrix, response_mask).detach().item()\n\n        # 2c. Rollout policy perplexity\n        mean_log_prob_rollout = verl_F.masked_mean(rollout_log_prob, response_mask, axis=-1)  # (batch_size,)\n        rollout_ppl = torch.exp(-mean_log_prob_rollout).mean()  # Batch mean of per-sequence PPL\n        metrics[\"rollout_ppl\"] = rollout_ppl.detach().item()\n        metrics[\"rollout_log_ppl\"] = (-mean_log_prob_rollout).mean().detach().item()\n\n        # 2d. Log PPL difference (sequence-level perplexity difference)\n        # log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training\n        # Since ppl = exp(-log_prob), we have:\n        #   log(ppl_ratio) = log(training_ppl/rollout_ppl) = log_ppl_diff\n        # Positive value means training assigns lower probability (higher PPL) than rollout\n        log_ppl_diff = mean_log_prob_rollout - mean_log_prob_training\n        metrics[\"log_ppl_diff\"] = log_ppl_diff.mean().detach().item()\n        metrics[\"log_ppl_abs_diff\"] = log_ppl_diff.abs().mean().detach().item()\n        metrics[\"log_ppl_diff_max\"] = log_ppl_diff.max().detach().item()\n        metrics[\"log_ppl_diff_min\"] = log_ppl_diff.min().detach().item()\n\n        # 2e. PPL ratio (how much higher is training PPL vs rollout PPL)\n        # Compute per-sequence ratio first, then average\n        # For numerical stability, compute in log space using log_ppl_diff\n        # Note: log_ppl_diff = log(ppl_ratio), so ppl_ratio = exp(log_ppl_diff)\n        # This is the inverse of geometric IS: ppl_ratio_i = 1 / geometric_is_i for each sequence\n        ppl_ratio = torch.exp(log_ppl_diff).mean()  # mean(exp(log_ppl_diff)) = mean(ppl_ratio_i)\n        metrics[\"ppl_ratio\"] = ppl_ratio.detach().item()\n\n        # 2f. Chi-squared divergence: χ²(π_training || π_rollout) = E_μ[ρ²] - 1\n        # where ρ = π_training / π_rollout and μ = π_rollout (rollout distribution)\n        # This measures the variance of importance sampling weights\n        # Token-level: E_token[ρ²] - 1 (averaged over all tokens)\n        log_ratio_safe = torch.clamp(log_ratio, min=-SAFETY_BOUND, max=SAFETY_BOUND)\n        rho_token = torch.exp(log_ratio_safe)  # ρ = π_training / π_rollout (token-level)\n        rho_squared_token = rho_token.square()\n        chi2_token = verl_F.masked_mean(rho_squared_token, response_mask) - 1.0\n        metrics[\"chi2_token\"] = chi2_token.detach().item()\n\n        # Sequence-level: E_seq[(Π ρ_t)²] - 1 = E_seq[exp(2 * Σ log ρ_t)] - 1\n        log_ratio_sum = verl_F.masked_sum(log_ratio, response_mask, axis=-1)  # Σ log ρ_t per sequence\n        log_ratio_sum_safe = torch.clamp(log_ratio_sum, min=-SAFETY_BOUND, max=SAFETY_BOUND)\n        rho_squared_seq = torch.exp(2.0 * log_ratio_sum_safe)  # (Π ρ_t)²\n        chi2_seq = rho_squared_seq.mean() - 1.0\n        metrics[\"chi2_seq\"] = chi2_seq.detach().item()\n\n    return metrics\n\n\ndef compute_rollout_correction_and_add_to_batch(\n    batch: DataProto, rollout_corr_config: RolloutCorrectionConfig\n) -> tuple[DataProto, dict]:\n    \"\"\"Compute rollout correction weights and apply rejection sampling.\n\n    Computes importance sampling weights to correct for off-policy issues between\n    rollout and training policies. Applies rejection sampling by modifying response_mask.\n    Always updates response_mask; conditionally adds IS weights.\n\n    Key behavior:\n    - response_mask: ALWAYS updated with rejection (RS exclusions removed from training)\n    - rollout_is_weights: Added to batch ONLY if rollout_is parameter is set\n\n    This separation ensures:\n    - Rejection works independently of IS weight application\n    - Metrics can be monitored before enabling IS weight correction\n\n    Args:\n        batch: DataProto with old_log_probs, rollout_log_probs, response_mask\n\n    Returns:\n        Tuple of (updated_batch, metrics):\n            updated_batch: Batch with modified response_mask (always) and rollout_is_weights (if enabled)\n            metrics: Dict of IS and off-policy metrics, all with \"rollout_corr/\" prefix\n\n    Note:\n        The implementation is copied from szrlee <szrlee@gmail.com>.\n    \"\"\"\n    # Get new API parameters directly from config\n    rollout_is = rollout_corr_config.get(\"rollout_is\", None)\n    rollout_is_threshold = rollout_corr_config.get(\"rollout_is_threshold\", 2.0)\n    rollout_is_batch_normalize = rollout_corr_config.get(\"rollout_is_batch_normalize\", False)\n    rollout_rs = rollout_corr_config.get(\"rollout_rs\", None)\n    rollout_rs_threshold = rollout_corr_config.get(\"rollout_rs_threshold\", None)\n\n    # Compute IS weights and get modified response_mask\n    rollout_is_weights, modified_response_mask, rollout_corr_metrics = compute_rollout_correction_and_rejection_mask(\n        old_log_prob=batch.batch[\"old_log_probs\"],\n        rollout_log_prob=batch.batch[\"rollout_log_probs\"],\n        response_mask=batch.batch[\"response_mask\"],\n        rollout_is=rollout_is,\n        rollout_is_threshold=rollout_is_threshold,\n        rollout_is_batch_normalize=rollout_is_batch_normalize,\n        rollout_rs=rollout_rs,\n        rollout_rs_threshold=rollout_rs_threshold,\n    )\n\n    # ALWAYS update response_mask with rejection applied\n    batch.batch[\"response_mask\"] = modified_response_mask\n\n    # Add IS weights to batch if computed\n    if rollout_is_weights is not None:\n        batch = batch.union(rollout_is_weights)\n\n    return batch, rollout_corr_metrics\n\n\ndef compute_rollout_corr_metrics_from_logprobs(\n    log_prob: torch.Tensor,\n    rollout_log_prob: torch.Tensor,\n    response_mask: torch.Tensor,\n) -> dict[str, float]:\n    \"\"\"Compute rollout correction metrics from log probabilities during training.\n\n    This function is used in the actor to compute metrics using the CURRENT policy\n    log probabilities versus rollout log probabilities, allowing tracking of the\n    off-policy gap as training progresses.\n\n    It computes off-policy diagnostic metrics (KL, PPL, χ²) from log probabilities.\n\n    Args:\n        log_prob: Current policy log probabilities, shape (batch_size, seq_length)\n        rollout_log_prob: Rollout policy log probabilities, shape (batch_size, seq_length)\n        response_mask: Valid token mask, shape (batch_size, seq_length)\n\n    Returns:\n        Dictionary of metrics with \"rollout_corr/\" prefix\n    \"\"\"\n    # Compute off-policy diagnostic metrics\n    offpolicy_metrics = compute_offpolicy_metrics(\n        old_log_prob=log_prob,\n        rollout_log_prob=rollout_log_prob,\n        response_mask=response_mask,\n    )\n\n    # Add rollout_corr/ prefix to all metrics\n    metrics_with_prefix = {}\n    for key, value in offpolicy_metrics.items():\n        if isinstance(value, torch.Tensor):\n            metrics_with_prefix[f\"rollout_corr/{key}\"] = value.item()\n        else:\n            metrics_with_prefix[f\"rollout_corr/{key}\"] = value\n\n    return metrics_with_prefix\n\n\ndef apply_bypass_mode(\n    batch: DataProto,\n    rollout_corr_config: Optional[RolloutCorrectionConfig] = None,\n    policy_loss_config: PolicyLossConfig = None,\n) -> None:\n    \"\"\"\n    Setup bypass mode: Use rollout_log_probs as old_log_probs.\n\n    Bypass mode skips expensive actor forward pass for old_log_prob computation\n    by setting old_log_probs = rollout_log_probs (2 policies instead of 3).\n\n    Uses compute_policy_loss_bypass_mode() which supports:\n    - loss_type=\"ppo_clip\" (default): PPO clipped objective (IS handled by ratio)\n    - loss_type=\"reinforce\": REINFORCE with explicit IS weights\n\n    Both loss types benefit from rejection sampling (RS) which masks out-of-distribution samples.\n\n    Note:\n        The implementation is copied from szrlee <szrlee@gmail.com>.\n    \"\"\"\n    from omegaconf import open_dict\n\n    if \"rollout_log_probs\" not in batch.batch:\n        raise ValueError(\n            \"bypass_mode=True requires rollout_log_probs in batch. \"\n            \"Ensure rollout worker is configured to calculate_log_probs=true.\"\n        )\n\n    # Use rollout log probs as old log probs (zero-cost substitution)\n    batch.batch[\"old_log_probs\"] = batch.batch[\"rollout_log_probs\"]\n\n    with open_dict(policy_loss_config):\n        # Pass rollout_correction config to actor for loss computation and metrics\n        policy_loss_config[\"rollout_correction\"] = rollout_corr_config\n        # Always use bypass_mode loss function which handles both loss_types\n        policy_loss_config[\"loss_mode\"] = \"bypass_mode\"\n"
  },
  {
    "path": "verl/trainer/ppo/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport warnings\nfrom enum import Enum\n\nfrom omegaconf import DictConfig\n\nfrom verl.single_controller.base import Worker\nfrom verl.trainer.ppo.core_algos import AdvantageEstimator\n\nWorkerType = type[Worker]\n\n\nclass Role(Enum):\n    \"\"\"\n    To create more roles dynamically, you can subclass Role and add new members\n    \"\"\"\n\n    Actor = 0\n    Rollout = 1\n    ActorRollout = 2\n    Critic = 3\n    RefPolicy = 4\n    RewardModel = 5\n    ActorRolloutRef = 6\n    Env = 7\n\n    def __str__(self):\n        return self._get_role_string()\n\n    def _get_role_string(self):\n        role_mapping = {\n            Role.Actor: \"actor\",\n            Role.Rollout: \"rollout\",\n            Role.ActorRollout: \"actor_rollout\",\n            Role.Critic: \"critic\",\n            Role.RefPolicy: \"ref\",\n            Role.RewardModel: \"rm\",\n            Role.ActorRolloutRef: \"actor_rollout_ref\",\n        }\n        return role_mapping.get(self, self.name.lower())\n\n    @classmethod\n    def from_string(cls, name: str):\n        string_mapping = {\n            \"actor\": cls.Actor,\n            \"rollout\": cls.Rollout,\n            \"actor_rollout\": cls.ActorRollout,\n            \"critic\": cls.Critic,\n            \"ref\": cls.RefPolicy,\n            \"rm\": cls.RewardModel,\n            \"actor_rollout_ref\": cls.ActorRolloutRef,\n        }\n        role = string_mapping.get(name.lower())\n        if role is None:\n            raise ValueError(f\"No Role found for string: {name}\")\n        return role\n\n\ndef need_reference_policy(\n    config: DictConfig,\n) -> bool:\n    \"\"\"Given the config, do we need ref policy.\"\"\"\n    return config.algorithm.use_kl_in_reward or config.actor_rollout_ref.actor.use_kl_loss\n\n\ndef need_reward_model(\n    config: DictConfig,\n) -> bool:\n    \"\"\"Given the config, do we need reward model.\"\"\"\n    return config.reward.reward_model.enable\n\n\ndef need_critic(config: DictConfig) -> bool:\n    \"\"\"Given a config, do we need critic.\"\"\"\n    if config.critic.enable is not None:\n        return bool(config.critic.enable)\n    elif config.algorithm.adv_estimator == AdvantageEstimator.GAE:\n        return True\n    else:\n        warnings.warn(\n            \"Disabled critic as algorithm.adv_estimator != gae. If it is not intended, please set critic.enable=True\",\n            stacklevel=2,\n        )\n        return False\n"
  },
  {
    "path": "verl/trainer/runtime_env.yaml",
    "content": "working_dir: ./\nexcludes: [\"/.git/\"]\nenv_vars:\n  TORCH_NCCL_AVOID_RECORD_STREAMS: \"1\"\n  CUDA_DEVICE_MAX_CONNECTIONS: \"1\"\n  HCCL_HOST_SOCKET_PORT_RANGE: \"auto\"\n  HCCL_NPU_SOCKET_PORT_RANGE: \"auto\""
  },
  {
    "path": "verl/trainer/sft_trainer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport os\nfrom functools import partial\n\nfrom tensordict.tensorclass import NonTensorData\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n\nimport logging\n\nimport hydra\nimport torch\nimport torch.distributed\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import DistributedSampler\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom tqdm import tqdm\n\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.checkpoint import CheckpointHandler\nfrom verl.utils.dataset.dataset_utils import SFTTensorCollator\nfrom verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset\nfrom verl.utils.device import auto_set_device, get_device_name\nfrom verl.utils.distributed import destroy_global_process_group\nfrom verl.utils.logger import log_with_rank\nfrom verl.utils.memory_utils import aggressive_empty_cache\nfrom verl.utils.profiler import log_gpu_memory_usage\nfrom verl.utils.tracking import Tracking\nfrom verl.workers.engine_workers import TrainingWorker\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_SFT_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass SFTTrainer:\n    def __init__(\n        self,\n        config,\n    ):\n        self.config = config\n\n        log_gpu_memory_usage(f\"rank {torch.distributed.get_rank()}: Before SFTTrainer init\", logger=logger)\n\n        self.rank = torch.distributed.get_rank()\n\n        self._build_config()\n        self._build_dataset()\n\n        self._build_engine()\n\n        self._build_dataloader()\n\n        self._init_engine()\n\n        self._build_ckpt_handler()\n\n        # Initialize resume-related variables\n        self.resume_global_step = self.ckpt_handler.load_checkpoint()\n\n        self.device_name = self.config.trainer.device\n\n        if self.rank == 0:\n            print(self.config)\n\n        log_gpu_memory_usage(f\"rank {self.rank}: After SFTTrainer init\", logger=logger)\n\n    def _build_ckpt_handler(self):\n        resume_mode = getattr(self.config.trainer, \"resume_mode\", \"auto\")\n        resume_from_path = getattr(self.config.trainer, \"resume_from_path\", None)\n        max_ckpt_to_keep = getattr(self.config.trainer, \"max_ckpt_to_keep\", None)\n        default_hdfs_dir = getattr(self.config.trainer, \"default_hdfs_dir\", None)\n        lora_train_meta = self._get_lora_train_meta()\n\n        self.ckpt_handler = CheckpointHandler(\n            engine=self.engine,\n            train_dataloader=self.train_dataloader,\n            default_local_dir=self.config.trainer.default_local_dir,\n            max_ckpt_to_keep=max_ckpt_to_keep,\n            default_hdfs_dir=default_hdfs_dir,\n            resume_mode=resume_mode,\n            resume_from_path=resume_from_path,\n            lora_train_meta=lora_train_meta,\n        )\n\n    def _get_lora_train_meta(self):\n        lora_adapter_path = self.config.model.get(\"lora_adapter_path\", None)\n        lora_rank = int(getattr(self.config.model, \"lora_rank\", 0) or 0)\n\n        if lora_adapter_path is None and lora_rank <= 0:\n            return None\n\n        raw_lora_alpha = self.config.model.get(\"lora_alpha\", None)\n        if raw_lora_alpha is None:\n            log_with_rank(\n                \"LoRA is enabled but `model.lora_alpha` is not set; fallback to 0 in checkpoint metadata.\",\n                logger=logger,\n                rank=self.rank,\n                level=logging.WARNING,\n                log_only_rank_0=True,\n            )\n            lora_alpha = 0\n        else:\n            lora_alpha = int(raw_lora_alpha)\n            if lora_alpha == 0:\n                log_with_rank(\n                    \"LoRA is enabled but `model.lora_alpha` is 0; this may lead to ineffective LoRA scaling.\",\n                    logger=logger,\n                    rank=self.rank,\n                    level=logging.WARNING,\n                    log_only_rank_0=True,\n                )\n\n        task_type = self.config.model.get(\"task_type\", None)\n        if task_type is None:\n            task_type = \"CAUSAL_LM\"\n\n        return {\n            \"r\": lora_rank,\n            \"lora_alpha\": int(lora_alpha or 0),\n            \"task_type\": str(task_type),\n        }\n\n    def _build_config(self):\n        from verl.utils.config import omega_conf_to_dataclass\n\n        self.model_config = omega_conf_to_dataclass(self.config.model)\n        self.engine_config = omega_conf_to_dataclass(self.config.engine)\n        self.optimizer_config = omega_conf_to_dataclass(self.config.optim)\n        self.checkpoint_config = omega_conf_to_dataclass(self.config.checkpoint)\n        self.profiler_config = omega_conf_to_dataclass(self.config.profiler)\n\n        # check profile interval\n        self.profiler_interval = self.config.trainer.profile_interval\n        self._validate_profiler_interval()\n\n    def _validate_profiler_interval(self):\n        assert len(self.profiler_interval) == 2\n        self.start_profile_step = self.profiler_interval[0]\n        self.end_profile_step = self.profiler_interval[1]\n        assert self.end_profile_step >= self.start_profile_step\n        if self.start_profile_step < 0:\n            assert self.end_profile_step < 0\n\n    def _build_engine(self):\n        from verl.workers.engine_workers import TrainingWorkerConfig\n        from verl.workers.utils.losses import sft_loss\n\n        self.loss_fn = partial(sft_loss, config=None)\n\n        config = TrainingWorkerConfig(\n            model_type=\"language_model\",\n            model_config=self.model_config,\n            engine_config=self.engine_config,\n            optimizer_config=self.optimizer_config,\n            checkpoint_config=self.checkpoint_config,\n            profiler_config=self.profiler_config,\n        )\n\n        self.training_client = TrainingWorker(config=config)\n        self.training_client.set_loss_fn(loss_fn=self.loss_fn)\n        # Note that in SPMD world, this abstraction has to break\n        self.engine = self.training_client.engine\n\n    def _init_engine(self):\n        # patch optimizer config\n        if self.config.trainer.total_training_steps is not None:\n            self.total_training_steps = self.config.trainer.total_training_steps\n        else:\n            self.total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs\n        self.optimizer_config.total_training_steps = self.total_training_steps\n\n        self.steps_per_epoch = len(self.train_dataloader)\n\n        # manage save and test frequency\n        self.save_freq = self.config.trainer.save_freq\n        if self.save_freq == \"after_each_epoch\":\n            self.save_freq = self.steps_per_epoch\n\n        self.test_freq = self.config.trainer.test_freq\n        if self.test_freq == \"after_each_epoch\":\n            self.test_freq = self.steps_per_epoch\n\n        self.training_client.reset()\n\n    def _build_dataset(self):\n        config = self.config\n        tokenizer = self.model_config.tokenizer\n        processor = self.model_config.processor\n        train_dataset = create_sft_dataset(\n            config.data.train_files,\n            config.data,\n            tokenizer,\n            processor,\n            max_samples=config.data.get(\"train_max_samples\", -1),\n        )\n        if config.data.val_files:\n            val_dataset = create_sft_dataset(\n                config.data.val_files,\n                config.data,\n                tokenizer,\n                processor,\n                max_samples=config.data.get(\"val_max_samples\", -1),\n            )\n        else:\n            val_dataset = None\n\n        self.train_dataset, self.val_dataset = train_dataset, val_dataset\n\n    def _build_dataloader(self):\n        # build dataset\n        config = self.config\n        # build dataloader\n        # Use data parallel rank and size instead of global rank and world size\n\n        # Set pin_memory_device when pin_memory is enabled.\n        device_name = get_device_name()\n\n        dp_rank = self.engine.get_data_parallel_rank()\n        dp_size = self.engine.get_data_parallel_size()\n\n        self.train_sampler = DistributedSampler(\n            self.train_dataset, shuffle=True, num_replicas=dp_size, rank=dp_rank, drop_last=True\n        )\n\n        self.global_batch_size = config.data.train_batch_size\n        self.train_batch_size_per_dp = self.global_batch_size // dp_size\n        self.collate_fn = SFTTensorCollator(config.data.pad_mode)\n\n        self.train_dataloader = StatefulDataLoader(\n            dataset=self.train_dataset,\n            batch_size=self.train_batch_size_per_dp,\n            sampler=self.train_sampler,\n            collate_fn=self.collate_fn,\n            num_workers=self.config.data.num_workers,\n            pin_memory=False,\n            drop_last=True,\n            pin_memory_device=device_name,\n        )\n\n        if self.val_dataset:\n            self.val_sampler = DistributedSampler(\n                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=True\n            )\n            self.val_dataloader = StatefulDataLoader(\n                dataset=self.val_dataset,\n                batch_size=self.train_batch_size_per_dp,\n                sampler=self.val_sampler,\n                collate_fn=self.collate_fn,\n                num_workers=self.config.data.num_workers,\n                pin_memory=False,\n                drop_last=True,\n                pin_memory_device=device_name,\n            )\n        else:\n            self.val_dataloader = None\n\n    def _get_batch_seqlens(self, data):\n        # mean over dp group\n        is_nested = data[\"input_ids\"].is_nested\n        if is_nested:\n            batch_seqlens: torch.Tensor = data[\"input_ids\"].offsets().diff()\n        else:\n            batch_seqlens: torch.Tensor = data[\"attention_mask\"].sum(dim=-1)\n        batch_seqlens = batch_seqlens.to(self.device_name)  # (global_bsz // dp)\n\n        dp_group = self.engine.get_data_parallel_group()\n        dp_size = self.engine.get_data_parallel_size()\n\n        if dp_size == 1 or dp_group is None:\n            return batch_seqlens.tolist()\n\n        output_tensor = torch.empty(\n            (batch_seqlens.shape[0] * dp_size,),\n            dtype=batch_seqlens.dtype,\n            device=self.device_name,\n        )  # (global_bsz,)\n\n        torch.distributed.all_gather_into_tensor(\n            output_tensor=output_tensor,\n            input_tensor=batch_seqlens,\n            group=dp_group,\n        )\n\n        batch_seqlens = output_tensor.tolist()\n        return batch_seqlens\n\n    def fit(self):\n        is_logging = self.engine.is_mp_src_rank_with_outputs() and self.engine.get_data_parallel_rank() == 0\n\n        # TODO: add a unified tracking\n        if is_logging:\n            tracking = Tracking(\n                project_name=self.config.trainer.project_name,\n                experiment_name=self.config.trainer.experiment_name,\n                default_backend=self.config.trainer.logger,\n                config=OmegaConf.to_container(self.config, resolve=True),\n            )\n\n        global_step = self.resume_global_step  # Start from resumed step\n        last_valid_metric = None\n\n        log_with_rank(\n            f\"Total training steps: {self.total_training_steps},\",\n            logger=logger,\n            rank=0,\n            log_only_rank_0=True,\n        )\n\n        # With StatefulDataLoader, we don't need to manually calculate epochs and steps\n        # The dataloader will automatically resume from where it left off\n        if global_step > 0:\n            log_with_rank(\n                f\"StatefulDataLoader will automatically resume from global step: {global_step}\",\n                logger=logger,\n                rank=0,\n                log_only_rank_0=True,\n            )\n\n        # Calculate which epoch we're starting from for sampler.set_epoch()\n        start_epoch = global_step // self.steps_per_epoch\n\n        meta_info = {\n            \"use_remove_padding\": self.config.model.use_remove_padding,\n            \"use_dynamic_bsz\": self.config.data.use_dynamic_bsz,\n            \"max_token_len_per_gpu\": self.config.data.max_token_len_per_gpu,\n            \"micro_batch_size_per_gpu\": self.config.data.micro_batch_size_per_gpu,\n            \"temperature\": 1.0,\n            \"global_batch_size\": self.global_batch_size,\n            \"pad_mode\": self.config.data.pad_mode,\n            \"pad_token_id\": self.model_config.tokenizer.pad_token_id,\n        }\n\n        train_time = 0\n        total_tokens = 0\n        for epoch in range(start_epoch, self.config.trainer.total_epochs):\n            self.train_sampler.set_epoch(epoch=epoch)\n\n            aggressive_empty_cache(force_sync=True)\n            log_gpu_memory_usage(f\"rank {self.rank}: At start of epoch {epoch}\", logger=logger)\n\n            for step_in_epoch, data in enumerate(\n                tqdm(\n                    self.train_dataloader,\n                    initial=global_step % self.steps_per_epoch if epoch == start_epoch else 0,\n                    total=self.steps_per_epoch,\n                    desc=f\"Epoch {epoch + 1}/{self.config.trainer.total_epochs}\",\n                    disable=not is_logging,\n                )\n            ):\n                global_step += 1\n\n                # construct tensordict\n                data = tu.get_tensordict(tensor_dict=data, non_tensor_dict=meta_info)\n                batch_seqlens = self._get_batch_seqlens(data=data)\n                # this is necessary. Otherwise, it is interpreted as NonTensorStack\n                batch_seqlens_ntd = NonTensorData(batch_seqlens)\n\n                tu.assign_non_tensor(data, update_lr_scheduler=True, global_token_num=batch_seqlens_ntd)\n\n                # start profile in SPMD mode\n                if global_step == self.start_profile_step:\n                    self.training_client.start_profile()\n                # train for on batch\n                output = self.training_client.train_batch(data=data)\n\n                if global_step == self.end_profile_step:\n                    self.training_client.stop_profile()\n\n                if self.engine.is_mp_src_rank_with_outputs():\n                    metrics = tu.get(output, \"metrics\")\n\n                    # TODO: we can actual accumulate metrics for N steps and perform aggregate metrics\n                    for k in [\"loss\", \"grad_norm\", \"lr\", \"mfu\"]:\n                        if k in metrics.keys():\n                            value = metrics.pop(k)\n                            metrics[f\"train/{k}\"] = value\n\n                    metrics[\"train/global_tokens\"] = torch.sum(\n                        torch.tensor(batch_seqlens, device=self.device_name)\n                    ).item()\n                    total_tokens += metrics[\"train/global_tokens\"]\n                    metrics[\"train/total_tokens(B)\"] = total_tokens / 1e9\n\n                    if self.engine.get_data_parallel_rank() == 0:\n                        tracking.log(data=metrics, step=global_step)\n\n                is_last_step = global_step >= self.total_training_steps\n                is_valid_step = global_step % self.test_freq == 0\n                is_save_step = global_step % self.save_freq == 0\n\n                # early exit or validation step\n                if is_last_step and self.val_dataloader is not None or (self.test_freq > 0 and is_valid_step):\n                    # Perform validation\n                    val_losses = []\n                    for val_data in self.val_dataloader:\n                        val_data = tu.get_tensordict(tensor_dict=val_data, non_tensor_dict=meta_info)\n                        output = self.training_client.infer_batch(val_data)\n\n                        if self.engine.is_mp_src_rank_with_outputs():\n                            metrics = tu.get(output, \"metrics\")\n                            val_losses.append(metrics[\"loss\"])\n\n                    if self.engine.is_mp_src_rank_with_outputs():\n                        val_loss = torch.mean(torch.tensor(val_losses, device=self.device_name))\n                        # average over data parallel group\n                        dp_group = self.engine.get_data_parallel_group()\n                        if dp_group is not None:\n                            torch.distributed.all_reduce(val_loss, op=torch.distributed.ReduceOp.AVG, group=dp_group)\n\n                    if is_logging:\n                        metric = {\"val/loss\": val_loss.detach().item()}\n                        tracking.log(data=metric, step=global_step)\n                        last_valid_metric = metric\n                    torch.distributed.barrier()\n\n                if is_last_step or (self.save_freq > 0 and is_save_step):\n                    aggressive_empty_cache(force_sync=True)\n                    self.ckpt_handler.save_checkpoint(step=global_step)\n\n                if is_last_step:\n                    if is_logging:\n                        print(f\"Total time for train steps: {train_time:.2f}s\")\n                        print(f\"Final validation metrics: {last_valid_metric}\")\n                    return\n\n\ndef run_sft(config):\n    from verl.utils.distributed import initialize_global_process_group\n\n    initialize_global_process_group()\n    trainer = SFTTrainer(config=config)\n    trainer.fit()\n    destroy_global_process_group()\n\n\n@hydra.main(config_path=\"config\", config_name=\"sft_trainer_engine\", version_base=None)\ndef main(config):\n    # Automatically set `config.trainer.device = npu` when running on Ascend NPU.\n    auto_set_device(config)\n    run_sft(config)\n\n\ndef create_sft_dataset(data_paths, data_config, tokenizer, processor, max_samples=-1):\n    \"\"\"Create a dataset.\"\"\"\n    # build dataset\n    # First check if a custom dataset class is specified\n    if data_config.custom_cls.get(\"path\", None):\n        from verl.utils.import_utils import load_extern_object\n\n        dataset_cls = load_extern_object(data_config.custom_cls.path, data_config.custom_cls.name)\n    else:\n        # Default to multi-turn dataset\n        dataset_cls = MultiTurnSFTDataset\n\n    # Create datasets based on the selected class\n    dataset = dataset_cls(\n        parquet_files=data_paths, tokenizer=tokenizer, config=data_config, processor=processor, max_samples=max_samples\n    )\n    return dataset\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/trainer/sft_trainer_ray.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport os\nfrom functools import partial\n\nfrom tensordict.tensorclass import NonTensorData\n\nos.environ[\"NCCL_DEBUG\"] = \"WARN\"\nos.environ[\"TOKENIZERS_PARALLELISM\"] = \"true\"\n\nimport logging\n\nimport hydra\nimport ray\nimport torch\nimport torch.distributed\nfrom omegaconf import OmegaConf\nfrom torch.utils.data import DistributedSampler\nfrom torchdata.stateful_dataloader import StatefulDataLoader\nfrom tqdm import tqdm\n\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.checkpoint import CheckpointHandler, OrchestrationMode\nfrom verl.utils.dataset.dataset_utils import SFTTensorCollator\nfrom verl.utils.dataset.multiturn_sft_dataset import MultiTurnSFTDataset\nfrom verl.utils.device import auto_set_device, get_device_name\nfrom verl.utils.logger import log_with_rank\nfrom verl.utils.tracking import Tracking\nfrom verl.workers.engine_workers import TrainingWorker\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_SFT_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass SFTTrainer:\n    def __init__(\n        self,\n        config,\n    ):\n        self.config = config\n\n        self._build_config()\n        self._build_dataset()\n        self._build_dataloader()\n\n        self._build_engine()\n        self._build_ckpt_handler()\n\n        # Initialize resume-related variables\n        self.resume_global_step = self.ckpt_handler.load_checkpoint()\n\n        self.device_name = self.config.trainer.device\n\n        print(self.config)\n\n    def _build_ckpt_handler(self):\n        resume_mode = getattr(self.config.trainer, \"resume_mode\", \"auto\")\n        resume_from_path = getattr(self.config.trainer, \"resume_from_path\", None)\n        max_ckpt_to_keep = getattr(self.config.trainer, \"max_ckpt_to_keep\", None)\n        default_hdfs_dir = getattr(self.config.trainer, \"default_hdfs_dir\", None)\n\n        self.ckpt_handler = CheckpointHandler(\n            engine=self.training_client,\n            train_dataloader=self.train_dataloader,\n            default_local_dir=self.config.trainer.default_local_dir,\n            max_ckpt_to_keep=max_ckpt_to_keep,\n            default_hdfs_dir=default_hdfs_dir,\n            resume_mode=resume_mode,\n            resume_from_path=resume_from_path,\n            mode=OrchestrationMode.RAY,\n        )\n\n    def _build_config(self):\n        from verl.utils.config import omega_conf_to_dataclass\n\n        self.model_config = omega_conf_to_dataclass(self.config.model)\n        self.engine_config = omega_conf_to_dataclass(self.config.engine)\n        self.optimizer_config = omega_conf_to_dataclass(self.config.optim)\n        self.checkpoint_config = omega_conf_to_dataclass(self.config.checkpoint)\n        self.profiler_config = omega_conf_to_dataclass(self.config.profiler)\n\n        # check profile interval\n        self.profiler_interval = self.config.trainer.profile_interval\n        self._validate_profiler_interval()\n\n    def _validate_profiler_interval(self):\n        assert len(self.profiler_interval) == 2\n        self.start_profile_step = self.profiler_interval[0]\n        self.end_profile_step = self.profiler_interval[1]\n        assert self.end_profile_step >= self.start_profile_step\n        if self.start_profile_step < 0:\n            assert self.end_profile_step < 0\n\n    def _build_engine(self):\n        from verl.workers.engine_workers import TrainingWorkerConfig\n        from verl.workers.utils.losses import sft_loss\n\n        self.loss_fn = partial(sft_loss, config=None)\n\n        config = TrainingWorkerConfig(\n            model_type=\"language_model\",\n            model_config=self.model_config,\n            engine_config=self.engine_config,\n            optimizer_config=self.optimizer_config,\n            checkpoint_config=self.checkpoint_config,\n            profiler_config=self.profiler_config,\n        )\n\n        wg_kwargs = {}\n        if self.start_profile_step != -1:\n            wg_kwargs[\"profile_steps\"] = list(range(self.start_profile_step, self.end_profile_step + 1))\n            # Only require nsight worker options when tool is nsys\n            if OmegaConf.select(self.config.profiler, \"tool\") == \"nsys\":\n                wg_kwargs[\"worker_nsight_options\"] = OmegaConf.to_container(\n                    OmegaConf.select(self.config.global_profiler.global_tool_config.nsys, \"worker_nsight_options\")\n                )\n\n        # create resource pool and worker group\n        from verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup\n\n        n_gpus_per_node = self.config.trainer.n_gpus_per_node\n        nnodes = self.config.trainer.nnodes\n        self.resource_pool = RayResourcePool(process_on_nodes=[n_gpus_per_node] * nnodes)\n        ray_cls_with_init = RayClassWithInitArgs(ray.remote(TrainingWorker), config=config)\n        self.training_client = RayWorkerGroup(\n            resource_pool=self.resource_pool,\n            ray_cls_with_init=ray_cls_with_init,\n            device_name=self.config.trainer.device,\n            **wg_kwargs,\n        )\n        self.training_client.set_loss_fn(loss_fn=self.loss_fn)\n        self.training_client.reset()\n\n    def _build_dataset(self):\n        config = self.config\n        tokenizer = self.model_config.tokenizer\n        processor = self.model_config.processor\n        train_dataset = create_sft_dataset(\n            config.data.train_files,\n            config.data,\n            tokenizer,\n            processor=processor,\n            max_samples=config.data.get(\"train_max_samples\", -1),\n        )\n        if config.data.val_files:\n            val_dataset = create_sft_dataset(\n                config.data.val_files,\n                config.data,\n                tokenizer,\n                processor=processor,\n                max_samples=config.data.get(\"val_max_samples\", -1),\n            )\n        else:\n            val_dataset = None\n\n        self.train_dataset, self.val_dataset = train_dataset, val_dataset\n\n    def _build_dataloader(self):\n        # build dataset\n        config = self.config\n        # build dataloader\n        # Use data parallel rank and size instead of global rank and world size\n\n        # Set pin_memory_device when pin_memory is enabled.\n        device_name = get_device_name()\n\n        dp_rank = 0\n        dp_size = 1\n\n        self.train_sampler = DistributedSampler(\n            self.train_dataset, shuffle=True, num_replicas=dp_size, rank=dp_rank, drop_last=True\n        )\n\n        self.global_batch_size = config.data.train_batch_size\n        self.train_batch_size_per_dp = self.global_batch_size // dp_size\n        self.collate_fn = SFTTensorCollator(config.data.pad_mode)\n\n        self.train_dataloader = StatefulDataLoader(\n            dataset=self.train_dataset,\n            batch_size=self.train_batch_size_per_dp,\n            sampler=self.train_sampler,\n            collate_fn=self.collate_fn,\n            num_workers=8,\n            pin_memory=False,\n            drop_last=True,\n            pin_memory_device=device_name,\n        )\n\n        if self.val_dataset:\n            self.val_sampler = DistributedSampler(\n                self.val_dataset, shuffle=False, num_replicas=dp_size, rank=dp_rank, drop_last=True\n            )\n            self.val_dataloader = StatefulDataLoader(\n                dataset=self.val_dataset,\n                batch_size=self.train_batch_size_per_dp,\n                sampler=self.val_sampler,\n                collate_fn=self.collate_fn,\n                num_workers=8,\n                pin_memory=False,\n                drop_last=True,\n                pin_memory_device=device_name,\n            )\n        else:\n            self.val_dataloader = None\n\n        # update\n        if self.config.trainer.total_training_steps is not None:\n            self.total_training_steps = self.config.trainer.total_training_steps\n        else:\n            self.total_training_steps = len(self.train_dataloader) * self.config.trainer.total_epochs\n        self.optimizer_config.total_training_steps = self.total_training_steps\n\n        self.steps_per_epoch = len(self.train_dataloader)\n\n        # manage save and test frequency\n        self.save_freq = self.config.trainer.save_freq\n        if self.save_freq == \"after_each_epoch\":\n            self.save_freq = self.steps_per_epoch\n\n        self.test_freq = self.config.trainer.test_freq\n        if self.test_freq == \"after_each_epoch\":\n            self.test_freq = self.steps_per_epoch\n\n    def _get_batch_seqlens(self, data):\n        # mean over dp group\n        is_nested = data[\"input_ids\"].is_nested\n        if is_nested:\n            batch_seqlens: torch.Tensor = data[\"input_ids\"].offsets().diff()\n        else:\n            batch_seqlens: torch.Tensor = data[\"attention_mask\"].sum(dim=-1)\n        return batch_seqlens\n\n    def fit(self):\n        tracking = Tracking(\n            project_name=self.config.trainer.project_name,\n            experiment_name=self.config.trainer.experiment_name,\n            default_backend=self.config.trainer.logger,\n            config=OmegaConf.to_container(self.config, resolve=True),\n        )\n\n        global_step = self.resume_global_step  # Start from resumed step\n        last_valid_metric = None\n\n        log_with_rank(\n            f\"Total training steps: {self.total_training_steps},\",\n            logger=logger,\n            rank=0,\n            log_only_rank_0=True,\n        )\n\n        # With StatefulDataLoader, we don't need to manually calculate epochs and steps\n        # The dataloader will automatically resume from where it left off\n        if global_step > 0:\n            log_with_rank(\n                f\"StatefulDataLoader will automatically resume from global step: {global_step}\",\n                logger=logger,\n                rank=0,\n                log_only_rank_0=True,\n            )\n\n        # Calculate which epoch we're starting from for sampler.set_epoch()\n        start_epoch = global_step // self.steps_per_epoch\n\n        meta_info = {\n            \"use_remove_padding\": self.config.model.use_remove_padding,\n            \"use_dynamic_bsz\": self.config.data.use_dynamic_bsz,\n            \"max_token_len_per_gpu\": self.config.data.max_token_len_per_gpu,\n            \"micro_batch_size_per_gpu\": self.config.data.micro_batch_size_per_gpu,\n            \"temperature\": 1.0,\n            \"global_batch_size\": self.global_batch_size,\n            \"pad_mode\": self.config.data.pad_mode,\n            \"pad_token_id\": self.model_config.tokenizer.pad_token_id,\n        }\n\n        train_time = 0\n        total_tokens = 0\n        for epoch in range(start_epoch, self.config.trainer.total_epochs):\n            self.train_sampler.set_epoch(epoch=epoch)\n\n            for step_in_epoch, data in enumerate(\n                tqdm(\n                    self.train_dataloader,\n                    initial=global_step % self.steps_per_epoch if epoch == start_epoch else 0,\n                    total=self.steps_per_epoch,\n                    desc=f\"Epoch {epoch + 1}/{self.config.trainer.total_epochs}\",\n                )\n            ):\n                global_step += 1\n                # construct tensordict\n                data = tu.get_tensordict(tensor_dict=data, non_tensor_dict=meta_info)\n                batch_seqlens = self._get_batch_seqlens(data=data).tolist()\n                # this is necessary. Otherwise, it is interpreted as NonTensorStack\n                batch_seqlens_ntd = NonTensorData(batch_seqlens)\n\n                tu.assign_non_tensor(data, update_lr_scheduler=True, global_token_num=batch_seqlens_ntd)\n\n                # start profile in SPMD mode\n                if global_step == self.start_profile_step:\n                    self.training_client.start_profile()\n\n                # train for on batch\n                output = self.training_client.train_batch(data)\n                output = output.get()\n\n                if global_step == self.end_profile_step:\n                    self.training_client.stop_profile()\n\n                metrics = tu.get(output, \"metrics\")\n\n                # TODO: we can actual accumulate metrics for N steps and perform aggregate metrics\n                metrics[\"train/loss\"] = metrics.pop(\"loss\")\n                metrics[\"train/grad_norm\"] = metrics.pop(\"grad_norm\")\n                metrics[\"train/lr\"] = metrics.pop(\"lr\")\n                metrics[\"train/mfu\"] = metrics.pop(\"mfu\")\n                metrics[\"train/global_tokens\"] = torch.sum(torch.tensor(batch_seqlens, device=self.device_name)).item()\n                total_tokens += metrics[\"train/global_tokens\"]\n                metrics[\"train/total_tokens(B)\"] = total_tokens / 1e9\n                tracking.log(data=metrics, step=global_step)\n\n                is_last_step = global_step >= self.total_training_steps\n                is_valid_step = global_step % self.test_freq == 0\n                is_save_step = global_step % self.save_freq == 0\n\n                # early exit or validation step\n                if is_last_step and self.val_dataloader is not None or (self.test_freq > 0 and is_valid_step):\n                    # Perform validation\n                    val_losses = []\n                    for val_data in self.val_dataloader:\n                        val_data = tu.get_tensordict(tensor_dict=val_data, non_tensor_dict=meta_info)\n                        output = self.training_client.infer_batch(val_data)\n                        output = output.get()\n                        metrics = tu.get(output, \"metrics\")\n                        val_losses.append(metrics[\"loss\"])\n\n                    val_loss = torch.mean(torch.tensor(val_losses, device=self.device_name))\n\n                    metric = {\"val/loss\": val_loss.detach().item()}\n                    tracking.log(data=metric, step=global_step)\n                    last_valid_metric = metric\n\n                if is_last_step or (self.save_freq > 0 and is_save_step):\n                    self.ckpt_handler.save_checkpoint(step=global_step)\n\n                if is_last_step:\n                    print(f\"Total time for train steps: {train_time:.2f}s\")\n                    print(f\"Final validation metrics: {last_valid_metric}\")\n                    return\n\n\ndef run_sft(config):\n    ray.init()\n    trainer = SFTTrainer(config=config)\n    trainer.fit()\n\n\n@hydra.main(config_path=\"config\", config_name=\"sft_trainer_engine\", version_base=None)\ndef main(config):\n    # Automatically set `config.trainer.device = npu` when running on Ascend NPU.\n    auto_set_device(config)\n    run_sft(config)\n\n\ndef create_sft_dataset(data_paths, data_config, tokenizer, processor, max_samples=-1):\n    \"\"\"Create a dataset.\"\"\"\n    # build dataset\n    # First check if a custom dataset class is specified\n    if data_config.custom_cls.get(\"path\", None):\n        from verl.utils.import_utils import load_extern_type\n\n        dataset_cls = load_extern_type(data_config.custom_cls.path, data_config.custom_cls.name)\n    else:\n        # Default to multi-turn dataset\n        dataset_cls = MultiTurnSFTDataset\n\n    # Create datasets based on the selected class\n    dataset = dataset_cls(\n        parquet_files=data_paths, tokenizer=tokenizer, config=data_config, processor=processor, max_samples=max_samples\n    )\n    return dataset\n\n\nif __name__ == \"__main__\":\n    main()\n"
  },
  {
    "path": "verl/utils/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom . import config, tokenizer\nfrom .config import omega_conf_to_dataclass, validate_config\nfrom .groupwise import as_torch_index, group_mean_std\nfrom .tokenizer import hf_processor, hf_tokenizer, normalize_token_ids\n\n__all__ = (\n    tokenizer.__all__\n    + config.__all__\n    + [\"hf_processor\", \"hf_tokenizer\", \"normalize_token_ids\", \"omega_conf_to_dataclass\", \"validate_config\"]\n    + [\"as_torch_index\", \"group_mean_std\"]\n)\n"
  },
  {
    "path": "verl/utils/activation_offload.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2022-2025, NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Functionality for CPU offloading of tensors saved for backward pass.\"\"\"\n\nfrom __future__ import annotations\n\nimport functools\nimport logging\nimport os\nfrom typing import Any, Optional\n\nimport torch\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n\nfrom verl.utils.device import get_torch_device\nfrom verl.utils.fsdp_utils import FSDPModule as FSDP2\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef _get_unique_tensor_key(tensor):\n    key = (tensor.untyped_storage().data_ptr() + tensor.storage_offset(), tensor.dtype)\n    return key\n\n\nclass FSDPParameterFilter:\n    def __init__(self):\n        self.model_parameters_storage = set()\n\n    def __call__(self, tensor):\n        return tensor.untyped_storage().data_ptr() not in self.model_parameters_storage\n\n    def update_model_parameters(self, model):\n        new_storage = set()\n        for p in model.parameters():\n            new_storage.add(p.data.untyped_storage().data_ptr())\n        self.model_parameters_storage = new_storage\n\n\nclass CpuOffloadHookWithOffloadHandler:\n    \"\"\"Context-manager that offloads/recovers tensors through an offload hander.\n\n    The hook just offloads/recovers the tensor object to the handler through `tensor_push`\n    and `tensor_pop` interface. How the offload-handler manages the offloading, recovering\n    or prefetching timing is transparent to this hook.\n    \"\"\"\n\n    def __init__(\n        self,\n        offload_handler: OffloadHandler,\n        handler_extra_kwargs: Optional[dict[str, Any]] = None,\n    ) -> None:\n        if handler_extra_kwargs is None:\n            handler_extra_kwargs = {}\n        self.offload_handler: OffloadHandler = offload_handler\n        self.handler_extra_kwargs: dict[str, Any] = handler_extra_kwargs\n        self.inside_context = False\n\n    def __enter__(self):\n        self.inside_context = True\n        torch._C._autograd._push_saved_tensors_default_hooks(self.on_save_for_backward, self.on_get_saved_tensor)\n\n    def __exit__(self, *args: Any):\n        self.inside_context = False\n        torch._C._autograd._pop_saved_tensors_default_hooks()\n\n    def on_save_for_backward(self, tensor: torch.Tensor) -> Any:\n        retrieve_identifier = self.offload_handler.tensor_push(tensor, **self.handler_extra_kwargs)\n        return retrieve_identifier\n\n    def on_get_saved_tensor(self, saved_state: Any) -> torch.Tensor:\n        tensor = self.offload_handler.tensor_pop(saved_state, **self.handler_extra_kwargs)\n        return tensor\n\n\nclass OffloadHandler:\n    \"\"\"A base class for CPU offload-handler.\"\"\"\n\n    def __init__(self) -> None:\n        pass\n\n    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:\n        \"\"\"Tensor push.\"\"\"\n        raise NotImplementedError(\n            \"`tensor_push is not implented in OffloadHandler class. Inherit this class and implement your \"\n            \"custom tensor_push.\"\n        )\n\n    def tensor_pop(self, tensor_tag: Any, **kwargs):\n        \"\"\"Tensor pop.\"\"\"\n        raise NotImplementedError(\n            \"`tensor_pop is not implented in OffloadHandler class. Inherit this class and implement your \"\n            \"custom tensor_pop.\"\n        )\n\n\nclass GroupCommitFunction(torch.autograd.Function):\n    \"\"\"this is a dummy op with output identical to input.\n    However, it is necessary for marking a timepoint for offload handler to\n    accomplish all synchronizations. Implementing it as a function is necessary\n    because we need to actions in both forward and backward.\n    \"\"\"\n\n    @staticmethod\n    def forward(ctx, tensor, cpu_offload_handler):\n        # pylint: disable=missing-function-docstring\n        cpu_offload_handler.on_group_commit_forward()\n        ctx.cpu_offload_handler = cpu_offload_handler\n        # return the identical tensor\n        return tensor\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        # pylint: disable=missing-function-docstring\n        cpu_offload_handler = ctx.cpu_offload_handler\n        cpu_offload_handler.on_group_commit_backward()\n        return grad_output, None\n\n\ngroup_prefetch_offload_commit = GroupCommitFunction.apply\n\n\nclass SynchronizedGroupOffloadHandler(OffloadHandler):\n    \"\"\"Offload Handler that offloads/reloads in a synchronized way.\n    The device-to-host and host-to-device copying happen in the same stream\n    as the computation kernels, thus the copying will block computation.\n    \"\"\"\n\n    def __init__(self, num_offload_group, tensor_need_offloading_checker=(lambda _: True)) -> None:\n        super().__init__()\n\n        self.num_offload_group = num_offload_group\n        self.tensor_need_offloading_checker = tensor_need_offloading_checker\n\n        self.groupid_reset()\n\n    def groupid_reset(self):\n        \"\"\"Groupid reset.\"\"\"\n        # Data structures to label saved tensors and book-keep their cpu copies.\n        # Currently, on push, create a new cpu tensor and copies; on pop, copies\n        # the tensor back to gpu and deletes the cpu tensor.\n        # These will increment whenever `group_commit()` is invoked\n        self.current_group, self.tensor_count_current_group = (0, 0)\n        self.torch_tensor_count = 0\n        self.tensor_tag_to_state = {}\n\n    def on_group_commit_forward(self):\n        \"\"\"On group commit forward.\"\"\"\n        # finishing up with updating current group and tensor count\n        self.current_group += 1  # increment\n        self.tensor_count_current_group = 0  # reset\n\n    def on_group_commit_backward(self):\n        \"\"\"On group commit backward.\"\"\"\n        self.current_group -= 1\n        assert self.current_group >= 0\n\n    @staticmethod\n    def offload(src_tensor, pin_memory=True):\n        \"\"\"Offload.\"\"\"\n\n        cpu_backup = torch.empty(\n            src_tensor.size(),\n            dtype=src_tensor.dtype,\n            layout=src_tensor.layout,\n            device=\"cpu\",\n            pin_memory=pin_memory,\n        )\n        cpu_backup.copy_(src_tensor, non_blocking=True)\n        state = (src_tensor.device, cpu_backup)\n        return state\n\n    @staticmethod\n    def reload(state, non_blocking=None):\n        \"\"\"Reload.\"\"\"\n        dev, cpu_backup = state\n        if non_blocking is None:\n            non_blocking = cpu_backup.is_pinned()\n        return cpu_backup.to(dev, non_blocking=non_blocking)\n\n    def tensor_push(self, tensor: torch.Tensor, **kwargs):\n        \"\"\"Tensor push.\"\"\"\n        # obtain a unique tensor tag\n        tensor_tag = (self.current_group, self.tensor_count_current_group)\n        self.tensor_count_current_group += 1\n        assert tensor_tag not in self.tensor_tag_to_state\n        if self.current_group < self.num_offload_group and self.tensor_need_offloading_checker(tensor):\n            state = SynchronizedGroupOffloadHandler.offload(tensor)\n            self.tensor_tag_to_state[tensor_tag] = state\n        else:\n            # will be offloaded together after group commit\n            self.tensor_tag_to_state[tensor_tag] = tensor\n\n        return tensor_tag\n\n    def tensor_pop(self, tensor_tag, **kwargs):\n        \"\"\"Tensor pop.\"\"\"\n        assert tensor_tag in self.tensor_tag_to_state\n        state = self.tensor_tag_to_state.pop(tensor_tag)\n        if isinstance(state, tuple):\n            tensor = SynchronizedGroupOffloadHandler.reload(state)\n        else:\n            tensor = state\n        return tensor\n\n\nclass AsyncDoubleBufferGroupOffloadHandler(SynchronizedGroupOffloadHandler):\n    \"\"\"Compared to synchronize, this uses more memory because of the buffer but\n    achieves better performance due to the overlapping. D2h and h2d copying are\n    completely hidden behind computation if computation time of a layer is longer\n    than host-device communication time. Bulk offloading with delay and bulk reloading\n    with prefetch are implemented.\"\"\"\n\n    def __init__(\n        self,\n        num_offload_group,  # must be <= actual number of groups (number of commits)\n        num_model_group,\n        tensor_need_offloading_checker=(lambda t: True),\n    ) -> None:\n        super().__init__(\n            num_offload_group=num_offload_group,\n            tensor_need_offloading_checker=tensor_need_offloading_checker,\n        )\n        # Number of layers in the model\n        self.num_layers = num_model_group\n        # Data Structure to maintain reference to activation tensors\n        self.tensor_tag_to_buf = {}\n        # Tracking the number of layers offloaded\n        self.offloaded_group_count = 0\n        # Core data structure that decides the window for offloading\n        self.layer_window_map = {}\n        self.group_offload_mapping = {}\n\n        # Logic to make offloading load balance across computation\n        # for optimal CPU/GPU interconnect usage\n        constant = 0\n        for i in range(self.num_offload_group):\n            self.layer_window_map[i] = ((self.num_layers // self.num_offload_group) * (i + 1)) - 1\n            if i < (self.num_layers % self.num_offload_group):\n                self.layer_window_map[i] += i + 1\n                constant = i + 1\n            else:\n                self.layer_window_map[i] += constant\n\n        # allocate streams and events for synchronization\n        self.d2h_stream = get_torch_device().Stream()\n        self.h2d_stream = get_torch_device().Stream()\n\n    def tensor_push(self, tensor: torch.Tensor, **kwargs) -> Any:\n        torch_stray_tensor = isinstance(\n            tensor,\n            torch._subclasses.fake_tensor.FakeTensor | torch._subclasses.functional_tensor.FunctionalTensor,\n        )\n        need_offload = not torch_stray_tensor\n        need_offload = need_offload and self.tensor_need_offloading_checker(tensor)\n\n        if need_offload:\n            # obtain a unique tensor tag\n            tensor_tag = (self.current_group, self.tensor_count_current_group)\n            self.tensor_count_current_group += 1\n\n            assert tensor_tag not in self.tensor_tag_to_state\n            self.tensor_tag_to_state[tensor_tag] = tensor\n\n            if self.current_group < self.num_offload_group:\n                self.tensor_tag_to_buf[tensor_tag] = tensor\n        else:\n            tensor_tag = tensor\n        return tensor_tag\n\n    def tensor_pop(self, tensor_tag, **kwargs):\n        \"\"\"Tensor pop.\"\"\"\n        if isinstance(tensor_tag, torch.Tensor):\n            return tensor_tag\n        assert tensor_tag in self.tensor_tag_to_state\n        tensor = self.tensor_tag_to_state.pop(tensor_tag)\n        self.tensor_tag_to_buf.pop(tensor_tag, None)\n\n        # the tensor should have been copied back in on_group_commit_backward()\n        # which invokes bulk_reload_group.\n        assert not isinstance(tensor, tuple)\n        return tensor\n\n    def bulk_offload_group(self, group_to_offload):\n        \"\"\"Bulk offload group.\"\"\"\n        offload_mapping = {}\n        offload_size = 0\n        with get_torch_device().stream(self.d2h_stream):\n            for tensor_tag, state in self.tensor_tag_to_state.items():\n                group_id, _ = tensor_tag\n                if group_id == group_to_offload:\n                    assert not isinstance(state, tuple)\n                    key = _get_unique_tensor_key(state)\n                    if key not in offload_mapping:\n                        offload_mapping[key] = state\n                    # if offload, return the reference to cpu copy\n                    self.tensor_tag_to_state[tensor_tag] = (key, state.shape)\n            for key, tensor in offload_mapping.items():\n                state = SynchronizedGroupOffloadHandler.offload(tensor)\n                offload_size += tensor.numel() * tensor.element_size()\n                offload_mapping[key] = state\n\n            self.group_offload_mapping[group_to_offload] = offload_mapping\n\n    def synchronize_on_group_commit_forward(self, current_group):\n        \"\"\"Synchronize on group commit forward.\"\"\"\n\n        # For the first group, kickstart the offload after we have\n        # the first compute completion\n        if current_group == 0:\n            self.d2h_stream.wait_stream(get_torch_device().current_stream())\n            self.bulk_offload_group(current_group)\n\n        # Window map data structure helps us synchronize based on number\n        # of layers offloaded\n        if self.layer_window_map[self.offloaded_group_count] == current_group:\n            # Stream synchronization both ways\n            self.d2h_stream.wait_stream(get_torch_device().current_stream())\n            get_torch_device().current_stream().wait_stream(self.d2h_stream)\n\n            # Time to free the activation memory after usage\n            for tensor_tag, _ in self.tensor_tag_to_buf.items():\n                if tensor_tag[0] == self.offloaded_group_count:\n                    self.tensor_tag_to_buf[tensor_tag] = None\n\n            # Time to offload the next group\n            if self.offloaded_group_count < (self.num_offload_group - 1):\n                self.bulk_offload_group(self.offloaded_group_count + 1)\n\n            # Increment the offload group count to keep track\n            self.offloaded_group_count += 1\n\n    def on_group_commit_forward(self):\n        \"\"\"This function will cause host device synchronization\"\"\"\n        # handle synchronization events\n        self.synchronize_on_group_commit_forward(self.current_group)\n\n        super().on_group_commit_forward()\n\n    @torch.no_grad\n    def bulk_reload_group(self, group_to_reload):\n        \"\"\"Bulk reload group.\"\"\"\n        assert group_to_reload < self.num_offload_group\n\n        with get_torch_device().stream(self.h2d_stream):\n            # move back tensors\n            offload_mapping = self.group_offload_mapping.pop(group_to_reload)\n            assert offload_mapping is not None\n            for key, state in offload_mapping.items():\n                offload_mapping[key] = SynchronizedGroupOffloadHandler.reload(state)\n            for tensor_label, state in self.tensor_tag_to_state.items():\n                group_id, _ = tensor_label\n                if group_id == group_to_reload and not isinstance(state, torch.Tensor):\n                    assert isinstance(state, tuple), f\"{group_id} {state}\"\n                    key, shape = state\n                    recovered_tensor = offload_mapping[key].view(shape)\n                    self.tensor_tag_to_state[tensor_label] = recovered_tensor\n\n    def on_group_commit_backward(self):\n        # first decrement the current group.\n        # after last commit in forward, the group will +1; in backward it -1.\n        # Finally it should be decremented to 0.\n        self.current_group -= 1\n        assert self.current_group >= 0\n\n        # Layer window data structure helps us to reload at right times\n        if self.layer_window_map[self.offloaded_group_count - 1] == self.current_group:\n            # Stream synchronization both ways\n            self.h2d_stream.wait_stream(get_torch_device().current_stream())\n            get_torch_device().current_stream().wait_stream(self.h2d_stream)\n\n            # Time to reload the next group\n            self.bulk_reload_group(self.offloaded_group_count - 1)\n\n            # Decrease the offloading group counter\n            self.offloaded_group_count -= 1 if self.offloaded_group_count > 1 else 0\n\n        # Last group computation needs to wait till all the reloads complete\n        if self.current_group == 0:\n            get_torch_device().current_stream().wait_stream(self.h2d_stream)\n            self.offloaded_group_count = 0\n\n\ndef get_activation_offload_context(\n    num_layers: int = 1, model_layers: int = 1, tensor_need_offloading_checker=(lambda t: True)\n):\n    cpu_offload_handler = AsyncDoubleBufferGroupOffloadHandler(\n        num_offload_group=num_layers,\n        num_model_group=model_layers,\n        tensor_need_offloading_checker=tensor_need_offloading_checker,\n    )\n\n    def group_prefetch_offload_commit_async(tensor):\n        return group_prefetch_offload_commit(tensor, cpu_offload_handler)\n\n    return (\n        CpuOffloadHookWithOffloadHandler(offload_handler=cpu_offload_handler),\n        group_prefetch_offload_commit_async,\n    )\n\n\nclass ActivationHandler:\n    def __init__(self, offload_ctx, sync_func, tensor_filter, enable_ckpt):\n        self._offload_ctx = offload_ctx\n        self._sync_func = sync_func\n        self._enable_ckpt = enable_ckpt\n        self._tensor_filter = tensor_filter\n        if enable_ckpt:\n            self.checkpoint_fn = functools.partial(\n                torch.utils.checkpoint.checkpoint,\n                use_reentrant=True,\n            )\n\n    def pre_forward(self, module):\n        if module.training:\n            self._offload_ctx.__enter__()\n            self._tensor_filter.update_model_parameters(module)\n\n    def post_forward(self, module):\n        if module.training:\n            self._offload_ctx.__exit__(None, None, None)\n\n    def _pack_kwargs(self, *args, **kwargs):\n        kwarg_keys = []\n        flat_args = list(args)\n        for k, v in kwargs.items():\n            kwarg_keys.append(k)\n            flat_args.append(v)\n\n        return tuple(flat_args), tuple(kwarg_keys)\n\n    def _unpack_kwargs(self, flat_args, kwarg_keys):\n        assert len(kwarg_keys) <= len(flat_args), f\"too many keys {len(kwarg_keys)} vs. {len(flat_args)}\"\n        if len(kwarg_keys) == 0:\n            return flat_args, {}\n        args = flat_args[: -len(kwarg_keys)]\n        kwargs = dict(zip(kwarg_keys, flat_args[-len(kwarg_keys) :], strict=True))\n        return args, kwargs\n\n    def _ckpt_forward(self, forward_method, *args, **kwargs):\n        flat_args, kwarg_keys = self._pack_kwargs(*args, **kwargs)\n\n        def my_function(*inputs):\n            # unpack back into args and kwargs\n            nonlocal forward_method, kwarg_keys\n            unpacked_args, unpacked_kwargs = self._unpack_kwargs(inputs, kwarg_keys)\n            # run original module\n            return forward_method(*unpacked_args, **unpacked_kwargs)\n\n        return self.checkpoint_fn(\n            my_function,\n            *flat_args,\n        )\n\n    def forward(self, module, forward_method, *args, **kwargs):\n        if not module.training:\n            return forward_method(*args, **kwargs)\n        if not self._enable_ckpt:\n            ret = forward_method(*args, **kwargs)\n        else:\n            ret = self._ckpt_forward(forward_method, *args, **kwargs)\n        binded_tensor = ret\n        if isinstance(ret, tuple):\n            binded_tensor = ret[0]\n        binded_tensor = self._sync_func(binded_tensor)\n        final_ret = binded_tensor\n        if isinstance(ret, tuple):\n            final_ret = (final_ret,) + ret[1:]\n        return final_ret\n\n    def wrap_module_forward_method(self, module):\n        orig_method = module.forward\n        handler = self\n\n        @functools.wraps(orig_method)\n        def wrapped_method(model_self, *args, **kwargs):\n            nonlocal handler\n            handler.pre_forward(model_self)\n            out = handler.forward(model_self, orig_method, *args, **kwargs)\n            handler.post_forward(model_self)\n            return out\n\n        module.forward = wrapped_method.__get__(module, type(module))\n\n\ndef enable_activation_offloading(model, strategy, enable_ckpt=False):\n    \"\"\"\n    Enable activation offloading for the model. It groups activations by TransformerLayer and offloads activation\n    groups asynchronously. This means that the offloading of the i-th activation group and the computation of the i+1-th\n    activation group happen at the same time, and there are at most two activation groups in GPU memory.\n\n    Args:\n        model: the model to enable activation offloading\n        strategy: the training strategy of the model, such as \"fsdp\"\n        enable_ckpt: whether activation checkpointing(also called gradient checkpointing) has been enabled for the model\n\n    Note:\n        For best efficiency, activation offloading is usually combined with activation checkpointing. However, this\n        implementation of activation offloading is conflicted with the implementation of activation checkpointing in\n        some training strategies. This function resolves this conflict, and therefore requires the \"strategy\" and\n        \"enable_ckpt\" arguments.\n\n    Returns:\n\n    \"\"\"\n\n    assert strategy == \"fsdp\" or strategy == \"fsdp2\", \"activation offloading only supports fsdp strategy\"\n    layers = []\n\n    def get_layers(module):\n        for name, child in module.named_children():\n            if not isinstance(child, FSDP | FSDP2):\n                get_layers(child)\n            else:\n                wrapped_module = child\n                if isinstance(child, FSDP):\n                    wrapped_module = child._fsdp_wrapped_module\n                # In some cases, torch.nn.Embedding is wrapped with FSDP alone. However, the activation\n                # size of torch.nn.Embedding is small, so it's not necessary to offload it.\n                if not isinstance(wrapped_module, torch.nn.Embedding):\n                    layers.append(child)\n\n    get_layers(model)\n    if len(layers) < 3:\n        logger.warning(f\"Find only {len(layers)} fsdp layers, not necessary to enable async activation offloading\")\n        return\n\n    tensor_filter = FSDPParameterFilter()\n    context, sync_func = get_activation_offload_context(len(layers) - 1, len(layers), tensor_filter)\n    if enable_ckpt:\n        # The implementation of activation checkpointing in transformers library is incompatible with\n        # activation offloading,\n        # so it will be disabled, but this implementation supports another version of activation checkpointing, so that\n        # these two features can be enabled at the same time.\n        for module in model.modules():\n            if hasattr(module, \"gradient_checkpointing_disable\"):\n                module.gradient_checkpointing_disable()\n\n    handler = ActivationHandler(context, sync_func, tensor_filter, enable_ckpt)\n    for layer in layers:\n        module = layer\n        if isinstance(layer, FSDP):\n            module = module._fsdp_wrapped_module\n        handler.wrap_module_forward_method(module)\n"
  },
  {
    "path": "verl/utils/attention_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable\n\n_index_first_axis, _pad_input, _rearrange, _unpad_input = None, None, None, None\n\n\ndef _get_attention_functions() -> tuple[Callable, Callable, Callable, Callable]:\n    \"\"\"Dynamically import attention functions based on available hardware.\"\"\"\n\n    from verl.utils.device import is_torch_npu_available\n\n    global _index_first_axis, _pad_input, _rearrange, _unpad_input\n\n    if is_torch_npu_available(check_device=False):\n        from verl.utils.npu_flash_attn_utils import index_first_axis, pad_input, rearrange, unpad_input\n    else:\n        from flash_attn.bert_padding import index_first_axis, pad_input, rearrange, unpad_input\n\n    _index_first_axis, _pad_input, _rearrange, _unpad_input = index_first_axis, pad_input, rearrange, unpad_input\n\n    return _index_first_axis, _pad_input, _rearrange, _unpad_input\n\n\ndef index_first_axis(*args, **kwargs):\n    \"\"\"\n    Unified entry point for `index_first_axis` across CUDA and NPU backends.\n\n    Dynamically dispatches to the appropriate device-specific implementation:\n      - On CUDA: `flash_attn.bert_padding.index_first_axis`\n      - On NPU: `transformers.integrations.npu_flash_attention.index_first_axis`\n        (falls back to `transformers.modeling_flash_attention_utils._index_first_axis`\n        in newer versions of transformers).\n\n    Users can call this function directly without worrying about the underlying device.\n    \"\"\"\n    func, *_ = _get_attention_functions()\n    return func(*args, **kwargs)\n\n\ndef pad_input(*args, **kwargs):\n    \"\"\"\n    Unified entry point for `pad_input` across CUDA and NPU backends.\n\n    Dynamically dispatches to the appropriate device-specific implementation:\n      - On CUDA: `flash_attn.bert_padding.pad_input`\n      - On NPU: `transformers.integrations.npu_flash_attention.pad_input`\n        (falls back to `transformers.modeling_flash_attention_utils._pad_input`\n        in newer versions of transformers).\n\n    Users can call this function directly without worrying about the underlying device.\n    \"\"\"\n    _, func, *_ = _get_attention_functions()\n    return func(*args, **kwargs)\n\n\ndef rearrange(*args, **kwargs):\n    \"\"\"\n    Unified entry point for `rearrange` across CUDA and NPU backends.\n\n    Dynamically dispatches to the appropriate device-specific implementation:\n      - On CUDA: `flash_attn.bert_padding.rearrange`\n      - On NPU: `transformers.integrations.npu_flash_attention.rearrange`\n        (falls back to `einops.rearrange` if no dedicated NPU implementation exists).\n\n    Users can call this function directly without worrying about the underlying device.\n    \"\"\"\n    *_, func, _ = _get_attention_functions()\n    return func(*args, **kwargs)\n\n\ndef unpad_input(*args, **kwargs):\n    \"\"\"\n    Unified entry point for `unpad_input` across CUDA and NPU backends.\n\n    Dynamically dispatches to the appropriate device-specific implementation:\n      - On CUDA: `flash_attn.bert_padding.unpad_input`\n      - On NPU: `transformers.integrations.npu_flash_attention.unpad_input`\n        (falls back to `transformers.modeling_flash_attention_utils._unpad_input`\n        in newer versions of transformers).\n\n    Users can call this function directly without worrying about the underlying device.\n    \"\"\"\n    *_, func = _get_attention_functions()\n    return func(*args, **kwargs)\n\n\n__all__ = [\"index_first_axis\", \"pad_input\", \"rearrange\", \"unpad_input\"]\n"
  },
  {
    "path": "verl/utils/chat_template.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\nimport logging\nimport os\n\nfrom transformers import PreTrainedTokenizerBase, ProcessorMixin\n\nfrom verl.utils.tokenizer import normalize_token_ids\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef initialize_system_prompt(tokenizer, **apply_chat_template_kwargs) -> list[int]:\n    \"\"\"\n    Initialize system prompt tokens for chat templates that support them.\n\n    Args:\n        tokenizer: The tokenizer with a chat template\n        **apply_chat_template_kwargs: Additional arguments for apply_chat_template\n\n    Returns:\n        List of token IDs for the system prompt, or empty list if not supported\n    \"\"\"\n    token1 = normalize_token_ids(\n        tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": \"\"}], add_generation_prompt=False, tokenize=True)\n    )\n    token2 = normalize_token_ids(\n        tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": \"\"}] * 2, add_generation_prompt=False, tokenize=True)\n    )\n    # get system prompt tokens\n    system_prompt = token1[: -(len(token2) - len(token1))]\n    return system_prompt\n\n\ndef extract_system_prompt_and_generation(tokenizer):\n    token1 = normalize_token_ids(\n        tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": \"\"}], add_generation_prompt=False, tokenize=True)\n    )\n    token2 = normalize_token_ids(\n        tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": \"\"}] * 2, add_generation_prompt=False, tokenize=True)\n    )\n    # get system prompt tokens\n    system_prompt = token1[: -(len(token2) - len(token1))]\n    # get generate prompt tokens\n    token3 = normalize_token_ids(\n        tokenizer.apply_chat_template([{\"role\": \"user\", \"content\": \"\"}], add_generation_prompt=True, tokenize=True)\n    )\n    generate_prompt = token3[len(token1) :]\n\n    return system_prompt, generate_prompt\n\n\ndef apply_chat_template(\n    processor: PreTrainedTokenizerBase | ProcessorMixin,\n    messages: list[dict],\n    *,\n    tokenize: bool = True,\n    add_generation_prompt: bool = True,\n    tools=None,\n    return_dict: bool = False,\n    **kwargs,\n) -> list[int] | str:\n    \"\"\"apply_chat_template to messages with special attention to template requiring\n    at least one user message, e.g. Qwen3.5.\n\n    Args:\n        processor: tokenizer or processor.\n        messages: list[dict], messages.\n        tokenize: bool, whether to tokenize the output.\n        add_generation_prompt: bool, whether to add generation prompt.\n        tools: list[dict], tools schema.\n        return_dict: bool, whether to return a dict.\n        **kwargs: additional arguments for apply_chat_template.\n\n    Returns:\n        list[int] | str: tokenized ids or text string.\n    \"\"\"\n    try:\n        return processor.apply_chat_template(\n            messages,\n            tokenize=tokenize,\n            add_generation_prompt=add_generation_prompt,\n            tools=tools,\n            return_dict=return_dict,\n            **kwargs,\n        )\n    except Exception:\n        # Qwen3.5 apply_chat_template needs messages with at least one user message\n        dummy_user_message = [{\"role\": \"user\", \"content\": [{\"type\": \"text\", \"text\": \"\"}]}]\n        dummy_user_prefix = processor.apply_chat_template(\n            dummy_user_message,\n            tokenize=tokenize,\n            add_generation_prompt=False,\n            tools=tools,\n            return_dict=return_dict,\n            **kwargs,\n        )\n        output = processor.apply_chat_template(\n            dummy_user_message + messages,\n            tokenize=tokenize,\n            add_generation_prompt=add_generation_prompt,\n            tools=tools,\n            return_dict=return_dict,\n            **kwargs,\n        )\n\n        if not tokenize:  # tokenize=False\n            return output[len(dummy_user_prefix) :]\n        elif not return_dict:  # tokenize=True and return_dict=False\n            if isinstance(output[0], list):  # transformers>=5\n                assert len(output) == 1, \"output must be a list[int] or list[list[int]]\"\n                dummy_user_prefix = dummy_user_prefix[0]\n                output = output[0]\n            return output[len(dummy_user_prefix) :]\n        else:  # tokenize=True and return_dict=True and return_tensors=\"pt\"\n            dummy_user_prefix = dict(dummy_user_prefix)\n            output = dict(output)\n            prefix_len = dummy_user_prefix[\"input_ids\"].shape[1]\n            output[\"input_ids\"] = output[\"input_ids\"][:, prefix_len:]\n            output[\"attention_mask\"] = output[\"attention_mask\"][:, prefix_len:]\n            if \"mm_token_type_ids\" in output:\n                output[\"mm_token_type_ids\"] = output[\"mm_token_type_ids\"][:, prefix_len:]\n            return output\n"
  },
  {
    "path": "verl/utils/checkpoint/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .checkpoint_handler import CheckpointHandler, OrchestrationMode\n\n__all__ = [\"CheckpointHandler\", \"OrchestrationMode\"]\n"
  },
  {
    "path": "verl/utils/checkpoint/checkpoint_handler.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n# TODO: add unit tests\n\nimport json\nimport logging\nimport os\nimport re\nfrom enum import Enum\n\nimport torch\n\nimport verl.utils.hdfs_io as hdfs_io\nfrom verl.single_controller import WorkerGroup\nfrom verl.utils.checkpoint.checkpoint_manager import find_latest_ckpt_path, get_checkpoint_tracker_filename\nfrom verl.utils.logger import log_with_rank\nfrom verl.workers.engine import BaseEngine\n\n\ndef extract_step(path):\n    match = re.search(r\"global_step_(\\d+)\", path)\n    if match:\n        return int(match.group(1))\n    return None\n\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_SFT_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass OrchestrationMode(Enum):\n    SPMD = 0\n    RAY = 1\n\n\nclass CheckpointHandler:\n    \"\"\"\n    Checkpoint handler handles the path, global_step of a checkpoint folder.\n    Currently, it only works with a single model.\n    We can expand it to support multiple models. It is expected to be used with SPMD style (e.g., torchrun)\n    \"\"\"\n\n    def __init__(\n        self,\n        engine: BaseEngine | WorkerGroup,\n        train_dataloader,\n        *,\n        default_local_dir,\n        max_ckpt_to_keep=None,\n        default_hdfs_dir=None,\n        resume_mode=\"auto\",\n        resume_from_path=None,\n        mode=OrchestrationMode.SPMD,\n        lora_train_meta=None,\n    ):\n        self.default_local_dir = default_local_dir\n        self.max_ckpt_to_keep = max_ckpt_to_keep\n        self.default_hdfs_dir = default_hdfs_dir\n        self.resume_mode = resume_mode\n        self.resume_from_path = resume_from_path\n        self.engine = engine\n        self.train_dataloader = train_dataloader\n        self.mode = mode\n        self.lora_train_meta = lora_train_meta\n\n        if self.mode == OrchestrationMode.SPMD:\n            self.rank = torch.distributed.get_rank()\n            self.is_mp_src_rank_with_outputs = self.engine.is_mp_src_rank_with_outputs()\n            self.dp_rank = self.engine.get_data_parallel_rank()\n        elif self.mode == OrchestrationMode.RAY:\n            self.rank = 0\n            self.is_mp_src_rank_with_outputs = True\n            self.dp_rank = 0\n        else:\n            raise ValueError(f\"Unknown {self.mode=}\")\n\n    def save_checkpoint(self, step):\n        \"\"\"Save checkpoint using FSDPCheckpointManager with improved tracking\"\"\"\n        from verl.utils.fs import local_mkdir_safe\n\n        # Determine checkpoint path\n        local_global_step_folder = os.path.join(self.default_local_dir, f\"global_step_{step}\")\n        if self.rank == 0:\n            print(f\"Saving checkpoint to: {local_global_step_folder}\")\n\n        # Get max checkpoints to keep\n        max_ckpt_to_keep = self.max_ckpt_to_keep\n\n        # Use checkpoint manager to save\n        self.engine.save_checkpoint(\n            local_path=local_global_step_folder, global_step=step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n\n        # Save dataloader state. Note that we only save the iterator in the train_dataloader.\n        # So it's identical in each dp rank.\n        if self.rank == 0 and self.lora_train_meta is not None:\n            local_mkdir_safe(local_global_step_folder)\n            lora_meta_path = os.path.join(local_global_step_folder, \"lora_train_meta.json\")\n            with open(lora_meta_path, \"w\", encoding=\"utf-8\") as f:\n                json.dump(self.lora_train_meta, f, ensure_ascii=False, indent=4)\n            print(f\"Saved LoRA rank/alpha metadata to: {lora_meta_path}\")\n\n        if self.is_mp_src_rank_with_outputs:\n            dp_rank = self.dp_rank\n            local_mkdir_safe(local_global_step_folder)\n            dataloader_local_path = os.path.join(local_global_step_folder, f\"data_{dp_rank}.pt\")\n\n            # Use StatefulDataLoader's built-in state dict functionality\n            dataloader_state_dict = self.train_dataloader.state_dict()\n            torch.save(dataloader_state_dict, dataloader_local_path)\n            print(f\"Saved dataloader state to: {dataloader_local_path}\")\n\n        if self.rank == 0:\n            # Update latest checkpoint tracker (atomic write)\n            tracker_file = get_checkpoint_tracker_filename(self.default_local_dir)\n            temp_tracker_file = tracker_file + \".tmp\"\n            with open(temp_tracker_file, \"w\") as f:\n                f.write(str(step))\n            os.rename(temp_tracker_file, tracker_file)\n            print(f\"Updated checkpoint tracker: {tracker_file}\")\n\n        # Copy to HDFS if configured\n        if self.rank == 0 and self.default_hdfs_dir:\n            hdfs_io.makedirs(self.default_hdfs_dir, exist_ok=True)\n            hdfs_io.copy(src=local_global_step_folder, dst=self.default_hdfs_dir, dirs_exist_ok=True)\n\n        if self.mode == OrchestrationMode.SPMD:\n            torch.distributed.barrier()\n\n    def load_checkpoint(self):\n        # Determine resume path based on configuration\n        checkpoint_path = self._determine_resume_path()\n\n        if checkpoint_path is None:\n            return 0\n\n        # extract resume step from checkpoint path\n        resume_step = extract_step(checkpoint_path)\n        if resume_step is None:\n            log_with_rank(\n                f\"Warning: Could not extract step number from {checkpoint_path}, starting from step 0\",\n                logger=logger,\n                rank=self.rank,\n                level=logging.WARNING,\n                log_only_rank_0=True,\n            )\n            return 0\n        self.resume_global_step = resume_step\n\n        # Use checkpoint manager to load model state\n        self.engine.load_checkpoint(checkpoint_path)\n        # Always load dataloader state for StatefulDataLoader\n        self._load_dataloader_state(checkpoint_path)\n\n        return resume_step\n\n    def _load_dataloader_state(self, checkpoint_path: str):\n        \"\"\"Load dataloader state from checkpoint\"\"\"\n        dp_rank = self.dp_rank\n        dataloader_path = os.path.join(checkpoint_path, f\"data_{dp_rank}.pt\")\n\n        if os.path.exists(dataloader_path):\n            # Use StatefulDataLoader's built-in state dict functionality\n            dataloader_state_dict = torch.load(dataloader_path, map_location=\"cpu\", weights_only=False)\n            self.train_dataloader.load_state_dict(dataloader_state_dict)\n\n            log_with_rank(\n                f\"Successfully loaded dataloader state from {dataloader_path}\",\n                logger=logger,\n                rank=self.rank,\n                log_only_rank_0=True,\n            )\n\n        else:\n            log_with_rank(\n                f\"Warning: No dataloader state found at {dataloader_path}, will start from scratch\",\n                logger=logger,\n                rank=self.rank,\n                level=logging.WARNING,\n                log_only_rank_0=True,\n            )\n\n    def _determine_resume_path(self):\n        \"\"\"Determine the path to resume from based on resume_mode configuration\"\"\"\n        resume_mode = self.resume_mode\n        resume_from_path = self.resume_from_path\n\n        if resume_mode == \"disable\":\n            return None\n        elif resume_mode == \"auto\":\n            if resume_from_path is not None:\n                assert os.path.exists(resume_from_path), (\n                    \"resume_from_path must be null or an existing path when resume_mode is 'auto'\"\n                )\n                assert \"global_step_\" in resume_from_path, \"resume_from_path must specify the global_steps\"\n                return resume_from_path\n            # Try to find the latest checkpoint in the default directory\n            return self._find_latest_checkpoint()\n        elif resume_mode == \"resume_path\":\n            assert os.path.exists(resume_from_path), (\n                \"resume_from_path must be an existing path when resume_mode is 'resume_path'\"\n            )\n            assert \"global_step_\" in resume_from_path, \"resume_from_path must specify the global_steps\"\n            return resume_from_path\n        else:\n            raise ValueError(f\"Invalid resume_mode: {resume_mode}. Must be 'auto', 'disable', or 'resume_path'\")\n\n    def _find_latest_checkpoint(self):\n        \"\"\"Find the latest checkpoint in the default local directory\"\"\"\n        checkpoint_dir = self.default_local_dir\n\n        if not os.path.exists(checkpoint_dir):\n            return None\n\n        latest_checkpoint = find_latest_ckpt_path(checkpoint_dir)\n\n        if latest_checkpoint and self.rank == 0:\n            step_num = extract_step(latest_checkpoint)\n            print(f\"Found latest checkpoint: {latest_checkpoint} (step {step_num})\")\n\n        return latest_checkpoint\n"
  },
  {
    "path": "verl/utils/checkpoint/checkpoint_manager.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport random\nimport shutil\n\nimport numpy as np\nimport torch\nimport torch.distributed\nfrom omegaconf import DictConfig\nfrom transformers import PreTrainedTokenizer, ProcessorMixin\n\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils.device import get_device_name, get_torch_device\n\n\nclass BaseCheckpointManager:\n    \"\"\"\n    A checkpoint manager that saves and loads the following states in a SPMD way:\n    - model\n    - optimizer\n    - lr_scheduler\n    - extra_states\n\n    We save\n    - sharded model states and optimizer states\n    - full lr_scheduler states\n    - huggingface tokenizer and config for ckpt merge\n    \"\"\"\n\n    def __init__(\n        self,\n        model,\n        optimizer: torch.optim.Optimizer,\n        lr_scheduler: torch.optim.lr_scheduler.LRScheduler = None,\n        processing_class: PreTrainedTokenizer | ProcessorMixin = None,\n        checkpoint_config: DictConfig | CheckpointConfig = None,\n    ):\n        self.checkpoint_config = checkpoint_config\n        checkpoint_load_contents = checkpoint_config.get(\"load_contents\", None) if checkpoint_config else None\n        checkpoint_save_contents = checkpoint_config.get(\"save_contents\", None) if checkpoint_config else None\n        if checkpoint_load_contents is None:\n            checkpoint_load_contents = [\"model\", \"optimizer\", \"extra\"]\n        if checkpoint_save_contents is None:\n            checkpoint_save_contents = [\"model\", \"optimizer\", \"extra\"]\n        self.previous_global_step = None\n        self.previous_saved_paths = []\n\n        self.model = model\n        self.optimizer = optimizer\n        self.lr_scheduler = lr_scheduler\n        self.processing_class = processing_class\n        self.checkpoint_load_contents = checkpoint_load_contents\n        self.checkpoint_save_contents = checkpoint_save_contents\n\n        self.rank = torch.distributed.get_rank()\n        self.world_size = torch.distributed.get_world_size()\n\n    @property\n    def should_save_model(self) -> bool:\n        \"\"\"\n        Returns True if 'model' is in checkpoint_save_contents, indicating the model state should be saved.\n        \"\"\"\n        return \"model\" in self.checkpoint_save_contents\n\n    @property\n    def should_save_optimizer(self) -> bool:\n        \"\"\"\n        Returns True if 'optimizer' is in checkpoint_save_contents, indicating the optimizer state should be saved.\n        \"\"\"\n        return \"optimizer\" in self.checkpoint_save_contents\n\n    @property\n    def should_save_extra(self) -> bool:\n        \"\"\"\n        Returns True if 'extra' is in checkpoint_save_contents, indicating the extra state should be saved.\n        \"\"\"\n        return \"extra\" in self.checkpoint_save_contents\n\n    @property\n    def should_save_hf_model(self) -> bool:\n        \"\"\"\n        Returns True if 'hf_model' is in checkpoint_save_contents, indicating the model should be converted to hf\n        model and saved.\n        \"\"\"\n        return \"hf_model\" in self.checkpoint_save_contents\n\n    @property\n    def should_load_model(self) -> bool:\n        \"\"\"\n        Returns True if 'model' is in checkpoint_load_contents, indicating the model state should be loaded.\n        \"\"\"\n        return \"model\" in self.checkpoint_load_contents\n\n    @property\n    def should_load_optimizer(self) -> bool:\n        \"\"\"\n        Returns True if 'optimizer' is in checkpoint_load_contents, indicating the optimizer state should be loaded.\n        \"\"\"\n        return \"optimizer\" in self.checkpoint_load_contents\n\n    @property\n    def should_load_extra(self) -> bool:\n        \"\"\"\n        Returns True if 'extra' is in checkpoint_load_contents, indicating the extra state should be loaded.\n        \"\"\"\n        return \"extra\" in self.checkpoint_load_contents\n\n    def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load: bool = False):\n        raise NotImplementedError\n\n    def save_checkpoint(\n        self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep: int = None\n    ):\n        raise NotImplementedError\n\n    @staticmethod\n    def checkpath(local_path: str, hdfs_path: str):\n        assert local_path is not None or hdfs_path is not None, \"local_path and hdfs_path cannot be both None\"\n        return local_path is not None, local_path if local_path is not None else hdfs_path\n\n    def remove_previous_save_local_path(self, path):\n        if isinstance(path, str):\n            path = [path]\n        for p in path:\n            abs_path = os.path.abspath(p)\n            print(f\"Checkpoint manager remove previous save local path: {abs_path}\")\n            if not os.path.exists(abs_path):\n                continue\n            shutil.rmtree(abs_path, ignore_errors=True)\n\n    def ensure_checkpoint_capacity(self, max_ckpt_to_keep: int):\n        \"\"\"\n        Remove old checkpoints to make room for a new one, keeping a safety buffer.\n\n        With max_ckpt_to_keep=1, this does nothing - we keep the existing checkpoint\n        until the new save completes successfully (handled by register_checkpoint).\n        For max_ckpt_to_keep >= 2, we keep (max_ckpt_to_keep - 1) checkpoints before save.\n        \"\"\"\n        if not (max_ckpt_to_keep and isinstance(max_ckpt_to_keep, int) and max_ckpt_to_keep > 1):\n            return\n        if len(self.previous_saved_paths) >= max_ckpt_to_keep:\n            keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep + 1\n            self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start])\n            self.previous_saved_paths = self.previous_saved_paths[keep_start:]\n\n    def register_checkpoint(self, new_path: str, max_ckpt_to_keep: int):\n        \"\"\"\n        Register a successfully saved checkpoint and enforce retention limit.\n\n        Adds the new checkpoint path to tracking and removes excess old\n        checkpoints beyond max_ckpt_to_keep.\n        \"\"\"\n        self.previous_saved_paths.append(new_path)\n        if not (max_ckpt_to_keep and isinstance(max_ckpt_to_keep, int) and max_ckpt_to_keep > 0):\n            return\n        if len(self.previous_saved_paths) > max_ckpt_to_keep:\n            keep_start = len(self.previous_saved_paths) - max_ckpt_to_keep\n            self.remove_previous_save_local_path(self.previous_saved_paths[:keep_start])\n            self.previous_saved_paths = self.previous_saved_paths[keep_start:]\n\n    @staticmethod\n    def get_rng_state():\n        rng_state = {\n            \"cpu\": torch.get_rng_state(),\n            \"numpy\": np.random.get_state(),\n            \"random\": random.getstate(),\n        }\n\n        if get_device_name() != \"cpu\":\n            rng_state[get_device_name()] = get_torch_device().get_rng_state()\n\n        return rng_state\n\n    @staticmethod\n    def load_rng_state(rng_state):\n        torch.set_rng_state(rng_state[\"cpu\"])\n        np.random.set_state(rng_state[\"numpy\"])\n        random.setstate(rng_state[\"random\"])\n\n        if get_device_name() != \"cpu\":\n            get_torch_device().set_rng_state(rng_state[get_device_name()])\n\n\ndef find_latest_ckpt_path(path, directory_format=\"global_step_{}\"):\n    \"\"\"\n    Return the most recent checkpoint directory based on a tracker file.\n\n    Args:\n        path (str): Base directory containing the checkpoint tracker.\n        directory_format (str): Template for checkpoint subfolders with one\n            placeholder for the iteration number (default \"global_step_{}\").\n\n    Returns:\n        str or None: Full path to the latest checkpoint directory, or\n        None if the tracker or checkpoint folder is missing.\n    \"\"\"\n    if path is None:\n        return None\n\n    tracker_file = get_checkpoint_tracker_filename(path)\n    if not os.path.exists(tracker_file):\n        if not torch.distributed.is_initialized() or torch.distributed.get_rank() == 0:\n            print(f\"Checkpoint tracker file does not exist: {tracker_file}\")\n        return None\n\n    with open(tracker_file, \"rb\") as f:\n        iteration = int(f.read().decode())\n    ckpt_path = os.path.join(path, directory_format.format(iteration))\n    if not os.path.exists(ckpt_path):\n        print(\"Checkpoint does not exist: %s\", ckpt_path)\n        return None\n\n    print(\"Found checkpoint: %s\", ckpt_path)\n    return ckpt_path\n\n\ndef get_checkpoint_tracker_filename(root_path: str):\n    \"\"\"\n    Tracker file rescords the latest chckpoint during training to restart from.\n    \"\"\"\n    return os.path.join(root_path, \"latest_checkpointed_iteration.txt\")\n\n\ndef should_save_ckpt_esi(max_steps_duration: float, save_ckpt_duration: float = 60, redundant_time: float = 0) -> bool:\n    \"\"\"\n    Determine if checkpoint should be saved based on capacity esi expiration.\n\n    Args:\n        max_steps_duration: Max estimated time (seconds) required to complete one training step\n        save_ckpt_duration: Estimated time (seconds) required to save checkpoint (default: 60)\n        redundant_time: Additional buffer time (seconds) for unexpected delays (default: 0)\n    \"\"\"\n    exp_ts_mlp = os.getenv(\"MLP_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\")  # vemlp\n    exp_ts_aws = os.getenv(\"SAGEMAKER_CURRENT_CAPACITY_BLOCK_EXPIRATION_TIMESTAMP\")  # aws\n    if exp_ts_mlp:\n        try:\n            import time\n\n            remaining = float(exp_ts_mlp) - time.time()\n        except ValueError:\n            return False\n        return (\n            remaining > 0\n            and max_steps_duration > 0\n            and remaining <= save_ckpt_duration + max_steps_duration + redundant_time\n        )\n    elif exp_ts_aws:\n        from datetime import datetime, timedelta\n\n        expiration_time = datetime.fromtimestamp(int(exp_ts_aws))\n        time_difference = expiration_time - datetime.now()\n        threshold_minutes = (save_ckpt_duration + max_steps_duration + redundant_time) / 60\n        return time_difference < timedelta(minutes=threshold_minutes)\n    else:\n        return False\n"
  },
  {
    "path": "verl/utils/checkpoint/fsdp_checkpoint_manager.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport logging\nimport os\nimport warnings\nfrom dataclasses import asdict, dataclass\nfrom typing import Optional\n\nimport torch\nimport torch.distributed\nfrom accelerate import init_empty_weights\nfrom omegaconf import DictConfig\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp import ShardedOptimStateDictConfig, ShardedStateDictConfig, StateDictType\nfrom transformers import GenerationConfig, PreTrainedTokenizer, ProcessorMixin\nfrom transformers.dynamic_module_utils import custom_object_save\n\nfrom verl.utils.device import is_cuda_available\nfrom verl.utils.fs import copy_to_local, is_non_local, local_mkdir_safe\nfrom verl.utils.fsdp_utils import fsdp_version, get_fsdp_full_state_dict, get_fsdp_state_ctx\nfrom verl.utils.logger import log_with_rank\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\n\nfrom .checkpoint_manager import BaseCheckpointManager\n\n# Setup logging\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\n\n\n@dataclass\nclass FSDPConfig:\n    \"\"\"Configuration for FSDP checkpointing.\n\n    Args:\n        FSDP_version (int): Version of FSDP being used.\n        world_size (int): Number of processes in the distributed training setup.\n    \"\"\"\n\n    FSDP_version: int\n    world_size: int\n\n\nclass FSDPCheckpointManager(BaseCheckpointManager):\n    \"\"\"\n    Manage FSDP checkpointing in SPMD training.\n\n    - Saves/loads per-rank sharded model & optimizer states\n    - Persists full lr_scheduler and RNG state\n    - Stores HF tokenizer/processor and model/config for unified restore\n\n    Args:\n        model (FSDP): Wrapped model instance.\n        optimizer (Optimizer): Training optimizer.\n        lr_scheduler (LRScheduler): Learning-rate scheduler.\n        processing_class (PreTrainedTokenizer or ProcessorMixin, optional):\n            Pre-/post-processing artifact handler.\n        checkpoint_contents DictConfig: Configuration for checkpoint contents.\n            - 'load': Components to load; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].\n            - 'save': Components to save; must contain 'model'. Defaults to ['model', 'optimizer', 'extra'].\n        trust_remote_code: Whether to trust_remote_code when loading the model configuration\n    \"\"\"\n\n    def __init__(\n        self,\n        model: FSDP,\n        optimizer: Optional[torch.optim.Optimizer] = None,\n        lr_scheduler: Optional[torch.optim.lr_scheduler.LRScheduler] = None,\n        processing_class: PreTrainedTokenizer | ProcessorMixin = None,\n        checkpoint_config: DictConfig = None,\n        trust_remote_code: bool = False,\n        **kwargs,\n    ):\n        if processing_class is None and \"tokenizer\" in kwargs:\n            warnings.warn(\n                \"`tokenizer` is deprecated. use `processing_class` instead.\", DeprecationWarning, stacklevel=2\n            )\n            processing_class = kwargs.pop(\"tokenizer\")\n\n        super().__init__(\n            model,\n            optimizer,\n            lr_scheduler=lr_scheduler,\n            processing_class=processing_class,\n            checkpoint_config=checkpoint_config,\n        )\n        self.trust_remote_code = trust_remote_code\n\n    def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False):\n        \"\"\"\n        Load an FSDP checkpoint for this rank.\n\n        Downloads and loads:\n          - model and optimizer shards\n          - extra state dict (scheduler + RNG)\n\n        Args:\n            local_path: Directory with per-rank checkpoint files.\n            hdfs_path: Unused (for API compatibility).\n            del_local_after_load: Remove local files after loading.\n        \"\"\"\n        if local_path is None:\n            return\n\n        # check if the checkpoint_load_contents is valid\n        if self.should_load_model:\n            assert self.model is not None, \"model must be provided when checkpoint_contents.load includes ['model']\"\n        if self.should_load_optimizer:\n            assert self.optimizer is not None, (\n                \"optimizer must be provided when checkpoint_contents.load includes ['optimizer']\"\n            )\n\n        # every rank download its own checkpoint\n        state_dict_cfg = (\n            ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)\n            if self.should_load_model\n            else None\n        )\n        optim_cfg = (\n            ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)\n            if self.should_load_optimizer\n            else None\n        )\n        with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):\n            if self.should_load_model:\n                remote_model_path = os.path.join(local_path, f\"model_world_size_{self.world_size}_rank_{self.rank}.pt\")\n                local_model_path = copy_to_local(remote_model_path)\n                model_state_dict = torch.load(local_model_path, weights_only=False)\n                self.model.load_state_dict(model_state_dict)\n                log_with_rank(f\"Loaded model from {remote_model_path}\", rank=self.rank, logger=logger)\n\n            if self.should_load_optimizer:\n                remote_optim_path = os.path.join(local_path, f\"optim_world_size_{self.world_size}_rank_{self.rank}.pt\")\n                local_optim_path = copy_to_local(remote_optim_path)\n                optimizer_state_dict = torch.load(local_optim_path, weights_only=False)\n                self.optimizer.load_state_dict(optimizer_state_dict)\n                log_with_rank(f\"Loaded optimizer from {remote_optim_path}\", rank=self.rank, logger=logger)\n\n        if self.should_load_extra:\n            remote_extra_state_path = os.path.join(\n                local_path, f\"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt\"\n            )\n            local_extra_state_path = copy_to_local(remote_extra_state_path)\n            extra_state_dict = torch.load(local_extra_state_path, weights_only=False)\n            # recover random state\n            if \"rng\" in extra_state_dict:\n                # 'rng' may not exist for backward compatibility\n                self.load_rng_state(extra_state_dict[\"rng\"])\n                log_with_rank(f\"Loaded rng from {remote_extra_state_path}\", rank=self.rank, logger=logger)\n\n            lr_scheduler_state_dict = extra_state_dict[\"lr_scheduler\"]\n            if lr_scheduler_state_dict is not None and self.lr_scheduler is not None:\n                self.lr_scheduler.load_state_dict(lr_scheduler_state_dict)\n                log_with_rank(f\"Loaded lr_scheduler from {remote_extra_state_path}\", rank=self.rank, logger=logger)\n\n        if self.rank == 0 and del_local_after_load:\n            try:\n                os.remove(local_model_path) if is_non_local(local_model_path) else None\n                os.remove(local_optim_path) if is_non_local(local_optim_path) else None\n                os.remove(local_extra_state_path) if is_non_local(local_extra_state_path) else None\n            except Exception as e:\n                log_with_rank(\n                    f\"remove local resume ckpt file after loading failed, exception {e} will be ignored\",\n                    rank=self.rank,\n                    logger=logger,\n                )\n\n        # wait for everyone to load checkpoints\n        torch.distributed.barrier()\n\n    def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None):\n        \"\"\"\n        Save an FSDP checkpoint for this rank.\n\n        Writes:\n          - model & optimizer shard files\n          - extra state dict (scheduler + RNG)\n          - HF tokenizer/processor and model/config on rank 0\n          - optional full HF model under 'huggingface/' if requested\n\n        Rotates old checkpoints, keeping at most `max_ckpt_to_keep`.\n\n        Args:\n            local_path: Target directory for checkpoint files.\n            hdfs_path: Unused (for API compatibility).\n            global_step: Current training step (used for bookkeeping).\n            max_ckpt_to_keep: Number of recent checkpoints to retain.\n        \"\"\"\n        if local_path is None:\n            return\n\n        # record the previous global step\n        self.previous_global_step = global_step\n\n        if self.rank == 0:\n            self.ensure_checkpoint_capacity(max_ckpt_to_keep)\n\n        local_path = local_mkdir_safe(local_path)\n        torch.distributed.barrier()\n\n        # check if the checkpoint_save_contents is valid\n        if self.should_save_model:\n            assert self.model is not None, \"model must be provided when checkpoint_contents.save includes ['model']\"\n        if self.should_save_optimizer:\n            assert self.optimizer is not None, (\n                \"optimizer must be provided when checkpoint_contents.save includes ['optimizer']\"\n            )\n\n        # every rank will save its own model and optim shard\n        state_dict_cfg = ShardedStateDictConfig(offload_to_cpu=True if is_cuda_available else False)\n        optim_cfg = ShardedOptimStateDictConfig(offload_to_cpu=True if is_cuda_available else False)\n        with warnings.catch_warnings():\n            warnings.simplefilter(\"ignore\")\n            with get_fsdp_state_ctx(self.model, StateDictType.SHARDED_STATE_DICT, state_dict_cfg, optim_cfg):\n                model_path = os.path.join(local_path, f\"model_world_size_{self.world_size}_rank_{self.rank}.pt\")\n                optim_path = os.path.join(local_path, f\"optim_world_size_{self.world_size}_rank_{self.rank}.pt\")\n                extra_path = os.path.join(local_path, f\"extra_state_world_size_{self.world_size}_rank_{self.rank}.pt\")\n\n                if self.should_save_model:\n                    model_state_dict = self.model.state_dict()\n                    torch.save(model_state_dict, model_path)\n                    log_with_rank(f\"Saved model to {os.path.abspath(model_path)}\", rank=self.rank, logger=logger)\n\n                if self.should_save_optimizer:\n                    optimizer_state_dict = self.optimizer.state_dict()\n                    torch.save(optimizer_state_dict, optim_path)\n                    log_with_rank(f\"Saved optim to {os.path.abspath(optim_path)}\", rank=self.rank, logger=logger)\n\n                if self.should_save_extra:\n                    lr_scheduler_state_dict = self.lr_scheduler.state_dict() if self.lr_scheduler is not None else None\n                    extra_state_dict = {\n                        \"lr_scheduler\": lr_scheduler_state_dict,\n                        \"rng\": self.get_rng_state(),\n                    }\n                    torch.save(extra_state_dict, extra_path)\n                    log_with_rank(f\"Saved extra_state to {os.path.abspath(extra_path)}\", rank=self.rank, logger=logger)\n\n        if self.rank == 0:\n            # Save HF tokenizer/processor and model config on rank 0 to huggingface/ directory, no matter whether\n            # huggingface model is requested to be saved or not.\n\n            if fsdp_version(self.model) == 1:\n                unwrap_model = self.model._fsdp_wrapped_module\n            else:\n                unwrap_model = self.model\n\n            hf_config_tokenizer_path = os.path.join(local_path, \"huggingface\")\n            local_mkdir_safe(hf_config_tokenizer_path)\n            model_config = unwrap_model.config\n            generation_config = None\n            if unwrap_model.can_generate() and hasattr(model_config, \"name_or_path\") and model_config.name_or_path:\n                try:\n                    # Some model's name_or_path is empty if not initialized from pretrained,\n                    # in this cases, we don't save generation config.\n                    generation_config = GenerationConfig.from_pretrained(model_config.name_or_path)\n                    generation_config.save_pretrained(hf_config_tokenizer_path)\n                except Exception:\n                    # if the generation config isn't available, we don't save it\n                    pass\n\n            if hasattr(model_config, \"auto_map\") and None in model_config.auto_map:\n                model_config.auto_map = {k: v for k, v in model_config.auto_map.items() if k is not None}\n\n            model_config.save_pretrained(hf_config_tokenizer_path)\n            if self.processing_class is not None:\n                self.processing_class.save_pretrained(hf_config_tokenizer_path)\n            log_with_rank(\n                f\"Saved model config and tokenizer class to {os.path.abspath(hf_config_tokenizer_path)}\",\n                rank=self.rank,\n                logger=logger,\n                log_only_rank_0=True,\n            )\n\n            # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be\n            # loaded from the Hub.\n            if hasattr(model_config, \"auto_map\"):\n                custom_object_save(unwrap_model, hf_config_tokenizer_path, config=model_config)\n\n            # Also save runtime FSDP config\n            fsdp_config_path = os.path.join(local_path, \"fsdp_config.json\")\n            fsdp_config = FSDPConfig(\n                FSDP_version=fsdp_version(self.model),\n                world_size=self.world_size,\n            )\n            with open(fsdp_config_path, \"w\") as f:\n                json.dump(asdict(fsdp_config), f, indent=4)\n\n        # wait for everyone to dump to local\n        torch.distributed.barrier()\n\n        if self.should_save_hf_model:\n            # Only rank 0 will save hf model and,\n            # offload to cpu to save LLMs which may be too large to fit in one GPU\n            state_dict = get_fsdp_full_state_dict(self.model, offload_to_cpu=True, rank0_only=True)\n\n            if self.rank == 0:\n                hf_local_path = os.path.join(local_path, \"huggingface\")\n                os.makedirs(hf_local_path, exist_ok=True)\n\n                if \"ForTokenClassification\" in model_config.architectures[0]:\n                    from transformers import AutoModelForTokenClassification\n\n                    auto_model_cls = AutoModelForTokenClassification\n                elif \"ForCausalLM\" in model_config.architectures[0]:\n                    from transformers import AutoModelForCausalLM\n\n                    auto_model_cls = AutoModelForCausalLM\n                elif \"ForConditionalGeneration\" in model_config.architectures[0]:\n                    auto_model_cls = get_auto_model_for_vision2seq()\n                else:\n                    raise NotImplementedError(f\"Unknown architecture {model_config['architectures']}\")\n\n                with init_empty_weights():\n                    save_model = auto_model_cls.from_config(\n                        model_config, torch_dtype=torch.bfloat16, trust_remote_code=self.trust_remote_code\n                    )\n\n                save_model.to_empty(device=\"cpu\")\n\n                if save_model.can_generate():\n                    if generation_config is not None:\n                        save_model.generation_config = generation_config\n                    else:\n                        print(\n                            f\"Warning: {self.__class__.__name__}.save_checkpoint: Generation config file not found \"\n                            f\"in, using a generation config created from the model config when saving hf_model.\"\n                        )\n\n                save_model.save_pretrained(hf_local_path, state_dict=state_dict)\n                log_with_rank(\n                    f\"Saved hf_model to {os.path.abspath(hf_local_path)}\",\n                    rank=self.rank,\n                    logger=logger,\n                    log_only_rank_0=True,\n                )\n                del state_dict\n                del save_model\n\n            # wait for rank0 to dump hf_model to local\n            torch.distributed.barrier()\n\n        if self.rank == 0:\n            self.register_checkpoint(local_path, max_ckpt_to_keep)\n"
  },
  {
    "path": "verl/utils/checkpoint/megatron_checkpoint_manager.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\nimport json\nimport logging\nimport os\nimport random\nfrom collections.abc import Callable\nfrom dataclasses import asdict\n\nimport megatron.core\nimport numpy as np\nimport torch\nimport torch.distributed\nfrom megatron.core import dist_checkpointing, mpu, tensor_parallel\nfrom megatron.core.dist_checkpointing.mapping import ShardedObject\nfrom megatron.core.transformer.enums import AttnBackend\nfrom packaging import version\nfrom transformers import GenerationConfig\n\nfrom verl.models.weight_loader_registry import get_weight_saver\nfrom verl.utils.device import get_device_name, get_torch_device\nfrom verl.utils.fs import is_non_local, local_mkdir_safe\nfrom verl.utils.logger import log_with_rank\nfrom verl.utils.megatron.dist_checkpointing import load_dist_checkpointing, save_dist_checkpointing\nfrom verl.utils.megatron_utils import (\n    get_dist_checkpoint_path,\n    get_hf_model_checkpoint_path,\n    get_transformer_config_checkpoint_path,\n)\n\nfrom .checkpoint_manager import BaseCheckpointManager\n\n# Setup logging\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\nmcore_ge_014 = version.parse(megatron.core.__version__) >= version.parse(\"0.14.0\")\nif not mcore_ge_014:\n    logger.warning(\n        \"Detected megatron.core %s, recommend upgrading to >= 0.14.0 for better checkpoint compatibility\",\n        megatron.core.__version__,\n    )\n\n\nclass MegatronCheckpointManager(BaseCheckpointManager):\n    \"\"\"\n    Checkpoint manager for Megatron-LM distributed training.\n\n    This class manages the saving and loading of model checkpoints in a Megatron-LM\n    distributed training environment. It handles various aspects of checkpointing\n    including model states, optimizer states, learning rate schedulers, and random\n    number generator states, ensuring compatibility with HuggingFace formats.\n\n    Key features:\n    - Distributed checkpoint saving and loading using Megatron's dist_checkpointing\n    - Support for tensor parallel, pipeline parallel, and data parallel configurations\n    - Automatic handling of model state dictionaries across multiple pipeline stages\n    - Integration with HuggingFace model configurations and tokenizers\n    - Random number generator state management for reproducibility\n    - Support for both synchronous and asynchronous checkpoint operations\n\n    The manager automatically handles:\n    - Directory structure creation based on global steps and process ranks\n    - Model configuration and tokenizer saving in HuggingFace format\n    - Optimizer and scheduler state persistence\n    - CUDA RNG state management for deterministic training\n    - Checkpoint cleanup and retention policies\n\n    Args:\n        model: The Megatron model instance to checkpoint\n        optimizer: The optimizer instance (optional)\n        lr_scheduler: The learning rate scheduler instance (optional)\n\n    Attributes:\n        model: Reference to the Megatron model being checkpointed\n        optimizer: Reference to the optimizer (if provided)\n        lr_scheduler: Reference to the learning rate scheduler (if provided)\n        rank: Current process rank in the distributed setup\n\n    Example:\n        ```python\n        checkpoint_manager = MegatronCheckpointManager(\n            model=megatron_model,\n            optimizer=optimizer,\n            lr_scheduler=scheduler\n        )\n\n        checkpoint_manager.save_checkpoint(\n            local_path=\"checkpoints/step_1000\",\n            global_step=1000\n        )\n\n        checkpoint_manager.load_checkpoint(\n            local_path=\"checkpoints/step_1000\"\n        )\n        ```\n    \"\"\"\n\n    def __init__(\n        self,\n        config,\n        checkpoint_config,\n        model_config,\n        transformer_config,\n        role,\n        model: torch.nn.ModuleList,\n        arch: str,\n        hf_config,\n        param_dtype: torch.dtype,\n        share_embeddings_and_output_weights: bool,\n        processing_class,\n        optimizer,\n        optimizer_scheduler,\n        use_distributed_optimizer: bool,\n        use_checkpoint_opt_param_scheduler: bool = False,\n        use_dist_checkpointing: bool = True,\n        bridge=None,\n        provider=None,\n        peft_cls=None,\n        **kwargs,\n    ):\n        super().__init__(\n            model,\n            optimizer=optimizer,\n            lr_scheduler=optimizer_scheduler,\n            processing_class=processing_class,\n            checkpoint_config=checkpoint_config,\n        )\n        self.arch = arch\n        self.config = config\n        self.transformer_config = transformer_config\n        self.role = role\n        self.is_value_model = False\n        if self.role in [\"reward\", \"critic\"]:\n            self.is_value_model = True\n        self.model_config = model_config\n        self.hf_config = hf_config\n        self.param_dtype = param_dtype\n        self.share_embeddings_and_output_weights = share_embeddings_and_output_weights\n        self.model_path = self.config.model.path\n        self.use_distributed_optimizer = use_distributed_optimizer\n        self.use_checkpoint_opt_param_scheduler = use_checkpoint_opt_param_scheduler\n        self.bridge = bridge\n        self.provider = provider\n        self.vanilla_bridge = self.provider is None\n        self.peft_cls = peft_cls\n        self.rank = torch.distributed.get_rank()\n        # Megatron-Bridge is Okay to load/save HF checkpoint for value model as well\n        self.use_dist_checkpointing = (\n            use_dist_checkpointing or not self.bridge or (self.vanilla_bridge and self.is_value_model)\n        )\n        self.use_hf_checkpoint = not self.use_dist_checkpointing\n\n        self.weight_saver = None\n        if self.bridge is None:\n            self.weight_saver = get_weight_saver(self.arch)\n\n    def get_rng_state(self, use_dist_ckpt: bool = True, data_parallel_random_init: bool = False):\n        \"\"\"collect rng state across data parallel ranks\"\"\"\n        rng_state = {\n            \"random_rng_state\": random.getstate(),\n            \"np_rng_state\": np.random.get_state(),\n            \"torch_rng_state\": torch.get_rng_state(),\n            \"rng_tracker_states\": tensor_parallel.get_cuda_rng_tracker().get_states(),\n        }\n\n        if get_device_name() != \"cpu\":\n            rng_state[f\"{get_device_name()}_rng_state\"] = get_torch_device().get_rng_state()\n\n        rng_state_list = None\n        if torch.distributed.is_initialized() and mpu.get_data_parallel_world_size() > 1 and data_parallel_random_init:\n            rng_state_list = [None for i in range(mpu.get_data_parallel_world_size())]\n            torch.distributed.all_gather_object(rng_state_list, rng_state, group=mpu.get_data_parallel_group())\n        else:\n            rng_state_list = [rng_state]\n\n        if use_dist_ckpt:\n            pp_rank = mpu.get_pipeline_model_parallel_rank()\n            pp_size = mpu.get_pipeline_model_parallel_world_size()\n            tp_rank = mpu.get_tensor_model_parallel_rank()\n            tp_size = mpu.get_tensor_model_parallel_world_size()\n            rng_state_list = ShardedObject(\n                \"rng_state\",\n                rng_state_list,\n                (pp_size, tp_size),\n                (pp_rank, tp_rank),\n                replica_id=mpu.get_data_parallel_rank(with_context_parallel=True),\n            )\n\n        return rng_state_list\n\n    def get_checkpoint_name(\n        self,\n        checkpoints_path,\n        pipeline_parallel=None,\n        tensor_rank=None,\n        pipeline_rank=None,\n        cp_rank=None,\n        expert_parallel=None,\n        expert_rank=None,\n        return_base_dir=True,\n        basename=\"model.pt\",\n    ):\n        \"\"\"Determine the directory name for this rank's checkpoint.\"\"\"\n        # Use both the tensor and pipeline MP rank.\n        if pipeline_parallel is None:\n            pipeline_parallel = mpu.get_pipeline_model_parallel_world_size() > 1\n        if tensor_rank is None:\n            tensor_rank = mpu.get_tensor_model_parallel_rank()\n        if pipeline_rank is None:\n            pipeline_rank = mpu.get_pipeline_model_parallel_rank()\n        if cp_rank is None:\n            cp_rank = mpu.get_context_parallel_rank()\n        if expert_parallel is None:\n            expert_parallel = mpu.get_expert_model_parallel_world_size() > 1\n        if expert_rank is None:\n            expert_rank = mpu.get_expert_model_parallel_rank()\n\n        # Use both the tensor and pipeline MP rank. If using the distributed\n        # optimizer, then the optimizer's path must additionally include the\n        # data parallel rank.\n\n        # due to the fact that models are identical across cp ranks, cp rank is not used in the checkpoint path\n        if not pipeline_parallel:\n            common_path = os.path.join(checkpoints_path, f\"mp_rank_{tensor_rank:02d}\")\n        else:\n            common_path = os.path.join(checkpoints_path, f\"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}\")\n\n        if expert_parallel:\n            common_path = common_path + f\"_{expert_rank:03d}\"\n\n        os.makedirs(common_path, exist_ok=True)\n\n        if return_base_dir:\n            return common_path\n        return os.path.join(common_path, basename)\n\n    def generate_state_dict(\n        self,\n        generate_model: bool = True,\n        generate_optimizer: bool = True,\n        generate_extra: bool = True,\n        is_loading: bool = False,\n        metadata: dict | None = None,\n    ):\n        # For save dist checkpointing\n        state_dict = {}\n        base_metadata = metadata or self._build_sharded_state_dict_metadata()\n\n        # Should always generate model state dict\n        # All ranks Save Model to reduce memory pressure\n        # Get sharded state dict, notice that state_dict will collect among dp groups, causing memory pressure\n        for vpp_rank, model in enumerate(self.model):\n            if len(self.model) > 1:\n                mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)\n                key = f\"model{vpp_rank}\" if len(self.model) > 1 else \"model\"\n            else:\n                key = \"model\"\n            if hasattr(model, \"module\"):\n                model = model.module\n\n            # GPTModel's sharded_state_dict function when having mtp requires metadata['dp_cp_group']\n            model_metadata = dict(base_metadata)\n            model_metadata[\"dp_cp_group\"] = mpu.get_data_parallel_group(with_context_parallel=True)\n            kwargs = {\"metadata\": model_metadata}\n            state_dict[key] = model.sharded_state_dict(**kwargs)\n\n        # Optimizer State Dict\n        if generate_optimizer:\n            torch.distributed.barrier()\n            sharded_state_dict_kwargs = {\"is_loading\": is_loading}\n            if base_metadata is not None:\n                # https://github.com/NVIDIA/Megatron-LM/blob/core_v0.14.0/megatron/core/optimizer/distrib_optimizer.py#L1109-L1123\n                if mcore_ge_014:\n                    sharded_state_dict_kwargs[\"metadata\"] = base_metadata\n            optimizer_sharded_states = self.optimizer.sharded_state_dict(state_dict, **sharded_state_dict_kwargs)\n            state_dict[\"optimizer\"] = optimizer_sharded_states\n\n            if self.lr_scheduler is not None:\n                lr_state_dict = self.lr_scheduler.state_dict()\n                state_dict[\"lr_scheduler\"] = lr_state_dict\n\n        if not generate_model:\n            state_dict.pop(\"model\", None)\n\n        # RNG States State Dict\n        if generate_extra:\n            torch.distributed.barrier()\n            rng_state = self.get_rng_state()\n            state_dict[\"rng_state\"] = rng_state\n\n        return state_dict\n\n    def _build_sharded_state_dict_metadata(self) -> dict:\n        \"\"\"Builds metadata used for sharded_state_dict versioning.\n\n\n        The whole content metadata is passed to ``sharded_state_dict`` model and optimizer methods\n        and therefore affects only the logic behind sharded_state_dict creation.\n        The content metadata should be minimalistic, ideally flat (or with a single nesting level)\n        and with semantically meaningful flag names (e.g. `distrib_optim_sharding_type`).\n        In particular, a simple integer (or SemVer) versioning flag (e.g. `metadata['version'] = 3.4`)\n        is discouraged, because the metadata serves for all models and optimizers and it's practically\n        impossible to enforce a linearly increasing versioning for this whole space.\n        \"\"\"\n        metadata: dict = {}\n\n        if not mcore_ge_014:\n            # For backward compatibility with Megatron core < v0.14.0\n            if self.use_distributed_optimizer:\n                metadata[\"distrib_optim_sharding_type\"] = \"fully_sharded_model_space\"\n            return metadata\n\n        if self.use_distributed_optimizer:\n            megatron_config = getattr(self.config, self.role, self.config).megatron\n            dist_ckpt_optim_fully_reshardable = megatron_config.dist_ckpt_optim_fully_reshardable\n            distrib_optim_fully_reshardable_mem_efficient = (\n                megatron_config.distrib_optim_fully_reshardable_mem_efficient\n            )\n            if dist_ckpt_optim_fully_reshardable:\n                metadata[\"distrib_optim_sharding_type\"] = \"fully_reshardable\"\n                metadata[\"distrib_optim_fully_reshardable_mem_efficient\"] = (\n                    distrib_optim_fully_reshardable_mem_efficient\n                )\n            else:\n                metadata[\"distrib_optim_sharding_type\"] = \"dp_reshardable\"\n\n        metadata[\"singleton_local_shards\"] = False\n        metadata[\"chained_optim_avoid_prefix\"] = True\n        return metadata\n\n    def load_rng_states(self, rng_states, data_parallel_random_init=False, use_dist_ckpt=True):\n        # access rng_state for data parallel rank\n        if data_parallel_random_init:\n            rng_states = rng_states[mpu.get_data_parallel_rank()]\n        else:\n            rng_states = rng_states[0]\n        random.setstate(rng_states[\"random_rng_state\"])\n        np.random.set_state(rng_states[\"np_rng_state\"])\n        torch.set_rng_state(rng_states[\"torch_rng_state\"])\n\n        if get_device_name() != \"cpu\":\n            get_torch_device().set_rng_state(rng_states[f\"{get_device_name()}_rng_state\"])\n\n        # Check for empty states array\n        if not rng_states[\"rng_tracker_states\"]:\n            raise KeyError\n        tensor_parallel.get_cuda_rng_tracker().set_states(rng_states[\"rng_tracker_states\"])\n\n    def load_checkpoint(self, local_path: str, hdfs_path: str = None, del_local_after_load=False):\n        if local_path is not None:\n            assert os.path.exists(local_path), f\"Checkpoint path {local_path} does not exist.\"\n\n        # For load optimizer dist_ckpt\n        try:\n            import transformer_engine\n\n            torch.serialization.add_safe_globals([torch.optim.AdamW])\n            torch.serialization.add_safe_globals([transformer_engine.pytorch.optimizers.fused_adam.FusedAdam])\n        except Exception:\n            pass\n\n        dist_checkpoint_path = get_dist_checkpoint_path(local_path)\n\n        load_content_metadata = getattr(dist_checkpointing, \"load_content_metadata\", None)\n        if load_content_metadata is None:\n            # For backward compatibility\n            sharded_sd_metadata = None\n        else:\n            sharded_sd_metadata = load_content_metadata(checkpoint_dir=dist_checkpoint_path)\n        if sharded_sd_metadata is None:\n            if self.use_distributed_optimizer:\n                # Backward-compatibility with old checkpoints which don't have content versioning\n                # Can be removed after ending support for MLM optimizer checkpoints with MCore < v0.13\n                # (for MCore v0.13+ checkpoints `sharded_sd_metadata is not None`)\n                sharded_sd_metadata = {\n                    \"distrib_optim_sharding_type\": \"fully_sharded_model_space\",\n                }\n            else:\n                sharded_sd_metadata = self._build_sharded_state_dict_metadata()\n\n        # Get State Dict for loading\n        sharded_state_dict = self.generate_state_dict(\n            self.should_load_model and self.use_dist_checkpointing,\n            self.should_load_optimizer,\n            self.should_load_extra,\n            is_loading=True,\n            metadata=sharded_sd_metadata,\n        )\n        log_with_rank(f\"Generated state dict for loading: {sharded_state_dict.keys()}\", rank=self.rank, logger=logger)\n\n        # Load Dist Checkpointing\n        state_dict = load_dist_checkpointing(\n            sharded_state_dict=sharded_state_dict,\n            ckpt_dir=dist_checkpoint_path,\n        )\n\n        if self.should_load_model and self.use_dist_checkpointing:\n            assert \"model\" in state_dict or any(\n                f\"model{vpp_rank}\" in state_dict for vpp_rank in range(len(self.model))\n            ), f\"Model state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}.\"\n            for vpp_rank, model in enumerate(self.model):\n                if len(self.model) == 1:\n                    model_state_dict = state_dict[\"model\"]\n                else:\n                    assert f\"model{vpp_rank}\" in state_dict, f\"model{vpp_rank} not found in state_dict\"\n                    model_state_dict = state_dict[f\"model{vpp_rank}\"]\n                mpu.set_virtual_pipeline_model_parallel_rank(vpp_rank)\n                self.model[vpp_rank].load_state_dict(model_state_dict)\n            log_with_rank(f\"Loaded sharded model checkpoint from {local_path}\", rank=self.rank, logger=logger)\n\n        # Skip HF checkpoint loading if PEFT is used\n        elif self.should_load_model and self.use_hf_checkpoint and self.peft_cls is None:\n            hf_model_path = get_hf_model_checkpoint_path(local_path)\n            if self.vanilla_bridge:\n                self.bridge.load_weights(self.model, hf_model_path)\n            else:\n                self.bridge.load_hf_weights(self.model, hf_model_path)\n            log_with_rank(f\"Loaded HF model checkpoint from {hf_model_path} with bridge\", rank=self.rank, logger=logger)\n        # Load PEFT adapter checkpoint if available\n        if self.should_load_model and self.peft_cls is not None:\n            adapter_ckpt_path = os.path.join(local_path, \"adapter_checkpoint\")\n            if os.path.exists(adapter_ckpt_path):\n                from verl.utils.megatron_peft_utils import load_adapter_checkpoint\n\n                # TODO: a better format for adapter checkpoint, waiting megatron-bridge support\n\n                load_adapter_checkpoint(\n                    self.model,\n                    adapter_ckpt_path,\n                )\n                log_with_rank(\n                    f\"Loaded adapter checkpoint from {adapter_ckpt_path}\",\n                    rank=self.rank,\n                    logger=logger,\n                )\n            else:\n                log_with_rank(\n                    f\"PEFT config is set but no adapter checkpoint found at {adapter_ckpt_path}\",\n                    rank=self.rank,\n                    logger=logger,\n                )\n\n        if self.should_load_optimizer:\n            assert \"optimizer\" in state_dict, (\n                f\"Optimizer state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}.\"\n            )\n            optimizer_state_dict = state_dict[\"optimizer\"]\n            self.optimizer.load_state_dict(optimizer_state_dict)\n            log_with_rank(f\"Loaded optimizer checkpoint from {local_path}\", rank=self.rank, logger=logger)\n            if self.use_checkpoint_opt_param_scheduler:\n                assert \"lr_scheduler\" in state_dict, (\n                    f\"LR scheduler state dict not found in {state_dict.keys()}. Please check the checkpoint file \"\n                    f\"{local_path}.\"\n                )\n                lr_scheduler_state_dict = state_dict[\"lr_scheduler\"]\n                if self.lr_scheduler is not None:\n                    self.lr_scheduler.load_state_dict(lr_scheduler_state_dict)\n                    log_with_rank(f\"Loaded LR scheduler checkpoint from {local_path}\", rank=self.rank, logger=logger)\n\n        if self.should_load_extra:\n            assert \"rng_state\" in state_dict, (\n                f\"RNG state dict not found in {state_dict.keys()}. Please check the checkpoint file {local_path}.\"\n            )\n            rng_state = state_dict[\"rng_state\"]\n            self.load_rng_states(rng_state)\n            log_with_rank(f\"Loaded RNG states from {local_path}\", rank=self.rank, logger=logger)\n\n        if del_local_after_load:\n            try:\n                os.remove(local_path) if is_non_local(local_path) else None\n            except Exception as e:\n                log_with_rank(\n                    f\"remove local resume ckpt file after loading failed, exception {e} will be ignored\",\n                    rank=self.rank,\n                    logger=logger,\n                )\n\n    def save_checkpoint(self, local_path: str, hdfs_path: str = None, global_step: int = 0, max_ckpt_to_keep=None):\n        # record the previous global step\n        self.previous_global_step = global_step\n\n        if not self.checkpoint_config.async_save:\n            self.ensure_checkpoint_capacity(max_ckpt_to_keep)\n\n        local_path = local_mkdir_safe(local_path)\n        dist_checkpoint_path = get_dist_checkpoint_path(local_path)\n\n        # Note that model weights, optimizer states, and extra states are generated\n        # together in a state dict, we save them in one time\n        if self.use_dist_checkpointing:\n            # Generate state dict for saving\n            sharded_sd_metadata = self._build_sharded_state_dict_metadata()\n            state_dict = self.generate_state_dict(\n                self.should_save_model,\n                self.should_save_optimizer,\n                self.should_save_extra,\n                metadata=sharded_sd_metadata,\n            )\n            log_with_rank(f\"Generated state dict for saving: {state_dict.keys()}\", rank=self.rank, logger=logger)\n            for vpp_rank, model in enumerate(self.model):\n                if len(self.model) > 1:\n                    model_i_keys = state_dict[f\"model{vpp_rank}\"].keys()\n                    log_with_rank(f\"Generated state dict for saving: {model_i_keys}\", rank=self.rank, logger=logger)\n                else:\n                    log_with_rank(\n                        f\"Generated state dict for saving: {state_dict['model'].keys()}\", rank=self.rank, logger=logger\n                    )\n            # Start Async save if enabled\n            async_save_request = save_dist_checkpointing(\n                sharded_state_dict=state_dict,\n                ckpt_path=dist_checkpoint_path,\n                async_save=self.checkpoint_config.async_save,\n                content_metadata=sharded_sd_metadata,\n            )\n\n            # Synchronize all async save requests\n            if not self.checkpoint_config.async_save:\n                assert async_save_request is None, \"Async save request should be None when not using async save.\"\n                torch.distributed.barrier()\n        else:\n            assert self.use_hf_checkpoint, \"When not using distributed checkpointing, use_hf_checkpoint should be True.\"\n            # Generate optimizer and exra state dicts\n            sharded_sd_metadata = self._build_sharded_state_dict_metadata()\n            state_dict = self.generate_state_dict(\n                generate_model=False,\n                generate_optimizer=self.should_save_optimizer,\n                generate_extra=self.should_save_extra,\n                metadata=sharded_sd_metadata,\n            )\n            # Save optimizer and extra states to local path\n            # Start Async save if enabled\n            async_save_request = save_dist_checkpointing(\n                sharded_state_dict=state_dict,\n                ckpt_path=dist_checkpoint_path,\n                async_save=self.checkpoint_config.async_save,\n                content_metadata=sharded_sd_metadata,\n            )\n\n            # Synchronize all async save requests\n            if not self.checkpoint_config.async_save:\n                assert async_save_request is None, \"Async save request should be None when not using async save.\"\n                torch.distributed.barrier()\n\n        if self.should_save_model:\n            # Save adapter-only checkpoint if PEFT is enabled\n            if self.peft_cls is not None:\n                from verl.utils.megatron_peft_utils import save_adapter_checkpoint\n\n                adapter_ckpt_path = os.path.join(local_path, \"adapter_checkpoint\")\n\n                # Save adapter weights only (much smaller than full model)\n                save_adapter_checkpoint(\n                    self.model,\n                    adapter_ckpt_path,\n                    self.rank,\n                )\n\n                log_with_rank(\n                    f\"Saved adapter-only checkpoint to {adapter_ckpt_path}\",\n                    rank=self.rank,\n                    logger=logger,\n                    log_only_rank_0=True,\n                )\n            elif self.use_hf_checkpoint:\n                # Use mbridge to save HF model checkpoint\n                log_with_rank(f\"Saving HF model checkpoint to {local_path} with bridge\", rank=self.rank, logger=logger)\n                hf_ckpt_path = get_hf_model_checkpoint_path(local_path)\n                if self.vanilla_bridge:\n                    extended_args = {}\n                    mbridge_config = getattr(self.checkpoint_config, \"mbridge_config\", None) or {}\n                    for sig in inspect.signature(self.bridge.save_weights).parameters:\n                        if sig == \"weights_path\" or sig == \"models\":\n                            continue\n                        if sig in mbridge_config:\n                            extended_args[sig] = mbridge_config[sig]\n                    self.bridge.save_weights(self.model, hf_ckpt_path, **extended_args)\n                else:\n                    self.bridge.save_hf_weights(self.model, hf_ckpt_path)\n\n                log_with_rank(f\"Saved bridge checkpoint to {hf_ckpt_path}\", rank=self.rank, logger=logger)\n\n            # Only rank 0 saves the hf config and tokenizer to huggingface path\n            # No matter whether we save hf model or not\n            if self.rank == 0:\n                # Save tokenizer\n                hf_config_tokenizer_path = get_hf_model_checkpoint_path(local_path)\n                if self.processing_class is not None:\n                    self.processing_class.save_pretrained(hf_config_tokenizer_path)\n                # Save huggingface config\n                self.hf_config.save_pretrained(hf_config_tokenizer_path)\n                if hasattr(self.hf_config, \"name_or_path\") and self.hf_config.name_or_path:\n                    try:\n                        generation_config = GenerationConfig.from_pretrained(self.hf_config.name_or_path)\n                        generation_config.save_pretrained(hf_config_tokenizer_path)\n                    except Exception:\n                        # if the generation config isn't available, we don't save it\n                        pass\n                log_with_rank(\n                    f\"Saved Huggingface config and tokenizer to {hf_config_tokenizer_path}\",\n                    rank=self.rank,\n                    logger=logger,\n                    log_only_rank_0=True,\n                )\n\n        if self.should_save_extra:\n            if self.rank == 0:\n                # Save transformer config\n                print(self.transformer_config)\n                bypass_keys = [\n                    \"finalize_model_grads_func\",\n                    \"grad_scale_func\",\n                    \"no_sync_func\",\n                    \"grad_sync_func\",\n                    \"param_sync_func\",\n                    \"generation_config\",\n                    \"_pg_collection\",\n                ]\n                backup = {}\n                for k in bypass_keys:\n                    if hasattr(self.transformer_config, k):\n                        backup[k] = getattr(self.transformer_config, k, None)\n                        delattr(self.transformer_config, k)\n                transformer_config_dict = asdict(self.transformer_config)\n                for k in backup:\n                    setattr(self.transformer_config, k, backup[k])\n                to_convert_types = {torch.dtype: str, AttnBackend: str}\n                ignore_types = [Callable]\n                pop_keys = []\n                for key, value in transformer_config_dict.items():\n                    if type(value) in to_convert_types:\n                        transformer_config_dict[key] = to_convert_types[type(value)](value)\n                    if type(value) in ignore_types:\n                        pop_keys.append(key)\n                    if callable(value):\n                        pop_keys.append(key)\n                for key in pop_keys:\n                    transformer_config_dict.pop(key)\n                transformer_config_path = get_transformer_config_checkpoint_path(local_path)\n                with open(transformer_config_path, \"w\") as f:\n                    json.dump(transformer_config_dict, f, indent=2)\n\n        if self.should_save_hf_model and not self.use_hf_checkpoint:\n            # wait for everyone to dump to local\n            if self.bridge is not None:\n                hf_model_ckpt_path = get_hf_model_checkpoint_path(local_path)\n                if self.vanilla_bridge:\n                    extended_args = {}\n                    mbridge_config = getattr(self.checkpoint_config, \"mbridge_config\", None) or {}\n                    for sig in inspect.signature(self.bridge.save_weights).parameters:\n                        if sig == \"weights_path\" or sig == \"models\":\n                            continue\n                        if sig in mbridge_config:\n                            extended_args[sig] = mbridge_config[sig]\n                    self.bridge.save_weights(self.model, hf_model_ckpt_path, **extended_args)\n                else:\n                    self.bridge.save_hf_weights(self.model, hf_model_ckpt_path)\n            else:\n                state_dict = self.weight_saver(\n                    self.model,\n                    self.hf_config,\n                    dtype=self.param_dtype,\n                    is_value_model=self.is_value_model,\n                    tie_word_embeddings=self.share_embeddings_and_output_weights,\n                )\n\n                torch.distributed.barrier()\n                if self.rank == 0:\n                    hf_model_ckpt_path = get_hf_model_checkpoint_path(local_path)\n                    import warnings\n\n                    from accelerate import init_empty_weights\n\n                    with init_empty_weights(), warnings.catch_warnings():\n                        warnings.simplefilter(\"ignore\")\n                        if \"mistral7b-rm\" in self.config.model.path:\n                            from transformers import MistralForSequenceClassification\n\n                            model = MistralForSequenceClassification.from_pretrained(\n                                self.config.model.path\n                            )  # use score head instead of lm_head\n                            state_dict[\"score.weight\"] = state_dict[\"score.weight\"]\n                        else:\n                            from transformers import AutoModelForCausalLM\n\n                            model = AutoModelForCausalLM.from_pretrained(self.config.model.path, torch_dtype=\"auto\")\n                    model.save_pretrained(hf_model_ckpt_path, state_dict=state_dict)\n                    log_with_rank(\n                        f\"Saved Huggingface config and tokenizer to {hf_model_ckpt_path}\",\n                        rank=self.rank,\n                        logger=logger,\n                        log_only_rank_0=True,\n                    )\n\n                    if hdfs_path is not None:\n                        log_with_rank(\n                            f\"Uploading checkpoint to {hdfs_path}\", rank=self.rank, logger=logger, log_only_rank_0=True\n                        )\n                        from verl.utils import hdfs_io\n\n                        hdfs_io.makedirs(hdfs_path, exist_ok=True)\n                        hdfs_io.copy(src=hf_model_ckpt_path, dst=hdfs_path, dirs_exist_ok=True)\n                        log_with_rank(\n                            f\"HDFS checkpoint uploaded to {hdfs_path}\",\n                            rank=self.rank,\n                            logger=logger,\n                            log_only_rank_0=True,\n                        )\n\n        def finalize_save_fn():\n            # Rank 0 uploads checkpoint to HDFS if hdfs_path is provided\n            log_with_rank(\n                f\"Dist checkpointing save completed for {dist_checkpoint_path}\", rank=self.rank, logger=logger\n            )\n            if self.rank == 0:\n                if hdfs_path is not None:\n                    log_with_rank(f\"Uploading checkpoint to {hdfs_path}\", rank=self.rank, logger=logger)\n                    from verl.utils import hdfs_io\n\n                    hdfs_io.makedirs(hdfs_path, exist_ok=True)\n                    hdfs_io.copy(src=dist_checkpoint_path, dst=hdfs_path, dirs_exist_ok=True)\n                    hdfs_io.copy(src=hf_config_tokenizer_path, dst=hdfs_path, dirs_exist_ok=True)\n\n            # update latest_checkpointed_iteration.txt when async_save is True\n            if self.checkpoint_config.async_save and self.rank == 0:\n                log_with_rank(\n                    f\"Update latest_checkpointed_iteration.txt to step {global_step}\",\n                    rank=self.rank,\n                    logger=logger,\n                )\n                local_latest_checkpointed_iteration = os.path.join(\n                    os.path.dirname(os.path.dirname(local_path)), \"latest_checkpointed_iteration.txt\"\n                )\n                with open(local_latest_checkpointed_iteration, \"w\") as f:\n                    f.write(str(global_step))\n\n            self.register_checkpoint(local_path, max_ckpt_to_keep)\n\n        if self.checkpoint_config.async_save:\n            assert async_save_request is not None, \"Async save request should not be None when using async save.\"\n            async_save_request.add_finalize_fn(finalize_save_fn)\n            from megatron.core.dist_checkpointing.strategies.base import async_calls\n\n            async_calls.schedule_async_request(async_save_request)\n        else:\n            finalize_save_fn()\n"
  },
  {
    "path": "verl/utils/config.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import is_dataclass\nfrom typing import Any, Optional\n\nfrom omegaconf import DictConfig, ListConfig, OmegaConf\n\n__all__ = [\"omega_conf_to_dataclass\", \"validate_config\"]\n\n\ndef omega_conf_to_dataclass(config: DictConfig | dict, dataclass_type: Optional[type[Any]] = None) -> Any:\n    \"\"\"\n    Convert an OmegaConf DictConfig to a dataclass.\n\n    Args:\n        config: The OmegaConf DictConfig or dict to convert.\n        dataclass_type: The dataclass type to convert to. When dataclass_type is None,\n            the DictConfig must contain _target_ to be instantiated via hydra.instantiate API.\n\n    Returns:\n        The dataclass instance.\n    \"\"\"\n    # Got an empty config\n    if not config:\n        return dataclass_type if dataclass_type is None else dataclass_type()\n    # Got an object\n    if not isinstance(config, DictConfig | ListConfig | dict | list):\n        return config\n\n    if dataclass_type is None:\n        assert \"_target_\" in config, (\n            \"When dataclass_type is not provided, config must contain _target_. \"\n            \"See trainer/config/ppo_trainer.yaml algorithm section for an example. \"\n            f\"Got config: {config}\"\n        )\n        from hydra.utils import instantiate\n\n        return instantiate(config, _convert_=\"partial\")\n\n    if not is_dataclass(dataclass_type):\n        raise ValueError(f\"{dataclass_type} must be a dataclass\")\n    cfg = OmegaConf.create(config)  # in case it's a dict\n    # pop _target_ to avoid hydra instantiate error, as most dataclass do not have _target_\n    # Updated (vermouth1992) We add _target_ to BaseConfig so that it is compatible.\n    # Otherwise, this code path can't support recursive instantiation.\n    # if \"_target_\" in cfg:\n    #     cfg.pop(\"_target_\")\n    cfg_from_dataclass = OmegaConf.structured(dataclass_type)\n    # let cfg override the existing vals in `cfg_from_dataclass`\n    cfg_merged = OmegaConf.merge(cfg_from_dataclass, cfg)\n    # now convert to `dataclass_type`\n    config_object = OmegaConf.to_object(cfg_merged)\n    return config_object\n\n\ndef update_dict_with_config(dictionary: dict, config: DictConfig):\n    for key in dictionary:\n        if hasattr(config, key):\n            dictionary[key] = getattr(config, key)\n\n\ndef validate_config(\n    config: DictConfig,\n    use_reference_policy: bool,\n    use_critic: bool,\n) -> None:\n    \"\"\"Validate an OmegaConf DictConfig.\n\n    Args:\n        config (DictConfig): The OmegaConf DictConfig to validate.\n        use_reference_policy (bool): is ref policy needed\n        use_critic (bool): is critic needed\n    \"\"\"\n    # number of GPUs total\n    n_gpus = config.trainer.n_gpus_per_node * config.trainer.nnodes\n\n    if not config.actor_rollout_ref.actor.use_dynamic_bsz:\n        if config.actor_rollout_ref.actor.strategy == \"megatron\":\n            model_parallel_size = (\n                config.actor_rollout_ref.actor.megatron.tensor_model_parallel_size\n                * config.actor_rollout_ref.actor.megatron.pipeline_model_parallel_size\n            )\n            assert (\n                n_gpus % (model_parallel_size * config.actor_rollout_ref.actor.megatron.context_parallel_size) == 0\n            ), (\n                f\"n_gpus ({n_gpus}) must be divisible by model_parallel_size ({model_parallel_size}) times \"\n                f\"context_parallel_size ({config.actor_rollout_ref.actor.megatron.context_parallel_size})\"\n            )\n            megatron_dp = n_gpus // (\n                model_parallel_size * config.actor_rollout_ref.actor.megatron.context_parallel_size\n            )\n            minimal_bsz = megatron_dp * config.actor_rollout_ref.actor.ppo_micro_batch_size_per_gpu\n        else:\n            minimal_bsz = n_gpus\n\n        # 1. Check total batch size for data correctness\n        real_train_batch_size = config.data.train_batch_size * config.actor_rollout_ref.rollout.n\n        assert real_train_batch_size % minimal_bsz == 0, (\n            f\"real_train_batch_size ({real_train_batch_size}) must be divisible by minimal possible batch size \"\n            f\"({minimal_bsz})\"\n        )\n\n    # A helper function to check \"micro_batch_size\" vs \"micro_batch_size_per_gpu\"\n    # We throw an error if the user sets both. The new convention is \"..._micro_batch_size_per_gpu\".\n    def check_mutually_exclusive(mbs, mbs_per_gpu, name: str):\n        \"\"\"Validate mutually exclusive micro batch size configuration options.\n\n        Ensures that users don't set both deprecated micro_batch_size and\n        the new micro_batch_size_per_gpu parameters simultaneously.\n\n        Args:\n            mbs: Deprecated micro batch size parameter value.\n            mbs_per_gpu: New micro batch size per GPU parameter value.\n            name (str): Configuration section name for error messages.\n\n        Raises:\n            ValueError: If both parameters are set or neither is set.\n        \"\"\"\n        settings = {\n            \"actor_rollout_ref.ref\": \"log_prob_micro_batch_size\",\n            \"actor_rollout_ref.rollout\": \"log_prob_micro_batch_size\",\n        }\n\n        if name in settings:\n            param = settings[name]\n            param_per_gpu = f\"{param}_per_gpu\"\n\n            if mbs is None and mbs_per_gpu is None:\n                raise ValueError(f\"[{name}] Please set at least one of '{name}.{param}' or '{name}.{param_per_gpu}'.\")\n\n            if mbs is not None and mbs_per_gpu is not None:\n                raise ValueError(\n                    f\"[{name}] You have set both '{name}.{param}' AND '{name}.{param_per_gpu}'. Please remove \"\n                    f\"'{name}.{param}' because only '*_{param_per_gpu}' is supported (the former is deprecated).\"\n                )\n\n    # Actor validation done in ActorConfig.__post_init__ and validate()\n    actor_config = omega_conf_to_dataclass(config.actor_rollout_ref.actor)\n    actor_config.validate(n_gpus, config.data.train_batch_size, config.actor_rollout_ref.model)\n\n    if not config.actor_rollout_ref.actor.use_dynamic_bsz:\n        if use_reference_policy:\n            # reference: log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu\n            check_mutually_exclusive(\n                config.actor_rollout_ref.ref.log_prob_micro_batch_size,\n                config.actor_rollout_ref.ref.log_prob_micro_batch_size_per_gpu,\n                \"actor_rollout_ref.ref\",\n            )\n\n        #  The rollout section also has log_prob_micro_batch_size vs. log_prob_micro_batch_size_per_gpu\n        check_mutually_exclusive(\n            config.actor_rollout_ref.rollout.log_prob_micro_batch_size,\n            config.actor_rollout_ref.rollout.log_prob_micro_batch_size_per_gpu,\n            \"actor_rollout_ref.rollout\",\n        )\n\n    if config.algorithm.use_kl_in_reward and config.actor_rollout_ref.actor.use_kl_loss:\n        print(\"NOTICE: You have both enabled in-reward kl and kl loss.\")\n\n    # critic\n    if use_critic:\n        critic_config = omega_conf_to_dataclass(config.critic)\n        critic_config.validate(n_gpus, config.data.train_batch_size)\n\n    if config.data.get(\"val_batch_size\", None) is not None:\n        print(\n            \"WARNING: val_batch_size is deprecated.\"\n            + \" Validation datasets are sent to inference engines as a whole batch,\"\n            + \" which will schedule the memory themselves.\"\n        )\n\n    # check eval config\n    if config.actor_rollout_ref.rollout.val_kwargs.do_sample:\n        assert config.actor_rollout_ref.rollout.temperature > 0, (\n            \"validation gen temperature should be greater than 0 when enabling do_sample\"\n        )\n\n    # check LoRA rank in vLLM\n    lora_config = config.actor_rollout_ref.model.get(\"lora\", {})\n    lora_rank = lora_config.get(\"rank\", 0)\n    if lora_rank <= 0:\n        lora_rank = config.actor_rollout_ref.model.get(\"lora_rank\", 0)\n    if lora_config.get(\"merge\", False):\n        lora_rank = 0\n    if lora_rank > 0 and config.actor_rollout_ref.rollout.name == \"vllm\":\n        from verl.workers.rollout.vllm_rollout.utils import get_vllm_max_lora_rank\n\n        get_vllm_max_lora_rank(lora_rank)\n\n    print(\"[validate_config] All configuration checks passed successfully!\")\n"
  },
  {
    "path": "verl/utils/dataset/README.md",
    "content": "# Dataset Format\n## RLHF dataset\nWe combine all the data sources into a single parquet files. We directly organize the prompt into the chat format so that multi-turn chats can be easily incorporated. In the prompt, we may add instruction following texts to guide the model output the answers in a particular format so that we can extract the answers.\n\nMath problems\n```json\n{\n    \"data_source\": \"openai/gsm8k\",\n    \"prompt\": [{\"role\": \"user\", \"content\": \"Natalia sold clips to 48 of her friends in April, and then she sold half as many clips in May. How many clips did Natalia sell altogether in April and May? Let's think step by step and output the final answer after \\\"####\\\"\"}],\n    \"ability\": \"math\",\n    \"reward_model\": {\n        \"style\": \"rule\",\n        \"ground_truth\": [\"72\"]\n    },\n}\n```\n"
  },
  {
    "path": "verl/utils/dataset/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .rl_dataset import RLHFDataset\nfrom .rm_dataset import RMDataset\n\n__all__ = [\"RLHFDataset\", \"RMDataset\"]\n"
  },
  {
    "path": "verl/utils/dataset/dataset_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom enum import Enum\n\nimport torch\nfrom tensordict.tensorclass import NonTensorData\n\n\nclass DatasetPadMode(str, Enum):\n    \"\"\"Padding mode for dataset\"\"\"\n\n    RIGHT = \"right\"\n    LEFT_RIGHT = \"left_right\"\n    NO_PADDING = \"no_padding\"\n\n\nclass SFTTensorCollator:\n    \"\"\"\n    A custom collate_fn that handles batching of sequences.\n    1. for variable-length sequences, convert them into NestedTensors.\n    2. for fixed-length sequences, use default_collate.\n    \"\"\"\n\n    def __init__(self, pad_mode: DatasetPadMode = DatasetPadMode.LEFT_RIGHT):\n        self.pad_mode = pad_mode\n\n    def __call__(self, batch: list[dict[str, any]]) -> dict[str, any]:\n        if self.pad_mode == DatasetPadMode.NO_PADDING:\n            return self.collate_variable_batch(batch)\n        elif self.pad_mode in [DatasetPadMode.RIGHT, DatasetPadMode.LEFT_RIGHT]:\n            from torch.utils.data import default_collate\n\n            return default_collate(batch)\n        else:\n            raise NotImplementedError(f\"pad_mode {self.pad_mode} not implemented\")\n\n    def collate_variable_batch(self, batch: list[dict[str, any]]) -> dict[str, any]:\n        \"\"\"\n        Collates a list of samples into a single batch.\n\n        Args:\n            batch: A list of dictionary samples from the dataset.\n\n        Returns:\n            A dictionary representing the batched data, with variable-length\n            sequences converted to NestedTensors.\n        \"\"\"\n\n        final_batch = {}\n\n        tensor_keys = set().union(*(d.keys() for d in batch))\n\n        # Handle tensor values by creating a NestedTensor.\n        for key in tensor_keys:\n            if isinstance(batch[0][key], torch.Tensor):\n                tensors = [item[key] for item in batch]\n                final_batch[key] = torch.nested.as_nested_tensor(tensors, layout=torch.jagged)\n            else:\n                tensors = [NonTensorData(item.get(key)) for item in batch]\n                final_batch[key] = torch.stack(tensors, dim=0)\n\n        return final_batch\n"
  },
  {
    "path": "verl/utils/dataset/multiturn_sft_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n\n#     http://www.apache.org/licenses/LICENSE-2.0\n\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nMulti-turn SFT dataset that supports training on conversation data with multiple turns\n\"\"\"\n\nimport logging\nimport os\nimport re\nfrom functools import wraps\nfrom typing import Any, Optional\n\nimport numpy as np\nimport pandas as pd\nimport torch\nimport torch.nn.functional as F\nfrom omegaconf import DictConfig, ListConfig\nfrom torch.utils.data import Dataset\nfrom transformers import PreTrainedTokenizer, ProcessorMixin\n\nfrom verl.models.transformers.qwen2_vl import get_rope_index\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.chat_template import apply_chat_template, extract_system_prompt_and_generation\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.dataset.vision_utils import process_image, process_video\nfrom verl.utils.fs import copy_local_path_from_hdfs\nfrom verl.utils.py_functional import convert_nested_value_to_list_recursive\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef once(func):\n    \"\"\"Decorator to ensure a function runs only once. Subsequent calls do nothing.\"\"\"\n\n    @wraps(func)\n    def wrapper(*args, **kwargs):\n        if not hasattr(wrapper, \"called\"):\n            wrapper.called = True\n            return func(*args, **kwargs)\n\n    return wrapper\n\n\n@once\ndef print_assembled_message(tokenizer, message_list, input_ids, loss_mask, attn_mask, tools):\n    \"\"\"\n    Print the message after applying the chat template\n    \"\"\"\n\n    tokenized = tokenizer.apply_chat_template(message_list, add_generation_prompt=False, tokenize=False, tools=tools)\n    sep = \"\\n\\n\"\n    str = f\"tokenized entire message:\\n{tokenized}\"\n    str += sep\n    decoded_ids = input_ids.tolist() if hasattr(input_ids, \"tolist\") else input_ids\n    str += f\"tokenized seperately    :\\n{tokenizer.decode(decoded_ids)}\"\n\n    logger.debug(str)\n\n\nclass MultiTurnSFTDataset(Dataset):\n    \"\"\"\n    Dataset for multi-turn conversations where each assistant response should be trained\n\n    Args:\n        data_files (str or list): Path(s) to Parquet file(s).\n        tokenizer (PreTrainedTokenizer): For the tokenization of text to token IDs.\n        config (DictConfig): Options like cache_dir, prompt_key, max_prompt_length, truncation, etc.\n        processor (ProcessorMixin, optional): Multimodal preprocessor for images/videos.\n        max_samples (int, optional): Limit the number of samples. Defaults to -1 (use all).\n    \"\"\"\n\n    def __init__(\n        self,\n        parquet_files: str | list[str],\n        tokenizer: PreTrainedTokenizer,\n        config: DictConfig,\n        processor: Optional[ProcessorMixin] = None,\n        max_samples: int = -1,\n    ):\n        # Set defaults and extract parameters from config if provided\n        config = config or {}\n        self.pad_mode = config.get(\"pad_mode\", \"right\")\n        assert self.pad_mode in [\"right\", \"no_padding\"], (\n            f\"Expect pad_mode to be 'right' or 'no_padding'. Got {self.pad_mode}\"\n        )\n        self.truncation = config.get(\"truncation\", \"error\")\n        # for right padding\n        self.max_length = config.get(\"max_length\", 1024)\n        # Get messages_key from the new multiturn config structure\n        self.messages_key = config.get(\"messages_key\", \"messages\")\n        self.image_key = config.get(\"image_key\", \"images\")\n        self.video_key = config.get(\"video_key\", \"videos\")\n        self.image_patch_size = config.get(\n            \"image_patch_size\", processor.image_processor.patch_size if processor else None\n        )\n        self.tools_key = config.get(\"tools_key\", \"tools\")\n        self.enable_thinking_key = config.get(\"enable_thinking_key\", \"enable_thinking\")\n        self.enable_thinking_default = config.get(\"enable_thinking_default\", None)\n        self.apply_chat_template_kwargs = config.get(\"apply_chat_template_kwargs\", {})\n        self.shuffle = config.get(\"shuffle\", False)\n        self.seed = config.get(\"seed\")\n        self.max_samples = max_samples\n        self.ignore_input_ids_mismatch = config.get(\"ignore_input_ids_mismatch\", False)\n        assert self.truncation in [\"error\", \"left\", \"right\"]\n\n        if not isinstance(parquet_files, list | ListConfig):\n            parquet_files = [parquet_files]\n\n        self.parquet_files = parquet_files\n        if isinstance(tokenizer, str):\n            tokenizer = hf_tokenizer(tokenizer)\n        self.tokenizer: PreTrainedTokenizer = tokenizer\n        self.processor = processor\n\n        self._download()\n        self._read_files_and_process()\n\n    def _download(self):\n        for i, parquet_file in enumerate(self.parquet_files):\n            self.parquet_files[i] = copy_local_path_from_hdfs(parquet_file, verbose=True)\n\n    def _read_files_and_process(self):\n        def series_to_item(ls):\n            import numpy\n            import pandas\n\n            while isinstance(ls, pandas.core.series.Series | numpy.ndarray) and len(ls) == 1:\n                ls = ls[0]\n            return ls\n\n        dataframes = []\n        for parquet_file in self.parquet_files:\n            # default loader loads some list as np.ndarray, which fails the tokenizer\n            dataframe = pd.read_parquet(parquet_file, dtype_backend=\"pyarrow\")\n            dataframes.append(dataframe)\n        self.dataframe = pd.concat(dataframes)\n\n        total = len(self.dataframe)\n        print(f\"dataset len: {len(self.dataframe)}\")\n\n        if self.max_samples > 0 and self.max_samples < total:\n            if self.shuffle:\n                rngs_args = (self.seed,) if self.seed is not None else ()\n                rng = np.random.default_rng(*rngs_args)\n                indices = rng.choice(total, size=self.max_samples, replace=False)\n            else:\n                indices = np.arange(self.max_samples)\n            self.dataframe = self.dataframe.iloc[indices.tolist()]\n            print(f\"selected {self.max_samples} random samples out of {total}\")\n\n        # Extract messages list from dataframe\n        self.messages = self.dataframe[self.messages_key].apply(convert_nested_value_to_list_recursive).tolist()\n\n        # Extract tools list from dataframe\n        if self.tools_key in self.dataframe.columns:\n            self.tools = self.dataframe[self.tools_key].apply(convert_nested_value_to_list_recursive).tolist()\n        else:\n            self.tools = None\n        # Extract enable_thinking list from dataframe\n        if self.enable_thinking_key in self.dataframe.columns:\n            self.enable_thinking = self.dataframe[self.enable_thinking_key].tolist()\n        else:\n            self.enable_thinking = None\n\n        # system prompt: <|im_start|>system\\nYou are a helpful assistant.<|im_end|>\\n\n        # generation prompt: <|im_start|>assistant\\n\n        self.system_prompt, self.generation_prompt = extract_system_prompt_and_generation(self.tokenizer)\n\n    def __len__(self):\n        return len(self.messages)\n\n    def _process_single_message(\n        self,\n        index: int,\n        message: dict[str, Any],\n        full_message: list,\n        tools: Optional[list[dict[str, Any]]] = None,\n        enable_thinking: Optional[bool] = None,\n    ) -> tuple[list[int], list[int], list[int]]:\n        \"\"\"\n        Process a single message and return its tokenized representation.\n\n        Args:\n            index: turn index in the conversation\n            message: A single message dictionary\n            images: List of images to be used\n            videos: List of videos to be used\n            tools: List of tools to be used\n            enable_thinking: Whether to enable thinking mode\n\n        Returns:\n            Tuple of (input_ids, loss_mask, attention_mask, dict[str, torch.Tensor])\n        \"\"\"\n        processor = self.processor if self.processor is not None else self.tokenizer\n        apply_chat_template_kwargs = {**self.apply_chat_template_kwargs}\n        if enable_thinking is not None:\n            apply_chat_template_kwargs[\"enable_thinking\"] = enable_thinking\n\n        inputs = apply_chat_template(\n            processor,\n            messages=[message],\n            tools=tools,\n            add_generation_prompt=False,\n            tokenize=True,\n            return_dict=True,\n            return_tensors=\"pt\",\n            **apply_chat_template_kwargs,\n        )\n\n        inputs = dict(inputs)\n        input_ids = inputs.pop(\"input_ids\")[0]\n        attention_mask = inputs.pop(\"attention_mask\")[0]\n\n        # remove system prompt if exists\n        if index != 0 and message[\"role\"] != \"system\":\n            input_ids = input_ids[len(self.system_prompt) :]\n            attention_mask = attention_mask[len(self.system_prompt) :]\n\n        if message[\"role\"] == \"assistant\":\n            loss_mask = torch.ones_like(attention_mask)\n            # mask out generation prompt if assistant message\n            loss_mask[: len(self.generation_prompt)] = 0\n        else:\n            loss_mask = torch.zeros_like(attention_mask)\n\n        return input_ids, loss_mask, attention_mask, inputs\n\n    def _build_messages(self, example: dict):\n        \"\"\"Replace <image> and <video> placeholder in messages with corresponding image and video\n        which is required by processor.apply_chat_template.\n        - <image>: {\"type\": \"image\", \"image\": image}\n        - <video>: {\"type\": \"video\", \"video\": video}\n\n        Args:\n            example: Row dictionary from dataframe.\n\n        Returns:\n            messages: List of messages with replaced placeholder.\n        \"\"\"\n        messages: list = example[self.messages_key]\n        images = example[self.image_key] if self.image_key in example else []\n        videos = example[self.video_key] if self.video_key in example else []\n\n        image_offset, video_offset = 0, 0\n        for message in messages:\n            content = message[\"content\"]\n            if not isinstance(content, str):\n                continue\n\n            if self.image_key not in example and self.video_key not in example:\n                if self.processor is not None:\n                    message[\"content\"] = [{\"type\": \"text\", \"text\": content}]\n                continue\n            assert self.processor is not None, \"processor is needed to process image and video\"\n\n            content_list = []\n            segments = re.split(\"(<image>|<video>)\", content)\n            segments = [item for item in segments if item != \"\"]\n            for segment in segments:\n                if segment == \"<image>\":\n                    image = process_image(images[image_offset], image_patch_size=self.image_patch_size)\n                    content_list.append({\"type\": \"image\", \"image\": image})\n                    image_offset += 1\n                elif segment == \"<video>\":\n                    video = process_video(videos[video_offset], image_patch_size=self.image_patch_size)\n                    content_list.append({\"type\": \"video\", \"video\": video})\n                    video_offset += 1\n                else:\n                    content_list.append({\"type\": \"text\", \"text\": segment})\n            message[\"content\"] = content_list\n\n        assert image_offset == len(images), f\"image_offset {image_offset} != len(images) {len(images)}\"\n        assert video_offset == len(videos), f\"video_offset {video_offset} != len(videos) {len(videos)}\"\n        return messages\n\n    def __getitem__(self, item):\n        row_dict: dict = self.dataframe.iloc[item].to_dict()\n        messages = self._build_messages(row_dict)\n        tools = self.tools[item] if self.tools is not None else None\n        enable_thinking = (\n            self.enable_thinking[item] if self.enable_thinking is not None else self.enable_thinking_default\n        )\n\n        # 1. tokenize each message\n        input_ids, loss_mask, attention_mask, multi_modal_inputs = [], [], [], {}\n        for i, message in enumerate(messages):\n            _input_ids, _loss_mask, _attention_mask, _inputs = self._process_single_message(\n                index=i,\n                message=message,\n                full_message=messages,\n                tools=tools if i == 0 else None,\n                enable_thinking=enable_thinking,\n            )\n            input_ids.append(_input_ids)\n            loss_mask.append(_loss_mask)\n            attention_mask.append(_attention_mask)\n            for k, v in _inputs.items():\n                multi_modal_inputs.setdefault(k, []).append(v)\n\n        input_ids = torch.cat(input_ids, dim=0)\n        loss_mask = torch.cat(loss_mask, dim=0)\n        attention_mask = torch.cat(attention_mask, dim=0)\n        assert input_ids.shape == loss_mask.shape == attention_mask.shape, (\n            f\"Shape mismatch: {input_ids.shape}, {loss_mask.shape}, {attention_mask.shape}\"\n        )\n\n        print_assembled_message(self.tokenizer, messages, input_ids, loss_mask, attention_mask, tools)\n        self.sanity_check(input_ids, messages, tools, enable_thinking)\n\n        # Since the tokenizer may return user-customized results, we need to filter out inconsistent tensor shapes\n        keys_to_remove = []\n        for k, v in multi_modal_inputs.items():\n            if len(v) > 0 and v[0] is not None and isinstance(v[0], torch.Tensor):\n                # Check if all tensors in the list have the same shape\n                first_shape = v[0].shape[1:]\n                if not all(tensor.shape[1:] == first_shape for tensor in v):\n                    keys_to_remove.append(k)\n\n        for k in keys_to_remove:\n            del multi_modal_inputs[k]\n\n        for k, v in multi_modal_inputs.items():\n            multi_modal_inputs[k] = torch.concat(v, dim=0)\n\n        # 2. handle position_ids for Qwen-VL series models\n        if self.processor is not None and \"Qwen2VLImageProcessor\" in self.processor.image_processor.__class__.__name__:\n            image_grid_thw = multi_modal_inputs.get(\"image_grid_thw\", None)\n            video_grid_thw = multi_modal_inputs.get(\"video_grid_thw\", None)\n            second_per_grid_ts = multi_modal_inputs.get(\"second_per_grid_ts\", None)\n\n            vision_position_ids = get_rope_index(\n                self.processor,\n                input_ids=input_ids,\n                image_grid_thw=image_grid_thw,\n                video_grid_thw=video_grid_thw,\n                second_per_grid_ts=second_per_grid_ts,\n                attention_mask=attention_mask,\n            )  # (3, seq_len)\n            text_position_ids = torch.arange(input_ids.shape[0], dtype=torch.long).unsqueeze(0)  # (1, seq_len)\n            position_ids = torch.cat((text_position_ids, vision_position_ids), dim=0)  # (4, seq_length)\n        else:\n            position_ids = torch.arange(input_ids.shape[0], dtype=torch.long)  # (seq_len,)\n\n        # 3. handle padding\n        sequence_length = input_ids.shape[0]\n        # Handle sequence length\n        if self.pad_mode == DatasetPadMode.RIGHT:\n            if sequence_length < self.max_length:\n                # Pad sequences\n                pad_token_id = self.tokenizer.pad_token_id if self.tokenizer.pad_token_id is not None else 0\n                padded_input_ids = torch.full((self.max_length - sequence_length,), pad_token_id, dtype=input_ids.dtype)\n                padded_attention_mask = torch.zeros((self.max_length - sequence_length,), dtype=attention_mask.dtype)\n                padded_loss_mask = torch.zeros((self.max_length - sequence_length,), dtype=loss_mask.dtype)\n\n                input_ids = torch.cat((input_ids, padded_input_ids))\n                attention_mask = torch.cat((attention_mask, padded_attention_mask))\n                loss_mask = torch.cat((loss_mask, padded_loss_mask))\n                position_ids = F.pad(position_ids, (0, self.max_length - sequence_length), value=0)\n            elif sequence_length > self.max_length:\n                if self.truncation == \"left\":\n                    input_ids = input_ids[-self.max_length :]\n                    attention_mask = attention_mask[-self.max_length :]\n                    loss_mask = loss_mask[-self.max_length :]\n                    position_ids = position_ids[..., -self.max_length :]\n                elif self.truncation == \"right\":\n                    input_ids = input_ids[: self.max_length]\n                    attention_mask = attention_mask[: self.max_length]\n                    loss_mask = loss_mask[: self.max_length]\n                    position_ids = position_ids[..., : self.max_length]\n                elif self.truncation == \"error\":\n                    raise ValueError(f\"{sequence_length=} is larger than {self.max_length=}\")\n                else:\n                    raise ValueError(f\"Unknown truncation method {self.truncation}\")\n\n            res = {\n                \"input_ids\": input_ids,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n                \"loss_mask\": loss_mask,\n            }\n            if len(multi_modal_inputs) > 0:\n                res[\"multi_modal_inputs\"] = multi_modal_inputs\n            return res\n        elif self.pad_mode == DatasetPadMode.NO_PADDING:\n            if sequence_length > self.max_length and self.truncation == \"error\":\n                raise ValueError(f\"{sequence_length=} is larger than {self.max_length=}\")\n            # truncate input_ids if it is longer than max_length\n            if len(input_ids) > self.max_length:\n                input_ids = input_ids[: self.max_length]\n                loss_mask = loss_mask[: self.max_length]\n                position_ids = position_ids[..., : self.max_length]\n\n            # return nested tensor with out padding\n            res = {\n                \"input_ids\": input_ids,\n                \"position_ids\": position_ids,\n                \"loss_mask\": loss_mask,\n            }\n            if len(multi_modal_inputs) > 0:\n                res[\"multi_modal_inputs\"] = multi_modal_inputs\n            return res\n        else:\n            raise ValueError(f\"Unknown pad mode {self.pad_mode}\")\n\n    def sanity_check(self, input_ids: torch.Tensor, messages: list[dict], tools: list[dict], enable_thinking: bool):\n        \"\"\"Check concatenated input_ids of apply_chat_template to each turn equals\n        apply_chat_template to whole messages.\n        \"\"\"\n        processor = self.processor if self.processor is not None else self.tokenizer\n        apply_chat_template_kwargs = {**self.apply_chat_template_kwargs}\n        if enable_thinking is not None:\n            apply_chat_template_kwargs[\"enable_thinking\"] = enable_thinking\n        inputs = processor.apply_chat_template(\n            messages,\n            tools=tools,\n            add_generation_prompt=False,\n            tokenize=True,\n            return_dict=True,\n            return_tensors=\"pt\",\n            **apply_chat_template_kwargs,\n        )\n\n        error_message = (\n            \"MultiTurnSFTDataset apply_chat_template to each turn separately and concat `input_ids` \"\n            \"as a whole sequence, which may not equal to apply_chat_template to whole messages at once.\\n\"\n            \"For example, Qwen Thinking series models add <think></think> tags to last turn, please check \"\n            \"your tokenizer chat template settings.\\n\"\n            \"Set `ignore_input_ids_mismatch=True` to ignore input_ids mismatch and use the concatenated \"\n            \"input_ids as the final input_ids. \"\n        )\n\n        if not torch.equal(input_ids, inputs[\"input_ids\"].squeeze(0)):\n            if self.ignore_input_ids_mismatch:\n                logger.warning_once(error_message)\n            else:\n                raise AssertionError(error_message)\n"
  },
  {
    "path": "verl/utils/dataset/rl_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport copy\nimport logging\nimport os\nimport re\nimport traceback\nfrom collections import defaultdict\nfrom io import BytesIO\nfrom typing import Optional\n\nimport datasets\nimport numpy as np\nimport torch\nfrom omegaconf import DictConfig, ListConfig\nfrom PIL import Image\nfrom torch.utils.data import Dataset\nfrom transformers import PreTrainedTokenizer, ProcessorMixin\n\nfrom verl.utils.import_utils import load_extern_object\nfrom verl.utils.tokenizer import normalize_token_ids\n\nlogger = logging.getLogger(__name__)\n\n\ndef collate_fn(data_list: list[dict]) -> dict:\n    \"\"\"\n    Collate a batch of sample dicts into batched tensors and arrays.\n\n    Args:\n        data_list: List of dicts mapping feature names to torch.Tensor or other values.\n\n    Returns:\n        Dict where tensor entries are stacked into a torch.Tensor of shape\n        (batch_size, \\\\*dims) and non-tensor entries are converted to\n        np.ndarray of dtype object with shape (batch_size,).\n    \"\"\"\n    tensors = defaultdict(list)\n    non_tensors = defaultdict(list)\n\n    for data in data_list:\n        for key, val in data.items():\n            if isinstance(val, torch.Tensor):\n                tensors[key].append(val)\n            else:\n                non_tensors[key].append(val)\n\n    for key, val in tensors.items():\n        tensors[key] = torch.stack(val, dim=0)\n\n    for key, val in non_tensors.items():\n        non_tensors[key] = np.fromiter(val, dtype=object, count=len(val))\n\n    return {**tensors, **non_tensors}\n\n\nclass RLHFDataset(Dataset):\n    \"\"\"\n    Load and preprocess RLHF data from Parquet files.\n\n    - Caches files locally.\n    - Reads into a HuggingFace Dataset and tokenizes prompts.\n    - Optionally handles images/videos via a ProcessorMixin.\n    - Filters prompts over a max length.\n    - Supports resuming from checkpoints.\n\n    Args:\n        data_files (str or list): Path(s) to Parquet file(s).\n        tokenizer (PreTrainedTokenizer): For the tokenization of text to token IDs.\n        config (DictConfig): Options like cache_dir, prompt_key, max_prompt_length, truncation, etc.\n        processor (ProcessorMixin, optional): Multimodal preprocessor for images/videos.\n    \"\"\"\n\n    def __init__(\n        self,\n        data_files: str | list[str],\n        tokenizer: PreTrainedTokenizer,\n        config: DictConfig,\n        processor: Optional[ProcessorMixin] = None,\n        max_samples: int = -1,\n    ):\n        if not isinstance(data_files, list | ListConfig):\n            data_files = [data_files]\n\n        self.data_files = copy.deepcopy(data_files)\n        self.original_data_files = copy.deepcopy(data_files)  # use for resume\n        self.tokenizer = tokenizer\n        self.processor = processor\n        self.max_samples = max_samples\n        self.config = config\n\n        self.cache_dir = os.path.expanduser(config.get(\"cache_dir\", \"~/.cache/verl/rlhf\"))\n        self.prompt_key = config.get(\"prompt_key\", \"prompt\")\n        self.image_key = config.get(\"image_key\", \"images\")\n        self.video_key = config.get(\"video_key\", \"videos\")\n        self.image_patch_size = config.get(\"image_patch_size\", 14)\n        self.max_prompt_length = config.get(\"max_prompt_length\", 1024)\n        self.return_raw_chat = config.get(\"return_raw_chat\", False)\n        self.return_full_prompt = config.get(\"return_full_prompt\", False)\n        self.truncation = config.get(\"truncation\", \"error\")\n        self.filter_overlong_prompts = config.get(\"filter_overlong_prompts\", True)\n        self.apply_chat_template_kwargs = config.get(\"apply_chat_template_kwargs\", {})\n\n        self.tool_config_path = config.get(\"tool_config_path\", None)\n        self.tool_schemas = None\n        if self.tool_config_path:\n            try:\n                from verl.tools.utils.tool_registry import initialize_tools_from_config\n\n                tool_list = initialize_tools_from_config(self.tool_config_path)\n                # match ToolAgentLoop behaviour: model_dump to plain dicts\n                self.tool_schemas = [\n                    tool.tool_schema.model_dump(exclude_unset=True, exclude_none=True) for tool in tool_list\n                ]\n            except Exception as e:\n                logger.warning(\"Failed to initialize tools from %s: %s\", self.tool_config_path, e)\n                self.tool_schemas = None\n\n        self.num_workers = config.get(\"filter_overlong_prompts_workers\", max(1, os.cpu_count() // 4))\n        self.num_workers = min(self.num_workers, os.cpu_count()) if self.num_workers is not None else None\n        self.use_shm = config.get(\"use_shm\", False)\n        self.chat_template_func = config.get(\"chat_template_func\", None)\n        self.need_tools_kwargs = config.get(\"need_tools_kwargs\", False)\n        self.filter_prompts = config.get(\"filter_prompts\", True)\n        self.serialize_dataset = False\n        self.return_multi_modal_inputs = config.get(\"return_multi_modal_inputs\", True)\n        self.shuffle = config.get(\"shuffle\", False)\n        self.seed = config.get(\"seed\")\n\n        self._download()\n        self._read_files_and_tokenize()\n\n    def _download(self, use_origin_parquet=False):\n        from verl.utils.fs import copy_to_local\n\n        data_files = self.data_files if not use_origin_parquet else self.original_data_files\n        for i, parquet_file in enumerate(data_files):\n            self.data_files[i] = copy_to_local(src=parquet_file, cache_dir=self.cache_dir, use_shm=self.use_shm)\n\n    def _read_files_and_tokenize(self):\n        dataframes = []\n        for parquet_file in self.data_files:\n            # read files and cache\n            if parquet_file.endswith(\".parquet\"):\n                dataframe = datasets.load_dataset(\"parquet\", data_files=parquet_file)[\"train\"]\n            elif parquet_file.endswith(\".json\") or parquet_file.endswith(\".jsonl\"):\n                dataframe = datasets.load_dataset(\"json\", data_files=parquet_file)[\"train\"]\n            else:\n                raise ValueError(f\"Unsupported file format: {parquet_file}\")\n            dataframes.append(dataframe)\n        self.dataframe: datasets.Dataset = datasets.concatenate_datasets(dataframes)\n\n        total = len(self.dataframe)\n        print(f\"dataset len: {len(self.dataframe)}\")\n\n        if self.max_samples > 0 and self.max_samples < total:\n            if self.shuffle:\n                rngs_args = (self.seed,) if self.seed is not None else ()\n                rng = np.random.default_rng(*rngs_args)\n                indices = rng.choice(total, size=self.max_samples, replace=False)\n            else:\n                indices = np.arange(self.max_samples)\n            self.dataframe = self.dataframe.select(indices.tolist())\n            print(f\"selected {self.max_samples} random samples out of {total}\")\n\n        self.dataframe = self.maybe_filter_out_long_prompts(self.dataframe)\n\n    def maybe_filter_out_long_prompts(self, dataframe: datasets.Dataset = None):\n        # filter out too long prompts\n        if self.filter_overlong_prompts:\n            tokenizer = self.tokenizer\n            processor = self.processor\n            prompt_key = self.prompt_key\n            image_key = self.image_key\n            video_key = self.video_key\n\n            if processor is not None:\n                from verl.utils.dataset.vision_utils import process_image, process_video\n\n                def doc2len(doc) -> int:\n                    try:\n                        messages = self._build_messages(doc)\n                        # pass tool schemas if available so the processor can format prompts\n                        apply_kwargs = dict(**self.apply_chat_template_kwargs)\n                        if self.tool_schemas is not None:\n                            apply_kwargs[\"tools\"] = self.tool_schemas\n\n                        raw_prompt = self.processor.apply_chat_template(\n                            messages, add_generation_prompt=True, tokenize=False, **apply_kwargs\n                        )\n                        if image_key in doc and doc[image_key]:\n                            images = [\n                                process_image(image, image_patch_size=self.image_patch_size) for image in doc[image_key]\n                            ]\n                        else:\n                            images = None\n\n                        if video_key in doc and doc[video_key]:\n                            videos, video_metadata = zip(\n                                *[\n                                    process_video(\n                                        video, image_patch_size=self.image_patch_size, return_video_metadata=True\n                                    )\n                                    for video in doc[video_key]\n                                ],\n                                strict=True,\n                            )\n                            videos = list(videos)\n                            video_metadata = list(video_metadata)\n                            videos_kwargs = {\"video_metadata\": video_metadata, \"do_sample_frames\": False}\n                        else:\n                            videos = None\n                            videos_kwargs = {}\n\n                        if images is None and videos is None:\n                            # only text prompt\n                            return len(\n                                processor.tokenizer(\n                                    text=raw_prompt,\n                                    add_special_tokens=False,  # avoid adding special tokens\n                                    return_attention_mask=False,\n                                )[\"input_ids\"]\n                            )\n                        else:\n                            # multi-modal prompt\n                            return len(\n                                processor(text=[raw_prompt], images=images, videos=videos, videos_kwargs=videos_kwargs)[\n                                    \"input_ids\"\n                                ][0]\n                            )\n                    except Exception:\n                        print(\"Error processing one of the samples, skipping...\")\n                        traceback.print_exc()\n                        return self.max_prompt_length + 1\n\n            else:\n\n                def doc2len(doc) -> int:\n                    try:\n                        apply_kwargs = dict(**self.apply_chat_template_kwargs)\n                        if self.tool_schemas is not None:\n                            apply_kwargs[\"tools\"] = self.tool_schemas\n\n                        # Keep explicit tokenization to avoid transformers version default changes.\n                        apply_kwargs.pop(\"tokenize\", None)\n                        apply_kwargs.pop(\"return_dict\", None)\n                        apply_kwargs.pop(\"return_tensors\", None)\n\n                        tokenized_prompt = tokenizer.apply_chat_template(\n                            doc[prompt_key], add_generation_prompt=True, tokenize=True, **apply_kwargs\n                        )\n                        return len(normalize_token_ids(tokenized_prompt))\n                    except Exception:\n                        print(\"Error processing one of the samples, skipping...\")\n                        traceback.print_exc()\n                        return self.max_prompt_length + 1\n\n            dataframe = dataframe.filter(\n                lambda doc: doc2len(doc) <= self.max_prompt_length,\n                num_proc=self.num_workers,\n                desc=f\"Filtering prompts longer than {self.max_prompt_length} tokens\",\n            )\n\n            print(f\"filter dataset len: {len(dataframe)}\")\n        return dataframe\n\n    def resume_dataset_state(self):\n        self.serialize_dataset = not hasattr(self, \"original_data_files\")\n        # resume dataframe if not it's serialized in data.pt\n        if not self.serialize_dataset:\n            self._download(use_origin_parquet=True)  # download and resume from original parquet files\n            self._read_files_and_tokenize()\n        else:\n            print(r\"old dataloader ckpt file is used, please train from scratch for better ckpt performance\")\n\n    def __getstate__(self):\n        if not self.serialize_dataset:\n            state = self.__dict__.copy()\n\n            if \"dataframe\" in state:\n                del state[\"dataframe\"]\n            return state\n\n        return self.__dict__.copy()\n\n    def __len__(self):\n        return len(self.dataframe)\n\n    def _build_messages(self, example: dict):\n        \"\"\"Replace <image> and <video> placeholder in messages with corresponding image and video\n        which is required by processor.apply_chat_template.\n        - <image>: {\"type\": \"image\", **image}\n        - <video>: {\"type\": \"video\", **video}\n\n        Args:\n            example: Row dictionary from dataframe.\n\n        Returns:\n            messages: List of messages with replaced placeholder.\n        \"\"\"\n        messages: list = example[self.prompt_key]\n        # When concatenating image and video datasets, pop will return None for image or video sample\n        images = example.pop(self.image_key, None) or []\n        videos = example.pop(self.video_key, None) or []\n\n        image_offset, video_offset = 0, 0\n        for message in messages:\n            if not images and not videos:\n                continue\n            assert self.processor is not None, \"processor is needed to process image and video\"\n\n            content = message[\"content\"]\n            if not isinstance(content, str):\n                continue\n\n            content_list = []\n            segments = re.split(\"(<image>|<video>)\", content)\n            segments = [item for item in segments if item != \"\"]\n            for segment in segments:\n                if segment == \"<image>\":\n                    assert image_offset < len(images), f\"image_offset {image_offset} >= len(images) {len(images)}\"\n                    image = images[image_offset]\n                    if isinstance(image, Image.Image):\n                        image = image.convert(\"RGB\")\n                        content_list.append({\"type\": \"image\", \"image\": image})\n                    elif isinstance(image, dict):\n                        if \"bytes\" in image:\n                            image[\"image\"] = Image.open(BytesIO(image[\"bytes\"]))\n                        content_list.append({\"type\": \"image\", **image})\n                    else:\n                        raise TypeError(f\"image must be dict or PIL.Image, unsupported image type: {type(image)}\")\n                    image_offset += 1\n                elif segment == \"<video>\":\n                    assert video_offset < len(videos), f\"video_offset {video_offset} >= len(videos) {len(videos)}\"\n                    content_list.append({\"type\": \"video\", **videos[video_offset]})\n                    video_offset += 1\n                else:\n                    content_list.append({\"type\": \"text\", \"text\": segment})\n            message[\"content\"] = content_list\n\n        assert image_offset == len(images), f\"image_offset {image_offset} != len(images) {len(images)}\"\n        assert video_offset == len(videos), f\"video_offset {video_offset} != len(videos) {len(videos)}\"\n        return messages\n\n    def __getitem__(self, item):\n        \"\"\"For rollout, apply_chat_template has been moved to AgentLoop, so we only return raw_prompt here.\"\"\"\n        row_dict: dict = self.dataframe[item]\n        row_dict[\"raw_prompt\"] = self._build_messages(row_dict)\n\n        # TODO(wuxibin): We still need a dummy tensor to make sure DataProto.batch is not empty.\n        # Remove this after deprecate DataProto by TensorDict.\n        row_dict[\"dummy_tensor\"] = torch.tensor([0], dtype=torch.uint8)\n\n        # add index for each prompt\n        if \"extra_info\" not in row_dict or row_dict[\"extra_info\"] is None:\n            row_dict[\"extra_info\"] = dict()\n        index = row_dict.get(\"extra_info\", {}).get(\"index\", 0)\n        tools_kwargs = row_dict.get(\"extra_info\", {}).get(\"tools_kwargs\", {})\n        interaction_kwargs = row_dict.get(\"extra_info\", {}).get(\"interaction_kwargs\", {})\n        need_tools_kwargs = row_dict.get(\"extra_info\", {}).get(\"need_tools_kwargs\", self.need_tools_kwargs)\n        if need_tools_kwargs and not tools_kwargs:\n            logger.warning(\"tools_kwargs is empty for index %s, data source: %s\", index, row_dict[\"data_source\"])\n        row_dict[\"index\"] = index\n        row_dict[\"tools_kwargs\"] = tools_kwargs\n        row_dict[\"interaction_kwargs\"] = interaction_kwargs\n        return row_dict\n\n    @classmethod\n    async def process_vision_info(\n        cls,\n        messages: list[dict],\n        image_patch_size,\n        config: DictConfig,\n    ) -> tuple[list[Image.Image], list[tuple[torch.Tensor, dict]]]:\n        \"\"\"Extract images and videos from messages.\n\n        This method is called by AgentLoop (e.g SingleTurnAgentLoop) before apply_chat_template to\n        the `raw_prompt` from dataset. User may customize RLHFDataset and override this method to\n        support custom vision extraction.\n\n        >>> messages = kwargs[\"raw_prompt\"]\n        >>> images, videos = RLHFDataset.process_vision_info(messages, image_patch_size)\n        >>> videos, video_metadatas = zip(*videos)\n        >>> raw_prompt = processor.apply_chat_template(messages, tokenize=False)\n        >>> inputs = processor(text=[raw_prompt], images=images, videos=videos,\n        ...                    video_metadata=video_metadatas, do_sample_frames=False)\n\n        Args:\n            messages: List of messages from dataset `raw_prompt`.\n            image_patch_size: Image patch size for processor.\n            config: Config for dataset.\n\n        Returns:\n            images: List of images.\n            videos: List of videos, each video is a tuple of (video_tensor, video_metadata).\n        \"\"\"\n        from qwen_vl_utils import process_vision_info\n\n        images, videos = process_vision_info(messages, image_patch_size=image_patch_size, return_video_metadata=True)\n        return images, videos\n\n    def split(self, num_splits: int):\n        \"\"\"\n        split the dataset into num_splits sub-datasets\n        Args:\n            num_splits: specified number of splits\n        Returns:\n            List[RLHFDataset]: list of RLHFDataset splits\n        Raises:\n            ValueError: if num_splits is not a positive integer\n        \"\"\"\n        if not isinstance(num_splits, int) or num_splits <= 0:\n            raise ValueError(f\"num_splits must be a positive integer, got {num_splits}\")\n\n        if not hasattr(self, \"dataframe\"):\n            raise AttributeError(\n                \"dataframe not found in RLHFDataset\\n\"\n                \"reason: _read_files_and_tokenize() not called or Parquet file loading failed\"\n            )\n        if self.dataframe is None:\n            raise ValueError(\"RLHFDataset dataframe 为 None!\")\n\n        total_samples = len(self.dataframe)\n        print(f\"total_samples: {total_samples}\")\n        if total_samples == 0:\n            raise ValueError(\"Cannot split an empty dataset\")\n\n        # Calculate effective sample count after dropping remainders if needed\n        if total_samples % num_splits != 0:\n            total_samples = total_samples - (total_samples % num_splits)\n            logging.warning(f\"Dropping {len(self.dataframe) % num_splits} samples, effective samples: {total_samples}\")\n\n        split_size = total_samples // num_splits\n        splits = []\n\n        for i in range(num_splits):\n            start_idx = i * split_size\n            end_idx = (i + 1) * split_size if i < num_splits - 1 else total_samples\n\n            split_dataframe = self.dataframe.select(range(start_idx, end_idx))\n\n            split_dataset = RLHFDataset(\n                data_files=self.data_files,\n                tokenizer=self.tokenizer,\n                config=self.config,\n                processor=self.processor,\n                max_samples=self.max_samples,\n            )\n            split_dataset.dataframe = split_dataframe\n            split_dataset.serialize_dataset = self.serialize_dataset\n            split_dataset.original_data_files = self.original_data_files\n\n            splits.append(split_dataset)\n\n        return splits\n\n\ndef get_dataset_class(data_config: DictConfig):\n    \"\"\"Get RLHF dataset class.\n\n    Args:\n        data_config: The data config.\n\n    Returns:\n        dataset_cls: The dataset class.\n    \"\"\"\n\n    # Check if a custom dataset class is specified in the data configuration\n    # and if the path to the custom class is provided\n    if \"custom_cls\" in data_config and data_config.custom_cls.get(\"path\", None) is not None:\n        # Dynamically load the custom dataset class\n        dataset_cls = load_extern_object(data_config.custom_cls.path, data_config.custom_cls.name)\n        # Verify that the custom dataset class inherits from torch.utils.data.Dataset\n        if not issubclass(dataset_cls, Dataset):\n            raise TypeError(\n                f\"The custom dataset class '{data_config.custom_cls.name}' from \"\n                f\"'{data_config.custom_cls.path}' must inherit from torch.utils.data.Dataset\"\n            )\n    else:\n        # Use the default RLHFDataset class if no custom class is specified\n        dataset_cls = RLHFDataset\n    print(f\"Using dataset class: {dataset_cls.__name__}\")\n\n    return dataset_cls\n"
  },
  {
    "path": "verl/utils/dataset/rm_dataset.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nfrom typing import Optional\n\nimport numpy as np\nimport pandas as pd\nimport torch\nfrom torch.utils.data import Dataset\n\nfrom verl.utils import hf_tokenizer\n\n\ndef download_files_distributed(download_fn):\n    import torch.distributed\n\n    if torch.distributed.is_initialized():\n        if torch.distributed.get_rank() == 0:\n            # download files\n            download_fn()\n\n        torch.distributed.barrier()\n    else:\n        # download anyway\n        download_fn()\n\n\nclass RMDataset(Dataset):\n    def __init__(\n        self,\n        parquet_files: str | list[str],\n        tokenizer,\n        prompt_key=\"prompt\",\n        chosen_key=\"chosen\",\n        rejected_key=\"rejected\",\n        max_length=1024,\n        add_eos=True,\n        cache_dir=\"~/.cache/verl/rm\",\n        max_samples: int = -1,\n        shuffle: bool = False,\n        seed: Optional[int] = None,\n    ):\n        if not isinstance(parquet_files, list):\n            parquet_files = [parquet_files]\n\n        self.parquet_files = parquet_files\n        self.max_samples = max_samples\n        self.shuffle = shuffle\n        self.seed = seed\n        self.cache_dir = os.path.expanduser(cache_dir)\n        if isinstance(tokenizer, str):\n            tokenizer = hf_tokenizer(tokenizer)\n        self.tokenizer = tokenizer\n\n        self.prompt_key = prompt_key\n        self.chosen_key = chosen_key\n        self.rejected_key = rejected_key\n\n        self.add_eos = add_eos\n        self.max_length = max_length\n\n        self._download()\n        self._read_files_and_tokenize()\n\n    def _download(self):\n        def _download_files():\n            from verl.utils.fs import copy, is_non_local\n\n            os.makedirs(self.cache_dir, exist_ok=True)\n            assert os.path.exists(self.cache_dir)\n            for i, parquet_file in enumerate(self.parquet_files):\n                if is_non_local(parquet_file):\n                    dst = os.path.join(self.cache_dir, os.path.basename(parquet_file))\n                    if not os.path.exists(dst):\n                        copy(src=parquet_file, dst=dst)\n                    self.parquet_files[i] = dst\n\n        download_files_distributed(_download_files)\n\n    def _read_files_and_tokenize(self):\n        dataframes = []\n        for parquet_file in self.parquet_files:\n            # read parquet files and cache\n            dataframe = pd.read_parquet(parquet_file)\n            dataframes.append(dataframe)\n        self.dataframe = pd.concat(dataframes)\n\n        total = len(self.dataframe)\n        print(f\"dataset len: {len(self.dataframe)}\")\n\n        if self.max_samples > 0 and self.max_samples < total:\n            if self.shuffle:\n                rngs_args = (self.seed,) if self.seed is not None else ()\n                rng = np.random.default_rng(*rngs_args)\n                indices = rng.choice(total, size=self.max_samples, replace=False)\n            else:\n                indices = np.arange(self.max_samples)\n            self.dataframe = self.dataframe.iloc[indices.tolist()]\n            print(f\"selected {self.max_samples} random samples out of {total}\")\n\n        self.prompts = self.dataframe[self.prompt_key].tolist()\n        self.chosen_responses = self.dataframe[self.chosen_key].tolist()\n        self.rejected_responses = self.dataframe[self.rejected_key].tolist()\n\n    def __len__(self):\n        return len(self.prompts)\n\n    def _pad_to_length(self, input_ids, attention_mask):\n        curr_length = input_ids.shape[-1]\n\n        if curr_length < self.max_length:\n            input_ids = torch.cat(\n                (input_ids, torch.zeros(size=(self.max_length - curr_length,), dtype=input_ids.dtype)), dim=-1\n            )\n            attention_mask = torch.cat(\n                (attention_mask, torch.zeros(size=(self.max_length - curr_length,), dtype=attention_mask.dtype)), dim=-1\n            )\n        elif curr_length > self.max_length:\n            input_ids = input_ids[: self.max_length]\n            attention_mask = attention_mask[: self.max_length]\n\n        return input_ids, attention_mask\n\n    def __getitem__(self, item):\n        prompt = self.prompts[item]\n        chosen_response = self.chosen_responses[item]\n        rejected_response = self.rejected_responses[item]\n\n        prompt_ids = self.tokenizer(prompt, return_tensors=\"pt\")[\"input_ids\"][0]\n        chosen_response_ids = self.tokenizer(chosen_response, return_tensors=\"pt\")[\"input_ids\"][0]\n        rejected_response_ids = self.tokenizer(rejected_response, return_tensors=\"pt\")[\"input_ids\"][0]\n\n        if self.add_eos:\n            chosen_response_ids = torch.cat((chosen_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1)\n            rejected_response_ids = torch.cat(\n                (rejected_response_ids, torch.tensor([self.tokenizer.eos_token_id])), dim=-1\n            )\n\n        chosen_input_ids = torch.cat((prompt_ids, chosen_response_ids), dim=-1)\n        chosen_attention_mask = torch.ones_like(chosen_input_ids)\n\n        rejected_input_ids = torch.cat((prompt_ids, rejected_response_ids), dim=-1)\n        rejected_attention_mask = torch.ones_like(rejected_input_ids)\n\n        chosen_input_ids, chosen_attention_mask = self._pad_to_length(chosen_input_ids, chosen_attention_mask)\n        rejected_input_ids, rejected_attention_mask = self._pad_to_length(rejected_input_ids, rejected_attention_mask)\n\n        input_ids = torch.stack((chosen_input_ids, rejected_input_ids), dim=0)\n        attention_mask = torch.stack((chosen_attention_mask, rejected_attention_mask), dim=0)\n\n        return {\n            \"input_ids\": input_ids,\n            \"attention_mask\": attention_mask,\n        }\n"
  },
  {
    "path": "verl/utils/dataset/vision_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom io import BytesIO\nfrom typing import Optional\n\nimport torch\nfrom PIL import Image\n\n\ndef process_image(image: dict | Image.Image, image_patch_size: int = 14) -> Image.Image:\n    from qwen_vl_utils import fetch_image\n\n    if isinstance(image, Image.Image):\n        return image.convert(\"RGB\")\n\n    if \"bytes\" in image:\n        assert \"image\" not in image, \"Cannot have both `bytes` and `image`\"\n        image[\"image\"] = Image.open(BytesIO(image[\"bytes\"]))\n\n    try:\n        ans = fetch_image(image, image_patch_size=image_patch_size)\n    except Exception:\n        ans = fetch_image(image)\n    return ans\n\n\nVIDEO_FORMAT_HELP = \"\"\"Currently, we only support the video formats introduced in qwen2-vl.\nRefer to https://github.com/QwenLM/Qwen2.5-VL?tab=readme-ov-file#using---transformers-to-chat.\n\neg.\n{\n    \"type\": \"video\",\n    \"video\": [\n        \"file:///path/to/frame1.jpg\",\n        \"file:///path/to/frame2.jpg\"\n    ]\n}\n\n{\n    \"type\": \"video\",\n    \"video\": \"file:///path/to/video.mp4\"\n}\n# Defaults to fps=2, min_frames=4, max_frames=768\n\n{\n    \"type\": \"video\",\n    \"video\": \"file:///path/to/video.mp4\",\n    \"fps\": 2,\n    \"min_frames\": 1,\n    \"max_frames\": 32\n}\n\"\"\"\n\n\ndef process_video(\n    video: dict,\n    image_patch_size: int = 14,\n    nframes: Optional[int] = None,\n    fps: Optional[float] = None,\n    fps_min_frames: Optional[int] = None,\n    fps_max_frames: Optional[int] = None,\n    return_video_sample_fps: bool = False,\n    return_video_metadata: bool = False,\n) -> torch.Tensor:\n    \"\"\"Converts a video dict into a [n_frames, 3, H, W] tensor\n\n    Add video sample FPS in a future MR\n    \"\"\"\n    from qwen_vl_utils import fetch_video\n\n    if not isinstance(video, dict) or \"video\" not in video:\n        raise NotImplementedError(VIDEO_FORMAT_HELP)\n    assert nframes is None or fps is None, \"Can't use both `nframes` or `fps`\"\n\n    # Shallow copy... since we might want to add some keys\n    video = dict(video)\n\n    contains_sampling_rules = \"nframes\" in video or \"fps\" in video\n    if not contains_sampling_rules:\n        if nframes is not None:\n            video[\"nframes\"] = nframes\n        elif fps is not None:\n            video[\"fps\"] = fps\n            if fps_min_frames is not None:\n                video[\"min_frames\"] = fps_min_frames\n            if fps_max_frames is not None:\n                video[\"max_frames\"] = fps_max_frames\n\n    return fetch_video(\n        video,\n        image_patch_size=image_patch_size,\n        return_video_sample_fps=return_video_sample_fps,\n        return_video_metadata=return_video_metadata,\n    )\n\n\ndef process_multi_modal_inputs_for_minicpmo(input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs):\n    # Adjust image bounds based on left padding and cumulative sequence lengths\n    # This is necessary for MiniCPM-o's vision-language alignment\n    left_padding_length = torch.argmax(attention_mask, dim=1)\n    image_bounds = []\n    for i in range(len(multi_modal_inputs[\"image_bound\"])):\n        image_bound = (\n            multi_modal_inputs[\"image_bound\"][i].to(left_padding_length.device) - left_padding_length[i] + cu_seqlens[i]\n        )\n        image_bounds.append(image_bound)\n\n    # Flatten pixel values list for MiniCPM-o processing\n    pixel_values = []\n    for i in range(len(multi_modal_inputs[\"pixel_values\"])):\n        pixel_values.extend([p for p in multi_modal_inputs[\"pixel_values\"][i]])\n\n    multi_modal_inputs[\"pixel_values\"] = [pixel_values]\n    multi_modal_inputs[\"image_bound\"] = [torch.vstack(image_bounds)]\n    multi_modal_inputs[\"tgt_sizes\"] = [torch.vstack(multi_modal_inputs[\"tgt_sizes\"])]\n    multi_modal_inputs[\"input_ids\"] = input_ids\n    multi_modal_inputs[\"attention_mask\"] = attention_mask\n    multi_modal_inputs[\"position_ids\"] = position_ids\n    return {\"data\": multi_modal_inputs}\n"
  },
  {
    "path": "verl/utils/debug/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# APIs kept for backward compatibility purpose\n# For new features please develop in verl/utils/profiler/\nfrom ..profiler import *  # noqa: F401\n"
  },
  {
    "path": "verl/utils/debug/metrics.py",
    "content": "# Copyright 2025 Individual Contributor: TomQunChaoA\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\n\nimport torch\n\nfrom verl.protocol import DataProto\n\nlogger = logging.getLogger(__file__)\n\n\ndef calculate_token_list_diff(tensor1: torch.Tensor, tensor2: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:\n    # verify inputs\n    if tensor1.numel() == 0 or tensor2.numel() == 0:\n        return torch.zeros(tensor1.shape[0], dtype=torch.long, device=tensor1.device)\n    if tensor1.shape != tensor2.shape or mask.shape != tensor1.shape or mask.shape != tensor2.shape:\n        print(\n            f\"<WARN> dim of tensor1, tensor2, mask is not equal, {(tensor1.shape)=},{(tensor2.shape)=}, {(mask.shape)=}\"\n        )\n        return torch.ones_like(tensor1)\n    # transfer to same device\n    if tensor2.device != tensor1.device:\n        tensor2 = tensor2.to(tensor1.device)\n    if mask.device != tensor1.device:\n        mask = mask.to(tensor1.device)\n\n    # calculate diff\n    diff_mask = tensor1 != tensor2\n\n    valid_diff_mask = diff_mask & (mask == 1)\n\n    diff_counts = valid_diff_mask.sum(dim=1)\n\n    return diff_counts\n\n\ndef pearson_correlation_coefficient(tensor1: torch.Tensor, tensor2: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:\n    # implemention of https://arxiv.org/pdf/2506.13585\n    if tensor1.shape != tensor2.shape or mask.shape != tensor1.shape or mask.shape != tensor2.shape:\n        return 0\n    mt1 = torch.masked_select(tensor1, mask)\n    mt2 = torch.masked_select(tensor2, mask)\n    result = torch.corrcoef(torch.stack([mt1, mt2], dim=0))\n    return result[0][1].detach().item()\n\n\ndef calculate_log_prob_diff(log_probs1: torch.Tensor, log_probs2: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:\n    full_diff = torch.abs(log_probs1 - log_probs2)\n    return torch.masked_select(full_diff, mask)\n\n\ndef calculate_debug_metrics(data: DataProto) -> dict:\n    \"\"\"\n    calculate rollout vs actor logprobs diff, for debugging purpose\n\n    Args:\n        data: DataProto\n            the data batch to calculate\n            rollout_log_probs: log_probs record when rollout forward tokens\n            old_log_probs(actor log probs): log_probs record when actor forward tokens\n            loss_mask or attention_mask: to mask unrelated token\n            responses: the response tokens, for calculating size\n    Returns:\n        dict: metrics\n            \"training/rollout_probs_diff_valid\": 1->input is valid, 0->input is invalid\n            \"training/rollout_probs_diff_max\": max value of logprob diff of rollout vs. actor\n            \"training/rollout_probs_diff_mean\": mean value of logprob diff of rollout vs. actor\n            \"training/rollout_probs_diff_std\": std value of logprob diff of rollout vs. actor\n            \"training/rollout_actor_probs_pearson_corr\": logprob's pearson corrcoef of rollout vs. actor, reference to https://arxiv.org/pdf/2506.13585\n    \"\"\"\n\n    rollout_old_log_probs = data.batch[\"rollout_log_probs\"]\n    actor_old_log_probs = data.batch[\"old_log_probs\"]\n    if \"response_mask\" in data.batch:\n        logger.debug(\"response mask found, use it to mask log probs\")\n        log_prob_mask = data.batch[\"response_mask\"]\n    elif \"attention_mask\" in data.batch:\n        log_prob_mask = data.batch[\"attention_mask\"]\n    else:\n        logger.warning(f\"no mask info found, use all log probs, {(data.batch.keys())=}\")\n        log_prob_mask = torch.ones_like(rollout_old_log_probs)\n    responses = data.batch[\"responses\"]\n    response_length = responses.size(1)\n\n    response_mask = log_prob_mask[:, -response_length:]\n    # calculate pearson corrcoef\n    actor_probs = torch.exp(actor_old_log_probs)\n    rollout_probs = torch.exp(rollout_old_log_probs)\n    response_mask_bool = response_mask.bool()\n    pearson_corrcoef = pearson_correlation_coefficient(actor_probs, rollout_probs, response_mask_bool)\n    rollout_probs_diff = calculate_log_prob_diff(actor_probs, rollout_probs, response_mask_bool)\n    return {\n        \"training/rollout_probs_diff_valid\": 1,\n        \"training/rollout_probs_diff_max\": torch.max(rollout_probs_diff).detach().item(),\n        \"training/rollout_probs_diff_mean\": torch.mean(rollout_probs_diff).detach().item(),\n        \"training/rollout_probs_diff_std\": torch.std(rollout_probs_diff).detach().item(),\n        \"training/rollout_actor_probs_pearson_corr\": pearson_corrcoef,\n    }\n"
  },
  {
    "path": "verl/utils/debug/performance.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# APIs kept for backward compatibility purpose\n# This file is deprecated, for new features please develop in profiler/performance.py\nfrom verl.utils.profiler.performance import reduce_timing, simple_timer  # noqa: F401\n"
  },
  {
    "path": "verl/utils/debug/trajectory_tracker.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nTrajectory tracker can be inserted into code to save the intermediate results.\nThe results will be dump to hdfs for offline comparison.\nEach process will have a client that first move all the tensors to CPU\n\"\"\"\n\nimport io\nimport os\nimport tempfile\nfrom collections import deque\n\nimport ray\nimport torch\n\nfrom verl.utils.hdfs_io import copy, makedirs\n\nremote_copy = ray.remote(copy)\n\n\n@ray.remote\ndef save_to_hdfs(data: io.BytesIO, name, hdfs_dir, verbose):\n    filename = name + \".pth\"\n    with tempfile.TemporaryDirectory() as tmpdirname:\n        local_filepath = os.path.join(tmpdirname, filename)\n        with open(local_filepath, \"wb\") as f:\n            f.write(data.getbuffer())\n        # upload to hdfs\n\n        if verbose:\n            print(f\"Saving {local_filepath} to {hdfs_dir}\")\n        try:\n            copy(local_filepath, hdfs_dir)\n        except Exception as e:\n            print(e)\n\n\n@ray.remote\nclass TrajectoryTracker:\n    def __init__(self, hdfs_dir, verbose) -> None:\n        self.hdfs_dir = hdfs_dir\n        makedirs(hdfs_dir)\n        self.verbose = verbose\n\n        self.handle = deque()\n\n    def dump(self, data: io.BytesIO, name):\n        # get a temp file and write to it\n        self.handle.append(save_to_hdfs.remote(data, name, self.hdfs_dir, self.verbose))\n\n    def wait_for_hdfs(self):\n        while len(self.handle) != 0:\n            future = self.handle.popleft()\n            ray.get(future)\n\n\ndef dump_data(data, name):\n    enable = os.getenv(\"VERL_ENABLE_TRACKER\", \"0\") == \"1\"\n    if not enable:\n        return\n    buffer = io.BytesIO()\n    torch.save(data, buffer)\n    tracker = get_trajectory_tracker()\n    ray.get(tracker.dump.remote(buffer, name))\n\n\ndef get_trajectory_tracker():\n    hdfs_dir = os.getenv(\"VERL_TRACKER_HDFS_DIR\", default=None)\n    verbose = os.getenv(\"VERL_TRACKER_VERBOSE\", default=\"0\") == \"1\"\n    assert hdfs_dir is not None\n    tracker = TrajectoryTracker.options(name=\"global_tracker\", get_if_exists=True, lifetime=\"detached\").remote(\n        hdfs_dir, verbose\n    )\n    return tracker\n\n\nif __name__ == \"__main__\":\n    # testing\n    os.environ[\"VERL_ENABLE_TRACKER\"] = \"1\"\n    os.environ[\"VERL_TRACKER_HDFS_DIR\"] = \"~/debug/test\"\n\n    @ray.remote\n    def process(iter):\n        data = {\"obs\": torch.randn(10, 20)}\n        dump_data(data, f\"process_{iter}_obs\")\n\n    ray.init()\n\n    output_lst = []\n\n    for i in range(10):\n        output_lst.append(process.remote(i))\n\n    out = ray.get(output_lst)\n\n    tracker = get_trajectory_tracker()\n    ray.get(tracker.wait_for_hdfs.remote())\n"
  },
  {
    "path": "verl/utils/device.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# This code is inspired by the torchtune.\n# https://github.com/pytorch/torchtune/blob/main/torchtune/utils/_device.py\n#\n# Copyright (c) Meta Platforms, Inc. and affiliates.\n# All rights reserved.\n#\n# This source code is licensed under the BSD-style license in https://github.com/pytorch/torchtune/blob/main/LICENSE\n\nimport logging\nimport os\nimport platform\nimport subprocess\n\nimport torch\nfrom packaging import version\n\nlogger = logging.getLogger(__name__)\n\n\ndef is_torch_npu_available(check_device=True) -> bool:\n    \"\"\"Check if Ascend NPU is available for PyTorch operations.\n\n    Attempts to detect NPU availability by checking for the torch.npu module\n    and its is_available() function.\n\n    Args:\n        check_device : only check torch_npu package or strictly check if NPU device is available\n\n    Returns:\n        bool: True if NPU is available, False otherwise.\n    \"\"\"\n    try:\n        if not hasattr(torch, \"npu\"):\n            return False\n\n        if check_device:\n            return torch.npu.is_available()\n        else:\n            return True\n    except ImportError:\n        return False\n\n\nis_cuda_available = torch.cuda.is_available()\nis_npu_available = is_torch_npu_available()\n\n\ndef get_resource_name() -> str:\n    \"\"\"Function that return ray resource name based on the device type.\n    Returns:\n        ray resource name string, either \"GPU\" or \"NPU\".\n    \"\"\"\n    return \"GPU\" if is_cuda_available else \"NPU\"\n\n\ndef get_visible_devices_keyword() -> str:\n    \"\"\"Get the environment variable name for visible device selection.\n\n    Returns the appropriate environment variable name based on the available\n    accelerator type (CUDA or Ascend NPU).\n\n    Returns:\n        str: 'CUDA_VISIBLE_DEVICES' if CUDA is available,\n            'ASCEND_RT_VISIBLE_DEVICES' otherwise.\n    \"\"\"\n    return \"CUDA_VISIBLE_DEVICES\" if not is_torch_npu_available(check_device=False) else \"ASCEND_RT_VISIBLE_DEVICES\"\n\n\ndef get_device_name() -> str:\n    \"\"\"Get the device type string based on available accelerators.\n\n    Detects the available accelerator and returns the corresponding PyTorch\n    device type string. Currently supports CUDA, Ascend NPU, and CPU.\n\n    Returns:\n        str: Device type string ('cuda', 'npu', or 'cpu').\n    \"\"\"\n    if is_cuda_available:\n        device = \"cuda\"\n    elif is_npu_available:\n        device = \"npu\"\n    else:\n        device = \"cpu\"\n    return device\n\n\ndef get_torch_device():\n    \"\"\"Get the PyTorch device module for the current accelerator.\n\n    Returns the torch device namespace (e.g., torch.cuda, torch.npu) based on\n    the detected accelerator type. Falls back to torch.cuda if the namespace\n    is not found.\n\n    Returns:\n        module: The PyTorch device module (torch.cuda, torch.npu, etc.).\n    \"\"\"\n    device_name = get_device_name()\n    try:\n        return getattr(torch, device_name)\n    except AttributeError:\n        logger.warning(f\"Device namespace '{device_name}' not found in torch, try to load torch.cuda.\")\n        return torch.cuda\n\n\ndef get_device_id() -> int:\n    \"\"\"Get the index of the current accelerator device.\n\n    Returns:\n        int: The current device index (e.g., 0 for 'cuda:0').\n    \"\"\"\n    return get_torch_device().current_device()\n\n\ndef get_nccl_backend() -> str:\n    \"\"\"Get the distributed communication backend based on device type.\n\n    Returns the appropriate collective communication backend for the\n    detected accelerator (HCCL for Ascend NPU, NCCL for CUDA).\n\n    Returns:\n        str: Backend name ('hccl' for NPU, 'nccl' for CUDA/default).\n    \"\"\"\n    if is_npu_available:\n        return \"hccl\"\n    else:\n        # default to nccl\n        return \"nccl\"\n\n\ndef set_expandable_segments(enable: bool) -> None:\n    \"\"\"Configure CUDA memory allocator expandable segments setting.\n\n    Expandable segments can help avoid out-of-memory (OOM) errors by allowing\n    the memory allocator to expand existing memory segments rather than\n    allocating new ones.\n\n    Args:\n        enable: If True, enable expandable segments. If False, disable them.\n\n    Note:\n        This function only has an effect when CUDA is available.\n    \"\"\"\n    if is_cuda_available:\n        torch.cuda.memory._set_allocator_settings(f\"expandable_segments:{enable}\")\n\n\ndef auto_set_device(config) -> None:\n    \"\"\"Automatically configure device name for different accelerators.\n\n    For example, on Ascend NPU, this function defaults the trainer device to \"npu\"\n    unless explicitly set to \"cpu\".\n\n    Args:\n        config: Configuration object with trainer.device attribute.\n    \"\"\"\n    if config and hasattr(config, \"trainer\") and hasattr(config.trainer, \"device\"):\n        if is_torch_npu_available():\n            if config.trainer.device not in [\"cpu\", \"npu\"]:\n                logger.warning(\n                    f\"Detect setting config.trainer.device to {config.trainer.device} for Ascend NPU, maybe\"\n                    f\"from default value in config file, automatically set to `npu` instead.\"\n                )\n\n            config.trainer.device = \"npu\"\n        # Other cases: set device to \"cuda\" via config file, no need to change.\n\n\ndef get_device_capability(device_id: int = 0) -> tuple[int | None, int | None]:\n    \"\"\"Get the compute capability of a CUDA device.\n\n    Args:\n        device_id: The CUDA device index to query. Defaults to 0.\n\n    Returns:\n        tuple: A tuple of (major, minor) compute capability version,\n            or (None, None) if CUDA is not available.\n    \"\"\"\n    major, minor = None, None\n    if is_cuda_available:\n        major, minor = torch.cuda.get_device_capability(device_id)\n\n    return major, minor\n\n\ndef get_npu_versions() -> tuple[str, str]:\n    \"\"\"Get the software version and CANN toolkit version for NPU devices.\n\n    Returns:\n        tuple[str, str]: A tuple of (software_version, cann_version)\n\n    Raises:\n        RuntimeError: If unable to retrieve version information\n    \"\"\"\n    # Check npu-smi software version\n    result = subprocess.run([\"npu-smi\", \"info\", \"-t\", \"board\", \"-i\", \"1\"], capture_output=True, text=True, check=True)\n\n    # Parse software version from output\n    software_version = None\n    for line in result.stdout.split(\"\\n\"):\n        if \"Software Version\" in line:\n            # Extract version from line like: \"Software Version : 25.3.rc1.2\"\n            parts = line.split(\":\")\n            if len(parts) > 1:\n                software_version = parts[1].strip().lower()\n            break\n\n    if not software_version:\n        raise RuntimeError(\"Could not find Software Version in npu-smi output\")\n\n    # Check CANN toolkit version\n    arch = platform.machine()\n    if arch not in [\"arm64\", \"aarch64\", \"x86_64\"]:\n        raise RuntimeError(f\"Unsupported architecture: {arch}\")\n\n    ascend_home = os.environ.get(\"ASCEND_HOME_PATH\", \"/usr/local/Ascend/ascend-toolkit/latest\")\n    cann_path = os.path.join(ascend_home, f\"{arch}-linux\")\n\n    if not os.path.exists(cann_path):\n        raise RuntimeError(f\"CANN toolkit path does not exist: {cann_path}\")\n\n    info_file = os.path.join(cann_path, \"ascend_toolkit_install.info\")\n    if not os.path.exists(info_file):\n        raise RuntimeError(f\"CANN toolkit info file does not exist: {info_file}\")\n\n    # Parse version from info file\n    cann_version = None\n    with open(info_file) as f:\n        for line in f:\n            if line.startswith(\"version=\"):\n                cann_version = line.split(\"=\", 1)[1].strip().lower()\n                break\n\n    if not cann_version:\n        raise RuntimeError(\"Could not find version in CANN toolkit info file\")\n\n    return software_version, cann_version\n\n\ndef check_ipc_version_support(software_version: str, cann_version: str) -> bool:\n    \"\"\"Check if the given software and CANN versions support IPC.\n\n    Compares the software version and CANN toolkit version against minimum\n    required versions for IPC support:\n    - Software Version should be >= 25.3.rc1\n    - CANN version should be >= 8.3.rc1\n\n    Args:\n        software_version: The software version string (e.g., \"25.5.0\", \"25.3.rc1.2\", \"25.5.t3.b001\")\n        cann_version: The CANN toolkit version string (e.g., \"8.3.0\", \"8.3.rc1\")\n\n    Returns:\n        bool: True if IPC is supported, False otherwise.\n\n    Raises:\n        RuntimeError: If version format is invalid\n    \"\"\"\n    # For software_version like \"25.3.rc1.2\", \"25.5.0\", or \"25.5.t3.b001\",\n    # we need to extract the base version\n    # Use regex to extract version with the following rules:\n    # - Standard version: 25.5.0 -> 25.5.0\n    # - RC version: 25.3.rc1.2 -> 25.3.rc1\n    # - t suffix version: 25.5.t3.b001 -> 25.5 (only first 2 parts if third part is lowercase t)\n    # - RC version: 25.3.rc1 -> 25.3.rc1\n    # For versions with more than 3 parts (e.g., 25.3.rc1.2), only match the first 3 parts\n    import re\n\n    # Match version with optional rc part or lowercase t suffix:\n    # - If version has lowercase t (e.g., 25.5.t3.b001), only match first 2 parts\n    # - Otherwise, match up to 3 parts (e.g., 25.5.0, 25.3.rc1.2)\n    ascend_version_pattern = r\"(\\d+\\.\\d+(?=\\.t))|(\\d+\\.\\d+(?:\\.(?:rc\\d+|\\d+))?)\"\n    software_match = re.match(ascend_version_pattern, software_version)\n    if not software_match:\n        raise RuntimeError(f\"Invalid software version format: {software_version}\")\n\n    # Select the matched group (either first 2 parts or up to 3 parts)\n    software_base = software_match.group(1) if software_match.group(1) else software_match.group(2)\n\n    cann_match = re.match(ascend_version_pattern, cann_version)\n    if not cann_match:\n        raise RuntimeError(f\"Invalid CANN version format: {cann_version}\")\n    else:\n        # Select the matched group (either first 2 parts or up to 3 parts)\n        cann_base = cann_match.group(1) if cann_match.group(1) else cann_match.group(2)\n\n    if version.parse(software_base) >= version.parse(\"25.3.rc1\"):\n        if version.parse(cann_base) >= version.parse(\"8.3.rc1\"):\n            return True\n        else:\n            logger.info(f\"CANN version {cann_version} is below 8.3.RC1\")\n    else:\n        logger.info(f\"Software version {software_version} is below 25.3.rc1\")\n\n    return False\n\n\ndef is_support_ipc() -> bool:\n    \"\"\"Check if the device supports IPC (Inter-Process Communication).\n\n    For GPU devices, always returns True.\n    For NPU devices, checks the software version and CANN toolkit version\n    to determine if IPC is supported.\n\n    Returns:\n        bool: True if IPC is supported, False otherwise.\n    \"\"\"\n    # If CUDA is available, it's a GPU device\n    if is_cuda_available:\n        return True\n\n    # For NPU devices, check the software version and CANN toolkit version\n    if is_npu_available:\n        try:\n            software_version, cann_version = get_npu_versions()\n            return check_ipc_version_support(software_version, cann_version)\n\n        except subprocess.CalledProcessError as e:\n            raise RuntimeError(f\"Failed to execute npu-smi command: {e}\") from e\n        except Exception as e:\n            raise RuntimeError(f\"Error checking IPC support: {e}\") from e\n\n    # For other devices (CPU), return False\n    return False\n"
  },
  {
    "path": "verl/utils/distributed.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Utilities for distributed training.\"\"\"\n\nimport ctypes\nimport os\nimport socket\nfrom datetime import timedelta\n\nimport ray\nimport torch.distributed\n\nfrom verl.utils.device import get_device_name, get_nccl_backend, get_torch_device, is_npu_available\nfrom verl.utils.net_utils import is_ipv6\n\n\ndef set_numa_affinity():\n    if is_npu_available:\n        # TODO (FightingZhen) libnuma.so is not available in e2e_ascend CI image, remove this code after image update.\n        return\n\n    initialized = False\n    try:\n        libnuma = ctypes.CDLL(\"libnuma.so\")\n        if libnuma.numa_available() < 0:\n            return\n\n        import pynvml\n\n        pynvml.nvmlInit()\n        initialized = True\n        device_name = \"NPU\" if is_npu_available else \"GPU\"\n        local_rank = int(ray.get_runtime_context().get_accelerator_ids()[device_name][0])\n        handle = pynvml.nvmlDeviceGetHandleByIndex(local_rank)\n        pynvml.nvmlDeviceSetCpuAffinity(handle)\n    except ImportError:\n        print(\"Warning: pynvml not available, skipping NUMA affinity setup\")\n    except Exception as e:\n        print(f\"Warning: Failed to set NUMA affinity: {e}\")\n    finally:\n        if initialized:\n            pynvml.nvmlShutdown()\n\n\ndef initialize_global_process_group(timeout_second=36000):\n    torch.distributed.init_process_group(\n        get_nccl_backend(),\n        timeout=timedelta(seconds=timeout_second),\n        init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n    )\n    local_rank = int(os.environ[\"LOCAL_RANK\"])\n    rank = int(os.environ[\"RANK\"])\n    world_size = int(os.environ[\"WORLD_SIZE\"])\n\n    if torch.distributed.is_initialized():\n        get_torch_device().set_device(local_rank)\n    return local_rank, rank, world_size\n\n\ndef destroy_global_process_group():\n    if torch.distributed.is_initialized():\n        torch.distributed.destroy_process_group()\n\n\ndef initialize_global_process_group_ray(timeout_second=None, backend=None):\n    # in current ray environment, LOCAL_RANK is always zero.\n\n    import torch.distributed\n\n    timeout = timedelta(seconds=timeout_second) if timeout_second is not None else None\n    backend = backend or f\"cpu:gloo,{get_device_name()}:{get_nccl_backend()}\"\n    if not torch.distributed.is_initialized():\n        rank = int(os.environ.get(\"RANK\", 0))\n        world_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n        torch.distributed.init_process_group(\n            backend=backend,\n            rank=rank,\n            world_size=world_size,\n            timeout=timeout,\n            init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n        )\n\n\ndef stateless_init_process_group(master_address, master_port, rank, world_size, device):\n    \"\"\"\n    vLLM provides `StatelessProcessGroup` to create a process group\n    without considering the global process group in torch.distributed.\n    It is recommended to create `StatelessProcessGroup`, and then initialize\n    the data-plane communication (NCCL) between external (train processes)\n    and vLLM workers.\n    \"\"\"\n    # NOTE: If it is necessary to support weight synchronization with the sglang backend in the future,\n    # the following can be used:\n    # from sglang.srt.distributed.device_communicators.pynccl import PyNcclCommunicator\n    # from sglang.srt.distributed.utils import statelessprocessgroup\n\n    from torch.distributed import TCPStore\n    from vllm.distributed.utils import StatelessProcessGroup\n\n    from verl.utils.device import is_npu_available\n\n    if is_npu_available:\n        from vllm_ascend.distributed.device_communicators.pyhccl import PyHcclCommunicator as PyNcclCommunicator\n    else:\n        from vllm.distributed.device_communicators.pynccl import PyNcclCommunicator\n\n    def create_process_group(\n        host: str,\n        port: int,\n        rank: int,\n        world_size: int,\n        data_expiration_seconds: int = 3600,\n        store_timeout: int = 300,\n    ) -> \"StatelessProcessGroup\":\n        \"\"\"\n        This is copied from vllm/distributed/utils.py:StatelessProcessGroup.create\n        Modified to support ipv6 stateless communication groups.\"\"\"\n        launch_server = rank == 0\n        if launch_server:\n            # listen on the specified interface (instead of 0.0.0.0)\n            if is_ipv6(master_address):\n                listen_socket = socket.socket(socket.AF_INET6, socket.SOCK_STREAM)\n            else:\n                listen_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM)\n            listen_socket.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n            listen_socket.bind((host, port))\n            listen_socket.listen()\n            listen_fd = listen_socket.fileno()\n        else:\n            listen_socket = None\n            listen_fd = None\n\n        store = TCPStore(\n            host_name=host,\n            port=port,\n            world_size=world_size,\n            is_master=launch_server,\n            timeout=timedelta(seconds=store_timeout),\n            use_libuv=False,  # for now: github.com/pytorch/pytorch/pull/150215\n            master_listen_fd=listen_fd,\n        )\n\n        return StatelessProcessGroup(\n            rank=rank,\n            world_size=world_size,\n            store=store,\n            socket=listen_socket,\n            data_expiration_seconds=data_expiration_seconds,\n        )\n\n    pg = create_process_group(host=master_address, port=master_port, rank=rank, world_size=world_size)\n\n    pynccl = PyNcclCommunicator(pg, device=device)\n    return pynccl\n"
  },
  {
    "path": "verl/utils/experimental/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/utils/experimental/torch_functional.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Optional\n\nimport torch\n\n\ndef _fused_linear_for_ppo_fwd(\n    hidden_states: torch.FloatTensor,\n    vocab_weights: torch.FloatTensor,\n    input_ids: torch.LongTensor,\n    temperature: float = 1.0,\n) -> tuple[torch.FloatTensor, torch.FloatTensor]:\n    logits = (hidden_states @ vocab_weights.t()) / temperature\n    orig_dtype = logits.dtype\n    logits = logits.to(torch.float32)\n\n    # Slower but more numerically stable to do log_softmax than probs.log()\n    probs = logits.softmax(dim=-1)\n    log_probs = logits.log_softmax(dim=-1)\n\n    token_log_probs = log_probs.gather(-1, input_ids.unsqueeze(-1)).squeeze(-1)\n    entropy = torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1)\n\n    return token_log_probs.to(orig_dtype), entropy.to(orig_dtype)\n\n\ndef _fused_linear_for_ppo_bwd(\n    dlog_probs: Optional[torch.FloatTensor],\n    dentropy: Optional[torch.FloatTensor],\n    hidden_states: torch.FloatTensor,\n    vocab_weights: torch.FloatTensor,\n    input_ids: torch.LongTensor,\n    temperature: float = 1.0,\n) -> tuple[torch.FloatTensor, torch.FloatTensor]:\n    logits = (hidden_states @ vocab_weights.t()) / temperature\n    orig_dtype = logits.dtype\n    logits = logits.to(torch.float32)\n\n    probs = logits.softmax(dim=-1)\n\n    dlogits = 0\n\n    # Gradient from log_probs\n    if dlog_probs is not None:\n        one_hot_input = torch.zeros_like(logits).scatter_(-1, input_ids.unsqueeze(-1), 1)\n        dlogits += dlog_probs.to(torch.float32).unsqueeze(-1) * (one_hot_input - probs)\n\n    # Gradient from entropy\n    if dentropy is not None:\n        log_probs = logits.log_softmax(dim=-1)\n        entropy = torch.logsumexp(logits, dim=-1) - torch.sum(probs * logits, dim=-1)\n        dlogits += probs * (log_probs + entropy.unsqueeze(-1)) * (-dentropy.unsqueeze(-1))\n\n    dlogits = dlogits.to(orig_dtype) / temperature\n\n    dhidden_states = dlogits @ vocab_weights\n    dvocab_weights = dlogits.t() @ hidden_states\n\n    return dhidden_states, dvocab_weights\n\n\nclass FusedLinearForPPOFunction(torch.autograd.Function):\n    @staticmethod\n    def forward(\n        ctx,\n        hidden_states: torch.FloatTensor,\n        vocab_weights: torch.FloatTensor,\n        input_ids: torch.LongTensor,\n        temperature: float = 1.0,\n        chunk_size: int = 512,\n    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:\n        ctx.set_materialize_grads(False)\n\n        # Cast to a 2D tensor of the shape [T, D] for ease of working\n        orig_ndim = hidden_states.ndim\n        assert orig_ndim in (2, 3), f\"Invalid hidden_states shape, received {hidden_states.shape}\"\n\n        orig_batch_size = -1\n        if orig_ndim == 3:\n            assert input_ids.ndim == 2, f\"input_ids shape doesn't match, {hidden_states.shape} {input_ids.shape}\"\n            orig_batch_size = hidden_states.shape[0]\n            hidden_states = hidden_states.flatten(0, 1)\n            input_ids = input_ids.flatten(0, 1)\n\n        T = hidden_states.shape[0]\n\n        # Allocate memory for outputs\n        output_requires_grad = hidden_states.requires_grad or vocab_weights.requires_grad\n        log_probs = hidden_states.new_zeros(T, requires_grad=output_requires_grad)\n        entropy = hidden_states.new_zeros(T, requires_grad=output_requires_grad)\n\n        # Perform forward one chunk at a time\n        for chunk_start in range(0, T, chunk_size):\n            chunk_end = min(chunk_start + chunk_size, T)\n\n            chunk_log_probs, chunk_entropy = _fused_linear_for_ppo_fwd(\n                hidden_states=hidden_states[chunk_start:chunk_end],\n                vocab_weights=vocab_weights,\n                input_ids=input_ids[chunk_start:chunk_end],\n                temperature=temperature,\n            )\n            log_probs[chunk_start:chunk_end] = chunk_log_probs\n            entropy[chunk_start:chunk_end] = chunk_entropy\n\n        # Cast the output back to the original input dimension\n        if orig_ndim == 3:\n            log_probs = log_probs.view(orig_batch_size, -1)\n            entropy = entropy.view(orig_batch_size, -1)\n\n        ctx.save_for_backward(hidden_states, vocab_weights, input_ids)\n        ctx.orig_batch_size = orig_batch_size\n        ctx.orig_ndim = orig_ndim\n        ctx.temperature = temperature\n        ctx.chunk_size = chunk_size\n\n        return log_probs, entropy\n\n    @staticmethod\n    def backward(ctx, dlog_probs: Optional[torch.FloatTensor], dentropy: Optional[torch.FloatTensor]):\n        assert dlog_probs is not None or dentropy is not None\n\n        hidden_states, vocab_weights, input_ids = ctx.saved_tensors\n        orig_batch_size = ctx.orig_batch_size\n        orig_ndim = ctx.orig_ndim\n        temperature = ctx.temperature\n        chunk_size = ctx.chunk_size\n\n        # Here orig_ndim refers to the orig_ndim of hidden_states\n        if orig_ndim == 3:\n            if dlog_probs is not None:\n                dlog_probs = dlog_probs.flatten()\n            if dentropy is not None:\n                dentropy = dentropy.flatten()\n\n        T = hidden_states.shape[0]\n\n        # Allocate memory for outputs\n        dhidden_states = None\n        if hidden_states.requires_grad:\n            dhidden_states = torch.zeros_like(hidden_states)\n        dvocab_weights = None\n        if vocab_weights.requires_grad:\n            dvocab_weights = torch.zeros_like(vocab_weights)\n\n        # Perform backward one chunk at a time\n        for chunk_start in range(0, T, chunk_size):\n            chunk_end = min(chunk_start + chunk_size, T)\n            chunk_dlog_probs = None\n            if dlog_probs is not None:\n                chunk_dlog_probs = dlog_probs[chunk_start:chunk_end]\n            chunk_dentropy = None\n            if dentropy is not None:\n                chunk_dentropy = dentropy[chunk_start:chunk_end]\n\n            h, v = _fused_linear_for_ppo_bwd(\n                dlog_probs=chunk_dlog_probs,\n                dentropy=chunk_dentropy,\n                hidden_states=hidden_states[chunk_start:chunk_end],\n                vocab_weights=vocab_weights,\n                input_ids=input_ids[chunk_start:chunk_end],\n                temperature=temperature,\n            )\n\n            if hidden_states.requires_grad:\n                dhidden_states[chunk_start:chunk_end] += h\n            if vocab_weights.requires_grad:\n                dvocab_weights += v\n\n        # Cast the output back to the original input dimension\n        if orig_ndim == 3 and hidden_states.requires_grad:\n            hidden_size = hidden_states.shape[-1]\n            dhidden_states = dhidden_states.view(orig_batch_size, -1, hidden_size)\n\n        return (\n            dhidden_states,  # hidden_states\n            dvocab_weights,  # vocab_weights\n            None,  # input_ids\n            None,  # temperature\n            None,  # chunk_size\n        )\n\n\nclass FusedLinearForPPO(torch.nn.Module):\n    def __init__(self, chunk_size: int = 512):\n        super().__init__()\n\n        self.chunk_size = chunk_size\n\n    def forward(\n        self,\n        hidden_states: torch.FloatTensor,\n        vocab_weights: torch.FloatTensor,\n        input_ids: torch.LongTensor,\n        temperature: float = 1.0,\n    ) -> tuple[torch.FloatTensor, torch.FloatTensor]:\n        input_ids = input_ids.to(torch.int64)\n        return FusedLinearForPPOFunction.apply(\n            hidden_states,\n            vocab_weights,\n            input_ids,\n            temperature,\n            self.chunk_size,\n        )\n"
  },
  {
    "path": "verl/utils/flops_counter.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport inspect\n\nimport torch\nfrom transformers import PretrainedConfig\n\nfrom verl.utils.device import get_torch_device\n\n_DEVICE_FLOPS = {\n    \"CPU\": 448e9,\n    \"GB200\": 2.5e15,\n    \"B200\": 2.25e15,\n    \"MI300X\": 1336e12,\n    \"H100\": 989e12,\n    \"H800\": 989e12,\n    \"H200\": 989e12,\n    \"A100\": 312e12,\n    \"A800\": 312e12,\n    \"L40S\": 362.05e12,\n    \"L40\": 181.05e12,\n    \"A40\": 149.7e12,\n    \"L20\": 119.5e12,\n    \"H20\": 148e12,\n    \"910B\": 354e12,\n    \"Ascend910\": 354e12,\n    \"RTX 3070 Ti\": 21.75e12,\n}\n\n\ndef get_device_flops(unit=\"T\", device_name=None):\n    \"\"\"Get the theoretical FLOPS (Floating Point Operations Per Second) capacity of the current device.\n\n    Args:\n        unit (str): The unit to return the FLOPS in. Supported values are:\n            \"B\" - Billion (1e9)\n            \"K\" - Thousand (1e3)\n            \"M\" - Million (1e6)\n            \"G\" - Giga (1e9)\n            \"T\" - Tera (1e12, default)\n            \"P\" - Peta (1e15)\n\n    Returns:\n        float: The theoretical FLOPS capacity of the current device in the specified unit.\n        Returns float('inf') for unknown GPU types.\n    \"\"\"\n\n    def unit_convert(number, level):\n        units = [\"B\", \"K\", \"M\", \"G\", \"T\", \"P\"]\n        if number <= 0:\n            return number\n        ptr = 0\n        while ptr < len(units) and units[ptr] != level:\n            number /= 1000\n            ptr += 1\n        return number\n\n    # pass device_name is for testing purpose only\n    if device_name is None:\n        device = get_torch_device()\n        if device == torch.cpu:\n            device_name = \"CPU\"\n        else:\n            device_name = get_torch_device().get_device_name()\n\n    flops = float(\"inf\")  # INF flops for unkown gpu type\n\n    for key, value in sorted(_DEVICE_FLOPS.items(), reverse=True):\n        if key in device_name:\n            flops = value\n            break\n    flops_unit = unit_convert(flops, unit)\n    return flops_unit\n\n\ndef _estimate_qwen2_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    num_hidden_layers = config.num_hidden_layers\n    num_key_value_heads = config.num_key_value_heads\n    num_attention_heads = config.num_attention_heads\n    intermediate_size = config.intermediate_size\n\n    head_dim = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # non-attn per layer parm\n    # Qwen2/LLama use SwiGelu, gate, having up and down linear layer in mlp\n    mlp_N = hidden_size * intermediate_size * 3\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_qwen3_vl_flops(config, tokens_sum, batch_seqlens, delta_time, **kargs):\n    # qwen3_vl uses text_config and vision_config to distinguish configs of different parts.\n    hidden_size = config.text_config.hidden_size\n    vocab_size = config.text_config.vocab_size\n    num_hidden_layers = config.text_config.num_hidden_layers\n    num_key_value_heads = config.text_config.num_key_value_heads\n    num_attention_heads = config.text_config.num_attention_heads\n    intermediate_size = config.text_config.intermediate_size\n\n    head_dim = hidden_size // num_attention_heads\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # non-attn per layer parm\n    mlp_N = hidden_size * intermediate_size * 3\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # qwen3_vl uses deepstack to merge visual embeds and text embeds, but it has no tensor operation.\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers\n\n    # vit flops\n    images_seqlens = kargs.get(\"images_seqlens\", None)\n    if images_seqlens is not None:\n        vit_flops = _estimate_qwen3_vit_flop(images_seqlens, config.vision_config)\n    else:\n        vit_flops = 0\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops + vit_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_qwen3_vl_moe_flops(config, tokens_sum, batch_seqlens, delta_time, **kargs):\n    # qwen3_vl uses text_config and vision_config to distinguish configs of different parts.\n    hidden_size = config.text_config.hidden_size\n    vocab_size = config.text_config.vocab_size\n    num_hidden_layers = config.text_config.num_hidden_layers\n    num_key_value_heads = config.text_config.num_key_value_heads\n    num_attention_heads = config.text_config.num_attention_heads\n    moe_intermediate_size = config.text_config.moe_intermediate_size\n    moe_num_expert = config.text_config.num_experts\n    moe_topk = config.text_config.num_experts_per_tok\n\n    head_dim = getattr(\n        config.text_config, \"head_dim\", config.text_config.hidden_size // config.text_config.num_attention_heads\n    )\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # non-attn per layer parm\n    moe_gata_N = hidden_size * moe_num_expert\n    # moe has gate_proj, up_proj and down_proj using SwiGLU in ExpertMlp layer & shared experts\n    moe_expertmlp_N = hidden_size * moe_intermediate_size * (moe_topk) * 3\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    moe_N = (moe_gata_N + moe_expertmlp_N + attn_linear_N) * (num_hidden_layers) + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * moe_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers\n\n    # vit flops\n    images_seqlens = kargs.get(\"images_seqlens\", None)\n    if images_seqlens is not None:\n        vit_flops = _estimate_qwen3_vit_flop(images_seqlens, config.vision_config)\n    else:\n        vit_flops = 0\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops + vit_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_qwen3_vit_flop(images_seqlens, config):\n    \"\"\"\n    Estimate the FLOPS of the vision encoder for Qwen3-VL\n    \"\"\"\n\n    if config is None:\n        return 0\n    tokens_sum = sum(images_seqlens)\n\n    num_heads = config.num_heads\n    depth = config.depth\n\n    dim = config.hidden_size\n    mlp_hidden_dim = config.intermediate_size\n    out_hidden_size = config.out_hidden_size\n\n    spatial_merge_size = config.spatial_merge_size\n\n    head_dim = dim // num_heads\n\n    # every vision token's patch_embed comes from a conv of (C, T, H, W) -> (dim,)\n    patch_embed_N = dim * config.in_channels * config.temporal_patch_size * config.patch_size * config.patch_size\n    # Qwen3 VL vision mlp does not use GLU, thus 2.\n    mlp_N = dim * mlp_hidden_dim * 2\n    attn_linear_N = dim * (4 * dim)  # qkv and output proj\n    merger_N = (out_hidden_size + (dim * (spatial_merge_size**2))) * (dim * (spatial_merge_size**2))\n\n    # Qwen3 VL uses deep stack, one merger for every deepstack layer\n    deepstack_merger_N = merger_N * len(config.deepstack_visual_indexes)\n    # non-attn all_layer parm\n    dense_N = patch_embed_N + (mlp_N + attn_linear_N) * depth + deepstack_merger_N + merger_N\n\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # In Qwen3 VL, full attention is used in all vision layers.\n    full_attn_layer_num = depth\n\n    # full attn layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in images_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 12 * seqlen_square_sum * head_dim * num_heads * full_attn_layer_num\n\n    vit_flops = dense_N_flops + attn_qkv_flops\n\n    return vit_flops\n\n\ndef _estimate_deepseek_v3_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    moe_intermediate_size = config.moe_intermediate_size\n    num_hidden_layers = config.num_hidden_layers\n    first_k_dense_replace = config.first_k_dense_replace\n    num_query_heads = config.num_attention_heads\n    moe_num_expert = config.n_routed_experts\n\n    moe_topk = config.num_experts_per_tok\n    share_expert_num = config.n_shared_experts\n\n    # non-attn per layer parm\n    moe_gata_N = hidden_size * moe_num_expert\n    # moe has fc1_1, fc1_2 and fc2 using SwiGLU in ExpertMlp layer & shared experts\n    moe_expertmlp_N = hidden_size * moe_intermediate_size * (moe_topk + share_expert_num) * 3\n    # MLA attn\n    attn_linear_N = 0\n    q_head_dim = config.qk_nope_head_dim + config.qk_rope_head_dim\n    if config.q_lora_rank is None:\n        attn_linear_N += hidden_size * num_query_heads * q_head_dim\n    else:\n        attn_linear_N += hidden_size * config.q_lora_rank\n        attn_linear_N += num_query_heads * q_head_dim * config.q_lora_rank\n\n    attn_linear_N += hidden_size * (config.kv_lora_rank + config.qk_rope_head_dim)\n    attn_linear_N += num_query_heads * (q_head_dim - config.qk_rope_head_dim + config.v_head_dim) * config.kv_lora_rank\n    attn_linear_N += num_query_heads * config.v_head_dim * hidden_size\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    moe_N = (\n        (moe_gata_N + moe_expertmlp_N + attn_linear_N) * (num_hidden_layers - first_k_dense_replace)\n        + (hidden_size * config.intermediate_size * 3 + attn_linear_N) * first_k_dense_replace\n        + emd_and_lm_head_N\n    )\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * moe_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen * num_hidden_layers\n\n    # Core attention FLOPS for MLA with causal mask:\n    # Q @ K^T: 3 * 2 * seq^2 * q_head_dim * num_heads / 2 (causal)\n    # attn @ V: 3 * 2 * seq^2 * v_head_dim * num_heads / 2 (causal)\n    attn_qkv_flops = 3 * seqlen_square_sum * (q_head_dim + config.v_head_dim) * num_query_heads\n    # all_layer & all_token fwd & bwk flops\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n\n    return flops_achieved\n\n\ndef _estimate_qwen2_moe_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    num_hidden_layers = config.num_hidden_layers\n    num_key_value_heads = config.num_key_value_heads\n    num_attention_heads = config.num_attention_heads\n    moe_intermediate_size = config.moe_intermediate_size\n    moe_topk = config.num_experts_per_tok\n    num_experts = config.num_experts\n\n    head_dim = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # non-attn per layer parm\n    # gate + moe export\n    moe_mlp_N = hidden_size * moe_topk * moe_intermediate_size * 3 + hidden_size * num_experts\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    dense_N = (moe_mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_gemma3_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    num_hidden_layers = config.num_hidden_layers\n    num_key_value_heads = config.num_key_value_heads\n    num_attention_heads = config.num_attention_heads\n    intermediate_size = config.intermediate_size\n\n    head_dim = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # non-attn per layer parm\n    # Gemma3 uses GeGLU (gelu_pytorch_tanh), having 3 matrices in MLP (inherited from Gemma2MLP)\n    mlp_N = hidden_size * intermediate_size * 3\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer parm\n    dense_N = (mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    # Gemma3 alternates between full and sliding window attention based on layer_types\n    seqlen_square_sum = 0\n\n    layer_types = getattr(config, \"layer_types\", None)\n    sliding_window = getattr(config, \"sliding_window\", 1024)  # default 1024\n    # default pattern: every 6th layer is full\n    sliding_window_pattern = getattr(config, \"sliding_window_pattern\", 6)\n\n    # If layer_types is not provided, generate it based on sliding_window_pattern\n    if layer_types is None and sliding_window is not None and sliding_window_pattern is not None:\n        layer_types = [\n            \"sliding_attention\" if bool((i + 1) % sliding_window_pattern) else \"full_attention\"\n            for i in range(num_hidden_layers)\n        ]\n\n    if layer_types:\n        # Calculate attention flops per layer based on attention type\n        for layer_idx in range(num_hidden_layers):\n            is_sliding = False\n            if layer_types and layer_idx < len(layer_types):\n                is_sliding = layer_types[layer_idx] == \"sliding_attention\"\n\n            for seqlen in batch_seqlens:\n                if is_sliding and sliding_window:\n                    # Sliding window limits each token to attend to at most window_size tokens\n                    effective_seqlen = min(seqlen, sliding_window)\n                    seqlen_square_sum += seqlen * effective_seqlen\n                else:\n                    # Full attention\n                    seqlen_square_sum += seqlen * seqlen\n    else:\n        # If no layer_types config, assume all layers use full attention\n        for seqlen in batch_seqlens:\n            seqlen_square_sum += seqlen * seqlen\n        seqlen_square_sum *= num_hidden_layers\n\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_apertus_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    num_hidden_layers = config.num_hidden_layers\n    num_key_value_heads = config.num_key_value_heads\n    num_attention_heads = config.num_attention_heads\n    intermediate_size = config.intermediate_size\n\n    head_dim = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # Apertus MLP with XIELU activation uses only 2 linear layers (up_proj, down_proj)\n    # No gate_proj for XIELU, unlike SwiGLU which has 3 layers\n    mlp_N = hidden_size * intermediate_size * 2\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n\n    # ApertusConfig has qk_norm defaulting to True.\n    # This adds params for q_norm (on H) and k_norm (on num_kv_heads * head_dim)\n    qk_norm_params_per_layer = hidden_size + num_key_value_heads * head_dim  # q_norm + k_norm\n\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n    # non-attn all_layer params\n    dense_N = (mlp_N + attn_linear_N + qk_norm_params_per_layer) * num_hidden_layers + emd_and_lm_head_N\n    # non-attn all_layer & all_token fwd & bwd flops\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # attn all_layer & all_token fwd & bwd flops\n    seqlen_square_sum = 0\n    for seqlen in batch_seqlens:\n        seqlen_square_sum += seqlen * seqlen\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads * num_hidden_layers\n\n    # all_layer & all_token fwd & bwd flops\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_gpt_oss_flops(config, tokens_sum, batch_seqlens, delta_time):\n    hidden_size = config.hidden_size\n    vocab_size = config.vocab_size\n    num_hidden_layers = config.num_hidden_layers\n    num_key_value_heads = config.num_key_value_heads\n    num_attention_heads = config.num_attention_heads\n\n    # MoE params\n    moe_intermediate_size = config.intermediate_size\n    num_experts = config.num_local_experts\n    num_experts_per_tok = config.num_experts_per_tok\n    mlp_matrices = 3\n\n    # Head dim\n    head_dim = getattr(config, \"head_dim\", hidden_size // num_attention_heads)\n    q_size = num_attention_heads * head_dim\n    k_size = num_key_value_heads * head_dim\n    v_size = num_key_value_heads * head_dim\n\n    # 1. Attention Block (GQA)\n    attn_linear_N = hidden_size * (q_size + k_size + v_size + num_attention_heads * head_dim)\n    # 2. MLP / MoE Block\n    # Gate network\n    moe_gate_N = hidden_size * num_experts\n    # Expert forward calculation, Active parameters: mlp_matrices * H * I * num_experts_per_tok\n    moe_expert_N = hidden_size * moe_intermediate_size * mlp_matrices * num_experts_per_tok\n\n    moe_mlp_N = moe_gate_N + moe_expert_N\n\n    emd_and_lm_head_N = vocab_size * hidden_size * 2\n\n    # Total non-attn params per layer * layers + embeddings\n    # (moe_mlp_N + attn_linear_N) * layers\n    dense_N = (moe_mlp_N + attn_linear_N) * num_hidden_layers + emd_and_lm_head_N\n\n    # FLOPs for dense part (fwd + bwd = 6 * N)\n    dense_N_flops = 6 * dense_N * tokens_sum\n\n    # 3. Attention Matrix FLOPs\n    seqlen_square_sum = 0\n\n    # Handle sliding window attention\n    layer_types = getattr(config, \"layer_types\", None)\n    sliding_window = getattr(config, \"sliding_window\", 128)\n\n    if layer_types:\n        for layer_type in layer_types:\n            is_sliding = layer_type == \"sliding_attention\"\n\n            for seqlen in batch_seqlens:\n                if is_sliding and sliding_window:\n                    # Sliding window limits each token to attend to at most window_size tokens\n                    effective_seqlen = min(seqlen, sliding_window)\n                    seqlen_square_sum += seqlen * effective_seqlen\n                else:\n                    # Full attention\n                    seqlen_square_sum += seqlen * seqlen\n    else:\n        # Default to full attention for all layers\n        for seqlen in batch_seqlens:\n            seqlen_square_sum += seqlen * seqlen\n        seqlen_square_sum *= num_hidden_layers\n\n    attn_qkv_flops = 6 * seqlen_square_sum * head_dim * num_attention_heads\n\n    # Total FLOPs\n    flops_all_token = dense_N_flops + attn_qkv_flops\n    flops_achieved = flops_all_token * (1.0 / delta_time) / 1e12\n    return flops_achieved\n\n\ndef _estimate_unknown_flops(config, tokens_sum, batch_seqlens, delta_time):\n    return 0\n\n\nESTIMATE_FUNC = {\n    \"qwen2\": _estimate_qwen2_flops,\n    \"llama\": _estimate_qwen2_flops,\n    \"qwen2_moe\": _estimate_qwen2_moe_flops,\n    \"qwen2_vl\": _estimate_qwen2_flops,\n    \"qwen2_5_vl\": _estimate_qwen2_flops,\n    \"qwen3\": _estimate_qwen2_flops,\n    \"qwen3_moe\": _estimate_qwen2_moe_flops,\n    \"qwen3_vl\": _estimate_qwen3_vl_flops,\n    \"qwen3_vl_moe\": _estimate_qwen3_vl_moe_flops,\n    \"deepseek_v3\": _estimate_deepseek_v3_flops,\n    \"minicpmv\": _estimate_qwen2_flops,\n    \"minicpmo\": _estimate_qwen2_flops,\n    \"mistral\": _estimate_qwen2_flops,\n    \"gemma3_text\": _estimate_gemma3_flops,\n    \"seed_oss\": _estimate_qwen2_flops,\n    \"apertus\": _estimate_apertus_flops,\n    \"glm4v\": _estimate_qwen2_flops,\n    \"gpt_oss\": _estimate_gpt_oss_flops,\n    \"mimo\": _estimate_qwen2_flops,\n}\n\n\nclass FlopsCounter:\n    \"\"\"\n    Used to count mfu during training loop\n\n    Example:\n        flops_counter = FlopsCounter(config)\n        flops_achieved, flops_promised = flops_counter.estimate_flops(tokens_list, delta_time)\n\n    \"\"\"\n\n    def __init__(self, config: PretrainedConfig):\n        VALID_CONFIG_TYPE = ESTIMATE_FUNC.keys()\n        if config.model_type not in VALID_CONFIG_TYPE:\n            print(\n                f\"Only support config type of {VALID_CONFIG_TYPE}, but got {config.model_type}. MFU will always be \"\n                f\"zero.\"\n            )\n\n        self.config = config\n\n    # TODO: actually we can make this a static method\n    def estimate_flops(self, batch_seqlens, delta_time, **kargs):\n        \"\"\"\n        Estimate the FLOPS based on the number of valid tokens in the current batch and the time taken.\n\n        Args:\n            batch_seqlens (List[int]): A list where each element represents the number of valid tokens in the\n                current batch.\n            delta_time (float): The time taken to process the batch, in seconds.\n\n        Returns:\n            estimated_flops (float): The estimated FLOPS based on the input tokens and time.\n            promised_flops (float): The expected FLOPS of the current device.\n        \"\"\"\n        tokens_sum = sum(batch_seqlens)\n        func = ESTIMATE_FUNC.get(self.config.model_type, _estimate_unknown_flops)\n        sig = inspect.signature(func)\n        if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):\n            estimated_flops = func(self.config, tokens_sum, batch_seqlens, delta_time, **kargs)\n        else:\n            estimated_flops = func(self.config, tokens_sum, batch_seqlens, delta_time)\n        promised_flops = get_device_flops()\n        return estimated_flops, promised_flops\n"
  },
  {
    "path": "verl/utils/fp8_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\n\nimport torch\n\nfrom verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\n\n\nclass FP8QuantizerHelper:\n    def __init__(self, quant_config):\n        self.quant_config = quant_config\n\n    def should_quantize_param(self, param_name):\n        \"\"\"Determine whether to quantize to FP8 based on parameter name\n\n        Quantization rules:\n        - Must end with .weight (exclude bias)\n        - Exclude embedding layers\n        - Exclude normalization layers\n        - Exclude output layer (lm_head)\n        \"\"\"\n        # Must be a weight parameter\n        if not param_name.endswith(\".weight\"):\n            return False\n\n        # Layer types to exclude\n        exclude_patterns = [\n            \"embed_tokens\",  # Embedding layer\n            \"lm_head\",  # Output layer\n            \"layernorm\",  # LayerNorm\n            \"norm\",  # Various Norm layers\n            \"ln_\",  # LayerNorm variants\n            \"embeddings\",  # Embeddings\n            \"mlp.gate.weight\",  # MoE router\n        ]\n\n        # Check if matches exclude patterns\n        param_lower = param_name.lower()\n        for pattern in exclude_patterns:\n            if pattern in param_lower:\n                return False\n\n        # Layer types to include (Linear layers)\n        include_patterns = [\n            \"q_proj\",  # Query projection\n            \"k_proj\",  # Key projection\n            \"v_proj\",  # Value projection\n            \"o_proj\",  # Output projection\n            \"gate_proj\",  # Gate projection (for MLP)\n            \"up_proj\",  # Up projection (for MLP)\n            \"down_proj\",  # Down projection (for MLP)\n            \"fc1\",  # Fully connected 1\n            \"fc2\",  # Fully connected 2\n            \"mlp\",  # MLP layers\n        ]\n\n        # Check if matches include patterns\n        for pattern in include_patterns:\n            if pattern in param_lower:\n                logger.debug(f\"Will quantize FP8: {param_name}\")\n                return True\n\n        # Do not quantize by default\n        logger.debug(f\"Skip quantization: {param_name}\")\n        return False\n\n    def quant_weights_by_name(self, weights, dtype=torch.bfloat16):\n        \"\"\"FP8 quantization based on parameter name using a memory-efficient generator.\n\n\n        Args:\n            weights: Generator or iterable of (name, tensor) pairs\n            dtype: Data type for intermediate computation\n\n        Yields:\n            Tuples of (name, tensor) for each weight and its scale\n        \"\"\"\n        if isinstance(self.quant_config, dict):\n            weight_block_size = self.quant_config.get(\"weight_block_size\")\n        else:\n            weight_block_size = getattr(self.quant_config, \"weight_block_size\", None)\n\n        if weight_block_size is None:\n            raise ValueError(\"weight_block_size not found in quant_config\")\n\n        for k, v in weights:\n            # Check if quantization is needed\n            if not self.should_quantize_param(k):\n                yield (k, v)\n                continue\n\n            # Quantize to FP8\n            try:\n                if torch.distributed.get_rank() == 0:\n                    logger.debug(f\"Quantizing to FP8 blockwise: {k}\")\n\n                param_lp, param_scale = scaled_fp8_blockwise(\n                    v.to(dtype),\n                    weight_block_size=weight_block_size,\n                )\n                param_scale = param_scale.squeeze(-1)\n\n                # Yield the quantized weight and scale\n                yield (k, param_lp)\n                yield (k + \"_scale_inv\", param_scale)\n\n                # Explicitly delete to help GC\n                del param_lp, param_scale\n\n            except Exception as e:\n                logger.error(f\"Failed to quantize {k}: {e}\")\n                # If quantization fails, use original weights\n                yield (k, v)\n"
  },
  {
    "path": "verl/utils/fs.py",
    "content": "#!/usr/bin/env python\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# -*- coding: utf-8 -*-\n\"\"\"File-system agnostic IO APIs\"\"\"\n\nimport hashlib\nimport os\nimport shutil\nimport tempfile\n\ntry:\n    from hdfs_io import copy, exists, makedirs  # for internal use only\nexcept ImportError:\n    from .hdfs_io import copy, exists, makedirs\n\n__all__ = [\"copy\", \"exists\", \"makedirs\"]\n\n_HDFS_PREFIX = \"hdfs://\"\n\n\ndef is_non_local(path):\n    \"\"\"Check if a path is a non-local (HDFS) path.\n\n    Args:\n        path (str): The path to check.\n\n    Returns:\n        bool: True if the path is an HDFS path, False otherwise.\n    \"\"\"\n    return path.startswith(_HDFS_PREFIX)\n\n\ndef md5_encode(path: str) -> str:\n    \"\"\"Generate an MD5 hash of a path string.\n\n    This function is used to create unique identifiers for paths, typically\n    for creating cache directories or lock files.\n\n    Args:\n        path (str): The path to encode.\n\n    Returns:\n        str: The hexadecimal MD5 hash of the path.\n    \"\"\"\n    return hashlib.md5(path.encode()).hexdigest()\n\n\ndef get_local_temp_path(hdfs_path: str, cache_dir: str) -> str:\n    \"\"\"Generate a unique local cache path for an HDFS resource.\n    Creates a MD5-hashed subdirectory in cache_dir to avoid name conflicts,\n    then returns path combining this subdirectory with the HDFS basename.\n\n    Args:\n        hdfs_path (str): Source HDFS path to be cached\n        cache_dir (str): Local directory for storing cached files\n\n    Returns:\n        str: Absolute local filesystem path in format:\n            {cache_dir}/{md5(hdfs_path)}/{basename(hdfs_path)}\n    \"\"\"\n    # make a base64 encoding of hdfs_path to avoid directory conflict\n    encoded_hdfs_path = md5_encode(hdfs_path)\n    temp_dir = os.path.join(cache_dir, encoded_hdfs_path)\n    os.makedirs(temp_dir, exist_ok=True)\n    dst = os.path.join(temp_dir, os.path.basename(hdfs_path))\n    return dst\n\n\ndef verify_copy(src: str, dest: str) -> bool:\n    \"\"\"\n    verify the copy of src to dest by comparing their sizes and file structures.\n\n    return:\n        bool: True if the copy is verified, False otherwise.\n    \"\"\"\n    if not os.path.exists(src):\n        return False\n    if not os.path.exists(dest):\n        return False\n\n    if os.path.isfile(src) != os.path.isfile(dest):\n        return False\n\n    if os.path.isfile(src):\n        src_size = os.path.getsize(src)\n        dest_size = os.path.getsize(dest)\n        if src_size != dest_size:\n            return False\n        return True\n\n    src_files = set()\n    dest_files = set()\n\n    for root, dirs, files in os.walk(src):\n        rel_path = os.path.relpath(root, src)\n        dest_root = os.path.join(dest, rel_path) if rel_path != \".\" else dest\n\n        if not os.path.exists(dest_root):\n            return False\n\n        for entry in os.listdir(root):\n            src_entry = os.path.join(root, entry)\n            src_files.add(os.path.relpath(src_entry, src))\n\n        for entry in os.listdir(dest_root):\n            dest_entry = os.path.join(dest_root, entry)\n            dest_files.add(os.path.relpath(dest_entry, dest))\n\n    if src_files != dest_files:\n        return False\n\n    for rel_path in src_files:\n        src_entry = os.path.join(src, rel_path)\n        dest_entry = os.path.join(dest, rel_path)\n\n        if os.path.isdir(src_entry) != os.path.isdir(dest_entry):\n            return False\n\n        if os.path.isfile(src_entry):\n            src_size = os.path.getsize(src_entry)\n            dest_size = os.path.getsize(dest_entry)\n            if src_size != dest_size:\n                return False\n\n    return True\n\n\ndef copy_to_shm(src: str):\n    \"\"\"\n    Load the model into   /dev/shm   to make the process of loading the model multiple times more efficient.\n    \"\"\"\n    shm_model_root = \"/dev/shm/verl-cache/\"\n    src_abs = os.path.abspath(os.path.normpath(src))\n    dest = os.path.join(shm_model_root, hashlib.md5(src_abs.encode(\"utf-8\")).hexdigest())\n    os.makedirs(dest, exist_ok=True)\n    dest = os.path.join(dest, os.path.basename(src_abs))\n    if os.path.exists(dest) and verify_copy(src, dest):\n        # inform user and depends on him\n        print(\n            f\"[WARNING]: The memory model path {dest} already exists. If it is not you want, please clear it and \"\n            f\"restart the task.\"\n        )\n    else:\n        if os.path.isdir(src):\n            shutil.copytree(src, dest, symlinks=False, dirs_exist_ok=True)\n        else:\n            shutil.copy2(src, dest)\n    return dest\n\n\ndef _record_directory_structure(folder_path):\n    record_file = os.path.join(folder_path, \".directory_record.txt\")\n    with open(record_file, \"w\") as f:\n        for root, dirs, files in os.walk(folder_path):\n            for dir_name in dirs:\n                relative_dir = os.path.relpath(os.path.join(root, dir_name), folder_path)\n                f.write(f\"dir:{relative_dir}\\n\")\n            for file_name in files:\n                if file_name != \".directory_record.txt\":\n                    relative_file = os.path.relpath(os.path.join(root, file_name), folder_path)\n                    f.write(f\"file:{relative_file}\\n\")\n    return record_file\n\n\ndef _check_directory_structure(folder_path, record_file):\n    if not os.path.exists(record_file):\n        return False\n    existing_entries = set()\n    for root, dirs, files in os.walk(folder_path):\n        for dir_name in dirs:\n            relative_dir = os.path.relpath(os.path.join(root, dir_name), folder_path)\n            existing_entries.add(f\"dir:{relative_dir}\")\n        for file_name in files:\n            if file_name != \".directory_record.txt\":\n                relative_file = os.path.relpath(os.path.join(root, file_name), folder_path)\n                existing_entries.add(f\"file:{relative_file}\")\n    with open(record_file) as f:\n        recorded_entries = set(f.read().splitlines())\n    return existing_entries == recorded_entries\n\n\ndef copy_to_local(\n    src: str, cache_dir=None, filelock=\".file.lock\", verbose=False, always_recopy=False, use_shm: bool = False\n) -> str:\n    \"\"\"Copy files/directories from HDFS to local cache with validation.\n\n    Args:\n        src (str): Source path - HDFS path (hdfs://...), local filesystem path, or Hugging Face model ID\n        cache_dir (str, optional): Local directory for cached files. Uses system tempdir if None\n        filelock (str): Base name for file lock. Defaults to \".file.lock\"\n        verbose (bool): Enable copy operation logging. Defaults to False\n        always_recopy (bool): Force fresh copy ignoring cache. Defaults to False\n        use_shm (bool): Enable shared memory copy. Defaults to False\n\n    Returns:\n        str: Local filesystem path to copied resource\n    \"\"\"\n    # Save to a local path for persistence.\n    local_path = copy_local_path_from_hdfs(src, cache_dir, filelock, verbose, always_recopy)\n\n    if use_shm and isinstance(local_path, str) and not os.path.exists(local_path):\n        try:\n            from huggingface_hub import snapshot_download\n\n            resolved = snapshot_download(local_path)\n            if isinstance(resolved, str) and os.path.exists(resolved):\n                local_path = resolved\n        except ImportError:\n            pass\n        except Exception as e:\n            print(f\"WARNING: Failed to download model from Hugging Face: {e}\")\n\n    # Load into shm to improve efficiency.\n    if use_shm:\n        return copy_to_shm(local_path)\n    return local_path\n\n\ndef copy_local_path_from_hdfs(\n    src: str, cache_dir=None, filelock=\".file.lock\", verbose=False, always_recopy=False\n) -> str:\n    \"\"\"Deprecated. Please use copy_to_local instead.\"\"\"\n    from filelock import FileLock\n\n    assert src[-1] != \"/\", f\"Make sure the last char in src is not / because it will cause error. Got {src}\"\n\n    if is_non_local(src):\n        # download from hdfs to local\n        if cache_dir is None:\n            # get a temp folder\n            cache_dir = tempfile.gettempdir()\n        os.makedirs(cache_dir, exist_ok=True)\n        assert os.path.exists(cache_dir)\n        local_path = get_local_temp_path(src, cache_dir)\n        # get a specific lock\n        filelock = md5_encode(src) + \".lock\"\n        lock_file = os.path.join(cache_dir, filelock)\n        with FileLock(lock_file=lock_file):\n            if always_recopy and os.path.exists(local_path):\n                if os.path.isdir(local_path):\n                    shutil.rmtree(local_path, ignore_errors=True)\n                else:\n                    os.remove(local_path)\n            if not os.path.exists(local_path):\n                if verbose:\n                    print(f\"Copy from {src} to {local_path}\")\n                copy(src, local_path)\n                if os.path.isdir(local_path):\n                    _record_directory_structure(local_path)\n            elif os.path.isdir(local_path):\n                # always_recopy=False, local path exists, and it is a folder: check whether there is anything missed\n                record_file = os.path.join(local_path, \".directory_record.txt\")\n                if not _check_directory_structure(local_path, record_file):\n                    if verbose:\n                        print(f\"Recopy from {src} to {local_path} due to missing files or directories.\")\n                    shutil.rmtree(local_path, ignore_errors=True)\n                    copy(src, local_path)\n                    _record_directory_structure(local_path)\n        return local_path\n    else:\n        return src\n\n\ndef local_mkdir_safe(path):\n    \"\"\"_summary_\n    Thread-safe directory creation function that ensures the directory is created\n    even if multiple processes attempt to create it simultaneously.\n\n    Args:\n        path (str): The path to create a directory at.\n    \"\"\"\n\n    from filelock import FileLock\n\n    if not os.path.isabs(path):\n        working_dir = os.getcwd()\n        path = os.path.join(working_dir, path)\n\n    # Using hash value of path as lock file name to avoid long file name\n    lock_filename = f\"ckpt_{hash(path) & 0xFFFFFFFF:08x}.lock\"\n    lock_path = os.path.join(tempfile.gettempdir(), lock_filename)\n\n    try:\n        with FileLock(lock_path, timeout=60):  # Add timeout\n            # make a new dir\n            os.makedirs(path, exist_ok=True)\n    except Exception as e:\n        print(f\"Warning: Failed to acquire lock for {path}: {e}\")\n        # Even if the lock is not acquired, try to create the directory\n        os.makedirs(path, exist_ok=True)\n\n    return path\n"
  },
  {
    "path": "verl/utils/fsdp_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nimport itertools\nimport json\nimport math\nimport os\nfrom abc import ABC\nfrom collections import OrderedDict\nfrom contextlib import contextmanager, nullcontext\nfrom typing import Optional, cast\n\nimport torch\nimport torch.distributed as dist\nimport torch.nn as nn\nfrom packaging import version\nfrom torch.distributed import DeviceMesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp._runtime_utils import _lazy_init\nfrom torch.distributed.fsdp.wrap import size_based_auto_wrap_policy, transformer_auto_wrap_policy\nfrom transformers.trainer_pt_utils import get_module_class_from_name\n\nfrom verl.utils.device import get_device_id, get_device_name, get_torch_device\nfrom verl.utils.model import check_exclude_modules, check_target_modules\n\nif version.parse(torch.__version__) >= version.parse(\"2.6\"):\n    from torch.distributed.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard\n    from torch.distributed.fsdp._fully_shard._fsdp_init import _get_post_forward_mesh_info\n    from torch.distributed.tensor import DTensor, Shard\n    from torch.distributed.tensor._dtensor_spec import DTensorSpec\n\n    fully_shard_module = torch.distributed.fsdp._fully_shard._fully_shard\nelif version.parse(torch.__version__) >= version.parse(\"2.4\"):\n    from torch.distributed._composable.fsdp import CPUOffloadPolicy, FSDPModule, MixedPrecisionPolicy, fully_shard\n\n    fully_shard_module = torch.distributed._composable.fsdp\nelse:\n    fully_shard, MixedPrecisionPolicy, FSDPModule, CPUOffloadPolicy, fully_shard_module = None, None, None, None, None\n\n\ndef init_fn(x: torch.nn.Module):\n    if torch.distributed.get_rank() != 0:\n        x = x.to_empty(device=get_device_id(), recurse=False)\n        get_torch_device().empty_cache()\n    return x\n\n\ndef get_init_weight_context_manager(use_meta_tensor=True, mesh: DeviceMesh = None):\n    from accelerate import init_empty_weights\n\n    cpu_init_weights = lambda: torch.device(\"cpu\")\n    if use_meta_tensor:\n        if mesh is None:\n            init_context = init_empty_weights if torch.distributed.get_rank() != 0 else cpu_init_weights\n        else:\n            init_context = init_empty_weights if mesh.get_coordinate()[-1] != 0 else cpu_init_weights\n    else:\n        init_context = cpu_init_weights\n    return init_context\n\n\n# Copyright 2020-present the HuggingFace Inc. team.\n# Adapted from https://github.com/huggingface/transformers/src/transformers/trainer.py\ndef get_fsdp_wrap_policy(module, config=None, is_lora=False):\n    \"\"\"Get FSDP wrap policy for the module.\n\n    Args:\n        module: The module to get wrap policy for\n        config: Configuration for wrap policy\n        is_lora: Whether to enable lambda policy for LoRA modules\n    \"\"\"\n    if config is None:\n        config = {}\n\n    # NOTE: This is a temporary workaround to be compatible with the OmegaConf & dataclass. We will remove this\n    # once we have make all config in verl from OmegaConf to data class.\n    def _get_attr(attr_name, default_value=None):\n        if hasattr(config, \"get\"):\n            return config.get(attr_name, default_value)\n        else:\n            return config.__getattribute__(attr_name)\n\n    if _get_attr(\"disable\", False):\n        return None\n\n    default_transformer_cls_names_to_wrap = getattr(module, \"_no_split_modules\", None)\n    fsdp_transformer_layer_cls_to_wrap = _get_attr(\n        \"transformer_layer_cls_to_wrap\", default_transformer_cls_names_to_wrap\n    )\n    min_num_params = _get_attr(\"min_num_params\", 0)\n    auto_wrap_policy = None\n\n    policies = []\n\n    from torch.distributed.fsdp.wrap import _or_policy, lambda_auto_wrap_policy\n\n    # Add lambda policy for LoRA modules if is_lora is True\n    if is_lora:\n\n        def lambda_policy_fn(module):\n            return bool(\n                len(list(module.named_children())) == 0\n                and getattr(module, \"weight\", None) is not None\n                and module.weight.requires_grad\n            )\n\n        lambda_policy = functools.partial(lambda_auto_wrap_policy, lambda_fn=lambda_policy_fn)\n        policies.append(lambda_policy)\n\n    if min_num_params > 0:\n        size_policy = functools.partial(size_based_auto_wrap_policy, min_num_params=min_num_params)\n        policies.append(size_policy)\n    elif fsdp_transformer_layer_cls_to_wrap is not None:\n        transformer_cls_to_wrap = set()\n        for layer_class in fsdp_transformer_layer_cls_to_wrap:\n            transformer_cls = get_module_class_from_name(module, layer_class)\n            if transformer_cls is None:\n                raise Exception(\"Could not find the transformer layer class to wrap in the model.\")\n            else:\n                transformer_cls_to_wrap.add(transformer_cls)\n\n        transformer_policy = functools.partial(\n            transformer_auto_wrap_policy,\n            transformer_layer_cls=transformer_cls_to_wrap,\n        )\n        policies.append(transformer_policy)\n\n    if len(policies) > 0:\n        auto_wrap_policy = functools.partial(_or_policy, policies=policies)\n\n    return auto_wrap_policy\n\n\n@torch.no_grad()\ndef offload_fsdp_model_to_cpu(model: FSDP, empty_cache: bool = True):\n    if fsdp_version(model) == 2 or fsdp_version(model) == 0:\n        offload_fsdp2_model_to_cpu(model, empty_cache)\n        return\n\n    assert isinstance(model, FSDP)\n    # lazy init FSDP model\n    _lazy_init(model, model)\n    assert model._is_root, \"Only support root model offloading to CPU\"\n    for handle in model._all_handles:\n        if handle._offload_params:\n            continue\n        flat_param = handle.flat_param\n        assert (\n            flat_param.data.data_ptr() == flat_param._local_shard.data_ptr()\n            and id(flat_param.data) != id(flat_param._local_shard)\n            and flat_param.data.size() == flat_param._local_shard.size()\n        )\n        handle.flat_param_to(torch.device(\"cpu\"), non_blocking=True)\n        # the following still keeps id(._local_shard) != id(.data)\n        flat_param._local_shard = flat_param.data\n        assert id(flat_param._local_shard) != id(flat_param.data)\n    if empty_cache:\n        get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef offload_fsdp2_model_to_cpu(model, empty_cache: bool = True):\n    model.cpu()\n    if empty_cache:\n        get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef load_fsdp_model_to_gpu(model: FSDP):\n    if fsdp_version(model) == 2 or fsdp_version(model) == 0:\n        load_fsdp2_model_to_gpu(model)\n        return\n\n    assert isinstance(model, FSDP)\n    # lazy init FSDP model\n    _lazy_init(model, model)\n    assert model._is_root, \"Only support root model loading to GPU\"\n    device_id = get_device_id()\n    for handle in model._all_handles:\n        if handle._offload_params:\n            continue\n        flat_param = handle.flat_param\n        handle.flat_param_to(torch.device(f\"{get_device_name()}:{device_id}\"), non_blocking=True)\n        # the following still keeps id(._local_shard) != id(.data)\n        flat_param._local_shard = flat_param.data\n\n\n@torch.no_grad()\ndef load_fsdp2_model_to_gpu(model):\n    device = get_device_id()\n    model.to(device)\n\n\n@torch.no_grad()\ndef offload_fsdp_optimizer(optimizer):\n    if not optimizer.state:\n        return\n    for param_group in optimizer.param_groups:\n        for param in param_group[\"params\"]:\n            state = optimizer.state[param]\n            for key, value in state.items():\n                if isinstance(value, torch.Tensor):\n                    state[key] = value.to(\"cpu\", non_blocking=True)\n\n\n@torch.no_grad()\ndef load_fsdp_optimizer(optimizer, device_id):\n    if not optimizer.state:\n        return\n    for param_group in optimizer.param_groups:\n        for param in param_group[\"params\"]:\n            state = optimizer.state[param]\n            for key, value in state.items():\n                if isinstance(value, torch.Tensor):\n                    state[key] = value.to(device_id, non_blocking=True)\n\n\n@contextmanager\ndef meta_device_init():\n    \"\"\"\n    Create model parameters with meta device.\n\n    Note buffers in model will still be initialized in default device (e.g., CPU),\n    since the buffers can be non-persistent and filled with expected values that can\n    NOT be captured in meta device.\n    \"\"\"\n    device = torch.device(\"meta\")\n    old_register_parameter = nn.Module.register_parameter\n    registered = set()\n\n    def register_empty_parameter(module, name, param):\n        old_register_parameter(module, name, param)\n        # we will skip register shared parameters as it\n        # is already registered previously\n        if param is not None and param not in registered:\n            param_cls = type(module._parameters[name])\n            kwargs = module._parameters[name].__dict__\n            kwargs[\"requires_grad\"] = param.requires_grad\n            module._parameters[name] = param_cls(module._parameters[name].to(device), **kwargs)\n            registered.add(module._parameters[name])\n\n    try:\n        nn.Module.register_parameter = register_empty_parameter\n        yield\n    finally:\n        registered.clear()\n        nn.Module.register_parameter = old_register_parameter\n\n\ndef parallel_load_safetensors(filepath):\n    \"\"\"\n    Parallel load safetensors from huggingface checkpoint\n\n    Huggingface checkpoint contains:\n\n    - config.json: a json file for model configuration\n    - model.safetensor.index.json: a json file for safetensors (parameters & buffers) index\n    - model-000x-of-ooxx.safetensors: a binary file for safetensors (parameters & buffers) chunks\n\n    Or (when model is small),\n\n    - model.safetensors: a binary file for all parameters and buffers\n\n    Each rank will own a part of model chunks and load them directly into GPU memory.\n    \"\"\"\n    from safetensors.torch import load_file\n\n    safetensors2param = {}\n\n    index_file = os.path.join(filepath, \"model.safetensors.index.json\")\n    if os.path.exists(index_file):\n        index = json.load(open(index_file, \"rb\"))\n        for param_name, filename in index[\"weight_map\"].items():\n            safetensors2param.setdefault(filename, []).append(param_name)\n    else:\n        # in this case, the model is small and we can load it all at once\n        param_file = os.path.join(filepath, \"model.safetensors\")\n        assert os.path.exists(param_file), f\"Cannot find {param_file}\"\n        states = load_file(param_file)\n        for param_name in states:\n            safetensors2param.setdefault(\"model.safetensors\", []).append(param_name)\n        del states\n\n    total_files = len(safetensors2param)\n    ckpt_chunks = sorted(safetensors2param.keys())\n    world_size = dist.get_world_size()\n    size = int(math.ceil(total_files / world_size))\n    ckpt_chunks = [ckpt_chunks[rank * size : rank * size + size] for rank in range(world_size)]\n\n    shard_states = {}\n    device = get_device_id()\n    for rank, files in enumerate(ckpt_chunks):\n        if rank == dist.get_rank():\n            for file in files:\n                file = os.path.join(filepath, file)\n                states = load_file(file, device=device)\n                # print(f\"rank {rank} loading {file}...\")\n                shard_states.update(states)\n        else:\n            for file in files:\n                for param_name in safetensors2param[file]:\n                    shard_states[param_name] = rank\n    return shard_states\n\n\ndef parallel_init_module_fn(module: torch.nn.Module, shard_states: dict[str, torch.nn.Parameter]):\n    \"\"\"\n    Generate a function to initialize sub-modules in the `module` with `shard_states`\n    from huggingface checkpoint.\n\n    Args:\n        module (torch.nn.Module): the global module to be initialized\n        shard_states (Dict[str, torch.nn.Parameter]): the shard states from huggingface checkpoint\n\n    Returns:\n        init_fn (Callable): a function to initialize sub-modules in the `module` with `shard_states`\n    \"\"\"\n\n    state2fqn = {}\n    for name, state in itertools.chain(\n        module.named_parameters(remove_duplicate=False), module.named_buffers(remove_duplicate=False)\n    ):\n        state2fqn.setdefault(state, []).append(name)\n    # remove standalone parameters and buffers\n    shared = {s for s, names in state2fqn.items() if len(names) > 1}\n    materialized_states = {}\n\n    @torch.no_grad()\n    def create_and_sync_state(param_name, state, is_param):\n        assert param_name in shard_states, f\"{param_name} not loaded\"\n        device = get_device_id()\n        if is_param:\n            param = torch.nn.Parameter(torch.empty_like(state.data, device=device), requires_grad=state.requires_grad)\n        else:  # buffer\n            param = torch.empty_like(state.data, device=device)\n        loaded = shard_states[param_name]\n        if isinstance(loaded, torch.nn.Parameter | torch.Tensor):\n            # NOTE: loaded.dtype can be different with param.dtype\n            param.data.copy_(loaded.data)\n            dist.broadcast(param.data, src=dist.get_rank())\n        else:\n            assert isinstance(loaded, int)  # the rank that holds the state\n            dist.broadcast(param.data, src=loaded)\n        shard_states.pop(param_name)\n        del loaded\n        return param\n\n    def init_fn(sub_mod: torch.nn.Module, recurse: bool = True):\n        param_and_buffers = tuple(sub_mod.named_parameters(recurse=False)) + tuple(sub_mod.named_buffers(recurse=False))\n        # param_and_buffers = sorted(sub_mod.named_parameters(recurse=False), key=lambda x: x[0])\n        for name, state in param_and_buffers:\n            if not state.is_meta:\n                continue\n            is_param = name in sub_mod._parameters\n            fqn = state2fqn[state].pop(0)\n            # non-persistent buffers will not be saved in state dict, we can safely skip it\n            if (not is_param) and fqn not in shard_states:\n                if state.is_meta:\n                    raise RuntimeError(\n                        f\"find a non-persistent buffer ({fqn}) initiated with device meta. Such buffer is not saved \"\n                        f\"in checkpoint and user should guarantee to init in CPU / GPU device.\"\n                    )\n                continue\n            # for shared parameter, we get it from the first time it is created\n            if state in shared:\n                if state not in materialized_states:\n                    materialized_states[state] = create_and_sync_state(fqn, state, is_param)\n                else:\n                    if fqn in shard_states:\n                        shard_states.pop(fqn)\n                materialize_state = materialized_states[state]\n            # for not shared parameter, we create it directly\n            else:\n                materialize_state = create_and_sync_state(fqn, state, is_param)\n            if is_param:\n                sub_mod._parameters[name] = materialize_state\n            else:\n                sub_mod._buffers[name] = materialize_state\n        if recurse:\n            for module in sub_mod.children():\n                init_fn(module, recurse=True)\n\n        # for debug\n        # if len(shard_states) == 0: print(\"clear\")\n        return sub_mod\n\n    return init_fn\n\n\ndef fsdp_version(model):\n    if isinstance(model, FSDP):\n        return 1\n    elif isinstance(model, FSDPModule):\n        return 2\n    else:\n        return 0\n\n\ndef get_fsdp_state_ctx(model, state_type, state_cfg, optim_cfg):\n    if fsdp_version(model) == 1:\n        return FSDP.state_dict_type(model, state_type, state_cfg, optim_cfg)\n    else:\n        return nullcontext()\n\n\ndef get_fsdp_full_state_dict(model: torch.nn.Module, offload_to_cpu: bool = True, rank0_only: bool = True):\n    \"\"\"\n    Get the full state dict from an FSDP model.\n\n    Args:\n        model (torch.nn.Module): The FSDP model to get state dict from\n        offload_to_cpu (bool, optional): Whether to offload the state dict to CPU. Defaults to True.\n        rank0_only (bool, optional): Whether to only get state dict on rank 0. Defaults to True.\n\n    Returns:\n        dict: The full state dict of the model\n\n    Raises:\n        NotImplementedError: If the FSDP version is unknown\n    \"\"\"\n    if fsdp_version(model) == 1:\n        from torch.distributed.fsdp import FullStateDictConfig, StateDictType\n\n        state_dict_config = FullStateDictConfig(offload_to_cpu=offload_to_cpu, rank0_only=rank0_only)\n        with get_fsdp_state_ctx(\n            model, state_type=StateDictType.FULL_STATE_DICT, state_cfg=state_dict_config, optim_cfg=None\n        ):\n            state_dict = model.state_dict()\n        return state_dict\n    elif fsdp_version(model) == 2 or fsdp_version(model) == 0:\n        from torch.distributed.checkpoint.state_dict import StateDictOptions, get_model_state_dict\n\n        state_dict_config = StateDictOptions(\n            full_state_dict=True, cpu_offload=offload_to_cpu, broadcast_from_rank0=not rank0_only\n        )\n        state_dict = get_model_state_dict(model, options=state_dict_config)\n        return state_dict\n    else:\n        raise NotImplementedError(f\"Unknown FSDP version {fsdp_version}\")\n\n\ndef fsdp2_load_full_state_dict(model: torch.nn.Module, full_state: dict, device_mesh=None, cpu_offload=None):\n    \"\"\"\n    Loads the full state dict (could be only on rank 0) into the sharded model. This is done by broadcasting the\n    parameters from rank 0 to all other ranks. This function modifies the model in-place.\n\n    Args:\n        model (`torch.nn.Module`): The model to load the state dict into\n        full_state (`dict`): The full state dict to load, can only be on rank 0\n    \"\"\"\n\n    if version.parse(torch.__version__) >= version.parse(\"2.7.0\"):\n        from torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict\n    else:\n        # official torch 2.6.0 set_model_state_dict API leads to OOM\n        # use torch 2.7.0 copy from verl/third_party/torch/distributed/checkpoint\n        from verl.third_party.torch.distributed.checkpoint.state_dict import StateDictOptions, set_model_state_dict\n\n    # To broadcast, it needs to be instantiated in the GPU.\n    if dist.get_rank() == 0:\n        model = model.to(device=get_device_id(), non_blocking=True)\n    else:\n        model = model.to_empty(device=get_device_id())\n\n    cpu_offload = cpu_offload is not None\n    options = StateDictOptions(full_state_dict=True, cpu_offload=cpu_offload, broadcast_from_rank0=True)\n    set_model_state_dict(model, full_state, options=options)\n\n    # rotary_emb is not in state_dict, so we need to broadcast it manually\n    for name, buf in model.named_buffers():\n        dist.broadcast(buf, src=0)\n\n    if cpu_offload:\n        model.to(\"cpu\", non_blocking=True)\n        for buf in model.buffers():\n            buf.data = buf.data.to(get_device_id())\n\n\n@contextmanager\ndef maybe_patch_fsdp_module(model):\n    if fully_shard_module is None:\n        yield\n        return\n\n    orig_fsdp_module = fully_shard_module.FSDPModule\n\n    class FSDPModuleABC(ABC, orig_fsdp_module):\n        pass\n\n    try:\n        if isinstance(model, ABC):\n            fully_shard_module.FSDPModule = FSDPModuleABC\n        yield\n    finally:\n        fully_shard_module.FSDPModule = orig_fsdp_module\n\n\ndef _select_fsdp2_wrap_targets(model, fsdp_transformer_layer_cls_to_wrap):\n    \"\"\"Select modules to wrap individually with fully_shard in FSDP2.\n\n    Matches transformer layers by class name, and embed_tokens/lm_head by name\n    (with isinstance fallback). Name-based matching is needed because peft wraps\n    embed_tokens in ModulesToSaveWrapper, breaking isinstance(module, nn.Embedding).\n    When tie_word_embeddings is True, embed_tokens and lm_head share weights and\n    must not be wrapped separately.\n    \"\"\"\n    _tie = getattr(model.config, \"tie_word_embeddings\", False)\n    _wrap_by_name = set() if _tie else {\"embed_tokens\", \"lm_head\"}\n\n    modules = []\n    for name, module in model.named_modules():\n        leaf_name = name.rsplit(\".\", 1)[-1] if \".\" in name else name\n        if (\n            module.__class__.__name__ in fsdp_transformer_layer_cls_to_wrap\n            or (isinstance(module, nn.Embedding) and not _tie)\n            or (leaf_name in _wrap_by_name and hasattr(module, \"weight\"))\n        ):\n            modules.append(module)\n    return modules\n\n\ndef apply_fsdp2(model, fsdp_kwargs, config):\n    \"\"\"model: AutoModelForCausalLM\"\"\"\n    assert CPUOffloadPolicy is not None, \"PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)\"\n\n    default_transformer_cls_names_to_wrap = getattr(model, \"_no_split_modules\", None)\n    fsdp_transformer_layer_cls_to_wrap = config.get(\"wrap_policy\", {}).get(\n        \"transformer_layer_cls_to_wrap\", default_transformer_cls_names_to_wrap\n    )\n\n    if isinstance(fsdp_transformer_layer_cls_to_wrap, str):\n        fsdp_transformer_layer_cls_to_wrap = [fsdp_transformer_layer_cls_to_wrap]\n\n    assert len(fsdp_transformer_layer_cls_to_wrap) > 0 and fsdp_transformer_layer_cls_to_wrap[0] is not None\n\n    modules = _select_fsdp2_wrap_targets(model, fsdp_transformer_layer_cls_to_wrap)\n\n    for idx, module in enumerate(modules):\n        # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:\n        #     print(f\"wrap module {module.__class__.__name__}\")\n        with maybe_patch_fsdp_module(module):\n            fully_shard(module, **fsdp_kwargs)\n\n    # if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:\n    #     print(f\"wrap module {model.__class__.__name__}\")\n    with maybe_patch_fsdp_module(model):\n        fully_shard(model, **fsdp_kwargs)  # fsdp2 will not reshard_after_forward for root module\n\n\ndef get_shard_placement_fn(fsdp_size):\n    \"\"\"Choose the dimension that can divide fsdp_size to avoid padding\"\"\"\n\n    def shard_placement_fn(param):\n        shape = list(param.shape)\n        for i in range(len(shape)):\n            if shape[i] % fsdp_size == 0:\n                return Shard(i)\n        return Shard(0)\n\n    return shard_placement_fn\n\n\ndef fsdp2_clip_grad_norm_(parameters, max_norm, norm_type=2.0, error_if_nonfinite=False, foreach=None):\n    \"\"\"torch.nn.utils.clip_grad_norm_ cann't run on cpu parameter DTensor\"\"\"\n    from torch.nn.utils.clip_grad import _clip_grads_with_norm_, _get_total_norm\n\n    if isinstance(parameters, torch.Tensor):\n        parameters = [parameters]\n    else:\n        # prevent generators from being exhausted\n        parameters = list(parameters)\n    grads = [p.grad for p in parameters if p.grad is not None]\n    total_norm = _get_total_norm(grads, norm_type, error_if_nonfinite, foreach)\n    total_norm = total_norm.to(get_device_id(), non_blocking=True)\n    _clip_grads_with_norm_(parameters, max_norm, total_norm, foreach)\n    return total_norm\n\n\ndef layered_summon_lora_params(fsdp_module) -> OrderedDict:\n    from peft.utils.save_and_load import get_peft_model_state_dict\n\n    def __prefix_submodules(module, prefix):\n        for name, submodule in module.named_modules():\n            if name.startswith(prefix) and \".\" not in name[len(prefix) :]:\n                yield name, submodule\n\n    lora_params = OrderedDict()\n    prefix_list = [\n        # fsdp\n        \"_fsdp_wrapped_module.base_model.model.\",\n        \"_fsdp_wrapped_module.base_model.model.model.\",\n        \"_fsdp_wrapped_module.base_model.model.model.layers.\",\n        \"_fsdp_wrapped_module.base_model.model.model.language_model.layers.\",\n        # fsdp2\n        \"base_model.model.\",\n        \"base_model.model.model.\",\n        \"base_model.model.model.layers.\",\n        \"base_model.model.model.language_model.layers.\",\n    ]\n    peft_model = getattr(fsdp_module, \"_fsdp_wrapped_module\", fsdp_module)\n    for prefix in prefix_list:\n        for name, submodule in __prefix_submodules(fsdp_module, prefix):\n            prefix = name.replace(\"_fsdp_wrapped_module.base_model.model.\", \"base_model.model.\")\n            if name.endswith(\".model\") or name.endswith(\".layers\"):\n                continue\n            if fsdp_version(submodule) > 0:\n                with FSDP.summon_full_params(submodule, writeback=False):\n                    sub_lora_params = get_peft_model_state_dict(peft_model, state_dict=submodule.state_dict())\n                    sub_lora_params = {\n                        f\"{prefix}.{name}\": param.full_tensor().detach().cpu()\n                        if hasattr(param, \"full_tensor\")\n                        else param.detach().cpu()\n                        for name, param in sub_lora_params.items()\n                    }\n                    lora_params.update(sub_lora_params)\n                    submodule._is_root = False\n                get_torch_device().empty_cache()\n    return lora_params\n\n\ndef collect_lora_params(module: FSDP, layered_summon: bool, base_sync_done: bool) -> OrderedDict:\n    \"\"\"\n    collect lora params or full params if base model is not ready in vllm\n    work with if isinstance(self.module._fsdp_wrapped_module, PeftModel)\n    \"\"\"\n    from peft.utils.save_and_load import get_peft_model_state_dict\n\n    lora_params = OrderedDict()\n    peft_model = getattr(module, \"_fsdp_wrapped_module\", module)\n    if fsdp_version(module) > 0:\n        if layered_summon:\n            if not base_sync_done:\n                raise ValueError(\n                    \"To use layered_summon, you must make sure base-model is preloaded in vllm, e.g. let \"\n                    \"rollout.load_format=safetensors\"\n                )\n            lora_params = layered_summon_lora_params(module)\n        else:\n            with FSDP.summon_full_params(module, writeback=False):\n                if base_sync_done:\n                    lora_params = get_peft_model_state_dict(peft_model)\n                    lora_params = {\n                        name: param.full_tensor().detach().cpu()\n                        if hasattr(param, \"full_tensor\")\n                        else param.detach().cpu()\n                        for name, param in lora_params.items()\n                    }\n                else:\n                    model = peft_model.base_model.model\n                    orig_dev = \"cpu\" if \"cpu\" in str(next(model.parameters()).device) else get_device_name()\n                    model = model.to(\"cpu\")\n                    for name, param in model.state_dict().items():\n                        if any(x in name for x in [\"_flat_param\", \"lora_\"]):\n                            continue\n                        name = name.replace(\"_fsdp_wrapped_module.\", \"\").replace(\".base_layer\", \"\")\n                        lora_params[name] = (\n                            param.full_tensor().detach().cpu()\n                            if hasattr(param, \"full_tensor\")\n                            else param.detach().cpu()\n                        )\n                    model = model.to(orig_dev)\n            get_torch_device().empty_cache()\n    else:\n        if base_sync_done:\n            lora_params = get_peft_model_state_dict(peft_model)\n        else:\n            model = peft_model.base_model.model\n            orig_dev = \"cpu\" if \"cpu\" in str(next(model.parameters()).device) else get_device_name()\n            model = model.to(\"cpu\")\n            for name, param in model.state_dict().items():\n                if any(x in name for x in [\"_flat_param\", \"lora_\"]):\n                    continue\n                name = name.replace(\"_fsdp_wrapped_module.\", \"\").replace(\".base_layer\", \"\")\n                lora_params[name] = param.detach().cpu()\n            model = model.to(orig_dev)\n    return lora_params\n\n\ndef replace_lora_wrapper(k, peft_config):\n    \"\"\"Replace LoRA parameter keys with base layer equivalents.\n\n    Transforms LoRA parameter names to their corresponding base layer\n    names for proper weight loading in vLLM when base model sync is not done.\n\n    Args:\n        k (str): Original parameter key name.\n\n    Returns:\n        str: Transformed parameter key for base layer.\n    \"\"\"\n    stacked_params = [\"q_proj\", \"k_proj\", \"v_proj\", \"o_proj\", \"gate_proj\", \"up_proj\", \"down_proj\"]\n    if k.endswith(\".weight\"):\n        module_k = k[: -len(\".weight\")]\n        if check_exclude_modules(peft_config, module_k):\n            return k\n        elif any([module_k.endswith(s) for s in stacked_params]) or check_target_modules(peft_config, module_k):\n            return f\"{module_k}.base_layer.weight\"\n    if k.endswith(\".bias\"):\n        module_k = k[: -len(\".bias\")]\n        if check_exclude_modules(peft_config, module_k):\n            return k\n        elif any([module_k.endswith(s) for s in stacked_params]) or check_target_modules(peft_config, module_k):\n            return f\"{module_k}.base_layer.bias\"\n    return k\n\n\ndef set_reshard_after_forward(module: FSDPModule, reshard_after_forward: bool, recurse: bool = True) -> None:\n    \"\"\"\n    Sets if the module should reshard parameters after forward. This can be\n    used to change the ``reshard_after_forward`` FSDP arg at runtime. For\n    example, this can be used to set the FSDP root module's value to\n    ``True`` (since it is otherwise specially set to ``False``), or it can\n    set an FSDP module's value to ``False`` for running evals and set back\n    to ``True`` for training.\n\n    Args:\n        reshard_after_forward (bool): Whether to reshard parameters after\n            forward.\n        recurse (bool): Whether to set for all FSDP submodules or just the\n            passed-in module.\n\n    ---\n    Copied from https://github.com/pytorch/pytorch/blob/main/torch/distributed/fsdp/_fully_shard/_fully_shard.py to\n    address the absence of the set_reshard_after_forward function in torch versions earlier than 2.8.0.\n    \"\"\"\n\n    if not isinstance(reshard_after_forward, bool):\n        raise ValueError(f\"reshard_after_forward should be a bool, got {type(reshard_after_forward)}\")\n    self_module = cast(nn.Module, module)\n    modules = list(self_module.modules()) if recurse else [self_module]\n    for module in modules:\n        if isinstance(module, FSDPModule):\n            state = module._get_fsdp_state()\n            state._auto_reshard_after_forward = False\n            if fsdp_param_group := state._fsdp_param_group:\n                fsdp_param_group.post_forward_mesh_info = _get_post_forward_mesh_info(\n                    reshard_after_forward, fsdp_param_group.mesh_info\n                )\n\n\ndef normalize_peft_param_name(params: dict) -> dict:\n    \"\"\"\n    Converts peft model parameter name to base parameter name\n    For example,\n        base_model.model.model.embed_tokens.weight -> model.embed_tokens.weight\n        base_model.model.model.layers.0.self_attn.q_proj.base_layer.weight -> model.layers.0.self_attn.q_proj.weight\n    and remove params such as base_model.model.model.layers.0.self_attn.q_proj.lora_A.default.weight,\n    base_model.model.model.layers.0.self_attn.q_proj.lora_B.default.weight\n    \"\"\"\n\n    def _normalize_peft_name(name: str) -> str:\n        return name.replace(\"base_model.model.\", \"\").replace(\"base_model.\", \"\").replace(\".base_layer\", \"\")\n\n    def _is_lora_key(name: str) -> bool:\n        # catch typical PEFT keys\n        return (\"lora_\" in name) or (\".adapter_\" in name)\n\n    params = [(_normalize_peft_name(k), v) for k, v in params.items()]\n    # strip any residual LoRA tensors\n    params = {k: v for k, v in params if not _is_lora_key(k)}\n    return params\n\n\ndef _merge_or_unmerge_lora_(module, merge: bool):\n    \"\"\"Merge or unmerge LoRA adapters in a module.\n\n    Args:\n        module: The module containing LoRA layers\n        merge: If True, merge LoRA into base model; if False, unmerge LoRA\n    \"\"\"\n    from peft.tuners.lora import LoraLayer\n\n    with torch.no_grad():\n        for m in module.modules():\n            if isinstance(m, LoraLayer):\n                is_merged = getattr(m, \"merged\", False)\n                if merge and not is_merged:\n                    m.merge()\n                elif (not merge) and is_merged:\n                    m.unmerge()\n\n\n# merged_adapters\ndef _clean_merged_lora_(module):\n    \"\"\"Cleans the merged lora adapters\"\"\"\n    from peft.tuners.lora import LoraLayer\n\n    with torch.no_grad():\n        for m in module.modules():\n            if isinstance(m, LoraLayer):\n                merged_adapters = getattr(m, \"merged_adapters\", False)\n                if merged_adapters:\n                    m.merged_adapters = []\n\n\ndef fsdp_merge_unmerge(module: nn.Module, do_merge: bool):\n    \"\"\"Merge or unmerge LoRA adapters in FSDP module.\n\n    For FSDP (v1), it gathers all model parameters to each device, which may cause OOM.\n    For FSDP2, it gathers model parameters layer-by-layer to reduce memory footprint.\n\n    Args:\n        module: The FSDP module to merge/unmerge LoRA adapters\n        do_merge: If True, merge LoRA into base model; if False, unmerge LoRA\n    \"\"\"\n    version = fsdp_version(module)\n    assert version in [1, 2], f\"fsdp_merge_unmerge requires FSDP module, got version {version}\"\n\n    if version == 1:\n        # Unshard → merge → Reshard\n        with FSDP.summon_full_params(module, writeback=True, with_grads=False):\n            _merge_or_unmerge_lora_(module, merge=do_merge)\n    else:\n        # FSDP2: Unshard → merge → Reshard layer-by-layer\n        for name, submodule in module.named_modules():\n            if isinstance(submodule, FSDPModule) and name != \"\":  # skip root model\n                with FSDP.summon_full_params(submodule, writeback=True, with_grads=False):\n                    _merge_or_unmerge_lora_(submodule, merge=do_merge)\n\n\ndef backup_base_model_weights(module):\n    \"\"\"Backup base model weights to CPU with LoRA temporarily disabled.\n\n    This function temporarily disables LoRA adapters, backs up the clean base model weights\n    to CPU, then re-enables the adapters.\n\n    Args:\n        module: The PEFT model with LoRA adapters\n\n    Returns:\n        dict: Dictionary mapping parameter name to CPU tensor backup of base model weights\n    \"\"\"\n    from peft import PeftModel\n\n    backup = {}\n    with torch.no_grad():\n        # Check if module is a PEFT model\n        if isinstance(module, PeftModel):\n            # Temporarily disable adapters to get clean base model weights\n            with module.disable_adapter():\n                # Backup base model weights (excluding lora parameters)\n                for name, param in module.named_parameters():\n                    if \"lora\" not in name.lower():\n                        backup[name] = param.data.clone().cpu()\n        else:\n            # For non-PEFT models, just backup all parameters\n            for name, param in module.named_parameters():\n                backup[name] = param.data.clone().cpu()\n    return backup\n\n\ndef restore_base_model_weights(module, backup):\n    \"\"\"Restore base model weights from CPU backup.\n\n    This function restores the base model weights from the CPU backup, effectively\n    undoing any LoRA merge operations.\n\n    Args:\n        module: The PEFT model with LoRA adapters\n        backup: Dictionary mapping parameter name to CPU tensor backup of base model weights\n    \"\"\"\n    with torch.no_grad():\n        for name, param in module.named_parameters():\n            if name in backup:\n                param.data.copy_(backup[name].to(param.device))\n\n\n@contextmanager\ndef merged_lora_context(actor, backup_adapters=False):\n    \"\"\"Context manager to temporarily merge LoRA adapters.\n\n    This context manager merges LoRA adapters into the base model weights,\n    performs operations (like syncing weights to vLLM), then restores the base model\n    weights from backup.\n\n    Args:\n        actor: The actor module with LoRA adapters to merge\n        backup_adapters: If True, backup base model weights (with LoRA disabled) before\n            merging and restore them after. This is more numerically stable than unmerging.\n\n    Yields:\n        None\n    \"\"\"\n    base_weights_backup = None\n    if backup_adapters:\n        # Backup base model weights with LoRA temporarily disabled\n        base_weights_backup = backup_base_model_weights(actor)\n\n    # Merge LoRA adapters into base model\n    fsdp_merge_unmerge(actor, do_merge=True)\n    try:\n        # Do work while merged (sync_to_vllm / generate / etc.)\n        yield\n    finally:\n        if backup_adapters and base_weights_backup is not None:\n            # Restore base model weights from CPU backup (effectively undoing the merge)\n            restore_base_model_weights(actor, base_weights_backup)\n            _clean_merged_lora_(actor)\n        else:\n            # Fall back to unmerge if no backup was made\n            fsdp_merge_unmerge(actor, do_merge=False)\n\n\ndef fsdp2_sharded_save_to_cpu(\n    model: torch.nn.Module,\n) -> tuple[dict[str, tuple[torch.Tensor, DTensorSpec]], DTensorSpec]:\n    \"\"\"\n    Sharded Save: Each process only saves the local DTensor shard from its own GPU to CPU memory.\n\n    Args:\n        model: FSDP2-wrapped model whose parameters are of DTensor type.\n\n    Returns:\n        cpu_sharded_state: Dictionary of CPU shards for the current process.\n                          Key = parameter name, Value = (CPU shard tensor, original DTensorSpec)\n        global_spec: DTensorSpec of the first parameter (used to verify global rules during loading)\n    \"\"\"\n    cpu_sharded_state = {}\n    global_spec = None  # Record global sharding rules (all parameters follow the same spec)\n\n    for param_name, param in model.named_parameters():\n        # Only process sharded parameters of DTensor type (core parameters of FSDP2)\n        if not isinstance(param, DTensor):\n            # Save non-sharded parameters (e.g., running_mean of BatchNorm) as local data\n            cpu_tensor = param.detach().cpu()\n            cpu_sharded_state[param_name] = (cpu_tensor, None)\n            continue\n\n        # Record global sharding rules (take spec of the first DTensor to ensure consistency)\n        if global_spec is None:\n            global_spec = param._spec\n            assert hasattr(global_spec, \"device_mesh\"), \"DTensorSpec must contain 'device_mesh' attribute\"\n            assert hasattr(global_spec, \"placements\"), \"DTensorSpec must contain 'placements' attribute\"\n\n        # 1. Extract local shard data from the current GPU (_local_tensor)\n        local_gpu_tensor = param._local_tensor  # Local shard attribute defined in your DTensor class\n        # 2. Move to CPU memory and detach from computation graph\n        local_cpu_tensor = local_gpu_tensor.detach().cpu()\n        # 3. Save CPU shard + original DTensorSpec (ensure sharding rules remain unchanged)\n        cpu_sharded_state[param_name] = (local_cpu_tensor, param._spec)\n\n    assert global_spec is not None, \"No DTensor-type parameters found in the model. FSDP2 sharding may not be enabled.\"\n    return cpu_sharded_state, global_spec\n\n\ndef fsdp2_sharded_load_from_cpu(\n    model: torch.nn.Module,\n    cpu_sharded_state: dict[str, tuple[torch.Tensor, Optional[DTensorSpec]]],\n    target_spec: DTensorSpec,\n) -> None:\n    \"\"\"\n    Sharded Load: Each process only loads the CPU shard it is responsible for to the GPU,\n                  keeping sharding rules unchanged.\n\n    Args:\n        model: FSDP2 model to be restored (must have the same structure as when saved)\n        cpu_sharded_state: Shard data read from CPU memory by the current process\n                          (from fsdp2_sharded_save_to_cpu)\n        target_spec: Global DTensorSpec from saving (used to verify sharding rule consistency)\n    \"\"\"\n    # Verify device_mesh consistency (core: ensure loaded shards map to original GPUs)\n    current_device_mesh = None\n    for param in model.parameters():\n        if isinstance(param, DTensor):\n            current_device_mesh = param._spec.device_mesh\n            break\n    assert current_device_mesh is not None, \"DTensor parameters not initialized in the model to be loaded\"\n    assert current_device_mesh == target_spec.device_mesh, (\n        f\"device_mesh mismatch during loading! Original: {target_spec.device_mesh}, Current: {current_device_mesh}\"\n    )\n\n    for param_name, param in model.named_parameters():\n        # Skip parameters not in the saved state (e.g., newly added parameters)\n        if param_name not in cpu_sharded_state:\n            continue\n\n        # Extract CPU shard data and original Spec\n        local_cpu_tensor, saved_spec = cpu_sharded_state[param_name]\n\n        # Handle different parameter types: DTensor sharded parameters vs. regular parameters\n        if isinstance(param, DTensor):\n            # 1. Verify sharding rule consistency (placements must match original Spec)\n            assert saved_spec is not None, f\"DTensorSpec missing in saved state for parameter {param_name}\"\n            assert saved_spec.placements == target_spec.placements, (\n                f\"Sharding strategy mismatch for parameter {param_name} (conflicts with global rules)!\"\n            )\n\n            # 2. Move CPU shard data to the current GPU (device of param._local_tensor)\n            target_device = param._local_tensor.device\n            local_gpu_tensor = local_cpu_tensor.to(target_device)\n\n            # 3. Restore to DTensor's local shard (directly copy to _local_tensor, keep spec unchanged)\n            param._local_tensor.copy_(local_gpu_tensor)\n\n        else:\n            # Regular parameters: load directly to original device\n            target_device = param.device\n            param.data.copy_(local_cpu_tensor.to(target_device))\n\n    # Process synchronization: ensure all processes complete loading before proceeding\n    dist.barrier()\n"
  },
  {
    "path": "verl/utils/groupwise.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nGroup-wise helpers for RL training utilities.\n\nPublic API:\n    - as_torch_index(index, device=None) -> torch.LongTensor\n    - group_mean_std(scores, gidx, eps=1e-6, device=None) -> (mean_g, std_g, count_g)\n\nDefault device policy:\n    - If `device` is None:\n        * In pytest (detected by env \"PYTEST_CURRENT_TEST\"): use CPU.\n        * Else if CUDA is available: use CUDA.\n        * Else: use CPU.\n    - You can override via env \"VERL_FORCE_DEVICE\" (e.g., \"cuda:0\" / \"cpu\").\n\nNotes:\n- as_torch_index: canonicalizes arbitrary group labels to a contiguous 1-D torch.long\n  tensor in range [0..G-1]. Robust to torch/numpy/list/tuple, ints/floats/bools,\n  numeric strings, UUIDs, mixed object arrays. Near-integer floats (|x-round(x)|<=1e-6)\n  are rounded; otherwise factorization is applied.\n- group_mean_std: pure-PyTorch per-group mean/std with Bessel correction for variance\n  (denominator max(count-1, 1)). Singleton groups fallback to mean=0, std=1 for\n  compatibility with common “native” conventions.\n\"\"\"\n\nfrom __future__ import annotations\n\nimport os\nfrom typing import Any, Optional\n\nimport numpy as np\nimport torch\n\nfrom verl.utils.device import get_device_name\n\n__all__ = [\"as_torch_index\", \"group_mean_std\"]\n\n\ndef _resolve_device(explicit: Optional[torch.device | str]) -> torch.device:\n    \"\"\"\n    Resolve device according to policy described in the module docstring.\n    Priority:\n      1) explicit argument\n      2) VERL_FORCE_DEVICE env\n      3) pytest detection -> cpu\n      4) cuda if available, else cpu\n    \"\"\"\n    if explicit is not None:\n        return torch.device(explicit)\n\n    forced = os.getenv(\"VERL_FORCE_DEVICE\")\n    if forced:\n        return torch.device(forced)\n\n    # Heuristic: pytest sets PYTEST_CURRENT_TEST\n    if \"PYTEST_CURRENT_TEST\" in os.environ:\n        return torch.device(\"cpu\")\n\n    return torch.device(get_device_name())\n\n\ndef _to_1d_numpy_object_array(x: Any) -> np.ndarray:\n    \"\"\"Best-effort: convert arbitrary input into a 1-D numpy array; fallback to object dtype.\"\"\"\n    try:\n        arr = np.asarray(x)\n    except Exception:\n        try:\n            arr = np.array(list(x), dtype=object)\n        except Exception:\n            arr = np.array([x], dtype=object)\n    if arr.ndim != 1:\n        arr = arr.reshape(-1)\n    return arr\n\n\ndef as_torch_index(index: Any, device: torch.device | str | None = None) -> torch.Tensor:\n    \"\"\"\n    Convert arbitrary group labels to a contiguous 1-D torch.long tensor (0..G-1).\n\n    Args:\n        index: Any iterable of labels or tensor/ndarray.\n        device: Target device; if None, resolved via _resolve_device().\n\n    Returns:\n        torch.LongTensor with shape (N,)\n    \"\"\"\n    target = _resolve_device(device)\n\n    # ---------- Fast path: torch.Tensor ----------\n    if isinstance(index, torch.Tensor):\n        t = index.reshape(-1)\n        if t.dtype in (\n            torch.int64,\n            torch.int32,\n            torch.int16,\n            torch.int8,\n            getattr(torch, \"uint8\", torch.uint8),\n            torch.bool,\n        ):\n            return t.to(device=target, dtype=torch.long)\n\n        if t.dtype in (torch.float16, torch.float32, torch.float64, torch.bfloat16):\n            t64 = t.to(dtype=torch.float64)\n            rounded = torch.round(t64)\n            if torch.allclose(t64, rounded, rtol=0.0, atol=1e-6):\n                return rounded.to(device=target, dtype=torch.long)\n            arr = np.array([str(x.item()) for x in t], dtype=object)\n        else:\n            arr = np.array([str(x.item()) if hasattr(x, \"item\") else str(x) for x in t], dtype=object)\n\n    else:\n        # ---------- Non-torch: go through numpy ----------\n        arr = _to_1d_numpy_object_array(index)\n\n        # Pure integers (incl. bool)\n        if arr.dtype != object and np.issubdtype(arr.dtype, np.integer):\n            return torch.from_numpy(arr.astype(np.int64, copy=False)).to(device=target)\n\n        # Floats nearly equal to integers\n        if arr.dtype != object and np.issubdtype(arr.dtype, np.floating):\n            arr64 = arr.astype(np.float64, copy=False)\n            rounded = np.rint(arr64)\n            if np.allclose(arr64, rounded, rtol=0.0, atol=1e-6):\n                return torch.from_numpy(rounded.astype(np.int64)).to(device=target)\n            # fall through\n\n        # Try numeric string coercion\n        try:\n            coerced = arr.astype(np.int64)\n            return torch.from_numpy(coerced).to(device=target)\n        except Exception:\n            pass\n\n        if arr.dtype != object:\n            arr = arr.astype(object)\n\n    # ---------- Factorization (UUIDs / mixed types / arbitrary labels) ----------\n    try:\n        _, inv = np.unique(arr, return_inverse=True)\n    except Exception:\n        sarr = np.array([str(x) for x in arr], dtype=object)\n        _, inv = np.unique(sarr, return_inverse=True)\n\n    inv = inv.astype(np.int64, copy=False)\n    return torch.from_numpy(inv).to(device=target)\n\n\n@torch.no_grad()\ndef group_mean_std(\n    scores: torch.Tensor,\n    gidx: torch.Tensor,\n    eps: float = 1e-6,\n    device: torch.device | str | None = None,\n) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:\n    \"\"\"\n    Compute per-group mean/std/count in pure PyTorch.\n\n    mean_g = sum / count\n    std_g  = sqrt( max( (sum2 - sum^2/count) / max(count-1, 1), eps ) )\n\n    Singleton groups fallback to mean=0, std=1.\n\n    Args:\n        scores: (N,) float tensor.\n        gidx  : (N,) long/int tensor with group indices (0..G-1).\n        eps   : Numerical floor for variance.\n        device: Target device; if None, resolved via _resolve_device().\n\n    Returns:\n        mean_g: (G,) float32\n        std_g : (G,) float32\n        count : (G,) float32\n    \"\"\"\n    target = _resolve_device(device)\n\n    scores = scores.reshape(-1).to(device=target, dtype=torch.float32)\n    gidx = gidx.reshape(-1).to(device=target, dtype=torch.long)\n\n    if scores.numel() != gidx.numel():\n        raise ValueError(f\"scores and gidx length mismatch: {scores.numel()} vs {gidx.numel()}\")\n\n    G = int(torch.max(gidx).item()) + 1 if gidx.numel() > 0 else 0\n    if G == 0:\n        # Return empty tensors on the selected device\n        empty = torch.empty(0, device=target, dtype=torch.float32)\n        return empty, empty, empty\n\n    ones = torch.ones_like(scores, dtype=torch.float32)\n\n    count = torch.zeros(G, device=target, dtype=torch.float32).index_add_(0, gidx, ones)\n    s1 = torch.zeros(G, device=target, dtype=torch.float32).index_add_(0, gidx, scores)\n    s2 = torch.zeros(G, device=target, dtype=torch.float32).index_add_(0, gidx, scores * scores)\n\n    mean = s1 / count.clamp_min(1.0)\n    var_num = s2 - (s1 * s1) / count.clamp_min(1.0)\n    denom = (count - 1.0).clamp_min(1.0)\n    var = var_num / denom\n    std = torch.sqrt(torch.clamp(var, min=eps))\n\n    # Singleton groups: mean=0, std=1\n    single = count <= 1.0\n    if torch.any(single):\n        mean = mean.clone()\n        std = std.clone()\n        mean[single] = 0.0\n        std[single] = 1.0\n\n    return mean, std, count\n"
  },
  {
    "path": "verl/utils/hdfs_io.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nimport shutil\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_SFT_LOGGING_LEVEL\", \"WARN\"))\n\n_HDFS_PREFIX = \"hdfs://\"\n\n_HDFS_BIN_PATH = shutil.which(\"hdfs\")\n\n\ndef exists(path: str, **kwargs) -> bool:\n    r\"\"\"Works like os.path.exists() but supports hdfs.\n\n    Test whether a path exists. Returns False for broken symbolic links.\n\n    Args:\n        path (str): path to test\n\n    Returns:\n        bool: True if the path exists, False otherwise\n    \"\"\"\n    if _is_non_local(path):\n        return _exists(path, **kwargs)\n    return os.path.exists(path)\n\n\ndef _exists(file_path: str):\n    \"\"\"hdfs capable to check whether a file_path is exists\"\"\"\n    if file_path.startswith(\"hdfs\"):\n        return _run_cmd(_hdfs_cmd(f\"-test -e {file_path}\")) == 0\n    return os.path.exists(file_path)\n\n\ndef makedirs(name, mode=0o777, exist_ok=False, **kwargs) -> None:\n    r\"\"\"Works like os.makedirs() but supports hdfs.\n\n    Super-mkdir; create a leaf directory and all intermediate ones.  Works like\n    mkdir, except that any intermediate path segment (not just the rightmost)\n    will be created if it does not exist. If the target directory already\n    exists, raise an OSError if exist_ok is False. Otherwise no exception is\n    raised.  This is recursive.\n\n    Args:\n        name (str): directory to create\n        mode (int): file mode bits\n        exist_ok (bool): if True, do not raise an exception if the directory already exists\n        kwargs: keyword arguments for hdfs\n\n    \"\"\"\n    if _is_non_local(name):\n        # TODO(haibin.lin):\n        # - handle OSError for hdfs(?)\n        # - support exist_ok for hdfs(?)\n        _mkdir(name, **kwargs)\n    else:\n        os.makedirs(name, mode=mode, exist_ok=exist_ok)\n\n\ndef _mkdir(file_path: str) -> bool:\n    \"\"\"hdfs mkdir\"\"\"\n    if file_path.startswith(\"hdfs\"):\n        _run_cmd(_hdfs_cmd(f\"-mkdir -p {file_path}\"))\n    else:\n        os.makedirs(file_path, exist_ok=True)\n    return True\n\n\ndef copy(src: str, dst: str, **kwargs) -> bool:\n    r\"\"\"Works like shutil.copy() for file, and shutil.copytree for dir, and supports hdfs.\n\n    Copy data and mode bits (\"cp src dst\"). Return the file's destination.\n    The destination may be a directory.\n    If source and destination are the same file, a SameFileError will be\n    raised.\n\n    Arg:\n        src (str): source file path\n        dst (str): destination file path\n        kwargs: keyword arguments for hdfs copy\n\n    Returns:\n        str: destination file path\n\n    \"\"\"\n    if _is_non_local(src) or _is_non_local(dst):\n        # TODO(haibin.lin):\n        # - handle SameFileError for hdfs files(?)\n        # - return file destination for hdfs files\n        return _copy(src, dst)\n    else:\n        if os.path.isdir(src):\n            return shutil.copytree(src, dst, **kwargs)\n        else:\n            return shutil.copy(src, dst, **kwargs)\n\n\ndef _copy(from_path: str, to_path: str, timeout: int = None) -> bool:\n    if to_path.startswith(\"hdfs\"):\n        if from_path.startswith(\"hdfs\"):\n            returncode = _run_cmd(_hdfs_cmd(f\"-cp -f {from_path} {to_path}\"), timeout=timeout)\n        else:\n            returncode = _run_cmd(_hdfs_cmd(f\"-put -f {from_path} {to_path}\"), timeout=timeout)\n    else:\n        if from_path.startswith(\"hdfs\"):\n            returncode = _run_cmd(\n                _hdfs_cmd(\n                    f\"-get \\\n                {from_path} {to_path}\"\n                ),\n                timeout=timeout,\n            )\n        else:\n            try:\n                shutil.copy(from_path, to_path)\n                returncode = 0\n            except shutil.SameFileError:\n                returncode = 0\n            except Exception as e:\n                logger.warning(f\"copy {from_path} {to_path} failed: {e}\")\n                returncode = -1\n    return returncode == 0\n\n\ndef _run_cmd(cmd: str, timeout=None):\n    return os.system(cmd)\n\n\ndef _hdfs_cmd(cmd: str) -> str:\n    return f\"{_HDFS_BIN_PATH} dfs {cmd}\"\n\n\ndef _is_non_local(path: str):\n    return path.startswith(_HDFS_PREFIX)\n"
  },
  {
    "path": "verl/utils/import_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUtilities to check if packages are available.\nWe assume package availability won't change during runtime.\n\"\"\"\n\nimport importlib\nimport importlib.util\nimport os\nimport warnings\nfrom functools import cache, wraps\nfrom typing import Optional\n\n\n@cache\ndef is_megatron_core_available():\n    try:\n        mcore_spec = importlib.util.find_spec(\"megatron.core\")\n    except ModuleNotFoundError:\n        mcore_spec = None\n    return mcore_spec is not None\n\n\n@cache\ndef is_vllm_available():\n    try:\n        vllm_spec = importlib.util.find_spec(\"vllm\")\n    except ModuleNotFoundError:\n        vllm_spec = None\n    return vllm_spec is not None\n\n\n@cache\ndef is_sglang_available():\n    try:\n        sglang_spec = importlib.util.find_spec(\"sglang\")\n    except ModuleNotFoundError:\n        sglang_spec = None\n    return sglang_spec is not None\n\n\n@cache\ndef is_nvtx_available():\n    try:\n        nvtx_spec = importlib.util.find_spec(\"nvtx\")\n    except ModuleNotFoundError:\n        nvtx_spec = None\n    return nvtx_spec is not None\n\n\n@cache\ndef is_trl_available():\n    try:\n        trl_spec = importlib.util.find_spec(\"trl\")\n    except ModuleNotFoundError:\n        trl_spec = None\n    return trl_spec is not None\n\n\ndef import_external_libs(external_libs=None):\n    if external_libs is None:\n        return\n    if not isinstance(external_libs, list):\n        external_libs = [external_libs]\n    import importlib\n\n    for external_lib in external_libs:\n        importlib.import_module(external_lib)\n\n\nPKG_PATH_PREFIX = \"pkg://\"\nFILE_PATH_PREFIX = \"file://\"\n\n\ndef load_module(module_path: str, module_name: Optional[str] = None) -> object:\n    \"\"\"Load a module from a path.\n\n    Args:\n        module_path (str):\n            The path to the module. Either\n                - `pkg_path`, e.g.,\n                    - \"pkg://verl.utils.dataset.rl_dataset\"\n                    - \"pkg://verl/utils/dataset/rl_dataset\"\n                - or `file_path` (absolute or relative), e.g.,\n                    - \"file://verl/utils/dataset/rl_dataset.py\"\n                    - \"/path/to/verl/utils/dataset/rl_dataset.py\"\n        module_name (str, optional):\n            The name of the module to added to ``sys.modules``. If not provided, the module will not be added,\n                thus will not be cached and directly ``import``able.\n    \"\"\"\n    if not module_path:\n        return None\n\n    if module_path.startswith(PKG_PATH_PREFIX):\n        module_name = module_path[len(PKG_PATH_PREFIX) :].replace(\"/\", \".\")\n        module = importlib.import_module(module_name)\n\n    else:\n        if module_path.startswith(FILE_PATH_PREFIX):\n            module_path = module_path[len(FILE_PATH_PREFIX) :]\n\n        if not os.path.exists(module_path):\n            raise FileNotFoundError(f\"Custom module file not found: {module_path=}\")\n\n        # Use the provided module_name for the spec, or derive a unique name to avoid collisions.\n        spec_name = module_name or f\"custom_module_{hash(os.path.abspath(module_path))}\"\n        spec = importlib.util.spec_from_file_location(spec_name, module_path)\n        if spec is None or spec.loader is None:\n            raise ImportError(f\"Could not load module from {module_path=}\")\n\n        module = importlib.util.module_from_spec(spec)\n        try:\n            spec.loader.exec_module(module)\n        except Exception as e:\n            raise RuntimeError(f\"Error loading module from {module_path=}\") from e\n\n        if module_name is not None:\n            import sys\n\n            # Avoid overwriting an existing module with a different object.\n            if module_name in sys.modules and sys.modules[module_name] is not module:\n                raise RuntimeError(\n                    f\"Module name '{module_name}' already in `sys.modules` and points to a different module.\"\n                )\n            sys.modules[module_name] = module\n\n    return module\n\n\ndef _get_qualified_name(func):\n    \"\"\"Get full qualified name including module and class (if any).\"\"\"\n    module = func.__module__\n    qualname = func.__qualname__\n    return f\"{module}.{qualname}\"\n\n\ndef deprecated(replacement: str = \"\"):\n    \"\"\"Decorator to mark functions or classes as deprecated.\"\"\"\n\n    def decorator(obj):\n        qualified_name = _get_qualified_name(obj)\n\n        if isinstance(obj, type):\n            original_init = obj.__init__\n\n            @wraps(original_init)\n            def wrapped_init(self, *args, **kwargs):\n                msg = f\"Warning: Class '{qualified_name}' is deprecated.\"\n                if replacement:\n                    msg += f\" Please use '{replacement}' instead.\"\n                warnings.warn(msg, category=FutureWarning, stacklevel=2)\n                return original_init(self, *args, **kwargs)\n\n            obj.__init__ = wrapped_init\n            return obj\n\n        else:\n\n            @wraps(obj)\n            def wrapped(*args, **kwargs):\n                msg = f\"Warning: Function '{qualified_name}' is deprecated.\"\n                if replacement:\n                    msg += f\" Please use '{replacement}' instead.\"\n                warnings.warn(msg, category=FutureWarning, stacklevel=2)\n                return obj(*args, **kwargs)\n\n            return wrapped\n\n    return decorator\n\n\ndef load_extern_object(module_path: str, object_name: str) -> object:\n    \"\"\"Load an object from a module path.\n\n    Args:\n        module_path (str): See :func:`load_module`.\n        object_name (str):\n            The name of the object to load with ``getattr(module, object_name)``.\n    \"\"\"\n    module = load_module(module_path)\n\n    if not hasattr(module, object_name):\n        raise AttributeError(f\"Object not found in module: {object_name=}, {module_path=}.\")\n\n    return getattr(module, object_name)\n\n\ndef load_class_from_fqn(fqn: str, description: str = \"class\") -> type:\n    \"\"\"Load a class from its fully qualified name.\n\n    Args:\n        fqn: Fully qualified class name (e.g., 'mypackage.module.ClassName').\n        description: Description for error messages (e.g., 'AgentLoopManager').\n\n    Returns:\n        The loaded class.\n\n    Raises:\n        ValueError: If fqn format is invalid (missing dot separator).\n        ImportError: If the module cannot be imported.\n        AttributeError: If the class is not found in the module.\n\n    Example:\n        >>> cls = load_class_from_fqn(\"verl.experimental.agent_loop.AgentLoopManager\")\n        >>> instance = cls(config=config, ...)\n    \"\"\"\n    if \".\" not in fqn:\n        raise ValueError(\n            f\"Invalid {description} '{fqn}'. Expected fully qualified class name (e.g., 'mypackage.module.ClassName').\"\n        )\n    try:\n        module_path, class_name = fqn.rsplit(\".\", 1)\n        module = importlib.import_module(module_path)\n        return getattr(module, class_name)\n    except ImportError as e:\n        raise ImportError(f\"Failed to import module '{module_path}' for {description}: {e}\") from e\n    except AttributeError as e:\n        raise AttributeError(f\"Class '{class_name}' not found in module '{module_path}': {e}\") from e\n\n\n@deprecated(replacement=\"load_module(file_path); getattr(module, type_name)\")\ndef load_extern_type(file_path: str, type_name: str) -> type:\n    \"\"\"DEPRECATED. Directly use `load_extern_object` instead.\"\"\"\n    return load_extern_object(file_path, type_name)\n"
  },
  {
    "path": "verl/utils/kernel/__init__.py",
    "content": "#\n# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n"
  },
  {
    "path": "verl/utils/kernel/fp8_kernel.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\n\nimport torch\n\nlogger = logging.getLogger(__name__)\n\n# Check if Triton is available\n_TRITON_AVAILABLE = False\ntry:\n    import triton\n    import triton.language as tl\n\n    _TRITON_AVAILABLE = True\nexcept ImportError:\n    logger.debug(\"Triton not available, FP8 Triton kernels will not be used\")\n\n# Environment variable to control Triton FP8 usage (set to \"1\" to disable)\n_DISABLE_TRITON_FP8 = os.environ.get(\"VERL_DISABLE_TRITON_FP8\", \"0\").lower() in (\"1\", \"true\", \"yes\")\n\n# FP8 constants\nFP8_DTYPE = torch.float8_e4m3fn\nFP8_MAX = torch.finfo(FP8_DTYPE).max\nFP8_MIN = -FP8_MAX\n\n\ndef ceil_div(x: int, y: int) -> int:\n    \"\"\"Perform ceiling division of two integers.\"\"\"\n    return (x + y - 1) // y\n\n\ndef is_triton_available() -> bool:\n    \"\"\"Check if Triton is available for FP8 kernels.\"\"\"\n    return _TRITON_AVAILABLE\n\n\nif _TRITON_AVAILABLE:\n\n    @triton.jit\n    def _blockwise_cast_to_fp8_kernel(\n        X,\n        Y,\n        S,\n        stride_xm,\n        stride_xn,\n        stride_ym,\n        stride_yn,\n        stride_sm,\n        stride_sn,\n        M,\n        N,\n        eps,\n        fp8_min,\n        fp8_max,\n        BLOCK_M: tl.constexpr = 128,\n        BLOCK_N: tl.constexpr = 128,\n    ):\n        \"\"\"Triton kernel for blockwise FP8 quantization.\n\n        Each program instance handles one block of size (BLOCK_M, BLOCK_N).\n        Computes per-block scale and quantizes to FP8 in a single pass.\n\n        Refer to https://github.com/THUDM/slime/blob/main/slime/backends/megatron_utils/kernels/fp8_kernel.py\n        \"\"\"\n        pid_m = tl.cast(tl.program_id(axis=0), tl.int64)\n        pid_n = tl.cast(tl.program_id(axis=1), tl.int64)\n\n        # Compute block offsets\n        off_m = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)\n        off_n = pid_n * BLOCK_N + tl.arange(0, BLOCK_N)\n\n        # Create masks for boundary handling\n        mask_m = off_m < M\n        mask_n = off_n < N\n        mask = mask_m[:, None] & mask_n[None, :]\n\n        # Load input block and convert to float32 for precision\n        x = tl.load(X + off_m[:, None] * stride_xm + off_n[None, :] * stride_xn, mask=mask, other=0.0).to(tl.float32)\n\n        # Compute block-wise absolute maximum with epsilon for numerical stability\n        _absmax = tl.maximum(tl.max(tl.abs(x)), eps)\n\n        # Compute scale: scale = absmax / fp8_max\n        x_s = _absmax / fp8_max\n\n        # Compute inverse scale for quantization\n        s_inv = 1.0 / x_s\n\n        # Quantize: clamp(x * s_inv, fp8_min, fp8_max)\n        y_q = tl.clamp(x * s_inv, fp8_min, fp8_max).to(Y.dtype.element_ty)\n\n        # Store quantized values and scale\n        tl.store(Y + off_m[:, None] * stride_ym + off_n[None, :] * stride_yn, y_q, mask=mask)\n        tl.store(S + pid_m * stride_sm + pid_n * stride_sn, x_s)\n\n    def blockwise_cast_to_fp8_triton(\n        x: torch.Tensor,\n        weight_block_size: list[int] | tuple[int, int] | None = None,\n    ) -> tuple[torch.Tensor, torch.Tensor]:\n        \"\"\"Quantize a 2D tensor to FP8 using blockwise quantization with Triton.\n\n        This function provides high-performance FP8 quantization with minimal memory overhead.\n        All computations (abs, max, scale, clamp) are performed in a single Triton kernel,\n        eliminating intermediate tensor allocations.\n\n        Args:\n            x: Input tensor of shape (M, N), must be 2D.\n            weight_block_size: Block size for quantization as [BLOCK_M, BLOCK_N].\n                              Defaults to [128, 128] if None.\n\n        Returns:\n            Tuple of (quantized_tensor, scale_tensor):\n                - quantized_tensor: FP8 quantized tensor of shape (M, N)\n                - scale_tensor: Per-block scale factors of shape (ceil(M/BLOCK_M), ceil(N/BLOCK_N))\n                               This is the inverse scale (multiply to dequantize).\n        \"\"\"\n        assert x.dim() == 2, f\"Expected 2D tensor, got {x.dim()}D\"\n\n        # Default block size\n        BLOCK_M, BLOCK_N = 128, 128\n        if weight_block_size is not None:\n            BLOCK_M, BLOCK_N = weight_block_size[0], weight_block_size[1]\n\n        M, N = x.shape\n\n        # Pre-allocate output tensors (only memory allocation in this function)\n        y = torch.empty(M, N, device=x.device, dtype=FP8_DTYPE)\n        s = torch.empty(ceil_div(M, BLOCK_M), ceil_div(N, BLOCK_N), dtype=torch.float32, device=x.device)\n\n        # Grid: one program per block\n        def grid(meta):\n            return (triton.cdiv(M, meta[\"BLOCK_M\"]), triton.cdiv(N, meta[\"BLOCK_N\"]))\n\n        # Tune kernel parameters based on memory layout\n        if x.is_contiguous():\n            kwargs = {\"BLOCK_M\": BLOCK_M, \"BLOCK_N\": BLOCK_N, \"num_warps\": 8, \"num_stages\": 2}\n        else:\n            kwargs = {\"BLOCK_M\": BLOCK_M, \"BLOCK_N\": BLOCK_N, \"num_warps\": 1, \"num_stages\": 4}\n\n        # Launch kernel\n        _blockwise_cast_to_fp8_kernel[grid](\n            x,\n            y,\n            s,\n            *x.stride(),\n            *y.stride(),\n            *s.stride(),\n            M,\n            N,\n            1e-10,  # eps for numerical stability\n            FP8_MIN,\n            FP8_MAX,\n            **kwargs,\n        )\n\n        return y, s\n\n\ndef scaled_fp8_blockwise_triton(\n    data_hp: torch.Tensor,\n    weight_block_size: list[int] | tuple[int, int],\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"High-performance FP8 blockwise quantization using Triton kernel.\n\n    This is the recommended function to use for FP8 quantization when Triton is available.\n    It handles padding automatically and returns results in the expected format.\n\n    Args:\n        data_hp: Input high-precision tensor of shape (M, N).\n        weight_block_size: Block size for quantization as [BLOCK_M, BLOCK_N].\n\n    Returns:\n        Tuple of (fp8_data, descale):\n            - fp8_data: FP8 quantized tensor of original shape\n            - descale: Per-block descale factors (inverse of scale, for dequantization)\n\n    Raises:\n        RuntimeError: If Triton is not available.\n    \"\"\"\n    if not _TRITON_AVAILABLE:\n        raise RuntimeError(\"Triton is required for scaled_fp8_blockwise_triton but is not available\")\n\n    block_size0 = weight_block_size[0]\n    block_size1 = weight_block_size[1]\n\n    # Save original shape for potential cropping\n    original_shape = data_hp.shape\n\n    # Pad dimensions to be multiples of block size if needed\n    pad_dim0 = (block_size0 - data_hp.shape[0] % block_size0) % block_size0\n    pad_dim1 = (block_size1 - data_hp.shape[1] % block_size1) % block_size1\n\n    if pad_dim0 > 0 or pad_dim1 > 0:\n        logger.debug(\n            f\"Padding weight from {data_hp.shape} to \"\n            f\"({data_hp.shape[0] + pad_dim0}, {data_hp.shape[1] + pad_dim1}) \"\n            f\"for blockwise FP8 quantization\"\n        )\n        data_hp = torch.nn.functional.pad(data_hp, (0, pad_dim1, 0, pad_dim0), mode=\"constant\", value=0)\n\n    # Call Triton kernel\n    fp_data, scale = blockwise_cast_to_fp8_triton(data_hp, weight_block_size)\n\n    # Remove padding to restore original shape\n    if pad_dim0 > 0 or pad_dim1 > 0:\n        fp_data = fp_data[: original_shape[0], : original_shape[1]].contiguous()\n\n    # Return scale as descale (the Triton kernel returns scale, we need to return it as-is\n    # since it's already the inverse scale format expected by vLLM/SGLang)\n    return fp_data, scale\n\n\ndef _scaled_fp8_blockwise_pytorch(\n    data_hp: torch.Tensor,\n    weight_block_size: list[int] | tuple[int, int],\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"PyTorch implementation of blockwise FP8 quantization.\n\n    Memory-optimized implementation that:\n    - Uses in-place operations where possible\n    - Explicitly deletes intermediate tensors\n    - Minimizes peak memory usage during quantization\n\n    Args:\n        data_hp: Input high-precision tensor of shape (M, N).\n        weight_block_size: Block size for quantization as [BLOCK_M, BLOCK_N].\n\n    Returns:\n        Tuple of (fp8_data, descale):\n            - fp8_data: FP8 quantized tensor\n            - descale: Per-block descale factors for dequantization\n    \"\"\"\n    block_size0 = weight_block_size[0]\n    block_size1 = weight_block_size[1]\n    assert block_size0 == block_size1, \"Block sizes must be equal\"\n\n    # Save unpadded shape for later cropping\n    original_shape = data_hp.shape\n\n    # Pad dimensions to be multiples of block size if needed\n    pad_dim0 = (block_size0 - data_hp.shape[0] % block_size0) % block_size0\n    pad_dim1 = (block_size1 - data_hp.shape[1] % block_size1) % block_size1\n\n    if pad_dim0 > 0 or pad_dim1 > 0:\n        logger.debug(\n            f\"Padding weight from {data_hp.shape} to \"\n            f\"({data_hp.shape[0] + pad_dim0}, {data_hp.shape[1] + pad_dim1}) \"\n            f\"for blockwise FP8 quantization\"\n        )\n        data_hp = torch.nn.functional.pad(data_hp, (0, pad_dim1, 0, pad_dim0), mode=\"constant\", value=0)\n\n    # FP8\n    max_dtype = FP8_MAX\n\n    padded_shape = data_hp.shape\n    blk_m, blk_n = data_hp.shape[0] // block_size0, data_hp.shape[1] // block_size1\n\n    # Reshape and permute - these are views, no memory allocation\n    data_hp = data_hp.reshape(blk_m, block_size0, blk_n, block_size1)\n    data_hp = data_hp.permute(0, 2, 1, 3).contiguous()\n\n    # Flatten to (BLK_M, BLK_N, BLOCK_SIZE_M * BLOCK_SIZE_N) in float32 for precision\n    data_hp = data_hp.to(torch.float32).flatten(start_dim=2)\n\n    # Calculate max absolute value per block - use fused abs+amax\n    max_abs = data_hp.abs().amax(dim=-1, keepdim=True)\n\n    # Compute scale in-place where possible\n    scale_fp = torch.empty_like(max_abs)\n    torch.div(max_dtype, max_abs, out=scale_fp)\n    # Handle edge cases: zero and inf\n    scale_fp = torch.where(max_abs == 0, torch.ones_like(scale_fp), scale_fp)\n    scale_fp = torch.where(max_abs == torch.inf, torch.ones_like(scale_fp), scale_fp)\n    del max_abs  # Free max_abs memory\n\n    # Compute descale before modifying data\n    descale_fp = torch.reciprocal(scale_fp)\n\n    # Scale and clamp in a memory-efficient way\n    data_hp.mul_(scale_fp)\n    del scale_fp  # Free scale memory\n    data_hp.clamp_(min=-max_dtype, max=max_dtype)\n\n    # Convert to FP8\n    fp_data = data_hp.to(FP8_DTYPE)\n    del data_hp  # Free float32 data\n\n    # Reshape back to original layout\n    fp_data = fp_data.reshape(blk_m, blk_n, block_size0, block_size1).permute(0, 2, 1, 3).reshape(padded_shape)\n\n    # Remove padding to restore original shape\n    if original_shape[0] != padded_shape[0] or original_shape[1] != padded_shape[1]:\n        fp_data = fp_data[: original_shape[0], : original_shape[1]].contiguous()\n\n    return fp_data, descale_fp\n\n\ndef scaled_fp8_blockwise(\n    data_hp: torch.Tensor,\n    weight_block_size: list[int] | tuple[int, int],\n) -> tuple[torch.Tensor, torch.Tensor]:\n    \"\"\"Cast tensor from high precision to FP8 with blockwise quantization.\n\n    This function automatically selects the best available implementation:\n    1. Triton kernel (if available): Highest performance, minimal memory overhead\n    2. PyTorch fallback: Memory-optimized implementation using in-place operations\n\n    To disable Triton and force PyTorch fallback, set environment variable:\n        VERL_DISABLE_TRITON_FP8=1\n\n    Args:\n        data_hp: Input tensor of shape (M, N) in high precision (bf16/fp16/fp32).\n        weight_block_size: Block size for quantization as [BLOCK_M, BLOCK_N].\n\n    Returns:\n        Tuple of (fp8_data, descale):\n            - fp8_data: FP8 quantized tensor\n            - descale: Per-block descale factors for dequantization\n    \"\"\"\n    assert len(data_hp.shape) == 2, \"Only 2d input tensor is supported\"\n\n    # Use Triton kernel if available and not disabled\n    if _TRITON_AVAILABLE and not _DISABLE_TRITON_FP8:\n        return scaled_fp8_blockwise_triton(data_hp, weight_block_size)\n\n    # PyTorch fallback implementation (memory-optimized)\n    return _scaled_fp8_blockwise_pytorch(data_hp, weight_block_size)\n"
  },
  {
    "path": "verl/utils/kernel/kernels.py",
    "content": "#\n# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nImplementations of the linear cross entropy with token entropy kernel.\n\"\"\"\n\nimport typing\nfrom dataclasses import dataclass\n\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_capability, get_device_name, is_cuda_available\n\ntry:\n    import triton\n    import triton.language as tl\n\n    HAVE_TRITON = True\n    SUPPORT_CUDA_TMA = is_cuda_available and get_device_capability()[0] >= 9 and hasattr(tl, \"make_tensor_descriptor\")\n\nexcept ImportError:\n    HAVE_TRITON = False\n    SUPPORT_CUDA_TMA = False\n\nfrom verl.utils.device import get_torch_device\n\nif not HAVE_TRITON:\n    from contextlib import contextmanager\n    from unittest.mock import MagicMock\n\n    @contextmanager\n    def null_decorator(*args, **kwargs):\n        if len(kwargs) == 0 and len(args) == 1 and callable(args[0]):\n            return args[0]\n        else:\n\n            def inner(func):\n                return func\n\n            return inner\n\n    triton = MagicMock()\n    triton.jit = null_decorator\n    triton.autotune = null_decorator\n    tl = MagicMock()\n\nelif SUPPORT_CUDA_TMA:\n    # TMA descriptors require a global memory allocation\n    def alloc_fn(size: int, alignment: int, stream: typing.Optional[int]):\n        return torch.empty(size, device=get_device_name(), dtype=torch.int8)\n\n    # https://github.com/triton-lang/triton/commit/43625fc968b693ab51884ca95adbcf3e43483fd0\n    # Triton 3.5.0 stores allocators in ContextVar; values do not propagate to new\n    # threads by default. Some execution paths in verl use thread pools (e.g.,\n    # concurrent.futures), so we set a ContextVar *default* to avoid falling\n    # back to NullAllocator in worker threads.\n    try:\n        import contextvars\n\n        import triton.runtime._allocation as _triton_allocation\n\n        if isinstance(getattr(_triton_allocation, \"_allocator\", None), contextvars.ContextVar):\n            _triton_allocation._allocator = contextvars.ContextVar(\n                _triton_allocation._allocator.name,\n                default=alloc_fn,\n            )\n    except (ImportError, AttributeError):\n        pass\n\n    triton.set_allocator(alloc_fn)\n\n\n@dataclass\nclass EntropyReductionEnum:\n    \"\"\"\n    Enum for the reduction method of cross entropy.\n    \"\"\"\n\n    _None = 0\n    _Sum = 1\n    _Mean = 2\n\n\ndef get_entropy_reduction_enum_number(reduction: str) -> int:\n    \"\"\"\n    Get the enum number for the reduction method of cross entropy.\n    \"\"\"\n    _enum = EntropyReductionEnum._None\n    if reduction == \"none\":\n        _enum = EntropyReductionEnum._None\n    elif reduction == \"sum\":\n        _enum = EntropyReductionEnum._Sum\n    elif reduction == \"mean\":\n        _enum = EntropyReductionEnum._Mean\n    else:\n        raise ValueError(f\"Invalid reduction: {reduction}\")\n    return _enum\n\n\ndef get_entropy_reduction_enum(ce_reduction: int) -> EntropyReductionEnum:\n    \"\"\"\n    Get the enum for the reduction method of cross entropy.\n    \"\"\"\n    _enum = EntropyReductionEnum._None\n    if ce_reduction == 0:\n        _enum = EntropyReductionEnum._None\n    elif ce_reduction == 1:\n        _enum = EntropyReductionEnum._Sum\n    elif ce_reduction == 2:\n        _enum = EntropyReductionEnum._Mean\n    else:\n        raise ValueError(f\"Invalid ce_reduction: {ce_reduction}\")\n    return _enum\n\n\n@dataclass\nclass BackwardEnum:\n    \"\"\"\n    Enum for the backward method.\n    \"\"\"\n\n    _Total_Fuse_MN = (\n        0  # Fuse d_logits & d_hidden & d_weight, no intermediate storage, requires fp32 for d_hidden & d_weight\n    )\n    _Total_Separate = 1  # Store d_logits, no special requirements for d_hidden & d_weight\n    _Split_Dlogits_N = 2  # split d_logits along its N dimension, aka. vocab_size\n    _Split_Dlogits_M = 3  # split d_logits along its M dimension, aka. num_tokens\n\n\n@dataclass\nclass Config:\n    \"\"\"Configuration for efficient entropy kernel operations.\n\n    Args:\n        _backward (BackwardEnum): Backward computation method. Defaults to BackwardEnum._Split_Dlogits_N.\n        _use_triton (bool): Whether to use Triton kernels for computation. Defaults to True.\n    \"\"\"\n\n    _backward: BackwardEnum = BackwardEnum._Split_Dlogits_N\n    _use_triton: bool = True\n\n\n_config = Config()\n\n\ndef set_backward_method(backward_method: BackwardEnum):\n    \"\"\"\n    Set the backward method.\n    \"\"\"\n    global _config\n    _config._backward = backward_method\n\n\n@triton.autotune(\n    configs=[triton.Config({\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 256, \"BLOCK_SIZE_K\": 32}, num_stages=3, num_warps=8)],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_kernel_general_mainloop(\n    rank,\n    hidden_ptr,\n    weight_ptr,\n    labels_ptr,\n    num_tokens,\n    hidden_size,\n    vocab_size,\n    vocab_per_split,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    max_ptr,\n    stride_max_m: tl.int64,\n    stride_max_n: tl.int64,\n    accu_ptr,\n    stride_accu_m: tl.int64,\n    stride_accu_n: tl.int64,\n    entropy_b_ptr,\n    stride_entropy_b_m: tl.int64,\n    stride_entropy_b_n: tl.int64,\n    global_logprobs_ptr,\n    stride_global_logprobs: tl.int64,\n    global_logprobs_scalar_ptr,\n    rcp_temperature: tl.float32,\n    # Meta-parameters\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    USE_TMA: tl.constexpr,\n):\n    \"\"\"\n    forward mainloop\n    \"\"\"\n    pid = tl.program_id(axis=0)\n    num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split\n    num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    num_pid_n = tl.cdiv(vocab_per_split, BLOCK_SIZE_N)\n    pid_m = pid % num_pid_m\n    pid_n = pid // num_pid_m\n\n    if pid_m == 0 and pid_n == 0:\n        tl.store(global_logprobs_scalar_ptr, 0.0)\n\n    # create pointers for the first blocks of hidden\n    start_offs_am = pid_m * BLOCK_SIZE_M\n    offs_am = start_offs_am + tl.arange(0, BLOCK_SIZE_M)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n\n    if USE_TMA:\n        # using TMA and device-side descriptor creation\n        hidden_desc = tl.make_tensor_descriptor(\n            hidden_ptr,\n            shape=[num_tokens, hidden_size],\n            strides=[stride_hidden_m, 1],\n            block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],\n        )\n\n        weight_desc = tl.make_tensor_descriptor(\n            weight_ptr,\n            shape=[vocab_size, hidden_size],\n            strides=[stride_weight_n, 1],\n            block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],\n        )\n\n    else:\n        hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n\n    # load labels for this block\n    labels = tl.load(labels_ptr + offs_am, mask=offs_am < num_tokens)\n\n    # traverse over N dimension\n    # _max = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    _max = tl.full((BLOCK_SIZE_M,), -float(\"inf\"), dtype=tl.float32)\n    _accu = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    _entropy_b = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    _logprobs = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    vocab_bound = min((pid_n + 1) * vocab_per_split, vocab_size)\n    for n in range(0, num_pid_n):\n        start_offs_bn = pid_n * vocab_per_split + n * BLOCK_SIZE_N\n        offs_bn = start_offs_bn + tl.arange(0, BLOCK_SIZE_N)\n\n        logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n        if not USE_TMA:\n            # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n)\n            weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n\n        # iterate over K dimension\n        for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n            if USE_TMA:\n                # load the next block of hidden and weight\n                start_offs_k = k * BLOCK_SIZE_K\n                _hidden = hidden_desc.load([start_offs_am, start_offs_k])\n                _weight = weight_desc.load([start_offs_bn, start_offs_k])\n            else:\n                # load the next block of hidden and weight\n                _hidden = tl.load(\n                    hidden_ptrs,\n                    mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n                    other=0.0,\n                )\n\n                _weight = tl.load(\n                    weight_ptrs,\n                    mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K)\n                    & (offs_bn[:, None] < (min((pid_n + 1) * vocab_per_split, vocab_size))),\n                    other=0.0,\n                )\n\n                # advance the ptrs to the next K block\n                hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n                weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n\n            # GEMM\n            logits = tl.dot(_hidden, _weight.trans(), logits)\n\n        if not USE_TMA:\n            # reset hidden_ptrs for next iteration\n            hidden_ptrs -= hidden_size * stride_hidden_k\n\n        # scale logits by temperature\n        logits *= rcp_temperature\n\n        logits_for_lse = tl.where(offs_bn[None, :] < vocab_bound, logits, float(\"-inf\"))\n\n        # update global maximum\n        _max_old = _max\n        m_pid_n = tl.max(logits_for_lse, axis=1)\n        _max = tl.maximum(_max_old, m_pid_n)\n\n        exp_logits = tl.exp(logits_for_lse - _max[:, None])\n        coeff = tl.exp(_max_old - _max)\n        _accu = coeff * _accu + tl.sum(exp_logits, axis=1)\n\n        _entropy_b = _entropy_b * coeff + tl.sum(logits * exp_logits, axis=1)\n\n        label_mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None]\n        _logprobs += tl.sum(logits * label_mask, axis=1)\n\n    # store maximum\n    offs_max_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n    offs_max_n = pid_n\n    maximum_ptrs = max_ptr + offs_max_n * stride_max_n + offs_max_m * stride_max_m\n    tl.store(maximum_ptrs, _max, mask=(offs_max_m < num_tokens) & (offs_max_n < num_splits))\n\n    # store entropy\n    accu_ptrs = accu_ptr + offs_max_n * stride_accu_n + offs_max_m * stride_accu_m\n    tl.store(accu_ptrs, _accu, mask=(offs_max_m < num_tokens) & (offs_max_n[None] < num_splits))\n    entropy_b_ptrs = entropy_b_ptr + offs_max_n * stride_entropy_b_n + offs_max_m * stride_entropy_b_m\n    tl.store(entropy_b_ptrs, _entropy_b, mask=(offs_max_m < num_tokens) & (offs_max_n < num_splits))\n    # store logprobs\n    vocab_left_idx = pid_n * vocab_per_split + rank * vocab_size\n    vocab_right_idx = min((pid_n + 1) * vocab_per_split, vocab_size) + rank * vocab_size\n    mask = (labels >= vocab_left_idx) & (labels < vocab_right_idx)\n    mask &= offs_am < num_tokens\n    global_logprobs_ptrs = global_logprobs_ptr + offs_am * stride_global_logprobs\n    # tl.atomic_add(global_logprobs_ptrs, _logprobs, mask=mask)\n    tl.store(global_logprobs_ptrs, _logprobs, mask=mask)\n\n\n@triton.autotune(configs=[triton.Config({\"BLOCK_SIZE_M\": 16, \"BLOCK_SIZE_N\": 64})], key=[\"num_tokens\", \"num_splits\"])\n@triton.jit\ndef efficient_entropy_triton_kernel_epilogue(\n    max_ptr,\n    stride_max_m: tl.int64,\n    stride_max_n: tl.int64,\n    num_tokens,\n    num_splits,\n    global_max_ptr,\n    stride_global_max: tl.int64,\n    accu_ptr,\n    stride_accu_m: tl.int64,\n    stride_accu_n: tl.int64,\n    global_accu_ptr,\n    stride_global_accu: tl.int64,\n    entropy_b_ptr,\n    stride_entropy_b_m: tl.int64,\n    stride_entropy_b_n: tl.int64,\n    global_entropy_b_ptr,\n    stride_global_entropy_b: tl.int64,\n    global_entropy_ptr,\n    stride_global_entropy: tl.int64,\n    global_logprobs_ptr,\n    stride_global_logprobs: tl.int64,\n    global_logprobs_scalar_ptr,\n    reduction: int,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n):\n    \"\"\"\n    foward epilogue\n    \"\"\"\n    pid_m = tl.program_id(axis=0)\n\n    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n    global_max = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    global_accu = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    global_entropy_b = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    for pid_n in range(0, tl.cdiv(num_splits, BLOCK_SIZE_N)):\n        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n        max_ptrs = max_ptr + offs_m[:, None] * stride_max_m + offs_n[None, :] * stride_max_n\n\n        _max = tl.load(max_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0)\n\n        accu_ptrs = accu_ptr + offs_m[:, None] * stride_accu_m + offs_n[None, :] * stride_accu_n\n        _accu = tl.load(accu_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0)\n\n        entropy_b_ptrs = entropy_b_ptr + offs_m[:, None] * stride_entropy_b_m + offs_n[None, :] * stride_entropy_b_n\n        _entropy_b = tl.load(\n            entropy_b_ptrs, mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits), other=0.0\n        )\n\n        # local reduction\n        _max_old = global_max\n        _local_max = tl.max(_max, axis=1)\n        global_max = tl.maximum(global_max, _local_max)\n\n        _scale = tl.exp(_max - global_max[:, None])\n        _coeff = tl.exp(_max_old - global_max)\n        global_accu = _coeff * global_accu + tl.sum(_scale * _accu, axis=1)\n        global_entropy_b = _coeff * global_entropy_b + tl.sum(_scale * _entropy_b, axis=1)\n\n    # store\n    maximum_ptrs = global_max_ptr + offs_m * stride_global_max\n    tl.store(maximum_ptrs, global_max, mask=offs_m < num_tokens)\n\n    # store entropy_b\n    global_entropy_b = tl.fdiv(global_entropy_b, global_accu)  # entropy_b\n    tl.store(global_entropy_b_ptr + offs_m * stride_global_entropy_b, global_entropy_b, mask=offs_m < num_tokens)\n\n    # store entropy\n    global_accu_ptrs = global_accu_ptr + offs_m * stride_global_accu\n    tl.store(global_accu_ptrs, global_accu, mask=offs_m < num_tokens)\n    global_entropy = tl.log(global_accu) + global_max - global_entropy_b  # entropy_a\n    global_entropy_ptrs = global_entropy_ptr + offs_m * stride_global_entropy\n    tl.store(global_entropy_ptrs, global_entropy, mask=offs_m < num_tokens)\n    # update logprobs\n    global_logprobs_ptrs = global_logprobs_ptr + offs_m * stride_global_logprobs\n    global_logprobs = tl.load(global_logprobs_ptrs, mask=offs_m < num_tokens)\n    global_logprobs = global_max + tl.log(global_accu) - global_logprobs\n\n    global_logprobs = -1 * global_logprobs\n    if reduction == 0:\n        tl.store(global_logprobs_ptrs, global_logprobs, mask=offs_m < num_tokens)\n    elif reduction == 1:\n        global_logprobs_scalar = tl.sum(global_logprobs, axis=0)\n        tl.atomic_add(global_logprobs_scalar_ptr, global_logprobs_scalar)\n    elif reduction == 2:\n        global_logprobs_scalar = tl.sum(global_logprobs, axis=0) / num_tokens.to(tl.float32)\n        tl.atomic_add(global_logprobs_scalar_ptr, global_logprobs_scalar)\n\n\n@triton.autotune(configs=[triton.Config({\"BLOCK_SIZE_M\": 16, \"BLOCK_SIZE_N\": 64})], key=[\"num_tokens\", \"num_splits\"])\n@triton.jit\ndef efficient_entropy_triton_kernel_epilogue_tp(\n    num_tokens,\n    num_splits,\n    reduced_max_ptr,\n    stride_reduced_max_m: tl.int64,\n    stride_reduced_max_n: tl.int64,\n    original_max_ptr,\n    stride_original_max_m: tl.int64,\n    stride_original_max_n: tl.int64,\n    accu_ptr,\n    stride_accu_m: tl.int64,\n    stride_accu_n: tl.int64,\n    entropy_b_ptr,\n    stride_entropy_b_m: tl.int64,\n    stride_entropy_b_n: tl.int64,\n    global_max_ptr,\n    stride_global_max: tl.int64,\n    global_accu_ptr,\n    stride_global_accu: tl.int64,\n    global_entropy_b_ptr,\n    stride_global_entropy_b: tl.int64,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n):\n    pid_m = tl.program_id(axis=0)\n\n    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n\n    global_max = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    global_accu = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    global_entropy_b = tl.zeros((BLOCK_SIZE_M,), dtype=tl.float32)\n    for pid_n in range(0, tl.cdiv(num_splits, BLOCK_SIZE_N)):\n        offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n\n        _reduced_max = tl.load(\n            reduced_max_ptr + offs_m[:, None] * stride_reduced_max_m + offs_n[None, :] * stride_reduced_max_n,\n            mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits),\n            other=0.0,\n        )\n        _original_max = tl.load(\n            original_max_ptr + offs_m[:, None] * stride_original_max_m + offs_n[None, :] * stride_original_max_n,\n            mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits),\n            other=0.0,\n        )\n        _accu = tl.load(\n            accu_ptr + offs_m[:, None] * stride_accu_m + offs_n[None, :] * stride_accu_n,\n            mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits),\n            other=0.0,\n        )\n\n        # local reduce-max\n        _max_old = global_max\n        _local_max = tl.max(_reduced_max, axis=1)\n        global_max = tl.maximum(global_max, _local_max)\n\n        # update accumulate\n        _coeff = tl.exp(_max_old - global_max)\n        _scale = tl.exp(_original_max - global_max[:, None])\n        global_accu = _coeff * global_accu + tl.sum(_scale * _accu, axis=1)\n\n        # update entropy_b\n        _entropy_b = tl.load(\n            entropy_b_ptr + offs_m[:, None] * stride_entropy_b_m + offs_n[None, :] * stride_entropy_b_n,\n            mask=(offs_m[:, None] < num_tokens) & (offs_n[None, :] < num_splits),\n            other=0.0,\n        )\n        global_entropy_b = _coeff * global_entropy_b + tl.sum(_scale * _entropy_b, axis=1)\n\n    # store\n    tl.store(global_max_ptr + offs_m * stride_global_max, global_max, mask=offs_m < num_tokens)\n    tl.store(global_accu_ptr + offs_m * stride_global_accu, global_accu, mask=offs_m < num_tokens)\n    tl.store(global_entropy_b_ptr + offs_m * stride_global_entropy_b, global_entropy_b, mask=offs_m < num_tokens)\n\n\n@triton.autotune(configs=[triton.Config({\"BLOCK_SIZE_M\": 16})], key=[\"num_tokens\"])\n@triton.jit\ndef efficient_entropy_triton_epilogue_tp_update(\n    num_tokens,\n    logprobs_ptr,\n    stride_logprobs: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accumulate_ptr,\n    stride_accumulate: tl.int64,\n    entropy_b_ptr,\n    stride_entropy_b: tl.int64,\n    entropy_ptr,\n    stride_entropy: tl.int64,\n    logprobs_scalar_ptr,\n    reduction: int,\n    BLOCK_SIZE_M: tl.constexpr,\n):\n    pid_m = tl.program_id(axis=0)\n\n    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n\n    maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens)\n    accumulate = tl.load(accumulate_ptr + offs_m * stride_accumulate, mask=offs_m < num_tokens)\n\n    entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens)\n    entropy_b = tl.fdiv(entropy_b, accumulate)\n    tl.store(entropy_b_ptr + offs_m * stride_entropy_b, entropy_b, mask=offs_m < num_tokens)\n\n    entropy = tl.log(accumulate) + maximum - entropy_b\n    tl.store(entropy_ptr + offs_m * stride_entropy, entropy, mask=offs_m < num_tokens)\n\n    logprobs = tl.load(logprobs_ptr + offs_m * stride_logprobs, mask=offs_m < num_tokens)\n    logprobs = maximum + tl.log(accumulate) - logprobs\n\n    logprobs = -1 * logprobs\n    if reduction == 0:\n        tl.store(logprobs_ptr + offs_m * stride_logprobs, logprobs, mask=offs_m < num_tokens)\n    elif reduction == 1:\n        logprobs_scalar = tl.sum(logprobs, axis=0)\n        tl.atomic_add(logprobs_scalar_ptr, logprobs_scalar)\n    elif reduction == 2:\n        logprobs_scalar = tl.sum(logprobs, axis=0) / num_tokens.to(tl.float32)\n        tl.atomic_add(logprobs_scalar_ptr, logprobs_scalar)\n\n\n_dedicated_stream, _dedicated_events = None, None\n\n\ndef efficient_entropy_forward(\n    hidden: torch.Tensor,\n    weight: torch.Tensor,\n    labels: torch.Tensor,\n    reduction: typing.Optional[int] = 2,\n    temperature: typing.Optional[float] = 1.0,\n    dist_process_group: typing.Optional[dist.ProcessGroup] = None,\n) -> list[torch.Tensor]:\n    \"\"\"\n    forward host function\n    \"\"\"\n    assert hidden.is_cuda and weight.is_cuda and labels.is_cuda\n    assert weight.device == hidden.device and labels.device == hidden.device\n    assert hidden.dim() == 2 and weight.dim() == 2 and labels.dim() == 1\n    assert hidden.is_contiguous() and weight.is_contiguous() and labels.is_contiguous()\n\n    assert hidden.shape[0] == labels.shape[0] and hidden.shape[1] == weight.shape[1]\n\n    _rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group)\n    _world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)\n\n    if dist_process_group is not None and not hasattr(efficient_entropy_forward, \"_initialized\"):\n        global _dedicated_stream, _dedicated_events\n        _dedicated_stream = get_torch_device().Stream(hidden.device)\n        _dedicated_events = [get_torch_device().Event() for _ in range(2)]\n        efficient_entropy_forward._initialized = True\n\n    num_tokens, hidden_size = hidden.shape\n    num_tokens = labels.shape[0]\n    vocab_size, hidden_size = weight.shape\n    assert hidden_size % 128 == 0\n\n    REDUCTION = get_entropy_reduction_enum(reduction)\n\n    if REDUCTION == EntropyReductionEnum._None:\n        if dist_process_group is None:\n            logprobs = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32)\n        else:\n            logprobs = torch.zeros((num_tokens,), device=hidden.device, dtype=torch.float32)\n    elif REDUCTION in (EntropyReductionEnum._Sum, EntropyReductionEnum._Mean):\n        logprobs = torch.empty((), device=hidden.device, dtype=torch.float32)\n    else:\n        raise ValueError(f\"Invalid reduction: {reduction}\")\n\n    entropy = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32)\n    assert logprobs.is_contiguous() and entropy.is_contiguous()\n\n    maximum = torch.empty_like(entropy)\n    accumulate_and_entropy_b = torch.empty((num_tokens * 2,), device=hidden.device, dtype=torch.float32)\n    accumulate_and_entropy_b_view = accumulate_and_entropy_b.view(2, num_tokens)\n    accumulate = accumulate_and_entropy_b_view[0, :]\n    entropy_b = accumulate_and_entropy_b_view[1, :]\n    assert maximum.is_contiguous() and accumulate.is_contiguous() and entropy_b.is_contiguous()\n\n    vocab_per_split = 1024\n    assert vocab_per_split % 128 == 0\n    num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split\n\n    _max = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32)\n    _accu = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32)\n    _entropy_b = torch.empty((num_tokens, num_splits), device=hidden.device, dtype=torch.float32)\n\n    if REDUCTION == EntropyReductionEnum._None:\n        _logprobs = logprobs\n    else:\n        _logprobs = torch.empty((num_tokens,), device=hidden.device, dtype=torch.float32)\n\n    assert _accu.is_contiguous() and _entropy_b.is_contiguous() and _max.is_contiguous()\n    assert _accu.is_cuda and _entropy_b.is_cuda and _max.is_cuda\n\n    if _config._use_triton:\n        # 1D kernel launch, then split the tile\n        def mainloop_grid(meta):\n            return (triton.cdiv(num_tokens, meta[\"BLOCK_SIZE_M\"]) * num_splits,)\n\n        efficient_entropy_kernel_general_mainloop[mainloop_grid](\n            _rank,\n            hidden,\n            weight,\n            labels,\n            num_tokens,\n            hidden_size,\n            vocab_size,\n            vocab_per_split,\n            hidden.stride(0),\n            hidden.stride(1),\n            weight.stride(0),\n            weight.stride(1),\n            _max,\n            _max.stride(0),\n            _max.stride(1),\n            _accu,\n            _accu.stride(0),\n            _accu.stride(1),\n            _entropy_b,\n            _entropy_b.stride(0),\n            _entropy_b.stride(1),\n            _logprobs,\n            _logprobs.stride(0),\n            logprobs,\n            1.0 / temperature,\n            USE_TMA=SUPPORT_CUDA_TMA and hidden.stride(1) == 1 and weight.stride(1) == 1,\n        )\n    else:\n        raise AssertionError(\"Triton is required for efficient entropy kernel\")\n\n    # reduction on maximum and maximum_indices\n    def epilogue_grid(meta):\n        return (triton.cdiv(num_tokens, meta[\"BLOCK_SIZE_M\"]),)\n\n    if dist_process_group is None:\n        efficient_entropy_triton_kernel_epilogue[epilogue_grid](\n            _max,\n            _max.stride(0),\n            _max.stride(1),\n            num_tokens,\n            num_splits,\n            maximum,\n            maximum.stride(0),\n            _accu,\n            _accu.stride(0),\n            _accu.stride(1),\n            accumulate,\n            accumulate.stride(0),\n            _entropy_b,\n            _entropy_b.stride(0),\n            _entropy_b.stride(1),\n            entropy_b,\n            entropy_b.stride(0),\n            entropy,\n            entropy.stride(0),\n            _logprobs,\n            _logprobs.stride(0),\n            logprobs,\n            REDUCTION,\n        )\n    else:\n        # tensor-parallel\n        _max_backup = _max.clone()\n        dist.all_reduce(_max, op=dist.ReduceOp.MAX, group=dist_process_group)\n\n        get_torch_device().current_stream().record_event(_dedicated_events[0])\n        with get_torch_device().stream(_dedicated_stream):\n            _dedicated_stream.wait_event(_dedicated_events[0])\n            dist.all_reduce(_logprobs, op=dist.ReduceOp.SUM, group=dist_process_group)\n            _dedicated_stream.record_event(_dedicated_events[1])\n\n        efficient_entropy_triton_kernel_epilogue_tp[epilogue_grid](\n            num_tokens,\n            num_splits,\n            _max,\n            _max.stride(0),\n            _max.stride(1),\n            _max_backup,\n            _max_backup.stride(0),\n            _max_backup.stride(1),\n            _accu,\n            _accu.stride(0),\n            _accu.stride(1),\n            _entropy_b,\n            _entropy_b.stride(0),\n            _entropy_b.stride(1),\n            maximum,\n            maximum.stride(0),\n            accumulate,\n            accumulate.stride(0),\n            entropy_b,\n            entropy_b.stride(0),\n        )\n        get_torch_device().current_stream().wait_event(_dedicated_events[1])\n\n        dist.all_reduce(accumulate_and_entropy_b, op=dist.ReduceOp.SUM, group=dist_process_group)\n\n        # update logprobs & entropy\n        efficient_entropy_triton_epilogue_tp_update[epilogue_grid](\n            num_tokens,\n            _logprobs,\n            _logprobs.stride(0),\n            maximum,\n            maximum.stride(0),\n            accumulate,\n            accumulate.stride(0),\n            entropy_b,\n            entropy_b.stride(0),\n            entropy,\n            entropy.stride(0),\n            logprobs,\n            REDUCTION,\n        )\n\n    return (logprobs, entropy, maximum, accumulate, entropy_b)\n\n\n# NOTE: merge d_weight & d_hidden here, split along M & N\n@triton.autotune(\n    configs=[\n        triton.Config(\n            {\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 128, \"BLOCK_SIZE_K\": 32, \"GROUP_SIZE_M\": 16},\n            num_stages=3,\n            num_warps=8,\n        )\n    ],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_backward_kernel_general_mainloop_MN(\n    num_tokens: int,\n    hidden_size: int,\n    vocab_size: int,\n    rank: int,\n    hidden_ptr,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    weight_ptr,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    labels_ptr,\n    stride_labels: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accu_ptr,\n    stride_accu: tl.int64,\n    d_entropy_ptr,\n    stride_d_entropy: tl.int64,\n    d_logprobs_ptr,\n    stride_d_logprobs: tl.int64,\n    reduction: int,\n    entropy_b_ptr,\n    stride_entropy_b: tl.int64,\n    d_hidden_ptr,\n    stride_d_hidden_m: tl.int64,\n    stride_d_hidden_k: tl.int64,\n    d_weight_ptr,\n    stride_d_weight_n: tl.int64,\n    stride_d_weight_k: tl.int64,\n    rcp_temperature: tl.float32,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr,\n    USE_TMA: tl.constexpr,\n):\n    \"\"\"\n    backward mainloop, where d_logits & d_hidden & d_weight are fused\n    \"\"\"\n    # block swizzling\n    # pid = tl.program_id(axis=0)\n    # num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    # pid_m = pid % num_pid_m\n    # pid_n = pid // num_pid_m\n\n    pid = tl.program_id(axis=0)\n    num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    num_pid_n = tl.cdiv(vocab_size, BLOCK_SIZE_N)\n    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n    group_id = pid // num_pid_in_group\n    first_pid_m = group_id * GROUP_SIZE_M\n    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)\n    pid_n = (pid % num_pid_in_group) // group_size_m\n\n    start_offs_am = pid_m * BLOCK_SIZE_M\n    offs_am = start_offs_am + tl.arange(0, BLOCK_SIZE_M)\n    start_offs_bn = pid_n * BLOCK_SIZE_N\n    offs_bn = start_offs_bn + tl.arange(0, BLOCK_SIZE_N)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n    if USE_TMA:\n        # using TMA and device-side descriptor creation\n        hidden_desc = tl.make_tensor_descriptor(\n            hidden_ptr,\n            shape=[num_tokens, hidden_size],\n            strides=[stride_hidden_m, 1],\n            block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],\n        )\n\n        weight_desc = tl.make_tensor_descriptor(\n            weight_ptr,\n            shape=[vocab_size, hidden_size],\n            strides=[stride_weight_n, 1],\n            block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],\n        )\n\n    maximum_ptrs = maximum_ptr + offs_am * stride_maximum\n    maximum = tl.load(maximum_ptrs, mask=offs_am < num_tokens, other=0.0)\n    accu_ptrs = accu_ptr + offs_am * stride_accu\n    accu = tl.load(accu_ptrs, mask=offs_am < num_tokens, other=1e-6)  # epsilon to avoid division by zero\n    accu_rcp = tl.fdiv(1.0, accu)\n\n    d_entropy_ptrs = d_entropy_ptr + offs_am * stride_d_entropy\n    d_entropy = tl.load(d_entropy_ptrs, mask=offs_am < num_tokens, other=0.0)\n    if reduction == 0:  # none\n        d_logprobs_ptrs = d_logprobs_ptr + offs_am * stride_d_logprobs\n        d_logprobs = tl.load(d_logprobs_ptrs, mask=offs_am < num_tokens, other=0.0)\n    elif reduction == 1:  # sum\n        d_logprobs = tl.load(d_logprobs_ptr)\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    else:  # mean\n        d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32))\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    d_logprobs = -1 * d_logprobs\n\n    entropy_b_ptrs = entropy_b_ptr + offs_am * stride_entropy_b\n    entropy_b = tl.load(entropy_b_ptrs, mask=offs_am < num_tokens, other=0.0)\n\n    if not USE_TMA:\n        hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n        # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n)\n        weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n    labels_ptrs = labels_ptr + offs_am * stride_labels\n    labels = tl.load(labels_ptrs, mask=offs_am < num_tokens, other=0)\n\n    d_hidden_ptrs = d_hidden_ptr + offs_am[:, None] * stride_d_hidden_m + offs_k[None, :] * stride_d_hidden_k\n    # d_weight_ptrs = d_weight_ptr + offs_k[:, None] * stride_d_weight_k + offs_bn[None, :] * stride_d_weight_n\n    d_weight_ptrs = d_weight_ptr + offs_bn[:, None] * stride_d_weight_n + offs_k[None, :] * stride_d_weight_k\n\n    logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n    for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n        if USE_TMA:\n            start_offs_k = k * BLOCK_SIZE_K\n            _hidden = hidden_desc.load([start_offs_am, start_offs_k])\n            _weight = weight_desc.load([start_offs_bn, start_offs_k])\n        else:\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n                other=0.0,\n            )\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size),\n                other=0.0,\n            )\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n\n        logits = tl.dot(_hidden, _weight.T, logits)\n\n    if not USE_TMA:\n        hidden_ptrs -= hidden_size * stride_hidden_k\n        weight_ptrs -= hidden_size * stride_weight_k\n\n    # scale logits by temperature\n    logits *= rcp_temperature\n\n    exp_logits = tl.exp(logits - maximum[:, None])\n\n    mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None]\n    d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask)\n    d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None])\n\n    # scale d_logits by temperature\n    d_logits *= rcp_temperature\n\n    # loop for d_weight & d_hidden\n    for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n        start_offs_k = k * BLOCK_SIZE_K\n        if USE_TMA:\n            _hidden = hidden_desc.load([start_offs_am, start_offs_k])\n        else:\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n                other=0.0,\n            )\n            # _d_weight = tl.dot(tl.trans(_hidden).to(tl.float32), d_logits)\n            # tl.atomic_add(d_weight_ptrs,\n            #               _d_weight,\n            #               mask=(offs_k[:, None] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[None, :] < vocab_size))\n        _d_weight = tl.dot(d_logits.trans(), _hidden.to(tl.float32))\n        tl.atomic_add(\n            d_weight_ptrs,\n            _d_weight,\n            mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size),\n        )\n\n        if USE_TMA:\n            _weight = weight_desc.load([start_offs_bn, start_offs_k])\n        else:\n            # _weight = tl.load(\n            #     weight_ptrs,\n            #     mask=(offs_k[:, None] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[None, :] < vocab_size),\n            #     other=0.0\n            # )\n            # _d_hidden = tl.dot(d_logits, tl.trans(_weight).to(tl.float32))\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size),\n                other=0.0,\n            )\n        _d_hidden = tl.dot(d_logits, _weight.to(tl.float32))\n        tl.atomic_add(\n            d_hidden_ptrs,\n            _d_hidden,\n            mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n        )\n\n        if not USE_TMA:\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n        d_hidden_ptrs += BLOCK_SIZE_K * stride_d_hidden_k\n        d_weight_ptrs += BLOCK_SIZE_K * stride_d_weight_k\n\n\n@triton.autotune(\n    configs=[\n        triton.Config(\n            {\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 128, \"BLOCK_SIZE_K\": 32, \"GROUP_SIZE_M\": 16},\n            num_stages=3,\n            num_warps=8,\n        ),\n    ],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_backward_kernel_d_hidden(\n    num_tokens: int,\n    hidden_size: int,\n    vocab_size: int,\n    rank: int,\n    hidden_ptr,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    weight_ptr,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    labels_ptr,\n    stride_labels: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accu_ptr,\n    stride_accu: tl.int64,\n    d_entropy_ptr,\n    stride_d_entropy: tl.int64,\n    d_logprobs_ptr,\n    stride_d_logprobs: tl.int64,\n    reduction: int,\n    entropy_b_ptr,\n    stride_entropy_b: tl.int64,\n    d_hidden_ptr,\n    stride_d_hidden_m: tl.int64,\n    stride_d_hidden_k: tl.int64,\n    rcp_temperature: tl.float32,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr,\n):\n    \"\"\"\n    backward d_hidden\n    \"\"\"\n    pid = tl.program_id(axis=0)\n    num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    pid_m = pid % num_pid_m\n    pid_k = pid // num_pid_m\n\n    offs_m = pid_m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n    result_offs_k = pid_k * BLOCK_SIZE_K + offs_k\n\n    maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens, other=0.0)\n    accu = tl.load(accu_ptr + offs_m * stride_accu, mask=offs_m < num_tokens, other=1e-6)\n    accu_rcp = tl.fdiv(1.0, accu)\n    d_entropy = tl.load(d_entropy_ptr + offs_m * stride_d_entropy, mask=offs_m < num_tokens, other=0.0)\n    if reduction == 0:\n        d_logprobs = tl.load(d_logprobs_ptr + offs_m * stride_d_logprobs, mask=offs_m < num_tokens, other=0.0)\n    elif reduction == 1:\n        d_logprobs = tl.load(d_logprobs_ptr)\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    else:\n        d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32))\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    d_logprobs = -1 * d_logprobs\n\n    entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens, other=0.0)\n    labels = tl.load(labels_ptr + offs_m * stride_labels, mask=offs_m < num_tokens, other=0)\n\n    # iterate over vocab_size\n    d_hidden = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_K), dtype=tl.float32)\n    for n in range(0, tl.cdiv(vocab_size, BLOCK_SIZE_N)):\n        offs_n = n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n\n        hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n        weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n\n        # iterate over hidden_size to get logits\n        logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n        for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_m[:, None] < num_tokens),\n                other=0.0,\n            )\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_n[:, None] < vocab_size),\n                other=0.0,\n            )\n\n            logits = tl.dot(_hidden, _weight.trans(), logits)\n\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n\n        # scale logits by temperature\n        logits *= rcp_temperature\n\n        exp_logits = tl.exp(logits - maximum[:, None])\n\n        mask = (offs_n + rank * vocab_size)[None, :] == labels[:, None]\n        d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask)\n        d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None])\n\n        # scale d_logits\n        d_logits *= rcp_temperature\n\n        # calculate d_hidden\n        weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + result_offs_k[None, :] * stride_weight_k)\n        _weight = tl.load(\n            weight_ptrs, mask=(result_offs_k[None, :] < hidden_size) & (offs_n[:, None] < vocab_size), other=0.0\n        )\n        d_hidden = tl.dot(d_logits.to(weight_ptr.dtype.element_ty), _weight, d_hidden)\n\n    # write back\n    tl.store(\n        d_hidden_ptr + offs_m[:, None] * stride_d_hidden_m + result_offs_k[None, :] * stride_d_hidden_k,\n        d_hidden,\n        mask=(offs_m[:, None] < num_tokens) & (result_offs_k[None, :] < hidden_size),\n    )\n\n\n@triton.autotune(\n    configs=[\n        triton.Config(\n            {\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 128, \"BLOCK_SIZE_K\": 32, \"GROUP_SIZE_M\": 16},\n            num_stages=3,\n            num_warps=8,\n        ),\n    ],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_backward_kernel_d_weight(\n    num_tokens: int,\n    hidden_size: int,\n    vocab_size: int,\n    rank: int,\n    hidden_ptr,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    weight_ptr,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    labels_ptr,\n    stride_labels: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accu_ptr,\n    stride_accu: tl.int64,\n    d_entropy_ptr,\n    stride_d_entropy: tl.int64,\n    d_logprobs_ptr,\n    stride_d_logprobs: tl.int64,\n    reduction: int,\n    entropy_b_ptr,\n    stride_entropy_b: tl.int64,\n    d_weight_ptr,\n    stride_d_weight_n: tl.int64,\n    stride_d_weight_k: tl.int64,\n    rcp_temperature: tl.float32,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr,\n):\n    pid = tl.program_id(axis=0)\n    num_pid_n = tl.cdiv(vocab_size, BLOCK_SIZE_N)\n    pid_n = pid % num_pid_n\n    pid_k = pid // num_pid_n\n\n    offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n    result_offs_k = pid_k * BLOCK_SIZE_K + offs_k\n\n    d_weight = tl.zeros((BLOCK_SIZE_N, BLOCK_SIZE_K), dtype=tl.float32)\n    for m in range(0, tl.cdiv(num_tokens, BLOCK_SIZE_M)):\n        offs_m = m * BLOCK_SIZE_M + tl.arange(0, BLOCK_SIZE_M)\n\n        maximum = tl.load(maximum_ptr + offs_m * stride_maximum, mask=offs_m < num_tokens, other=0.0)\n        accu = tl.load(accu_ptr + offs_m * stride_accu, mask=offs_m < num_tokens, other=1e-6)\n        accu_rcp = tl.fdiv(1.0, accu)\n        d_entropy = tl.load(d_entropy_ptr + offs_m * stride_d_entropy, mask=offs_m < num_tokens, other=0.0)\n        if reduction == 0:\n            d_logprobs = tl.load(d_logprobs_ptr + offs_m * stride_d_logprobs, mask=offs_m < num_tokens, other=0.0)\n        elif reduction == 1:\n            d_logprobs = tl.load(d_logprobs_ptr)\n            d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n        else:\n            d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32))\n            d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n        d_logprobs = -1 * d_logprobs\n\n        entropy_b = tl.load(entropy_b_ptr + offs_m * stride_entropy_b, mask=offs_m < num_tokens, other=0.0)\n        labels = tl.load(labels_ptr + offs_m * stride_labels, mask=offs_m < num_tokens, other=0)\n\n        hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n        weight_ptrs = weight_ptr + (offs_n[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n\n        logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n        for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_m[:, None] < num_tokens),\n                other=0.0,\n            )\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_n[:, None] < vocab_size),\n                other=0.0,\n            )\n\n            logits = tl.dot(_hidden, _weight.trans(), logits)\n\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n\n        logits *= rcp_temperature\n\n        exp_logits = tl.exp(logits - maximum[:, None])\n\n        mask = (offs_n + rank * vocab_size)[None, :] == labels[:, None]\n        d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask)\n        d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None])\n\n        d_logits *= rcp_temperature\n\n        hidden_ptrs = hidden_ptr + (offs_m[:, None] * stride_hidden_m + result_offs_k[None, :] * stride_hidden_k)\n        _hidden = tl.load(\n            hidden_ptrs, mask=(result_offs_k[None, :] < hidden_size) & (offs_m[:, None] < num_tokens), other=0.0\n        )\n        d_weight = tl.dot(d_logits.to(d_weight_ptr.dtype.element_ty).trans(), _hidden, d_weight)\n\n    # write back\n    tl.store(\n        d_weight_ptr + offs_n[:, None] * stride_d_weight_n + result_offs_k[None, :] * stride_d_weight_k,\n        d_weight,\n        mask=(offs_n[:, None] < vocab_size) & (result_offs_k[None, :] < hidden_size),\n    )\n\n\n# NOTE: split tile from d_logits' perspective\n@triton.autotune(\n    configs=[\n        triton.Config(\n            {\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 256, \"BLOCK_SIZE_K\": 32, \"GROUP_SIZE_M\": 16},\n            num_stages=3,\n            num_warps=8,\n        ),\n    ],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_backward_kernel_general_d_logits(\n    num_tokens: int,\n    hidden_size: int,\n    vocab_size: int,\n    rank: int,\n    hidden_ptr,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    weight_ptr,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    labels_ptr,\n    stride_labels: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accu_ptr,\n    stride_accu: tl.int64,\n    d_entropy_ptr,\n    stride_d_entropy: tl.int64,\n    d_logprobs_ptr,\n    stride_d_logprobs: tl.int64,\n    reduction: int,\n    entropy_b_ptr,\n    stride_entropy_b,\n    d_logits_ptr,\n    stride_d_logits_m: tl.int64,\n    stride_d_logits_n: tl.int64,\n    rcp_temperature: tl.float32,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr,\n    USE_TMA: tl.constexpr,\n):\n    \"\"\"\n    backward d_logits\n    \"\"\"\n    # block swizzling\n    # pid = tl.program_id(axis=0)\n    # num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    # pid_m = pid % num_pid_m\n    # pid_n = pid // num_pid_m\n\n    pid = tl.program_id(axis=0)\n    num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    num_pid_n = tl.cdiv(vocab_size, BLOCK_SIZE_N)\n    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n    group_id = pid // num_pid_in_group\n    first_pid_m = group_id * GROUP_SIZE_M\n    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)\n    pid_n = (pid % num_pid_in_group) // group_size_m\n\n    start_offs_am = pid_m * BLOCK_SIZE_M\n    offs_am = start_offs_am + tl.arange(0, BLOCK_SIZE_M)\n    start_offs_bn = pid_n * BLOCK_SIZE_N\n    offs_bn = start_offs_bn + tl.arange(0, BLOCK_SIZE_N)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n\n    maximum_ptrs = maximum_ptr + offs_am * stride_maximum\n    maximum = tl.load(maximum_ptrs, mask=offs_am < num_tokens, other=0.0)\n    accu_ptrs = accu_ptr + offs_am * stride_accu\n    accu = tl.load(accu_ptrs, mask=offs_am < num_tokens, other=1e-6)  # epsilon to avoid division by zero\n    accu_rcp = tl.fdiv(1.0, accu)\n\n    d_entropy_ptrs = d_entropy_ptr + offs_am * stride_d_entropy\n    d_entropy = tl.load(d_entropy_ptrs, mask=offs_am < num_tokens, other=0.0)\n    if reduction == 0:  # none\n        d_logprobs_ptrs = d_logprobs_ptr + offs_am * stride_d_logprobs\n        d_logprobs = tl.load(d_logprobs_ptrs, mask=offs_am < num_tokens, other=0.0)\n    elif reduction == 1:  # sum\n        d_logprobs = tl.load(d_logprobs_ptr)\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    else:  # mean\n        d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32))\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    d_logprobs = -1 * d_logprobs\n\n    entropy_b_ptrs = entropy_b_ptr + offs_am * stride_entropy_b\n    entropy_b = tl.load(entropy_b_ptrs, mask=offs_am < num_tokens, other=0.0)\n\n    labels_ptrs = labels_ptr + offs_am * stride_labels\n    labels = tl.load(labels_ptrs, mask=offs_am < num_tokens, other=0)\n\n    logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n\n    if USE_TMA:\n        # using TMA and device-side descriptor creation\n        hidden_desc = tl.make_tensor_descriptor(\n            hidden_ptr,\n            shape=[num_tokens, hidden_size],\n            strides=[stride_hidden_m, 1],\n            block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],\n        )\n        weight_desc = tl.make_tensor_descriptor(\n            weight_ptr,\n            shape=[vocab_size, hidden_size],\n            strides=[stride_weight_n, 1],\n            block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],\n        )\n    else:\n        hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n        # weight_ptrs = weight_ptr + (offs_k[:, None] * stride_weight_k + offs_bn[None, :] * stride_weight_n)\n        weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n\n    for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n        if USE_TMA:\n            start_offs_k = k * BLOCK_SIZE_K\n            _hidden = hidden_desc.load([start_offs_am, start_offs_k])\n            _weight = weight_desc.load([start_offs_bn, start_offs_k])\n        else:\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n                other=0.0,\n            )\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_size),\n                other=0.0,\n            )\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n        logits = tl.dot(_hidden, _weight.T, logits)\n\n    if not USE_TMA:\n        hidden_ptrs -= hidden_size * stride_hidden_k\n        weight_ptrs -= hidden_size * stride_weight_k\n\n    # scale logits by temperature\n    logits *= rcp_temperature\n\n    exp_logits = tl.exp(logits - maximum[:, None])\n\n    mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None]\n    d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask)\n    d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None])\n\n    # scale d_logits by temperature\n    d_logits *= rcp_temperature\n\n    # store d_logits\n    d_logits_ptrs = d_logits_ptr + offs_am[:, None] * stride_d_logits_m + offs_bn[None, :] * stride_d_logits_n\n    tl.store(\n        d_logits_ptrs,\n        d_logits,  # will be implicitly converted to d_logits_ptrs.dtype.element_ty\n        mask=(offs_am[:, None] < num_tokens) & (offs_bn[None, :] < vocab_size),\n    )\n\n\n@triton.autotune(\n    configs=[\n        triton.Config(\n            {\"BLOCK_SIZE_M\": 128, \"BLOCK_SIZE_N\": 256, \"BLOCK_SIZE_K\": 32, \"GROUP_SIZE_M\": 16},\n            num_stages=3,\n            num_warps=8,\n        ),\n    ],\n    key=[\"num_tokens\", \"hidden_size\", \"vocab_size\"],\n)\n@triton.jit\ndef efficient_entropy_backward_kernel_general_d_logits_split_N(\n    split_idx: int,\n    num_tokens: int,\n    hidden_size: int,\n    vocab_size: int,\n    vocab_per_split: int,\n    rank: int,\n    hidden_ptr,\n    stride_hidden_m: tl.int64,\n    stride_hidden_k: tl.int64,\n    weight_ptr,\n    stride_weight_n: tl.int64,\n    stride_weight_k: tl.int64,\n    labels_ptr,\n    stride_labels: tl.int64,\n    maximum_ptr,\n    stride_maximum: tl.int64,\n    accu_ptr,\n    stride_accu: tl.int64,\n    d_entropy_ptr,\n    stride_d_entropy: tl.int64,\n    d_logprobs_ptr,\n    stride_d_logprobs: tl.int64,\n    reduction: int,\n    entropy_b_ptr,\n    stride_entropy_b,\n    d_logits_ptr,\n    stride_d_logits_m: tl.int64,\n    stride_d_logits_n: tl.int64,\n    rcp_temperature: tl.float32,\n    BLOCK_SIZE_M: tl.constexpr,\n    BLOCK_SIZE_N: tl.constexpr,\n    BLOCK_SIZE_K: tl.constexpr,\n    GROUP_SIZE_M: tl.constexpr,\n    USE_TMA: tl.constexpr,\n):\n    pid = tl.program_id(axis=0)\n    num_pid_m = tl.cdiv(num_tokens, BLOCK_SIZE_M)\n    num_pid_n = tl.cdiv(vocab_per_split, BLOCK_SIZE_N)\n    num_pid_in_group = GROUP_SIZE_M * num_pid_n\n    group_id = pid // num_pid_in_group\n    first_pid_m = group_id * GROUP_SIZE_M\n    group_size_m = min(num_pid_m - first_pid_m, GROUP_SIZE_M)\n    pid_m = first_pid_m + ((pid % num_pid_in_group) % group_size_m)\n    pid_n = (pid % num_pid_in_group) // group_size_m\n\n    start_offs_am = pid_m * BLOCK_SIZE_M\n    offs_am = start_offs_am + tl.arange(0, BLOCK_SIZE_M)\n    start_offs_bn = split_idx * vocab_per_split + pid_n * BLOCK_SIZE_N\n    offs_bn = start_offs_bn + tl.arange(0, BLOCK_SIZE_N)\n    offs_k = tl.arange(0, BLOCK_SIZE_K)\n\n    maximum = tl.load(maximum_ptr + offs_am * stride_maximum, mask=offs_am < num_tokens, other=0.0)\n    accu = tl.load(accu_ptr + offs_am * stride_accu, mask=offs_am < num_tokens, other=1e-6)\n    accu_rcp = tl.fdiv(1.0, accu)\n    d_entropy = tl.load(d_entropy_ptr + offs_am * stride_d_entropy, mask=offs_am < num_tokens, other=0.0)\n    if reduction == 0:\n        d_logprobs = tl.load(d_logprobs_ptr + offs_am * stride_d_logprobs, mask=offs_am < num_tokens, other=0.0)\n    elif reduction == 1:\n        d_logprobs = tl.load(d_logprobs_ptr)\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    else:\n        d_logprobs = tl.fdiv(tl.load(d_logprobs_ptr), num_tokens.to(tl.float32))\n        d_logprobs = tl.broadcast_to(d_logprobs, (BLOCK_SIZE_M,))\n    d_logprobs = -1 * d_logprobs\n    entropy_b = tl.load(entropy_b_ptr + offs_am * stride_entropy_b, mask=offs_am < num_tokens, other=0.0)\n    labels = tl.load(labels_ptr + offs_am * stride_labels, mask=offs_am < num_tokens, other=0)\n\n    logits = tl.zeros((BLOCK_SIZE_M, BLOCK_SIZE_N), dtype=tl.float32)\n\n    if USE_TMA:\n        # using TMA and device-side descriptor creation\n        hidden_desc = tl.make_tensor_descriptor(\n            hidden_ptr,\n            shape=[num_tokens, hidden_size],\n            strides=[stride_hidden_m, 1],\n            block_shape=[BLOCK_SIZE_M, BLOCK_SIZE_K],\n        )\n        weight_desc = tl.make_tensor_descriptor(\n            weight_ptr,\n            shape=[vocab_size, hidden_size],\n            strides=[stride_weight_n, 1],\n            block_shape=[BLOCK_SIZE_N, BLOCK_SIZE_K],\n        )\n    else:\n        hidden_ptrs = hidden_ptr + (offs_am[:, None] * stride_hidden_m + offs_k[None, :] * stride_hidden_k)\n        weight_ptrs = weight_ptr + (offs_bn[:, None] * stride_weight_n + offs_k[None, :] * stride_weight_k)\n        vocab_right_bound = min((split_idx + 1) * vocab_per_split, vocab_size)\n\n    for k in range(0, tl.cdiv(hidden_size, BLOCK_SIZE_K)):\n        if USE_TMA:\n            start_offs_k = k * BLOCK_SIZE_K\n            _hidden = hidden_desc.load([start_offs_am, start_offs_k])\n            _weight = weight_desc.load([start_offs_bn, start_offs_k])\n        else:\n            _hidden = tl.load(\n                hidden_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_am[:, None] < num_tokens),\n                other=0.0,\n            )\n            _weight = tl.load(\n                weight_ptrs,\n                mask=(offs_k[None, :] < hidden_size - k * BLOCK_SIZE_K) & (offs_bn[:, None] < vocab_right_bound),\n                other=0.0,\n            )\n            hidden_ptrs += BLOCK_SIZE_K * stride_hidden_k\n            weight_ptrs += BLOCK_SIZE_K * stride_weight_k\n        logits = tl.dot(_hidden, _weight.T, logits)\n\n    logits *= rcp_temperature\n    exp_logits = tl.exp(logits - maximum[:, None])\n\n    mask = (offs_bn + rank * vocab_size)[None, :] == labels[:, None]\n    d_logits = d_logprobs[:, None] * (exp_logits * accu_rcp[:, None] - mask)\n    d_logits += d_entropy[:, None] * (-exp_logits * accu_rcp[:, None]) * (logits - entropy_b[:, None])\n\n    d_logits *= rcp_temperature\n\n    # filter d_logits with mask\n    result_offs_n = pid_n * BLOCK_SIZE_N + tl.arange(0, BLOCK_SIZE_N)\n    mask = (offs_am[:, None] < num_tokens) & (result_offs_n[None, :] < vocab_per_split)\n\n    tl.store(\n        d_logits_ptr + offs_am[:, None] * stride_d_logits_m + result_offs_n[None, :] * stride_d_logits_n, d_logits, mask\n    )\n\n\ndef efficient_entropy_backward(\n    dlogprobs: torch.Tensor,\n    dentropy: torch.Tensor,\n    hidden: torch.Tensor,\n    weight: torch.Tensor,\n    labels: torch.Tensor,\n    maximum: torch.Tensor,\n    acc: torch.Tensor,\n    entropy_b: torch.Tensor,\n    reduction: typing.Optional[int] = 2,\n    should_return_fp32_grad: bool = False,\n    temperature: typing.Optional[float] = 1.0,\n    dist_process_group: typing.Optional[dist.ProcessGroup] = None,\n) -> list[torch.Tensor]:\n    \"\"\"\n    backward host function\n    \"\"\"\n    assert hidden.is_cuda and weight.is_cuda and labels.is_cuda\n    assert weight.device == hidden.device and labels.device == hidden.device\n    assert hidden.dim() == 2 and weight.dim() == 2 and labels.dim() == 1\n    assert hidden.is_contiguous() and weight.is_contiguous() and labels.is_contiguous()\n    assert hidden.shape[0] == labels.shape[0] and hidden.shape[1] == weight.shape[1]\n\n    _rank = 0 if dist_process_group is None else dist.get_rank(dist_process_group)\n    _world_size = 1 if dist_process_group is None else dist.get_world_size(dist_process_group)\n\n    num_tokens, hidden_size = hidden.shape\n    num_tokens = labels.shape[0]\n    vocab_size, hidden_size = weight.shape\n    assert hidden_size % 128 == 0\n\n    REDUCTION = get_entropy_reduction_enum(reduction)\n\n    if REDUCTION == EntropyReductionEnum._None:\n        assert dlogprobs.shape == (num_tokens,)\n    else:\n        assert dlogprobs.dim() == 0\n\n    assert dlogprobs.is_contiguous() and dentropy.is_contiguous()\n    assert dlogprobs.is_cuda and dentropy.is_cuda\n    assert dlogprobs.device == hidden.device and dlogprobs.device == dentropy.device\n    assert dentropy.shape == (num_tokens,)\n\n    d_hidden, d_weight = None, None\n    if _config._backward == BackwardEnum._Total_Fuse_MN or should_return_fp32_grad:\n        d_hidden = torch.zeros_like(hidden, dtype=torch.float32, device=hidden.device)\n        d_weight = torch.zeros_like(weight, dtype=torch.float32, device=weight.device)\n    else:\n        d_hidden = torch.empty_like(hidden, dtype=hidden.dtype, device=hidden.device)\n        d_weight = torch.empty_like(weight, dtype=hidden.dtype, device=weight.device)\n    assert d_hidden.is_contiguous() and d_weight.is_contiguous()\n\n    assert maximum.is_contiguous() and acc.is_contiguous()\n    assert maximum.device == hidden.device and acc.device == hidden.device\n    assert maximum.shape == labels.shape == acc.shape\n    assert maximum.is_cuda and acc.is_cuda\n\n    vocab_per_split = 1024\n    assert vocab_per_split % 128 == 0\n    num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split\n\n    assert entropy_b.is_contiguous() and entropy_b.is_cuda\n    assert entropy_b.shape == (num_tokens,)\n\n    if _config._backward == BackwardEnum._Total_Fuse_MN:\n        # --- Triton doesn't materialize d_logits at all. Split tiles at the perspective of d_logits.\n        def mainloop_grid(meta):\n            return (triton.cdiv(num_tokens, meta[\"BLOCK_SIZE_M\"]) * triton.cdiv(vocab_size, meta[\"BLOCK_SIZE_N\"]),)\n\n        efficient_entropy_backward_kernel_general_mainloop_MN[mainloop_grid](\n            num_tokens,\n            hidden_size,\n            vocab_size,\n            _rank,\n            hidden,\n            hidden.stride(0),\n            hidden.stride(1),\n            weight,\n            weight.stride(0),\n            weight.stride(1),\n            labels,\n            labels.stride(0),\n            maximum,\n            maximum.stride(0),\n            acc,\n            acc.stride(0),\n            dentropy,\n            dentropy.stride(0),\n            dlogprobs,\n            dlogprobs.stride(0) if REDUCTION == EntropyReductionEnum._None else 0,\n            REDUCTION,\n            entropy_b,\n            entropy_b.stride(0),\n            d_hidden,\n            d_hidden.stride(0),\n            d_hidden.stride(1),\n            d_weight,\n            d_weight.stride(0),\n            d_weight.stride(1),\n            1.0 / temperature,\n            USE_TMA=SUPPORT_CUDA_TMA and hidden.stride(1) == 1 and weight.stride(1) == 1,\n        )\n\n    elif _config._backward == BackwardEnum._Total_Separate:\n        _d_logits = torch.empty((num_tokens, vocab_size), device=hidden.device, dtype=hidden.dtype).contiguous()\n        assert _d_logits.is_contiguous()\n\n        if _config._use_triton:\n\n            def d_logits_grid(meta):\n                return (triton.cdiv(num_tokens, meta[\"BLOCK_SIZE_M\"]) * triton.cdiv(vocab_size, meta[\"BLOCK_SIZE_N\"]),)\n\n            efficient_entropy_backward_kernel_general_d_logits[d_logits_grid](\n                num_tokens,\n                hidden_size,\n                vocab_size,\n                _rank,\n                hidden,\n                hidden.stride(0),\n                hidden.stride(1),\n                weight,\n                weight.stride(0),\n                weight.stride(1),\n                labels,\n                labels.stride(0),\n                maximum,\n                maximum.stride(0),\n                acc,\n                acc.stride(0),\n                dentropy,\n                dentropy.stride(0),\n                dlogprobs,\n                dlogprobs.stride(0) if REDUCTION == EntropyReductionEnum._None else 0,\n                REDUCTION,\n                entropy_b,\n                entropy_b.stride(0),\n                _d_logits,\n                _d_logits.stride(0),\n                _d_logits.stride(1),\n                1.0 / temperature,\n                USE_TMA=SUPPORT_CUDA_TMA and hidden.stride(1) == 1 and weight.stride(1) == 1,\n            )\n\n            torch.matmul(_d_logits, weight, out=d_hidden)\n            torch.matmul(_d_logits.T, hidden, out=d_weight)\n        else:\n            raise AssertionError(\"Triton is required for efficient entropy kernel\")\n\n    elif _config._backward == BackwardEnum._Split_Dlogits_N:\n        vocab_per_split = 9504\n        num_splits = (vocab_size + vocab_per_split - 1) // vocab_per_split\n\n        _d_logits = torch.empty((num_tokens, vocab_per_split), device=hidden.device, dtype=hidden.dtype).contiguous()\n        assert _d_logits.is_contiguous()\n\n        def d_logits_grid(meta):\n            return (triton.cdiv(num_tokens, meta[\"BLOCK_SIZE_M\"]) * triton.cdiv(vocab_per_split, meta[\"BLOCK_SIZE_N\"]),)\n\n        for split_idx in range(num_splits):\n            efficient_entropy_backward_kernel_general_d_logits_split_N[d_logits_grid](\n                split_idx,\n                num_tokens,\n                hidden_size,\n                vocab_size,\n                vocab_per_split,\n                _rank,\n                hidden,\n                hidden.stride(0),\n                hidden.stride(1),\n                weight,\n                weight.stride(0),\n                weight.stride(1),\n                labels,\n                labels.stride(0),\n                maximum,\n                maximum.stride(0),\n                acc,\n                acc.stride(0),\n                dentropy,\n                dentropy.stride(0),\n                dlogprobs,\n                dlogprobs.stride(0) if REDUCTION == EntropyReductionEnum._None else 0,\n                REDUCTION,\n                entropy_b,\n                entropy_b.stride(0),\n                _d_logits,\n                _d_logits.stride(0),\n                _d_logits.stride(1),\n                1.0 / temperature,\n                USE_TMA=SUPPORT_CUDA_TMA and hidden.stride(1) == 1 and weight.stride(1) == 1,\n            )\n\n            if split_idx == (num_splits - 1):\n                vocab_right_bound = min((split_idx + 1) * vocab_per_split, vocab_size) - split_idx * vocab_per_split\n                _d_logits = _d_logits[:, :vocab_right_bound].contiguous()\n\n            if split_idx == 0:\n                torch.matmul(\n                    _d_logits, weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :], out=d_hidden\n                )\n            else:\n                d_hidden += torch.matmul(\n                    _d_logits, weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :]\n                )\n            torch.matmul(\n                _d_logits.T, hidden, out=d_weight[split_idx * vocab_per_split : (split_idx + 1) * vocab_per_split, :]\n            )\n\n    elif _config._backward == BackwardEnum._Split_Dlogits_M:\n        raise NotImplementedError(\"BackwardEnum._Split_Dlogits_M is not implemented yet\")\n\n    return d_hidden, d_weight\n"
  },
  {
    "path": "verl/utils/kernel/linear_cross_entropy.py",
    "content": "#\n# SPDX-FileCopyrightText: Copyright (c) 2025 NVIDIA CORPORATION & AFFILIATES. All rights reserved.\n# SPDX-License-Identifier: Apache-2.0\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n# http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport typing\n\nimport torch\nimport torch.distributed as dist\n\n\nclass LinearCrossEntropy(torch.autograd.Function):\n    @staticmethod\n    def forward(\n        ctx,\n        hidden: torch.Tensor,\n        weight: torch.Tensor,\n        labels: torch.Tensor,\n        temperature: typing.Optional[float] = 1.0,\n        reduction: typing.Optional[str] = \"none\",\n        dist_process_group: typing.Optional[dist.ProcessGroup] = None,\n    ) -> list[torch.Tensor]:\n        \"\"\"_summary_\n\n        Args:\n            ctx (_type_): _description_\n            hidden (torch.Tensor): (batch_size, num_tokens, hidden_size) -> (batch_size * num_tokens, hidden_size)\n            weight (torch.Tensor): (vocab_size, hidden_size)\n            labels (torch.Tensor): (batch_size, num_tokens) -> (batch_size * num_tokens, )\n            temperature (typing.Optional[float], optional): _description_. Defaults to 1.0.\n            reduction (typing.Optional[str], optional): _description_. Defaults to \"none\".\n            dist_process_group (typing.Optional[dist.ProcessGroup], optional): _description_. Defaults to None.\n\n        Returns:\n            typing.List[torch.Tensor]: _description_\n        \"\"\"\n\n        assert isinstance(temperature, float), f\"temperature must be a float, but got {type(temperature)}\"\n        assert isinstance(reduction, str), f\"reduction must be a str, but got {type(reduction)}\"\n        with torch.cuda.nvtx.range(\"LinearCrossEntropy-forward\"):\n            from . import kernels\n\n            REDUCTION = kernels.get_entropy_reduction_enum_number(reduction.lower())\n\n            original_hidden_shape = hidden.shape\n            if len(hidden.shape) != 2:\n                hidden = hidden.view(-1, hidden.shape[-1])  # (batch_size * num_tokens, hidden_size)\n            if len(labels.shape) != 1:\n                labels = labels.view(-1)\n\n            logprobs, entropy, _maximum, _accumulate, _entropy_b = kernels.efficient_entropy_forward(\n                hidden, weight, labels, REDUCTION, temperature, dist_process_group\n            )\n\n            ctx.save_for_backward(hidden, weight, labels, _maximum, _accumulate, _entropy_b)\n            ctx.original_hidden_shape = original_hidden_shape\n            ctx.REDUCTION = REDUCTION\n            ctx.dist_process_group = dist_process_group\n            ctx.should_return_fp32_grad = False\n            ctx.temperature = temperature\n        return logprobs, entropy\n\n    @staticmethod\n    def backward(ctx, dlogprobs: torch.Tensor, dentropy: torch.Tensor) -> list[torch.Tensor]:\n        from . import kernels\n\n        with torch.cuda.nvtx.range(\"LinearCrossEntropy-backward\"):\n            (hidden, weight, labels, _maximum, _accumulate, _entropy_b) = ctx.saved_tensors\n            REDUCTION = ctx.REDUCTION\n            dist_process_group = ctx.dist_process_group\n            should_return_fp32_grad = ctx.should_return_fp32_grad\n            temperature = ctx.temperature\n\n            d_hidden, d_weight = kernels.efficient_entropy_backward(\n                dlogprobs,\n                dentropy,\n                hidden,\n                weight,\n                labels,\n                _maximum,\n                _accumulate,\n                _entropy_b,\n                REDUCTION,\n                should_return_fp32_grad,\n                temperature,\n                dist_process_group,\n            )\n            d_hidden = d_hidden.view(ctx.original_hidden_shape)\n\n        return (d_hidden, d_weight, None, None, None, None)\n\n\nlinear_cross_entropy = LinearCrossEntropy.apply\n"
  },
  {
    "path": "verl/utils/logger/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom .aggregate_logger import (\n    DecoratorLoggerBase,\n    LocalLogger,\n    log_with_rank,\n    print_rank_0,\n    print_with_rank,\n    print_with_rank_and_timer,\n)\n\n__all__ = [\n    \"LocalLogger\",\n    \"DecoratorLoggerBase\",\n    \"print_rank_0\",\n    \"print_with_rank\",\n    \"print_with_rank_and_timer\",\n    \"log_with_rank\",\n]\n"
  },
  {
    "path": "verl/utils/logger/aggregate_logger.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nA Ray logger will receive logging info from different processes.\n\"\"\"\n\nimport datetime\nimport logging\nimport numbers\nimport pprint\n\nimport torch\n\n\ndef concat_dict_to_str(dict: dict, step):\n    output = [f\"step:{step}\"]\n    for k, v in dict.items():\n        if isinstance(v, numbers.Number):\n            output.append(f\"{k}:{pprint.pformat(v)}\")\n    output_str = \" - \".join(output)\n    return output_str\n\n\nclass LocalLogger:\n    \"\"\"\n    A local logger that logs messages to the console.\n\n    Args:\n        print_to_console (bool): Whether to print to the console.\n    \"\"\"\n\n    def __init__(self, print_to_console=True):\n        self.print_to_console = print_to_console\n\n    def flush(self):\n        pass\n\n    def log(self, data, step):\n        if self.print_to_console:\n            print(concat_dict_to_str(data, step=step), flush=True)\n\n\nclass DecoratorLoggerBase:\n    \"\"\"\n    Base class for all decorators that log messages.\n\n    Args:\n        role (str): The role (the name) of the logger.\n        logger (logging.Logger): The logger instance to use for logging.\n        level (int): The logging level.\n        rank (int): The rank of the process.\n        log_only_rank_0 (bool): If True, only log for rank 0.\n    \"\"\"\n\n    def __init__(\n        self, role: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0, log_only_rank_0: bool = True\n    ):\n        self.role = role\n        self.logger = logger\n        self.level = level\n        self.rank = rank\n        self.log_only_rank_0 = log_only_rank_0\n        self.logging_function = self.log_by_logging\n        if logger is None:\n            self.logging_function = self.log_by_print\n\n    def log_by_print(self, log_str):\n        if not self.log_only_rank_0 or self.rank == 0:\n            print(f\"{self.role} {log_str}\", flush=True)\n\n    def log_by_logging(self, log_str):\n        if self.logger is None:\n            raise ValueError(\"Logger is not initialized\")\n        if not self.log_only_rank_0 or self.rank == 0:\n            self.logger.log(self.level, f\"{self.role} {log_str}\")\n\n\ndef print_rank_0(message):\n    \"\"\"If distributed is initialized, print only on rank 0.\"\"\"\n    if torch.distributed.is_initialized():\n        if torch.distributed.get_rank() == 0:\n            print(message, flush=True)\n    else:\n        print(message, flush=True)\n\n\ndef print_with_rank(message: str, rank: int = 0, log_only_rank_0: bool = False):\n    \"\"\"_summary_\n    Print a message with rank information.\n    This function prints the message only if `log_only_rank_0` is False or if the rank is 0.\n\n    Args:\n        message (str): _description_\n        rank (int, optional): _description_. Defaults to 0.\n        log_only_rank_0 (bool, optional): _description_. Defaults to False.\n    \"\"\"\n    if not log_only_rank_0 or rank == 0:\n        print(f\"[Rank {rank}] {message}\", flush=True)\n\n\ndef print_with_rank_and_timer(message: str, rank: int = 0, log_only_rank_0: bool = False):\n    \"\"\"_summary_\n    Print a message with rank information and a timestamp.\n    This function prints the message only if `log_only_rank_0` is False or if the rank is 0.\n\n    Args:\n        message (str): _description_\n        rank (int, optional): _description_. Defaults to 0.\n        log_only_rank_0 (bool, optional): _description_. Defaults to False.\n    \"\"\"\n    now = datetime.datetime.now()\n    message = f\"[{now.strftime('%Y-%m-%d %H:%M:%S')}] [Rank {rank}] {message}\"\n    if not log_only_rank_0 or rank == 0:\n        print(message, flush=True)\n\n\ndef log_with_rank(message: str, rank, logger: logging.Logger, level=logging.INFO, log_only_rank_0: bool = False):\n    \"\"\"_summary_\n    Log a message with rank information using a logger.\n    This function logs the message only if `log_only_rank_0` is False or if the rank is 0.\n    Args:\n        message (str): The message to log.\n        rank (int): The rank of the process.\n        logger (logging.Logger): The logger instance to use for logging.\n        level (int, optional): The logging level. Defaults to logging.INFO.\n        log_only_rank_0 (bool, optional): If True, only log for rank 0. Defaults to False.\n    \"\"\"\n    if not log_only_rank_0 or rank == 0:\n        logger.log(level, f\"[Rank {rank}] {message}\")\n"
  },
  {
    "path": "verl/utils/logging_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\n\nimport torch\n\n\ndef set_basic_config(level):\n    \"\"\"\n    This function sets the global logging format and level. It will be called when import verl\n    \"\"\"\n    logging.basicConfig(format=\"%(levelname)s:%(asctime)s:%(message)s\", level=level)\n\n\ndef log_to_file(string):\n    print(string)\n    if os.path.isdir(\"logs\"):\n        with open(f\"logs/log_{torch.distributed.get_rank()}\", \"a+\") as f:\n            f.write(string + \"\\n\")\n"
  },
  {
    "path": "verl/utils/megatron/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/utils/megatron/dist_checkpointing.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport megatron.core\nimport torch\nfrom megatron.core import dist_checkpointing, mpu\nfrom megatron.core.dist_checkpointing.serialization import (\n    get_default_load_sharded_strategy,\n    get_default_save_sharded_strategy,\n)\nfrom megatron.core.dist_checkpointing.strategies.fully_parallel import (\n    FullyParallelLoadStrategyWrapper,\n    FullyParallelSaveStrategyWrapper,\n)\nfrom packaging import version\n\n\ndef save_dist_checkpointing(\n    sharded_state_dict,\n    ckpt_path,\n    async_save=False,\n    content_metadata=None,\n):\n    validate_sharding_integrity = True\n    # Get checkpointing strategies\n    save_strategy = get_default_save_sharded_strategy(\"torch_dist\")\n    save_strategy = FullyParallelSaveStrategyWrapper(\n        save_strategy, mpu.get_data_parallel_group(with_context_parallel=True)\n    )\n\n    # https://github.com/NVIDIA/Megatron-LM/blob/core_v0.14.0/megatron/core/optimizer/distrib_optimizer.py#L1109-L1123\n    mcore_ge_014 = version.parse(megatron.core.__version__) >= version.parse(\"0.14.0\")\n    # Save model sharded state dicts\n    save_kwargs = dict(\n        sharded_strategy=save_strategy,\n        async_sharded_save=async_save,\n        validate_access_integrity=validate_sharding_integrity,\n    )\n    if content_metadata is not None:\n        if mcore_ge_014:\n            save_kwargs[\"content_metadata\"] = content_metadata\n    return dist_checkpointing.save(sharded_state_dict, ckpt_path, **save_kwargs)\n\n\ndef load_dist_checkpointing(sharded_state_dict, ckpt_dir):\n    # Get checkpointing strategies\n    load_strategy = get_default_load_sharded_strategy(ckpt_dir)\n    load_strategy = FullyParallelLoadStrategyWrapper(\n        load_strategy, mpu.get_data_parallel_group(with_context_parallel=True)\n    )\n\n    # Fix torch.load weights only error\n    try:\n        import transformer_engine as te\n\n        torch.serialization.add_safe_globals([torch.optim.AdamW])\n        torch.serialization.add_safe_globals([te.pytorch.optimizers.fused_adam.FusedAdam])\n    except Exception:\n        pass\n\n    # Load model sharded state dicts\n    state_dict = dist_checkpointing.load(sharded_state_dict, ckpt_dir, sharded_strategy=load_strategy)\n\n    return state_dict\n"
  },
  {
    "path": "verl/utils/megatron/memory.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\n\nfrom verl.utils.device import get_device_id\n\n\nclass MemoryBuffer:\n    def __init__(self, numel, numel_padded, dtype):\n        self.numel = numel\n        self.numel_padded = numel_padded\n        self.dtype = dtype\n        self.data = torch.zeros(self.numel_padded, dtype=self.dtype, device=get_device_id(), requires_grad=False)\n\n    def zero(self):\n        \"\"\"Reset the buffer to zero.\"\"\"\n        self.data.zero_()\n\n    def get(self, shape, start_index):\n        \"\"\"Return a tensor with the input `shape` as a view into the\n        1-D data starting at `start_index`.\"\"\"\n        end_index = start_index + shape.numel()\n        assert end_index <= self.numel, \"requested tensor is out of the buffer range.\"\n        buffer_tensor = self.data[start_index:end_index]\n        buffer_tensor = buffer_tensor.view(shape)\n        return buffer_tensor\n"
  },
  {
    "path": "verl/utils/megatron/optimizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom megatron.core.optimizer import OptimizerConfig\nfrom megatron.core.optimizer import get_megatron_optimizer as get_megatron_optimizer_native\nfrom megatron.core.optimizer_param_scheduler import OptimizerParamScheduler\n\nfrom verl.utils.logger import print_rank_0\n\n\ndef init_megatron_optim_config(\n    optim_config: dict, use_distributed_optimizer: bool = True, fp16: bool = False\n) -> OptimizerConfig:\n    optim_args = {\n        \"optimizer\": optim_config.optimizer,\n        \"lr\": optim_config.lr,\n        \"min_lr\": optim_config.min_lr,\n        \"clip_grad\": optim_config.clip_grad,\n        \"weight_decay\": optim_config.weight_decay,\n        \"use_distributed_optimizer\": use_distributed_optimizer,\n    }\n    if fp16:\n        optim_args.update(\n            {\n                \"bf16\": False,\n                \"fp16\": True,\n                \"params_dtype\": torch.float16,\n                \"initial_loss_scale\": 32768,\n                \"min_loss_scale\": 1,\n                \"use_precision_aware_optimizer\": True,\n                \"store_param_remainders\": False,\n            }\n        )\n    else:  # bf16 mode\n        optim_args.update(\n            {\n                \"bf16\": True,\n                \"params_dtype\": torch.bfloat16,\n            }\n        )\n    override_config = optim_config.get(\"override_optimizer_config\", {})\n    if override_config:\n        for k, v in override_config.items():\n            optim_args[k] = v\n\n    print_rank_0(f\"optimizer config after override: {optim_args}\")\n\n    config = OptimizerConfig(**optim_args)\n    return config\n\n\ndef get_megatron_optimizer(\n    model,\n    config: OptimizerConfig,\n):\n    # Base optimizer.\n    return get_megatron_optimizer_native(\n        config=config,\n        model_chunks=model,\n    )\n\n\ndef get_megatron_optimizer_param_scheduler(\n    optimizer,\n    config,\n):\n    \"\"\"\n    Get the optimizer parameter scheduler for Megatron.\n    \"\"\"\n    lr_decay_steps = config.lr_decay_steps\n    lr_warmup_steps = config.lr_warmup_steps\n    if config.get(\"lr_decay_steps\", None) is None:\n        lr_decay_steps = config.total_training_steps\n    wsd_decay_steps = None\n    if config.get(\"lr_wsd_decay_steps\", None) is not None:\n        wsd_decay_steps = config.lr_wsd_decay_steps\n    if config.get(\"lr_warmup_steps_ratio\", None) is not None and (\n        config.get(\"lr_warmup_steps\", None) is None or config.lr_warmup_steps <= 0\n    ):\n        lr_warmup_steps = int(config.lr_warmup_steps_ratio * lr_decay_steps)\n\n    opt_param_scheduler = OptimizerParamScheduler(\n        optimizer,\n        init_lr=config.lr_warmup_init,\n        max_lr=config.lr,\n        min_lr=config.min_lr,\n        lr_warmup_steps=lr_warmup_steps,\n        lr_decay_steps=lr_decay_steps,\n        lr_decay_style=config.lr_decay_style,\n        start_wd=config.weight_decay,\n        end_wd=config.weight_decay,\n        wd_incr_steps=config.total_training_steps,\n        wd_incr_style=config.weight_decay_incr_style,\n        use_checkpoint_opt_param_scheduler=config.use_checkpoint_opt_param_scheduler,\n        override_opt_param_scheduler=(not config.use_checkpoint_opt_param_scheduler),\n        wsd_decay_steps=wsd_decay_steps,\n        lr_wsd_decay_style=config.lr_wsd_decay_style,\n    )\n\n    return opt_param_scheduler\n\n\ndef get_megatron_last_lr(optimizer):\n    \"\"\"\n    Get the last learning rate from the optimizer parameter scheduler.\n    \"\"\"\n    return optimizer.param_groups[0][\"lr\"]\n"
  },
  {
    "path": "verl/utils/megatron/pipeline_parallel.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nfrom megatron.core import parallel_state as mpu\n\nfrom .sequence_parallel import pad_to_sequence_parallel\n\n\ndef compute_transformers_input_shapes(batches, meta_info):\n    from flash_attn.bert_padding import unpad_input  # flash 2 is a must for Megatron\n\n    # pre-compute input shapes for each micro-batch at each pp stage\n    input_shapes = []\n    for model_inputs in batches:\n        input_ids = model_inputs[\"input_ids\"]\n        attention_mask = model_inputs[\"attention_mask\"]\n        input_ids_rmpad = unpad_input(input_ids.unsqueeze(dim=-1), attention_mask)[0]  # (total_nnz, 1)\n        if meta_info[\"sequence_parallel\"]:\n            input_ids_rmpad = pad_to_sequence_parallel(input_ids_rmpad)\n            # compute shapes for model_inputs\n            input_shapes.append(\n                torch.Size(\n                    [\n                        input_ids_rmpad.shape[0] // mpu.get_tensor_model_parallel_world_size(),\n                        1,\n                        meta_info[\"hidden_size\"],\n                    ]\n                )\n            )\n        else:\n            # compute shapes for model_inputs\n            input_shapes.append(torch.Size([input_ids_rmpad.shape[0], 1, meta_info[\"hidden_size\"]]))\n    return input_shapes\n\n\ndef make_batch_generator(batches, vpp_size):\n    \"\"\"\n    Creates a batch generator suitable for Megatron pipeline parallelism,\n    handling virtual pipeline parallelism (VPP).\n\n    If VPP is used (vpp_size > 1), it duplicates the batch iterator for each\n    virtual pipeline stage. Otherwise, it returns a single iterator.\n\n    Args:\n        batches: An iterable (e.g., list) of micro-batches.\n        vpp_size (int): The virtual pipeline model parallel size.\n\n    Returns:\n        An iterator or a list of iterators over the micro-batches.\n    \"\"\"\n    if vpp_size > 1:\n        # has vpp\n        batch_generator = [batches] * vpp_size  # number of vpp chunks\n        batch_generator = [iter(b) for b in batch_generator]\n    else:\n        # no vpp\n        batch_generator = iter(batches)\n    return batch_generator\n"
  },
  {
    "path": "verl/utils/megatron/router_replay_patch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport inspect\nimport types\nimport warnings\nfrom enum import Enum\nfrom functools import wraps\n\nimport torch\n\ntry:\n    from megatron.core.transformer.moe.moe_utils import (\n        apply_router_token_dropping,\n        compute_routing_scores_for_aux_loss,\n        group_limited_topk,\n    )\n    from megatron.core.transformer.moe.token_dispatcher import MoEAlltoAllTokenDispatcher\nexcept ImportError:\n    warnings.warn(\"NPU not support router replay for now.\", stacklevel=2)\n    MoEAlltoAllTokenDispatcher = None\nfrom megatron.core.transformer.moe.router import TopKRouter\nfrom megatron.core.transformer.transformer_config import TransformerConfig\n\n# https://github.com/THUDM/slime/blob/main/slime/utils/routing_replay.py\n\n\nclass RouterReplayAction(Enum):\n    RECORD = \"record\"\n    REPLAY_FORWARD = \"replay_forward\"\n    REPLAY_BACKWARD = \"replay_backward\"\n\n\nclass RouterReplay:\n    \"\"\"\n    A class to manage the recording and replaying of MoE routing decisions.\n    It holds all router instances and provides static methods to globally\n    control recording and replaying.\n    \"\"\"\n\n    # Static variable to hold all router instances, one per MoE layer.\n    router_instances = []\n\n    @staticmethod\n    def set_replay_data(all_layers_topk_indices: list):\n        \"\"\"\n        Distributes the topk indices for all layers to their respective RouterReplay instances.\n        :param all_layers_topk_indices: A list of tensors, where each tensor contains the\n                                        topk indices for a specific layer. The order\n                                        must match the instantiation order of the routers.\n        \"\"\"\n        if len(all_layers_topk_indices) != len(RouterReplay.router_instances):\n            raise ValueError(\n                f\"The number of replay tensors ({len(all_layers_topk_indices)}) \"\n                f\"does not match the number of router instances ({len(RouterReplay.router_instances)}).\"\n            )\n        for i, router_instance in enumerate(RouterReplay.router_instances):\n            router_instance.set_target_indices(all_layers_topk_indices[i])\n\n    @staticmethod\n    def get_recorded_data() -> list:\n        \"\"\"\n        Collects the recorded topk indices from all RouterReplay instances.\n        :return: A list of tensors, each containing the recorded topk indices for a layer.\n        \"\"\"\n        return [router.get_recorded_indices() for router in RouterReplay.router_instances]\n\n    @staticmethod\n    def clear_global_indices():\n        \"\"\"Clears the recorded and target topk indices in all instances.\"\"\"\n        for router in RouterReplay.router_instances:\n            router.clear_indices()\n\n    def __init__(self):\n        \"\"\"Initializes a RouterReplay instance for a specific layer.\"\"\"\n        self.target_topk_idx = None  # For replay\n        self.recorded_topk_idx = None  # For recording\n        self.router_replay_action = None  # Router replay action for this layer\n        self.replay_backward_list = []  # List of tensors for backward pass replay\n        RouterReplay.router_instances.append(self)\n\n    def set_target_indices(self, topk_indices: torch.Tensor):\n        \"\"\"Sets the target topk indices for replay.\"\"\"\n        self.target_topk_idx = topk_indices\n        self.replay_backward_list.append(topk_indices)\n\n    def get_recorded_indices(self):\n        \"\"\"Returns the recorded topk indices.\"\"\"\n        return self.recorded_topk_idx\n\n    def record_indices(self, topk_indices: torch.Tensor):\n        \"\"\"Records the topk indices.\"\"\"\n        self.recorded_topk_idx = topk_indices\n\n    def clear_indices(self):\n        \"\"\"Clears the recorded and target topk indices.\"\"\"\n        self.recorded_topk_idx = None\n        self.target_topk_idx = None\n        self.replay_backward_list = []\n\n    def set_router_replay_action(self, router_replay_action: RouterReplayAction):\n        \"\"\"Sets the router replay action for this layer.\"\"\"\n        self.router_replay_action = router_replay_action\n\n    def clear_router_replay_action(self):\n        \"\"\"Clears the router replay action for this layer.\"\"\"\n        self.router_replay_action = None\n\n    @staticmethod\n    def set_global_router_replay_action(router_replay_action: RouterReplayAction):\n        \"\"\"Sets the router replay action for all router instances.\"\"\"\n        for router in RouterReplay.router_instances:\n            router.set_router_replay_action(router_replay_action)\n\n    @staticmethod\n    def clear_global_router_replay_action():\n        \"\"\"Clears the router replay action for all router instances.\"\"\"\n        for router in RouterReplay.router_instances:\n            router.clear_router_replay_action()\n\n\ndef _patched_topk_routing_with_score_function(\n    logits: torch.Tensor,\n    topk: int,\n    use_pre_softmax: bool,\n    num_groups: int,\n    group_topk: int,\n    score_function: str,\n    expert_bias: torch.Tensor,\n    fused: bool,\n    router_replay: RouterReplay,\n    scaling_factor: float,\n):\n    \"\"\"\n    Patched version of topk_routing_with_score_function that supports router replay.\n    \"\"\"\n    num_tokens, num_experts = logits.shape\n\n    def _compute_topk(scores, topk, num_groups=None, group_topk=None):\n        if group_topk:\n            return group_limited_topk(\n                scores=scores,\n                topk=topk,\n                num_tokens=num_tokens,\n                num_experts=num_experts,\n                num_groups=num_groups,\n                group_topk=group_topk,\n            )\n        else:\n            return torch.topk(scores, k=topk, dim=1)\n\n    def compute_topk(scores, topk, num_groups=None, group_topk=None):\n        # Default behavior if no replay is active\n\n        routing_action = router_replay.router_replay_action if router_replay is not None else None\n\n        if routing_action is None:\n            return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)\n\n        if routing_action == RouterReplayAction.RECORD:\n            probs, top_indices = _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)\n            if router_replay is not None:\n                router_replay.record_indices(top_indices)\n            return probs, top_indices\n\n        elif routing_action == RouterReplayAction.REPLAY_FORWARD:\n            if router_replay is None or router_replay.target_topk_idx is None:\n                # Fallback if replay data is not available\n                return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)\n\n            # Use the provided indices for replay\n            top_indices = router_replay.target_topk_idx\n            # Ensure indices are on the correct device\n            top_indices = top_indices.to(scores.device)\n            # Gather the scores for the replayed indices to get the probabilities\n            probs = scores.gather(1, top_indices)\n            return probs, top_indices\n        elif routing_action == RouterReplayAction.REPLAY_BACKWARD:\n            if router_replay is None or not router_replay.replay_backward_list:\n                # Fallback if replay data is not available\n                return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)\n\n            # Use the last recorded indices for backward replay\n            top_indices = router_replay.replay_backward_list.pop(0)\n            # Ensure indices are on the correct device\n            top_indices = top_indices.to(scores.device)\n            # Gather the scores for the replayed indices to get the probabilities\n            probs = scores.gather(1, top_indices)\n            return probs, top_indices\n        else:  # Unknown action, fallback\n            return _compute_topk(scores, topk, num_groups=num_groups, group_topk=group_topk)\n\n    if score_function == \"softmax\":\n        if use_pre_softmax:\n            scores = torch.softmax(logits, dim=-1, dtype=torch.float32).type_as(logits)\n            probs, top_indices = compute_topk(scores, topk, num_groups, group_topk)\n        else:\n            scores, top_indices = compute_topk(logits, topk, num_groups, group_topk)\n            probs = torch.softmax(scores, dim=-1, dtype=torch.float32).type_as(logits)\n    elif score_function == \"sigmoid\":\n        scores = torch.sigmoid(logits.float()).type_as(logits)\n        if expert_bias is not None:\n            scores_for_routing = scores + expert_bias\n            _, top_indices = compute_topk(scores_for_routing, topk, num_groups, group_topk)\n            scores = torch.gather(scores, dim=1, index=top_indices).type_as(logits)\n        else:\n            scores, top_indices = compute_topk(scores, topk, num_groups, group_topk)\n        probs = scores / (scores.sum(dim=-1, keepdim=True) + 1e-20) if topk > 1 else scores\n    else:\n        raise ValueError(f\"Invalid score_function: {score_function}\")\n\n    if scaling_factor:\n        probs = probs * scaling_factor\n\n    if torch.are_deterministic_algorithms_enabled():\n        # build [num_tokens, num_experts] from [num_tokens, topk]\n        routing_probs = torch.zeros_like(logits)\n        rows = torch.arange(num_tokens, device=logits.device).unsqueeze(1)\n        routing_probs.index_put_((rows, top_indices), probs, accumulate=False)\n\n        routing_map = torch.zeros_like(logits, dtype=logits.dtype)\n        routing_map.index_put_((rows, top_indices), torch.ones_like(probs, dtype=routing_map.dtype), accumulate=False)\n        routing_map = routing_map.bool()\n    else:\n        # TODO Try using element-wise operations instead of scatter?\n        routing_probs = torch.zeros_like(logits).scatter(1, top_indices, probs)\n        routing_map = torch.zeros_like(logits).int().scatter(1, top_indices, 1).bool()\n\n    return routing_probs, routing_map\n\n\ndef _get_aux_loss_coeff(_self, aux_loss_type: str) -> float:\n    \"\"\"Return the aux loss coeff for the given auxiliary loss type.\n    If the auxiliary loss type is not found, return 0.0.\n    \"\"\"\n    if isinstance(_self.routing_type, str):\n        if _self.routing_type == aux_loss_type:\n            return _self.config.moe_aux_loss_coeff\n    if isinstance(_self.routing_type, list):\n        try:\n            idx = _self.routing_type.index(aux_loss_type)\n            return _self.config.moe_aux_loss_coeff[idx]\n        except (ValueError, IndexError):\n            return 0.0\n    return 0.0\n\n\ndef _is_aux_loss_enabled(_self) -> bool:\n    \"\"\"Check if the auxiliary loss is enabled.\"\"\"\n    for aux_loss_type in [\"aux_loss\", \"seq_aux_loss\", \"global_aux_loss\"]:\n        if _get_aux_loss_coeff(_self, aux_loss_type) > 0:\n            return True\n    return False\n\n\ndef patched_routing(self, logits: torch.Tensor, *args, **kwargs):\n    \"\"\"Top-k routing function\n\n    Args:\n        logits (torch.Tensor): Logits tensor after gating.\n\n    Returns:\n        probs (torch.Tensor): The probabilities of token to experts assignment.\n        routing_map (torch.Tensor): The mapping of token to experts assignment,\n            with shape [num_tokens, num_experts].\n    \"\"\"\n    seq_length, bsz = logits.shape[:2]\n    logits = logits.view(-1, self.config.num_moe_experts)\n\n    # Apply Z-Loss\n    logits = self.apply_z_loss(logits)\n\n    # Megatron versions before 0.14.0 do not have 'moe_router_fusion' in TransformerConfig.\n    # We use getattr with a default value of False to ensure compatibility across different\n    # versions of Megatron-LM and MindSpeed.\n    moe_router_fusion = getattr(self.config, \"moe_router_fusion\", False)\n\n    # Calculate probs and routing_map for token dispatching\n    if self.routing_type == \"sinkhorn\":\n        probs, routing_map = self.sinkhorn_load_balancing(logits)\n    else:\n        probs, routing_map = _patched_topk_routing_with_score_function(\n            logits=logits,\n            topk=self.topk,\n            use_pre_softmax=self.config.moe_router_pre_softmax,\n            num_groups=self.config.moe_router_num_groups,\n            group_topk=self.config.moe_router_group_topk,\n            scaling_factor=self.config.moe_router_topk_scaling_factor,\n            score_function=self.score_function,\n            expert_bias=self.expert_bias,\n            fused=moe_router_fusion,\n            router_replay=getattr(self, \"router_replay\", None),\n        )\n\n    # Apply token dropping to probs and routing_map.\n    if self.config.moe_expert_capacity_factor is not None:\n        probs, routing_map = apply_router_token_dropping(\n            probs,\n            routing_map,\n            router_topk=self.topk,\n            capacity_factor=self.config.moe_expert_capacity_factor,\n            drop_policy=self.config.moe_token_drop_policy,\n            pad_to_capacity=self.config.moe_pad_expert_input_to_capacity,\n        )\n\n    if not hasattr(self, \"is_aux_loss_enabled\"):\n        self.is_aux_loss_enabled = types.MethodType(_is_aux_loss_enabled, self)\n    # Apply each aux loss type and attach aux loss autograd function to probs\n    if self.training and torch.is_grad_enabled() and self.is_aux_loss_enabled():\n        # Calculate scores and routing_map for aux loss\n        routing_map_for_aux_loss, scores_for_aux_loss = compute_routing_scores_for_aux_loss(\n            logits, self.topk, self.score_function, fused=self.config.moe_router_fusion\n        )\n        probs = self._apply_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)\n        probs = self._apply_seq_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss, seq_length, bsz)\n        probs = self._apply_global_aux_loss(probs, scores_for_aux_loss, routing_map_for_aux_loss)\n\n    # Update expert bias and tokens_per_expert\n    # Prevent extra local tokens accumulation on evaluation or activation recomputation\n    if self.enable_expert_bias and torch.is_grad_enabled():\n        with torch.no_grad():\n            self.local_tokens_per_expert += routing_map.sum(dim=0)\n\n    return probs, routing_map\n\n\ndef apply_router_replay_patch():\n    \"\"\"\n    Applies the monkey patch for MoE Router Replay functionality.\n    This patch dynamically adds the 'enable_routing_replay' attribute to TransformerConfig\n    and modifies the TopKRouter to support recording and replaying of routing decisions.\n    \"\"\"\n    print(\"Applying Router Replay Patch...\")\n    # Clear router instances to avoid state leakage between model initializations.\n    RouterReplay.router_instances.clear()\n    # Step 1: Patch TransformerConfig to include the feature flag\n\n    try:\n        sig = inspect.signature(TransformerConfig.__init__)\n        native_params = sig.parameters\n        params = list(sig.parameters.values())\n    except Exception:\n        sig = None\n        native_params = {}\n        params = []\n\n    ext_attrs = [\"enable_routing_replay\"]\n\n    # Update __signature__ to prevent NPU/MindSpeed wrappers from filtering out or blocking custom parameters.\n    for attr in ext_attrs:\n        if attr not in native_params:\n            if sig:\n                new_param = inspect.Parameter(attr, inspect.Parameter.KEYWORD_ONLY, default=False)\n                if params and params[-1].kind == inspect.Parameter.VAR_KEYWORD:\n                    params.insert(-1, new_param)\n                else:\n                    params.append(new_param)\n\n    if sig:\n        try:\n            TransformerConfig.__init__.__signature__ = sig.replace(parameters=params)\n        except Exception as e:\n            print(f\"Failed to update signature metadata: {e}\")\n\n    if not hasattr(TransformerConfig, \"_verl_router_patched\"):\n        # Store original __init__ method\n        original_tf_config_init = TransformerConfig.__init__\n\n        # Define new __init__ method that safely handles enable_routing_replay parameter\n        @wraps(original_tf_config_init)\n        def patched_tf_config_init(self, *args, **kwargs):\n            # Simple solution: remove the unknown parameter before calling original constructor\n            enable_routing_replay = kwargs.get(\"enable_routing_replay\", False)\n            if \"enable_routing_replay\" not in native_params:\n                enable_routing_replay = kwargs.pop(\"enable_routing_replay\", False)\n\n            # Call original constructor with remaining kwargs\n            original_tf_config_init(self, *args, **kwargs)\n\n            # Set the instance attribute\n            self.enable_routing_replay = enable_routing_replay\n\n        # Apply the patch\n        TransformerConfig.__init__ = patched_tf_config_init\n        TransformerConfig._verl_router_patched = True\n\n    # Step 2: Patch TopKRouter only once to ensure idempotency.\n    if hasattr(TopKRouter, \"_router_replay_patched\"):\n        return\n\n    original_init = TopKRouter.__init__\n\n    # Step 3: Define the new __init__ method\n    def patched_init(self, *args, **kwargs):\n        original_init(self, *args, **kwargs)\n        self.router_replay = None\n        if getattr(self.config, \"enable_routing_replay\", False):\n            self.router_replay = RouterReplay()\n\n    # Step 4: Patch MoEAlltoAllTokenDispatcher.preprocess to handle router replay\n    # When router replay is enabled, duplicate indices in top_indices can cause\n    # routing_map.sum() < num_tokens * topk, leading to split size mismatch in alltoall.\n    if MoEAlltoAllTokenDispatcher is not None and not hasattr(MoEAlltoAllTokenDispatcher, \"_preprocess_patched\"):\n        original_preprocess = MoEAlltoAllTokenDispatcher.preprocess\n\n        def patched_preprocess(self, routing_map):\n            \"\"\"Patched preprocess that handles router replay correctly for alltoall dispatcher.\"\"\"\n            # Call original preprocess\n            result = original_preprocess(self, routing_map)\n\n            # Fix num_out_tokens when router replay is enabled\n            if (\n                getattr(self.config, \"enable_routing_replay\", False)\n                and not self.drop_and_pad\n                and self.config.moe_expert_capacity_factor is None\n                and not (\n                    getattr(self.config, \"moe_router_padding_for_quantization\", None)\n                    or getattr(self.config, \"moe_router_padding_for_fp8\", None)\n                )\n            ):\n                # With router replay, duplicate indices can reduce the actual routed\n                # token count, so derive it from the routing map instead.\n                self.num_out_tokens = int(routing_map.sum().item())\n\n            return result\n\n        MoEAlltoAllTokenDispatcher.preprocess = patched_preprocess\n        MoEAlltoAllTokenDispatcher._preprocess_patched = True\n\n    # Step 5: Apply the patches\n    TopKRouter.__init__ = patched_init\n    TopKRouter.routing = patched_routing\n    TopKRouter._router_replay_patched = True\n"
  },
  {
    "path": "verl/utils/megatron/router_replay_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nRouter Replay Utilities\nUtilities for handling router replay functionality in Megatron models.\n\"\"\"\n\nimport inspect\nimport warnings\nfrom typing import Optional\n\nimport torch\n\ntry:\n    from megatron.core.pipeline_parallel.utils import is_vp_first_stage, is_vp_last_stage\nexcept ImportError:\n    warnings.warn(\"NPU not support router replay for now.\", stacklevel=2)\n    pass\n\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.pipeline_parallel.schedules import get_schedule_table\nfrom megatron.core.tensor_parallel import gather_from_sequence_parallel_region, scatter_to_sequence_parallel_region\nfrom megatron.core.transformer.transformer_config import TransformerConfig\nfrom megatron.core.transformer.transformer_layer import get_transformer_layer_offset\n\nfrom verl.models.mcore.util import (\n    postprocess_packed_seqs,\n    postprocess_thd_no_padding,\n    preprocess_packed_seqs,\n    preprocess_thd_no_padding,\n)\nfrom verl.utils.device import get_device_name\nfrom verl.utils.megatron.router_replay_patch import RouterReplay, RouterReplayAction\n\ndevice_name = get_device_name()\n\n\n# from megatron.core.transformer.transformer_block import get_num_layers_to_build\ndef get_num_layers_to_build(\n    config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None\n) -> int:\n    \"\"\"\n    Determine the number of transformer layers to build for the current pipeline stage.\n    Args:\n        config (TransformerConfig): Configuration object containing transformer model parameters.\n        vp_stage (Optional[int]): Virtual pipeline stage number.\n        pp_rank (Optional[int]): Pipeline parallel rank.\n\n    Returns:\n        int: The number of layers to be built for the current pipeline stage.\n    \"\"\"\n    # If we have a custom PP layout, straightforwardly\n    # return the number of decoders in the layout array.\n    if hasattr(config, \"pipeline_model_parallel_layout\") and config.pipeline_model_parallel_layout is not None:\n        from megatron.core.transformer.enums import LayerType\n\n        return config.pipeline_model_parallel_layout.get_num_layers_to_build(\n            layer_type=LayerType.decoder, vp_stage=vp_stage\n        )\n\n    # Fallback for legacy tests.\n    if pp_rank is None:\n        pp_rank = mpu.get_pipeline_model_parallel_rank()\n\n    is_first_pp_stage = pp_rank == 0\n    is_last_pp_stage = pp_rank == config.pipeline_model_parallel_size - 1\n\n    if config.num_layers_in_first_pipeline_stage is not None or config.num_layers_in_last_pipeline_stage is not None:\n        assert not (config.account_for_embedding_in_pipeline_split or config.account_for_loss_in_pipeline_split), (\n            \" \\\n        Does not support standalone embedding stage and standalone loss stage with uneven pp\"\n        )\n        # Number of layers to distribute over rest of pipeline stages\n        layers_to_distribute = config.num_layers\n        # Number of pipeline stages left for distributing transformer layers\n        pipeline_stages_left = config.pipeline_model_parallel_size\n\n        # If the uneven first (last) pipeline stage is enabled, remove the specified number\n        # of layers to calculate the number of layers on each middle pipeline stage.\n        if config.num_layers_in_first_pipeline_stage is not None:\n            layers_to_distribute -= config.num_layers_in_first_pipeline_stage\n            pipeline_stages_left -= 1\n\n        if config.num_layers_in_last_pipeline_stage is not None:\n            layers_to_distribute -= config.num_layers_in_last_pipeline_stage\n            pipeline_stages_left -= 1\n\n        # If pp_size <= 2, we do not have any intermediate pipeline stages, and we do not\n        # need to check if the left over layers are divisible by the left over stages.\n        if pipeline_stages_left > 0:\n            assert layers_to_distribute % pipeline_stages_left == 0, (\n                \"With uneven pipelineing the left over layers must be divisible by left over stages\"\n            )\n            num_layers_per_pipeline_rank = layers_to_distribute // pipeline_stages_left\n        else:\n            num_layers_per_pipeline_rank = 0\n\n        # If the uneven first (last) pipeline stage is enabled, return the specified number\n        # of layers for all virtual pipeline parallel stages within the first (last) pipeline\n        # parallel stage.\n\n        if is_first_pp_stage and config.num_layers_in_first_pipeline_stage is not None:\n            num_layers_per_pipeline_rank = config.num_layers_in_first_pipeline_stage\n\n        if is_last_pp_stage and config.num_layers_in_last_pipeline_stage is not None:\n            num_layers_per_pipeline_rank = config.num_layers_in_last_pipeline_stage\n    else:\n        # Include the embedding layer and loss layer into pipeline parallelism partition\n        num_layers = config.num_layers\n        if config.account_for_embedding_in_pipeline_split:\n            num_layers += 1\n\n        if config.account_for_loss_in_pipeline_split:\n            num_layers += 1\n\n        assert num_layers % config.pipeline_model_parallel_size == 0, (\n            \"num_layers should be divisible by pipeline_model_parallel_size\"\n        )\n        num_layers_per_pipeline_rank = num_layers // config.pipeline_model_parallel_size\n\n    vp_size = config.virtual_pipeline_model_parallel_size\n    if vp_size is not None and config.pipeline_model_parallel_size > 1:\n        # Interleaved pipeline parallelism:\n        # Number of layers in each model chunk is the number of layers in the stage,\n        # divided by the number of model chunks in a stage.\n        # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of\n        # layers to stages like (each list is a model chunk):\n        # Stage 0: [0]  [2]  [4]  [6]\n        # Stage 1: [1]  [3]  [5]  [7]\n        # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of\n        # layers to stages like (each list is a model chunk):\n        # Stage 0: [0, 1]  [4, 5]\n        # Stage 1: [2, 3]  [6, 7]\n\n        assert num_layers_per_pipeline_rank % vp_size == 0, (\n            f\"num_layers_per_pipeline_rank {num_layers_per_pipeline_rank} \\\n            should be divisible by vp_size {vp_size}\"\n        )\n        num_layers_per_virtual_stage = num_layers_per_pipeline_rank // vp_size\n\n        num_layers_to_build = num_layers_per_virtual_stage\n\n    else:\n        # Non-interleaved pipeline parallelism:\n        # Each stage gets a contiguous set of layers.\n        num_layers_to_build = num_layers_per_pipeline_rank\n\n    # The embedding (or loss) layer cannot function as a standalone transformer layer\n    # Reduce the number of layers to construct by 1 on the first (or last) stage if the\n    # embedding (or loss) layer is included in the pipeline parallelism partition and placement.\n    if config.account_for_embedding_in_pipeline_split:\n        if is_vp_first_stage(vp_stage, vp_size) and is_first_pp_stage:\n            num_layers_to_build -= 1\n            assert num_layers_to_build >= 0, \"Not enough layers in the first virtual pipeline stage\"\n\n    if config.account_for_loss_in_pipeline_split:\n        if is_vp_last_stage(vp_stage, vp_size) and is_last_pp_stage:\n            num_layers_to_build -= 1\n            assert num_layers_to_build >= 0, \"Not enough layers in the last virtual pipeline stage\"\n\n    return num_layers_to_build\n\n\ndef is_moe_layer(tf_config, layer_idx):\n    moe_layer_freq = getattr(tf_config, \"moe_layer_freq\", None)\n\n    if isinstance(moe_layer_freq, int):\n        return layer_idx % moe_layer_freq == 0\n    elif isinstance(moe_layer_freq, list):\n        return moe_layer_freq[layer_idx] == 1\n    else:\n        raise ValueError(f\"Unsupported moe_layer_freq type: {type(moe_layer_freq)}\")\n\n\ndef get_moe_num_layers_to_build(\n    config: TransformerConfig, vp_stage: Optional[int] = None, pp_rank: Optional[int] = None\n) -> int:\n    \"\"\"Count the number of MoE layers assigned to the current rank.\n    When ``moe_layer_freq`` is 1 or unset, every transformer layer is an MoE\n    layer, so the count equals the total layer count. Otherwise only layers\n    whose global index satisfies the frequency predicate are counted.\n    Args:\n        config: Megatron TransformerConfig providing layer layout information.\n        vp_stage: Virtual-pipeline stage index (None defaults to current).\n        pp_rank: Pipeline-parallel rank (None defaults to current).\n    Returns:\n        Number of MoE layers on the specified rank/stage.\n    \"\"\"\n    total_layers = get_num_layers_to_build(config, vp_stage=vp_stage, pp_rank=pp_rank)\n\n    sig = inspect.signature(get_transformer_layer_offset)\n    # core 0.12.1 is not support vp_stage and pp_rank as parameters\n    if \"vp_stage\" in sig.parameters and \"pp_rank\" in sig.parameters:\n        layer_offset = get_transformer_layer_offset(config, vp_stage=vp_stage, pp_rank=pp_rank)\n    elif \"pp_rank\" in sig.parameters:\n        layer_offset = get_transformer_layer_offset(config, pp_rank=pp_rank)\n    else:\n        layer_offset = get_transformer_layer_offset(config)\n\n    local_global_indices = range(layer_offset, layer_offset + total_layers)\n\n    num_moe_layers = sum(1 for idx in local_global_indices if is_moe_layer(config, idx))\n\n    return num_moe_layers\n\n\ndef merge_router_topk_indices(attention_mask, input_ids, mini_layer_topk_idx_list, tf_config, vp_rank=None):\n    \"\"\"\n    Merge recorded router top-k indices across sequence-parallel ranks for all router instances,\n    then pack/unpack them to align with the original (batch, seq_len) layout and append the result.\n\n    Args:\n        attention_mask (torch.Tensor): Attention mask of shape [batch_size, seq_len]. Used to determine\n            the valid token positions during pack/unpack.\n        input_ids (torch.Tensor): Input token IDs of shape [batch_size, seq_len]. Used together with\n            attention_mask for sequence packing/unpacking.\n        mini_layer_topk_idx_list (list): A Python list to which the merged top-k indices tensor will be appended.\n        tf_config: Megatron/Transformer engine configuration object. Used to locate router instances for\n            the current micro-batch.\n        vp_rank (Optional[int]): Virtual pipeline stage rank override. If None, the current VP rank from\n            Megatron parallel state will be used.\n\n    Returns:\n        None: The function has side effects only; it appends a tensor of shape\n        [1, dynamic_bs_all, layer_num, topk] to mini_layer_topk_idx_list.\n    \"\"\"\n    with torch.no_grad():\n        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)\n        layers_topk_idx = []\n        for router in router_instances_list:\n            layers_topk_idx.append(router.recorded_topk_idx.to(torch.uint8))  # dynamic_bs, topk\n\n        # layer_num, dynamic_bs, topk  -> dynamic_bs, layer_num, topk\n        layers_topk_idx = torch.stack(layers_topk_idx).permute(1, 0, 2).to(device_name)\n        # dynamic_bs, layer_num, topk -> 1, dynamic_bs_all, layer_num, topk\n        layers_topk_idx = (\n            gather_from_sequence_parallel_region(layers_topk_idx, tensor_parallel_output_grad=False)\n            .unsqueeze(0)\n            .contiguous()\n        )\n\n        if input_ids.is_nested:\n            batch_size = input_ids.shape[0]\n            _, packed_seq_params = preprocess_thd_no_padding(input_ids, pre_process=True)\n            layers_topk_idx = postprocess_thd_no_padding(\n                layers_topk_idx, packed_seq_params, input_ids, batch_size, post_process=True\n            )\n        else:\n            batch_size, seq_len = attention_mask.shape[:2]\n            _, packed_seq_params = preprocess_packed_seqs(input_ids, attention_mask, pre_process=True)\n            layers_topk_idx = postprocess_packed_seqs(\n                layers_topk_idx, packed_seq_params, attention_mask, batch_size, seq_len, post_process=True\n            )\n        mini_layer_topk_idx_list.append(layers_topk_idx.cpu())\n\n\ndef set_router_replay_data(layers_topk_idx, attention_mask, tf_config, vp_rank=None):\n    \"\"\"\n    Scatter the packed router top-k indices back to sequence-parallel ranks and update each local\n    RouterReplay instance with target indices for replay mode.\n\n    This function prepares the per-layer, per-sample top-k routing decisions (recorded during an earlier\n    forward) so that subsequent replay passes can follow exactly the same routing.\n\n    Args:\n        layers_topk_idx (torch.Tensor): Router top-k indices with shape [bs, max_seq_len, layer_num, topk].\n            This should be the merged output produced by merge_router_topk_indices.\n        attention_mask (torch.Tensor): Attention mask [batch_size, seq_len] used for pack/unpack alignment.\n        tf_config: Megatron/Transformer engine configuration object.\n        vp_rank (Optional[int]): Virtual pipeline stage rank override. If None, the current VP rank from\n            Megatron parallel state will be used.\n\n    Returns:\n        None: The function updates internal RouterReplay instances in-place.\n    \"\"\"\n    with torch.no_grad():\n        if layers_topk_idx.is_nested:\n            layers_topk_idx_rmpad, _, _ = preprocess_thd_no_padding(layers_topk_idx, pre_process=True)\n        else:\n            layers_topk_idx_rmpad, _ = preprocess_packed_seqs(layers_topk_idx, attention_mask, pre_process=True)\n        layers_topk_idx_rmpad = layers_topk_idx_rmpad.contiguous()  # 1, dynamic_bs_all, layer_num, topk\n\n        # 1, dynamic_bs_split, layer_num, topk\n        layers_topk_idx_rmpad_split = scatter_to_sequence_parallel_region(\n            layers_topk_idx_rmpad.to(device_name).squeeze(dim=0)\n        ).unsqueeze(dim=0)\n\n        # dynamic_bs_split, layer_num, topk -> layer_num, dynamic_bs_split, topk\n        layers_topk_idx_reshape = layers_topk_idx_rmpad_split.permute(0, 2, 1, 3).squeeze(\n            dim=0\n        )  # layer_num, dynamic_bs_all, topk\n        local_rank_info = get_current_rank_layer_info(tf_config, vp_rank)\n        offset, end = local_rank_info[\"start\"], local_rank_info[\"end\"]\n        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)\n\n        # When dim-0 covers all layers (e.g. R3, or R2 with all-MoE models),\n        # index by absolute layer_idx; otherwise (R2 with mixed dense/MoE),\n        # dim-0 only contains MoE layers, index by MoE-layer ordinal.\n        index_by_layer = len(layers_topk_idx_reshape) == tf_config.num_layers\n\n        # For R2: count MoE layers before `offset` as the starting position.\n        moe_idx = sum(1 for i in range(offset) if is_moe_layer(tf_config, i))\n\n        router_offset = 0\n        for layer_idx in range(offset, end):\n            if not is_moe_layer(tf_config, layer_idx):\n                continue\n            router = router_instances_list[router_offset]\n            idx = layer_idx if index_by_layer else moe_idx\n            router.set_target_indices(layers_topk_idx_reshape[idx].to(torch.int64))\n            router_offset += 1\n            moe_idx += 1\n\n\ndef reorder_and_merge_vpp_layers(\n    micro_batch_tensor_list,\n    num_microbatches: int,\n    vpp_size: int,\n    microbatch_group_size_per_vp_stage: int,\n) -> torch.Tensor:\n    \"\"\"\n    Reorder and merge per-VPP layer blocks into a contiguous layer dimension.\n\n    Given a tensor shaped as [bs*vpp_size, max_token_len, layer_num_per_vpp, topk], this function:\n    1) Builds the schedule table for virtual microbatches and reorders the first dimension so that entries\n       belonging to the same model chunk (VPP stage) become contiguous.\n    2) Reshapes and merges the (vpp_size, layer_num_per_vpp) into a single layer dimension, producing\n       [bs, max_token_len, layer_num, topk].\n\n    Args:\n        micro_batch_tensor_list : the list of Input tensor.\n        num_microbatches (int): Number of microbatches per pipeline stage (bs).\n        vpp_size (int): Virtual pipeline parallel size (number of model chunks).\n        microbatch_group_size_per_vp_stage (int): Number of consecutive microbatches processed per VPP stage.\n\n    Returns:\n        torch.Tensor: Output tensor of shape [bs, max_token_len, layer_num, topk].\n\n    Raises:\n        ValueError: If input tensor dimensionality or expected sizes do not match.\n        RuntimeError: If the computed output shape is unexpected or the schedule length mismatches.\n    \"\"\"\n    # 1) Build schedule table: map each virtual_microbatch_id -> (microbatch_id, model_chunk_id)\n    schedule_table = get_schedule_table(num_microbatches, vpp_size, microbatch_group_size_per_vp_stage)\n\n    # 2) Group by model_chunk_id to build reorder indices so entries of the same chunk become contiguous along dim 0\n    tensor_by_chunk = [[] for _ in range(vpp_size)]\n    mini_tensor_list = []\n\n    for vidx, (_mb, chunk_id) in enumerate(schedule_table):\n        tensor_by_chunk[chunk_id].append(micro_batch_tensor_list[vidx])\n\n    if micro_batch_tensor_list[0].is_nested:\n        for chunk_id in range(vpp_size):\n            tensors = [tensor for nt in tensor_by_chunk[chunk_id] for tensor in nt.unbind()]\n            mini_tensor_list.append(torch.nested.as_nested_tensor(tensors, layout=torch.jagged))\n    else:\n        for chunk_id in range(vpp_size):\n            mini_tensor_list.append(torch.cat(tensor_by_chunk[chunk_id], dim=0))\n\n    out = torch.cat(mini_tensor_list, dim=2)\n\n    return out\n\n\ndef get_current_rank_layer_info(tf_config, vp_rank=None):\n    # When vp_rank is None, default to the current VP rank (or 0 if VP is disabled).\n    \"\"\"Return the local layer range/count for the current process and the full assignment table.\n\n    Args:\n        tf_config: Configuration object used by compute_pipeline_layer_assignment.\n        vp_rank (Optional[int]): Explicit virtual pipeline stage rank to query. If None, uses\n            mpu.get_virtual_pipeline_model_parallel_rank() when VP is enabled; otherwise 0.\n\n    Returns:\n        Tuple[dict, dict]: A tuple of (local_assignment, all_assignments) where local_assignment contains\n        keys {\"start\", \"end\", \"count\"} for the current (pp_rank, vp_stage).\n    \"\"\"\n    if vp_rank is None:\n        vp_rank = 0\n    num_layers_to_build = get_num_layers_to_build(tf_config, vp_stage=vp_rank)\n\n    sig = inspect.signature(get_transformer_layer_offset)\n\n    if \"vp_stage\" in sig.parameters:\n        offset = get_transformer_layer_offset(tf_config, vp_stage=vp_rank)\n    else:\n        offset = get_transformer_layer_offset(tf_config)\n    local = {}\n    local[\"start\"] = offset\n    local[\"end\"] = offset + num_layers_to_build\n    local[\"count\"] = num_layers_to_build\n    return local\n\n\ndef pp_gather(local_layers_router_map, tf_config):\n    # TODO: Consider non-uniform layer allocation cases.\n    \"\"\"\n    Gather local router maps from all PP ranks into a global router map.\n\n    Args:\n        local_layers_router_map (torch.Tensor): Local router map of shape\n            [bs, max_seq_len, local_num_layers, topk].\n        tf_config: Configuration providing pipeline_model_parallel_size.\n\n    Returns:\n        torch.Tensor: Global router map of shape [bs, max_seq_len, num_layers, topk] placed on CPU.\n    \"\"\"\n    pp_size = tf_config.pipeline_model_parallel_size\n    if pp_size <= 1:\n        return local_layers_router_map\n\n    pp_group = mpu.get_pipeline_model_parallel_group()\n    world_size = torch.distributed.get_world_size(pp_group)\n    local_layers_router_map = local_layers_router_map.to(device_name)\n    if local_layers_router_map.is_nested:\n        layers_topk_idx_global_list = [None] * world_size\n        torch.distributed.all_gather_object(layers_topk_idx_global_list, local_layers_router_map, pp_group)\n    else:\n        layers_topk_idx_global_list = [\n            torch.empty(\n                size=local_layers_router_map.shape,\n                dtype=local_layers_router_map.dtype,\n                device=local_layers_router_map.device,\n            )\n            for _ in range(world_size)\n        ]\n        torch.distributed.all_gather(\n            tensor=local_layers_router_map,\n            tensor_list=layers_topk_idx_global_list,\n            group=pp_group,\n            async_op=False,\n        )\n    vp_size = tf_config.virtual_pipeline_model_parallel_size\n    if vp_size is not None:\n        vpp_router_map_offset = [[] for _ in range(pp_size)]\n        for pp_stage in range(pp_size):\n            vpp_router_map_offset[pp_stage].append(0)\n            for vp_stage in range(vp_size):\n                num_layers_to_build = get_moe_num_layers_to_build(tf_config, vp_stage, pp_stage)\n                vpp_router_map_offset[pp_stage].append(num_layers_to_build + vpp_router_map_offset[pp_stage][-1])\n        layers_topk_idx_global = []\n        for vp_stage in range(vp_size):\n            for pp_stage in range(pp_size):\n                piece = slice(vpp_router_map_offset[pp_stage][vp_stage], vpp_router_map_offset[pp_stage][vp_stage + 1])\n                if layers_topk_idx_global_list[pp_stage].is_nested:\n                    nested_item = layers_topk_idx_global_list[pp_stage]\n                    sliced_tensors = [t[:, piece, :] for t in nested_item.unbind()]\n                    sliced_nested = torch.nested.as_nested_tensor(sliced_tensors, layout=torch.jagged)\n                    layers_topk_idx_global.append(sliced_nested)\n                else:\n                    layers_topk_idx_global.append(layers_topk_idx_global_list[pp_stage][:, :, piece, :])\n        global_router_map = torch.cat(layers_topk_idx_global, dim=2).to(\"cpu\")\n    else:\n        global_router_map = torch.cat(layers_topk_idx_global_list, dim=2).to(\"cpu\")\n\n    return global_router_map\n\n\nclass RouterReplayHelper:\n    \"\"\"Helper class to query router replay state and locate local RouterReplay instances.\"\"\"\n\n    @staticmethod\n    def get_micro_batch_router_list(tf_config, vp_rank=None):\n        \"\"\"\n        Return the list of RouterReplay instances corresponding to the current micro-batch and local\n        (pp_rank, vp_stage) layer range.\n\n        When virtual pipeline (VPP) is enabled, the local range for the PP rank is expanded to include\n        all VP stages by multiplying the per-VP count by vp_size. The returned slice is taken from the\n        global RouterReplay.router_instances list.\n\n        Args:\n            tf_config: Configuration object used to compute layer assignments.\n            vp_rank (Optional[int]): Explicit virtual pipeline stage to query. If None, the current VP\n                rank from Megatron parallel state is used when available.\n        Returns:\n            list: A contiguous sublist of RouterReplay.router_instances for the local layer range.\n        \"\"\"\n        vp_size = tf_config.virtual_pipeline_model_parallel_size\n        if vp_size is not None:\n            vp_rank = 0 if vp_rank is None else vp_rank\n            offset = 0\n            for pre_vp_stage in range(vp_size):\n                if pre_vp_stage == vp_rank:\n                    break\n                offset += get_moe_num_layers_to_build(tf_config, pre_vp_stage)\n        else:\n            offset = 0\n\n        num_layers_to_build = get_moe_num_layers_to_build(tf_config, vp_rank)\n        router_instances_list = RouterReplay.router_instances[offset : offset + num_layers_to_build]\n        return router_instances_list\n\n    @staticmethod\n    def is_r2_record_action(tf_config, vp_rank=None) -> bool:\n        \"\"\"Return True if the current router_replay_action is RECORD (R2) for the local router instances.\n\n        This inspects the first local RouterReplay instance's router_replay_action and compares it to\n        RouterReplayAction.RECORD.\n        \"\"\"\n        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)\n        return router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.RECORD\n\n    @staticmethod\n    def is_replay_forward_action(tf_config, vp_rank=None) -> bool:\n        \"\"\"Return True if the current router_replay_action is REPLAY_FORWARD for the local router instances.\n\n        This inspects the first local RouterReplay instance's router_replay_action and compares it to\n        RouterReplayAction.REPLAY_FORWARD.\n        \"\"\"\n        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)\n        return (\n            router_instances_list and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_FORWARD\n        )\n\n    @staticmethod\n    def is_replay_backward_action(tf_config, vp_rank=None) -> bool:\n        \"\"\"Return True if the current router_replay_action is REPLAY_BACKWARD for the local router instances.\n\n        This inspects the first local RouterReplay instance's router_replay_action and compares it to\n        RouterReplayAction.REPLAY_BACKWARD.\n        \"\"\"\n        router_instances_list = RouterReplayHelper.get_micro_batch_router_list(tf_config, vp_rank)\n        return (\n            router_instances_list\n            and router_instances_list[0].router_replay_action == RouterReplayAction.REPLAY_BACKWARD\n        )\n"
  },
  {
    "path": "verl/utils/megatron/sequence_parallel.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nimport torch.nn.functional as F\nfrom megatron.core import parallel_state as mpu\n\n\ndef mark_parameter_as_sequence_parallel(parameter):\n    parameter.sequence_parallel = True\n\n\ndef is_sequence_parallel_param(param):\n    return hasattr(param, \"sequence_parallel\") and param.sequence_parallel\n\n\ndef pad_to_sequence_parallel(unpad_tokens: torch.Tensor):\n    \"\"\"pad the tokens such that the total length is a multiple of sp world size\n\n    Args:\n        unpad_tokens: (total_nnz, ...). Tokens after removing padding\n\n    Returns:\n        the padded tokens: (total_nnz + pad_size,...)\n\n    \"\"\"\n    total_nnz = unpad_tokens.shape[0]\n    sp_world_size = mpu.get_tensor_model_parallel_world_size()\n\n    pad_size = 0 if total_nnz % sp_world_size == 0 else sp_world_size - total_nnz % sp_world_size\n\n    if pad_size > 0:\n        if unpad_tokens.ndim == 1:\n            unpad_tokens = F.pad(unpad_tokens, (0, pad_size))\n        elif unpad_tokens.ndim == 2:\n            unpad_tokens = F.pad(unpad_tokens, (0, 0, 0, pad_size))\n        else:\n            raise NotImplementedError(f\"Padding dim {unpad_tokens.ndim()} is not supported\")\n\n    return unpad_tokens\n"
  },
  {
    "path": "verl/utils/megatron/tensor_parallel.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUtilities for using tensor_parallel in megatron\n\"\"\"\n\nfrom typing import TYPE_CHECKING\n\nimport torch\nimport torch.distributed as dist\nfrom megatron.core import parallel_state as mpu\nfrom torch.nn import init\n\nif TYPE_CHECKING:\n    from megatron.core import ModelParallelConfig\n\n\ndef update_kwargs_with_config(dictionary: dict, config: \"ModelParallelConfig\"):\n    dictionary[\"config\"] = config\n    return dictionary\n\n\ndef get_default_kwargs_for_model_parallel_config():\n    model_parallel_config_kwargs = {\n        \"params_dtype\": torch.float32,\n        \"use_cpu_initialization\": False,\n        \"perform_initialization\": True,\n        \"gradient_accumulation_fusion\": False,\n        \"sequence_parallel\": False,\n    }\n    return model_parallel_config_kwargs\n\n\ndef get_default_model_parallel_config():\n    from megatron.core import ModelParallelConfig\n\n    return ModelParallelConfig(**get_default_kwargs_for_model_parallel_config())\n\n\ndef get_common_default_kwargs_for_parallel_linear():\n    default_model_parallel_config = get_default_model_parallel_config()\n    common_default_kwargs = {\n        \"init_method\": init.xavier_normal_,\n        \"stride\": 1,\n        \"keep_master_weight_for_test\": False,\n        \"config\": default_model_parallel_config,\n    }\n    return common_default_kwargs\n\n\ndef get_default_kwargs_for_column_parallel_linear():\n    from megatron.core import ModelParallelConfig\n\n    model_parallel_config_kwargs = get_default_kwargs_for_model_parallel_config()\n    column_parallel_config_kwargs = {\n        \"async_tensor_model_parallel_allreduce\": False,\n    }\n    model_parallel_config_kwargs.update(column_parallel_config_kwargs)\n    column_default_kwargs = {\n        \"config\": ModelParallelConfig(**model_parallel_config_kwargs),\n    }\n    common_default_kwargs = get_common_default_kwargs_for_parallel_linear()\n    common_default_kwargs.update(column_default_kwargs)\n    return common_default_kwargs\n\n\ndef get_default_kwargs_for_row_parallel_linear():\n    common_default_kwargs = get_common_default_kwargs_for_parallel_linear()\n    return common_default_kwargs\n\n\ndef get_default_kwargs_for_parallel_embedding():\n    from megatron.core import ModelParallelConfig\n\n    model_parallel_config_kwargs = get_default_kwargs_for_model_parallel_config()\n    embedding_default_kwargs = {\n        \"init_method\": init.xavier_normal_,\n        \"config\": ModelParallelConfig(**model_parallel_config_kwargs),\n    }\n    return embedding_default_kwargs\n\n\ndef is_tensor_parallel_param(param):\n    return hasattr(param, \"tensor_model_parallel\") and param.tensor_model_parallel\n\n\ndef get_tensor_parallel_partition_dim(param):\n    assert is_tensor_parallel_param(param)\n    return param.partition_dim\n\n\ndef get_tensor_parallel_partition_stride(param):\n    assert is_tensor_parallel_param(param)\n    return param.partition_stride\n\n\nclass _VocabParallelEntropy(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, vocab_parallel_logits: torch.Tensor) -> torch.Tensor:\n        @torch.compile(dynamic=True)\n        def mul_reduce(a, b):\n            return (a * b).sum(dim=-1, keepdim=True)\n\n        logits_max = vocab_parallel_logits.max(dim=-1, keepdim=True).values\n        dist.all_reduce(logits_max, op=dist.ReduceOp.MAX, group=mpu.get_tensor_model_parallel_group())\n        normalized_vocab_parallel_logits = vocab_parallel_logits - logits_max\n        normalized_exp_logits = normalized_vocab_parallel_logits.exp_()\n        normalized_sum_exp_logits = normalized_exp_logits.sum(dim=-1, keepdim=True)\n        dist.all_reduce(normalized_sum_exp_logits, group=mpu.get_tensor_model_parallel_group())\n        softmax_logits = normalized_exp_logits.div_(normalized_sum_exp_logits)\n        sum_softmax_times_logits = mul_reduce(softmax_logits, vocab_parallel_logits)\n        dist.all_reduce(sum_softmax_times_logits, group=mpu.get_tensor_model_parallel_group())\n        entropy = logits_max + normalized_sum_exp_logits.log() - sum_softmax_times_logits\n        ctx.save_for_backward(vocab_parallel_logits, softmax_logits, sum_softmax_times_logits)\n        return entropy.squeeze(dim=-1)\n\n    @staticmethod\n    def backward(ctx, grad_output: torch.Tensor) -> torch.Tensor:\n        vocab_parallel_logits, softmax_logits, sum_softmax_times_logits = ctx.saved_tensors\n        # reuse softmax_logits as grad\n        vocab_parallel_logits.sub_(sum_softmax_times_logits)\n        softmax_logits.mul_(vocab_parallel_logits)\n        softmax_logits.mul_(grad_output.unsqueeze(dim=-1))\n        # recover vocab_parallel_logits\n        vocab_parallel_logits.add_(sum_softmax_times_logits)\n        softmax_logits.mul_(-1)\n        return softmax_logits\n\n\ndef vocab_parallel_entropy(vocab_parallel_logits: torch.Tensor) -> torch.Tensor:\n    \"\"\"Compute entropy when the logits are sharded in tp ranks\n\n    Args:\n        vocab_parallel_logits: (total_nnz, vocab_size // tp_size)\n\n    Returns: (total_nnz,)\n\n    \"\"\"\n    return _VocabParallelEntropy.apply(vocab_parallel_logits)\n\n\ndef vocab_parallel_log_probs_from_logits(logits, labels):\n    \"\"\"TODO(zhangchi.usc1992): We may change the implementation later\"\"\"\n    from megatron.core import tensor_parallel\n\n    return -tensor_parallel.vocab_parallel_cross_entropy(vocab_parallel_logits=logits, target=labels)\n\n\ndef vocab_parallel_log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length):\n    \"\"\"Similar to log_probs_from_logits_response_rmpad, but the logits_rmpad is now spliited across tensor parallel\n    region.\n    This will further reduce the peak memory usage during training\n\n    Args:\n        input_ids: [batch_size, seqlen]\n        attention_mask: [batch_size, seqlen]\n        logits_rmpad: [total_nnz, vocab_size // tp_size]\n        response_length: int\n\n    \"\"\"\n    from flash_attn.bert_padding import pad_input, unpad_input\n\n    batch_size, seqlen = input_ids.shape\n    input_ids_rmpad, indices, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask=attention_mask)\n    input_ids_rmpad = input_ids_rmpad.squeeze(-1)\n    input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=0)\n    full_log_probs_rmpad = vocab_parallel_log_probs_from_logits(\n        logits=logits_rmpad, labels=input_ids_rmpad_rolled\n    )  # (total_nnz,)\n    full_output = pad_input(\n        hidden_states=full_log_probs_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen\n    )\n    output = full_output.squeeze(-1)[:, -response_length - 1 : -1]  # [batch_size, response_length]\n    return output\n"
  },
  {
    "path": "verl/utils/megatron_peft_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Utilities for PEFT (Parameter-Efficient Fine-Tuning) of Megatron in VERL.\"\"\"\n\nimport os\nfrom pathlib import Path\nfrom typing import Iterator\n\nimport torch\n\n# Map megatron lora target modules to HF-style module names for vLLM\nMEGATRON_TO_HF_MODULES = {\n    \"linear_qkv\": [\"q_proj\", \"k_proj\", \"v_proj\"],\n    \"linear_proj\": [\"o_proj\"],\n    \"linear_fc1\": [\"gate_proj\", \"up_proj\"],\n    \"linear_fc2\": [\"down_proj\"],\n    \"router\": [\"gate\"],\n    # Canonical LoRA mappings\n    \"linear_q\": [\"q_proj\"],\n    \"linear_k\": [\"k_proj\"],\n    \"linear_v\": [\"v_proj\"],\n    \"linear_fc1_up\": [\"up_proj\"],\n    \"linear_fc1_gate\": [\"gate_proj\"],\n    # MLA mappings\n    \"linear_kv_down_proj\": [\"kv_a_proj_with_mqa\"],\n    \"linear_kv_up_proj\": [\"kv_b_proj\"],\n    \"linear_q_down_proj\": [\"q_a_proj\"],\n    \"linear_q_up_proj\": [\"q_b_proj\"],\n    \"linear_q_proj\": [\"q_proj\"],\n    # DSA indexer mappings\n    \"linear_wq_b\": [\"wq_b\"],\n    \"linear_wk\": [\"wk\"],\n    \"linear_weights_proj\": [\"weights_proj\"],\n}\n\n# Modules with stacked parameters that need .base_layer suffix in vLLM\nSTACKED_PARAMS = [\n    \".q_proj.weight\",\n    \".q_proj.bias\",\n    \".k_proj.weight\",\n    \".k_proj.bias\",\n    \".v_proj.weight\",\n    \".v_proj.bias\",\n    \".o_proj.weight\",\n    \".o_proj.bias\",\n    \".gate_proj.weight\",\n    \".up_proj.weight\",\n    \".down_proj.weight\",\n    \".mlp.gate.weight\",\n    \".mlp.gate.bias\",\n    \".mlp.gate.e_score_correction_bias\",\n    \".kv_a_proj_with_mqa.weight\",\n    \".kv_b_proj.weight\",\n    \".q_a_proj.weight\",\n    \".q_b_proj.weight\",\n    \".wq_b.weight\",\n    \".wk.weight\",\n    \".weights_proj.weight\",\n]\n\n\ndef _get_rank_checkpoint_path(base_path: str) -> str:\n    \"\"\"Get rank-specific checkpoint path following Megatron's convention.\n\n    Returns path like: base_path/mp_rank_{tp:02d}_{pp:03d}_{ep:03d}/\n\n    Args:\n        base_path: Base checkpoint directory\n\n    Returns:\n        Rank-specific subdirectory path\n    \"\"\"\n    from megatron.core import mpu\n\n    tensor_rank = mpu.get_tensor_model_parallel_rank()\n    pipeline_rank = mpu.get_pipeline_model_parallel_rank()\n    expert_rank = mpu.get_expert_model_parallel_rank()\n\n    pipeline_parallel = mpu.get_pipeline_model_parallel_world_size() > 1\n    expert_parallel = mpu.get_expert_model_parallel_world_size() > 1\n\n    if not pipeline_parallel:\n        rank_path = os.path.join(base_path, f\"mp_rank_{tensor_rank:02d}\")\n    else:\n        rank_path = os.path.join(base_path, f\"mp_rank_{tensor_rank:02d}_{pipeline_rank:03d}\")\n\n    if expert_parallel:\n        rank_path = rank_path + f\"_{expert_rank:03d}\"\n\n    return rank_path\n\n\ndef get_adapter_state_dict(model):\n    \"\"\"Extract only adapter parameters from a model.\n\n    Args:\n        model: PyTorch model (possibly wrapped in DDP/Float16Module)\n\n    Returns:\n        Dict of adapter parameter names to tensors\n    \"\"\"\n    from verl.utils.megatron_utils import unwrap_model\n\n    # Unwrap model from DDP/Float16Module\n    unwrapped = unwrap_model(model)\n    if isinstance(unwrapped, list):\n        unwrapped = unwrapped[0]\n\n    adapter_state = {}\n    for name, param in unwrapped.named_parameters():\n        if \".adapter.\" in name.lower():\n            adapter_state[name] = param.data.clone()\n\n    return adapter_state\n\n\ndef save_adapter_checkpoint(\n    model: torch.nn.Module | list[torch.nn.Module],\n    checkpoint_path: str,\n    rank: int = 0,\n):\n    \"\"\"Save only adapter parameters to checkpoint.\n\n    This is much more efficient than saving the full model when using PEFT,\n    as adapters typically represent <1% of total parameters.\n\n    Uses Megatron's distributed checkpoint structure: each rank saves to\n    checkpoint_path/mp_rank_{tp:02d}_{pp:03d}/adapter.pt\n\n    Args:\n        model: Model or list of models\n        checkpoint_path: Base path to save checkpoint (rank-specific subdirs created)\n        rank: Process rank (used for logging only)\n    \"\"\"\n\n    if isinstance(model, list):\n        models = model\n    else:\n        models = [model]\n\n    # Get adapter state from first model\n    adapter_state = get_adapter_state_dict(models[0])\n\n    if not adapter_state:\n        if rank == 0:\n            print(\"Warning: No adapter parameters found to save\")\n        return\n\n    # Get rank-specific directory path\n    Path(checkpoint_path).mkdir(parents=True, exist_ok=True)\n    rank_path = _get_rank_checkpoint_path(checkpoint_path)\n    adapter_file = rank_path + \"_adapter.pt\"\n\n    torch.save(\n        {\n            \"adapter_state_dict\": adapter_state,\n        },\n        adapter_file,\n    )\n\n    if rank == 0:\n        print(f\"Saved {len(adapter_state)} adapter parameters to {checkpoint_path} (distributed)\")\n\n\ndef load_adapter_checkpoint(\n    model: torch.nn.Module | list[torch.nn.Module],\n    checkpoint_path: str,\n    strict: bool = True,\n):\n    \"\"\"Load adapter parameters from checkpoint.\n\n    Loads from Megatron's distributed checkpoint structure: reads from\n    checkpoint_path/mp_rank_{tp:02d}_{pp:03d}/adapter.pt for each rank.\n\n    Args:\n        model: Model or list of models\n        checkpoint_path: Base path to checkpoint directory\n        strict: Whether to strictly enforce parameter name matching\n    \"\"\"\n    from megatron.core import mpu\n\n    from verl.utils.megatron_utils import unwrap_model\n\n    # Get rank-specific path\n    rank_path = _get_rank_checkpoint_path(checkpoint_path)\n    adapter_file = rank_path + \"_adapter.pt\"\n\n    if not os.path.isfile(adapter_file):\n        raise FileNotFoundError(f\"Adapter checkpoint not found: {adapter_file}\")\n\n    checkpoint = torch.load(adapter_file, map_location=\"cpu\")\n    adapter_state = checkpoint.get(\"adapter_state_dict\", {})\n\n    if not adapter_state:\n        print(\"Warning: No adapter parameters found in checkpoint\")\n        return\n\n    if isinstance(model, list):\n        models = model\n    else:\n        models = [model]\n\n    # Load adapter parameters into each model (for VPP, models may have multiple chunks)\n    loaded_count = 0\n    for m in models:\n        unwrapped = unwrap_model(m)\n        if isinstance(unwrapped, list):\n            unwrapped = unwrapped[0]\n\n        # Load parameters\n        _, unexpected = unwrapped.load_state_dict(adapter_state, strict=False)\n\n        if strict and unexpected:\n            raise RuntimeError(f\"Error loading adapter checkpoint:\\nUnexpected keys: {unexpected}\")\n\n        loaded_count += len(adapter_state)\n\n    if (\n        mpu.get_data_parallel_rank() == 0\n        and mpu.get_tensor_model_parallel_rank() == 0\n        and mpu.get_pipeline_model_parallel_rank() == 0\n    ):\n        print(f\"Loaded {len(adapter_state)} adapter parameters from {checkpoint_path}\")\n\n\ndef count_adapter_parameters(model):\n    \"\"\"Count the number of trainable adapter parameters.\n\n    Args:\n        model: PyTorch model\n\n    Returns:\n        Tuple of (adapter_params, total_params, percentage)\n    \"\"\"\n    from verl.utils.megatron_utils import unwrap_model\n\n    unwrapped = unwrap_model(model)\n    if isinstance(unwrapped, list):\n        unwrapped = unwrapped[0]\n\n    adapter_params = 0\n    total_params = 0\n\n    for name, param in unwrapped.named_parameters():\n        total_params += param.numel()\n        if \"lora\" in name.lower() or \"adapter\" in name.lower():\n            if param.requires_grad:\n                adapter_params += param.numel()\n\n    percentage = 100 * adapter_params / total_params if total_params > 0 else 0\n\n    return adapter_params, total_params, percentage\n\n\ndef print_adapter_info(model):\n    \"\"\"Print information about adapter parameters in the model.\"\"\"\n    adapter_params, total_params, percentage = count_adapter_parameters(model)\n\n    print(f\"\\n{'=' * 60}\")\n    print(\"PEFT Adapter Information:\")\n    print(f\"  Total parameters:     {total_params:,}\")\n    print(f\"  Adapter parameters:   {adapter_params:,}\")\n    print(f\"  Trainable percentage: {percentage:.2f}%\")\n    print(f\"{'=' * 60}\\n\")\n\n\ndef convert_megatron_to_hf_target_modules(megatron_modules: list[str]) -> list[str]:\n    \"\"\"Convert megatron lora target modules to HF-style module names.\n\n    Args:\n        megatron_modules: List of megatron-style module names.\n\n    Returns:\n        List of HF-style module names with duplicates removed.\n    \"\"\"\n    hf_target_modules = []\n    for module in megatron_modules:\n        if module in MEGATRON_TO_HF_MODULES:\n            hf_target_modules.extend(MEGATRON_TO_HF_MODULES[module])\n        else:\n            hf_target_modules.append(module)\n    # Remove duplicates while preserving order\n    return list(dict.fromkeys(hf_target_modules))\n\n\ndef build_peft_config_for_vllm(lora_config: dict) -> dict:\n    \"\"\"Build a peft_config dict compatible with vLLM's PEFTHelper from megatron lora config.\n\n    Args:\n        lora_config: Megatron lora configuration dictionary.\n\n    Returns:\n        A dictionary compatible with vLLM's PEFTHelper.from_dict().\n    \"\"\"\n    from peft import TaskType\n\n    target_modules = lora_config.get(\"target_modules\", [\"linear_qkv\", \"linear_proj\", \"linear_fc1\", \"linear_fc2\"])\n    exclude_modules = lora_config.get(\"exclude_modules\", [])\n    hf_target_modules = convert_megatron_to_hf_target_modules(target_modules)\n    hf_exclude_modules = convert_megatron_to_hf_target_modules(exclude_modules)\n\n    return {\n        \"task_type\": TaskType.CAUSAL_LM,\n        \"r\": lora_config.get(\"rank\", 0),\n        \"lora_alpha\": lora_config.get(\"alpha\", 32),\n        \"target_modules\": hf_target_modules,\n        \"exclude_modules\": hf_exclude_modules,\n        \"bias\": \"none\",\n        \"lora_dropout\": lora_config.get(\"dropout\", 0.0),\n    }\n\n\n# vLLM needs to target all-linear no matter about specific LoRA config\ndef add_base_layer_suffix(\n    params: Iterator[tuple[str, torch.Tensor]],\n    model_type: str,\n) -> Iterator[tuple[str, torch.Tensor]]:\n    \"\"\"Yield param pairs with a base-layer suffix added to the param name.\n\n    Args:\n        params: Iterator of (param_name, tensor)\n        model_type: The type of the model (e.g., \"llama\").\n    \"\"\"\n    stacked_params = STACKED_PARAMS\n    # TODO: other models may have more special treatment, or integrate this into Megatron-Bridge\n    if model_type == \"llama\":\n        stacked_params = [\".embed_tokens.weight\", *STACKED_PARAMS]\n    for name, param in params:\n        ending_suffix = \"\"\n        for suffix in stacked_params:\n            if name.endswith(suffix):\n                ending_suffix = suffix\n                break\n        if ending_suffix:\n            suffix = ending_suffix.rsplit(\".\", 1)[-1]\n            name = f\"{name[: -len(suffix)]}base_layer.{suffix}\"\n        yield name, param\n\n\n__all__ = [\n    \"get_adapter_state_dict\",\n    \"save_adapter_checkpoint\",\n    \"load_adapter_checkpoint\",\n    \"count_adapter_parameters\",\n    \"print_adapter_info\",\n    \"convert_megatron_to_hf_target_modules\",\n    \"build_peft_config_for_vllm\",\n    \"add_base_layer_suffix\",\n]\n"
  },
  {
    "path": "verl/utils/megatron_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Pretrain utilities.\"\"\"\n\nimport gc\nimport inspect\nimport logging\nimport os\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport torch\nimport torch.nn.functional as F\nfrom megatron.core import ModelParallelConfig, mpu, parallel_state, tensor_parallel\nfrom megatron.core.distributed import DistributedDataParallel as DDP\nfrom megatron.core.distributed import DistributedDataParallelConfig\nfrom megatron.core.enums import ModelType\nfrom megatron.core.optimizer import ChainedOptimizer\nfrom megatron.core.parallel_state import get_global_memory_buffer\nfrom megatron.core.transformer import MLATransformerConfig, TransformerConfig\nfrom megatron.core.transformer.module import Float16Module\nfrom megatron.core.transformer.multi_token_prediction import MTPLossLoggingHelper\nfrom megatron.core.utils import get_attr_wrapped_model\nfrom transformers import PretrainedConfig\n\nimport verl.utils.megatron.tensor_parallel as tp_utils\nfrom verl.utils.device import get_device_id, get_device_name, get_torch_device\nfrom verl.utils.fs import local_mkdir_safe\nfrom verl.utils.model import normalize_model_name\nfrom verl.utils.torch_dtypes import PrecisionType\nfrom verl.workers.config import HFModelConfig, McoreEngineConfig\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef get_model_config(model):\n    return get_attr_wrapped_model(model, \"config\", allow_none=False)\n\n\ndef get_model(\n    model_provider_func,\n    model_type=ModelType.encoder_or_decoder,\n    wrap_with_ddp=True,\n    use_distributed_optimizer=True,\n    transformer_config=None,\n    override_ddp_config=None,\n):\n    \"\"\"Build the model.\"\"\"\n    # Build model.\n    if (\n        mpu.get_pipeline_model_parallel_world_size() > 1\n        and mpu.get_virtual_pipeline_model_parallel_world_size() is not None\n    ):\n        assert model_type != ModelType.encoder_and_decoder, (\n            \"Interleaved schedule not supported for model with both encoder and decoder\"\n        )\n        model = []\n        has_vp_stage = inspect.signature(mpu.is_pipeline_first_stage).parameters.get(\"vp_stage\", None) is not None\n        for i in range(mpu.get_virtual_pipeline_model_parallel_world_size()):\n            mpu.set_virtual_pipeline_model_parallel_rank(i)\n            # Set pre_process and post_process only after virtual rank is set.\n            extra_kwargs = {} if not has_vp_stage else {\"ignore_virtual\": False, \"vp_stage\": i}\n            pre_process = mpu.is_pipeline_first_stage(**extra_kwargs)\n            post_process = mpu.is_pipeline_last_stage(**extra_kwargs)\n            this_model = model_provider_func(pre_process=pre_process, post_process=post_process, vp_stage=i)\n            this_model.model_type = model_type\n            model.append(this_model)\n        mpu.set_virtual_pipeline_model_parallel_rank(0)\n    else:\n        pre_process = mpu.is_pipeline_first_stage()\n        post_process = mpu.is_pipeline_last_stage()\n        add_encoder = True\n        add_decoder = True\n        assert model_type != ModelType.encoder_and_decoder, \"Model type encoder_and_decoder is not supported\"\n        if model_type == ModelType.encoder_and_decoder:\n            if mpu.get_pipeline_model_parallel_world_size() > 1:\n                assert mpu.get_pipeline_model_parallel_split_rank() is not None, (\n                    \"Split rank needs to be specified for model with both encoder and decoder\"\n                )\n                rank = mpu.get_pipeline_model_parallel_rank()\n                split_rank = mpu.get_pipeline_model_parallel_split_rank()\n                world_size = mpu.get_pipeline_model_parallel_world_size()\n                pre_process = rank == 0 or rank == split_rank\n                post_process = (rank == (split_rank - 1)) or (rank == (world_size - 1))\n                add_encoder = mpu.is_pipeline_stage_before_split()\n                add_decoder = mpu.is_pipeline_stage_after_split()\n            model = model_provider_func(\n                pre_process=pre_process, post_process=post_process, add_encoder=add_encoder, add_decoder=add_decoder\n            )\n        else:\n            model = model_provider_func(pre_process=pre_process, post_process=post_process)\n        model.model_type = model_type\n\n    if not isinstance(model, list):\n        model = [model]\n\n    # Set tensor model parallel attributes if not set.\n    # Only parameters that are already tensor model parallel have these\n    # attributes set for them. We should make sure the default attributes\n    # are set for all params so the optimizer can use them.\n    for model_module in model:\n        for param in model_module.parameters():\n            tensor_parallel.set_defaults_if_not_set_tensor_model_parallel_attributes(param)\n\n    # Print number of parameters.\n    if mpu.get_data_parallel_rank() == 0:\n        print(\n            \" > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}\".format(\n                mpu.get_tensor_model_parallel_rank(),\n                mpu.get_pipeline_model_parallel_rank(),\n                sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model]),\n            ),\n            flush=True,\n        )\n\n    # GPU allocation.\n    if transformer_config is None or (not transformer_config.use_cpu_initialization):\n        for model_module in model:\n            model_module.to(f\"{get_device_name()}:{get_device_id()}\")\n\n    # Fp16 conversion.\n    config: TransformerConfig = get_model_config(model[0])\n    config.fp8 = None\n    tfconfig: TransformerConfig = model[0].config\n    if config.fp16 or config.bf16:  # the ModelParallelConfig in GPTModel\n        model = [Float16Module(config, model_module) for model_module in model]\n\n    if wrap_with_ddp:\n        ddp_models = []\n        ddp_config_dict = {\n            \"use_distributed_optimizer\": use_distributed_optimizer,\n            \"grad_reduce_in_fp32\": True,\n            \"overlap_grad_reduce\": False,\n        }\n        if override_ddp_config is not None:\n            ddp_config_dict.update(override_ddp_config)\n        ddp_config = DistributedDataParallelConfig(**ddp_config_dict)\n        for model_chunk_idx, model_chunk in enumerate(model):\n            ddp_model = DDP(\n                config=tfconfig,\n                module=model_chunk,\n                disable_bucketing=(model_chunk_idx > 0),\n                ddp_config=ddp_config,\n            )\n            ddp_models.append(ddp_model)\n        model = ddp_models\n        # # Broadcast params from data parallel src rank to other data parallel ranks.\n        # # if args.data_parallel_random_init:\n        for model_module in model:\n            model_module.broadcast_params()\n    return model\n\n\n@dataclass\nclass McoreModuleWrapperConfig:\n    \"\"\"Configuration for Mcore module wrapper.\"\"\"\n\n    is_value_model: bool = False\n    share_embeddings_and_output_weights: bool = False\n    wrap_with_ddp: bool = True\n    use_distributed_optimizer: bool = True\n\n\ndef make_megatron_module(\n    wrap_config: McoreModuleWrapperConfig,\n    tf_config: TransformerConfig,\n    hf_config: PretrainedConfig,\n    bridge: Any = None,\n    provider: Any = None,\n    override_model_config: dict[str, Any] = None,\n    override_ddp_config: dict[str, Any] = None,\n    peft_cls: Any = None,\n    peft_config: Any = None,\n):\n    if override_model_config is None:\n        override_model_config = {}\n\n    if bridge is not None:\n        if provider is None:\n            from verl.models.mcore.mbridge import freeze_moe_router, make_value_model\n\n            value_model_hook = make_value_model\n        else:\n            from verl.models.mcore.bridge import freeze_moe_router, make_value_model\n\n            hidden_size = (\n                hf_config.text_config.hidden_size if hasattr(hf_config, \"text_config\") else hf_config.hidden_size\n            )\n            value_model_hook = make_value_model(hidden_size, provider.sequence_parallel)\n\n        post_model_creation_callbacks = []\n        if wrap_config.is_value_model:\n            post_model_creation_callbacks.append(value_model_hook)\n        if override_model_config.get(\"moe_config\", {}).get(\"freeze_moe_router\", False):\n            post_model_creation_callbacks.append(freeze_moe_router)\n        if provider is not None:\n            # When using PEFT with Megatron-Bridge, we must apply PEFT transformation\n            # BEFORE wrapping the model in DDP. This is required because:\n            # 1. PEFT freezes base model parameters (requires_grad=False)\n            # 2. DDP must be aware of which parameters are trainable when building gradient buckets\n            # 3. The distributed optimizer must only track trainable (adapter) parameters\n            # See Megatron-Bridge docs: training/peft.md\n\n            # Register PEFT transformation as pre-wrap hook if peft_cls is specified\n            # This must happen BEFORE DDP wrapping to avoid KeyError with frozen parameters\n            if peft_cls is not None:\n                from verl.utils.megatron_peft_utils import load_adapter_checkpoint, print_adapter_info\n\n                def peft_pre_wrap_hook(model):\n                    \"\"\"Pre-wrap hook that applies PEFT transformation.\"\"\"\n                    # Apply PEFT transformation - this will freeze base model and add adapters\n                    # The PEFT callable handles both freezing and transformation\n                    transformed_model = peft_cls(model, training=True)\n\n                    # Set parameters to save (adapter-only checkpointing)\n                    peft_cls.set_params_to_save(transformed_model)\n\n                    # Load adapter weights if adapter_path is specified\n                    adapter_path = getattr(peft_config, \"adapter_path\", None)\n                    if adapter_path is not None and adapter_path:\n                        print(f\"Loading adapter weights from: {adapter_path}\")\n                        load_adapter_checkpoint(transformed_model, adapter_path)\n\n                    # Print PEFT statistics\n                    if torch.distributed.get_rank() == 0:\n                        print_adapter_info(transformed_model)\n\n                    return transformed_model\n\n                provider.register_pre_wrap_hook(peft_pre_wrap_hook)\n\n            # Register post-creation callbacks (make_value_model, freeze_moe_router) as pre-wrap hooks\n            for callback in post_model_creation_callbacks:\n                provider.register_pre_wrap_hook(callback)\n\n            # Create DDP config if needed\n            ddp_config = None\n            if wrap_config.wrap_with_ddp:\n                from megatron.bridge.training.config import DistributedDataParallelConfig\n\n                ddp_config_dict = {\n                    \"use_distributed_optimizer\": wrap_config.use_distributed_optimizer,\n                }\n                # Apply any DDP config overrides\n                if override_ddp_config is not None:\n                    ddp_config_dict.update(override_ddp_config)\n\n                ddp_config = DistributedDataParallelConfig(**ddp_config_dict)\n                ddp_config.finalize()\n\n            # Now call provide_distributed_model with all hooks registered\n            # Hooks will be applied automatically before DDP wrapping\n            model = provider.provide_distributed_model(\n                wrap_with_ddp=wrap_config.wrap_with_ddp,\n                ddp_config=ddp_config,\n                fp16=provider.fp16,\n                bf16=provider.bf16,\n            )\n\n            # Extract TransformerConfig from the created model\n            tf_config = get_model_config(model[0] if isinstance(model, list) else model)\n        else:\n            model = bridge.get_model(\n                post_model_creation_callbacks=post_model_creation_callbacks,\n                wrap_with_ddp=wrap_config.wrap_with_ddp,\n                fp16=tf_config.fp16,\n                bf16=tf_config.bf16,\n                ddp_config=override_ddp_config,\n            )\n\n        if isinstance(tf_config, MLATransformerConfig):\n            # Keep the same behavior as hf_to_mcore_config_dpskv3\n            from verl.models.mcore.patch import apply_patch\n\n            apply_patch()\n    else:\n\n        def megatron_model_provider(pre_process, post_process, vp_stage=None):\n            from verl.models.mcore import init_mcore_model\n\n            parallel_model = init_mcore_model(\n                tf_config,\n                hf_config,\n                pre_process,\n                post_process,\n                share_embeddings_and_output_weights=wrap_config.share_embeddings_and_output_weights,\n                value=wrap_config.is_value_model,\n                freeze_moe_router=override_model_config.get(\"moe_config\", {}).get(\"freeze_moe_router\", False),\n                vp_stage=vp_stage,\n            )\n            parallel_model.to(get_device_name())\n            return parallel_model\n\n        model = get_model(\n            megatron_model_provider,\n            wrap_with_ddp=wrap_config.wrap_with_ddp,\n            use_distributed_optimizer=wrap_config.use_distributed_optimizer,\n            override_ddp_config=override_ddp_config,\n        )\n    return model, tf_config\n\n\nALL_MODULE_WRAPPER_CLASSNAMES = (DDP, Float16Module)\n\n\ndef unwrap_model(model, module_instances=ALL_MODULE_WRAPPER_CLASSNAMES):\n    return_list = True\n    if not isinstance(model, list):\n        model = [model]\n        return_list = False\n    unwrapped_model = []\n    for model_module in model:\n        while isinstance(model_module, module_instances):\n            model_module = model_module.module\n        unwrapped_model.append(model_module)\n    if not return_list:\n        return unwrapped_model[0]\n    return unwrapped_model\n\n\ndef convert_config(hf_config: PretrainedConfig, megatron_config) -> TransformerConfig:\n    \"\"\"[Deprecated] convert config\n\n    Args:\n        hf_config (PretrainedConfig): _description_\n        megatron_config (_type_): _description_\n\n    Returns:\n        TransformerConfig: _description_\n    \"\"\"\n\n    warnings.warn(\"[deprecated] use config converter for more model support\", stacklevel=2)\n    print(f\"megatron config {megatron_config}\")\n    dt = PrecisionType.to_dtype(megatron_config.params_dtype)\n    print(f\"pipeline_dtype=megatron_config {dt}\")\n    qkv_bias = True if \"Qwen2ForCausalLM\" in hf_config.architectures else getattr(hf_config, \"attention_bias\", False)\n    overlap_p2p_comm = (\n        mpu.get_virtual_pipeline_model_parallel_world_size() is not None\n        and mpu.get_virtual_pipeline_model_parallel_world_size() > 1\n    )\n    batch_p2p_comm = False\n    transformer_config = TransformerConfig(\n        num_layers=hf_config.num_hidden_layers,\n        hidden_size=hf_config.hidden_size,\n        num_attention_heads=hf_config.num_attention_heads,\n        num_query_groups=hf_config.num_key_value_heads,\n        ffn_hidden_size=hf_config.intermediate_size,\n        #    max_position_embeddings=hf_config.max_position_embeddings,\n        activation_func=F.silu,\n        normalization=\"RMSNorm\",\n        #    rotary_percent=False, # default,\n        gated_linear_unit=True,  # for llama\n        use_cpu_initialization=True,\n        apply_residual_connection_post_layernorm=False,  # check what's this mean\n        add_bias_linear=False,\n        tensor_model_parallel_size=mpu.get_tensor_model_parallel_world_size(),\n        pipeline_model_parallel_size=mpu.get_pipeline_model_parallel_world_size(),\n        virtual_pipeline_model_parallel_size=mpu.get_virtual_pipeline_model_parallel_world_size(),\n        context_parallel_size=mpu.get_context_parallel_world_size(),\n        overlap_p2p_comm=overlap_p2p_comm,\n        batch_p2p_comm=batch_p2p_comm,\n        pipeline_dtype=dt,\n        params_dtype=dt,\n        sequence_parallel=mpu.get_tensor_model_parallel_world_size() > 1,\n        variable_seq_lengths=True,\n        masked_softmax_fusion=True,\n        moe_token_dispatcher_type=\"alltoall\",\n        attention_dropout=hf_config.attention_dropout,\n        hidden_dropout=getattr(hf_config, \"hidden_dropout\", 0.0),\n        add_qkv_bias=qkv_bias,\n        bf16=dt is torch.bfloat16,\n    )\n\n    return transformer_config\n\n\ndef mcore_model_parallel_config(\n    sequence_parallel: bool,\n    params_dtype: torch.dtype,\n) -> ModelParallelConfig:\n    # WARNING: Code should not reach this point. This function is deprecated and will be removed.\n    # Please use hf_to_mcore_config_dense() from verl.models.mcore.config_converter instead.\n    warnings.warn(\n        \"Code should not reach this point. This function is deprecated and will be removed. Please use \"\n        \"hf_to_mcore_config_dense() from verl.models.mcore.config_converter instead.\",\n        DeprecationWarning,\n        stacklevel=2,\n    )\n    return ModelParallelConfig(\n        tensor_model_parallel_size=mpu.get_tensor_model_parallel_world_size(),\n        pipeline_model_parallel_size=mpu.get_pipeline_model_parallel_world_size(),\n        virtual_pipeline_model_parallel_size=mpu.get_virtual_pipeline_model_parallel_world_size(),\n        context_parallel_size=mpu.get_context_parallel_world_size(),\n        sequence_parallel=sequence_parallel,\n        params_dtype=params_dtype,\n        pipeline_dtype=params_dtype,\n        bf16=True,\n        fp16=False,\n        timers=None,\n    )\n\n\n@torch.no_grad()\ndef offload_megatron_model_to_cpu(models):\n    \"\"\"\n    In megatron, the model and optimizer storage are:\n    - bf16 parameter data chunked in model parallel group\n    - fp32 grad chunked in model parallel group\n    - fp32 main_parameter chunked in model and dp group\n    - fp32 optimizer state chunked in model and dp group\n    \"\"\"\n    for model_chunk in models:\n        if isinstance(model_chunk, DDP):\n            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]\n            for buffers in model_chunk_all_buffers:\n                for buffer in buffers:\n                    # offload parameters\n                    if buffer.param_data.storage().size() > 0:\n                        buffer.param_data.cpu_data = buffer.param_data.data.cpu().pin_memory()\n                        buffer.param_data_size = buffer.param_data.storage().size()\n                        buffer.param_data.storage().resize_(0)\n\n                    assert buffer.param_data_size == buffer.param_data.cpu_data.storage().size()\n\n                    if buffer.grad_data.storage().size() > 0:\n                        # if the grad_data size is already zero, we assume that it is already offloaded\n                        buffer.grad_data_size = buffer.grad_data.storage().size()\n                        buffer.grad_data.storage().resize_(0)\n            # Offload frozen parameters not in DDP buffers (e.g. base model in LoRA/PEFT)\n            # DDP buffers only contain requires_grad=True params, so frozen params must be offloaded separately.\n            for param in model_chunk.module.parameters():\n                if not param.requires_grad and param.device.type != \"cpu\":\n                    param.data = param.data.to(\"cpu\", non_blocking=True)\n        else:\n            # we need this for ref module\n            for _, param in model_chunk.named_parameters():\n                param.data = param.data.to(\"cpu\", non_blocking=True)\n                if param.grad is not None:\n                    param.grad = param.grad.to(\"cpu\", non_blocking=True)\n    gc.collect()\n    get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef load_megatron_model_to_gpu(models, load_grad=True, load_frozen_params=True):\n    \"\"\"\n    Load megatron model to GPU.\n    Args:\n        models: The model to load.\n        load_grad: Whether to load gradients.\n        load_frozen_params: Whether to load frozen parameters.\n    \"\"\"\n    for model_chunk in models:\n        if isinstance(model_chunk, DDP):\n            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]\n            for buffers in model_chunk_all_buffers:\n                for buffer in buffers:\n                    # sometimes, we don't want to load grad for pure inference\n                    if load_grad and hasattr(buffer, \"grad_data_size\"):\n                        current_storage_size = buffer.grad_data.storage().size()\n                        if current_storage_size == 0 or current_storage_size == buffer.grad_data_size:\n                            buffer.grad_data.storage().resize_(buffer.grad_data_size)\n                            buffer.grad_data.zero_()\n                        else:\n                            # Non-standard layers (e.g. GatedDeltaNet) may have grad\n                            # buffers with mismatched storage size; skip resize and\n                            # zero in-place with current storage.\n                            buffer.grad_data.zero_()\n\n                    if buffer.param_data.storage().size() == 0:\n                        buffer.param_data.storage().resize_(buffer.param_data_size)\n                        # copy data from cpu to cuda\n                        buffer.param_data.copy_(buffer.param_data.cpu_data, non_blocking=True)\n\n            # Load frozen parameters that were offloaded (e.g. base model in LoRA/PEFT)\n            if load_frozen_params:\n                device_id = get_device_id()\n                for param in model_chunk.module.parameters():\n                    if not param.requires_grad and param.device.type == \"cpu\":\n                        param.data = param.data.to(device_id, non_blocking=True)\n        else:\n            # we need this for ref module\n            device_id = get_device_id()\n            for _, param in model_chunk.named_parameters():\n                param.data = param.data.to(device_id, non_blocking=True)\n                if param.grad is not None:\n                    param.grad = param.grad.to(device_id, non_blocking=True)\n    gc.collect()\n    get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef offload_megatron_copy_params(optimizers):\n    \"\"\"\n    Offload optimizer parameters to CPU. Supports both Megatron optimizers\n    and `ChainedOptimizer`, which wraps a list of underlying optimizers.\n\n    Args:\n        optimizers: The optimizer or ChainedOptimizer instance.\n    \"\"\"\n\n    def _iter_opts(opt):\n        if isinstance(opt, ChainedOptimizer):\n            return opt.chained_optimizers\n        return [opt]\n\n    def offload_tensor_to_cpu(tensor):\n        if tensor is None:\n            return\n        tensor.data = tensor.data.to(\"cpu\", non_blocking=True)\n\n    def offload_group_to_cpu(group):\n        if group is None:\n            return\n\n        if isinstance(group, list):\n            for param_group in group:\n                if isinstance(param_group, list):\n                    for param in param_group:\n                        offload_tensor_to_cpu(param)\n                else:\n                    offload_tensor_to_cpu(param_group)\n        else:\n            offload_tensor_to_cpu(group)\n\n    # Offload all parameter groups to CPU for each underlying optimizer\n\n    for _opt in _iter_opts(optimizers):\n        if hasattr(_opt, \"shard_fp32_from_float16_groups\"):\n            offload_group_to_cpu(_opt.shard_fp32_from_float16_groups)\n\n\n@torch.no_grad()\ndef load_megatron_copy_params(optimizers):\n    \"\"\"\n    Load optimizer parameters back to GPU. Handles ChainedOptimizer.\n\n    Args:\n        optimizers: Optimizer or ChainedOptimizer instance.\n    \"\"\"\n\n    def _iter_opts(opt):\n        if isinstance(opt, ChainedOptimizer):\n            return opt.chained_optimizers\n        return [opt]\n\n    def load_tensor_to_gpu(tensor):\n        if tensor is None:\n            return\n        device_id = get_device_id()\n        tensor.data = tensor.data.to(device_id, non_blocking=True)\n\n    def load_group_to_gpu(group):\n        if group is None:\n            return\n\n        if isinstance(group, list):\n            for param_group in group:\n                if isinstance(param_group, list):\n                    for param in param_group:\n                        load_tensor_to_gpu(param)\n                else:\n                    load_tensor_to_gpu(param_group)\n        else:\n            load_tensor_to_gpu(group)\n\n    # Load all parameter groups to GPU for each underlying optimizer\n\n    for _opt in _iter_opts(optimizers):\n        if hasattr(_opt, \"shard_fp32_from_float16_groups\"):\n            load_group_to_gpu(_opt.shard_fp32_from_float16_groups)\n\n\n@torch.no_grad()\ndef offload_megatron_optimizer(optimizers):\n    def _iter_opts(opt):\n        if isinstance(opt, ChainedOptimizer):\n            return opt.chained_optimizers\n        return [opt]\n\n    for _opt in _iter_opts(optimizers):\n        offload_megatron_copy_params(_opt)\n        ## worker may hold zero parameter when enabling custom pipeline layout\n        if _opt.optimizer is not None:\n            # HybridDeviceOptimizer: offload all sub-optimizer states to CPU\n            # TODO: this should be a method in Megatron-LM's HybridDeviceOptimizer\n            hdo = _opt.optimizer\n            if all(hasattr(hdo, attr) for attr in (\"sub_optimizers\", \"inner_param_to_orig_param\", \"state\")):\n                for optimizer in hdo.sub_optimizers:\n                    for param, state in optimizer.state.items():\n                        for k, v in state.items():\n                            if not isinstance(v, torch.Tensor):\n                                continue\n                            orig_param = hdo.inner_param_to_orig_param.get(param, param)\n                            hdo.state[orig_param][k] = state[k] = v.to(\"cpu\")\n            else:\n                opt_state_dict_values = _opt.optimizer.state.values()\n                for v in opt_state_dict_values:\n                    if \"exp_avg\" in v:\n                        v[\"exp_avg\"] = v[\"exp_avg\"].to(\"cpu\", non_blocking=True)\n                    if \"exp_avg_sq\" in v:\n                        v[\"exp_avg_sq\"] = v[\"exp_avg_sq\"].to(\"cpu\", non_blocking=True)\n\n        try:\n            # Free TransformerEngine's dummy weight gradients cache\n            # https://github.com/NVIDIA/TransformerEngine/blob/release_v2.10/transformer_engine/pytorch/module/base.py#L64\n            from transformer_engine.pytorch.module.base import _dummy_wgrads\n\n            _dummy_wgrads.clear()\n        except ImportError:\n            pass\n\n        # Free Megatron-LM's global memory buffer\n        get_global_memory_buffer().buffer.clear()\n\n        gc.collect()\n        get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef load_megatron_optimizer(optimizers):\n    def _iter_opts(opt):\n        if isinstance(opt, ChainedOptimizer):\n            return opt.chained_optimizers\n        return [opt]\n\n    for _opt in _iter_opts(optimizers):\n        load_megatron_copy_params(_opt)\n        ## worker may hold zero parameter when enabling custom pipeline layout\n        if _opt.optimizer is not None:\n            # if we are using HybridDeviceOptimizer, we need to only move gpu optimizer state to gpu\n            if hasattr(_opt.optimizer, \"_move_new_state_to_right_device\"):\n                _opt.optimizer._move_new_state_to_right_device()\n            else:\n                opt_state_dict_values = _opt.optimizer.state.values()\n                for v in opt_state_dict_values:\n                    if \"exp_avg\" in v:\n                        v[\"exp_avg\"] = v[\"exp_avg\"].to(get_device_id(), non_blocking=True)\n                    if \"exp_avg_sq\" in v:\n                        v[\"exp_avg_sq\"] = v[\"exp_avg_sq\"].to(get_device_id(), non_blocking=True)\n        gc.collect()\n        get_torch_device().empty_cache()\n\n\ndef get_dist_checkpoint_path(checkpoint_path):\n    local_mkdir_safe(checkpoint_path)\n    local_mkdir_safe(os.path.join(checkpoint_path, \"dist_ckpt\"))\n    return os.path.join(checkpoint_path, \"dist_ckpt\")\n\n\ndef get_hf_model_checkpoint_path(checkpoint_path):\n    local_mkdir_safe(checkpoint_path)\n    local_mkdir_safe(os.path.join(checkpoint_path, \"huggingface\"))\n    return os.path.join(checkpoint_path, \"huggingface\")\n\n\ndef get_transformer_config_checkpoint_path(checkpoint_path):\n    os.makedirs(checkpoint_path, exist_ok=True)\n    return os.path.join(checkpoint_path, \"transformer_config.json\")\n\n\ndef convert_megatron_model_to_transformers_model(\n    name,\n    param,\n    config: PretrainedConfig,\n    tp_size: int,\n    num_query_groups: int,\n    convert_qkv_gate_up_by_trunk_concat=False,\n):\n    \"\"\"Convert megatron model to transformers model.\"\"\"\n    new_params = {}\n\n    def convert_qkv_shard(full_tensor, q_name, k_name, v_name):\n        nonlocal config\n        nonlocal tp_size\n        nonlocal num_query_groups\n\n        q_shard_list = []\n        k_shard_list = []\n        v_shard_list = []\n        hidden_size_per_head = getattr(config, \"head_dim\", config.hidden_size // config.num_attention_heads)\n\n        if config.num_key_value_heads >= tp_size:\n            q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n            kv_size_tp = hidden_size_per_head * config.num_key_value_heads // tp_size\n            total_size = q_size_tp + 2 * kv_size_tp\n            for i in range(tp_size):\n                num_query_groups_per_partition = num_query_groups // tp_size\n                qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                q_size_chunk = q_size_tp // num_query_groups_per_partition\n                kv_size_chunk = kv_size_tp // num_query_groups_per_partition\n                for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition):\n                    q_part = qkv_part_chunk[:q_size_chunk]\n                    k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk]\n                    v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :]\n                    q_shard_list.append(q_part)\n                    k_shard_list.append(k_part)\n                    v_shard_list.append(v_part)\n        else:\n            q_size_tp = hidden_size_per_head * config.num_attention_heads // tp_size\n            kv_size_tp = hidden_size_per_head\n            total_size = q_size_tp + 2 * kv_size_tp\n            for i in range(tp_size):\n                num_query_groups_per_partition = num_query_groups // tp_size\n                qkv_part = full_tensor[i * total_size : (i + 1) * total_size]\n                q_size_chunk = q_size_tp // num_query_groups_per_partition\n                kv_size_chunk = kv_size_tp // num_query_groups_per_partition\n                for qkv_part_chunk in qkv_part.chunk(num_query_groups_per_partition):\n                    q_part = qkv_part_chunk[:q_size_chunk]\n                    k_part = qkv_part_chunk[q_size_chunk : q_size_chunk + kv_size_chunk]\n                    v_part = qkv_part_chunk[q_size_chunk + kv_size_chunk :]\n                    q_shard_list.append(q_part)\n                    if i * config.num_key_value_heads % tp_size == 0:\n                        k_shard_list.append(k_part)\n                        v_shard_list.append(v_part)\n\n        new_params[q_name] = torch.cat(q_shard_list, dim=0)\n        new_params[k_name] = torch.cat(k_shard_list, dim=0)\n        new_params[v_name] = torch.cat(v_shard_list, dim=0)\n\n    def convert_gate_up_shard(full_tensor, gate_name, up_name):\n        nonlocal config\n        nonlocal tp_size\n\n        intermediate_size_tp = config.intermediate_size // tp_size\n        gate_weight_list = []\n        up_weight_list = []\n        for i in range(tp_size):\n            gate_up_weight_tp = full_tensor[intermediate_size_tp * 2 * i : intermediate_size_tp * 2 * (i + 1)]\n            gate_weight_tp = gate_up_weight_tp[:intermediate_size_tp]\n            up_weight_tp = gate_up_weight_tp[intermediate_size_tp:]\n            gate_weight_list.append(gate_weight_tp)\n            up_weight_list.append(up_weight_tp)\n\n        new_params[gate_name] = torch.cat(gate_weight_list, dim=0)\n        new_params[up_name] = torch.cat(up_weight_list, dim=0)\n\n    if name == \"embedding.word_embeddings.weight\":\n        new_params[\"model.embed_tokens.weight\"] = param\n    elif \"self_attention\" in name:\n        splitted_name = name.split(\".\")\n        layer_number = splitted_name[2]\n        component = splitted_name[4]\n        param_type = splitted_name[5]\n        if component == \"linear_proj\":\n            new_params[f\"model.layers.{layer_number}.self_attn.o_proj.weight\"] = param\n        elif component == \"linear_qkv\" and not isinstance(param, list):\n            if param_type == \"layer_norm_weight\":\n                new_params[f\"model.layers.{layer_number}.input_layernorm.weight\"] = param\n            else:\n                if convert_qkv_gate_up_by_trunk_concat:\n                    convert_qkv_shard(\n                        param,\n                        f\"model.layers.{layer_number}.self_attn.q_proj.{param_type}\",\n                        f\"model.layers.{layer_number}.self_attn.k_proj.{param_type}\",\n                        f\"model.layers.{layer_number}.self_attn.v_proj.{param_type}\",\n                    )\n                else:\n                    new_params[f\"model.layers.{layer_number}.self_attn.qkv_proj.{param_type}\"] = param\n        elif component == \"q_layernorm\" or component == \"k_layernorm\":\n            hf_component = component.replace(\"layer\", \"\")\n            new_params[f\"model.layers.{layer_number}.self_attn.{hf_component}.weight\"] = param\n        else:\n            assert isinstance(param, list) and len(param) == 3\n            assert param_type == \"weight\" or param_type == \"bias\"\n            new_params[f\"model.layers.{layer_number}.self_attn.q_proj.{param_type}\"] = param[0]\n            new_params[f\"model.layers.{layer_number}.self_attn.k_proj.{param_type}\"] = param[1]\n            new_params[f\"model.layers.{layer_number}.self_attn.v_proj.{param_type}\"] = param[2]\n    elif \"mlp\" in name:\n        splitted_name = name.split(\".\")\n        layer_number = splitted_name[2]\n        component = splitted_name[4]\n        param_type = splitted_name[5]\n        if component == \"linear_fc1\" and not isinstance(param, list):\n            if param_type == \"layer_norm_weight\":\n                new_params[f\"model.layers.{layer_number}.post_attention_layernorm.weight\"] = param\n            elif param_type == \"weight\":\n                if convert_qkv_gate_up_by_trunk_concat:\n                    convert_gate_up_shard(\n                        param,\n                        f\"model.layers.{layer_number}.mlp.gate_proj.weight\",\n                        f\"model.layers.{layer_number}.mlp.up_proj.weight\",\n                    )\n                else:\n                    new_params[f\"model.layers.{layer_number}.mlp.gate_up_proj.weight\"] = param\n        elif component == \"linear_fc1\" and isinstance(param, list):\n            assert len(param) == 2\n            assert param_type == \"weight\" or param_type == \"bias\"\n            new_params[f\"model.layers.{layer_number}.mlp.gate_proj.weight\"] = param[0]\n            new_params[f\"model.layers.{layer_number}.mlp.up_proj.weight\"] = param[1]\n        elif component == \"linear_fc2\":\n            new_params[f\"model.layers.{layer_number}.mlp.down_proj.weight\"] = param\n    elif name == \"decoder.final_layernorm.weight\":\n        new_params[\"model.norm.weight\"] = param\n    elif name == \"output_layer.weight\":\n        new_params[\"lm_head.weight\"] = param\n    else:\n        raise ValueError(f\"Unknown param name: {name}\")\n    return new_params.keys(), new_params.values()\n\n\ndef broadcast_from_megatron_pp(tensor: torch.Tensor):\n    # tensor is not None only in one of the pp ranks\n    if tensor is not None:\n        shape = tensor.shape\n        dtype = tensor.dtype\n        tensor_parallel = getattr(tensor, \"tensor_model_parallel\", None)\n        partition_dim = getattr(tensor, \"partition_dim\", None)\n        tensor_spec = (shape, dtype, tensor_parallel, partition_dim)\n    else:\n        tensor_spec = None\n    tensor_spec_output = [None] * mpu.get_pipeline_model_parallel_world_size()\n    torch.distributed.all_gather_object(\n        object_list=tensor_spec_output, obj=tensor_spec, group=mpu.get_pipeline_model_parallel_group()\n    )\n    # find the src rank\n    target_tensor_spec = None\n    src_rank = None\n    for rank, tensor_spec in enumerate(tensor_spec_output):\n        if tensor_spec is not None:\n            if target_tensor_spec is None:\n                target_tensor_spec = tensor_spec\n            else:\n                raise ValueError(\"A tensor exists on two pp ranks\")\n            src_rank = rank\n    assert target_tensor_spec is not None\n    if tensor is None:\n        tensor = torch.empty(size=target_tensor_spec[0], dtype=target_tensor_spec[1], device=get_device_id())\n        if target_tensor_spec[2] is not None:\n            tensor.tensor_model_parallel = target_tensor_spec[2]\n        if target_tensor_spec[3] is not None:\n            tensor.partition_dim = target_tensor_spec[3]\n\n    global_rank = torch.distributed.get_global_rank(group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank)\n    torch.distributed.broadcast(tensor=tensor, src=global_rank, group=mpu.get_pipeline_model_parallel_group())\n    return tensor\n\n\ndef broadcast_str_from_megatron_pp(obj: Any):\n    obj_output = [None] * mpu.get_pipeline_model_parallel_world_size()\n    torch.distributed.all_gather_object(object_list=obj_output, obj=obj, group=mpu.get_pipeline_model_parallel_group())\n\n    src_rank = None\n    target_obj = None\n    for rank, item in enumerate(obj_output):\n        if item is not None:\n            if target_obj is not None:\n                raise ValueError(\"An object exists on two pp ranks\")\n            target_obj = item\n            src_rank = rank\n\n    assert target_obj is not None, \"No valid object found to broadcast.\"\n\n    global_rank = torch.distributed.get_global_rank(group=mpu.get_pipeline_model_parallel_group(), group_rank=src_rank)\n\n    obj_output = [None] * torch.distributed.get_world_size(group=mpu.get_pipeline_model_parallel_group())\n    obj_output[0] = target_obj\n    torch.distributed.broadcast_object_list(\n        object_list=obj_output, src=global_rank, group=mpu.get_pipeline_model_parallel_group()\n    )\n\n    return obj_output[0]\n\n\ndef default_tp_concat_fn(\n    layer_name_mapping,\n    name,\n    train_params,\n    infer_params,\n    model_config,\n    hf_config=None,\n    convert_qkv_gate_up_by_simple_split=False,\n):\n    \"\"\"\n    name: name of the parameter\n    train_params: training parameters\n    infer_params (Iterable[torch.Tensor]): a iterator towards list of parameters all-gathered from micro_dp_group\n    model_config: huggingface model_config\n    TODO(zhangchi.usc1992): currently, the implementation is adhoc. We can move this function to the model\n    definition so that it is model-agnostic. If the model doesn't implement this function,\n    we can throw an error to force user disable TP HybridEngine.\n    \"\"\"\n    from megatron.core import mpu\n\n    train_tp_size = mpu.get_tensor_model_parallel_world_size()\n    if layer_name_mapping.get(\"qkv_layer_name\") in name and \"layer_norm\" not in name:\n        # if the tensor is qkv, for each param on tp, split into q, k, v\n        # concat q, k, v separately.\n        q_lst = []\n        k_lst = []\n        v_lst = []\n        num_attention_heads = model_config.num_attention_heads\n        num_key_value_heads = model_config.num_key_value_heads\n        if \"vision_model\" in name:\n            num_attention_heads = hf_config.vision_config.num_heads\n            num_key_value_heads = hf_config.vision_config.num_heads\n        assert num_attention_heads % num_key_value_heads == 0\n        num_q_per_kv = num_attention_heads // num_key_value_heads\n        assert infer_params[0].shape[0] % (num_q_per_kv + 2) == 0, (\n            f\"param '{name}' shape '{infer_params[0].shape}' dim0 is not divisible by {num_q_per_kv + 2}\"\n        )\n        kv_size_per_tp = infer_params[0].shape[0] // (num_q_per_kv + 2)\n        split_size = [kv_size_per_tp * num_q_per_kv, kv_size_per_tp, kv_size_per_tp]\n        for infer_param in infer_params:\n            num_query_groups_per_partition = num_key_value_heads // train_tp_size\n            for chunk in infer_param.chunk(num_query_groups_per_partition):\n                split_size = [\n                    kv_size_per_tp * num_q_per_kv // num_query_groups_per_partition,\n                    kv_size_per_tp // num_query_groups_per_partition,\n                    kv_size_per_tp // num_query_groups_per_partition,\n                ]\n                q, k, v = chunk.split(split_size)\n                q_lst.append(q)\n                k_lst.append(k)\n                v_lst.append(v)\n        q = torch.cat(q_lst, dim=0)\n        k = torch.cat(k_lst, dim=0)\n        v = torch.cat(v_lst, dim=0)\n        infer_params = torch.cat((q, k, v), dim=0) if not convert_qkv_gate_up_by_simple_split else [q, k, v]\n\n    elif (\n        layer_name_mapping.get(\"gate_proj_layer_name\") in name\n        and \"layer_norm\" not in name\n        and \"vision_model.projection\" not in name\n    ):\n        # if the tensor is gate and proj\n        gate_lst = []\n        up_lst = []\n        for infer_param in infer_params:\n            gate, up = infer_param.chunk(2)\n            gate_lst.append(gate)\n            up_lst.append(up)\n        gate = torch.cat(gate_lst, dim=0)\n        up = torch.cat(up_lst, dim=0)\n        infer_params = torch.cat((gate, up), dim=0) if not convert_qkv_gate_up_by_simple_split else [gate, up]\n\n    elif \"mlp.experts.linear_fc2.weight\" in name:  # moe\n        infer_params = torch.cat(infer_params, dim=1)\n\n    else:\n        # concat tensor\n        infer_params = torch.cat(infer_params, dim=tp_utils.get_tensor_parallel_partition_dim(train_params))\n\n    return infer_params\n\n\ndef per_tensor_generator(\n    actor_module,\n    model_config,\n    weight_converter,\n    transformer_config,\n    layer_name_mapping,\n    convert_qkv_gate_up_by_simple_split=True,\n):\n    from megatron.core import parallel_state as mpu\n\n    pp_rank = mpu.get_pipeline_model_parallel_rank()\n    ep_size = mpu.get_expert_model_parallel_world_size()\n    etp_size = mpu.get_expert_tensor_parallel_world_size()\n    ep_group = mpu.get_expert_model_parallel_group()\n    etp_group = mpu.get_expert_tensor_parallel_group()\n    vpp_size = len(actor_module)\n    all_gather_group = mpu.get_tensor_model_parallel_group()\n    all_gather_group_size = torch.distributed.get_world_size(group=all_gather_group)\n\n    def tensor_generator():\n        for scan_vpp_idx in range(vpp_size):\n            existing_keys = set()\n            model = unwrap_model(actor_module[scan_vpp_idx])\n            for name, param in model.named_parameters():\n                existing_keys.add(name)\n                yield name, param\n            # note\n            # there is a bug in megatron GPTModel\n            # decoder.layers[n].mlp.router.expert_bias\" in GPTModel is not registered in named_parameter, but in\n            # state_dict(). for now we patch it by adding those keys to extra_keys.\n            extra_keys = [x for x in model.state_dict().keys() if \"_extra_state\" not in x and x not in existing_keys]\n            for name in extra_keys:\n                yield name, model.state_dict()[name].to(get_device_id())\n\n    # we need first make all rank get full model information\n    meta_info = []\n    for scan_vpp_idx in range(vpp_size):\n        existing_keys = set()\n        model = unwrap_model(actor_module[scan_vpp_idx])\n        for idx, (name, _) in enumerate(model.named_parameters()):\n            existing_keys.add(name)\n            meta_info.append((pp_rank, scan_vpp_idx, idx, name))\n        extra_keys = [x for x in model.state_dict().keys() if \"_extra_state\" not in x and x not in existing_keys]\n        for name in extra_keys:\n            meta_info.append((pp_rank, scan_vpp_idx, idx, name))\n\n    obj_spec_output = [None] * mpu.get_pipeline_model_parallel_world_size()\n    torch.distributed.all_gather_object(\n        object_list=obj_spec_output, obj=meta_info, group=mpu.get_pipeline_model_parallel_group()\n    )\n    layer_list_meta = [item for sublist in obj_spec_output for item in sublist]\n\n    gen_func = tensor_generator()\n\n    # lazy load tensor for full model\n    for cur_pp_rank, scan_vpp_idx, idx, name in layer_list_meta:\n        if model_config.tie_word_embeddings and (\"output_layers\" in name):\n            import warnings\n\n            warnings.warn(\n                \"Current model sharing word and embedding weights, skip output layer conversion\", stacklevel=2\n            )\n            continue\n\n        if cur_pp_rank == pp_rank:\n            try:\n                cur_name, cur_tensor = next(gen_func)\n            except StopIteration:\n                cur_name, cur_tensor = None, None\n            cur_name = normalize_model_name(name, cur_pp_rank, scan_vpp_idx, transformer_config)\n        else:\n            cur_tensor, cur_name = None, None\n\n        # pp broadcast model tensor and name\n        cur_name = broadcast_str_from_megatron_pp(cur_name)\n        broad_pp_tensor = broadcast_from_megatron_pp(cur_tensor)\n\n        # (xya): this is a hack to fix the name of the parameters\n        while cur_name.startswith(\"module.\"):\n            cur_name = cur_name[len(\"module.\") :]\n\n        # EP\n        if \".mlp.experts.linear_fc\" in cur_name and ep_size > 1:\n            num_experts = weight_converter.mcore_config.num_moe_experts\n            num_experts_per_rank = num_experts // ep_size\n            infer_params = [torch.empty_like(broad_pp_tensor) for _ in range(ep_size)]\n            torch.distributed.all_gather(infer_params, broad_pp_tensor, group=ep_group)\n\n            name_prefix, local_expert_id = cur_name.split(\".weight\")\n            local_expert_id = int(local_expert_id)\n            global_expert_ids = [num_experts_per_rank * ep_rank + local_expert_id for ep_rank in range(ep_size)]\n            global_expert_names = [f\"{name_prefix}.weight{expert_id}\" for expert_id in global_expert_ids]\n\n            for name, param in zip(global_expert_names, infer_params, strict=True):\n                if etp_size > 1:\n                    # gather etp\n                    etp_params = [torch.empty_like(param) for _ in range(etp_size)]\n                    torch.distributed.all_gather(etp_params, param, group=etp_group)\n                    params = etp_params\n                else:\n                    params = [param]\n\n                merge_params = default_tp_concat_fn(\n                    layer_name_mapping,\n                    name,\n                    broad_pp_tensor,\n                    params,\n                    model_config,\n                    weight_converter.hf_config,\n                    convert_qkv_gate_up_by_simple_split,\n                )\n                if not isinstance(merge_params, list):\n                    merge_params = [merge_params]\n                converted_names, converted_params = weight_converter.convert_param(name, merge_params)\n\n                yield from zip(converted_names, [param.detach() for param in converted_params], strict=True)\n            continue\n\n        # tp all gather\n        if tp_utils.is_tensor_parallel_param(broad_pp_tensor):\n            # allocate a new tensor with proper size\n            if all_gather_group_size <= 1:\n                infer_params = [broad_pp_tensor]\n            else:\n                infer_params = [torch.empty_like(broad_pp_tensor) for _ in range(all_gather_group_size)]\n                torch.distributed.all_gather(infer_params, broad_pp_tensor, group=mpu.get_tensor_model_parallel_group())\n            infer_params = default_tp_concat_fn(\n                layer_name_mapping,\n                cur_name,\n                broad_pp_tensor,\n                infer_params,\n                model_config,\n                weight_converter.hf_config,\n                convert_qkv_gate_up_by_simple_split,\n            )\n        else:\n            infer_params = broad_pp_tensor\n\n        if not isinstance(infer_params, list):\n            infer_params = [infer_params]\n        converted_names, converted_params = weight_converter.convert_param(cur_name, infer_params)\n\n        yield from zip(converted_names, [param.detach() for param in converted_params], strict=True)\n\n\ndef get_transformer_layer_offset(pipeline_rank, vp_stage, config: TransformerConfig):\n    \"\"\"\n    Get the index offset of any pipeline stage, given the level of pipelining.\n\n    Make pipeline_rank and vp_stage as two arguments to make it more flexible,\n    which is able to fetch layer offset for any pipeline stage.\n    The original function only returns the layer offset for current pipeline stage.\n\n    Extension to https://github.com/NVIDIA/Megatron-LM/blob/main/megatron/core/transformer/transformer_layer.py::get_transformer_layer_offset\n    \"\"\"\n\n    has_vp_stage = (\n        inspect.signature(parallel_state.is_pipeline_first_stage).parameters.get(\"vp_stage\", None) is not None\n    )\n    extra_kwargs = {} if not has_vp_stage else {\"ignore_virtual\": False, \"vp_stage\": vp_stage}\n\n    if config.pipeline_model_parallel_size > 1:\n        if hasattr(config, \"pipeline_model_parallel_layout\") and config.pipeline_model_parallel_layout:\n            from megatron.core.transformer.enums import LayerType\n\n            offset = config.pipeline_model_parallel_layout.get_layer_offset(\n                layer_type=LayerType.decoder, vp_stage=vp_stage\n            )\n        elif (\n            config.num_layers_in_first_pipeline_stage is not None\n            or config.num_layers_in_last_pipeline_stage is not None\n        ):\n            # Calculate number of pipeline stages to distribute the remaining Transformer\n            # layers after deducting the Transformer layers in the first or the last stages\n            middle_pipeline_stages = config.pipeline_model_parallel_size\n            middle_pipeline_stages -= sum(\n                [\n                    1 if x is not None else 0\n                    for x in (\n                        config.num_layers_in_first_pipeline_stage,\n                        config.num_layers_in_last_pipeline_stage,\n                    )\n                ]\n            )\n\n            # Calculate layers to distribute in each pipeline stage. If the\n            # num_layers_in_first_pipeline_stage and num_layers_in_last_pipeline_stage\n            # are not set, we will not enable uneven pipeline. All layers will be treated\n            # as middle layers.\n            num_layers_in_first_pipeline_stage = (\n                0 if config.num_layers_in_first_pipeline_stage is None else config.num_layers_in_first_pipeline_stage\n            )\n            num_layers_in_last_pipeline_stage = (\n                0 if config.num_layers_in_last_pipeline_stage is None else config.num_layers_in_last_pipeline_stage\n            )\n\n            middle_num_layers = (\n                config.num_layers - num_layers_in_first_pipeline_stage - num_layers_in_last_pipeline_stage\n            )\n\n            if (vp_size := config.virtual_pipeline_model_parallel_size) is not None:\n                assert vp_stage is not None, \"vp_stage must be provided if virtual pipeline model parallel size is set\"\n\n                # Calculate number of layers in each virtual model chunk\n                # If the num_layers_in_first_pipeline_stage and\n                # num_layers_in_last_pipeline_stage are not set, all pipeline stages\n                # will be treated as middle pipeline stages in the calculation\n                num_layers_per_virtual_model_chunk_in_first_pipeline_stage = (\n                    0\n                    if config.num_layers_in_first_pipeline_stage is None\n                    else config.num_layers_in_first_pipeline_stage // vp_size\n                )\n\n                num_layers_per_virtual_model_chunk_in_last_pipeline_stage = (\n                    0\n                    if config.num_layers_in_last_pipeline_stage is None\n                    else config.num_layers_in_last_pipeline_stage // vp_size\n                )\n\n                num_layers_per_vritual_model_chunk_in_middle_pipeline_stage = middle_num_layers // vp_size\n\n                # First stage + middle stage + last stage\n                total_virtual_chunks = (\n                    num_layers_per_virtual_model_chunk_in_first_pipeline_stage\n                    + num_layers_per_vritual_model_chunk_in_middle_pipeline_stage\n                    + num_layers_per_virtual_model_chunk_in_last_pipeline_stage\n                )\n\n                # Calculate the layer offset with interleaved uneven pipeline parallelism\n                if pipeline_rank == 0:\n                    offset = vp_stage * total_virtual_chunks\n                else:\n                    offset = (\n                        vp_stage * total_virtual_chunks\n                        + num_layers_per_virtual_model_chunk_in_first_pipeline_stage\n                        + (pipeline_rank - 1)\n                        * (num_layers_per_vritual_model_chunk_in_middle_pipeline_stage // middle_pipeline_stages)\n                    )\n            else:\n                if middle_pipeline_stages > 0:\n                    num_layers_per_pipeline_rank = middle_num_layers // middle_pipeline_stages\n                else:\n                    num_layers_per_pipeline_rank = 0\n\n                middle_pipeline_rank = (\n                    pipeline_rank if config.num_layers_in_first_pipeline_stage is None else pipeline_rank - 1\n                )\n\n                if pipeline_rank == 0:\n                    offset = 0\n                else:\n                    offset = (middle_pipeline_rank * num_layers_per_pipeline_rank) + num_layers_in_first_pipeline_stage\n        else:\n            num_layers = config.num_layers\n\n            # Increase the number of layers by one if we include the embedding (loss)\n            # layer into pipeline parallelism partition and placement\n            if config.account_for_embedding_in_pipeline_split:\n                num_layers += 1\n\n            if config.account_for_loss_in_pipeline_split:\n                num_layers += 1\n\n            num_layers_per_pipeline_rank = num_layers // config.pipeline_model_parallel_size\n\n            if (vp_size := config.virtual_pipeline_model_parallel_size) is not None:\n                assert vp_stage is not None, \"vp_stage must be provided if virtual pipeline model parallel size is set\"\n\n                num_layers_per_virtual_rank = num_layers_per_pipeline_rank // vp_size\n                total_virtual_chunks = num_layers // vp_size\n                offset = vp_stage * total_virtual_chunks + (pipeline_rank * num_layers_per_virtual_rank)\n\n                # Reduce the offset of embedding layer from the total layer number\n                if config.account_for_embedding_in_pipeline_split and not parallel_state.is_pipeline_first_stage(\n                    **extra_kwargs\n                ):\n                    offset -= 1\n            else:\n                offset = pipeline_rank * num_layers_per_pipeline_rank\n\n                # Reduce the offset of embedding layer from the total layer number\n                if config.account_for_embedding_in_pipeline_split and not parallel_state.is_pipeline_first_stage(\n                    **extra_kwargs\n                ):\n                    offset -= 1\n    else:\n        offset = 0\n    return offset\n\n\ndef register_megatron_training_hooks(model: list[torch.nn.Module], optimizer):\n    from megatron.core.distributed import finalize_model_grads\n    from megatron.core.utils import get_model_config\n\n    try:\n        from megatron.core.distributed.fsdp.mcore_fsdp_adapter import FullyShardedDataParallel as megatron_FSDP\n    except ImportError:\n        megatron_FSDP = DDP\n\n    # register some callbacks for megatron training, following https://github.com/NVIDIA/Megatron-LM/blob/core_v0.15.0rc7/megatron/training/training.py#L2039-L2057\n    for one_model in model:\n        config = get_model_config(one_model)\n        config.grad_scale_func = optimizer.scale_loss\n        config.finalize_model_grads_func = finalize_model_grads\n\n        overlap_param_gather = getattr(optimizer.config, \"overlap_param_gather\", False)\n        overlap_grad_reduce = getattr(one_model.ddp_config, \"overlap_grad_reduce\", False)\n        align_grad_reduce = True  # default to True, seldom to be false\n        align_param_gather = getattr(one_model.ddp_config, \"align_param_gather\", False)\n\n        if isinstance(model[0], megatron_FSDP | DDP) and overlap_grad_reduce:\n            assert config.no_sync_func is None, (\n                \"When overlap_grad_reduce is True, config.no_sync_func must be None; \"\n                \"a custom no_sync_func is not supported when overlapping grad-reduce\"\n            )\n            config.no_sync_func = [model_chunk.no_sync for model_chunk in model]\n            if len(model) == 1:\n                config.no_sync_func = config.no_sync_func[0]\n            if align_grad_reduce:\n                config.grad_sync_func = [model_chunk.start_grad_sync for model_chunk in model]\n                if len(model) == 1:\n                    config.grad_sync_func = config.grad_sync_func[0]\n        if overlap_param_gather and align_param_gather:\n            config.param_sync_func = [model_chunk.start_param_sync for model_chunk in model]\n            if len(model) == 1:\n                config.param_sync_func = config.param_sync_func[0]\n\n\ndef mapping_string_to_attn_backend(args: dict) -> dict:\n    if \"attention_backend\" in args and isinstance(args[\"attention_backend\"], str):\n        from megatron.core.transformer.enums import AttnBackend\n\n        args[\"attention_backend\"] = AttnBackend[args[\"attention_backend\"]]\n    return args\n\n\ndef get_megatron_mtp_loss(n_micro_batch):\n    # Calculate MTP loss scale similar to Megatron-LM implementation\n    mtp_loss_scale = 1.0 / n_micro_batch\n\n    # Create a dummy total_loss_dict to collect MTP metrics\n    total_loss_dict = {}\n\n    # Track MTP metrics - this will populate total_loss_dict with MTP losses\n    MTPLossLoggingHelper.track_mtp_metrics(\n        loss_scale=mtp_loss_scale, iteration=0, writer=None, wandb_writer=None, total_loss_dict=total_loss_dict\n    )\n    # Add MTP metrics to losses_reduced if any were collected\n    # total_loss_dict: {'mtp_1 loss': tensor(value, device='cuda:0')}\n    output = {}\n    if total_loss_dict:\n        for key, value in total_loss_dict.items():\n            # Convert key to have proper prefix and format\n            formatted_key = f\"mtp_losses/{key.replace(' ', '_')}\"\n            # only added to the 0th batch, the MTP loss obtained is a global value, and will be the same for every batch\n            output[formatted_key] = value.cpu().item()\n    return output\n\n\ndef get_megatron_module_device(models: list[Any]) -> str:\n    if not models:\n        return \"cpu\"\n\n    model_chunk = models[0]\n    if not model_chunk.buffers:\n        try:\n            return next(model_chunk.module.parameters()).device.type\n        except StopIteration:\n            return \"cpu\"\n\n    buffer = model_chunk.buffers[0]\n    if buffer.param_data.storage().size() == 0:\n        return \"cpu\"\n    else:\n        return get_device_name()\n\n\ndef check_mtp_config(model_config: HFModelConfig, engine_config: McoreEngineConfig):\n    \"\"\"\n    Check and configure MTP (Multi-Token Prediction) settings.\n\n    Cases:\n        - mtp.enable == False and no MTP layers: return directly\n        - mtp.enable == False and has MTP layers: set num_nextn_predict_layers = 0\n        - mtp.enable == True and has MTP layers: configure override_transformer_config\n        - mtp.enable == True and no MTP layers: raise ValueError\n    \"\"\"\n    has_mtp = (\n        model_config.hf_config.num_nextn_predict_layers > 0\n        if hasattr(model_config.hf_config, \"num_nextn_predict_layers\")\n        else False\n    )\n    enable_mtp = model_config.mtp.enable\n\n    if not enable_mtp and not has_mtp:\n        return\n    elif not enable_mtp and has_mtp:\n        model_config.hf_config.num_nextn_predict_layers = 0\n    elif enable_mtp and not has_mtp:\n        raise ValueError(\"enable mtp while model has no mtp layer, please use a model with mtp layer\")\n    elif enable_mtp and has_mtp:\n        if \"mtp_loss_scaling_factor\" not in engine_config.override_transformer_config:\n            engine_config.override_transformer_config[\"mtp_loss_scaling_factor\"] = (\n                model_config.mtp.mtp_loss_scaling_factor\n            )\n    return\n\n\ndef patch_engine_mtp(module, model_config):\n    \"\"\"\n    Apply MTP patches to the model module.\n\n    Args:\n        module: The model module to patch. Can be a single module or a list of modules.\n        model_config: The model configuration containing MTP settings.\n    \"\"\"\n    logger.warning(\"Applying mtp patch...\")\n    from verl.models.mcore.mtp_patch import patch_mtp_layer_get_embeddings, patch_postprocess\n\n    print(module)\n\n    modules = module if isinstance(module, list) else [module]\n    for m in modules:\n        patch_postprocess(m)\n        if model_config.mtp.detach_encoder:\n            patch_mtp_layer_get_embeddings(m)\n\n\n@torch.no_grad()\ndef copy_megatron_model_to_cpu(models):\n    \"\"\"\n    Copy Megatron model parameters to CPU memory (non-destructive copy).\n    Unlike offload_megatron_model_to_cpu which moves data, this function creates\n    independent copies on CPU while keeping GPU data intact.\n\n    Args:\n        models: List of model chunks (DDP-wrapped or unwrapped)\n\n    Returns:\n        dict: CPU state containing copied parameters and buffers\n    \"\"\"\n    cpu_state = {}\n\n    for model_idx, model_chunk in enumerate(models):\n        if isinstance(model_chunk, DDP):\n            # Handle DDP-wrapped models\n            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]\n            buffer_states = []\n\n            for buffers in model_chunk_all_buffers:\n                buffer_list = []\n                for buffer in buffers:\n                    buffer_state = {}\n\n                    # Copy parameter data to CPU\n                    if buffer.param_data.storage().size() > 0:\n                        buffer_state[\"param_data\"] = buffer.param_data.data.cpu().clone().pin_memory()\n\n                    buffer_list.append(buffer_state)\n                buffer_states.append(buffer_list)\n\n            cpu_state[f\"model_chunk_{model_idx}\"] = {\"buffer_states\": buffer_states, \"is_ddp\": True}\n        else:\n            # Handle non-DDP models (ref module)\n            model_state = {}\n            for name, param in model_chunk.named_parameters():\n                param_state = {\"data\": param.data.cpu().clone().pin_memory()}\n                model_state[name] = param_state\n\n            cpu_state[f\"model_chunk_{model_idx}\"] = {\"model_state\": model_state, \"is_ddp\": False}\n\n    return cpu_state\n\n\n@torch.no_grad()\ndef restore_megatron_model_from_cpu(models, cpu_state):\n    \"\"\"\n    Restore Megatron model parameters from CPU memory back to GPU.\n\n    Args:\n        models: List of model chunks to restore to\n        cpu_state: CPU state dict returned from copy_megatron_model_to_cpu\n    \"\"\"\n    for model_idx, model_chunk in enumerate(models):\n        chunk_key = f\"model_chunk_{model_idx}\"\n        if chunk_key not in cpu_state:\n            continue\n\n        chunk_state = cpu_state[chunk_key]\n\n        if chunk_state[\"is_ddp\"] and isinstance(model_chunk, DDP):\n            # Restore DDP buffers\n            model_chunk_all_buffers = [model_chunk.buffers, model_chunk.expert_parallel_buffers]\n            buffer_states = chunk_state[\"buffer_states\"]\n\n            for buffers, buffer_list in zip(model_chunk_all_buffers, buffer_states, strict=False):\n                for buffer, buffer_state in zip(buffers, buffer_list, strict=False):\n                    # Restore parameter data\n                    if \"param_data\" in buffer_state:\n                        buffer.param_data.data.copy_(buffer_state[\"param_data\"].to(buffer.param_data.device))\n\n        elif not chunk_state[\"is_ddp\"] and not isinstance(model_chunk, DDP):\n            # Restore non-DDP models\n            model_state = chunk_state[\"model_state\"]\n            for name, param in model_chunk.named_parameters():\n                if name in model_state:\n                    param_state = model_state[name]\n                    param.data.copy_(param_state[\"data\"].to(param.device))\n"
  },
  {
    "path": "verl/utils/memory_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport gc\nimport inspect\nimport logging\nimport os\nfrom datetime import datetime\nfrom pathlib import Path\n\nimport torch\n\nfrom verl.utils.device import get_torch_device, is_cuda_available\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef aggressive_empty_cache(force_sync: bool = True, max_retries: int = 3) -> None:\n    \"\"\"\n    More aggressive GPU memory cleanup function, tries to release PyTorch reserved\n    but unallocated memory.\n\n    Args:\n        force_sync: Whether to force device synchronization\n        max_retries: Maximum number of retries\n    \"\"\"\n    device = get_torch_device()\n    if not device.is_available():\n        return\n\n    for attempt in range(max_retries):\n        # Record memory status before cleanup\n        before_reserved = device.memory_reserved()\n        before_allocated = device.memory_allocated()\n\n        # Run garbage collection\n        gc.collect()\n\n        # Clear PyTorch cache\n        device.empty_cache()\n\n        # Force synchronization (optional)\n        if force_sync:\n            device.synchronize()\n\n        # Record memory status after cleanup\n        after_reserved = device.memory_reserved()\n        after_allocated = device.memory_allocated()\n\n        # Calculate freed memory\n        reserved_freed = before_reserved - after_reserved\n        allocated_freed = before_allocated - after_allocated\n\n        logger.info(\n            f\"Memory cleanup attempt {attempt + 1}: Freed {reserved_freed / 1024**3:.2f} GB reserved, \"\n            f\"{allocated_freed / 1024**3:.2f} GB allocated\"\n        )\n\n        # Stop retrying if little memory was freed\n        if reserved_freed < 1024**3:  # less than 1GB\n            break\n\n\ndef reset_memory_stats() -> None:\n    \"\"\"Reset GPU memory statistics\"\"\"\n    if get_torch_device().is_available():\n        device = get_torch_device()\n        device.reset_peak_memory_stats()\n        device.reset_accumulated_memory_stats()\n\n\ndef get_memory_info() -> dict:\n    \"\"\"Get detailed GPU memory information\"\"\"\n    if not get_torch_device().is_available():\n        return {}\n\n    device = get_torch_device()\n    device_id = device.current_device()\n\n    return {\n        \"total_memory_gb\": device.get_device_properties(device_id).total_memory / 1024**3,\n        \"reserved_memory_gb\": device.memory_reserved() / 1024**3,\n        \"allocated_memory_gb\": device.memory_allocated() / 1024**3,\n        \"cached_memory_gb\": (device.memory_reserved() - device.memory_allocated()) / 1024**3,\n        \"max_memory_allocated_gb\": device.max_memory_allocated() / 1024**3,\n        \"max_memory_reserved_gb\": device.max_memory_reserved() / 1024**3,\n    }\n\n\ndef log_memory_usage(stage: str = \"current\") -> None:\n    \"\"\"Log GPU memory usage\"\"\"\n    if not get_torch_device().is_available():\n        return\n\n    info = get_memory_info()\n    logger.info(\n        f\"Memory usage [{stage}]: \"\n        f\"Total: {info['total_memory_gb']:.2f} GB, \"\n        f\"Allocated: {info['allocated_memory_gb']:.2f} GB, \"\n        f\"Reserved: {info['reserved_memory_gb']:.2f} GB, \"\n        f\"Cached: {info['cached_memory_gb']:.2f} GB\"\n    )\n\n\ndef optimize_memory_for_inference() -> None:\n    \"\"\"Optimize GPU memory usage for inference\"\"\"\n    if not get_torch_device().is_available():\n        return\n\n    # Set a more aggressive memory allocation policy\n    get_torch_device().set_per_process_memory_fraction(0.95)  # Use 95% of GPU memory\n\n    # Clear cache\n    aggressive_empty_cache(force_sync=True)\n\n    logger.info(\"Optimized GPU memory usage for inference\")\n\n\ndef optimize_memory_for_training() -> None:\n    \"\"\"Optimize GPU memory usage for training\"\"\"\n    if not get_torch_device().is_available():\n        return\n\n    # Set a moderate memory allocation policy\n    get_torch_device().set_per_process_memory_fraction(0.9)  # Use 90% of GPU memory\n\n    # Clear cache\n    aggressive_empty_cache(force_sync=False)\n\n    logger.info(\"Optimized GPU memory usage for training\")\n\n\ndef enable_memory_visualize(\n    trace_alloc_max_entries: int = 200_000,\n    stack_depth: int = 32,\n    context: str = \"all\",\n    stacks: str = \"all\",\n    devices=None,\n    record_context: bool = True,\n):\n    \"\"\"\n    Enables memory history recording for CUDA allocations. This function\n    should be called before any large-scale CUDA allocations. For DDP or\n    multi-process setups, it must be called on each rank.\n\n    Args:\n        trace_alloc_max_entries (int): Maximum number of allocation entries\n            to record.\n        stack_depth (int): The depth of the call stack to capture for each\n            allocation. (Supported by some PyTorch versions).\n        context (str): The type of memory events to record.\n            'alloc': records only allocation events.\n            'state': records memory state changes.\n            'all': records both.\n        stacks (str): The type of call stacks to record.\n            'python': records Python stacks.\n            'cpp': records C++ stacks (available in some versions).\n            'all': records both.\n        devices (Union[int, list[int], None]): The device for which to enable\n            memory history. `None` enables it for the current default device.\n        record_context (bool): Whether to record context information for\n            allocations. Required by older PyTorch versions.\n    \"\"\"\n    # Memory history recording is CUDA-specific functionality\n    if not is_cuda_available:\n        logger.warning(\"[memory_visualize] Memory history recording is only available on CUDA devices\")\n        return\n\n    f = get_torch_device().memory._record_memory_history\n    params = set(inspect.signature(f).parameters.keys())\n\n    def _one_call(dev_kw=None):\n        kwargs = {}\n        if \"context\" in params:\n            kwargs[\"context\"] = context\n        if \"stacks\" in params:\n            kwargs[\"stacks\"] = stacks\n        if \"max_entries\" in params:\n            kwargs[\"max_entries\"] = trace_alloc_max_entries\n        elif \"trace_alloc_max_entries\" in params:\n            kwargs[\"trace_alloc_max_entries\"] = trace_alloc_max_entries\n        if \"stack_depth\" in params:\n            kwargs[\"stack_depth\"] = stack_depth\n        if dev_kw is not None:\n            if \"device\" in params:\n                kwargs[\"device\"] = dev_kw\n            elif \"devices\" in params:\n                kwargs[\"devices\"] = dev_kw if isinstance(dev_kw, list) else [dev_kw]\n        if \"record_context\" in params:\n            kwargs[\"record_context\"] = record_context\n\n        try:\n            f(**kwargs)\n            return \"native\", kwargs\n        except TypeError:\n            try:\n                if \"trace_alloc_max_entries\" in params and \"record_context\" in params:\n                    f(enabled=True, trace_alloc_max_entries=trace_alloc_max_entries, record_context=True)\n                    return \"legacy\", {\n                        \"enabled\": True,\n                        \"trace_alloc_max_entries\": trace_alloc_max_entries,\n                        \"record_context\": True,\n                    }\n                else:\n                    f(enabled=True)\n                    return \"legacy-min\", {\"enabled\": True}\n            except Exception:\n                raise\n\n    if devices is None or isinstance(devices, str | int | torch.device):\n        mode, used = _one_call(devices if devices is not None else None)\n    else:\n        mode, used = \"multi-device\", {}\n        for d in list(devices):\n            _mode, _used = _one_call(d)\n            used[f\"dev{d}\"] = _used\n\n    device = get_torch_device()\n    if device.is_available():\n        device.reset_peak_memory_stats()\n        device.synchronize()\n\n    rank = int(os.environ.get(\"RANK\", \"0\") or 0)\n    logger.info(f\"[memory_visualize][rank {rank}] recording enabled ({mode}); args={used}\")\n\n\nclass MemorySnapshotSampler:\n    \"\"\"\n    A utility class that dumps GPU memory snapshots.\n    This is useful for monitoring memory usage over a long-running process.\n\n    The dumped files can be visualized with https://docs.pytorch.org/memory_viz\n\n    Args:\n        out_dir (str): The directory where the snapshots will be saved.\n        tag (str): A tag for the snapshot filenames.\n    \"\"\"\n\n    def __init__(self, out_dir: str = \"./mem_snapshots\", tag: str = \"periodic\"):\n        self.out_dir = out_dir\n        self.tag = tag\n\n    def dump_memory_snapshot(self, out_dir: str = \"./mem_snapshots\", tag: str = \"snapshot\", sub_dir: str = None):\n        \"\"\"\n        Generates a memory snapshot and saves it as a pickle file in a specified directory.\n        The files are organized by timestamp in subdirectories, with all ranks' files\n        placed in the same timestamp subdirectory.\n\n        Args:\n            out_dir (str): The directory where the snapshot file will be saved.\n                The directory is created if it does not exist.\n            tag (str): A string tag to prepend to the filename for easier identification.\n            sub_dir (str): A subdirectory to place the snapshot file in.\n        \"\"\"\n        if sub_dir is None:\n            timestamp = datetime.now().strftime(\"%Y%m%d-%H%M\")\n            out_path = Path(out_dir) / timestamp\n        else:\n            out_path = Path(out_dir) / sub_dir\n        out_path.mkdir(parents=True, exist_ok=True)\n\n        # get the GPU rank on the current process\n        rank = os.environ.get(\"RANK\", \"0\")\n        pid = os.getpid()\n        # todo(chenyang): check wether we need to sync all ranks before dump\n        fname = f\"{tag}_rank{rank}_pid{pid}.pickle\"\n        path = out_path / fname\n\n        device = get_torch_device()\n        if not device.is_available():\n            logger.warning(\"[memory_visualize] is only available on CUDA devices.\")\n            return\n        try:\n            device.synchronize()\n            # Memory snapshot is CUDA-specific functionality\n            device.memory._dump_snapshot(str(path))\n            logger.info(f\"[memory_visualize] dumped: {path}\")\n        except Exception as e:\n            logger.info(f\"[memory_visualize][warn] dump failed: {e}\")\n"
  },
  {
    "path": "verl/utils/metric/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .utils import AggregationType, Metric, reduce_metrics\n\n__all__ = [\"reduce_metrics\", \"AggregationType\", \"Metric\"]\n"
  },
  {
    "path": "verl/utils/metric/utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nMetrics utils.\n\"\"\"\n\nfrom enum import Enum\nfrom typing import Any, Optional, Union\n\nimport numpy as np\nimport torch\n\n\ndef reduce_metrics(metrics: dict[str, Union[\"Metric\", list[Any]]]) -> dict[str, Any]:\n    \"\"\"\n    Reduces a dictionary of metric lists by computing the mean, max, or min of each list.\n    The reduce operation is determined by the key name:\n    - If the key contains \"max\", np.max is used\n    - If the key contains \"min\", np.min is used\n    - Otherwise, np.mean is used\n\n    Args:\n        metrics: A dictionary mapping metric names to lists of metric values.\n\n    Returns:\n        A dictionary with the same keys but with each list replaced by its reduced value.\n\n    Example:\n        >>> metrics = {\n        ...     \"loss\": [1.0, 2.0, 3.0],\n        ...     \"accuracy\": [0.8, 0.9, 0.7],\n        ...     \"max_reward\": [5.0, 8.0, 6.0],\n        ...     \"min_error\": [0.1, 0.05, 0.2]\n        ... }\n        >>> reduce_metrics(metrics)\n        {\"loss\": 2.0, \"accuracy\": 0.8, \"max_reward\": 8.0, \"min_error\": 0.05}\n    \"\"\"\n    for key, val in metrics.items():\n        if isinstance(val, Metric):\n            metrics[key] = val.aggregate()\n        elif \"max\" in key:\n            metrics[key] = np.max(val)\n        elif \"min\" in key:\n            metrics[key] = np.min(val)\n        else:\n            metrics[key] = np.mean(val)\n    return metrics\n\n\nclass AggregationType(Enum):\n    MEAN = \"mean\"\n    SUM = \"sum\"\n    MIN = \"min\"\n    MAX = \"max\"\n\n\nNumericType = int, float, torch.Tensor, np.ndarray\nNumeric = int | float | torch.Tensor | np.ndarray\n\n\nclass Metric:\n    \"\"\"\n    A metric aggregator for collecting and aggregating numeric values.\n\n    This class accumulates numeric values (int, float, or scalar tensors) and computes\n    an aggregate statistic based on the specified aggregation type (MEAN, SUM, MIN, or MAX).\n\n    Args:\n        aggregation: The aggregation method to use. Can be a string (\"mean\", \"sum\", \"min\", \"max\")\n            or an AggregationType enum value.\n        value: Optional initial value(s) to add. Can be a single numeric value or a list of values.\n\n    Example:\n        >>> metric = Metric(aggregation=\"mean\", value=1.0)\n        >>> metric.append(2.0)\n        >>> metric.append(3.0)\n        >>> metric.aggregate()\n        2.0\n    \"\"\"\n\n    def __init__(self, aggregation: str | AggregationType, value: Optional[Numeric | list[Numeric]] = None) -> None:\n        if isinstance(aggregation, str):\n            self.aggregation = AggregationType(aggregation)\n        else:\n            self.aggregation = aggregation\n        if not isinstance(self.aggregation, AggregationType):\n            raise ValueError(f\"Unsupported aggregation type: {aggregation}\")\n        self.values = []\n        if value is not None:\n            self.append(value)\n\n    def append(self, value: Union[Numeric, \"Metric\"]) -> None:\n        if isinstance(value, Metric):\n            self.extend(value)\n            return\n        if isinstance(value, torch.Tensor):\n            if value.numel() != 1:\n                raise ValueError(\"Only scalar tensors can be converted to float\")\n            value = value.detach().item()\n        if not isinstance(value, NumericType):\n            raise ValueError(f\"Unsupported value type: {type(value)}\")\n        self.values.append(value)\n\n    def extend(self, values: Union[\"Metric\", list[Numeric]]) -> None:\n        if isinstance(values, Metric):\n            if values.aggregation != self.aggregation:\n                raise ValueError(f\"Aggregation type mismatch: {self.aggregation} != {values.aggregation}\")\n            values = values.values\n        for value in values:\n            self.append(value)\n\n    def aggregate(self) -> float:\n        return self._aggregate(self.values, self.aggregation)\n\n    @classmethod\n    def _aggregate(cls, values: list[Numeric], aggregation: AggregationType) -> float:\n        match aggregation:\n            case AggregationType.MEAN:\n                return np.mean(values)\n            case AggregationType.SUM:\n                return np.sum(values)\n            case AggregationType.MIN:\n                return np.min(values)\n            case AggregationType.MAX:\n                return np.max(values)\n\n    @classmethod\n    def aggregate_dp(cls, metric_lists: list[\"Metric\"]) -> float:\n        if not metric_lists:\n            raise ValueError(\"Cannot aggregate an empty list of metrics.\")\n        value_lists = [ml.values for ml in metric_lists]\n        if not all(len(ls) == len(value_lists[0]) for ls in value_lists):\n            raise ValueError(\n                f\"All Metric instances must have the same number of values \"\n                f\"for dp aggregation: {[len(ls) for ls in value_lists]}\"\n            )\n        value_arrays = np.array(value_lists)  # [num_dp, num_grad_accumulation]\n        aggregation = metric_lists[0].aggregation\n        match aggregation:\n            case AggregationType.SUM | AggregationType.MEAN:\n                return cls._aggregate(\n                    values=np.mean(value_arrays, axis=0), aggregation=aggregation\n                )  # mean over dp ranks\n            case AggregationType.MIN | AggregationType.MAX:\n                return cls._aggregate(values=value_arrays.flatten(), aggregation=aggregation)  # min/max over all values\n\n    @classmethod\n    def from_dict(cls, data: dict[str, Numeric], aggregation: str | AggregationType) -> dict[str, \"Metric\"]:\n        return {key: cls(value=value, aggregation=aggregation) for key, value in data.items()}\n\n    def init_list(self) -> \"Metric\":\n        return Metric(aggregation=self.aggregation)\n"
  },
  {
    "path": "verl/utils/model.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUtilities to create common models from huggingface\n\"\"\"\n\nimport json\nimport os\nimport re\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nimport numpy as np\nimport torch\nfrom tensordict.tensorclass import NonTensorData\nfrom torch import nn\nfrom transformers import (\n    AutoConfig,\n    AutoModel,\n    AutoModelForCausalLM,\n    AutoModelForSequenceClassification,\n    AutoModelForTokenClassification,\n    GenerationConfig,\n    MistralForSequenceClassification,\n    PretrainedConfig,\n    PreTrainedModel,\n)\n\ntry:\n    from transformers import AutoModelForVision2Seq\nexcept ImportError:\n    AutoModelForVision2Seq = None\n\ntry:\n    from transformers import AutoModelForImageTextToText\nexcept ImportError:\n    AutoModelForImageTextToText = AutoModelForVision2Seq\n\nfrom transformers.modeling_outputs import CausalLMOutputWithPast\n\nfrom verl.models.registry import ModelRegistry\nfrom verl.utils.import_utils import is_trl_available\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\n\nAutoModelForVision2Seq = get_auto_model_for_vision2seq()\n\n\nclass LambdaLayer(nn.Module):\n    def __init__(self, fn):\n        super().__init__()\n        self.fn = fn\n\n    def forward(self, *args, **kwargs):\n        return self.fn(*args, **kwargs)\n\n\ndef squeeze(x):\n    return torch.squeeze(x, dim=-1)\n\n\ndef update_model_config(module_config, override_config_kwargs):\n    \"\"\"Update the module config with the override_config_kwargs.\n    Args:\n        module_config: The module config from Huggingface Transformers.\n        override_config_kwargs: The kwargs to override the module config.\n    \"\"\"\n    for key, val in override_config_kwargs.items():\n        if isinstance(val, dict):\n            update_model_config(getattr(module_config, key), val)\n        else:\n            setattr(module_config, key, val)\n\n\ndef get_huggingface_actor_config(model_name: str, override_config_kwargs=None, trust_remote_code=False) -> dict:\n    if override_config_kwargs is None:\n        override_config_kwargs = {}\n    assert isinstance(override_config_kwargs, dict), (\n        f\"override_config_kwargs must be a dict, got {type(override_config_kwargs)}\"\n    )\n    module_config = AutoConfig.from_pretrained(model_name, trust_remote_code=trust_remote_code)\n    update_model_config(module_config, override_config_kwargs)\n\n    return module_config\n\n\ndef get_generation_config(\n    model: str,\n    trust_remote_code: bool = False,\n) -> Optional[GenerationConfig]:\n    try:\n        return GenerationConfig.from_pretrained(model)\n    except OSError:  # Not found\n        try:\n            config = get_huggingface_actor_config(\n                model,\n                trust_remote_code=trust_remote_code,\n            )\n            return GenerationConfig.from_model_config(config)\n        except OSError:  # Not found\n            return None\n\n\ndef create_huggingface_actor(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module:\n    \"\"\"\n\n    Args:\n        model_name:\n        override_config_kwargs:\n\n    Returns:\n\n    \"\"\"\n    if override_config_kwargs is None:\n        override_config_kwargs = {}\n    if automodel_kwargs is None:\n        automodel_kwargs = {}\n    assert isinstance(override_config_kwargs, dict), (\n        f\"override_config_kwargs must be a dict, got {type(override_config_kwargs)}\"\n    )\n    module_config = get_huggingface_actor_config(\n        model_name, override_config_kwargs, trust_remote_code=automodel_kwargs.get(\"trust_remote_code\", False)\n    )\n    module: nn.Module = AutoModelForCausalLM.from_config(module_config, **automodel_kwargs)\n    return module\n\n\ndef create_huggingface_critic(model_name: str, override_config_kwargs=None, automodel_kwargs=None) -> nn.Module:\n    \"\"\"\n\n    Args:\n        model_name:\n        override_config_kwargs:\n\n    Returns:\n\n    \"\"\"\n    critic_module: nn.Module = create_huggingface_actor(\n        model_name, override_config_kwargs=override_config_kwargs, automodel_kwargs=automodel_kwargs\n    )\n    if automodel_kwargs is None:\n        automodel_kwargs = {}\n    torch_dtype = automodel_kwargs.get(\"torch_dtype\", torch.float32)\n    critic_module.lm_head = nn.Sequential(\n        nn.Linear(critic_module.config.hidden_size, 1, dtype=torch_dtype), LambdaLayer(fn=squeeze)\n    )\n    return critic_module\n\n\ndef get_model_size(model: nn.Module, scale=\"auto\"):\n    n_params = sum(p.numel() for p in model.parameters())\n\n    if scale == \"auto\":\n        if n_params > 1e9:\n            scale = \"B\"\n        elif n_params > 1e6:\n            scale = \"M\"\n        elif n_params > 1e3:\n            scale = \"K\"\n        else:\n            scale = \"\"\n\n    if scale == \"B\":\n        n_params = n_params / 1e9\n    elif scale == \"M\":\n        n_params = n_params / 1e6\n    elif scale == \"K\":\n        n_params = n_params / 1e3\n    elif scale == \"\":\n        pass\n    else:\n        raise NotImplementedError(f\"Unknown scale {scale}\")\n\n    return n_params, scale\n\n\ndef print_model_size(model: nn.Module, name: str = None):\n    n_params, scale = get_model_size(model, scale=\"auto\")\n    if name is None:\n        name = model.__class__.__name__\n    print(f\"{name} contains {n_params:.2f}{scale} parameters\")\n\n\ndef create_random_mask(\n    input_ids: torch.Tensor,\n    max_ratio_of_valid_token: float,\n    max_ratio_of_left_padding: float,\n    min_ratio_of_valid_token: float = 0,\n):\n    \"\"\"Create a random mask given input_ids. Support left padding and right padding.\n    Process:\n    - Sample valid token length\n    - Sample left_padding length\n    - Generate padding\n\n    Args:\n        input_ids:\n            shape (batch_size, seq_len)\n\n    Returns:\n\n    \"\"\"\n    assert max_ratio_of_valid_token > 0 and max_ratio_of_valid_token <= 1.0\n    assert max_ratio_of_left_padding >= 0 and max_ratio_of_left_padding < 1.0\n    assert min_ratio_of_valid_token <= max_ratio_of_valid_token\n\n    batch_size, sequence_length = input_ids.shape\n    max_num_valid_tokens = int(sequence_length * max_ratio_of_valid_token)\n    min_num_valid_tokens = max(1, int(sequence_length * min_ratio_of_valid_token))\n    max_left_padding = int(sequence_length * max_ratio_of_left_padding)\n    assert max_num_valid_tokens + max_left_padding <= sequence_length\n    assert max_num_valid_tokens > 0 and max_ratio_of_valid_token <= sequence_length\n    masks = torch.ones_like(input_ids, dtype=torch.int64)\n    # TODO: we can make this faster\n    for i in range(batch_size):\n        num_left_padding = np.random.randint(low=0, high=max_left_padding + 1, dtype=np.int64)\n        num_valid = np.random.randint(low=min_num_valid_tokens, high=max_num_valid_tokens + 1, dtype=np.int64)\n\n        for index in range(num_left_padding):\n            masks[i, index] = 0\n\n        for index in range(num_left_padding + num_valid, sequence_length):\n            masks[i, index] = 0\n    return masks\n\n\ndef compute_position_id_with_mask(mask):\n    return torch.clip(torch.cumsum(mask, dim=-1) - 1, min=0, max=None)\n\n\ndef convert_weight_keys(state_dict: dict[str, torch.Tensor], model: PreTrainedModel):\n    # convert state dict keys: https://github.com/huggingface/transformers/pull/38385\n    if not hasattr(model, \"_checkpoint_conversion_mapping\"):\n        return state_dict\n\n    reverse_key_mapping = {v: k for k, v in model._checkpoint_conversion_mapping.items()}\n    original_weights = {}\n    for key, value in state_dict.items():\n        for pattern, replacement in reverse_key_mapping.items():\n            replacement = replacement.lstrip(\"^\")  # strip off un-needed chars and patterns\n            replacement = re.sub(r\"\\(.*\\)\", \"\", replacement)\n            key, n_replace = re.subn(pattern, replacement, key)\n            # Early exit of the loop\n            if n_replace > 0:\n                break\n\n        original_weights[key] = value\n\n    return original_weights\n\n\ndef check_exclude_modules(config, key: str) -> bool:\n    \"\"\"\n    A helper method to check if the passed module's key name matches any of the exclude modules in the adapter_config.\n    Adapted from https://github.com/huggingface/peft/blob/main/src/peft/tuners/tuners_utils.py\n\n    Args:\n        config (`LoraConfig` | `LycorisConfig`): A config to match exclude modules from\n        key (`str`): A key to search any matches in config\n\n    Returns:\n        True of match object if key matches any exclude modules from config, False if no match found\n    \"\"\"\n    if hasattr(config, \"exclude_modules\") and config.exclude_modules:\n        if isinstance(config.exclude_modules, str):\n            if re.fullmatch(config.exclude_modules, key):\n                return True\n        elif key in config.exclude_modules:\n            return True\n        elif any(key.endswith(f\".{exclude_key}\") for exclude_key in config.exclude_modules):\n            return True\n    return False\n\n\ndef check_target_modules(config, key: str) -> bool:\n    \"\"\"\n    A helper method to check if the passed module's key name matches any of the target modules in the adapter_config.\n    Adapted from https://github.com/huggingface/peft/blob/main/src/peft/tuners/tuners_utils.py\n\n    Args:\n        config (`LoraConfig` | `LycorisConfig`): A config to match target modules from\n        key (`str`): A key to search any matches in config\n\n    Returns:\n        True of match object if key matches any target modules from config, False if no match found\n    \"\"\"\n    if isinstance(config.target_modules, str):\n        target_module_found = re.fullmatch(config.target_modules, key)\n    elif key in config.target_modules:\n        # this module is specified directly in target_modules\n        target_module_found = True\n    else:\n        target_module_found = any(key.endswith(f\".{target_key}\") for target_key in config.target_modules)\n\n        layer_indexes = getattr(config, \"layers_to_transform\", None)\n        layers_pattern = getattr(config, \"layers_pattern\", None)\n\n        is_using_layer_indexes = layer_indexes is not None and (\n            len(layer_indexes) != 0 if isinstance(layer_indexes, list) else True\n        )\n        if is_using_layer_indexes and target_module_found:\n            layer_index = None\n            # TODO: It's still unclear how empty layers_pattern (None, [], or \"\") should behave\n            # For now, empty layers_pattern means any layer pattern is ok\n            if layers_pattern is None or len(layers_pattern) == 0:\n                layer_index = re.match(r\".*\\.[^.]*\\.(\\d+)\\.\", key)\n            else:\n                layers_pattern = [layers_pattern] if isinstance(layers_pattern, str) else layers_pattern\n                for pattern in layers_pattern:\n                    layer_index = re.match(rf\".*\\.{pattern}\\.(\\d+)\\.\", key)\n                    if layer_index is not None:\n                        break\n\n            if layer_index is None:\n                target_module_found = False\n            else:\n                layer_index = int(layer_index.group(1))\n                if isinstance(layer_indexes, int):\n                    target_module_found = layer_index == layer_indexes\n                else:\n                    target_module_found = layer_index in layer_indexes\n\n    return target_module_found\n\n\ndef normalize_model_name(name, pp_rank, vpp_rank, transformer_config, layer_name=\"layers\"):\n    \"\"\"\n    Transform the model name in each model_chunk in each pp stage into the name in inference engine\n    \"\"\"\n    from verl.utils.megatron_utils import get_transformer_layer_offset\n\n    layer_offset = get_transformer_layer_offset(pp_rank, vpp_rank, transformer_config)\n\n    if layer_name in name:  # belong to an intermediate layer\n        split_name = name.split(\".\")\n        # find the num next to split_name\n        for i, name in enumerate(split_name):\n            if name == layer_name:\n                break\n        layer_num_idx = i + 1\n        # check the name\n        assert len(split_name) >= layer_num_idx + 1, f\"split_name = {split_name}\"\n        assert split_name[layer_num_idx].isdigit(), f\"split_name = {split_name}\"\n        # increment layer_num_idx by layer_offset\n        split_name[layer_num_idx] = str(int(split_name[layer_num_idx]) + layer_offset)\n        name = \".\".join(split_name)  # weight name in inference_tp_model\n    return name\n\n\ndef normalize_pp_vpp_params(params, num_hidden_layers, layer_name=\"layers\"):\n    \"\"\"\n    Normalize the pp vpp params into a complete named parameters.\n    This is useful when gather parameters from pp ranks and passed to a model without pp\n\n    params: Iterable[List[Dict[str, param]]]\n        params contains a list of pp, with a list of vpp named_parameters in each vpp chunk.\n    output: Dict[str, param]\n\n    \"\"\"\n    pp_size = len(params)\n    for pp_rank in range(len(params)):\n        vpp_size = len(params[pp_rank])\n        for vpp_rank in range(vpp_size):\n            for name, param in params[pp_rank][vpp_rank].items():\n                normalized_name = normalize_model_name(\n                    name, pp_rank, vpp_rank, pp_size, vpp_size, num_hidden_layers, layer_name=layer_name\n                )\n                yield normalized_name, param\n\n\ndef get_parallel_model_from_config(\n    config, megatron_config, pre_process=None, post_process=None, share_embeddings_and_output_weights=False, value=False\n):\n    from megatron.core import ModelParallelConfig\n\n    assert isinstance(megatron_config, ModelParallelConfig)\n    model_class = _get_parallel_model_architecture_from_config(config, value)\n\n    model = model_class(\n        config,\n        megatron_config,\n        pre_process=pre_process,\n        post_process=post_process,\n        share_embeddings_and_output_weights=share_embeddings_and_output_weights,\n    )\n    return model\n\n\ndef _get_parallel_model_architecture_from_config(config: PretrainedConfig, value=False) -> type[nn.Module]:\n    architectures = getattr(config, \"architectures\", [])\n    for arch in architectures:\n        model_cls = ModelRegistry.load_model_cls(arch, value)\n        print(\"after load model cls\")\n        if model_cls is not None:\n            return model_cls\n    raise ValueError(\n        f\"Model architectures {architectures} are not supported for now. Supported architectures: \"\n        f\"{ModelRegistry.get_supported_archs()}\"\n    )\n\n\ndef _load_hf_model(config, model_config, is_value_model):\n    \"\"\"Helper function containing the loading hf model logic\"\"\"\n    from accelerate import init_empty_weights\n    from megatron.core import parallel_state as mpu\n\n    from verl.models.mcore.saver import _megatron_calc_global_rank\n\n    assert hasattr(model_config, \"architectures\"), \"architectures cannot be empty when load weight!\"\n    architectures = getattr(model_config, \"architectures\", [])\n\n    # get auto class\n    auto_cls = get_hf_auto_model_class(model_config)\n\n    if config.model.path.startswith(\"hdfs:\"):\n        from verl.utils.fs import copy_to_local\n\n        print(f\"start download from {config.model.path}\")\n        local_model_path = copy_to_local(src=config.model.path, use_shm=config.model.get(\"use_shm\", False))\n        print(\"finish download\")\n    else:\n        local_model_path = config.model.path\n        print(f\"load from local dir {local_model_path}\")\n\n    src_rank = _megatron_calc_global_rank(tp_rank=0, dp_rank=0, pp_rank=0, cp_rank=mpu.get_context_parallel_rank())\n    cpu_init_weights = lambda: torch.device(\"cpu\")\n    init_context = init_empty_weights if torch.distributed.get_rank() != src_rank else cpu_init_weights\n    with init_context(), warnings.catch_warnings():\n        warnings.simplefilter(\"ignore\")\n        # TODO: to find a better way to load mistral7b-rm lm_head\n        if \"mistral7b-rm\" in config.model.path:\n            model = MistralForSequenceClassification.from_pretrained(\n                local_model_path,\n                torch_dtype=\"auto\",\n                # device_map=\"auto\",  # disable auto device_map, the HF weight is only loaded to CPU in src_rank\n                # low_cpu_mem_usage=True\n            )  # use score head instead of lm_head\n            state_dict = model.state_dict()\n            state_dict[\"lm_head.weight\"] = state_dict[\"score.weight\"]\n            state_dict[\"model.embed_tokens.weight\"] = state_dict[\"model.embed_tokens.weight\"][\n                :32000\n            ]  # workaround, 32001 -> 32000\n            is_value_model = True\n        else:\n            model = auto_cls.from_pretrained(\n                local_model_path,\n                torch_dtype=\"auto\",\n                # device_map=\"auto\", # disable auto device_map, the HF weight is only loaded to CPU in src_rank\n                # low_cpu_mem_usage=True\n            )\n            state_dict = model.state_dict()\n\n    return architectures, model, state_dict, is_value_model\n\n\ndef get_hf_model_path(config):\n    if config.model.path.startswith(\"hdfs:\"):\n        from verl.utils.fs import copy_to_local\n\n        local_model_path = copy_to_local(src=config.model.path, use_shm=config.model.get(\"use_shm\", False))\n    else:\n        local_model_path = config.model.path\n    return local_model_path\n\n\ndef load_megatron_model_weights(config, model_config, parallel_model, params_dtype, is_value_model=False):\n    \"\"\"Load weights for verl customized model.\"\"\"\n    architectures, model, state_dict, is_value_model = _load_hf_model(config, model_config, is_value_model)\n\n    from verl.models.weight_loader_registry import get_weight_loader\n\n    print(f\"before weight loader: architectures = {architectures}...\")\n    for arch in architectures:\n        print(f\"call weight loader arch = {arch}, model config = {model.config}\")\n        weight_loader = get_weight_loader(arch)\n        weight_loader(\n            state_dict=state_dict,\n            wrapped_models=parallel_model,\n            config=model.config,\n            params_dtype=params_dtype,\n            is_value_model=is_value_model,\n            tie_word_embeddings=model_config.tie_word_embeddings,\n        )\n    return model.config\n\n\ndef load_megatron_gptmodel_weights(config, model_config, parallel_model, params_dtype, is_value_model=False):\n    \"\"\"Load weights for mcore GPT model.\"\"\"\n    _, model, state_dict, is_value_model = _load_hf_model(config, model_config, is_value_model)\n\n    from verl.models.mcore.loader import load_state_dict_to_megatron_gptmodel\n\n    load_state_dict_to_megatron_gptmodel(\n        state_dict=state_dict,\n        wrapped_models=parallel_model,\n        config=model.config,\n        params_dtype=params_dtype,\n        is_value_model=is_value_model,\n    )\n    del state_dict, model\n\n\n# pad input_ids_rmpad, cu_seqlens and max_seqlen_in_batch to be divisible by tp\ndef pad_packed_inputs(unpad_tokens: torch.Tensor, cu_seqlens, max_seqlen_in_batch, size):\n    \"\"\"pad the tokens such that the total length is a multiple of size.\n    This function is useful when applying sequence parallel and context parallel\n\n    Args:\n        unpad_tokens: (total_nnz, ...). Tokens after removing padding\n        cu_seqlens: (total_nnz + 1,)\n        max_seqlen_in_batch: int\n\n    Returns:\n\n    \"\"\"\n    F = nn.functional\n\n    total_nnz = unpad_tokens.shape[0]\n\n    pad_size = 0 if total_nnz % size == 0 else size - total_nnz % size\n\n    # we assume adding a new data in the batch with seqlen pad_size\n    if pad_size > 0:\n        if unpad_tokens.ndim == 1:\n            unpad_tokens = F.pad(unpad_tokens, (0, pad_size))\n        elif unpad_tokens.ndim == 2:\n            unpad_tokens = F.pad(unpad_tokens, (0, 0, 0, pad_size))\n        else:\n            raise NotImplementedError(f\"Padding dim {unpad_tokens.ndim()} is not supported\")\n\n        cu_seqlens = F.pad(cu_seqlens, (0, 1), value=pad_size + cu_seqlens[-1])\n        max_seqlen_in_batch = max(max_seqlen_in_batch, pad_size)\n\n    return unpad_tokens, cu_seqlens, max_seqlen_in_batch\n\n\ndef load_mcore_dist_weights(parallel_model, dist_weight_path, is_value_model=False, prefix=\"\"):\n    from megatron.core import dist_checkpointing\n    from megatron.core.dist_checkpointing.serialization import StrictHandling\n\n    from verl.utils.megatron_utils import unwrap_model\n\n    # strict = StrictHandling.IGNORE_ALL if is_value_model else StrictHandling.ASSUME_OK_UNEXPECTED\n    strict = StrictHandling.ASSUME_OK_UNEXPECTED\n    for model in parallel_model:\n        ssd = unwrap_model(model).sharded_state_dict(prefix=prefix)\n        if is_value_model:\n            for k in list(ssd.keys()):\n                if \"output_layer\" in k:\n                    ssd.pop(k)\n        dist_checkpointing.load(ssd, dist_weight_path, strict=strict)\n\n    return\n\n\ndef get_parallel_gptmodel_from_config(\n    tfconfig, hf_config, pre_process=None, post_process=None, share_embeddings_and_output_weights=False, value=False\n):\n    from megatron.core.models.gpt.gpt_layer_specs import get_gpt_decoder_block_spec\n    from megatron.core.models.gpt.gpt_model import GPTModel\n\n    use_te = True\n    assert tfconfig.normalization == \"RMSNorm\", \"only RMSNorm is supported for now\"\n    transformer_layer_spec = get_gpt_decoder_block_spec(tfconfig, use_transformer_engine=use_te)\n    rope_scaling_args = {}\n    if hf_config.rope_scaling is not None:\n        assert hf_config.rope_scaling[\"type\"] == \"linear\", \"only linear scaling is supported for now\"\n        rope_scaling_args[\"seq_len_interpolation_factor\"] = hf_config.rope_scaling[\"factor\"]\n    parallel_model = GPTModel(\n        config=tfconfig,\n        transformer_layer_spec=transformer_layer_spec,\n        vocab_size=hf_config.vocab_size,\n        max_sequence_length=hf_config.max_position_embeddings,\n        pre_process=pre_process,\n        post_process=post_process,\n        share_embeddings_and_output_weights=share_embeddings_and_output_weights,\n        position_embedding_type=\"rope\",\n        rotary_base=hf_config.rope_theta,\n        **rope_scaling_args,\n    )\n    # # for layer in parallel_model.decoder.layers:\n    # layer.self_attention.core_attention.flash_attention.softmax_scale = None\n    if post_process and value:\n        from verl.models.llama.megatron.layers.parallel_linear import LinearForLastLayer\n\n        parallel_model.output_layer = LinearForLastLayer(\n            input_size=tfconfig.hidden_size, output_size=1, config=tfconfig\n        )\n    return parallel_model\n\n\ndef patch_valuehead_model(model) -> None:\n    from types import MethodType\n\n    from transformers import PreTrainedModel\n    from trl import AutoModelForCausalLMWithValueHead\n\n    def tie_weights(self: \"AutoModelForCausalLMWithValueHead\") -> None:\n        if isinstance(self.pretrained_model, PreTrainedModel):\n            self.pretrained_model.tie_weights()\n\n    def get_input_embeddings(self: \"AutoModelForCausalLMWithValueHead\") -> torch.nn.Module:\n        if isinstance(self.pretrained_model, PreTrainedModel):\n            return self.pretrained_model.get_input_embeddings()\n\n    def get_output_embeddings(self: \"AutoModelForCausalLMWithValueHead\") -> torch.nn.Module:\n        if isinstance(self.pretrained_model, PreTrainedModel):\n            return self.pretrained_model.get_output_embeddings()\n\n    def can_generate(self):\n        return False\n\n    ignore_modules = [name for name, _ in model.named_parameters() if \"pretrained_model\" in name]\n    model._keys_to_ignore_on_save = ignore_modules\n    model.tie_weights = MethodType(tie_weights, model)\n    model.get_input_embeddings = MethodType(get_input_embeddings, model)\n    model.get_output_embeddings = MethodType(get_output_embeddings, model)\n    model.can_generate = MethodType(can_generate, model)\n    model._no_split_modules = getattr(model.pretrained_model, \"_no_split_modules\", [])\n\n\ndef load_valuehead_model(local_path, torch_dtype, model_config, trust_remote_code):\n    from transformers import AutoModelForCausalLM, AutoModelForTokenClassification\n\n    try:\n        model = AutoModelForTokenClassification.from_pretrained(\n            pretrained_model_name_or_path=local_path,\n            torch_dtype=torch_dtype,\n            config=model_config,\n            attn_implementation=\"flash_attention_2\",\n            trust_remote_code=trust_remote_code,\n        )\n        return model\n    except BaseException as e:\n        if not is_trl_available():\n            raise RuntimeError(\n                f\"model({local_path}) is not a value head model, please install trl to make it valid\"\n            ) from e\n\n    assert is_trl_available()\n\n    from trl import AutoModelForCausalLMWithValueHead\n\n    if type(model_config) in AutoModelForVision2Seq._model_mapping.keys():\n        module_class = AutoModelForVision2Seq\n    else:\n        module_class = AutoModelForCausalLM\n    ori_model = module_class.from_pretrained(\n        pretrained_model_name_or_path=local_path,\n        torch_dtype=torch_dtype,\n        config=model_config,\n        attn_implementation=\"flash_attention_2\",\n        trust_remote_code=trust_remote_code,\n    )\n    model = AutoModelForCausalLMWithValueHead.from_pretrained(ori_model)\n    patch_valuehead_model(model)\n    return model\n\n\n_architecture_to_auto_class = {\n    \"ForCausalLM\": AutoModelForCausalLM,\n    \"ForVision2Seq\": AutoModelForVision2Seq,\n    \"ForTokenClassification\": AutoModelForTokenClassification,\n    \"ForSequenceClassification\": AutoModelForSequenceClassification,\n}\n\n\ndef get_hf_auto_model_class(hf_config):\n    has_remote_code = hasattr(hf_config, \"auto_map\") and any(\n        hf_config.architectures[0] in val for val in hf_config.auto_map.values()\n    )\n    if has_remote_code:\n        auto_class = next(k for k, v in hf_config.auto_map.items() if hf_config.architectures[0] in v)\n        match auto_class:\n            case \"AutoModelForVision2Seq\":\n                actor_module_class = AutoModelForVision2Seq\n            case \"AutoModelForCausalLM\":\n                actor_module_class = AutoModelForCausalLM\n            case \"AutoModelForImageTextToText\":\n                actor_module_class = AutoModelForImageTextToText\n            case _:\n                actor_module_class = AutoModel\n    else:\n        actor_module_class = AutoModel\n        # For VLM models, we use type to check instead of architecture\n        if type(hf_config) in AutoModelForImageTextToText._model_mapping.keys():\n            actor_module_class = AutoModelForImageTextToText\n        else:\n            for key, cls in _architecture_to_auto_class.items():\n                if key in hf_config.architectures[0]:\n                    actor_module_class = cls\n                    break\n\n    return actor_module_class\n\n\ndef extract_multi_modal_inputs(\n    batch_data: list[dict[str, torch.Tensor]],\n    indices: Optional[list[int]] = None,\n) -> dict[str, torch.Tensor | list[torch.Tensor]]:\n    \"\"\"\n    Extract and process multi-modal inputs from a batch.\n\n    Args:\n        batch_data (list[dict[str, torch.Tensor]]): The batch containing potential multi-modal inputs\n        indices (Optional[list[int]]): If provided, only extract inputs at these indices\n\n    Returns:\n        dict[str, torch.Tensor | list[torch.Tensor]]: Processed multi-modal inputs ready for model consumption\n\n    \"\"\"\n    multi_modal_inputs = {}\n    multi_modal_inputs_collected = {}\n    has_image_bound = False\n\n    selected_batch_data = batch_data\n    if indices is not None:\n        selected_batch_data = [batch_data[i] for i in indices if i < len(batch_data)]\n\n    for inputs in selected_batch_data:\n        inputs = inputs.data if isinstance(inputs, NonTensorData) else inputs\n        # Mixed pure text and multi-modal dataset.\n        if inputs is None:\n            continue\n        if \"image_bound\" in inputs:\n            has_image_bound = True\n        for key, value in inputs.items():\n            if value is not None:\n                if key not in multi_modal_inputs_collected:\n                    multi_modal_inputs_collected[key] = []\n                multi_modal_inputs_collected[key].append(value)\n\n    for key, values in multi_modal_inputs_collected.items():\n        if has_image_bound:  # minicpm-o logic\n            multi_modal_inputs[key] = values\n        else:\n            multi_modal_inputs[key] = torch.cat(values, dim=0)\n\n    return multi_modal_inputs\n\n\ndef get_lora_rank_from_adapter(adapter_path: str | os.PathLike) -> int:\n    \"\"\"\n    Extract LoRA rank from adapter configuration file.\n\n    Args:\n        adapter_path: Path to LoRA adapter directory\n\n    Returns:\n        LoRA rank value from adapter_config.json\n\n    Raises:\n        FileNotFoundError: If adapter path or config file doesn't exist\n        ValueError: If config file is invalid or missing rank\n    \"\"\"\n    adapter_path = os.path.abspath(os.path.expanduser(str(adapter_path)))\n\n    if not os.path.exists(adapter_path):\n        raise FileNotFoundError(f\"LoRA adapter path not found: {adapter_path}\")\n\n    config_path = os.path.join(adapter_path, \"adapter_config.json\")\n    if not os.path.exists(config_path):\n        raise FileNotFoundError(f\"adapter_config.json not found in {adapter_path}\")\n\n    try:\n        with open(config_path, encoding=\"utf-8\") as f:\n            config = json.load(f)\n            if \"r\" not in config:\n                raise ValueError(f\"LoRA rank 'r' not found in {config_path}\")\n            return int(config[\"r\"])\n    except json.JSONDecodeError as e:\n        raise ValueError(f\"Invalid JSON in {config_path}: {e}\") from e\n    except (KeyError, ValueError) as e:\n        raise ValueError(f\"Cannot parse LoRA rank from {config_path}: {e}\") from e\n\n\n@dataclass\nclass CausalLMOutputForPPO(CausalLMOutputWithPast):\n    log_probs: Optional[torch.FloatTensor] = None\n    entropy: Optional[torch.FloatTensor] = None\n"
  },
  {
    "path": "verl/utils/net_utils.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# ==============================================================================\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport ipaddress\nimport socket\n\n\ndef is_ipv4(ip_str: str) -> bool:\n    \"\"\"\n    Check if the given string is an IPv4 address\n\n    Args:\n        ip_str: The IP address string to check\n\n    Returns:\n        bool: Returns True if it's an IPv4 address, False otherwise\n    \"\"\"\n    try:\n        ipaddress.IPv4Address(ip_str)\n        return True\n    except ipaddress.AddressValueError:\n        return False\n\n\ndef is_ipv6(ip_str: str) -> bool:\n    \"\"\"\n    Check if the given string is an IPv6 address\n\n    Args:\n        ip_str: The IP address string to check\n\n    Returns:\n        bool: Returns True if it's an IPv6 address, False otherwise\n    \"\"\"\n    try:\n        ipaddress.IPv6Address(ip_str)\n        return True\n    except ipaddress.AddressValueError:\n        return False\n\n\ndef is_valid_ipv6_address(address: str) -> bool:\n    try:\n        ipaddress.IPv6Address(address)\n        return True\n    except ValueError:\n        return False\n\n\ndef get_free_port(address: str, with_alive_sock: bool = False) -> tuple[int, socket.socket | None]:\n    \"\"\"Find a free port on the given address.\n\n    By default the socket is closed internally, suitable for immediate use.\n    Set with_alive_sock=True to keep the socket open as a port reservation,\n    preventing other calls from getting the same port. The caller is\n    responsible for closing the socket before the port is actually bound\n    by the target service (e.g. NCCL, uvicorn).\n    \"\"\"\n    family = socket.AF_INET6 if is_valid_ipv6_address(address) else socket.AF_INET\n\n    sock = socket.socket(family=family, type=socket.SOCK_STREAM)\n    sock.setsockopt(socket.SOL_SOCKET, socket.SO_REUSEADDR, 1)\n    sock.bind((address, 0))\n    port = sock.getsockname()[1]\n    if with_alive_sock:\n        return port, sock\n    sock.close()\n    return port, None\n"
  },
  {
    "path": "verl/utils/npu_flash_attn_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport torch\nimport torch.nn.functional as F\nfrom einops import rearrange, repeat\n\n\n# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py\nclass IndexFirstAxis(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, input, indices):\n        ctx.save_for_backward(indices)\n        assert input.ndim >= 2\n        ctx.first_axis_dim, other_shape = input.shape[0], input.shape[1:]\n        second_dim = other_shape.numel()\n        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.\n        # return input[indices]\n        return torch.gather(rearrange(input, \"b ... -> b (...)\"), 0, repeat(indices, \"z -> z d\", d=second_dim)).reshape(\n            -1, *other_shape\n        )\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        (indices,) = ctx.saved_tensors\n        assert grad_output.ndim >= 2\n        other_shape = grad_output.shape[1:]\n        grad_output = rearrange(grad_output, \"b ... -> b (...)\")\n        grad_input = torch.zeros(\n            [ctx.first_axis_dim, grad_output.shape[1]],\n            device=grad_output.device,\n            dtype=grad_output.dtype,\n        )\n        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.\n        # grad_input[indices] = grad_output\n        grad_input.scatter_(0, repeat(indices, \"z -> z d\", d=grad_output.shape[1]), grad_output)\n        return grad_input.reshape(ctx.first_axis_dim, *other_shape), None\n\n\nindex_first_axis = IndexFirstAxis.apply\n\n\n# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py\nclass IndexPutFirstAxis(torch.autograd.Function):\n    @staticmethod\n    def forward(ctx, values, indices, first_axis_dim):\n        ctx.save_for_backward(indices)\n        assert indices.ndim == 1\n        assert values.ndim >= 2\n        output = torch.zeros(first_axis_dim, *values.shape[1:], device=values.device, dtype=values.dtype)\n        # TD [2022-03-04] For some reason torch.scatter is a bit faster than indexing.\n        output[indices] = values\n        # output.scatter_(0, repeat(indices, 'z -> z d', d=values.shape[1]), values)\n        return output\n\n    @staticmethod\n    def backward(ctx, grad_output):\n        (indices,) = ctx.saved_tensors\n        # TD [2022-03-04] For some reason torch.gather is a bit faster than indexing.\n        grad_values = grad_output[indices]\n        # grad_values = torch.gather(grad_output, 0, repeat(indices, 'z -> z d', d=grad_output.shape[1]))\n        return grad_values, None, None\n\n\nindex_put_first_axis = IndexPutFirstAxis.apply\n\n\n# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py\ndef pad_input(hidden_states, indices, batch, seqlen):\n    \"\"\"\n    Arguments:\n        hidden_states: (total_nnz, ...), where total_nnz = number of tokens in selected in attention_mask.\n        indices: (total_nnz), the indices that represent the non-masked tokens of the original padded input sequence.\n        batch: int, batch size for the padded sequence.\n        seqlen: int, maximum sequence length for the padded sequence.\n    Return:\n        hidden_states: (batch, seqlen, ...)\n    \"\"\"\n    # dim = hidden_states.shape[-1]\n    # output = torch.zeros((batch * seqlen), dim, device=hidden_states.device, dtype=hidden_states.dtype)\n    # output[indices] = hidden_states\n    output = index_put_first_axis(hidden_states, indices, batch * seqlen)\n    return rearrange(output, \"(b s) ... -> b s ...\", b=batch)\n\n\n# Copied from https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py\ndef unpad_input(hidden_states, attention_mask, unused_mask=None):\n    \"\"\"\n    Arguments:\n        hidden_states: (batch, seqlen, ...)\n        attention_mask: (batch, seqlen), bool / int, 1 means valid and 0 means not valid.\n        unused_mask: (batch, seqlen), bool / int, 1 means the element is allocated but unused.\n    Return:\n        hidden_states: (total_nnz, ...), where total_nnz = number of tokens selected in attention_mask + unused_mask.\n        indices: (total_nnz), the indices of masked tokens from the flattened input sequence.\n        cu_seqlens: (batch + 1), the cumulative sequence lengths, used to index into hidden_states.\n        max_seqlen_in_batch: int\n        seqused: (batch), returns the number of tokens selected in attention_mask + unused_mask.\n    \"\"\"\n    all_masks = (attention_mask + unused_mask) if unused_mask is not None else attention_mask\n    seqlens_in_batch = all_masks.sum(dim=-1, dtype=torch.int32)\n    used_seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)\n    indices = torch.nonzero(all_masks.flatten(), as_tuple=False).flatten()\n    max_seqlen_in_batch = seqlens_in_batch.max().item()\n    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))\n    # TD [2022-03-04] We don't want to index with a bool mask, because Pytorch will expand the\n    # bool mask, then call nonzero to get the indices, then index with those. The indices is @dim\n    # times larger than it needs to be, wasting memory. It's faster and more memory-efficient to\n    # index with integer indices. Moreover, torch's index is a bit slower than it needs to be,\n    # so we write custom forward and backward to make it a bit faster.\n    return (\n        index_first_axis(rearrange(hidden_states, \"b s ... -> (b s) ...\"), indices),\n        indices,\n        cu_seqlens,\n        max_seqlen_in_batch,\n        used_seqlens_in_batch,\n    )\n"
  },
  {
    "path": "verl/utils/profiler/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom ..device import is_npu_available\nfrom ..import_utils import is_nvtx_available\nfrom .config import build_sglang_profiler_args, build_vllm_profiler_args\nfrom .performance import GPUMemoryLogger, log_gpu_memory_usage, simple_timer\nfrom .profile import DistProfiler, DistProfilerExtension, ProfilerConfig\n\n# Select marker implementations by availability, but keep DistProfiler as our dispatcher\nif is_nvtx_available():\n    from .nvtx_profile import mark_annotate, mark_end_range, mark_start_range, marked_timer\nelif is_npu_available:\n    from .mstx_profile import mark_annotate, mark_end_range, mark_start_range, marked_timer\nelse:\n    from .performance import marked_timer\n    from .profile import mark_annotate, mark_end_range, mark_start_range\n\n__all__ = [\n    \"GPUMemoryLogger\",\n    \"log_gpu_memory_usage\",\n    \"mark_start_range\",\n    \"mark_end_range\",\n    \"mark_annotate\",\n    \"DistProfiler\",\n    \"DistProfilerExtension\",\n    \"ProfilerConfig\",\n    \"simple_timer\",\n    \"marked_timer\",\n    \"build_vllm_profiler_args\",\n    \"build_sglang_profiler_args\",\n]\n"
  },
  {
    "path": "verl/utils/profiler/config.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport os\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nfrom omegaconf import MISSING\n\nfrom verl.base_config import BaseConfig\n\n\n@dataclass\nclass NsightToolConfig(BaseConfig):\n    \"\"\"Nsight tool config.\"\"\"\n\n    \"True for each task has its own database, False for all tasks in one training step share one database.\"\n    discrete: bool = False\n    name: str = \"nsight\"\n\n    def __post_init__(self) -> None:\n        pass\n\n\n@dataclass\nclass TorchProfilerToolConfig(BaseConfig):\n    \"\"\"Torch profiler tool config.\"\"\"\n\n    # options: cuda, cpu, memory, shapes, stack\n    contents: list[str] = field(default_factory=list)\n    discrete: bool = False\n    name: str = \"torch\"\n\n    def __post_init__(self) -> None:\n        \"\"\"config validation logics go here\"\"\"\n        __support_contents = [\"cuda\", \"cpu\", \"memory\", \"shapes\", \"stack\"]\n        for content in self.contents:\n            assert content in __support_contents, (\n                f\"Profiler contents only supports {__support_contents}, but gets {content}\"\n            )\n        assert isinstance(self.contents, list), f\"Profiler contents must be of type list, got {type(self.contents)}\"\n\n\n@dataclass\nclass TorchMemoryToolConfig(BaseConfig):\n    \"\"\"Torch memory profiler tool config.\n\n    Args:\n        trace_alloc_max_entries (int): Maximum number of memory allocation entries to track.\n        stack_depth (int): Stack trace depth for memory allocations.\n    \"\"\"\n\n    trace_alloc_max_entries: int = 100_000\n    stack_depth: int = 32\n    name: str = \"torch_memory\"\n\n    def __post_init__(self) -> None:\n        \"\"\"config validation logics go here\"\"\"\n        assert isinstance(self.trace_alloc_max_entries, int), (\n            f\"trace_alloc_max_entries must be int, got {type(self.trace_alloc_max_entries)}\"\n        )\n        assert isinstance(self.stack_depth, int), f\"stack_depth must be int, got {type(self.stack_depth)}\"\n        assert self.trace_alloc_max_entries > 0, (\n            f\"trace_alloc_max_entries must be positive, got {self.trace_alloc_max_entries}\"\n        )\n        assert self.stack_depth > 0, f\"stack_depth must be positive, got {self.stack_depth}\"\n\n\n@dataclass\nclass NPUToolConfig(NsightToolConfig):\n    \"\"\"NPU profiler too; config.\"\"\"\n\n    # options: npu, cpu, memory, shapes, module, stack\n    contents: list[str] = field(default_factory=list)\n\n    # Collection level, optional values: level_none, level0, level1, level2.\n    level: str = \"level0\"\n\n    # Whether to automatically parse the data.\n    analysis: bool = False\n\n    name: str = \"npu\"\n\n    def __post_init__(self) -> None:\n        \"\"\"config validation logics go here\"\"\"\n        assert isinstance(self.contents, list), f\"Profiler contents must be of type list, got {type(self.contents)}\"\n        assert isinstance(self.level, str), f\"Profiler level must be of type str, got {type(self.level)}\"\n        assert isinstance(self.analysis, bool), f\"Profiler analysis must be of type bool, got {type(self.analysis)}\"\n        for content in self.contents:\n            assert content in [\"npu\", \"cpu\", \"memory\", \"shapes\", \"module\", \"stack\"], (\n                f\"Profiler contents only supports npu, cpu, memory, shapes, module, stack, but gets {content}\"\n            )\n        assert self.level in [\"level_none\", \"level0\", \"level1\", \"level2\"], (\n            f\"Profiler level only supports level0, 1, 2, and level_none, but gets {self.level}\"\n        )\n\n\n@dataclass\nclass ProfilerConfig(BaseConfig):\n    \"\"\"Worker profiler config.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        discrete (bool): True for each task has its own database, False for all tasks in one training step\n          share one database.\n        all_ranks (bool): Whether to profile all ranks.\n        ranks (list[int]): The ranks that will be profiled. Defaults to [].\n        global_tool_config (Any): Global tool configuration for all profiling tools.\n    \"\"\"\n\n    tool: Optional[str] = MISSING\n    enable: bool = False\n    all_ranks: bool = False\n    ranks: list[int] = field(default_factory=list)\n    save_path: Optional[str] = MISSING\n    tool_config: Any = MISSING  # Just a placeholder, will use configs above directly\n    global_tool_config: Optional[Any] = None  # Global tool configuration for all profiling tools\n\n    def union(self, other: \"ProfilerConfig\") -> \"ProfilerConfig\":\n        assert self.tool == other.tool, f\"Cannot union ProfilerConfig with different tools: {self.tool} vs {other.tool}\"\n        return ProfilerConfig(\n            tool=self.tool,\n            enable=self.enable or other.enable,\n            all_ranks=self.all_ranks or other.all_ranks,\n            ranks=list(set(self.ranks or []) | set(other.ranks or [])),\n            save_path=self.save_path,\n            tool_config=self.tool_config,\n            global_tool_config=self.global_tool_config or other.global_tool_config,\n        )\n\n    def intersect(self, other: \"ProfilerConfig\") -> \"ProfilerConfig\":\n        assert self.tool == other.tool, (\n            f\"Cannot intersect ProfilerConfig with different tools: {self.tool} vs {other.tool}\"\n        )\n        return ProfilerConfig(\n            tool=self.tool,\n            enable=self.enable and other.enable,\n            all_ranks=self.all_ranks and other.all_ranks,\n            ranks=list(set(self.ranks or []) & set(other.ranks or [])),\n            save_path=self.save_path,\n            tool_config=self.tool_config,\n            global_tool_config=self.global_tool_config if self.global_tool_config else other.global_tool_config,\n        )\n\n    def __post_init__(self) -> None:\n        \"\"\"config validation logics go here\"\"\"\n        assert isinstance(self.ranks, set | list | tuple), (\n            f\"Profiler ranks must be of type list, got {type(self.ranks)}\"\n        )\n\n\ndef build_vllm_profiler_args(profiler_config: ProfilerConfig, tool_config: BaseConfig, rank: int) -> dict:\n    \"\"\"\n    Build arguments and environment variables for vLLM profiler.\n\n    Acts as an adapter to bridge verl's unified profiler config and vLLM's specific requirements.\n    It sets environment variables for compatibility and constructs arguments for vLLM >= 0.13.0.\n\n    Args:\n        profiler_config (ProfilerConfig): The unified profiler configuration.\n        tool_config (BaseConfig): The tool configuration.\n        rank (int): The rank of the replica.\n\n    Returns:\n        dict: A dictionary of arguments to be passed to vLLM's start_profile method.\n    \"\"\"\n    if not profiler_config or not tool_config or not hasattr(tool_config, \"contents\"):\n        return {}\n\n    contents = tool_config.contents\n    with_stack = True if \"stack\" in contents or \"module\" in contents else False\n    record_shapes = True if \"shapes\" in contents else False\n    with_memory = True if \"memory\" in contents else False\n    save_path = os.path.join(profiler_config.save_path, f\"agent_loop_rollout_replica_{rank}\")\n\n    # vLLM < 0.13.0 supports controlling profiler via environment variables\n    os.environ[\"VLLM_TORCH_PROFILER_DIR\"] = save_path\n    os.environ[\"VLLM_TORCH_PROFILER_WITH_STACK\"] = \"1\" if with_stack else \"0\"\n    os.environ[\"VLLM_TORCH_PROFILER_RECORD_SHAPES\"] = \"1\" if record_shapes else \"0\"\n    os.environ[\"VLLM_TORCH_PROFILER_WITH_PROFILE_MEMORY\"] = \"1\" if with_memory else \"0\"\n\n    # vLLM >= 0.13.0 supports controlling profiler via arguments.\n    # While it maintains backward compatibility with environment variables,\n    # we provide arguments explicitly to align with the new API style.\n    return {\n        \"profiler_config\": json.dumps(\n            {\n                \"profiler\": \"torch\",\n                \"torch_profiler_dir\": save_path,\n                \"torch_profiler_with_memory\": with_memory,\n                \"torch_profiler_with_stack\": with_stack,\n                \"torch_profiler_record_shapes\": record_shapes,\n            }\n        )\n    }\n\n\ndef build_sglang_profiler_args(profiler_config: ProfilerConfig, tool_config: BaseConfig, rank: int) -> dict:\n    \"\"\"\n    Build arguments for SGLang profiler.\n\n    Args:\n        profiler_config (ProfilerConfig): The unified profiler configuration.\n        tool_config (BaseConfig): The tool configuration.\n        rank (int): The rank of the replica.\n\n    Returns:\n        dict: A dictionary of arguments suitable for starting the SGLang profiler.\n    \"\"\"\n    if not profiler_config or not tool_config or not hasattr(tool_config, \"contents\"):\n        return {}\n\n    contents = tool_config.contents\n    if \"memory\" in contents:\n        warnings.warn(\"SGLang profiler does not support memory profiling. Ignoring memory content.\", stacklevel=2)\n\n    return {\n        \"output_dir\": os.path.join(profiler_config.save_path, f\"agent_loop_rollout_replica_{rank}\"),\n        \"with_stack\": \"stack\" in contents or \"module\" in contents,\n        \"record_shapes\": \"shapes\" in contents,\n    }\n"
  },
  {
    "path": "verl/utils/profiler/empty_annotations.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable, Optional\n\n\ndef mark_start_range(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> None:\n    pass\n\n\ndef mark_end_range(range_id: str) -> None:\n    pass\n\n\ndef mark_annotate(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> Callable:\n    def decorator(func):\n        return func\n\n    return decorator\n"
  },
  {
    "path": "verl/utils/profiler/mstx_profile.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Inspired from https://gitee.com/ascend/MindSpeed-RL/blob/master/mindspeed_rl/utils/utils.py\nimport functools\nimport logging\nimport os\nfrom contextlib import contextmanager\nfrom typing import Any, Callable, Optional\n\nimport torch_npu\nfrom packaging import version\nfrom torch_npu.npu import mstx\n\nfrom .config import NPUToolConfig\nfrom .profile import DistProfiler, ProfilerConfig\n\n\ndef mark_start_range(message: Optional[str] = None) -> None:\n    \"\"\"Start a mark range in the profiler.\n\n    Args:\n        message (str, optional):\n            The message to be displayed in the profiler. Defaults to None.\n    \"\"\"\n    return mstx.range_start(message=message)\n\n\ndef mark_end_range(range_id: str) -> None:\n    \"\"\"End a mark range in the profiler.\n\n    Args:\n        range_id (str):\n            The id of the mark range to end.\n    \"\"\"\n    return mstx.range_end(range_id)\n\n\ndef mark_annotate(message: Optional[str] = None) -> Callable:\n    \"\"\"Decorate a function to annotate a mark range along with the function life cycle.\n\n    Args:\n        message (str, optional):\n            The message to be displayed in the profiler. Defaults to None.\n    \"\"\"\n\n    def decorator(func):\n        profile_message = message or func.__name__\n        return mstx.mstx_range(profile_message)(func)\n\n    return decorator\n\n\n@contextmanager\ndef marked_timer(name: str, timing_raw: dict[str, float], *args: Any, **kwargs: Any) -> None:\n    \"\"\"Context manager for timing with MSTX markers.\n\n    This utility function measures the execution time of code within its context,\n    accumulates the timing information, and adds MSTX markers for profiling.\n\n    Args:\n        name (str): The name/identifier for this timing measurement.\n        timing_raw (Dict[str, float]): Dictionary to store timing information.\n\n    Yields:\n        None: This is a context manager that yields control back to the code block.\n    \"\"\"\n    if args:\n        logging.warning(f\"Args are not supported in mstx_profile, but received: {args}\")\n    if kwargs:\n        logging.warning(f\"Kwargs are not supported in mstx_profile, but received: {kwargs}\")\n    mark_range = mark_start_range(message=name)\n    from .performance import _timer\n\n    yield from _timer(name, timing_raw)\n    mark_end_range(mark_range)\n\n\ndef get_npu_profiler(\n    contents: list[str],\n    profile_level: str,\n    profile_save_path: str,\n    analysis: bool,\n    role: Optional[str] = None,\n    profile_step: Optional[str] = None,\n):\n    \"\"\"Generate and return an NPU profiler object.\n\n    Args:\n        contents (list[str]):\n            A list of options to control the collection content,\n            such as npu, cpu, memory, shapes, module, stack.\n        profile_level (str):\n            The collection level, which can be set to level_none,\n            level0, level1 and level2.\n        profile_save_path (str):\n            The path to save the collected data.\n        analysis (bool):\n            Whether to enables automatic data parsing.\n        role (str, optional):\n            The role of the current data collection. Defaults to None.\n        profile_step(str, optional):\n            The current training step. Defaults to None.\n    \"\"\"\n    if profile_level == \"level_none\":\n        level = torch_npu.profiler.ProfilerLevel.Level_none\n    elif profile_level == \"level0\":\n        level = torch_npu.profiler.ProfilerLevel.Level0\n    elif profile_level == \"level1\":\n        level = torch_npu.profiler.ProfilerLevel.Level1\n    elif profile_level == \"level2\":\n        level = torch_npu.profiler.ProfilerLevel.Level2\n    else:\n        raise ValueError(f\"level only supports level0, 1, 2, and level_none, but gets {profile_level}\")\n\n    if profile_step:\n        profile_save_path = os.path.join(profile_save_path, profile_step)\n    if role:\n        profile_save_path = os.path.join(profile_save_path, role)\n\n    # The ability to filter communication via mstx_domain_exclude requires torch_npu==2.1 or higher.\n    if version.parse(torch_npu.__version__) < version.parse(\"2.1\"):\n        raise RuntimeError(\"torch_npu==2.1 or higher is required to use mstx_domain_exclude\")\n\n    experimental_config = torch_npu.profiler._ExperimentalConfig(\n        profiler_level=level,\n        export_type=torch_npu.profiler.ExportType.Db,\n        data_simplification=True,\n        msprof_tx=True,\n        mstx_domain_exclude=[\"communication\"],\n    )\n\n    activites = []\n    if contents is None or \"npu\" in contents:\n        activites.append(torch_npu.profiler.ProfilerActivity.NPU)\n    if contents is None or \"cpu\" in contents:\n        activites.append(torch_npu.profiler.ProfilerActivity.CPU)\n\n    prof = torch_npu.profiler.profile(\n        with_modules=contents is None or \"module\" in contents,\n        with_stack=contents is None or \"stack\" in contents,\n        record_shapes=contents is None or \"shapes\" in contents,\n        profile_memory=contents is None or \"memory\" in contents,\n        activities=activites,\n        on_trace_ready=torch_npu.profiler.tensorboard_trace_handler(profile_save_path, analyse_flag=analysis),\n        experimental_config=experimental_config,\n    )\n    return prof\n\n\nclass NPUProfiler(DistProfiler):\n    \"\"\"\n    NPU profiler. Initialized in a worker to control the NPU profiler.\n    \"\"\"\n\n    _define_count = 0\n\n    def __init__(self, rank: int, config: ProfilerConfig, tool_config: NPUToolConfig, **kwargs):\n        \"\"\"Initialize the NsightSystemsProfiler.\n\n        Args:\n            rank (int): The rank of the current process.\n            config (Optional[ProfilerConfig]): Configuration for the profiler. If None, a default configuration is used.\n            tool_config (NPUToolConfig): The config to control npu profiler behavior.\n        \"\"\"\n        if not config:\n            config = ProfilerConfig(ranks=[], enable=False)\n        if not tool_config:\n            assert not config.enable, \"tool_config must be set when profiler is enabled\"\n        self.discrete: bool = tool_config.discrete\n        self.profile_npu = None\n        self.profile_contents = tool_config.contents\n        self.profile_level = tool_config.level\n        self.profile_save_path = config.save_path\n        self.analysis = tool_config.analysis\n\n    def start(self, **kwargs):\n        role = kwargs.get(\"role\", None)\n        if not self.discrete and NPUProfiler._define_count == 0:\n            self.profile_npu = get_npu_profiler(\n                contents=self.profile_contents,\n                profile_level=self.profile_level,\n                profile_save_path=self.profile_save_path,\n                analysis=self.analysis,\n                role=role,\n            )\n            self.profile_npu.start()\n            NPUProfiler._define_count += 1\n\n    def stop(self):\n        if not self.discrete and NPUProfiler._define_count == 1:\n            self.profile_npu.step()\n            self.profile_npu.stop()\n            NPUProfiler._define_count -= 1\n\n    def annotate(self, message: Optional[str] = None, role: Optional[str] = None, **kwargs_outer) -> Callable:\n        \"\"\"Decorate a Worker member function to profile the current rank in the current training step.\n\n        Requires the target function to be a member function of a Worker,\n        which has a member field `profiler` with NPUProfiler type.\n\n        Args:\n            message (str, optional):\n                The message to be displayed in the profiler. Defaults to None.\n            role (str, optional):\n                The role of the current data collection. Defaults to None.\n        \"\"\"\n\n        def decorator(func):\n            @functools.wraps(func)\n            def wrapper(*args, **kwargs_inner):\n                profile_name = message or func.__name__\n                discrete_mode = self.discrete\n\n                if not discrete_mode:\n                    mark_range = mark_start_range(message=profile_name)\n                else:\n                    profile_npu = get_npu_profiler(\n                        contents=self.profile_contents,\n                        profile_level=self.profile_level,\n                        profile_save_path=self.profile_save_path,\n                        analysis=self.analysis,\n                        role=role,\n                    )\n                    profile_npu.start()\n                    mark_range = mark_start_range(message=profile_name)\n\n                result = func(*args, **kwargs_inner)\n\n                if not discrete_mode:\n                    mark_end_range(mark_range)\n                else:\n                    mark_end_range(mark_range)\n                    profile_npu.step()\n                    profile_npu.stop()\n\n                return result\n\n            return wrapper\n\n        return decorator\n"
  },
  {
    "path": "verl/utils/profiler/nvtx_profile.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2024, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nfrom contextlib import contextmanager\nfrom typing import Callable, Optional\n\nimport nvtx\nimport torch\n\nfrom .config import NsightToolConfig\nfrom .profile import DistProfiler, ProfilerConfig\n\n\ndef mark_start_range(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> None:\n    \"\"\"Start a mark range in the profiler.\n\n    Args:\n        message (str, optional):\n            The message to be displayed in the profiler. Defaults to None.\n        color (str, optional):\n            The color of the range. Defaults to None.\n        domain (str, optional):\n            The domain of the range. Defaults to None.\n        category (str, optional):\n            The category of the range. Defaults to None.\n    \"\"\"\n    return nvtx.start_range(message=message, color=color, domain=domain, category=category)\n\n\ndef mark_end_range(range_id: str) -> None:\n    \"\"\"End a mark range in the profiler.\n\n    Args:\n        range_id (str):\n            The id of the mark range to end.\n    \"\"\"\n    return nvtx.end_range(range_id)\n\n\ndef mark_annotate(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> Callable:\n    \"\"\"Decorate a function to annotate a mark range along with the function life cycle.\n\n    Args:\n        message (str, optional):\n            The message to be displayed in the profiler. Defaults to None.\n        color (str, optional):\n            The color of the range. Defaults to None.\n        domain (str, optional):\n            The domain of the range. Defaults to None.\n        category (str, optional):\n            The category of the range. Defaults to None.\n    \"\"\"\n\n    def decorator(func):\n        profile_message = message or func.__name__\n        return nvtx.annotate(profile_message, color=color, domain=domain, category=category)(func)\n\n    return decorator\n\n\n@contextmanager\ndef marked_timer(\n    name: str,\n    timing_raw: dict[str, float],\n    color: str = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n):\n    \"\"\"Context manager for timing with NVTX markers.\n\n    This utility function measures the execution time of code within its context,\n    accumulates the timing information, and adds NVTX markers for profiling.\n\n    Args:\n        name (str): The name/identifier for this timing measurement.\n        timing_raw (Dict[str, float]): Dictionary to store timing information.\n        color (Optional[str]): Color for the NVTX marker. Defaults to None.\n        domain (Optional[str]): Domain for the NVTX marker. Defaults to None.\n        category (Optional[str]): Category for the NVTX marker. Defaults to None.\n\n    Yields:\n        None: This is a context manager that yields control back to the code block.\n    \"\"\"\n    mark_range = mark_start_range(message=name, color=color, domain=domain, category=category)\n    from .performance import _timer\n\n    yield from _timer(name, timing_raw)\n    mark_end_range(mark_range)\n\n\nclass NsightSystemsProfiler(DistProfiler):\n    \"\"\"Nsight system profiler. Installed in a worker to control the Nsight system profiler.\"\"\"\n\n    def __init__(self, rank: int, config: Optional[ProfilerConfig], tool_config: Optional[NsightToolConfig], **kwargs):\n        \"\"\"Initialize the NsightSystemsProfiler.\n\n        Args:\n            rank (int): The rank of the current process.\n            config (Optional[ProfilerConfig]): Configuration for the profiler. If None, a default configuration is used.\n        \"\"\"\n        # If no configuration is provided, create a default ProfilerConfig with an empty list of ranks\n        if not config:\n            config = ProfilerConfig(ranks=[])\n        if not tool_config:\n            assert not config.enable, \"tool_config must be provided when profiler is enabled\"\n        self.discrete: bool = tool_config.discrete\n\n    def start(self, **kwargs):\n        if not self.discrete:\n            torch.cuda.profiler.start()\n\n    def stop(self):\n        if not self.discrete:\n            torch.cuda.profiler.stop()\n\n    def annotate(\n        self,\n        message: Optional[str] = None,\n        color: Optional[str] = None,\n        domain: Optional[str] = None,\n        category: Optional[str] = None,\n        **kwargs_outer,\n    ) -> Callable:\n        \"\"\"Decorate a Worker member function to profile the current rank in the current training step.\n\n        Requires the target function to be a member function of a Worker, which has a member field `profiler` with\n        NightSystemsProfiler type.\n\n        Args:\n            message (str, optional):\n                The message to be displayed in the profiler. Defaults to None.\n            color (str, optional):\n                The color of the range. Defaults to None.\n            domain (str, optional):\n                The domain of the range. Defaults to None.\n            category (str, optional):\n                The category of the range. Defaults to None.\n        \"\"\"\n\n        def decorator(func):\n            @functools.wraps(func)\n            def wrapper(*args, **kwargs_inner):\n                profile_name = message or func.__name__\n\n                if self.discrete:\n                    torch.cuda.profiler.start()\n                mark_range = mark_start_range(message=profile_name, color=color, domain=domain, category=category)\n\n                result = func(*args, **kwargs_inner)\n\n                mark_end_range(mark_range)\n                if self.discrete:\n                    torch.cuda.profiler.stop()\n\n                return result\n\n            return wrapper\n\n        return decorator\n"
  },
  {
    "path": "verl/utils/profiler/performance.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport datetime\nimport inspect\nimport logging\nfrom contextlib import contextmanager\nfrom typing import Any, Optional\n\nimport torch\nimport torch.distributed as dist\nfrom codetiming import Timer\n\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.logger import DecoratorLoggerBase\n\n\ndef _get_current_mem_info(unit: str = \"GB\", precision: int = 2) -> tuple[str]:\n    \"\"\"Get current memory usage.\n\n    Note that CPU device memory info is always 0.\n\n    Args:\n        unit (str, optional): The unit of memory measurement. Defaults to \"GB\".\n        precision (int, optional): The number of decimal places to round memory values. Defaults to 2.\n\n    Returns:\n        tuple[str]: A tuple containing memory allocated, memory reserved, memory used, and memory total\n        in the specified unit.\n    \"\"\"\n    assert unit in [\"GB\", \"MB\", \"KB\"]\n    device = get_torch_device()\n    # torch.cpu.memory_allocated() does not exist\n    if device == torch.cpu:\n        return \"0.00\", \"0.00\", \"0.00\", \"0.00\"\n\n    divisor = 1024**3 if unit == \"GB\" else 1024**2 if unit == \"MB\" else 1024\n    mem_allocated = get_torch_device().memory_allocated()\n    mem_reserved = get_torch_device().memory_reserved()\n    # use get_torch_device().mem_get_info to profile device memory\n    # since vllm's sleep mode works below pytorch\n    # see https://github.com/vllm-project/vllm/pull/11743#issuecomment-2754338119\n    mem_free, mem_total = get_torch_device().mem_get_info()\n    mem_used = mem_total - mem_free\n    mem_allocated = f\"{mem_allocated / divisor:.{precision}f}\"\n    mem_reserved = f\"{mem_reserved / divisor:.{precision}f}\"\n    mem_used = f\"{mem_used / divisor:.{precision}f}\"\n    mem_total = f\"{mem_total / divisor:.{precision}f}\"\n    return mem_allocated, mem_reserved, mem_used, mem_total\n\n\ndef log_gpu_memory_usage(head: str, logger: logging.Logger = None, level=logging.DEBUG, rank: int = 0):\n    \"\"\"Log GPU memory usage information.\n\n    Args:\n        head (str): A descriptive header for the memory usage log message.\n        logger (logging.Logger, optional): Logger instance to use for logging. If None, prints to stdout.\n        level: Logging level to use. Defaults to logging.DEBUG.\n        rank (int): The rank of the process to log memory for. Defaults to 0.\n    \"\"\"\n    if (not dist.is_initialized()) or (rank is None) or (dist.get_rank() == rank):\n        mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()\n        message = (\n            f\"{head}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, \"\n            f\"device memory used/total (GB): {mem_used}/{mem_total}\"\n        )\n\n        if logger is None:\n            print(message)\n        else:\n            logger.log(msg=message, level=level)\n\n\nclass GPUMemoryLogger(DecoratorLoggerBase):\n    \"\"\"A decorator class to log GPU memory usage.\n\n    Example:\n        >>> from verl.utils.profiler.performance import GPUMemoryLogger\n        >>> @GPUMemoryLogger(role=\"actor\")\n        >>> def update_actor(self, batch):\n        ...     # real actor update logics\n        ...     return\n    \"\"\"\n\n    def __init__(self, role: str, logger: logging.Logger = None, level=logging.DEBUG, log_only_rank_0: bool = True):\n        if dist.is_initialized() and dist.get_world_size() > 1:\n            rank = dist.get_rank()\n        else:\n            rank = 0\n        super().__init__(role, logger, level, rank, log_only_rank_0)\n\n    def __call__(self, decorated_function: callable):\n        def f(*args, **kwargs):\n            return self.log(decorated_function, *args, **kwargs)\n\n        return f\n\n    def log(self, func, *args, **kwargs):\n        name = func.__name__\n        mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()\n        message = (\n            f\"Before {name}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, \"\n            f\"device memory used/total (GB): {mem_used}/{mem_total}\"\n        )\n        self.logging_function(message)\n\n        output = func(*args, **kwargs)\n\n        mem_allocated, mem_reserved, mem_used, mem_total = _get_current_mem_info()\n        message = (\n            f\"After {name}, memory allocated (GB): {mem_allocated}, memory reserved (GB): {mem_reserved}, \"\n            f\"device memory used/total (GB): {mem_used}/{mem_total}\"\n        )\n\n        self.logging_function(message)\n        return output\n\n\ndef log_print(ctn: Any):\n    current_time = datetime.datetime.now().strftime(\"%Y-%m-%d %H:%M:%S\")\n\n    frame = inspect.currentframe().f_back\n    function_name = frame.f_code.co_name\n    line_number = frame.f_lineno\n    file_name = frame.f_code.co_filename.split(\"/\")[-1]\n    print(f\"[{current_time}-{file_name}:{line_number}:{function_name}]: {ctn}\")\n\n\ndef _timer(name: str, timing_raw: dict[str, float]):\n    \"\"\"Inner function that handles the core timing logic.\n\n    Args:\n        name (str): The name/identifier for this timing measurement.\n        timing_raw (Dict[str, float]): Dictionary to store timing information.\n    \"\"\"\n    with Timer(name=name, logger=None) as timer:\n        yield\n    if name not in timing_raw:\n        timing_raw[name] = 0\n    timing_raw[name] += timer.last\n\n\n@contextmanager\ndef simple_timer(name: str, timing_raw: dict[str, float]):\n    \"\"\"Context manager for basic timing without NVTX markers.\n\n    This utility function measures the execution time of code within its context\n    and accumulates the timing information in the provided dictionary.\n\n    Args:\n        name (str): The name/identifier for this timing measurement.\n        timing_raw (Dict[str, float]): Dictionary to store timing information.\n\n    Yields:\n        None: This is a context manager that yields control back to the code block.\n    \"\"\"\n    yield from _timer(name, timing_raw)\n\n\n@contextmanager\ndef marked_timer(\n    name: str,\n    timing_raw: dict[str, float],\n    color: str = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n):\n    \"\"\"Context manager for timing with platform markers.\n\n    This utility function measures the execution time of code within its context,\n    accumulates the timing information, and adds platform markers for profiling.\n    This function is a default implementation when hardware profiler is not available.\n\n    Args:\n        name (str): The name/identifier for this timing measurement.\n        timing_raw (Dict[str, float]): Dictionary to store timing information.\n        color (Optional[str]): Color for the marker. Defaults to None.\n        domain (Optional[str]): Domain for the marker. Defaults to None.\n        category (Optional[str]): Category for the marker. Defaults to None.\n\n    Yields:\n        None: This is a context manager that yields control back to the code block.\n    \"\"\"\n    yield from _timer(name, timing_raw)\n\n\ndef reduce_timing(\n    timing_raw: dict[str, float], reduce_op: torch.distributed.ReduceOp = torch.distributed.ReduceOp.AVG\n) -> dict[str, float]:\n    \"\"\"Reduce timing information across all processes.\n\n    This function uses distributed communication to gather and sum the timing\n    information from all processes in a distributed environment.\n\n    Args:\n        timing_raw (Dict[str, float]): Dictionary containing timing information.\n\n    Returns:\n        Dict[str, float]: Reduced timing information.\n    \"\"\"\n    if not dist.is_initialized():\n        return timing_raw\n\n    key_list, timing_list = [], []\n    for key in sorted(timing_raw.keys()):\n        key_list.append(key)\n        timing_list.append(timing_raw[key])\n    timing_list = torch.tensor(timing_list, dtype=torch.float32, device=get_device_id())\n    torch.distributed.all_reduce(timing_list, op=reduce_op)\n    timing_list = [tensor.item() for tensor in timing_list.to(\"cpu\")]\n    timing_generate = {key_list[i]: timing_list[i] for i in range(len(key_list))}\n    return timing_generate\n\n\ndef topk_reduce_ratio_min_max(timing: float, k: int = 10) -> tuple[float, float, float]:\n    \"\"\"Calculate topk items take-up ratio, and min/max timing across all ranks.\"\"\"\n    if not dist.is_initialized():\n        return -1.0, -1.0, -1.0\n\n    world_size = dist.get_world_size()\n    timing_tensor = torch.tensor(timing, dtype=torch.float32, device=get_device_id())\n    tensor_list = [torch.zeros(1, dtype=torch.float32, device=get_device_id()) for _ in range(world_size)]\n    torch.distributed.all_gather(tensor_list, timing_tensor)\n    tensor_stack = torch.stack(tensor_list)\n    timing_min = tensor_stack.min().cpu().item()\n    timing_max = tensor_stack.max().cpu().item()\n    top_k_percentile = torch.quantile(tensor_stack, 1 - k / 100)\n    tail_ratio = torch.mean((tensor_stack > top_k_percentile).float()).cpu().item()\n    return tail_ratio, timing_min, timing_max\n\n\ndef gather_timing(timing_raw: dict[str, float]) -> dict[str, list[float]]:\n    if not dist.is_initialized():\n        return {k: [v] for k, v in timing_raw.items()}\n\n    key_list, timing_list = [], []\n    for key in sorted(timing_raw.keys()):\n        key_list.append(key)\n        timing_list.append(timing_raw[key])\n\n    world_size = torch.distributed.get_world_size()\n\n    object_gather_list = [None] * world_size\n\n    torch.distributed.all_gather_object(object_gather_list, timing_list)\n\n    timing_generate = {\n        key_list[i]: [timing_list[i] for timing_list in object_gather_list] for i in range(len(key_list))\n    }\n\n    return timing_generate\n"
  },
  {
    "path": "verl/utils/profiler/profile.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nfrom typing import Callable, Optional\n\nfrom ..memory_utils import MemorySnapshotSampler, enable_memory_visualize\nfrom .config import ProfilerConfig, TorchMemoryToolConfig\n\n\ndef mark_start_range(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> None:\n    \"\"\"Start a profiling range marker (no-op implementation).\n\n    Args:\n        message (Optional[str]): Message to associate with the range marker.\n        color (Optional[str]): Color for the marker visualization.\n        domain (Optional[str]): Domain for the marker.\n        category (Optional[str]): Category for the marker.\n    \"\"\"\n    pass\n\n\ndef mark_end_range(range_id: str) -> None:\n    \"\"\"End a profiling range marker (no-op implementation).\n\n    Args:\n        range_id (str): Identifier of the range to end.\n    \"\"\"\n    pass\n\n\ndef mark_annotate(\n    message: Optional[str] = None,\n    color: Optional[str] = None,\n    domain: Optional[str] = None,\n    category: Optional[str] = None,\n) -> Callable:\n    \"\"\"Decorator to annotate a function with profiling markers (no-op implementation).\n\n    Args:\n        message (Optional[str]): Message to associate with the annotation.\n        color (Optional[str]): Color for the marker visualization.\n        domain (Optional[str]): Domain for the marker.\n        category (Optional[str]): Category for the marker.\n\n    Returns:\n        Callable: Decorator function that returns the original function unchanged.\n    \"\"\"\n\n    def decorator(func):\n        return func\n\n    return decorator\n\n\nclass DistProfiler:\n    \"\"\"A dispatcher that delegates to specific profilers based on config.tool.\n\n    Supported tools:\n    - nsys: NsightSystemsProfiler\n    - npu: NPUProfiler (Ascend)\n    - torch: PyTorch torch.profiler wrapper\n    - torch_memory: Torch CUDA memory snapshot dump\n    \"\"\"\n\n    def __init__(\n        self, rank: int, config: Optional[ProfilerConfig] = None, tool_config: Optional[object] = None, **kwargs\n    ):\n        # Default config\n        if not config:\n            config = ProfilerConfig(ranks=[], enable=False, tool_config=None)\n\n        if tool_config is None:\n            tool_config = config.tool_config\n\n        self.config = config\n        self.tool_config = tool_config\n\n        self._impl = None\n        self._tool = getattr(config, \"tool\", None)\n        self._enable = config.enable\n        self._this_step = False\n\n        # Normalize rank selection\n        self._this_rank = False\n        if config.all_ranks:\n            self._this_rank = True\n        elif config.ranks:\n            self._this_rank = rank in config.ranks\n        else:\n            # default rank 0 if enabled but ranks unspecified\n            self._this_rank = (rank == 0) if self._enable else False\n\n        # TorchMemoryProfiler currently do not support discrete mode.\n        self._discrete = getattr(tool_config, \"discrete\", False) if tool_config else False\n\n        # Lazy import to avoid circular deps\n        if self._tool == \"nsys\":\n            from .nvtx_profile import NsightSystemsProfiler as _Nsight\n\n            self._impl = _Nsight(rank=rank, config=config, tool_config=tool_config, **kwargs)\n        elif self._tool == \"npu\":\n            from .mstx_profile import NPUProfiler as _Npu\n\n            self._impl = _Npu(rank=rank, config=config, tool_config=tool_config, **kwargs)\n        elif self._tool == \"torch\":\n            from .torch_profile import Profiler as _Torch\n\n            self._impl = _Torch(rank=rank, config=config, tool_config=tool_config)\n        elif self._tool == \"torch_memory\":\n            self._impl = TorchMemoryProfiler(rank=rank, config=config, tool_config=tool_config)\n        else:\n            # Fallback to a no-op impl\n            self._impl = _NoOpProfiler()\n\n    def check_enable(self):\n        return self._enable\n\n    def check_this_rank(self):\n        return self._this_rank\n\n    def check_this_step(self):\n        return self._this_step\n\n    def is_discrete_mode(self):\n        return self._discrete\n\n    def start(self, **kwargs):\n        if self.check_enable() and self.check_this_rank():\n            self._this_step = True\n            return getattr(self._impl, \"start\", lambda **_: None)(**kwargs)\n\n    def stop(self):\n        if self.check_enable() and self.check_this_rank():\n            self._this_step = False\n            return getattr(self._impl, \"stop\", lambda: None)()\n\n    @classmethod\n    def annotate(\n        cls,\n        message: Optional[str] = None,\n        color: Optional[str] = None,\n        domain: Optional[str] = None,\n        category: Optional[str] = None,\n        **kwargs_outer,\n    ) -> Callable:\n        def decorator(func):\n            @functools.wraps(func)\n            def wrapper(self_instance, *args, **kwargs_inner):\n                profiler = getattr(self_instance, \"profiler\", None)\n                if (\n                    not profiler\n                    or not profiler.check_enable()\n                    or not profiler.check_this_step()\n                    or not profiler.check_this_rank()\n                ):\n                    return func(self_instance, *args, **kwargs_inner)\n\n                impl = profiler._impl\n                if hasattr(impl, \"annotate\"):\n                    try:\n                        actual_decorator = impl.annotate(\n                            message=message, color=color, domain=domain, category=category, **kwargs_outer\n                        )\n\n                        return actual_decorator(func)(self_instance, *args, **kwargs_inner)\n                    except Exception:\n                        return func(self_instance, *args, **kwargs_inner)\n                return func(self_instance, *args, **kwargs_inner)\n\n            return wrapper\n\n        return decorator\n\n\nclass _NoOpProfiler:\n    def start(self, **kwargs):\n        return\n\n    def stop(self):\n        return\n\n\nclass TorchMemoryProfiler:\n    \"\"\"Profiler that dumps CUDA memory snapshots at step boundaries.\n\n    Behavior:\n    - On first construction (per process), enable memory history recording if CUDA is available\n    - On start(step=X), remember sub_dir for this step\n    - On stop(), dump a memory snapshot into config.save_path under the remembered sub_dir\n    \"\"\"\n\n    _memory_history_enabled: bool = False\n\n    def __init__(\n        self, rank: int, config: Optional[ProfilerConfig], tool_config: Optional[TorchMemoryToolConfig] = None\n    ):\n        # Always respond to explicit start/stop calls for torch_memory tool,\n        # regardless of per-role enable flag, to align with global step control.\n        self.enable = True\n        if not config:\n            config = ProfilerConfig(ranks=[])\n        self.config = config\n        self.rank = rank\n        self.this_step = False\n        self.sub_dir = None\n        self.sampler = MemorySnapshotSampler()\n\n        # Get parameters from tool_config, with fallback to defaults\n        if tool_config:\n            trace_alloc_max_entries = tool_config.trace_alloc_max_entries\n            stack_depth = tool_config.stack_depth\n        else:\n            trace_alloc_max_entries = 100_000\n            stack_depth = 32\n\n        # Best-effort enable memory history once\n        if not TorchMemoryProfiler._memory_history_enabled:\n            try:\n                enable_memory_visualize(trace_alloc_max_entries=trace_alloc_max_entries, stack_depth=stack_depth)\n            except Exception:\n                # silently ignore if not supported\n                pass\n            TorchMemoryProfiler._memory_history_enabled = True\n\n    def start(self, **kwargs):\n        if not self.enable:\n            return\n        if not self._should_profile_this_rank():\n            return\n        profile_step = kwargs.get(\"profile_step\", None)\n        # Keep ranks aligned under same folder name\n        self.sub_dir = f\"step{profile_step}\" if profile_step is not None else None\n        self.this_step = True\n\n    def stop(self):\n        if not self.enable or not self.this_step:\n            return\n        self.this_step = False\n        if not self._should_profile_this_rank():\n            return\n        out_dir = self.config.save_path or \"outputs/profile\"\n        tag = \"torch_memory\"\n        # Dump snapshot; all ranks write into same sub_dir\n        try:\n            self.sampler.dump_memory_snapshot(out_dir=out_dir, tag=tag, sub_dir=self.sub_dir)\n        except Exception:\n            pass\n\n    def _should_profile_this_rank(self) -> bool:\n        if self.config.all_ranks:\n            return True\n        if self.config.ranks:\n            return self.rank in self.config.ranks\n        # default rank 0\n        return self.rank == 0\n\n\nclass DistProfilerExtension:\n    \"\"\"An extension class for DistProfiler that provides distributed profiling capabilities.\n\n    It is intended for workers in verl that single controller invokes.\n\n    This class wraps a DistProfiler instance and provides methods to start/stop profiling\n    that can be dispatched across multiple ranks in a distributed training environment.\n\n    Args:\n        profiler (DistProfiler): The base distributed profiler instance to extend\n    \"\"\"\n\n    def __init__(self, profiler: DistProfiler):\n        self.profiler = profiler\n\n    from verl.single_controller.base.decorator import Dispatch, register\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def start_profile(self, **kwargs) -> None:\n        \"\"\"Start profiling for the current rank in the current training step.\"\"\"\n        self.profiler.start(**kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def stop_profile(self) -> None:\n        \"\"\"Stop profiling for the current rank in the current training step.\"\"\"\n        self.profiler.stop()\n"
  },
  {
    "path": "verl/utils/profiler/torch_profile.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport functools\nimport os\nfrom datetime import datetime, timezone\nfrom typing import Callable, Optional\n\nimport torch\n\nfrom .config import ProfilerConfig, TorchProfilerToolConfig\nfrom .profile import DistProfiler\n\n\ndef get_torch_profiler(\n    contents: list[str],\n    save_path: str,\n    role: Optional[str] = None,\n    save_file_prefix: Optional[str] = None,\n    rank: int = 0,\n):\n    if role:\n        save_path = os.path.join(save_path, role)\n\n    os.makedirs(save_path, exist_ok=True)\n\n    current_time = datetime.now(tz=timezone.utc).astimezone()\n    timestamp = current_time.strftime(\"%Y%m%d%H%M%S%f\")[:-3]\n    pid = os.getpid()\n\n    save_file_name = f\"prof_rank-{rank}_{pid}_{timestamp}.json.gz\"\n    if save_file_prefix:\n        save_file_name = f\"{save_file_prefix}_{save_file_name}\"\n    save_path = os.path.join(save_path, save_file_name)\n\n    def _trace_handler(prof):\n        print(f\"[Profiler] Saving trace to {save_path}\")\n        prof.export_chrome_trace(save_path)\n\n    contents = set(contents) if contents else set()\n    activities = []\n    if not contents or \"cpu\" in contents:\n        activities.append(torch.profiler.ProfilerActivity.CPU)\n    if not contents or \"cuda\" in contents:\n        activities.append(torch.profiler.ProfilerActivity.CUDA)\n\n    return torch.profiler.profile(\n        activities=activities,\n        with_stack=\"stack\" in contents,\n        record_shapes=\"shapes\" in contents,\n        profile_memory=\"memory\" in contents,\n        on_trace_ready=_trace_handler,\n    )\n\n\nclass Profiler(DistProfiler):\n    \"\"\"A PyTorch profiler wrapper class for collecting performance metrics.\n\n    This profiler provides a convenient interface for profiling PyTorch operations,\n    with support for:\n\n    - CPU and CUDA activity profiling\n    - Configurable profiling schedule (wait/warmup/active steps)\n    - Multi-rank profiling support\n    - Chrome trace export\n\n    Args:\n        config: Configuration object containing profiling parameters\n    \"\"\"\n\n    _define_count = 0\n\n    def __init__(\n        self,\n        rank,\n        config: ProfilerConfig,\n        tool_config: Optional[TorchProfilerToolConfig] = None,\n        save_file_prefix=None,\n    ):\n        # note : if we do not set use_profile, it will be set as None, so that all function will be skip\n        config = config or ProfilerConfig(ranks=[], enable=False)\n        self.save_file_prefix = save_file_prefix\n\n        if not tool_config:\n            assert not config.enable, \"tool_config must be provided when profiler is enabled\"\n\n        self.prof = None\n        self.rank = rank\n        self.config = config\n        self.tool_config = tool_config\n        self.contents = self.tool_config.contents\n        self.save_path = self.config.save_path\n        # Align with other profilers: read discrete mode, default to False for torch profiler\n        self.discrete = getattr(self.tool_config, \"discrete\", False)\n\n    def check(self):\n        return self.prof is not None\n\n    def start(self, **kwargs):\n        role = kwargs.get(\"role\", None)\n        if not self.discrete and Profiler._define_count == 0:\n            self.prof = get_torch_profiler(\n                contents=self.contents,\n                save_path=self.save_path,\n                role=role,\n                save_file_prefix=self.save_file_prefix,\n                rank=self.rank,\n            )\n            print(f\"[Profiler] started for rank {self.rank}\")\n            self.prof.start()\n            Profiler._define_count += 1\n\n    def step(self):\n        if self.check():\n            self.prof.step()\n\n    def stop(self):\n        if not self.discrete and Profiler._define_count == 1:\n            self.step()\n            print(f\"[Profiler] stopped for rank {self.rank}\")\n            self.prof.stop()\n            Profiler._define_count -= 1\n\n    def annotate(self, message: Optional[str] = None, role: Optional[str] = None, **kwargs_outer) -> Callable:\n        \"\"\"Decorate a Worker member function to profile the current rank in the current training step.\n\n        Requires the target function to be a member function of a Worker,\n        which has a member field `profiler` with Profiler type.\n\n        Args:\n            message (str, optional):\n                The message to be displayed in the profiler. Defaults to None.\n            role (str, optional):\n                The role of the current data collection. Defaults to None.\n        \"\"\"\n\n        def decorator(func):\n            @functools.wraps(func)\n            def wrapper(*args, **kwargs_inner):\n                profile_name = message or func.__name__\n\n                if not self.discrete:\n                    # In continuous mode, we just record function, profiler started globally\n                    with torch.profiler.record_function(profile_name):\n                        return func(*args, **kwargs_inner)\n\n                # In discrete mode, we start/stop profiler around the function\n                prof = get_torch_profiler(\n                    contents=self.contents,\n                    save_path=self.save_path,\n                    role=role,\n                    save_file_prefix=self.save_file_prefix,\n                    rank=self.rank,\n                )\n                prof.start()\n                with torch.profiler.record_function(profile_name):\n                    result = func(*args, **kwargs_inner)\n                prof.stop()\n                return result\n\n            return wrapper\n\n        return decorator\n"
  },
  {
    "path": "verl/utils/py_functional.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nContain small python utility functions\n\"\"\"\n\nimport importlib\nimport multiprocessing\nimport os\nimport queue  # Import the queue module for exception type hint\nimport signal\nfrom contextlib import contextmanager\nfrom functools import wraps\nfrom types import SimpleNamespace\nfrom typing import Any, Callable, Iterator, Optional\n\nimport numpy as np\n\nfrom verl.utils.metric import Metric\n\n\n# --- Top-level helper for multiprocessing timeout ---\n# This function MUST be defined at the top level to be pickleable\ndef _mp_target_wrapper(target_func: Callable, mp_queue: multiprocessing.Queue, args: tuple, kwargs: dict[str, Any]):\n    \"\"\"\n    Internal wrapper function executed in the child process.\n    Calls the original target function and puts the result or exception into the queue.\n    \"\"\"\n    try:\n        result = target_func(*args, **kwargs)\n        mp_queue.put((True, result))  # Indicate success and put result\n    except Exception as e:\n        # Ensure the exception is pickleable for the queue\n        try:\n            import pickle\n\n            pickle.dumps(e)  # Test if the exception is pickleable\n            mp_queue.put((False, e))  # Indicate failure and put exception\n        except (pickle.PicklingError, TypeError):\n            # Fallback if the original exception cannot be pickled\n            mp_queue.put((False, RuntimeError(f\"Original exception type {type(e).__name__} not pickleable: {e}\")))\n\n\n# Renamed the function from timeout to timeout_limit\ndef timeout_limit(seconds: float, use_signals: bool = False):\n    \"\"\"\n    Decorator to add a timeout to a function.\n\n    Args:\n        seconds: The timeout duration in seconds.\n        use_signals: (Deprecated)  This is deprecated because signals only work reliably in the main thread\n                     and can cause issues in multiprocessing or multithreading contexts.\n                     Defaults to False, which uses the more robust multiprocessing approach.\n\n    Returns:\n        A decorated function with timeout.\n\n    Raises:\n        TimeoutError: If the function execution exceeds the specified time.\n        RuntimeError: If the child process exits with an error (multiprocessing mode).\n        NotImplementedError: If the OS is not POSIX (signals are only supported on POSIX).\n    \"\"\"\n\n    def decorator(func):\n        if use_signals:\n            if os.name != \"posix\":\n                raise NotImplementedError(f\"Unsupported OS: {os.name}\")\n            # Issue deprecation warning if use_signals is explicitly True\n            print(\n                \"WARN: The 'use_signals=True' option in the timeout decorator is deprecated. \\\n                Signals are unreliable outside the main thread. \\\n                Please use the default multiprocessing-based timeout (use_signals=False).\"\n            )\n\n            @wraps(func)\n            def wrapper_signal(*args, **kwargs):\n                def handler(signum, frame):\n                    # Update function name in error message if needed (optional but good practice)\n                    raise TimeoutError(f\"Function {func.__name__} timed out after {seconds} seconds (signal)!\")\n\n                old_handler = signal.getsignal(signal.SIGALRM)\n                signal.signal(signal.SIGALRM, handler)\n                # Use setitimer for float seconds support, alarm only supports integers\n                signal.setitimer(signal.ITIMER_REAL, seconds)\n\n                try:\n                    result = func(*args, **kwargs)\n                finally:\n                    # Reset timer and handler\n                    signal.setitimer(signal.ITIMER_REAL, 0)\n                    signal.signal(signal.SIGALRM, old_handler)\n                return result\n\n            return wrapper_signal\n        else:\n            # --- Multiprocessing based timeout (existing logic) ---\n            @wraps(func)\n            def wrapper_mp(*args, **kwargs):\n                q = multiprocessing.Queue(maxsize=1)\n                process = multiprocessing.Process(target=_mp_target_wrapper, args=(func, q, args, kwargs))\n                process.start()\n                process.join(timeout=seconds)\n\n                if process.is_alive():\n                    process.terminate()\n                    process.join(timeout=0.5)  # Give it a moment to terminate\n                    if process.is_alive():\n                        print(f\"Warning: Process {process.pid} did not terminate gracefully after timeout.\")\n                    # Update function name in error message if needed (optional but good practice)\n                    raise TimeoutError(f\"Function {func.__name__} timed out after {seconds} seconds (multiprocessing)!\")\n\n                try:\n                    success, result_or_exc = q.get(timeout=0.1)  # Small timeout for queue read\n                    if success:\n                        return result_or_exc\n                    else:\n                        raise result_or_exc  # Reraise exception from child\n                except queue.Empty as err:\n                    exitcode = process.exitcode\n                    if exitcode is not None and exitcode != 0:\n                        raise RuntimeError(\n                            f\"Child process exited with error (exitcode: {exitcode}) before returning result.\"\n                        ) from err\n                    else:\n                        # Should have timed out if queue is empty after join unless process died unexpectedly\n                        # Update function name in error message if needed (optional but good practice)\n                        raise TimeoutError(\n                            f\"Operation timed out or process finished unexpectedly without result \"\n                            f\"(exitcode: {exitcode}).\"\n                        ) from err\n                finally:\n                    q.close()\n                    q.join_thread()\n\n            return wrapper_mp\n\n    return decorator\n\n\ndef union_two_dict(dict1: dict, dict2: dict):\n    \"\"\"Union two dict. Will throw an error if there is an item not the same object with the same key.\n\n    Args:\n        dict1:\n        dict2:\n\n    Returns:\n\n    \"\"\"\n    for key, val in dict2.items():\n        if key in dict1:\n            assert dict2[key] == dict1[key], f\"{key} in meta_dict1 and meta_dict2 are not the same object\"\n        dict1[key] = val\n\n    return dict1\n\n\ndef rename_dict(data: dict, prefix: str = \"\") -> dict:\n    \"\"\"Add a prefix to all the keys in the data dict if it's name is not started with prefix\n\n    Args:\n        data: a dictionary\n        prefix: prefix\n\n    Returns:\n        dictionary with modified name\n\n    \"\"\"\n    new_data = {}\n    for key, val in data.items():\n        new_key = f\"{prefix}{key}\" if not key.startswith(prefix) else key\n        new_data[new_key] = val\n    return new_data\n\n\ndef append_to_dict(data: dict, new_data: dict, prefix: str = \"\"):\n    \"\"\"Append values from new_data to lists in data.\n\n    For each key in new_data, this function appends the corresponding value to a list\n    stored under the same key in data. If the key doesn't exist in data, a new list is created.\n\n    Args:\n        data (Dict): The target dictionary containing lists as values.\n        new_data (Dict): The source dictionary with values to append.\n\n    Returns:\n        None: The function modifies data in-place.\n    \"\"\"\n    for key, val in new_data.items():\n        new_key = f\"{prefix}{key}\" if not key.startswith(prefix) else key\n        if new_key not in data:\n            data[new_key] = val.init_list() if isinstance(val, Metric) else []\n        if isinstance(val, list):\n            data[new_key].extend(val)\n        else:\n            data[new_key].append(val)\n\n\nclass NestedNamespace(SimpleNamespace):\n    \"\"\"A nested version of SimpleNamespace that recursively converts dictionaries to namespaces.\n\n    This class allows for dot notation access to nested dictionary structures by recursively\n    converting dictionaries to NestedNamespace objects.\n\n    Example:\n        config_dict = {\"a\": 1, \"b\": {\"c\": 2, \"d\": 3}}\n        config = NestedNamespace(config_dict)\n        # Access with: config.a, config.b.c, config.b.d\n\n    Args:\n        dictionary: The dictionary to convert to a nested namespace.\n        **kwargs: Additional attributes to set on the namespace.\n    \"\"\"\n\n    def __init__(self, dictionary, **kwargs):\n        super().__init__(**kwargs)\n        for key, value in dictionary.items():\n            if isinstance(value, dict):\n                self.__setattr__(key, NestedNamespace(value))\n            else:\n                self.__setattr__(key, value)\n\n\nclass DynamicEnumMeta(type):\n    def __iter__(cls) -> Iterator[Any]:\n        return iter(cls._registry.values())\n\n    def __contains__(cls, item: Any) -> bool:\n        # allow `name in EnumClass` or `member in EnumClass`\n        if isinstance(item, str):\n            return item in cls._registry\n        return item in cls._registry.values()\n\n    def __getitem__(cls, name: str) -> Any:\n        return cls._registry[name]\n\n    def __reduce_ex__(cls, protocol):\n        # Always load the existing module and grab the class\n        return getattr, (importlib.import_module(cls.__module__), cls.__name__)\n\n    def names(cls):\n        return list(cls._registry.keys())\n\n    def values(cls):\n        return list(cls._registry.values())\n\n\nclass DynamicEnum(metaclass=DynamicEnumMeta):\n    _registry: dict[str, \"DynamicEnum\"] = {}\n    _next_value: int = 0\n\n    def __init__(self, name: str, value: int):\n        self.name = name\n        self.value = value\n\n    def __repr__(self):\n        return f\"<{self.__class__.__name__}.{self.name}: {self.value}>\"\n\n    def __reduce_ex__(self, protocol):\n        \"\"\"\n        Unpickle via: getattr(import_module(module).Dispatch, 'ONE_TO_ALL')\n        so the existing class is reused instead of re-executed.\n        \"\"\"\n        module = importlib.import_module(self.__class__.__module__)\n        enum_cls = getattr(module, self.__class__.__name__)\n        return getattr, (enum_cls, self.name)\n\n    @classmethod\n    def register(cls, name: str) -> \"DynamicEnum\":\n        key = name.upper()\n        if key in cls._registry:\n            raise ValueError(f\"{key} already registered\")\n        member = cls(key, cls._next_value)\n        cls._registry[key] = member\n        setattr(cls, key, member)\n        cls._next_value += 1\n        return member\n\n    @classmethod\n    def remove(cls, name: str):\n        key = name.upper()\n        member = cls._registry.pop(key)\n        delattr(cls, key)\n        return member\n\n    @classmethod\n    def from_name(cls, name: str) -> Optional[\"DynamicEnum\"]:\n        return cls._registry.get(name.upper())\n\n\n@contextmanager\ndef temp_env_var(key: str, value: str):\n    \"\"\"Context manager for temporarily setting an environment variable.\n\n    This context manager ensures that environment variables are properly set and restored,\n    even if an exception occurs during the execution of the code block.\n\n    Args:\n        key: Environment variable name to set\n        value: Value to set the environment variable to\n\n    Yields:\n        None\n\n    Example:\n        >>> with temp_env_var(\"MY_VAR\", \"test_value\"):\n        ...     # MY_VAR is set to \"test_value\"\n        ...     do_something()\n        ... # MY_VAR is restored to its original value or removed if it didn't exist\n    \"\"\"\n    original = os.environ.get(key)\n    os.environ[key] = value\n    try:\n        yield\n    finally:\n        if original is None:\n            os.environ.pop(key, None)\n        else:\n            os.environ[key] = original\n\n\ndef convert_to_regular_types(obj):\n    \"\"\"Convert Hydra configs and other special types to regular Python types.\"\"\"\n    from omegaconf import DictConfig, ListConfig\n\n    if isinstance(obj, ListConfig | DictConfig):\n        return {k: convert_to_regular_types(v) for k, v in obj.items()} if isinstance(obj, DictConfig) else list(obj)\n    elif isinstance(obj, list | tuple):\n        return [convert_to_regular_types(x) for x in obj]\n    elif isinstance(obj, dict):\n        return {k: convert_to_regular_types(v) for k, v in obj.items()}\n    return obj\n\n\ndef convert_nested_value_to_list_recursive(data_item):\n    if isinstance(data_item, dict):\n        return {k: convert_nested_value_to_list_recursive(v) for k, v in data_item.items()}\n    elif isinstance(data_item, list):\n        return [convert_nested_value_to_list_recursive(elem) for elem in data_item]\n    elif isinstance(data_item, np.ndarray):\n        # Convert to list, then recursively process the elements of the new list\n        return convert_nested_value_to_list_recursive(data_item.tolist())\n    else:\n        # Base case: item is already a primitive type (int, str, float, bool, etc.)\n        return data_item\n\n\ndef list_of_dict_to_dict_of_list(list_of_dict: list[dict]):\n    if len(list_of_dict) == 0:\n        return {}\n    keys = list_of_dict[0].keys()\n    output = {key: [] for key in keys}\n    for data in list_of_dict:\n        for key, item in data.items():\n            assert key in output, f\"Key '{key}' is not present in the keys of the first dictionary in the list.\"\n            output[key].append(item)\n    return output\n"
  },
  {
    "path": "verl/utils/qat/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nQAT (Quantization-Aware Training) module for verl.\n\nSupports NVFP4 (W4A4 and W4A16) quantization modes for FSDP training.\n\nModule Structure:\n- core.py: QATConfig, apply_qat, enable_qat_fuse (training setup)\n- linear.py: QATLinear layer with Triton kernels for fake quantization\n- quantizer.py: QATQuantizer for true quantization + scale computation utilities\n- vllm_patch.py: Patches for vLLM dynamic weight loading\n\nUsage:\n    from verl.utils.qat import apply_qat, QATConfig\n\n    config = QATConfig(enable=True, mode=\"w4a16\")\n    model = apply_qat(model, config)  # Before FSDP wrapping\n\"\"\"\n\nfrom verl.utils.qat.core import (\n    QATConfig,\n    apply_qat,\n    enable_qat_fuse,\n    invalidate_all_scales,\n    load_quantization_config,\n)\nfrom verl.utils.qat.vllm_patch import (\n    apply_qat_patches,\n    manual_process_weights_after_loading,\n    prepare_qat_for_load_weights,\n)\n\n__all__ = [\n    # Core\n    \"QATConfig\",\n    \"apply_qat\",\n    \"load_quantization_config\",\n    \"enable_qat_fuse\",\n    \"invalidate_all_scales\",\n    # vLLM Patch\n    \"apply_qat_patches\",\n    \"manual_process_weights_after_loading\",\n    \"prepare_qat_for_load_weights\",\n]\n"
  },
  {
    "path": "verl/utils/qat/core.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"QAT (Quantization-Aware Training) utilities for verl FSDP training.\"\"\"\n\nimport json\nimport logging\nimport re\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nimport torch.nn as nn\n\nfrom verl.base_config import BaseConfig\n\nlogger = logging.getLogger(__name__)\n\n\n@dataclass\nclass QATConfig(BaseConfig):\n    \"\"\"Unified configuration for QAT (Quantization-Aware Training).\"\"\"\n\n    enable: bool = False\n    mode: str = \"w4a16\"\n    group_size: int = 16\n    ignore_patterns: list[str] = field(default_factory=lambda: [\"lm_head\", \"embed_tokens\", \"re:.*mlp.gate$\"])\n    activation_observer: str = \"static_minmax\"\n    quantization_config_path: Optional[str] = None\n\n\ndef load_quantization_config(qat_config: QATConfig) -> dict[str, Any]:\n    \"\"\"Load quantization config JSON file from QATConfig.\"\"\"\n    if not qat_config.quantization_config_path:\n        raise ValueError(\"quantization_config_path is required when QAT is enabled\")\n\n    logger.info(f\"Loading QAT quantization config from: {qat_config.quantization_config_path}\")\n\n    with open(qat_config.quantization_config_path) as f:\n        quant_config = json.load(f)\n\n    if qat_config.ignore_patterns:\n        original_ignore = quant_config.get(\"ignore\", [])\n        quant_config[\"ignore\"] = qat_config.ignore_patterns\n        if original_ignore != qat_config.ignore_patterns:\n            logger.info(f\"Overriding JSON 'ignore' field: {original_ignore} -> {qat_config.ignore_patterns}\")\n\n    logger.info(\"Successfully loaded QAT quantization config\")\n    return quant_config\n\n\ndef _should_quantize(name: str, module: nn.Module, config: QATConfig) -> bool:\n    \"\"\"Check if a module should be quantized.\"\"\"\n    if not isinstance(module, nn.Linear):\n        return False\n\n    for pattern in config.ignore_patterns:\n        if pattern.startswith(\"re:\"):\n            regex = pattern[3:]\n            if re.match(regex, name):\n                logger.debug(f\"Ignoring {name} due to regex pattern: {regex}\")\n                return False\n        else:\n            if pattern in name:\n                logger.debug(f\"Ignoring {name} due to pattern: {pattern}\")\n                return False\n\n    if module.in_features % config.group_size != 0:\n        logger.warning(\n            f\"Skipping {name}: in_features={module.in_features} not divisible by group_size={config.group_size}\"\n        )\n        return False\n\n    return True\n\n\ndef apply_qat(\n    model: nn.Module,\n    config: QATConfig | dict[str, Any],\n) -> nn.Module:\n    \"\"\"Apply QAT to a model by replacing nn.Linear with QATLinear.\"\"\"\n    from verl.utils.qat.linear import QATLinear, QATMode\n\n    if not isinstance(config, QATConfig):\n        config = QATConfig(**config)\n\n    if not config.enable:\n        logger.info(\"QAT is disabled, returning original model\")\n        return model\n\n    mode = QATMode(config.mode.lower())\n    logger.info(f\"Applying QAT with mode={mode.value}, group_size={config.group_size}\")\n\n    modules_to_replace = []\n    for name, module in model.named_modules():\n        if _should_quantize(name, module, config):\n            modules_to_replace.append((name, module))\n\n    logger.info(f\"Found {len(modules_to_replace)} Linear layers to convert to QAT\")\n\n    converted_count = 0\n    for name, module in modules_to_replace:\n        if isinstance(module, QATLinear):\n            continue\n\n        fake_quant_module = QATLinear.from_linear(\n            module,\n            mode=mode,\n            group_size=config.group_size,\n            activation_observer=config.activation_observer,\n        )\n\n        _set_module(model, name, fake_quant_module)\n        converted_count += 1\n\n    logger.info(f\"Successfully applied QAT to {converted_count} layers\")\n\n    return model\n\n\ndef _set_module(model: nn.Module, name: str, new_module: nn.Module):\n    \"\"\"Set a module in the model by its full name.\"\"\"\n    parts = name.split(\".\")\n    parent = model\n    for part in parts[:-1]:\n        parent = getattr(parent, part)\n    setattr(parent, parts[-1], new_module)\n\n\nFUSION_PATTERNS = {\n    \"qkv\": [\"q_proj\", \"k_proj\", \"v_proj\"],\n    \"gate_up\": [\"gate_proj\", \"up_proj\"],\n}\n\n\ndef setup_fusion_siblings(model: nn.Module):\n    \"\"\"Setup fusion siblings for QKV and GateUp layers.\"\"\"\n    import weakref\n\n    from verl.utils.qat.linear import QATLinear\n\n    qat_modules = {name: m for name, m in model.named_modules() if isinstance(m, QATLinear)}\n\n    counts = {}\n    for group_name, suffixes in FUSION_PATTERNS.items():\n        groups: dict[str, dict[str, nn.Module]] = {}\n        for name, module in qat_modules.items():\n            for suffix in suffixes:\n                if name.endswith(suffix):\n                    parent = name.rsplit(\".\", 1)[0]\n                    groups.setdefault(parent, {})[suffix] = module\n\n        count = 0\n        for parent, projs in groups.items():\n            if len(projs) >= 2:\n                modules = list(projs.values())\n                for i, m in enumerate(modules):\n                    siblings = modules[:i] + modules[i + 1 :]\n                    m._fusion_siblings_ref = [weakref.ref(s) for s in siblings]\n                count += 1\n        counts[group_name] = count\n\n    logger.info(f\"[QAT Fuse] Setup fusion siblings: {counts}\")\n    return counts\n\n\ndef enable_qat_fuse(model: nn.Module):\n    \"\"\"Enable QAT fuse mode: sets up fusion siblings for weight scale fusion.\"\"\"\n    setup_fusion_siblings(model)\n    model._qat_fuse_enabled = True\n    logger.info(\"[QAT Fuse] Enabled QAT fuse mode\")\n\n\ndef invalidate_all_scales(model: nn.Module):\n    \"\"\"Clear all cached weight scales after optimizer.step().\"\"\"\n    from verl.utils.qat.linear import QATLinear\n\n    count = 0\n    for module in model.modules():\n        if isinstance(module, QATLinear):\n            module._weight_blockwise_scale = None\n            module._weight_global_scale = None\n            module._cached_weight_amax = None\n            count += 1\n\n    logger.debug(f\"[QAT Fuse] Invalidated scales for {count} QATLinear layers\")\n"
  },
  {
    "path": "verl/utils/qat/linear.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"QAT FakeQuantized Linear module for NVFP4 (W4A4/W4A16) with FSDP compatibility.\n\nIncludes Triton kernels for high-performance FP4 quantization.\n\"\"\"\n\nfrom enum import Enum\nfrom typing import Optional\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\n\n__all__ = [\"QATLinear\", \"QATMode\"]\n\n\nimport triton\nimport triton.language as tl\n\n_TORCH_TO_TL_DTYPE = {\n    torch.float32: tl.float32,\n    torch.float16: tl.float16,\n    torch.bfloat16: tl.bfloat16,\n}\nFP4_E2M1_MAX: float = 6.0\nFP8_E4M3_MAX: float = 448.0\n\n\n@triton.jit\ndef _fp4_fake_quant_kernel(\n    x_ptr,\n    y_ptr,\n    M,\n    N,\n    global_scale_ptr,\n    stride_xm,\n    stride_xn,\n    stride_ym,\n    stride_yn,\n    BLOCK_SIZE: tl.constexpr,\n    TILE_M: tl.constexpr,\n    TILE_N: tl.constexpr,\n    NUM_FP4_BLOCKS: tl.constexpr,\n    OUT_DTYPE: tl.constexpr,\n    FP4_MAX: tl.constexpr,\n    FP8_MAX: tl.constexpr,\n):\n    pid_m = tl.program_id(axis=0)\n    pid_n = tl.program_id(axis=1)\n    row_start = pid_m * TILE_M\n    col_start = pid_n * TILE_N\n\n    x_block_ptr = tl.make_block_ptr(\n        base=x_ptr,\n        shape=(M, N),\n        strides=(stride_xm, stride_xn),\n        offsets=(row_start, col_start),\n        block_shape=(TILE_M, TILE_N),\n        order=(1, 0),\n    )\n    y_block_ptr = tl.make_block_ptr(\n        base=y_ptr,\n        shape=(M, N),\n        strides=(stride_ym, stride_yn),\n        offsets=(row_start, col_start),\n        block_shape=(TILE_M, TILE_N),\n        order=(1, 0),\n    )\n\n    global_scale = tl.load(global_scale_ptr).to(tl.float32)\n    global_scale_safe = tl.where(global_scale > 0.0, global_scale, 1e-12)\n\n    tile = tl.load(x_block_ptr, boundary_check=(0, 1), padding_option=\"zero\").to(tl.float32)\n    tile_reshaped = tl.reshape(tile, (TILE_M, NUM_FP4_BLOCKS, BLOCK_SIZE))\n    x_abs = tl.abs(tile_reshaped)\n\n    block_max = tl.max(x_abs, axis=2, keep_dims=True)\n    block_max_scaled = block_max / (FP4_MAX * global_scale_safe)\n    block_max_scaled = tl.minimum(block_max_scaled, FP8_MAX)\n    block_max_quant = block_max_scaled.to(tl.float8e4nv).to(tl.float32) * global_scale\n    block_max_quant = tl.where(block_max_quant >= 1e-5, block_max_quant, 1.0)\n\n    block_max_quant_broadcast = tl.broadcast_to(block_max_quant, (TILE_M, NUM_FP4_BLOCKS, BLOCK_SIZE))\n    abs_scaled = x_abs / block_max_quant_broadcast\n\n    q_val = tl.where(\n        abs_scaled <= 0.25,\n        0.0,\n        tl.where(\n            abs_scaled < 0.75,\n            0.5,\n            tl.where(\n                abs_scaled <= 1.25,\n                1.0,\n                tl.where(\n                    abs_scaled < 1.75,\n                    1.5,\n                    tl.where(\n                        abs_scaled <= 2.5,\n                        2.0,\n                        tl.where(abs_scaled < 3.5, 3.0, tl.where(abs_scaled <= 5.0, 4.0, FP4_MAX)),\n                    ),\n                ),\n            ),\n        ),\n    )\n\n    x_rescaled = q_val * block_max_quant_broadcast\n    x_rescaled = tl.where(tile_reshaped >= 0, x_rescaled, -x_rescaled)\n    tile_quant = tl.reshape(x_rescaled, (TILE_M, TILE_N))\n\n    tl.store(y_block_ptr, tile_quant.to(OUT_DTYPE), boundary_check=(0, 1))\n\n\ndef fp4_fake_quant_weight(\n    weight: torch.Tensor,\n    global_amax: torch.Tensor = None,\n    block_size: int = 16,\n    tile_rows: int = 16,\n    tile_cols: int = 64,\n) -> torch.Tensor:\n    \"\"\"Apply FP4 fake quantization using Triton kernel.\"\"\"\n    x_shape = weight.shape\n    x_dtype = weight.dtype\n    x = weight.reshape(-1, x_shape[-1]).contiguous()\n    M, N = x.shape\n    y = torch.empty_like(x)\n\n    stride_xm, stride_xn = x.stride()\n    stride_ym, stride_yn = y.stride()\n\n    tile_cols = max(tile_cols, block_size)\n    tile_cols_aligned = ((tile_cols + block_size - 1) // block_size) * block_size\n    num_fp4_blocks = tile_cols_aligned // block_size\n\n    if global_amax is None:\n        global_amax = weight.abs().max().to(torch.float32)\n    global_scale = global_amax.float() / (FP4_E2M1_MAX * FP8_E4M3_MAX)\n\n    grid = (triton.cdiv(M, tile_rows), triton.cdiv(N, tile_cols_aligned))\n\n    _fp4_fake_quant_kernel[grid](\n        x,\n        y,\n        M,\n        N,\n        global_scale,\n        stride_xm,\n        stride_xn,\n        stride_ym,\n        stride_yn,\n        BLOCK_SIZE=block_size,\n        TILE_M=tile_rows,\n        TILE_N=tile_cols_aligned,\n        NUM_FP4_BLOCKS=num_fp4_blocks,\n        OUT_DTYPE=_TORCH_TO_TL_DTYPE[x_dtype],\n        FP4_MAX=FP4_E2M1_MAX,\n        FP8_MAX=FP8_E4M3_MAX,\n    )\n    return y.view(*x_shape)\n\n\nclass STEFP4QuantTriton(torch.autograd.Function):\n    \"\"\"Straight-Through Estimator wrapper for Triton FP4 quantization kernel.\"\"\"\n\n    @staticmethod\n    def forward(ctx, x: torch.Tensor, global_amax: torch.Tensor, block_size: int) -> torch.Tensor:\n        return fp4_fake_quant_weight(x, global_amax=global_amax, block_size=block_size)\n\n    @staticmethod\n    def backward(ctx, grad_output: torch.Tensor) -> tuple:\n        return grad_output, None, None\n\n\nclass QATMode(str, Enum):\n    \"\"\"QAT quantization mode.\"\"\"\n\n    W4A4 = \"w4a4\"  # Weight 4-bit, Activation 4-bit (dynamic)\n    W4A16 = \"w4a16\"  # Weight 4-bit, Activation 16-bit (weight only)\n\n\nclass QATLinear(nn.Linear):\n    \"\"\"QAT FakeQuantized Linear layer with FSDP compatibility.\"\"\"\n\n    _UNINITIALIZED_SCALE = -1.0\n\n    def __init__(\n        self,\n        in_features: int,\n        out_features: int,\n        bias: bool = True,\n        mode: QATMode = QATMode.W4A4,\n        group_size: int = 16,\n        activation_observer: str = \"static_minmax\",  # Observer strategy for activation global_scale\n        device: Optional[torch.device] = None,\n        dtype: Optional[torch.dtype] = None,\n    ):\n        super().__init__(in_features, out_features, bias, device=device, dtype=dtype)\n\n        self.mode = mode\n        self.group_size = group_size\n        self.activation_observer = activation_observer\n\n        self._weight_blockwise_scale: Optional[torch.Tensor] = None\n        self._weight_global_scale: Optional[torch.Tensor] = None\n        self._cached_weight_amax: Optional[torch.Tensor] = None\n        self._fusion_siblings_ref = None\n\n        if mode == QATMode.W4A4:\n            self.register_buffer(\n                \"input_global_scale\", torch.tensor([self._UNINITIALIZED_SCALE], dtype=torch.float32), persistent=True\n            )\n\n            self.register_buffer(\n                \"input_amax\", torch.tensor([self._UNINITIALIZED_SCALE], dtype=torch.float32), persistent=True\n            )\n\n            self._ema_decay: float = 0.01\n\n        self.fake_quant_enabled = True\n\n    @classmethod\n    def from_linear(\n        cls,\n        linear: nn.Linear,\n        mode: QATMode = QATMode.W4A4,\n        group_size: int = 16,\n        activation_observer: str = \"static_minmax\",\n    ) -> \"QATLinear\":\n        \"\"\"Create QATLinear from an existing nn.Linear.\"\"\"\n        has_bias = linear.bias is not None\n\n        new_linear = cls(\n            in_features=linear.in_features,\n            out_features=linear.out_features,\n            bias=has_bias,\n            mode=mode,\n            group_size=group_size,\n            activation_observer=activation_observer,\n            device=linear.weight.device,\n            dtype=linear.weight.dtype,\n        )\n\n        if linear.weight.device != torch.device(\"meta\"):\n            new_linear.weight = nn.Parameter(linear.weight.clone())\n            if has_bias:\n                new_linear.bias = nn.Parameter(linear.bias.clone())\n\n        return new_linear\n\n    def _is_amax_initialized(self) -> bool:\n        \"\"\"Check if input_amax has been initialized.\"\"\"\n        if not hasattr(self, \"input_amax\"):\n            return False\n        return self.input_amax.item() != self._UNINITIALIZED_SCALE\n\n    def _update_input_global_scale(self, x: torch.Tensor):\n        \"\"\"Update static input_global_scale based on observer strategy.\"\"\"\n        assert self.mode == QATMode.W4A4, \"_update_input_global_scale should only be called in W4A4 mode\"\n\n        current_amax = torch.amax(torch.abs(x)).detach().to(torch.float32)\n\n        if torch.distributed.is_initialized() and torch.distributed.get_world_size() > 1:\n            torch.distributed.all_reduce(current_amax, op=torch.distributed.ReduceOp.MAX)\n\n        scale_factor = FP8_E4M3_MAX * FP4_E2M1_MAX\n\n        if self.activation_observer == \"memoryless_minmax\":\n            new_scale = (scale_factor / (current_amax + 1e-12)).view(1)\n            self.input_global_scale.copy_(new_scale.to(self.input_global_scale.device))\n\n        elif self.activation_observer == \"static_minmax\":\n            if not self._is_amax_initialized():\n                self.input_amax.copy_(current_amax.view(1).to(self.input_amax.device))\n            else:\n                new_amax = torch.maximum(self.input_amax, current_amax.view(1).to(self.input_amax.device))\n                self.input_amax.copy_(new_amax)\n            amax_f32 = self.input_amax.to(torch.float32)\n            new_scale = (scale_factor / (amax_f32 + 1e-12)).float().view(1)\n            self.input_global_scale.copy_(new_scale.to(self.input_global_scale.device))\n\n        elif self.activation_observer == \"minmax\":\n            if not self._is_amax_initialized():\n                self.input_amax.copy_(current_amax.view(1).to(self.input_amax.device))\n            else:\n                new_amax = (1 - self._ema_decay) * self.input_amax + self._ema_decay * current_amax.view(1).to(\n                    self.input_amax.device\n                )\n                self.input_amax.copy_(new_amax)\n            amax_f32 = self.input_amax.to(torch.float32)\n            new_scale = (scale_factor / (amax_f32 + 1e-12)).float().view(1)\n            self.input_global_scale.copy_(new_scale.to(self.input_global_scale.device))\n\n        else:\n            raise ValueError(f\"Unknown activation_observer: {self.activation_observer}\")\n\n    def _fake_quantize_weight(self, weight: torch.Tensor) -> torch.Tensor:\n        \"\"\"Apply fake quantization to weight tensor using Triton kernel.\"\"\"\n        with torch.no_grad():\n            if self._cached_weight_amax is not None:\n                global_amax = self._cached_weight_amax\n            else:\n                siblings_ref = getattr(self, \"_fusion_siblings_ref\", None)\n\n                if siblings_ref is not None:\n                    siblings = [ref() for ref in siblings_ref if ref() is not None]\n                    siblings = [s for s in siblings if s.weight.device != torch.device(\"meta\")]\n\n                    for sibling in siblings:\n                        sibling_amax = getattr(sibling, \"_cached_weight_amax\", None)\n                        if sibling_amax is not None:\n                            global_amax = sibling_amax\n                            self._cached_weight_amax = global_amax\n                            break\n                    else:\n                        all_modules = [self] + siblings\n                        amaxes = [m.weight.abs().max().to(torch.float32) for m in all_modules]\n                        global_amax = torch.max(torch.stack(amaxes))\n\n                        self._cached_weight_amax = global_amax\n                        for sibling in siblings:\n                            sibling._cached_weight_amax = global_amax\n                else:\n                    global_amax = weight.abs().max().to(torch.float32)\n                    self._cached_weight_amax = global_amax\n\n            if self._weight_global_scale is None:\n                self._weight_global_scale = global_amax.float() / (FP4_E2M1_MAX * FP8_E4M3_MAX)\n\n        result = STEFP4QuantTriton.apply(weight, global_amax, self.group_size)\n\n        return result\n\n    def _fake_quantize_activation(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Apply fake quantization to activation tensor (W4A4 mode only).\"\"\"\n        original_shape = x.shape\n\n        if x.dim() == 3:\n            x_2d = x.view(-1, x.shape[-1])\n        else:\n            x_2d = x\n\n        if self.training:\n            self._update_input_global_scale(x_2d)\n\n        if self.input_global_scale.item() == self._UNINITIALIZED_SCALE:\n            raise RuntimeError(\"W4A4 input_global_scale uninitialized. Load PTQ model first.\")\n\n        global_amax = (FP4_E2M1_MAX * FP8_E4M3_MAX) / self.input_global_scale.to(x.device)\n        result = STEFP4QuantTriton.apply(x_2d, global_amax, self.group_size)\n        return result.view(original_shape)\n\n    def forward(self, x: torch.Tensor) -> torch.Tensor:\n        \"\"\"Forward pass with fake quantization.\"\"\"\n        if not self.fake_quant_enabled:\n            return F.linear(x, self.weight, self.bias)\n\n        weight_fq = self._fake_quantize_weight(self.weight)\n\n        if self.mode == QATMode.W4A4:\n            x_fq = self._fake_quantize_activation(x)\n        else:\n            x_fq = x\n\n        return F.linear(x_fq, weight_fq, self.bias)\n\n    def extra_repr(self) -> str:\n        return (\n            f\"in_features={self.in_features}, out_features={self.out_features}, \"\n            f\"bias={self.bias is not None}, mode={self.mode.value}, \"\n            f\"group_size={self.group_size}, fake_quant_enabled={self.fake_quant_enabled}\"\n        )\n"
  },
  {
    "path": "verl/utils/qat/quantizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nFast NVFP4 Quantizer for verl FSDP training.\n\nDirectly computes scales and quantizes weights using compressed_tensors APIs.\nIncludes scale computation utilities for weight quantization.\n\"\"\"\n\nimport logging\nimport os\nimport re\nfrom typing import Generator, Iterable, Optional\n\nimport torch\nfrom compressed_tensors.compressors.quantized_compressors.fp4_quantized import NVFP4PackedCompressor\nfrom compressed_tensors.quantization.quant_args import (\n    FP4_E2M1_DATA,\n    FP8_E4M3_DATA,\n    QuantizationArgs,\n    QuantizationStrategy,\n    QuantizationType,\n)\nfrom compressed_tensors.quantization.utils.helpers import generate_gparam\n\nfrom verl.utils.device import get_device_name, get_torch_device\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n_LAYER_IDX_RE = re.compile(r\"layers\\.(\\d+)\\.\")\n\n\ndef compute_blockwise_scale(\n    weight: torch.Tensor,\n    global_scale: torch.Tensor,\n    group_size: int = 16,\n) -> torch.Tensor:\n    \"\"\"Compute blockwise scale using pre-computed global_scale (for fusion).\n    Returns FP8 E4M3 blockwise scale tensor.\n    \"\"\"\n    out_features, in_features = weight.shape\n    num_groups = in_features // group_size\n    weight_reshaped = weight.view(out_features, num_groups, group_size)\n    block_max = torch.amax(torch.abs(weight_reshaped), dim=-1).to(torch.float32)\n\n    local_scale = block_max / FP4_E2M1_DATA.max\n    blockwise_scale_f32 = torch.clamp(\n        global_scale * local_scale,\n        min=-FP8_E4M3_DATA.max,\n        max=FP8_E4M3_DATA.max,\n    )\n\n    blockwise_scale = blockwise_scale_f32.to(torch.float8_e4m3fn)\n    eps = torch.finfo(torch.float8_e4m3fn).eps\n    blockwise_scale = torch.where(\n        blockwise_scale == 0,\n        torch.tensor(eps, dtype=blockwise_scale.dtype, device=weight.device),\n        blockwise_scale,\n    )\n\n    return blockwise_scale\n\n\n# Fusion patterns for transformer models\nFUSE_PATTERNS = {\n    \"qkv\": [\"q_proj\", \"k_proj\", \"v_proj\"],\n    \"gate_up\": [\"gate_proj\", \"up_proj\"],\n}\n\n\ndef fuse_global_scales(\n    layer_global_scales: dict[str, torch.Tensor],\n    strategy: str = \"min\",\n) -> dict[str, torch.Tensor]:\n    \"\"\"Fuse global scales for QKV/GateUp groups (take min across group).\"\"\"\n    if not layer_global_scales:\n        return {}\n\n    # Group by parent module\n    parent_to_children: dict[str, dict[str, str]] = {}\n    for name in layer_global_scales:\n        parent, child = name.rsplit(\".\", 1) if \".\" in name else (\"\", name)\n        parent_to_children.setdefault(parent, {})[child] = name\n\n    fused_scales = {}\n    processed = set()\n\n    for parent, children in parent_to_children.items():\n        for _, patterns in FUSE_PATTERNS.items():\n            matched = [children[p] for p in patterns if p in children]\n            if len(matched) == len(patterns):\n                group_scales = [layer_global_scales[n] for n in matched]\n                if strategy == \"min\":\n                    fused_scale = torch.min(torch.cat(group_scales)).reshape([1])\n                else:\n                    raise ValueError(f\"Unknown fuse strategy: {strategy}\")\n                for layer_name in matched:\n                    fused_scales[layer_name] = fused_scale.clone()\n                    processed.add(layer_name)\n\n    for name, scale in layer_global_scales.items():\n        if name not in processed:\n            fused_scales[name] = scale\n\n    return fused_scales\n\n\nclass QATQuantizer:\n    \"\"\"Quantizer for QAT-trained weights using compressed_tensors APIs.\"\"\"\n\n    def __init__(\n        self,\n        mode: str = \"w4a16\",\n        group_size: int = 16,\n        ignore_patterns: Optional[list] = None,\n        device: Optional[torch.device] = None,\n        param_dtype: Optional[torch.dtype] = None,\n    ):\n        self.mode = mode.lower()\n        self._is_w4a4 = self.mode == \"w4a4\"  # W4A4 needs input_global_scale\n        self.group_size = group_size\n        self.ignore_patterns = ignore_patterns or [\"lm_head\", \"embed_tokens\", \"re:.*mlp.gate$\"]\n        self.device = device or torch.device(get_device_name())\n        self.param_dtype = param_dtype\n\n        self._compressor = NVFP4PackedCompressor()\n        self._quant_args = QuantizationArgs(\n            num_bits=4,\n            type=QuantizationType.FLOAT,\n            symmetric=True,\n            strategy=QuantizationStrategy.TENSOR_GROUP,\n            group_size=group_size,\n            scale_dtype=FP8_E4M3_DATA.dtype,\n        )\n\n    def _should_quantize(self, name: str, tensor: torch.Tensor) -> bool:\n        \"\"\"Check if parameter should be quantized.\"\"\"\n        if not name.endswith(\".weight\"):\n            return False\n        if tensor.dim() != 2:\n            return False\n        if tensor.shape[1] % self.group_size != 0:\n            return False\n\n        module_name = name.rsplit(\".weight\", 1)[0]\n\n        for pattern in self.ignore_patterns:\n            if pattern.startswith(\"re:\"):\n                # Regex pattern - use re.match like vLLM does\n                regex = pattern[3:]\n                if re.match(regex, module_name):\n                    return False\n            else:\n                if pattern in module_name:\n                    return False\n        return True\n\n    @staticmethod\n    def _extract_layer_idx(name: str) -> Optional[int]:\n        \"\"\"Extract decoder layer index from parameter name.\"\"\"\n        match = _LAYER_IDX_RE.search(name)\n        return int(match.group(1)) if match else None\n\n    def _process_layer_group(\n        self,\n        layer_idx: Optional[int],\n        layer_params: dict[str, torch.Tensor],\n        input_global_scales: dict[str, torch.Tensor],\n        output_device: torch.device,\n    ) -> list[tuple[str, torch.Tensor]]:\n        \"\"\"Quantize one decoder layer's buffered params. Returns list of (name, tensor).\"\"\"\n        layer_weights = {}\n        layer_passthrough = {}\n\n        for name, tensor in layer_params.items():\n            if \"input_global_scale\" in name or \"input_amax\" in name:\n                continue\n\n            if self._should_quantize(name, tensor):\n                layer_name = name.rsplit(\".weight\", 1)[0]\n                layer_weights[layer_name] = (name, tensor)\n            else:\n                layer_passthrough[name] = tensor\n\n        if layer_idx is None and layer_weights:\n            raise RuntimeError(\n                f\"[QAT Quantizer] Unexpected quantizable weights outside decoder layers: \"\n                f\"{list(layer_weights.keys())}. These should be in ignore_patterns.\"\n            )\n\n        if not layer_weights:\n            return [(name, tensor.to(output_device)) for name, tensor in layer_passthrough.items()]\n\n        # Move weights to GPU, compute global scales\n        weights_on_gpu = {}\n        layer_global_scales = {}\n\n        for layer_name, (_, tensor) in layer_weights.items():\n            weight_gpu = tensor.to(device=self.device, dtype=self.param_dtype)\n            weights_on_gpu[layer_name] = weight_gpu\n            amax = torch.amax(torch.abs(weight_gpu)).to(torch.float32)\n            layer_global_scales[layer_name] = generate_gparam(\n                -amax.unsqueeze(0),\n                amax.unsqueeze(0),\n                scale_data=FP8_E4M3_DATA,\n                quant_data=FP4_E2M1_DATA,\n                dtype=torch.float32,\n            )\n\n        fused_global_scales = fuse_global_scales(layer_global_scales, strategy=\"min\")\n\n        results = []\n\n        for layer_name, weight_gpu in weights_on_gpu.items():\n            fused_global_scale = fused_global_scales[layer_name]\n            weight_scale = compute_blockwise_scale(weight_gpu, fused_global_scale, self.group_size)\n            weight_packed = self._compressor.compress_weight(\n                weight=weight_gpu,\n                scale=weight_scale.float(),\n                global_scale=fused_global_scale,\n                quantization_args=self._quant_args,\n            )[\"weight_packed\"]\n\n            results.append((f\"{layer_name}.weight_packed\", weight_packed.to(output_device)))\n            results.append((f\"{layer_name}.weight_scale\", weight_scale.to(output_device)))\n            results.append((f\"{layer_name}.weight_global_scale\", fused_global_scale.to(output_device)))\n\n            if self._is_w4a4:\n                if layer_name in input_global_scales:\n                    results.append(\n                        (\n                            f\"{layer_name}.input_global_scale\",\n                            input_global_scales[layer_name].float().to(output_device),\n                        )\n                    )\n                else:\n                    raise ValueError(\n                        f\"W4A4 mode requires input_global_scale for layer '{layer_name}', \"\n                        f\"but it's not found or uninitialized (-1.0).\"\n                    )\n\n        del weights_on_gpu, layer_global_scales, fused_global_scales\n\n        for name, tensor in layer_passthrough.items():\n            results.append((name, tensor.to(output_device)))\n\n        return results\n\n    def quantize_with_fusion(\n        self,\n        params: dict[str, torch.Tensor] | Iterable[tuple[str, torch.Tensor]],\n        target_device: Optional[torch.device] = None,\n    ) -> Generator[tuple[str, torch.Tensor], None, None]:\n        \"\"\"Streaming quantize: consume input layer by layer, yield (name, tensor) pairs.\"\"\"\n        if isinstance(params, dict):\n            params = params.items()\n\n        output_device = target_device or torch.device(\"cpu\")\n\n        _sentinel = object()\n        current_layer_idx = _sentinel\n        layer_buffer: dict[str, torch.Tensor] = {}\n        input_global_scales: dict[str, torch.Tensor] = {}\n        for name, tensor in params:\n            tensor_cpu = tensor.to(\"cpu\") if tensor.is_cuda else tensor\n            layer_idx = self._extract_layer_idx(name)\n\n            # Collect input_global_scales for W4A4 as we go\n            if self._is_w4a4 and \"input_global_scale\" in name:\n                scale_layer_name = name.replace(\".input_global_scale\", \"\")\n                if tensor_cpu.numel() == 1 and tensor_cpu.item() == -1.0:\n                    logger.warning(f\"W4A4: {scale_layer_name} input_global_scale is uninitialized\")\n                else:\n                    input_global_scales[scale_layer_name] = tensor_cpu\n\n            # Layer boundary: flush previous layer\n            if layer_idx != current_layer_idx and current_layer_idx is not _sentinel and layer_buffer:\n                yield from self._process_layer_group(\n                    current_layer_idx, layer_buffer, input_global_scales, output_device\n                )\n                layer_buffer = {}\n\n            current_layer_idx = layer_idx\n            layer_buffer[name] = tensor_cpu\n\n        # Flush last buffered layer\n        if layer_buffer:\n            yield from self._process_layer_group(current_layer_idx, layer_buffer, input_global_scales, output_device)\n\n        get_torch_device().empty_cache()\n\n\n__all__ = [\n    \"QATQuantizer\",\n]\n"
  },
  {
    "path": "verl/utils/qat/vllm_patch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nvLLM NVFP4 Patches for Dynamic Weight Updates.\n\nEnables dynamic weight reloading for NVFP4 quantized models in vLLM.\n\nSupported schemes:\n- Dense: W4A16-FP4, W4A4-FP4\n- MoE: NVFP4-MoE\n\"\"\"\n\nimport logging\nimport os\nfrom typing import Optional\nfrom unittest.mock import patch\n\nimport torch\nfrom torch.nn import Parameter\n\nfrom verl.utils.device import get_device_name\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass ParamMetaDict(dict):\n    \"\"\"\n    Dict-like class for parameter management with metadata-based rebuild and tensor swap.\n\n    Supports:\n    - Rebuild of deleted parameters from saved metadata\n    - Tensor Swap for parameters with shape changes (address stability for CUDA Graph)\n    \"\"\"\n\n    def __init__(self, model: torch.nn.Module, device: Optional[torch.device] = None):\n        \"\"\"\n        Initialize ParamMetaDict from a model.\n\n        Args:\n            model: vLLM model (may be wrapped in ModelRunner)\n            device: Device for created parameters\n        \"\"\"\n        super().__init__()\n        self.device = device\n\n        # Get the actual model (handle vLLM's wrapper structure)\n        actual_model = model\n        if hasattr(model, \"model\"):\n            actual_model = model.model\n        self._model = actual_model\n\n        # Build mappings by scanning all modules\n        self._layer_meta_cache: dict[str, dict] = {}  # Cache of _hf_param_meta\n        self._tensor_swap_layers: dict[str, dict] = {}  # Layers needing tensor swap\n\n        self._build_mappings()\n\n        # Initialize with current parameters\n        for name, param in actual_model.named_parameters():\n            self[name] = param\n\n    def _build_mappings(self):\n        \"\"\"Build layer metadata cache for rebuild and tensor swap.\"\"\"\n        for layer_name, module in self._model.named_modules():\n            # Check for _hf_param_meta which indicates this layer has HF format params\n            if hasattr(module, \"_hf_param_meta\"):\n                self._layer_meta_cache[layer_name] = {\n                    \"module\": module,\n                    \"meta\": module._hf_param_meta,\n                }\n\n                # Check for tensor swap layers (weight_scale with shape change)\n                if \"weight_scale\" in module._hf_param_meta:\n                    marlin_refs = getattr(module, \"_marlin_tensor_refs\", {})\n                    if \"weight_scale\" in marlin_refs:\n                        self._tensor_swap_layers[layer_name] = {\n                            \"module\": module,\n                            \"marlin_ref\": marlin_refs[\"weight_scale\"],\n                            \"hf_meta\": module._hf_param_meta[\"weight_scale\"],\n                        }\n\n                # MoE layers (w13_weight_scale, w2_weight_scale)\n                if \"w13_weight_scale\" in module._hf_param_meta:\n                    marlin_refs = getattr(module, \"_marlin_tensor_refs\", {})\n                    if \"w13_weight_scale\" in marlin_refs:\n                        self._tensor_swap_layers[f\"{layer_name}.w13\"] = {\n                            \"module\": module,\n                            \"param_name\": \"w13_weight_scale\",\n                            \"marlin_ref\": marlin_refs[\"w13_weight_scale\"],\n                            \"hf_meta\": module._hf_param_meta[\"w13_weight_scale\"],\n                        }\n                    if \"w2_weight_scale\" in marlin_refs:\n                        self._tensor_swap_layers[f\"{layer_name}.w2\"] = {\n                            \"module\": module,\n                            \"param_name\": \"w2_weight_scale\",\n                            \"marlin_ref\": marlin_refs[\"w2_weight_scale\"],\n                            \"hf_meta\": module._hf_param_meta[\"w2_weight_scale\"],\n                        }\n\n    def _try_rebuild(self, key: str) -> Optional[Parameter]:\n        \"\"\"\n        Try to rebuild a parameter from metadata if it was deleted.\n\n        Args:\n            key: Full parameter name\n\n        Returns:\n            Rebuilt parameter or None if cannot rebuild\n        \"\"\"\n        # Extract layer name and param name\n        parts = key.rsplit(\".\", 1)\n        if len(parts) != 2:\n            return None\n\n        layer_name, param_name = parts\n\n        # Check if we have metadata for this layer\n        if layer_name not in self._layer_meta_cache:\n            return None\n\n        cache_entry = self._layer_meta_cache[layer_name]\n        module = cache_entry[\"module\"]\n        meta = cache_entry[\"meta\"]\n\n        # Check if this param needs rebuild\n        if param_name not in meta:\n            return None\n\n        # Already exists on module?\n        if hasattr(module, param_name):\n            param = getattr(module, param_name)\n            if param is not None:\n                return param\n\n        # Rebuild from metadata\n        new_param = _create_param_from_meta(module, param_name, meta[param_name], self.device)\n        module.register_parameter(param_name, new_param)\n        return new_param\n\n    def prepare_for_reload(self) -> None:\n        \"\"\"Replace Marlin-format tensors with HF-shape tensors for reload.\"\"\"\n        for layer_name, swap_info in self._tensor_swap_layers.items():\n            module = swap_info[\"module\"]\n            param_name = swap_info.get(\"param_name\", \"weight_scale\")\n            hf_meta = swap_info[\"hf_meta\"]\n            if hasattr(module, param_name):\n                new_param = _create_param_from_meta(module, param_name, hf_meta, self.device)\n                setattr(module, param_name, new_param)\n\n    def __getitem__(self, key: str) -> Parameter:\n        \"\"\"Get parameter with rebuild support.\"\"\"\n        # Try standard lookup first\n        if key in dict.keys(self):\n            return super().__getitem__(key)\n\n        # Try rebuild from metadata\n        param = self._try_rebuild(key)\n        if param is not None:\n            self[key] = param\n            return param\n\n        raise KeyError(f\"Parameter not found: {key}\")\n\n    def __contains__(self, key: str) -> bool:\n        \"\"\"Check if parameter exists (with rebuild check).\"\"\"\n        if super().__contains__(key):\n            return True\n\n        # Check if can rebuild from metadata\n        parts = key.rsplit(\".\", 1)\n        if len(parts) == 2:\n            layer_name, param_name = parts\n            if layer_name in self._layer_meta_cache:\n                meta = self._layer_meta_cache[layer_name][\"meta\"]\n                if param_name in meta:\n                    return True\n\n        return False\n\n    def get(self, key: str, default=None):\n        \"\"\"Get parameter with default.\"\"\"\n        try:\n            return self[key]\n        except KeyError:\n            return default\n\n\ndef _create_param_from_meta(\n    module: torch.nn.Module,\n    param_name: str,\n    meta: dict,\n    device: Optional[torch.device] = None,\n) -> Parameter:\n    \"\"\"Create a Parameter from saved metadata. Used by rebuild and tensor swap.\"\"\"\n    shape = meta[\"shape\"]\n    dtype = meta[\"dtype\"]\n    dev = device or meta.get(\"device\", get_device_name())\n    param_class = meta.get(\"param_class\", Parameter)\n\n    weight_loaders = getattr(module, \"_weight_loaders\", {})\n    weight_loader = weight_loaders.get(param_name)\n\n    data = torch.empty(shape, dtype=dtype, device=dev)\n\n    try:\n        if param_class is not Parameter and weight_loader is not None:\n            kwargs = {\"data\": data, \"weight_loader\": weight_loader}\n            if \"input_dim\" in meta:\n                kwargs[\"input_dim\"] = meta[\"input_dim\"]\n            if \"output_dim\" in meta:\n                kwargs[\"output_dim\"] = meta[\"output_dim\"]\n            new_param = param_class(**kwargs)\n        else:\n            new_param = Parameter(data, requires_grad=False)\n            if weight_loader is not None:\n                new_param.weight_loader = weight_loader\n    except Exception as e:\n        logger.warning(f\"Failed to create param {param_name} with class {param_class}: {e}, using Parameter\")\n        new_param = Parameter(data, requires_grad=False)\n        if weight_loader is not None:\n            new_param.weight_loader = weight_loader\n\n    if \"quant_method\" in meta:\n        new_param.quant_method = meta[\"quant_method\"]\n\n    return new_param\n\n\ndef save_param_meta(layer: torch.nn.Module, param_name: str):\n    \"\"\"Save parameter metadata for rebuild.\"\"\"\n    if not hasattr(layer, \"_hf_param_meta\"):\n        layer._hf_param_meta = {}\n\n    param = getattr(layer, param_name, None)\n    if param is None:\n        return\n\n    meta = {\n        \"shape\": tuple(param.shape),\n        \"dtype\": param.dtype,\n        \"device\": str(param.device),\n        \"param_class\": type(param),  # Save the actual parameter class\n    }\n\n    # Save vLLM-specific attributes needed for reconstruction\n    if hasattr(param, \"_input_dim\"):\n        meta[\"input_dim\"] = param._input_dim\n    if hasattr(param, \"_output_dim\"):\n        meta[\"output_dim\"] = param._output_dim\n\n    # Save MoE-specific attributes (quant_method is required by weight_loader)\n    if hasattr(param, \"quant_method\"):\n        meta[\"quant_method\"] = param.quant_method\n\n    layer._hf_param_meta[param_name] = meta\n\n\ndef _check_first_call(layer: torch.nn.Module) -> bool:\n    \"\"\"Check if this is the first process_weights call, and increment counter.\"\"\"\n    count = getattr(layer, \"_process_weights_call_count\", 0)\n    layer._process_weights_call_count = count + 1\n    return count == 0\n\n\n# Dense W4A16 Patches\ndef patched_w4a16_process_weights_after_loading(self, layer: torch.nn.Module) -> None:\n    \"\"\"Patched process_weights_after_loading for W4A16 Dense layer.\"\"\"\n    import vllm._custom_ops as ops\n    from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (\n        marlin_make_workspace_new,\n        marlin_permute_scales,\n        nvfp4_marlin_process_global_scale,\n        nvfp4_marlin_process_scales,\n    )\n\n    is_first_call = _check_first_call(layer)\n\n    group_size = 16\n    part_size_n = layer.output_size_per_partition\n    part_size_k = layer.input_size_per_partition\n    device = layer.weight_packed.device\n    param_dtype = getattr(layer, \"params_dtype\", torch.float16)\n\n    # Save metadata (first call only)\n    if is_first_call:\n        save_param_meta(layer, \"weight_packed\")\n        save_param_meta(layer, \"weight_global_scale\")\n        save_param_meta(layer, \"weight_scale\")\n        if not hasattr(layer, \"_weight_loaders\"):\n            layer._weight_loaders = {}\n        for pname in [\"weight_packed\", \"weight_global_scale\", \"weight_scale\"]:\n            param = getattr(layer, pname, None)\n            if param is not None and hasattr(param, \"weight_loader\"):\n                layer._weight_loaders[pname] = param.weight_loader\n\n    # Get HF format data\n    weight_packed_hf = layer.weight_packed.data\n    weight_global_scale_hf = layer.weight_global_scale.data\n    weight_scale_hf = layer.weight_scale.data\n\n    # Create workspace (first call only)\n    if is_first_call:\n        layer.workspace = marlin_make_workspace_new(device)\n\n    # Convert to Marlin format\n    perm = torch.empty(0, dtype=torch.int, device=device)\n    qweight = weight_packed_hf.view(torch.int32).T.contiguous()\n    marlin_weight = ops.gptq_marlin_repack(\n        b_q_weight=qweight,\n        perm=perm,\n        size_k=part_size_k,\n        size_n=part_size_n,\n        num_bits=4,\n        is_a_8bit=False,\n    )\n\n    weight_scale = weight_scale_hf.T.contiguous().to(param_dtype)\n    weight_scale_permuted = marlin_permute_scales(\n        s=weight_scale,\n        size_k=part_size_k,\n        size_n=part_size_n,\n        group_size=group_size,\n        is_a_8bit=False,\n    )\n    marlin_weight_scale = nvfp4_marlin_process_scales(weight_scale_permuted)\n\n    weight_scale_2_raw = (1.0 / weight_global_scale_hf.max()).to(param_dtype)\n    marlin_weight_scale_2 = nvfp4_marlin_process_global_scale(weight_scale_2_raw)\n\n    # Update compute parameters\n    if is_first_call:\n        layer.weight = Parameter(marlin_weight, requires_grad=False)\n        layer.weight_scale = Parameter(marlin_weight_scale, requires_grad=False)\n        layer.weight_scale_2 = Parameter(marlin_weight_scale_2, requires_grad=False)\n        if not hasattr(layer, \"_marlin_tensor_refs\"):\n            layer._marlin_tensor_refs = {}\n        layer._marlin_tensor_refs[\"weight_scale\"] = layer.weight_scale.data\n    else:\n        layer.weight.data.copy_(marlin_weight)\n        layer.weight_scale_2.data.copy_(marlin_weight_scale_2)\n        marlin_scale_ref = layer._marlin_tensor_refs.get(\"weight_scale\")\n        if marlin_scale_ref is not None:\n            marlin_scale_ref.copy_(marlin_weight_scale)\n            layer.weight_scale = Parameter(marlin_scale_ref, requires_grad=False)\n        else:\n            logger.warning(\"W4A16: _marlin_tensor_refs['weight_scale'] not found\")\n            layer.weight_scale = Parameter(marlin_weight_scale, requires_grad=False)\n\n    # Delete HF parameters\n    if hasattr(layer, \"weight_packed\"):\n        delattr(layer, \"weight_packed\")\n    if hasattr(layer, \"weight_global_scale\"):\n        delattr(layer, \"weight_global_scale\")\n\n\ndef patched_w4a4_process_weights_after_loading(self, layer: torch.nn.Module) -> None:\n    \"\"\"Patched process_weights_after_loading for W4A4 Dense (all backends).\"\"\"\n    from vllm.model_executor.layers.quantization.utils.quant_utils import swizzle_blockscale\n\n    is_first_call = _check_first_call(layer)\n\n    _W4A4_HF_PARAMS = [\"weight_packed\", \"weight_scale\", \"weight_global_scale\", \"input_global_scale\"]\n\n    if is_first_call:\n        for pname in _W4A4_HF_PARAMS:\n            save_param_meta(layer, pname)\n        if not hasattr(layer, \"_weight_loaders\"):\n            layer._weight_loaders = {}\n        for pname in _W4A4_HF_PARAMS:\n            param = getattr(layer, pname, None)\n            if param is not None and hasattr(param, \"weight_loader\"):\n                layer._weight_loaders[pname] = param.weight_loader\n\n    weight_packed_data = layer.weight_packed.data\n    weight_scale_data = layer.weight_scale.data\n    input_global_scale_data = layer.input_global_scale.data\n    weight_global_scale_data = layer.weight_global_scale.data\n\n    global_input_scale = input_global_scale_data.max().to(torch.float32)\n    global_weight_scale = weight_global_scale_data.max().to(torch.float32)\n\n    if self.backend == \"flashinfer-trtllm\":\n        from flashinfer import shuffle_matrix_a, shuffle_matrix_sf_a\n\n        epilogue_tile_m = 128\n        processed_weight = shuffle_matrix_a(weight_packed_data.view(torch.uint8), epilogue_tile_m)\n        processed_weight_scale = (\n            shuffle_matrix_sf_a(weight_scale_data.view(torch.uint8), epilogue_tile_m)\n            .reshape(weight_scale_data.shape)\n            .view(torch.float8_e4m3fn)\n        )\n    elif self.backend == \"fbgemm\":\n        processed_weight_scale = swizzle_blockscale(weight_scale_data).view(-1).view(torch.uint8)\n        processed_weight = weight_packed_data\n    else:\n        # cutlass / flashinfer-cutlass\n        processed_weight_scale = swizzle_blockscale(weight_scale_data)\n        processed_weight = weight_packed_data\n\n    alpha = 1.0 / (global_input_scale * global_weight_scale)\n\n    if is_first_call:\n        layer.weight_packed = Parameter(processed_weight, requires_grad=False)\n        layer.weight_scale = Parameter(processed_weight_scale, requires_grad=False)\n        layer.input_global_scale = Parameter(global_input_scale, requires_grad=False)\n        layer.weight_global_scale = Parameter(global_weight_scale, requires_grad=False)\n        layer.alpha = Parameter(alpha, requires_grad=False)\n\n        if not hasattr(layer, \"_marlin_tensor_refs\"):\n            layer._marlin_tensor_refs = {}\n        layer._marlin_tensor_refs[\"weight_packed\"] = layer.weight_packed.data\n        layer._marlin_tensor_refs[\"weight_scale\"] = layer.weight_scale.data\n        layer._marlin_tensor_refs[\"input_global_scale\"] = layer.input_global_scale.data\n        layer._marlin_tensor_refs[\"weight_global_scale\"] = layer.weight_global_scale.data\n        layer._marlin_tensor_refs[\"alpha\"] = layer.alpha.data\n    else:\n        refs = layer._marlin_tensor_refs\n        for ref_name, new_data in [\n            (\"weight_packed\", processed_weight),\n            (\"weight_scale\", processed_weight_scale),\n            (\"input_global_scale\", global_input_scale),\n            (\"weight_global_scale\", global_weight_scale),\n            (\"alpha\", alpha),\n        ]:\n            ref = refs.get(ref_name)\n            if ref is not None:\n                ref.copy_(new_data)\n                setattr(layer, ref_name, Parameter(ref, requires_grad=False))\n            else:\n                logger.warning(f\"W4A4: _marlin_tensor_refs['{ref_name}'] not found, creating new Parameter\")\n                setattr(\n                    layer,\n                    ref_name,\n                    Parameter(\n                        new_data.clone() if isinstance(new_data, torch.Tensor) else torch.tensor(new_data),\n                        requires_grad=False,\n                    ),\n                )\n\n\ndef _marlin_repack_experts(packed, perm, size_k, size_n, num_experts):\n    \"\"\"Repack weight for each expert into Marlin format and stack.\"\"\"\n    import vllm._custom_ops as ops\n\n    result = []\n    for i in range(num_experts):\n        qweight = packed[i].view(torch.int32).T.contiguous()\n        result.append(\n            ops.gptq_marlin_repack(\n                b_q_weight=qweight,\n                perm=perm,\n                size_k=size_k,\n                size_n=size_n,\n                num_bits=4,\n                is_a_8bit=False,\n            )\n        )\n    return torch.stack(result)\n\n\ndef _marlin_process_scales_experts(scale_hf, param_dtype, size_k, size_n, group_size, num_experts):\n    \"\"\"Process scales for each expert into Marlin format and stack.\"\"\"\n    from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (\n        marlin_permute_scales,\n        nvfp4_marlin_process_scales,\n    )\n\n    result = []\n    scales = scale_hf.to(param_dtype)\n    for i in range(num_experts):\n        s = marlin_permute_scales(\n            s=scales[i].T,\n            size_k=size_k,\n            size_n=size_n,\n            group_size=group_size,\n            is_a_8bit=False,\n        )\n        result.append(nvfp4_marlin_process_scales(s))\n    return torch.stack(result)\n\n\ndef _process_nvfp4_moe_marlin(self, layer: torch.nn.Module, is_first_call: bool) -> None:\n    \"\"\"Process MoE layer with MARLIN backend (W4A16).\"\"\"\n    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import make_nvfp4_moe_kernel\n    from vllm.model_executor.layers.quantization.utils.marlin_utils_fp4 import (\n        marlin_make_workspace_new,\n        nvfp4_marlin_process_global_scale,\n    )\n\n    group_size = 16\n    e = layer.num_experts\n    k = layer.hidden_size\n    n = layer.intermediate_size_per_partition\n    device = layer.w13_weight_packed.device\n    param_dtype = layer.params_dtype\n    w13_num_shards = 2 if self.moe.is_act_and_mul else 1\n\n    if is_first_call:\n        layer.workspace = marlin_make_workspace_new(device, 4)\n\n    perm = torch.empty(0, dtype=torch.int, device=device)\n\n    if self.moe.is_act_and_mul and not torch.allclose(\n        layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]\n    ):\n        logger.warning(\"w1_weight_global_scale must match w3_weight_global_scale. Accuracy may be affected.\")\n\n    size_n_w13, size_k_w13 = n * w13_num_shards, k\n    size_n_w2, size_k_w2 = k, n\n\n    w13_weight_marlin = _marlin_repack_experts(layer.w13_weight_packed.data, perm, size_k_w13, size_n_w13, e)\n    w2_weight_marlin = _marlin_repack_experts(layer.w2_weight_packed.data, perm, size_k_w2, size_n_w2, e)\n    w13_weight_scale_marlin = _marlin_process_scales_experts(\n        layer.w13_weight_scale.data, param_dtype, size_k_w13, size_n_w13, group_size, e\n    )\n    w2_weight_scale_marlin = _marlin_process_scales_experts(\n        layer.w2_weight_scale.data, param_dtype, size_k_w2, size_n_w2, group_size, e\n    )\n\n    # Process global scales\n    w13_scale_2 = 1.0 / layer.w13_weight_global_scale[:, 0]\n    w2_scale_2 = 1.0 / layer.w2_weight_global_scale.data\n    w13_scale_2_processed = nvfp4_marlin_process_global_scale(w13_scale_2.to(param_dtype))\n    w2_scale_2_processed = nvfp4_marlin_process_global_scale(w2_scale_2.to(param_dtype))\n\n    # Update parameters\n    if is_first_call:\n        layer.w13_weight = Parameter(w13_weight_marlin, requires_grad=False)\n        layer.w2_weight = Parameter(w2_weight_marlin, requires_grad=False)\n        layer.w13_weight_scale = Parameter(w13_weight_scale_marlin, requires_grad=False)\n        layer.w2_weight_scale = Parameter(w2_weight_scale_marlin, requires_grad=False)\n        layer.w13_weight_scale_2 = Parameter(w13_scale_2_processed, requires_grad=False)\n        layer.w2_weight_scale_2 = Parameter(w2_scale_2_processed, requires_grad=False)\n        if not hasattr(layer, \"_marlin_tensor_refs\"):\n            layer._marlin_tensor_refs = {}\n        layer._marlin_tensor_refs[\"w13_weight_scale\"] = layer.w13_weight_scale.data\n        layer._marlin_tensor_refs[\"w2_weight_scale\"] = layer.w2_weight_scale.data\n    else:\n        layer.w13_weight.data.copy_(w13_weight_marlin)\n        layer.w2_weight.data.copy_(w2_weight_marlin)\n        layer.w13_weight_scale_2.data.copy_(w13_scale_2_processed)\n        layer.w2_weight_scale_2.data.copy_(w2_scale_2_processed)\n        w13_marlin_ref = layer._marlin_tensor_refs.get(\"w13_weight_scale\")\n        w2_marlin_ref = layer._marlin_tensor_refs.get(\"w2_weight_scale\")\n        if w13_marlin_ref is not None:\n            w13_marlin_ref.copy_(w13_weight_scale_marlin)\n            layer.w13_weight_scale = Parameter(w13_marlin_ref, requires_grad=False)\n        else:\n            logger.warning(\"MoE: _marlin_tensor_refs['w13_weight_scale'] not found\")\n            layer.w13_weight_scale.data.copy_(w13_weight_scale_marlin)\n        if w2_marlin_ref is not None:\n            w2_marlin_ref.copy_(w2_weight_scale_marlin)\n            layer.w2_weight_scale = Parameter(w2_marlin_ref, requires_grad=False)\n        else:\n            logger.warning(\"MoE: _marlin_tensor_refs['w2_weight_scale'] not found\")\n            layer.w2_weight_scale.data.copy_(w2_weight_scale_marlin)\n\n    layer.w13_input_scale = None\n    layer.w2_input_scale = None\n\n    # Initialize kernel\n    self.moe_quant_config = self.get_fused_moe_quant_config(layer)\n    if self.moe_quant_config is not None and (\n        (not self.moe.moe_parallel_config.use_all2all_kernels) or self.moe.moe_parallel_config.use_naive_all2all_kernels\n    ):\n        self.kernel = make_nvfp4_moe_kernel(\n            moe_quant_config=self.moe_quant_config,\n            moe_config=self.moe,\n            experts_cls=self.experts_cls,\n        )\n\n\ndef _process_nvfp4_moe_flashinfer_cutlass(self, layer: torch.nn.Module, is_first_call: bool) -> None:\n    \"\"\"Process MoE layer with FlashInfer/CUTLASS backend (W4A4).\"\"\"\n    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import (\n        convert_to_nvfp4_moe_kernel_format,\n        make_nvfp4_moe_kernel,\n    )\n    from vllm.model_executor.utils import replace_parameter\n\n    w13_packed = layer.w13_weight_packed.data\n    w2_packed = layer.w2_weight_packed.data\n    w13_scale_hf = layer.w13_weight_scale.data\n    w2_scale_hf = layer.w2_weight_scale.data\n\n    if self.moe.is_act_and_mul and not torch.allclose(\n        layer.w13_weight_global_scale[:, 0], layer.w13_weight_global_scale[:, 1]\n    ):\n        logger.warning(\"w1_weight_global_scale must match w3_weight_global_scale. Accuracy may be affected.\")\n    w13_weight_global_scale = layer.w13_weight_global_scale[:, 0].contiguous()\n\n    w13_temp = Parameter(w13_packed.clone(), requires_grad=False)\n    w2_temp = Parameter(w2_packed.clone(), requires_grad=False)\n\n    if is_first_call:\n        layer.w13_weight = w13_temp\n        layer.w2_weight = w2_temp\n\n    (\n        w13,\n        w13_scale,\n        w13_scale_2,\n        a13_scale,\n        w2,\n        w2_scale,\n        w2_scale_2,\n        a2_scale,\n    ) = convert_to_nvfp4_moe_kernel_format(\n        nvfp4_backend=self.nvfp4_backend,\n        layer=layer,\n        w13=w13_temp,\n        w13_scale=w13_scale_hf,\n        w13_scale_2=(1.0 / w13_weight_global_scale),\n        a13_scale=(1.0 / layer.w13_input_global_scale),\n        w2=w2_temp,\n        w2_scale=w2_scale_hf,\n        w2_scale_2=(1.0 / layer.w2_weight_global_scale),\n        a2_scale=(1.0 / layer.w2_input_global_scale),\n        is_act_and_mul=self.moe.is_act_and_mul,\n    )\n\n    # Update parameters\n    if is_first_call:\n        replace_parameter(layer, \"w13_weight\", w13)\n        replace_parameter(layer, \"w2_weight\", w2)\n        layer.w13_weight_scale = Parameter(w13_scale, requires_grad=False)\n        layer.w2_weight_scale = Parameter(w2_scale, requires_grad=False)\n        if not hasattr(layer, \"_marlin_tensor_refs\"):\n            layer._marlin_tensor_refs = {}\n        layer._marlin_tensor_refs[\"w13_weight_scale\"] = layer.w13_weight_scale.data\n        layer._marlin_tensor_refs[\"w2_weight_scale\"] = layer.w2_weight_scale.data\n    else:\n        layer.w13_weight.data.copy_(w13.data)\n        layer.w2_weight.data.copy_(w2.data)\n        w13_scale_ref = layer._marlin_tensor_refs.get(\"w13_weight_scale\")\n        w2_scale_ref = layer._marlin_tensor_refs.get(\"w2_weight_scale\")\n        if w13_scale_ref is not None:\n            w13_scale_ref.copy_(w13_scale)\n            layer.w13_weight_scale = Parameter(w13_scale_ref, requires_grad=False)\n        else:\n            logger.warning(\"MoE W4A4: _marlin_tensor_refs['w13_weight_scale'] not found\")\n            layer.w13_weight_scale.data.copy_(w13_scale)\n        if w2_scale_ref is not None:\n            w2_scale_ref.copy_(w2_scale)\n            layer.w2_weight_scale = Parameter(w2_scale_ref, requires_grad=False)\n        else:\n            logger.warning(\"MoE W4A4: _marlin_tensor_refs['w2_weight_scale'] not found\")\n            layer.w2_weight_scale.data.copy_(w2_scale)\n\n    layer.w13_weight_scale_2 = w13_scale_2\n    layer.w2_weight_scale_2 = w2_scale_2\n    layer.w13_input_scale = a13_scale\n    layer.w2_input_scale = a2_scale\n\n    # Initialize kernel\n    self.moe_quant_config = self.get_fused_moe_quant_config(layer)\n    if self.moe_quant_config is not None and (\n        (not self.moe.moe_parallel_config.use_all2all_kernels) or self.moe.moe_parallel_config.use_naive_all2all_kernels\n    ):\n        self.kernel = make_nvfp4_moe_kernel(\n            moe_quant_config=self.moe_quant_config,\n            moe_config=self.moe,\n            experts_cls=self.experts_cls,\n        )\n\n\n# MoE NVFP4 Patches (entry points)\ndef patched_nvfp4_moe_process_weights_after_loading(self, layer: torch.nn.Module) -> None:\n    \"\"\"Patched process_weights_after_loading for NVFP4 MoE layer.\"\"\"\n    from vllm.model_executor.layers.fused_moe.oracle.nvfp4 import NvFp4MoeBackend\n\n    is_first_call = _check_first_call(layer)\n\n    # Save metadata (first call only)\n    if is_first_call:\n        save_param_meta(layer, \"w13_weight_packed\")\n        save_param_meta(layer, \"w2_weight_packed\")\n        save_param_meta(layer, \"w13_weight_scale\")\n        save_param_meta(layer, \"w2_weight_scale\")\n        if not hasattr(layer, \"_weight_loaders\"):\n            layer._weight_loaders = {}\n        for pname in [\"w13_weight_packed\", \"w2_weight_packed\", \"w13_weight_scale\", \"w2_weight_scale\"]:\n            param = getattr(layer, pname, None)\n            if param is not None and hasattr(param, \"weight_loader\"):\n                layer._weight_loaders[pname] = param.weight_loader\n\n    is_marlin = self.nvfp4_backend == NvFp4MoeBackend.MARLIN\n    if is_marlin:\n        _process_nvfp4_moe_marlin(self, layer, is_first_call)\n    else:\n        _process_nvfp4_moe_flashinfer_cutlass(self, layer, is_first_call)\n\n    # Delete HF parameters\n    if hasattr(layer, \"w13_weight_packed\"):\n        delattr(layer, \"w13_weight_packed\")\n    if hasattr(layer, \"w2_weight_packed\"):\n        delattr(layer, \"w2_weight_packed\")\n\n\n_PATCH_TARGETS = [\n    # Dense W4A16\n    (\n        \"vllm.model_executor.layers.quantization.compressed_tensors.schemes.\"\n        \"compressed_tensors_w4a16_nvfp4.CompressedTensorsW4A16Fp4.process_weights_after_loading\",\n        patched_w4a16_process_weights_after_loading,\n    ),\n    # Dense W4A4\n    (\n        \"vllm.model_executor.layers.quantization.compressed_tensors.schemes.\"\n        \"compressed_tensors_w4a4_nvfp4.CompressedTensorsW4A4Fp4.process_weights_after_loading\",\n        patched_w4a4_process_weights_after_loading,\n    ),\n    # MoE NVFP4\n    (\n        \"vllm.model_executor.layers.quantization.compressed_tensors.\"\n        \"compressed_tensors_moe.CompressedTensorsW4A4Nvfp4MoEMethod.process_weights_after_loading\",\n        patched_nvfp4_moe_process_weights_after_loading,\n    ),\n]\n\n_applied_patches = []\n\n\ndef apply_qat_patches():\n    \"\"\"Apply NVFP4 patches to support dynamic weight updates. Call before model loading.\"\"\"\n    global _applied_patches\n\n    if _applied_patches:\n        logger.warning(\"QAT patches already applied, skipping\")\n        return _applied_patches\n\n    logger.info(\"Applying NVFP4 patches for dynamic weight loading...\")\n\n    for target, replacement in _PATCH_TARGETS:\n        p = patch(target, replacement)\n        _applied_patches.append(p)\n        p.start()\n\n    logger.info(f\"Applied {len(_applied_patches)} NVFP4 patches for dynamic weight loading\")\n    return _applied_patches\n\n\ndef prepare_qat_for_load_weights(model, device=None):\n    \"\"\"\n    Prepare QAT model for weight loading. Call ONCE before multi-bucket weight loading.\n\n    Args:\n        model: vLLM model\n        device: Device for created parameters\n    \"\"\"\n    inner_model = model\n    if hasattr(model, \"model\"):\n        inner_model = model.model\n\n    param_meta = ParamMetaDict(inner_model, device=device)\n\n    param_meta.prepare_for_reload()\n    logger.info(f\"[prepare_qat] Tensor swap prepared for {len(param_meta._tensor_swap_layers)} layers\")\n\n    # Rebuild deleted (W4A16) or overwritten (W4A4) params back to HF format\n    rebuilt_count = 0\n    for layer_name, cache_entry in param_meta._layer_meta_cache.items():\n        module = cache_entry[\"module\"]\n        for param_name, pm in cache_entry[\"meta\"].items():\n            existing = getattr(module, param_name, None)\n            if existing is not None:\n                hf_shape = tuple(pm[\"shape\"])\n                hf_dtype = pm[\"dtype\"]\n                if (\n                    tuple(existing.shape) == hf_shape\n                    and existing.dtype == hf_dtype\n                    and hasattr(existing, \"weight_loader\")\n                ):\n                    continue\n            new_param = _create_param_from_meta(module, param_name, pm, device)\n            module.register_parameter(param_name, new_param)\n            rebuilt_count += 1\n\n    logger.info(f\"[prepare_qat] Rebuilt {rebuilt_count} parameters\")\n    inner_model._param_meta_for_restore = param_meta\n    return param_meta\n\n\ndef manual_process_weights_after_loading(model):\n    \"\"\"Trigger weight post-processing for all quantized layers after load_weights.\"\"\"\n    dense_count = 0\n    moe_count = 0\n\n    actual_model = model\n    if hasattr(model, \"model\"):\n        actual_model = model.model\n\n    for module in actual_model.modules():\n        if hasattr(module, \"scheme\"):\n            module.scheme.process_weights_after_loading(module)\n            dense_count += 1\n\n        quant_method = getattr(module, \"quant_method\", None)\n        if quant_method is not None and not hasattr(module, \"scheme\"):\n            if hasattr(quant_method, \"process_weights_after_loading\"):\n                # Skip KV cache quantization methods\n                if \"KVCache\" in quant_method.__class__.__name__:\n                    continue\n                quant_method.process_weights_after_loading(module)\n                moe_count += 1\n\n    logger.debug(f\"Processed {dense_count} dense layers, {moe_count} MoE layers\")\n    return dense_count + moe_count\n\n\n__all__ = [\n    \"apply_qat_patches\",\n    \"prepare_qat_for_load_weights\",\n    \"manual_process_weights_after_loading\",\n]\n"
  },
  {
    "path": "verl/utils/ray_utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nContains commonly used utilities for ray\n\"\"\"\n\nimport asyncio\nimport concurrent.futures\nimport functools\nimport inspect\nimport os\nfrom typing import Any, Optional\n\nimport ray\n\n\ndef ray_noset_visible_devices(env_vars=os.environ):\n    # Refer to\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/nvidia_gpu.py#L95-L96\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/amd_gpu.py#L102-L103\n    # https://github.com/ray-project/ray/blob/3b9e729f6a669ffd85190f901f5e262af79771b0/python/ray/_private/accelerators/amd_gpu.py#L114-L115\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/npu.py#L94-L95\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/hpu.py#L116-L117\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/neuron.py#L108-L109\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/tpu.py#L171-L172\n    # https://github.com/ray-project/ray/blob/161849364a784442cc659fb9780f1a6adee85fce/python/ray/_private/accelerators/intel_gpu.py#L97-L98\n    NOSET_VISIBLE_DEVICES_ENV_VARS_LIST = [\n        \"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES\",\n        \"RAY_EXPERIMENTAL_NOSET_ROCR_VISIBLE_DEVICES\",\n        \"RAY_EXPERIMENTAL_NOSET_HIP_VISIBLE_DEVICES\",\n        \"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES\",\n        \"RAY_EXPERIMENTAL_NOSET_HABANA_VISIBLE_MODULES\",\n        \"RAY_EXPERIMENTAL_NOSET_NEURON_RT_VISIBLE_CORES\",\n        \"RAY_EXPERIMENTAL_NOSET_TPU_VISIBLE_CHIPS\",\n        \"RAY_EXPERIMENTAL_NOSET_ONEAPI_DEVICE_SELECTOR\",\n    ]\n    return any(env_vars.get(env_var) for env_var in NOSET_VISIBLE_DEVICES_ENV_VARS_LIST)\n\n\ndef parallel_put(data_list: list[Any], max_workers: Optional[int] = None):\n    \"\"\"\n    Puts a list of data into the Ray object store in parallel using a thread pool.\n\n    Args:\n        data_list (List[Any]): A list of Python objects to be put into the Ray object store.\n        max_workers (int, optional): The maximum number of worker threads to use.\n                                     Defaults to min(len(data_list), 16).\n\n    Returns:\n        List[ray.ObjectRef]: A list of Ray object references corresponding to the input data_list,\n                             maintaining the original order.\n    \"\"\"\n    assert len(data_list) > 0, \"data_list must not be empty\"\n\n    def put_data(index, data):\n        return index, ray.put(data)\n\n    if max_workers is None:\n        max_workers = min(len(data_list), 16)\n\n    with concurrent.futures.ThreadPoolExecutor(max_workers=max_workers) as executor:\n        data_list_f = [executor.submit(put_data, i, data) for i, data in enumerate(data_list)]\n        res_lst = []\n        for future in concurrent.futures.as_completed(data_list_f):\n            res_lst.append(future.result())\n\n        # reorder based on index\n        output = [None for _ in range(len(data_list))]\n        for res in res_lst:\n            index, data_ref = res\n            output[index] = data_ref\n\n    return output\n\n\ndef get_event_loop():\n    try:\n        loop = asyncio.get_event_loop()\n    except RuntimeError:\n        loop = asyncio.new_event_loop()\n        asyncio.set_event_loop(loop)\n\n    return loop\n\n\ndef auto_await(func):\n    \"\"\"Auto await a coroutine function.\n\n    Handles three cases:\n    1. When the decorated function is called with await: returns the coroutine\n       so the caller can await it.\n    2. When called directly and there is no running event loop: runs the\n       coroutine with asyncio.run() and returns the result.\n    3. When called directly and the event loop is already running: runs the\n       coroutine (e.g. in a thread pool to avoid deadlock) and returns the result.\n    \"\"\"\n\n    @functools.wraps(func)\n    def wrapper(*args, **kwargs):\n        coro = func(*args, **kwargs)\n\n        if not inspect.iscoroutine(coro):\n            return coro\n\n        try:\n            loop = asyncio.get_running_loop()\n        except RuntimeError:\n            loop = None\n\n        # Case 1: No running loop -> run with asyncio.run()\n        if loop is None:\n            return asyncio.run(coro)\n\n        # Case 2: Running loop -> return coro if caller will await\n        caller_frame = inspect.currentframe()\n        if caller_frame is not None:\n            caller_frame = caller_frame.f_back\n        caller_is_async = caller_frame is not None and (caller_frame.f_code.co_flags & inspect.CO_COROUTINE) != 0\n        if caller_is_async:\n            return coro\n\n        # Case 3: Running loop -> run coro in thread pool\n        # (cannot block the loop thread without deadlock)\n        with concurrent.futures.ThreadPoolExecutor(max_workers=1) as pool:\n            future = pool.submit(asyncio.run, coro)\n            return future.result()\n\n    return wrapper\n"
  },
  {
    "path": "verl/utils/rendezvous/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/utils/rendezvous/ray_backend.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport time\n\nimport ray\nfrom cupy.cuda.nccl import NcclCommunicator, get_unique_id\nfrom ray.util import list_named_actors\n\n\n@ray.remote\nclass NCCLIDStore:\n    def __init__(self, nccl_id):\n        self._nccl_id = nccl_id\n\n    def get(self):\n        return self._nccl_id\n\n\ndef get_nccl_id_store_by_name(name):\n    all_actors = list_named_actors(all_namespaces=True)\n    matched_actors = [actor for actor in all_actors if actor.get(\"name\", None) == name]\n    if len(matched_actors) == 1:\n        actor = matched_actors[0]\n        return ray.get_actor(**actor)\n    elif len(matched_actors) > 1:\n        logging.warning(\"multiple actors with same name found: %s\", matched_actors)\n    elif len(matched_actors) == 0:\n        logging.info(\"failed to get any actor named %s\", name)\n    return None\n\n\ndef create_nccl_communicator_in_ray(\n    rank: int, world_size: int, group_name: str, max_retries: int = 100, interval_s: int = 5\n):\n    if rank == 0:\n        nccl_id = get_unique_id()\n        nccl_id_store = NCCLIDStore.options(name=group_name).remote(nccl_id)\n\n        assert ray.get(nccl_id_store.get.remote()) == nccl_id\n        communicator = NcclCommunicator(\n            ndev=world_size,\n            commId=nccl_id,\n            rank=0,\n        )\n        return communicator\n    else:\n        for i in range(max_retries):\n            nccl_id_store = get_nccl_id_store_by_name(group_name)\n            if nccl_id_store is not None:\n                logging.info(\"nccl_id_store %s got\", group_name)\n                nccl_id = ray.get(nccl_id_store.get.remote())\n                logging.info(\"nccl id for %s got: %s\", group_name, nccl_id)\n                communicator = NcclCommunicator(\n                    ndev=world_size,\n                    commId=nccl_id,\n                    rank=rank,\n                )\n                return communicator\n            logging.info(\"failed to get nccl_id for %d time, sleep for %d seconds\", i + 1, interval_s)\n            time.sleep(interval_s)\n"
  },
  {
    "path": "verl/utils/reward_score/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# from . import gsm8k, math, prime_math, prime_code\n\nfrom verl.utils.import_utils import deprecated\n\n\ndef default_compute_score(\n    data_source,\n    solution_str,\n    ground_truth,\n    extra_info=None,\n    sandbox_fusion_url=None,\n    concurrent_semaphore=None,\n    memory_limit_mb=None,\n    **kwargs,\n):\n    \"\"\"Compute the score for a given solution based on the data source.\n\n    Args:\n        data_source (str): The source dataset identifier which determines the scoring method.\n        solution_str (str): The solution string to be evaluated.\n        ground_truth (str): The ground truth answer for comparison.\n        extra_info (dict, optional): Additional information that might be needed for scoring. Defaults to None.\n\n    Returns:\n        float: The computed score as a floating point number. If the result is a dictionary,\n               it returns the dictionary instead.\n\n    Raises:\n        NotImplementedError: If the reward function is not implemented for the given data source.\n    \"\"\"\n    if data_source == \"openai/gsm8k\":\n        from . import gsm8k\n\n        res = gsm8k.compute_score(solution_str, ground_truth)\n    elif data_source in [\"lighteval/MATH\", \"DigitalLearningGmbH/MATH-lighteval\", \"HuggingFaceH4/MATH-500\"]:\n        from . import math_reward\n\n        res = math_reward.compute_score(solution_str, ground_truth)\n        # [Optional] Math-Verify Integration\n        # For enhanced accuracy, consider utilizing Math-Verify (https://github.com/huggingface/Math-Verify).\n        # Note: Math-Verify needs to be manually installed via pip: `pip install math-verify`.\n        # To use it, override the `compute_score` function with the following implementation:\n\n        # from . import math_verify\n        # res = math_verify.compute_score(solution_str, ground_truth)\n    elif data_source in [\"math_dapo\", \"math\", \"math_dapo_reasoning\"] or data_source.startswith(\"aime\"):\n        from . import math_dapo\n\n        res = math_dapo.compute_score(solution_str, ground_truth)\n    elif data_source in [\n        \"numina_aops_forum\",\n        \"numina_synthetic_math\",\n        \"numina_amc_aime\",\n        \"numina_synthetic_amc\",\n        \"numina_cn_k12\",\n        \"numina_olympiads\",\n    ]:\n        from . import prime_math\n\n        res = prime_math.compute_score(solution_str, ground_truth)\n    elif data_source in [\"codecontests\", \"apps\", \"codeforces\", \"taco\"]:\n        # Use the passed sandbox_fusion_url if available\n        if sandbox_fusion_url:\n            from . import sandbox_fusion\n\n            # Pass the URL directly, ground_truth likely contains test cases here\n            res = sandbox_fusion.compute_score(\n                sandbox_fusion_url, concurrent_semaphore, memory_limit_mb, solution_str, ground_truth, continuous=True\n            )\n        else:\n            # If no sandbox URL is provided, fall back to prime_code or raise error\n            from . import prime_code\n\n            # Assuming prime_code doesn't need the URL\n            res = prime_code.compute_score(solution_str, ground_truth, continuous=True)\n    elif data_source in [\"hiyouga/geometry3k\"]:\n        from . import geo3k\n\n        res = geo3k.compute_score(solution_str, ground_truth)\n    elif data_source in [\n        \"searchR1_nq\",\n        \"searchR1_triviaqa\",\n        \"searchR1_popqa\",\n        \"searchR1_hotpotqa\",\n        \"searchR1_2wikimultihopqa\",\n        \"searchR1_musique\",\n        \"searchR1_bamboogle\",\n    ]:\n        from . import search_r1_like_qa_em\n\n        res = search_r1_like_qa_em.compute_score(solution_str, ground_truth)\n\n    else:\n        raise NotImplementedError(f\"Reward function is not implemented for {data_source=}\")\n\n    if isinstance(res, dict):\n        return res\n    elif isinstance(res, int | float | bool):\n        return float(res)\n    else:\n        return float(res[0])\n\n\n@deprecated(\"verl.utils.reward_score.default_compute_score\")\ndef _default_compute_score(\n    data_source,\n    solution_str,\n    ground_truth,\n    extra_info=None,\n    sandbox_fusion_url=None,\n    concurrent_semaphore=None,\n    memory_limit_mb=None,\n):\n    \"\"\"\n    Legacy function API to be deprecated. Please use `default_compute_score` instead.\n    \"\"\"\n    return default_compute_score(\n        data_source, solution_str, ground_truth, extra_info, sandbox_fusion_url, concurrent_semaphore, memory_limit_mb\n    )\n\n\n__all__ = [\"default_compute_score\"]\n"
  },
  {
    "path": "verl/utils/reward_score/geo3k.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport re\n\nfrom mathruler.grader import extract_boxed_content, grade_answer\n\n\ndef format_reward(predict_str: str) -> float:\n    pattern = re.compile(r\"<think>.*</think>.*\\\\boxed\\{.*\\}.*\", re.DOTALL)\n    match_result = re.fullmatch(pattern, predict_str)\n    return 1.0 if match_result else 0.0\n\n\ndef acc_reward(predict_str: str, ground_truth: str, use_boxed: bool = True) -> float:\n    if use_boxed:\n        answer = extract_boxed_content(predict_str)\n    else:\n        answer = predict_str\n    return 1.0 if grade_answer(answer, ground_truth) else 0.0\n\n\ndef compute_score(predict_str: str, ground_truth: str, use_boxed: bool = True, format_score: float = 0.1) -> float:\n    return (1.0 - format_score) * acc_reward(predict_str, ground_truth, use_boxed) + format_score * format_reward(\n        predict_str\n    )\n"
  },
  {
    "path": "verl/utils/reward_score/gsm8k.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport re\n\n_SOLUTION_CLIP_CHARS = 300\n\n\ndef extract_solution(solution_str, method=\"strict\"):\n    assert method in [\"strict\", \"flexible\"]\n\n    # Optimization: Regular expression matching on very long strings can be slow.\n    # For math problems, the final answer is usually at the end.\n    # We only match on the last 300 characters, which is a safe approximation for 300 tokens.\n    if len(solution_str) > _SOLUTION_CLIP_CHARS:\n        solution_str = solution_str[-_SOLUTION_CLIP_CHARS:]\n\n    if method == \"strict\":\n        # this also tests the formatting of the model\n        solutions = re.findall(\"#### (\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n        if len(solutions) == 0:\n            final_answer = None\n        else:\n            # take the last solution\n            final_answer = solutions[-1].replace(\",\", \"\").replace(\"$\", \"\")\n    elif method == \"flexible\":\n        answer = re.findall(\"(\\\\-?[0-9\\\\.\\\\,]+)\", solution_str)\n        final_answer = None\n        if len(answer) == 0:\n            # no reward is there is no answer\n            pass\n        else:\n            invalid_str = [\"\", \".\"]\n            # find the last number that is not '.'\n            for final_answer in reversed(answer):\n                if final_answer not in invalid_str:\n                    break\n    return final_answer\n\n\ndef compute_score(solution_str, ground_truth, method=\"strict\", format_score=0.0, score=1.0):\n    \"\"\"The scoring function for GSM8k.\n\n    Reference: Trung, Luong, et al. \"Reft: Reasoning with reinforced fine-tuning.\" Proceedings of the 62nd Annual\n    Meeting of the Association for Computational Linguistics (Volume 1: Long Papers). 2024.\n\n    Args:\n        solution_str: the solution text\n        ground_truth: the ground truth\n        method: the method to extract the solution, choices are 'strict' and 'flexible'\n        format_score: the score for the format\n        score: the score for the correct answer\n    \"\"\"\n    answer = extract_solution(solution_str=solution_str, method=method)\n    if answer is None:\n        return 0\n    else:\n        if answer == ground_truth:\n            return score\n        else:\n            return format_score\n"
  },
  {
    "path": "verl/utils/reward_score/math_batch.py",
    "content": "# Copyright 2025 Individual Contributor: Mert Unsal\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .math_reward import compute_score\n\n\ndef compute_score_batched(data_sources, solution_strs, ground_truths, extra_infos):\n    \"\"\"\n    This is a demonstration of how the batched reward function should look like.\n    Typically, you want to use batched reward to speed up the process with parallelization\n    \"\"\"\n    return [\n        compute_score(solution_str, ground_truth)\n        for solution_str, ground_truth in zip(solution_strs, ground_truths, strict=True)\n    ]\n"
  },
  {
    "path": "verl/utils/reward_score/math_dapo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py\n\nimport re\nfrom typing import Optional\n\n\ndef last_boxed_only_string(string: str) -> Optional[str]:\n    \"\"\"Extract the last LaTeX boxed expression from a string.\n\n    Args:\n        string: Input string containing LaTeX code\n\n    Returns:\n        The last boxed expression or None if not found\n    \"\"\"\n    idx = string.rfind(\"\\\\boxed{\")\n    if idx < 0:\n        return None\n\n    i = idx\n    right_brace_idx = None\n    num_left_braces_open = 0\n\n    while i < len(string):\n        if string[i] == \"{\":\n            num_left_braces_open += 1\n        if string[i] == \"}\":\n            num_left_braces_open -= 1\n            if num_left_braces_open == 0:\n                right_brace_idx = i\n                break\n        i += 1\n\n    return string[idx : right_brace_idx + 1] if right_brace_idx is not None else None\n\n\ndef remove_boxed(s: str) -> str:\n    \"\"\"Remove the LaTeX boxed command from a string.\n\n    Args:\n        s: String with format \"\\\\boxed{content}\"\n\n    Returns:\n        The content inside the boxed command\n    \"\"\"\n    left = \"\\\\boxed{\"\n    assert s[: len(left)] == left, f\"box error: {s}\"\n    assert s[-1] == \"}\", f\"box error: {s}\"\n    return s[len(left) : -1]\n\n\n# Constants for normalization\nSUBSTITUTIONS = [\n    (\"an \", \"\"),\n    (\"a \", \"\"),\n    (\".$\", \"$\"),\n    (\"\\\\$\", \"\"),\n    (r\"\\ \", \"\"),\n    (\" \", \"\"),\n    (\"mbox\", \"text\"),\n    (\",\\\\text{and}\", \",\"),\n    (\"\\\\text{and}\", \",\"),\n    (\"\\\\text{m}\", \"\\\\text{}\"),\n]\n\nREMOVED_EXPRESSIONS = [\n    \"square\",\n    \"ways\",\n    \"integers\",\n    \"dollars\",\n    \"mph\",\n    \"inches\",\n    \"hours\",\n    \"km\",\n    \"units\",\n    \"\\\\ldots\",\n    \"sue\",\n    \"points\",\n    \"feet\",\n    \"minutes\",\n    \"digits\",\n    \"cents\",\n    \"degrees\",\n    \"cm\",\n    \"gm\",\n    \"pounds\",\n    \"meters\",\n    \"meals\",\n    \"edges\",\n    \"students\",\n    \"childrentickets\",\n    \"multiples\",\n    \"\\\\text{s}\",\n    \"\\\\text{.}\",\n    \"\\\\text{\\ns}\",\n    \"\\\\text{}^2\",\n    \"\\\\text{}^3\",\n    \"\\\\text{\\n}\",\n    \"\\\\text{}\",\n    r\"\\mathrm{th}\",\n    r\"^\\circ\",\n    r\"^{\\circ}\",\n    r\"\\;\",\n    r\",\\!\",\n    \"{,}\",\n    '\"',\n    \"\\\\dots\",\n]\n\n\ndef normalize_final_answer(final_answer: str) -> str:\n    \"\"\"Normalize a final answer to a quantitative reasoning question.\n\n    Args:\n        final_answer: The answer string to normalize\n\n    Returns:\n        Normalized answer string\n    \"\"\"\n    final_answer = final_answer.split(\"=\")[-1]\n\n    # Apply substitutions and removals\n    for before, after in SUBSTITUTIONS:\n        final_answer = final_answer.replace(before, after)\n    for expr in REMOVED_EXPRESSIONS:\n        final_answer = final_answer.replace(expr, \"\")\n\n    # Extract and normalize LaTeX math\n    final_answer = re.sub(r\"(.*?)(\\$)(.*?)(\\$)(.*)\", \"$\\\\3$\", final_answer)\n    final_answer = re.sub(r\"(\\\\text\\{)(.*?)(\\})\", \"\\\\2\", final_answer)\n    final_answer = re.sub(r\"(\\\\textbf\\{)(.*?)(\\})\", \"\\\\2\", final_answer)\n    final_answer = re.sub(r\"(\\\\overline\\{)(.*?)(\\})\", \"\\\\2\", final_answer)\n    final_answer = re.sub(r\"(\\\\boxed\\{)(.*)(\\})\", \"\\\\2\", final_answer)\n\n    # Normalize shorthand TeX:\n    #  \\fracab -> \\frac{a}{b}\n    #  \\frac{abc}{bef} -> \\frac{abc}{bef}\n    #  \\fracabc -> \\frac{a}{b}c\n    #  \\sqrta -> \\sqrt{a}\n    #  \\sqrtab -> sqrt{a}b\n    final_answer = re.sub(r\"(frac)([^{])(.)\", \"frac{\\\\2}{\\\\3}\", final_answer)\n    final_answer = re.sub(r\"(sqrt)([^{])\", \"sqrt{\\\\2}\", final_answer)\n    final_answer = final_answer.replace(\"$\", \"\")\n\n    # Normalize numbers\n    if final_answer.replace(\",\", \"\").isdigit():\n        final_answer = final_answer.replace(\",\", \"\")\n\n    return final_answer.strip()\n\n\ndef is_correct_minerva(\n    solution_str: str, gt: str, gt_need_extract: bool = False, answer_pattern: str = r\"(?i)Answer\\s*:\\s*([^\\n]+)\"\n) -> tuple[bool, str]:\n    \"\"\"Check if the solution is correct according to Minerva criteria.\n\n    Args:\n        solution_str: The solution string to check\n        gt: The ground truth answer\n        gt_need_extract: Whether the ground truth needs extraction\n        answer_pattern: Regex pattern to extract the answer\n\n    Returns:\n        Tuple of (is_correct, normalized_prediction)\n    \"\"\"\n    # Extract answer from solution\n    match = re.findall(answer_pattern, solution_str)\n    extracted_answer = match[-1] if match else \"[INVALID]\"\n    pred = normalize_final_answer(extracted_answer)\n\n    # Process ground truth\n    if gt_need_extract:\n        gt = normalize_final_answer(remove_boxed(last_boxed_only_string(gt)))\n    else:\n        gt = normalize_final_answer(gt)\n\n    return (pred == gt), pred\n\n\ndef is_correct_strict_box(\n    pred: str, gt: str, pause_tokens_index: Optional[list[int]] = None\n) -> tuple[int, Optional[str]]:\n    \"\"\"Check if the prediction is correct using strict boxed answer criteria.\n\n    Args:\n        pred: The prediction string\n        gt: The ground truth answer\n        pause_tokens_index: Indices of pause tokens\n\n    Returns:\n        Tuple of (score, extracted_prediction)\n    \"\"\"\n    # Extract the relevant part of the prediction\n    if pause_tokens_index is not None:\n        assert len(pause_tokens_index) == 4\n        pred = pred[pause_tokens_index[-1] - 100 :]\n    else:\n        pred = pred[-100:]\n\n    # Extract and check the boxed answer\n    boxed_pred = last_boxed_only_string(pred)\n    extracted_pred = remove_boxed(boxed_pred) if boxed_pred is not None else None\n\n    return 1 if (extracted_pred == gt) else -1, extracted_pred\n\n\ndef verify(\n    solution_str: str, answer: str, strict_box_verify: bool = False, pause_tokens_index: Optional[list[int]] = None\n) -> bool:\n    \"\"\"Verify if the solution is correct.\n\n    Args:\n        solution_str: The solution string to verify\n        answer: The ground truth answer\n        strict_box_verify: Whether to use strict box verification\n        pause_tokens_index: Indices of pause tokens\n\n    Returns:\n        True if the solution is correct, False otherwise\n    \"\"\"\n    if strict_box_verify:\n        correct, pred = is_correct_strict_box(solution_str, answer, pause_tokens_index)\n        return correct == 1, pred\n\n    correct, pred = is_correct_minerva(solution_str, answer)\n    return correct, pred\n\n\ndef compute_score(\n    solution_str: str,\n    ground_truth: str,\n    strict_box_verify: bool = False,\n    pause_tokens_index: Optional[list[int]] = None,\n) -> float:\n    \"\"\"Compute the reward score for a solution.\n\n    Args:\n        solution_str: The solution string\n        ground_truth: The ground truth answer\n        strict_box_verify: Whether to use strict box verification\n        pause_tokens_index: Indices of pause tokens\n\n    Returns:\n        Reward score (1.0 for correct, -1.0 for incorrect)\n    \"\"\"\n    # Limit solution length for efficiency\n    solution_str = solution_str[-300:]  # The longest answer in MATH-500 has 159 characters\n\n    # Verify the solution\n    correct, pred = verify(solution_str, ground_truth, strict_box_verify, pause_tokens_index)\n\n    reward = 1.0 if correct else -1.0\n    acc = correct\n\n    return {\n        \"score\": reward,\n        \"acc\": acc,\n        \"pred\": pred,\n    }\n"
  },
  {
    "path": "verl/utils/reward_score/math_reward.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2022 EleutherAI and the HuggingFace Inc. team. All rights reserved.\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n# Adapted from https://github.com/EleutherAI/lm-evaluation-harness/blob/main/lm_eval/tasks/hendrycks_math/utils.py\n\n\ndef compute_score(solution_str, ground_truth) -> float:\n    retval = 0.0\n    try:\n        string_in_last_boxed = last_boxed_only_string(solution_str)\n        if string_in_last_boxed is not None:\n            answer = remove_boxed(string_in_last_boxed)\n            if is_equiv(answer, ground_truth):\n                retval = 1.0\n    except Exception as e:\n        print(e)\n\n    return retval\n\n\n# string normalization from https://github.com/EleutherAI/lm-evaluation-harness/blob/master/lm_eval/tasks/hendrycks_math.py\ndef is_equiv(str1, str2, verbose=False):\n    if str1 is None and str2 is None:\n        print(\"WARNING: Both None\")\n        return True\n    if str1 is None or str2 is None:\n        return False\n\n    try:\n        ss1 = strip_string(str1)\n        ss2 = strip_string(str2)\n        if verbose:\n            print(ss1, ss2)\n        return ss1 == ss2\n    except Exception:\n        return str1 == str2\n\n\ndef remove_boxed(s):\n    if \"\\\\boxed \" in s:\n        left = \"\\\\boxed \"\n        assert s[: len(left)] == left\n        return s[len(left) :]\n\n    left = \"\\\\boxed{\"\n\n    assert s[: len(left)] == left\n    assert s[-1] == \"}\"\n\n    return s[len(left) : -1]\n\n\ndef last_boxed_only_string(string):\n    idx = string.rfind(\"\\\\boxed\")\n    if \"\\\\boxed \" in string:\n        return \"\\\\boxed \" + string.split(\"\\\\boxed \")[-1].split(\"$\")[0]\n    if idx < 0:\n        idx = string.rfind(\"\\\\fbox\")\n        if idx < 0:\n            return None\n\n    i = idx\n    right_brace_idx = None\n    num_left_braces_open = 0\n    while i < len(string):\n        if string[i] == \"{\":\n            num_left_braces_open += 1\n        if string[i] == \"}\":\n            num_left_braces_open -= 1\n            if num_left_braces_open == 0:\n                right_brace_idx = i\n                break\n        i += 1\n\n    retval = None if right_brace_idx is None else string[idx : right_brace_idx + 1]\n\n    return retval\n\n\ndef fix_fracs(string):\n    substrs = string.split(\"\\\\frac\")\n    new_str = substrs[0]\n    if len(substrs) > 1:\n        substrs = substrs[1:]\n        for substr in substrs:\n            new_str += \"\\\\frac\"\n            if substr[0] == \"{\":\n                new_str += substr\n            else:\n                try:\n                    assert len(substr) >= 2\n                except Exception:\n                    return string\n                a = substr[0]\n                b = substr[1]\n                if b != \"{\":\n                    if len(substr) > 2:\n                        post_substr = substr[2:]\n                        new_str += \"{\" + a + \"}{\" + b + \"}\" + post_substr\n                    else:\n                        new_str += \"{\" + a + \"}{\" + b + \"}\"\n                else:\n                    if len(substr) > 2:\n                        post_substr = substr[2:]\n                        new_str += \"{\" + a + \"}\" + b + post_substr\n                    else:\n                        new_str += \"{\" + a + \"}\" + b\n    string = new_str\n    return string\n\n\ndef fix_a_slash_b(string):\n    if len(string.split(\"/\")) != 2:\n        return string\n    a = string.split(\"/\")[0]\n    b = string.split(\"/\")[1]\n    try:\n        a = int(a)\n        b = int(b)\n        assert string == \"{}/{}\".format(a, b)\n        new_string = \"\\\\frac{\" + str(a) + \"}{\" + str(b) + \"}\"\n        return new_string\n    except Exception:\n        return string\n\n\ndef remove_right_units(string):\n    # \"\\\\text{ \" only ever occurs (at least in the val set) when describing units\n    if \"\\\\text{ \" in string:\n        splits = string.split(\"\\\\text{ \")\n        assert len(splits) == 2\n        return splits[0]\n    else:\n        return string\n\n\ndef fix_sqrt(string):\n    if \"\\\\sqrt\" not in string:\n        return string\n    splits = string.split(\"\\\\sqrt\")\n    new_string = splits[0]\n    for split in splits[1:]:\n        if split[0] != \"{\":\n            a = split[0]\n            new_substr = \"\\\\sqrt{\" + a + \"}\" + split[1:]\n        else:\n            new_substr = \"\\\\sqrt\" + split\n        new_string += new_substr\n    return new_string\n\n\ndef strip_string(string):\n    # linebreaks\n    string = string.replace(\"\\n\", \"\")\n\n    # remove inverse spaces\n    string = string.replace(\"\\\\!\", \"\")\n\n    # replace \\\\ with \\\n    string = string.replace(\"\\\\\\\\\", \"\\\\\")\n\n    # replace tfrac and dfrac with frac\n    string = string.replace(\"tfrac\", \"frac\")\n    string = string.replace(\"dfrac\", \"frac\")\n\n    # remove \\left and \\right\n    string = string.replace(\"\\\\left\", \"\")\n    string = string.replace(\"\\\\right\", \"\")\n\n    # Remove circ (degrees)\n    string = string.replace(\"^{\\\\circ}\", \"\")\n    string = string.replace(\"^\\\\circ\", \"\")\n\n    # remove dollar signs\n    string = string.replace(\"\\\\$\", \"\")\n\n    # remove units (on the right)\n    string = remove_right_units(string)\n\n    # remove percentage\n    string = string.replace(\"\\\\\\\\%\", \"\")\n    string = string.replace(\"\\\\%\", \"\")\n\n    # \" 0.\" equivalent to \" .\" and \"{0.\" equivalent to \"{.\" Alternatively, add \"0\" if \".\" is the start of the string\n    string = string.replace(\" .\", \" 0.\")\n    string = string.replace(\"{.\", \"{0.\")\n    # if empty, return empty string\n    if len(string) == 0:\n        return string\n    if string[0] == \".\":\n        string = \"0\" + string\n\n    # to consider: get rid of e.g. \"k = \" or \"q = \" at beginning\n    if len(string.split(\"=\")) == 2 and len(string.split(\"=\")[0]) <= 2:\n        string = string.split(\"=\")[1]\n\n    # fix sqrt3 --> sqrt{3}\n    string = fix_sqrt(string)\n\n    # remove spaces\n    string = string.replace(\" \", \"\")\n\n    # \\frac1b or \\frac12 --> \\frac{1}{b} and \\frac{1}{2}, etc. Even works with \\frac1{72} (but not \\frac{72}1).\n    # Also does a/b --> \\\\frac{a}{b}\n    string = fix_fracs(string)\n\n    # manually change 0.5 --> \\frac{1}{2}\n    if string == \"0.5\":\n        string = \"\\\\frac{1}{2}\"\n\n    # NOTE: X/Y changed to \\frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y\n    string = fix_a_slash_b(string)\n\n    return string\n"
  },
  {
    "path": "verl/utils/reward_score/math_verify.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\ntry:\n    from math_verify.errors import TimeoutException\n    from math_verify.metric import math_metric\n    from math_verify.parser import ExprExtractionConfig, LatexExtractionConfig\nexcept ImportError:\n    print(\"To use Math-Verify, please install it first by running `pip install math-verify`.\")\n\n\ndef compute_score(model_output: str, ground_truth: str, timeout_score: float = 0) -> bool:\n    verify_func = math_metric(\n        gold_extraction_target=(LatexExtractionConfig(),),\n        pred_extraction_target=(ExprExtractionConfig(), LatexExtractionConfig()),\n    )\n    ret_score = 0.0\n\n    # Wrap the ground truth in \\boxed{} format for verification\n    ground_truth_boxed = \"\\\\boxed{\" + ground_truth + \"}\"\n    try:\n        ret_score, _ = verify_func([ground_truth_boxed], [model_output])\n    except Exception:\n        pass\n    except TimeoutException:\n        ret_score = timeout_score\n\n    return ret_score\n"
  },
  {
    "path": "verl/utils/reward_score/prime_code/README.md",
    "content": "## LiveCodeBench\n\n### Introduction\n[LiveCodeBench](https://github.com/LiveCodeBench/LiveCodeBench) provides holistic and contamination-free evaluation of coding capabilities of LLMs. Particularly, LiveCodeBench continuously collects new problems over time from contests across three competition platforms -- LeetCode, AtCoder, and CodeForces. \n\n### How to reproduce\nOur evaluation is grounded on the version found in LiveCodeBench.\n> **Installation**\n```bash\n# Make sure the CUDA version > 12.0.\npip install -r requirements.txt\npip install flash-attn --no-build-isolation\n```\n\n### Acknowleage\nThank you to the [LiveCodeBench](https://livecodebench.github.io/leaderboard.html) team for their contributions to the open-source community."
  },
  {
    "path": "verl/utils/reward_score/prime_code/__init__.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport traceback\n\nfrom .utils import check_correctness as apps_check_correctness\n\n\ndef compute_score(completion, test_cases, continuous=False):\n    # try to get code solution from completion. if the completion is pure code, this will not take effect.\n    solution = completion.split(\"```python\")[-1].split(\"```\")[0]\n    try:\n        try:\n            if not isinstance(test_cases, dict):\n                test_cases = json.loads(test_cases)\n        except Exception as e:\n            print(f\"Error:{e}\")\n\n        # Complete check on all in-out pairs first. If there is no failure, per-sample test can be skipped.\n        try:\n            res, metadata = apps_check_correctness(in_outs=test_cases, generation=solution, timeout=5, debug=False)\n            metadata = dict(enumerate(metadata))[0]\n            success = all(map(lambda x: x is True, res))\n            if success:\n                return success, metadata\n        except Exception:\n            pass\n\n        test_cases_list = []\n        inputs = test_cases[\"inputs\"]\n        outputs = test_cases[\"outputs\"]\n        for i in range(len(inputs)):\n            test_cases_list.append({\"inputs\": [inputs[i]], \"outputs\": [outputs[i]]})\n\n        if continuous:\n            # per sample test: if continuous score is needed, test first 10 samples regardless of failures\n            # do not test all samples cuz some problems have enormous test cases\n            metadata_list = []\n            res_list = []\n            for test_case_id, test_case in enumerate(test_cases_list):\n                res, metadata = apps_check_correctness(in_outs=test_case, generation=solution, timeout=10, debug=False)\n                try:\n                    metadata = dict(enumerate(metadata))[0]  # metadata can be empty occasionally\n                except Exception:\n                    metadata = {}\n                metadata[\"test_case\"] = {}\n                metadata[\"test_case\"][\"input\"] = str(test_case[\"inputs\"][0])\n                metadata[\"test_case\"][\"output\"] = str(test_case[\"outputs\"][0])\n                metadata[\"test_case\"][\"res\"] = str(res)\n                metadata_list.append(metadata)\n                res_list.extend(res)\n\n                if test_case_id >= 9:\n                    break\n            res_count = len(res_list) if len(res_list) > 0 else 1\n            success = sum(map(lambda x: x is True, res_list)) / res_count\n    except Exception:\n        traceback.print_exc(10)\n        success = False\n        metadata_list = None\n    return success, metadata_list\n"
  },
  {
    "path": "verl/utils/reward_score/prime_code/testing_util.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport ast\nimport faulthandler\nimport json\nimport platform\n\n# to run the solution files we're using a timing based approach\nimport signal\nimport sys\nimport traceback\n\n# used for debugging to time steps\nfrom datetime import datetime\nfrom enum import Enum\n\n# for capturing the stdout\nfrom io import StringIO\n\n# used for testing the code that reads from input\nfrom unittest.mock import mock_open, patch\n\nimport numpy as np\nfrom pyext import RuntimeModule\n\n\ndef truncatefn(s, length=300):\n    assert isinstance(s, str)\n    if len(s) <= length:\n        return s\n\n    return s[: length // 2] + \"...(truncated) ...\" + s[-length // 2 :]\n\n\nclass CODE_TYPE(Enum):\n    call_based = 0\n    standard_input = 1\n\n\n# used to capture stdout as a list\n# from https://stackoverflow.com/a/16571630/6416660\n# alternative use redirect_stdout() from contextlib\nclass Capturing(list):\n    def __enter__(self):\n        self._stdout = sys.stdout\n        sys.stdout = self._stringio = StringIO()\n        # Make closing the StringIO a no-op\n        self._stringio.close = lambda x: 1\n        return self\n\n    def __exit__(self, *args):\n        self.append(self._stringio.getvalue())\n        del self._stringio  # free up some memory\n        sys.stdout = self._stdout\n\n\ndef only_int_check(val):\n    return isinstance(val, int)\n\n\ndef string_int_check(val):\n    return isinstance(val, str) and val.isdigit()\n\n\ndef combined_int_check(val):\n    return only_int_check(val) or string_int_check(val)\n\n\ndef clean_traceback(error_traceback):\n    file_start = error_traceback.find('File \"<string>\"')\n    # print(file_start)\n    error_traceback = \"Traceback (most recent call last):\\n  \" + error_traceback[file_start:]\n    return error_traceback\n\n\ndef run_test(in_outs, test=None, debug=False, timeout=15):\n    \"\"\"\n    if test(generated_code) is not None it'll try to run the code.\n    otherwise it'll just return an input and output pair.\n    \"\"\"\n    # Disable functionalities that can make destructive changes to the test.\n    reliability_guard()\n\n    if debug:\n        print(f\"start = {datetime.now().time()}\")\n\n    if in_outs:\n        if in_outs.get(\"fn_name\") is None:\n            which_type = CODE_TYPE.standard_input  # Standard input\n            method_name = None\n        else:\n            which_type = CODE_TYPE.call_based  # Call-based\n            method_name = in_outs[\"fn_name\"]\n\n    if debug:\n        print(f\"loaded input_output = {datetime.now().time()}\")\n\n    if test is None:\n        raise AssertionError(\"should not happen: test code is none\")\n    elif test is not None:\n        results = []\n        sol = \"from string import *\\nfrom re import *\\nfrom datetime import *\\nfrom collections import *\\nfrom heapq import *\\nfrom bisect import *\\nfrom copy import *\\nfrom math import *\\nfrom random import *\\nfrom statistics import *\\nfrom itertools import *\\nfrom functools import *\\nfrom operator import *\\nfrom io import *\\nfrom sys import *\\nfrom json import *\\nfrom builtins import *\\nfrom typing import *\\nimport string\\nimport re\\nimport datetime\\nimport collections\\nimport heapq\\nimport bisect\\nimport copy\\nimport math\\nimport random\\nimport statistics\\nimport itertools\\nimport functools\\nimport operator\\nimport io\\nimport sys\\nimport json\\nsys.setrecursionlimit(6*10**5)\\n\"  # noqa: E501\n        if debug:\n            print(f\"loading test code = {datetime.now().time()}\")\n\n        if which_type == CODE_TYPE.call_based:\n            sol += test\n            if debug:\n                print(f\"sol = {sol}\")\n            signal.alarm(timeout)\n            try:\n                tmp_sol = RuntimeModule.from_string(\"tmp_sol\", \"\", sol)\n                tmp = tmp_sol if \"class Solution\" not in test else tmp_sol.Solution()\n                signal.alarm(0)\n            except Exception as e:\n                signal.alarm(0)\n                error_traceback = traceback.format_exc()\n                if debug:\n                    print(f\"type 0 compilation error = {e}\")\n                results.append(-2)\n                return results, {\n                    \"error\": repr(e),\n                    # \"error_code\": -1,\n                    # \"error_message\": \"Compilation Error\",\n                    \"traceback\": clean_traceback(error_traceback),\n                }\n            signal.alarm(0)\n\n        elif which_type == CODE_TYPE.standard_input:\n            # sol\n            # if code has if __name__ == \"__main__\": then remove it\n            try:\n                astree = ast.parse(test)\n                last_block = astree.body[-1]\n                if isinstance(last_block, ast.If):\n                    condition = last_block.test\n                    if ast.unparse(condition).strip() == \"__name__ == '__main__'\":\n                        test = ast.unparse(astree.body[:-1]) + \"\\n\" + ast.unparse(last_block.body)\n            except Exception:\n                pass\n\n            tmp_test = test.split(\"\\n\")\n\n            new_test = []\n            for x in tmp_test:\n                if (not x.startswith(\"from \")) and (not x.startswith(\"import \")):\n                    new_test.append(\"\\t\" + x + \"\\n\")\n                else:\n                    new_test.append(x + \"\\n\")\n            tmp_test = new_test\n\n            new_test = \"\"\n            started = False\n            for i in tmp_test:\n                if i.startswith(\"\\t\") and not started:\n                    new_test += \"stdin = sys.stdin\\nstdout = sys.stdout\\n\"\n                    new_test += \"def code():\\n\"\n                    new_test += i\n                    started = True\n                elif started and ((i.startswith(\"from \")) or (i.startswith(\"import \"))):\n                    new_test += \"\\t\" + i\n                else:\n                    new_test += i\n            tmp_test = new_test\n\n            sol += tmp_test\n            if debug:\n                print(f\"sol = {sol}\")\n            method_name = \"code\"\n            signal.alarm(timeout)\n            try:\n                tmp_sol = RuntimeModule.from_string(\"tmp_sol\", \"\", sol)\n                tmp = tmp_sol\n                signal.alarm(0)\n            except Exception as e:\n                signal.alarm(0)\n                error_traceback = traceback.format_exc()\n                if debug:\n                    print(f\"type 1 compilation error = {e}\")\n                results.append(-2)\n                return results, {\n                    \"error\": repr(e),\n                    # \"error_code\": -1,\n                    # \"error_message\": \"Compilation Error\",\n                    \"traceback\": clean_traceback(error_traceback),\n                }\n            signal.alarm(0)\n        if debug:\n            print(f\"get method = {datetime.now().time()}\")\n\n        try:\n            method = getattr(tmp, method_name)  # get_attr second arg must be str\n        except Exception:\n            signal.alarm(0)\n            error_traceback = traceback.format_exc()\n            error_info = sys.exc_info()\n            print(f\"unable to get function error = {error_info}\")\n            results.append(-2)\n            return results, {\n                \"error\": repr(error_info),\n                # \"error_code\": -1,\n                # \"error_message\": \"Unable to extract code\",\n                \"traceback\": clean_traceback(error_traceback),\n            }\n\n        for index, inputs in enumerate(in_outs[\"inputs\"]):\n            raw_inputs = inputs\n            raw_outputs = in_outs[\"outputs\"][index]\n            if which_type == CODE_TYPE.call_based:\n                inputs = [json.loads(line) for line in inputs.split(\"\\n\")]\n                in_outs[\"outputs\"][index] = json.loads(in_outs[\"outputs\"][index])\n\n                truncate_line_size = 300 // (raw_inputs.count(\"\\n\") + 1)\n                raw_inputs = \"\\n\".join(\n                    [truncatefn(line, truncate_line_size) for line in raw_inputs.strip().split(\"\\n\")]\n                )\n                raw_outputs = truncatefn(raw_outputs, 200)\n            else:\n                raw_inputs = truncatefn(raw_inputs)\n                raw_outputs = truncatefn(raw_outputs, 200)\n            # JSON forces dictionaries to have string keys; this undoes this (assuming a singleton list)\n            try:\n                if isinstance(inputs[0], dict):\n                    inputs = [{int(k): v for k, v in inputs[0].items()}]\n            except Exception:\n                pass\n            try:\n                if isinstance(in_outs[\"outputs\"][index], dict):\n                    in_outs[\"outputs\"][index] = [{int(k): v for k, v in in_outs[\"outputs\"][index].items()}]\n            except Exception:\n                pass\n            try:\n                if isinstance(in_outs[\"outputs\"][index][0], dict):\n                    in_outs[\"outputs\"][index] = [{int(k): v for k, v in in_outs[\"outputs\"][index][0].items()}]\n            except Exception:\n                pass\n\n            if debug:\n                print(\n                    f\"time: {datetime.now().time()} testing index = {index}  inputs = {inputs}, {type(inputs)}. \"\n                    f\"type = {which_type}\"\n                )\n            if which_type == CODE_TYPE.call_based:  # Call-based\n                signal.alarm(timeout)\n                faulthandler.enable()\n                try:\n                    output = method(*inputs)\n                    raw_true_output = output\n\n                    raw_true_output_copy = json.dumps(output)\n                    raw_true_output_copy = truncatefn(raw_true_output_copy, 200)\n\n                    # ground truth sequences are not tuples\n                    if isinstance(output, tuple):\n                        output = list(output)\n\n                    tmp_result = output == in_outs[\"outputs\"][index]\n                    if isinstance(in_outs[\"outputs\"][index], list) and in_outs[\"outputs\"][index]:\n                        tmp_result = tmp_result or (output == in_outs[\"outputs\"][index][0])\n\n                    # ground truth sequences are not tuples\n                    try:\n                        if isinstance(output[0], tuple):\n                            tmp_result = tmp_result or ([list(x) for x in output] == in_outs[\"outputs\"][index][0])\n                    except Exception:\n                        pass\n                    results.append(tmp_result)\n                    if tmp_result is not True:\n                        return results, {\n                            \"output\": raw_true_output_copy,\n                            \"expected\": raw_outputs,\n                            \"inputs\": raw_inputs,\n                            # \"error_code\": -2,\n                            \"error_message\": \"Wrong Answer\",\n                        }\n                    # reset the alarm\n                    signal.alarm(0)\n                except Exception as e:\n                    signal.alarm(0)\n                    error_traceback = traceback.format_exc()\n                    faulthandler.disable()\n                    if debug:\n                        print(f\"Standard input runtime error or time limit exceeded error = {e}\")\n                    results.append(-1)\n                    return results, {\n                        \"error\": repr(e),\n                        \"traceback\": clean_traceback(error_traceback),\n                    }\n                faulthandler.disable()\n                signal.alarm(0)\n                if debug:\n                    print(\n                        f\"outputs = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, \"\n                        f\"{type(inputs)}, {output == [in_outs['outputs'][index]]}\"\n                    )\n            elif which_type == CODE_TYPE.standard_input:  # Standard input\n                faulthandler.enable()\n                passed = False\n\n                if isinstance(inputs, list):\n                    inputs = \"\\n\".join(inputs)\n                if isinstance(in_outs[\"outputs\"][index], list):\n                    in_outs[\"outputs\"][index] = \"\\n\".join(in_outs[\"outputs\"][index])\n\n                signal.alarm(timeout)\n                with Capturing() as output:\n                    try:\n                        call_method(method, inputs)\n                        # reset the alarm\n                        signal.alarm(0)\n                        passed = True\n                    except Exception as e:\n                        # runtime error or took too long\n                        signal.alarm(0)\n                        error_traceback = traceback.format_exc()\n                        print(f\"Call-based runtime error or time limit exceeded error = {repr(e)}{e}\")\n                        results.append(-1)\n                        return results, {\n                            \"error\": repr(e),\n                            \"traceback\": clean_traceback(error_traceback),\n                        }\n                    signal.alarm(0)\n                raw_true_output = output[0]\n                raw_true_output_copy = truncatefn(raw_true_output, 200)\n                output = raw_true_output.splitlines()\n                if not passed:\n                    if debug:\n                        nl = \"\\n\"\n                        if not isinstance(inputs, list):\n                            print(\n                                f\"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, \"\n                                f\"inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, \"\n                                f\"{output == [in_outs['outputs'][index]]}\"\n                            )\n                        else:\n                            print(\n                                f\"not passed output = {output}, test outputs = {in_outs['outputs'][index]}, \"\n                                f\"inputs = {inputs}, {type(inputs)}, {output == [in_outs['outputs'][index]]}\"\n                            )\n                    continue\n\n                if passed and debug:\n                    print(f\"==> output = {output}, test outputs = {in_outs['outputs'][index]}\")\n\n                if custom_compare_(output, in_outs[\"outputs\"][index]):\n                    tmp_result = True\n                    results.append(tmp_result)\n                    continue\n\n                # ground truth sequences are expressed as lists not tuples\n                if isinstance(output, tuple):\n                    output = list(output)\n\n                tmp_result = False\n                try:\n                    tmp_result = output == [in_outs[\"outputs\"][index]]\n                    if isinstance(in_outs[\"outputs\"][index], list):\n                        tmp_result = tmp_result or (output == in_outs[\"outputs\"][index])\n                        if isinstance(output[0], str):\n                            tmp_result = tmp_result or ([e.strip() for e in output] == in_outs[\"outputs\"][index])\n                except Exception as e:\n                    if debug:\n                        print(f\"Failed check1 exception = {e}\")\n                    pass\n\n                if tmp_result is True:\n                    results.append(tmp_result)\n                    continue\n\n                # try one more time without \\n\n                if isinstance(in_outs[\"outputs\"][index], list):\n                    for tmp_index, i in enumerate(in_outs[\"outputs\"][index]):\n                        in_outs[\"outputs\"][index][tmp_index] = i.split(\"\\n\")\n                        in_outs[\"outputs\"][index][tmp_index] = [\n                            x.strip() for x in in_outs[\"outputs\"][index][tmp_index] if x\n                        ]\n                else:\n                    in_outs[\"outputs\"][index] = in_outs[\"outputs\"][index].split(\"\\n\")\n                    in_outs[\"outputs\"][index] = list(filter(len, in_outs[\"outputs\"][index]))\n                    in_outs[\"outputs\"][index] = list(map(lambda x: x.strip(), in_outs[\"outputs\"][index]))\n\n                try:\n                    tmp_result = output == [in_outs[\"outputs\"][index]]\n                    if isinstance(in_outs[\"outputs\"][index], list):\n                        tmp_result = tmp_result or (output == in_outs[\"outputs\"][index])\n                except Exception as e:\n                    if debug:\n                        print(f\"Failed check2 exception = {e}\")\n                    pass\n\n                if tmp_result is True:\n                    results.append(tmp_result)\n                    continue\n\n                # try by converting the output into a split up list too\n                if isinstance(output, list):\n                    output = list(filter(len, output))\n\n                if debug:\n                    nl = \"\\n\"\n                    if not isinstance(inputs, list):\n                        print(\n                            f\"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, \"\n                            f\"inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, \"\n                            f\"{output == [in_outs['outputs'][index]]} {tmp_result=}\"\n                        )\n                    else:\n                        print(\n                            f\"@1 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, \"\n                            f\"{type(inputs)}, {output == [in_outs['outputs'][index]]} {tmp_result=}\"\n                        )\n\n                if debug:\n                    print(f\"{tmp_result=} @a\")\n\n                try:\n                    tmp_result = output == [in_outs[\"outputs\"][index]]\n                    if isinstance(in_outs[\"outputs\"][index], list):\n                        tmp_result = tmp_result or (output == in_outs[\"outputs\"][index])\n                except Exception as e:\n                    if debug:\n                        print(f\"Failed check3 exception = {e}\")\n                    pass\n\n                if debug:\n                    print(f\"{tmp_result=} @b\")\n\n                try:\n                    all_ints = all(\n                        combined_int_check(e1) and combined_int_check(e2)\n                        for e1, e2 in zip(output, in_outs[\"outputs\"][index], strict=True)\n                    )\n                    if not all_ints:\n                        if debug:\n                            print(\n                                [\n                                    combined_int_check(e1) and combined_int_check(e2)\n                                    for e1, e2 in zip(output, in_outs[\"outputs\"][index], strict=True)\n                                ]\n                            )\n                        output_float = [float(e) for e in output]\n                        gt_float = [float(e) for e in in_outs[\"outputs\"][index]]\n                        tmp_result = tmp_result or (\n                            (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float)\n                        )\n                except Exception:\n                    pass\n\n                if debug:\n                    print(f\"{tmp_result=} @c\")\n\n                try:\n                    if isinstance(output[0], list):\n                        all_ints = all(\n                            combined_int_check(e1) and combined_int_check(e2)\n                            for e1, e2 in zip(output[0], in_outs[\"outputs\"][index], strict=True)\n                        )\n                        if not all_ints:\n                            output_float = [float(e) for e in output[0]]\n                            gt_float = [float(e) for e in in_outs[\"outputs\"][index][0]]\n                            tmp_result = tmp_result or (\n                                (len(output_float) == len(gt_float)) and np.allclose(output_float, gt_float)\n                            )\n                except Exception:\n                    pass\n\n                if tmp_result is True:\n                    results.append(tmp_result)\n                    continue\n\n                if debug:\n                    print(f\"{tmp_result=} @d\")\n                # try by converting the stuff into split up list\n                if isinstance(in_outs[\"outputs\"][index], list):\n                    for tmp_index, i in enumerate(in_outs[\"outputs\"][index]):\n                        in_outs[\"outputs\"][index][tmp_index] = set(i.split())\n                else:\n                    in_outs[\"outputs\"][index] = set(in_outs[\"outputs\"][index].split())\n\n                if debug:\n                    print(f\"{tmp_result=} @e\")\n\n                try:\n                    tmp_result = output == in_outs[\"outputs\"][index]\n                except Exception as e:\n                    if debug:\n                        print(f\"Failed check4 exception = {e}\")\n                    continue\n\n                if tmp_result is True:\n                    results.append(tmp_result)\n                    continue\n\n                if debug:\n                    print(f\"{tmp_result=} @f\")\n\n                # try by converting the output into a split up list too\n                if isinstance(output, list):\n                    for tmp_index, i in enumerate(output):\n                        output[tmp_index] = i.split()\n                    output = list(filter(len, output))\n                    for tmp_index, i in enumerate(output):\n                        output[tmp_index] = set(i)\n                else:\n                    output = output.split()\n                    output = list(filter(len, output))\n                    output = set(output)\n\n                if debug:\n                    print(f\"{tmp_result=} @g\")\n\n                if tmp_result is True and debug:\n                    print(\"PASSED\")\n\n                results.append(tmp_result)\n                if tmp_result is not True:\n                    return results, {\n                        \"output\": raw_true_output_copy,\n                        \"expected\": raw_outputs,\n                        \"inputs\": raw_inputs,\n                        # \"error_code\": -2,\n                        \"error_message\": \"Wrong Answer\",\n                    }\n\n                if debug:\n                    nl = \"\\n\"\n                    if not isinstance(inputs, list):\n                        print(\n                            f\"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, \"\n                            f\"inputs = {inputs.replace(nl, ' new-line ')}, {type(inputs)}, \"\n                            f\"{output == [in_outs['outputs'][index]]}\"\n                        )\n                    else:\n                        print(\n                            f\"@2 output = {output}, test outputs = {in_outs['outputs'][index]}, inputs = {inputs}, \"\n                            f\"{type(inputs)}, {output == [in_outs['outputs'][index]]}\"\n                        )\n\n                    print(f\"results = {results}\")\n\n    return results, {}\n\n\ndef custom_compare_(output, ground_truth):\n    if isinstance(output, list):\n        output_1 = \"\\n\".join(output)\n        if stripped_string_compare(output_1, ground_truth):\n            return True\n\n    if isinstance(output, list):\n        output_2 = [o.lstrip().rstrip() for o in output]\n        output_2 = \"\\n\".join(output_2)\n        if stripped_string_compare(output_2, ground_truth):\n            return True\n\n    return False\n\n\ndef stripped_string_compare(s1, s2):\n    s1 = s1.lstrip().rstrip()\n    s2 = s2.lstrip().rstrip()\n    return s1 == s2\n\n\ndef call_method(method, inputs):\n    if isinstance(inputs, list):\n        inputs = \"\\n\".join(inputs)\n\n    inputs_line_iterator = iter(inputs.split(\"\\n\"))\n\n    # sys.setrecursionlimit(10000)\n\n    # @patch('builtins.input', side_effect=inputs.split(\"\\n\"))\n    @patch(\"builtins.open\", mock_open(read_data=inputs))\n    @patch(\"sys.stdin\", StringIO(inputs))\n    @patch(\"sys.stdin.readline\", lambda *args: next(inputs_line_iterator))\n    @patch(\"sys.stdin.readlines\", lambda *args: inputs.split(\"\\n\"))\n    @patch(\"sys.stdin.read\", lambda *args: inputs)\n    # @patch('sys.stdout.write', print)\n    def _inner_call_method(_method):\n        try:\n            return _method()\n        except SystemExit:\n            pass\n        finally:\n            pass\n\n    return _inner_call_method(method)\n\n\ndef reliability_guard(maximum_memory_bytes=None):\n    \"\"\"\n    This disables various destructive functions and prevents the generated code\n    from interfering with the test (e.g. fork bomb, killing other processes,\n    removing filesystem files, etc.)\n    WARNING\n    This function is NOT a security sandbox. Untrusted code, including, model-\n    generated code, should not be blindly executed outside of one. See the\n    Codex paper for more information about OpenAI's code sandbox, and proceed\n    with caution.\n    \"\"\"\n\n    if maximum_memory_bytes is not None:\n        import resource\n\n        resource.setrlimit(resource.RLIMIT_AS, (maximum_memory_bytes, maximum_memory_bytes))\n        resource.setrlimit(resource.RLIMIT_DATA, (maximum_memory_bytes, maximum_memory_bytes))\n        if platform.uname().system != \"Darwin\":\n            resource.setrlimit(resource.RLIMIT_STACK, (maximum_memory_bytes, maximum_memory_bytes))\n\n    faulthandler.disable()\n\n    import builtins\n\n    builtins.exit = None\n    builtins.quit = None\n\n    import os\n\n    os.environ[\"OMP_NUM_THREADS\"] = \"1\"\n\n    os.kill = None\n    os.system = None  # 防止干扰repl评测\n    os.putenv = None\n    os.remove = None\n    os.removedirs = None\n    os.rmdir = None\n    os.fchdir = None\n    os.setuid = None\n    os.fork = None\n    os.forkpty = None\n    os.killpg = None\n    os.rename = None\n    os.renames = None\n    os.truncate = None\n    os.replace = None\n    os.unlink = None\n    os.fchmod = None\n    os.fchown = None\n    os.chmod = None\n    os.chown = None\n    os.chroot = None\n    os.lchflags = None\n    os.lchmod = None\n    os.lchown = None\n    os.getcwd = None\n    os.chdir = None\n\n    import shutil\n\n    shutil.rmtree = None\n    shutil.move = None\n    shutil.chown = None\n\n    import subprocess\n\n    subprocess.Popen = None  # type: ignore\n\n    __builtins__[\"help\"] = None\n\n    import sys\n\n    sys.modules[\"ipdb\"] = None\n    sys.modules[\"joblib\"] = None\n    sys.modules[\"resource\"] = None\n    sys.modules[\"psutil\"] = None\n    sys.modules[\"tkinter\"] = None\n\n    # Disable some built-in functions that can be destructive\n    for mod in [\"subprocess\", \"ctypes\"]:\n        sys.modules[mod] = None\n"
  },
  {
    "path": "verl/utils/reward_score/prime_code/utils.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Borrowed from: https://huggingface.co/spaces/codeparrot/apps_metric/blob/main/utils.py\n\nimport multiprocessing\nimport os\nimport sys\nimport traceback\nfrom typing import Optional\n\nfrom .testing_util import run_test\n\n\ndef _temp_run(sample, generation, debug, result, metadata_list, timeout):\n    with open(os.devnull, \"w\") as devnull:\n        sys.stdout = devnull\n        sys.stderr = devnull\n        try:\n            res, metadata = run_test(in_outs=sample, test=generation, debug=debug, timeout=timeout)\n            result.append(res)\n            metadata_list.append(metadata)\n        except Exception:\n            # print(e) # some tracebacks are extremely long.\n            traceback.print_exc(10)\n            result.append([-1 for i in range(len(sample[\"inputs\"]))])\n            metadata_list.append({})\n\n\ndef check_correctness(in_outs: Optional[dict], generation, timeout=10, debug=True):\n    \"\"\"Check correctness of code generation with a global timeout.\n    The global timeout is to catch some extreme/rare cases not handled by the timeouts\n    inside `run_test`\"\"\"\n\n    manager = multiprocessing.Manager()\n    result = manager.list()\n    metadata_list = manager.list()\n    p = multiprocessing.Process(target=_temp_run, args=(in_outs, generation, debug, result, metadata_list, timeout))\n    p.start()\n    p.join(timeout=timeout + 1)\n    if p.is_alive():\n        p.kill()\n        # p.terminate()\n    if not result:\n        # consider that all tests failed\n        result = [[-1 for i in range(len(in_outs[\"inputs\"]))]]\n        if debug:\n            print(\"global timeout\")\n    return result[0], metadata_list\n"
  },
  {
    "path": "verl/utils/reward_score/prime_math/__init__.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nAnswer checker API that uses sympy to simplify expressions and check for equality.\n\nCall grade_answer(given_answer: str, ground_truth: str).\n\nFROM: https://github.com/openai/prm800k/blob/main/prm800k/grading/grader.py\n\"\"\"\n\nimport contextlib\nimport math\nimport re\n\nimport sympy\nfrom pylatexenc import latex2text\nfrom sympy.parsing import sympy_parser\n\nfrom verl.utils.py_functional import timeout_limit\n\nfrom . import math_normalize\nfrom .grader import math_equal\n\n# import math_normalize\n# from grader import math_equal\n\n# sympy might hang -- we don't care about trying to be lenient in these cases\nBAD_SUBSTRINGS = [\"^{\", \"^(\"]\nBAD_REGEXES = [r\"\\^[0-9]+\\^\", r\"\\^[0-9][0-9]+\"]\nTUPLE_CHARS = \"()[]\"\n\n\ndef _sympy_parse(expr: str):\n    \"\"\"Parses an expression with sympy.\"\"\"\n    py_expr = expr.replace(\"^\", \"**\")\n    return sympy_parser.parse_expr(\n        py_expr,\n        transformations=(sympy_parser.standard_transformations + (sympy_parser.implicit_multiplication_application,)),\n    )\n\n\ndef _parse_latex(expr: str) -> str:\n    \"\"\"Attempts to parse latex to an expression sympy can read.\"\"\"\n    expr = expr.replace(\"\\\\tfrac\", \"\\\\frac\")\n    expr = expr.replace(\"\\\\dfrac\", \"\\\\frac\")\n    expr = expr.replace(\"\\\\frac\", \" \\\\frac\")  # Play nice with mixed numbers.\n    expr = latex2text.LatexNodes2Text().latex_to_text(expr)\n\n    # Replace the specific characters that this parser uses.\n    expr = expr.replace(\"√\", \"sqrt\")\n    expr = expr.replace(\"π\", \"pi\")\n    expr = expr.replace(\"∞\", \"inf\")\n    expr = expr.replace(\"∪\", \"U\")\n    expr = expr.replace(\"·\", \"*\")\n    expr = expr.replace(\"×\", \"*\")\n\n    return expr.strip()\n\n\ndef _is_float(num: str) -> bool:\n    try:\n        float(num)\n        return True\n    except ValueError:\n        return False\n\n\ndef _is_int(x: float) -> bool:\n    try:\n        return abs(x - int(round(x))) <= 1e-7\n    except Exception:\n        return False\n\n\ndef _is_frac(expr: str) -> bool:\n    return bool(re.search(r\"^-?[0-9]+.?/0*[1-9][0-9]*.?$\", expr))\n\n\ndef _str_is_int(x: str) -> bool:\n    try:\n        x = _strip_properly_formatted_commas(x)\n        x = float(x)\n        return abs(x - int(round(x))) <= 1e-7\n    except Exception:\n        return False\n\n\ndef _str_to_int(x: str) -> bool:\n    x = x.replace(\",\", \"\")\n    x = float(x)\n    return int(x)\n\n\ndef _inject_implicit_mixed_number(step: str):\n    \"\"\"\n    Automatically make a mixed number evalable\n    e.g. 7 3/4 => 7+3/4\n    \"\"\"\n    p1 = re.compile(r\"([0-9]) +([0-9])\")\n    step = p1.sub(r\"\\1+\\2\", step)  ## implicit mults\n    return step\n\n\ndef _strip_properly_formatted_commas(expr: str):\n    # We want to be careful because we don't want to strip tuple commas\n    p1 = re.compile(r\"(\\d)(,)(\\d\\d\\d)($|\\D)\")\n    while True:\n        next_expr = p1.sub(r\"\\1\\3\\4\", expr)\n        if next_expr == expr:\n            break\n        expr = next_expr\n    return next_expr\n\n\ndef _normalize(expr: str) -> str:\n    \"\"\"Normalize answer expressions.\"\"\"\n    if expr is None:\n        return None\n\n    # Remove enclosing `\\text{}`.\n    m = re.search(r\"^\\\\text\\{(?P<text>.+?)\\}$\", expr)\n    if m is not None:\n        expr = m.group(\"text\")\n\n    expr = expr.replace(\"\\\\%\", \"%\")\n    expr = expr.replace(\"\\\\$\", \"$\")\n    expr = expr.replace(\"$\", \"\")\n    expr = expr.replace(\"%\", \"\")\n    expr = expr.replace(\" or \", \" , \")\n    expr = expr.replace(\" and \", \" , \")\n\n    expr = expr.replace(\"million\", \"*10^6\")\n    expr = expr.replace(\"billion\", \"*10^9\")\n    expr = expr.replace(\"trillion\", \"*10^12\")\n\n    for unit in [\n        \"degree\",\n        \"cm\",\n        \"centimeter\",\n        \"meter\",\n        \"mile\",\n        \"second\",\n        \"minute\",\n        \"hour\",\n        \"day\",\n        \"week\",\n        \"month\",\n        \"year\",\n        \"foot\",\n        \"feet\",\n        \"inch\",\n        \"yard\",\n        \"liter\",\n    ]:\n        expr = re.sub(f\"{unit}(es)?(s)? *(\\\\^[0-9]+)?\", \"\", expr)\n    expr = re.sub(r\"\\^ *\\\\circ\", \"\", expr)\n\n    if len(expr) > 0 and expr[0] == \"{\" and expr[-1] == \"}\":\n        expr = expr[1:-1]\n\n    expr = re.sub(\",\\\\\\\\! *\", \"\", expr)\n    if _is_float(expr) and _is_int(float(expr)):\n        expr = str(int(round(float(expr))))\n    if \"\\\\\" in expr:\n        with contextlib.suppress(Exception):\n            expr = _parse_latex(expr)\n\n    # edge case with mixed numbers and negative signs\n    expr = re.sub(\"- *\", \"-\", expr)\n\n    expr = _inject_implicit_mixed_number(expr)\n\n    # don't be case sensitive for text answers\n    expr = expr.lower()\n\n    if _str_is_int(expr):\n        expr = str(_str_to_int(expr))\n\n    return expr\n\n\ndef count_unknown_letters_in_expr(expr: str):\n    expr = expr.replace(\"sqrt\", \"\")\n    expr = expr.replace(\"frac\", \"\")\n    letters_in_expr = set([x for x in expr if x.isalpha()])\n    return len(letters_in_expr)\n\n\ndef should_allow_eval(expr: str):\n    # we don't want to try parsing unknown text or functions of more than two variables\n    if count_unknown_letters_in_expr(expr) > 2:\n        return False\n\n    for bad_string in BAD_SUBSTRINGS:\n        if bad_string in expr:\n            return False\n\n    return all(re.search(bad_regex, expr) is None for bad_regex in BAD_REGEXES)\n\n\n@timeout_limit(seconds=10)\ndef are_equal_under_sympy(ground_truth_normalized: str, given_normalized: str):\n    are_equal = False\n    try:\n        expr = f\"({ground_truth_normalized})-({given_normalized})\"\n        if should_allow_eval(expr):\n            sympy_diff = _sympy_parse(expr)\n            simplified = sympy.simplify(sympy_diff)\n            if simplified == 0:\n                are_equal = True\n    except Exception:\n        pass\n    return are_equal\n\n\ndef split_tuple(expr: str):\n    \"\"\"\n    Split the elements in a tuple/interval, while handling well-formatted commas in large numbers\n    \"\"\"\n    expr = _strip_properly_formatted_commas(expr)\n    if len(expr) == 0:\n        return []\n    if (\n        len(expr) > 2\n        and expr[0] in TUPLE_CHARS\n        and expr[-1] in TUPLE_CHARS\n        and all([ch not in expr[1:-1] for ch in TUPLE_CHARS])\n    ):\n        elems = [elem.strip() for elem in expr[1:-1].split(\",\")]\n    else:\n        elems = [expr]\n    return elems\n\n\ndef grade_answer(given_answer: str, ground_truth: str) -> bool:\n    \"\"\"\n    The answer will be considered correct if:\n    (a) it normalizes to the same string as the ground truth answer\n    OR\n    (b) sympy can simplify the difference between the expressions to 0\n    \"\"\"\n    if given_answer is None:\n        return False\n\n    ground_truth_normalized_mathd = math_normalize.normalize_answer(ground_truth)\n    given_answer_normalized_mathd = math_normalize.normalize_answer(given_answer)\n\n    # be at least as lenient as mathd\n    if ground_truth_normalized_mathd == given_answer_normalized_mathd:\n        return True\n\n    ground_truth_normalized = _normalize(ground_truth)\n    given_normalized = _normalize(given_answer)\n\n    if ground_truth_normalized is None:\n        return False\n\n    if ground_truth_normalized == given_normalized:\n        return True\n\n    if len(given_normalized) == 0:\n        return False\n\n    ground_truth_elems = split_tuple(ground_truth_normalized)\n    given_elems = split_tuple(given_normalized)\n\n    if (\n        len(ground_truth_elems) > 1\n        and (ground_truth_normalized[0] != given_normalized[0] or ground_truth_normalized[-1] != given_normalized[-1])\n        or len(ground_truth_elems) != len(given_elems)\n    ):\n        is_correct = False\n    else:\n        for ground_truth_elem, given_elem in zip(ground_truth_elems, given_elems, strict=True):\n            if _is_frac(ground_truth_elem) and _is_frac(given_elem):\n                # if fractions aren't reduced, then shouldn't be marked as correct\n                # so, we don't want to allow sympy.simplify in this case\n                is_correct = ground_truth_elem == given_elem\n            elif _str_is_int(ground_truth_elem) != _str_is_int(given_elem):\n                # if the ground truth answer is an integer, we require the given answer to be a strict match\n                # (no sympy.simplify)\n                is_correct = False\n            else:\n                try:\n                    is_correct = are_equal_under_sympy(ground_truth_elem, given_elem)\n                except Exception as e:\n                    # if there's an error, we'll just say it's not correct\n                    is_correct = False\n                    print(f\"Error: {e} from are_equal_under_sympy, {ground_truth_elem}, {given_elem}\")\n            if not is_correct:\n                break\n\n    return is_correct\n\n\ndef remove_boxed(s):\n    left = \"\\\\boxed{\"\n    try:\n        assert s[: len(left)] == left\n        assert s[-1] == \"}\"\n        return s[len(left) : -1]\n    except Exception:\n        return None\n\n\ndef _last_boxed_only_string(string):\n    idx = string.rfind(\"\\\\boxed\")\n    if idx < 0:\n        idx = string.rfind(\"\\\\fbox\")\n        if idx < 0:\n            return None\n\n    i = idx\n    left_brace_idx = None\n    right_brace_idx = None\n    num_left_braces_open = 0\n    while i < len(string):\n        if string[i] == \"{\":\n            num_left_braces_open += 1\n            if left_brace_idx is None:\n                left_brace_idx = i\n        elif string[i] == \"}\":\n            num_left_braces_open -= 1\n            if num_left_braces_open == 0:\n                right_brace_idx = i\n                break\n\n        i += 1\n\n    if left_brace_idx is None or right_brace_idx is None:\n        return None\n\n    return string[left_brace_idx + 1 : right_brace_idx].strip()\n\n\ndef match_answer(response):\n    is_matched = False\n    for ans_marker in [\"answer:\", \"answer is\", \"answers are\"]:\n        ans_idx = response.lower().rfind(ans_marker)\n        if ans_idx != -1:\n            is_matched = True\n            response = response[ans_idx + len(ans_marker) :].strip()\n            if response.endswith(\"\\n\"):\n                response = response[:-2]\n\n    for ans_marker in [\"is answer\", \"is the answer\", \"are answers\", \"are the answers\"]:\n        ans_idx = response.lower().rfind(ans_marker)\n        if ans_idx != -1:\n            is_matched = True\n            response = response[:ans_idx].strip()\n            if response.endswith(\"\\n\"):\n                response = response[:-2]\n\n    # Find boxed\n    ans_boxed = _last_boxed_only_string(response)\n    if ans_boxed:\n        is_matched = True\n        response = ans_boxed\n\n    if \". \" in response:\n        dot_idx = response.lower().rfind(\". \")\n        if dot_idx != -1:\n            response = response[:dot_idx].strip()\n\n    for ans_marker in [\"be \", \"is \", \"are \", \"=\", \": \", \"get \", \"be\\n\", \"is\\n\", \"are\\n\", \":\\n\", \"get\\n\"]:\n        ans_idx = response.lower().rfind(ans_marker)\n        if ans_idx != -1:\n            is_matched = True\n            response = response[ans_idx + len(ans_marker) :].strip()\n            if response.endswith(\"\\n\"):\n                response = response[:-2]\n\n    is_matched = is_matched if any([c.isdigit() for c in response]) else False  # answer must have a digit\n    # Grade\n    return is_matched, response\n\n\ndef compute_score(model_output: str, ground_truth: str) -> bool:\n    model_output = str(model_output)\n    ground_truth = str(ground_truth)\n\n    is_matched, extracted_model_output = match_answer(model_output)\n    format_correctness = \"Step 2:\" in model_output and \"\\\\box\" in model_output\n\n    # grade simple algebra questions. if succeeded, return; otherwise, proceed to more complex grading\n    if grade_answer(extracted_model_output, ground_truth):\n        return True, True, extracted_model_output\n\n    try:\n        if \"\\\\pi\" in extracted_model_output or \"\\\\pi\" in ground_truth:\n            equivs = []\n            for pi in [math.pi, 3.14]:\n                equivs.append(math_equal(extracted_model_output, ground_truth, timeout=True, pi=pi))\n            is_correct = any(equivs)\n        else:\n            is_correct = math_equal(extracted_model_output, ground_truth, timeout=True)\n    except Exception:\n        is_correct = False\n\n    return is_correct, format_correctness, extracted_model_output\n"
  },
  {
    "path": "verl/utils/reward_score/prime_math/grader.py",
    "content": "# Copyright (c) 2024, NVIDIA CORPORATION.  All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Copyright (c) Microsoft Corporation.\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE\n\n# Copyright (c) 2023 OpenAI\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\n# Copyright (c) 2021 Dan Hendrycks\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\n# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThis logic is largely copied from the Hendrycks' MATH release (math_equivalence), and borrowed from:\n- https://github.com/microsoft/ToRA/blob/main/src/eval/grader.py\n- https://github.com/microsoft/ProphetNet/tree/master/CRITIC\n- https://github.com/openai/prm800k\n\"\"\"\n\nimport contextlib\nimport math\nimport re\nfrom math import isclose\n\n# sympy related\nfrom sympy import N, simplify\nfrom sympy.parsing.latex import parse_latex\nfrom sympy.parsing.sympy_parser import parse_expr\n\n# verl related\nfrom verl.utils.py_functional import timeout_limit\n\n\ndef is_digit(s):\n    try:\n        if \"{,}\" in str(s):\n            num = float(str(s).replace(\"{,}\", \"\"))\n            return True, num\n\n        num = float(str(s).replace(\",\", \"\"))\n        return True, num\n    except ValueError:\n        return False, None\n\n\ndef normalize(answer, pi) -> str:\n    # checking if answer is $<number> and removing $ in that case to compare\n    if isinstance(answer, str) and bool(re.match(r\"\\$\\d+(\\.\\d+)?\", answer)):\n        return answer[1:]\n\n    # checking if answer is <number>% or <number>\\\\% and removing %\n    if isinstance(answer, str) and (\n        bool(re.match(r\"^\\d+(\\.\\d+)?%$\", answer)) or bool(re.match(r\"^\\d+(\\.\\d+)?\\\\%$\", answer))\n    ):\n        return answer.replace(\"\\\\%\", \"\").replace(\"%\", \"\")\n\n    # handle base\n    answer = handle_base(answer)\n\n    # handle pi\n    answer = handle_pi(answer, pi)\n\n    return answer\n\n\ndef handle_base(x) -> str:\n    if isinstance(x, str) and \"_\" in x:\n        # Due to base\n        x = x.split(\"_\")[0]\n        x = float(x)\n        return int(x)\n    return x\n\n\ndef handle_pi(string, pi):\n    if isinstance(string, str) and \"\\\\pi\" in string:\n        # Find the first occurrence of \"\\pi\"\n        idx = string.find(\"\\\\pi\")\n\n        # Iterate over the string and find all occurrences of \"\\pi\" with a valid previous character\n        while idx != -1:\n            if idx > 0 and string[idx - 1].isdigit():\n                # Replace \"\\pi\" with \"*math.pi\" if the previous character is a digit\n                string = string[:idx] + f\"*{pi}\" + string[idx + 3 :]\n            else:\n                # Replace \"\\pi\" with \"1*math.pi\" if the previous character is not a digit\n                string = string[:idx] + f\"1*{pi}\" + string[idx + 3 :]\n\n            # Find the next occurrence of \"\\pi\"\n            idx = string.find(\"\\\\pi\", idx + 1)\n\n        # Evaluate the expression using eval() function\n        with contextlib.suppress(Exception):\n            string = eval(string)\n\n    return string\n\n\ndef math_equal(\n    prediction: bool | float | str,\n    reference: float | str,\n    include_percentage: bool = True,\n    tolerance: float = 1e-4,\n    timeout: float = 10.0,\n    pi: float = math.pi,\n) -> bool:\n    \"\"\"\n    Exact match of math if and only if:\n    1. numerical equal: both can convert to float and are equal\n    2. symbolic equal: both can convert to sympy expression and are equal\n    \"\"\"\n\n    prediction = normalize(prediction, pi)\n    reference = normalize(reference, pi)\n\n    if isinstance(prediction, str) and len(prediction) > 1000:  # handling weird corner-cases\n        prediction = prediction[:1000]\n\n    # 0. string comparison\n    if isinstance(prediction, str) and isinstance(reference, str):\n        if prediction.strip().lower() == reference.strip().lower():\n            return True\n        if prediction.replace(\" \", \"\") == reference.replace(\" \", \"\"):\n            return True\n\n    try:  # 1. numerical equal\n        if is_digit(prediction)[0] and is_digit(reference)[0]:\n            prediction = is_digit(prediction)[1]\n            reference = is_digit(reference)[1]\n            # number questions\n            gt_result = [reference / 100, reference, reference * 100] if include_percentage else [reference]\n            for item in gt_result:\n                try:\n                    if isclose(item, prediction, rel_tol=tolerance):\n                        return True\n                except Exception:\n                    continue\n            return False\n    except Exception:\n        pass\n\n    if not prediction and prediction not in [0, False]:\n        return False\n\n    # 2. symbolic equal\n    reference = str(reference).strip()\n    prediction = str(prediction).strip()\n\n    ## deal with [], (), {}\n    prediction = format_intervals(prediction)\n\n    pred_str, ref_str = prediction, reference\n    if (prediction.startswith(\"[\") and prediction.endswith(\"]\") and not reference.startswith(\"(\")) or (\n        prediction.startswith(\"(\") and prediction.endswith(\")\") and not reference.startswith(\"[\")\n    ):\n        pred_str = pred_str.strip(\"[]()\")\n        ref_str = ref_str.strip(\"[]()\")\n    for s in [\"{\", \"}\", \"(\", \")\"]:\n        ref_str = ref_str.replace(s, \"\")\n        pred_str = pred_str.replace(s, \"\")\n    if pred_str == ref_str:\n        return True\n\n    ## [a, b] vs. [c, d], return a==c and b==d\n    if (\n        prediction\n        and reference\n        and prediction[0] in \"([\"\n        and prediction[-1] in \")]\"\n        and prediction[0] == reference[0]\n        and prediction[-1] == reference[-1]\n    ):\n        pred_parts = prediction[1:-1].split(\",\")\n        ref_parts = reference[1:-1].split(\",\")\n        if len(pred_parts) == len(ref_parts) and all(\n            [\n                math_equal(pred_pt, ref_pt, include_percentage, tolerance)\n                for pred_pt, ref_pt in zip(pred_parts, ref_parts, strict=True)\n            ]\n        ):\n            return True\n\n    if \",\" in prediction and \",\" in reference:\n        pred_parts = [item.strip() for item in prediction.split(\",\")]\n        ref_parts = [item.strip() for item in reference.split(\",\")]\n\n        if len(pred_parts) == len(ref_parts):\n            return bool(\n                all(\n                    [\n                        math_equal(pred_parts[i], ref_parts[i], include_percentage, tolerance)\n                        for i in range(len(pred_parts))\n                    ]\n                )\n            )\n\n    # if we have point == tuple of values\n    if prediction.startswith(\"Point\") and reference[0] == \"(\" and reference[-1] == \")\":\n        pred_parts = prediction[prediction.find(\"(\") + 1 : -1].split(\",\")\n        ref_parts = reference[1:-1].split(\",\")\n        if len(pred_parts) == len(ref_parts) and all(\n            [\n                math_equal(pred_pt, ref_pt, include_percentage, tolerance)\n                for pred_pt, ref_pt in zip(pred_parts, ref_parts, strict=False)\n            ]\n        ):\n            return True\n\n    # if reference is a matrix\n    if r\"\\begin{pmatrix}\" in reference and prediction.startswith(\"Matrix\"):\n        try:\n            pred_matrix = parse_expr(prediction)\n            ref_matrix_items = reference.split()[1:-1:2]\n            if len(pred_matrix) == len(ref_matrix_items) and all(\n                [\n                    math_equal(pred, ref, include_percentage, tolerance)\n                    for ref, pred in zip(ref_matrix_items, pred_matrix, strict=False)\n                ]\n            ):\n                return True\n        except Exception:\n            pass\n    elif r\"\\begin{pmatrix}\" in reference and prediction.startswith(\"[\") and prediction.endswith(\"]\"):\n        if isinstance(eval(prediction), list):\n            try:\n                pred_matrix = eval(prediction)\n                # ref_matrix_items = reference.split()[1:-1:2]\n                ref_matrix_items = (\n                    reference.removeprefix(r\"\\\\begin{pmatrix}\")\n                    .removeprefix(r\"\\begin{pmatrix}\")\n                    .removesuffix(r\"\\\\end{pmatrix}\")\n                    .removesuffix(r\"\\end{pmatrix}\")\n                )\n                ref_matrix_items = ref_matrix_items.split(\"\\\\\")\n                ref_matrix_items = [row.split(\"&\") if \"&\" in row else row for row in ref_matrix_items]\n                if len(pred_matrix) == len(ref_matrix_items) and all(\n                    [\n                        math_equal(pred, ref, include_percentage, tolerance)\n                        for ref, pred in zip(ref_matrix_items, pred_matrix, strict=False)\n                    ]\n                ):\n                    return True\n            except Exception:\n                pass\n\n    return symbolic_equal(prediction, reference, tolerance, timeout)\n\n\ndef symbolic_equal(a, b, tolerance, timeout=10.0):\n    def _parse(s):\n        for f in [parse_expr, parse_latex]:\n            try:\n                with timeout_limit(seconds=timeout):\n                    return f(s)\n            except TimeoutError:\n                print(f\"Parsing timed out for {s}\")\n                continue\n            except Exception:\n                continue\n        return s\n\n    a = _parse(a)\n    b = _parse(b)\n\n    try:\n        with timeout_limit(seconds=timeout):\n            if simplify(a - b) == 0:\n                return True\n    except TimeoutError:\n        print(f\"Simplification timed out for {a} - {b}\")\n        pass\n    except Exception:\n        pass\n\n    try:\n        with timeout_limit(seconds=timeout):\n            if isclose(N(a), N(b), rel_tol=tolerance):\n                return True\n    except TimeoutError:\n        print(f\"Numerical evaluation timed out for {a}, {b}\")\n        pass\n    except Exception:\n        pass\n    return False\n\n\ndef format_intervals(prediction):\n    patterns = {\n        \"Interval(\": r\"^Interval\\((.*)\\)$\",\n        \"Interval.Ropen(\": r\"^Interval\\.Ropen\\((.*)\\)$\",\n        \"Interval.Lopen(\": r\"^Interval\\.Lopen\\((.*)\\)$\",\n        \"Interval.open(\": r\"^Interval\\.open\\((.*)\\)$\",\n    }\n\n    for key, pattern in patterns.items():\n        match = re.match(pattern, prediction)\n        if match:\n            inner_content = match.group(1)\n\n            if key == \"Interval(\":  # Intarval(a, b) == [a, b]\n                return f\"[{inner_content}]\"\n            elif key == \"Interval.Ropen(\":  # Intarval.Ropen(a, b) == [a, b)\n                return f\"[{inner_content})\"\n            elif key == \"Interval.Lopen(\":  # Intarval.Lopen(a, b) == (a, b]\n                return f\"({inner_content}]\"\n            elif key == \"Interval.open(\":  # Intarval.open(a, b) == (a, b)\n                return f\"({inner_content})\"\n\n    return prediction\n"
  },
  {
    "path": "verl/utils/reward_score/prime_math/math_normalize.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# Copyright (c) 2021 Dan Hendrycks\n#\n# Permission is hereby granted, free of charge, to any person obtaining a copy\n# of this software and associated documentation files (the \"Software\"), to deal\n# in the Software without restriction, including without limitation the rights\n# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell\n# copies of the Software, and to permit persons to whom the Software is\n# furnished to do so, subject to the following conditions:\n#\n# The above copyright notice and this permission notice shall be included in all\n# copies or substantial portions of the Software.\n#\n# THE SOFTWARE IS PROVIDED \"AS IS\", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR\n# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,\n# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE\n# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER\n# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,\n# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE\n# SOFTWARE.\n\"\"\"\nThis logic is largely copied from the Hendrycks' MATH release (math_equivalence).\n\nFrom: https://github.com/openai/prm800k/blob/main/prm800k/grading/math_normalize.py\n\"\"\"\n\nimport re\nfrom typing import Optional\n\n\ndef normalize_answer(answer: Optional[str]) -> Optional[str]:\n    if answer is None:\n        return None\n    answer = answer.strip()\n    try:\n        # Remove enclosing `\\text{}`.\n        m = re.search(r\"^\\\\text\\{(?P<text>.+?)\\}$\", answer)\n        if m is not None:\n            answer = m.group(\"text\").strip()\n        return _strip_string(answer)\n    except Exception:\n        return answer\n\n\ndef _fix_fracs(string):\n    substrs = string.split(\"\\\\frac\")\n    new_str = substrs[0]\n    if len(substrs) > 1:\n        substrs = substrs[1:]\n        for substr in substrs:\n            new_str += \"\\\\frac\"\n            if substr[0] == \"{\":\n                new_str += substr\n            else:\n                try:\n                    assert len(substr) >= 2\n                except Exception:\n                    return string\n                a = substr[0]\n                b = substr[1]\n                if b != \"{\":\n                    if len(substr) > 2:\n                        post_substr = substr[2:]\n                        new_str += \"{\" + a + \"}{\" + b + \"}\" + post_substr\n                    else:\n                        new_str += \"{\" + a + \"}{\" + b + \"}\"\n                else:\n                    if len(substr) > 2:\n                        post_substr = substr[2:]\n                        new_str += \"{\" + a + \"}\" + b + post_substr\n                    else:\n                        new_str += \"{\" + a + \"}\" + b\n    string = new_str\n    return string\n\n\ndef _fix_a_slash_b(string):\n    if len(string.split(\"/\")) != 2:\n        return string\n    a = string.split(\"/\")[0]\n    b = string.split(\"/\")[1]\n    try:\n        a = int(a)\n        b = int(b)\n        assert string == \"{}/{}\".format(a, b)\n        new_string = \"\\\\frac{\" + str(a) + \"}{\" + str(b) + \"}\"\n        return new_string\n    except Exception:\n        return string\n\n\ndef _remove_right_units(string):\n    # \"\\\\text{ \" only ever occurs (at least in the val set) when describing units\n    if \"\\\\text{ \" in string:\n        splits = string.split(\"\\\\text{ \")\n        assert len(splits) == 2\n        return splits[0]\n    else:\n        return string\n\n\ndef _fix_sqrt(string):\n    if \"\\\\sqrt\" not in string:\n        return string\n    splits = string.split(\"\\\\sqrt\")\n    new_string = splits[0]\n    for split in splits[1:]:\n        if split[0] != \"{\":\n            a = split[0]\n            new_substr = \"\\\\sqrt{\" + a + \"}\" + split[1:]\n        else:\n            new_substr = \"\\\\sqrt\" + split\n        new_string += new_substr\n    return new_string\n\n\ndef _strip_string(string):\n    # linebreaks\n    string = string.replace(\"\\n\", \"\")\n\n    # remove inverse spaces\n    string = string.replace(\"\\\\!\", \"\")\n\n    # replace \\\\ with \\\n    string = string.replace(\"\\\\\\\\\", \"\\\\\")\n\n    # replace tfrac and dfrac with frac\n    string = string.replace(\"tfrac\", \"frac\")\n    string = string.replace(\"dfrac\", \"frac\")\n\n    # remove \\left and \\right\n    string = string.replace(\"\\\\left\", \"\")\n    string = string.replace(\"\\\\right\", \"\")\n\n    # Remove circ (degrees)\n    string = string.replace(\"^{\\\\circ}\", \"\")\n    string = string.replace(\"^\\\\circ\", \"\")\n\n    # remove dollar signs\n    string = string.replace(\"\\\\$\", \"\")\n\n    # remove units (on the right)\n    string = _remove_right_units(string)\n\n    # remove percentage\n    string = string.replace(\"\\\\\\\\%\", \"\")\n    string = string.replace(\"\\\\%\", \"\")\n\n    # \" 0.\" equivalent to \" .\" and \"{0.\" equivalent to \"{.\" Alternatively, add \"0\" if \".\" is the start of the string\n    string = string.replace(\" .\", \" 0.\")\n    string = string.replace(\"{.\", \"{0.\")\n    # if empty, return empty string\n    if len(string) == 0:\n        return string\n    if string[0] == \".\":\n        string = \"0\" + string\n\n    # to consider: get rid of e.g. \"k = \" or \"q = \" at beginning\n    if len(string.split(\"=\")) == 2 and len(string.split(\"=\")[0]) <= 2:\n        string = string.split(\"=\")[1]\n\n    # fix sqrt3 --> sqrt{3}\n    string = _fix_sqrt(string)\n\n    # remove spaces\n    string = string.replace(\" \", \"\")\n\n    # \\frac1b or \\frac12 --> \\frac{1}{b} and \\frac{1}{2}, etc. Even works with \\frac1{72} (but not \\frac{72}1).\n    # Also does a/b --> \\\\frac{a}{b}\n    string = _fix_fracs(string)\n\n    # manually change 0.5 --> \\frac{1}{2}\n    if string == \"0.5\":\n        string = \"\\\\frac{1}{2}\"\n\n    # NOTE: X/Y changed to \\frac{X}{Y} in dataset, but in simple cases fix in case the model output is X/Y\n    string = _fix_a_slash_b(string)\n\n    return string\n"
  },
  {
    "path": "verl/utils/reward_score/rlla.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport json\nimport random\nimport re\nfrom collections import Counter\n\n\ndef match_score(list1, list2):\n    \"\"\"Compute a similarity score considering element frequency, ignoring order.\n\n    Reference: Liu S Y, Dong X, Lu X, et al. \"Gdpo: Group reward-decoupled normalization policy\n    optimization for multi-reward rl optimization.\"\n    arXiv preprint arXiv:2601.05242, 2026.\n    \"\"\"\n    if list1 == list2:\n        return 1.0\n\n    if not list1 or not list2:\n        return 0.0\n\n    count1 = Counter(list1)  # Frequency count for list1\n    count2 = Counter(list2)  # Frequency count for list2\n\n    intersection = sum(min(count1[k], count2[k]) for k in count1.keys() & count2.keys())\n    max_possible = len(list1) + len(list2) - intersection\n\n    return intersection / max_possible if max_possible > 0 else 0.0\n\n\n# custoimzed reward functions: format\ndef customize_format_reward_func(\n    completions, answer, step, max_possible_reward, min_possible_reward, do_print, **kwargs\n):\n    rewards = []\n    responses = [completion[0][\"content\"] for completion in completions]\n\n    if do_print:\n        print(\"\\n======= Answer ======= \")\n        print(answer[0])\n        print(\"\\n======= Responses ======= \")\n        for idx, response in enumerate(responses):\n            print(f\"*** Response {idx + 1}***\\n{response}\")\n\n    for response, ans in zip(responses, answer, strict=False):\n        reward = min_possible_reward\n        if \"<response>\" in ans and \"<tool_call>\" not in ans:\n            pattern = r\"^<think>.*?</think>\\n<response>.*?</response>$\"\n            if (\n                re.search(pattern, response, re.DOTALL)\n                and response.count(\"<response>\") == 1\n                and response.count(\"</response>\") == 1\n            ):\n                reward = max_possible_reward\n        elif \"<response>\" not in ans and \"<tool_call>\" in ans:\n            pattern = r\"^<think>.*?</think>\\n<tool_call>\\n.*?\\n</tool_call>$\"\n            if (\n                re.search(pattern, response, re.DOTALL)\n                and response.count(\"<tool_call>\") == 1\n                and response.count(\"</tool_call>\") == 1\n            ):\n                reward = max_possible_reward\n        elif \"<response>\" in ans and \"<tool_call>\" in ans:\n            pattern = r\"^<think>.*?</think>\\n<tool_call>\\n.*?\\n</tool_call>\\n<response>.*?</response>$\"\n            if (\n                re.search(pattern, response, re.DOTALL)\n                and response.count(\"<tool_call>\") == 1\n                and response.count(\"</tool_call>\") == 1\n                and response.count(\"<response>\") == 1\n                and response.count(\"</response>\") == 1\n            ):\n                reward = max_possible_reward\n        else:\n            pattern = r\"^<think>.*?</think>$\"\n            if re.search(pattern, response, re.DOTALL):\n                reward = max_possible_reward\n\n        rewards.append(reward)\n\n    if do_print:\n        print(\"\\n======= Reward for <format> =======\")\n        print(\"Reward function for <format> is called ...\")\n        print(rewards)\n\n    return rewards\n\n\ndef compute_tool_call_reward(gt_tools, pd_tools, max_possible_reward, min_possible_reward, do_print):\n    if gt_tools == pd_tools:\n        if do_print:\n            print(\"Max possible score:\", \"Exact Match!\")\n            print(\"Score:\", max_possible_reward)\n        return max_possible_reward\n\n    gt_names = [tool[\"name\"] for tool in gt_tools]\n    pd_names = [tool[\"name\"] for tool in pd_tools]\n    score = match_score(list(gt_names), list(pd_names))\n\n    local_max_possible = 1.0\n    used_pd_indices = set()  # Keep track of matched pd_tools\n\n    for gt_tool in gt_tools:\n        gt_name = gt_tool[\"name\"]\n        gt_params = gt_tool[\"parameters\"]\n\n        local_max_possible += 1.0 + len(gt_params)\n\n        best_match = None\n        best_match_score = 0.0\n        best_match_index = -1\n\n        # Find the best matching unused pd_tool\n        for i, pd_tool in enumerate(pd_tools):\n            if i in used_pd_indices or pd_tool[\"name\"] != gt_name:\n                continue\n\n            pd_params = pd_tool[\"parameters\"]\n            param_score = match_score(list(gt_params.keys()), list(pd_params.keys()))\n\n            # Calculate correctness score for parameter values\n            correctness_score = sum(1.0 for k, v in gt_params.items() if k in pd_params and pd_params[k] == v)\n\n            total_score = param_score + correctness_score\n\n            if total_score > best_match_score:\n                best_match_score = total_score\n                best_match = pd_tool\n                best_match_index = i\n\n        if best_match:\n            used_pd_indices.add(best_match_index)\n            score += best_match_score\n\n    if do_print:\n        print()\n        print(\"Max possible score:\", local_max_possible)\n        print(\"Score:\", score)\n\n    return (max_possible_reward - min_possible_reward) * score / local_max_possible + min_possible_reward\n\n\n# custoimzed reward functions: tool call correctness\ndef customize_correctness_reward_tool(\n    completions, answer, step, max_possible_reward, min_possible_reward, do_print, **kwargs\n):\n    responses = [completion[0][\"content\"] for completion in completions]\n    rewards = []\n\n    for response, ans in zip(responses, answer, strict=False):\n        reward = 0.0\n\n        if \"<tool_call>\" not in ans:\n            # if \"<tool_call>\" not in response and \"</tool_call>\" not in response:\n            #     reward = max_possible_reward\n            # else:\n            #     reward = min_possible_reward\n            rewards.append(reward)\n            continue\n\n        gt_tool_call = ans.split(\"<tool_call>\")[1].split(\"</tool_call>\")[0].strip()\n        gt_tools = gt_tool_call.split(\"\\n\")\n        gt_tools = [json.loads(tool) for tool in gt_tools]  # each diction contains \"name\" and \"parameter\"\n\n        try:\n            # Change here as a constrint in training: if the format is not correct,\n            # directly give the lowest possible score\n            assert \"<tool_call>\" in response\n            assert \"</tool_call>\" in response\n            pd_tools = response.split(\"<tool_call>\")[1].split(\"</tool_call>\")[0].strip().split(\"\\n\")\n            pd_tools = [json.loads(tool) for tool in pd_tools]\n            reward = compute_tool_call_reward(\n                gt_tools, pd_tools, max_possible_reward, min_possible_reward, do_print\n            )  # top reward is 2\n        except Exception:\n            reward = min_possible_reward\n\n        rewards.append(reward)\n\n    if do_print:\n        print(\"\\n======= Reward for <tool call> =======\")\n        print(\"Reward function for <tool call> correctness is called ...\")\n        print(rewards)\n    return rewards\n\n\ndef compute_score(data_source, solution_str, ground_truth, extra_info, step=0):\n    \"\"\"The scoring function for GSM8k.\n\n    Reference: Trung, Luong, et al. \"Reft: Reasoning with reinforced fine-tuning.\"\n    Proceedings of the 62nd Annual Meeting of the Association for\n    Computational Linguistics (Volume 1: Long Papers). 2024.\n\n    Args:\n        solution_str: the solution text\n        ground_truth: the ground truth\n        method: the method to extract the solution, choices are 'strict' and 'flexible'\n        format_score: the score for the format\n        score: the score for the correct answer\n    \"\"\"\n    exp_name = extra_info.get(\"experiment_name\", \"\")\n    if \"llama\" in exp_name:\n        predict_str = (\n            solution_str.split(\"<|start_header_id|>assistant<|end_header_id|>\")[-1].split(\"<|eot_id|>\")[0].strip()\n        )\n    elif \"qwen\" in exp_name:\n        predict_str = solution_str.split(\"<|im_start|>assistant\")[-1].split(\"<|im_end|>\")[0].strip()\n    else:\n        raise NotImplementedError(f\"Unknown model name: {exp_name}\")\n\n    tool_max_possible = 3.0\n    tool_min_possible = -3.0\n\n    format_max_possible = 1.0\n    format_min_possible = 0.0\n\n    completions = [[{\"role\": \"assistant\", \"content\": predict_str}]]\n    answer = [ground_truth]\n\n    do_print = random.randint(1, 64) == 1\n\n    fomrat_score = customize_format_reward_func(\n        completions, answer, step, format_max_possible, format_min_possible, do_print\n    )[0]\n    correctness_score = customize_correctness_reward_tool(\n        completions, answer, step, tool_max_possible, tool_min_possible, do_print\n    )[0]\n\n    score = fomrat_score + correctness_score\n\n    result = {\n        \"score\": score,\n        \"format_reward\": fomrat_score,\n        \"accuracy_reward\": correctness_score,\n    }\n\n    return result\n"
  },
  {
    "path": "verl/utils/reward_score/sandbox_fusion/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport json\nimport logging\nimport traceback\n\nfrom .utils import check_correctness\n\n\"\"\"\nVerify code correctness using the Sandbox Fusion (https://github.com/bytedance/SandboxFusion).\nYou can either deploy the sandbox_fusion service yourself or use the\nFaaS service provided by public cloud, eg: volcengine.com.\n\"\"\"\nlogger = logging.getLogger(__name__)\n\n\ndef compute_score(\n    sandbox_fusion_url, concurrent_semaphore, memory_limit_mb, completion, test_cases, continuous=False, timeout=10\n):\n    \"\"\"\n    Computes the code score using the remote sandbox API.\n\n    Args:\n        sandbox_fusion_url: The URL of the sandbox_fusion service, eg: \"https://<your service endpoint>/run_code\"\n\n        completion: The completion string containing the code.\n        test_cases: JSON string or dictionary containing \"inputs\" and \"outputs\".\n        continuous: Whether to compute a continuous score (based on the first N test cases).\n        timeout: Timeout for each test case.\n\n    Returns:\n        A tuple (score, metadata_list).\n        score: Float score (0.0 to 1.0).\n        metadata_list: List containing execution metadata for each test case.\n    \"\"\"\n    solution = completion\n    if \"```python\" in completion:\n        solution = completion.split(\"```python\")[-1].split(\"```\")[0]\n    elif \"```\" in completion:\n        # Handle cases like ```\\ncode\\n```\n        parts = completion.split(\"```\")\n        if len(parts) >= 2:\n            solution = parts[1]\n            # Remove potential language specifier like 'python\\n'\n            if \"\\n\" in solution:\n                first_line, rest = solution.split(\"\\n\", 1)\n                if first_line.strip().isalpha():  # Simple check for language name\n                    solution = rest\n    else:\n        return 0.0, [{\"error\": \"Invalid completion (missing code block)\"}]\n\n    try:\n        if not isinstance(test_cases, dict):\n            try:\n                test_cases = json.loads(test_cases)\n            except json.JSONDecodeError as e:\n                logger.error(f\"Failed to parse test_cases JSON: {e}\")\n                return 0.0, [{\"error\": \"Invalid test_cases JSON format\"}]\n\n        if test_cases is not None and \"assert_case\" in test_cases and isinstance(test_cases.get(\"assert_case\"), list):\n            assert_cases = test_cases.get(\"assert_case\")\n            test_cases.setdefault(\"inputs\", [\"\" for _ in assert_cases])\n            test_cases.setdefault(\"outputs\", [None for _ in assert_cases])\n        elif not test_cases or \"inputs\" not in test_cases or \"outputs\" not in test_cases:\n            logger.error(\"Invalid test_cases structure.\")\n            return 0.0, [{\"error\": \"Invalid test_cases structure (missing inputs/outputs)\"}]\n\n        # Check all test cases\n        # Note: The return value of check_correctness might need adaptation here\n        # Assume check_correctness returns (results_list, metadata_list)\n        # results_list contains True, False, or error codes (-1, -2, -3, etc.)\n        res_list, metadata_list = check_correctness(\n            sandbox_fusion_url=sandbox_fusion_url,\n            in_outs=test_cases,\n            generation=solution,\n            timeout=timeout,\n            concurrent_semaphore=concurrent_semaphore,\n            memory_limit_mb=memory_limit_mb,\n        )\n\n        # Calculate score\n        if not res_list:  # If there are no results (e.g., invalid input)\n            return 0.0, metadata_list\n\n        if continuous:\n            # Calculate pass rate for the first N (e.g., 10) test cases\n            num_to_consider = min(len(res_list), 10)\n            if num_to_consider == 0:\n                score = 0.0\n            else:\n                passed_count = sum(1 for r in res_list[:num_to_consider] if r is True)\n                score = passed_count / num_to_consider\n            # Return all metadata, even if score is based on the first N\n            final_metadata = metadata_list\n        else:\n            # Calculate pass rate for all test cases\n            passed_count = sum(1 for r in res_list if r is True)\n            total_cases = len(res_list)\n            score = passed_count / total_cases if total_cases > 0 else 0.0\n            final_metadata = metadata_list\n\n    except Exception as e:\n        logger.error(f\"Error during compute_score: {e}\")\n        traceback.print_exc()\n        score = 0.0\n        # Try to return partial metadata if available, otherwise return error info\n        final_metadata = metadata_list if \"metadata_list\" in locals() else [{\"error\": f\"Unhandled exception: {e}\"}]\n\n    # Ensure float and list are returned\n    return float(score), final_metadata if isinstance(final_metadata, list) else [final_metadata]\n"
  },
  {
    "path": "verl/utils/reward_score/sandbox_fusion/utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport concurrent.futures  # <-- Import concurrent.futures\nimport json\nimport logging\nimport os\nimport threading\nimport time\nimport traceback\nimport uuid\nfrom typing import Any, Optional\n\nimport requests\n\nDEFAULT_TIMEOUT = 10  # Default compile and run timeout\nMAX_RETRIES = 3\nINITIAL_RETRY_DELAY = 1\nAPI_TIMEOUT = 10\n\nlogger = logging.getLogger(__name__)\n\n# Define supported languages list (optional, for documentation or validation)\nSUPPORTED_LANGUAGES = [\n    \"python\",\n    \"cpp\",\n    \"nodejs\",\n    \"go\",\n    \"go_test\",\n    \"java\",\n    \"php\",\n    \"csharp\",\n    \"bash\",\n    \"typescript\",\n    \"sql\",\n    \"rust\",\n    \"cuda\",\n    \"lua\",\n    \"R\",\n    \"perl\",\n    \"D_ut\",\n    \"ruby\",\n    \"scala\",\n    \"julia\",\n    \"pytest\",\n    \"junit\",\n    \"kotlin_script\",\n    \"jest\",\n    \"verilog\",\n    \"python_gpu\",\n    \"lean\",\n    \"swift\",\n    \"racket\",\n]\n\n\ndef call_sandbox_api(\n    sandbox_fusion_url: str,\n    code: str,\n    stdin: Optional[str],\n    compile_timeout: int,\n    run_timeout: int,\n    memory_limit_mb: int,\n    language: str = \"python\",\n) -> tuple[Optional[dict[str, Any]], Optional[str]]:  # <-- Remove request_id parameter\n    \"\"\"\n    Calls the remote sandbox API to execute code with retry logic for Gateway Timeout,\n    using increasing delay between retries. Logs internal calls with a unique ID.\n\n    Args:\n        sandbox_fusion_url: The URL of the sandbox fusion API.\n        code: The code string to execute.\n        stdin: The standard input string.\n        compile_timeout: Compile timeout in seconds.\n        run_timeout: Run timeout in seconds.\n        language: The programming language of the code (e.g., \"python\", \"cpp\", \"java\"). Defaults to \"python\".\n\n    Returns:\n        A tuple (response_json, error_message).\n        If successful, response_json is the API's returned JSON object, error_message is None.\n        If failed after retries, response_json is None, error_message contains the error information.\n    \"\"\"\n    request_id = str(uuid.uuid4())  # <-- Generate request_id internally\n    log_prefix = f\"[Request ID: {request_id}] \"  # <-- Create log prefix\n\n    if language not in SUPPORTED_LANGUAGES:\n        error_msg = f\"{log_prefix}Unsupported language: {language}\"\n        logger.error(error_msg)\n        return None, error_msg\n\n    payload = json.dumps(\n        {\n            \"compile_timeout\": compile_timeout,\n            \"run_timeout\": run_timeout,\n            \"code\": code,\n            \"stdin\": stdin,\n            \"memory_limit_MB\": memory_limit_mb,\n            \"language\": language,  # Use the passed language parameter\n            \"files\": {},\n            \"fetch_files\": [],\n        }\n    )\n    headers = {\"Content-Type\": \"application/json\", \"Accept\": \"application/json\"}\n    # Calculate a reasonable request timeout based on compile/run timeouts plus a buffer\n    request_timeout = compile_timeout + run_timeout + API_TIMEOUT\n\n    last_error = None  # Store the last error encountered\n\n    for attempt in range(MAX_RETRIES):\n        try:\n            logger.info(\n                f\"{log_prefix}Attempt {attempt + 1}/{MAX_RETRIES}: Calling sandbox API at {sandbox_fusion_url}\"\n            )  # <-- Use internal log_prefix\n            response = requests.post(\n                sandbox_fusion_url,\n                headers=headers,\n                data=payload,\n                timeout=request_timeout,  # Use the calculated timeout\n            )\n\n            # Check for Gateway Timeout (504) specifically for retrying\n            if response.status_code == 504:\n                last_error = (\n                    f\"{log_prefix}API Request Error: Gateway Timeout (504) on attempt \"\n                    f\"{attempt + 1}/{MAX_RETRIES}\"\n                )  # <-- Use internal log_prefix\n                logger.warning(last_error)\n                if attempt < MAX_RETRIES - 1:  # Don't sleep after the last attempt\n                    # Calculate increasing delay (e.g., 1s, 2s, 4s, ...) or (1s, 2s, 3s, ...)\n                    # Simple linear increase: delay = INITIAL_RETRY_DELAY * (attempt + 1)\n                    # Exponential backoff: delay = INITIAL_RETRY_DELAY * (2 ** attempt)\n                    delay = INITIAL_RETRY_DELAY * (attempt + 1)  # Using linear increase for simplicity\n                    logger.info(f\"{log_prefix}Retrying after {delay} seconds...\")  # <-- Use internal log_prefix\n                    time.sleep(delay)\n                continue  # Go to the next retry attempt\n\n            # Check for other HTTP errors (e.g., 4xx, other 5xx)\n            response.raise_for_status()\n\n            # If successful (status code 2xx)\n            logger.info(\n                f\"{log_prefix}Sandbox API call successful on attempt {attempt + 1}\"\n            )  # <-- Use internal log_prefix\n            return response.json(), None\n\n        except requests.exceptions.RequestException as e:\n            last_error = f\"{log_prefix}API Request Error: {e}\"  # <-- Use internal log_prefix\n            break  # Exit retry loop on non-504 request errors\n        except json.JSONDecodeError as e:\n            raw_response_text = response.text if \"response\" in locals() else \"N/A\"\n            last_error = f\"{log_prefix}API Response JSON Decode Error: {e}\"  # <-- Use internal log_prefix\n            break  # Exit retry loop on JSON decode errors\n        except Exception as e:\n            last_error = f\"{log_prefix}Unexpected Error: {e}\"  # <-- Use internal log_prefix\n            break  # Exit retry loop on other unexpected errors\n\n    # If loop finishes without returning success, return the last recorded error\n    logger.error(f\"{log_prefix}Sandbox API call failed. Last error: {last_error}\")  # <-- Use internal log_prefix\n    # Return the error message without the prefix, as the caller doesn't need the internal ID\n    # Ensure API call failure returns error message, leading to -1 in check_correctness\n    return None, last_error.replace(log_prefix, \"API Call Failed: \") if last_error else \"API Call Failed after retries\"\n\n\ndef _process_single_case(\n    case_index: int,\n    stdin_data: Any,\n    expected_output: Any,\n    sandbox_fusion_url: str,\n    generation: str,\n    timeout: int,\n    memory_limit_mb: int,\n    language: str,\n    concurrent_semaphore: Optional[threading.Semaphore] = None,\n    fn_name: Optional[str] = None,\n) -> tuple[int, dict[str, Any]]:\n    \"\"\"Helper function to process a single test case.\"\"\"\n    api_response = None\n    error_msg = None\n    logger.info(f\"Processing test case {case_index + 1}.\")\n\n    current_generation_code = generation\n\n    if fn_name and language == \"python\":\n        # Wrapper assumes stdin_data is a JSON string for function arguments.\n        wrapper_code = f\"\"\"\nimport traceback\nfrom string import *\nfrom re import *\nfrom datetime import *\nfrom collections import *\nfrom heapq import *\nfrom bisect import *\nfrom copy import *\nfrom math import *\nfrom random import *\nfrom statistics import *\nfrom itertools import *\nfrom functools import *\nfrom operator import *\nfrom io import *\nfrom sys import *\nfrom json import *\nfrom builtins import *\nfrom typing import *\nimport string\nimport re\nimport datetime\nimport collections\nimport heapq\nimport bisect\nimport copy\nimport math\nimport random\nimport statistics\nimport itertools\nimport functools\nimport operator\nimport io\nimport sys\nimport json\n\n# === User's Original Code START ===\n{generation}\n# === User's Original Code END ===\n\n_SANDBOX_FN_NAME = \"{fn_name}\"\n\ndef _execute_user_function():\n    # --- Input Parsing ---\n    _raw_input_str = sys.stdin.read()\n    _args = []\n    if _raw_input_str.strip(): # If there's input\n        try:\n            _args = [json.loads(line) for line in _raw_input_str.split('\\\\n')]\n        except json.JSONDecodeError as _je:\n            sys.stderr.write(f\"WrapperError: Invalid JSON input for '{{_SANDBOX_FN_NAME}}': {{_je}}\\\\nInput was: \"\n                              f\"{{_raw_input_str[:200]}}\\\\n\")\n            return None, True # result, error_occurred\n\n    # --- Function Location and Execution ---\n    try:\n        _target_callable = None\n        # Try global scope first\n        if _SANDBOX_FN_NAME in globals():\n            _target_callable = globals()[_SANDBOX_FN_NAME]\n        # Else, if 'Solution' class exists, try to get its method\n        elif 'Solution' in globals():\n            _Solution_class = globals()['Solution']\n            # Attempt to instantiate and get method.\n            # Errors (e.g., Solution not a class, instantiation fails, method missing)\n            # will be caught by the broad except block below.\n            _solution_instance = _Solution_class()\n            _target_callable = getattr(_solution_instance, _SANDBOX_FN_NAME)\n\n        if not _target_callable:\n            sys.stderr.write(f\"WrapperError: Function or method '{{_SANDBOX_FN_NAME}}' not found.\\\\n\")\n            return None, True # result, error_occurred\n\n        _fn_result = _target_callable(*_args)\n        return _fn_result, False # result, no_error\n    except Exception: # Catches errors from Solution instantiation, getattr, or function call\n        sys.stderr.write(f\"Error during setup or execution of '{{_SANDBOX_FN_NAME}}':\\\\n{{traceback.format_exc()}}\\\\n\")\n        return None, True # result, error_occurred\n\nif __name__ == '__main__':\n    _result, _error_occurred = _execute_user_function()\n\n    if not _error_occurred:\n        # Serialize result to stdout\n        if isinstance(_result, (dict, list, tuple)) or _result is None or isinstance(_result, bool):\n            print(json.dumps(_result))\n        elif isinstance(_result, (int, float, str)):\n            print(str(_result)) # Ensure string conversion for print\n        else:\n            # For other types, default to string representation.\n            print(str(_result))\n    # Optional: To explicitly exit with an error code if the sandbox relies on it\n    # else:\n    #    sys.exit(1)\n\"\"\"\n        current_generation_code = wrapper_code\n\n    stdin = None if stdin_data is None else str(stdin_data)\n    try:\n        if concurrent_semaphore:\n            # logger.debug(f\"Case {case_index + 1}: Attempting to acquire semaphore.\")\n            with concurrent_semaphore:\n                # logger.debug(f\"Case {case_index + 1}: Semaphore acquired. Calling API.\")\n                api_response, error_msg = call_sandbox_api(\n                    sandbox_fusion_url=sandbox_fusion_url,\n                    code=current_generation_code,\n                    stdin=stdin,\n                    compile_timeout=timeout,\n                    run_timeout=timeout,\n                    memory_limit_mb=memory_limit_mb,\n                    language=language,\n                )\n            # logger.debug(f\"Case {case_index + 1}: Semaphore released.\")\n        else:\n            api_response, error_msg = call_sandbox_api(\n                sandbox_fusion_url=sandbox_fusion_url,\n                code=current_generation_code,\n                stdin=stdin,\n                compile_timeout=timeout,\n                run_timeout=timeout,\n                memory_limit_mb=memory_limit_mb,\n                language=language,\n            )\n    except Exception as e:\n        error_msg = f\"API Request Exception during check_correctness for case {case_index + 1}: {e}\"\n        logger.error(f\"Case {case_index + 1}: {error_msg}\")\n        traceback.print_exc()\n\n    metadata = {\n        \"case_index\": case_index,\n        \"input\": stdin,\n        \"expected_output\": str(expected_output) if expected_output else None,\n        \"api_request_error\": error_msg,\n        \"api_response\": None,\n        \"status\": \"unknown\",\n        \"stdout\": None,\n        \"stderr\": None,\n        \"exit_code\": None,\n        \"duration\": None,\n        \"compile_duration\": None,\n        \"compile_stderr\": None,\n        \"api_status\": None,\n        \"compile_status\": None,\n        \"run_status\": None,\n    }\n    result_status = -1  # Default error: API request error or unknown sandbox error\n\n    if error_msg:\n        metadata[\"status\"] = \"api_error\"\n        result_status = -1  # API request itself failed (includes timeout after retries)\n        logger.error(f\"Case {case_index}: API error occurred: {error_msg}\")\n        # Log code and input only on error for brevity\n        generation_to_log = generation[:200] + \"...\" if len(generation) > 200 else generation\n        logger.error(f\"Case {case_index}: code: {generation_to_log}\")\n        logger.error(f\"Case {case_index}: input: {stdin}\")\n    elif api_response:\n        # --- Add debug logging ---\n        logger.debug(f\"Case {case_index}: API Response: {api_response}\")\n        metadata[\"api_response\"] = api_response\n        metadata[\"api_status\"] = api_response.get(\"status\")\n        compile_result = api_response.get(\"compile_result\")\n        run_result = api_response.get(\"run_result\")\n\n        # Extract compile information\n        if compile_result:\n            metadata[\"compile_status\"] = compile_result.get(\"status\")\n            metadata[\"compile_duration\"] = compile_result.get(\"execution_time\")\n            metadata[\"compile_stderr\"] = compile_result.get(\"stderr\")\n\n        # Extract run information\n        if run_result:\n            metadata[\"run_status\"] = run_result.get(\"status\")\n            metadata[\"stdout\"] = run_result.get(\"stdout\")\n            metadata[\"stderr\"] = run_result.get(\"stderr\")  # stderr during runtime\n            metadata[\"exit_code\"] = run_result.get(\"return_code\")\n            metadata[\"duration\"] = run_result.get(\"execution_time\")\n\n        # --- Determine status based on API response ---\n        api_status = metadata[\"api_status\"]\n\n        if api_status == \"SandboxError\":\n            metadata[\"status\"] = \"sandbox_error\"\n            result_status = -1  # Internal sandbox error\n        elif api_status == \"Failed\":\n            # --- Add debug logging ---\n            logger.debug(f\"API returned Failed status. Response: {api_response}\")\n            logger.debug(f\"Compile Result: {compile_result}\")\n            logger.debug(f\"Run Result: {run_result}\")\n            # --- Check the logic here ---\n            # Compile failed or timed out\n            is_compile_error = compile_result and (\n                metadata[\"compile_status\"] in [\"Error\", \"TimeLimitExceeded\"]\n                or (metadata[\"compile_status\"] == \"Finished\" and compile_result.get(\"return_code\") != 0)\n            )\n            if is_compile_error:\n                # Differentiate between compile_error and compile_timeout based on specific status\n                if metadata[\"compile_status\"] == \"TimeLimitExceeded\":\n                    metadata[\"status\"] = \"compile_timeout\"\n                else:  # Includes Error and Finished but return_code != 0 cases\n                    metadata[\"status\"] = \"compile_error\"\n                result_status = -4\n            # Run failed or timed out\n            elif run_result:\n                # Modified condition: Check for TimeLimitExceeded OR (Finished with non-zero exit code) OR Error status\n                is_runtime_error = (\n                    metadata[\"run_status\"] == \"TimeLimitExceeded\"\n                    or metadata[\"run_status\"] == \"Error\"\n                    or (metadata[\"run_status\"] == \"Finished\" and run_result.get(\"return_code\") != 0)\n                )\n                if is_runtime_error:\n                    if metadata[\"run_status\"] == \"TimeLimitExceeded\":\n                        metadata[\"status\"] = \"timeout\"  # Runtime timeout\n                        result_status = -3\n                    else:  # Includes Error and Finished with non-zero return_code\n                        metadata[\"status\"] = \"runtime_error\"\n                        result_status = -2\n                else:\n                    # Other Failed status with run_result, classify as unknown failure\n                    logger.warning(f\"Unknown run_status '{metadata['run_status']}' or state within Failed API status.\")\n                    metadata[\"status\"] = \"unknown_failure\"\n                    result_status = -1  # Default to -1\n            else:\n                # Status is Failed but neither a clear compile error nor run_result exists\n                logger.warning(\"API status Failed but cannot determine specific error type (compile/run).\")\n                metadata[\"status\"] = \"unknown_failure_state\"\n                result_status = -1  # Default to -1\n        elif api_status == \"Success\":\n            # Run completed successfully, now check the answer\n            if run_result and metadata[\"run_status\"] == \"Finished\":\n                actual_output = metadata[\"stdout\"] if metadata[\"stdout\"] is not None else \"\"\n                # Note: Output might contain trailing newlines, need normalization\n                if expected_output is None or str(actual_output).rstrip(\"\\n\") == str(expected_output).rstrip(\"\\n\"):\n                    result_status = True\n                    metadata[\"status\"] = \"success\"\n                else:\n                    result_status = False\n                    metadata[\"status\"] = \"wrong_answer\"\n            else:\n                # Status is Success but run_result status is not Finished, this is unexpected\n                metadata[\"status\"] = \"unexpected_success_state\"\n                result_status = -1  # Classify as unknown error\n        else:\n            # API returned an unknown top-level status\n            logger.warning(f\"Unknown API status received: {api_status}\")\n            metadata[\"status\"] = f\"unknown_api_status_{api_status}\"\n            result_status = -1  # Default to -1\n    else:  # api_response is None and no error_msg (Should not happen with current call_sandbox_api logic)\n        metadata[\"status\"] = \"unknown_api_state\"\n        result_status = -1\n        logger.error(f\"Case {case_index}: Unknown API state (no response and no error message).\")\n    return result_status, metadata\n\n\ndef check_correctness(\n    sandbox_fusion_url: str,\n    in_outs: Optional[dict],\n    generation: str,\n    timeout: int = DEFAULT_TIMEOUT,\n    memory_limit_mb: int = 1024,\n    language: str = \"python\",\n    concurrent_semaphore: Optional[threading.Semaphore] = None,\n) -> tuple[list[Any], list[dict[str, Any]]]:\n    \"\"\"\n    Checks the correctness of code generation using the remote sandbox API,\n    processing test cases concurrently.\n\n    Args:\n        sandbox_fusion_url: The URL of the sandbox fusion API.\n        in_outs: Dictionary containing \"inputs\" and \"outputs\" lists.\n        generation: The generated code string.\n        timeout: Timeout for each test case (compile and run share this timeout).\n        language: The programming language of the code.\n\n    Returns:\n        A tuple (results, metadata_list).\n        results: A list containing the test result for each input/output pair\n                 (True/False/-1 api/sandbox err, -2 runtime err, -3 timeout, -4 compile err).\n                 Results are ordered corresponding to the inputs.\n        metadata_list: A list containing metadata dictionaries for each test case,\n                       ordered corresponding to the inputs.\n    \"\"\"\n    logger.info(\"Starting correctness check for generation.\")\n\n    if not in_outs or \"inputs\" not in in_outs or \"outputs\" not in in_outs:\n        logger.warning(\"Invalid in_outs format provided.\")\n        return [-1], [{\"error\": \"Invalid input/output data\"}]\n\n    inputs = in_outs[\"inputs\"]\n    expected_outputs = in_outs[\"outputs\"]\n    fn_name = in_outs.get(\"fn_name\")\n    num_cases = len(inputs)\n    assert_cases = in_outs.get(\"assert_case\", [\"\"] * num_cases)  # Default to empty strings if not provided\n    results = [None] * num_cases  # Initialize with placeholders\n    metadata_list = [None] * num_cases  # Initialize with placeholders\n\n    if num_cases == 0:\n        logger.warning(\"Empty inputs provided.\")\n        return [], []\n\n    if len(inputs) != len(expected_outputs):\n        logger.warning(f\"Mismatch between number of inputs ({len(inputs)}) and outputs ({len(expected_outputs)}).\")\n        # Return error based on the number of inputs provided\n        return [-1] * num_cases, [{\"error\": \"Input/output count mismatch\", \"case_index\": i} for i in range(num_cases)]\n\n    # If assert_cases is provided, it overrides inputs and outputs\n    if len(assert_cases) != num_cases:\n        logger.warning(\n            f\"Mismatch between number of assert cases ({len(assert_cases)}) and inputs/outputs ({num_cases}).\"\n        )\n        return [-1] * num_cases, [{\"error\": \"Input/output count mismatch\", \"case_index\": i} for i in range(num_cases)]\n\n    first_compile_error_index = -1\n\n    # max_workers is limited by sandbox_fusion_max_concurrent from concurrent_semaphore\n    with concurrent.futures.ThreadPoolExecutor(max_workers=max(32, os.cpu_count() * 5)) as executor:\n        # Submit all tasks, passing the concurrent_semaphore to _process_single_case\n        future_to_index = {\n            executor.submit(\n                _process_single_case,\n                i,\n                stdin_data,\n                expected_outputs[i],\n                sandbox_fusion_url,\n                generation + \"\\n\\n\" + assert_cases[i],  # Append assert case to generation\n                timeout,\n                memory_limit_mb,\n                language,\n                concurrent_semaphore,\n                fn_name,\n            ): i\n            for i, stdin_data in enumerate(inputs)\n        }\n\n        # Process results as they complete\n        for future in concurrent.futures.as_completed(future_to_index):\n            index = future_to_index[future]\n            try:\n                result_status, metadata = future.result()\n                results[index] = result_status\n                metadata_list[index] = metadata\n\n                # Check for compile error (-4)\n                if result_status == -4:\n                    if first_compile_error_index == -1 or index < first_compile_error_index:\n                        first_compile_error_index = index\n                    # Optimization: could potentially cancel futures for index > first_compile_error_index\n                    # However, cancellation is not guaranteed. Post-processing is safer.\n\n            except Exception as exc:\n                logger.error(f\"Test case {index} generated an exception: {exc}\")\n                traceback.print_exc()\n                results[index] = -1  # Mark as API/internal error\n                metadata_list[index] = {\n                    \"case_index\": index,\n                    \"input\": str(inputs[index]),\n                    \"expected_output\": str(expected_outputs[index]) if expected_outputs[index] else None,\n                    \"api_request_error\": f\"Internal execution error: {exc}\",\n                    \"status\": \"internal_error\",\n                }\n\n    # Post-processing for compile errors\n    if first_compile_error_index != -1:\n        logger.warning(\n            f\"Compile error detected in case {first_compile_error_index}. Marking subsequent cases as compile errors.\"\n        )\n        for i in range(first_compile_error_index + 1, num_cases):\n            # Only update if not already processed (though it should be None or have a result)\n            if results[i] != -4:  # Avoid overwriting if it somehow already got -4\n                results[i] = -4\n                # Update or create metadata for skipped cases due to compile error\n                if metadata_list[i] is None:  # If future failed before returning metadata\n                    metadata_list[i] = {\n                        \"case_index\": i,\n                        \"input\": str(inputs[i]),\n                        \"expected_output\": str(expected_outputs[i]) if expected_outputs[i] else None,\n                        \"api_request_error\": None,\n                        \"status\": \"compile_error_skipped\",  # Indicate skipped due to prior compile error\n                    }\n                else:  # If future completed but result is overridden\n                    metadata_list[i][\"status\"] = \"compile_error_skipped\"\n\n    logger.info(f\"Correctness check finished. Results: {results}\")\n    return results, metadata_list\n"
  },
  {
    "path": "verl/utils/reward_score/search_r1_like_qa_em.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\r\n# Copyright 2023-2024 SGLang Team\r\n# Copyright 2025 Search-R1 Contributors\r\n#\r\n# Licensed under the Apache License, Version 2.0 (the \"License\");\r\n# you may not use this file except in compliance with the License.\r\n# You may obtain a copy of the License at\r\n#\r\n#     http://www.apache.org/licenses/LICENSE-2.0\r\n#\r\n# Unless required by applicable law or agreed to in writing, software\r\n# distributed under the License is distributed on an \"AS IS\" BASIS,\r\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\r\n# See the License for the specific language governing permissions and\r\n# limitations under the License.\r\n# Adapted from https://github.com/PeterGriffinJin/Search-R1/blob/main/verl/utils/reward_score/qa_em.py\r\n\r\nimport random\r\nimport re\r\nimport string\r\n\r\n\r\ndef normalize_answer(s):\r\n    def remove_articles(text):\r\n        return re.sub(r\"\\b(a|an|the)\\b\", \" \", text)\r\n\r\n    def white_space_fix(text):\r\n        return \" \".join(text.split())\r\n\r\n    def remove_punc(text):\r\n        exclude = set(string.punctuation)\r\n        return \"\".join(ch for ch in text if ch not in exclude)\r\n\r\n    def lower(text):\r\n        return text.lower()\r\n\r\n    return white_space_fix(remove_articles(remove_punc(lower(s))))\r\n\r\n\r\ndef em_check(prediction, golden_answers):\r\n    if isinstance(golden_answers, str):\r\n        golden_answers = [golden_answers]\r\n    normalized_prediction = normalize_answer(prediction)\r\n    score = 0\r\n    for golden_answer in golden_answers:\r\n        golden_answer = normalize_answer(golden_answer)\r\n        if golden_answer == normalized_prediction:\r\n            score = 1\r\n            break\r\n    return score\r\n\r\n\r\ndef subem_check(prediction, golden_answers):\r\n    if isinstance(golden_answers, str):\r\n        golden_answers = [golden_answers]\r\n    normalized_prediction = normalize_answer(prediction)\r\n    score = 0\r\n    for golden_answer in golden_answers:\r\n        golden_answer = normalize_answer(golden_answer)\r\n        if golden_answer in normalized_prediction:\r\n            score = 1\r\n            break\r\n    return score\r\n\r\n\r\ndef extract_solution(solution_str):\r\n    \"\"\"Extract the equation from the solution string.\"\"\"\r\n    # Remove everything before the first \"Assistant:\"\r\n    # if \"Assistant:\" in solution_str:\r\n    #     solution_str = solution_str.split(\"Assistant:\", 1)[1]\r\n    # elif \"<|im_start|>assistant\" in solution_str:\r\n    #     solution_str = solution_str.split(\"<|im_start|>assistant\", 1)[1]\r\n    # else:\r\n    #     return None\r\n    # solution_str = solution_str.split('\\n')[-1]\r\n\r\n    answer_pattern = r\"<answer>(.*?)</answer>\"\r\n    match = re.finditer(answer_pattern, solution_str, re.DOTALL)\r\n    matches = list(match)\r\n\r\n    # If there are 0  matches, return None\r\n    if len(matches) < 1:\r\n        return None\r\n\r\n    # If there are 2 or more matches, return the last one\r\n    return matches[-1].group(1).strip()\r\n\r\n\r\ndef count_answer_tags(text):\r\n    opening_tags = text.count(\"<answer>\")\r\n    closing_tags = text.count(\"</answer>\")\r\n\r\n    return opening_tags, closing_tags\r\n\r\n\r\ndef compute_score(solution_str, ground_truth, method=\"strict\", format_score=0.0, score=1.0):\r\n    \"\"\"The scoring function for exact match (EM).\r\n\r\n    Args:\r\n        solution_str: the solution text\r\n        ground_truth: the ground truth\r\n        method: the method to extract the solution, choices are 'strict' and 'flexible'\r\n        format_score: the score for the format\r\n        score: the score for the correct answer\r\n    \"\"\"\r\n    answer = extract_solution(solution_str=solution_str)\r\n    open_count, close_count = count_answer_tags(solution_str)\r\n    do_print = random.randint(1, 64) == 1\r\n\r\n    if do_print:\r\n        print(\"--------------------------------\")\r\n        print(f\"Golden answers: {ground_truth['target']}\")\r\n        if answer is not None:\r\n            print(f\"Extracted answer is not None: {answer}\")\r\n        else:\r\n            print(\"Extracted answer: None!\")\r\n        print(f\"Solution string: {solution_str}\")\r\n\r\n    if answer is None:\r\n        return 0\r\n    else:\r\n        if em_check(answer, ground_truth[\"target\"]):\r\n            if open_count > 10 or close_count > 10:  # prevent output a lot of </answer>\r\n                score = score / 4\r\n                return score\r\n            return score\r\n        else:\r\n            return format_score\r\n\r\n\r\ndef compute_score_subem(solution_str, ground_truth, method=\"strict\", format_score=0.0, score=1.0):\r\n    \"\"\"The scoring function for substring exact match (EM).\r\n\r\n    Args:\r\n        solution_str: the solution text\r\n        ground_truth: the ground truth\r\n        method: the method to extract the solution, choices are 'strict' and 'flexible'\r\n        format_score: the score for the format\r\n        score: the score for the correct answer\r\n    \"\"\"\r\n    answer = extract_solution(solution_str=solution_str)\r\n    do_print = random.randint(1, 64) == 1\r\n\r\n    if do_print:\r\n        print(\"--------------------------------\")\r\n        print(f\"Golden answers: {ground_truth['target']}\")\r\n        print(f\"Extracted answer: {answer}\")\r\n        print(f\"Solution string: {solution_str}\")\r\n\r\n    if answer is None:\r\n        return 0\r\n    else:\r\n        if subem_check(answer, ground_truth[\"target\"]):\r\n            return score\r\n        else:\r\n            return format_score\r\n"
  },
  {
    "path": "verl/utils/rollout_skip.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom pathlib import Path\n\nfrom verl.protocol import DataProto\n\n\nclass RolloutSkip:\n    \"\"\"\n    RolloutSkip skips sequence generation during rollout by attempting to load previously dumped data.\n    If no dumped data is found, it generates new sequences and saves them to disk.\n\n    Args:\n        config: The configuration object containing rollout settings.\n        rollout_wg: The worker group that handles the rollout process.\n\n    Note:\n        When rollout.n or rollout.gen_batch_size differ from previous runs,\n        new sequences will be generated and saved with different filenames.\n    \"\"\"\n\n    print_mark = \"[RolloutSkip()]\"\n\n    def __init__(self, config, rollout_wg):\n        self.rollout_config = config.actor_rollout_ref.rollout\n        self.exp_name = config.data.get(\"experiment_name\", \"\")\n        self.project_name = config.data.get(\"project_name\", \"\")\n\n        self.n = int(self.rollout_config.get(\"n\", 0))\n        self.gbs = int(config.data.get(\"gen_batch_size\", config.data.get(\"train_batch_size\", 0)))\n\n        self.dumped_dir = Path(self.rollout_config.get(\"skip_dump_dir\", \"/tmp/verl/rollout_dump\"))\n        self.dumped_dir.mkdir(parents=True, exist_ok=True)\n\n        # Check if path is in Ray temporary directory\n        if str(self.dumped_dir.absolute()).startswith(\"/tmp/ray/session\"):\n            print(\n                f\"\\033[33m{self.print_mark} Warning: \\nUsing dump path \",\n                f\"'{self.dumped_dir.absolute()}' is not recommended \",\n                \"as it's located in /tmp/ray/session*\\033[0m\",\n                flush=True,\n            )\n\n        print(\n            f\"{self.print_mark} Rollout skip dump path set to: \",\n            f\"{self.dumped_dir.absolute()}\",\n            flush=True,\n        )\n\n        self._rollout_wg = rollout_wg\n\n    @property\n    def curr_path_dump(self):\n        return self.dumped_dir.joinpath(f\"{self.exp_name}_{self.project_name}_GBS{self.gbs}__N{self.n}\").absolute()\n\n    def wrap_generate_sequences(self):\n        try:\n            self._rollout_wg.generate_sequences = wrap_generate_sequences(self, self._rollout_wg)\n            print(\n                f\"{self.print_mark} Successfully patched `actor_rollout_wg.generate_sequences()`\",\n                flush=True,\n            )\n        except Exception as e:\n            raise RuntimeError(\n                \"{self.print_mark} Failed to patch `actor_rollout_wg.generate_sequences()`\",\n                flush=True,\n            ) from e\n\n    def try_load(self):\n        if not self.curr_path_dump.exists():\n            print(\n                f\"{self.print_mark} No data dump found at {self.curr_path_dump}.\",\n                \"The trainer will generate and automatically dump the data for this first run.\",\n                flush=True,\n            )\n            return None\n\n        try:\n            # * Load\n            ret_batch = DataProto.load_from_disk(self.curr_path_dump)\n            print(\n                f\"\\033[32m{self.print_mark} Successfully load pre-generated data from {self.curr_path_dump}\\033[0m\",\n                flush=True,\n            )\n            return ret_batch\n        except Exception as e:\n            print(\n                f\"\\033[31m{self.print_mark} Failed to load pre-generated data from {self.curr_path_dump}\",\n                f\"Error: {str(e)}\\033[0m\",\n                flush=True,\n            )\n            return None\n\n    def dump(self, outputs: DataProto):\n        try:\n            outputs.save_to_disk(self.curr_path_dump)\n            print(\n                f\"\\033[32m{self.print_mark} Successfully dump data in {self.curr_path_dump}\\033[0m\",\n                flush=True,\n            )\n        except Exception as e:\n            print(\n                f\"\\033[31m{self.print_mark} Failed to dump data in {self.curr_path_dump}: {e}\\033[0m\",\n                flush=True,\n            )\n\n\ndef wrap_generate_sequences(rolloutskip: RolloutSkip, rollout_wg):\n    generate_sequences = rollout_wg.generate_sequences\n\n    def warp_fn(batch, **kwargs):\n        gen_batch_output = rolloutskip.try_load()\n\n        if gen_batch_output is None:\n            # * 1. Generation\n            gen_batch_output = generate_sequences(batch, **kwargs)\n            # * 2. Dump\n            rolloutskip.dump(gen_batch_output)\n        return gen_batch_output\n\n    return warp_fn\n"
  },
  {
    "path": "verl/utils/rollout_trace.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport contextlib\nimport functools\nimport inspect\nimport os\nfrom contextvars import ContextVar\nfrom typing import Optional\n\nfrom pydantic import BaseModel\n\nfrom verl.utils.ray_utils import get_event_loop\n\n_trace_enabled: ContextVar[bool] = ContextVar(\"_trace_enabled\", default=True)\n\n\nclass RolloutTraceConfig:\n    \"\"\"Configuration for rollout tracing with various backends.\n\n    Singleton configuration class for managing rollout trace settings across different\n    tracing backends like Weave and MLflow.\n\n    Args:\n        backend (Optional[str]): Tracing backend to use ('weave', 'mlflow', or None).\n        client (Optional[object]): Client instance for the selected backend.\n        token2text (bool): Whether to convert tokens to text in traces. Defaults to False.\n        project_name (str): Name of the project for tracing.\n        experiment_name (str): Name of the experiment for tracing.\n        max_samples_per_step_per_worker (Optional[int]): Maximum number of unique samples to trace\n            per worker per step. If None, all samples are traced. If set, each worker will randomly\n            select up to this many unique samples to trace (including all their rollouts for GRPO).\n            Total traces = max_samples_per_step_per_worker * num_workers * n_rollouts_per_sample.\n    \"\"\"\n\n    _instance: Optional[\"RolloutTraceConfig\"] = None\n    backend: Optional[str] = None\n    client: Optional[object] = None\n    token2text: bool = False\n    _initialized: bool = False\n    project_name: str = None\n    experiment_name: str = None\n    max_samples_per_step_per_worker: Optional[int] = None\n\n    def __new__(cls, *args, **kwargs):\n        if cls._instance is None:\n            cls._instance = super().__new__(cls)\n            cls._instance._initialized = False\n        return cls._instance\n\n    @classmethod\n    def get_instance(cls) -> \"RolloutTraceConfig\":\n        if cls._instance is None:\n            cls._instance = cls()\n        return cls._instance\n\n    @classmethod\n    def init(\n        cls,\n        project_name: str,\n        experiment_name: str,\n        backend: str,\n        token2text: bool = False,\n        max_samples_per_step_per_worker: Optional[int] = None,\n    ):\n        config = cls.get_instance()\n        if config._initialized:\n            return\n\n        config.backend = backend\n        config.token2text = token2text\n        config.project_name = project_name\n        config.experiment_name = experiment_name\n        config.max_samples_per_step_per_worker = max_samples_per_step_per_worker\n\n        if backend == \"weave\":\n            import weave\n\n            config.client = weave.init(project_name)\n        elif backend == \"mlflow\":\n            import mlflow\n\n            mlflow.config.enable_async_logging()\n            config.client = mlflow\n\n            MLFLOW_TRACKING_URI = os.environ.get(\"MLFLOW_TRACKING_URI\", \"sqlite:////tmp/mlruns.db\")\n            mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)\n\n            mlflow.set_experiment(project_name)\n        else:\n            config.client = None\n\n        config._initialized = True\n\n    @classmethod\n    def get_backend(cls) -> Optional[str]:\n        return cls.get_instance().backend\n\n    @classmethod\n    def get_client(cls) -> Optional[object]:\n        return cls.get_instance().client\n\n    @classmethod\n    def enable_token2text(cls) -> Optional[bool]:\n        return cls.get_instance().token2text\n\n    @classmethod\n    def reset(cls):\n        cls._instance = None\n\n\n@contextlib.contextmanager\ndef rollout_trace_attr(\n    sample_index=None, step=None, rollout_n=None, name=\"rollout_trace\", validate=False, trace: bool = True\n):\n    \"\"\"A context manager to add attributes to a trace for the configured backend.\n\n    Args:\n        sample_index: Sample index for the trace.\n        step: Training step number.\n        rollout_n: Rollout number (for GRPO with multiple rollouts per sample).\n        name: Name for the trace span (used by mlflow backend).\n        validate: Whether this is a validation run.\n        trace: If False, disables tracing for the duration of the context.\n    \"\"\"\n    backend = RolloutTraceConfig.get_backend()\n\n    should_skip = backend is not None and not trace\n\n    if should_skip:\n        token = _trace_enabled.set(False)\n        try:\n            yield\n        finally:\n            _trace_enabled.reset(token)\n        return\n\n    # Build attributes for the trace\n    attributes = {}\n    if backend:\n        if sample_index is not None:\n            attributes[\"sample_index\"] = sample_index\n        if step is not None:\n            attributes[\"step\"] = step\n        if rollout_n is not None:\n            attributes[\"rollout_n\"] = rollout_n\n        attributes[\"validate\"] = validate\n        attributes[\"experiment_name\"] = RolloutTraceConfig.get_instance().experiment_name\n\n    if not attributes or backend is None:\n        yield\n        return\n\n    if backend == \"weave\":\n        import weave\n\n        with weave.attributes(attributes):\n            yield\n    elif backend == \"mlflow\":\n        import mlflow\n\n        with mlflow.start_span(name=name) as span:\n            trace_id = span.trace_id\n            for key, value in attributes.items():\n                mlflow.set_trace_tag(trace_id, str(key), str(value))\n            yield\n    else:\n        yield\n\n\ndef rollout_trace_op(func):\n    @functools.wraps(func)\n    async def async_wrapper(self, *args, **kwargs):\n        if not _trace_enabled.get():\n            return await func(self, *args, **kwargs)\n\n        backend = RolloutTraceConfig.get_backend()\n        enable_token2text = RolloutTraceConfig.enable_token2text()\n        if backend is None:\n            return await func(self, *args, **kwargs)\n\n        sig = inspect.signature(func)\n        bound_args = sig.bind(self, *args, **kwargs)\n        bound_args.apply_defaults()\n        inputs = dict(bound_args.arguments)\n        del inputs[\"self\"]\n\n        async def add_token2text(self, result):\n            if hasattr(result, \"prompt_ids\") and hasattr(self, \"tokenizer\") and hasattr(self.tokenizer, \"decode\"):\n                # Use model_dump() for Pydantic models to get a proper copy,\n                # otherwise vars() returns a reference to internal __dict__ which\n                # can cause serialization issues with MLflow\n                if isinstance(result, BaseModel):\n                    _result = result.model_dump()\n                else:\n                    _result = dict(vars(result))\n                loop = get_event_loop()\n                if hasattr(result, \"prompt_ids\"):\n                    prompt_text = await loop.run_in_executor(None, self.tokenizer.decode, result.prompt_ids)\n                    _result[\"prompt_text\"] = prompt_text\n\n                if hasattr(result, \"response_ids\"):\n                    response_text = await loop.run_in_executor(None, self.tokenizer.decode, result.response_ids)\n                    _result[\"response_text\"] = response_text\n                return _result\n            return result\n\n        if backend == \"weave\":\n            tracer = RolloutTraceConfig.get_client()\n            from weave.trace.context import call_context\n\n            cur_attributes = {**call_context.call_attributes.get()}\n            call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes)\n            try:\n                result = await func(self, *args, **kwargs)\n\n                if enable_token2text:\n                    _result = await add_token2text(self, result)\n                    tracer.finish_call(call, output=_result)\n                else:\n                    tracer.finish_call(call, output=result)\n\n                return result\n\n            except Exception as e:\n                tracer.finish_call(call, exception=e)\n                raise e\n        elif backend == \"mlflow\":\n            import mlflow\n\n            with mlflow.start_span(name=func.__qualname__) as span:\n                span.set_inputs(inputs)\n                result = await func(self, *args, **kwargs)\n                if enable_token2text:\n                    _result = await add_token2text(self, result)\n                    span.set_outputs(_result)\n                else:\n                    span.set_outputs(result)\n\n            return result\n\n        else:\n            return await func(self, *args, **kwargs)\n\n    @functools.wraps(func)\n    def wrapper(self, *args, **kwargs):\n        if not _trace_enabled.get():\n            return func(self, *args, **kwargs)\n\n        backend = RolloutTraceConfig.get_backend()\n        if backend is None:\n            return func(self, *args, **kwargs)\n\n        sig = inspect.signature(func)\n        bound_args = sig.bind(self, *args, **kwargs)\n        bound_args.apply_defaults()\n        inputs = dict(bound_args.arguments)\n        del inputs[\"self\"]\n\n        if backend == \"weave\":\n            tracer = RolloutTraceConfig.get_client()\n            from weave.trace.context import call_context\n\n            cur_attributes = {**call_context.call_attributes.get()}\n            call = tracer.create_call(op=func.__qualname__, inputs=inputs, attributes=cur_attributes)\n            try:\n                result = func(self, *args, **kwargs)\n                tracer.finish_call(call, output=result)\n                return result\n            except Exception as e:\n                tracer.finish_call(call, exception=e)\n                raise e\n        elif backend == \"mlflow\":\n            import mlflow\n\n            return mlflow.trace(func)(self, *args, **kwargs)\n        else:\n            return func(self, *args, **kwargs)\n\n    return async_wrapper if inspect.iscoroutinefunction(func) else wrapper\n"
  },
  {
    "path": "verl/utils/seqlen_balancing.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport copy\nimport heapq\nfrom itertools import chain\n\nimport torch\nfrom torch import distributed as dist\n\nfrom verl.protocol import DataProto\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.device import get_device_name\n\n\ndef calculate_workload(seqlen_list: torch.Tensor) -> torch.Tensor:\n    \"\"\"Calculate approximate computational workload for transformer attention.\n\n    Estimates FLOPs for dense transformer blocks based on sequence length using\n    the formula: FLOPs ≈ 12 * hidden_size² * seqlen + 2 * hidden_size * seqlen²\n\n    The constants are calibrated for a 7B model (hidden_size=4096), yielding:\n    workload ∝ 24576 * seqlen + seqlen²\n\n    Args:\n        seqlen_list: Sequence lengths as a tensor.\n\n    Returns:\n        torch.Tensor: Estimated workload values proportional to actual FLOPs.\n\n    Note:\n        The returned values are relative workloads, not actual FLOP counts.\n        Useful for balancing computation across data parallel ranks.\n    \"\"\"\n    return 24576 * seqlen_list + seqlen_list**2\n\n\ndef karmarkar_karp(seqlen_list: list[int], k_partitions: int, equal_size: bool) -> list[list[int]]:\n    \"\"\"Partition items into k groups using the Karmarkar-Karp differencing method.\n\n    Implements the Largest Differencing Method (LDM) algorithm for balanced\n    multi-way number partitioning. This heuristic produces near-optimal partitions\n    by iteratively combining the sets with the largest difference.\n\n    Args:\n        seqlen_list: Values to partition (typically sequence lengths or workloads).\n        k_partitions: Number of partitions to create.\n        equal_size: If True, each partition will have exactly len(seqlen_list) / k_partitions\n            items. If False, partitions may have different sizes.\n\n    Returns:\n        list[list[int]]: List of k partitions, each containing indices into seqlen_list.\n\n    See Also:\n        https://en.wikipedia.org/wiki/Largest_differencing_method\n\n    Note:\n        When equal_size=True, len(seqlen_list) must be divisible by k_partitions.\n    \"\"\"\n\n    # see: https://en.wikipedia.org/wiki/Largest_differencing_method\n    class Set:\n        def __init__(self) -> None:\n            self.sum = 0\n            self.items = []\n\n        def add(self, idx: int, val: int):\n            self.items.append((idx, val))\n            self.sum += val\n\n        def merge(self, other):\n            for idx, val in other.items:\n                self.items.append((idx, val))\n                self.sum += val\n\n        def __lt__(self, other):\n            if self.sum != other.sum:\n                return self.sum < other.sum\n            if len(self.items) != len(other.items):\n                return len(self.items) < len(other.items)\n            return self.items < other.items\n\n    class State:\n        def __init__(self, items: list[tuple[int, int]], k: int) -> None:\n            self.k = k\n            # sets should always be decreasing order\n            self.sets = [Set() for _ in range(k)]\n            assert len(items) in [1, k], f\"{len(items)} not in [1, {k}]\"\n            for i, (idx, seqlen) in enumerate(items):\n                self.sets[i].add(idx=idx, val=seqlen)\n            self.sets = sorted(self.sets, reverse=True)\n\n        def get_partitions(self):\n            partitions = []\n            for i in range(len(self.sets)):\n                cur_partition = []\n                for idx, _ in self.sets[i].items:\n                    cur_partition.append(idx)\n                partitions.append(cur_partition)\n            return partitions\n\n        def merge(self, other):\n            for i in range(self.k):\n                self.sets[i].merge(other.sets[self.k - 1 - i])\n            self.sets = sorted(self.sets, reverse=True)\n\n        @property\n        def spread(self) -> int:\n            return self.sets[0].sum - self.sets[-1].sum\n\n        def __lt__(self, other):\n            # least heap, let the state with largest spread to be popped first,\n            # if the spread is the same, let the state who has the largest set\n            # to be popped first.\n            if self.spread != other.spread:\n                return self.spread > other.spread\n            return self.sets[0] > other.sets[0]\n\n        def __repr__(self) -> str:\n            repr_str = \"[\"\n            for i in range(self.k):\n                if i > 0:\n                    repr_str += \",\"\n                repr_str += \"{\"\n                for j, (_, seqlen) in enumerate(self.sets[i].items):\n                    if j > 0:\n                        repr_str += \",\"\n                    repr_str += str(seqlen)\n                repr_str += \"}\"\n            repr_str += \"]\"\n            return repr_str\n\n    sorted_seqlen_list = sorted([(seqlen, i) for i, seqlen in enumerate(seqlen_list)])\n    states_pq = []\n    if equal_size:\n        assert len(seqlen_list) % k_partitions == 0, f\"{len(seqlen_list)} % {k_partitions} != 0\"\n        for offset in range(0, len(sorted_seqlen_list), k_partitions):\n            items = []\n            for i in range(k_partitions):\n                seqlen, idx = sorted_seqlen_list[offset + i]\n                items.append((idx, seqlen))\n            heapq.heappush(states_pq, State(items=items, k=k_partitions))\n    else:\n        for seqlen, idx in sorted_seqlen_list:\n            heapq.heappush(states_pq, State(items=[(idx, seqlen)], k=k_partitions))\n\n    while len(states_pq) > 1:\n        state0 = heapq.heappop(states_pq)\n        state1 = heapq.heappop(states_pq)\n        # merge states\n        state0.merge(state1)\n        heapq.heappush(states_pq, state0)\n\n    final_state = states_pq[0]\n    partitions = final_state.get_partitions()\n    if equal_size:\n        for i, partition in enumerate(partitions):\n            assert len(partition) * k_partitions == len(seqlen_list), (\n                f\"{len(partition)} * {k_partitions} != {len(seqlen_list)}\"\n            )\n    return partitions\n\n\ndef greedy_partition(seqlen_list: list[int], k_partitions: int, equal_size: bool) -> list[list[int]]:\n    \"\"\"Partition items into k groups using a greedy assignment strategy.\n\n    Assigns each item to the partition with the smallest current sum, iterating\n    through items in order. Simpler but typically less optimal than Karmarkar-Karp.\n\n    Args:\n        seqlen_list: Values to partition (typically sequence lengths or workloads).\n        k_partitions: Number of partitions to create.\n        equal_size: If True, adds a bias to ensure equal partition sizes.\n            Requires len(seqlen_list) to be divisible by k_partitions.\n\n    Returns:\n        list[list[int]]: List of k partitions, each containing indices into seqlen_list.\n\n    Note:\n        When equal_size=True, a large bias is added to encourage equal distribution\n        of items before considering the actual values.\n    \"\"\"\n    bias = sum(seqlen_list) + 1 if equal_size else 0\n    sorted_seqlen = [(seqlen + bias, i) for i, seqlen in enumerate(seqlen_list)]\n    partitions = [[] for _ in range(k_partitions)]\n    partition_sums = [0 for _ in range(k_partitions)]\n    for seqlen, i in sorted_seqlen:\n        min_idx = None\n        for j in range(k_partitions):\n            if min_idx is None or partition_sums[j] < partition_sums[min_idx]:\n                min_idx = j\n        partitions[min_idx].append(i)\n        partition_sums[min_idx] += seqlen\n    if equal_size:\n        for i, partition in enumerate(partitions):\n            assert len(partition) * k_partitions == len(seqlen_list), (\n                f\"{len(partition)} * {k_partitions} != {len(seqlen_list)}\"\n            )\n    return partitions\n\n\ndef get_seqlen_balanced_partitions(seqlen_list: list[int], k_partitions: int, equal_size: bool):\n    \"\"\"\n    Calculates partitions of indices from seqlen_list such that the sum of sequence lengths\n    in each partition is balanced. Uses the Karmarkar-Karp differencing method.\n\n    This is useful for balancing workload across devices or batches, especially when\n    dealing with variable sequence lengths.\n\n    Args:\n        seqlen_list (List[int]): A list of sequence lengths for each item.\n        k_partitions (int): The desired number of partitions.\n        equal_size (bool): If True, ensures that each partition has the same number of items.\n                           Requires len(seqlen_list) to be divisible by k_partitions.\n                           If False, partitions can have varying numbers of items, focusing\n                           only on balancing the sum of sequence lengths.\n\n    Returns:\n        List[List[int]]: A list containing k_partitions lists. Each inner list contains the\n                         original indices of the items assigned to that partition. The indices\n                         within each partition list are sorted.\n\n    Raises:\n        AssertionError: If len(seqlen_list) < k_partitions.\n        AssertionError: If equal_size is True and len(seqlen_list) is not divisible by k_partitions.\n        AssertionError: If any resulting partition is empty.\n    \"\"\"\n    assert len(seqlen_list) >= k_partitions, f\"number of items:[{len(seqlen_list)}] < k_partitions:[{k_partitions}]\"\n\n    def _check_and_sort_partitions(partitions):\n        assert len(partitions) == k_partitions, f\"{len(partitions)} != {k_partitions}\"\n        seen_idx = set()\n        sorted_partitions = [None] * k_partitions\n        for i, partition in enumerate(partitions):\n            assert len(partition) > 0, f\"the {i}-th partition is empty\"\n            for idx in partition:\n                seen_idx.add(idx)\n            sorted_partitions[i] = sorted(partition)\n        assert seen_idx == set(range(len(seqlen_list)))\n        return sorted_partitions\n\n    partitions = karmarkar_karp(seqlen_list=seqlen_list, k_partitions=k_partitions, equal_size=equal_size)\n    return _check_and_sort_partitions(partitions)\n\n\ndef log_seqlen_unbalance(seqlen_list: list[int], partitions: list[list[int]], prefix):\n    \"\"\"\n    Calculate and log metrics related to sequence length imbalance before and after partitioning.\n\n    Args:\n        seqlen_list (List[int]): A list of sequence lengths for each item.\n        partitions (List[List[int]]): A list of partitions, where each inner list contains indices\n                                      from seqlen_list assigned to that partition.\n        prefix (str): A prefix to be added to each metric key in the returned dictionary.\n\n    Returns:\n        dict: A dictionary containing metrics related to sequence length imbalance.\n    \"\"\"\n    # Get the number of partitions\n    k_partition = len(partitions)\n    # assert len(seqlen_list) % k_partition == 0\n    batch_size = len(seqlen_list) // k_partition\n    min_sum_seqlen = None\n    max_sum_seqlen = None\n    total_sum_seqlen = 0\n\n    # Iterate over each batch of sequence lengths\n    for offset in range(0, len(seqlen_list), batch_size):\n        cur_sum_seqlen = sum(seqlen_list[offset : offset + batch_size])\n        if min_sum_seqlen is None or cur_sum_seqlen < min_sum_seqlen:\n            min_sum_seqlen = cur_sum_seqlen\n        if max_sum_seqlen is None or cur_sum_seqlen > max_sum_seqlen:\n            max_sum_seqlen = cur_sum_seqlen\n        total_sum_seqlen += cur_sum_seqlen\n\n    balanced_sum_seqlen_list = []\n    for partition in partitions:\n        cur_sum_seqlen_balanced = sum([seqlen_list[i] for i in partition])\n        balanced_sum_seqlen_list.append(cur_sum_seqlen_balanced)\n    # print(\"balanced_sum_seqlen_list: \", balanced_sum_seqlen_list)\n    min_sum_seqlen_balanced = min(balanced_sum_seqlen_list)\n    max_sum_seqlen_balanced = max(balanced_sum_seqlen_list)\n\n    return {\n        f\"{prefix}/min\": min_sum_seqlen,\n        f\"{prefix}/max\": max_sum_seqlen,\n        f\"{prefix}/minmax_diff\": max_sum_seqlen - min_sum_seqlen,\n        f\"{prefix}/balanced_min\": min_sum_seqlen_balanced,\n        f\"{prefix}/balanced_max\": max_sum_seqlen_balanced,\n        f\"{prefix}/mean\": total_sum_seqlen / len(partitions),\n    }\n\n\ndef ceildiv(a: int, b: int) -> int:\n    \"\"\"Compute ceiling division of a by b.\n\n    Returns the smallest integer greater than or equal to a/b.\n    Uses the identity: ceil(a/b) = floor((a + b - 1) / b) = -(-a // b)\n\n    Args:\n        a: Dividend (numerator).\n        b: Divisor (denominator), must be non-zero.\n\n    Returns:\n        int: Ceiling of a divided by b.\n\n    Example:\n        >>> ceildiv(7, 3)  # ceil(7/3) = ceil(2.33) = 3\n        3\n        >>> ceildiv(6, 3)  # ceil(6/3) = ceil(2.0) = 2\n        2\n    \"\"\"\n    return -(a // -b)\n\n\ndef roundup_divisible(a: int, b: int) -> int:\n    \"\"\"Round up a to the nearest multiple of b.\n\n    Returns the smallest multiple of b that is >= a.\n\n    Args:\n        a: Value to round up.\n        b: Divisor to round to (must be positive).\n\n    Returns:\n        int: Smallest multiple of b that is >= a.\n\n    Example:\n        >>> roundup_divisible(7, 4)  # nearest multiple of 4 >= 7 is 8\n        8\n        >>> roundup_divisible(8, 4)  # 8 is already a multiple of 4\n        8\n    \"\"\"\n    return ((a + b - 1) // b) * b\n\n\ndef rearrange_micro_batches(\n    batch,\n    max_token_len,\n    dp_group=None,\n    num_batches_divided_by=None,\n    same_micro_num_in_dp=True,\n    min_num_micro_batch=None,\n    use_dynamic_bsz_balance=True,\n    force_group_size=1,\n):\n    \"\"\"\n    Split a batch into micro-batches by total token count, with optional DP sync and padding.\n\n    Args:\n        batch (TensorDict): must include \"attention_mask\" (B*S); other fields are sliced similarly.\n        max_token_len (int): max sum of attention_mask per micro-batch.\n        dp_group (optional): torch.distributed group for data-parallel sync.\n        num_batches_divided_by (optional): virtual pipeline parallel size, for megatron.\n        same_micro_num_in_dp (bool): if True and dp_group set, pad all ranks to the same count.\n        min_num_micro_batch (int, optional): force at least this many splits (pads empty ones).\n        use_dynamic_bsz_balance (bool, optional): balance the computational workload between micro-batches\n        force_group_size (int, optional): force consecutive samples to be in the same micro-batch (for RM training).\n\n    Returns:\n        List[TensorDict]: the micro-batches.\n        List[List[int]]: index lists mapping each micro-batch back to original positions.\n    \"\"\"\n    # this is per local micro_bsz\n    input_ids = batch[\"input_ids\"]\n    if input_ids.is_nested:\n        seq_len_effective: torch.Tensor = input_ids.offsets().diff()\n        max_seq_len = max(seq_len_effective)\n    else:\n        max_seq_len = batch[\"attention_mask\"].shape[-1]\n        seq_len_effective: torch.Tensor = batch[\"attention_mask\"].sum(dim=1)\n\n    assert max_token_len >= max_seq_len, (\n        f\"max_token_len must be greater than the sequence length. Got {max_token_len=} and {max_seq_len=}\"\n    )\n\n    # Validate force_group_size\n    batch_size = len(seq_len_effective)\n    assert batch_size % force_group_size == 0, (\n        f\"Batch size {batch_size} must be divisible by force_group_size {force_group_size}\"\n    )\n\n    total_seqlen = seq_len_effective.sum().item()\n    # NOTE: num_microbatches <= batch_size, so take the min of this two.\n    # When force_group_size > 1, we work with groups instead of individual samples\n    num_groups = batch_size // force_group_size\n    num_micro_batches = min(num_groups, ceildiv(total_seqlen, max_token_len))\n    if min_num_micro_batch is not None:\n        # used to support pp\n        num_micro_batches = max(min_num_micro_batch, num_micro_batches)\n    if dist.is_initialized() and same_micro_num_in_dp and dp_group is not None:\n        num_micro_batches = torch.tensor([num_micro_batches], device=get_device_name())\n        dist.all_reduce(num_micro_batches, op=dist.ReduceOp.MAX, group=dp_group)\n        num_micro_batches = num_micro_batches.cpu().item()\n    if num_batches_divided_by is not None:\n        num_micro_batches = roundup_divisible(num_micro_batches, num_batches_divided_by)\n\n    assert num_micro_batches <= num_groups\n\n    # upcast to int64 to avoid potential overflow im `calculate_workload` computation.\n    seq_len_effective = seq_len_effective.long()\n\n    # When force_group_size > 1, aggregate workloads by groups\n    if force_group_size > 1:\n        # Calculate workload for each group (sum of workloads of samples in the group)\n        workloads_per_sample = calculate_workload(seq_len_effective)\n        workloads_per_sample_grouped = workloads_per_sample.view(num_groups, force_group_size)\n        group_workloads = workloads_per_sample_grouped.sum(dim=1).cpu().tolist()\n\n        # Partition groups instead of individual samples\n        micro_bsz_group_idx = get_seqlen_balanced_partitions(group_workloads, num_micro_batches, equal_size=False)\n\n        # Convert group indices back to sample indices\n        micro_bsz_idx = []\n        for group_partition in micro_bsz_group_idx:\n            sample_partition = []\n            for group_idx in group_partition:\n                start_idx = group_idx * force_group_size\n                sample_partition.extend(range(start_idx, start_idx + force_group_size))\n            micro_bsz_idx.append(sample_partition)\n\n        workloads = group_workloads\n    else:\n        # Original logic for force_group_size == 1\n        # note that seq_len_effective is a GPU tensor. We need to make it a list to avoid D2H!\n        workloads = calculate_workload(seq_len_effective).cpu().tolist()\n        micro_bsz_idx = get_seqlen_balanced_partitions(workloads, num_micro_batches, equal_size=False)\n\n    if use_dynamic_bsz_balance:\n        # Use the sum of squared sequence lengths to approximate attention computation workload\n        if force_group_size > 1:\n            # For grouped samples, use group workloads for sorting\n            micro_bsz_idx.sort(\n                key=lambda partition: (\n                    sum(workloads[idx // force_group_size] for idx in partition[::force_group_size]),\n                    partition[0] if partition else 0,\n                ),\n                reverse=True,\n            )\n        else:\n            micro_bsz_idx.sort(\n                key=lambda partition: (\n                    sum(workloads[idx] for idx in partition),\n                    partition[0] if partition else 0,\n                ),\n                reverse=True,\n            )\n        # Place smaller micro-batches at both ends to reduce the bubbles exposed during the warm-up and cool-down.\n        micro_bsz_idx = micro_bsz_idx[::2][::-1] + micro_bsz_idx[1::2]\n\n    micro_batches = []\n\n    for partition in micro_bsz_idx:\n        curr_micro_batch = tu.index_select_tensor_dict(batch, partition)\n        micro_batches.append(curr_micro_batch)\n\n    return micro_batches, micro_bsz_idx\n\n\ndef get_reverse_idx(idx_map):\n    \"\"\"\n    Build the inverse of an index mapping.\n\n    Args:\n        idx_map (Sequence[int]): Sequence where idx_map[i] = j.\n\n    Returns:\n        List[int]: Inverse mapping list such that output[j] = i for each i.\n    \"\"\"\n    reverse_idx_map = copy.deepcopy(idx_map)\n\n    for i, idx in enumerate(idx_map):\n        reverse_idx_map[idx] = i\n\n    return reverse_idx_map\n\n\ndef prepare_dynamic_batch(\n    data: DataProto,\n    max_token_len: int,\n    dp_group=None,\n    num_batches_divided_by=None,\n    same_micro_num_in_dp=True,\n    min_num_micro_batch=None,\n    use_dynamic_bsz_balance=True,\n) -> tuple[list[DataProto], list[list[int]]]:\n    \"\"\"\n    Prepare a batch for dynamic batching.\n\n    Args:\n        data (DataProto): The input data.\n        max_token_len (int): The maximum token length for dynamic batching.\n\n    Returns:\n        Tuple[List[DataProto], List[List[int]]]: A tuple containing a list of DataProto objects\n        and a list of index lists.\n    \"\"\"\n    batch, batch_idx_list = rearrange_micro_batches(\n        data.batch,\n        max_token_len=max_token_len,\n        dp_group=dp_group,\n        num_batches_divided_by=num_batches_divided_by,\n        same_micro_num_in_dp=same_micro_num_in_dp,\n        min_num_micro_batch=min_num_micro_batch,\n        use_dynamic_bsz_balance=use_dynamic_bsz_balance,\n    )\n    micro_batches = []\n    for i, batch_idx in enumerate(batch_idx_list):\n        tensors = dict(batch[i])\n        non_tensors = {key: value[batch_idx] for key, value in data.non_tensor_batch.items()}\n        meta_info = copy.deepcopy(data.meta_info)\n        micro_batches.append(DataProto.from_dict(tensors, non_tensors, meta_info=meta_info))\n\n    return micro_batches, batch_idx_list\n\n\ndef restore_dynamic_batch(data: torch.Tensor, batch_idx_list: list[list[int]]) -> torch.Tensor:\n    \"\"\"\n    Restore a batch from dynamic batching.\n\n    Args:\n        data (torch.Tensor): The input data.\n        batch_idx_list (List[List[int]]): The list of index lists.\n\n    Returns:\n        torch.Tensor: The restored data.\n    \"\"\"\n    indices = list(chain.from_iterable(batch_idx_list))\n    batch_size = data.shape[0]\n    assert len(indices) == batch_size, f\"{len(indices)} vs. {batch_size}\"\n    revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long)\n\n    if data.is_nested:\n        data_lst = data.unbind()\n        tensors = [data_lst[i] for i in revert_indices]\n        reverted_data = torch.nested.as_nested_tensor(tensors, layout=torch.jagged)\n    else:\n        reverted_data = data[revert_indices]\n\n    return reverted_data\n\n\ndef get_group_balanced_partitions(\n    seqlen_list: list[int],\n    uid_list: list,\n    k_partitions: int,\n) -> list[list[int]]:\n    \"\"\"\n    Partition samples into k groups while keeping samples with the same uid together.\n\n    Args:\n        seqlen_list: List of sequence lengths for each sample.\n        uid_list: List of uids identifying which samples share the same prefix.\n                  Samples with the same uid will be kept together.\n        k_partitions: Number of partitions (typically world_size).\n\n    Returns:\n        List of k lists, each containing sample indices assigned to that partition.\n        Samples with the same uid are guaranteed to be in the same partition.\n    \"\"\"\n    assert len(seqlen_list) == len(uid_list), \"seqlen_list and uid_list must have same length\"\n\n    # Build groups: each group contains indices of samples with the same uid\n    # Assumes samples with same uid are contiguous\n    groups = []  # List of (group_indices, group_total_seqlen)\n    current_uid = None\n    current_indices = []\n    current_seqlen = 0\n\n    for i, (seqlen, uid) in enumerate(zip(seqlen_list, uid_list, strict=False)):\n        if uid != current_uid:\n            if current_indices:\n                groups.append((current_indices, current_seqlen))\n            current_uid = uid\n            current_indices = [i]\n            current_seqlen = seqlen\n        else:\n            current_indices.append(i)\n            current_seqlen += seqlen\n\n    # Don't forget the last group\n    if current_indices:\n        groups.append((current_indices, current_seqlen))\n\n    num_groups = len(groups)\n    assert num_groups >= k_partitions, (\n        f\"Number of uid groups ({num_groups}) must be >= k_partitions ({k_partitions}). \"\n        f\"Consider reducing world_size or increasing batch_size.\"\n    )\n\n    # Calculate workload for each group (as integers for partitioning)\n    group_workloads = []\n    for indices, total_seqlen in groups:\n        # Use sum of individual workloads for more accurate estimation\n        workload = sum(int(calculate_workload(torch.tensor([seqlen_list[i]])).item()) for i in indices)\n        group_workloads.append(workload)\n\n    # Use Karmarkar-Karp to partition groups\n    # equal_size=True ensures each partition gets the same number of groups,\n    # which is required when each group has the same number of samples (rollout.n)\n    group_partitions = get_seqlen_balanced_partitions(\n        seqlen_list=group_workloads,\n        k_partitions=k_partitions,\n        equal_size=True,\n    )\n\n    # Convert group partitions to sample partitions\n    sample_partitions = []\n    for group_partition in group_partitions:\n        sample_indices = []\n        for group_idx in group_partition:\n            sample_indices.extend(groups[group_idx][0])\n        sample_partitions.append(sorted(sample_indices))\n\n    return sample_partitions\n"
  },
  {
    "path": "verl/utils/sglang/sglang_fp8_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom verl.utils.fp8_utils import FP8QuantizerHelper\n\n\nclass SGLangFP8QuantizerHelper(FP8QuantizerHelper):\n    def __init__(self, quant_config):\n        super().__init__(quant_config)\n"
  },
  {
    "path": "verl/utils/tensordict_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nfrom typing import Any, Iterable\n\nimport torch\nfrom tensordict import TensorDict\nfrom tensordict.tensorclass import NonTensorData, NonTensorStack\n\n\ndef assign_non_tensor_data(tensor_dict: TensorDict, key, val):\n    \"\"\"Assign a single non-tensor value to a TensorDict.\n\n    Wraps the value in NonTensorData so it can be stored alongside tensors\n    in the TensorDict. Use this for scalar metadata or simple non-tensor values.\n\n    Args:\n        tensor_dict: The TensorDict to assign to.\n        key: The key under which to store the value.\n        val: Any non-tensor value to store (e.g., string, int, dict).\n\n    Raises:\n        AssertionError: If tensor_dict is not a TensorDict.\n\n    Example:\n        >>> td = TensorDict({\"obs\": torch.randn(3, 4)}, batch_size=[3])\n        >>> assign_non_tensor_data(td, \"experiment_name\", \"run_001\")\n    \"\"\"\n    assert isinstance(tensor_dict, TensorDict), \"input dict must be a TensorDict\"\n    tensor_dict[key] = NonTensorData(val)\n\n\ndef assign_non_tensor_stack(tensor_dict: TensorDict, key, val: list):\n    \"\"\"Assign a list with potentially nested structures (lists, dicts, etc.) to TensorDict.\n\n    This function handles complex nested data structures like:\n    - Lists of lists: [[], [0.5, 0.8], [0.9]]\n    - Lists of dicts: [{\"acc\": 1.0}, {\"acc\": 0.0}]\n    - Lists of lists of dicts: [[{\"content\": \"...\", \"role\": \"user\"}]]\n\n    These structures are wrapped in NonTensorStack so TensorDict can handle them correctly.\n\n    Args:\n        tensor_dict: The TensorDict to assign to\n        key: The key to assign the value under\n        val: A list containing potentially nested structures\n\n    Example:\n        >>> td = TensorDict({}, batch_size=[])\n        >>> turn_scores = [[], [0.5, 0.8], [0.9]]\n        >>> assign_non_tensor_stack(td, \"turn_scores\", turn_scores)\n        >>> # Now td[\"turn_scores\"] contains the nested data\n    \"\"\"\n    # Convert list to NonTensorStack to handle nested structures\n    # This wraps each item in NonTensorData to preserve complex objects\n    # TODO(petersh6): can convert back to val directly if we are not accessing .data from the NonTensorStack\n    assert isinstance(tensor_dict, TensorDict), \"input dict must be a TensorDict\"\n    tensor_dict[key] = NonTensorStack.from_list([NonTensorData(item) for item in val])\n\n\ndef assign_non_tensor(tensor_dict: TensorDict, **kwargs):\n    \"\"\"Assign non-tensor data to a TensorDict.\n\n    Automatically detects if the value is a list with nested structures and uses\n    the appropriate assignment method (NonTensorData for simple values,\n    NonTensorStack for lists with nested structures).\n\n    Args:\n        tensor_dict: The TensorDict to assign to\n        **kwargs: Key-value pairs where values can be:\n            - Simple values (stored as NonTensorData)\n            - Lists with nested structures (stored as NonTensorStack)\n\n    Example:\n        >>> td = TensorDict({\"obs\": torch.randn(3, 4)}, batch_size=[3])\n        >>> assign_non_tensor(\n        ...     tensor_dict=td,\n        ...     metadata=\"experiment_1\",  # Simple value\n        ...     turn_scores=[[], [0.5, 0.8], [0.9]]  # Nested list\n        ... )\n    \"\"\"\n    assert isinstance(tensor_dict, TensorDict), \"input dict must be a TensorDict\"\n    for key, val in kwargs.items():\n        if isinstance(val, (NonTensorData | NonTensorStack)):\n            tensor_dict[key] = val\n        elif isinstance(val, list):\n            # For lists, use NonTensorStack\n            assign_non_tensor_stack(tensor_dict=tensor_dict, key=key, val=val)\n        else:\n            # For non-list values, use NonTensorData\n            assign_non_tensor_data(tensor_dict=tensor_dict, key=key, val=val)\n    return tensor_dict\n\n\ndef unwrap_non_tensor_data(data):\n    \"\"\"Unwrap a NonTensorData object to get the underlying value.\n\n    If the input is a NonTensorData wrapper, extracts and returns the\n    underlying data. Otherwise, returns the input unchanged.\n\n    Args:\n        data: Either a NonTensorData object or any other value.\n\n    Returns:\n        The unwrapped data if input was NonTensorData, otherwise the\n        original input unchanged.\n\n    Example:\n        >>> wrapped = NonTensorData(\"hello\")\n        >>> unwrap_non_tensor_data(wrapped)\n        'hello'\n        >>> unwrap_non_tensor_data(42)  # Non-wrapped value\n        42\n    \"\"\"\n    if isinstance(data, NonTensorData):\n        return data.data\n    return data\n\n\ndef get_non_tensor_data(data: TensorDict, key: str, default):\n    \"\"\"Retrieve and unwrap non-tensor data from a TensorDict.\n\n    Fetches the value for the given key from the TensorDict and automatically\n    unwraps it if it's stored as NonTensorData.\n\n    Args:\n        data: The TensorDict to retrieve from.\n        key: The key to look up.\n        default: Value to return if the key is not found.\n\n    Returns:\n        The unwrapped value if the key exists and was wrapped in NonTensorData,\n        the raw value if it wasn't wrapped, or the default if key not found.\n\n    Example:\n        >>> td = TensorDict({}, batch_size=[])\n        >>> assign_non_tensor_data(td, \"config\", {\"lr\": 0.01})\n        >>> get_non_tensor_data(td, \"config\", None)\n        {'lr': 0.01}\n        >>> get_non_tensor_data(td, \"missing\", \"default_value\")\n        'default_value'\n    \"\"\"\n    output = data.get(key, default)\n    return unwrap_non_tensor_data(output)\n\n\ndef concat_nested_tensors(tensors: list[torch.Tensor]) -> torch.Tensor:\n    \"\"\"Concatenate multiple nested tensors along the batch dimension.\n\n    Takes a list of nested tensors with jagged layout and concatenates them\n    into a single nested tensor. Each input tensor must have 2 or more dimensions and be contiguous.\n\n    Args:\n        tensors: List of nested tensors to concatenate. All tensors must\n            be nested, contiguous, and have 2 or more dimensions.\n\n    Returns:\n        A new nested tensor with jagged layout containing all rows from\n        the input tensors concatenated along dimension 0.\n\n    Raises:\n        AssertionError: If any tensor is not nested, not contiguous, or\n            doesn't have 2 or more dimensions.\n\n    Example:\n        >>> t1 = torch.nested.as_nested_tensor([torch.randn(3), torch.randn(5)], layout=torch.jagged)\n        >>> t2 = torch.nested.as_nested_tensor([torch.randn(2), torch.randn(4)], layout=torch.jagged)\n        >>> result = concat_nested_tensors([t1, t2])\n        >>> # result contains 4 rows: lengths [3, 5, 2, 4]\n    \"\"\"\n    for tensor in tensors:\n        assert tensor.is_nested and tensor.is_contiguous()\n    unbind_tensors = []\n    for tensor in tensors:\n        assert len(tensor.shape) >= 2, f\"nested tensor must have 2 or more dimensions. Got {tensor.shape}\"\n        unbind_tensor = tensor.unbind(0)\n        unbind_tensors.extend(list(unbind_tensor))\n\n    tensor = torch.nested.as_nested_tensor(unbind_tensors, layout=torch.jagged)\n    return tensor\n\n\ndef concat_tensordict_with_none_bsz(data: list[TensorDict]):\n    \"\"\"Handle concatenation of TensorDicts with empty batch size.\n\n    For TensorDicts that contain only metadata (NonTensorData) with no batch\n    dimension, returns the first TensorDict as the concatenation result.\n\n    Args:\n        data: List of TensorDicts, each with empty batch_size (batch_size=[]).\n\n    Returns:\n        The first TensorDict from the list, as metadata concatenation\n        simply preserves the first instance.\n\n    Raises:\n        AssertionError: If any TensorDict has a non-empty batch_size.\n\n    Note:\n        This is used internally by concat_tensordict when handling\n        TensorDicts that contain only non-tensor metadata.\n    \"\"\"\n    for d in data:\n        assert len(d.batch_size) == 0\n    # directly return the first meta info\n    return data[0]\n\n\ndef concat_tensordict(data: list[TensorDict]) -> TensorDict:\n    \"\"\"Concatenate multiple TensorDicts along dimension zero.\n\n    Combines a list of TensorDicts into a single TensorDict by concatenating\n    all tensors along the batch dimension (dim=0). Handles nested tensors\n    specially by unbinding and rebinding them.\n\n    Args:\n        data: List of TensorDicts to concatenate. All TensorDicts must have\n            the same keys and the same set of nested tensor keys.\n\n    Returns:\n        A new TensorDict containing concatenated tensors from all inputs.\n\n    Raises:\n        AssertionError: If data is empty or if TensorDicts have inconsistent\n            nested tensor keys.\n\n    Note:\n        - For TensorDicts with empty batch_size, returns the first one\n        - Nested tensors are handled specially via concat_nested_tensors\n        - Regular tensors use TensorDict.cat for efficient concatenation\n    \"\"\"\n    assert len(data) > 0, \"Must have at least one tensordict\"\n\n    # Find nested tensor keys from the first tensordict\n    nested_tensor_keys = {key for key, value in data[0].items() if isinstance(value, torch.Tensor) and value.is_nested}\n\n    if not nested_tensor_keys:\n        if len(data[0].batch_size) == 0:\n            return concat_tensordict_with_none_bsz(data)\n        # if batch size is None (only contain NonTensorData)\n        return TensorDict.cat(data, dim=0)\n\n    # Create a list of tensordicts containing only non-nested tensors for concatenation\n    regular_tds = []\n    for td in data:\n        current_nested_keys = {k for k, v in td.items() if isinstance(v, torch.Tensor) and v.is_nested}\n        assert current_nested_keys == nested_tensor_keys, \"All tensordicts must have the same set of nested tensors.\"\n\n        # Create a new TensorDict with non-nested items without modifying the original\n        regular_items = {k: v for k, v in td.items() if k not in nested_tensor_keys}\n        regular_tds.append(TensorDict(regular_items, batch_size=td.batch_size, device=td.device))\n\n    # Concatenate the regular tensordicts\n    output = TensorDict.cat(regular_tds, dim=0)\n\n    # Concatenate and add nested tensors to the output\n    for key in nested_tensor_keys:\n        nested_tensors_to_concat = [td[key] for td in data]\n        output[key] = concat_nested_tensors(nested_tensors_to_concat)\n\n    return output\n\n\ndef chunk_tensordict(td: TensorDict, chunks: int) -> list[TensorDict]:\n    \"\"\"Split a TensorDict into equal-sized chunks with special nested tensor handling.\n\n    Divides a TensorDict into the specified number of chunks along the batch\n    dimension. Handles NestedTensors specially since TensorDict.chunk() doesn't\n    support jagged tensors.\n\n    Args:\n        td: The TensorDict to split.\n        chunks: Number of chunks to create. Must evenly divide len(td).\n\n    Returns:\n        List of TensorDicts, each containing a portion of the original data.\n\n    Raises:\n        AssertionError: If td is not a TensorDict or if its length is not\n            evenly divisible by chunks.\n\n    Note:\n        PyTorch ``unbind(dim=0)`` on 3D+ jagged NestedTensors has a bug where\n        ``split_with_sizes`` is applied to the wrong dimension of the internal\n        ``_values`` tensor.  For example, mRoPE ``position_ids`` with per-sample\n        shape ``(4, seq_len)`` becomes a 3D jagged NestedTensor\n        ``[B, *(ragged=4), seq_len]``; ``_values`` is ``[B*4, seq_len]`` and\n        ``unbind`` erroneously splits dimension 1 (``seq_len``) instead of\n        dimension 0, causing::\n\n            RuntimeError: split_with_sizes expects split_sizes to sum exactly\n            to <seq_len>, but got split_sizes=[4, 4, ...]\n\n        2D jagged NestedTensors (e.g. ``input_ids``, ``loss_mask``) are\n        unaffected — ``unbind(dim=0)`` works correctly for them.\n\n        The workaround: try ``unbind`` first (fast path for 2D); on failure,\n        fall back to ``to_padded_tensor`` → ``chunk`` → reconstruct per-chunk\n        NestedTensors using the original ragged lengths from ``offsets``.\n\n        See https://github.com/pytorch/pytorch/issues/153238\n    \"\"\"\n    assert isinstance(td, TensorDict) and len(td) % chunks == 0, (\n        f\"expecting td with length divisible by chunks, but got {len(td)} and {chunks}\"\n    )\n    chunk_size = len(td) // chunks\n    nested_keys = {key for key, val in td.items() if isinstance(val, torch.Tensor) and val.is_nested}\n    new_td = TensorDict(\n        {k: v for k, v in td.items() if k not in nested_keys}, batch_size=td.batch_size, device=td.device\n    )\n\n    tds = new_td.chunk(chunks=chunks)\n    for key in nested_keys:\n        nt = td[key]\n        try:\n            tensors = nt.unbind(dim=0)\n        except RuntimeError:\n            padded = nt.to_padded_tensor(0)\n            padded_chunks = padded.chunk(chunks, dim=0)\n            offsets = nt.offsets()\n            lengths = offsets.diff().tolist()\n            for i, chunk_td in enumerate(tds):\n                chunk_lengths = lengths[i * chunk_size : (i + 1) * chunk_size]\n                chunk_tensors = [padded_chunks[i][j, :seq_len] for j, seq_len in enumerate(chunk_lengths)]\n                chunk_td[key] = torch.nested.as_nested_tensor(chunk_tensors, layout=torch.jagged)\n            continue\n\n        for i, chunk_td in enumerate(tds):\n            chunk_td[key] = torch.nested.as_nested_tensor(\n                tensors[i * chunk_size : (i + 1) * chunk_size], layout=torch.jagged\n            )\n\n    return tds\n\n\ndef get_tensordict(tensor_dict: dict[str, torch.Tensor | list], non_tensor_dict: dict = None) -> TensorDict:\n    \"\"\"Create a TensorDict from tensors and non-tensor data.\n\n    Automatically handles nested structures in lists by converting them to NonTensorStack.\n    This enables support for:\n    - Lists of lists: [[], [0.5, 0.8], [0.9]]\n    - Lists of dicts: [{\"acc\": 1.0}, {\"acc\": 0.0}]\n    - Lists of lists of dicts: [[{\"content\": \"...\", \"role\": \"user\"}]]\n\n    Args:\n        tensor_dict: Dictionary of tensors and lists to include in the TensorDict\n        non_tensor_dict: Dictionary of metadata to store as NonTensorData\n\n    Returns:\n        TensorDict with proper handling of nested structures\n\n    Example:\n        >>> td = get_tensordict(\n        ...     tensor_dict={\n        ...         \"obs\": torch.randn(3, 4),\n        ...         \"turn_scores\": [[], [0.5, 0.8], [0.9]]  # Nested list\n        ...     },\n        ...     non_tensor_dict={\"experiment\": \"test\"}\n        ... )\n    \"\"\"\n    tensor_dict = tensor_dict.copy()\n    if non_tensor_dict is None:\n        non_tensor_dict = {}\n\n    batch_size = None\n\n    for key, val in tensor_dict.items():\n        if isinstance(val, torch.Tensor) and val.is_nested:\n            assert val.is_contiguous(), \"Nested tensors must be contiguous. Try setting layout=torch.jagged\"\n            assert val.layout == torch.jagged, \"Nested tensors must be jagged.\"\n\n        # Skip validation for NonTensorStack as it's already properly formatted\n        if isinstance(val, NonTensorStack):\n            if batch_size is None:\n                batch_size = len(val)\n            else:\n                assert len(val) == batch_size, (\n                    f\"Batch size of NonTensorStack {key} is not consistent with other tensors. \"\n                    f\"Expected {batch_size}, got {len(val)}\"\n                )\n            continue\n\n        if isinstance(val, list):\n            for v in val:\n                assert not isinstance(v, torch.Tensor), (\n                    \"Passing a list makes the data NonTensorStack, \"\n                    \"which doesn't support torch.Tensor. Please convert to numpy first\"\n                )\n            # Convert to NonTensorStack to handle nested structures\n            tensor_dict[key] = NonTensorStack.from_list([NonTensorData(item) for item in val])\n\n        assert isinstance(val, torch.Tensor | list)\n\n        if batch_size is None:\n            batch_size = val.size(0) if isinstance(val, torch.Tensor) else len(val)\n        else:\n            val_batch_size = val.size(0) if isinstance(val, torch.Tensor) else len(val)\n            assert val_batch_size == batch_size, (\n                f\"Batch size of tensor {key} is not consistent with other tensors. \"\n                f\"Expected {batch_size}, got {val_batch_size}\"\n            )\n\n    if batch_size is None:\n        batch_size = []\n    else:\n        batch_size = [batch_size]\n\n    for key, val in non_tensor_dict.items():\n        assert key not in tensor_dict\n        tensor_dict[key] = NonTensorData(val)\n\n    return TensorDict(source=tensor_dict, batch_size=batch_size)\n\n\ndef index_select_tensor_dict(batch: TensorDict, indices: torch.Tensor | list[int]) -> TensorDict:\n    \"\"\"Select rows from a TensorDict using indices.\n\n    Creates a new TensorDict containing only the rows specified by indices.\n    Handles regular tensors, nested tensors, NonTensorStack, and NonTensorData\n    appropriately.\n\n    Args:\n        batch: The TensorDict to index into. Can be None.\n        indices: 1D tensor or list of integers specifying which rows to select.\n\n    Returns:\n        A new TensorDict containing only the selected rows, or None if\n        batch was None.\n\n    Raises:\n        AssertionError: If indices is not 1-dimensional.\n\n    Note:\n        - Regular tensors are indexed directly\n        - Nested tensors are unbound, indexed, and rebound\n        - NonTensorStack is indexed by batch dimension\n        - NonTensorData (scalar metadata) is preserved unchanged\n    \"\"\"\n    if isinstance(indices, list):\n        indices = torch.tensor(indices)\n\n    assert indices.dim() == 1, \"indices must be a 1D tensor\"\n\n    data_dict = {}\n    batch_size = indices.shape[0]\n\n    if batch is not None:\n        for key, tensor in batch.items():\n            if isinstance(tensor, torch.Tensor) and not tensor.is_nested:\n                data_dict[key] = tensor[indices]\n            elif isinstance(tensor, torch.Tensor) and tensor.is_nested:\n                tensor_lst = tensor.unbind()  # for performance\n                data_dict[key] = torch.nested.as_nested_tensor(\n                    [tensor_lst[idx] for idx in indices], layout=torch.jagged\n                )\n            else:\n                # This handles NonTensorStack (indexable by batch dim) and NonTensorData (scalar metadata).\n                if tensor.shape:\n                    data_dict[key] = tensor[indices]\n                else:\n                    data_dict[key] = tensor\n        selected_batch = TensorDict(source=data_dict, batch_size=batch_size)\n    else:\n        selected_batch = None\n\n    return selected_batch\n\n\ndef union_tensor_dict(tensor_dict1: TensorDict, tensor_dict2: TensorDict) -> TensorDict:\n    \"\"\"Merge two TensorDicts, adding keys from the second to the first.\n\n    Performs an in-place union of two TensorDicts. Keys from tensor_dict2\n    that don't exist in tensor_dict1 are added. Keys that exist in both\n    must have identical values.\n\n    Args:\n        tensor_dict1: The base TensorDict to merge into (modified in-place).\n        tensor_dict2: The TensorDict whose keys will be added to tensor_dict1.\n\n    Returns:\n        The modified tensor_dict1 containing the union of both TensorDicts.\n\n    Raises:\n        AssertionError: If batch sizes don't match, or if a key exists in\n            both TensorDicts with different values.\n\n    Example:\n        >>> td1 = TensorDict({\"a\": torch.tensor([1, 2])}, batch_size=[2])\n        >>> td2 = TensorDict({\"b\": torch.tensor([3, 4])}, batch_size=[2])\n        >>> result = union_tensor_dict(td1, td2)\n        >>> list(result.keys())\n        ['a', 'b']\n    \"\"\"\n    assert tensor_dict1.batch_size == tensor_dict2.batch_size, (\n        f\"Two tensor dict must have identical batch size. Got {tensor_dict1.batch_size} and {tensor_dict2.batch_size}\"\n    )\n    for key in tensor_dict2.keys():\n        if key not in tensor_dict1.keys():\n            # Note that there is a difference between tensor_dict2[key] and tensor_dict2.get(key)\n            tensor_dict1[key] = tensor_dict2.get(key)\n        else:\n            if isinstance(tensor_dict2[key], torch.Tensor):\n                assert tensor_dict1[key].equal(tensor_dict2[key]), (\n                    f\"{key} in tensor_dict1 and tensor_dict2 are not the same object\"\n                )\n            else:\n                # non-tensor\n                assert tensor_dict1[key] == tensor_dict2[key], (\n                    f\"{key} in tensor_dict1 and tensor_dict2 are not the same object\"\n                )\n\n    return tensor_dict1\n\n\ndef make_iterator(tensordict: TensorDict, mini_batch_size, epochs, seed=None, dataloader_kwargs=None):\n    \"\"\"Create an iterator that yields mini-batches from a TensorDict.\n\n    Wraps a TensorDict in a DataLoader-style iterator that yields mini-batches\n    for the specified number of epochs. Useful for training loops.\n\n    Args:\n        tensordict: The TensorDict to iterate over.\n        mini_batch_size: Size of each mini-batch. Must evenly divide the\n            TensorDict's batch size.\n        epochs: Number of times to iterate through the entire dataset.\n        seed: Optional random seed for reproducible shuffling.\n        dataloader_kwargs: Optional dict of additional kwargs to pass to\n            the underlying DataLoader (e.g., shuffle=True, num_workers=4).\n\n    Returns:\n        An iterator that yields TensorDict mini-batches.\n\n    Raises:\n        AssertionError: If batch size is not divisible by mini_batch_size.\n\n    Example:\n        >>> td = TensorDict({\"obs\": torch.randn(100, 4)}, batch_size=[100])\n        >>> for batch in make_iterator(td, mini_batch_size=10, epochs=2):\n        ...     # batch is a TensorDict with batch_size=[10]\n        ...     pass\n    \"\"\"\n    from torch.utils.data import DataLoader\n\n    assert tensordict.batch_size[0] % mini_batch_size == 0, f\"{tensordict.batch_size[0]} % {mini_batch_size} != 0\"\n    # we can directly create a dataloader from TensorDict\n    if dataloader_kwargs is None:\n        dataloader_kwargs = {}\n\n    if seed is not None:\n        generator = torch.Generator()\n        generator.manual_seed(seed)\n    else:\n        generator = None\n\n    assert isinstance(dataloader_kwargs, dict)\n\n    idx_lst = torch.arange(tensordict.shape[0])\n\n    train_dataloader = DataLoader(\n        dataset=idx_lst, batch_size=mini_batch_size, collate_fn=lambda x: x, generator=generator, **dataloader_kwargs\n    )\n\n    def get_data():\n        for _ in range(epochs):\n            for idx in train_dataloader:\n                yield index_select_tensor_dict(tensordict, idx)\n\n    return iter(get_data())\n\n\ndef assert_tensordict_eq(tensordict1: TensorDict, tensordict2: TensorDict):\n    \"\"\"Assert that two TensorDicts are equal.\n\n    Performs a deep equality check between two TensorDicts, verifying that\n    they have the same keys with identical values. Handles nested tensors\n    by comparing their unbound components.\n\n    Args:\n        tensordict1: First TensorDict to compare.\n        tensordict2: Second TensorDict to compare.\n\n    Raises:\n        AssertionError: If the TensorDicts differ in keys, value types, or\n            value contents. The error message indicates what differs.\n\n    Note:\n        - Regular tensors are compared element-wise\n        - Nested tensors are unbound and compared component by component\n        - Non-tensor values are compared with standard equality\n    \"\"\"\n    tensordict1_key_set = set(tensordict1.keys())\n    tensordict2_key_set = set(tensordict2.keys())\n    assert tensordict1_key_set == tensordict2_key_set, (\n        f\"key set diffs. Got {tensordict2_key_set=} vs {tensordict1_key_set=}\"\n    )\n\n    for key in tensordict1.keys():\n        val = tensordict1[key]\n        val2 = tensordict2[key]\n\n        assert type(val) is type(val2), f\"The type of {key} must be the same. Got {type(val)} vs {type(val2)}\"\n\n        if isinstance(val, torch.Tensor):\n            if val.is_nested:\n                assert val.is_nested and val2.is_nested, (\n                    f\"Both tensors must be nested tensors. {val.is_nested=}, {val2.is_nested=}\"\n                )\n                t1, t2 = val.unbind(), val2.unbind()\n                assert len(t1) == len(t2), f\"Nested tensor should have the same lengths. {len(t1)=} vs {len(t2)=}\"\n                for c1, c2 in zip(t1, t2, strict=True):\n                    assert torch.equal(c1, c2), f\"Nested tensor components have different values. {c1=} vs {c2=}\"\n            else:\n                assert torch.all(torch.eq(val, val2)).item()\n        else:\n            assert val == val2\n\n\ndef get(tensordict: TensorDict, key: str, default=None) -> Any:\n    \"\"\"Get a value from a TensorDict with automatic unwrapping.\n\n    Retrieves a value from the TensorDict and automatically converts it\n    to a Python-native format:\n    - Tensors are returned as-is\n    - NonTensorStack is converted to a Python list\n    - NonTensorData is unwrapped to its underlying value\n\n    Args:\n        tensordict: The TensorDict to retrieve from.\n        key: The key to look up.\n        default: Value to return if the key doesn't exist. Defaults to None.\n\n    Returns:\n        The value for the key in its native format, or default if not found.\n\n    Example:\n        >>> td = get_tensordict({\"obs\": torch.randn(3, 4), \"labels\": [\"a\", \"b\", \"c\"]})\n        >>> get(td, \"obs\")  # Returns torch.Tensor\n        >>> get(td, \"labels\")  # Returns [\"a\", \"b\", \"c\"] as a list\n        >>> get(td, \"missing\", \"default\")  # Returns \"default\"\n    \"\"\"\n    if key not in tensordict:\n        return default\n\n    output = tensordict.get(key)\n    if isinstance(output, torch.Tensor):\n        return output\n    elif isinstance(output, NonTensorStack):\n        return output.tolist()\n    else:\n        assert isinstance(output, NonTensorData)\n        return output.data\n\n\ndef get_keys(tensordict: TensorDict, keys: Iterable[str]) -> TensorDict:\n    \"\"\"Extract a subset of keys from a TensorDict into a new TensorDict.\n\n    Creates a new TensorDict containing only the specified keys. Values\n    are properly categorized as tensor or non-tensor data.\n\n    Args:\n        tensordict: The source TensorDict.\n        keys: Iterable of key names to extract.\n\n    Returns:\n        A new TensorDict containing only the specified keys with their values.\n\n    Raises:\n        KeyError: If any key in keys doesn't exist in the tensordict.\n\n    Example:\n        >>> td = get_tensordict({\"a\": torch.randn(3), \"b\": torch.randn(3), \"c\": torch.randn(3)})\n        >>> subset = get_keys(td, [\"a\", \"c\"])\n        >>> list(subset.keys())\n        ['a', 'c']\n    \"\"\"\n    tensor_output = {}\n    non_tensor_output = {}\n    for key in keys:\n        if key not in tensordict.keys():\n            raise KeyError(f\"key {key} not in tensordict\")\n        output = tensordict.get(key)\n        if isinstance(output, torch.Tensor):\n            tensor_output[key] = output\n        elif isinstance(output, NonTensorStack):\n            tensor_output[key] = output.tolist()\n        else:\n            assert isinstance(output, NonTensorData)\n            non_tensor_output[key] = output.data\n\n    return get_tensordict(tensor_output, non_tensor_output)\n\n\ndef pop(tensordict: TensorDict, key: str, default=None) -> Any:\n    \"\"\"Remove and return a value from a TensorDict with automatic unwrapping.\n\n    Removes the specified key from the TensorDict and returns its value,\n    automatically converting to Python-native format (same as get()).\n\n    Args:\n        tensordict: The TensorDict to pop from.\n        key: The key to remove and return.\n        default: Value to return if the key doesn't exist. Defaults to None.\n\n    Returns:\n        The value for the key in its native format, or default if not found.\n        The key is removed from the TensorDict.\n\n    Example:\n        >>> td = get_tensordict({\"obs\": torch.randn(3, 4), \"labels\": [\"a\", \"b\", \"c\"]})\n        >>> labels = pop(td, \"labels\")  # Returns [\"a\", \"b\", \"c\"], removes from td\n        >>> \"labels\" in td.keys()\n        False\n    \"\"\"\n    _sentinel = object()\n    output = tensordict.pop(key, _sentinel)\n    if output is _sentinel:\n        return default\n\n    if isinstance(output, torch.Tensor):\n        return output\n    elif isinstance(output, NonTensorStack):\n        return output.tolist()\n    else:\n        assert isinstance(output, NonTensorData)\n        return output.data\n\n\ndef pop_keys(tensordict: TensorDict, keys: Iterable[str]) -> TensorDict:\n    \"\"\"Remove multiple keys from a TensorDict and return them as a new TensorDict.\n\n    Removes the specified keys from the source TensorDict and creates a new\n    TensorDict containing those keys and their values.\n\n    Args:\n        tensordict: The source TensorDict to pop from (modified in-place).\n        keys: Iterable of key names to remove and return.\n\n    Returns:\n        A new TensorDict containing the popped keys and their values.\n\n    Raises:\n        KeyError: If any key in keys doesn't exist in the tensordict.\n\n    Example:\n        >>> td = get_tensordict({\"a\": torch.randn(3), \"b\": torch.randn(3), \"c\": torch.randn(3)})\n        >>> popped = pop_keys(td, [\"a\", \"c\"])\n        >>> list(td.keys())  # Only 'b' remains\n        ['b']\n        >>> list(popped.keys())\n        ['a', 'c']\n    \"\"\"\n    tensor_output = {}\n    non_tensor_output = {}\n    for key in keys:\n        if key not in tensordict.keys():\n            raise KeyError(f\"key {key} not in tensordict\")\n        output = tensordict.get(key)\n        if isinstance(output, torch.Tensor):\n            tensor_output[key] = tensordict.pop(key)\n        elif isinstance(output, NonTensorStack):\n            tensor_output[key] = tensordict.pop(key).tolist()\n        else:\n            assert isinstance(output, NonTensorData)\n            non_tensor_output[key] = tensordict.pop(key)\n\n    return get_tensordict(tensor_output, non_tensor_output)\n\n\ndef pad_to_divisor(data: TensorDict, size_divisor: int):\n    \"\"\"Pad a TensorDict's batch dimension to be divisible by a given divisor.\n\n    If the TensorDict's length is not evenly divisible by size_divisor,\n    pads the batch dimension by repeating elements from the beginning.\n    Useful for ensuring even distribution across workers in distributed training.\n\n    Args:\n        data: The TensorDict to pad.\n        size_divisor: The divisor that the padded length must be divisible by.\n\n    Returns:\n        tuple: A tuple containing:\n            - data (TensorDict): The padded TensorDict (or original if no padding needed)\n            - pad_size (int): Number of elements added as padding (0 if none)\n\n    Raises:\n        AssertionError: If data is not a TensorDict.\n\n    Example:\n        >>> td = TensorDict({\"obs\": torch.randn(10, 4)}, batch_size=[10])\n        >>> padded, pad_size = pad_to_divisor(td, 4)\n        >>> len(padded)  # 12 (next multiple of 4 after 10)\n        12\n        >>> pad_size\n        2\n    \"\"\"\n    assert isinstance(data, TensorDict), \"data must be a TensorDict\"\n    if len(data) % size_divisor != 0:\n        pad_size = size_divisor - len(data) % size_divisor\n        padding_protos = []\n        remaining_pad = pad_size\n        while remaining_pad > 0:\n            take_size = min(remaining_pad, len(data))\n            padding_protos.append(data[:take_size])\n            remaining_pad -= take_size\n        data_padded = torch.cat([data] + padding_protos)\n    else:\n        if len(data) == 0:\n            logging.warning(\"padding a DataProto with no item, no changed made\")\n        pad_size = 0\n        data_padded = data\n    return data_padded, pad_size\n\n\ndef unpad(data: TensorDict, pad_size):\n    \"\"\"Remove padding from a TensorDict.\n\n    Reverses the effect of pad_to_divisor by removing the specified number\n    of elements from the end of the TensorDict.\n\n    Args:\n        data: The padded TensorDict.\n        pad_size: Number of padding elements to remove. If 0, returns\n            data unchanged.\n\n    Returns:\n        The TensorDict with padding removed, equivalent to data[:-pad_size].\n\n    Example:\n        >>> td = TensorDict({\"obs\": torch.randn(12, 4)}, batch_size=[12])\n        >>> unpadded = unpad(td, pad_size=2)\n        >>> len(unpadded)\n        10\n    \"\"\"\n    if pad_size != 0:\n        data = data[:-pad_size]\n    return data\n\n\ndef contiguous(data: TensorDict) -> TensorDict:\n    \"\"\"Call contiguous on a tensor dict. The contiguous function of tensordict lib will make NonTensorStack.\n    This function will always return a new tensordict\n\n    Args:\n        data: The input tensordict\n\n    Returns:\n        a tensordict that is contiguous\n\n    \"\"\"\n    tensor_dict = {}\n    non_tensor_dict = {}\n\n    for key in data.keys():\n        val = data.get(key)\n        if isinstance(val, NonTensorData):\n            non_tensor_dict[key] = val\n        elif isinstance(val, NonTensorStack):\n            tensor_dict[key] = val\n        else:\n            assert isinstance(val, torch.Tensor), f\"Expect val to be a torch.Tensor. Got {type(val)}\"\n            tensor_dict[key] = val.contiguous()\n\n    return get_tensordict(tensor_dict=tensor_dict, non_tensor_dict=non_tensor_dict)\n\n\ndef maybe_fix_3d_position_ids(data: TensorDict):\n    # note for tensordict with pickle/unpickle. nested tensor in tensordict after consolidate and pickle/unpickle\n    # will incur indexing error for ragged tensor. This only happens when using 3D position ids in VLMs.\n    # This is likely a bug in tensordict. As a workaround, we manually set _ragged_index.\n    if \"position_ids\" in data.keys() and data[\"position_ids\"].dim() == 3 and data[\"position_ids\"].is_nested:\n        data[\"position_ids\"]._ragged_idx = 2\n"
  },
  {
    "path": "verl/utils/tokenizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"Utils for tokenization.\"\"\"\n\nimport types\nimport warnings\n\n__all__ = [\"hf_tokenizer\", \"hf_processor\", \"normalize_token_ids\"]\n\n\ndef normalize_token_ids(tokenized_output) -> list[int]:\n    \"\"\"Normalize tokenizer outputs into a flat ``list[int]``.\n\n    This handles Transformers 4/5 differences where ``apply_chat_template(tokenize=True)``\n    may return either ``list[int]`` or a ``BatchEncoding``/mapping with ``input_ids``.\n    \"\"\"\n\n    token_ids = tokenized_output\n    if isinstance(tokenized_output, dict):\n        if \"input_ids\" in tokenized_output:\n            token_ids = tokenized_output[\"input_ids\"]\n    elif hasattr(tokenized_output, \"input_ids\"):\n        token_ids = tokenized_output.input_ids\n\n    if hasattr(token_ids, \"tolist\"):\n        token_ids = token_ids.tolist()\n\n    if isinstance(token_ids, tuple):\n        token_ids = list(token_ids)\n\n    if isinstance(token_ids, list) and len(token_ids) == 1 and isinstance(token_ids[0], list | tuple):\n        token_ids = list(token_ids[0])\n\n    if not isinstance(token_ids, list):\n        raise TypeError(f\"token_ids must be list-like token ids, got {type(token_ids).__name__}: {token_ids!r}\")\n\n    normalized_ids = []\n    for idx, token_id in enumerate(token_ids):\n        if hasattr(token_id, \"item\"):\n            token_id = token_id.item()\n        try:\n            normalized_ids.append(int(token_id))\n        except (TypeError, ValueError) as e:\n            raise TypeError(f\"token_id must be int-convertible, got {type(token_id).__name__}: {token_id!r}\") from e\n    return normalized_ids\n\n\ndef set_pad_token_id(tokenizer):\n    \"\"\"Set pad_token_id to eos_token_id if it is None.\n\n    Args:\n        tokenizer (transformers.PreTrainedTokenizer): The tokenizer to be set.\n\n    \"\"\"\n    if tokenizer.pad_token_id is None:\n        tokenizer.pad_token_id = tokenizer.eos_token_id\n        warnings.warn(f\"tokenizer.pad_token_id is None. Now set to {tokenizer.eos_token_id}\", stacklevel=1)\n    if tokenizer.pad_token is None:\n        tokenizer.pad_token = tokenizer.eos_token\n        warnings.warn(f\"tokenizer.pad_token is None. Now set to {tokenizer.eos_token}\", stacklevel=1)\n\n\ndef hf_tokenizer(name_or_path, correct_pad_token=True, correct_gemma2=True, **kwargs):\n    \"\"\"Create a huggingface pretrained tokenizer which correctness handles eos and pad tokens.\n\n    Args:\n\n        name (str): The name of the tokenizer.\n        correct_pad_token (bool): Whether to correct the pad token id.\n        correct_gemma2 (bool): Whether to correct the gemma2 tokenizer.\n\n    Returns:\n\n        transformers.PreTrainedTokenizer: The pretrained tokenizer.\n\n    \"\"\"\n    from transformers import AutoTokenizer\n\n    if correct_gemma2 and isinstance(name_or_path, str) and \"gemma-2-2b-it\" in name_or_path:\n        # the EOS token in gemma2 is ambiguious, which may worsen RL performance.\n        # https://huggingface.co/google/gemma-2-2b-it/commit/17a01657f5c87135bcdd0ec7abb4b2dece04408a\n        warnings.warn(\n            \"Found gemma-2-2b-it tokenizer. Set eos_token and eos_token_id to <end_of_turn> and 107.\", stacklevel=1\n        )\n        kwargs[\"eos_token\"] = \"<end_of_turn>\"\n        kwargs[\"eos_token_id\"] = 107\n    tokenizer = AutoTokenizer.from_pretrained(name_or_path, **kwargs)\n    if correct_pad_token:\n        set_pad_token_id(tokenizer)\n    return tokenizer\n\n\ndef hf_processor(name_or_path, **kwargs):\n    \"\"\"Create a huggingface processor to process multimodal data.\n\n    Args:\n        name_or_path (str): The name of the processor.\n\n    Returns:\n        Optional[transformers.ProcessorMixin]: The pretrained multimodal processor.\n        Returns ``None`` for text-only models (including AutoProcessor fallbacks to\n        tokenizer backends such as ``TokenizersBackend``).\n    \"\"\"\n    from transformers import AutoConfig, AutoProcessor, PreTrainedTokenizerBase\n\n    try:\n        processor = AutoProcessor.from_pretrained(name_or_path, **kwargs)\n        # In newer transformers, AutoProcessor may legitimately fall back to a\n        # tokenizer backend (e.g. TokenizersBackend) for text-only models.\n        # Treat it as \"no multimodal processor\" and let callers use hf_tokenizer.\n        if isinstance(processor, PreTrainedTokenizerBase):\n            return None\n\n        config = AutoConfig.from_pretrained(name_or_path, **kwargs)\n\n        # Bind vlm model's get_rope_index method to processor\n        processor.config = config\n        model_class = None\n        match processor.__class__.__name__:\n            case \"Qwen2VLProcessor\":\n                from transformers.models.qwen2_vl import Qwen2VLModel\n\n                model_class = Qwen2VLModel\n            case \"Qwen2_5_VLProcessor\":\n                from transformers.models.qwen2_5_vl import Qwen2_5_VLModel\n\n                model_class = Qwen2_5_VLModel\n            case \"Qwen3VLProcessor\":\n                from transformers.models.qwen3_vl import Qwen3VLModel\n\n                model_class = Qwen3VLModel\n            case \"Glm4vImageProcessor\":\n                from transformers.models.glm4v import Glm4vModel\n\n                model_class = Glm4vModel\n            case \"MllamaProcessor\":\n                pass  # MllamaProcessor and MllamaModel doesn't have get_rope_index property\n            case _:\n                raise ValueError(f\"Unsupported processor type: {processor.__class__.__name__}\")\n\n        if model_class is not None:\n            processor.get_rope_index = types.MethodType(model_class.get_rope_index, processor)\n            if hasattr(model_class, \"get_vision_position_ids\"):\n                processor.get_vision_position_ids = types.MethodType(model_class.get_vision_position_ids, processor)\n    except Exception as e:\n        processor = None\n        # TODO(haibin.lin): try-catch should be removed after adding transformer version req to setup.py to avoid\n        # silent failure\n        warnings.warn(f\"Failed to create processor: {e}. This may affect multimodal processing\", stacklevel=1)\n    # Avoid load tokenizer, see:\n    # https://github.com/huggingface/transformers/blob/v4.49.0/src/transformers/models/auto/processing_auto.py#L344\n    if processor is not None and \"Processor\" not in processor.__class__.__name__:\n        processor = None\n    return processor\n"
  },
  {
    "path": "verl/utils/torch_dtypes.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nAdapted from Cruise.\n\"\"\"\n\nimport torch\n\nHALF_LIST = [16, \"16\", \"fp16\", \"float16\", torch.float16]\nFLOAT_LIST = [32, \"32\", \"fp32\", \"float32\", torch.float32]\nBFLOAT_LIST = [\"bf16\", \"bfloat16\", torch.bfloat16]\n\n\nclass PrecisionType:\n    \"\"\"Type of precision used.\n\n    >>> PrecisionType.HALF == 16\n    True\n    >>> PrecisionType.HALF in (16, \"16\")\n    True\n    \"\"\"\n\n    HALF = \"16\"\n    FLOAT = \"32\"\n    FULL = \"64\"\n    BFLOAT = \"bf16\"\n    MIXED = \"mixed\"\n\n    @staticmethod\n    def supported_type(precision: str | int) -> bool:\n        return any(x == precision for x in PrecisionType)\n\n    @staticmethod\n    def supported_types() -> list[str]:\n        return [x.value for x in PrecisionType]\n\n    @staticmethod\n    def is_fp16(precision):\n        return precision in HALF_LIST\n\n    @staticmethod\n    def is_fp32(precision):\n        return precision in FLOAT_LIST\n\n    @staticmethod\n    def is_bf16(precision):\n        return precision in BFLOAT_LIST\n\n    @staticmethod\n    def to_dtype(precision):\n        if precision in HALF_LIST:\n            return torch.float16\n        elif precision in FLOAT_LIST:\n            return torch.float32\n        elif precision in BFLOAT_LIST:\n            return torch.bfloat16\n        else:\n            raise RuntimeError(f\"unexpected precision: {precision}\")\n\n    @staticmethod\n    def to_str(precision):\n        if precision == torch.float16:\n            return \"fp16\"\n        elif precision == torch.float32:\n            return \"fp32\"\n        elif precision == torch.bfloat16:\n            return \"bf16\"\n        else:\n            raise RuntimeError(f\"unexpected precision: {precision}\")\n"
  },
  {
    "path": "verl/utils/torch_functional.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nContain small torch utilities\n\"\"\"\n\nimport math\nfrom contextlib import contextmanager\nfrom typing import Optional\n\nimport torch\nimport torch.distributed\nimport torch.nn.functional as F\nfrom tensordict import TensorDict\nfrom torch import nn\nfrom torch.optim import Optimizer\nfrom torch.optim.lr_scheduler import LambdaLR\nfrom transformers import PreTrainedTokenizer\n\nfrom verl.utils.device import get_device_name, get_torch_device\n\ntry:\n    from flash_attn.ops.triton.cross_entropy import cross_entropy_loss\n\n    FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE = True\nexcept ImportError:\n    FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE = False\n\n\ntry:\n    import torch_npu\n\n    NPU_CROSS_ENTROPY_LOSS_AVAILABLE = hasattr(torch_npu, \"npu_cross_entropy_loss\")\nexcept ImportError:\n    NPU_CROSS_ENTROPY_LOSS_AVAILABLE = False\n\n\ndef gather_from_labels(data: torch.Tensor, label: torch.Tensor) -> torch.Tensor:\n    \"\"\"Gather values from data tensor at positions specified by label indices.\n\n    Selects elements from the last dimension of `data` based on indices in `label`.\n    Commonly used to extract log-probabilities for specific token IDs from a\n    vocabulary distribution.\n\n    Args:\n        data: Input tensor of shape (..., vocab_size) containing values to gather from.\n        label: Index tensor of shape (...,) with values in range [0, vocab_size).\n\n    Returns:\n        torch.Tensor: Gathered values with shape (...,), same as label shape.\n\n    Example:\n        >>> logits = torch.randn(2, 3, 100)  # [batch, seq, vocab]\n        >>> labels = torch.randint(0, 100, (2, 3))  # [batch, seq]\n        >>> gathered = gather_from_labels(logits, labels)  # [batch, seq]\n    \"\"\"\n    output = torch.gather(data, -1, label.unsqueeze(-1)).squeeze(-1)\n    return output\n\n\ndef logprobs_from_logits(logits, labels, inplace_backward=True):\n    \"\"\"\n    Compute per-token log-probabilities for the given labels.\n\n    Uses a Flash-Attention–based cross-entropy (if available) for efficient backward,\n    otherwise falls back to a standard log-softmax+gather approach.\n\n    See: https://github.com/pytorch/pytorch/issues/563#issuecomment-330103591\n\n    Args:\n        logits (Tensor): Model outputs of shape (..., vocab_size).\n        labels (LongTensor): True class indices of shape matching logits[..., :-1].\n        inplace_backward (bool): If True and Flash-Attn is available, perform backward in-place.\n\n    Returns:\n        Tensor: Log-probabilities of the target labels, shape logits.shape[:-1].\n    \"\"\"\n    if FLAH_ATTN_CROSS_ENTROPY_LOSS_AVAILABLE:\n        batch_dim = logits.shape[:-1]\n        last_dim = logits.shape[-1]\n        logits = logits.reshape(-1, last_dim)\n        labels = labels.reshape(-1)\n        output = logprobs_from_logits_flash_attn(logits, labels, inplace_backward=inplace_backward)\n        output = output.view(*batch_dim)\n    elif NPU_CROSS_ENTROPY_LOSS_AVAILABLE:\n        output = logprobs_from_logits_torch_npu(logits, labels)\n    else:\n        output = logprobs_from_logits_v2(logits, labels)\n    return output\n\n\ndef logprobs_from_logits_flash_attn(\n    logits: torch.Tensor, labels: torch.Tensor, inplace_backward: bool = True\n) -> torch.Tensor:\n    \"\"\"Compute log-probabilities using Flash Attention's optimized cross-entropy.\n\n    Uses the Flash Attention library's Triton-based cross-entropy implementation\n    for efficient computation on NVIDIA GPUs.\n\n    Args:\n        logits: Model output logits of shape (batch_size, vocab_size).\n        labels: Target token indices of shape (batch_size,).\n        inplace_backward: If True, perform backward pass in-place for memory efficiency.\n\n    Returns:\n        torch.Tensor: Log-probabilities for target labels, shape (batch_size,).\n\n    Raises:\n        AssertionError: If flash-attn version < 2.4.3 (different return format).\n    \"\"\"\n    output = cross_entropy_loss(logits, labels, inplace_backward=inplace_backward)\n    assert isinstance(output, tuple), (\n        \"please make sure flash-attn>=2.4.3 where cross_entropy_loss returns Tuple[losses, z_losses].\"\n    )\n    return -output[0]\n\n\ndef logprobs_from_logits_torch_npu(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n    \"\"\"Compute log-probabilities using Ascend NPU's optimized cross-entropy.\n\n    Uses torch_npu's native cross-entropy implementation for efficient\n    computation on Huawei Ascend NPU devices.\n\n    Args:\n        logits: Model output logits of shape (..., vocab_size).\n        labels: Target token indices of shape (...,).\n\n    Returns:\n        torch.Tensor: Log-probabilities for target labels, same shape as labels.\n    \"\"\"\n    batch_dim = logits.shape[:-1]\n    logits = logits.reshape(-1, logits.shape[-1])\n    loss, _, _, _ = torch_npu.npu_cross_entropy_loss(logits, labels.reshape(-1), reduction=\"none\")\n    return -loss.view(*batch_dim)\n\n\ndef logprobs_from_logits_naive(logits: torch.Tensor, labels: torch.Tensor) -> torch.Tensor:\n    \"\"\"Compute log-probabilities using standard log-softmax approach.\n\n    Simple implementation using PyTorch's log_softmax followed by gathering.\n    Less memory-efficient than specialized implementations but works on all devices.\n\n    Args:\n        logits: Model output logits of shape (..., vocab_size).\n        labels: Target token indices of shape (...,).\n\n    Returns:\n        torch.Tensor: Log-probabilities for target labels, same shape as labels.\n    \"\"\"\n    logp = F.log_softmax(logits, dim=-1)\n    logpy = gather_from_labels(logp, labels)\n    return logpy\n\n\ndef logprobs_from_logits_v2(logits: torch.FloatTensor, labels: torch.Tensor) -> torch.Tensor:\n    \"\"\"Memory-efficient log-probability computation using row-wise processing.\n\n    Computes log-probabilities by processing one row at a time to reduce peak\n    memory consumption. Uses logsumexp for float32/float64, falls back to\n    log_softmax for bfloat16 due to numerical stability concerns.\n\n    The mathematical identity used is: log_softmax(x_i) = x_i - logsumexp(x)\n\n    Args:\n        logits: Model output logits of shape (batch_size, seq_len, vocab_size)\n            or (batch_size, vocab_size).\n        labels: Target token indices matching logits shape without vocab dimension.\n\n    Returns:\n        torch.Tensor: Log-probabilities for target labels.\n\n    Note:\n        This implementation trades compute for memory by iterating over batch\n        dimension, making it suitable for large vocabulary sizes.\n    \"\"\"\n    if logits.dtype in [torch.float32, torch.float64]:\n        logits_labels = torch.gather(logits, dim=-1, index=labels.unsqueeze(-1)).squeeze(-1)\n        # loop to reduce peak mem consumption\n        logsumexp_values = torch.stack([torch.logsumexp(logit, dim=-1) for logit in logits])\n        logprobs_labels = logits_labels - logsumexp_values  # log_softmax(x_i) = x_i - logsumexp(x)\n    else:\n        # logsumexp approach is unstable with bfloat16, fall back to slightly less efficent approach\n        logprobs_labels = []\n        for row_logits, row_labels in zip(logits, labels, strict=True):  # loop to reduce peak mem consumption\n            row_logprobs = F.log_softmax(row_logits, dim=-1)\n            row_logprobs_labels = row_logprobs.gather(dim=-1, index=row_labels.unsqueeze(-1)).squeeze(-1)\n            logprobs_labels.append(row_logprobs_labels)\n        logprobs_labels = torch.stack(logprobs_labels)\n    return logprobs_labels\n\n\ndef clip_by_value(x: torch.Tensor, tensor_min: torch.Tensor, tensor_max: torch.Tensor) -> torch.Tensor:\n    \"\"\"Clip tensor values to a range defined by tensor bounds.\n\n    Extension of torch.clamp that supports tensor-valued min/max bounds\n    instead of only scalar bounds.\n\n    Args:\n        x: Input tensor to clip.\n        tensor_min: Minimum bound tensor (broadcastable to x).\n        tensor_max: Maximum bound tensor (broadcastable to x).\n\n    Returns:\n        torch.Tensor: Clipped tensor with values in [tensor_min, tensor_max].\n\n    See Also:\n        https://github.com/pytorch/pytorch/issues/2793#issuecomment-428784713\n    \"\"\"\n    clipped = torch.max(torch.min(x, tensor_max), tensor_min)\n    return clipped\n\n\ndef entropy_from_logits(logits: torch.Tensor) -> torch.Tensor:\n    \"\"\"Calculate Shannon entropy from unnormalized logits.\n\n    Computes H(p) = -sum(p * log(p)) using the numerically stable formula:\n    entropy = logsumexp(logits) - sum(softmax(logits) * logits)\n\n    Args:\n        logits: Unnormalized log-probabilities of shape (..., vocab_size).\n\n    Returns:\n        torch.Tensor: Entropy values with shape (...,), one per distribution.\n    \"\"\"\n    pd = torch.nn.functional.softmax(logits, dim=-1)\n    entropy = torch.logsumexp(logits, dim=-1) - torch.sum(pd * logits, dim=-1)\n    return entropy\n\n\ndef entropy_from_logits_with_chunking(logits: torch.Tensor, chunk_size: int = 2048) -> torch.Tensor:\n    \"\"\"Memory-efficient entropy calculation using chunked processing.\n\n    Computes entropy by processing the batch in chunks to reduce peak memory\n    usage. Useful for large batch sizes or when memory is constrained.\n\n    Args:\n        logits: Unnormalized log-probabilities of shape (batch_size, vocab_size).\n        chunk_size: Number of samples to process at once. Defaults to 2048.\n\n    Returns:\n        torch.Tensor: Entropy values with shape (batch_size,).\n\n    Note:\n        Converts chunks to float32 for numerical stability during computation.\n    \"\"\"\n    entropy = torch.zeros(logits.shape[0], device=logits.device)\n    for i in range(0, logits.shape[0], chunk_size):\n        logits_chunk = logits[i : i + chunk_size].float()\n        pd_chunk = torch.nn.functional.softmax(logits_chunk, dim=-1)\n        entropy_chunk = torch.logsumexp(logits_chunk, dim=-1) - torch.sum(pd_chunk * logits_chunk, dim=-1)\n        entropy[i : i + chunk_size] = entropy_chunk\n    return entropy\n\n\ndef masked_sum(values: torch.Tensor, mask: torch.Tensor, axis: int | tuple[int, ...] | None = None) -> torch.Tensor:\n    \"\"\"Compute sum of tensor values where mask is True.\n\n    NaN values outside the mask are replaced with zeros to prevent\n    contaminating the sum.\n\n    Args:\n        values: Input tensor containing values to sum.\n        mask: Boolean or numeric mask tensor (same shape as values).\n            Non-zero values indicate elements to include.\n        axis: Dimension(s) along which to sum. None sums all elements.\n\n    Returns:\n        torch.Tensor: Sum of masked values, reduced along specified axis.\n    \"\"\"\n    # If NaNs exist out of mask, replace NaNs in values with a value that\n    # won't affect the sum (e.g., 0 for masked regions)\n    valid_values = torch.where(mask.bool(), values, 0.0)\n    return (valid_values * mask).sum(axis=axis)\n\n\ndef masked_mean(values, mask, axis=None):\n    \"\"\"\n    Compute the mean of `values` over elements selected by `mask`.\n\n    Args:\n        values (Tensor): Input tensor.\n        mask (Tensor): Boolean or numeric mask of the same shape as `values`.\n        axis (int or tuple of int, optional): Dimension(s) along which to compute the mean.\n            Defaults to None (over all elements).\n\n    Returns:\n        Tensor: Masked mean, with shape equal to `values` reduced over `axis`.\n    \"\"\"\n    s = masked_sum(values, mask, axis)\n    return s / (mask.sum(axis=axis) + 1e-8)\n\n\ndef masked_var(values, mask, unbiased=True):\n    \"\"\"Compute variance of tensor with masked values.\"\"\"\n    mean = masked_mean(values, mask)\n    centered_values = values - mean\n    variance = masked_mean(centered_values**2, mask)\n    if unbiased:\n        mask_sum = mask.sum()\n        if mask_sum == 0:\n            raise ValueError(\"At least one element in the mask has to be 1.\")\n        # note that if mask_sum == 1, then there is a division by zero issue\n        # to avoid it you just need to use a larger minibatch_size\n        if mask_sum == 1:\n            raise ValueError(\"The sum of the mask is one, which can cause a division by zero.\")\n        bessel_correction = mask_sum / (mask_sum - 1)\n        variance = variance * bessel_correction\n    return variance\n\n\ndef masked_whiten(values, mask, shift_mean=True):\n    \"\"\"\n    Whiten `values` by normalizing with mean and variance computed over `mask`.\n\n    Args:\n        values (torch.Tensor): Input tensor.\n        mask (torch.Tensor): Boolean tensor of same shape, selects elements for stats.\n        shift_mean (bool): If True (default), output is zero-mean;\n                           if False, the original mean is re-added after scaling.\n\n    Returns:\n        torch.Tensor: Whitened tensor of same shape as `values`.\n    \"\"\"\n    mean, var = masked_mean(values, mask), masked_var(values, mask)\n    whitened = (values - mean) * torch.rsqrt(var + 1e-8)\n    if not shift_mean:\n        whitened += mean\n    return whitened\n\n\ndef get_response_mask(response_id: torch.Tensor, eos_token: int | list[int] = 2, dtype=torch.int64):\n    \"\"\"\n    end of sentence token can be int or list: 1 or [1, 2]\n    e.g.\n    response_id = torch.tensor([[20, 10, 34, 1, 0, 0, 0],\n                                [78, 0, 76, 2, 1, 0, 0],\n                                [23, 98, 1, 0, 0, 0, 0],\n                                [33, 3, 98, 45, 1, 0, 0]])\n    #eos_token=1\n    response_mask:  tensor([[1, 1, 1, 1, 0, 0, 0],\n                            [1, 1, 1, 1, 1, 0, 0],\n                            [1, 1, 1, 0, 0, 0, 0],\n                            [1, 1, 1, 1, 1, 0, 0]])\n    #eos_token=[1,2]\n    response_mask:  tensor([[1, 1, 1, 1, 0, 0, 0],\n                            [1, 1, 1, 1, 0, 0, 0],\n                            [1, 1, 1, 0, 0, 0, 0],\n                            [1, 1, 1, 1, 1, 0, 0]])\n    \"\"\"\n    eos_mask = torch.isin(response_id, torch.tensor(eos_token, device=response_id.device)).int()\n    return (eos_mask.cumsum(dim=1) - eos_mask).eq(0).to(dtype)\n\n\ndef compute_grad_norm(model: nn.Module) -> float:\n    \"\"\"Compute the squared L2 norm of all gradients in a model.\n\n    Sums the squared values of all gradient tensors across all parameters.\n    Useful for monitoring gradient magnitudes during training.\n\n    Args:\n        model: PyTorch model with computed gradients.\n\n    Returns:\n        float: Sum of squared gradient values (not the square root).\n\n    Note:\n        Returns the squared norm, not the norm itself. To get the actual\n        L2 norm, take the square root of the returned value.\n    \"\"\"\n    total_grad_square = 0\n    for param in model.parameters():\n        if param.grad is not None:\n            total_grad_square += torch.sum(torch.square(param.grad.detach())).item()\n    return total_grad_square\n\n\ndef broadcast_dict_tensor(tensors: dict[str, torch.Tensor] | TensorDict, src: int, group) -> None:\n    \"\"\"Broadcast all tensors in a dictionary from source rank to all ranks.\n\n    Iterates over all tensors in the dictionary and broadcasts each one\n    from the source rank to all other ranks in the process group.\n\n    Args:\n        tensors: Dictionary or TensorDict containing tensors to broadcast.\n        src: Source rank from which to broadcast.\n        group: Process group for the broadcast operation.\n\n    Note:\n        This implementation broadcasts tensors one at a time. Could be optimized\n        to use a single broadcast with packed tensors.\n    \"\"\"\n    for key in tensors.sorted_keys:\n        torch.distributed.broadcast(tensors[key], src=src, group=group, async_op=False)\n\n\ndef allgather_dict_tensors(\n    tensors: dict[str, torch.Tensor] | TensorDict, size: int, group, dim: int = 0\n) -> dict[str, torch.Tensor] | TensorDict:\n    \"\"\"Gather tensors from all ranks and concatenate them.\n\n    Performs all_gather on each tensor in the dictionary and concatenates\n    the results along the specified dimension.\n\n    Args:\n        tensors: Dictionary or TensorDict containing tensors to gather.\n        size: Number of ranks in the process group.\n        group: Process group for the all_gather operation.\n        dim: Dimension along which to concatenate gathered tensors. Defaults to 0.\n\n    Returns:\n        Dictionary or TensorDict (matching input type) with gathered and\n        concatenated tensors. Each tensor's size along `dim` is multiplied by `size`.\n\n    Note:\n        This implementation gathers tensors one at a time synchronously.\n        Could be optimized using async ops or packed all_gather.\n    \"\"\"\n    if isinstance(tensors, TensorDict):\n        is_tensor_dict = True\n        tensors_as_dict = tensors.to_dict()\n    else:\n        tensors_as_dict = tensors\n        is_tensor_dict = False\n\n    output = {}\n    sorted_keys = sorted(tensors_as_dict.keys())\n    for key in sorted_keys:\n        val = tensors_as_dict[key]\n        output[key] = [torch.empty_like(val) for _ in range(size)]\n        torch.distributed.all_gather(output[key], val, group=group, async_op=False)\n        output[key] = torch.cat(output[key], dim=dim)\n\n    if is_tensor_dict:\n        output = TensorDict(source=output, batch_size=tensors.batch_size[0] * size)\n\n    return output\n\n\ndef allgather_dict_into_dict(data: dict, group=None) -> dict:\n    \"\"\"allgather a dict into a dict of list\n\n    Args:\n        data: a dict\n        group: the process group to allgather\n\n    Returns: dict containing a list of the results from allgather\n\n    \"\"\"\n    assert isinstance(data, dict), f\"Expect data to be a dictionary, Got {type(data)}\"\n\n    group_size = torch.distributed.get_world_size(group=group)\n\n    final_metrics = {}\n    all_metrics_lst = [None for _ in range(group_size)]\n    torch.distributed.all_gather_object(all_metrics_lst, data, group=group)\n\n    for all_metrics in all_metrics_lst:\n        for key, val in all_metrics.items():\n            if key not in final_metrics:\n                final_metrics[key] = []\n            final_metrics[key].append(val)\n    return final_metrics\n\n\ndef split_dict_tensor_into_batches(tensors: TensorDict, batch_size) -> list[TensorDict]:\n    assert tensors.batch_size[0] % batch_size == 0, (\n        f\"input data batch size: {tensors.batch_size[0]}, split batch size: {batch_size}\"\n    )\n    return tensors.split(batch_size)\n\n\ndef pad_2d_list_to_length(response, pad_token_id, max_length=None):\n    \"\"\"\n    pad a 2D list (e.g. responses, logprobs) to a 2D tensor.\n    \"\"\"\n    response_length = max(len(sub_list) for sub_list in response)\n    target_length = max_length if max_length is not None and max_length > response_length else response_length\n    padded_response = [tuple(sub_list) + (pad_token_id,) * (target_length - len(sub_list)) for sub_list in response]\n    tensor = torch.tensor(padded_response)\n    return tensor\n\n\ndef pad_sequence_to_length(tensors, max_seq_len, pad_token_id, left_pad=False):\n    \"\"\"\n    pad a 2D tensors (e.g. responses, logprobs) in the last dim to max_seq_length.\n    input shape: [bs, seq_length]\n    output shape: [bs, max_seq_length]\n    \"\"\"\n    if tensors.shape[-1] >= max_seq_len:\n        return tensors\n    # (0, max_seq_len - tensors.shape[-1]) means right pad to max_seq_length and no left pad\n    pad_tuple = (max_seq_len - tensors.shape[-1], 0) if left_pad else (0, max_seq_len - tensors.shape[-1])\n    return F.pad(tensors, pad_tuple, \"constant\", pad_token_id)\n\n\ndef postprocess_data(\n    input_ids: torch.Tensor,\n    attention_mask: torch.Tensor,\n    max_length: int,\n    pad_token_id: int,\n    left_pad=True,\n    truncation=\"error\",\n):\n    \"\"\"Process tokenizer outputs to consistent shapes via padding/truncation.\n\n    Args:\n        input_ids: Token indices [batch_size, seq_len]\n        attention_mask: Mask [batch_size, seq_len]\n        max_length: Target sequence length\n        pad_token_id: Padding token ID\n        left_pad: Pad left if True\n        truncation: \"left\", \"right\", \"middle\" or \"error\"\n\n    Returns:\n        (input_ids, attention_mask) padded/truncated to max_length\n    \"\"\"\n    assert truncation in [\"left\", \"right\", \"middle\", \"error\"]\n    assert input_ids.ndim == 2\n\n    sequence_length = input_ids.shape[-1]\n    if sequence_length < max_length:\n        input_ids = pad_sequence_to_length(\n            input_ids, max_seq_len=max_length, pad_token_id=pad_token_id, left_pad=left_pad\n        )\n        attention_mask = pad_sequence_to_length(\n            attention_mask, max_seq_len=max_length, pad_token_id=0, left_pad=left_pad\n        )\n    elif sequence_length > max_length:\n        if truncation == \"left\":\n            # actually, left truncation may not be reasonable\n            input_ids = input_ids[:, -max_length:]\n            attention_mask = attention_mask[:, -max_length:]\n        elif truncation == \"right\":\n            input_ids = input_ids[:, :max_length]\n            attention_mask = attention_mask[:, :max_length]\n        elif truncation == \"middle\":\n            left_half = max_length // 2\n            right_half = max_length - left_half\n            input_ids = torch.cat([input_ids[:, :left_half], input_ids[:, -right_half:]], dim=-1)\n            attention_mask = torch.cat([attention_mask[:, :left_half], attention_mask[:, -right_half:]], dim=-1)\n        elif truncation == \"error\":\n            raise NotImplementedError(f\"{sequence_length=} is larger than {max_length=}\")\n        else:\n            raise NotImplementedError(f\"Unknown truncation method {truncation}\")\n\n    return input_ids, attention_mask\n\n\ndef tokenize_and_postprocess_data(\n    prompt: str, tokenizer: PreTrainedTokenizer, max_length: int, pad_token_id: int, left_pad=True, truncation=\"error\"\n):\n    \"\"\"Tokenize text and process outputs to consistent tensor shapes.\n\n    Args:\n        prompt: Input text to tokenize\n        tokenizer: HuggingFace tokenizer instance\n        max_length: Target sequence length\n        pad_token_id: Padding token ID\n        left_pad: Pad left if True\n        truncation: Truncation strategy (\"left\"/\"right\"/\"error\")\n\n    Returns:\n        Tuple of (input_ids, attention_mask) from postprocess_data\n    \"\"\"\n    input_data = tokenizer(prompt, return_tensors=\"pt\", add_special_tokens=False)\n    input_ids = input_data[\"input_ids\"]\n    attention_mask = input_data[\"attention_mask\"]\n\n    return postprocess_data(input_ids, attention_mask, max_length, pad_token_id, left_pad, truncation)\n\n\ndef remove_pad_token(input_ids: torch.Tensor, attention_mask: torch.Tensor):\n    \"\"\"Remove the pad token.\n\n    Args:\n        input_ids shape: [bs, seq_length]\n        attention_mask shape: [bs, seq_length]\n    Returns:\n        no_padding_batch(List[List[int]]): contains the rmpad token ids per query.\n    \"\"\"\n    no_padding_batch = []\n    for ids, mask in zip(input_ids, attention_mask, strict=True):\n        no_padding_batch.append((ids[len(ids) - mask.sum() :]).cpu().numpy().tolist())\n    return no_padding_batch\n\n\ndef log_probs_from_logits_response(input_ids, logits, response_length):\n    \"\"\"Compute the response log_probs from full logits. Note that logits = model(input_ids)\n\n    Args:\n        input_ids: [batch_size, seqlen]\n        logits: [batch_size, seqlen, vocab_size]\n\n    Returns:\n        response_log_prob:\n    \"\"\"\n    response_logits = logits[:, -response_length - 1 : -1]\n    response = input_ids[:, -response_length:]\n    response_log_prob = logprobs_from_logits(logits=response_logits, labels=response)\n    return response_log_prob\n\n\ndef log_probs_from_logits_response_rmpad(input_ids, attention_mask, logits_rmpad, response_length):\n    \"\"\"Compute the log_probs from logits with rmpad logits and pad input. Note that\n    logits_rmpad = model(input_ids_rmpad). For each sentences, there is a shift between\n    logits and input_ids.\n    The reason for this function to is to compute logprobs_from_logits in rmpad mode because it is memory-intensive\n    for large vocab_size\n\n    Args:\n        input_ids: [batch_size, seqlen]\n        attention_mask: [batch_size, seqlen]\n        logits_rmpad: [total_nnz, vocab_size]\n        response_length: int\n    \"\"\"\n    from flash_attn.bert_padding import pad_input, unpad_input\n\n    batch_size, seqlen = input_ids.shape\n    input_ids_rmpad, indices, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask=attention_mask)\n    input_ids_rmpad = input_ids_rmpad.squeeze(-1)\n    input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=0)\n    full_log_probs_rmpad = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled)  # (total_nnz,)\n    full_output = pad_input(\n        hidden_states=full_log_probs_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen\n    )\n    output = full_output.squeeze(-1)[:, -response_length - 1 : -1]  # [batch_size, response_length]\n    return output\n\n\ndef log_probs_from_logits_all_rmpad(input_ids_rmpad, logits_rmpad, indices, batch_size, seqlen, response_length):\n    \"\"\"Compute the log_probs from logits with rmpad input_ids and logits. Note that\n    logits_rmpad = model(input_ids_rmpad). For each sentences, there is a shift between\n    logits and input_ids.\n    The reason for this function to is to compute logprobs_from_logits in rmpad mode because it is memory-intensive\n    for large vocab_size\n\n    Args:\n        input_ids_rmpad: [1, total_nnz]\n        logits_rmpad: [total_nnz, vocab_size]\n        indices: [total_nnz]\n        batch_size: int\n        seqlen: int\n        response_length: int\n    \"\"\"\n    if get_device_name() == \"cuda\":\n        from flash_attn.bert_padding import pad_input\n    elif get_device_name() == \"npu\":\n        from verl.utils.attention_utils import pad_input\n\n    input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # transpose back to [total_nnz, 1]\n    input_ids_rmpad = input_ids_rmpad.squeeze(-1)\n    input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=0)\n    full_log_probs_rmpad = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled)  # (total_nnz,)\n    full_output = pad_input(\n        hidden_states=full_log_probs_rmpad.unsqueeze(-1), indices=indices, batch=batch_size, seqlen=seqlen\n    )\n    output = full_output.squeeze(-1)[:, -response_length - 1 : -1]  # [batch_size, response_length]\n    return output\n\n\ndef post_process_logits(input_ids, logits, temperature, top_k, top_p):\n    if temperature != 1.0:\n        logits = logits.div_(temperature)  # inplace operation to avoid OOM\n    # TODO: add them back\n    # if top_k is not None and top_k > 0:\n    #     logits = TopKLogitsWarper(top_k=top_k)(input_ids, logits)\n    # if top_p is not None and top_p < 1.0 and top_p > 0.0:\n    #     logits = TopPLogitsWarper(top_p=top_p)(input_ids, logits)\n    return logits\n\n\ndef calculate_sum_pi_squared_from_logits(logits: torch.Tensor):\n    \"\"\"\n    Compute exact sum of squared probabilities from logits.\n    Formula: Σπ² = exp(logsumexp(2*logits) - 2*logsumexp(logits))\n\n    Used for optimal baseline variance reduction as described in\n    \"What Matters for Model Merging at Scale?\" (arXiv:2410.03617)\n\n    Args:\n        logits: Logits tensor (..., vocab_size).\n\n    Returns:\n        Sum of squared probabilities tensor (...).\n    \"\"\"\n    return torch.exp(torch.logsumexp(2.0 * logits, dim=-1) - 2.0 * torch.logsumexp(logits, dim=-1))\n\n\n\"\"\"\nOptimizer related\n\"\"\"\n\n\ndef get_cosine_schedule_with_warmup(\n    optimizer: Optimizer,\n    num_warmup_steps: int,\n    num_training_steps: int,\n    min_lr_ratio: float = 0.0,\n    num_cycles: float = 0.5,\n    last_epoch: int = -1,\n    init_lr_ratio: float = None,\n    zero_indexed_step: bool = True,\n):\n    \"\"\"\n    Create a schedule with a learning rate that decreases following the values of the cosine function between the\n    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the\n    initial lr set in the optimizer.\n    Args:\n        optimizer (:class:`~torch.optim.Optimizer`):\n            The optimizer for which to schedule the learning rate.\n        num_warmup_steps (:obj:`int`):\n            The number of steps for the warmup phase.\n        num_training_steps (:obj:`int`):\n            The total number of training steps.\n        min_lr_ratio (:obj:`float`, `optional`, defaults to 0.0):\n            The minimum lr ratio w.r.t the maximum.\n        num_cycles (:obj:`float`, `optional`, defaults to 0.5):\n            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0\n            following a half-cosine).\n        last_epoch (:obj:`int`, `optional`, defaults to -1):\n            The index of the last epoch when resuming training.\n        init_lr_ratio (:obj:`float`, `optional`, defaults to None):\n            The initial lr ratio w.r.t the maximum.\n        zero_indexed_step (:obj:`bool`, `optional`, defaults to True):\n            Whether the LR schedule uses 0-indexed steps. If True (default), step counting starts at 0.\n            If False (used by torchtitan), step counting starts at 1.\n    Return:\n        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n    \"\"\"\n    min_lr_ratio = 0.0 if min_lr_ratio is None else min_lr_ratio\n    assert min_lr_ratio >= 0 and min_lr_ratio <= 1.0\n    coef = (1 - min_lr_ratio) * 0.5\n    intercept = (1 + min_lr_ratio) * 0.5\n\n    init_lr_ratio = 0.0 if init_lr_ratio is None else init_lr_ratio\n    assert init_lr_ratio >= 0 and init_lr_ratio <= 1.0\n\n    def lr_lambda(current_step):\n        if not zero_indexed_step:\n            current_step += 1\n        if current_step < num_warmup_steps:\n            return init_lr_ratio + (1.0 - init_lr_ratio) * (float(current_step) / float(max(1, num_warmup_steps)))\n        progress = float(current_step - num_warmup_steps) / float(max(1, num_training_steps - num_warmup_steps))\n        x = math.cos(math.pi * float(num_cycles) * 2.0 * progress)\n        return max(min_lr_ratio, x * coef + intercept)\n\n    return LambdaLR(optimizer, lr_lambda, last_epoch)\n\n\ndef get_constant_schedule_with_warmup(\n    optimizer: Optimizer,\n    num_warmup_steps: int,\n    last_epoch: int = -1,\n):\n    \"\"\"\n    Create a constant LR schedule with a linear warmup phase.\n\n    Args:\n        optimizer (Optimizer): Wrapped optimizer.\n        num_warmup_steps (int): Number of steps to ramp up the LR from 0 to initial value.\n        last_epoch (int, optional): The index of the last epoch when resuming training. Defaults to -1.\n\n    Returns:\n        LambdaLR: Scheduler that increases LR linearly during warmup, then holds it constant.\n    \"\"\"\n\n    def lr_lambda(current_step):\n        if current_step < num_warmup_steps:\n            return float(current_step) / float(max(1.0, num_warmup_steps))\n        return 1.0\n\n    return LambdaLR(optimizer, lr_lambda, last_epoch)\n\n\ndef prepare_decoder_attention_mask(attention_mask, input_shape, inputs_embeds):\n    # create causal mask\n    # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n    combined_attention_mask = None\n    if input_shape[-1] > 1:\n        combined_attention_mask = _make_causal_mask(\n            input_shape,\n            inputs_embeds.dtype,\n            device=inputs_embeds.device,\n        )\n\n    if attention_mask is not None:\n        # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]\n        expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(\n            inputs_embeds.device\n        )\n        combined_attention_mask = (\n            expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask\n        )\n\n    return combined_attention_mask\n\n\n# Copied from transformers.models.bart.modeling_bart._make_causal_mask\ndef _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, device: torch.device):\n    \"\"\"\n    Make causal mask used for bi-directional self-attention.\n    \"\"\"\n    bsz, tgt_len = input_ids_shape\n    mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)\n    mask_cond = torch.arange(mask.size(-1), device=device)\n    mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)\n    mask = mask.to(dtype)\n    return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len)\n\n\n# Copied from transformers.models.bart.modeling_bart._expand_mask\ndef _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):\n    \"\"\"\n    Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.\n    \"\"\"\n    bsz, src_len = mask.size()\n    tgt_len = tgt_len if tgt_len is not None else src_len\n\n    expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)\n\n    inverted_mask = 1.0 - expanded_mask\n\n    return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)\n\n\ndef get_unpad_data(attention_mask):\n    seqlens_in_batch = attention_mask.sum(dim=-1, dtype=torch.int32)\n    indices = torch.nonzero(attention_mask.flatten(), as_tuple=False).flatten()\n    max_seqlen_in_batch = seqlens_in_batch.max().item()\n    cu_seqlens = F.pad(torch.cumsum(seqlens_in_batch, dim=0, dtype=torch.int32), (1, 0))\n    return (\n        indices,\n        cu_seqlens,\n        max_seqlen_in_batch,\n    )\n\n\ndef get_wsd_schedule_with_warmup(\n    optimizer: Optimizer,\n    num_warmup_steps: int,\n    num_training_steps: int,\n    min_lr_ratio: float = 0.0,\n    num_cycles: float = 0.5,\n    last_epoch: int = -1,\n    stable_ratio: float = 0.9,\n):\n    \"\"\"\n    Create a Warmup-Stable-Decay learning rate scheduler.\n\n    The schedule follows three phases:\n    1. Warmup: Learning rate increases linearly from 0 to the initial LR\n    2. Stable: Learning rate remains constant at the initial LR\n    3. Decay: Learning rate decreases following a cosine curve to min_lr_ratio * initial LR\n\n    Args:\n        optimizer (:class:`~torch.optim.Optimizer`):\n            The optimizer for which to schedule the learning rate.\n        num_warmup_steps (:obj:`int`):\n            The number of steps for the warmup phase.\n        num_training_steps (:obj:`int`):\n            The total number of training steps.\n        min_lr_ratio (:obj:`float`, `optional`, defaults to 0.0):\n            The minimum learning rate ratio w.r.t the initial learning rate.\n        num_cycles (:obj:`float`, `optional`, defaults to 0.5):\n            The number of waves in the cosine schedule during decay phase.\n        last_epoch (:obj:`int`, `optional`, defaults to -1):\n            The index of the last epoch when resuming training.\n        stable_ratio (:obj:`float`, `optional`, defaults to 0.0):\n            The ratio of non-warmup steps that should maintain a constant learning rate.\n            Set to 0.0 to behave exactly like cosine schedule.\n\n    Return:\n        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.\n    \"\"\"\n    remaining_steps = max(0, num_training_steps - num_warmup_steps)\n    num_stable_steps = int(remaining_steps * stable_ratio)\n    num_decay_steps = remaining_steps - num_stable_steps\n\n    def lr_lambda(current_step):\n        if current_step < num_warmup_steps:\n            return float(current_step) / float(max(1, num_warmup_steps))\n        if current_step < num_warmup_steps + num_stable_steps:\n            return 1.0\n        if current_step < num_training_steps:\n            progress = float(current_step - num_warmup_steps - num_stable_steps) / float(max(1, num_decay_steps))\n            value = max(0.0, 0.5 * (1.0 + math.cos(math.pi * float(num_cycles) * 2.0 * progress)))\n            return (1.0 - min_lr_ratio) * value + min_lr_ratio\n        return min_lr_ratio\n\n    return LambdaLR(optimizer, lr_lambda, last_epoch)\n\n\n@contextmanager\ndef check_device_is_available():\n    \"\"\"\n    Some modules must be imported after CUDA is initialized. Such as sglang's sharding manager.\n\n    This context manager checks if CUDA is available and raises an error if it is not.\n    \"\"\"\n    if not get_torch_device().is_available():\n        raise RuntimeError(\"Device {} must be initialized before importing this module.\".format(get_device_name()))\n\n    yield\n\n\ndef distributed_mean_max_min_std(local_tensor, compute_max=True, compute_min=True, compute_std=True):\n    \"\"\"Compute distributed statistics across all processes.\n\n    Args:\n        local_tensor: Tensor containing local values\n        compute_max: Include maximum value calculation\n        compute_min: Include minimum value calculation\n        compute_std: Include standard deviation calculation\n\n    Returns:\n        Tuple containing (mean, max, min, std) in this order. None for disabled metrics.\n    \"\"\"\n    # Sum the local tensor across all processes\n    local_sum = torch.sum(local_tensor)\n    local_num = torch.tensor(torch.numel(local_tensor), device=get_device_name())\n\n    torch.distributed.all_reduce(local_sum, op=torch.distributed.ReduceOp.SUM)\n    torch.distributed.all_reduce(local_num, op=torch.distributed.ReduceOp.SUM)\n\n    global_mean = local_sum / local_num\n\n    if compute_max:\n        local_max = torch.max(local_tensor)\n        torch.distributed.all_reduce(local_max, op=torch.distributed.ReduceOp.MAX)\n    else:\n        local_max = None\n\n    if compute_min:\n        local_min = torch.min(local_tensor)\n        torch.distributed.all_reduce(local_min, op=torch.distributed.ReduceOp.MIN)\n    else:\n        local_min = None\n\n    if compute_std:\n        square_diff = torch.sum(torch.pow(local_tensor - global_mean, 2))\n        torch.distributed.all_reduce(square_diff, op=torch.distributed.ReduceOp.SUM)\n        global_std = torch.sqrt(square_diff / (local_num - 1))\n    else:\n        global_std = None\n\n    return global_mean, local_max, local_min, global_std\n\n\ndef distributed_masked_mean(local_tensor, local_mask):\n    \"\"\"Compute global mean of non-masked elements across distributed processes.\n\n    Args:\n        local_tensor (torch.Tensor): Input tensor with local values\n        local_mask (torch.Tensor): Binary mask (1=valid, 0=ignore) matching local_tensor shape\n\n    Returns:\n        torch.Tensor: Global mean of all valid elements across processes\n    \"\"\"\n    local_tensor = local_tensor * local_mask\n\n    local_sum = torch.sum(local_tensor)\n    local_num = torch.sum(local_mask)\n\n    torch.distributed.all_reduce(local_sum, op=torch.distributed.ReduceOp.SUM)\n    torch.distributed.all_reduce(local_num, op=torch.distributed.ReduceOp.SUM)\n\n    global_mean = local_sum / local_num\n    return global_mean\n\n\ndef expand_as_nested(tensor: torch.Tensor, nested_tensor: torch.Tensor) -> torch.Tensor:\n    \"\"\"\n\n    Args:\n        tensor: a tensor with shape (bsz,)\n        nested_tensor: a nested tensor with shape (bsz, xxx)\n\n    Returns:\n        a tensor with the same shape as nested_tensor\n\n    \"\"\"\n    assert nested_tensor.is_nested, \"nested_tensor must be nested\"\n    assert tensor.shape[0] == nested_tensor.shape[0], (\n        f\"The batch shape must be the same. Got {tensor.shape[0]} vs {nested_tensor.shape[0]}\"\n    )\n    assert len(tensor.shape) == 1, \"The ndim of tensor must be 1\"\n    assert len(nested_tensor.shape) == 2, \"The ndim of nested_tensor must be 2\"\n\n    offsets = nested_tensor.offsets()\n    seqlens = offsets.diff()\n    output = torch.repeat_interleave(tensor, seqlens, dim=0)\n    output = torch.nested.nested_tensor_from_jagged(values=output, offsets=offsets)\n    return output\n\n\n@contextmanager\ndef use_original_torch_compile():\n    \"\"\"torch.compile might be replaced by mindspeed on NPU, this contextmanager\n    can revert torch.compile temporarily.\n    \"\"\"\n    try:\n        from mindspeed.patch_utils import MindSpeedPatchesManager\n\n        compile_patch = None\n        for patch in MindSpeedPatchesManager.patches_info.values():\n            if patch.orig_module_name == \"torch\" and patch.orig_func_name == \"compile\":\n                if patch.is_applied():\n                    compile_patch = patch\n                break\n        if compile_patch is not None:\n            compile_patch.remove_patch()\n            yield\n            compile_patch.apply_patch()\n        else:\n            yield\n    except Exception:\n        yield\n"
  },
  {
    "path": "verl/utils/tracking.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nA unified tracking interface that supports logging data to different backend\n\"\"\"\n\nimport dataclasses\nimport json\nimport logging\nimport os\nfrom enum import Enum\nfrom functools import partial\nfrom pathlib import Path\nfrom typing import Any\n\nimport orjson\n\nlogger = logging.getLogger(__name__)\n\nMLFLOW_MAX_ATTEMPTS = 3\nMLFLOW_SLEEP_SECONDS = 5\n\n\nclass Tracking:\n    \"\"\"A unified tracking interface for logging experiment data to multiple backends.\n\n    This class provides a centralized way to log experiment metrics, parameters, and artifacts\n    to various tracking backends including WandB, MLflow, SwanLab, TensorBoard, and console.\n\n    Attributes:\n        supported_backend: List of supported tracking backends.\n        logger: Dictionary of initialized logger instances for each backend.\n    \"\"\"\n\n    supported_backend = [\n        \"wandb\",\n        \"mlflow\",\n        \"swanlab\",\n        \"vemlp_wandb\",\n        \"tensorboard\",\n        \"console\",\n        \"clearml\",\n        \"trackio\",\n        \"file\",\n    ]\n\n    def __init__(self, project_name, experiment_name, default_backend: str | list[str] = \"console\", config=None):\n        if isinstance(default_backend, str):\n            default_backend = [default_backend]\n        for backend in default_backend:\n            if backend == \"tracking\":\n                import warnings\n\n                warnings.warn(\"`tracking` logger is deprecated. use `wandb` instead.\", DeprecationWarning, stacklevel=2)\n            else:\n                assert backend in self.supported_backend, f\"{backend} is not supported\"\n\n        self.logger = {}\n\n        if \"tracking\" in default_backend or \"wandb\" in default_backend:\n            import os\n\n            import wandb\n\n            settings = None\n            if config and config[\"trainer\"].get(\"wandb_proxy\", None):\n                settings = wandb.Settings(https_proxy=config[\"trainer\"][\"wandb_proxy\"])\n            entity = os.environ.get(\"WANDB_ENTITY\", None)\n            wandb.init(project=project_name, name=experiment_name, entity=entity, config=config, settings=settings)\n            self.logger[\"wandb\"] = wandb\n\n        if \"trackio\" in default_backend:\n            import trackio\n\n            trackio.init(project=project_name, name=experiment_name, config=config)\n            self.logger[\"trackio\"] = trackio\n\n        if \"mlflow\" in default_backend:\n            import os\n            import time\n\n            import mlflow\n\n            for _mlflow_attempt in range(1, MLFLOW_MAX_ATTEMPTS + 1):\n                try:\n                    MLFLOW_TRACKING_URI = os.environ.get(\"MLFLOW_TRACKING_URI\", \"sqlite:////tmp/mlruns.db\")\n                    logger.info(\"Using MLFlow tracking URI: %s\", MLFLOW_TRACKING_URI)\n                    mlflow.set_tracking_uri(MLFLOW_TRACKING_URI)\n\n                    # Some cloud providers like Azure ML or Databricks automatically set MLFLOW_RUN_ID\n                    # If set, attach to the existing run instead of creating a new one\n                    run_id = os.environ.get(\"MLFLOW_RUN_ID\")\n                    if run_id:\n                        mlflow.start_run(run_id=run_id)\n                    else:\n                        # Project_name is actually experiment_name in MLFlow\n                        # If experiment does not exist, will create a new experiment\n                        experiment = mlflow.set_experiment(project_name)\n                        mlflow.start_run(experiment_id=experiment.experiment_id, run_name=experiment_name)\n\n                    mlflow.log_params(_compute_mlflow_params_from_objects(config))\n                    self.logger[\"mlflow\"] = _MlflowLoggingAdapter()\n                    break  # Success\n                except Exception as e:\n                    logger.warning(\n                        \"MLflow initialization attempt %d/%d failed: %s\", _mlflow_attempt, MLFLOW_MAX_ATTEMPTS, e\n                    )\n                    if _mlflow_attempt < MLFLOW_MAX_ATTEMPTS:\n                        time.sleep(MLFLOW_SLEEP_SECONDS)\n                    else:\n                        logger.warning(\"All MLflow initialization attempts failed. Proceeding without MLflow tracking.\")\n\n        if \"swanlab\" in default_backend:\n            import os\n\n            import swanlab\n\n            SWANLAB_API_KEY = os.environ.get(\"SWANLAB_API_KEY\", None)\n            SWANLAB_LOG_DIR = os.environ.get(\"SWANLAB_LOG_DIR\", \"swanlog\")\n            SWANLAB_MODE = os.environ.get(\"SWANLAB_MODE\", \"cloud\")\n            if SWANLAB_API_KEY:\n                swanlab.login(SWANLAB_API_KEY)  # NOTE: previous login information will be overwritten\n\n            if config is None:\n                config = {}  # make sure config is not None, otherwise **config will raise error\n            swanlab.init(\n                project=project_name,\n                experiment_name=experiment_name,\n                config={\"FRAMEWORK\": \"verl\", **config},\n                logdir=SWANLAB_LOG_DIR,\n                mode=SWANLAB_MODE,\n            )\n            self.logger[\"swanlab\"] = swanlab\n\n        if \"vemlp_wandb\" in default_backend:\n            import os\n\n            import volcengine_ml_platform\n            from volcengine_ml_platform import wandb as vemlp_wandb\n\n            volcengine_ml_platform.init(\n                ak=os.environ[\"VOLC_ACCESS_KEY_ID\"],\n                sk=os.environ[\"VOLC_SECRET_ACCESS_KEY\"],\n                region=os.environ[\"MLP_TRACKING_REGION\"],\n            )\n\n            vemlp_wandb.init(\n                project=project_name,\n                name=experiment_name,\n                config=config,\n                sync_tensorboard=True,\n            )\n            self.logger[\"vemlp_wandb\"] = vemlp_wandb\n\n        if \"tensorboard\" in default_backend:\n            self.logger[\"tensorboard\"] = _TensorboardAdapter(project_name, experiment_name)\n\n        if \"console\" in default_backend:\n            from verl.utils.logger import LocalLogger\n\n            self.console_logger = LocalLogger(print_to_console=True)\n            self.logger[\"console\"] = self.console_logger\n\n        if \"clearml\" in default_backend:\n            self.logger[\"clearml\"] = ClearMLLogger(project_name, experiment_name, config)\n\n        if \"file\" in default_backend:\n            self.logger[\"file\"] = FileLogger(project_name, experiment_name)\n\n    def log(self, data, step, backend=None):\n        for default_backend, logger_instance in self.logger.items():\n            if backend is None or default_backend in backend:\n                logger_instance.log(data=data, step=step)\n\n    def __del__(self):\n        if \"wandb\" in self.logger:\n            self.logger[\"wandb\"].finish(exit_code=0)\n        if \"swanlab\" in self.logger:\n            self.logger[\"swanlab\"].finish()\n        if \"vemlp_wandb\" in self.logger:\n            self.logger[\"vemlp_wandb\"].finish(exit_code=0)\n        if \"tensorboard\" in self.logger:\n            self.logger[\"tensorboard\"].finish()\n        if \"clearml\" in self.logger:\n            self.logger[\"clearml\"].finish()\n        if \"trackio\" in self.logger:\n            self.logger[\"trackio\"].finish()\n        if \"file\" in self.logger:\n            self.logger[\"file\"].finish()\n\n\nclass ClearMLLogger:\n    def __init__(self, project_name: str, experiment_name: str, config):\n        self.project_name = project_name\n        self.experiment_name = experiment_name\n\n        import clearml\n\n        self._task: clearml.Task = clearml.Task.init(\n            task_name=experiment_name,\n            project_name=project_name,\n            continue_last_task=True,\n            output_uri=False,\n        )\n\n        self._task.connect_configuration(config, name=\"Hyperparameters\")\n\n    def _get_logger(self):\n        return self._task.get_logger()\n\n    def log(self, data, step):\n        import numpy as np\n        import pandas as pd\n\n        # logs = self._rewrite_logs(data)\n        logger = self._get_logger()\n        for k, v in data.items():\n            title, series = k.split(\"/\", 1)\n\n            if isinstance(v, int | float | np.floating | np.integer):\n                logger.report_scalar(\n                    title=title,\n                    series=series,\n                    value=v,\n                    iteration=step,\n                )\n            elif isinstance(v, pd.DataFrame):\n                logger.report_table(\n                    title=title,\n                    series=series,\n                    table_plot=v,\n                    iteration=step,\n                )\n            else:\n                logger.warning(\n                    f'Trainer is attempting to log a value of \"{v}\" of type {type(v)} for key \"{k}\". This '\n                    f\"invocation of ClearML logger's function is incorrect so this attribute was dropped. \"\n                )\n\n    def finish(self):\n        self._task.close()\n\n\nclass FileLogger:\n    def __init__(self, project_name: str, experiment_name: str):\n        self.project_name = project_name\n        self.experiment_name = experiment_name\n\n        self.filepath = os.getenv(\"VERL_FILE_LOGGER_PATH\", None)\n        if self.filepath is None:\n            root_path = os.path.expanduser(os.getenv(\"VERL_FILE_LOGGER_ROOT\", \".\"))\n            directory = os.path.join(root_path, self.project_name)\n            os.makedirs(directory, exist_ok=True)\n            self.filepath = os.path.join(directory, f\"{self.experiment_name}.jsonl\")\n            print(f\"Creating file logger at {self.filepath}\")\n        self.fp = open(self.filepath, \"wb\", buffering=0)\n\n    def log(self, data, step):\n        data = {\"step\": step, \"data\": data}\n        self.fp.write(orjson.dumps(data, option=orjson.OPT_SERIALIZE_NUMPY) + b\"\\n\")\n\n    def finish(self):\n        self.fp.close()\n\n\nclass _TensorboardAdapter:\n    def __init__(self, project_name, experiment_name):\n        import os\n\n        from torch.utils.tensorboard import SummaryWriter\n\n        tensorboard_dir = os.environ.get(\"TENSORBOARD_DIR\", f\"tensorboard_log/{project_name}/{experiment_name}\")\n        os.makedirs(tensorboard_dir, exist_ok=True)\n        print(f\"Saving tensorboard log to {tensorboard_dir}.\")\n        self.writer = SummaryWriter(tensorboard_dir)\n\n    def log(self, data, step):\n        for key in data:\n            self.writer.add_scalar(key, data[key], step)\n\n    def finish(self):\n        self.writer.close()\n\n\nclass _MlflowLoggingAdapter:\n    def __init__(self):\n        import logging\n        import re\n\n        self.logger = logging.getLogger(__name__)\n        # Suppress noisy \"Found credentials from IAM Role\" on every MLflow request\n        logging.getLogger(\"botocore.credentials\").setLevel(logging.WARNING)\n        # MLflow metric key validation logic:\n        # https://github.com/mlflow/mlflow/blob/master/mlflow/utils/validation.py#L157C12-L157C44\n        # Only characters allowed: slashes, alphanumerics, underscores, periods, dashes, colons,\n        # and spaces.\n        self._invalid_chars_pattern = re.compile(\n            r\"[^/\\w.\\- :]\"\n        )  # Allowed: slashes, alphanumerics, underscores, periods, dashes, colons, and spaces.\n        self._consecutive_slashes_pattern = re.compile(r\"/+\")\n        self._sanitized_key_cache = {}\n\n    def _sanitize_key(self, key):\n        if key in self._sanitized_key_cache:\n            return self._sanitized_key_cache[key] or key\n        # First replace @ with _at_ for backward compatibility\n        sanitized = key.replace(\"@\", \"_at_\")\n        # Replace consecutive slashes with a single slash (MLflow treats them as file paths)\n        sanitized = self._consecutive_slashes_pattern.sub(\"/\", sanitized)\n        # Then replace any other invalid characters with _\n        sanitized = self._invalid_chars_pattern.sub(\"_\", sanitized)\n        if sanitized == key:\n            self._sanitized_key_cache[key] = None\n        else:\n            self.logger.warning(\"[MLflow] Metric key '%s' sanitized to '%s' due to invalid characters.\", key, sanitized)\n            self._sanitized_key_cache[key] = sanitized\n        return sanitized\n\n    def log(self, data, step):\n        import mlflow\n\n        results = {self._sanitize_key(k): v for k, v in data.items()}\n        mlflow.log_metrics(metrics=results, step=step)\n\n\ndef _compute_mlflow_params_from_objects(params) -> dict[str, Any]:\n    if params is None:\n        return {}\n\n    return _flatten_dict(_transform_params_to_json_serializable(params, convert_list_to_dict=True), sep=\"/\")\n\n\ndef _transform_params_to_json_serializable(x, convert_list_to_dict: bool):\n    _transform = partial(_transform_params_to_json_serializable, convert_list_to_dict=convert_list_to_dict)\n\n    if dataclasses.is_dataclass(x):\n        return _transform(dataclasses.asdict(x))\n    if isinstance(x, dict):\n        return {k: _transform(v) for k, v in x.items()}\n    if isinstance(x, list):\n        if convert_list_to_dict:\n            return {\"list_len\": len(x)} | {f\"{i}\": _transform(v) for i, v in enumerate(x)}\n        else:\n            return [_transform(v) for v in x]\n    if isinstance(x, Path):\n        return str(x)\n    if isinstance(x, Enum):\n        return x.value\n\n    return x\n\n\ndef _flatten_dict(raw: dict[str, Any], *, sep: str) -> dict[str, Any]:\n    import pandas as pd\n\n    ans = pd.json_normalize(raw, sep=sep).to_dict(orient=\"records\")[0]\n    assert isinstance(ans, dict)\n    return ans\n\n\n@dataclasses.dataclass\nclass ValidationGenerationsLogger:\n    project_name: str = None\n    experiment_name: str = None\n\n    def log(self, loggers, samples, step):\n        if \"wandb\" in loggers:\n            self.log_generations_to_wandb(samples, step)\n        if \"swanlab\" in loggers:\n            self.log_generations_to_swanlab(samples, step)\n        if \"mlflow\" in loggers:\n            self.log_generations_to_mlflow(samples, step)\n\n        if \"clearml\" in loggers:\n            self.log_generations_to_clearml(samples, step)\n        if \"tensorboard\" in loggers:\n            self.log_generations_to_tensorboard(samples, step)\n\n        if \"vemlp_wandb\" in loggers:\n            self.log_generations_to_vemlp_wandb(samples, step)\n\n    def log_generations_to_vemlp_wandb(self, samples, step):\n        from volcengine_ml_platform import wandb as vemlp_wandb\n\n        self._log_generations_to_wandb(samples, step, vemlp_wandb)\n\n    def log_generations_to_wandb(self, samples, step):\n        import wandb\n\n        self._log_generations_to_wandb(samples, step, wandb)\n\n    def _log_generations_to_wandb(self, samples, step, wandb):\n        \"\"\"Log samples to wandb as a table\"\"\"\n\n        # Create column names for all samples\n        columns = [\"step\"] + sum(\n            [[f\"input_{i + 1}\", f\"output_{i + 1}\", f\"score_{i + 1}\"] for i in range(len(samples))], []\n        )\n\n        if not hasattr(self, \"validation_table\"):\n            # Initialize the table on first call\n            self.validation_table = wandb.Table(columns=columns)\n\n        # Create a new table with same columns and existing data\n        # Workaround for https://github.com/wandb/wandb/issues/2981#issuecomment-1997445737\n        new_table = wandb.Table(columns=columns, data=self.validation_table.data)\n\n        # Add new row with all data\n        row_data = []\n        row_data.append(step)\n        for sample in samples:\n            row_data.extend(sample)\n\n        new_table.add_data(*row_data)\n\n        # Update reference and log\n        if wandb.run is not None:\n            wandb.log({\"val/generations\": new_table}, step=step)\n        self.validation_table = new_table\n\n    def log_generations_to_swanlab(self, samples, step):\n        \"\"\"Log samples to swanlab as text\"\"\"\n        import swanlab\n\n        swanlab_table = swanlab.echarts.Table()\n\n        # Create column names\n        headers = [\"step\", \"input\", \"output\", \"score\"]\n\n        swanlab_row_list = [[step, *sample] for sample in samples]\n        swanlab_table.add(headers=headers, rows=swanlab_row_list)\n\n        # Log to swanlab\n        swanlab.log({\"val/generations\": swanlab_table}, step=step)\n\n    def log_generations_to_mlflow(self, samples, step):\n        \"\"\"Log validation generation to mlflow as artifacts\"\"\"\n        # https://mlflow.org/docs/latest/api_reference/python_api/mlflow.html?highlight=log_artifact#mlflow.log_artifact\n\n        import tempfile\n\n        import mlflow\n\n        try:\n            with tempfile.TemporaryDirectory() as tmp_dir:\n                validation_gen_step_file = Path(tmp_dir, f\"val_step{step}.json\")\n                row_data = []\n                for sample in samples:\n                    data = {\"input\": sample[0], \"output\": sample[1], \"score\": sample[2]}\n                    row_data.append(data)\n                with open(validation_gen_step_file, \"w\") as file:\n                    json.dump(row_data, file)\n                mlflow.log_artifact(validation_gen_step_file)\n        except Exception as e:\n            print(f\"WARNING: save validation generation file to mlflow failed with error {e}\")\n\n    def log_generations_to_clearml(self, samples, step):\n        \"\"\"Log validation generation to clearml as table\"\"\"\n\n        import clearml\n        import pandas as pd\n\n        task: clearml.Task | None = clearml.Task.current_task()\n        if task is None:\n            return\n\n        table = [\n            {\n                \"step\": step,\n                \"input\": sample[0],\n                \"output\": sample[1],\n                \"score\": sample[2],\n            }\n            for sample in samples\n        ]\n\n        logger = task.get_logger()\n        logger.report_table(\n            series=\"Validation generations\",\n            title=\"Validation\",\n            table_plot=pd.DataFrame.from_records(table),\n            iteration=step,\n        )\n\n    def log_generations_to_tensorboard(self, samples, step):\n        \"\"\"Log samples to tensorboard as text\"\"\"\n        # Initialize tensorboard writer if not exists\n        if not hasattr(self, \"writer\"):\n            from torch.utils.tensorboard import SummaryWriter\n\n            # Use the same directory structure as _TensorboardAdapter\n            if self.project_name and self.experiment_name:\n                default_dir = os.path.join(\"tensorboard_log\", self.project_name, self.experiment_name)\n            else:\n                default_dir = \"tensorboard_log\"\n\n            tensorboard_dir = os.environ.get(\"TENSORBOARD_DIR\", default_dir)\n            os.makedirs(tensorboard_dir, exist_ok=True)\n            self.writer = SummaryWriter(log_dir=tensorboard_dir)\n\n        # Format the samples data into readable text\n        text_content = f\"**Generation Results - Step {step}**\\n\\n\"\n\n        for i, sample in enumerate(samples):\n            text_content += f\"### Sample {i + 1}\\n\"\n\n            # Assuming sample contains [input, output, score]\n            if len(sample) >= 3:\n                input_text, output_text, score = sample[0], sample[1], sample[2]\n\n                text_content += f\"**Input:** {input_text}\\n\\n\"\n                text_content += f\"**Output:** {output_text}\\n\\n\"\n                text_content += f\"**Score:** {score}\\n\\n\"\n            else:\n                # Handle cases where sample format might be different\n                text_content += f\"**Data:** {sample}\\n\\n\"\n\n            text_content += \"---\\n\\n\"\n\n        # Log to tensorboard as text\n        self.writer.add_text(\"val/generations\", text_content, step)\n        # Flush to ensure data is written\n        self.writer.flush()\n"
  },
  {
    "path": "verl/utils/transformers_compat.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"\nCompatibility utilities for different versions of transformers library.\n\"\"\"\n\nimport importlib.metadata\nfrom functools import lru_cache\nfrom typing import Optional\n\nfrom packaging import version\n\n# Handle version compatibility for flash_attn_supports_top_left_mask\n# This function was added in newer versions of transformers\ntry:\n    from transformers.modeling_flash_attention_utils import flash_attn_supports_top_left_mask\nexcept ImportError:\n    # For older versions of transformers that don't have this function\n    # Default to False as a safe fallback for older versions\n    def flash_attn_supports_top_left_mask():\n        \"\"\"Fallback implementation for older transformers versions.\n        Returns False to disable features that require this function.\n        \"\"\"\n        return False\n\n\n@lru_cache\ndef is_transformers_version_in_range(min_version: Optional[str] = None, max_version: Optional[str] = None) -> bool:\n    try:\n        # Get the installed version of the transformers library\n        transformers_version_str = importlib.metadata.version(\"transformers\")\n    except importlib.metadata.PackageNotFoundError as e:\n        raise ModuleNotFoundError(\"The `transformers` package is not installed.\") from e\n\n    transformers_version = version.parse(transformers_version_str)\n\n    lower_bound_check = True\n    if min_version is not None:\n        lower_bound_check = version.parse(min_version) <= transformers_version\n\n    upper_bound_check = True\n    if max_version is not None:\n        upper_bound_check = transformers_version <= version.parse(max_version)\n\n    return lower_bound_check and upper_bound_check\n\n\n@lru_cache\ndef get_auto_model_for_vision2seq():\n    \"\"\"Return the available VL auto model class across transformers versions.\"\"\"\n\n    try:\n        # Prefer the newer class when available. In transformers 4.x this class has\n        # a broader mapping than AutoModelForVision2Seq, and AutoModelForVision2Seq\n        # is deprecated for removal in v5.\n        from transformers import AutoModelForImageTextToText\n    except ImportError:\n        from transformers import AutoModelForVision2Seq\n\n        return AutoModelForVision2Seq\n\n    return AutoModelForImageTextToText\n"
  },
  {
    "path": "verl/utils/trtllm/trtllm_fp8_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom verl.utils.fp8_utils import FP8QuantizerHelper\n\n\nclass TRTLLMFP8QuantizerHelper(FP8QuantizerHelper):\n    def __init__(self, quant_config):\n        super().__init__(quant_config)\n"
  },
  {
    "path": "verl/utils/ulysses.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nUtilities for DeepSpeed Ulysses Sequence Parallelism.\nDeepSpeed Ulysses Paper: https://arxiv.org/abs/2309.14509\nInspired from: https://github.com/deepspeedai/DeepSpeed/blob/master/deepspeed/sequence/layer.py\n\"\"\"\n\nfrom typing import Any, Optional\n\nimport torch\nimport torch.distributed as dist\nfrom torch import Tensor\nfrom torch.distributed import ProcessGroup\n\n_ULYSSES_SEQUENCE_PARALLEL_GROUP = None\n\n\ndef set_ulysses_sequence_parallel_group(group: dist.ProcessGroup):\n    \"\"\"\n    Set ulysses sequence parallel process group.\n    \"\"\"\n    global _ULYSSES_SEQUENCE_PARALLEL_GROUP\n    _ULYSSES_SEQUENCE_PARALLEL_GROUP = group\n\n\ndef get_ulysses_sequence_parallel_group() -> Optional[dist.ProcessGroup]:\n    \"\"\"\n    Get ulysses sequence parallel process group.\n    \"\"\"\n    global _ULYSSES_SEQUENCE_PARALLEL_GROUP\n    return _ULYSSES_SEQUENCE_PARALLEL_GROUP\n\n\ndef get_ulysses_sequence_parallel_world_size(group: ProcessGroup = None) -> int:\n    \"\"\"\n    Get ulysses sequence parallel world size.\n    \"\"\"\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    return dist.get_world_size(group) if group else 1\n\n\ndef get_ulysses_sequence_parallel_rank(group: ProcessGroup = None) -> int:\n    \"\"\"\n    Get ulysses sequence parallel rank.\n    \"\"\"\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    return dist.get_rank(group) if group else 0\n\n\ndef gather_seq_scatter_heads(\n    x: Tensor,\n    seq_dim: int,\n    head_dim: int,\n    unpadded_dim_size: int = 0,\n    group: ProcessGroup = None,\n) -> Tensor:\n    \"\"\"\n    A func to sync embedding input with alltoall in sequence parallel\n    gather sequence dimension and scatter head dim:\n    e.g. seq_dim: 1, head_dim: 2\n    [bsz, seq/n, h, ...] -> [bsz, seq, h/n, ...]\n    \"\"\"\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    if not group:\n        return x\n    sp_world = get_ulysses_sequence_parallel_world_size(group)\n    x = SeqAllToAll.apply(group, x, head_dim, seq_dim)\n    if unpadded_dim_size and unpadded_dim_size % sp_world != 0:\n        padding_size = x.size(seq_dim) - unpadded_dim_size\n        x = _unpad_tensor(x, seq_dim, padding_size)\n    return x\n\n\ndef gather_heads_scatter_seq(x: Tensor, head_dim: int, seq_dim: int, group: ProcessGroup = None) -> Tensor:\n    \"\"\"\n    A func to sync attention result with alltoall in sequence parallel\n    gather head dimension and scatter seq dim:\n    e.g. seq_dim: 1, head_dim: 2\n    [bsz, seq, h/n, ...] -> [bsz, seq/n, h, ...]\n    \"\"\"\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    if not group:\n        return x\n    dim_size = x.size(seq_dim)\n    sp_world = get_ulysses_sequence_parallel_world_size(group)\n    if dim_size % sp_world != 0:\n        padding_size = sp_world - (dim_size % sp_world)\n        x = _pad_tensor(x, seq_dim, padding_size)\n    return SeqAllToAll.apply(group, x, seq_dim, head_dim, False)\n\n\ndef _pad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor:\n    shape = list(x.shape)\n    shape[dim] = padding_size\n    pad = torch.zeros(shape, dtype=x.dtype, device=x.device)\n    return torch.cat([x, pad], dim=dim)\n\n\ndef _unpad_tensor(x: Tensor, dim: int, padding_size: int) -> Tensor:\n    slc = [slice(None)] * len(x.shape)\n    slc[dim] = slice(0, -padding_size)\n    return x[tuple(slc)]\n\n\ndef slice_input_tensor(x: Tensor, dim: int, padding: bool = True, group: ProcessGroup = None) -> Tensor:\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    sp_world_size = dist.get_world_size(group)\n    sp_rank = get_ulysses_sequence_parallel_rank()\n    dim_size = x.size(dim)\n    # pad before slice\n    if padding and dim_size % sp_world_size:\n        padding_size = sp_world_size - (dim_size % sp_world_size)\n        x = _pad_tensor(x, dim, padding_size)\n    # slice the input tensor\n    parts = x.size(dim) // sp_world_size\n    slc = [slice(None)] * len(x.shape)\n    slc[dim] = slice(sp_rank * parts, (sp_rank + 1) * parts)\n    return x[tuple(slc)].contiguous()\n\n\ndef all_to_all_tensor(\n    local_input: Tensor,\n    scatter_dim: int,\n    gather_dim: int,\n    group: Optional[dist.ProcessGroup] = None,\n    async_op: bool = False,\n):\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    seq_world_size = dist.get_world_size(group)\n    input_list = [t.contiguous() for t in torch.tensor_split(local_input, seq_world_size, scatter_dim)]\n    output_list = [torch.empty_like(input_list[0]) for _ in range(seq_world_size)]\n    comm = dist.all_to_all(output_list, input_list, group=group, async_op=async_op)\n    if async_op:\n\n        def wait():\n            comm.wait()\n            return torch.cat(output_list, dim=gather_dim).contiguous()\n\n        return wait\n    return torch.cat(output_list, dim=gather_dim).contiguous()\n\n\ndef all_gather_tensor(local_tensor: Tensor, group: Optional[dist.ProcessGroup] = None, async_op: bool = False):\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    sp_world_size = dist.get_world_size(group=group)\n    output_shape = list(local_tensor.shape)\n    output_shape[0] = output_shape[0] * sp_world_size\n    output = torch.empty(output_shape, dtype=local_tensor.dtype, device=local_tensor.device)\n    dist.all_gather_into_tensor(output, local_tensor, group=group, async_op=async_op)\n    return output\n\n\nclass SeqAllToAll(torch.autograd.Function):\n    @staticmethod\n    def forward(\n        ctx: Any,\n        group: dist.ProcessGroup,\n        local_input: Tensor,\n        scatter_dim: int,\n        gather_dim: int,\n        async_op: bool = False,\n    ) -> Tensor:\n        ctx.group = group\n        ctx.scatter_dim = scatter_dim\n        ctx.gather_dim = gather_dim\n        ctx.async_op = async_op\n        return all_to_all_tensor(local_input, scatter_dim, gather_dim, group, async_op)\n\n    @staticmethod\n    def backward(ctx: Any, *grad_output: Tensor) -> tuple[None, Tensor, None, None]:\n        input_t = torch.cat(grad_output[1:], dim=ctx.gather_dim).contiguous() if ctx.async_op else grad_output[0]\n        return (\n            None,\n            all_to_all_tensor(input_t, ctx.gather_dim, ctx.scatter_dim, ctx.group, False),\n            None,\n            None,\n            None,\n            None,\n        )\n\n\nclass Gather(torch.autograd.Function):\n    @staticmethod\n    def forward(\n        ctx: Any,\n        group: dist.ProcessGroup,\n        local_tensor: Tensor,\n        gather_dim: int,\n        grad_scaler: bool = True,\n        async_op=False,\n    ) -> Tensor:\n        ctx.group = group\n        ctx.gather_dim = gather_dim\n        ctx.grad_scaler = grad_scaler\n        ctx.async_op = async_op\n\n        sp_world_size = dist.get_world_size(group=group)\n        ctx.sp_world_size = sp_world_size\n\n        sp_rank = dist.get_rank(group=group)\n        ctx.sp_rank = sp_rank\n\n        local_shape = list(local_tensor.size())\n        split_size = local_shape[0]\n        part_size = local_shape[gather_dim]  # store original size\n        ctx.part_size = part_size\n\n        output = all_gather_tensor(local_tensor, group, async_op)\n        return torch.cat(output.split(split_size, dim=0), dim=gather_dim)\n\n    @staticmethod\n    def backward(ctx: Any, grad_output: Tensor) -> Any:\n        if ctx.grad_scaler:\n            grad_output = grad_output * ctx.sp_world_size\n        return (\n            None,\n            grad_output.split(ctx.part_size, dim=ctx.gather_dim)[ctx.sp_rank].contiguous(),\n            None,\n            None,\n            None,\n            None,\n        )\n\n\ndef gather_outpus_and_unpad(*args, **kwargs):\n    raise RuntimeError(\n        \"please use verl.utils.ulysses.gather_outputs_and_unpad instead of verl.utils.ulysses.gather_outpus_and_unpad\"\n    )\n\n\ndef gather_outputs_and_unpad(\n    x: Tensor,\n    gather_dim: int,\n    unpad_dim: int = None,\n    padding_size: int = 0,\n    grad_scaler: bool = True,\n    group: Optional[dist.ProcessGroup] = None,\n):\n    \"\"\"\n    Gather a tensor across a process group and optionally unpad its padded elements.\n\n    Args:\n        x (Tensor): Input tensor to gather.\n        gather_dim (int): Dimension along which to gather across ranks.\n        unpad_dim (int, optional): Dimension from which to remove padding. If None, no unpadding.\n        padding_size (int): Number of padding elements to remove on `unpad_dim`. Defaults to 0.\n        grad_scaler (bool): Whether to apply gradient scaling during gather. Defaults to True.\n        group (ProcessGroup, optional): Process group for gathering. If None, uses\n            `get_ulysses_sequence_parallel_group()`. If still None, returns `x` unchanged.\n\n    Returns:\n        Tensor: The gathered tensor, with padding removed if requested.\n    \"\"\"\n    group = get_ulysses_sequence_parallel_group() if group is None else group\n    if group is None:\n        return x\n    x = Gather.apply(group, x, gather_dim, grad_scaler)\n    if unpad_dim is not None:\n        assert isinstance(padding_size, int), \"padding size is not given or is not an integer\"\n        if padding_size == 0:\n            return x\n        x = _unpad_tensor(x, unpad_dim, padding_size)\n    return x\n\n\ndef ulysses_pad(\n    input_ids_rmpad: torch.Tensor, position_ids_rmpad: Optional[torch.Tensor] = None, sp_size: int = 1, pad_value=0\n):\n    if position_ids_rmpad is not None:\n        assert position_ids_rmpad.size(-2) == 1\n        assert input_ids_rmpad.size(-1) == position_ids_rmpad.size(-1)\n    if sp_size <= 1:\n        return input_ids_rmpad, position_ids_rmpad, 0\n    _, total_seq_len = input_ids_rmpad.shape\n    pad_size = (sp_size - total_seq_len % sp_size) % sp_size\n    if pad_size > 0:\n        input_ids_rmpad = torch.nn.functional.pad(input_ids_rmpad, (0, pad_size), value=pad_value)\n        if position_ids_rmpad is not None:\n            pad_pos_ids = torch.arange(pad_size, device=position_ids_rmpad.device).unsqueeze(0)\n            if position_ids_rmpad.dim() == 3:\n                pad_pos_ids = pad_pos_ids.unsqueeze(0).repeat(position_ids_rmpad.size(0), 1, 1)\n            position_ids_rmpad = torch.cat((position_ids_rmpad, pad_pos_ids), dim=-1)\n    return input_ids_rmpad, position_ids_rmpad, pad_size\n\n\ndef ulysses_pad_and_slice_inputs(\n    input_ids_rmpad: torch.Tensor,\n    position_ids_rmpad: Optional[torch.Tensor] = None,\n    sp_size: int = 1,\n    skip_position_ids_rmpad: bool = False,\n    pad_value=0,\n):\n    \"\"\"\n    Pad and slice input_ids to be divisible by sp_size\n    Pad position_ids to be divisible by sp_size.\n\n    Note both input_ids_rmpad and position_ids_rmpad will be padded and sliced.\n\n    The is the utility of pre-forward for ulysses sequence parallelism\n\n    Args:\n        input_ids_rmpad: shape of [bsz, seqlen]\n        position_ids_rmpad: shape of [bsz, seqlen], where bsz must be 1\n        sp_size (int): ulysses sequence parallelism size\n        skip_position_ids_rmpad: whether to skip position_ids_rmpad for VeOmniEngine\n\n    Returns:\n        torch.Tensor: padded and sliced input_ids\n        torch.Tensor: padded and sliced position_ids\n        int: pad size\n    \"\"\"\n    input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad(\n        input_ids_rmpad, position_ids_rmpad, sp_size, pad_value=pad_value\n    )\n    input_ids_rmpad = slice_input_tensor(input_ids_rmpad, dim=1, padding=False)\n    if position_ids_rmpad is not None and not skip_position_ids_rmpad:\n        position_ids_rmpad = slice_input_tensor(position_ids_rmpad, dim=1, padding=False)\n    return input_ids_rmpad, position_ids_rmpad, pad_size\n\n\ndef validate_ulysses_config(num_heads, ulysses_sequence_size):\n    if ulysses_sequence_size > 1:\n        assert num_heads % ulysses_sequence_size == 0, (\n            f\"num_heads ({num_heads}) must be divisible by ulysses sequence size({ulysses_sequence_size})\"\n        )\n"
  },
  {
    "path": "verl/utils/vllm/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .npu_vllm_patch import check_vllm_ascend_before_server_launch\nfrom .utils import TensorLoRARequest, VLLMHijack, is_version_ge\n\n# The contents of vllm/patch.py should not be imported here, because the contents of\n# patch.py should be imported after the vllm LLM instance is created. Therefore,\n# wait until you actually start using it before importing the contents of\n# patch.py separately.\n\n__all__ = [\n    \"TensorLoRARequest\",\n    \"VLLMHijack\",\n    \"is_version_ge\",\n    \"check_vllm_ascend_before_server_launch\",\n]\n"
  },
  {
    "path": "verl/utils/vllm/npu_vllm_patch.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Copyright 2025 The Qwen Team and The HuggingFace Inc. team\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport os\nfrom functools import wraps\n\nfrom verl.utils.device import is_torch_npu_available\n\n\ndef vllm_ascend_v011_select_moe_comm_method_wrapper(fn):\n    @wraps(fn)\n    def wrapper(self, num_tokens, with_prefill):\n        moe_comm_method = fn(self, num_tokens, with_prefill)\n        from vllm_ascend.ascend_forward_context import MoECommType\n        from vllm_ascend.utils import AscendSocVersion, enable_sp, get_ascend_soc_version\n\n        soc_version = get_ascend_soc_version()\n\n        # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now.\n        if soc_version in {AscendSocVersion.A2} and moe_comm_method == MoECommType.MC2:\n            quant_type = getattr(self.vllm_config.model_config.hf_config, \"moe_quantize\", None)\n            # Currently, w4a8_dynamic does not support allgatherep\n            if quant_type == \"w4a8_dynamic\":\n                moe_comm_method = MoECommType.ALLTOALL\n            else:\n                moe_comm_method = MoECommType.ALLGATHER\n\n        if with_prefill:\n            if enable_sp():\n                moe_comm_method = MoECommType.ALLGATHER\n            else:\n                moe_comm_method = MoECommType.NAIVE_MULTICAST\n\n        return moe_comm_method\n\n    return wrapper\n\n\ndef vllm_ascend_v011_matmul_and_reduce_wrapper(fn):\n    @wraps(fn)\n    def wrapper(self, *args, **kwargs):\n        from vllm_ascend.utils import AscendSocVersion, get_ascend_soc_version\n\n        soc_version = get_ascend_soc_version()\n        # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now.\n        if soc_version in {AscendSocVersion.A2}:\n            from vllm.forward_context import get_forward_context\n\n            try:\n                forward_context = get_forward_context()\n                forward_context.mmrs_fusion = False\n            except AssertionError:\n                # forward_context.mmrs_fusion will be false in matmul_and_reduce func.\n                pass\n        return fn(self, *args, **kwargs)\n\n    return wrapper\n\n\ndef check_vllm_ascend_before_server_launch():\n    import torch_npu\n    import vllm\n\n    def _is_ascend_soc_version_A2_v011_local():\n        from vllm_ascend.utils import AscendSocVersion\n\n        soc_version = torch_npu.npu.get_soc_version()\n        if 220 <= soc_version <= 225:\n            _ascend_soc_version = AscendSocVersion.A2\n        elif 250 <= soc_version <= 255:\n            _ascend_soc_version = AscendSocVersion.A3\n        else:\n            _ascend_soc_version = AscendSocVersion.UNDEFINED\n\n        return _ascend_soc_version == AscendSocVersion.A2\n\n    def _is_ascend_soc_version_A2_v013_local():\n        from vllm_ascend.utils import AscendDeviceType\n\n        soc_version = torch_npu.npu.get_soc_version()\n        if 220 <= soc_version <= 225:\n            cur_device_type = AscendDeviceType.A2\n        elif 250 <= soc_version <= 255:\n            cur_device_type = AscendDeviceType.A3\n        elif 200 <= soc_version <= 205:\n            cur_device_type = AscendDeviceType._310P\n        elif soc_version == 260:\n            cur_device_type = AscendDeviceType.A5\n        else:\n            raise RuntimeError(f\"Can not support soc_version: {soc_version}.\")\n\n        return cur_device_type == AscendDeviceType.A2\n\n    if vllm.__version__ == \"0.11.0\":\n        is_A2 = _is_ascend_soc_version_A2_v011_local()\n    elif vllm.__version__ == \"0.13.0\":\n        is_A2 = _is_ascend_soc_version_A2_v013_local()\n    else:\n        is_A2 = False\n\n    if is_A2:\n        VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE = bool(int(os.getenv(\"VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE\", \"0\")))\n        if VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE:\n            raise AssertionError(\n                \"AscendSocVersion.A2 is not support VLLM_ASCEND_ENABLE_MATMUL_ALLREDUCE \\\n                in Single-card multi-process scenario now. \"\n            )\n\n\ndef vllm_ascend_v013_select_moe_comm_method_wrapper(fn):\n    @wraps(fn)\n    def wrapper(*args, **kwargs):\n        moe_comm_method = fn(*args, **kwargs)\n        from vllm_ascend.ascend_forward_context import MoECommType\n        from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type\n\n        ascend_device_type = get_ascend_device_type()\n\n        # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now.\n        if ascend_device_type in {AscendDeviceType.A2} and moe_comm_method == MoECommType.MC2:\n            moe_comm_method = MoECommType.ALLGATHER\n\n        return moe_comm_method\n\n    return wrapper\n\n\ndef vllm_ascend_v013_matmul_and_reduce_wrapper(fn):\n    @wraps(fn)\n    def wrapper(self, *args, **kwargs):\n        from vllm_ascend.utils import AscendDeviceType, get_ascend_device_type\n\n        ascend_device_type = get_ascend_device_type()\n        # AscendSocVersion.A2 is not support MC2 in Single-card multi-process scenario now.\n        if ascend_device_type in {AscendDeviceType.A2}:\n            from vllm.forward_context import get_forward_context\n\n            try:\n                forward_context = get_forward_context()\n                forward_context.mmrs_fusion = False\n            except AssertionError:\n                # forward_context.mmrs_fusion will be false in matmul_and_reduce func.\n                pass\n        return fn(self, *args, **kwargs)\n\n    return wrapper\n\n\ndef patch_vllm013_rotary_emb():\n    from vllm.model_executor.layers.rotary_embedding.common import ApplyRotaryEmb\n    def vllm013_npu_rotary_embedding_init_impl(\n            self,\n            enforce_enable: bool = False,\n            is_neox_style: bool = True,\n            enable_fp32_compute: bool = False,\n        ) -> None:\n        super(ApplyRotaryEmb, self).__init__(enforce_enable)\n        self.is_neox_style = is_neox_style\n        self.enable_fp32_compute = enable_fp32_compute\n        self.apply_rotary_emb_flash_attn = None\n\n    ApplyRotaryEmb.__init__ = vllm013_npu_rotary_embedding_init_impl\n\n\nif is_torch_npu_available(check_device=False):\n    import vllm\n    from packaging import version\n    _VLLM_VERSION = version.parse(vllm.__version__)\n    if _VLLM_VERSION >= version.parse(\"0.13.0\"):\n        # Disable flash_attn in RotaryEmbedding (NPU) when VLLM >= 0.13\n        patch_vllm013_rotary_emb()\n\n    VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2 = bool(int(os.getenv(\"VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2\", \"1\")))\n    if VERL_NPU_ENABLE_A2_PATCH_VLLM_ASCEND_MC2:\n        # only support vllm 0.13 and 0.11 now.\n        if _VLLM_VERSION >= version.parse(\"0.13.0\"):\n            from vllm_ascend import ascend_forward_context\n            from vllm_ascend.ops.linear_op import SequenceRowParallelOp\n\n            ascend_forward_context.select_moe_comm_method = vllm_ascend_v013_select_moe_comm_method_wrapper(\n                ascend_forward_context.select_moe_comm_method\n            )\n            SequenceRowParallelOp.matmul_and_reduce = vllm_ascend_v013_matmul_and_reduce_wrapper(\n                SequenceRowParallelOp.matmul_and_reduce\n            )\n        elif _VLLM_VERSION >= version.parse(\"0.11.0\"):\n            from vllm_ascend.ops.linear_op import SequenceRowParallelOp\n            from vllm_ascend.worker.model_runner_v1 import NPUModelRunner\n\n            NPUModelRunner._select_moe_comm_method = vllm_ascend_v011_select_moe_comm_method_wrapper(\n                NPUModelRunner._select_moe_comm_method\n            )\n            SequenceRowParallelOp.matmul_and_reduce = vllm_ascend_v011_matmul_and_reduce_wrapper(\n                SequenceRowParallelOp.matmul_and_reduce\n            )\n"
  },
  {
    "path": "verl/utils/vllm/patch.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n# To support different vLLM versions, we add the model into SUPPORTED_MOE_MODELS separately to avoid triggering\n# unsupported issues.\nSUPPORTED_MOE_MODELS = []\n\ntry:\n    from vllm.model_executor.models.deepseek_v2 import DeepseekV2ForCausalLM, DeepseekV3ForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(DeepseekV2ForCausalLM)\n    SUPPORTED_MOE_MODELS.append(DeepseekV3ForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.mixtral import MixtralForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(MixtralForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.qwen2_moe import Qwen2MoeForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(Qwen2MoeForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.qwen3_moe import Qwen3MoeForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(Qwen3MoeForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.qwen3_vl_moe import Qwen3MoeLLMForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(Qwen3MoeLLMForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.qwen3_next import Qwen3NextForCausalLM\n\n    SUPPORTED_MOE_MODELS.append(Qwen3NextForCausalLM)\nexcept ImportError:\n    pass\n\ntry:\n    from vllm.model_executor.models.kimi_vl import KimiVLForConditionalGeneration\n\n    SUPPORTED_MOE_MODELS.append(KimiVLForConditionalGeneration)\nexcept ImportError:\n    pass\n\n\ndef patch_vllm_moe_model_weight_loader(model):\n    # this is a work around to load the weight of vllm fused moe model\n    # it is from a bug from vllm 0.8.2\n    # all the weights are supposed to have a weight_loader, but the moe weights\n    # do not have a weight_loader, so we need to patch it\n    # (True, 'model.embed_tokens.weight')\n    # (True, 'model.layers.0.self_attn.qkv_proj.weight')\n    # (True, 'model.layers.0.self_attn.qkv_proj.bias')\n    # (True, 'model.layers.0.self_attn.o_proj.weight')\n    # (True, 'model.layers.0.mlp.gate.weight')\n    # (True, 'model.layers.0.mlp.shared_expert.gate_up_proj.weight')\n    # (True, 'model.layers.0.mlp.shared_expert.down_proj.weight')\n    # (False, 'model.layers.0.mlp.shared_expert_gate.weight')   use default\n    # (False, 'model.layers.0.input_layernorm.weight')          use default\n    # (False, 'model.layers.0.post_attention_layernorm.weight') use default\n    # (False, 'model.layers.0.mlp.experts.w13_weight')          use mlp.experts.weight_loader\n    # (False, 'model.layers.0.mlp.experts.w2_weight')          use mlp.experts.weight_loader\n\n    # Early return if no MOE models are supported\n    if not SUPPORTED_MOE_MODELS:\n        return\n\n    original_model_type = type(model)\n    if hasattr(model, \"runnable\") and \"ACLGraphWrapper\" in str(original_model_type):\n        model = model.runnable\n        original_model_type = type(model)\n\n    # Define MLP attribute mapping for different model types\n    MLP_ATTR_MAPPING = {}\n    try:\n        from vllm.model_executor.models.mixtral import MixtralForCausalLM\n\n        MLP_ATTR_MAPPING[MixtralForCausalLM] = \"block_sparse_moe\"\n    except ImportError:\n        pass\n\n    DEFAULT_MLP_ATTR = \"mlp\"\n\n    # Get inner model (either model.model or model.language_model)\n    inner_model = getattr(model, \"model\", None) or getattr(model, \"language_model\", None)\n    if inner_model is None:\n        raise ValueError(\"The provided model does not have a valid 'model' or 'language_model' attribute.\")\n\n    if not isinstance(model, tuple(SUPPORTED_MOE_MODELS)) and not isinstance(inner_model, tuple(SUPPORTED_MOE_MODELS)):\n        return\n\n    # TODO(@leisuzz): class Qwen3MoeLLMForCausalLM is not available if VLLM version < 0.11.0,\n    # will update the 'if statement' with 'isinstance' when verl commonly use VLLM version >= 0.11.0\n    if type(inner_model).__name__ == \"Qwen3MoeLLMForCausalLM\":\n        inner_model = inner_model.model  # Reassign inner_model in Qwen3-vl\n\n    for layer_idx, layer in enumerate(inner_model.layers):\n        mlp_attr = MLP_ATTR_MAPPING.get(original_model_type, DEFAULT_MLP_ATTR)\n\n        mlp = getattr(layer, mlp_attr, None)\n        if not mlp:\n            continue\n\n        experts = getattr(mlp, \"experts\", None)\n        if not experts or not hasattr(experts, \"weight_loader\"):\n            continue\n\n        # Patch the weight loaders\n        for name, param in mlp.named_parameters():\n            if \"w13_weight\" in name or \"w2_weight\" in name:\n                param.weight_loader = experts.weight_loader\n"
  },
  {
    "path": "verl/utils/vllm/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom msgspec import field\nfrom packaging import version as vs\n\ntry:\n    from vllm.lora.lora_model import LoRAModel\nexcept ImportError:\n    from vllm.lora.models import LoRAModel\n\nfrom vllm.lora.request import LoRARequest\nfrom vllm.lora.utils import get_adapter_absolute_path\nfrom vllm.lora.worker_manager import LRUCacheWorkerLoRAManager\n\nfrom verl.third_party.vllm import get_version\n\n\nclass TensorLoRARequest(LoRARequest):\n    peft_config: dict = field(default=None)\n    lora_tensors: dict = field(default=None)\n\n\nclass VLLMHijack:\n    @staticmethod\n    def hijack():\n        def hijack__load_adapter(self, lora_request: TensorLoRARequest) -> LoRAModel:\n            \"\"\"\n            based on vllm.lora.worker_manager.WorkerLoRAManager._load_adapter, support load adapter with lora tensors\n\n            Reason:\n            VLLM does not support adding LoRA from tensors directly. It only supports adding LoRA via file paths.\n            To synchronize the LoRA tensors of the actor model, we need to find a workaround to enable VLLM to\n            load memory-based LoRA tensors.\n            \"\"\"\n            try:\n                supported_lora_modules = self._adapter_manager.supported_lora_modules\n                packed_modules_mapping = self._adapter_manager.packed_modules_mapping\n                expected_lora_modules: list[str] = []\n                for module in supported_lora_modules:\n                    if module in packed_modules_mapping:\n                        expected_lora_modules.extend(packed_modules_mapping[module])\n                    else:\n                        expected_lora_modules.append(module)\n\n                expected_lora_modules = list(set(expected_lora_modules))\n\n                lora_tensors = None\n                from vllm.lora.peft_helper import PEFTHelper\n\n                if isinstance(lora_request, TensorLoRARequest):\n                    peft_config = lora_request.peft_config\n                    lora_tensors = lora_request.lora_tensors\n                    peft_helper = PEFTHelper.from_dict(peft_config)\n                else:\n                    lora_path = get_adapter_absolute_path(lora_request.lora_path)\n\n                    peft_helper = PEFTHelper.from_local_dir(lora_path, self.max_position_embeddings)\n\n                # Validates the LoRA configuration against requirements before\n                # loading weights, throwing an exception if validation fails.\n                peft_helper.validate_legal(self.lora_config)\n\n                # For some models like Qwen2VL, we need to use hf_to_vllm_mapper\n                # to ensure correct loading of lora weights.\n                model = self._adapter_manager.model\n                hf_to_vllm_mapper = None\n                if hasattr(model, \"hf_to_vllm_mapper\") and model.hf_to_vllm_mapper is not None:\n                    hf_to_vllm_mapper = model.hf_to_vllm_mapper\n\n                lora_request_kwargs = {\n                    \"peft_helper\": peft_helper,\n                    \"lora_model_id\": lora_request.lora_int_id,\n                    \"device\": \"cpu\",\n                    \"dtype\": self.lora_config.lora_dtype,\n                    \"weights_mapper\": hf_to_vllm_mapper,\n                }\n                if hasattr(self, \"embedding_padding_modules\"):\n                    lora_request_kwargs[\"embedding_modules\"] = self.embedding_modules\n                    lora_request_kwargs[\"embedding_padding_modules\"] = self.embedding_padding_modules\n                else:\n                    lora_request_kwargs[\"model_vocab_size\"] = self.vocab_size\n                if hasattr(self.lora_config, \"lora_extra_vocab_size\"):\n                    lora_request_kwargs[\"target_embedding_padding\"] = (\n                        self.vocab_size + self.lora_config.lora_extra_vocab_size\n                    )\n                if isinstance(lora_request, TensorLoRARequest):\n                    lora = self._lora_model_cls.from_lora_tensors(\n                        tensors=lora_tensors,\n                        **lora_request_kwargs,\n                    )\n                else:\n                    lora = self._lora_model_cls.from_local_checkpoint(\n                        lora_path,\n                        expected_lora_modules,\n                        **lora_request_kwargs,\n                    )\n            except Exception:\n                raise\n\n            if getattr(lora, \"extra_vocab_size\", 0) > getattr(self.lora_config, \"lora_extra_vocab_size\", 0):\n                raise ValueError(\n                    f\"LoRA added vocab size {lora.extra_vocab_size} is greater than lora_extra_vocab_size \"\n                    f\"{self.lora_config.lora_extra_vocab_size}.\"\n                )\n            return lora\n\n        def do_hijack(target_cls, target_method_name, hooking_method):\n            setattr(target_cls, target_method_name, hooking_method)\n\n        do_hijack(LRUCacheWorkerLoRAManager, \"_load_adapter\", hijack__load_adapter)\n\n\ndef is_version_ge(pkg: str = \"vllm\", minver: str = \"0.7.3\"):\n    \"\"\"check if the package version is greater than or equal to the minimum version\"\"\"\n    return vs.parse(get_version(pkg)) >= vs.parse(minver)\n"
  },
  {
    "path": "verl/utils/vllm/vllm_fp8_utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nfrom dataclasses import dataclass, field\nfrom unittest.mock import patch\n\nimport torch\nimport vllm\nfrom packaging import version\n\ntry:\n    from vllm.model_executor.layers.fused_moe.layer import FusedMoE\n    from vllm.model_executor.layers.linear import LinearBase\nexcept ImportError as e:\n    raise ImportError(\"FP8 quantization not available\") from e\n\nfrom verl.utils.kernel.fp8_kernel import scaled_fp8_blockwise\n\nlogger = logging.getLogger(__name__)\n\n\n# Ref: https://github.com/NVIDIA-NeMo/RL/commit/bc24887c72a6e1b2699a228bc87c588546dfe6b7\n@dataclass()\nclass FP8State:\n    # A cache of fp8 parameter names, we can check this cache to see if a\n    # param name corresponds to a fp8 weight\n    seen_params: set = field(default_factory=lambda: set())\n    fp8_param_names: set = field(default_factory=lambda: set())\n    vllm_patches: list = field(default_factory=lambda: [])\n\n\nfp8_state: FP8State = FP8State()\n\n\ndef is_fp8_model(vllm_config):\n    from vllm.model_executor.layers.quantization.fp8 import Fp8Config\n\n    if hasattr(vllm_config, \"quant_config\") and isinstance(vllm_config.quant_config, Fp8Config):\n        return True\n\n    return False\n\n\ndef get_module_from_param_name(model, name: str):\n    # Split the name into parts (e.g., 'layers', '0', 'self_attn', 'q_proj', 'weight')\n    # The module path is all but the last part (the parameter's own name)\n    path_parts = name.split(\".\")\n    module_path = path_parts[:-1]\n    # Replace with the fused model name\n    packed_modules_mapping = model.packed_modules_mapping\n    reversed_mapping = {\n        original_name: fused_name\n        for fused_name, original_names_list in packed_modules_mapping.items()\n        for original_name in original_names_list\n    }\n    if module_path[-1] in reversed_mapping.keys():\n        module_path[-1] = reversed_mapping[module_path[-1]]\n\n    current_module = model\n    try:\n        # Traverse the model hierarchy\n        for part in module_path:\n            if isinstance(current_module, FusedMoE):\n                return current_module\n            elif isinstance(current_module, torch.nn.ModuleList):\n                current_module = current_module[int(part)]\n            else:\n                current_module = getattr(current_module, part)\n    except (AttributeError, IndexError, ValueError) as e:\n        print(f\"Warning: Could not find module for parameter '{name}'. Error: {e}\")\n    return current_module\n\n\ndef is_fp8_weight(name, model):\n    if name not in fp8_state.seen_params:\n        fp8_state.seen_params.add(name)\n        # Filter out bias params\n        if name.endswith(\"weight\"):\n            module = get_module_from_param_name(model, name)\n            # We currently only quantize linear layers\n\n            if (isinstance(module, LinearBase) and module.weight.dtype == torch.float8_e4m3fn) or (\n                isinstance(module, FusedMoE)\n                and module.w13_weight.dtype == torch.float8_e4m3fn\n                and module.w2_weight.dtype == torch.float8_e4m3fn\n            ):\n                fp8_state.fp8_param_names.add(name)\n    return name in fp8_state.fp8_param_names\n\n\ndef quant_weights(weights, model, quant_config, dtype=torch.bfloat16):\n    \"\"\"Quantize weights to FP8 format using a memory-efficient generator.\n\n\n    Args:\n        weights: Generator or iterable of (name, tensor) pairs\n        model: The model to check for FP8 weight names\n        quant_config: Quantization configuration with weight_block_size\n        dtype: Data type for intermediate computation (default: bfloat16)\n\n    Yields:\n        Tuples of (name, tensor) for each weight and its scale\n    \"\"\"\n    if quant_config.weight_block_size is None:\n        raise ValueError(\"Currently only support blockwise quantization, please set weight_block_size in quant_config\")\n\n    # vLLM v0.11-v0.12 renamed weight_scale_inv → weight_scale in process_weights_after_loading,\n    # so load_weights expects \"_scale\" suffix. v0.14+ keeps weight_scale_inv, so expects \"_scale_inv\".\n    _use_scale_not_scale_inv = version.parse(\"0.11.0\") <= version.parse(vllm.__version__) < version.parse(\"0.14.0\")\n\n    for k, v in weights:\n        if not is_fp8_weight(k, model):\n            yield (k, v)\n            continue\n\n        # Cast the weight into fp8 and its scale factor\n        if torch.distributed.get_rank() == 0:\n            logger.debug(f\"Quantizing to FP8 blockwise: {k}\")\n\n        param_lp, param_scale = scaled_fp8_blockwise(\n            v.to(dtype),\n            weight_block_size=quant_config.weight_block_size,\n        )\n        param_scale = param_scale.squeeze(-1)\n\n        # Yield the quantized weight\n        yield (k, param_lp)\n\n        # Yield the scale with appropriate naming based on vLLM version\n        if _use_scale_not_scale_inv and \"expert\" not in k:\n            yield (k + \"_scale\", param_scale)\n        else:\n            yield (k + \"_scale_inv\", param_scale)\n\n        # Explicitly delete original tensor reference to help GC\n        del v, param_lp, param_scale\n\n\ndef load_quanted_weights(weights, model_runner):\n    model = model_runner.model\n    quant_config = model_runner.vllm_config.quant_config\n    vllm_dtype = model_runner.vllm_config.model_config.dtype\n\n    weights_quantized = quant_weights(weights, model, quant_config, dtype=vllm_dtype)\n\n    # Monkey patch the param class to their subclass, as certain models\n    # will check the param type to call the proper weightloader\n    for name, param in model.named_parameters():\n        if hasattr(param, \"subclass_type\"):\n            param.orig_type = param.__class__\n            param.__class__ = param.subclass_type\n    # Finally load the weights into vllm\n    loaded_params = model.load_weights(weights_quantized)\n    # Undo the type change above to the original type\n    for name, param in model.named_parameters():\n        if hasattr(param, \"subclass_type\"):\n            param.__class__ = param.orig_type\n    return loaded_params\n\n\ndef process_weights_after_loading_for_vllm10(self, layer) -> None:\n    \"\"\"This function is used to process the weights after loading for a Linear layer, it is used for vllm v0.10\n\n    Compared to the original process_weights_after_loading in vllm, we just avoid creation of\n    new torch.nn.Parameter objects, because that removes the weight_loader attribute which we need for refit.\n    \"\"\"\n    logger.debug(\"Applying patch process_weights_after_loading\")\n    try:\n        from vllm.model_executor.parameter import (\n            BlockQuantScaleParameter,\n            ModelWeightParameter,\n        )\n    except Exception:\n        print(\"error\")\n    from torch.nn import Parameter\n\n    def _create_param_from_subclass_attributes(custom_param):\n        param = Parameter(custom_param.data, requires_grad=False)\n        base_param_dir = dir(torch.nn.Parameter)\n        custom_param_dir = dir(custom_param)\n        # Find the attributes that are unique to the custom parameter\n        custom_attributes = [\n            attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith(\"__\")\n        ]\n        # Set the custom attributes into the base parameter object\n        for attr in custom_attributes:\n            setattr(param, attr, getattr(custom_param, attr))\n\n        param.subclass_type = type(custom_param)\n        return param\n\n    assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized\n    assert self.quant_config.activation_scheme == \"dynamic\"\n    weight = layer.weight.data\n    weight_scale_inv = layer.weight_scale_inv.data\n    weight = self._maybe_pad_weight(weight)\n\n    layer.weight = _create_param_from_subclass_attributes(\n        ModelWeightParameter(\n            data=weight,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight.weight_loader,\n        )\n    )\n    layer.weight_scale_inv = _create_param_from_subclass_attributes(\n        BlockQuantScaleParameter(\n            data=weight_scale_inv,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight_scale_inv.weight_loader,\n        )\n    )\n\n\ndef process_weights_after_loading_for_vllm11(self, layer) -> None:\n    \"\"\"This function is used to process the weights after loading for a Linear layer, it is used for vllm 0.11\n\n    Compared to the original process_weights_after_loading in vllm, we just avoid creation of\n    new torch.nn.Parameter objects, because that removes the weight_loader attribute which we need for refit.\n    \"\"\"\n    from torch.nn import Parameter\n    from vllm.model_executor.layers.quantization.utils.fp8_utils import (\n        maybe_post_process_fp8_weight_block,\n        process_fp8_weight_block_strategy,\n    )\n    from vllm.model_executor.parameter import (\n        BlockQuantScaleParameter,\n        ModelWeightParameter,\n    )\n\n    assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized\n    assert self.quant_config.activation_scheme == \"dynamic\"\n\n    def _create_param_from_subclass_attributes(custom_param):\n        param = Parameter(custom_param.data, requires_grad=False)\n        base_param_dir = dir(torch.nn.Parameter)\n        custom_param_dir = dir(custom_param)\n        # Find the attributes that are unique to the custom parameter\n        custom_attributes = [\n            attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith(\"__\")\n        ]\n        # Set the custom attributes into the base parameter object\n        for attr in custom_attributes:\n            setattr(param, attr, getattr(custom_param, attr))\n\n        param.subclass_type = type(custom_param)\n        return param\n\n    weight_scale = layer.weight_scale_inv if hasattr(layer, \"weight_scale_inv\") else layer.weight_scale\n    weight, weight_scale = process_fp8_weight_block_strategy(layer.weight, weight_scale)\n\n    layer.weight = _create_param_from_subclass_attributes(\n        ModelWeightParameter(\n            data=weight.data,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight.weight_loader,\n        )\n    )\n    layer.weight_scale = _create_param_from_subclass_attributes(\n        BlockQuantScaleParameter(\n            data=weight_scale.data,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight_scale_inv.weight_loader,\n        )\n    )\n\n    del layer.weight_scale_inv\n\n    if version.parse(vllm.__version__) == version.parse(\"0.11.0\"):\n        maybe_post_process_fp8_weight_block(layer, self.cutlass_block_fp8_supported)\n    else:\n        maybe_post_process_fp8_weight_block(layer)\n\n\ndef process_weights_after_loading_for_vllm14(self, layer) -> None:\n    \"\"\"process_weights_after_loading for vLLM >= 0.14.\n\n    Starting from v0.14, vLLM keeps the scale parameter as `weight_scale_inv`\n    (instead of renaming it to `weight_scale` like v0.11-v0.12), and `apply()`\n    accesses `layer.weight_scale_inv`. We preserve `weight_loader` and\n    `subclass_type` attributes so that refit (repeated weight sync) works.\n    \"\"\"\n    from torch.nn import Parameter\n    from vllm.model_executor.layers.quantization.utils.fp8_utils import (\n        maybe_post_process_fp8_weight_block,\n        process_fp8_weight_block_strategy,\n    )\n    from vllm.model_executor.parameter import (\n        BlockQuantScaleParameter,\n        ModelWeightParameter,\n    )\n\n    assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized\n    assert self.quant_config.activation_scheme == \"dynamic\"\n\n    def _create_param_from_subclass_attributes(custom_param):\n        param = Parameter(custom_param.data, requires_grad=False)\n        base_param_dir = dir(torch.nn.Parameter)\n        custom_param_dir = dir(custom_param)\n        custom_attributes = [\n            attr for attr in custom_param_dir if attr not in base_param_dir and not attr.startswith(\"__\")\n        ]\n        for attr in custom_attributes:\n            setattr(param, attr, getattr(custom_param, attr))\n\n        param.subclass_type = type(custom_param)\n        return param\n\n    weight, weight_scale_inv = process_fp8_weight_block_strategy(layer.weight, layer.weight_scale_inv)\n\n    layer.weight = _create_param_from_subclass_attributes(\n        ModelWeightParameter(\n            data=weight.data,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight.weight_loader,\n        )\n    )\n    layer.weight_scale_inv = _create_param_from_subclass_attributes(\n        BlockQuantScaleParameter(\n            data=weight_scale_inv.data,\n            output_dim=0,\n            input_dim=1,\n            weight_loader=layer.weight_scale_inv.weight_loader,\n        )\n    )\n\n    # vLLM v0.17 removed the `else: register_parameter(\"input_scale\", None)` from\n    # create_weights() for dynamic activation, but apply() still accesses layer.input_scale.\n    # Since block_quant always uses dynamic activation, ensure the attribute exists.\n    if not hasattr(layer, \"input_scale\"):\n        layer.input_scale = None\n\n    maybe_post_process_fp8_weight_block(layer)\n\n\ndef process_weights_after_loading_moe_for_vllm10(self, layer) -> None:\n    \"\"\"This function is used to process the weights after loading for a FusedMoE layer, it is used for vllm v0.10\"\"\"\n    from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled\n    from vllm.model_executor.layers.quantization.fp8 import _is_col_major, _swap_w13_to_w31\n    from vllm.model_executor.layers.quantization.utils.fp8_utils import (\n        get_col_major_tma_aligned_tensor,\n        requant_weight_ue8m0_inplace,\n    )\n    from vllm.utils.deep_gemm import is_blackwell_deep_gemm_used\n\n    self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()\n    assert self.quant_config.activation_scheme == \"dynamic\"\n    if self.flashinfer_moe_enabled:\n        w13_weight = _swap_w13_to_w31(layer.w13_weight.data)\n        w13_weight_scale_inv = _swap_w13_to_w31(layer.w13_weight_scale_inv.data)\n        w2_weight = layer.w2_weight.data\n        w2_weight_scale_inv = layer.w2_weight_scale_inv.data\n    else:\n        w13_weight = layer.w13_weight.data\n        w13_weight_scale_inv = layer.w13_weight_scale_inv.data\n        w2_weight = layer.w2_weight\n        w2_weight_scale_inv = layer.w2_weight_scale_inv\n\n    from torch.nn import Parameter\n\n    def _create_param_from_subclass_attributes(custom_data, custom_weight):\n        param = Parameter(custom_data, requires_grad=False)\n        base_param_dir = dir(torch.nn.Parameter)\n        custom_weight_dir = dir(custom_weight)\n        # Find the attributes that are unique to the custom parameter\n        custom_attributes = [\n            attr for attr in custom_weight_dir if attr not in base_param_dir and not attr.startswith(\"__\")\n        ]\n        # Set the custom attributes into the base parameter object\n        for attr in custom_attributes:\n            setattr(param, attr, getattr(custom_weight, attr))\n\n        return param\n\n    layer.w13_weight = _create_param_from_subclass_attributes(w13_weight, layer.w13_weight)\n    layer.w13_weight_scale_inv = _create_param_from_subclass_attributes(\n        w13_weight_scale_inv, layer.w13_weight_scale_inv\n    )\n    layer.w2_weight = _create_param_from_subclass_attributes(w2_weight, layer.w2_weight)\n    layer.w2_weight_scale_inv = _create_param_from_subclass_attributes(w2_weight_scale_inv, layer.w2_weight_scale_inv)\n\n    # DeepGemm scales need to be transposed and aligned.  We try to do\n    # it ahead of time for performance reasons.\n    if self.allow_deep_gemm and not is_blackwell_deep_gemm_used():\n        # Lazy import to avoid CUDA initialization problems.\n        if _is_col_major(layer.w13_weight_scale_inv):\n            layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()\n        if _is_col_major(layer.w2_weight_scale_inv):\n            layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()\n\n    if is_blackwell_deep_gemm_used():\n        assert layer.weight_block_size is not None\n        # Re-quantise the expert weights so their scales are UE8M0.\n        block_sz = tuple(layer.weight_block_size)\n        requant_weight_ue8m0_inplace(\n            layer.w13_weight.data,\n            layer.w13_weight_scale_inv.data,\n            block_sz,\n        )\n        requant_weight_ue8m0_inplace(\n            layer.w2_weight.data,\n            layer.w2_weight_scale_inv.data,\n            block_sz,\n        )\n\n        if _is_col_major(layer.w13_weight_scale_inv):\n            layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv).contiguous()\n        if _is_col_major(layer.w2_weight_scale_inv):\n            layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv).contiguous()\n\n\ndef process_weights_after_loading_moe_for_vllm11(self, layer) -> None:\n    \"\"\"This function is used to process the weights after loading for a FusedMoE layer, it is used for vllm 0.11\"\"\"\n    from vllm.model_executor.layers.quantization.utils.flashinfer_utils import (\n        swap_w13_to_w31,\n    )\n    from vllm.model_executor.layers.quantization.utils.fp8_utils import (\n        expert_weight_is_col_major,\n        requant_weight_ue8m0_inplace,\n    )\n    from vllm.utils.deep_gemm import (\n        get_col_major_tma_aligned_tensor,\n        is_deep_gemm_e8m0_used,\n    )\n\n    try:\n        from vllm.model_executor.layers.fused_moe.rocm_aiter_fused_moe import is_rocm_aiter_moe_enabled\n\n        self.rocm_aiter_moe_enabled = is_rocm_aiter_moe_enabled()\n    except ImportError:\n        from vllm._aiter_ops import rocm_aiter_ops\n\n        self.rocm_aiter_moe_enabled = rocm_aiter_ops.is_fused_moe_enabled()\n\n    assert self.block_quant and self.quant_config.is_checkpoint_fp8_serialized\n    assert self.quant_config.activation_scheme == \"dynamic\"\n\n    if self.flashinfer_moe_backend is not None:\n        layer.w13_weight.data = swap_w13_to_w31(layer.w13_weight.data)\n        layer.w13_weight_scale_inv.data = swap_w13_to_w31(layer.w13_weight_scale_inv.data)\n\n    if self.allow_deep_gemm and not is_deep_gemm_e8m0_used():\n        if expert_weight_is_col_major(layer.w13_weight_scale_inv):\n            layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv)\n        if expert_weight_is_col_major(layer.w2_weight_scale_inv):\n            layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv)\n\n    if is_deep_gemm_e8m0_used():\n        assert layer.weight_block_size is not None\n        # Re-quantise the expert weights so their scales are UE8M0.\n        block_sz = tuple(layer.weight_block_size)\n        requant_weight_ue8m0_inplace(\n            layer.w13_weight.data,\n            layer.w13_weight_scale_inv.data,\n            block_sz,\n        )\n        requant_weight_ue8m0_inplace(\n            layer.w2_weight.data,\n            layer.w2_weight_scale_inv.data,\n            block_sz,\n        )\n\n        # Ensure column-major TMA alignment expected by DeepGEMM.\n        if expert_weight_is_col_major(layer.w13_weight_scale_inv):\n            layer.w13_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w13_weight_scale_inv)\n        if expert_weight_is_col_major(layer.w2_weight_scale_inv):\n            layer.w2_weight_scale_inv = get_col_major_tma_aligned_tensor(layer.w2_weight_scale_inv)\n\n\ndef process_weights_after_loading_moe_for_vllm14(self, layer) -> None:\n    # removed the reentrancy guard here for refit\n    from vllm.model_executor.layers.fused_moe.oracle.fp8 import (\n        convert_to_fp8_moe_kernel_format,\n        make_fp8_moe_kernel,\n    )\n\n    # Allow for accessing weights and scales in standard way.\n    w13 = layer.w13_weight\n    w2 = layer.w2_weight\n    w13_scale = getattr(layer, f\"w13_{self.weight_scale_name}\")\n    w2_scale = getattr(layer, f\"w2_{self.weight_scale_name}\")\n    w13_input_scale = layer.w13_input_scale\n    w2_input_scale = layer.w2_input_scale\n\n    # Shuffle weights to runtime format and setup kernel.\n    w13, w2, w13_scale, w2_scale = convert_to_fp8_moe_kernel_format(\n        fp8_backend=self.fp8_backend,\n        layer=layer,\n        w13=w13,\n        w2=w2,\n        w13_scale=w13_scale,\n        w2_scale=w2_scale,\n        w13_input_scale=w13_input_scale,\n        w2_input_scale=w2_input_scale,\n    )\n    from torch.nn import Parameter\n\n    def _create_param_from_subclass_attributes(custom_data, custom_weight):\n        param = Parameter(custom_data, requires_grad=False)\n        base_param_dir = dir(torch.nn.Parameter)\n        custom_weight_dir = dir(custom_weight)\n        # Find the attributes that are unique to the custom parameter\n        custom_attributes = [\n            attr for attr in custom_weight_dir if attr not in base_param_dir and not attr.startswith(\"__\")\n        ]\n        # Set the custom attributes into the base parameter object\n        for attr in custom_attributes:\n            setattr(param, attr, getattr(custom_weight, attr))\n\n        return param\n\n    # Replace parameters with updated versions. Note that this helper\n    # function ensures the replacement is compatible with RL weight reloads.\n    layer.w13_weight = _create_param_from_subclass_attributes(w13, layer.w13_weight)\n    layer.w2_weight = _create_param_from_subclass_attributes(w2, layer.w2_weight)\n    layer.w13_weight_scale_inv = _create_param_from_subclass_attributes(w13_scale, layer.w13_weight_scale_inv)\n    layer.w2_weight_scale_inv = _create_param_from_subclass_attributes(w2_scale, layer.w2_weight_scale_inv)\n\n    self.moe_quant_config = self.get_fused_moe_quant_config(layer)\n    if self.moe_quant_config:\n        assert self.experts_cls is not None\n\n        self.moe_kernel = make_fp8_moe_kernel(\n            moe_quant_config=self.moe_quant_config,\n            moe_config=self.moe,\n            fp8_backend=self.fp8_backend,\n            experts_cls=self.experts_cls,\n            routing_tables=layer._maybe_init_expert_routing_tables(),\n            shared_experts=layer.shared_experts,\n        )\n\n\ndef apply_vllm_fp8_patches():\n    logger.info(\"Applying vllm fp8 patches for blockwise quantization\")\n    vllm_ver = version.parse(vllm.__version__)\n\n    # Linear patch: v0.14+ keeps weight_scale_inv, v0.11-v0.12 renames to weight_scale\n    func1_path = \"vllm.model_executor.layers.quantization.fp8.Fp8LinearMethod.process_weights_after_loading\"\n    if vllm_ver >= version.parse(\"0.14.0\"):\n        linear_patch_fn = process_weights_after_loading_for_vllm14\n    elif vllm_ver >= version.parse(\"0.11.0\"):\n        linear_patch_fn = process_weights_after_loading_for_vllm11\n    else:\n        linear_patch_fn = process_weights_after_loading_for_vllm10\n    patcher1 = patch(func1_path, linear_patch_fn)\n    patcher1.start()\n\n    # MoE patch\n    func2_path = \"vllm.model_executor.layers.quantization.fp8.Fp8MoEMethod.process_weights_after_loading\"\n    if vllm_ver >= version.parse(\"0.14.0\"):\n        moe_patch_fn = process_weights_after_loading_moe_for_vllm14\n    elif vllm_ver >= version.parse(\"0.11.0\"):\n        moe_patch_fn = process_weights_after_loading_moe_for_vllm11\n    else:\n        moe_patch_fn = process_weights_after_loading_moe_for_vllm10\n    patcher2 = patch(func2_path, moe_patch_fn)\n    patcher2.start()\n"
  },
  {
    "path": "verl/version/version",
    "content": "0.8.0.dev\n"
  },
  {
    "path": "verl/workers/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/workers/actor/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .base import BasePPOActor\nfrom .dp_actor import DataParallelPPOActor\n\n__all__ = [\"BasePPOActor\", \"DataParallelPPOActor\"]\n"
  },
  {
    "path": "verl/workers/actor/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe base class for Actor\n\"\"\"\n\nfrom abc import ABC, abstractmethod\n\nimport torch\n\nfrom verl import DataProto\n\n__all__ = [\"BasePPOActor\"]\n\n\nclass BasePPOActor(ABC):\n    def __init__(self, config):\n        \"\"\"The base class for PPO actor\n\n        Args:\n            config (DictConfig): a config passed to the PPOActor. We expect the type to be\n                DictConfig (https://omegaconf.readthedocs.io/), but it can be any namedtuple in general.\n        \"\"\"\n        super().__init__()\n        self.config = config\n\n    @abstractmethod\n    def compute_log_prob(self, data: DataProto) -> torch.Tensor:\n        \"\"\"Compute logits given a batch of data.\n\n        Args:\n            data (DataProto): a batch of data represented by DataProto. It must contain key ```input_ids```,\n                ```attention_mask``` and ```position_ids```.\n\n        Returns:\n            DataProto: a DataProto containing the key ```log_probs```\n\n\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def update_policy(self, data: DataProto) -> dict:\n        \"\"\"Update the policy with an iterator of DataProto\n\n        Args:\n            data (DataProto): an iterator over the DataProto that returns by\n                ```make_minibatch_iterator```\n\n        Returns:\n            Dict: a dictionary contains anything. Typically, it contains the statistics during updating the model\n            such as ```loss```, ```grad_norm```, etc,.\n\n        \"\"\"\n        pass\n"
  },
  {
    "path": "verl/workers/actor/dp_actor.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nSingle Process Actor\n\"\"\"\n\nimport logging\nimport os\n\nimport torch\nfrom torch import nn\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.tensor import DTensor\n\nimport verl.utils.torch_functional as verl_F\nfrom verl import DataProto\nfrom verl.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn, kl_penalty\nfrom verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_\nfrom verl.utils.profiler import GPUMemoryLogger\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import prepare_dynamic_batch, restore_dynamic_batch\nfrom verl.utils.torch_dtypes import PrecisionType\nfrom verl.utils.torch_functional import logprobs_from_logits\nfrom verl.utils.ulysses import gather_outputs_and_unpad, ulysses_pad, ulysses_pad_and_slice_inputs\nfrom verl.workers.actor import BasePPOActor\nfrom verl.workers.config import ActorConfig\n\n__all__ = [\"DataParallelPPOActor\"]\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass DataParallelPPOActor(BasePPOActor):\n    \"\"\"FSDP DataParallel PPO Actor or Ref worker\n\n    Args:\n        config (ActorConfig): Actor config\n        actor_module (nn.Module): Actor or ref module\n        actor_optimizer (torch.optim.Optimizer, optional): Actor optimizer. Defaults to None.\n    \"\"\"\n\n    def __init__(self, config: ActorConfig, actor_module: nn.Module, actor_optimizer: torch.optim.Optimizer = None):\n        \"\"\"When optimizer is None, it is Reference Policy\"\"\"\n        super().__init__(config)\n        self.actor_module = actor_module\n        self.actor_optimizer = actor_optimizer\n        role = \"Ref\" if actor_optimizer is None else \"Actor\"\n\n        self.use_remove_padding = self.config.get(\"use_remove_padding\", False)\n        if torch.distributed.get_rank() == 0:\n            print(f\"{role} use_remove_padding={self.use_remove_padding}\")\n        self.use_fused_kernels = self.config.get(\"use_fused_kernels\", False)\n        if torch.distributed.get_rank() == 0:\n            print(f\"{role} use_fused_kernels={self.use_fused_kernels}\")\n\n        self.ulysses_sequence_parallel_size = self.config.ulysses_sequence_parallel_size\n        self.use_ulysses_sp = self.ulysses_sequence_parallel_size > 1\n\n        self.use_dynamic_bsz = self.config.get(\"use_dynamic_bsz\", False)\n\n        self.use_prefix_grouper = self.config.get(\"use_prefix_grouper\", False)\n        if torch.distributed.get_rank() == 0:\n            print(f\"{role} use_prefix_grouper={self.use_prefix_grouper}\")\n\n        if self.config.entropy_from_logits_with_chunking:\n            entropy_from_logits = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_from_logits = verl_F.entropy_from_logits\n\n        self.compute_entropy_from_logits = (\n            torch.compile(entropy_from_logits, dynamic=True)\n            if self.config.get(\"use_torch_compile\", True)  # use torch compile by default\n            else entropy_from_logits\n        )\n        self.device_name = get_device_name()\n        self.param_dtype = PrecisionType.to_dtype(self.config.fsdp_config.get(\"dtype\", \"bfloat16\"))\n        if self.param_dtype == torch.float16:\n            from torch.distributed.fsdp.sharded_grad_scaler import ShardedGradScaler\n\n            self.scaler = ShardedGradScaler(growth_interval=400)\n        else:\n            self.scaler = None\n\n        # Sum of squared probabilities computation (for optimal_token_baseline)\n        # Only initialize if calculate_sum_pi_squared config is enabled\n        if self.config.get(\"calculate_sum_pi_squared\", False):\n            self.calculate_sum_pi_squared_from_logits = (\n                torch.compile(verl_F.calculate_sum_pi_squared_from_logits, dynamic=True)\n                if self.config.get(\"use_torch_compile\", True)\n                else verl_F.calculate_sum_pi_squared_from_logits\n            )\n            assert not (self.use_fused_kernels or self.use_prefix_grouper), (\n                \"calculate_sum_pi_squared is not supported with \"\n                f\"{self.use_fused_kernels=} or {self.use_prefix_grouper=} for now.\"\n            )\n\n    def _forward_micro_batch(\n        self, micro_batch: dict[str, torch.Tensor], temperature: float, calculate_entropy: bool = False\n    ) -> dict[str, torch.Tensor]:\n        \"\"\"\n        Returns:\n            dict[str, torch.Tensor]:\n                log_probs: (bs, response_len)\n                if calculate_entropy is True:\n                    entropys: (bs, response_len)\n                if calculate_sum_pi_squared is False:\n                    sum_pi_squared: (bs, response_len)\n        \"\"\"\n        calculate_sum_pi_squared = self.config.get(\"calculate_sum_pi_squared\", False)\n        sum_pi_squared_checkpointing = self.config.get(\"sum_pi_squared_checkpointing\", False)\n        # PrefixGrouper path for shared-prefix optimization\n        if self.use_prefix_grouper:\n            can_use_pg = (\n                not self.use_remove_padding\n                and not self.use_ulysses_sp\n                and not self.use_fused_kernels\n                and not self.use_dynamic_bsz\n            )\n            if can_use_pg and \"response_mask\" in micro_batch and \"uid\" in micro_batch:\n                from verl.trainer.ppo.prefix_grouper_utils import forward_micro_batch_with_prefix_grouper\n\n                return forward_micro_batch_with_prefix_grouper(\n                    micro_batch=micro_batch,\n                    model=self.actor_module,\n                    temperature=temperature,\n                    calculate_entropy=calculate_entropy,\n                    device_name=self.device_name,\n                    param_dtype=self.param_dtype,\n                    use_chunking_entropy=self.config.get(\"entropy_from_logits_with_chunking\", False),\n                )\n\n        response_length = micro_batch[\"responses\"].size(-1)\n        multi_modal_inputs = {}\n        if \"multi_modal_inputs\" in micro_batch.keys():\n            from verl.utils.model import extract_multi_modal_inputs\n\n            multi_modal_inputs = extract_multi_modal_inputs(micro_batch[\"multi_modal_inputs\"])\n\n        with torch.autocast(device_type=self.device_name, dtype=self.param_dtype):\n            input_ids = micro_batch[\"input_ids\"]\n            batch_size, seqlen = input_ids.shape\n            attention_mask = micro_batch[\"attention_mask\"]\n            position_ids = micro_batch[\"position_ids\"]\n            entropy = None\n            if position_ids.dim() == 3:  # qwen2vl mrope\n                position_ids = position_ids.transpose(0, 1)  # (bsz, 4, seqlen) -> (4, bsz, seqlen)\n\n            if self.use_remove_padding:\n                input_ids_rmpad, indices, cu_seqlens, *_ = unpad_input(\n                    input_ids.unsqueeze(-1), attention_mask\n                )  # input_ids_rmpad (total_nnz, ...)\n                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n\n                # unpad the position_ids to align the rotary\n                if position_ids.dim() == 3:\n                    position_ids_rmpad = (\n                        index_first_axis(rearrange(position_ids, \"c b s ... -> (b s) c ...\"), indices)\n                        .transpose(0, 1)\n                        .unsqueeze(1)\n                    )  # (4, bsz, seqlen) -> (4, 1, bsz * seqlen)\n                else:\n                    position_ids_rmpad = index_first_axis(\n                        rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n                    ).transpose(0, 1)\n\n                is_mask_all_zero = attention_mask.sum() == 0\n                if is_mask_all_zero:\n                    input_ids_rmpad = torch.zeros(\n                        (1, self.ulysses_sequence_parallel_size),\n                        device=input_ids.device,\n                        dtype=input_ids.dtype,\n                    )\n                    if position_ids.dim() == 3:\n                        position_ids_rmpad = torch.zeros(\n                            (position_ids.shape[0], 1, self.ulysses_sequence_parallel_size),\n                            device=position_ids.device,\n                            dtype=position_ids.dtype,\n                        )\n                    else:\n                        position_ids_rmpad = torch.zeros(\n                            (1, self.ulysses_sequence_parallel_size),\n                            device=position_ids.device,\n                            dtype=position_ids.dtype,\n                        )\n\n                if \"image_bound\" in multi_modal_inputs:\n                    from verl.utils.dataset.vision_utils import process_multi_modal_inputs_for_minicpmo\n\n                    multi_modal_inputs = process_multi_modal_inputs_for_minicpmo(\n                        input_ids, attention_mask, position_ids, cu_seqlens, multi_modal_inputs\n                    )\n\n                # for compute the log_prob\n                input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1)  # (1, total_nnz)\n\n                # pad and slice the inputs if sp > 1\n                if self.use_ulysses_sp:\n                    is_vlm_model = hasattr(\n                        getattr(self.actor_module, \"module\", self.actor_module).config, \"vision_config\"\n                    )\n                    if is_vlm_model:\n                        # vlm model's inputs will be sliced after embedding\n                        input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad(\n                            input_ids_rmpad,\n                            position_ids_rmpad=position_ids_rmpad,\n                            sp_size=self.ulysses_sequence_parallel_size,\n                        )\n                    else:\n                        input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(\n                            input_ids_rmpad,\n                            position_ids_rmpad=position_ids_rmpad,\n                            sp_size=self.ulysses_sequence_parallel_size,\n                        )\n                    input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs(\n                        input_ids_rmpad_rolled,\n                        position_ids_rmpad=None,\n                        sp_size=self.ulysses_sequence_parallel_size,\n                    )\n\n                input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0)  # ((total_nnz / sp) + pad)\n\n                # only pass input_ids and position_ids to enable flash_attn_varlen\n                extra_args = {}\n                if self.use_fused_kernels:\n                    extra_args[\"temperature\"] = temperature\n                    extra_args[\"return_dict\"] = True\n\n                output = self.actor_module(\n                    input_ids=input_ids_rmpad,\n                    attention_mask=None,\n                    position_ids=position_ids_rmpad,\n                    **multi_modal_inputs,\n                    use_cache=False,\n                    **extra_args,\n                )  # prevent model thinks we are generating\n\n                if self.use_fused_kernels:\n                    log_probs = output.log_probs.squeeze(0)  # (total_nnz,)\n                    entropy_rmpad = output.entropy.squeeze(0)  # (total_nnz,)\n\n                else:\n                    logits_rmpad = output.logits.squeeze(0)  # (total_nnz, vocab_size)\n                    logits_rmpad.div_(temperature)\n\n                    # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen)\n                    inplace_backward = True\n                    if calculate_entropy:\n                        inplace_backward = False\n                    log_probs = logprobs_from_logits(\n                        logits=logits_rmpad,\n                        labels=input_ids_rmpad_rolled,\n                        inplace_backward=inplace_backward,\n                    )\n\n                    # compute entropy\n                    if calculate_entropy:\n                        # ((total_nnz / sp) + pad)\n                        entropy_rmpad = (\n                            self.compute_entropy_from_logits(logits_rmpad)\n                            if not self.config.entropy_checkpointing\n                            else torch.utils.checkpoint.checkpoint(self.compute_entropy_from_logits, logits_rmpad)\n                        )\n\n                    # Compute sum_pi_squared if requested (for optimal_token_baseline)\n                    if calculate_sum_pi_squared:\n                        sum_pi_squared_rmpad = (\n                            self.calculate_sum_pi_squared_from_logits(logits_rmpad)\n                            if not sum_pi_squared_checkpointing\n                            else torch.utils.checkpoint.checkpoint(\n                                self.calculate_sum_pi_squared_from_logits, logits_rmpad\n                            )\n                        )\n\n                # gather log_prob if sp > 1\n                if self.use_ulysses_sp:\n                    # gather and unpad for the ulysses sp\n                    log_probs = gather_outputs_and_unpad(\n                        log_probs,\n                        gather_dim=0,\n                        unpad_dim=0,\n                        padding_size=pad_size,\n                    )\n                    if calculate_entropy:\n                        entropy_rmpad = gather_outputs_and_unpad(\n                            entropy_rmpad,\n                            gather_dim=0,\n                            unpad_dim=0,\n                            padding_size=pad_size,\n                        )\n                    if calculate_sum_pi_squared:\n                        sum_pi_squared_rmpad = gather_outputs_and_unpad(\n                            sum_pi_squared_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size\n                        )\n\n                if is_mask_all_zero:\n                    log_probs = log_probs[:0]\n                    if calculate_entropy:\n                        entropy_rmpad = entropy_rmpad[:0]\n\n                # pad back to (bsz, seqlen)\n                if calculate_entropy:\n                    full_entropy = pad_input(\n                        hidden_states=entropy_rmpad.unsqueeze(-1),\n                        indices=indices,\n                        batch=batch_size,\n                        seqlen=seqlen,\n                    )\n                if calculate_sum_pi_squared:\n                    full_sum_pi_squared = pad_input(\n                        hidden_states=sum_pi_squared_rmpad.unsqueeze(-1),\n                        indices=indices,\n                        batch=batch_size,\n                        seqlen=seqlen,\n                    )\n                full_log_probs = pad_input(\n                    hidden_states=log_probs.unsqueeze(-1),\n                    indices=indices,\n                    batch=batch_size,\n                    seqlen=seqlen,\n                )\n\n                # only return response part:\n                if calculate_entropy:\n                    entropy = full_entropy.squeeze(-1)[:, -response_length - 1 : -1]  # (bsz, response_length)\n                if calculate_sum_pi_squared:\n                    # (bsz, response_length)\n                    sum_pi_squared = full_sum_pi_squared.squeeze(-1)[:, -response_length - 1 : -1]\n                log_probs = full_log_probs.squeeze(-1)[:, -response_length - 1 : -1]  # (bsz, response_length)\n\n            else:  # not using rmpad and no ulysses sp\n                extra_args = {}\n                if self.use_fused_kernels:\n                    extra_args[\"temperature\"] = temperature\n                    extra_args[\"return_dict\"] = True\n\n                output = self.actor_module(\n                    input_ids=input_ids,\n                    attention_mask=attention_mask,\n                    position_ids=position_ids,\n                    **multi_modal_inputs,\n                    use_cache=False,\n                    **extra_args,\n                )  # prevent model thinks we are generating\n\n                if self.use_fused_kernels:\n                    log_probs = output.log_probs[:, -response_length - 1 : -1]\n                    entropy = output.entropy[:, -response_length - 1 : -1]  # (bsz, response_length)\n\n                else:\n                    logits = output.logits\n\n                    logits.div_(temperature)\n                    logits = logits[:, -response_length - 1 : -1, :]  # (bsz, response_length, vocab_size)\n                    log_probs = logprobs_from_logits(logits, micro_batch[\"responses\"])\n                    if calculate_entropy:\n                        if not self.config.entropy_checkpointing:\n                            entropy = verl_F.entropy_from_logits(logits)  # (bsz, response_length)\n                        else:\n                            entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits)\n                    # Compute sum_pi_squared if requested (for optimal_token_baseline)\n                    if calculate_sum_pi_squared:\n                        sum_pi_squared = (\n                            self.calculate_sum_pi_squared_from_logits(logits)\n                            if not sum_pi_squared_checkpointing\n                            else torch.utils.checkpoint.checkpoint(self.calculate_sum_pi_squared_from_logits, logits)\n                        )\n\n            outputs = {\"log_probs\": log_probs}\n            if calculate_entropy:\n                outputs[\"entropys\"] = entropy\n            if calculate_sum_pi_squared:\n                outputs[\"sum_pi_squared\"] = sum_pi_squared\n            return outputs\n\n    def _optimizer_step(self):\n        assert self.config.grad_clip is not None\n        if self.scaler is not None:\n            self.scaler.unscale_(self.actor_optimizer)\n        if isinstance(self.actor_module, FSDP):\n            grad_norm = self.actor_module.clip_grad_norm_(max_norm=self.config.grad_clip)\n        elif isinstance(self.actor_module, FSDPModule):\n            grad_norm = fsdp2_clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(self.actor_module.parameters(), max_norm=self.config.grad_clip)\n\n        if isinstance(grad_norm, DTensor):\n            grad_norm = grad_norm.full_tensor()\n\n        # if grad_norm is not finite, skip the update\n        if self.scaler is not None:\n            self.scaler.step(self.actor_optimizer)\n            self.scaler.update()\n        else:\n            if not torch.isfinite(grad_norm):\n                print(f\"WARN: rank {torch.distributed.get_rank()} grad_norm is not finite: {grad_norm}\")\n                self.actor_optimizer.zero_grad()\n            else:\n                self.actor_optimizer.step()\n\n        # Clear cached weight scales for QAT (weights changed)\n        if getattr(self.actor_module, \"_qat_fuse_enabled\", False):\n            from verl.utils.qat import invalidate_all_scales\n\n            invalidate_all_scales(self.actor_module)\n\n        return grad_norm\n\n    @GPUMemoryLogger(role=\"dp actor\", logger=logger)\n    def compute_log_prob(self, data: DataProto, calculate_entropy: bool = False) -> dict[str, torch.Tensor]:\n        \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids\n\n        Args:\n            data (DataProto): a DataProto containing keys\n\n                ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the\n                concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``.\n\n                ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``responses``:  tensor of shape [batch_size, response_length]. torch.int64.\n\n        Returns:\n            dict[str, torch.Tensor]: a dict containing keys\n                - ``log_probs``: tensor of shape [batch_size, response_length]. torch.float32.\n                - ``entropys``: tensor of shape [batch_size, response_length]. torch.float32.\n                - ``sum_pi_squared``: tensor of shape [batch_size, response_length]. torch.float32.\n        \"\"\"\n        calculate_sum_pi_squared = self.config.get(\"calculate_sum_pi_squared\", False)\n\n        # set to eval\n        self.actor_module.eval()\n\n        micro_batch_size = data.meta_info[\"micro_batch_size\"]\n        temperature = data.meta_info[\"temperature\"]  # temperature must be in the data.meta_info to avoid silent error\n        use_dynamic_bsz = data.meta_info[\"use_dynamic_bsz\"]\n        pad_token_id = data.meta_info.get(\"pad_token_id\", 0)\n        has_multi_modal_inputs = \"multi_modal_inputs\" in data.non_tensor_batch.keys()\n\n        select_keys = [\"responses\", \"input_ids\", \"attention_mask\", \"position_ids\"]\n        non_tensor_select_keys = [\"multi_modal_inputs\"] if has_multi_modal_inputs else []\n        if self.use_prefix_grouper:\n            select_keys += [k for k in [\"prompts\", \"response_mask\"] if k in data.batch]\n            if \"uid\" in data.non_tensor_batch:\n                non_tensor_select_keys.append(\"uid\")\n\n        data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys)\n\n        if use_dynamic_bsz:\n            max_token_len = data.meta_info[\"max_token_len\"] * self.ulysses_sequence_parallel_size\n            micro_batches, batch_idx_list = prepare_dynamic_batch(data, max_token_len=max_token_len)\n        else:\n            micro_batches = data.split(micro_batch_size)\n\n        log_probs_lst = []\n        entropy_lst = []\n        sum_pi_squared_lst = []\n        for micro_batch in micro_batches:\n            micro_batch = micro_batch.to(get_device_id())\n            model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch, \"pad_token_id\": pad_token_id}\n            with torch.no_grad():\n                outputs = self._forward_micro_batch(\n                    model_inputs, temperature=temperature, calculate_entropy=calculate_entropy\n                )\n            log_probs_lst.append(outputs[\"log_probs\"])\n            if calculate_entropy:\n                entropy_lst.append(outputs[\"entropys\"])\n            if calculate_sum_pi_squared:\n                sum_pi_squared_lst.append(outputs[\"sum_pi_squared\"])\n\n        log_probs = torch.concat(log_probs_lst, dim=0)\n        if calculate_entropy:\n            entropys = torch.concat(entropy_lst, dim=0)\n        if calculate_sum_pi_squared:\n            sum_pi_squared = torch.concat(sum_pi_squared_lst, dim=0)\n\n        if use_dynamic_bsz:\n            log_probs = restore_dynamic_batch(log_probs, batch_idx_list)\n            if calculate_entropy:\n                entropys = restore_dynamic_batch(entropys, batch_idx_list)\n            if calculate_sum_pi_squared:\n                sum_pi_squared = restore_dynamic_batch(sum_pi_squared, batch_idx_list)\n\n        outputs = {\"log_probs\": log_probs}\n        if calculate_entropy:\n            outputs[\"entropys\"] = entropys\n        if calculate_sum_pi_squared:\n            outputs[\"sum_pi_squared\"] = sum_pi_squared\n        return outputs\n\n    @GPUMemoryLogger(role=\"dp actor\", logger=logger)\n    def update_policy(self, data: DataProto):\n        # make sure we are in training mode\n        self.actor_module.train()\n\n        temperature = data.meta_info[\"temperature\"]  # temperature must be in the data.meta_info to avoid silent error\n        pad_token_id = data.meta_info.get(\"pad_token_id\", 0)\n\n        select_keys = [\n            \"responses\",\n            \"response_mask\",\n            \"input_ids\",\n            \"attention_mask\",\n            \"position_ids\",\n            \"old_log_probs\",\n            \"advantages\",\n        ]\n        if self.use_prefix_grouper and \"prompts\" in data.batch.keys():\n            select_keys.append(\"prompts\")\n        if self.config.use_kl_loss:\n            select_keys.append(\"ref_log_prob\")\n        # Include pre-computed IS weights if present in batch\n        # Weights are computed centrally in trainer and added to batch when algorithm.rollout_is=True\n        if \"rollout_is_weights\" in data.batch.keys():\n            select_keys.append(\"rollout_is_weights\")\n        # Include rollout_log_probs for computing rollout_corr metrics in bypass mode\n        if \"rollout_log_probs\" in data.batch.keys():\n            select_keys.append(\"rollout_log_probs\")\n\n        has_multi_modal_inputs = \"multi_modal_inputs\" in data.non_tensor_batch.keys()\n        non_tensor_select_keys = []\n        if has_multi_modal_inputs:\n            non_tensor_select_keys.append(\"multi_modal_inputs\")\n        if self.use_prefix_grouper and \"uid\" in data.non_tensor_batch.keys():\n            non_tensor_select_keys.append(\"uid\")\n\n        data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys)\n\n        # Split to make minibatch iterator for updating the actor\n        # See PPO paper for details. https://arxiv.org/abs/1707.06347\n        mini_batches = data.split(self.config.ppo_mini_batch_size)\n\n        on_policy = len(mini_batches) == 1 and self.config.ppo_epochs == 1\n\n        metrics = {\n            \"actor/pg_loss\": 0.0,\n            \"actor/kl_loss\": 0.0,\n        }\n        for _ in range(self.config.ppo_epochs):\n            for batch_idx, mini_batch in enumerate(mini_batches):\n                if self.config.use_dynamic_bsz:\n                    max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size\n                    micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len)\n                else:\n                    self.gradient_accumulation = (\n                        self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu\n                    )\n                    micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)\n\n                self.actor_optimizer.zero_grad()\n\n                for micro_batch in micro_batches:\n                    micro_batch = micro_batch.to(get_device_id())\n                    micro_batch_metrics = {}\n                    model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch, \"pad_token_id\": pad_token_id}\n                    response_mask = model_inputs[\"response_mask\"]\n                    old_log_prob = model_inputs[\"old_log_probs\"]\n                    advantages = model_inputs[\"advantages\"]\n\n                    entropy_coeff = self.config.entropy_coeff\n                    loss_agg_mode = self.config.loss_agg_mode\n\n                    calculate_entropy = self.config.calculate_entropy or (entropy_coeff != 0)\n\n                    if self.config.use_dynamic_bsz:\n                        loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size\n                    else:\n                        loss_scale_factor = 1 / self.gradient_accumulation\n\n                    # all return: (bsz, response_length)\n                    outputs = self._forward_micro_batch(\n                        model_inputs, temperature=temperature, calculate_entropy=calculate_entropy\n                    )\n                    log_prob = outputs[\"log_probs\"]\n                    entropy = outputs[\"entropys\"] if calculate_entropy else None\n\n                    # for fully_async_policy\n                    if hasattr(self.config, \"use_rollout_log_probs\") and self.config.use_rollout_log_probs:\n                        old_log_prob = model_inputs[\"old_log_probs\"]\n                    else:\n                        if on_policy:\n                            old_log_prob = log_prob.detach()\n                        else:\n                            old_log_prob = model_inputs[\"old_log_probs\"]\n\n                    loss_mode = self.config.policy_loss.get(\"loss_mode\", \"vanilla\")\n                    # vanilla -> verl.trainer.ppo.core_algos.compute_policy_loss_vanilla\n\n                    # Extract pre-computed rollout correction weights if present\n                    # Weights are computed centrally in trainer and added when algorithm.rollout_is=True\n                    rollout_is_weights = model_inputs.get(\"rollout_is_weights\", None)\n\n                    # gpg -> verl.trainer.ppo.core_algos.compute_policy_loss_gpg\n                    # clip_cov -> verl.trainer.ppo.core_algos.compute_policy_loss_clip_cov\n                    policy_loss_fn = get_policy_loss_fn(loss_mode)\n\n                    # Compute policy loss (any function is expected to return 2 values)\n                    pg_loss, pg_metrics = policy_loss_fn(\n                        old_log_prob=old_log_prob,\n                        log_prob=log_prob,\n                        advantages=advantages,\n                        response_mask=response_mask,\n                        loss_agg_mode=loss_agg_mode,\n                        config=self.config,\n                        rollout_is_weights=rollout_is_weights,\n                    )\n                    micro_batch_metrics.update(pg_metrics)\n\n                    # Skip if using bypass_mode loss (metrics already computed in pg_metrics)\n                    rollout_log_prob = model_inputs.get(\"rollout_log_probs\", None)\n                    if loss_mode != \"bypass_mode\" and rollout_log_prob is not None:\n                        # Compute metrics using CURRENT policy π_θ vs π_rollout\n                        # Tracks evolving off-policy gap as π_θ updates during mini-batch training\n                        from verl.trainer.ppo.rollout_corr_helper import compute_rollout_corr_metrics_from_logprobs\n\n                        rollout_corr_metrics = compute_rollout_corr_metrics_from_logprobs(\n                            log_prob=log_prob,\n                            rollout_log_prob=rollout_log_prob,\n                            response_mask=response_mask,\n                        )\n                        micro_batch_metrics.update(rollout_corr_metrics)\n\n                    policy_loss = pg_loss\n                    if calculate_entropy and entropy is not None:\n                        entropy_agg = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n                        micro_batch_metrics[\"actor/entropy\"] = entropy_agg.detach().item()\n                        if entropy_coeff != 0:\n                            policy_loss -= entropy_agg * entropy_coeff\n\n                    if self.config.use_kl_loss:\n                        ref_log_prob = model_inputs[\"ref_log_prob\"]\n                        # compute kl loss\n                        kld = kl_penalty(\n                            logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type\n                        )\n                        kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n\n                        policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef\n                        metrics[\"actor/kl_loss\"] += kl_loss.detach().item() * loss_scale_factor\n                        micro_batch_metrics[\"actor/kl_coef\"] = self.config.kl_loss_coef\n\n                    if self.config.use_dynamic_bsz:\n                        # relative to the dynamic bsz\n                        loss = policy_loss * loss_scale_factor\n                    else:\n                        loss = policy_loss * loss_scale_factor\n                    if self.scaler is not None:\n                        self.scaler.scale(loss).backward()\n                    else:\n                        loss.backward()\n\n                    metrics[\"actor/pg_loss\"] += pg_loss.detach().item() * loss_scale_factor\n                    append_to_dict(metrics, micro_batch_metrics)\n\n                grad_norm = self._optimizer_step()\n                mini_batch_metrics = {\"actor/grad_norm\": grad_norm.detach().item()}\n                append_to_dict(metrics, mini_batch_metrics)\n        self.actor_optimizer.zero_grad()\n        return metrics\n"
  },
  {
    "path": "verl/workers/actor/megatron_actor.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nMegatron Actor.\nIn megatron actor, the differences are:\n1. We only make minibatch\n\nNote that our model doesn't have to be `MegatronModule` because we don't share embedding in the last layer\n\"\"\"\n\nimport itertools\nimport logging\nimport os\nfrom functools import partial\nfrom typing import Iterable\n\nimport torch\nimport torch.distributed\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.distributed import finalize_model_grads\n\n# from megatron.core.optimizer import DistributedOptimizer\nfrom megatron.core.optimizer import DistributedOptimizer\nfrom megatron.core.pipeline_parallel import get_forward_backward_func\nfrom omegaconf import OmegaConf\nfrom torch import nn\n\nfrom verl import DataProto\nfrom verl.trainer.ppo.core_algos import agg_loss, get_policy_loss_fn, kl_penalty\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.megatron.pipeline_parallel import make_batch_generator\nfrom verl.utils.megatron.router_replay_patch import RouterReplay, RouterReplayAction\nfrom verl.utils.megatron.router_replay_utils import (\n    RouterReplayHelper,\n    merge_router_topk_indices,\n    pp_gather,\n    reorder_and_merge_vpp_layers,\n    set_router_replay_data,\n)\nfrom verl.utils.megatron.tensor_parallel import vocab_parallel_entropy, vocab_parallel_log_probs_from_logits\nfrom verl.utils.megatron_utils import get_megatron_mtp_loss, get_model_config, unwrap_model\nfrom verl.utils.profiler import GPUMemoryLogger\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import get_reverse_idx, rearrange_micro_batches\nfrom verl.utils.torch_functional import broadcast_dict_tensor\nfrom verl.workers.actor import BasePPOActor\nfrom verl.workers.config import MtpConfig\n\n__all__ = [\"MegatronPPOActor\"]\n\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass MegatronPPOActor(BasePPOActor):\n    def __init__(\n        self,\n        config,\n        model_config,\n        hf_config,\n        tf_config,\n        actor_module: nn.ModuleList,\n        actor_optimizer: DistributedOptimizer,\n        mtp_config: MtpConfig = None,\n    ):\n        \"\"\"MeagtronPPOActor class. This class implements the simple PPO logics when the model is built with Megatron.\n\n        Args:\n            config (OmegaConf): the basic config that contains the hyper-parameters of PPO Actor. It must contain\n\n                ``ppo_micro_batch_size_per_gpu``: micro batch size when updating ppo.\n\n                ``ppo_mini_batch_size``: minibatch size when updating ppo using the batch data.\n\n                ``ppo_epochs``: number of epochs to update the actor using the batch data.\n\n                ``shuffle``: whether to shuffle the data after each ppo epoch.\n\n                ``clip_ratio``: clip ratio of the ppo algorithm. See https://arxiv.org/abs/1707.06347.\n\n                ``entropy_coeff``: entropy coefficient of the PPO loss. See https://arxiv.org/abs/1707.06347.\n            model_config (OmegaConf): model configuration. It must contains ``model_config.vocab_size`` and\n                ``model_config.hidden_size``\n            hf_config (PretrainedConfig): huggingface config\n            tf_config (TransformerConfig): mcore transformer config\n            mtp_config (MtpConfig): mtp config, default None\n            actor_module (nn.ModuleList): actor module is a ModuleList that contains a list of nn.Module in this\n                pp stage.\n                each nn.Module in this rank holds a vpp module chunk. See https://arxiv.org/pdf/2104.04473.pdf for\n                more details.\n                The actor module has some constraints to follow in order to use the updating logics implemented here\n\n                1. It must implement unpad_input before any computation and pad_input after all the computation.\n                Remove padding is an\n                optimization that removes the padding tokens. See unpad_input and pad_input function in flash-attn\n                (https://github.com/Dao-AILab/flash-attention/blob/main/flash_attn/bert_padding.py).\n\n                2. Each pp stage must return the hidden state with the same shape [total_nnz, 1, hidden_size],\n                where total_nnz is the number of valid tokens in this batch. If sequence parallel is enabled, the size\n                of the hidden state is [total_nnz // tp, 1, hidden_size].\n            actor_optimizer (DistributedOptimizer): currently, we only support DistributedOptimizer in Megatron.\n                It implements\n                zero1 optimizer that shards the optimizer state across dp ranks.\n\n        >>> from megatron.training import get_model\n        >>> from megatron.optimizer import get_megatron_optimizer\n        >>> actor_module = get_model(megatron_actor_model_provider, wrap_with_ddp=True)\n        >>> actor_module = nn.ModuleList(actor_module)\n        >>> actor_optimizer = get_megatron_optimizer(actor_module)\n        >>> actor = MegatronPPOActor(config=config,\n        >>>                          model_config=actor_model_config,\n        >>>                          hf_config=hf_config,\n        >>>                          tf_config=tf_config,\n        >>>                          actor_module=actor_module,\n        >>>                          actor_optimizer=actor_optimizer)\n        \"\"\"\n        super().__init__(config)\n        self._validate_config(config)\n        self.model_config = model_config\n        self.hf_config = hf_config\n        self.tf_config = tf_config\n        self.mtp_config = mtp_config\n        self.actor_module = actor_module\n        self.actor_optimizer: DistributedOptimizer = actor_optimizer\n\n        if self.mtp_config:\n            assert self.mtp_config.enable, \"MTP requires mtp_config.enable to be True\"\n\n        self.use_fused_kernels = self.config.get(\"use_fused_kernels\", False)\n        if getattr(self.mtp_config, \"enable\", False) and self.use_fused_kernels:\n            self.use_fused_kernels = False\n            logger.warning_once(\n                \"MTP is not compatible with fused kernels for now. Automatically disable use_fused_kernels.\"\n            )\n        if self.use_fused_kernels and not getattr(self.config, \"overlap_moe_expert_parallel_comm\", False):\n            # do not patch if overlap_moe_expert_parallel_comm is enabled\n            logger.warning_once(\n                \"Recommend to disable use_fused_kernels since the fused kernel's performance is broken for triton>=3.3\"\n                \"Unless you are using a very old version of triton < 3.3\"\n            )\n            from verl.models.mcore.model_forward_fused import patch_fused_forward\n\n            for model in self.actor_module:\n                patch_fused_forward(model)\n        else:\n            from verl.models.mcore.mtp_patch import patch_postprocess\n\n            for model in self.actor_module:\n                if self.mtp_config:\n                    from verl.models.mcore.mtp_patch import patch_mtp_layer_get_embeddings\n\n                    patch_postprocess(model)\n\n                    if self.mtp_config.detach_encoder:\n                        patch_mtp_layer_get_embeddings(model)\n\n        self.optimizer_step_args = OmegaConf.create(\n            {\n                \"skip_grad\": None,\n                \"overlap_dp_param_comm\": False,\n                \"overlap_dp_grad_comm\": False,\n                \"gradient_accumulation_steps\": 1,\n                \"sequence_parallel\": self.tf_config.sequence_parallel,\n                \"DDP_impl\": \"local\",\n                \"layernorm_allreduce_bucket_threshold\": 0,\n                \"reduce_grads_use_alltoall\": False,\n            }\n        )\n\n        self.router_replay = self.config.router_replay\n        self.enable_routing_replay = self.router_replay.mode != \"disabled\"\n        if self.enable_routing_replay:\n            self.mini_layer_topk_idx_list = []\n\n        config = get_model_config(self.actor_module[0])\n        print(config)\n        config.finalize_model_grads_func = finalize_model_grads\n\n    def _validate_config(self, config) -> None:\n        \"\"\"Validate config options not implemented for Megatron backend\"\"\"\n        assert config.get(\"ulysses_sequence_parallel_size\", 1) == 1\n        if config.get(\"shuffle\", False):\n            assert config.data_loader_seed is not None, \"If shuffle dataloader, seed must be manually set\"\n        if config.megatron.tensor_model_parallel_size == 1:\n            print(\"[Warining] Because actor tp size == 1, set sp to False\")\n            config.megatron.sequence_parallel = False\n        self.config = config\n\n    @GPUMemoryLogger(role=\"megatron actor\", logger=logger)\n    def compute_log_prob(self, data: DataProto, calculate_entropy=False) -> torch.Tensor:\n        \"\"\"Compute the log probability of the responses given input_ids, attention_mask and position_ids\n\n        Args:\n            data (DataProto): a DataProto containing keys\n\n                ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64. Note that input_ids is the\n                concatenation of prompt and response. Note that ``sequence_length = prompt_length + response_length``.\n\n                ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64.\n\n                ``responses``:  tensor of shape [batch_size, response_length]. torch.int64.\n\n        Returns:\n            DataProto: torch.Tensor: the log_prob tensor\n        \"\"\"\n        prev_modes = [m.training for m in self.actor_module]\n        for module in self.actor_module:\n            module.eval()\n        use_dynamic_bsz = data.meta_info.get(\"use_dynamic_bsz\", False)\n        micro_batch_size = data.meta_info.get(\"micro_batch_size\", None)\n        max_token_len = data.meta_info.get(\"max_token_len\", None)\n        if use_dynamic_bsz:\n            assert max_token_len is not None, \"max_token_len must be set when use_dynamic_bsz is True\"\n            max_token_len = max_token_len * self.config.megatron.context_parallel_size\n        else:\n            assert micro_batch_size is not None, (\n                \"micro batch size is needed for forward compute when use_dynamic_bsz is False\"\n            )\n\n        def compute_logprobs_fn(output, data, use_dynamic_bsz=False, indices=None):\n            response = data[\"responses\"]\n            response_length = response.size(1)\n            log_probs = output[\"log_probs\"][:, -response_length - 1 : -1].contiguous()\n            return {\"log_probs\": log_probs}\n\n        # We make recompute_old_log_prob by default here.\n        # TODO (zhangchi.usc1992): actually, this function should only return log_prob and this logic should be\n        # handled by user outside\n        recompute_old_log_prob = self.config.get(\"recompute_old_log_prob\", True)\n\n        entropys = torch.Tensor()\n        if recompute_old_log_prob:\n            select_keys = [\"responses\", \"input_ids\", \"attention_mask\", \"position_ids\"]\n\n            if self.enable_routing_replay and self.config.router_replay.mode == \"R3\":\n                assert \"routed_experts\" in data.batch.keys(), \"routed_experts must be in data.batch.keys()\"\n                select_keys.append(\"routed_experts\")\n\n            batch = data.select(batch_keys=select_keys).batch\n            input_ids = batch[\"input_ids\"]\n            batch_size = input_ids.size(0)\n            response = batch[\"responses\"]\n            response_length = response.size(1)\n            with torch.no_grad():\n                output = self.forward_backward_batch(\n                    data,\n                    forward_only=True,\n                    post_process_fn=compute_logprobs_fn,\n                    calculate_entropy=calculate_entropy,\n                    use_dynamic_bsz=use_dynamic_bsz,\n                    micro_batch_size=micro_batch_size,\n                    max_token_len=max_token_len,\n                )\n                if mpu.is_pipeline_last_stage(ignore_virtual=True):\n                    # only on last rank. It should be on every tp rank\n                    if calculate_entropy:\n                        log_probs = [o[0][\"log_probs\"] for o in output[\"output\"]]  # (bs, seq_size)\n                    else:\n                        log_probs = [o[\"log_probs\"] for o in output[\"output\"]]  # (bs, seq_size)\n                    log_probs = torch.cat(log_probs, dim=0).to(torch.float32)\n                    if use_dynamic_bsz:\n                        indices = output[\"indices\"]\n                        indices = list(itertools.chain.from_iterable(indices))\n                        assert len(indices) == log_probs.size(0), f\"{len(indices)} vs. {log_probs.size()}\"\n                        revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long)\n                        log_probs = log_probs[revert_indices]\n                else:\n                    log_probs = torch.empty(\n                        size=(batch_size, response_length), dtype=torch.float32, device=input_ids.device\n                    )\n                log_probs = log_probs.to(get_device_id())\n                # broadcast across pp ranks\n                torch.distributed.broadcast(\n                    tensor=log_probs,\n                    src=mpu.get_pipeline_model_parallel_last_rank(),\n                    group=mpu.get_pipeline_model_parallel_group(),\n                    async_op=False,\n                )\n                log_probs = log_probs.to(\"cpu\")\n                if calculate_entropy:\n                    # Note that o[0] is metrics, o[1] is entropy\n                    if mpu.is_pipeline_last_stage(ignore_virtual=True):\n                        entropys = torch.cat([o[1] for o in output[\"output\"]], dim=0)\n                        entropys = entropys.to(torch.float32)\n                        if use_dynamic_bsz:\n                            indices = output[\"indices\"]\n                            indices = list(itertools.chain.from_iterable(indices))\n                            assert len(indices) == entropys.size(0), f\"{len(indices)} vs. {entropys.size()}\"\n                            revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long)\n                            entropys = entropys[revert_indices]\n                    else:\n                        entropys = torch.empty(\n                            size=(batch_size, response_length), dtype=torch.float32, device=input_ids.device\n                        )\n                    # broadcast across pp ranks\n                    entropys = entropys.to(get_device_id())\n                    torch.distributed.broadcast(\n                        tensor=entropys,\n                        src=mpu.get_pipeline_model_parallel_last_rank(),\n                        group=mpu.get_pipeline_model_parallel_group(),\n                        async_op=False,\n                    )\n                    entropys = entropys.to(\"cpu\")\n                layers_topk_idx = None\n\n                if RouterReplayHelper.is_r2_record_action(self.tf_config):\n                    # (bs, max_seq_len/response_len,local_layer_num,topk)\n                    layers_topk_idx = output[\"mini_layer_topk_idx_tensor\"].to(torch.uint8)\n                    if use_dynamic_bsz:\n                        indices = output[\"indices\"]\n                        indices = list(itertools.chain.from_iterable(indices))\n                        assert len(indices) == layers_topk_idx.size(0), f\"{len(indices)} vs. {layers_topk_idx.size()}\"\n                        revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long)\n                        layers_topk_idx = layers_topk_idx[revert_indices]\n                    layers_topk_idx = pp_gather(layers_topk_idx, self.tf_config)\n        # add empty cache after each compute\n        get_torch_device().empty_cache()\n\n        for module, mode in zip(self.actor_module, prev_modes, strict=False):\n            module.train(mode)\n        return log_probs, entropys, layers_topk_idx\n\n    def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]:\n        \"\"\"Make minibatch iterator for updating the actor\n\n        Args:\n            data (DataProto): a DataProto containing keys\n\n                ``input_ids``: tensor of shape [batch_size, sequence_length]. torch.int64, where\n                ``sequence_length = prompt_length + response_length``\n\n                ``attention_mask``: tensor of shape [batch_size, sequence_length]. torch.int64\n\n                ``position_ids``: tensor of shape [batch_size, sequence_length]. torch.int64\n\n                ``responses``: tensor of shape [batch_size, response_length]. torch.int64. Note that\n                responses = input_ids[:, -response_length:]\n\n                ``old_log_probs``: tensor of shape [batch_size, response_length]. torch.float32. The log probability\n                of responses.\n\n                ``advantages``: tensor of shape [batch_size, response_length]. torch.float32. The advantages of\n                responses.\n                See PPO paper for details. https://arxiv.org/abs/1707.06347\n\n        Returns:\n\n        \"\"\"\n        select_keys = [\n            \"responses\",\n            \"input_ids\",\n            \"attention_mask\",\n            \"response_mask\",\n            \"position_ids\",\n            \"old_log_probs\",\n            \"advantages\",\n        ]\n        if self.config.use_kl_loss:\n            select_keys.append(\"ref_log_prob\")\n        # Include pre-computed IS weights if present in batch\n        # Weights are computed centrally in trainer and added to batch when algorithm.rollout_is=True\n        if \"rollout_is_weights\" in data.batch.keys():\n            select_keys.append(\"rollout_is_weights\")\n        # Include rollout_log_probs for computing rollout_corr metrics in bypass mode\n        if \"rollout_log_probs\" in data.batch.keys():\n            select_keys.append(\"rollout_log_probs\")\n        self.has_multi_modal_inputs = \"multi_modal_inputs\" in data.non_tensor_batch.keys()\n        # router replay\n        if self.enable_routing_replay:\n            select_keys.append(\"routed_experts\")\n        if self.has_multi_modal_inputs:\n            data = data.select(select_keys, [\"multi_modal_inputs\"])\n        else:\n            data = data.select(batch_keys=select_keys)\n\n        return data.make_iterator(\n            mini_batch_size=self.config.ppo_mini_batch_size,\n            epochs=self.config.ppo_epochs,\n            seed=self.config.data_loader_seed,\n            dataloader_kwargs={\"shuffle\": self.config.shuffle},\n        )\n\n    def forward_backward_batch(\n        self,\n        data: DataProto,\n        forward_only=False,\n        post_process_fn=None,\n        calculate_entropy=False,\n        use_dynamic_bsz=False,\n        micro_batch_size=None,\n        max_token_len=None,\n        mini_batch_size=None,\n    ):\n        \"\"\"\n        We assume:\n        - The model takes input: (input_ids, attention_mask, position_ids). No rmpad for the input\n        - The communication shape is (total_nnz_pad_to_sp // tp_size, 1, hidden_size) if sequence parallel is enabled\n        \"\"\"\n        # broadcast from last pp rank to all other pp ranks\n        # TODO: actually, we just need to control the sampling order.\n        data.to(get_device_id())\n        data.batch = data.batch.contiguous()\n        mini_batch = data\n        broadcast_dict_tensor(\n            mini_batch.batch,\n            src=mpu.get_pipeline_model_parallel_last_rank(),\n            group=mpu.get_pipeline_model_parallel_group(),\n        )\n        mini_batch.to(\"cpu\")\n        # split into micro-batches\n        mini_batch.batch[\"attention_mask\"] = mini_batch.batch[\"attention_mask\"].to(bool)\n        self.has_multi_modal_inputs = \"multi_modal_inputs\" in mini_batch.non_tensor_batch.keys()\n        if self.has_multi_modal_inputs:\n            mini_batch.batch[\"multi_modal_inputs\"] = mini_batch.non_tensor_batch[\"multi_modal_inputs\"]\n            mini_batch.batch[\"multi_modal_inputs_idx\"] = torch.Tensor(\n                list(range(len(mini_batch.non_tensor_batch[\"multi_modal_inputs\"])))\n            ).to(torch.int64)\n\n        if mini_batch.batch[\"position_ids\"].dim() == 3:  # qwen2vl mrope [bs, 3, seq_len]\n            mini_batch.batch[\"position_ids\"] = mini_batch.batch[\"position_ids\"][\n                :, 0\n            ]  # mcore patch recompute qwen2vl's pos ids during forward\n\n        indices = None\n        temperature = data.meta_info[\"temperature\"]\n        if use_dynamic_bsz:\n            assert max_token_len is not None, \"max_token_len must be set when use_dynamic_bsz is True\"\n            dp_group = mpu.get_data_parallel_group()\n            vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()\n            if vpp_size is not None and vpp_size > 1:\n                microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage\n                micro_batches, indices = rearrange_micro_batches(\n                    batch=mini_batch.batch,\n                    num_batches_divided_by=microbatch_group_size_per_vp_stage,\n                    max_token_len=max_token_len,\n                    dp_group=dp_group,\n                )\n                assert len(micro_batches) % self.tf_config.microbatch_group_size_per_vp_stage == 0, (\n                    f\"micro_batches {micro_batches} must be divisible by microbatch_group_size_per_vp_stage \"\n                    f\"{microbatch_group_size_per_vp_stage} for megatron backend\"\n                )\n            else:\n                micro_batches, indices = rearrange_micro_batches(\n                    batch=mini_batch.batch, max_token_len=max_token_len, dp_group=dp_group\n                )\n            total_seqlen = max_token_len\n        else:\n            assert micro_batch_size is not None, (\n                \"micro_batch_size is needed to be passed in when not using dynamic batch size\"\n            )\n            micro_batches = mini_batch.batch.split(micro_batch_size)\n            seq_len = micro_batches[0][\"input_ids\"].shape[1]\n            total_seqlen = micro_batch_size * seq_len\n        # compute input shapes for pp stages\n        n_micro_batch = len(micro_batches)\n\n        forward_backward_func = get_forward_backward_func()\n\n        def loss_func(output, data, meta_info):\n            # For memory efficiency\n            # We move calculation of entropy to compute_log_probs, forward_only == True\n            log_probs = None\n            entropy = None\n            if isinstance(output, dict):\n                log_probs = output[\"log_probs\"]\n                if \"entropy\" in output:\n                    entropy = output[\"entropy\"]\n            else:\n                assert isinstance(output, torch.Tensor)\n                log_probs = output\n\n            device = log_probs.device\n            metrics = {}\n            if forward_only:\n                if post_process_fn is None:\n                    pass\n                    # metrics[\"logits\"] = output\n                else:\n                    stats = post_process_fn(output, data)\n                    metrics.update(stats)\n                if not calculate_entropy:\n                    return torch.tensor(1.0, device=device), metrics\n\n            responses = data[\"responses\"]\n            response_length = responses.size(1)\n            response_mask = data[\"response_mask\"].to(bool)\n            loss_agg_mode = self.config.loss_agg_mode\n            # compute policy loss\n            log_prob = log_probs[:, -response_length - 1 : -1].contiguous()\n            ret_entropy = None\n            stats = {}\n            if not forward_only:\n                old_log_prob = data[\"old_log_probs\"]\n                advantages = data[\"advantages\"]\n\n                entropy_coeff = self.config.entropy_coeff\n                loss_agg_mode = self.config.loss_agg_mode\n\n                loss_mode = self.config.policy_loss.get(\"loss_mode\", \"vanilla\")\n\n                policy_loss_fn = get_policy_loss_fn(loss_mode)\n\n                # Extract pre-computed rollout correction weights if present\n                # Weights are computed centrally in trainer and added when algorithm.rollout_is=True\n                rollout_is_weights = data.get(\"rollout_is_weights\", None)\n                pg_loss, pg_metrics = policy_loss_fn(\n                    old_log_prob=old_log_prob,\n                    log_prob=log_prob,\n                    advantages=advantages,\n                    response_mask=response_mask,\n                    loss_agg_mode=loss_agg_mode,\n                    config=self.config,\n                    rollout_is_weights=rollout_is_weights,\n                )\n                stats.update(pg_metrics)\n\n                # Skip if using bypass_mode loss (metrics already computed in pg_metrics)\n                rollout_log_prob = data.get(\"rollout_log_probs\", None)\n                if loss_mode != \"bypass_mode\" and rollout_log_prob is not None:\n                    # Compute metrics using CURRENT policy π_θ vs π_rollout\n                    # Tracks evolving off-policy gap as π_θ updates during mini-batch training\n                    from verl.trainer.ppo.rollout_corr_helper import compute_rollout_corr_metrics_from_logprobs\n\n                    rollout_corr_metrics = compute_rollout_corr_metrics_from_logprobs(\n                        log_prob=log_prob,\n                        rollout_log_prob=rollout_log_prob,\n                        response_mask=response_mask,\n                    )\n                    stats.update(rollout_corr_metrics)\n\n                stats[\"actor/pg_loss\"] = pg_loss.detach().item()\n                policy_loss = pg_loss\n\n            if calculate_entropy:\n                entropy = output[\"entropy\"][:, -response_length - 1 : -1].contiguous()\n                if not forward_only:\n                    entropy_loss = agg_loss(loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode)\n                    entropy_coeff = meta_info[\"entropy_coeff\"]\n                    policy_loss = pg_loss - entropy_coeff * entropy_loss\n                else:\n                    ret_entropy = entropy\n\n            if forward_only:\n                policy_loss = torch.tensor(1.0, device=device)\n            else:\n                if self.config.use_kl_loss:\n                    ref_log_prob = data[\"ref_log_prob\"]\n                    # compute kl loss\n                    kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=self.config.kl_loss_type)\n                    kl_loss = agg_loss(loss_mat=kld, loss_mask=response_mask, loss_agg_mode=self.config.loss_agg_mode)\n\n                    policy_loss = policy_loss + kl_loss * self.config.kl_loss_coef\n                    metrics[\"actor/kl_loss\"] = kl_loss.detach().item()\n                    metrics[\"actor/kl_coef\"] = self.config.kl_loss_coef\n\n                # return loss and stats\n\n            append_to_dict(metrics, stats)\n            return policy_loss, [metrics, ret_entropy]\n\n        def forward_step(batch_iter, model, return_schedule_plan: bool = False):\n            \"\"\"\n            Args:\n                batch_iter: the batch iterator\n                model: the model\n                return_schedule_plan: whether to return the schedule plan, for 1f1b overlap\n            \"\"\"\n            if return_schedule_plan:\n                assert self.tf_config.overlap_moe_expert_parallel_comm, (\n                    \"overlap_moe_expert_parallel_comm must be enabled to return the schedule plan\"\n                )\n                # TODO: Fix this\n                assert not calculate_entropy, \"calculate_entropy must be disabled to return the schedule plan\"\n                from megatron.core.models.gpt.gpt_model import GPTModel\n\n                assert isinstance(model, GPTModel), \"model must be a GPTModel\"\n                assert self.use_fused_kernels, \"use_fused_kernels must be enabled to return the schedule plan\"\n                # TODO: support VLM with MoE\n                from verl.models.mcore.model_forward_1f1b_overlap import gptmodel_forward_1f1b_overlap\n\n            batch = next(batch_iter)\n            batch = batch.to(get_device_id())\n            batch = batch.contiguous()\n\n            input_ids = batch[\"input_ids\"]\n            attention_mask = batch[\"attention_mask\"].to(bool)\n            position_ids = batch[\"position_ids\"]\n\n            unwrapped_model = unwrap_model(model)\n            if hasattr(unwrapped_model, \"vp_stage\"):\n                vp_rank = unwrapped_model.vp_stage\n            else:\n                vp_rank = 0\n\n            multi_modal_inputs = {}\n            if \"multi_modal_inputs\" in batch:\n                from verl.utils.model import extract_multi_modal_inputs\n\n                indices = batch.get(\"multi_modal_inputs_idx\", None)\n                multi_modal_inputs = extract_multi_modal_inputs(batch[\"multi_modal_inputs\"], indices)\n            responses = batch[\"responses\"]\n            response_length = responses.size(1)\n            label = position_ids.clone()\n            label[:, -response_length - 1 : -1] = responses\n            label_mask = attention_mask.clone()\n            label_mask[:, : -response_length - 1] = False\n            label_mask[:, -1] = False\n\n            if RouterReplayHelper.is_replay_backward_action(self.tf_config, vp_rank):\n                router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank)\n                for router in router_instance_list:\n                    router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n\n            if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank):\n                layers_topk_idx = batch[\"routed_experts\"]\n                set_router_replay_data(layers_topk_idx, attention_mask, self.tf_config, vp_rank)\n\n            from verl.models.mcore import get_mcore_forward_fn, get_mcore_forward_fused_fn\n\n            if self.use_fused_kernels:\n                forward_fn = get_mcore_forward_fused_fn(self.hf_config)\n                if return_schedule_plan:\n                    forward_fn = gptmodel_forward_1f1b_overlap\n                # return dict of [logits, entropy]\n                output = forward_fn(\n                    model=model,\n                    input_ids=input_ids,\n                    position_ids=position_ids,\n                    attention_mask=attention_mask,\n                    labels=label,\n                    labels_mask=label_mask,\n                    temperature=temperature,\n                    multi_modal_inputs=multi_modal_inputs,\n                )\n            else:\n                forward_fn = get_mcore_forward_fn(self.hf_config)\n\n                def logits_processor(logits, label, label_mask):\n                    assert logits.shape[:2] == label.shape[:2]\n                    assert label.shape == label_mask.shape\n                    logits.div_(temperature)\n                    ret = {}\n                    if calculate_entropy:\n                        logits_bak = logits.clone()\n                        # # disable the hint until the fused_kernel is optimized for triton>=3.3\n                        # logger.warning_once(\n                        #     \"For memory-efficient computation, enable fused kernels via \"\n                        #     \"`actor_rollout_ref.model.use_fused_kernels=True`. \"\n                        #     \"The current `clone()` operation ensures correctness but increases memory usage.\"\n                        # )\n                        entropy = vocab_parallel_entropy(logits)\n                        ret[\"entropy\"] = entropy\n                    else:\n                        logits_bak = logits\n                    log_probs = vocab_parallel_log_probs_from_logits(logits_bak, label)\n                    log_probs = log_probs.masked_fill(~label_mask, 0.0)\n                    ret[\"log_probs\"] = log_probs\n                    return ret\n\n                logits_processor_args = {\"label\": label, \"label_mask\": label_mask}\n                output = forward_fn(\n                    model=model,\n                    input_ids=input_ids,\n                    attention_mask=attention_mask,\n                    position_ids=position_ids,\n                    multi_modal_inputs=multi_modal_inputs,\n                    logits_processor=logits_processor,\n                    logits_processor_args=logits_processor_args,\n                    data_format=\"thd\" if self.config.megatron.use_remove_padding else \"bshd\",\n                    mtp_config=None if forward_only else self.mtp_config,\n                )\n\n            if forward_only:\n                meta_info = None\n            else:\n                clip_ratio_c = self.config.get(\"clip_ratio_c\", 3.0)\n                meta_info = {\n                    \"clip_ratio\": self.config.clip_ratio,\n                    \"entropy_coeff\": self.config.entropy_coeff,\n                    \"clip_ratio_c\": clip_ratio_c,\n                }\n\n            if RouterReplayHelper.is_r2_record_action(self.tf_config, vp_rank):\n                merge_router_topk_indices(\n                    attention_mask, input_ids, self.mini_layer_topk_idx_list, self.tf_config, vp_rank\n                )\n\n            if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank):\n                router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank)\n                for router in router_instance_list:\n                    router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD)\n\n            return output, partial(loss_func, data=batch, meta_info=meta_info)\n\n        # batch should be a list of batches inside micro-batches\n        batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.actor_module))\n\n        # TODO: we may use the new schedule instead\n        # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size)\n        if mpu.get_pipeline_model_parallel_world_size() > 1:\n            losses_reduced = forward_backward_func(\n                forward_step_func=forward_step,\n                data_iterator=batch_generator,\n                model=self.actor_module,\n                num_microbatches=n_micro_batch,\n                seq_length=total_seqlen,  # no use when input_shapes was set\n                micro_batch_size=1,  # no use when input_shapes was set\n                forward_only=forward_only,\n            )\n        else:\n            losses_reduced = forward_backward_func(\n                forward_step_func=forward_step,\n                data_iterator=batch_generator,\n                model=self.actor_module,\n                num_microbatches=n_micro_batch,\n                seq_length=total_seqlen,  # in use for pp = 1\n                micro_batch_size=1,  # in use for pp = 1\n                forward_only=forward_only,\n            )\n        # loss_reduces contains the stats returned from loss_func\n\n        if self.has_multi_modal_inputs:\n            data.batch.pop(\"multi_modal_inputs\")\n            data.batch.pop(\"multi_modal_inputs_idx\")\n            data.non_tensor_batch.pop(\"multi_modal_inputs\")\n\n        losses_reduced = {\"output\": losses_reduced}\n        if use_dynamic_bsz:\n            losses_reduced[\"indices\"] = indices\n        if RouterReplayHelper.is_r2_record_action(self.tf_config):\n            if self.tf_config.virtual_pipeline_model_parallel_size is not None:\n                # config = self.actor_module[0].module.module.config\n                vp_size = len(self.actor_module)\n                microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage\n                bs = n_micro_batch\n                losses_reduced[\"mini_layer_topk_idx_tensor\"] = reorder_and_merge_vpp_layers(\n                    self.mini_layer_topk_idx_list, bs, vp_size, microbatch_group_size_per_vp_stage\n                )\n            else:\n                losses_reduced[\"mini_layer_topk_idx_tensor\"] = torch.cat(self.mini_layer_topk_idx_list, dim=0)\n            self.mini_layer_topk_idx_list = []\n\n        # Collect and pass MTP metrics to losses_reduced\n        if not forward_only and self.mtp_config and self.mtp_config.enable_train:\n            metrics = get_megatron_mtp_loss(n_micro_batch)\n            losses_reduced[\"mtp_losses\"] = [metrics]\n\n        return losses_reduced\n\n    @GPUMemoryLogger(role=\"megatron actor\", logger=logger)\n    def update_policy(self, dataloader: Iterable[DataProto], enable_mtp: bool = False) -> dict:\n        \"\"\"Update the policy with an iterator of DataProto\n\n        Args:\n            dataloader (Iterable[DataProto]): an iterator over the DataProto that returns by ``make_minibatch_iterator``\n                The keys of each data batch is described in the make_minibatch_iterator.\n\n            enable_mtp (bool, optional): whether to enable MTP communication\n\n        Returns:\n            Dict: a dictionary containing the statistics. Note that the statistics are only valid in the last pp stage\n            and users have to combine the output in each dp rank manually.\n\n        \"\"\"\n        metrics = {}\n        for data in dataloader:\n            if self.config.router_replay.mode in [\"R2\", \"R3\"]:\n                RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n            self.actor_optimizer.zero_grad()\n            # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm\n            for chunk in self.actor_module:\n                # if use distributed optimizer, zero grad buffer will be handled by optimizer\n                chunk.zero_grad_buffer()\n\n            calculate_entropy = self.config.entropy_coeff != 0\n            if data.meta_info.get(\"micro_batch_size\", None) is not None:\n                micro_batch_size = data.meta_info[\"micro_batch_size\"]\n            else:\n                micro_batch_size = self.config.ppo_micro_batch_size_per_gpu\n            max_token_len = None\n            if self.config.use_dynamic_bsz:\n                max_token_len = self.config.ppo_max_token_len_per_gpu * self.config.megatron.context_parallel_size\n            metric_micro_batch = self.forward_backward_batch(\n                data,\n                calculate_entropy=calculate_entropy,\n                use_dynamic_bsz=self.config.use_dynamic_bsz,\n                micro_batch_size=micro_batch_size,\n                max_token_len=max_token_len,\n                mini_batch_size=self.config.ppo_mini_batch_size,\n            )\n\n            mtp_losses = metric_micro_batch.get(\"mtp_losses\", None)\n            if mtp_losses is not None:\n                # mtp_losses is now in format: [{\"mtp_losses/mtp_1_loss\": [value1], \"mtp_losses/mtp_2_loss\": [value2]}]\n                for mtp_metrics_dict in mtp_losses:\n                    append_to_dict(metrics, mtp_metrics_dict)\n\n            metric_micro_batch = metric_micro_batch[\"output\"]\n            for metric in metric_micro_batch:\n                # Note that o[0] is metrics, o[1] is entropy, o[2] is response_mask\n                append_to_dict(metrics, metric[0])  # append the metric from this micro-batch to global metrics.\n\n            update_successful, grad_norm, num_zeros_in_grad = self.actor_optimizer.step()\n            data = {\"actor/grad_norm\": grad_norm}\n            append_to_dict(metrics, data)\n\n            if update_successful:\n                # allgather already execute in optimizer.step in new megatron\n                pass\n            else:\n                raise NotImplementedError\n\n            if self.config.router_replay.mode in [\"R2\", \"R3\"]:\n                RouterReplay.clear_global_router_replay_action()\n                RouterReplay.clear_global_indices()\n\n        self.actor_optimizer.zero_grad()\n        get_torch_device().empty_cache()\n        return metrics\n"
  },
  {
    "path": "verl/workers/config/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom . import actor, critic, engine, model, optimizer, reward, rollout\nfrom .actor import *  # noqa: F401\nfrom .critic import *  # noqa: F401\nfrom .engine import *  # noqa: F401\nfrom .model import *  # noqa: F401\nfrom .optimizer import *  # noqa: F401\nfrom .reward import *  # noqa: F401\nfrom .rollout import *  # noqa: F401\n\n__all__ = (\n    actor.__all__\n    + critic.__all__\n    + reward.__all__\n    + engine.__all__\n    + optimizer.__all__\n    + rollout.__all__\n    + model.__all__\n)\n"
  },
  {
    "path": "verl/workers/config/actor.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nfrom omegaconf import MISSING\n\nfrom verl.base_config import BaseConfig\nfrom verl.trainer.config import CheckpointConfig, RolloutCorrectionConfig\nfrom verl.utils.profiler.config import ProfilerConfig\nfrom verl.utils.qat import QATConfig\n\nfrom .engine import FSDPEngineConfig, McoreEngineConfig, TorchtitanEngineConfig, VeOmniEngineConfig\nfrom .model import HFModelConfig\nfrom .optimizer import OptimizerConfig\n\n__all__ = [\n    \"PolicyLossConfig\",\n    \"RouterReplayConfig\",\n    \"ActorConfig\",\n    \"FSDPActorConfig\",\n    \"McoreActorConfig\",\n    \"VeOmniActorConfig\",\n    \"QATConfig\",\n    \"TorchTitanActorConfig\",\n]\n\n\n@dataclass\nclass RouterReplayConfig(BaseConfig):\n    \"\"\"Configuration for router replay in MoE models.\n\n    This configuration controls the routing behavior for Mixture of Experts (MoE) models,\n    allowing for deterministic training through route recording and replay.\n\n    Args:\n        mode (str): Router replay mode. Options: 'disabled', 'R2', 'R3'.\n            - 'disabled': No router replay functionality\n            - 'R2': Use Router Replay routing strategy\n            - 'R3': Use Rollout Router Replay routing strategy\n        record_file (Optional[str]): File path to save recorded routing decisions.\n            Required when mode is 'record', 'R2', or 'R3'.\n        replay_file (Optional[str]): File path to load recorded routing decisions for replay.\n            Required when mode is 'replay'.\n    \"\"\"\n\n    mode: str = \"disabled\"\n    record_file: Optional[str] = None\n    replay_file: Optional[str] = None\n\n    def __post_init__(self):\n        \"\"\"Validate router replay configuration.\"\"\"\n        valid_modes = [\"disabled\", \"R2\", \"R3\"]\n        if self.mode not in valid_modes:\n            raise ValueError(f\"Invalid router_replay mode: {self.mode}. Must be one of {valid_modes}\")\n\n\n@dataclass\nclass PolicyLossConfig(BaseConfig):\n    \"\"\"Configuration for policy loss computation.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        loss_mode (str): Loss function mode. Options: 'vanilla', 'clip-cov', 'kl-cov', 'gpg'.\n        clip_cov_ratio (float): Ratio of tokens to be clipped for clip-cov loss.\n        clip_cov_lb (float): Lower bound for clip-cov loss.\n        clip_cov_ub (float): Upper bound for clip-cov loss.\n        kl_cov_ratio (float): Ratio of tokens to be applied KL penalty for kl-cov loss.\n        ppo_kl_coef (float): KL divergence penalty coefficient.\n        rollout_correction (RolloutCorrectionConfig): Configuration for rollout correction.\n    \"\"\"\n\n    loss_mode: str = \"vanilla\"\n    clip_cov_ratio: float = 0.0002\n    clip_cov_lb: float = 1.0\n    clip_cov_ub: float = 5.0\n    kl_cov_ratio: float = 0.0002\n    ppo_kl_coef: float = 0.1\n    rollout_correction: RolloutCorrectionConfig = field(default_factory=RolloutCorrectionConfig)\n\n\n@dataclass\nclass ActorConfig(BaseConfig):\n    \"\"\"Configuration for actor model training.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Training strategy. Must be specified.\n        ppo_mini_batch_size (int): Mini-batch size for PPO training.\n        ppo_micro_batch_size (Optional[int]): Micro-batch size for PPO training.\n            If None, uses ppo_micro_batch_size_per_gpu.\n        ppo_micro_batch_size_per_gpu (Optional[int]): Micro-batch size per GPU for PPO training.\n        use_dynamic_bsz (bool): Whether to use dynamic batch sizing.\n        ppo_max_token_len_per_gpu (int): Maximum token length per GPU for PPO training.\n        clip_ratio (float): PPO clipping ratio for policy loss.\n        clip_ratio_low (float): Lower bound for PPO clipping ratio.\n        clip_ratio_high (float): Upper bound for PPO clipping ratio.\n        policy_loss (PolicyLossConfig): Configuration for policy loss computation.\n        clip_ratio_c (float): Clipping ratio for critic loss.\n        loss_agg_mode (str): Loss aggregation mode. Options: 'token-mean', 'sample-mean'.\n        loss_scale_factor (Optional[int]): Scale factor for 'seq-mean-token-sum-norm' loss aggregation mode.\n            If None, uses response_length. Set to a constant to ensure consistent normalization.\n        entropy_coeff (float): Entropy coefficient for regularization.\n        tau_pos (float): Positive tau for SAPO smoothing (>= 1.0 keeps rewards stable).\n        tau_neg (float): Negative tau for SAPO smoothing (> tau_pos for asymmetry).\n        use_kl_loss (bool): Whether to use KL divergence loss.\n        use_torch_compile (bool): Whether to use torch.compile for optimization.\n        kl_loss_coef (float): KL divergence loss coefficient.\n        kl_loss_type (str): Type of KL loss to use.\n        ppo_epochs (int): Number of PPO epochs per training step.\n        shuffle (bool): Whether to shuffle data during training.\n        checkpoint (CheckpointConfig): Configuration for checkpointing.\n        optim (OptimizerConfig): Configuration for optimizer.\n        use_fused_kernels (bool): Whether to use custom fused kernels (e.g., FlashAttention, fused MLP).\n        data_loader_seed (int): Seed for data loader. If None, uses global seed.\n        router_replay (RouterReplayConfig): Configuration for router replay in MoE models.\n    \"\"\"\n\n    _mutable_fields = BaseConfig._mutable_fields | {\n        \"ppo_mini_batch_size\",\n        \"ppo_micro_batch_size\",\n        \"ppo_micro_batch_size_per_gpu\",\n        \"ppo_infer_micro_batch_size_per_gpu\",\n        \"engine\",\n        \"model_config\",\n    }\n\n    strategy: str = MISSING\n    ppo_mini_batch_size: int = 256\n    ppo_micro_batch_size: Optional[int] = None  # deprecate\n    ppo_micro_batch_size_per_gpu: Optional[int] = None\n    ppo_infer_micro_batch_size_per_gpu: Optional[int] = None\n    use_dynamic_bsz: bool = False\n    ppo_max_token_len_per_gpu: int = 16384\n    ppo_infer_max_token_len_per_gpu: int = 16384\n    clip_ratio: float = 0.2\n    clip_ratio_low: float = 0.2\n    clip_ratio_high: float = 0.2\n    freeze_vision_tower: bool = False\n    policy_loss: PolicyLossConfig = field(default_factory=PolicyLossConfig)\n    clip_ratio_c: float = 3.0\n    loss_agg_mode: str = \"token-mean\"\n    loss_scale_factor: Optional[int] = None\n    entropy_coeff: float = 0\n    tau_pos: float = 1.0\n    tau_neg: float = 1.05\n    calculate_entropy: bool = False\n    use_kl_loss: bool = False\n    # Whether to enable PrefixGrouper-based shared-prefix forward\n    use_prefix_grouper: bool = False\n    use_torch_compile: bool = True\n    kl_loss_coef: float = 0.001\n    kl_loss_type: str = \"low_var_kl\"\n    ppo_epochs: int = 1\n    shuffle: bool = False\n    data_loader_seed: int = 1\n    checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)\n    optim: OptimizerConfig = field(default_factory=OptimizerConfig)\n    use_fused_kernels: bool = False\n    profiler: ProfilerConfig = field(default_factory=ProfilerConfig)\n    engine: BaseConfig = field(default_factory=BaseConfig)\n    rollout_n: int = MISSING  # must be override by sampling config\n    model_config: HFModelConfig = field(default_factory=BaseConfig)\n    router_replay: RouterReplayConfig = field(default_factory=RouterReplayConfig)\n\n    # Store global batch info for loss aggregation:\n    # dp_size: data parallel size\n    # batch_num_tokens: number of valid tokens in global batch\n    # global_batch_size: global batch size\n    global_batch_info: dict = field(default_factory=dict)\n\n    def __post_init__(self):\n        \"\"\"Validate actor configuration parameters.\"\"\"\n        assert self.strategy != MISSING\n        assert self.rollout_n != MISSING\n        if not self.use_dynamic_bsz:\n            if self.ppo_micro_batch_size is not None and self.ppo_micro_batch_size_per_gpu is not None:\n                raise ValueError(\n                    \"[actor] You have set both 'actor.ppo_micro_batch_size' AND 'actor.ppo_micro_batch_size_per_gpu'. \"\n                    \"Please remove 'actor.ppo_micro_batch_size' because only '*_ppo_micro_batch_size_per_gpu' is \"\n                    \"supported (the former is deprecated).\"\n                )\n            else:\n                assert not (self.ppo_micro_batch_size is None and self.ppo_micro_batch_size_per_gpu is None), (\n                    \"[actor] Please set at least one of 'actor.ppo_micro_batch_size' or \"\n                    \"'actor.ppo_micro_batch_size_per_gpu' if use_dynamic_bsz is not enabled.\"\n                )\n\n        valid_loss_agg_modes = [\n            \"token-mean\",\n            \"seq-mean-token-sum\",\n            \"seq-mean-token-mean\",\n            \"seq-mean-token-sum-norm\",\n        ]\n        if self.loss_agg_mode not in valid_loss_agg_modes:\n            raise ValueError(f\"Invalid loss_agg_mode: {self.loss_agg_mode}\")\n\n    def validate(self, n_gpus: int, train_batch_size: int, model_config: dict = None):\n        \"\"\"Validate actor configuration with runtime parameters.\"\"\"\n        if not self.use_dynamic_bsz:\n            if train_batch_size < self.ppo_mini_batch_size:\n                raise ValueError(\n                    f\"train_batch_size ({train_batch_size}) must be >= \"\n                    f\"actor.ppo_mini_batch_size ({self.ppo_mini_batch_size})\"\n                )\n\n            sp_size = getattr(self, \"ulysses_sequence_parallel_size\", 1)\n            if self.ppo_micro_batch_size is not None:\n                if self.ppo_mini_batch_size % self.ppo_micro_batch_size != 0:\n                    raise ValueError(\n                        f\"ppo_mini_batch_size ({self.ppo_mini_batch_size}) must be divisible by \"\n                        f\"ppo_micro_batch_size ({self.ppo_micro_batch_size})\"\n                    )\n                if self.ppo_micro_batch_size * sp_size < n_gpus:\n                    raise ValueError(\n                        f\"ppo_micro_batch_size ({self.ppo_micro_batch_size}) * \"\n                        f\"ulysses_sequence_parallel_size ({sp_size}) must be >= n_gpus ({n_gpus})\"\n                    )\n\n    @staticmethod\n    def _check_mutually_exclusive(mbs, mbs_per_gpu, name: str):\n        \"\"\"Validate mutually exclusive micro batch size configuration options.\"\"\"\n        param = \"ppo_micro_batch_size\"\n        param_per_gpu = f\"{param}_per_gpu\"\n\n        if mbs is None and mbs_per_gpu is None:\n            raise ValueError(f\"[{name}] Please set at least one of '{name}.{param}' or '{name}.{param_per_gpu}'.\")\n\n        if mbs is not None and mbs_per_gpu is not None:\n            raise ValueError(\n                f\"[{name}] You have set both '{name}.{param}' AND '{name}.{param_per_gpu}'. Please remove \"\n                f\"'{name}.{param}' because only '*_{param_per_gpu}' is supported (the former is deprecated).\"\n            )\n\n\n@dataclass\nclass McoreActorConfig(ActorConfig):\n    \"\"\"Configuration for Megatron actor models.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Training strategy set to 'megatron' for Megatron parallelism.\n        load_weight (bool): Whether to load model weights from checkpoint.\n        megatron (dict[str, Any]): Configuration for Megatron parallelism settings.\n        profile (dict[str, Any]): Configuration for profiling settings.\n    \"\"\"\n\n    strategy: str = \"megatron\"\n    load_weight: bool = True\n    megatron: McoreEngineConfig = field(default_factory=McoreEngineConfig)\n    profile: dict[str, Any] = field(default_factory=dict)\n    use_rollout_log_probs: bool = False\n\n    def __post_init__(self):\n        \"\"\"Validate FSDP actor configuration parameters.\"\"\"\n        super().__post_init__()\n        self.engine = self.megatron\n\n\n@dataclass\nclass FSDPActorConfig(ActorConfig):\n    \"\"\"Configuration for FSDP actor models.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Training strategy set to 'fsdp' for Fully Sharded Data Parallel.\n        grad_clip (float): Gradient clipping threshold.\n        ulysses_sequence_parallel_size (int): [DEPRECATED] Ulysses sequence parallel size for long sequences.\n        entropy_from_logits_with_chunking (bool): Whether to compute entropy from logits\n            with chunking for memory efficiency.\n        entropy_checkpointing (bool): Whether to use gradient checkpointing for entropy computation.\n        fsdp_config (dict[str, Any]): Configuration for FSDP settings.\n        use_remove_padding (bool): Whether to remove padding tokens in inputs during training\n    \"\"\"\n\n    strategy: str = \"fsdp\"\n    grad_clip: float = 1.0\n    ulysses_sequence_parallel_size: int = 1\n    entropy_from_logits_with_chunking: bool = False\n    entropy_checkpointing: bool = False\n    fsdp_config: FSDPEngineConfig = field(default_factory=FSDPEngineConfig)\n    use_remove_padding: bool = False\n    use_rollout_log_probs: bool = False\n    calculate_sum_pi_squared: bool = False\n    sum_pi_squared_checkpointing: bool = False\n    qat: QATConfig = field(default_factory=QATConfig)\n\n    def __post_init__(self):\n        \"\"\"Validate FSDP actor configuration parameters.\"\"\"\n        super().__post_init__()\n        self.engine = self.fsdp_config\n\n        # backward compatibility\n        if self.ulysses_sequence_parallel_size > 1:\n            self.fsdp_config.ulysses_sequence_parallel_size = self.ulysses_sequence_parallel_size\n\n    def validate(self, n_gpus: int, train_batch_size: int, model_config: dict = None):\n        \"\"\"Validate FSDP actor configuration with runtime parameters.\"\"\"\n        super().validate(n_gpus, train_batch_size, model_config)\n\n        if self.strategy in {\"fsdp\", \"fsdp2\"} and self.ulysses_sequence_parallel_size > 1:\n            if model_config and not model_config.get(\"use_remove_padding\", False):\n                raise ValueError(\n                    \"When using sequence parallelism for actor/ref policy, you must enable `use_remove_padding`.\"\n                )\n\n\n@dataclass\nclass VeOmniActorConfig(ActorConfig):\n    \"\"\"Configuration for VeOmni actor models.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Training strategy set to 'veomni' for VeOmni parallelism.\n        veomni (dict[str, Any]): Configuration for VeOmni settings.\n        use_remove_padding (bool): Whether to remove padding tokens in inputs during training\n    \"\"\"\n\n    strategy: str = \"veomni\"\n    veomni: VeOmniEngineConfig = field(default_factory=VeOmniEngineConfig)\n    use_remove_padding: bool = False\n    use_rollout_log_probs: bool = False\n\n    def __post_init__(self):\n        \"\"\"Validate VeOmni actor configuration parameters.\"\"\"\n        super().__post_init__()\n        self.engine = self.veomni\n\n\n@dataclass\nclass TorchTitanActorConfig(ActorConfig):\n    \"\"\"Configuration for TorchTitan actor models.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Training strategy set to 'torchtitan' for TorchTitan parallelism.\n        torchtitan (TorchtitanEngineConfig): Configuration for TorchTitan engine settings.\n        use_remove_padding (bool): Whether to remove padding tokens in inputs during training\n        use_rollout_log_probs (bool): Whether to use log probabilities from rollout engine\n    \"\"\"\n\n    strategy: str = \"torchtitan\"\n    torchtitan: TorchtitanEngineConfig = field(default_factory=TorchtitanEngineConfig)\n    use_remove_padding: bool = False\n    use_rollout_log_probs: bool = False\n\n    def __post_init__(self):\n        \"\"\"Validate TorchTitan actor configuration parameters.\"\"\"\n        super().__post_init__()\n        self.engine = self.torchtitan\n"
  },
  {
    "path": "verl/workers/config/critic.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nfrom omegaconf import MISSING\n\nfrom verl.base_config import BaseConfig\nfrom verl.trainer.config import BaseModelConfig, CheckpointConfig\nfrom verl.utils.profiler import ProfilerConfig\n\nfrom .engine import FSDPEngineConfig, McoreEngineConfig, TorchtitanEngineConfig\nfrom .model import HFModelConfig\nfrom .optimizer import OptimizerConfig\n\n__all__ = [\"CriticConfig\", \"FSDPCriticConfig\", \"McoreCriticConfig\", \"TorchTitanCriticConfig\", \"FSDPCriticModelCfg\"]\n\n\n@dataclass\nclass CriticConfig(BaseConfig):\n    \"\"\"Configuration for critic model training.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        strategy (str): Strategy used for critic model training (fsdp, fsdp2, megatron).\n        ppo_micro_batch_size_per_gpu (int): Local per-GPU micro batch size.\n        rollout_n (int): Number of rollouts per update (mirrors actor rollout_n).\n        optim (Dict[str, Any]): Optimizer configuration including lr, weight_decay, etc.\n        model (Dict[str, Any]): Model configuration including path, tokenizer_path, etc.\n        ppo_mini_batch_size (int): PPO mini-batch size per update.\n        ppo_micro_batch_size (Optional[int]): Global micro batch size (deprecated).\n        use_dynamic_bsz (bool): Whether to automatically adjust batch size at runtime.\n        ppo_max_token_len_per_gpu (int): Max tokens per GPU in one PPO batch.\n        forward_max_token_len_per_gpu (int): Max token length per GPU in forward pass.\n        ppo_epochs (int): Number of PPO epochs per batch.\n        shuffle (bool): Shuffle training data across PPO epochs.\n        cliprange_value (float): PPO value function clipping range.\n        loss_agg_mode (str): Loss aggregation mode.\n        checkpoint (Dict[str, Any]): Checkpoint configuration.\n        profiler (Dict[str, Any]): Profiler configuration.\n        enable (Optional[bool]): Whether to enable the critic.\n    \"\"\"\n\n    _mutable_fields = BaseConfig._mutable_fields | {\n        \"ppo_micro_batch_size_per_gpu\",\n        \"ppo_mini_batch_size\",\n        \"ppo_micro_batch_size\",\n        \"model_config\",\n    }\n\n    strategy: str = MISSING\n    ppo_micro_batch_size_per_gpu: Optional[int] = None\n    enable: Optional[bool] = None\n    rollout_n: int = 1\n    ppo_mini_batch_size: int = 1\n    use_dynamic_bsz: bool = False\n    ppo_max_token_len_per_gpu: int = 32768\n    # deprecate this\n    forward_max_token_len_per_gpu: int = 32768\n    ppo_infer_micro_batch_size_per_gpu: Optional[int] = None\n    ppo_infer_max_token_len_per_gpu: int = 32768\n    ppo_epochs: int = 1\n    data_loader_seed: int = 1\n    shuffle: bool = True\n    cliprange_value: float = 0.5\n    loss_agg_mode: str = \"token-mean\"\n    ppo_micro_batch_size: Optional[int] = None\n    engine: BaseConfig = field(default_factory=BaseConfig)\n    optim: OptimizerConfig = field(default_factory=OptimizerConfig)\n    # deprecate model to favor model_config\n    model: BaseModelConfig = field(default_factory=BaseModelConfig)\n    model_config: HFModelConfig = None\n    checkpoint: CheckpointConfig = field(default_factory=CheckpointConfig)\n    profiler: ProfilerConfig = field(default_factory=ProfilerConfig)\n\n    def __post_init__(self):\n        \"\"\"Validate critic configuration parameters.\"\"\"\n        assert self.strategy != MISSING\n\n        if self.model_config is None:\n            warnings.warn(\"using model in Critic Config is deprecated, please use model_config instead\", stacklevel=2)\n            self.model_config = HFModelConfig(\n                path=self.model.path,\n                tokenizer_path=self.model.tokenizer_path,\n                override_config=self.model.override_config,\n                external_lib=self.model.external_lib,\n                trust_remote_code=self.model.trust_remote_code,\n            )\n\n        if not self.use_dynamic_bsz:\n            self._check_mutually_exclusive(self.ppo_micro_batch_size, self.ppo_micro_batch_size_per_gpu, \"critic\")\n\n            if self.ppo_micro_batch_size is not None:\n                if self.ppo_mini_batch_size % self.ppo_micro_batch_size != 0:\n                    raise ValueError(\n                        f\"[critic] ppo_mini_batch_size ({self.ppo_mini_batch_size}) must be divisible by \"\n                        f\"ppo_micro_batch_size ({self.ppo_micro_batch_size})\"\n                    )\n\n    def validate(self, n_gpus: int, train_batch_size: int):\n        \"\"\"Validate critic configuration with runtime parameters.\n\n        Args:\n            n_gpus: Total number of GPUs available\n            train_batch_size: Training batch size from data config\n        \"\"\"\n        if not self.use_dynamic_bsz:\n            if train_batch_size < self.ppo_mini_batch_size:\n                raise ValueError(\n                    f\"train_batch_size ({train_batch_size}) must be >= \"\n                    f\"critic.ppo_mini_batch_size ({self.ppo_mini_batch_size})\"\n                )\n\n    @staticmethod\n    def _check_mutually_exclusive(mbs, mbs_per_gpu, name: str):\n        \"\"\"Validate mutually exclusive micro batch size configuration options.\n\n        Ensures that users don't set both deprecated micro_batch_size and\n        the new micro_batch_size_per_gpu parameters simultaneously.\n\n        Args:\n            mbs: Deprecated micro batch size parameter value.\n            mbs_per_gpu: New micro batch size per GPU parameter value.\n            name (str): Configuration section name for error messages.\n\n        Raises:\n            ValueError: If both parameters are set or neither is set.\n        \"\"\"\n        param = \"micro_batch_size\"\n        param_per_gpu = f\"{param}_per_gpu\"\n\n        if mbs is None and mbs_per_gpu is None:\n            raise ValueError(f\"[{name}] Please set at least one of '{name}.{param}' or '{name}.{param_per_gpu}'.\")\n\n        if mbs is not None and mbs_per_gpu is not None:\n            raise ValueError(\n                f\"[{name}] You have set both '{name}.{param}' AND '{name}.{param_per_gpu}'. Please remove \"\n                f\"'{name}.{param}' because only '*_{param_per_gpu}' is supported (the former is deprecated).\"\n            )\n\n\n@dataclass\nclass McoreCriticConfig(CriticConfig):\n    \"\"\"Configuration for Megatron-based critic model training.\n\n    The inheritance from CriticConfig provides all base critic configuration plus Megatron-specific settings.\n\n    Args:\n        nccl_timeout (int): NCCL timeout in seconds for distributed operations.\n        megatron (Dict[str, Any]): Megatron-specific parallelism settings.\n        load_weight (bool): Whether to load initial weights.\n    \"\"\"\n\n    strategy: str = \"megatron\"\n    nccl_timeout: int = 600\n    megatron: McoreEngineConfig = field(default_factory=McoreEngineConfig)\n    load_weight: bool = True\n\n    def validate(self, n_gpus: int, train_batch_size: int):\n        \"\"\"Validate Megatron critic configuration with runtime parameters.\"\"\"\n        super().validate(n_gpus, train_batch_size)\n\n\n@dataclass\nclass FSDPCriticConfig(CriticConfig):\n    \"\"\"Configuration for FSDP-based critic model training.\n\n    The inheritance from CriticConfig provides all base critic configuration plus FSDP-specific settings.\n\n    Args:\n        forward_micro_batch_size (int): Forward-only batch size during inference (global).\n        forward_micro_batch_size_per_gpu (int): Forward-only batch size during inference (per GPU).\n        ulysses_sequence_parallel_size (int): [DEPRECATED] Ulysses sequence parallel size for long sequences.\n        grad_clip (float): Gradient clipping for critic updates.\n    \"\"\"\n\n    _mutable_fields = CriticConfig._mutable_fields | {\n        \"forward_micro_batch_size\",\n        \"forward_micro_batch_size_per_gpu\",\n    }\n\n    strategy: str = \"fsdp\"\n    forward_micro_batch_size: int = 1\n    forward_micro_batch_size_per_gpu: int = 1\n    ulysses_sequence_parallel_size: int = 1\n    grad_clip: float = 1.0\n\n    def __post_init__(self):\n        \"\"\"Validate FSDP critic configuration parameters.\"\"\"\n        super().__post_init__()\n\n        if self.strategy in {\"fsdp\", \"fsdp2\"}:\n            if self.ulysses_sequence_parallel_size > 1:\n                if not self.model.get(\"use_remove_padding\", False):\n                    raise ValueError(\n                        \"When using sequence parallelism for critic, you must enable `use_remove_padding`.\"\n                    )\n\n    def validate(self, n_gpus: int, train_batch_size: int):\n        \"\"\"Validate FSDP critic configuration with runtime parameters.\"\"\"\n        super().validate(n_gpus, train_batch_size)\n\n        if not self.use_dynamic_bsz:\n            sp_size = self.ulysses_sequence_parallel_size\n            if self.ppo_micro_batch_size is not None:\n                if self.ppo_micro_batch_size * sp_size < n_gpus:\n                    raise ValueError(\n                        f\"critic.ppo_micro_batch_size ({self.ppo_micro_batch_size}) * \"\n                        f\"ulysses_sequence_parallel_size ({sp_size}) must be >= n_gpus ({n_gpus})\"\n                    )\n\n\n@dataclass\nclass TorchTitanCriticConfig(CriticConfig):\n    \"\"\"Configuration for TorchTitan-based critic model training.\n\n    The inheritance from CriticConfig provides all base critic configuration plus TorchTitan-specific settings.\n\n    Args:\n        strategy (str): Training strategy set to 'torchtitan' for TorchTitan parallelism.\n        torchtitan (TorchtitanEngineConfig): Configuration for TorchTitan engine settings.\n    \"\"\"\n\n    strategy: str = \"torchtitan\"\n    torchtitan: TorchtitanEngineConfig = field(default_factory=TorchtitanEngineConfig)\n\n    def __post_init__(self):\n        \"\"\"Validate TorchTitan critic configuration parameters.\"\"\"\n        super().__post_init__()\n        self.engine = self.torchtitan\n\n\n@dataclass\nclass FSDPCriticModelCfg(BaseModelConfig):\n    \"\"\"FSDP-enabled critic model configuration.\n    Inherits base critic settings and adds distributed-memory and LoRA options.\n\n    Args:\n        use_shm (bool): Whether to use shared memory for loading the model.\n        enable_activation_offload (bool): Offload activations to CPU to reduce GPU memory usage.\n        use_remove_padding (bool): Use remove-padding optimization (saves compute).\n        enable_gradient_checkpointing (bool): Enable gradient checkpointing for memory efficiency.\n        fsdp_config (FSDPEngineConfig): FSDP-specific configuration block.\n        lora_rank (int): Set to positive value to enable LoRA (e.g., 32).\n        lora_alpha (int): LoRA scaling factor.\n        target_modules (Union[str, List[str]]): LoRA target modules: \"all-linear\" or list of layer names.\n    \"\"\"\n\n    use_shm: bool = False\n    enable_activation_offload: bool = False\n    use_remove_padding: bool = False\n    enable_gradient_checkpointing: bool = True\n    fsdp_config: FSDPEngineConfig = field(default_factory=FSDPEngineConfig)\n    lora_rank: int = 0\n    lora_alpha: int = 16\n    target_modules: str | list[str] = \"all-linear\"\n    # TiledMLP configuration for memory-efficient MLP computation\n    tiled_mlp: dict = field(default_factory=lambda: {\"enabled\": False, \"num_shards\": 4})\n"
  },
  {
    "path": "verl/workers/config/engine.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Any, Callable, Literal, Optional\n\nfrom verl.base_config import BaseConfig\nfrom verl.trainer.config import CheckpointConfig\n\nfrom ...utils.profiler import ProfilerConfig\nfrom .model import HFModelConfig\nfrom .optimizer import OptimizerConfig\n\n__all__ = [\n    \"FSDPEngineConfig\",\n    \"McoreEngineConfig\",\n    \"TrainingWorkerConfig\",\n    \"TorchtitanEngineConfig\",\n    \"VeOmniEngineConfig\",\n    \"AutomodelEngineConfig\",\n    \"EngineConfig\",\n    \"EngineRouterReplayConfig\",\n    \"QATEngineConfig\",\n]\n\n\n# TODO: rename to RouterReplayConfig after removing the legacy implementation\n@dataclass\nclass EngineRouterReplayConfig(BaseConfig):\n    \"\"\"Configuration for router replay in MoE models.\n\n    This configuration controls the routing behavior for Mixture of Experts (MoE) models,\n    allowing for deterministic training through route recording and replay.\n\n    Args:\n        mode (str): Router replay mode. Options: 'disabled', 'R2', 'R3'.\n            - 'disabled': No router replay functionality\n            - 'R2': Use Router Replay routing strategy\n            - 'R3': Use Rollout Router Replay routing strategy\n        record_file (Optional[str]): File path to save recorded routing decisions.\n            Required when mode is 'record', 'R2', or 'R3'.\n        replay_file (Optional[str]): File path to load recorded routing decisions for replay.\n            Required when mode is 'replay'.\n    \"\"\"\n\n    mode: str = \"disabled\"\n    record_file: Optional[str] = None\n    replay_file: Optional[str] = None\n\n    def __post_init__(self):\n        \"\"\"Validate router replay configuration.\"\"\"\n        valid_modes = [\"disabled\", \"R2\", \"R3\"]\n        if self.mode not in valid_modes:\n            raise ValueError(f\"Invalid router_replay mode: {self.mode}. Must be one of {valid_modes}\")\n\n\n@dataclass\nclass EngineConfig(BaseConfig):\n    _mutable_fields = BaseConfig._mutable_fields | {\n        \"use_dynamic_bsz\",\n        \"max_token_len_per_gpu\",\n        \"micro_batch_size_per_gpu\",\n        \"infer_max_token_len_per_gpu\",\n        \"infer_micro_batch_size_per_gpu\",\n        \"use_fused_kernels\",\n        \"use_remove_padding\",\n        \"forward_only\",\n        \"param_offload\",\n    }\n    # whether to offload param\n    param_offload: bool = False\n    # whether to offload optimizer\n    optimizer_offload: bool = False\n    # whether to offload grad\n    grad_offload: bool = False\n    # whether the engine is forward only (e.g., ref policy)\n    forward_only: bool = False\n    # the strategy (backend)\n    strategy: str = None\n    # model dtype\n    dtype: str = \"bfloat16\"  # [\"bfloat16\", \"float16\"]\n    # whether to use dynamic bsz\n    use_dynamic_bsz: bool = True\n    # for training\n    max_token_len_per_gpu: int = None\n    micro_batch_size_per_gpu: int = None\n    # for inference\n    infer_max_token_len_per_gpu: int = None\n    infer_micro_batch_size_per_gpu: int = None\n    # whether use fuse lm head kernel\n    use_fused_kernels: bool = False\n    # TODO (this may conflict with the one in model config)\n    use_remove_padding: bool = True\n\n    seed: int = 42\n\n    full_determinism: bool = False\n    router_replay: EngineRouterReplayConfig = field(default_factory=EngineRouterReplayConfig)\n\n    def __post_init__(self):\n        pass\n        # TODO: turn on this check after we reorg config\n        # if self.use_dynamic_bsz:\n        #     assert self.max_token_len_per_gpu is not None\n        # else:\n        #     assert self.micro_batch_size_per_gpu is not None\n\n\n@dataclass\nclass McoreEngineConfig(EngineConfig):\n    \"\"\"Configuration for Megatron parallelism.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        param_offload (bool): Whether to offload parameters to CPU.\n        grad_offload (bool): Whether to offload gradients to CPU.\n        optimizer_offload (bool): Whether to offload optimizer states to CPU.\n        tensor_model_parallel_size (int): Tensor model parallel size.\n        expert_model_parallel_size (int): Expert model parallel size for MoE models.\n        expert_tensor_parallel_size (Optional[int]): Expert tensor parallel size for MoE models.\n        pipeline_model_parallel_size (int): Pipeline model parallel size.\n        virtual_pipeline_model_parallel_size (Optional[int]): Virtual pipeline model parallel size\n            for interleaved scheduling.\n        context_parallel_size (int): Context parallel size for long sequences.\n        sequence_parallel (bool): Whether to enable sequence parallelism.\n        use_distributed_optimizer (bool): Whether to use distributed optimizer.\n        use_dist_checkpointing (bool): Whether to use distributed checkpointing.\n        dist_checkpointing_path (Optional[str]): Path for distributed checkpointing.\n        dist_ckpt_optim_fully_reshardable (bool): Use fully reshardable optimizer checkpoints.\n        distrib_optim_fully_reshardable_mem_efficient (bool): Use memory-efficient fully reshardable format.\n        seed (int): Random seed for reproducibility.\n        override_ddp_config (dict[str, Any]): Override configuration for DDP.\n        override_transformer_config (dict[str, Any]): Override configuration for transformer.\n        use_mbridge (bool): Whether to use MBridge for communication.\n        dtype (str): Mixed precision training param dtype, default \"bfloat16\"\n    \"\"\"\n\n    # sequence_parallel is not listed as a frozen field for auto-correction purpose\n    _mutable_fields = EngineConfig._mutable_fields | {\"sequence_parallel\"}\n    # mcore parallelism\n    tensor_model_parallel_size: int = 1\n    expert_model_parallel_size: int = 1\n    expert_tensor_parallel_size: Optional[int] = None\n    pipeline_model_parallel_size: int = 1\n    virtual_pipeline_model_parallel_size: Optional[int] = None\n    context_parallel_size: int = 1\n    sequence_parallel: bool = True\n    use_distributed_optimizer: bool = True\n    use_dist_checkpointing: bool = False\n    dist_checkpointing_path: Optional[str] = None\n    dist_checkpointing_prefix: str = \"\"\n    dist_ckpt_optim_fully_reshardable: bool = False\n    distrib_optim_fully_reshardable_mem_efficient: bool = False\n    override_ddp_config: dict[str, Any] = field(default_factory=dict)\n    override_transformer_config: dict[str, Any] = field(default_factory=dict)\n    override_mcore_model_config: dict[str, Any] = field(default_factory=dict)\n    use_mbridge: bool = True\n    vanilla_mbridge: bool = True\n    strategy: str = \"megatron\"\n\n    def __post_init__(self) -> None:\n        super().__post_init__()\n        \"\"\"config validation logics go here\"\"\"\n        assert self.strategy == \"megatron\"\n        assert self.dtype in [\"bfloat16\", \"float16\"], f\"dtype {self.dtype} not supported\"\n        if self.tensor_model_parallel_size == 1:\n            warnings.warn(\"set sequence parallel to false as TP size is 1\", stacklevel=2)\n            self.sequence_parallel = False\n\n\n@dataclass\nclass QATEngineConfig(BaseConfig):\n    \"\"\"Configuration for QAT (Quantization-Aware Training) within an engine.\n\n    Args:\n        enable (bool): Whether to enable QAT, default False\n        mode (str): Quantization mode, \"w4a16\" or \"w4a4\", default \"w4a16\"\n        group_size (int): Group size for blockwise quantization, default 16\n        ignore_patterns (list[str]): Module name patterns to exclude from quantization\n        activation_observer (str): Observer strategy for activation global_scale (W4A4 only)\n        quantization_config_path (Optional[str]): Path to quantization config JSON for vLLM\n    \"\"\"\n\n    enable: bool = False\n    mode: str = \"w4a16\"\n    group_size: int = 16\n    ignore_patterns: list[str] = field(default_factory=lambda: [\"lm_head\", \"embed_tokens\", \"re:.*mlp.gate$\"])\n    activation_observer: str = \"static_minmax\"\n    quantization_config_path: Optional[str] = None\n\n\n@dataclass\nclass FSDPEngineConfig(EngineConfig):\n    \"\"\"Configuration for FSDP (Fully Sharded Data Parallel).\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy.\n        param_offload (bool): Whether to offload parameters to CPU, default False\n        optimizer_offload (bool): Whether to offload optimizer states to CPU, default False\n        offload_policy (bool): Whether to offload policy model parameters, default False\n        reshard_after_forward (bool): Whether to reshard parameters after forward pass, default True\n        fsdp_size (int): FSDP group size. -1 means use all available GPUs.\n        forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False\n        model_dtype (str): Model data type used to initialize the transformers model. default \"fp32\"\n        use_orig_params (bool): Whether to use original parameters when initialize FSDP1, default False\n        seed (int): Random seed for reproducibility.\n        full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results\n            in distributed training. Important: this will negatively impact performance, so only use it for\n            debugging.\n        mixed_precision (Optional[dict[str, Any]]): Mixed precision configuration for FSDP, default None\n        dtype (str): Mixed precision training param dtype, default \"bfloat16\"\n        qat (QATEngineConfig): QAT configuration, default disabled\n    \"\"\"\n\n    # ulysses_sequence_parallel_size is mutable for backward compatibility\n    _mutable_fields = EngineConfig._mutable_fields | {\"ulysses_sequence_parallel_size\"}\n\n    # fsdp specific flags\n    wrap_policy: dict[str, Any] = field(default_factory=dict)\n    offload_policy: bool = False\n    reshard_after_forward: bool = True\n    fsdp_size: int = -1\n    forward_prefetch: bool = False\n    model_dtype: str = \"fp32\"\n    use_orig_params: bool = False\n    mixed_precision: Optional[dict[str, Any]] = None\n    ulysses_sequence_parallel_size: int = 1\n    entropy_from_logits_with_chunking: bool = False\n    use_torch_compile: bool = True\n    entropy_checkpointing: bool = False\n    strategy: str = \"fsdp\"\n    qat: QATEngineConfig = field(default_factory=QATEngineConfig)\n\n    def __post_init__(self):\n        super().__post_init__()\n        assert self.strategy in [\"fsdp\", \"fsdp2\"], f\"strategy {self.strategy} not supported\"\n\n\n@dataclass\nclass VeOmniEngineConfig(EngineConfig):\n    \"\"\"Configuration for VeOmni.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy.\n        param_offload (bool): Whether to offload parameters to CPU, default False\n        optimizer_offload (bool): Whether to offload optimizer states to CPU, default False\n        offload_policy (bool): Whether to offload policy model parameters, default False\n        reshard_after_forward (bool): Whether to reshard parameters after forward pass, default True\n        fsdp_size (int): FSDP group size. -1 means use all available GPUs, default -1\n        ulysses_parallel_size (int): Ulysses sequence parallel size, default 1\n        expert_parallel_size (int): Expert parallel size, default 1\n        init_device (str): Device to initialize model weights.\n            1. `cpu`: Init parameters on CPU in rank0 only.\n            2. `cuda`: Init parameters on GPU.\n            3. `meta`: Init parameters on meta.\n            4. `npu`: Init parameters on Ascend NPU.\n            default \"meta\"\n        enable_full_shard (bool): Enable fully shard for FSDP training (ZeRO-3), default False\n        enable_fsdp_offload (bool): Enable CPU offload for FSDP1, default False\n        enable_reentrant (bool): Use reentrant gradient checkpointing, default False\n        attn_implementation (str): Attention implementation to use.\n            1. `eager`\n            2. `sdpa`\n            3. `flash_attention_2`\n            4. `flash_attention_3`\n            5. `veomni_flash_attention_2_with_sp`\n            6. `veomni_flash_attention_3_with_sp`\n            7. `native-sparse`\n            default \"flash_attention_2\"\n            Note: In case VeOmni add more attn_implementation, please check https://github.com/ByteDance-Seed/VeOmni/\n        moe_implementation (str): MoE implementation to use.\n            1. `eager`\n            2. `fused`\n            default \"fused\"\n            Note: In case VeOmni add more moe_implementation, please check https://github.com/ByteDance-Seed/VeOmni/\n        force_use_huggingface (bool): Force loading model from huggingface, default False\n        activation_gpu_limit (float): When enabling activation offload, `activation_gpu_limit` GB\n            activations are allowed to reserve on GPU, default 0.0\n        basic_modules (list[str]): List of basic modules to use, default None\n        forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False\n        model_dtype (str): Model data type used to initialize the transformers model. default \"fp32\"\n        use_orig_params (bool): Whether to use original parameters when initialize FSDP1, default False\n        seed (int): Random seed for reproducibility.\n        full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results\n            in distributed training. Important: this will negatively impact performance, so only use it for\n            debugging.\n        mixed_precision (Optional[dict[str, Any]]): Mixed precision configuration for FSDP, default None\n\n    \"\"\"\n\n    wrap_policy: dict[str, Any] = field(default_factory=dict)\n    offload_policy: bool = False\n    reshard_after_forward: bool = True\n    forward_prefetch: bool = False\n    use_orig_params: bool = False\n    entropy_from_logits_with_chunking: bool = False\n    use_torch_compile: bool = True\n    entropy_checkpointing: bool = False\n    strategy: str = \"veomni\"\n    fsdp_size: int = -1\n    ulysses_parallel_size: int = 1\n    expert_parallel_size: int = 1\n    seed: int = 42\n    full_determinism: bool = False\n    mixed_precision: bool = False\n    init_device: str = \"meta\"\n    enable_full_shard: bool = False\n    ckpt_manager: Literal[\"dcp\"] = \"dcp\"\n    load_checkpoint_path: Optional[str] = None\n    enable_fsdp_offload: bool = False\n    enable_reentrant: bool = False\n    attn_implementation: str = \"flash_attention_2\"\n    moe_implementation: str = \"fused\"\n    force_use_huggingface: bool = False\n    activation_gpu_limit: float = 0.0\n    basic_modules: Optional[list[str]] = field(default_factory=list)\n\n    def __post_init__(self):\n        super().__post_init__()\n        assert self.strategy in [\"veomni\"], f\"strategy {self.strategy} not supported\"\n\n\n@dataclass\nclass TorchtitanEngineConfig(EngineConfig):\n    \"\"\"Configuration for Torchtitan.\n\n    The inheritance from BaseConfig provides omegaconf.DictConfig-like interface for a dataclass config.\n\n    Args:\n        wrap_policy (Dict[str, Any]): Configuration for FSDP wrap policy.\n        reshard_after_forward (Literal[\"default\", \"always\", \"never\"]): The policy for applying\n            `reshard_after_forward` within an FSDP setup, default \"default\"\n        forward_prefetch (bool): Whether to prefetch parameters for next forward pass, default False\n        use_orig_params (bool): Whether to use original parameters when initialize FSDP, default False\n        mixed_precision (bool): Mixed precision configuration for FSDP, default False\n        offload_policy (bool): Whether to offload policy model parameters, default False\n        data_parallel_size (int): Data parallel group size, default 1\n        data_parallel_replicate_size (int): Data parallel replicate size, default 1\n        data_parallel_shard_size (int): Data parallel shard degree, default 1\n        tensor_parallel_size (int): Tensor parallel size, default 1\n        expert_parallel_size (int): Expert parallel size, default 1\n        expert_tensor_parallel_size (int): Expert tensor parallel size, default 1\n        pipeline_parallel_size (int): Pipeline parallel size, default 1\n        context_parallel_size (int): Context parallel size, default 1\n        attn_type (str): Attention type for torchtitan's model (e.g., \"sdpa\", \"flex\", \"varlen\"),\n            default \"flex\"\n        strategy (str): Strategy to use for distributed training, default \"torchtitan\"\n        seed (int): Random seed for reproducibility.\n        full_determinism (bool): If true, enable_full_determinism is called to ensure reproducible results\n            in distributed training. Important: this will negatively impact performance, so only use it for\n            debugging.\n\n    \"\"\"\n\n    wrap_policy: dict[str, Any] = field(default_factory=dict)\n    reshard_after_forward: Literal[\"default\", \"always\", \"never\"] = \"default\"\n    forward_prefetch: bool = False\n    use_orig_params: bool = False\n    mixed_precision: bool = False\n    offload_policy: bool = False\n    use_torch_compile: bool = True\n    entropy_from_logits_with_chunking: bool = False\n    entropy_checkpointing: bool = False\n    data_parallel_size: int = 1\n    data_parallel_replicate_size: int = 1\n    data_parallel_shard_size: int = 1\n    tensor_parallel_size: int = 1\n    expert_parallel_size: int = 1\n    expert_tensor_parallel_size: int = 1\n    pipeline_parallel_size: int = 1\n    context_parallel_size: int = 1\n    attn_type: str = \"flex\"\n    max_seq_len: Optional[int] = None\n    strategy: str = \"torchtitan\"\n    seed: int = 42\n    full_determinism: bool = False\n\n    def __post_init__(self):\n        super().__post_init__()\n        assert self.strategy in [\"torchtitan\"], f\"strategy {self.strategy} not supported\"\n\n\n@dataclass\nclass AutomodelEngineConfig(EngineConfig):\n    \"\"\"Configuration for Automodel (nemo_automodel) backend.\n\n    The Automodel backend uses NeMoAutoModelForCausalLM for model loading and\n    supports FSDP2, MegatronFSDP, and DDP distributed strategies with optional\n    TP, CP, and EP parallelism.\n\n    Args:\n        strategy (str): Backend strategy identifier, must be \"automodel\".\n        distributed_strategy (str): Distributed training strategy: \"fsdp2\", \"megatron_fsdp\", or \"ddp\".\n        tp_size (int): Tensor parallel size.\n        pp_size (int): Pipeline parallel size (only pp_size=1 supported initially).\n        cp_size (int): Context parallel size.\n        ep_size (int): Expert parallel size for MoE models.\n        dp_replicate_size (int): Data-parallel replicate size for HSDP. 1 = pure sharding.\n        sequence_parallel (bool): Enable sequence parallelism in the TP plan.\n        defer_fsdp_grad_sync (bool): Defer FSDP gradient sync to the final micro-batch.\n        activation_checkpointing (bool): Whether to enable activation checkpointing.\n        enable_fp8 (bool): Whether to enable FP8 training.\n        enable_compile (bool): Whether to enable torch.compile for the model.\n        model_dtype (str): Model data type for loading weights. \"fp32\" loads in float32\n            (matching FSDP golden), \"auto\" uses the dtype from the model config.\n        attn_implementation (str): Attention implementation to use (\"sdpa\", \"flash_attention_2\", \"eager\", \"te\").\n\n    Backend settings (nemo_automodel BackendConfig):\n        backend_config (dict): Dict of kwargs passed directly to\n            nemo_automodel.components.models.common.BackendConfig(**backend_config).\n            Controls how model layers are implemented (TE vs PyTorch) and MoE dispatch.\n            See automodel.yaml for all predefined keys with defaults.\n            Key fields:\n                attn (str): Attention backend. \"te\" = TransformerEngine fused attention,\n                    \"sdpa\" = PyTorch scaled dot-product attention. Default: \"sdpa\".\n                linear (str): Linear layer backend. \"te\" = TE fused linear (with FP8 support),\n                    \"torch\" = standard PyTorch linear. Default: \"te\".\n                rms_norm (str): RMSNorm backend. \"te\" = TE fused RMSNorm, \"torch\" = PyTorch,\n                    \"torch_fp32\" = PyTorch in FP32 (better numerical stability for MoE).\n                    Default: \"torch_fp32\".\n                rope_fusion (bool): Enable fused RoPE kernel (requires CP=1). Default: true.\n                experts (str): MoE expert computation backend.\n                    \"gmm\" = grouped_gemm (requires pip install grouped_gemm),\n                    \"torch_mm\" = torch._grouped_mm (no external dependency),\n                    \"te\" = TE GroupedLinear. Default: \"gmm\".\n                dispatcher (str): MoE token dispatch strategy.\n                    \"torch\" = standard all-gather + local compute,\n                    \"deepep\" = DeepEP optimized all-to-all (higher throughput).\n                    Default: \"torch\".\n                    Note: \"deepep\" with experts=\"gmm\" matches the legacy enable_deepep=True behavior.\n                enable_fsdp_optimizations (bool): Enable FSDP-specific optimizations in Automodel.\n                    Default: false.\n                enable_hf_state_dict_adapter (bool): Enable HuggingFace state dict adapter for\n                    checkpoint compatibility. Default: true.\n                fake_balanced_gate (bool): Use fake balanced gating for debugging. Default: false.\n                fake_gate_noise (float): Noise added to fake balanced gate. Default: 0.0.\n                gate_precision: Gate computation precision. Default: null (auto).\n            Full reference: nemo_automodel/components/models/common/backend_config.py\n\n    MoE / Expert Parallelism settings:\n        moe_config (dict): Dict of kwargs passed directly to\n            nemo_automodel.components.moe.parallelizer.MoEParallelizerConfig(**moe_config).\n            Controls MoE parallelization behavior within FSDP2.\n            See automodel.yaml for all predefined keys with defaults.\n            Key fields:\n                ignore_router_for_ac (bool): Exclude router from activation checkpointing.\n                    Default: false.\n                reshard_after_forward (bool): Reshard expert params after forward pass\n                    (trades compute for memory). Default: false.\n                lm_head_precision: Precision for the LM head. Default: null (auto).\n                wrap_outer_model (bool): Whether to FSDP-wrap the outermost model module.\n                    Default: true.\n            Full reference: nemo_automodel/components/moe/parallelizer.py\n\n    Mixed precision policy (FSDP2):\n        mp_param_dtype (str): Parameter dtype for FSDP2 mixed precision policy.\n        mp_reduce_dtype (str): Reduce dtype for FSDP2 mixed precision policy.\n        mp_output_dtype (str): Output dtype for FSDP2 mixed precision policy.\n\n    Entropy computation:\n        entropy_from_logits_with_chunking (bool): Whether to use chunked entropy computation.\n        use_torch_compile (bool): Whether to use torch.compile for entropy computation.\n        entropy_checkpointing (bool): Whether to use checkpointing for entropy computation.\n    \"\"\"\n\n    strategy: str = \"automodel\"\n    distributed_strategy: str = \"fsdp2\"\n    # Parallelism sizes\n    tp_size: int = 1\n    pp_size: int = 1\n    cp_size: int = 1\n    ep_size: int = 1\n    dp_replicate_size: int = 1\n    sequence_parallel: bool = False\n    defer_fsdp_grad_sync: bool = True\n    # Model settings\n    activation_checkpointing: bool = False\n    enable_fp8: bool = False\n    enable_compile: bool = False\n    model_dtype: str = \"fp32\"\n    attn_implementation: str = \"flash_attention_2\"\n    # Backend settings\n    backend_config: dict = field(default_factory=dict)\n    # MoE settings\n    moe_config: dict = field(default_factory=dict)\n    # Mixed precision policy\n    mp_param_dtype: str = \"bf16\"\n    mp_reduce_dtype: str = \"fp32\"\n    mp_output_dtype: str = \"bf16\"\n    # Entropy computation\n    entropy_from_logits_with_chunking: bool = False\n    use_torch_compile: bool = True\n    entropy_checkpointing: bool = False\n\n    def __post_init__(self):\n        super().__post_init__()\n        assert self.strategy == \"automodel\", f\"strategy must be 'automodel', got {self.strategy}\"\n        assert self.distributed_strategy in [\"fsdp2\", \"megatron_fsdp\", \"ddp\"], (\n            f\"distributed_strategy {self.distributed_strategy} not supported\"\n        )\n        assert self.pp_size == 1, \"Pipeline parallelism (pp_size > 1) is not yet supported for automodel backend\"\n\n\n@dataclass\nclass TrainingWorkerConfig(BaseConfig):\n    model_type: str = None  # model type (language_model/value_model)\n    model_config: HFModelConfig = None\n    engine_config: EngineConfig = None\n    optimizer_config: OptimizerConfig = None\n    checkpoint_config: CheckpointConfig = None\n    profiler_config: ProfilerConfig = None\n    # automatically select engine and optimizer function.\n    # This function takes model config and the device name as parameter.\n    # Users can pass in a higher-order function to take more parameters\n    auto_select_engine_optim_fn: Callable[[\"HFModelConfig\", str], tuple[\"EngineConfig\", \"OptimizerConfig\"]] = None\n"
  },
  {
    "path": "verl/workers/config/megatron_peft.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"PEFT configuration of Megatron for VERL.\"\"\"\n\n\ndef get_peft_cls(model_config, bridge, provider, dtype=None):\n    \"\"\"Get PEFT class from model config.\n\n    Args:\n        model_config: Model configuration object.\n        bridge: Megatron-Bridge AutoBridge instance.\n        provider: Provider instance.\n\n    Returns:\n        PEFT configuration object (LoRAConfig, CanonicalLoRAConfig, DoRAConfig) or None.\n    \"\"\"\n\n    peft_cls = None\n    if not hasattr(model_config, \"lora\"):\n        return peft_cls\n\n    lora_cfg = model_config.lora\n    # Only enable if rank > 0\n    if lora_cfg.get(\"rank\", 0) <= 0:\n        return peft_cls\n\n    assert bridge is not None and provider is not None, \"LoRA/PEFT only supported via Megatron-Bridge\"\n\n    from verl.models.mcore.bridge import CanonicalLoRA, DoRA, LoRA, VLMLoRA\n\n    lora_dtype = lora_cfg.get(\"dtype\", dtype)\n    if lora_dtype is not None:\n        from verl.utils.torch_dtypes import PrecisionType\n\n        lora_dtype = PrecisionType.to_dtype(lora_dtype)\n\n    lora_type = lora_cfg.get(\"type\", \"lora\")\n    if lora_type == \"lora\":\n        peft_cls = LoRA(\n            target_modules=lora_cfg.get(\"target_modules\", [\"linear_qkv\", \"linear_proj\", \"linear_fc1\", \"linear_fc2\"]),\n            dim=lora_cfg.get(\"rank\"),\n            alpha=lora_cfg.get(\"alpha\", 32),\n            dropout=lora_cfg.get(\"dropout\", 0.0),\n            dropout_position=lora_cfg.get(\"dropout_position\", \"pre\"),\n            lora_A_init_method=lora_cfg.get(\"lora_A_init_method\", \"xavier\"),\n            lora_B_init_method=lora_cfg.get(\"lora_B_init_method\", \"zero\"),\n            a2a_experimental=lora_cfg.get(\"a2a_experimental\", False),\n            lora_dtype=lora_dtype,\n            exclude_modules=lora_cfg.get(\"exclude_modules\", []),\n        )\n    if lora_type == \"vlm_lora\":\n        peft_cls = VLMLoRA(\n            target_modules=lora_cfg.get(\"target_modules\", [\"linear_qkv\", \"linear_proj\", \"linear_fc1\", \"linear_fc2\"]),\n            dim=lora_cfg.get(\"rank\"),\n            alpha=lora_cfg.get(\"alpha\", 32),\n            dropout=lora_cfg.get(\"dropout\", 0.0),\n            dropout_position=lora_cfg.get(\"dropout_position\", \"pre\"),\n            lora_A_init_method=lora_cfg.get(\"lora_A_init_method\", \"xavier\"),\n            lora_B_init_method=lora_cfg.get(\"lora_B_init_method\", \"zero\"),\n            a2a_experimental=lora_cfg.get(\"a2a_experimental\", False),\n            lora_dtype=lora_dtype,\n            freeze_vision_model=lora_cfg.get(\"freeze_vision_model\", True),\n            freeze_vision_projection=lora_cfg.get(\"freeze_vision_projection\", True),\n            freeze_language_model=lora_cfg.get(\"freeze_language_model\", True),\n            exclude_modules=lora_cfg.get(\"exclude_modules\", []),\n        )\n    elif lora_type == \"canonical_lora\":\n        peft_cls = CanonicalLoRA(\n            target_modules=lora_cfg.get(\n                \"target_modules\",\n                [\n                    \"linear_q\",\n                    \"linear_k\",\n                    \"linear_v\",\n                    \"linear_proj\",\n                    \"linear_fc1_up\",\n                    \"linear_fc1_gate\",\n                    \"linear_fc2\",\n                ],\n            ),\n            dim=lora_cfg.get(\"rank\"),\n            alpha=lora_cfg.get(\"alpha\", 32),\n            dropout=lora_cfg.get(\"dropout\", 0.0),\n            dropout_position=lora_cfg.get(\"dropout_position\", \"pre\"),\n            lora_A_init_method=lora_cfg.get(\"lora_A_init_method\", \"xavier\"),\n            lora_B_init_method=lora_cfg.get(\"lora_B_init_method\", \"zero\"),\n            exclude_modules=lora_cfg.get(\"exclude_modules\", []),\n        )\n    elif lora_type == \"dora\":\n        peft_cls = DoRA(\n            target_modules=lora_cfg.get(\"target_modules\", [\"linear_qkv\", \"linear_proj\", \"linear_fc1\", \"linear_fc2\"]),\n            dim=lora_cfg.get(\"rank\"),\n            alpha=lora_cfg.get(\"alpha\", 32),\n            dropout=lora_cfg.get(\"dropout\", 0.0),\n            dropout_position=lora_cfg.get(\"dropout_position\", \"pre\"),\n            lora_A_init_method=lora_cfg.get(\"lora_A_init_method\", \"xavier\"),\n            lora_B_init_method=lora_cfg.get(\"lora_B_init_method\", \"zero\"),\n            exclude_modules=lora_cfg.get(\"exclude_modules\", []),\n        )\n\n    print(\n        f\"Enabling {lora_type.upper()} with rank={lora_cfg.get('rank')}, \"\n        f\"alpha={lora_cfg.get('alpha')}, dropout={lora_cfg.get('dropout')}\"\n    )\n    return peft_cls\n\n\n__all__ = [\n    \"get_peft_cls\",\n]\n"
  },
  {
    "path": "verl/workers/config/model.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom dataclasses import dataclass, field\nfrom typing import Any, Optional\n\nfrom omegaconf import MISSING\nfrom transformers import AutoConfig\n\nfrom verl.base_config import BaseConfig\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.fs import copy_to_local\nfrom verl.utils.import_utils import import_external_libs\nfrom verl.utils.model import get_generation_config, update_model_config\n\n__all__ = [\"HFModelConfig\", \"MtpConfig\"]\n\n\n@dataclass\nclass MtpConfig(BaseConfig):\n    \"\"\"\n    Configuration for MTP model.\n\n    enable: Enable loading and saving of MTP parameters, but do not use them\n\n    enable_train: Whether to enable using MTP parameters during training\n    enable_rollout: Whether to enable using MTP parameters during rollout\n\n    Training parameters:\n        detach_encoder: Whether to detach encoder parameters during MTP training\n        mtp_loss_scaling_factor: Loss scaling factor during MTP training\n\n    vLLM rollout parameters:\n        method: \"mtp\"\n        num-speculative-tokens: 1\n\n    SGLang rollout parameters:\n        speculative-algorithm: EAGLE\n        speculative-num-steps: 3\n        speculative-eagle-topk: 1\n        speculative-num-draft-tokens: 4\n    \"\"\"\n\n    enable: bool = False\n    enable_train: bool = False\n    enable_rollout: bool = False\n\n    detach_encoder: bool = False\n    mtp_loss_scaling_factor: float = 0.1\n\n    speculative_algorithm: str = \"EAGLE\"\n    speculative_num_steps: int = 3\n    speculative_eagle_topk: int = 1\n    speculative_num_draft_tokens: int = 4\n\n    method: str = \"mtp\"\n    num_speculative_tokens: int = 1\n\n\n@dataclass\nclass HFModelConfig(BaseConfig):\n    # note that we separate model_path, model_config_path and tokenizer_path in case they are different\n    _mutable_fields = {\n        \"hf_config_path\",\n        \"tokenizer_path\",\n        \"hf_config\",\n        \"generation_config\",\n        \"tokenizer\",\n        \"processor\",\n        \"local_path\",\n        \"architectures\",\n        \"local_hf_config_path\",\n        \"local_tokenizer_path\",\n        \"mtp\",\n    }\n\n    path: str = MISSING\n    local_path: Optional[str] = None\n    hf_config_path: Optional[str] = None\n    local_hf_config_path: Optional[str] = None\n    tokenizer_path: Optional[str] = None\n    local_tokenizer_path: Optional[str] = None\n\n    # whether to load tokenizer. This is useful when we only want to load model config\n    load_tokenizer: bool = True\n\n    hf_config: Any = None\n    generation_config: Any = None\n    tokenizer: Any = None\n    processor: Any = None\n\n    # whether to use shared memory\n    use_shm: bool = False\n    trust_remote_code: bool = False\n\n    # custom chat template for the model\n    custom_chat_template: Optional[str] = None\n\n    external_lib: Optional[str] = None\n\n    override_config: dict = field(default_factory=dict)\n\n    enable_gradient_checkpointing: bool = True\n    enable_activation_offload: bool = False\n\n    use_remove_padding: bool = True\n\n    # TODO: unify fsdp and megatron lora config\n    # fsdp lora related. We may setup a separate config later\n    lora_rank: int = 0\n    lora_alpha: int = 16\n    target_modules: Optional[Any] = \"all-linear\"  # allow both \"all-linear\" and [\"q_proj\",\"k_proj\"]\n    target_parameters: Optional[list[str]] = None  # for lora adapter on nn.Parameter\n\n    exclude_modules: Optional[str] = None\n\n    # megatron lora config\n    lora: dict[str, Any] = field(default_factory=dict)\n\n    # path to pre-trained LoRA adapter to load for continued training\n    lora_adapter_path: Optional[str] = None\n    use_liger: bool = False\n\n    use_fused_kernels: bool = False\n    fused_kernel_options: dict = field(default_factory=dict)\n\n    # TiledMLP configuration for memory-efficient MLP computation\n    tiled_mlp: dict = field(default_factory=lambda: {\"enabled\": False, \"num_shards\": 4})\n\n    architectures: Optional[list[str]] = None\n\n    mtp: MtpConfig = field(default_factory=MtpConfig)\n\n    def __post_init__(self):\n        import_external_libs(self.external_lib)\n\n        if self.hf_config_path is None:\n            self.hf_config_path = self.path\n        if self.tokenizer_path is None:\n            self.tokenizer_path = self.path\n\n        self.local_path = copy_to_local(self.path, use_shm=self.use_shm)\n\n        # construct tokenizer\n        if self.load_tokenizer:\n            self.local_tokenizer_path = copy_to_local(self.tokenizer_path, use_shm=self.use_shm)\n            self.tokenizer = hf_tokenizer(self.local_tokenizer_path, trust_remote_code=self.trust_remote_code)\n            self.processor = hf_processor(self.local_tokenizer_path, trust_remote_code=self.trust_remote_code)\n\n        # For base models (e.g. Qwen3.5-2b-Base), the processor may not have a chat_template\n        # while the tokenizer does. Sync it so that processor.apply_chat_template() works.\n        if (\n            self.processor is not None\n            and not getattr(self.processor, \"chat_template\", None)\n            and getattr(self.tokenizer, \"chat_template\", None)\n        ):\n            self.processor.chat_template = self.tokenizer.chat_template\n\n        if self.custom_chat_template is not None:\n            if self.processor is not None:\n                self.processor.chat_template = self.custom_chat_template\n            else:\n                self.tokenizer.chat_template = self.custom_chat_template\n\n        self.local_hf_config_path = copy_to_local(self.hf_config_path, use_shm=self.use_shm)\n        self.generation_config = get_generation_config(\n            self.local_hf_config_path, trust_remote_code=self.trust_remote_code\n        )\n\n        # construct hf_config\n        attn_implementation = self.override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        self.hf_config = AutoConfig.from_pretrained(\n            self.local_hf_config_path, trust_remote_code=self.trust_remote_code, attn_implementation=attn_implementation\n        )\n\n        override_config_kwargs = {}\n\n        if self.tokenizer is not None:\n            override_config_kwargs.update(\n                {\n                    \"bos_token_id\": self.tokenizer.bos_token_id,\n                    \"eos_token_id\": self.tokenizer.eos_token_id,\n                    \"pad_token_id\": self.tokenizer.pad_token_id,\n                }\n            )\n\n        # TODO: (vermouth1992). self.config.model in megatron differs from that of fsdp in the override_config.\n        override_config = (\n            self.override_config[\"model_config\"] if \"model_config\" in self.override_config else self.override_config\n        )\n        override_config_kwargs.update(override_config)\n        update_model_config(self.hf_config, override_config_kwargs=override_config_kwargs)\n\n        self.share_embeddings_and_output_weights = getattr(self.hf_config, \"tie_word_embeddings\", False)\n\n        # get model architectures\n        self.architectures = getattr(self.hf_config, \"architectures\", None)\n        assert self.architectures is not None and len(self.architectures) == 1, (\n            \"Expect only one architecture, got {}\".format(self.architectures)\n        )\n\n        # per model patch\n        if getattr(self.hf_config, \"model_type\", None) == \"kimi_vl\":\n            self.hf_config.text_config.topk_method = \"greedy\"\n\n        # Ensure target_modules is a str or list[str] (only if not None)\n        if self.target_modules is not None:\n            if not isinstance(self.target_modules, (str | list)):\n                raise TypeError(\n                    \"target_modules must be a string or a list of strings, \"\n                    f\"but got {type(self.target_modules).__name__}\"\n                )\n            if isinstance(self.target_modules, list):\n                for x in self.target_modules:\n                    if not isinstance(x, str):\n                        raise TypeError(\n                            f\"All elements in target_modules list must be strings, but found {type(x).__name__}\"\n                        )\n\n    def get_processor(self):\n        return self.processor if self.processor is not None else self.tokenizer\n"
  },
  {
    "path": "verl/workers/config/optimizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport warnings\nfrom dataclasses import dataclass\nfrom typing import Optional\n\nfrom omegaconf import MISSING\n\nfrom verl.base_config import BaseConfig\n\n__all__ = [\n    \"OptimizerConfig\",\n    \"FSDPOptimizerConfig\",\n    \"McoreOptimizerConfig\",\n    \"build_optimizer\",\n    \"VeOmniOptimizerConfig\",\n    \"TorchtitanOptimizerConfig\",\n    \"AutomodelOptimizerConfig\",\n]\n\n\n@dataclass\nclass OptimizerConfig(BaseConfig):\n    \"\"\"Base optimizer configuration.\n\n    Args:\n        lr (float): learning rate. Must be specified.\n        lr_warmup_steps_ratio (float): Warmup steps ratio; total steps will be injected at runtime.\n        total_training_steps (int): Total training steps (must be overridden at runtime).\n        weight_decay (float): Weight decay factor.\n        lr_warmup_steps (Optional[int]): Number of warmup steps; None delegates to lr_warmup_steps_ratio.\n    \"\"\"\n\n    _mutable_fields = {\"clip_grad\", \"total_training_steps\", \"lr_warmup_steps\"}\n\n    lr: float = 1e-3\n    lr_warmup_steps_ratio: float = 0.0\n    total_training_steps: int = -1\n    weight_decay: float = 0.01\n    lr_warmup_steps: Optional[int] = -1\n    betas: tuple[float, float] = (0.9, 0.999)\n    clip_grad: float = 1.0\n    # deprecate grad_clip\n    grad_clip: Optional[float] = None\n\n    def __post_init__(self):\n        assert self.lr != MISSING\n        if self.grad_clip is not None:\n            warnings.warn(\"`grad_clip` is deprecated, use `clip_grad` instead.\", DeprecationWarning, stacklevel=2)\n            self.clip_grad = self.grad_clip\n\n\n@dataclass\nclass VeOmniOptimizerConfig(OptimizerConfig):\n    \"\"\"VeOmni optimizer configuration extending base OptimizerConfig.\n\n    Args:\n        optimizer (str): Optimizer name; default is \"adamw\".\n        lr (float): Learning rate.\n        lr_min (float): Minimum learning rate.\n        lr_start (float): Starting learning rate for warmup.\n        lr_decay_ratio (float): LR decay ratio.\n        lr_scheduler_type (str): LR scheduler type: \"constant\" or \"cosine\".\n    \"\"\"\n\n    _mutable_fields = OptimizerConfig._mutable_fields.copy()\n\n    optimizer: str = \"adamw\"\n    lr_min: float = 0.0\n    lr_start: float = 0.0\n    lr_decay_ratio: float = 1.0\n    lr_scheduler_type: str = \"constant\"\n    override_optimizer_config: Optional[dict] = None\n\n\n@dataclass\nclass FSDPOptimizerConfig(OptimizerConfig):\n    \"\"\"FSDP optimizer configuration extending base OptimizerConfig.\n\n    Args:\n        optimizer (str): Optimizer class name (e.g., \"AdamW\", \"AdamW8bit\", \"_AdamW\").\n        optimizer_impl (str): Module path to import optimizer from (e.g., \"torch.optim\", \"torchao.optim\",\n            \"bitsandbytes.optim\").\n        lr (float): Learning rate.\n        min_lr_ratio (Optional[float]): Minimum LR ratio for cosine schedule.\n        lr_scheduler_type (str): LR scheduler type: \"constant\" or \"cosine\".\n        num_cycles (float): Number of cosine cycles in LR schedule.\n        zero_indexed_step (bool): Whether the LR schedule uses 0-indexed steps. If True (default),\n            step counting starts at 0. If False, step counting starts at 1.\n    \"\"\"\n\n    _mutable_fields = OptimizerConfig._mutable_fields.copy()\n    _mutable_fields.add(\"lr_scheduler_type\")\n\n    optimizer: str = \"AdamW\"\n    optimizer_impl: str = \"torch.optim\"\n    min_lr_ratio: Optional[float] = None\n    # deprecate warmup_style\n    warmup_style: Optional[str] = None\n    lr_scheduler_type: str = \"constant\"\n    num_cycles: float = 0.5\n    override_optimizer_config: Optional[dict] = None\n    zero_indexed_step: bool = True\n\n    def __post_init__(self):\n        if self.warmup_style is not None:\n            assert self.warmup_style in [\"constant\", \"cosine\"]\n            warnings.warn(\n                \"`warmup_style` is deprecated, use `lr_scheduler_type` instead.\", DeprecationWarning, stacklevel=2\n            )\n            self.lr_scheduler_type = self.warmup_style\n        assert self.lr_scheduler_type in [\"constant\", \"cosine\"]\n        return super().__post_init__()\n\n\n@dataclass\nclass McoreOptimizerConfig(OptimizerConfig):\n    \"\"\"Mcore optimizer configuration extending base OptimizerConfig.\n\n    Args:\n        optimizer (str): Optimizer name; default is \"adam\".\n        lr (float): Learning rate.\n        clip_grad (float): Gradient clipping norm.\n        lr_warmup_init (float): Initial learning rate for warmup; defaults to 0.0.\n        lr_decay_steps (Optional[int]): Number of decay steps.\n        lr_decay_style (str): LR decay style: \"constant\", \"linear\", \"cosine\", or \"inverse_square_root\".\n        min_lr (float): Minimum learning rate.\n        weight_decay_incr_style (str): Weight decay increment style: \"constant\" or \"cosine\".\n        lr_wsd_decay_style (str): Weight-standard-deviation decay style: \"constant\", \"exponential\", or \"cosine\".\n        lr_wsd_decay_steps (Optional[int]): Number of steps for weight-standard-deviation decay.\n        use_checkpoint_opt_param_scheduler (bool): Whether to use checkpoint optimizer parameter scheduler.\n    \"\"\"\n\n    optimizer: str = \"adam\"\n    lr_warmup_init: float = 0.0\n    lr_decay_steps: Optional[int] = None\n    lr_decay_style: str = \"linear\"\n    min_lr: float = 0.0\n    weight_decay_incr_style: str = \"constant\"\n    lr_wsd_decay_style: str = \"exponential\"\n    lr_wsd_decay_steps: Optional[int] = None\n    use_checkpoint_opt_param_scheduler: bool = False\n    override_optimizer_config: Optional[dict] = None\n\n\n@dataclass\nclass TorchtitanOptimizerConfig(OptimizerConfig):\n    \"\"\"Torchtitan optimizer configuration extending base OptimizerConfig.\n\n    Args:\n        name (str): Optimizer name; default is \"AdamW\".\n        eps (float): Epsilon value for AdamW optimizer, default 1e-8.\n        decay_type (str): Weight decay type: \"linear\", \"sqrt\", or \"cosine\".\n        min_lr_factor (float): Minimum learning rate factor.\n    \"\"\"\n\n    name: str = \"AdamW\"\n    eps: float = 1e-8\n    decay_type: str = \"linear\"\n    min_lr_factor: float = 0.0\n\n\n@dataclass\nclass AutomodelOptimizerConfig(OptimizerConfig):\n    \"\"\"Automodel optimizer configuration extending base OptimizerConfig.\n\n    Uses the same optimizer building mechanism as FSDP (dynamic import from optimizer_impl).\n    LR scheduling is handled by Automodel's OptimizerParamScheduler.\n\n    Args:\n        optimizer (str): Optimizer class name (e.g., \"AdamW\").\n        optimizer_impl (str): Module path to import optimizer from (e.g., \"torch.optim\").\n        lr (float): Learning rate (maps to max_lr in OptimizerParamScheduler).\n        init_lr_ratio (Optional[float]): Initial LR ratio for warmup start (init_lr = lr * init_lr_ratio).\n        min_lr_ratio (Optional[float]): Minimum LR ratio after decay (min_lr = lr * min_lr_ratio).\n        lr_scheduler_type (str): LR decay style: \"constant\", \"cosine\", \"linear\", or \"inverse-square-root\".\n        wd_incr_style (str): Weight decay increment style: \"constant\", \"linear\", or \"cosine\".\n        num_cycles (float): Kept for backward compatibility (unused by Automodel scheduler).\n        zero_indexed_step (bool): Kept for backward compatibility (unused by Automodel scheduler).\n    \"\"\"\n\n    _mutable_fields = OptimizerConfig._mutable_fields.copy()\n    _mutable_fields.add(\"lr_scheduler_type\")\n\n    optimizer: str = \"AdamW\"\n    optimizer_impl: str = \"torch.optim\"\n    init_lr_ratio: Optional[float] = 0.1\n    min_lr_ratio: Optional[float] = 0.01\n    lr_scheduler_type: str = \"cosine\"\n    wd_incr_style: str = \"constant\"\n    num_cycles: float = 0.5\n    zero_indexed_step: bool = True\n    # Common optimizer kwargs\n    eps: float = 1e-8\n    master_weights: bool = False\n    store_param_remainders: bool = False\n    exp_avg_dtype: Optional[str] = None  # \"fp32\", \"bf16\", \"fp16\", or \"torch.float32\" etc.\n    exp_avg_sq_dtype: Optional[str] = None  # \"fp32\", \"bf16\", \"fp16\", or \"torch.float32\" etc.\n    master_weight_dtype: Optional[str] = None  # \"fp32\", \"bf16\", \"fp16\", or \"torch.float32\" etc.\n    override_optimizer_config: Optional[dict] = None\n\n    def __post_init__(self):\n        assert self.lr_scheduler_type in [\"constant\", \"cosine\", \"linear\", \"inverse-square-root\"]\n        return super().__post_init__()\n\n\ndef build_optimizer(parameters, config: FSDPOptimizerConfig):\n    \"\"\"Build an optimizer based on the configuration.\n\n    Dynamically imports and instantiates an optimizer class from the specified module.\n\n    Args:\n        parameters: Model parameters to optimize\n        config: FSDPOptimizerConfig with optimizer settings\n\n    Returns:\n        Optimizer instance\n\n    Examples:\n        # PyTorch AdamW\n        config.optimizer_impl = \"torch.optim\"\n        config.optimizer = \"AdamW\"\n\n        # TorchAO AdamW with bf16 stochastic rounding\n        config.optimizer_impl = \"torchao.optim\"\n        config.optimizer = \"_AdamW\"\n        config.override_optimizer_config = {\"bf16_stochastic_round\": True}\n\n        # BitsAndBytes AdamW 8bit\n        config.optimizer_impl = \"bitsandbytes.optim\"\n        config.optimizer = \"AdamW8bit\"\n    \"\"\"\n    import importlib\n\n    optimizer_args = {\n        \"lr\": config.lr,\n        \"weight_decay\": config.weight_decay,\n    }\n\n    optimizer_name_lower = config.optimizer.lower()\n    if \"adam\" in optimizer_name_lower or \"ademamix\" in optimizer_name_lower:\n        optimizer_args[\"betas\"] = config.betas\n\n    if config.override_optimizer_config is not None:\n        optimizer_args.update(config.override_optimizer_config)\n\n    try:\n        module = importlib.import_module(config.optimizer_impl)\n        optimizer_cls = getattr(module, config.optimizer)\n    except ImportError as e:\n        raise ImportError(\n            f\"Failed to import module '{config.optimizer_impl}'. Make sure the package is installed. Error: {e}\"\n        ) from e\n    except AttributeError as e:\n        raise AttributeError(\n            f\"Optimizer '{config.optimizer}' not found in module '{config.optimizer_impl}'. \"\n            f\"Available optimizers: {dir(module)}\"\n        ) from e\n\n    return optimizer_cls(parameters, **optimizer_args)\n"
  },
  {
    "path": "verl/workers/config/reward.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nfrom verl.base_config import BaseConfig\nfrom verl.trainer.config.config import ModuleConfig\n\nfrom .rollout import RolloutConfig\n\n__all__ = [\"SandboxFusionConfig\", \"RewardConfig\", \"RewardModelConfig\"]\n\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@dataclass\nclass RewardManagerConfig(BaseConfig):\n    \"\"\"Configuration for reward manager.\n\n        A reward manager defines the mechanism of computing rule-based reward and handling different reward sources.\n\n    Args:\n        source (str): Source of the reward manager. Options: ``\"register\"``, ``\"importlib\"``. Default: ``\"register\"``.\n        name (str, optional):\n            - When ``source`` is ``\"register\"``, the name is used in `get_reward_manager_cls(name)``.\n                See ``verl/experimental/reward/reward_manager.py`` for options. Default: ``\"naive\"``.\n            - When ``source`` is ``\"importlib\"``, the name is used in ``getattr(module, name)``,\n                e.g., ``\"DAPORewardManager\"``.\n        module (ModuleConfig, optional): Optional configuration for the external module defining the reward manager,\n    \"\"\"\n\n    source: str = \"register\"\n    name: str = \"naive\"\n    module: Optional[ModuleConfig] = field(default_factory=ModuleConfig)\n\n    def __post_init__(self):\n        super().__post_init__()\n        if self.source == \"register\":\n            from verl.experimental.reward_loop.reward_manager.registry import REWARD_MANAGER\n\n            assert self.name in REWARD_MANAGER, (\n                f\"Reward manager is not registered: {self.name=} ,{REWARD_MANAGER.keys()=}\"\n            )\n        elif self.source == \"importlib\":\n            # NOTE: The existence is not checked since it depends on which machine the config is initialized on.\n            assert self.module is not None and self.module.path is not None, (\n                \"When source is importlib, module.path should be set.\"\n            )\n\n\n@dataclass\nclass SandboxFusionConfig(BaseConfig):\n    \"\"\"Configuration for cloud/local sandbox fusion.\n\n    Args:\n        url (Optional[str]): Cloud/local function URL for sandbox execution.\n        max_concurrent (int): Max concurrent requests allowed to sandbox.\n        memory_limit_mb (int): Max memory limit for each sandbox process in MB.\n    \"\"\"\n\n    url: Optional[str] = None\n    max_concurrent: int = 64\n    memory_limit_mb: int = 1024\n\n\n@dataclass\nclass RewardModelConfig(BaseConfig):\n    _mutable_fields = BaseConfig._mutable_fields\n\n    enable: bool = False\n    enable_resource_pool: bool = False\n    n_gpus_per_node: int = 0\n    nnodes: int = 0\n    model_path: Optional[str] = None\n    inference: RolloutConfig = field(default_factory=RolloutConfig)\n\n\n@dataclass\nclass RewardConfig(BaseConfig):\n    _mutable_fields = BaseConfig._mutable_fields\n\n    # reward manager args\n    num_workers: int = 8\n    reward_manager: RewardManagerConfig = field(default_factory=RewardManagerConfig)\n\n    # reward model args\n    reward_model: RewardModelConfig = field(default_factory=RewardModelConfig)\n\n    # sandbox fusion args\n    sandbox_fusion: SandboxFusionConfig = field(default_factory=SandboxFusionConfig)\n"
  },
  {
    "path": "verl/workers/config/rollout.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport warnings\nfrom dataclasses import dataclass, field\nfrom typing import Optional\n\nfrom omegaconf import MISSING\n\nfrom verl.base_config import BaseConfig\nfrom verl.utils.profiler import ProfilerConfig\nfrom verl.workers.config.model import MtpConfig\n\n__all__ = [\n    \"SamplingConfig\",\n    \"MultiTurnConfig\",\n    \"CustomAsyncServerConfig\",\n    \"AgentLoopConfig\",\n    \"TraceConfig\",\n    \"ServerConfig\",\n    \"PrometheusConfig\",\n    \"RolloutConfig\",\n    \"CheckpointEngineConfig\",\n]\n\n\n@dataclass\nclass SamplingConfig(BaseConfig):\n    temperature: float = 1.0\n    top_k: int = -1\n    top_p: float = 1.0\n    do_sample: bool = True\n    n: int = 1\n\n\n@dataclass\nclass MultiTurnConfig(BaseConfig):\n    _mutable_fields = {\"max_assistant_turns\", \"max_user_turns\"}\n\n    enable: bool = False\n    max_assistant_turns: Optional[int] = None\n    tool_config_path: Optional[str] = None\n    max_user_turns: Optional[int] = None\n    max_parallel_calls: int = 1\n    max_tool_response_length: int = 256\n    tool_response_truncate_side: str = \"middle\"\n    interaction_config_path: Optional[str] = None\n    use_inference_chat_template: bool = False\n    tokenization_sanity_check_mode: str = \"strict\"\n    format: str = \"hermes\"\n    num_repeat_rollouts: Optional[int] = None\n\n\n@dataclass\nclass CustomAsyncServerConfig(BaseConfig):\n    path: Optional[str] = None\n    name: Optional[str] = None\n\n\n@dataclass\nclass AgentLoopConfig(BaseConfig):\n    num_workers: int = 8\n    default_agent_loop: str = \"single_turn_agent\"\n    agent_loop_config_path: Optional[str] = None\n    custom_async_server: CustomAsyncServerConfig = field(default_factory=CustomAsyncServerConfig)\n    # Fully qualified class name for custom AgentLoopManager (e.g., \"mypackage.module.MyManager\").\n    # Security: This class will be dynamically imported via importlib. Only use trusted class paths.\n    agent_loop_manager_class: Optional[str] = None\n\n\n@dataclass\nclass TraceConfig(BaseConfig):\n    project_name: Optional[str] = None\n    experiment_name: Optional[str] = None\n    backend: Optional[str] = None\n    token2text: bool = False\n    max_samples_per_step_per_worker: Optional[int] = None\n\n    def __post_init__(self):\n        if self.max_samples_per_step_per_worker is not None and self.max_samples_per_step_per_worker < 0:\n            raise ValueError(\"`max_samples_per_step_per_worker` must be a non-negative integer or null.\")\n\n\n@dataclass\nclass ServerConfig(BaseConfig):\n    \"\"\"\n    Configuration for SGLang server when running in server mode\n    \"\"\"\n\n    timeout: float = 60.0\n    max_attempts: int = 3\n    retry_delay: float = 2.0\n    max_connections: int = 1000\n    max_start_wait_time: float = 300.0\n\n\n@dataclass\nclass PrometheusConfig(BaseConfig):\n    \"\"\"\n    Configuration for Prometheus server\n    \"\"\"\n\n    # whether enable prometheus on server mode rollout\n    enable: bool = False\n    # Port number that Prometheus listens on, default is 9090\n    port: int = 9090\n    # Path to Prometheus configuration file\n    file: str = \"/tmp/ray/session_latest/metrics/prometheus/prometheus.yml\"\n    # Specify served_model_name to avoid displaying overly long model paths in Grafana\n    served_model_name: Optional[str] = None\n\n\n@dataclass\nclass CheckpointEngineConfig(BaseConfig):\n    \"\"\"\n    Configuration for checkpoint engine to update weights from trainer to rollout\n    \"\"\"\n\n    # Backend for checkpoint engine: naive, nccl, nixl, hccl\n    backend: Optional[str] = \"naive\"\n    # Bucket size in MB to transfer multiple weights at one time\n    update_weights_bucket_megabytes: int = 2048\n    # Additional keyword arguments for checkpoint engine\n    engine_kwargs: dict = field(default_factory=dict)\n\n\n@dataclass\nclass RolloutConfig(BaseConfig):\n    _mutable_fields = {\"max_model_len\", \"load_format\", \"expert_parallel_size\", \"moe_tensor_parallel_size\"}\n\n    name: Optional[str] = MISSING\n    mode: str = \"async\"\n    nnodes: int = 0\n    n_gpus_per_node: int = 8\n\n    temperature: float = 1.0\n    top_k: int = -1\n    top_p: float = 1.0\n    do_sample: bool = True\n    n: int = 1\n    repetition_penalty: float = 1.0\n\n    # Early termination threshold for multi-turn rollout in sglang.\n    # Abort remaining requests when (1 - over_sample_rate) * total_requests are completed.\n    over_sample_rate: float = 0.0\n\n    prompt_length: int = 512\n    response_length: int = 512\n\n    dtype: str = \"bfloat16\"\n    gpu_memory_utilization: float = 0.5\n    ignore_eos: bool = False\n    enforce_eager: bool = True\n    cudagraph_capture_sizes: Optional[list] = None\n    free_cache_engine: bool = True\n    data_parallel_size: int = 1\n    expert_parallel_size: int = 1\n    tensor_model_parallel_size: int = 2\n    pipeline_model_parallel_size: int = 1\n    moe_tensor_parallel_size: int = 1\n    max_num_batched_tokens: int = 8192\n    logprobs_mode: Optional[str] = \"processed_logprobs\"\n    scheduling_policy: Optional[str] = \"fcfs\"\n\n    # TODO: enable train_kwargs\n    # train_sampling_config: SamplingConfig = field(default_factory=SamplingConfig)\n\n    val_kwargs: SamplingConfig = field(default_factory=SamplingConfig)\n\n    max_model_len: Optional[int] = None\n    max_num_seqs: int = 1024\n\n    # note that the logprob computation should belong to the actor\n    log_prob_micro_batch_size: Optional[int] = None\n    log_prob_micro_batch_size_per_gpu: Optional[int] = None\n    log_prob_use_dynamic_bsz: bool = False\n    log_prob_max_token_len_per_gpu: int = 16384\n\n    disable_log_stats: bool = True\n\n    multi_stage_wake_up: bool = False\n    engine_kwargs: dict = field(default_factory=dict)\n\n    calculate_log_probs: bool = False\n\n    agent: AgentLoopConfig = field(default_factory=AgentLoopConfig)\n\n    trace: TraceConfig = field(default_factory=TraceConfig)\n\n    multi_turn: MultiTurnConfig = field(default_factory=MultiTurnConfig)\n\n    # Server configuration for sglang server mode\n    server: ServerConfig = field(default_factory=ServerConfig)\n\n    # Use Prometheus to collect and monitor rollout statistics\n    prometheus: PrometheusConfig = field(default_factory=PrometheusConfig)\n\n    # Extension point for custom configurations\n    custom: Optional[dict] = None\n\n    # Checkpoint Engine config for update weights from trainer to rollout\n    checkpoint_engine: CheckpointEngineConfig = field(default_factory=CheckpointEngineConfig)\n\n    skip_rollout: bool = False\n\n    skip_dump_dir: str = \"/tmp/rollout_dump\"\n\n    profiler: Optional[ProfilerConfig] = None\n\n    enable_chunked_prefill: bool = True\n\n    enable_prefix_caching: bool = True\n\n    load_format: str = \"dummy\"\n\n    layered_summon: bool = False\n\n    layer_name_map: dict = field(default_factory=dict)\n\n    sglang_engine_mode: str = \"local\"\n\n    limit_images: Optional[int] = None\n\n    skip_tokenizer_init: bool = False\n\n    quantization: Optional[str] = None\n\n    quantization_config_file: Optional[str] = None\n\n    enable_rollout_routing_replay: bool = False\n\n    enable_sleep_mode: bool = True\n\n    mtp: MtpConfig = field(default_factory=MtpConfig)\n\n    qat: Optional[dict] = None\n\n    def __post_init__(self):\n        \"\"\"Validate the rollout config\"\"\"\n        # Deprecation warning for mode field - only async mode is supported\n        if self.mode == \"sync\":\n            raise ValueError(\n                \"Rollout mode 'sync' has been removed. Please set \"\n                \"`actor_rollout_ref.rollout.mode=async` or remove the mode setting entirely.\"\n            )\n        if self.mode != \"async\":\n            warnings.warn(\n                f\"Unknown rollout mode '{self.mode}'. Only 'async' mode is supported. \"\n                \"The 'mode' field is deprecated and will be removed in a future version.\",\n                DeprecationWarning,\n                stacklevel=2,\n            )\n\n        if self.name != \"trtllm\" and self.expert_parallel_size > 1:\n            assert self.expert_parallel_size == (self.tensor_model_parallel_size * self.data_parallel_size), (\n                \"expert_parallel_size must be equal to tensor_model_parallel_size * data_parallel_size\"\n            )\n\n        if self.moe_tensor_parallel_size is not None and self.moe_tensor_parallel_size > 1:\n            assert self.name == \"trtllm\", \"moe_tensor_parallel_size is only supported for trtllm\"\n\n        if self.name == \"trtllm\":\n            # If either expert_parallel_size or moe_tensor_parallel_size is at default 1,\n            # convert to None so TensorRT-LLM treats it as unspecified.\n            # When both unspecified: moe_ep_size=1, moe_tp_size=moe_world_size (no EP, all TP).\n            # When only one set: the other is auto-derived from tensor_model_parallel_size.\n            if self.expert_parallel_size is not None and self.expert_parallel_size == 1:\n                self.expert_parallel_size = None\n            if self.moe_tensor_parallel_size is not None and self.moe_tensor_parallel_size == 1:\n                self.moe_tensor_parallel_size = None\n            if self.expert_parallel_size is not None and self.moe_tensor_parallel_size is not None:\n                assert self.moe_tensor_parallel_size * self.expert_parallel_size == self.tensor_model_parallel_size, (\n                    \"moe_tensor_parallel_size * expert_parallel_size must equal tensor_model_parallel_size \"\n                    f\"(got {self.moe_tensor_parallel_size} * {self.expert_parallel_size} = \"\n                    f\"{self.moe_tensor_parallel_size * self.expert_parallel_size}, \"\n                    f\"tensor_model_parallel_size={self.tensor_model_parallel_size})\"\n                )\n\n        if self.pipeline_model_parallel_size > 1:\n            if self.name == \"vllm\" or self.name == \"sglang\" or self.name == \"trtllm\":\n                raise NotImplementedError(\n                    f\"Current rollout {self.name=} not implemented pipeline_model_parallel_size > 1 yet.\"\n                )\n"
  },
  {
    "path": "verl/workers/critic/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .base import BasePPOCritic\nfrom .dp_critic import DataParallelPPOCritic\n\n__all__ = [\"BasePPOCritic\", \"DataParallelPPOCritic\"]\n"
  },
  {
    "path": "verl/workers/critic/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nBase class for a critic\n\"\"\"\n\nfrom abc import ABC, abstractmethod\n\nimport torch\n\nfrom verl import DataProto\n\n__all__ = [\"BasePPOCritic\"]\n\n\nclass BasePPOCritic(ABC):\n    def __init__(self, config):\n        super().__init__()\n        self.config = config\n\n    @abstractmethod\n    def compute_values(self, data: DataProto) -> torch.Tensor:\n        \"\"\"Compute values\"\"\"\n        pass\n\n    @abstractmethod\n    def update_critic(self, data: DataProto):\n        \"\"\"Update the critic\"\"\"\n        pass\n"
  },
  {
    "path": "verl/workers/critic/dp_critic.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nImplement a multiprocess PPOCritic\n\"\"\"\n\nimport logging\nimport os\n\nimport torch\nimport torch.distributed\nfrom torch import nn, optim\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\n\nfrom verl import DataProto\nfrom verl.trainer.ppo import core_algos\nfrom verl.utils.attention_utils import index_first_axis, pad_input, rearrange, unpad_input\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.fsdp_utils import FSDPModule, fsdp2_clip_grad_norm_\nfrom verl.utils.profiler import GPUMemoryLogger\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import prepare_dynamic_batch, restore_dynamic_batch\nfrom verl.utils.torch_functional import masked_mean\nfrom verl.utils.ulysses import gather_outputs_and_unpad, ulysses_pad_and_slice_inputs\nfrom verl.workers.critic import BasePPOCritic\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass DataParallelPPOCritic(BasePPOCritic):\n    def __init__(self, config, critic_module: nn.Module, critic_optimizer: optim.Optimizer):\n        super().__init__(config=config)\n        self.critic_module = critic_module\n        self.critic_optimizer = critic_optimizer\n        self.use_remove_padding = self.config.model.get(\"use_remove_padding\", False)\n        print(f\"Critic use_remove_padding={self.use_remove_padding}\")\n\n        self.ulysses_sequence_parallel_size = self.config.get(\"ulysses_sequence_parallel_size\", 1)\n        self.device_name = get_device_name()\n\n    def _forward_micro_batch(self, micro_batch):\n        response_length = micro_batch[\"responses\"].size(-1)\n        multi_modal_inputs = {}\n        if \"multi_modal_inputs\" in micro_batch.keys():\n            from verl.utils.model import extract_multi_modal_inputs\n\n            multi_modal_inputs = extract_multi_modal_inputs(micro_batch[\"multi_modal_inputs\"])\n\n        with torch.autocast(device_type=self.device_name, dtype=torch.bfloat16):\n            input_ids = micro_batch[\"input_ids\"]\n            batch, seqlen = input_ids.shape\n            attention_mask = micro_batch[\"attention_mask\"]\n            position_ids = micro_batch[\"position_ids\"]\n            if position_ids.dim() == 3:  # qwen2vl mrope\n                position_ids = position_ids.transpose(0, 1)\n\n            if self.use_remove_padding:\n                input_ids_rmpad, indices, *_ = unpad_input(\n                    input_ids.unsqueeze(-1), attention_mask\n                )  # input_ids_rmpad (total_nnz, ...)\n                input_ids_rmpad = input_ids_rmpad.transpose(0, 1)  # (1, total_nnz)\n\n                # unpad the position_ids to align the rotary\n                if position_ids.dim() == 3:\n                    position_ids_rmpad = (\n                        index_first_axis(rearrange(position_ids, \"c b s ... -> (b s) c ...\"), indices)\n                        .transpose(0, 1)\n                        .unsqueeze(1)\n                    )  # (4, bsz, seqlen) -> (4, 1, bsz * seqlen)\n                else:\n                    position_ids_rmpad = index_first_axis(\n                        rearrange(position_ids.unsqueeze(-1), \"b s ... -> (b s) ...\"), indices\n                    ).transpose(0, 1)\n\n                # pad and slice the inputs if sp > 1\n                if self.ulysses_sequence_parallel_size > 1:\n                    input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(\n                        input_ids_rmpad, position_ids_rmpad, sp_size=self.ulysses_sequence_parallel_size\n                    )\n\n                # only pass input_ids and position_ids to enable flash_attn_varlen\n                output = self.critic_module(\n                    input_ids=input_ids_rmpad,\n                    attention_mask=None,\n                    position_ids=position_ids_rmpad,\n                    **multi_modal_inputs,\n                    use_cache=False,\n                )  # prevent model thinks we are generating\n\n                if hasattr(self.critic_module, \"v_head\"):\n                    # For trl.AutoModelForCausalLMWithValueHead\n                    values_rmpad = output[2].squeeze(0).unsqueeze(-1)\n                else:\n                    values_rmpad = output.logits\n                    values_rmpad = values_rmpad.squeeze(0)  # (total_nnz)\n\n                # gather output if sp > 1\n                if self.ulysses_sequence_parallel_size > 1:\n                    values_rmpad = gather_outputs_and_unpad(\n                        values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size\n                    )\n\n                # pad it back\n                values = pad_input(values_rmpad, indices=indices, batch=batch, seqlen=seqlen).squeeze(-1)\n                values = values[:, -response_length - 1 : -1]\n            else:\n                output = self.critic_module(\n                    input_ids=input_ids,\n                    attention_mask=attention_mask,\n                    position_ids=position_ids,\n                    **multi_modal_inputs,\n                    use_cache=False,\n                )  # prevent model thinks we are generating\n                if hasattr(self.critic_module, \"v_head\"):\n                    # For trl.AutoModelForCausalLMWithValueHead\n                    values = output[2]\n                else:\n                    values = output.logits\n                values = values[:, -response_length - 1 : -1].squeeze(-1)\n            return values\n\n    def _optimizer_step(self):\n        assert self.config.grad_clip is not None\n\n        if isinstance(self.critic_module, FSDP):\n            grad_norm = self.critic_module.clip_grad_norm_(self.config.grad_clip)\n        elif isinstance(self.critic_module, FSDPModule):\n            grad_norm = fsdp2_clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(self.critic_module.parameters(), max_norm=self.config.grad_clip)\n\n        # if grad_norm is not finite, skip the update\n        if not torch.isfinite(grad_norm):\n            print(f\"WARN: grad_norm is not finite: {grad_norm}\")\n            self.critic_optimizer.zero_grad()\n        else:\n            self.critic_optimizer.step()\n        return grad_norm\n\n    @GPUMemoryLogger(role=\"dp critic\", logger=logger)\n    def compute_values(self, data: DataProto) -> torch.Tensor:\n        self.critic_module.eval()\n        micro_batch_size = data.meta_info[\"micro_batch_size\"]\n        use_dynamic_bsz = data.meta_info[\"use_dynamic_bsz\"]\n        has_multi_modal_inputs = \"multi_modal_inputs\" in data.non_tensor_batch.keys()\n        select_keys = (\n            [\"responses\", \"input_ids\", \"response_mask\", \"attention_mask\", \"position_ids\"]\n            if \"response_mask\" in data.batch\n            else [\"responses\", \"input_ids\", \"attention_mask\", \"position_ids\"]\n        )\n        non_tensor_select_keys = [\"multi_modal_inputs\"] if has_multi_modal_inputs else []\n\n        data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys)\n\n        if use_dynamic_bsz:\n            max_token_len = data.meta_info[\"max_token_len\"] * self.ulysses_sequence_parallel_size\n            micro_batches, batch_idx_list = prepare_dynamic_batch(data, max_token_len=max_token_len)\n        else:\n            micro_batches = data.split(micro_batch_size)\n\n        values_lst = []\n        for micro_batch in micro_batches:\n            micro_batch = micro_batch.to(get_device_id())\n            model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}\n            with torch.no_grad():\n                values = self._forward_micro_batch(model_inputs)\n            values_lst.append(values)\n        values = torch.concat(values_lst, dim=0)\n\n        if use_dynamic_bsz:\n            values = restore_dynamic_batch(values, batch_idx_list)\n\n        if \"response_mask\" in data.batch:\n            response_mask = data.batch[\"response_mask\"]\n            response_mask = response_mask.to(values.device)\n            values = values * response_mask  # Only action tokens have values\n        return values\n\n    @GPUMemoryLogger(role=\"dp critic\", logger=logger)\n    def update_critic(self, data: DataProto):\n        # make sure we are in training mode\n        self.critic_module.train()\n        metrics = {\n            \"critic/vf_loss\": 0.0,\n        }\n\n        select_keys = [\"input_ids\", \"responses\", \"response_mask\", \"attention_mask\", \"position_ids\", \"values\", \"returns\"]\n        has_multi_modal_inputs = \"multi_modal_inputs\" in data.non_tensor_batch.keys()\n        non_tensor_select_keys = [\"multi_modal_inputs\"] if has_multi_modal_inputs else []\n\n        data = data.select(batch_keys=select_keys, non_tensor_batch_keys=non_tensor_select_keys)\n\n        # Split to make minibatch iterator for updating the actor\n        # See PPO paper for details. https://arxiv.org/abs/1707.06347\n        mini_batches = data.split(self.config.ppo_mini_batch_size)\n\n        for _ in range(self.config.ppo_epochs):\n            for batch_idx, mini_batch in enumerate(mini_batches):\n                if self.config.use_dynamic_bsz:\n                    max_token_len = self.config.ppo_max_token_len_per_gpu * self.ulysses_sequence_parallel_size\n                    micro_batches, _ = prepare_dynamic_batch(mini_batch, max_token_len=max_token_len)\n                else:\n                    self.gradient_accumulation = (\n                        self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu\n                    )\n                    micro_batches = mini_batch.split(self.config.ppo_micro_batch_size_per_gpu)\n\n                self.critic_optimizer.zero_grad()\n\n                for micro_batch in micro_batches:\n                    micro_batch = micro_batch.to(get_device_id())\n                    micro_batch_metrics = {}\n                    model_inputs = {**micro_batch.batch, **micro_batch.non_tensor_batch}\n                    response_mask = model_inputs[\"response_mask\"]\n                    values = model_inputs[\"values\"]\n                    returns = model_inputs[\"returns\"]\n\n                    vpreds = self._forward_micro_batch(model_inputs)\n                    vf_loss, vf_clipfrac = core_algos.compute_value_loss(\n                        vpreds=vpreds,\n                        values=values,\n                        returns=returns,\n                        response_mask=response_mask,\n                        cliprange_value=self.config.cliprange_value,\n                        loss_agg_mode=self.config.loss_agg_mode,\n                    )\n                    if self.config.use_dynamic_bsz:\n                        # relative to the dynamic bsz\n                        loss_scale_factor = response_mask.shape[0] / self.config.ppo_mini_batch_size\n                        loss = vf_loss * loss_scale_factor\n                    else:\n                        loss_scale_factor = 1 / self.gradient_accumulation\n                        loss = vf_loss * loss_scale_factor\n\n                    loss.backward()\n\n                    micro_batch_metrics.update(\n                        {\n                            \"critic/vf_clipfrac\": vf_clipfrac.detach().item(),\n                            \"critic/vpred_mean\": masked_mean(vpreds, response_mask).detach().item(),\n                        }\n                    )\n\n                    metrics[\"critic/vf_loss\"] += vf_loss.detach().item() * loss_scale_factor\n                    append_to_dict(metrics, micro_batch_metrics)\n\n                grad_norm = self._optimizer_step()\n                mini_batch_metrics = {\"critic/grad_norm\": grad_norm.detach().item()}\n                append_to_dict(metrics, mini_batch_metrics)\n        self.critic_optimizer.zero_grad()\n        return metrics\n"
  },
  {
    "path": "verl/workers/critic/megatron_critic.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nImplement a multiprocess PPOCritic\n\"\"\"\n\nimport itertools\nimport logging\nimport os\nfrom functools import partial\nfrom typing import Iterable\n\nimport torch\nimport torch.distributed\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.optimizer import DistributedOptimizer, OptimizerConfig\nfrom megatron.core.pipeline_parallel import get_forward_backward_func\nfrom omegaconf import OmegaConf\nfrom torch import nn\n\nfrom verl import DataProto\nfrom verl.trainer.ppo import core_algos\nfrom verl.utils.device import get_device_id, get_torch_device\nfrom verl.utils.megatron.pipeline_parallel import make_batch_generator\nfrom verl.utils.profiler import GPUMemoryLogger\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import get_reverse_idx, rearrange_micro_batches\nfrom verl.utils.torch_functional import broadcast_dict_tensor, masked_mean\nfrom verl.workers.critic import BasePPOCritic\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass MegatronPPOCritic(BasePPOCritic):\n    def __init__(\n        self,\n        config,\n        model_config,\n        hf_config,\n        tf_config,\n        critic_module: nn.ModuleList,\n        critic_optimizer: DistributedOptimizer,\n        critic_optimizer_config: OptimizerConfig,\n    ):\n        super().__init__(config=config)\n        self._validate_config(config)\n        self.model_config = model_config\n        self.hf_config = hf_config  # huggingface config\n        self.tf_config = tf_config  # mcore transformer config\n\n        self.critic_module = critic_module\n        self.critic_optimizer = critic_optimizer\n        self.critic_optimizer_config = critic_optimizer_config\n\n        # we create a separate nametuple for optimizer step so that global args won't affect it.\n        self.optimizer_step_args = OmegaConf.create(\n            {\n                \"skip_grad\": None,\n                \"overlap_dp_param_comm\": False,\n                \"overlap_dp_grad_comm\": False,\n                \"gradient_accumulation_steps\": 1,\n                \"sequence_parallel\": self.tf_config.sequence_parallel,\n                \"DDP_impl\": \"local\",\n                \"layernorm_allreduce_bucket_threshold\": 0,\n                \"reduce_grads_use_alltoall\": False,\n            }\n        )\n\n    def _validate_config(self, config) -> None:\n        \"\"\"Validate config options not implemented for Megatron backend\"\"\"\n        assert config.get(\"ulysses_sequence_parallel_size\", 1) == 1\n        if config.shuffle:\n            assert config.data_loader_seed is not None, \"If shuffle dataloader, seed must be manually set\"\n        self.config = config\n\n    @GPUMemoryLogger(\"megatron critic\", logger=logger)\n    def compute_values(self, data: DataProto) -> DataProto:\n        prev_modes = [m.training for m in self.critic_module]\n        for module in self.critic_module:\n            module.eval()\n        responses = data.batch[\"responses\"]\n        attention_mask = data.batch[\"attention_mask\"]\n        use_dynamic_bsz = data.meta_info.get(\"use_dynamic_bsz\", False)\n        micro_batch_size = data.meta_info.get(\"micro_batch_size\", None)\n        max_token_len = data.meta_info.get(\"max_token_len\", None)\n        assert micro_batch_size is not None, \"micro batch size is needed for forward compute\"\n        if use_dynamic_bsz:\n            assert max_token_len is not None, \"max_token_len must be set when use_dynamic_bsz is True\"\n            max_token_len = max_token_len * self.config.megatron.context_parallel_size\n        response_length = responses.size(1)\n        with torch.no_grad():\n            output = self.forward_backward_batch(\n                data=data,\n                forward_only=True,\n                use_dynamic_bsz=use_dynamic_bsz,\n                micro_batch_size=micro_batch_size,\n                max_token_len=max_token_len,\n                mini_batch_size=None,\n            )\n            if mpu.is_pipeline_last_stage(ignore_virtual=True):\n                # only on last rank. It should be on every tp rank\n                values = [o[\"vpreds\"] for o in output[\"output\"]]  # (bs, seq_size, vocal_size)\n                values = torch.cat(values, dim=0).to(torch.float32)\n                if use_dynamic_bsz:\n                    indices = output[\"indices\"]\n                    indices = list(itertools.chain.from_iterable(indices))\n                    assert len(indices) == values.size(0), f\"{len(indices)} vs. {values.size()}\"\n                    revert_indices = torch.tensor(get_reverse_idx(indices), dtype=torch.long)\n                    values = values[revert_indices]\n            else:\n                values = torch.empty_like(attention_mask, dtype=torch.float32)\n\n            # each tp ranks should contain the same value\n            values = values[\n                :, -response_length - 1 : -1\n            ]  # Values are predicted at the ends of prefixes, e.g., the last prompt token\n            response_mask = attention_mask[:, -response_length:]\n            values = values * response_mask  # Only action tokens have values\n            values = values.contiguous()\n\n            # sync among pp ranks\n            values = values.to(get_device_id())\n            torch.distributed.broadcast(\n                tensor=values,\n                src=mpu.get_pipeline_model_parallel_last_rank(),\n                group=mpu.get_pipeline_model_parallel_group(),\n            )\n            values = values.to(\"cpu\")\n\n        # add empty cache after each compute\n        get_torch_device().empty_cache()\n\n        for module, mode in zip(self.critic_module, prev_modes, strict=False):\n            module.train(mode)\n        return values\n\n    def make_minibatch_iterator(self, data: DataProto) -> Iterable[DataProto]:\n        select_keys = [\"input_ids\", \"responses\", \"attention_mask\", \"position_ids\", \"values\", \"returns\"]\n        data = data.select(batch_keys=select_keys)\n        return data.make_iterator(\n            mini_batch_size=self.config.ppo_mini_batch_size,\n            epochs=self.config.ppo_epochs,\n            seed=self.config.data_loader_seed,\n            dataloader_kwargs={\"shuffle\": self.config.shuffle},\n        )\n\n    def forward_backward_batch(\n        self,\n        data: DataProto,\n        forward_only=False,\n        use_dynamic_bsz=False,\n        micro_batch_size=None,\n        max_token_len=None,\n        mini_batch_size=None,\n    ):\n        # broadcast from last pp rank to all other pp ranks\n        data.to(get_device_id())\n        mini_batch = data\n        mini_batch.batch = mini_batch.batch.contiguous()\n        broadcast_dict_tensor(\n            mini_batch.batch,\n            src=mpu.get_pipeline_model_parallel_last_rank(),\n            group=mpu.get_pipeline_model_parallel_group(),\n        )\n        mini_batch.to(\"cpu\")\n        # split into micro-batches\n        mini_batch.batch[\"attention_mask\"] = mini_batch.batch[\"attention_mask\"].to(bool)\n\n        indices = None\n        if use_dynamic_bsz:\n            assert max_token_len is not None, \"max_token_len must be set when use_dynamic_bsz is True\"\n            dp_group = mpu.get_data_parallel_group()\n            vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()\n            if vpp_size is not None and vpp_size > 1:\n                microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage\n                micro_batches, indices = rearrange_micro_batches(\n                    batch=mini_batch.batch,\n                    num_batches_divided_by=microbatch_group_size_per_vp_stage,\n                    max_token_len=max_token_len,\n                    dp_group=dp_group,\n                )\n                assert len(micro_batches) % self.tf_config.microbatch_group_size_per_vp_stage == 0, (\n                    f\"micro_batches {micro_batches} must be divisible by microbatch_group_size_per_vp_stage \"\n                    f\"{microbatch_group_size_per_vp_stage} for megatron backend\"\n                )\n            else:\n                micro_batches, indices = rearrange_micro_batches(\n                    batch=mini_batch.batch, max_token_len=max_token_len, dp_group=dp_group\n                )\n            total_seqlen = max_token_len\n        else:\n            assert micro_batch_size is not None, (\n                \"micro_batch_size is needed to be passed in when not using dynamic batch size\"\n            )\n            micro_batches = mini_batch.batch.split(micro_batch_size)\n            seq_len = micro_batches[0][\"input_ids\"].shape[1]\n            total_seqlen = micro_batch_size * seq_len\n        n_micro_batch = len(micro_batches)\n\n        forward_backward_func = get_forward_backward_func()\n\n        def loss_func(output, data, meta_info):\n            nonlocal use_dynamic_bsz\n\n            if forward_only:\n                return torch.tensor(1.0, device=output.device), {\"vpreds\": output}\n\n            responses = data[\"responses\"]\n            attention_mask = data[\"attention_mask\"]\n            values = data[\"values\"]\n            returns = data[\"returns\"]\n            response_length = responses.size(1)\n\n            response_mask = attention_mask[:, -response_length:]\n\n            cliprange_value = self.config.cliprange_value\n\n            vpreds = output  # (bs, sequence_length)\n            vpreds = vpreds[:, -response_length - 1 : -1]\n\n            vf_loss, vf_clipfrac = core_algos.compute_value_loss(\n                vpreds=vpreds,\n                values=values,\n                returns=returns,\n                response_mask=response_mask,\n                cliprange_value=cliprange_value,\n                loss_agg_mode=self.config.loss_agg_mode,\n            )\n\n            stats = {\n                \"critic/vf_loss\": vf_loss.detach().item(),\n                \"critic/vf_clipfrac\": vf_clipfrac.detach().item(),\n                \"critic/vpred_mean\": masked_mean(vpreds, response_mask).detach().item(),\n            }\n\n            return vf_loss, stats\n\n        def forward_step(batch_iter, model):\n            batch = next(batch_iter)\n            batch = batch.to(get_device_id())\n            batch = batch.contiguous()\n\n            input_ids = batch[\"input_ids\"]\n            attention_mask = batch[\"attention_mask\"]\n            position_ids = batch[\"position_ids\"]\n            from verl.models.mcore import get_mcore_forward_fn\n\n            forward_fn = get_mcore_forward_fn(self.hf_config)\n\n            output = forward_fn(\n                model,\n                input_ids,\n                attention_mask,\n                position_ids,\n                {},  # multi_modal_inputs\n                value_model=True,\n            )\n\n            return output, partial(loss_func, data=batch, meta_info={})\n\n        # batch should be a list of batches inside micro-batches\n        batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.critic_module))\n\n        # TODO: we may use the new schedule instead\n        # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size)\n        if mpu.get_pipeline_model_parallel_world_size() > 1:\n            losses_reduced = forward_backward_func(\n                forward_step_func=forward_step,\n                data_iterator=batch_generator,\n                model=self.critic_module,\n                num_microbatches=n_micro_batch,\n                seq_length=total_seqlen,  # no use when input_shapes was set\n                micro_batch_size=1,  # no use when input_shapes was set\n                forward_only=forward_only,\n            )\n        else:\n            losses_reduced = forward_backward_func(\n                forward_step_func=forward_step,\n                data_iterator=batch_generator,\n                model=self.critic_module,\n                num_microbatches=n_micro_batch,\n                seq_length=total_seqlen,  # in use for pp = 1\n                micro_batch_size=1,  # in use for pp = 1\n                forward_only=forward_only,\n            )\n        # loss_reduces contains the stats returned from loss_func\n        losses_reduced = {\"output\": losses_reduced}\n        if use_dynamic_bsz:\n            losses_reduced[\"indices\"] = indices\n        return losses_reduced\n\n    @GPUMemoryLogger(\"megatron critic\", logger=logger)\n    def update_critic(self, dataloader: Iterable[DataProto]):\n        metrics = {}\n\n        for data in dataloader:\n            self.critic_optimizer.zero_grad()\n            # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm\n            for chunk in self.critic_module:\n                chunk.zero_grad_buffer()\n\n            micro_batch_size = self.config.ppo_micro_batch_size_per_gpu\n            max_token_len = None\n            if self.config.use_dynamic_bsz:\n                max_token_len = self.config.ppo_max_token_len_per_gpu * self.config.megatron.context_parallel_size\n            metric_micro_batch = self.forward_backward_batch(\n                data,\n                forward_only=False,\n                use_dynamic_bsz=self.config.use_dynamic_bsz,\n                micro_batch_size=micro_batch_size,\n                max_token_len=max_token_len,\n                mini_batch_size=self.config.ppo_mini_batch_size,\n            )\n            metric_micro_batch = metric_micro_batch[\"output\"]\n            update_successful, grad_norm, num_zeros_in_grad = self.critic_optimizer.step()\n            learning_rate = self.critic_optimizer.param_groups[-1][\"lr\"]\n            data = {\"critic/grad_norm\": grad_norm, \"critic/lr\": learning_rate}\n            append_to_dict(metrics, data)\n\n            if update_successful:\n                # allgather already execute in optimizer.step in new megatron\n                pass\n            else:\n                raise NotImplementedError\n\n            for metric in metric_micro_batch:\n                append_to_dict(metrics, metric)  # append the metric from this micro-batch to global metrics.\n\n        # add empty cache after each compute\n        get_torch_device().empty_cache()\n        return metrics\n"
  },
  {
    "path": "verl/workers/engine/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom .base import BaseEngine, EngineRegistry\nfrom .fsdp import FSDPEngine, FSDPEngineWithLMHead\n\n__all__ = [\n    \"BaseEngine\",\n    \"EngineRegistry\",\n    \"FSDPEngine\",\n    \"FSDPEngineWithLMHead\",\n]\n\ntry:\n    from .torchtitan import TorchTitanEngine, TorchTitanEngineWithLMHead\n\n    __all__ += [\"TorchTitanEngine\", \"TorchTitanEngineWithLMHead\"]\nexcept ImportError:\n    TorchTitanEngine = None\n    TorchTitanEngineWithLMHead = None\n\ntry:\n    from .veomni import VeOmniEngine, VeOmniEngineWithLMHead\n\n    __all__ += [\"VeOmniEngine\", \"VeOmniEngineWithLMHead\"]\nexcept ImportError:\n    VeOmniEngine = None\n    VeOmniEngineWithLMHead = None\n\ntry:\n    from .automodel import AutomodelEngine, AutomodelEngineWithLMHead\n\n    __all__ += [\"AutomodelEngine\", \"AutomodelEngineWithLMHead\"]\nexcept ImportError:\n    AutomodelEngine = None\n    AutomodelEngineWithLMHead = None\n\n# Mindspeed must be imported before Megatron to ensure the related monkey patches take effect as expected\ntry:\n    from .mindspeed import MindspeedEngineWithLMHead\n\n    __all__ += [\"MindspeedEngineWithLMHead\"]\nexcept ImportError:\n    MindspeedEngineWithLMHead = None\n\ntry:\n    from .megatron import MegatronEngine, MegatronEngineWithLMHead\n\n    __all__ += [\"MegatronEngine\", \"MegatronEngineWithLMHead\"]\nexcept ImportError:\n    MegatronEngine = None\n    MegatronEngineWithLMHead = None\n"
  },
  {
    "path": "verl/workers/engine/automodel/__init__.py",
    "content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .transformer_impl import AutomodelEngine, AutomodelEngineWithLMHead\n\n__all__ = [\n    \"AutomodelEngine\",\n    \"AutomodelEngineWithLMHead\",\n]\n"
  },
  {
    "path": "verl/workers/engine/automodel/transformer_impl.py",
    "content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Automodel (nemo_automodel) engine for verl SFT training.\n\nThis engine delegates model building, parallelization, optimizer sharding,\nLR scheduling, gradient clipping, and checkpointing to Automodel's\ninfrastructure while using verl's training loop, data pipeline, and loss function.\n\"\"\"\n\nimport gc\nimport logging\nimport os\nfrom contextlib import nullcontext\nfrom typing import Any, Callable, Optional\n\nimport torch\nimport torch.distributed\nfrom huggingface_hub.constants import HF_HUB_CACHE\nfrom nemo_automodel.components.checkpoint.checkpointing import Checkpointer, CheckpointingConfig\nfrom nemo_automodel.components.optim.scheduler import OptimizerParamScheduler\nfrom nemo_automodel.components.training.utils import (\n    prepare_for_final_backward,\n    prepare_for_grad_accumulation,\n    scale_grads_and_clip_grad_norm,\n)\nfrom tensordict import TensorDict\nfrom torch.distributed.tensor import DTensor\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.debug import log_gpu_memory_usage\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.model import convert_weight_keys, extract_multi_modal_inputs\nfrom verl.utils.torch_functional import logprobs_from_logits\nfrom verl.workers.config import AutomodelEngineConfig, AutomodelOptimizerConfig, HFModelConfig\n\nfrom ..base import BaseEngine, BaseEngineCtx, EngineRegistry\nfrom ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches\nfrom .utils import (\n    build_automodel_model,\n    build_distributed_config_from_engine_config,\n    get_dp_group_size,\n    get_dp_rank,\n    get_pp_rank,\n    get_tp_rank,\n    load_automodel_model_to_gpu,\n    load_automodel_optimizer,\n    maybe_fully_shard_optimizer,\n    offload_automodel_model_to_cpu,\n    offload_automodel_optimizer,\n)\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass AutomodelEngine(BaseEngine):\n    \"\"\"Engine implementation using Automodel for distributed training.\"\"\"\n\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: AutomodelEngineConfig,\n        optimizer_config: AutomodelOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n        **kwargs,\n    ):\n        super().__init__()\n\n        self.model_config = model_config\n        self.engine_config = engine_config\n        self.optimizer_config = optimizer_config\n        self.checkpoint_config = checkpoint_config\n\n        self.mode = None\n        self.rank = torch.distributed.get_rank()\n\n        # Apply compatibility patches early in the process\n        from nemo_automodel._transformers.utils import apply_cache_compatibility_patches\n        from nemo_automodel.shared.te_patches import apply_te_patches\n\n        apply_cache_compatibility_patches()\n        apply_te_patches()\n\n        world_size = torch.distributed.get_world_size()\n        self.distributed_config, self.device_mesh, self.moe_mesh = build_distributed_config_from_engine_config(\n            self.engine_config, world_size\n        )\n\n        if self.engine_config.full_determinism:\n            enable_full_determinism(seed=self.engine_config.seed)\n\n        self._is_offload_param = self.engine_config.param_offload\n        self._is_offload_optimizer = self.engine_config.optimizer_offload\n\n        if self.engine_config.entropy_from_logits_with_chunking:\n            entropy_from_logits = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_from_logits = verl_F.entropy_from_logits\n\n        self.compute_entropy_from_logits = (\n            torch.compile(entropy_from_logits, dynamic=True)\n            if self.engine_config.use_torch_compile\n            else entropy_from_logits\n        )\n\n    @property\n    def is_param_offload_enabled(self) -> bool:\n        return self._is_offload_param\n\n    @property\n    def is_optimizer_offload_enabled(self) -> bool:\n        return self._is_offload_optimizer\n\n    def initialize(self):\n        \"\"\"Build the model, optimizer, LR scheduler, and checkpointer using Automodel infrastructure.\"\"\"\n        self.module = build_automodel_model(\n            self.model_config, self.engine_config, self.distributed_config, self.device_mesh, self.moe_mesh\n        )\n        log_gpu_memory_usage(\"After Automodel model build\", logger=logger)\n\n        if not self.engine_config.forward_only:\n            self.optimizer = self._build_optimizer(self.module)\n            # maybe shard optimizer for MegatronFSDP\n            maybe_fully_shard_optimizer(self.module, self.optimizer, self.distributed_config)\n            self.lr_scheduler = self._build_lr_scheduler(self.optimizer)\n        else:\n            self.optimizer = None\n            self.lr_scheduler = None\n        self._build_checkpointer()\n\n        self.to(\n            device=\"cpu\",\n            model=self._is_offload_param,\n            optimizer=self._is_offload_optimizer,\n            grad=self._is_offload_param,\n        )\n\n        log_gpu_memory_usage(\"After offload model/optimizer/grad during init\", logger=logger)\n        torch.cuda.empty_cache()\n\n    def _build_optimizer(self, module):\n        \"\"\"Build optimizer via Automodel's build_optimizer.\"\"\"\n        from nemo_automodel.components.config.loader import ConfigNode\n        from nemo_automodel.recipes.llm.train_ft import build_optimizer as automodel_build_optimizer\n\n        config = self.optimizer_config\n\n        opt_dict = {\n            \"_target_\": f\"{config.optimizer_impl}.{config.optimizer}\",\n            \"lr\": config.lr,\n            \"weight_decay\": config.weight_decay,\n            \"eps\": config.eps,\n            \"betas\": list(config.betas),\n        }\n\n        if config.master_weights:\n            opt_dict[\"master_weights\"] = config.master_weights\n        if config.store_param_remainders:\n            opt_dict[\"store_param_remainders\"] = config.store_param_remainders\n\n        _short_to_torch = {\"bf16\": \"torch.bfloat16\", \"fp32\": \"torch.float32\", \"fp16\": \"torch.float16\"}\n        for attr in (\"exp_avg_dtype\", \"exp_avg_sq_dtype\", \"master_weight_dtype\"):\n            val = getattr(config, attr, None)\n            if val is not None:\n                opt_dict[attr] = _short_to_torch.get(val, val)\n\n        if config.override_optimizer_config:\n            opt_dict.update(config.override_optimizer_config)\n\n        cfg_opt = ConfigNode(opt_dict)\n        optimizers = automodel_build_optimizer(module, cfg_opt, self.distributed_config, self.device_mesh)\n        assert len(optimizers) == 1, f\"Expected 1 optimizer, got {len(optimizers)}\"\n        return optimizers[0]\n\n    def _build_lr_scheduler(self, optimizer):\n        cfg = self.optimizer_config\n        total_steps = cfg.total_training_steps\n        num_warmup_steps = cfg.lr_warmup_steps\n\n        if num_warmup_steps <= 0:\n            num_warmup_steps = int(cfg.lr_warmup_steps_ratio * total_steps)\n\n        base_lr = cfg.lr\n        init_lr_ratio = cfg.init_lr_ratio if cfg.init_lr_ratio is not None else 0.1\n        min_lr_ratio = cfg.min_lr_ratio if cfg.min_lr_ratio is not None else 0.01\n\n        if self.rank == 0:\n            print(\n                f\"Automodel LR Scheduler: total_steps={total_steps}, warmup={num_warmup_steps}, \"\n                f\"decay_style={cfg.lr_scheduler_type}, init_lr={base_lr * init_lr_ratio:.2e}, \"\n                f\"max_lr={base_lr:.2e}, min_lr={base_lr * min_lr_ratio:.2e}\"\n            )\n\n        scheduler = OptimizerParamScheduler(\n            optimizer=optimizer,\n            init_lr=base_lr * init_lr_ratio,\n            max_lr=base_lr,\n            min_lr=base_lr * min_lr_ratio,\n            lr_warmup_steps=num_warmup_steps,\n            lr_decay_steps=total_steps,\n            lr_decay_style=cfg.lr_scheduler_type,\n            start_wd=cfg.weight_decay,\n            end_wd=cfg.weight_decay,\n            wd_incr_steps=total_steps,\n            wd_incr_style=getattr(cfg, \"wd_incr_style\", \"constant\"),\n        )\n        return scheduler\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:\n        batch_num_tokens = data[\"loss_mask\"].sum().to(get_device_id())\n        torch.distributed.all_reduce(\n            batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group()\n        )\n        tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())\n        tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())\n\n        micro_batches, indices = prepare_micro_batches(\n            data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True\n        )\n\n        output_lst = []\n        ctx = torch.no_grad() if forward_only else nullcontext()\n\n        if not forward_only:\n            prepare_for_grad_accumulation([self.module])\n\n            # Set MoE aux loss backward scale to counteract FSDP's gradient allreduce.\n            if self.engine_config.ep_size > 1:\n                from nemo_automodel.components.moe.megatron.moe_utils import MoEAuxLossAutoScaler\n\n                MoEAuxLossAutoScaler.main_loss_backward_scale = torch.tensor(\n                    float(get_dp_group_size(self.device_mesh, include_cp=True))\n                )\n\n        num_micro_batches = len(micro_batches)\n        for i, micro_batch in enumerate(micro_batches):\n            # Signal final backward for MoE\n            if not forward_only and i == num_micro_batches - 1:\n                prepare_for_final_backward([self.module])\n\n            with ctx:\n                loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only)\n                if not forward_only:\n                    loss.backward()\n            output_lst.append(meta_info)\n\n        return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data)\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        raise NotImplementedError(\"forward_step must be implemented in subclass\")\n\n    def optimizer_zero_grad(self):\n        self.optimizer.zero_grad()\n\n    def optimizer_step(self):\n        grad_norm = scale_grads_and_clip_grad_norm(\n            max_grad_norm=self.optimizer_config.clip_grad,\n            model_parts=[self.module],\n            norm_type=2.0,\n            pp_enabled=False,\n            device_mesh=self.device_mesh,\n            moe_mesh=self.moe_mesh,\n            ep_axis_name=\"ep\" if self.moe_mesh is not None and \"ep\" in self.moe_mesh.mesh_dim_names else None,\n            pp_axis_name=None,\n            foreach=True,\n            num_label_tokens=None,\n            dp_group_size=get_dp_group_size(self.device_mesh, include_cp=True),\n        )\n\n        if isinstance(grad_norm, torch.Tensor):\n            grad_norm_val = grad_norm.item()\n        else:\n            grad_norm_val = float(grad_norm)\n\n        # If grad_norm is not finite, skip the update\n        if not torch.isfinite(torch.tensor(grad_norm_val)):\n            print(f\"WARN: grad_norm is not finite: {grad_norm_val}\")\n            self.optimizer.zero_grad()\n        else:\n            self.optimizer.step()\n            if hasattr(self.module, \"update_moe_gate_bias\"):\n                self.module.update_moe_gate_bias()\n\n        return grad_norm_val\n\n    def lr_scheduler_step(self):\n        \"\"\"Step Automodel's OptimizerParamScheduler and return current LR.\"\"\"\n        self.lr_scheduler.step(increment=1)\n        lr = self.optimizer.param_groups[0][\"lr\"]\n        return lr\n\n    def get_data_parallel_rank(self):\n        if self.device_mesh is not None:\n            return self.device_mesh.get_local_rank(\"dp\")\n        return torch.distributed.get_rank()\n\n    def get_data_parallel_size(self):\n        if self.device_mesh is not None:\n            return self.device_mesh[\"dp\"].size()\n        return torch.distributed.get_world_size()\n\n    def get_data_parallel_group(self):\n        if self.device_mesh is not None:\n            return self.device_mesh.get_group(mesh_dim=\"dp\")\n        return torch.distributed.group.WORLD\n\n    def is_mp_src_rank_with_outputs(self):\n        if self.device_mesh is not None and \"tp\" in self.device_mesh.mesh_dim_names:\n            if self.device_mesh[\"tp\"].size() > 1:\n                return self.device_mesh.get_local_rank(\"tp\") == 0\n        return True\n\n    def train_mode(self, **kwargs):\n        return AutomodelTrainModeCtx(self, **kwargs)\n\n    def eval_mode(self, **kwargs):\n        return AutomodelEvalModeCtx(self, **kwargs)\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        super().to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n        if self.engine_config.forward_only:\n            return\n\n        device_name = get_device_name()\n        assert device in (device_name, \"cpu\")\n\n        if device == device_name:\n            if model:\n                load_automodel_model_to_gpu(self.module)\n            if optimizer and self.optimizer is not None:\n                load_automodel_optimizer(self.optimizer, get_device_id())\n            gc.collect()\n        elif device == \"cpu\":\n            if model:\n                offload_automodel_model_to_cpu(self.module)\n            if optimizer and self.optimizer is not None:\n                offload_automodel_optimizer(self.optimizer)\n        else:\n            raise ValueError(f\"Invalid device type: {device}\")\n\n    def _build_checkpointer(self):\n        ckpt_config = CheckpointingConfig(\n            enabled=True,\n            checkpoint_dir=\"checkpoints/\",\n            model_save_format=\"safetensors\",\n            model_cache_dir=HF_HUB_CACHE,\n            model_repo_id=self.model_config.path,\n            save_consolidated=True,\n            is_peft=False,\n        )\n        self.checkpointer = Checkpointer(\n            config=ckpt_config,\n            dp_rank=get_dp_rank(self.device_mesh, include_cp=True),\n            tp_rank=get_tp_rank(self.device_mesh),\n            pp_rank=get_pp_rank(self.device_mesh),\n            moe_mesh=self.moe_mesh,\n        )\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"Save model, optimizer, and LR scheduler using Automodel's Checkpointer.\"\"\"\n        origin_module_device = next(self.module.parameters()).device.type\n        if self._is_offload_param or origin_module_device == \"cpu\":\n            load_automodel_model_to_gpu(self.module)\n\n        # Save model weights\n        self.checkpointer.save_model(self.module, local_path)\n\n        # Save optimizer and LR scheduler state\n        if self.optimizer is not None:\n            scheduler_list = [self.lr_scheduler] if self.lr_scheduler is not None else None\n            self.checkpointer.save_optimizer(self.optimizer, self.module, local_path, scheduler=scheduler_list)\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_automodel_model_to_cpu(self.module)\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs\n    ) -> None:\n        \"\"\"Load model, optimizer, and LR scheduler using Automodel's Checkpointer.\"\"\"\n        if self._is_offload_param:\n            load_automodel_model_to_gpu(self.module)\n\n        model_path = os.path.join(local_path, \"model\")\n        if not os.path.isdir(model_path):\n            model_path = local_path\n        self.checkpointer.load_model(self.module, model_path)\n\n        if self.optimizer is not None:\n            scheduler_list = [self.lr_scheduler] if self.lr_scheduler is not None else None\n            self.checkpointer.load_optimizer(self.optimizer, self.module, local_path, scheduler=scheduler_list)\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_automodel_model_to_cpu(self.module)\n\n        if self._is_offload_optimizer and self.optimizer is not None:\n            offload_automodel_optimizer(self.optimizer)\n\n    def get_per_tensor_param(self, **kwargs):\n        load_automodel_model_to_gpu(self.module)\n\n        params = self.module.state_dict()\n        params = convert_weight_keys(params, getattr(self.module, \"_fsdp_wrapped_module\", self.module))\n\n        if self._is_offload_param:\n            offload_automodel_model_to_cpu(self.module)\n\n        def param_generator():\n            for name, param in params.items():\n                unsharded_tensor = param.full_tensor() if isinstance(param, DTensor) else param\n                yield name, unsharded_tensor\n\n        return param_generator(), None\n\n\nclass AutomodelEvalModeCtx(BaseEngineCtx):\n    def __init__(self, engine: AutomodelEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"eval\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, AutomodelEngine)\n        super().__enter__()\n        self.engine.module.eval()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, AutomodelEngine)\n        # Reshard the root FSDP module\n        if hasattr(self.engine.module, \"reshard\"):\n            self.engine.module.reshard()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\nclass AutomodelTrainModeCtx(BaseEngineCtx):\n    def __init__(self, engine: AutomodelEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"train\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, AutomodelEngine)\n        super().__enter__()\n        self.engine.module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, AutomodelEngine)\n        self.engine.optimizer_zero_grad()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=[\"automodel\"], device=[\"cuda\"])\nclass AutomodelEngineWithLMHead(AutomodelEngine):\n    \"\"\"Automodel engine for language model with LM head training.\"\"\"\n\n    def prepare_model_inputs(self, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key=\"use_fused_kernels\", default=False)\n        temperature = micro_batch[\"temperature\"]\n        temperature_item = temperature\n        if use_fused_kernels:\n            assert not isinstance(temperature, torch.Tensor), (\n                \"use_fused_kernels does not support per sample temperature yet\"\n            )\n        assert pad_mode == DatasetPadMode.NO_PADDING, f\"pad_mode {pad_mode} not supported\"\n\n        multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get(\"multi_modal_inputs\", []))\n        input_ids = micro_batch[\"input_ids\"]\n        position_ids = micro_batch[\"position_ids\"]\n\n        if not isinstance(temperature, torch.Tensor):\n            temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device)\n\n        temperature = temperature.to(torch.float32)\n        assert temperature.shape[0] == input_ids.shape[0]\n\n        output_args = {}\n\n        if use_remove_padding:\n            temperature_rmpad = verl_F.expand_as_nested(temperature, input_ids).values()\n            temperature_rmpad = temperature_rmpad.unsqueeze(0)\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                input_ids_rmpad = input_ids.values().unsqueeze(0)\n                if position_ids.dim() == 3:\n                    position_ids_rmpad = position_ids.values().unsqueeze(1)\n                else:\n                    position_ids_rmpad = position_ids.values().unsqueeze(0)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n            input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1)\n\n            input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0)\n            temperature_rmpad = temperature_rmpad.squeeze(0)\n            output_args[\"input_ids_rmpad_rolled\"] = input_ids_rmpad_rolled\n            output_args[\"temperature_rmpad\"] = temperature_rmpad\n\n            model_inputs = {\n                \"input_ids\": input_ids_rmpad,\n                \"attention_mask\": None,\n                \"position_ids\": position_ids_rmpad,\n            }\n\n            # For TE attention backend, pass cu_seqlens\n            if self.engine_config.attn_implementation == \"te\":\n                cu_seqlens = input_ids.offsets().to(torch.int32)\n                max_seqlen = cu_seqlens.diff().max().item()\n                model_inputs[\"qkv_format\"] = \"thd\"\n                model_inputs[\"cu_seqlens\"] = cu_seqlens.unsqueeze(0)\n                model_inputs[\"max_seqlen\"] = max_seqlen\n\n        else:\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                input_ids = micro_batch[\"input_ids\"]\n                position_ids = micro_batch[\"position_ids\"]\n                loss_mask = micro_batch[\"loss_mask\"]\n\n                pad_token_id = tu.get_non_tensor_data(data=micro_batch, key=\"pad_token_id\", default=0)\n                batch_size = micro_batch.batch_size[0]\n                seq_len_effective = input_ids.offsets().diff()\n                max_seq_len = max(seq_len_effective)\n\n                input_ids_rmpad_rolled = torch.roll(input_ids.values(), shifts=-1, dims=0)\n                output_args[\"input_ids_rmpad_rolled\"] = input_ids_rmpad_rolled\n                output_args[\"temperature\"] = temperature\n\n                input_ids = torch.nested.to_padded_tensor(\n                    input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len)\n                )\n\n                if position_ids.dim() == 3:\n                    position_ids = torch.nested.to_padded_tensor(\n                        position_ids, padding=0, output_size=(batch_size, 4, max_seq_len)\n                    ).transpose(0, 1)\n                else:\n                    position_ids = torch.nested.to_padded_tensor(\n                        position_ids, padding=0, output_size=(batch_size, max_seq_len)\n                    )\n\n                attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask]\n                attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged)\n                attention_mask = torch.nested.to_padded_tensor(\n                    attention_mask, padding=0, output_size=(batch_size, max_seq_len)\n                )\n\n                model_inputs = {\n                    \"input_ids\": input_ids,\n                    \"attention_mask\": attention_mask,\n                    \"position_ids\": position_ids,\n                }\n\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        extra_args = {}\n        if use_fused_kernels:\n            extra_args[\"temperature\"] = temperature_item\n            extra_args[\"return_dict\"] = True\n\n        model_inputs.update(multi_modal_inputs)\n        model_inputs.update(extra_args)\n\n        return model_inputs, output_args\n\n    def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key=\"use_fused_kernels\", default=False)\n        calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key=\"calculate_entropy\", default=False)\n\n        if isinstance(output, torch.Tensor):\n            from types import SimpleNamespace\n\n            output = SimpleNamespace(logits=output)\n\n        model_output = {}\n        input_ids = micro_batch[\"input_ids\"]\n\n        if use_remove_padding:\n            input_ids_rmpad_rolled = output_args[\"input_ids_rmpad_rolled\"]\n            temperature_rmpad = output_args[\"temperature_rmpad\"]\n\n            if use_fused_kernels:\n                log_probs = output.log_probs.squeeze(0)\n                entropy_rmpad = output.entropy.squeeze(0)\n            else:\n                logits_rmpad = output.logits.squeeze(0)\n                # With TP, logits are DTensors sharded on vocab dim; gather for log_softmax.\n                if isinstance(logits_rmpad, DTensor):\n                    logits_rmpad = logits_rmpad.full_tensor()\n                logits_rmpad = logits_rmpad / temperature_rmpad.clamp(min=1e-8).unsqueeze(-1).to(logits_rmpad.dtype)\n\n                inplace_backward = True\n                if calculate_entropy:\n                    inplace_backward = False\n                log_probs = logprobs_from_logits(\n                    logits=logits_rmpad,\n                    labels=input_ids_rmpad_rolled,\n                    inplace_backward=inplace_backward,\n                )\n\n                if calculate_entropy:\n                    if not self.engine_config.entropy_checkpointing:\n                        entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad)\n                    else:\n                        entropy_rmpad = torch.utils.checkpoint.checkpoint(\n                            self.compute_entropy_from_logits, logits_rmpad\n                        )\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                cu_seqlens = input_ids.offsets()\n                log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens)\n                if calculate_entropy:\n                    entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        else:\n            response_length = tu.get_non_tensor_data(data=micro_batch, key=\"max_response_length\", default=1024)\n            if use_fused_kernels:\n                log_probs = output.log_probs[:, -response_length - 1 : -1]\n                entropy = output.entropy[:, -response_length - 1 : -1]\n            else:\n                logits = output.logits\n                # With TP, logits are DTensors sharded on vocab dim; gather for log_softmax.\n                if isinstance(logits, DTensor):\n                    logits = logits.full_tensor()\n                temperature = output_args[\"temperature\"]\n                temperature = temperature.unsqueeze(-1).unsqueeze(-1)\n                logits = logits / temperature.clamp(min=1e-8).to(logits.dtype)\n\n                if calculate_entropy:\n                    if not self.engine_config.entropy_checkpointing:\n                        entropy = verl_F.entropy_from_logits(logits)\n                    else:\n                        entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits)\n\n                if pad_mode == DatasetPadMode.NO_PADDING:\n                    cu_seqlens = input_ids.offsets()\n                    seq_lengths = cu_seqlens.diff()\n                    starts = torch.zeros_like(seq_lengths, dtype=torch.int64)\n                    logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged)\n                    logits_rmpad = torch.cat([t for t in logits.unbind()])\n                    input_ids_rmpad_rolled = output_args[\"input_ids_rmpad_rolled\"]\n                    log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled)\n                    log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens)\n                    if calculate_entropy:\n                        entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged)\n                        entropy_rmpad = torch.cat([t for t in entropy.unbind()])\n                        entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n                else:\n                    raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        model_output[\"log_probs\"] = log_probs\n        if calculate_entropy:\n            model_output[\"entropy\"] = entropy\n\n        return model_output\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        \"\"\"Run forward pass, compute loss, and return outputs.\"\"\"\n        device_name = get_device_name()\n        micro_batch = micro_batch.to(get_device_id())\n        model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch)\n\n        with torch.autocast(device_type=device_name, dtype=torch.bfloat16):\n            raw_output = self.module(\n                **model_inputs,\n                use_cache=False,\n            )\n\n            model_output = self.prepare_model_outputs(\n                output=raw_output, output_args=output_args, micro_batch=micro_batch\n            )\n\n            if loss_function is not None:\n                loss, metrics = loss_function(\n                    model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group()\n                )\n            else:\n                assert forward_only, \"forward_only must be True when loss_function is None\"\n                loss = torch.tensor(1.0, device=device_name)\n                metrics = {}\n\n            output = {\n                \"model_output\": model_output,\n                \"loss\": loss.detach().item(),\n                \"metrics\": metrics,\n            }\n\n            return loss, output\n"
  },
  {
    "path": "verl/workers/engine/automodel/utils.py",
    "content": "# Copyright (c) 2025, NVIDIA CORPORATION. All rights reserved.\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\"\"\"Utility functions for the Automodel engine integration.\"\"\"\n\nimport torch\nimport torch.distributed\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\n\ndef get_dp_rank(device_mesh, include_cp=False):\n    \"\"\"Get data-parallel rank from device mesh.\"\"\"\n    if device_mesh is None:\n        return 0\n    if include_cp and \"cp\" in device_mesh.mesh_dim_names and device_mesh[\"cp\"].size() > 1:\n        return device_mesh.get_local_rank(\"dp_cp\")\n    return device_mesh.get_local_rank(\"dp\")\n\n\ndef get_tp_rank(device_mesh):\n    \"\"\"Get tensor-parallel rank from device mesh.\"\"\"\n    if device_mesh is None or \"tp\" not in device_mesh.mesh_dim_names or device_mesh[\"tp\"].size() == 1:\n        return 0\n    return device_mesh.get_local_rank(\"tp\")\n\n\ndef get_pp_rank(device_mesh):\n    \"\"\"Get pipeline-parallel rank from device mesh.\"\"\"\n    if device_mesh is None or \"pp\" not in device_mesh.mesh_dim_names or device_mesh[\"pp\"].size() == 1:\n        return 0\n    return device_mesh.get_local_rank(\"pp\")\n\n\ndef get_dp_group_size(device_mesh, include_cp=False):\n    \"\"\"Get data-parallel group size from device mesh.\"\"\"\n    if device_mesh is None:\n        return torch.distributed.get_world_size()\n    if include_cp and \"cp\" in device_mesh.mesh_dim_names and device_mesh[\"cp\"].size() > 1:\n        return device_mesh[\"dp_cp\"].size()\n    if \"dp\" in device_mesh.mesh_dim_names:\n        return device_mesh[\"dp\"].size()\n    return torch.distributed.get_world_size()\n\n\ndef maybe_fully_shard_optimizer(model, optimizer, distributed_config):\n    \"\"\"Call fully_shard_optimizer for MegatronFSDP strategy.\"\"\"\n    from nemo_automodel.components.distributed.config import MegatronFSDPConfig\n\n    if isinstance(distributed_config, MegatronFSDPConfig) and torch.distributed.get_world_size() > 1:\n        from megatron_fsdp.fully_shard import fully_shard_optimizer\n\n        fully_shard_optimizer(model, optimizer)\n\n\ndef build_distributed_config_from_engine_config(engine_config, world_size):\n    \"\"\"Build v5 distributed config, device_mesh, and moe_mesh from engine config.\n\n    Args:\n        engine_config: AutomodelEngineConfig instance.\n        world_size: Total number of processes in the job.\n\n    Returns:\n        Tuple of (distributed_config, device_mesh, moe_mesh).\n    \"\"\"\n    from nemo_automodel.components.distributed.config import DDPConfig, FSDP2Config, MegatronFSDPConfig\n    from nemo_automodel.components.distributed.mesh_utils import create_device_mesh\n\n    strategy = engine_config.distributed_strategy\n\n    if strategy == \"fsdp2\":\n        from torch.distributed.fsdp import MixedPrecisionPolicy\n\n        from verl.utils.torch_dtypes import PrecisionType\n\n        mp_policy = MixedPrecisionPolicy(\n            param_dtype=PrecisionType.to_dtype(engine_config.mp_param_dtype),\n            reduce_dtype=PrecisionType.to_dtype(engine_config.mp_reduce_dtype),\n            output_dtype=PrecisionType.to_dtype(engine_config.mp_output_dtype),\n            cast_forward_inputs=True,\n        )\n\n        distributed_config = FSDP2Config(\n            sequence_parallel=engine_config.sequence_parallel,\n            mp_policy=mp_policy,\n            activation_checkpointing=engine_config.activation_checkpointing,\n            defer_fsdp_grad_sync=engine_config.defer_fsdp_grad_sync,\n        )\n\n    elif strategy == \"megatron_fsdp\":\n        distributed_config = MegatronFSDPConfig(\n            activation_checkpointing=engine_config.activation_checkpointing,\n        )\n\n    elif strategy == \"ddp\":\n        distributed_config = DDPConfig(\n            activation_checkpointing=engine_config.activation_checkpointing,\n        )\n\n    else:\n        raise ValueError(f\"Unsupported distributed_strategy: {strategy}\")\n\n    device_mesh, moe_mesh = create_device_mesh(\n        distributed_config,\n        tp_size=engine_config.tp_size,\n        pp_size=engine_config.pp_size,\n        cp_size=engine_config.cp_size,\n        ep_size=engine_config.ep_size,\n        dp_replicate_size=engine_config.dp_replicate_size,\n        world_size=world_size,\n    )\n\n    return distributed_config, device_mesh, moe_mesh\n\n\ndef build_automodel_model(model_config, engine_config, distributed_config, device_mesh, moe_mesh):\n    \"\"\"Build a model using NeMoAutoModelForCausalLM.from_pretrained().\n\n    Args:\n        model_config: HFModelConfig with model path and settings.\n        engine_config: AutomodelEngineConfig with distributed settings.\n        distributed_config: FSDP2Config, MegatronFSDPConfig, or DDPConfig instance.\n        device_mesh: Pre-created device mesh (or None for DDP).\n        moe_mesh: Pre-created MoE mesh (or None).\n\n    Returns:\n        A HuggingFace model with Automodel's distributed infrastructure applied.\n    \"\"\"\n    from nemo_automodel._transformers.auto_model import NeMoAutoModelForCausalLM\n\n    kwargs = {}\n\n    if engine_config.enable_fp8:\n        from nemo_automodel.components.quantization.fp8 import FP8Config\n\n        kwargs[\"fp8_config\"] = FP8Config()\n\n    if engine_config.enable_compile:\n        from nemo_automodel.components.utils.compile_utils import CompileConfig\n\n        kwargs[\"compile_config\"] = CompileConfig()\n\n    # Qwen/Llama with ep_size<=1: use HF implementation.\n    from transformers import AutoConfig\n\n    _cfg = AutoConfig.from_pretrained(model_config.path, trust_remote_code=model_config.trust_remote_code)\n    _arch = (getattr(_cfg, \"architectures\", None) or [\"\"])[0].lower()\n    if engine_config.ep_size <= 1 and (\"qwen\" in _arch or \"llama\" in _arch):\n        kwargs[\"force_hf\"] = True\n\n    if engine_config.backend_config and not kwargs.get(\"force_hf\", False):\n        from nemo_automodel.components.models.common.utils import BackendConfig\n\n        backend_kwargs = dict(engine_config.backend_config)\n        kwargs[\"backend\"] = BackendConfig(**backend_kwargs)\n\n    # MoE config for MoEParallelizerConfig\n    if engine_config.ep_size > 1:\n        from nemo_automodel.components.moe.config import MoEParallelizerConfig\n\n        moe_kwargs = dict(engine_config.moe_config) if engine_config.moe_config else {}\n        if hasattr(distributed_config, \"mp_policy\"):\n            moe_kwargs.setdefault(\"mp_policy\", distributed_config.mp_policy)\n\n        kwargs[\"moe_config\"] = MoEParallelizerConfig(**moe_kwargs)\n\n    kwargs[\"attn_implementation\"] = engine_config.attn_implementation\n\n    from verl.utils.torch_dtypes import PrecisionType\n\n    kwargs[\"torch_dtype\"] = PrecisionType.to_dtype(engine_config.model_dtype)\n\n    model = NeMoAutoModelForCausalLM.from_pretrained(\n        pretrained_model_name_or_path=model_config.path,\n        device_mesh=device_mesh,\n        moe_mesh=moe_mesh,\n        distributed_config=distributed_config,\n        activation_checkpointing=engine_config.activation_checkpointing,\n        trust_remote_code=model_config.trust_remote_code,\n        **kwargs,\n    )\n\n    return model\n\n\n@torch.no_grad()\ndef offload_automodel_model_to_cpu(model, empty_cache=True):\n    \"\"\"Offload an FSDP2-wrapped model to CPU (reshard, move to CPU, optional cache clear).\"\"\"\n    from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState\n    from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state\n\n    for module in model.modules():\n        state = _get_module_fsdp_state(module)\n        if state is None:\n            continue\n        fsdp_param_group = state._fsdp_param_group\n\n        if fsdp_param_group is None:\n            continue\n\n        fsdp_param_group._training_state = TrainingState.IDLE\n\n    model.reshard()\n    model.cpu()\n    if empty_cache:\n        get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef load_automodel_model_to_gpu(model):\n    \"\"\"Load model back to GPU.\"\"\"\n    device = get_device_id()\n    model.to(device, non_blocking=True)\n\n\n@torch.no_grad()\ndef offload_automodel_optimizer(optimizer):\n    \"\"\"Offload optimizer state to CPU.\"\"\"\n    if not optimizer.state:\n        return\n    for param_group in optimizer.param_groups:\n        for param in param_group[\"params\"]:\n            state = optimizer.state[param]\n            for key, value in state.items():\n                if isinstance(value, torch.Tensor):\n                    state[key] = value.to(\"cpu\", non_blocking=True)\n\n\n@torch.no_grad()\ndef load_automodel_optimizer(optimizer, device_id):\n    \"\"\"Load optimizer state back to GPU.\"\"\"\n    if not optimizer.state:\n        return\n    for param_group in optimizer.param_groups:\n        for param in param_group[\"params\"]:\n            state = optimizer.state[param]\n            for key, value in state.items():\n                if isinstance(value, torch.Tensor):\n                    state[key] = value.to(device_id, non_blocking=True)\n"
  },
  {
    "path": "verl/workers/engine/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe abstract base class defining the interface for model training engines.\n\"\"\"\n\nfrom abc import abstractmethod\nfrom contextlib import nullcontext\nfrom typing import Any, Callable, ContextManager, Generator, Optional\n\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl.utils.device import get_device_name\nfrom verl.utils.tensordict_utils import maybe_fix_3d_position_ids\n\n\nclass BaseEngine:\n    \"\"\"\n    Abstract base class defining the interface for model training engines. Interface is subject to\n    change before release.\n\n    Engine implementations must subclass BaseEngine and provide concrete behavior for all methods.\n    \"\"\"\n\n    def initialize(self):\n        \"\"\"\n        Instantiate or load the model, optimizer, and learning rate scheduler.\n\n        Should prepare all components necessary for training or evaluation.\n        \"\"\"\n        raise NotImplementedError\n\n    @property\n    @abstractmethod\n    def is_param_offload_enabled(self) -> bool:\n        \"\"\"Whether parameter offloading is enabled.\"\"\"\n        raise NotImplementedError\n\n    @property\n    @abstractmethod\n    def is_optimizer_offload_enabled(self) -> bool:\n        \"\"\"Whether optimizer offloading is enabled.\"\"\"\n        raise NotImplementedError\n\n    def train_mode(self, **kwargs):\n        \"\"\"\n        Context manager entry for switching the engine and model into training mode.\n\n        Usage:\n            with engine.train_mode():\n                # runs in training mode\n        \"\"\"\n        raise NotImplementedError\n\n    def eval_mode(self, **kwargs):\n        \"\"\"\n        Context manager entry for switching the engine and model into evaluation mode.\n\n        Usage:\n            with engine.eval_mode():\n                # runs in evaluation mode\n        \"\"\"\n        raise NotImplementedError\n\n    def optimizer_zero_grad(self):\n        \"\"\"\n        Zero the gradients of the optimizer.\n        \"\"\"\n        raise NotImplementedError\n\n    def optimizer_step(self):\n        \"\"\"\n        Perform an optimization step using the optimizer.\n        \"\"\"\n        raise NotImplementedError\n\n    def lr_scheduler_step(self):\n        \"\"\"\n        Advance the learning rate scheduler by one step.\n\n        Returns:\n            current_lr (float or list[float]): Updated learning rate(s).\n        \"\"\"\n        raise NotImplementedError\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:\n        \"\"\"\n        Perform a forward pass and optionally a backward pass on a batch of data.\n\n        Args:\n            data: The input data for the forward pass, typically containing tensors and metadata.\n            loss_function: The loss function to optimize. See `verl.workers.roles.utils.losses` for examples.\n            forward_only: If True, perform only the forward pass. If False, perform forward and backward pass.\n\n        Returns:\n            Any: The output of the forward pass, which can be used for loss computation or other purposes.\n        \"\"\"\n        raise NotImplementedError\n\n    def train_batch(self, data: TensorDict, loss_function: Callable) -> Any:\n        \"\"\"\n        Perform a training step on a batch of data.\n\n        Args:\n            data: The input data for training, typically containing tensors and metadata.\n            loss_function: A function that computes the loss and metrics given a batch and predictions.\n\n        Returns:\n            dict[str, torch.Tensor]: A dictionary containing the aggregated training metrics for the batch.\n        \"\"\"\n        maybe_fix_3d_position_ids(data)\n\n        self.optimizer_zero_grad()\n        outputs = self.forward_backward_batch(data, loss_function, forward_only=False)\n        grad_norm = self.optimizer_step()\n        if self.is_mp_src_rank_with_outputs():\n            assert \"grad_norm\" not in outputs[\"metrics\"]\n            outputs[\"metrics\"][\"grad_norm\"] = grad_norm\n        return outputs\n\n    def infer_batch(self, data: TensorDict, loss_function: Optional[Callable] = None) -> Any:\n        \"\"\"\n        Perform inference on a batch of data.\n\n        Args:\n            data: The input data for inference, typically containing tensors and metadata.\n\n        Returns:\n            Any: The output of the inference, which can be used for predictions or other purposes.\n        \"\"\"\n        # see comments from train_batch\n        maybe_fix_3d_position_ids(data)\n\n        with torch.no_grad():\n            outputs = self.forward_backward_batch(data, loss_function, forward_only=True)\n        return outputs\n\n    def get_per_tensor_param(self) -> tuple[Generator[tuple[str, torch.Tensor], None, None], Optional[dict]]:\n        \"\"\"\n        Get a generator that yields per-tensor parameters and optional peft config.\n\n        Returns:\n            Generator[tuple[str, torch.Tensor]]: A generator that yields tuples of parameter names and tensors.\n            Optional[dict]: Optional peft config.\n        \"\"\"\n        raise NotImplementedError\n\n    def get_data_parallel_size(self):\n        raise NotImplementedError\n\n    def get_data_parallel_rank(self):\n        raise NotImplementedError\n\n    def get_data_parallel_group(self):\n        raise NotImplementedError\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"\n        Move model parameters, optimizer states, or both to the specified device.\n\n        Args:\n            device: Target device identifier.\n            model: If True, move the model.\n            optimizer: If True, move the optimizer states.\n            grad: If True, move the gradient buffer.\n        \"\"\"\n        if not model:\n            assert not optimizer and not grad, \"Model must be moved to device along with optimizer and grad\"\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Save model, optimizer, and scheduler states to a checkpoint.\n\n        Args:\n            local_path: Local filesystem path to save checkpoint.\n            hdfs_path: Optional HDFS path to copy checkpoint.\n            global_step: Integer training step number for naming.\n            max_ckpt_to_keep: Maximum number of recent checkpoints to retain.\n            **kwargs: Arbitrary keyword arguments.\n        \"\"\"\n        raise NotImplementedError\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: bool = True, **kwargs\n    ) -> None:\n        \"\"\"\n        Load model, optimizer, and scheduler states from a checkpoint.\n\n        Args:\n            local_path: Local filesystem path of the checkpoint.\n            hdfs_path: Optional HDFS path where checkpoint is stored.\n            del_local_after_load: Whether to delete local copy after loading.\n            **kwargs: Arbitrary keyword arguments.\n        \"\"\"\n        raise NotImplementedError\n\n    def is_mp_src_rank_with_outputs(self):\n        \"\"\"\n        Whether the current rank is the first rank in model parallel group that contains model outputs\n        \"\"\"\n        raise NotImplementedError\n\n    def disable_adapter(self) -> ContextManager:\n        \"\"\"\n        Disable all adapters temporarily under the context in the model for LoRA\n        \"\"\"\n        return nullcontext()\n\n\nclass BaseEngineCtx:\n    def __init__(self, engine: BaseEngine, mode, **kwargs):\n        \"\"\"Base Engine context that handles load and offload\n\n        Args:\n            engine:\n            **kwargs:\n        \"\"\"\n        self.engine = engine\n        self.mode = mode\n        assert self.mode in (\"train\", \"eval\")\n        self.disable_auto_offload = kwargs.pop(\"disable_auto_offload\", False)\n\n    def _context_switch(self, device):\n        if self.disable_auto_offload:\n            return\n        if self.mode == \"eval\":\n            self.engine.to(device=device, model=self.engine.is_param_offload_enabled, optimizer=False, grad=False)\n        elif self.mode == \"train\":\n            self.engine.to(\n                device=device,\n                model=self.engine.is_param_offload_enabled,\n                optimizer=self.engine.is_optimizer_offload_enabled,\n                grad=self.engine.is_param_offload_enabled,\n            )\n\n    def __enter__(self):\n        self._context_switch(get_device_name())\n        self.engine.mode = self.mode\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        self._context_switch(\"cpu\")\n        self.engine.mode = None\n\n\nclass EngineRegistry:\n    \"\"\"\n    A registry for managing and instantiating different types of training engines.\n\n    This class uses a dictionary to store engine classes, mapping a string key to each class.\n    It provides a decorator `register` to add new engines to the registry and a `new` method\n    to create an instance of a registered engine.\n    \"\"\"\n\n    _engines = {}\n\n    @classmethod\n    def register(cls, model_type: str, backend: list[str] | str, device: list[str] | str = \"cuda\"):\n        \"\"\"\n        A class method decorator that registers an engine class with a given key.\n\n        This allows for dynamic instantiation of engine classes by their registered key.\n\n        Args:\n            model_type (str): The type of the model\n            backend (list[str] | str): The backend to use for the model type\n            device (list[str] | str): The device type (e.g., \"cuda\", \"npu\", \"cpu\") this engine supports,\n                default is \"cuda\"\n\n        Returns:\n            A decorator function that takes an engine class and registers it.\n        \"\"\"\n\n        def decorator(engine_class):\n            assert issubclass(engine_class, BaseEngine)\n            if model_type not in cls._engines:\n                cls._engines[model_type] = {}\n\n            backends = backend if isinstance(backend, list) else [backend]\n            devices = device if isinstance(device, list) else [device]\n            for current_backend in backends:\n                for current_device in devices:\n                    if current_backend not in cls._engines[model_type]:\n                        cls._engines[model_type][current_backend] = {}\n                    if current_device not in cls._engines[model_type][current_backend]:\n                        cls._engines[model_type][current_backend][current_device] = engine_class\n\n            return engine_class\n\n        return decorator\n\n    @classmethod\n    def get_engine_cls(cls, model_type: str, backend: str):\n        assert model_type in cls._engines, f\"Unknown model_type: {model_type}\"\n        assert backend in cls._engines[model_type], f\"Unknown backend: {backend}\"\n        device = get_device_name()\n        assert device in cls._engines[model_type][backend], (\n            f\"Unknown device: {device} for model_type: {model_type} and backend: {backend}\"\n        )\n        return cls._engines[model_type][backend][device]\n\n    @classmethod\n    def new(cls, model_type, backend, *args, **kwargs):\n        \"\"\"\n        Function to create a new training engine instance based on the provided config.\n        Args:\n            key: A configuration object containing the engine key and other settings.\n            *args: Variable length argument list.\n            **kwargs: Arbitrary keyword arguments.\n        Returns:\n            engine: An instance of the training engine corresponding to the config.\n        Raises:\n            NotImplementedError: If the engine key in the config does not match any known engines.\n        \"\"\"\n        engine_cls = cls.get_engine_cls(model_type, backend)\n        return engine_cls(*args, **kwargs)\n"
  },
  {
    "path": "verl/workers/engine/fsdp/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom .transformer_impl import FSDPEngine, FSDPEngineWithLMHead\n\n__all__ = [\"FSDPEngine\", \"FSDPEngineWithLMHead\"]\n"
  },
  {
    "path": "verl/workers/engine/fsdp/transformer_impl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe concrete Engine implementation using PyTorch FullyShardedDataParallel (FSDP)\n\"\"\"\n\nimport gc\nimport logging\nimport os\nimport warnings\nfrom contextlib import nullcontext\nfrom typing import Callable, ContextManager, Optional\n\nimport torch\nimport torch.distributed\nfrom peft import LoraConfig, TaskType, get_peft_model\nfrom tensordict import TensorDict\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType\nfrom torch.distributed.tensor import DTensor\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.models.transformers.monkey_patch import apply_monkey_patch\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.activation_offload import enable_activation_offloading\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.debug import log_gpu_memory_usage\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.fsdp_utils import (\n    CPUOffloadPolicy,\n    FSDPModule,\n    MixedPrecisionPolicy,\n    apply_fsdp2,\n    collect_lora_params,\n    fsdp2_clip_grad_norm_,\n    fsdp2_load_full_state_dict,\n    fsdp_version,\n    get_fsdp_wrap_policy,\n    get_init_weight_context_manager,\n    init_fn,\n    load_fsdp_model_to_gpu,\n    load_fsdp_optimizer,\n    merged_lora_context,\n    normalize_peft_param_name,\n    offload_fsdp_model_to_cpu,\n    offload_fsdp_optimizer,\n    replace_lora_wrapper,\n)\nfrom verl.utils.model import convert_weight_keys, extract_multi_modal_inputs\nfrom verl.utils.py_functional import convert_to_regular_types\nfrom verl.utils.torch_functional import logprobs_from_logits\nfrom verl.utils.ulysses import (\n    gather_outputs_and_unpad,\n    get_ulysses_sequence_parallel_group,\n    set_ulysses_sequence_parallel_group,\n    ulysses_pad,\n    ulysses_pad_and_slice_inputs,\n)\nfrom verl.workers.config import FSDPEngineConfig, FSDPOptimizerConfig, HFModelConfig\n\nfrom ..base import BaseEngine, BaseEngineCtx, EngineRegistry\nfrom ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches\nfrom .utils import create_device_mesh, get_sharding_strategy\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\ndevice_name = get_device_name()\n\n\nclass FSDPEngine(BaseEngine):\n    \"\"\"\n    Concrete Engine implementation using PyTorch FullyShardedDataParallel (FSDP).\n\n    Supports model sharding, activation/optimizer offloading, LoRA, and sequence parallelism.\n    \"\"\"\n\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: FSDPEngineConfig,\n        optimizer_config: FSDPOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n    ):\n        \"\"\"\n        Initialize the FSDPEngine.\n\n        Sets up distributed device meshes, LoRA, and offload policies based on config.\n\n        Args:\n            config: Configuration object with FSDP and model settings.\n        \"\"\"\n        super().__init__()\n\n        self.model_config = model_config\n        self.engine_config = engine_config\n        self.optimizer_config = optimizer_config\n        self.checkpoint_config = checkpoint_config\n\n        self.mode = None\n\n        self.rank = torch.distributed.get_rank()\n\n        # Apply NPU patches for FSDP backend\n        from .utils import apply_npu_fsdp_patches\n\n        apply_npu_fsdp_patches()\n\n        # build device mesh for Ulysses Sequence Parallel\n\n        self.use_remove_padding = self.model_config.use_remove_padding\n\n        self._init_device_mesh()\n\n        if self.engine_config.full_determinism:\n            enable_full_determinism(seed=self.engine_config.seed)\n\n        # set FSDP offload params\n        self._is_offload_param = self.engine_config.param_offload\n        self._is_offload_optimizer = self.engine_config.optimizer_offload\n        self._is_lora = self.model_config.lora_rank > 0\n\n        # QAT (Quantization-Aware Training)\n        self._qat_config = getattr(self.engine_config, \"qat\", None)\n        self._qat_enabled = self._qat_config is not None and getattr(self._qat_config, \"enable\", False)\n        if self._qat_enabled:\n            logger.info(f\"QAT enabled: mode={self._qat_config.mode}, group_size={self._qat_config.group_size}\")\n\n        if self.engine_config.entropy_from_logits_with_chunking:\n            entropy_from_logits = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_from_logits = verl_F.entropy_from_logits\n\n        self.compute_entropy_from_logits = (\n            torch.compile(entropy_from_logits, dynamic=True)\n            if self.engine_config.use_torch_compile  #  use torch compile by default\n            else entropy_from_logits\n        )\n\n    @property\n    def is_param_offload_enabled(self) -> bool:\n        return self._is_offload_param\n\n    @property\n    def is_optimizer_offload_enabled(self) -> bool:\n        return self._is_offload_optimizer\n\n    def is_mp_src_rank_with_outputs(self):\n        if self.ulysses_device_mesh is not None:\n            is_collect = self.ulysses_device_mesh[\"sp\"].get_local_rank() == 0\n        else:\n            is_collect = True\n        return is_collect\n\n    def initialize(self):\n        \"\"\"\n        Build the model, optimizer, and learning rate scheduler under FSDP.\n\n        Applies device, dtype, and precision configurations, including mixed precision.\n        Sets up checkpoint manager and FLOPs counter.\n        \"\"\"\n        # This is used to import external_lib into the huggingface systems\n        self._build_model_optimizer()\n\n        self.checkpoint_manager = FSDPCheckpointManager(\n            model=self.module,\n            optimizer=self.optimizer,\n            lr_scheduler=self.lr_scheduler,\n            processing_class=self.model_config.get_processor(),\n            checkpoint_config=self.checkpoint_config,\n            trust_remote_code=self.model_config.trust_remote_code,\n        )\n\n        self.to(\n            device=\"cpu\",\n            model=self._is_offload_param,\n            optimizer=self._is_offload_optimizer,\n            grad=self._is_offload_param,\n        )\n\n        log_gpu_memory_usage(\"After offload model/optimizer/grad during init\", logger=logger)\n\n    def _init_device_mesh(self):\n        world_size = torch.distributed.get_world_size()\n        from torch.distributed.device_mesh import init_device_mesh\n\n        fsdp_size = self.engine_config.fsdp_size\n\n        self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size)\n        self.ulysses_device_mesh = None\n        self.ulysses_parallel_group = None\n        self.ulysses_sequence_parallel_size = self.engine_config.ulysses_sequence_parallel_size\n        dp_size = self.get_data_parallel_size()\n        if self.ulysses_sequence_parallel_size > 1:\n            self.ulysses_device_mesh = init_device_mesh(\n                device_name, mesh_shape=(dp_size, self.ulysses_sequence_parallel_size), mesh_dim_names=[\"dp\", \"sp\"]\n            )\n            self.ulysses_parallel_group = self.ulysses_device_mesh[\"sp\"].get_group()\n\n        self.use_ulysses_sp = self.ulysses_sequence_parallel_size > 1\n\n    def _build_module(self):\n        from verl.utils.model import get_hf_auto_model_class\n        from verl.utils.torch_dtypes import PrecisionType\n\n        torch_dtype = self.engine_config.model_dtype\n\n        if torch_dtype is None:\n            # if it is training, we force torch_dtype to fp32\n            torch_dtype = torch.float32 if not self.engine_config.forward_only else torch.bfloat16\n\n        torch_dtype = PrecisionType.to_dtype(torch_dtype)\n\n        init_context = get_init_weight_context_manager(\n            use_meta_tensor=not self.model_config.hf_config.tie_word_embeddings, mesh=self.device_mesh\n        )\n\n        with init_context(), warnings.catch_warnings():\n            warnings.simplefilter(\"ignore\")\n\n            auto_class = get_hf_auto_model_class(hf_config=self.model_config.hf_config)\n\n            module = auto_class.from_pretrained(\n                pretrained_model_name_or_path=self.model_config.local_path,\n                torch_dtype=torch_dtype,\n                config=self.model_config.hf_config,\n                trust_remote_code=self.model_config.trust_remote_code,\n            )\n\n            use_liger = self.model_config.use_liger\n            # Apply Liger kernel to the model if use_liger is set to True\n            if use_liger:\n                from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance\n\n                _apply_liger_kernel_to_instance(model=module)\n\n            fused_kernel_options = self.model_config.fused_kernel_options\n            fused_kernels_backend = (\n                fused_kernel_options.get(\"impl_backend\", None) if fused_kernel_options is not None else None\n            )\n\n            use_fused_kernels = self.model_config.use_fused_kernels\n            apply_monkey_patch(\n                model=module,\n                use_remove_padding=self.use_remove_padding,\n                ulysses_sp_size=self.ulysses_sequence_parallel_size,\n                use_fused_kernels=use_fused_kernels,\n                fused_kernels_backend=fused_kernels_backend,\n            )\n\n            # some parameters may not in torch_dtype\n            module.to(torch_dtype)\n\n            if self.model_config.enable_gradient_checkpointing:\n                module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n        return module\n\n    def _build_lora_module(self, module):\n        module.enable_input_require_grads()\n\n        lora_adapter_path = getattr(self.model_config, \"lora_adapter_path\", None)\n        if lora_adapter_path is not None:\n            from peft import PeftModel\n\n            from verl.utils.fs import copy_to_local\n\n            print(f\"Loading pre-trained LoRA adapter to from: {lora_adapter_path}\")\n            # Copy adapter to local if needed\n            local_adapter_path = copy_to_local(lora_adapter_path, use_shm=self.model_config.use_shm)\n\n            module = PeftModel.from_pretrained(module, local_adapter_path, is_trainable=True)\n            peft_config = module.peft_config[\"default\"]\n            # Ensure task_type is TaskType enum, not string\n            if isinstance(peft_config.task_type, str):\n                peft_config.task_type = TaskType.CAUSAL_LM\n        else:\n            # Convert config to regular Python types before creating PEFT model\n            lora_config = {\n                \"task_type\": TaskType.CAUSAL_LM,\n                \"r\": self.model_config.lora_rank,\n                \"lora_alpha\": self.model_config.lora_alpha,\n                \"target_modules\": convert_to_regular_types(self.model_config.target_modules),\n                \"target_parameters\": convert_to_regular_types(self.model_config.target_parameters),\n                \"exclude_modules\": convert_to_regular_types(self.model_config.exclude_modules),\n                \"bias\": \"none\",\n            }\n            module = get_peft_model(module, LoraConfig(**lora_config))\n\n        return module\n\n    def _build_fsdp_module(self, module):\n        # TODO(ziheng): need to improve\n        from torch.distributed.fsdp import CPUOffload, MixedPrecision\n\n        from verl.utils.torch_dtypes import PrecisionType\n\n        mixed_precision_config = self.engine_config.mixed_precision\n        if mixed_precision_config is not None:\n            param_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"param_dtype\", \"bf16\"))\n            reduce_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"reduce_dtype\", \"fp32\"))\n            buffer_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"buffer_dtype\", \"fp32\"))\n        else:\n            param_dtype = torch.bfloat16\n            reduce_dtype = torch.float32\n            buffer_dtype = torch.float32\n\n        mixed_precision = MixedPrecision(param_dtype=param_dtype, reduce_dtype=reduce_dtype, buffer_dtype=buffer_dtype)\n\n        auto_wrap_policy = get_fsdp_wrap_policy(\n            module=module,\n            config=self.engine_config.wrap_policy,\n            is_lora=self.model_config.lora_rank > 0,\n        )\n\n        fsdp_mesh = self.device_mesh\n        sharding_strategy = get_sharding_strategy(fsdp_mesh)\n\n        # Note: We force turn off CPUOffload because it causes incorrect results when using grad accumulation\n        if self.engine_config.strategy == \"fsdp\":\n            # cpu_offload:\n            # - actor: None\n            # - critic: None\n            # - ref: CPUOffload(offload_params=True)\n\n            # We force reference policy to use CPUOffload to save memory.\n            # We force turn off CPUOffload for actor because it causes incorrect results when using grad accumulation\n            cpu_offload = None\n            if self.engine_config.forward_only:\n                cpu_offload = CPUOffload(offload_params=True)\n                self._is_offload_param = False\n                self._is_offload_optimizer = False\n\n            module = FSDP(\n                module,\n                param_init_fn=init_fn,\n                auto_wrap_policy=auto_wrap_policy,\n                device_id=get_device_id(),\n                sharding_strategy=sharding_strategy,\n                mixed_precision=mixed_precision,\n                sync_module_states=True,\n                device_mesh=self.device_mesh,\n                forward_prefetch=self.engine_config.forward_prefetch,\n                use_orig_params=self.engine_config.use_orig_params,\n                cpu_offload=cpu_offload,\n            )\n        elif self.engine_config.strategy == \"fsdp2\":\n            # - actor: offload_policy\n            # - critic: offload_policy\n            # - ref: CPUOffloadPolicy(pin_memory=True)\n            assert CPUOffloadPolicy is not None, \"PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)\"\n            mp_policy = MixedPrecisionPolicy(\n                param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True\n            )\n            offload_policy = None\n            if self.engine_config.offload_policy or self.engine_config.forward_only:\n                self._is_offload_param = False\n                self._is_offload_optimizer = False\n                offload_policy = CPUOffloadPolicy(pin_memory=True)\n\n            fsdp_kwargs = {\n                \"mesh\": fsdp_mesh,\n                \"mp_policy\": mp_policy,\n                \"offload_policy\": offload_policy,\n                \"reshard_after_forward\": self.engine_config.reshard_after_forward,\n            }\n            full_state = module.state_dict()\n            apply_fsdp2(module, fsdp_kwargs, self.engine_config)\n            fsdp2_load_full_state_dict(module, full_state, fsdp_mesh, offload_policy)\n        else:\n            raise NotImplementedError(f\"Unknown strategy {self.engine_config.strategy}\")\n\n        if self.model_config.enable_activation_offload:\n            enable_gradient_checkpointing = self.model_config.enable_gradient_checkpointing\n            enable_activation_offloading(module, self.engine_config.strategy, enable_gradient_checkpointing)\n\n        if torch.distributed.get_world_size() == 1 and fsdp_version(module) == 1:\n            FSDP.set_state_dict_type(\n                module,\n                state_dict_type=StateDictType.FULL_STATE_DICT,\n                state_dict_config=FullStateDictConfig(),\n            )\n        elif fsdp_version(module) == 1:\n            FSDP.set_state_dict_type(\n                module,\n                state_dict_type=StateDictType.SHARDED_STATE_DICT,\n                state_dict_config=ShardedStateDictConfig(),\n            )\n\n        return module\n\n    def _build_optimizer(self, module):\n        from verl.workers.config.optimizer import build_optimizer\n\n        optimizer = build_optimizer(module.parameters(), self.optimizer_config)\n\n        return optimizer\n\n    def _build_lr_scheduler(self, optimizer):\n        from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup\n\n        optim_config = self.optimizer_config\n\n        total_steps = optim_config.total_training_steps\n        num_warmup_steps = optim_config.lr_warmup_steps\n        lr_scheduler_type = optim_config.lr_scheduler_type\n        min_lr_ratio = optim_config.min_lr_ratio\n        num_cycles = optim_config.num_cycles\n        zero_indexed_step = optim_config.zero_indexed_step\n        if num_warmup_steps <= 0:\n            num_warmup_steps_ratio = optim_config.lr_warmup_steps_ratio\n            num_warmup_steps = int(num_warmup_steps_ratio * total_steps)\n\n        if self.rank == 0:\n            print(f\"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}\")\n\n        if lr_scheduler_type == \"constant\":\n            lr_scheduler = get_constant_schedule_with_warmup(optimizer=optimizer, num_warmup_steps=num_warmup_steps)\n        elif lr_scheduler_type == \"cosine\":\n            lr_scheduler = get_cosine_schedule_with_warmup(\n                optimizer=optimizer,\n                num_warmup_steps=num_warmup_steps,\n                num_training_steps=total_steps,\n                min_lr_ratio=min_lr_ratio,\n                num_cycles=num_cycles,\n                zero_indexed_step=zero_indexed_step,\n            )\n        else:\n            raise NotImplementedError(f\"LR scheduler type {lr_scheduler_type} is not supported\")\n        return lr_scheduler\n\n    def _apply_qat(self, module):\n        \"\"\"Apply QAT transformations to the model before FSDP wrapping.\"\"\"\n        from verl.utils.qat.core import apply_qat, enable_qat_fuse\n\n        module = apply_qat(\n            module,\n            {\n                \"enable\": self._qat_config.enable,\n                \"mode\": self._qat_config.mode,\n                \"group_size\": self._qat_config.group_size,\n                \"ignore_patterns\": list(self._qat_config.ignore_patterns),\n                \"activation_observer\": self._qat_config.activation_observer,\n            },\n        )\n        enable_qat_fuse(module)\n\n        if self._qat_config.mode == \"w4a4\":\n            self._restore_w4a4_input_scales(module, self.model_config.local_path)\n\n        return module\n\n    def _restore_w4a4_input_scales(self, model, model_path):\n        \"\"\"Restore input_global_scale and input_amax from checkpoint for W4A4 mode.\"\"\"\n        import glob\n\n        from safetensors import safe_open\n\n        safetensor_files = glob.glob(f\"{model_path}/model*.safetensors\")\n        loaded_count = 0\n\n        for sf_path in safetensor_files:\n            with safe_open(sf_path, framework=\"pt\") as f:\n                for key in f.keys():\n                    if \"input_global_scale\" in key:\n                        module_path = key.replace(\".input_global_scale\", \"\")\n                        amax_key = f\"{module_path}.input_amax\"\n\n                        module = model\n                        for part in module_path.split(\".\"):\n                            module = module[int(part)] if part.isdigit() else getattr(module, part)\n\n                        scale_val = f.get_tensor(key)\n                        val = scale_val.item() if scale_val.numel() == 1 else scale_val.max().item()\n                        module.input_global_scale.fill_(val)\n\n                        amax_val = f.get_tensor(amax_key)\n                        amax = amax_val.item() if amax_val.numel() == 1 else amax_val.max().item()\n                        module.input_amax.fill_(amax)\n                        loaded_count += 1\n\n        logger.info(f\"[QAT W4A4] Restored {loaded_count} input_global_scale/input_amax from {model_path}\")\n\n    def _build_model_optimizer(self):\n        from verl.utils.model import print_model_size\n\n        # Load base model with specified configuration and dtype\n        module = self._build_module()\n        # Apply LoRA adapters if low-rank adaptation is enabled\n        if self._is_lora:\n            module = self._build_lora_module(module)\n\n        # Apply QAT before FSDP wrapping (training only)\n        if self._qat_enabled and not self.engine_config.forward_only:\n            module = self._apply_qat(module)\n\n        # Synchronize all distributed processes before proceeding\n        torch.distributed.barrier()\n        if self.rank == 0:\n            print_model_size(module)\n        log_gpu_memory_usage(\"After init model from HF AutoModel\", logger=logger)\n\n        # Wrap model with FSDP for distributed training (sharding, mixed precision, etc.)\n        log_gpu_memory_usage(\"Before FSDP\", logger=None)\n        module = self._build_fsdp_module(module)\n        log_gpu_memory_usage(\"After FSDP\", logger=None)\n\n        if not self.engine_config.forward_only:\n            # Initialize optimizer with model parameters and config settings\n            optimizer = self._build_optimizer(module)\n            # Create learning rate scheduler with warmup and decay settings\n            lr_scheduler = self._build_lr_scheduler(optimizer)\n        else:\n            optimizer = None\n            lr_scheduler = None\n\n        self.module = module\n        self.optimizer = optimizer\n        self.lr_scheduler = lr_scheduler\n\n    def train_mode(self, **kwargs):\n        \"\"\"\n        Return a context manager that switches to training mode with FSDP-specific handling.\n\n        Includes parameter and optimizer offload entry/exit.\n        \"\"\"\n        return EngineTrainModeCtx(self, **kwargs)\n\n    def eval_mode(self, **kwargs):\n        \"\"\"\n        Return a context manager that switches to evaluation mode with FSDP-specific handling.\n\n        Includes activation offload entry/exit.\n        \"\"\"\n        return EngineEvalModeCtx(self, **kwargs)\n\n    def get_data_parallel_rank(self):\n        if self.ulysses_device_mesh is not None:\n            return self.ulysses_device_mesh[\"dp\"].get_local_rank()\n        else:\n            return torch.distributed.get_rank()\n\n    def get_data_parallel_size(self):\n        return torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size\n\n    def get_data_parallel_group(self):\n        if self.ulysses_device_mesh is not None:\n            return self.ulysses_device_mesh.get_group(mesh_dim=\"dp\")\n        else:\n            return torch.distributed.group.WORLD\n\n    def get_model_parallel_group(self):\n        raise NotImplementedError\n\n    def get_context_parallel_group(self):\n        raise NotImplementedError\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> list[TensorDict]:\n        # note that the global_batch_size should include data on all the dp\n        tu.assign_non_tensor(data, sp_size=self.ulysses_sequence_parallel_size)\n\n        # compute num_tokens in global batch for loss normalization\n        batch_num_tokens = data[\"loss_mask\"].sum().to(get_device_id())\n        torch.distributed.all_reduce(\n            batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group()\n        )\n        tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())\n        tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())\n\n        micro_batches, indices = prepare_micro_batches(\n            data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True\n        )\n\n        output_lst = []\n\n        ctx = torch.no_grad() if forward_only else nullcontext()\n\n        for micro_batch in micro_batches:\n            with ctx:\n                loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only)\n\n                if not forward_only:\n                    loss.backward()\n\n            output_lst.append(meta_info)\n\n        # postprocess and return\n        return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data)\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        raise NotImplementedError(\"forward_step must be implemented in subclass\")\n\n    def optimizer_zero_grad(self):\n        \"\"\"\n        Zero gradients and enforce FSDP grad-clipping logic.\n        \"\"\"\n        self.optimizer.zero_grad()\n\n    def optimizer_step(self):\n        \"\"\"\n        Clip gradients, skip update if non-finite, and step optimizer.\n\n        Returns:\n            grad_norm (float): Norm of gradients before clipping.\n        \"\"\"\n        assert self.optimizer_config.clip_grad is not None\n\n        if isinstance(self.module, FSDP):\n            grad_norm = self.module.clip_grad_norm_(self.optimizer_config.clip_grad)\n        elif isinstance(self.module, FSDPModule):\n            grad_norm = fsdp2_clip_grad_norm_(self.module.parameters(), max_norm=self.optimizer_config.clip_grad)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(\n                self.module.parameters(), max_norm=self.optimizer_config.clip_grad\n            )\n\n        if isinstance(grad_norm, DTensor):\n            grad_norm = grad_norm.full_tensor()\n\n        # if grad_norm is not finite, skip the update\n        if not torch.isfinite(grad_norm):\n            print(f\"WARN: grad_norm is not finite: {grad_norm}\")\n            self.optimizer.zero_grad()\n        else:\n            self.optimizer.step()\n\n        if self._qat_enabled:\n            from verl.utils.qat.core import invalidate_all_scales\n\n            invalidate_all_scales(self.module)\n\n        return grad_norm.item()\n\n    def lr_scheduler_step(self):\n        \"\"\"\n        Advance FSDP scheduler and return updated learning rate.\n        \"\"\"\n        self.lr_scheduler.step()\n        lr = self.lr_scheduler.get_last_lr()[0]  # only return the first group\n        return lr\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"\n        Move FSDP model and/or optimizer to CPU or GPU with offload support.\n        Note that this function executes irrespective of offload config. It serves as manual control\n        \"\"\"\n        super().to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n        if self.engine_config.forward_only:\n            # force cpu_offload\n            return\n\n        device_name = get_device_name()\n\n        assert device in (device_name, \"cpu\")\n        if device == device_name:\n            if model:\n                load_fsdp_model_to_gpu(self.module)\n            if optimizer and self.optimizer is not None:\n                load_fsdp_optimizer(self.optimizer, device)\n            gc.collect()\n        elif device == \"cpu\":\n            if model:\n                offload_fsdp_model_to_cpu(self.module)\n            if optimizer and self.optimizer is not None:\n                offload_fsdp_optimizer(self.optimizer)\n        else:\n            raise ValueError(f\"Invalid device type: {device}\")\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Save FSDP checkpoint, handling parameter offload as needed.\n        \"\"\"\n        origin_module_device = next(self.module.parameters()).device.type\n        if self._is_offload_param or origin_module_device == \"cpu\":\n            load_fsdp_model_to_gpu(self.module)\n\n        self.checkpoint_manager.save_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.module)\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs\n    ) -> None:\n        \"\"\"\n        Load FSDP checkpoint, restoring parameters and optimizer state.\n        \"\"\"\n        import torch\n\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.module)\n\n        self.checkpoint_manager.load_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.module)\n\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(self.optimizer)\n\n    def get_per_tensor_param(self, layered_summon=False, base_sync_done=False, **kwargs):\n        log_gpu_memory_usage(\"Before load_fsdp_model_to_gpu\", logger=logger)\n\n        load_fsdp_model_to_gpu(self.module)\n\n        log_gpu_memory_usage(\"After load_fsdp_model_to_gpu\", logger=logger)\n\n        peft_config = None\n        merge_lora = self.model_config.lora.get(\"merge\", False)\n\n        peft_model = getattr(self.module, \"_fsdp_wrapped_module\", self.module)\n        if hasattr(peft_model, \"peft_config\"):  # LoRA\n            if not merge_lora:\n                peft_config = peft_model.peft_config.get(\"default\", None)\n                params = collect_lora_params(\n                    module=self.module,\n                    layered_summon=layered_summon,\n                    base_sync_done=base_sync_done,\n                )\n                if not base_sync_done:\n                    params = {replace_lora_wrapper(k, peft_config): v for k, v in params.items()}\n            else:  # merge lora\n                with merged_lora_context(self.module, backup_adapters=True):\n                    params = self.module.state_dict()\n                    params = normalize_peft_param_name(params)\n        else:\n            params = self.module.state_dict()\n\n        params = convert_weight_keys(params, getattr(self.module, \"_fsdp_wrapped_module\", self.module))\n\n        log_gpu_memory_usage(\"Before offload_fsdp_model_to_cpu\", logger=logger)\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.module)\n        log_gpu_memory_usage(\"After offload_fsdp_model_to_cpu\", logger=logger)\n\n        if peft_config is not None and base_sync_done:\n            per_tensor_param = params.items()\n        else:\n            device = get_device_id()  # used when fsdp2 set cpu_offload_policy\n            # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate\n            per_tensor_param = (\n                (\n                    name,\n                    param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)\n                    if isinstance(param, DTensor)\n                    else param,\n                )\n                for name, param in params.items()\n            )\n\n        if self._qat_enabled:\n            from verl.utils.qat.quantizer import QATQuantizer\n            from verl.utils.torch_dtypes import PrecisionType\n\n            mixed_precision_config = self.engine_config.mixed_precision\n            if mixed_precision_config is not None:\n                param_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"param_dtype\", \"bf16\"))\n            else:\n                param_dtype = torch.bfloat16\n\n            quantizer = QATQuantizer(\n                mode=self._qat_config.mode,\n                group_size=self._qat_config.group_size,\n                ignore_patterns=list(self._qat_config.ignore_patterns),\n                device=torch.device(get_device_id()),\n                param_dtype=param_dtype,\n            )\n            per_tensor_param = quantizer.quantize_with_fusion(\n                per_tensor_param,\n                target_device=torch.device(\"cpu\"),\n            )\n\n        peft_config_dict = peft_config.to_dict() if peft_config is not None else None\n        return per_tensor_param, peft_config_dict\n\n    def disable_adapter(self) -> ContextManager:\n        return self.module.disable_adapter()\n\n\nclass EngineEvalModeCtx(BaseEngineCtx):\n    def __init__(self, engine: FSDPEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"eval\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, FSDPEngine)\n        super().__enter__()\n        self.prev_sp_group = get_ulysses_sequence_parallel_group()\n        set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group)\n        self.engine.module.eval()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, FSDPEngine)\n        set_ulysses_sequence_parallel_group(self.prev_sp_group)\n\n        # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes\n        # unshard the root FSDP module\n        if self.engine.engine_config.fsdp_size > 1:\n            if fsdp_version(self.engine.module) == 1:\n                self.engine.module._handle.reshard(True)\n            elif fsdp_version(self.engine.module) == 2:\n                self.engine.module.reshard()\n\n        super().__exit__(exc_type, exc_value, traceback)\n\n\nclass EngineTrainModeCtx(BaseEngineCtx):\n    def __init__(self, engine: FSDPEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"train\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, FSDPEngine)\n        super().__enter__()\n        self.prev_sp_group = get_ulysses_sequence_parallel_group()\n        set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group)\n        self.engine.module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, FSDPEngine)\n        set_ulysses_sequence_parallel_group(self.prev_sp_group)\n        self.engine.optimizer_zero_grad()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=[\"fsdp\", \"fsdp2\"], device=[\"cuda\", \"npu\"])\nclass FSDPEngineWithLMHead(FSDPEngine):\n    def prepare_model_inputs(self, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key=\"use_fused_kernels\", default=False)\n        temperature = micro_batch[\"temperature\"]\n        temperature_item = temperature\n        if use_fused_kernels:\n            assert not isinstance(temperature, torch.Tensor), (\n                \"use_fused_kernels does not support per sample temperature yet\"\n            )\n        assert pad_mode == DatasetPadMode.NO_PADDING, f\"pad_mode {pad_mode} not supported\"\n\n        multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get(\"multi_modal_inputs\", []))\n        input_ids = micro_batch[\"input_ids\"]\n        position_ids = micro_batch[\"position_ids\"]\n\n        if not isinstance(temperature, torch.Tensor):\n            temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device)\n\n        temperature = temperature.to(torch.float32)\n        assert temperature.shape[0] == input_ids.shape[0]\n\n        # args used to get outputs\n        output_args = {}\n\n        if use_remove_padding:\n            # support per sample temperature\n            # temperature (bsz,)\n            # input_ids (bsz, j1)\n            temperature_rmpad = verl_F.expand_as_nested(temperature, input_ids).values()  # (total_nnz,)\n            temperature_rmpad = temperature_rmpad.unsqueeze(0)  # (1, total_nnz)\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                input_ids_rmpad = input_ids.values().unsqueeze(0)  # (1, total_nnz)\n                if position_ids.dim() == 3:\n                    position_ids_rmpad = position_ids.values().unsqueeze(1)  # (4, 1, total_nnz)\n                else:\n                    position_ids_rmpad = position_ids.values().unsqueeze(0)  # (1, total_nnz)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n            # for compute the log_prob\n            input_ids_rmpad_rolled = torch.roll(input_ids_rmpad, shifts=-1, dims=1)  # (1, total_nnz)\n\n            # pad and slice the inputs if sp > 1\n            if self.use_ulysses_sp:\n                is_vlm_model = hasattr(getattr(self.module, \"module\", self.module).config, \"vision_config\")\n                if is_vlm_model:\n                    # vlm model's inputs will be sliced after embedding\n                    input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad(\n                        input_ids_rmpad,\n                        position_ids_rmpad=position_ids_rmpad,\n                        sp_size=self.ulysses_sequence_parallel_size,\n                    )\n                else:\n                    input_ids_rmpad, position_ids_rmpad, pad_size = ulysses_pad_and_slice_inputs(\n                        input_ids_rmpad,\n                        position_ids_rmpad=position_ids_rmpad,\n                        sp_size=self.ulysses_sequence_parallel_size,\n                        skip_position_ids_rmpad=True if self.__class__.__name__ == \"VeOmniEngineWithLMHead\" else False,\n                    )\n                input_ids_rmpad_rolled, _, _ = ulysses_pad_and_slice_inputs(\n                    input_ids_rmpad_rolled,\n                    position_ids_rmpad=None,\n                    sp_size=self.ulysses_sequence_parallel_size,\n                )\n\n                temperature_rmpad, _, _ = ulysses_pad_and_slice_inputs(\n                    temperature_rmpad, position_ids_rmpad=None, sp_size=self.ulysses_sequence_parallel_size, pad_value=1\n                )\n\n                output_args[\"pad_size\"] = pad_size\n\n            input_ids_rmpad_rolled = input_ids_rmpad_rolled.squeeze(0)  # ((total_nnz / sp) + pad)\n            temperature_rmpad = temperature_rmpad.squeeze(0)\n            output_args[\"input_ids_rmpad_rolled\"] = input_ids_rmpad_rolled\n            output_args[\"temperature_rmpad\"] = temperature_rmpad\n\n            # only pass input_ids and position_ids to enable flash_attn_varlen\n\n            model_inputs = {\n                \"input_ids\": input_ids_rmpad,\n                \"attention_mask\": None,\n                \"position_ids\": position_ids_rmpad,\n            }\n\n        else:\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                input_ids = micro_batch[\"input_ids\"]\n                position_ids = micro_batch[\"position_ids\"]\n                loss_mask = micro_batch[\"loss_mask\"]\n\n                pad_token_id = tu.get_non_tensor_data(data=micro_batch, key=\"pad_token_id\", default=0)\n                batch_size = micro_batch.batch_size[0]\n                seq_len_effective = input_ids.offsets().diff()\n                max_seq_len = max(seq_len_effective)\n\n                input_ids_rmpad_rolled = torch.roll(input_ids.values(), shifts=-1, dims=0)\n                output_args[\"input_ids_rmpad_rolled\"] = input_ids_rmpad_rolled\n                # we store the per sample temperature\n                output_args[\"temperature\"] = temperature\n\n                input_ids = torch.nested.to_padded_tensor(\n                    input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len)\n                )\n\n                if position_ids.dim() == 3:\n                    position_ids = torch.nested.to_padded_tensor(\n                        position_ids, padding=0, output_size=(batch_size, 4, max_seq_len)\n                    ).transpose(0, 1)  # (4, batch_size, max_seq_len)\n                else:\n                    position_ids = torch.nested.to_padded_tensor(\n                        position_ids, padding=0, output_size=(batch_size, max_seq_len)\n                    )\n\n                attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask]\n                attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged)\n                attention_mask = torch.nested.to_padded_tensor(\n                    attention_mask, padding=0, output_size=(batch_size, max_seq_len)\n                )\n\n                model_inputs = {\n                    \"input_ids\": input_ids,\n                    \"attention_mask\": attention_mask,\n                    \"position_ids\": position_ids,\n                }\n\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        extra_args = {}\n        if use_fused_kernels:\n            extra_args[\"temperature\"] = temperature_item\n            extra_args[\"return_dict\"] = True\n\n        model_inputs.update(multi_modal_inputs)\n        model_inputs.update(extra_args)\n\n        return model_inputs, output_args\n\n    def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        use_fused_kernels = tu.get_non_tensor_data(data=micro_batch, key=\"use_fused_kernels\", default=False)\n        calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key=\"calculate_entropy\", default=False)\n\n        model_output = {}\n\n        input_ids = micro_batch[\"input_ids\"]\n\n        if use_remove_padding:\n            input_ids_rmpad_rolled = output_args[\"input_ids_rmpad_rolled\"]\n            temperature_rmpad = output_args[\"temperature_rmpad\"]\n\n            if use_fused_kernels:\n                # temperature is singleton\n                log_probs = output.log_probs.squeeze(0)  # (total_nnz,)\n                entropy_rmpad = output.entropy.squeeze(0)  # (total_nnz,)\n            else:\n                logits_rmpad = output.logits.squeeze(0)  # (total_nnz, vocab_size)\n                logits_rmpad.div_(temperature_rmpad.clamp(min=1e-8).unsqueeze(-1).to(logits_rmpad.dtype))\n\n                # if use_sp: ((total_nnz / sp) + pad) ; if not use_sp: (batch, seqlen)\n                inplace_backward = True\n                if calculate_entropy:\n                    inplace_backward = False\n                log_probs = logprobs_from_logits(\n                    logits=logits_rmpad,\n                    labels=input_ids_rmpad_rolled,\n                    inplace_backward=inplace_backward,\n                )\n\n                # compute entropy\n                if calculate_entropy:\n                    if not self.engine_config.entropy_checkpointing:\n                        entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad)  # ((total_nnz / sp) + pad)\n                    else:\n                        entropy_rmpad = torch.utils.checkpoint.checkpoint(\n                            self.compute_entropy_from_logits, logits_rmpad\n                        )\n\n            # gather log_prob if sp > 1\n            if self.use_ulysses_sp:\n                pad_size = output_args[\"pad_size\"]\n\n                # gather and unpad for the ulysses sp\n                log_probs = gather_outputs_and_unpad(\n                    log_probs,\n                    gather_dim=0,\n                    unpad_dim=0,\n                    padding_size=pad_size,\n                )\n                if calculate_entropy:\n                    entropy_rmpad = gather_outputs_and_unpad(\n                        entropy_rmpad,\n                        gather_dim=0,\n                        unpad_dim=0,\n                        padding_size=pad_size,\n                    )\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                cu_seqlens = input_ids.offsets()\n                # (bsz, j1), for each sample, is the length of each sample: [real_prompt length + real_response length]\n                log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens)\n                if calculate_entropy:\n                    entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        else:  # not using rmpad and no ulysses sp\n            response_length = tu.get_non_tensor_data(data=micro_batch, key=\"max_response_length\", default=1024)\n            if use_fused_kernels:\n                log_probs = output.log_probs[:, -response_length - 1 : -1]\n                entropy = output.entropy[:, -response_length - 1 : -1]  # (bsz, response_length)\n\n            else:\n                logits = output.logits  # (bsz, response_length, vocab_size)\n                temperature = output_args[\"temperature\"]  # (bsz,)\n                temperature = temperature.unsqueeze(-1).unsqueeze(-1)\n                logits.div_(temperature.clamp(min=1e-8).to(logits.dtype))\n\n                if calculate_entropy:\n                    if not self.engine_config.entropy_checkpointing:\n                        entropy = verl_F.entropy_from_logits(logits)\n                    else:\n                        entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits)\n\n                if pad_mode == DatasetPadMode.NO_PADDING:\n                    cu_seqlens = input_ids.offsets()\n                    seq_lengths = cu_seqlens.diff()\n                    starts = torch.zeros_like(seq_lengths, dtype=torch.int64)\n                    logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged)\n                    logits_rmpad = torch.cat([t for t in logits.unbind()])\n                    input_ids_rmpad_rolled = output_args[\"input_ids_rmpad_rolled\"]\n                    log_probs = logprobs_from_logits(logits=logits_rmpad, labels=input_ids_rmpad_rolled)\n                    # (bsz, j1), for each sample, length of each sample: [real_prompt_length + real_response_length]\n                    log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens)\n                    if calculate_entropy:\n                        entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged)\n                        entropy_rmpad = torch.cat([t for t in entropy.unbind()])\n                        entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n                else:\n                    raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        model_output[\"log_probs\"] = log_probs\n        if calculate_entropy:\n            model_output[\"entropy\"] = entropy\n\n        return model_output\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        device_name = get_device_name()\n        # actually, we should avoid assigning like this...\n        micro_batch = micro_batch.to(get_device_id())\n        model_inputs, output_args = self.prepare_model_inputs(micro_batch=micro_batch)\n\n        with torch.autocast(device_type=device_name, dtype=torch.bfloat16):\n            raw_output = self.module(\n                **model_inputs,\n                use_cache=False,\n            )  # prevent model thinks we are generating\n\n            model_output = self.prepare_model_outputs(\n                output=raw_output, output_args=output_args, micro_batch=micro_batch\n            )\n\n            if loss_function is not None:\n                loss, metrics = loss_function(\n                    model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group()\n                )\n            else:\n                assert forward_only, \"forward_only must be True when loss_function is None\"\n                loss = torch.tensor(1.0, device=device_name)\n                metrics = {}\n\n            output = {\n                \"model_output\": model_output,\n                \"loss\": loss.detach().item(),\n                \"metrics\": metrics,\n            }\n\n            return loss, output\n\n\n@EngineRegistry.register(model_type=\"value_model\", backend=[\"fsdp\", \"fsdp2\"], device=[\"cuda\", \"npu\"])\nclass FSDPEngineWithValueHead(FSDPEngineWithLMHead):\n    \"\"\"\n    The only difference between critic and actor is how the raw model output is processed\n    \"\"\"\n\n    def prepare_model_outputs(self, output, output_args, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n\n        input_ids = micro_batch[\"input_ids\"]\n        if use_remove_padding:\n            if hasattr(self.module, \"v_head\"):\n                # For trl.AutoModelForCausalLMWithValueHead\n                values_rmpad = output[2].squeeze(0).unsqueeze(-1)\n            else:\n                values_rmpad = output.logits\n                values_rmpad = values_rmpad.squeeze(0)  # (total_nnz, 1)\n                # critic model arch is like Qwen3ForTokenClassfication and num_labels=1\n                # so we squeeze the last dimension here to get the value for each token\n                values_rmpad = values_rmpad.squeeze(-1)\n\n            # gather output if sp > 1\n            if self.use_ulysses_sp:\n                pad_size = output_args[\"pad_size\"]\n                values_rmpad = gather_outputs_and_unpad(values_rmpad, gather_dim=0, unpad_dim=0, padding_size=pad_size)\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                cu_seqlens = input_ids.offsets()\n                # (bsz, j1), for each sample, is the length of each sample: [real_prompt length + real_response length]\n                values = torch.nested.nested_tensor_from_jagged(values_rmpad, cu_seqlens)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        else:\n            if hasattr(self.module, \"v_head\"):\n                # For trl.AutoModelForCausalLMWithValueHead\n                values = output[2]\n            else:\n                values = output.logits\n\n            if pad_mode == DatasetPadMode.NO_PADDING:\n                cu_seqlens = input_ids.offsets()\n                seq_lengths = cu_seqlens.diff()\n                starts = torch.zeros_like(seq_lengths, dtype=torch.int64)\n                values = torch.nested.narrow(values, 1, starts, seq_lengths, layout=torch.jagged)\n                values_rmpad = torch.cat([t for t in values.unbind()])\n                # (bsz, j1), for each sample, length of each sample: [real_prompt_length + real_response_length]\n                values = torch.nested.nested_tensor_from_jagged(values_rmpad, cu_seqlens)\n            else:\n                raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        return {\"values\": values}\n"
  },
  {
    "path": "verl/workers/engine/fsdp/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\n\nimport torch\nfrom torch.distributed.device_mesh import init_device_mesh\n\nfrom verl.utils.device import get_device_name, is_npu_available\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef apply_npu_fsdp_patches():\n    \"\"\"Apply NPU patches for FSDP backend if NPU is available.\"\"\"\n    if is_npu_available:\n        try:\n            import verl.models.transformers.npu_patch  # noqa\n\n            if torch.distributed.is_initialized() and torch.distributed.get_rank() == 0:\n                logger.info(\"Applied NPU patches for FSDP backend\")\n        except Exception as e:\n            logger.warning(f\"Failed to apply NPU patches: {e}\")\n\n\ndef create_device_mesh(world_size, fsdp_size):\n    \"\"\"\n    Create a device mesh for distributed training based on the world size and FSDP size.\n\n    Args:\n        world_size (int): Total number of processes in the distributed training setup.\n        fsdp_size (int): Size of the Fully Sharded Data Parallel (FSDP) group.\n\n    Returns:\n        torch.distributed.device_mesh.DeviceMesh: The initialized device mesh.\n    \"\"\"\n    device_name = get_device_name()\n    if fsdp_size < 0 or fsdp_size >= world_size:\n        device_mesh = init_device_mesh(device_name, mesh_shape=(world_size,), mesh_dim_names=[\"fsdp\"])\n    else:\n        device_mesh = init_device_mesh(\n            device_name, mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=[\"ddp\", \"fsdp\"]\n        )\n    return device_mesh\n\n\ndef get_sharding_strategy(device_mesh):\n    \"\"\"\n    Determine the appropriate sharding strategy based on the number of dimensions of the device mesh.\n\n    Args:\n        device_mesh (torch.distributed.device_mesh.DeviceMesh): The device mesh used for distributed training.\n\n    Returns:\n        torch.distributed.fsdp.ShardingStrategy: The sharding strategy to be used with FSDP.\n\n    Raises:\n        NotImplementedError: If the number of dimensions of the device mesh is neither 1 nor 2.\n    \"\"\"\n    from torch.distributed.fsdp import ShardingStrategy\n\n    if device_mesh.ndim == 1:\n        sharding_strategy = ShardingStrategy.FULL_SHARD\n    elif device_mesh.ndim == 2:\n        sharding_strategy = ShardingStrategy.HYBRID_SHARD\n    else:\n        raise NotImplementedError(f\"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2\")\n    return sharding_strategy\n"
  },
  {
    "path": "verl/workers/engine/megatron/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\n# HACK Avoid cpu worker trigger cuda jit error\nimport os\n\nfrom verl.utils.device import is_cuda_available\n\nif not is_cuda_available and \"TORCH_CUDA_ARCH_LIST\" not in os.environ:\n    os.environ[\"TORCH_CUDA_ARCH_LIST\"] = \"8.0\"\n\nfrom .transformer_impl import MegatronEngine, MegatronEngineWithLMHead  # noqa: E402\n\nif not is_cuda_available:\n    del os.environ[\"TORCH_CUDA_ARCH_LIST\"]\n\n__all__ = [\"MegatronEngine\", \"MegatronEngineWithLMHead\"]\n"
  },
  {
    "path": "verl/workers/engine/megatron/transformer_impl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport logging\nimport os\nfrom functools import partial\nfrom typing import Any, Callable, ContextManager, Iterator, Optional\n\nimport torch\nimport torch.distributed\nfrom megatron.core import parallel_state as mpu\nfrom megatron.core.pipeline_parallel import get_forward_backward_func\nfrom omegaconf import OmegaConf\nfrom tensordict import TensorDict\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.models.mcore import get_mcore_forward_fused_no_padding_fn, get_mcore_weight_converter\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.checkpoint.megatron_checkpoint_manager import MegatronCheckpointManager\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.debug import log_gpu_memory_usage\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.megatron.pipeline_parallel import make_batch_generator\nfrom verl.utils.megatron.router_replay_patch import RouterReplay, RouterReplayAction, apply_router_replay_patch\nfrom verl.utils.megatron.router_replay_utils import (\n    RouterReplayHelper,\n    merge_router_topk_indices,\n    pp_gather,\n    reorder_and_merge_vpp_layers,\n    set_router_replay_data,\n)\nfrom verl.utils.megatron.tensor_parallel import vocab_parallel_entropy, vocab_parallel_log_probs_from_logits\nfrom verl.utils.megatron_peft_utils import add_base_layer_suffix, build_peft_config_for_vllm\nfrom verl.utils.megatron_utils import (\n    check_mtp_config,\n    get_megatron_module_device,\n    get_megatron_mtp_loss,\n    load_megatron_model_to_gpu,\n    load_megatron_optimizer,\n    offload_megatron_model_to_cpu,\n    offload_megatron_optimizer,\n    patch_engine_mtp,\n    register_megatron_training_hooks,\n    unwrap_model,\n)\nfrom verl.utils.model import extract_multi_modal_inputs, load_mcore_dist_weights\nfrom verl.utils.seqlen_balancing import restore_dynamic_batch\nfrom verl.workers.config import HFModelConfig, McoreEngineConfig, McoreOptimizerConfig\n\nfrom ..base import BaseEngine, BaseEngineCtx, EngineRegistry\nfrom ..utils import postprocess_batch_func, prepare_micro_batches\nfrom .utils import set_random_seed\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\nclass MegatronEngine(BaseEngine):\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: McoreEngineConfig,\n        optimizer_config: McoreOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n    ):\n        super().__init__()\n\n        self.model_config = model_config\n        self.engine_config = engine_config\n        self.optimizer_config = optimizer_config\n        self.checkpoint_config = checkpoint_config\n        assert self.engine_config.use_mbridge, \"use_mbridge must be True\"\n        self._init_device_mesh()\n\n        set_random_seed(seed=self.engine_config.seed)\n\n        self._is_offload_param = self.engine_config.param_offload\n        self._is_offload_grad = self.engine_config.grad_offload\n        self._is_offload_optimizer = self.engine_config.optimizer_offload\n\n        self.mode = None\n\n        self.layer_name_mapping = {\n            \"qkv_layer_name\": \"self_attention.linear_qkv.\",\n            \"gate_proj_layer_name\": \"linear_fc1.\",\n        }\n        self.weight_converter = None\n\n        # Router replay configuration for MoE models\n        self.enable_routing_replay = self.engine_config.router_replay.mode != \"disabled\"\n        logger.info(f\"enable_routing_replay in MegatronEngine: {self.enable_routing_replay}\")\n        if self.enable_routing_replay:\n            apply_router_replay_patch()\n            self.mini_layer_topk_idx_list = []\n\n    def _init_device_mesh(self):\n        # TODO: set different parallelism for actor, critic, ref\n        if mpu.is_initialized():\n            return\n\n        mpu.initialize_model_parallel(\n            tensor_model_parallel_size=self.engine_config.tensor_model_parallel_size,\n            pipeline_model_parallel_size=self.engine_config.pipeline_model_parallel_size,\n            virtual_pipeline_model_parallel_size=self.engine_config.virtual_pipeline_model_parallel_size,\n            use_sharp=False,\n            context_parallel_size=self.engine_config.context_parallel_size,\n            expert_model_parallel_size=self.engine_config.expert_model_parallel_size,\n            expert_tensor_parallel_size=self.engine_config.expert_tensor_parallel_size,\n            nccl_communicator_config_path=None,\n        )\n\n    def _build_tf_config(self):\n        from verl.utils.megatron_utils import mapping_string_to_attn_backend\n        from verl.utils.torch_dtypes import PrecisionType\n\n        check_mtp_config(self.model_config, self.engine_config)\n\n        self.param_dtype = PrecisionType.to_dtype(self.engine_config.dtype)\n        self.dtype = PrecisionType.to_dtype(self.param_dtype)\n\n        override_transformer_config = mapping_string_to_attn_backend({**self.engine_config.override_transformer_config})\n        if self.enable_routing_replay:\n            override_transformer_config[\"enable_routing_replay\"] = True\n\n        self.provider = None\n        self.vanilla_bridge = self.engine_config.vanilla_mbridge\n\n        if self.vanilla_bridge:\n            from verl.models.mcore.mbridge import AutoBridge\n\n            bridge = AutoBridge.from_config(self.model_config.hf_config, dtype=self.param_dtype)\n            bridge.set_extra_args(**override_transformer_config)\n            tf_config = bridge.config\n            tf_config.fp16 = self.param_dtype == torch.float16\n            tf_config.bf16 = self.param_dtype == torch.bfloat16\n        else:\n            from verl.models.mcore.bridge import AutoBridge\n\n            # Use Megatron-Bridge to convert HF config to Megatron config\n            bridge = AutoBridge.from_hf_pretrained(\n                self.model_config.local_path, trust_remote_code=self.model_config.trust_remote_code\n            )\n            # Get Megatron provider and configure it\n            provider = bridge.to_megatron_provider(load_weights=False)\n\n            # In case of invalid overrides, we need to make sure some critical params are set correctly\n            provider.params_dtype = self.param_dtype\n\n            # Ensure dtype settings propagate to Megatron-Bridge/TE\n            provider.fp16 = self.param_dtype == torch.float16\n            provider.bf16 = self.param_dtype == torch.bfloat16\n\n            # Pass distributed info\n            provider.tensor_model_parallel_size = self.engine_config.tensor_model_parallel_size\n            provider.pipeline_model_parallel_size = self.engine_config.pipeline_model_parallel_size\n            provider.expert_model_parallel_size = self.engine_config.expert_model_parallel_size\n            provider.expert_tensor_parallel_size = self.engine_config.expert_tensor_parallel_size\n            provider.virtual_pipeline_model_parallel_size = self.engine_config.virtual_pipeline_model_parallel_size\n            provider.context_parallel_size = self.engine_config.context_parallel_size\n            provider.sequence_parallel = self.engine_config.sequence_parallel\n\n            # Match verl implementation (need variable_seq_lengths)\n            from megatron.core.transformer.enums import AttnBackend\n\n            provider.attention_backend = AttnBackend.flash\n            provider.variable_seq_lengths = True\n            provider.moe_token_dispatcher_type = \"alltoall\"\n            provider.moe_router_load_balancing_type = \"none\"\n\n            # Apply transformer config overrides\n            for key, value in override_transformer_config.items():\n                setattr(provider, key, value)\n\n            provider.finalize()\n            self.provider = provider\n            tf_config = None  # Will be set after model creation\n        self.bridge = bridge\n\n        if not self.bridge:\n            self.weight_converter = get_mcore_weight_converter(self.model_config.hf_config, self.dtype)\n\n        if torch.distributed.get_rank() == 0:\n            if tf_config is not None:\n                print(f\"TF config: {tf_config}\")\n        self.tf_config = tf_config\n\n        from verl.workers.config.megatron_peft import get_peft_cls\n\n        self.peft_cls = get_peft_cls(\n            model_config=self.model_config, bridge=self.bridge, provider=self.provider, dtype=self.param_dtype\n        )\n\n    def _build_megatron_module(self):\n        from verl.utils.megatron_utils import McoreModuleWrapperConfig, make_megatron_module\n        from verl.utils.model import print_model_size\n\n        # TODO: add more cases\n        is_value_model = (\n            \"ForTokenClassification\" in self.model_config.architectures[0]\n            or \"ForSequenceClassification\" in self.model_config.architectures[0]\n        )\n\n        self.is_value_model = is_value_model\n\n        if self.engine_config.forward_only:\n            wrap_with_ddp = False\n        else:\n            wrap_with_ddp = True\n\n        wrap_config = McoreModuleWrapperConfig(\n            is_value_model=is_value_model,  # actor is not value model\n            share_embeddings_and_output_weights=self.model_config.share_embeddings_and_output_weights,\n            wrap_with_ddp=wrap_with_ddp,\n            use_distributed_optimizer=self.engine_config.use_distributed_optimizer,\n        )\n        module, updated_tf_config = make_megatron_module(\n            wrap_config=wrap_config,\n            tf_config=self.tf_config,\n            hf_config=self.model_config.hf_config,\n            bridge=self.bridge,\n            provider=self.provider,\n            override_model_config=self.engine_config.override_mcore_model_config,\n            override_ddp_config=self.engine_config.override_ddp_config,\n            peft_cls=self.peft_cls,\n            peft_config=self.model_config.get(\"lora\", None),\n        )\n        self.tf_config = updated_tf_config\n        print(f\"module: {len(module)}\")\n\n        if self.engine_config.use_dist_checkpointing:\n            load_mcore_dist_weights(module, self.engine_config.dist_checkpointing_path, is_value_model=is_value_model)\n        else:\n            if self.vanilla_bridge:\n                self.bridge.load_weights(module, self.model_config.local_path)\n            else:\n                allowed_mismatched_params = []\n                if self.is_value_model:\n                    allowed_mismatched_params = [\"output_layer.weight\"]\n                self.bridge.load_hf_weights(\n                    module, self.model_config.local_path, allowed_mismatched_params=allowed_mismatched_params\n                )\n\n        if torch.distributed.get_rank() == 0:\n            print_model_size(module[0])\n\n        if self.enable_routing_replay:\n            print(f\"routing replay layers: {len(RouterReplay.router_instances)}\")\n\n        return module\n\n    def _maybe_enable_fused_kernels(self):\n        if not self.engine_config.use_fused_kernels:\n            return\n\n        if self.is_value_model or self.model_config.mtp.enable:\n            logger.warning_once(\n                \"Fused kernels are not supported for value models or when MTP is enabled in Megatron engine; disabling.\"\n            )\n            self.engine_config.use_fused_kernels = False\n            return\n\n        from verl.models.mcore.model_forward_fused import patch_fused_forward\n\n        for model in self.module:\n            patch_fused_forward(model)\n\n    def _build_optimizer(self):\n        from verl.utils.megatron.optimizer import get_megatron_optimizer, init_megatron_optim_config\n\n        optim_config_megatron = init_megatron_optim_config(\n            self.optimizer_config,\n            use_distributed_optimizer=self.engine_config.use_distributed_optimizer,\n            fp16=self.param_dtype == torch.float16,\n        )\n        optimizer = get_megatron_optimizer(model=self.module, config=optim_config_megatron)\n        register_megatron_training_hooks(self.module, optimizer)\n        return optimizer\n\n    def _build_lr_scheduler(self):\n        from verl.utils.megatron.optimizer import get_megatron_optimizer_param_scheduler\n\n        optimizer_scheduler = get_megatron_optimizer_param_scheduler(\n            optimizer=self.optimizer, config=self.optimizer_config\n        )\n        return optimizer_scheduler\n\n    @property\n    def is_param_offload_enabled(self) -> bool:\n        return self._is_offload_param\n\n    @property\n    def is_optimizer_offload_enabled(self) -> bool:\n        return self._is_offload_optimizer\n\n    def is_mp_src_rank_with_outputs(self):\n        return (\n            mpu.get_tensor_model_parallel_rank() == 0\n            and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1\n            and mpu.get_context_parallel_rank() == 0\n        )\n\n    def initialize(self):\n        self._build_tf_config()\n\n        self.module = self._build_megatron_module()\n\n        self._maybe_enable_fused_kernels()\n\n        if self.model_config.mtp.enable:\n            patch_engine_mtp(self.module, self.model_config)\n\n        # For forward_only, we don't need optimizer, lr_scheduler, checkpoint_mananager\n        if self.engine_config.forward_only:\n            self.optimizer = None\n            self.lr_scheduler = None\n            self.to(device=\"cpu\", model=self._is_offload_param, optimizer=False, grad=False)\n            log_gpu_memory_usage(\"After offload model during init (forward_only)\", logger=logger)\n            return\n\n        self.optimizer = self._build_optimizer()\n        self.lr_scheduler = self._build_lr_scheduler()\n\n        full_reshardable = self.engine_config.dist_ckpt_optim_fully_reshardable\n        mem_eff = self.engine_config.distrib_optim_fully_reshardable_mem_efficient\n\n        tmp_config = OmegaConf.create(\n            {\n                \"model\": {\"path\": self.model_config.local_path},\n                \"megatron\": {\n                    \"dist_ckpt_optim_fully_reshardable\": full_reshardable,\n                    \"distrib_optim_fully_reshardable_mem_efficient\": mem_eff,\n                },\n            }\n        )\n\n        role = \"actor\" if not self.is_value_model else \"critic\"\n\n        self.checkpoint_mananager = MegatronCheckpointManager(\n            config=tmp_config,\n            checkpoint_config=self.checkpoint_config,\n            model_config=self.model_config.hf_config,\n            transformer_config=self.tf_config,\n            role=role,\n            model=self.module,\n            arch=self.model_config.architectures[0],\n            hf_config=self.model_config.hf_config,\n            param_dtype=self.param_dtype,\n            share_embeddings_and_output_weights=self.model_config.share_embeddings_and_output_weights,\n            processing_class=self.model_config.get_processor(),\n            optimizer=self.optimizer,\n            optimizer_scheduler=self.lr_scheduler,\n            use_distributed_optimizer=self.engine_config.use_distributed_optimizer,\n            use_checkpoint_opt_param_scheduler=self.optimizer_config.use_checkpoint_opt_param_scheduler,\n            bridge=self.bridge,\n            provider=self.provider,\n            peft_cls=self.peft_cls,\n            use_dist_checkpointing=self.engine_config.use_dist_checkpointing,\n        )\n\n        self.to(\n            device=\"cpu\",\n            model=self._is_offload_param,\n            optimizer=self._is_offload_optimizer,\n            grad=self._is_offload_param,\n        )\n\n        log_gpu_memory_usage(\"After offload model/optimizer/grad during init\", logger=logger)\n\n    def train_mode(self, **kwargs):\n        \"\"\"\n        Context manager entry for switching the engine and model into training mode.\n\n        Usage:\n            with engine.train_mode():\n                # runs in training mode\n        \"\"\"\n        return EngineTrainModeCtx(self, **kwargs)\n\n    def eval_mode(self, **kwargs):\n        \"\"\"\n        Context manager entry for switching the engine and model into evaluation mode.\n\n        Usage:\n            with engine.eval_mode():\n                # runs in evaluation mode\n        \"\"\"\n        return EngineEvalModeCtx(self, **kwargs)\n\n    def optimizer_zero_grad(self):\n        \"\"\"\n        Zero out gradients of all parameters before starting a new backward pass.\n        \"\"\"\n        self.optimizer.zero_grad()\n        # use use_contiguous_buffers_in_local_ddp and no overlap_dp_param_comm\n        for chunk in self.module:\n            # if use distributed optimizer, zero grad buffer will be handled by optimizer\n            chunk.zero_grad_buffer()\n\n    def optimizer_step(self):\n        \"\"\"\n        Perform an optimization step to update model parameters based on accumulated gradients.\n\n        Returns:\n            grad_norm (float): The norm of the gradients before clipping or update.\n        \"\"\"\n        update_successful, grad_norm, num_zeros_in_grad = self.optimizer.step()\n\n        if update_successful:\n            # allgather already execute in optimizer.step in new megatron\n            pass\n        else:\n            raise NotImplementedError(\"Megatron optimizer step failed. This should not happen\")\n\n        return grad_norm\n\n    def lr_scheduler_step(self):\n        \"\"\"\n        Advance the learning rate scheduler by one step.\n\n        Returns:\n            current_lr (float or list[float]): Updated learning rate(s).\n        \"\"\"\n        from verl.utils.megatron.optimizer import get_megatron_last_lr\n\n        self.lr_scheduler.step(1)\n        return get_megatron_last_lr(self.optimizer)\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"\n        Move model parameters, optimizer states, or both to the specified device.\n        Note that this function executes irrespective of offload config. It serves as manual control\n\n        Args:\n            device: Target device identifier.\n            model: If True, move the model.\n            optimizer: If True, move the optimizer states.\n        \"\"\"\n        super().to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n        device_name = get_device_name()\n\n        assert device in (device_name, \"cpu\")\n        if device == device_name:\n            if model:\n                load_megatron_model_to_gpu(self.module, load_grad=grad)\n            if optimizer and self.optimizer is not None:\n                load_megatron_optimizer(self.optimizer)\n        elif device == \"cpu\":\n            if model:\n                offload_megatron_model_to_cpu(self.module)\n            if optimizer and self.optimizer is not None:\n                offload_megatron_optimizer(self.optimizer)\n        else:\n            raise ValueError(f\"Invalid device type: {device}\")\n\n    def get_data_parallel_rank(self):\n        return mpu.get_data_parallel_rank()\n\n    def get_data_parallel_size(self):\n        return mpu.get_data_parallel_world_size()\n\n    def get_data_parallel_group(self):\n        return mpu.get_data_parallel_group()\n\n    def get_model_parallel_group(self):\n        return mpu.get_model_parallel_group()\n\n    def get_context_parallel_group(self):\n        return mpu.get_context_parallel_group()\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Save model, optimizer, and scheduler states to a checkpoint.\n\n        Args:\n            local_path: Local filesystem path to save checkpoint.\n            hdfs_path: Optional HDFS path to copy checkpoint.\n            global_step: Integer training step number for naming.\n            max_ckpt_to_keep: Maximum number of recent checkpoints to retain.\n        \"\"\"\n        origin_module_device = get_megatron_module_device(self.module)\n        if self._is_offload_param or origin_module_device == \"cpu\":\n            load_megatron_model_to_gpu(self.module, load_grad=True)\n        self.checkpoint_mananager.save_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.module)\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: bool = True, **kwargs\n    ) -> None:\n        \"\"\"\n        Load model, optimizer, and scheduler states from a checkpoint.\n\n        Args:\n            local_path: Local filesystem path of the checkpoint.\n            hdfs_path: Optional HDFS path where checkpoint is stored.\n            del_local_after_load: Whether to delete local copy after loading.\n        \"\"\"\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.module)\n        self.checkpoint_mananager.load_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.module)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.optimizer)\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:\n        tu.assign_non_tensor(data, sp_size=self.engine_config.context_parallel_size)\n\n        # compute num_tokens in global batch for loss normalization\n        batch_num_tokens = data[\"loss_mask\"].sum().to(get_device_id())\n        torch.distributed.all_reduce(\n            batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group()\n        )\n        tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())\n        tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())\n\n        vpp_size = mpu.get_virtual_pipeline_model_parallel_world_size()\n        if vpp_size is not None and vpp_size > 1:\n            num_batches_divided_by = self.tf_config.microbatch_group_size_per_vp_stage\n        else:\n            num_batches_divided_by = None\n\n        micro_batches, indices = prepare_micro_batches(\n            data=data,\n            dp_group=self.get_data_parallel_group(),\n            num_batches_divided_by=num_batches_divided_by,\n            same_micro_num_in_dp=True,\n            min_num_micro_batch=None,\n        )\n\n        if num_batches_divided_by is not None:\n            assert len(micro_batches) % num_batches_divided_by == 0, (\n                f\"micro_batches {micro_batches} must be divisible by num_batches_divided_by \"\n                f\"{num_batches_divided_by} for megatron backend\"\n            )\n\n        # compute input shapes for pp stages\n        n_micro_batch = len(micro_batches)\n\n        for micro_batch in micro_batches:\n            tu.assign_non_tensor(micro_batch, num_micro_batch=n_micro_batch)\n\n        forward_backward_func = get_forward_backward_func()\n\n        postprocess_micro_batch_func = partial(\n            self.postprocess_micro_batch_func,\n            forward_only=forward_only,\n            loss_function=loss_function,\n        )\n\n        tu.assign_non_tensor(data, num_micro_batch=n_micro_batch)\n\n        forward_step = partial(self.forward_step, postprocess_micro_batch_func=postprocess_micro_batch_func)\n\n        enable_routing_replay = tu.get_non_tensor_data(data, key=\"enable_routing_replay\", default=False)\n\n        if enable_routing_replay:\n            # Set to REPLAY mode: for R3 mode or actor update phase in R2 mode\n            RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n            if forward_only and self.engine_config.router_replay.mode == \"R2\":\n                # In R2 mode, forward_only calls (e.g., compute_log_probs) need to record routing information\n                RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)\n\n        # batch should be a list of batches inside micro-batches\n        batch_generator = make_batch_generator(micro_batches, vpp_size=len(self.module))\n\n        # TODO: we may use the new schedule instead\n        # for flash-attn: (seq_len, batch_size, hidden_size) = (mbs*seq_len, 1, hidden_size)\n        losses_reduced = forward_backward_func(\n            forward_step_func=forward_step,\n            data_iterator=batch_generator,\n            model=self.module,\n            num_microbatches=n_micro_batch,\n            seq_length=1,  # the communication shape is obtained via p2p comm\n            micro_batch_size=1,  # the communication shape is obtained via p2p comm\n            forward_only=forward_only,\n        )\n\n        if self.model_config.mtp.enable and self.is_mp_src_rank_with_outputs():\n            # add mtp_losses\n            metrics = get_megatron_mtp_loss(n_micro_batch)\n            if \"metrics\" not in losses_reduced[0]:\n                losses_reduced[0][\"metrics\"] = {}\n            losses_reduced[0][\"metrics\"].update(metrics)\n\n        if RouterReplayHelper.is_r2_record_action(self.tf_config):\n            if self.tf_config.virtual_pipeline_model_parallel_size is not None:\n                # config = self.actor_module[0].module.module.config\n                vp_size = len(self.module)\n                microbatch_group_size_per_vp_stage = self.tf_config.microbatch_group_size_per_vp_stage\n                bs = n_micro_batch\n                topk_idx_td = reorder_and_merge_vpp_layers(\n                    self.mini_layer_topk_idx_list, bs, vp_size, microbatch_group_size_per_vp_stage\n                )\n            else:\n                tensors = [tensor for nt in self.mini_layer_topk_idx_list for tensor in nt.unbind()]\n                topk_idx_td = torch.nested.as_nested_tensor(tensors, layout=torch.jagged)\n            self.mini_layer_topk_idx_list = []\n\n            layers_topk_idx = pp_gather(topk_idx_td.to(torch.uint8), self.tf_config)\n            use_dynamic_bsz = tu.get_non_tensor_data(data=data, key=\"use_dynamic_bsz\", default=True)\n            if use_dynamic_bsz and indices is not None:\n                layers_topk_idx = restore_dynamic_batch(layers_topk_idx, indices)\n\n        output = {}\n        if mpu.is_pipeline_last_stage(ignore_virtual=True):\n            output = postprocess_batch_func(output_lst=losses_reduced, indices=indices, data=data)\n            if RouterReplayHelper.is_r2_record_action(self.tf_config):\n                output[\"model_output\"][\"routed_experts\"] = layers_topk_idx\n        if enable_routing_replay:\n            RouterReplay.clear_global_indices()\n            RouterReplay.clear_global_router_replay_action()\n        return output\n\n    def get_per_tensor_param(self, base_sync_done=False, **kwargs):\n        peft_config = None\n        non_merge_lora_sync = self.peft_cls is not None and not self.model_config.lora.get(\"merge\", False)\n        adapter_only = base_sync_done and non_merge_lora_sync\n        # when lora adapter only, we only load adapter weights when base sync is done, otherwise load all weights\n        load_megatron_model_to_gpu(self.module, load_grad=False, load_frozen_params=not adapter_only)\n        if self.vanilla_bridge:\n            per_tensor_param = self.bridge.export_weights(self.module)\n        elif adapter_only:\n            # Only export adapter weights\n            peft_config = build_peft_config_for_vllm(self.model_config.lora)\n            per_tensor_param = self.bridge.export_adapter_weights(self.module)\n        else:\n            per_tensor_param = self.bridge.export_hf_weights(self.module)\n            if non_merge_lora_sync:\n                per_tensor_param = add_base_layer_suffix(\n                    per_tensor_param, model_type=self.model_config.hf_config.model_type\n                )\n        return per_tensor_param, peft_config\n\n    def disable_adapter(self) -> ContextManager:\n        return self.peft_cls.disable_adapter(self.module)\n\n    def forward_step(self, batch_iter, model, postprocess_micro_batch_func):\n        raise NotImplementedError(\"forward_step must be implemented in subclass\")\n\n    def postprocess_micro_batch_func(self, output, data: TensorDict, forward_only: bool, loss_function):\n        raise NotImplementedError(\"postprocess_micro_batch_func must be implemented in subclass\")\n\n\nclass EngineEvalModeCtx(BaseEngineCtx):\n    def __init__(self, engine: MegatronEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"eval\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, MegatronEngine)\n        super().__enter__()\n        # mcore module is a list of model chunk in each vpp stage\n        for module in self.engine.module:\n            module.eval()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, MegatronEngine)\n        super().__exit__(exc_type, exc_value, traceback)\n\n\nclass EngineTrainModeCtx(BaseEngineCtx):\n    def __init__(self, engine: MegatronEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"train\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, MegatronEngine)\n        super().__enter__()\n        # mcore module is a list of model chunk in each vpp stage\n        for module in self.engine.module:\n            module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, MegatronEngine)\n        self.engine.optimizer_zero_grad()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=\"megatron\")\nclass MegatronEngineWithLMHead(MegatronEngine):\n    def prepare_model_inputs(self, batch: TensorDict):\n        input_ids = batch[\"input_ids\"]\n        loss_mask = batch[\"loss_mask\"].to(bool)\n        multi_modal_inputs = extract_multi_modal_inputs(batch.get(\"multi_modal_inputs\", []))\n\n        routed_experts = batch.get(\"routed_experts\", None)\n\n        return {\n            \"input_ids\": input_ids,\n            \"loss_mask\": loss_mask,\n            \"multi_modal_inputs\": multi_modal_inputs,\n            \"routed_experts\": routed_experts,\n        }\n\n    def prepare_model_outputs(self, output: dict, data: TensorDict):\n        calculate_entropy = tu.get_non_tensor_data(data, key=\"calculate_entropy\", default=False)\n\n        log_prob = output[\"log_probs\"]\n        model_output = {\"log_probs\": log_prob}\n        if calculate_entropy:\n            entropy = output[\"entropy\"]\n            model_output[\"entropy\"] = entropy\n\n        return model_output\n\n    def forward_step(self, batch_iter: Iterator[TensorDict], model, postprocess_micro_batch_func):\n        batch: TensorDict = next(batch_iter)\n        batch = batch.to(get_device_id())\n        use_fused_kernels = tu.get_non_tensor_data(batch, key=\"use_fused_kernels\", default=False)\n        calculate_entropy = tu.get_non_tensor_data(batch, key=\"calculate_entropy\", default=False)\n        pad_mode = tu.get_non_tensor_data(batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        temperature = batch[\"temperature\"]\n        model_inputs = self.prepare_model_inputs(batch)\n        input_ids = model_inputs[\"input_ids\"]\n        multi_modal_inputs = model_inputs[\"multi_modal_inputs\"]\n        loss_mask = model_inputs[\"loss_mask\"]\n\n        unwrapped_model = unwrap_model(model)\n        if hasattr(unwrapped_model, \"vp_stage\"):\n            vp_rank = unwrapped_model.vp_stage\n        else:\n            vp_rank = 0\n\n        if RouterReplayHelper.is_replay_backward_action(self.tf_config, vp_rank):\n            router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank)\n            for router in router_instance_list:\n                router.set_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n\n        if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank):\n            layers_topk_idx = model_inputs[\"routed_experts\"]\n            set_router_replay_data(layers_topk_idx, None, self.tf_config, vp_rank)\n\n        if pad_mode == DatasetPadMode.NO_PADDING:\n            label = input_ids.clone()\n        else:\n            raise NotImplementedError(f\"Pad mode {pad_mode} is not supported for megatron engine\")\n\n        from verl.models.mcore import get_mcore_forward_no_padding_fn\n\n        if use_fused_kernels:\n            if not self.engine_config.use_remove_padding:\n                logger.warning_once(\n                    \"Fused kernels require `use_remove_padding=True` for Megatron engine. Falling back to non-fused.\"\n                )\n                use_fused_kernels = False\n            elif isinstance(temperature, torch.Tensor):\n                if temperature.numel() != 1:\n                    logger.warning_once(\n                        \"Fused kernels do not support per-sample temperature. Falling back to non-fused.\"\n                    )\n                    use_fused_kernels = False\n                else:\n                    temperature_value = float(temperature.item())\n            else:\n                temperature_value = float(temperature)\n\n        if use_fused_kernels:\n            fused_forward_fn = get_mcore_forward_fused_no_padding_fn(self.model_config.hf_config)\n            output = fused_forward_fn(\n                model=model,\n                input_ids=input_ids,\n                labels=label,\n                multi_modal_inputs=multi_modal_inputs,\n                temperature=temperature_value,\n                calculate_entropy=calculate_entropy,\n                pad_token_id=self.model_config.tokenizer.pad_token_id,\n            )\n        else:\n            if not isinstance(temperature, torch.Tensor):\n                temperature = torch.tensor([temperature] * input_ids.shape[0], device=input_ids.device)\n\n            temperature = temperature.to(torch.float32)\n            assert temperature.shape[0] == input_ids.shape[0]\n            temperature = verl_F.expand_as_nested(temperature, input_ids)  # (bsz, j1)\n\n            forward_fn = get_mcore_forward_no_padding_fn(self.model_config.hf_config)\n\n            def logits_processor(logits, label, temperature):\n                assert logits.shape[:2] == label.shape[:2]\n                # avoid non-positive temperature such as padding\n                temperature[temperature <= 0] = 1e-8\n                assert torch.all(temperature > 0).item(), f\"temperature tensor must be positive. Got {temperature}\"\n                logits.div_(temperature.unsqueeze(dim=-1).to(logits.dtype))\n                ret = {}\n                if calculate_entropy:\n                    logits_bak = logits.clone()\n                    # # disable the hint until the fused_kernel is optimized for triton>=3.3\n                    # if torch.distributed.get_rank() == 0:\n                    #     logger.warning_once(\n                    #         \"For memory-efficient computation, enable fused kernels via \"\n                    #         \"`actor_rollout_ref.model.use_fused_kernels=True`. \"\n                    #         \"The current `clone()` operation ensures correctness but increases memory usage.\"\n                    #     )\n                    entropy = vocab_parallel_entropy(logits)\n                    ret[\"entropy\"] = entropy\n                else:\n                    logits_bak = logits\n\n                log_probs = vocab_parallel_log_probs_from_logits(logits_bak, label)\n                ret[\"log_probs\"] = log_probs\n                return ret\n\n            logits_processor_args = {\"label\": label, \"temperature\": temperature, \"loss_mask\": loss_mask}\n\n            output = forward_fn(\n                model,\n                input_ids,\n                multi_modal_inputs,\n                logits_processor=logits_processor,\n                logits_processor_args=logits_processor_args,\n                vision_model=hasattr(self.model_config.hf_config, \"vision_config\"),\n                pad_token_id=self.model_config.tokenizer.pad_token_id,\n                data_format=\"thd\" if self.engine_config.use_remove_padding else \"bshd\",\n                mtp_enable_train=self.model_config.mtp.enable and self.model_config.mtp.enable_train,\n            )\n\n        # Router replay: record routing decisions for R2 mode\n        if RouterReplayHelper.is_r2_record_action(self.tf_config, vp_rank):\n            merge_router_topk_indices(None, input_ids, self.mini_layer_topk_idx_list, self.tf_config, vp_rank)\n\n        # Router replay: switch to backward replay mode for next backward pass\n        if RouterReplayHelper.is_replay_forward_action(self.tf_config, vp_rank):\n            router_instance_list = RouterReplayHelper.get_micro_batch_router_list(self.tf_config, vp_rank)\n            for router in router_instance_list:\n                router.set_router_replay_action(RouterReplayAction.REPLAY_BACKWARD)\n\n        return output, partial(postprocess_micro_batch_func, data=batch)\n\n    def postprocess_micro_batch_func(self, output, data: TensorDict, forward_only: bool, loss_function):\n        # For memory efficiency\n        # We move calculation of entropy to compute_log_probs, forward_only == True\n        device = data[\"input_ids\"].device\n        model_output = self.prepare_model_outputs(output, data)\n\n        if loss_function is not None:\n            loss, metrics = loss_function(model_output=model_output, data=data, dp_group=self.get_data_parallel_group())\n            # scale loss by num_micro_batch because megatron will scale loss\n            # by n_micro_batch inside pp schedule\n            scaled_loss = loss * data[\"num_micro_batch\"]\n        else:\n            assert forward_only, \"forward_only must be True when loss_function is None\"\n            loss = torch.tensor(1.0, device=device)\n            scaled_loss = loss\n            metrics = {}\n\n        output = {\n            \"model_output\": model_output,\n            \"loss\": loss.detach().item(),\n            \"metrics\": metrics,\n        }\n\n        # return loss and stats\n        return scaled_loss, output\n\n\n@EngineRegistry.register(model_type=\"value_model\", backend=\"megatron\")\nclass MegatronEngineWithValueHead(MegatronEngineWithLMHead):\n    # for value head\n    def forward_step(self, batch_iter, model, postprocess_micro_batch_func):\n        batch: TensorDict = next(batch_iter)\n        batch = batch.to(get_device_id())\n        model_inputs = self.prepare_model_inputs(batch)\n        input_ids = model_inputs[\"input_ids\"]\n        multi_modal_inputs = model_inputs[\"multi_modal_inputs\"]\n\n        from verl.models.mcore import get_mcore_forward_no_padding_fn\n\n        forward_fn = get_mcore_forward_no_padding_fn(self.model_config.hf_config)\n\n        output = forward_fn(\n            model,\n            input_ids,\n            multi_modal_inputs,\n            value_model=True,\n            vision_model=hasattr(self.model_config.hf_config, \"vision_config\"),\n            pad_token_id=self.model_config.tokenizer.pad_token_id,\n            enable_mtp=self.model_config.mtp.enable_train,\n        )\n\n        return output, partial(postprocess_micro_batch_func, data=batch)\n\n    def prepare_model_outputs(self, output: dict | torch.Tensor, data: TensorDict):\n        return {\"values\": output}\n"
  },
  {
    "path": "verl/workers/engine/megatron/utils.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nfrom verl.utils.device import get_torch_device\n\n\ndef set_random_seed(seed):\n    import random\n\n    import numpy as np\n    import torch\n\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    if get_torch_device().device_count() > 0:\n        from megatron.core import tensor_parallel\n\n        tensor_parallel.model_parallel_cuda_manual_seed(seed)\n    # FIXME: torch cumsum not support deterministic (used in vllm sampler),\n    # https://github.com/pytorch/pytorch/issues/89492\n    # torch.use_deterministic_algorithms(True, warn_only=True)\n    # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'\n"
  },
  {
    "path": "verl/workers/engine/mindspeed/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .transformer_impl import MindspeedEngineWithLMHead\n\n__all__ = [\"MindspeedEngineWithLMHead\"]\n"
  },
  {
    "path": "verl/workers/engine/mindspeed/transformer_impl.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport logging\nimport os\n\ntry:\n    from mindspeed.megatron_adaptor import repatch\nexcept ImportError:\n    repatch = None\n\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.workers.config import HFModelConfig, McoreEngineConfig, McoreOptimizerConfig\n\nfrom ..base import EngineRegistry\nfrom ..megatron import MegatronEngineWithLMHead\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=\"megatron\", device=\"npu\")\nclass MindspeedEngineWithLMHead(MegatronEngineWithLMHead):\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: McoreEngineConfig,\n        optimizer_config: McoreOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n    ):\n        super().__init__(model_config, engine_config, optimizer_config, checkpoint_config)\n\n        repatch_config = self.engine_config.get(\"override_transformer_config\", {})\n        repatch_config[\"use_flash_attn\"] = True\n        if self.engine_config.context_parallel_size > 1:\n            repatch_config[\"context_parallel_size\"] = self.engine_config.context_parallel_size\n\n        repatch(repatch_config)\n"
  },
  {
    "path": "verl/workers/engine/torchtitan/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom .transformer_impl import TorchTitanEngine, TorchTitanEngineWithLMHead\n\n__all__ = [\"TorchTitanEngine\", \"TorchTitanEngineWithLMHead\"]\n"
  },
  {
    "path": "verl/workers/engine/torchtitan/transformer_impl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe concrete Engine implementation using PyTorch TorchTitan parallelism (FSDP2 + TP + PP)\n\"\"\"\n\nimport gc\nimport importlib\nimport logging\nimport os\nimport re\nfrom contextlib import nullcontext\nfrom typing import Any, Callable, Optional\n\nimport torch\nimport torch.distributed\nfrom tensordict import TensorDict\nfrom torch.distributed.checkpoint.state_dict import get_model_state_dict\nfrom torch.distributed.tensor import DTensor\nfrom torchtitan.components.checkpoint import CheckpointManager\nfrom torchtitan.components.lr_scheduler import LRSchedulersContainer\nfrom torchtitan.components.optimizer import OptimizersContainer\nfrom torchtitan.config import CompileConfig, ParallelismConfig, TrainingConfig\nfrom torchtitan.distributed import utils as dist_utils\nfrom torchtitan.distributed.context_parallel import prepare_context_parallel_input\nfrom torchtitan.distributed.parallel_dims import ParallelDims\nfrom torchtitan.train import Trainer\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.debug import log_gpu_memory_usage\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.fsdp_utils import (\n    load_fsdp_model_to_gpu,\n    load_fsdp_optimizer,\n    offload_fsdp_model_to_cpu,\n    offload_fsdp_optimizer,\n)\nfrom verl.utils.model import extract_multi_modal_inputs\nfrom verl.utils.torch_functional import logprobs_from_logits\nfrom verl.workers.config import HFModelConfig, TorchtitanEngineConfig, TorchtitanOptimizerConfig\nfrom verl.workers.engine.torchtitan.utils import (\n    NoOpDataLoader,\n    derive_torchtitan_name_and_flavor,\n    enable_fsdp_gradient_division,\n    get_attention_masks,\n    iter_per_tensor_params_ep,\n)\n\nfrom ..base import BaseEngine, BaseEngineCtx, EngineRegistry\nfrom ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\ndevice_name = get_device_name()\n\n\nclass TorchTitanEngine(BaseEngine):\n    \"\"\"\n    Concrete Engine implementation using PyTorch TorchTitan parallelism.\n\n    Supports model sharding with FSDP2, tensor parallelism, activation/optimizer offloading,\n    LoRA, and sequence parallelism following the TorchTitan design.\n    \"\"\"\n\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: TorchtitanEngineConfig,\n        optimizer_config: TorchtitanOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n    ):\n        \"\"\"\n        Initialize the TorchTitanEngine.\n\n        Sets up distributed device meshes for tensor and data parallelism, LoRA, and offload policies.\n\n        Args:\n            model_config: Configuration for HuggingFace model.\n            engine_config: Configuration for FSDP/TorchTitan engine (uses FSDP2).\n            optimizer_config: Configuration for optimizer.\n            checkpoint_config: Configuration for checkpointing.\n        \"\"\"\n        super().__init__()\n\n        self.model_config = model_config\n        self.engine_config = engine_config\n        self.optimizer_config = optimizer_config\n        self.checkpoint_config = checkpoint_config\n\n        # Derive torchtitan model name and flavor from HF config\n        torchtitan_name, torchtitan_flavor = derive_torchtitan_name_and_flavor(self.model_config.hf_config)\n\n        # Get ModelSpec from model registry\n        model_module = importlib.import_module(f\"torchtitan.models.{torchtitan_name}\")\n        model_spec = model_module.model_registry(torchtitan_flavor)\n\n        # Override attn_backend on the model config if needed\n        attn_type = self.engine_config.attn_type\n        if hasattr(model_spec.model, \"layer\") and hasattr(model_spec.model.layer, \"attention\"):\n            model_spec.model.layer.attention.attn_backend = attn_type\n\n        optimizer = OptimizersContainer.Config(\n            name=self.optimizer_config.name,\n            lr=self.optimizer_config.lr,\n            eps=self.optimizer_config.eps,\n            beta1=self.optimizer_config.betas[0],\n            beta2=self.optimizer_config.betas[1],\n            weight_decay=self.optimizer_config.weight_decay,\n        )\n\n        total_steps = self.optimizer_config.total_training_steps\n        lr_warmup_steps = self.optimizer_config.lr_warmup_steps\n        if lr_warmup_steps is None or lr_warmup_steps <= 0:\n            lr_warmup_steps = int(self.optimizer_config.lr_warmup_steps_ratio * total_steps)\n\n        lr_scheduler = LRSchedulersContainer.Config(\n            warmup_steps=lr_warmup_steps,\n            decay_type=self.optimizer_config.decay_type,\n            min_lr_factor=self.optimizer_config.min_lr_factor,\n        )\n        parallelism = ParallelismConfig(\n            data_parallel_replicate_degree=self.engine_config.data_parallel_replicate_size,\n            data_parallel_shard_degree=self.engine_config.data_parallel_shard_size,\n            fsdp_reshard_after_forward=self.engine_config.reshard_after_forward,\n            tensor_parallel_degree=self.engine_config.tensor_parallel_size,\n            pipeline_parallel_degree=self.engine_config.pipeline_parallel_size,\n            context_parallel_degree=self.engine_config.context_parallel_size,\n            expert_parallel_degree=self.engine_config.expert_parallel_size,\n            expert_tensor_parallel_degree=self.engine_config.expert_tensor_parallel_size,\n        )\n        checkpoint = CheckpointManager.Config(\n            enable=True,\n            initial_load_in_hf=True,\n            initial_load_model_only=True,\n            initial_load_path=model_config.path,\n        )\n        compile_config = CompileConfig(enable=self.engine_config.use_torch_compile)\n        training_kwargs = {}\n        if self.engine_config.max_seq_len is not None:\n            training_kwargs[\"seq_len\"] = self.engine_config.max_seq_len\n        if self.engine_config.offload_policy or self.engine_config.forward_only:\n            training = TrainingConfig(enable_cpu_offload=True, **training_kwargs)\n        else:\n            training = TrainingConfig(**training_kwargs)\n\n        # Construct Torchtitan's Trainer.Config\n        self.config = Trainer.Config(\n            model_spec=model_spec,\n            hf_assets_path=self.model_config.path,\n            optimizer=optimizer,\n            lr_scheduler=lr_scheduler,\n            parallelism=parallelism,\n            checkpoint=checkpoint,\n            compile=compile_config,\n            training=training,\n            # Use a no-op dataloader since verl has its own data loading\n            dataloader=NoOpDataLoader.Config(),\n        )\n        self.trainer = Trainer(self.config)\n\n        self._init_device_mesh()\n\n        # Re-enable FSDP's gradient division for verl's loss scaling.\n        # TorchTitan disables gradient division by default (for global token normalization),\n        # but verl's loss function multiplies by dp_size to compensate for gradient averaging.\n        if self.engine_config.data_parallel_shard_size > 1:\n            dp_size = self.get_data_parallel_size()\n            for model_part in self.trainer.model_parts:\n                enable_fsdp_gradient_division(model_part, dp_size)\n\n        if self.engine_config.full_determinism:\n            enable_full_determinism(seed=self.engine_config.seed)\n\n        # set FSDP offload params\n        self._is_offload_param = self.engine_config.param_offload\n        self._is_offload_optimizer = self.engine_config.optimizer_offload\n\n        if self.engine_config.entropy_from_logits_with_chunking:\n            entropy_from_logits = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_from_logits = verl_F.entropy_from_logits\n\n        self.compute_entropy_from_logits = (\n            torch.compile(entropy_from_logits, dynamic=True)\n            if self.engine_config.use_torch_compile\n            else entropy_from_logits\n        )\n\n    @property\n    def is_param_offload_enabled(self) -> bool:\n        return self._is_offload_param\n\n    @property\n    def is_optimizer_offload_enabled(self) -> bool:\n        return self._is_offload_optimizer\n\n    def is_mp_src_rank_with_outputs(self):\n        \"\"\"\n        Whether the current rank is the first rank in model parallel group that contains model outputs\n        \"\"\"\n        is_collect = True\n        # TP: outputs are on TP rank 0\n        if self.parallel_dims.tp > 1:\n            tp_mesh = self.parallel_dims.get_optional_mesh(\"tp\")\n            is_collect = is_collect and (tp_mesh.get_local_rank() == 0)\n        # PP: outputs are on the last PP rank\n        if self.parallel_dims.pp > 1:\n            pp_mesh = self.parallel_dims.get_optional_mesh(\"pp\")\n            is_collect = is_collect and (pp_mesh.get_local_rank() == self.parallel_dims.pp - 1)\n        # CP: outputs are on CP rank 0\n        if self.parallel_dims.cp > 1:\n            cp_mesh = self.parallel_dims.get_optional_mesh(\"cp\")\n            is_collect = is_collect and (cp_mesh.get_local_rank() == 0)\n        return is_collect\n\n    def initialize(self):\n        \"\"\"\n        Build the model, optimizer, and learning rate scheduler with TorchTitan parallelism.\n\n        Applies device, dtype, and precision configurations, including mixed precision.\n        Sets up checkpoint manager.\n        \"\"\"\n        self.module = self.trainer.model_parts\n        self.checkpointer = self.trainer.checkpointer\n        # load initial HF weights\n        self.checkpointer.load()\n\n        if not self.engine_config.forward_only:\n            self.optimizer = self.trainer.optimizers\n            self.lr_scheduler = self.trainer.lr_schedulers\n        else:\n            self.optimizer = None\n            self.lr_scheduler = None\n\n        self.to(\n            device=\"cpu\",\n            model=self._is_offload_param,\n            optimizer=self._is_offload_optimizer,\n            grad=self._is_offload_param,\n        )\n\n        log_gpu_memory_usage(\"After offload model/optimizer/grad during init\", logger=logger)\n\n    def _init_device_mesh(self):\n        \"\"\"Initialize the device mesh for TorchTitan style parallelism.\"\"\"\n        world_size = torch.distributed.get_world_size()\n        self.parallel_dims = ParallelDims(\n            dp_shard=self.engine_config.data_parallel_shard_size,\n            dp_replicate=self.engine_config.data_parallel_replicate_size,\n            cp=self.engine_config.context_parallel_size,\n            tp=self.engine_config.tensor_parallel_size,\n            pp=self.engine_config.pipeline_parallel_size,\n            ep=self.engine_config.expert_parallel_size,\n            etp=self.engine_config.expert_tensor_parallel_size,\n            world_size=world_size,\n        )\n        self.device_mesh = self.parallel_dims.build_mesh()\n\n    def train_mode(self, **kwargs):\n        \"\"\"Return a context manager for training mode.\"\"\"\n        return EngineTrainModeCtx(self, **kwargs)\n\n    def eval_mode(self, **kwargs):\n        \"\"\"Return a context manager for evaluation mode.\"\"\"\n        return EngineEvalModeCtx(self, **kwargs)\n\n    def get_data_parallel_rank(self):\n        mesh = self._get_data_parallel_mesh()\n        return 0 if mesh is None else mesh.get_local_rank()\n\n    def get_data_parallel_size(self):\n        return self.engine_config.data_parallel_shard_size * self.engine_config.data_parallel_replicate_size\n\n    def get_data_parallel_group(self):\n        mesh = self._get_data_parallel_mesh()\n        if mesh is not None:\n            return mesh.get_group()\n        # If world_size == dp_size (e.g. single GPU, or all ranks are DP),\n        # return WORLD so that collective ops in _postprocess_output\n        # (allgather_dict_into_dict, all_reduce) still run and produce the\n        # correct metric aggregation format.\n        if torch.distributed.get_world_size() == self.get_data_parallel_size():\n            return torch.distributed.group.WORLD\n        return None\n\n    def get_model_parallel_group(self):\n        raise NotImplementedError\n\n    def get_context_parallel_group(self):\n        raise NotImplementedError\n\n    def _get_data_parallel_mesh(self):\n        \"\"\"Get the data parallel mesh, handling hybrid/fully/replicate shard modes.\"\"\"\n        mesh = self.parallel_dims.get_optional_mesh([\"dp_replicate\", \"fsdp\"])\n        if mesh is None:\n            mesh = self.parallel_dims.get_optional_mesh(\"fsdp\")\n        if mesh is None:\n            mesh = self.parallel_dims.get_optional_mesh(\"dp_replicate\")\n        return mesh\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False):\n        \"\"\"Perform forward and optionally backward pass on a batch.\"\"\"\n        tu.assign_non_tensor(data, sp_size=self.engine_config.tensor_parallel_size)\n\n        # Compute num_tokens in global batch for loss normalization\n        batch_num_tokens = data[\"loss_mask\"].sum().to(get_device_id())\n        dp_group = self.get_data_parallel_group()\n        if dp_group is not None:\n            torch.distributed.all_reduce(batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=dp_group)\n        tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())\n        tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())\n\n        micro_batches, indices = prepare_micro_batches(\n            data=data,\n            dp_group=self.get_data_parallel_group(),\n            same_micro_num_in_dp=True,\n        )\n\n        output_lst = []\n\n        ctx = torch.no_grad() if forward_only else nullcontext()\n\n        for micro_batch in micro_batches:\n            with ctx:\n                loss, output = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only)\n                if not forward_only:\n                    loss.backward()\n            output_lst.append(output)\n\n        return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data)\n\n    def model_forward_step(\n        self,\n        *,\n        inputs: torch.Tensor,\n        extra_inputs: dict[str, torch.Tensor] | None = None,\n        extra_kwargs: dict[str, torch.Tensor] | None = None,\n    ) -> torch.Tensor:\n        \"\"\"\n        Perform a forward pass through the trainer model without backward.\n        \"\"\"\n        model_parts = self.module\n        parallel_dims = self.parallel_dims\n\n        if parallel_dims.pp_enabled:\n            raise NotImplementedError(\n                \"Pipeline parallelism is not yet supported in model_forward_step. \"\n                \"This will be implemented in a follow-up PR.\"\n            )\n        else:\n            # Non-PP forward\n            assert len(model_parts) == 1\n            with self.trainer.train_context():\n                with self.trainer.maybe_enable_amp:\n                    pred = model_parts[0](inputs, **extra_inputs, **extra_kwargs)\n\n        if isinstance(pred, DTensor):\n            pred = pred.full_tensor()\n        return pred\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        raise NotImplementedError(\"forward_step must be implemented in subclass\")\n\n    def optimizer_zero_grad(self):\n        \"\"\"Zero gradients.\"\"\"\n        self.optimizer.zero_grad()\n\n    def optimizer_step(self):\n        \"\"\"Perform optimizer step with gradient clipping.\"\"\"\n        grad_norm = dist_utils.clip_grad_norm_(\n            [p for m in self.module for p in m.parameters()],\n            self.config.training.max_norm,\n            foreach=True,\n            pp_mesh=self.parallel_dims.get_optional_mesh(\"pp\"),\n            ep_enabled=self.parallel_dims.ep_enabled,\n        )\n\n        # if grad_norm is not finite, skip the update\n        if not torch.isfinite(grad_norm):\n            logger.warning(f\"grad_norm is not finite: {grad_norm}\")\n            self.optimizer.zero_grad()\n        else:\n            self.optimizer.step()\n        return grad_norm.item()\n\n    def lr_scheduler_step(self):\n        \"\"\"Advance learning rate scheduler.\"\"\"\n        self.lr_scheduler.step()\n        lr = self.lr_scheduler.schedulers[0].get_last_lr()[0]\n        return lr\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"Move model and/or optimizer to CPU or GPU.\"\"\"\n        super().to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n        if self.engine_config.forward_only:\n            return\n\n        device_name = get_device_name()\n        assert device in (device_name, \"cpu\")\n        if device == device_name:\n            if model:\n                for module in self.module:\n                    load_fsdp_model_to_gpu(module)\n            if optimizer and self.optimizer is not None:\n                load_fsdp_optimizer(self.optimizer, device)\n            gc.collect()\n        elif device == \"cpu\":\n            if model:\n                for module in self.module:\n                    offload_fsdp_model_to_cpu(module)\n            if optimizer and self.optimizer is not None:\n                offload_fsdp_optimizer(self.optimizer)\n        else:\n            raise ValueError(f\"Invalid device type: {device}\")\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"Save checkpoint.\"\"\"\n        if self._is_offload_param:\n            for module in self.module:\n                load_fsdp_model_to_gpu(module)\n\n        # Override TorchTitan's folder to use verl's path\n        parent_dir = os.path.dirname(local_path)\n        self.checkpointer.folder = parent_dir\n\n        if max_ckpt_to_keep is not None:\n            self.checkpointer.keep_latest_k = max_ckpt_to_keep\n\n        self.checkpointer.save(curr_step=global_step)\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            for module in self.module:\n                offload_fsdp_model_to_cpu(module)\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs\n    ) -> None:\n        \"\"\"Load checkpoint.\"\"\"\n        if self._is_offload_param:\n            for module in self.module:\n                load_fsdp_model_to_gpu(module)\n\n        # Override TorchTitan's folder to use verl's path\n        parent_dir = os.path.dirname(local_path)\n        self.checkpointer.folder = parent_dir\n\n        # Extract step number from path (verl uses global_step_N format)\n        match = re.search(r\"global_step_(\\d+)\", local_path)\n        if match:\n            step = int(match.group(1))\n            self.checkpointer.load(step=step)\n        else:\n            # Fallback to latest\n            self.checkpointer.load(step=-1)\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            for module in self.module:\n                offload_fsdp_model_to_cpu(module)\n\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(self.optimizer)\n\n    def get_per_tensor_param(self, **kwargs):\n        for module in self.module:\n            load_fsdp_model_to_gpu(module)\n\n        # Collect state dicts from all model parts\n        params = {}\n        for module in self.module:\n            module_params = get_model_state_dict(module)\n            params.update(module_params)\n\n        if self._is_offload_param:\n            for module in self.module:\n                offload_fsdp_model_to_cpu(module)\n\n        # Convert TorchTitan key names to HuggingFace key names (expected by vLLM)\n        sd_adapter = self.checkpointer.sd_adapter\n        if sd_adapter is not None:\n            params = sd_adapter.to_hf(params)\n\n        # When weight tying is enabled, the sd_adapter skips lm_head.weight during\n        # to_hf() conversion (since it's the same tensor as embed_tokens.weight in\n        # the torchtitan model). But vLLM needs lm_head.weight explicitly, so we\n        # add it back as a reference to embed_tokens.weight.\n        if \"model.embed_tokens.weight\" in params and \"lm_head.weight\" not in params:\n            params[\"lm_head.weight\"] = params[\"model.embed_tokens.weight\"]\n\n        device = get_device_id()  # used when fsdp2 set cpu_offload_policy\n\n        # When Expert Parallel (EP) is used, sd_adapter.to_hf() only produces\n        # individual expert weights for the locally-owned experts (e.g., 16 out of\n        # 128 with EP=8). vLLM needs ALL experts. We gather the missing experts\n        # by all-gathering each expert weight across the EP process group.\n        if self.parallel_dims.ep_enabled:\n            ep_mesh = self.parallel_dims.get_optional_mesh(\"ep\")\n            ep_group = ep_mesh.get_group()\n            ep_size = self.parallel_dims.ep\n            per_tensor_param = iter_per_tensor_params_ep(params, device, ep_group, ep_size)\n        else:\n            # TODO: cast fp32 to bf16 to reduce weight sync overhead, need more fine-grained control, e.g MoE gate\n            per_tensor_param = (\n                (\n                    name,\n                    param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)\n                    if isinstance(param, DTensor)\n                    else param,\n                )\n                for name, param in params.items()\n            )\n        # TODO: support Torchtitan PEFT\n        return per_tensor_param, None\n\n\nclass EngineEvalModeCtx(BaseEngineCtx):\n    def __init__(self, engine: TorchTitanEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"eval\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, TorchTitanEngine)\n        super().__enter__()\n        for module in self.engine.module:\n            module.eval()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, TorchTitanEngine)\n\n        # Reshard the root FSDP module\n        if self.engine.engine_config.data_parallel_shard_size > 1:\n            for module in self.engine.module:\n                module.reshard()\n\n        super().__exit__(exc_type, exc_value, traceback)\n\n\nclass EngineTrainModeCtx(BaseEngineCtx):\n    def __init__(self, engine: TorchTitanEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"train\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, TorchTitanEngine)\n        super().__enter__()\n        for module in self.engine.module:\n            module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, TorchTitanEngine)\n        self.engine.optimizer_zero_grad()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=[\"torchtitan\"], device=[\"cuda\", \"npu\"])\nclass TorchTitanEngineWithLMHead(TorchTitanEngine):\n    \"\"\"TorchTitan engine implementation for language models with LM head.\"\"\"\n\n    def prepare_model_inputs(self, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        assert pad_mode == DatasetPadMode.NO_PADDING, f\"pad_mode {pad_mode} not supported\"\n\n        multi_modal_inputs = extract_multi_modal_inputs(micro_batch.get(\"multi_modal_inputs\", []))\n        input_ids = micro_batch[\"input_ids\"]\n        position_ids = micro_batch[\"position_ids\"]\n        output_args = {}\n\n        if use_remove_padding:\n            input_ids = input_ids.values().unsqueeze(0)\n            if position_ids.dim() == 3:\n                position_ids = position_ids.values().unsqueeze(1)\n            else:\n                position_ids = position_ids.values().unsqueeze(0)\n\n            labels = torch.roll(input_ids, shifts=-1, dims=1)\n            attn_type = self.trainer.model_config.layer.attention.attn_backend\n            attention_mask = get_attention_masks(\n                input_batch=input_ids,\n                positions=position_ids,\n                attn_type=attn_type,\n            )\n        else:\n            loss_mask = micro_batch[\"loss_mask\"]\n            pad_token_id = tu.get_non_tensor_data(data=micro_batch, key=\"pad_token_id\", default=0)\n            batch_size = micro_batch.batch_size[0]\n            max_seq_len = max(input_ids.offsets().diff())\n\n            labels = torch.roll(input_ids.values(), shifts=-1, dims=0)\n            input_ids = torch.nested.to_padded_tensor(\n                input_ids, padding=pad_token_id, output_size=(batch_size, max_seq_len)\n            )\n\n            if position_ids.dim() == 3:\n                position_ids = torch.nested.to_padded_tensor(\n                    position_ids, padding=0, output_size=(batch_size, 4, max_seq_len)\n                ).transpose(0, 1)\n            else:\n                position_ids = torch.nested.to_padded_tensor(\n                    position_ids, padding=0, output_size=(batch_size, max_seq_len)\n                )\n\n            attention_mask_list = [torch.ones_like(t, dtype=torch.int32) for t in loss_mask]\n            attention_mask = torch.nested.as_nested_tensor(attention_mask_list, layout=torch.jagged)\n            attention_mask = torch.nested.to_padded_tensor(\n                attention_mask, padding=0, output_size=(batch_size, max_seq_len)\n            )\n\n        extra_inputs = {\n            \"positions\": position_ids,\n        }\n        # For arguments, like attention_masks, we have to put them in a separate\n        # dict as extra_inputs are not forwarded to other stages in PP, but\n        # extra_kwargs are.\n        extra_kwargs: dict[str, Any] = {\"attention_masks\": attention_mask}\n        if self.parallel_dims.cp_enabled:\n            input_ids, labels, extra_kwargs = prepare_context_parallel_input(\n                input_ids,\n                labels,\n                extra_kwargs,\n                self.parallel_dims.get_mesh(\"cp\"),\n                self.trainer.device,\n                self.trainer.config.parallelism.context_parallel_load_balancer,\n            )\n\n        # TODO(jessicazhong): multimodal is not yet supported for Torchtitan engine\n        extra_inputs.update(multi_modal_inputs)\n        output_args[\"labels\"] = labels\n        return input_ids, extra_inputs, extra_kwargs, output_args\n\n    def prepare_model_outputs(self, logits, output_args, micro_batch: TensorDict):\n        use_remove_padding = tu.get_non_tensor_data(data=micro_batch, key=\"use_remove_padding\", default=True)\n        pad_mode = tu.get_non_tensor_data(data=micro_batch, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n        assert pad_mode == DatasetPadMode.NO_PADDING, f\"pad_mode {pad_mode} not supported\"\n\n        temperature = micro_batch[\"temperature\"]\n        calculate_entropy = tu.get_non_tensor_data(data=micro_batch, key=\"calculate_entropy\", default=False)\n        labels = output_args[\"labels\"]\n        model_output = {}\n\n        input_ids = micro_batch[\"input_ids\"]\n        cu_seqlens = input_ids.offsets()\n        if use_remove_padding:\n            labels = labels.squeeze(0)\n            logits_rmpad = logits.squeeze(0)\n            # PyTorch's autograd doesn't allow in-place modification of views when gradients need to flow back\n            logits_rmpad = logits_rmpad / temperature\n\n            inplace_backward = True\n            if calculate_entropy:\n                inplace_backward = False\n            log_probs = logprobs_from_logits(\n                logits=logits_rmpad,\n                labels=labels,\n                inplace_backward=inplace_backward,\n            )\n\n            if calculate_entropy:\n                if not self.engine_config.entropy_checkpointing:\n                    entropy_rmpad = self.compute_entropy_from_logits(logits_rmpad)\n                else:\n                    entropy_rmpad = torch.utils.checkpoint.checkpoint(self.compute_entropy_from_logits, logits_rmpad)\n\n            log_probs = torch.nested.nested_tensor_from_jagged(log_probs.squeeze(0), cu_seqlens)\n            if calculate_entropy:\n                entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n        else:\n            logits.div_(temperature)\n            if calculate_entropy:\n                if not self.engine_config.entropy_checkpointing:\n                    entropy = verl_F.entropy_from_logits(logits)\n                else:\n                    entropy = torch.utils.checkpoint.checkpoint(verl_F.entropy_from_logits, logits)\n\n            seq_lengths = cu_seqlens.diff()\n            starts = torch.zeros_like(seq_lengths, dtype=torch.int64)\n            logits = torch.nested.narrow(logits, 1, starts, seq_lengths, layout=torch.jagged)\n            logits_rmpad = torch.cat([t for t in logits.unbind()])\n            log_probs = logprobs_from_logits(logits=logits_rmpad, labels=output_args[\"labels\"])\n            log_probs = torch.nested.nested_tensor_from_jagged(log_probs, cu_seqlens)\n            if calculate_entropy:\n                entropy = torch.nested.narrow(entropy, 1, starts, seq_lengths, layout=torch.jagged)\n                entropy_rmpad = torch.cat([t for t in entropy.unbind()])\n                entropy = torch.nested.nested_tensor_from_jagged(entropy_rmpad, cu_seqlens)\n\n        model_output[\"log_probs\"] = log_probs\n        if calculate_entropy:\n            model_output[\"entropy\"] = entropy\n\n        return model_output\n\n    def forward_step(self, micro_batch: TensorDict, loss_function, forward_only):\n        device_name = get_device_name()\n        micro_batch = micro_batch.to(get_device_id())\n        input_ids, extra_inputs, extra_kwargs, output_args = self.prepare_model_inputs(micro_batch=micro_batch)\n\n        with torch.autocast(device_type=device_name, dtype=torch.bfloat16):\n            logits = self.model_forward_step(inputs=input_ids, extra_inputs=extra_inputs, extra_kwargs=extra_kwargs)\n\n            model_output = self.prepare_model_outputs(logits=logits, output_args=output_args, micro_batch=micro_batch)\n\n            if loss_function is not None:\n                loss, metrics = loss_function(\n                    model_output=model_output, data=micro_batch, dp_group=self.get_data_parallel_group()\n                )\n            else:\n                assert forward_only, \"forward_only must be True when loss_function is None\"\n                loss = torch.tensor(1.0, device=device_name)\n                metrics = {}\n\n            output = {\n                \"model_output\": model_output,\n                \"loss\": loss.detach().item(),\n                \"metrics\": metrics,\n            }\n\n            return loss, output\n"
  },
  {
    "path": "verl/workers/engine/torchtitan/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport importlib\nimport logging\nimport re\nfrom collections import defaultdict\nfrom collections.abc import Generator, Iterator\nfrom dataclasses import dataclass\nfrom typing import Any\n\nimport torch\nimport torch.distributed\nimport torch.nn as nn\nfrom torch.distributed._composable.fsdp import FSDPModule\nfrom torch.distributed.tensor import DTensor\nfrom torch.nn.attention.flex_attention import _mask_mod_signature, and_masks\nfrom torchtitan.components.dataloader import BaseDataLoader\nfrom torchtitan.models.common.attention import (\n    AttentionMasksType,\n    VarlenMetadata,\n    create_attention_mask,\n    get_causal_mask_mod,\n)\n\nlogger = logging.getLogger(__name__)\n\n\nclass NoOpDataLoader(BaseDataLoader):\n    \"\"\"A no-op dataloader for use when verl manages its own data loading.\n\n    Satisfies the BaseDataLoader interface required by torchtitan's Trainer\n    but does nothing. Its __iter__ yields nothing, and state_dict /\n    load_state_dict are no-ops.\n    \"\"\"\n\n    @dataclass(kw_only=True, slots=True)\n    class Config(BaseDataLoader.Config):\n        pass\n\n    def __init__(self, **kwargs):\n        pass\n\n    def __iter__(self) -> Iterator[tuple[dict[str, torch.Tensor], torch.Tensor]]:\n        return iter([])\n\n    def state_dict(self):\n        return {}\n\n    def load_state_dict(self, state_dict):\n        pass\n\n\n# Mapping from HuggingFace model_type to torchtitan model name.\n# Torchtitan models not mapped here:\n#   - flux: diffusion model, not applicable to verl's RL/SFT workflows\n#   - llama3_ft: fault-tolerant variant of llama3, same HF models (mapped via \"llama\")\n_HF_MODEL_TYPE_TO_TORCHTITAN_NAME = {\n    \"qwen2\": \"qwen3\",\n    \"qwen3\": \"qwen3\",\n    \"qwen2_moe\": \"qwen3\",\n    \"qwen3_moe\": \"qwen3\",\n    \"llama\": \"llama3\",\n    \"llama4\": \"llama4\",\n    \"deepseek_v3\": \"deepseek_v3\",\n    \"gpt_oss\": \"gpt_oss\",\n}\n\n\ndef derive_torchtitan_name_and_flavor(hf_config) -> tuple[str, str]:\n    \"\"\"Derive torchtitan model name and flavor from a HuggingFace config.\n\n    The name is mapped from ``hf_config.model_type``. The flavor is found by\n    matching architecture parameters (dim, n_layers, vocab_size) against the\n    known flavors registered in the torchtitan model package.\n\n    Args:\n        hf_config: A HuggingFace AutoConfig object.\n\n    Returns:\n        A ``(name, flavor)`` tuple.\n\n    Raises:\n        ValueError: If model_type is unsupported or no matching flavor is found.\n    \"\"\"\n    model_type = getattr(hf_config, \"model_type\", None)\n    if model_type is None:\n        raise ValueError(\"HuggingFace config does not have 'model_type' field\")\n\n    name = _HF_MODEL_TYPE_TO_TORCHTITAN_NAME.get(model_type)\n    if name is None:\n        raise ValueError(\n            f\"Cannot derive torchtitan model name from HF model_type '{model_type}'. \"\n            f\"Supported types: {list(_HF_MODEL_TYPE_TO_TORCHTITAN_NAME.keys())}.\"\n        )\n\n    # Import the model package and find the configs dict\n    model_module = importlib.import_module(f\"torchtitan.models.{name}\")\n    model_configs = None\n    for attr in dir(model_module):\n        obj = getattr(model_module, attr)\n        if isinstance(obj, dict) and attr.endswith(\"_configs\"):\n            model_configs = obj\n            break\n\n    if model_configs is None:\n        raise ValueError(\n            f\"Could not find model configs dict in torchtitan.models.{name}. \"\n            f\"Expected a dict attribute ending with '_configs'.\"\n        )\n\n    hidden_size = hf_config.hidden_size\n    num_layers = hf_config.num_hidden_layers\n    vocab_size = hf_config.vocab_size\n\n    for flavor_name, model_cfg in model_configs.items():\n        if (\n            getattr(model_cfg, \"dim\", None) == hidden_size\n            and getattr(model_cfg, \"n_layers\", None) == num_layers\n            and getattr(model_cfg, \"vocab_size\", None) == vocab_size\n        ):\n            logger.info(\n                f\"Auto-derived torchtitan name='{name}', flavor='{flavor_name}' from HF model_type='{model_type}'\"\n            )\n            return name, flavor_name\n\n    raise ValueError(\n        f\"No matching torchtitan flavor found for model_type='{model_type}' \"\n        f\"(hidden_size={hidden_size}, num_hidden_layers={num_layers}, \"\n        f\"vocab_size={vocab_size}). \"\n        f\"Available flavors for '{name}': {list(model_configs.keys())}.\"\n    )\n\n\ndef enable_fsdp_gradient_division(model: nn.Module, dp_size: int) -> None:\n    \"\"\"\n    Re-enable FSDP's automatic gradient division.\n\n    TorchTitan calls disable_fsdp_gradient_division() which sets gradient_divide_factor=1.0.\n    This re-enables it by setting the factor to the specified dp_size, so gradients are\n    averaged across FSDP ranks. This is needed for verl's loss scaling (loss * dp_size)\n    to work correctly.\n\n    Args:\n        model: The model (or model part) to enable gradient division on.\n        dp_size: The data parallel size to use as the gradient divide factor.\n    \"\"\"\n\n    for module in model.modules():\n        if isinstance(module, FSDPModule):\n            module.set_gradient_divide_factor(float(dp_size))\n\n\ndef get_attention_masks(\n    input_batch: torch.Tensor,\n    positions: torch.Tensor,\n    attn_type: str,\n) -> AttentionMasksType:\n    match attn_type:\n        case \"flex\":\n            return _get_flex_attention_masks(\n                input_batch,\n                positions,\n            )\n        case \"varlen\":\n            return _create_varlen_metadata_for_document(\n                input_batch,\n                positions,\n            )\n        case _:\n            raise TypeError(\"Only varlen and flex attn masks are supported\")\n\n\ndef _get_document_mask_mod(positions: torch.Tensor) -> _mask_mod_signature:\n    # Detect boundaries from position resets\n    first_dummy_value = positions[:, :1] - 1\n    position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)\n    sequence_indices = (position_diff != 1).cumsum(-1)  # [batch, seq]\n\n    def document_mask(b: torch.Tensor, h: torch.Tensor, q_idx: torch.Tensor, kv_idx: torch.Tensor) -> torch.Tensor:\n        return sequence_indices[b, q_idx] == sequence_indices[b, kv_idx]\n\n    return document_mask\n\n\ndef _get_flex_attention_masks(\n    input_batch: torch.Tensor,\n    positions: torch.Tensor,\n) -> AttentionMasksType:\n    mask_mods = [get_causal_mask_mod()]\n    B = input_batch.shape[0]\n    mask_mods.append(_get_document_mask_mod(positions=positions))\n    return create_attention_mask(and_masks(*mask_mods), B, None, input_batch.shape[1], input_batch.shape[1])\n\n\ndef _create_varlen_metadata_for_document(input_batch: torch.Tensor, positions: torch.Tensor) -> VarlenMetadata:\n    \"\"\"\n    Creates cumulative sequence length indices needed for variable length attention\n\n    Args:\n        input_batch: Input token IDs with shape [batch, seq].\n        positions: Position IDs with shape [batch, seq]. Boundaries detected where\n            position diff != 1 (i.e., position resets).\n\n    Returns:\n        VarlenMetadata containing cumulative sequence length indices for q, k, and max_seq_len\n    \"\"\"\n    batch_size, seq_len = input_batch.shape\n    device = input_batch.device\n\n    # Detect boundaries from position resets (where diff != 1)\n    first_dummy_value = positions[:, :1] - 1\n    position_diff = torch.diff(positions, prepend=first_dummy_value, dim=-1)\n    # boundary_mask[b, i] is True if position i starts a new document\n    boundary_mask = position_diff != 1  # [batch, seq]\n    boundary_mask[:, 0] = True\n\n    cu_seqlens_list, all_seq_lengths = [], []\n    offset = 0\n\n    for b in range(batch_size):\n        # Find positions where new documents start\n        boundary_positions = boundary_mask[b].nonzero(as_tuple=True)[0].to(torch.int32)\n        sample_cu_seqlens = torch.cat(\n            [\n                boundary_positions,\n                torch.tensor([seq_len], dtype=torch.int32, device=device),\n            ]\n        )\n        sample_cu_seqlens = torch.unique_consecutive(sample_cu_seqlens)\n\n        seq_lengths = torch.diff(sample_cu_seqlens)\n        all_seq_lengths.append(seq_lengths)\n\n        cu_seqlens_adjusted = sample_cu_seqlens[:-1] + offset\n        cu_seqlens_list.append(cu_seqlens_adjusted)\n\n        offset += seq_len\n\n    packed_cu_seqlens = torch.cat(cu_seqlens_list + [torch.tensor([offset], dtype=torch.int32, device=device)])\n\n    max_seqlen = 0\n    if len(all_seq_lengths) > 0:\n        all_seq_lengths = torch.cat(all_seq_lengths)\n        # device to host sync but only done once per model forward\n        max_seqlen = all_seq_lengths.max().item()\n\n    return VarlenMetadata(\n        cu_seq_q=packed_cu_seqlens,\n        cu_seq_k=packed_cu_seqlens,\n        max_q=max_seqlen,\n        max_k=max_seqlen,\n    )\n\n\n# Regex to parse: model.layers.{L}.mlp.experts.{E}.{weight_suffix}\n_EXPERT_PATTERN = re.compile(r\"\\.layers\\.(\\d+)\\..*\\.experts\\.(\\d+)\\.(.*)\")\n\n\ndef _parse_expert_name(name: str) -> tuple[int, int, str] | None:\n    \"\"\"Parse layer_id, expert_id, weight_suffix from expert param name.\"\"\"\n    match = _EXPERT_PATTERN.search(name)\n    if match:\n        return int(match.group(1)), int(match.group(2)), match.group(3)\n    return None\n\n\ndef _make_expert_name_template(name: str) -> str:\n    \"\"\"Convert 'model.layers.0.mlp.experts.3.w1' -> 'model.layers.0.mlp.experts.{}.w1'\"\"\"\n    return _EXPERT_PATTERN.sub(lambda m: f\".layers.{m.group(1)}.mlp.experts.{{}}.{m.group(3)}\", name)\n\n\ndef iter_per_tensor_params_ep(\n    params: dict[str, Any],\n    device: int,\n    ep_group: torch.distributed.ProcessGroup,\n    ep_size: int,\n) -> Generator[tuple[str, torch.Tensor], None, None]:\n    \"\"\"Yield (name, tensor) pairs for weight sync with Expert Parallel.\n\n    Gathers expert weights across EP ranks one (layer, weight_type) group\n    at a time to avoid OOM from materializing all experts simultaneously.\n\n    Non-expert params are yielded first (with FSDP full_tensor() if needed),\n    then expert params are all-gathered per group and yielded individually.\n\n    Args:\n        params: HF-format state dict with per-expert keys. Expert keys must\n            follow the pattern ``model.layers.{L}.mlp.experts.{E}.{suffix}``.\n        device: device ID to place tensors on.\n        ep_group: The EP process group for all-gather.\n        ep_size: Number of EP ranks.\n    \"\"\"\n    expert_params: dict[tuple[int, str], dict[int, tuple[str, Any]]] = defaultdict(dict)\n    non_expert_params: list[tuple[str, Any]] = []\n\n    for name, param in params.items():\n        parsed = _parse_expert_name(name) if \"mlp.experts.\" in name else None\n        if parsed is None:\n            non_expert_params.append((name, param))\n        else:\n            layer_id, expert_id, weight_suffix = parsed\n            expert_params[(layer_id, weight_suffix)][expert_id] = (name, param)\n\n    params.clear()\n\n    # Yield non-expert params\n    for name, param in non_expert_params:\n        if isinstance(param, DTensor):\n            yield name, param.to(device, non_blocking=True).full_tensor().to(torch.bfloat16, non_blocking=True)\n        else:\n            yield name, param\n    del non_expert_params\n\n    # Yield expert params with all-gather\n    for (layer_id, weight_suffix), experts_dict in sorted(expert_params.items()):\n        sorted_expert_ids = sorted(experts_dict.keys())\n\n        # Stack local expert weights\n        local_weights = []\n        for eid in sorted_expert_ids:\n            _, param = experts_dict[eid]\n            if isinstance(param, DTensor):\n                param = param.to(device, non_blocking=True).full_tensor()\n            else:\n                param = param.to(device, non_blocking=True)\n            local_weights.append(param)\n\n        name_template = _make_expert_name_template(experts_dict[sorted_expert_ids[0]][0])\n        local_stacked = torch.stack(local_weights, dim=0)\n\n        # All-gather across EP ranks\n        gathered_list = [torch.empty_like(local_stacked) for _ in range(ep_size)]\n        torch.distributed.all_gather(gathered_list, local_stacked, group=ep_group)\n        all_experts = torch.cat(gathered_list, dim=0)\n\n        for expert_id in range(all_experts.shape[0]):\n            yield name_template.format(expert_id), all_experts[expert_id].to(torch.bfloat16).clone()\n\n        del local_weights, local_stacked, gathered_list, all_experts\n        torch.cuda.empty_cache()\n"
  },
  {
    "path": "verl/workers/engine/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport os\nimport random\n\nimport numpy as np\nimport torch\nfrom tensordict import TensorDict\n\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.device import is_npu_available\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.seqlen_balancing import rearrange_micro_batches, restore_dynamic_batch\n\n\ndef enable_full_determinism(seed: int):\n    \"\"\"\n    Helper function for reproducibility in distributed training.\n    See https://pytorch.org/docs/stable/notes/randomness.html for details.\n    \"\"\"\n\n    os.environ[\"PYTHONHASHSEED\"] = str(seed)\n    os.environ[\"CUBLAS_WORKSPACE_CONFIG\"] = \":16:8\"\n    os.environ[\"NCCL_DETERMINISTIC\"] = \"1\"\n    os.environ[\"FLASH_ATTENTION_DETERMINISTIC\"] = \"1\"\n    if is_npu_available:\n        # The environment variable required to enable deterministic mode on Ascend NPUs.\n        os.environ[\"NCCL_DETERMINISTIC\"] = \"true\"\n        os.environ[\"CLOSE_MATMUL_K_SHIFT\"] = \"1\"\n\n    random.seed(seed)\n    np.random.seed(seed)\n    torch.manual_seed(seed)\n    torch.cuda.manual_seed(seed)\n    torch.cuda.manual_seed_all(seed)\n    torch.use_deterministic_algorithms(True, warn_only=True)\n    # Enable CUDNN deterministic mode\n    torch.backends.cudnn.deterministic = True\n    torch.backends.cudnn.benchmark = False\n    torch.backends.cudnn.enabled = False\n    if is_npu_available:\n        torch.npu.manual_seed(seed)\n        torch.npu.manual_seed_all(seed)\n\n\ndef prepare_micro_batches(\n    data: TensorDict,\n    dp_group=None,\n    num_batches_divided_by=None,\n    same_micro_num_in_dp=True,\n    min_num_micro_batch=None,\n    use_dynamic_bsz_balance=True,\n):\n    \"\"\"\n    Prepare micro batches from data.\n    \"\"\"\n    use_dynamic_bsz = tu.get_non_tensor_data(data=data, key=\"use_dynamic_bsz\", default=True)\n    sp_size = tu.get_non_tensor_data(data=data, key=\"sp_size\", default=1)\n\n    force_group_size = tu.get_non_tensor_data(data=data, key=\"force_group_size\", default=1)\n\n    if use_dynamic_bsz:\n        assert \"max_token_len_per_gpu\" in data.keys(), \"max_token_len_per_gpu must be set when use_dynamic_bsz is True\"\n        max_token_len_per_gpu = data[\"max_token_len_per_gpu\"]\n        max_token_len = max_token_len_per_gpu * sp_size\n        micro_batches, batch_idx_list = rearrange_micro_batches(\n            data,\n            max_token_len=max_token_len,\n            dp_group=dp_group,\n            num_batches_divided_by=num_batches_divided_by,\n            same_micro_num_in_dp=same_micro_num_in_dp,\n            min_num_micro_batch=min_num_micro_batch,\n            use_dynamic_bsz_balance=use_dynamic_bsz_balance,\n            force_group_size=force_group_size,\n        )\n    else:\n        total_data_size = len(data)\n        micro_batch_size_per_gpu = data[\"micro_batch_size_per_gpu\"]\n        assert total_data_size % (force_group_size * micro_batch_size_per_gpu) == 0, (\n            \"data size must be divisible by force_group_size * micro_batch_size_per_gpu\"\n        )\n        micro_batches = tu.chunk_tensordict(data, total_data_size // (micro_batch_size_per_gpu * force_group_size))\n        batch_idx_list = None\n    return micro_batches, batch_idx_list\n\n\ndef postprocess_batch_func(output_lst, indices, data: TensorDict):\n    \"\"\"postprocess the output of a forward_backward_batch.\n    output_lst is a list of dict containing outputs for each micro-batch\n    reorder entropy and outputs. Return None for other pp ranks\n    only on last rank. It should be on every tp rank\n\n    each losses_reduced contains 1. model_output, 2. loss, 3. metrics.\n    \"\"\"\n\n    use_dynamic_bsz = tu.get_non_tensor_data(data=data, key=\"use_dynamic_bsz\", default=True)\n    pad_mode = tu.get_non_tensor_data(data=data, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n    assert pad_mode == DatasetPadMode.NO_PADDING, \"postprocess_batch_func only support NO_PADDING pad_mode\"\n\n    # losses_reduced is a list of dict containing outputs for each micro-batch\n    # reorder entropy and outputs. Return None for other pp ranks\n    # only on last rank. It should be on every tp rank\n\n    # losses_reduced contains 1. model_output, 2. loss, 3. metrics.\n    # We perform reverse\n\n    model_output = {}\n    losses = []\n    aggregated_metrics = {}\n\n    # model output\n    for o in output_lst:\n        if \"model_output\" in o:\n            for key, val in o[\"model_output\"].items():\n                if key not in model_output:\n                    model_output[key] = []\n                model_output[key].append(val)\n\n    # concat results from micro batches\n    for key, val in model_output.items():\n        if pad_mode == DatasetPadMode.NO_PADDING:\n            tensors = [tensor for nt in model_output[key] for tensor in nt.unbind()]\n            model_output[key] = torch.nested.as_nested_tensor(tensors, layout=torch.jagged)\n        else:\n            raise NotImplementedError(f\"pad_mode {pad_mode} not implemented\")\n\n        # reverse with dynamic bsz\n        if use_dynamic_bsz:\n            model_output[key] = restore_dynamic_batch(model_output[key], indices)\n\n    # loss\n    for o in output_lst:\n        if \"loss\" in o:\n            losses.append(o[\"loss\"])\n\n    # metrics\n    for o in output_lst:\n        if \"metrics\" in o:\n            metrics = o[\"metrics\"]\n            append_to_dict(aggregated_metrics, metrics)\n\n    output = {\n        \"model_output\": model_output,\n        \"loss\": losses,\n        \"metrics\": aggregated_metrics,\n    }\n\n    return output\n"
  },
  {
    "path": "verl/workers/engine/veomni/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .transformer_impl import VeOmniEngine, VeOmniEngineWithLMHead\n\n__all__ = [\"VeOmniEngine\", \"VeOmniEngineWithLMHead\"]\n"
  },
  {
    "path": "verl/workers/engine/veomni/transformer_impl.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport logging\nfrom dataclasses import dataclass, field\nfrom typing import Any, Callable, Optional, Sequence\n\nimport torch\nimport torch.distributed as dist\nfrom tensordict import TensorDict\nfrom torch.distributed.tensor import DTensor\nfrom veomni.distributed import parallel_state\nfrom veomni.distributed.offloading import build_activation_offloading_context\nfrom veomni.distributed.torch_parallelize import build_parallelize_model\nfrom veomni.models.auto import build_foundation_model\nfrom veomni.optim import build_lr_scheduler, build_optimizer\n\nimport verl.utils.torch_functional as verl_F\nfrom verl.trainer.config import CheckpointConfig\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.device import get_device_id, get_device_name\nfrom verl.utils.fsdp_utils import fsdp_version\nfrom verl.utils.model import convert_weight_keys\nfrom verl.utils.profiler import log_gpu_memory_usage\nfrom verl.utils.ulysses import (\n    get_ulysses_sequence_parallel_group,\n    set_ulysses_sequence_parallel_group,\n)\nfrom verl.workers.config import HFModelConfig, VeOmniEngineConfig, VeOmniOptimizerConfig\n\nfrom ..base import BaseEngineCtx, EngineRegistry\nfrom ..fsdp.transformer_impl import FSDPEngine, FSDPEngineWithLMHead\nfrom ..utils import enable_full_determinism, postprocess_batch_func, prepare_micro_batches\nfrom .utils import (\n    MOE_PARAM_HANDERS,\n    VL_TYPE2INDEX,\n    load_veomni_model_to_gpu,\n    load_veomni_optimizer,\n    offload_veomni_model_to_cpu,\n    offload_veomni_optimizer,\n)\n\nlogger = logging.getLogger(__file__)\n\n\nclass VeOmniEngine(FSDPEngine):\n    def __init__(\n        self,\n        model_config: HFModelConfig,\n        engine_config: VeOmniEngineConfig,\n        optimizer_config: VeOmniOptimizerConfig,\n        checkpoint_config: CheckpointConfig,\n        **kwargs,\n    ):\n        \"\"\"\n        Initialize the VeOmniEngine.\n\n        Sets up distributed device meshes, LoRA, and offload policies based on config.\n\n        Args:\n            config: Configuration object with VeOmni and model settings.\n        \"\"\"\n\n        self.model_config = model_config\n        self.engine_config = engine_config\n        self.optimizer_config = optimizer_config\n        self.checkpoint_config = checkpoint_config\n        # VeOmniEngine only supports fsdp2.\n        self.data_parallel_mode = \"fsdp2\"\n        self.rank = dist.get_rank()\n\n        fsdp_size = self.engine_config.fsdp_size\n        world_size = dist.get_world_size()\n        dp_size = world_size // self.engine_config.ulysses_parallel_size\n\n        if fsdp_size < 0 or fsdp_size >= dp_size:\n            data_parallel_replicate_size = 1\n            data_parallel_shard_size = dp_size\n        else:\n            if dp_size % fsdp_size != 0:\n                raise ValueError(\n                    f\"Data parallel size ({dp_size}) must be divisible by fsdp_size ({fsdp_size}). \"\n                    \"Please adjust your parallel configuration.\"\n                )\n            data_parallel_replicate_size = dp_size // fsdp_size\n            data_parallel_shard_size = fsdp_size\n\n        parallel_state.init_parallel_state(\n            dp_size=dp_size,\n            dp_replicate_size=data_parallel_replicate_size,\n            dp_shard_size=data_parallel_shard_size,\n            ep_size=self.engine_config.expert_parallel_size,\n            ulysses_size=self.engine_config.ulysses_parallel_size,\n            dp_mode=self.data_parallel_mode,\n        )\n\n        if self.engine_config.full_determinism:\n            enable_full_determinism(seed=self.engine_config.seed)\n\n        self.use_remove_padding = self.model_config.use_remove_padding\n\n        self._is_offload_param = self.engine_config.param_offload\n        self._is_offload_optimizer = self.engine_config.optimizer_offload\n        self._is_lora = self.model_config.lora_rank > 0\n\n        self.use_ulysses_sp = parallel_state.get_parallel_state().sp_enabled\n        self.ulysses_sequence_parallel_size = self.engine_config.ulysses_parallel_size\n\n        if self.use_ulysses_sp:\n            self.ulysses_parallel_group = parallel_state.get_parallel_state().device_mesh[\"sp\"].get_group()\n        else:\n            self.ulysses_parallel_group = None\n\n        if self.engine_config.entropy_from_logits_with_chunking:\n            entropy_from_logits = verl_F.entropy_from_logits_with_chunking\n        else:\n            entropy_from_logits = verl_F.entropy_from_logits\n\n        self.compute_entropy_from_logits = (\n            torch.compile(entropy_from_logits, dynamic=True)\n            if self.engine_config.use_torch_compile  #  use torch compile by default\n            else entropy_from_logits\n        )\n\n    def initialize(self):\n        \"\"\"\n        Build the model, optimizer, and learning rate scheduler under VeOmni.\n\n        Applies device, dtype, and precision configurations, including mixed precision.\n        Sets up checkpoint manager and FLOPs counter.\n        \"\"\"\n        self._build_model_optimizer()\n\n        self.checkpoint_manager = FSDPCheckpointManager(\n            model=self.module,\n            optimizer=self.optimizer,\n            lr_scheduler=self.lr_scheduler,\n            processing_class=self.model_config.get_processor(),\n            checkpoint_config=self.checkpoint_config,\n            trust_remote_code=self.model_config.trust_remote_code,\n        )\n\n        self.to(\n            device=\"cpu\",\n            model=self._is_offload_param,\n            optimizer=self._is_offload_optimizer,\n            grad=self._is_offload_optimizer,\n        )\n\n        log_gpu_memory_usage(\"After offload model/optimizer/grad during init\", logger=logger)\n\n    def _build_optimizer(self, module):\n        optimizer = build_optimizer(\n            module,\n            lr=self.optimizer_config.lr,\n            betas=self.optimizer_config.betas,\n            weight_decay=self.optimizer_config.weight_decay,\n            optimizer_type=self.optimizer_config.optimizer,\n        )\n        get_optimizer_pre_hook = getattr(module, \"get_optimizer_pre_hook\", None)\n        if get_optimizer_pre_hook is not None:\n            optimizer_pre_hook = get_optimizer_pre_hook(module, module.config, self.data_parallel_mode)\n            optimizer.register_step_pre_hook(optimizer_pre_hook)\n\n        return optimizer\n\n    def _build_lr_scheduler(self, optimizer):\n        optim_config = self.optimizer_config\n        lr_scheduler = build_lr_scheduler(\n            optimizer,\n            train_steps=optim_config.total_training_steps,\n            lr=optim_config.lr,\n            lr_min=optim_config.lr_min,\n            lr_decay_style=optim_config.lr_scheduler_type,\n            lr_decay_ratio=optim_config.lr_decay_ratio,\n            lr_warmup_ratio=optim_config.lr_warmup_steps_ratio,\n            lr_start=optim_config.lr_start,\n        )\n\n        return lr_scheduler\n\n    def _build_model_optimizer(self):\n        # Load base model with specified configuration and dtype\n        module = build_foundation_model(\n            config_path=self.model_config.hf_config_path,\n            weights_path=self.model_config.path,\n            torch_dtype=\"float32\" if self.engine_config.mixed_precision else \"bfloat16\",\n            attn_implementation=self.engine_config.attn_implementation,\n            moe_implementation=self.engine_config.moe_implementation,\n            init_device=self.engine_config.init_device,\n        )\n        log_gpu_memory_usage(\"After load base model\", logger=logger)\n\n        # Applies parallel strategies to the model.\n        log_gpu_memory_usage(\"Before parallelize model\", logger=logger)\n        module = build_parallelize_model(\n            module,\n            init_device=self.engine_config.init_device,\n            weights_path=self.model_config.path,\n            enable_full_shard=self.engine_config.enable_full_shard,\n            enable_mixed_precision=self.engine_config.mixed_precision,\n            enable_gradient_checkpointing=self.model_config.enable_gradient_checkpointing,\n            enable_fsdp_offload=self.engine_config.enable_fsdp_offload,\n            basic_modules=module._no_split_modules + self.engine_config.basic_modules,\n            enable_reentrant=self.engine_config.enable_reentrant,\n            enable_forward_prefetch=self.engine_config.forward_prefetch,\n        )\n        log_gpu_memory_usage(\"After parallelize model\", logger=logger)\n\n        if not self.engine_config.forward_only:\n            # Initialize optimizer with model parameters and config settings\n            optimizer = self._build_optimizer(module)\n            # Create learning rate scheduler with warmup and decay settings\n            lr_scheduler = self._build_lr_scheduler(optimizer)\n        else:\n            optimizer = None\n            lr_scheduler = None\n\n        self.module = module\n        self.optimizer = optimizer\n        self.lr_scheduler = lr_scheduler\n        self.model_fwd_context, self.model_bwd_context = build_activation_offloading_context(\n            self.model_config.enable_activation_offload,\n            self.model_config.enable_gradient_checkpointing,\n            self.engine_config.activation_gpu_limit,\n        )\n\n    def optimizer_step(self):\n        \"\"\"\n        Perform an optimization step using the optimizer.\n        \"\"\"\n        if hasattr(self.module, \"clip_grad_norm_\"):\n            grad_norm = self.module.clip_grad_norm_(self.optimizer_config.clip_grad)\n        else:\n            grad_norm = torch.nn.utils.clip_grad_norm_(self.module.parameters(), self.optimizer_config.clip_grad)\n\n        if isinstance(grad_norm, DTensor):\n            grad_norm = grad_norm.full_tensor()\n\n        # if grad_norm is not finite, skip the update\n        if not torch.isfinite(grad_norm):\n            print(f\"WARN: grad_norm is not finite: {grad_norm}\")\n            self.optimizer.zero_grad()\n        else:\n            self.optimizer.step()\n        return grad_norm.item()\n\n    def forward_backward_batch(self, data: TensorDict, loss_function: Callable, forward_only=False) -> Any:\n        \"\"\"\n        Perform a forward pass and optionally a backward pass on a batch of data.\n\n        Args:\n            data: The input data for the forward pass, typically containing tensors and metadata.\n            loss_function: The loss function to optimize. See `verl.workers.roles.utils.losses` for examples.\n            forward_only: If True, perform only the forward pass. If False, perform forward and backward pass.\n\n        Returns:\n            Any: The output of the forward pass, which can be used for loss computation or other purposes.\n        \"\"\"\n        tu.assign_non_tensor(data, sp_size=parallel_state.get_parallel_state().ulysses_size)\n\n        # compute num_tokens in global batch for loss normalization\n        batch_num_tokens = data[\"loss_mask\"].sum().to(get_device_id())\n        torch.distributed.all_reduce(\n            batch_num_tokens, op=torch.distributed.ReduceOp.SUM, group=self.get_data_parallel_group()\n        )\n        tu.assign_non_tensor(data, batch_num_tokens=batch_num_tokens.item())\n        tu.assign_non_tensor(data, dp_size=self.get_data_parallel_size())\n\n        micro_batches, indices = prepare_micro_batches(\n            data=data, dp_group=self.get_data_parallel_group(), same_micro_num_in_dp=True\n        )\n\n        output_lst = []\n\n        for micro_batch in micro_batches:\n            with self.model_fwd_context:\n                loss, meta_info = self.forward_step(micro_batch, loss_function=loss_function, forward_only=forward_only)\n            if not forward_only:\n                with self.model_bwd_context:\n                    loss.backward()\n\n            output_lst.append(meta_info)\n\n        return postprocess_batch_func(output_lst=output_lst, indices=indices, data=data)\n\n    def get_data_parallel_rank(self):\n        return parallel_state.get_parallel_state().device_mesh.get_local_rank(\"dp\")\n\n    def get_data_parallel_size(self):\n        return torch.distributed.get_world_size() // parallel_state.get_parallel_state().ulysses_size\n\n    def get_data_parallel_group(self):\n        if parallel_state.get_parallel_state().ulysses_size > 1:\n            return parallel_state.get_parallel_state().device_mesh.get_group(mesh_dim=\"dp\")\n        else:\n            return torch.distributed.group.WORLD\n\n    def get_model_parallel_group(self):\n        raise NotImplementedError\n\n    def get_context_parallel_group(self):\n        raise NotImplementedError\n\n    def is_mp_src_rank_with_outputs(self):\n        \"\"\"\n        Whether the current rank is the first rank in model parallel group that contains model outputs\n        \"\"\"\n        if parallel_state.get_parallel_state().ulysses_size > 1:\n            is_collect = parallel_state.get_parallel_state().device_mesh[\"ulysses\"].get_local_rank() == 0\n        else:\n            is_collect = True\n        return is_collect\n\n    def train_mode(self, **kwargs):\n        \"\"\"\n        Return a context manager that switches to training mode with VeOmni-specific handling.\n\n        Includes parameter and optimizer offload entry/exit.\n        \"\"\"\n        return EngineTrainModeCtx(self, **kwargs)\n\n    def eval_mode(self, **kwargs):\n        \"\"\"\n        Return a context manager that switches to evaluation mode with VeOmni-specific handling.\n\n        Includes activation offload entry/exit.\n        \"\"\"\n        return EngineEvalModeCtx(self, **kwargs)\n\n    def to(self, device: str, model: bool = True, optimizer: bool = True, grad: bool = True):\n        \"\"\"\n        Move model parameters, optimizer states, or both to the specified device.\n        Note that this function executes irrespective of offload config. It serves as manual control.\n\n        Args:\n            device: Target device identifier.\n            model: If True, move the model.\n            optimizer: If True, move the optimizer states.\n        \"\"\"\n        super(FSDPEngine, self).to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n        device_name = get_device_name()\n\n        assert device in (device_name, \"cpu\")\n        if device == device_name:\n            if model:\n                load_veomni_model_to_gpu(self.module)\n            if optimizer and self.optimizer is not None:\n                load_veomni_optimizer(self.optimizer, device)\n        elif device == \"cpu\":\n            if model:\n                offload_veomni_model_to_cpu(self.module)\n            if optimizer and self.optimizer is not None:\n                offload_veomni_optimizer(self.optimizer)\n        else:\n            raise ValueError(f\"Invalid device type: {device}\")\n\n    def save_checkpoint(\n        self,\n        local_path: str,\n        hdfs_path: Optional[str] = None,\n        global_step: int = 0,\n        max_ckpt_to_keep: Optional[int] = None,\n        **kwargs,\n    ) -> None:\n        \"\"\"\n        Save VeOmni checkpoint, handling parameter offload as needed.\n        \"\"\"\n        origin_module_device = next(self.module.parameters()).device.type\n        if self._is_offload_param or origin_module_device == \"cpu\":\n            load_veomni_model_to_gpu(self.module)\n\n        self.checkpoint_manager.save_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_veomni_model_to_cpu(self.module)\n\n    def load_checkpoint(\n        self, local_path: str, hdfs_path: Optional[str] = None, del_local_after_load: int = True, **kwargs\n    ) -> None:\n        \"\"\"\n        Load VeOmni checkpoint, restoring parameters and optimizer state.\n        \"\"\"\n        if self._is_offload_param:\n            load_veomni_model_to_gpu(self.module)\n\n        self.checkpoint_manager.load_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_veomni_model_to_cpu(self.module)\n\n        if self._is_offload_optimizer:\n            offload_veomni_optimizer(self.optimizer)\n\n    def get_per_tensor_param(self, **kwargs):\n        load_veomni_model_to_gpu(self.module)\n\n        params = self.module.state_dict()\n        params = convert_weight_keys(params, getattr(self.module, \"_fsdp_wrapped_module\", self.module))\n\n        if self._is_offload_param:\n            offload_veomni_model_to_cpu(self.module)\n\n        device = get_device_id()\n        ps = parallel_state.get_parallel_state()\n        model_type = getattr(self.module.config, \"model_type\", \"default\")\n        process_func = MOE_PARAM_HANDERS.get(model_type, lambda n, t: iter([(n, t)]))\n\n        def param_generator():\n            for name, param in params.items():\n                unsharded_tensor = param.full_tensor() if isinstance(param, DTensor) else param\n\n                is_expert_layer = \"mlp.experts.\" in name\n                is_proj = any(p in name for p in [\"down_proj\", \"gate_proj\", \"up_proj\", \"gate_up_proj\"])\n\n                if is_expert_layer and is_proj and ps.ep_enabled:\n                    output_shape = list(unsharded_tensor.shape)\n                    output_shape[0] *= ps.ep_size\n                    stacked_tensor = torch.empty(output_shape, dtype=unsharded_tensor.dtype, device=device)\n\n                    # all gather expert tensors [32, H, I] -> [128, H, I]\n                    torch.distributed.all_gather_into_tensor(stacked_tensor, unsharded_tensor, group=ps.ep_group)\n                    yield from process_func(name, stacked_tensor)\n\n                    del stacked_tensor\n                else:\n                    if is_expert_layer:\n                        yield from process_func(name, unsharded_tensor)\n                    else:\n                        yield name, unsharded_tensor\n\n        # TODO: support VeOmni LoRA\n        return param_generator(), None\n\n\nclass EngineEvalModeCtx(BaseEngineCtx):\n    def __init__(self, engine: VeOmniEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"eval\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, VeOmniEngine)\n        super().__enter__()\n        self.prev_sp_group = get_ulysses_sequence_parallel_group()\n        set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group)\n        self.engine.module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, VeOmniEngine)\n        set_ulysses_sequence_parallel_group(self.prev_sp_group)\n\n        # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes\n        # unshard the root FSDP module\n        if parallel_state.get_parallel_state().dp_shard_size > 1:\n            if fsdp_version(self.engine.module) == 1:\n                self.engine.module._handle.reshard(True)\n            elif fsdp_version(self.engine.module) == 2:\n                self.engine.module.reshard()\n\n        super().__exit__(exc_type, exc_value, traceback)\n\n\nclass EngineTrainModeCtx(BaseEngineCtx):\n    def __init__(self, engine: VeOmniEngine, **kwargs):\n        super().__init__(engine=engine, mode=\"train\", **kwargs)\n\n    def __enter__(self):\n        assert isinstance(self.engine, VeOmniEngine)\n        super().__enter__()\n        self.prev_sp_group = get_ulysses_sequence_parallel_group()\n        set_ulysses_sequence_parallel_group(self.engine.ulysses_parallel_group)\n        # TODO: Switch to eval mode after Integrating the CI environment\n        # VeOmni (ref: https://github.com/ByteDance-Seed/VeOmni/pull/421)\n        self.engine.module.train()\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        assert isinstance(self.engine, VeOmniEngine)\n        set_ulysses_sequence_parallel_group(self.prev_sp_group)\n        self.engine.optimizer_zero_grad()\n        super().__exit__(exc_type, exc_value, traceback)\n\n\n@dataclass\nclass OmniSequenceShardCollator:\n    \"\"\"\n    Data collator to chunk inputs along the sequence length.\n    \"\"\"\n\n    # features to slice sequence dimension\n    sp_slice_features: dict[str, int] = field(\n        default_factory=lambda: {\n            \"input_ids\": -1,\n            \"labels\": -1,\n            \"pixel_values\": 0,\n            \"pixel_values_videos\": 0,\n        },\n        metadata={\"help\": \"features to slice sequence dimension.\"},\n    )\n\n    # features to padding sequence dimension\n    padding_features: dict[str, int] = field(\n        default_factory=lambda: {\n            \"pixel_values\": 0,\n        },\n        metadata={\"help\": \"features to padding sequence dimension.\"},\n    )\n\n    # padding scale for padding features\n    padding_scale: dict[str, int] = field(\n        default_factory=lambda: {\"pixel_values\": 4}, metadata={\"help\": \"padding scale for padding features.\"}\n    )\n\n    def __post_init__(self):\n        self.sp_size = parallel_state.get_parallel_state().sp_size\n        self.sp_rank = parallel_state.get_parallel_state().sp_rank\n\n    def sp_slice(self, feature: torch.Tensor, dim: int = -1) -> dict[str, \"torch.Tensor\"]:\n        seq_length = feature.size(dim)\n        sp_chunk_size = (seq_length + self.sp_size - 1) // self.sp_size\n        return feature.narrow(dim, self.sp_rank * sp_chunk_size, sp_chunk_size)\n\n    def sp_padding(\n        self, tensor: \"torch.Tensor\", dim: int = -1, pad_value: int = 0, pad_scale: int = 1\n    ) -> \"torch.Tensor\":\n        \"\"\"\n        Pads a tensor with pad_length to aligns tensor with sp size.\n        \"\"\"\n        seq_length = tensor.size(dim)\n        scale_sp_size = self.sp_size * pad_scale\n\n        sp_chunk_size = (seq_length + scale_sp_size - 1) // scale_sp_size\n        pad_size = sp_chunk_size * scale_sp_size - seq_length\n        if pad_size == 0:\n            return tensor\n\n        pad_shape = list(tensor.shape)\n        pad_shape[dim] = pad_size\n        pad = torch.full(pad_shape, fill_value=pad_value, dtype=tensor.dtype, device=tensor.device)\n        return torch.cat((tensor, pad), dim=dim)\n\n    def __call__(self, batch: Sequence[dict[str, \"torch.Tensor\"]]) -> dict[str, \"torch.Tensor\"]:\n        for key in batch.keys():\n            if key in self.padding_features.keys():\n                batch[key] = self.sp_padding(\n                    batch[key],\n                    dim=self.sp_slice_features.get(key, -1),\n                    pad_value=self.padding_features[key],\n                    pad_scale=self.padding_scale.get(key, 1),\n                )\n\n        # sp slice\n        for key in batch.keys():\n            if key in self.sp_slice_features.keys():\n                batch[key] = self.sp_slice(batch[key], dim=self.sp_slice_features[key])\n\n        return batch\n\n\n@EngineRegistry.register(model_type=\"language_model\", backend=[\"veomni\"], device=[\"cuda\", \"npu\"])\nclass VeOmniEngineWithLMHead(VeOmniEngine, FSDPEngineWithLMHead):\n    def prepare_model_inputs(self, micro_batch: TensorDict):\n        # TODO: Cannot work properly for qwen_vl ulysses\n        model_inputs, output_args = super().prepare_model_inputs(micro_batch)\n        input_ids_rmpad = model_inputs[\"input_ids\"]\n        if self.module.config.model_type in VL_TYPE2INDEX.keys():\n            image_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type][\"IMAGE_INPUT_INDEX\"]\n            video_mask = input_ids_rmpad == VL_TYPE2INDEX[self.module.config.model_type][\"VIDEO_INPUT_INDEX\"]\n            model_inputs.update({\"image_mask\": image_mask, \"video_mask\": video_mask})\n\n            if parallel_state.get_parallel_state().sp_enabled:\n                omni_sequence_shard_collator = OmniSequenceShardCollator()\n                omni_sequence_shard_collator(model_inputs)\n\n        return model_inputs, output_args\n"
  },
  {
    "path": "verl/workers/engine/veomni/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\n\nfrom verl.utils.device import get_device_id, get_torch_device\n\nVL_TYPE2INDEX = {\n    \"qwen2_5_vl\": {\n        \"IMAGE_INPUT_INDEX\": 151655,\n        \"VIDEO_INPUT_INDEX\": 151656,\n    },\n    \"qwen3_vl\": {\n        \"IMAGE_INPUT_INDEX\": 151655,\n        \"VIDEO_INPUT_INDEX\": 151656,\n    },\n    \"qwen3_vl_moe\": {\n        \"IMAGE_INPUT_INDEX\": 151655,\n        \"VIDEO_INPUT_INDEX\": 151656,\n    },\n}\n\n\n@torch.no_grad()\ndef offload_veomni_model_to_cpu(model, empty_cache: bool = True):\n    from torch.distributed.fsdp._fully_shard._fsdp_common import TrainingState\n    from torch.distributed.fsdp._fully_shard._fsdp_state import _get_module_fsdp_state\n\n    for module in model.modules():\n        state = _get_module_fsdp_state(module)\n        if state is None:\n            continue\n        fsdp_param_group = state._fsdp_param_group\n\n        if fsdp_param_group is None:\n            continue\n\n        fsdp_param_group._training_state = TrainingState.IDLE\n\n    model.reshard()\n    model.cpu()\n    if empty_cache:\n        get_torch_device().empty_cache()\n\n\n@torch.no_grad()\ndef load_veomni_model_to_gpu(model):\n    device = get_device_id()\n    model.to(device)\n\n\n@torch.no_grad()\ndef offload_veomni_optimizer(optimizer):\n    optimizers = []\n    # Check if this is a MultiOptimizer (for ep and non-ep parameters when ep+fsdp2 is enabled)\n    if hasattr(optimizer, \"_is_multi_optimizer\") and optimizer._is_multi_optimizer:\n        optimizers.extend(optimizer.optimizers_dict.values())\n    else:\n        optimizers.append(optimizer)\n\n    for opt in optimizers:\n        if not opt.state:\n            continue\n        for param_group in opt.param_groups:\n            for param in param_group[\"params\"]:\n                state = opt.state[param]\n                for key, value in state.items():\n                    if isinstance(value, torch.Tensor):\n                        state[key] = value.to(\"cpu\", non_blocking=True)\n\n\n@torch.no_grad()\ndef load_veomni_optimizer(optimizer, device_id):\n    optimizers = []\n    # Check if this is a MultiOptimizer (for ep and non-ep parameters when ep+fsdp2 is enabled)\n    if hasattr(optimizer, \"_is_multi_optimizer\") and optimizer._is_multi_optimizer:\n        optimizers.extend(optimizer.optimizers_dict.values())\n    else:\n        optimizers.append(optimizer)\n\n    for opt in optimizers:\n        if not opt.state:\n            continue\n        for param_group in opt.param_groups:\n            for param in param_group[\"params\"]:\n                state = opt.state[param]\n                for key, value in state.items():\n                    if isinstance(value, torch.Tensor):\n                        state[key] = value.to(device_id, non_blocking=True)\n\n\ndef _map_moe_params_qwen3_moe(name, tensor):\n    for i in range(tensor.size(0)):\n        new_key = name.replace(\"mlp.experts.\", f\"mlp.experts.{i}.\") + \".weight\"\n        yield new_key, tensor[i].to(get_device_id(), non_blocking=True)\n\n\nMOE_PARAM_HANDERS = {\n    \"qwen3_moe\": _map_moe_params_qwen3_moe,\n}\n"
  },
  {
    "path": "verl/workers/engine_workers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport functools\nimport logging\nimport os\nfrom contextlib import nullcontext\nfrom copy import deepcopy\nfrom functools import partial\nfrom itertools import chain\n\nimport torch\nfrom codetiming import Timer\nfrom omegaconf import DictConfig, open_dict\nfrom tensordict import NonTensorData, TensorDict\nfrom torch.distributed.device_mesh import init_device_mesh\n\nfrom verl.checkpoint_engine import CheckpointEngineRegistry\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import get_device_name, set_expandable_segments\nfrom verl.utils.distributed import initialize_global_process_group_ray\nfrom verl.utils.flops_counter import FlopsCounter\nfrom verl.utils.memory_utils import aggressive_empty_cache\nfrom verl.utils.metric.utils import Metric\nfrom verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig, log_gpu_memory_usage\nfrom verl.utils.py_functional import append_to_dict\nfrom verl.utils.tensordict_utils import maybe_fix_3d_position_ids\nfrom verl.utils.torch_functional import allgather_dict_into_dict\nfrom verl.workers.config import ActorConfig, HFModelConfig, MtpConfig, RolloutConfig, TrainingWorkerConfig\nfrom verl.workers.rollout.base import BaseRollout, get_rollout_class\nfrom verl.workers.utils.losses import ppo_loss\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef _with_routing_replay_flag(enabled: bool):\n    \"\"\"Decorator to set 'enable_routing_replay' flag on the data TensorDict.\"\"\"\n\n    def decorator(func):\n        @functools.wraps(func)\n        def wrapper(self, data: TensorDict, *args, **kwargs):\n            if self.enable_routing_replay:\n                tu.assign_non_tensor_data(data, \"enable_routing_replay\", enabled)\n            return func(self, data, *args, **kwargs)\n\n        return wrapper\n\n    return decorator\n\n\nclass TrainingWorker(Worker, DistProfilerExtension):\n    \"\"\"\n    TrainingWorker provides a Tinker-like API (https://thinkingmachines.ai/tinker/) as a RayWorkerGroup\n    to a single controller. Currently, we only provide more coarse grained APIs,\n    and do not provide exact APIs as Tinker does. But this can be added in the future.\n    \"\"\"\n\n    def __init__(self, config: TrainingWorkerConfig):\n        Worker.__init__(self)\n\n        from verl.workers.engine import BaseEngine, EngineRegistry\n\n        initialize_global_process_group_ray(timeout_second=None)\n\n        self.config = config\n        self.model_config = self.config.model_config\n        self.engine_config = self.config.engine_config\n        self.optimizer_config = self.config.optimizer_config\n        self.checkpoint_config = self.config.checkpoint_config\n        self.device_name = get_device_name()\n\n        if self.engine_config is None:\n            assert self.optimizer_config is None\n            if self.config.auto_select_engine_optim_fn is None:\n                raise ValueError(\n                    \"engine_config is not provided and auto_select_engine_optim_fn is not set. \"\n                    \"Cannot determine engine backend.\"\n                )\n            # Support automatically select engine backend given model config\n            self.engine_config, self.optimizer_config = self.config.auto_select_engine_optim_fn(\n                self.model_config, self.device_name\n            )\n\n        # we use the one defined in model\n        # TODO: this is not elegant and should refactor later\n        self.engine_config.use_remove_padding = self.model_config.use_remove_padding\n        self.engine_config.use_fused_kernels = self.model_config.use_fused_kernels\n\n        # TODO: add DistProfilerExtension\n        self.profiler_config = self.config.profiler_config\n        if self.profiler_config is not None:\n            self.profiler_tool_config = self.profiler_config.tool_config.get(self.profiler_config.tool, {})\n        else:\n            self.profiler_tool_config = None\n\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=self.profiler_config, tool_config=self.profiler_tool_config)\n        )\n\n        self.engine: BaseEngine = EngineRegistry.new(\n            model_type=self.config.model_type,\n            backend=self.engine_config.strategy,\n            model_config=self.model_config,\n            engine_config=self.engine_config,\n            optimizer_config=self.optimizer_config,\n            checkpoint_config=self.checkpoint_config,\n        )\n\n        # build dispatch info\n        self._register_dispatch_collect_info(\n            mesh_name=\"train\",\n            dp_rank=self.engine.get_data_parallel_rank(),\n            is_collect=self.engine.is_mp_src_rank_with_outputs(),\n        )\n\n        self.flops_counter = FlopsCounter(self.model_config.hf_config)\n\n        self.loss_fn = None\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def to(self, device, model=True, optimizer=True, grad=True):\n        \"\"\"Manual control of load/offload\"\"\"\n        assert device in [\"cpu\", \"device\"]\n\n        if device == \"device\":\n            device = get_device_name()\n\n        self.engine.to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def set_loss_fn(self, loss_fn):\n        self.loss_fn = loss_fn\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def reset(self):\n        \"\"\"\n        Reset the model engine to the initial state. If the engine is not initialized,\n        we initialize it. Otherwise, reload ckpt and reset states\n        \"\"\"\n        self.engine.initialize()\n\n    def _postprocess_output(self, output, *, global_token_num, delta_time, forward_only, images_seqlens):\n        \"\"\"\n\n        Args:\n            output: a dictionary containing loss, model_outputs and metrics\n\n        Returns:\n\n        \"\"\"\n        # TODO: whether to log memory\n        # metrics[\"perf/max_memory_allocated_gb\"] = get_torch_device().max_memory_allocated() / (1024 ** 3)\n        # metrics[\"perf/max_memory_reserved_gb\"] = get_torch_device().max_memory_reserved() / (1024 ** 3)\n        # metrics[\"perf/cpu_memory_used_gb\"] = psutil.virtual_memory().used / (1024 ** 3)\n\n        metrics: dict = output.pop(\"metrics\")\n        # perform all gather in dp group to ensure that it's correct.\n        # Here each metric in metrics can be a list (micro-batch metrics) or a singleton\n        # we should always sum the loss of each micro-batch as we scale by global_bsz/global_token\n        loss = torch.sum(torch.tensor(output.pop(\"loss\"), device=self.device_name))\n        dp_group = self.engine.get_data_parallel_group()\n        if dp_group is not None:\n            torch.distributed.all_reduce(loss, op=torch.distributed.ReduceOp.AVG, group=dp_group)\n        loss = loss.item()\n\n        # For grad_norm, we do not perform all reduce because it is already been done when clipping grad\n        grad_norm = metrics.pop(\"grad_norm\", None)\n        lr = metrics.pop(\"lr\", None)\n\n        # For other metrics, we perform all gather in dp group (only if DP > 1)\n        if dp_group is not None:\n            final_metrics = allgather_dict_into_dict(data=metrics, group=dp_group)\n        else:\n            final_metrics = metrics\n        final_metrics[\"loss\"] = loss\n        if grad_norm is not None:\n            final_metrics[\"grad_norm\"] = grad_norm\n        if lr is not None:\n            final_metrics[\"lr\"] = lr\n\n        # TODO: confirm the mtp loss IS same across dp\n        for k, v in final_metrics.items():\n            if k.startswith(\"mtp_losses\"):\n                flatten_v = [sublist[0] for sublist in v]  # sublist should be single element\n                final_metrics[k] = sum(flatten_v) / len(flatten_v)\n        # compute mfu\n        if global_token_num is not None:\n            estimated_flops, promised_flops = self.flops_counter.estimate_flops(\n                global_token_num, delta_time, images_seqlens=images_seqlens\n            )\n            final_metrics[\"mfu\"] = estimated_flops / promised_flops / torch.distributed.get_world_size()\n            if forward_only:\n                final_metrics[\"mfu\"] /= 3.0\n        # model outputs\n        model_output = output.pop(\"model_output\", {})\n        # We only return final_metrics\n        final_output = tu.get_tensordict(tensor_dict=model_output, non_tensor_dict={\"metrics\": final_metrics})\n        return final_output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"), blocking=False)\n    def train_mini_batch(self, data: TensorDict) -> TensorDict:\n        \"\"\"Split a batch into N mini-batches run for multiple epochs\n\n        Args:\n            data:\n\n        Returns:\n\n        \"\"\"\n        maybe_fix_3d_position_ids(data)\n        batch_size_per_dp = data.shape[0]\n        disable_auto_offload = tu.pop(data, key=\"disable_auto_offload\", default=False)\n        mini_batch_size = tu.pop(data, key=\"mini_batch_size\", default=None)\n        num_mini_batch = tu.pop(data, key=\"num_mini_batch\", default=None)\n        epochs = tu.pop(data, key=\"epochs\", default=1)\n        seed = tu.pop(data, key=\"seed\", default=42)\n        dataloader_kwargs = tu.pop(data, key=\"dataloader_kwargs\", default={})\n\n        assert mini_batch_size is not None or num_mini_batch is not None\n\n        if mini_batch_size is None:\n            assert batch_size_per_dp % num_mini_batch == 0, f\"Got {batch_size_per_dp=} and {num_mini_batch=}\"\n            mini_batch_size_per_gpu = batch_size_per_dp // num_mini_batch\n        else:\n            assert mini_batch_size % self.engine.get_data_parallel_size() == 0, (\n                f\"Got {mini_batch_size=} and {self.engine.get_data_parallel_size()=}\"\n            )\n            mini_batch_size_per_gpu = mini_batch_size // self.engine.get_data_parallel_size()\n\n        # make iterator\n        dataloader = tu.make_iterator(\n            data,\n            mini_batch_size=mini_batch_size_per_gpu,\n            epochs=epochs,\n            seed=seed + self.engine.get_data_parallel_rank(),\n            dataloader_kwargs=dataloader_kwargs,\n        )\n\n        with (\n            self.engine.train_mode(disable_auto_offload=disable_auto_offload),\n            Timer(name=\"train_batch\", logger=None),\n        ):\n            # update\n            output_lst = []\n            total_num_iterations = data.shape[0] // mini_batch_size_per_gpu * epochs\n\n            for batch_idx, mini_batch_td in enumerate(dataloader):\n                # add global token num\n                global_token_num = mini_batch_td[\"input_ids\"].offsets().diff().tolist()  # (total_nnz,)\n                # allgather from dp rank\n                global_token_num_output = [None] * self.engine.get_data_parallel_size()\n                torch.distributed.all_gather_object(\n                    global_token_num_output, global_token_num, self.engine.get_data_parallel_group()\n                )\n                global_token_num = [x for xs in global_token_num_output for x in xs]\n                tu.assign_non_tensor(\n                    mini_batch_td,\n                    global_token_num=NonTensorData(global_token_num),\n                    update_lr_scheduler=batch_idx == total_num_iterations - 1,\n                    disable_auto_offload=True,\n                )\n                actor_output = self.train_batch(mini_batch_td)\n                output_lst.append(actor_output)\n\n            if self.engine.is_mp_src_rank_with_outputs():\n                actor_output = [tu.get(output, \"metrics\") for output in output_lst]\n                metrics = {}\n                for output in actor_output:\n                    for key, val in output.items():\n                        # flattn dp and micro batch\n                        if isinstance(val, list):\n                            output[key] = (\n                                Metric.aggregate_dp(val)\n                                if isinstance(val[0], Metric)\n                                else list(chain.from_iterable(val))\n                            )\n                    append_to_dict(metrics, output)\n\n                output = tu.get_tensordict(tensor_dict={}, non_tensor_dict={\"metrics\": metrics}).cpu()\n            else:\n                output = None\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"), blocking=False)\n    def train_batch(self, data: TensorDict) -> TensorDict:\n        assert self.loss_fn is not None, \"loss function can't be None when calling train_batch\"\n        assert not self.engine_config.forward_only, \"Can't run `train_batch` when forward_only is in the engine config.\"\n        # global_token_num should be a list of number of tokens of each seq in this batch\n        global_token_num = tu.get(data, key=\"global_token_num\")\n        disable_auto_offload = tu.get(data, key=\"disable_auto_offload\", default=False)\n        images_seqlens = tu.get(data, key=\"images_seqlens\", default=None)\n\n        # inject engineering parameters if not specified\n        default_keys = dict(\n            use_remove_padding=self.model_config.use_remove_padding,\n            use_dynamic_bsz=self.engine_config.use_dynamic_bsz,\n            max_token_len_per_gpu=self.engine_config.max_token_len_per_gpu,\n            micro_batch_size_per_gpu=self.engine_config.micro_batch_size_per_gpu,\n            use_fused_kernels=self.engine_config.use_fused_kernels,\n        )\n\n        for key, val in default_keys.items():\n            if key not in data.keys():\n                tu.assign_non_tensor(data, **{key: val})\n\n        with (\n            self.engine.train_mode(disable_auto_offload=disable_auto_offload),\n            Timer(name=\"train_batch\", logger=None) as timer,\n        ):\n            output = self.engine.train_batch(data, loss_function=self.loss_fn)\n            # containing loss, model_output and metrics\n            # for training, we only care about loss and metrics\n        delta_time = timer.last\n\n        update_lr_scheduler = tu.get(data, key=\"update_lr_scheduler\", default=False)\n        # update lr scheduler\n        if update_lr_scheduler:\n            lr = self.engine.lr_scheduler_step()\n        else:\n            lr = None\n\n        if self.engine.is_mp_src_rank_with_outputs():\n            # we don't need model_output in training. Maybe we change out mind later\n            output.pop(\"model_output\")\n            if lr is not None:\n                output[\"metrics\"][\"lr\"] = lr\n            final_output = self._postprocess_output(\n                output,\n                global_token_num=global_token_num,\n                delta_time=delta_time,\n                forward_only=False,\n                images_seqlens=images_seqlens,\n            ).cpu()\n        else:\n            final_output = None\n\n        return final_output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"train\"), blocking=False)\n    def infer_batch(self, data: TensorDict) -> TensorDict:\n        # add mfu calculator\n        global_token_num = tu.get(data, key=\"global_token_num\")\n        compute_loss = tu.get(data, key=\"compute_loss\", default=True)\n        disable_auto_offload = tu.get(data, key=\"disable_auto_offload\", default=False)\n        no_lora_adapter = tu.pop(data, key=\"no_lora_adapter\", default=False)\n        images_seqlens = tu.get(data, key=\"images_seqlens\", default=None)\n\n        default_keys = dict(\n            use_remove_padding=self.model_config.use_remove_padding,\n            use_dynamic_bsz=self.engine_config.use_dynamic_bsz,\n            max_token_len_per_gpu=self.engine_config.infer_max_token_len_per_gpu,\n            micro_batch_size_per_gpu=self.engine_config.infer_micro_batch_size_per_gpu,\n            use_fused_kernels=self.engine_config.use_fused_kernels,\n        )\n\n        for key, val in default_keys.items():\n            if key not in data.keys():\n                tu.assign_non_tensor(data, **{key: val})\n\n        # for sft training, we need to compute loss in eval\n        loss_function = self.loss_fn if compute_loss else None\n\n        with (\n            self.engine.eval_mode(disable_auto_offload=disable_auto_offload),\n            Timer(name=\"eval_batch\", logger=None) as timer,\n        ):\n            adapter_ctx = self.engine.disable_adapter() if no_lora_adapter else nullcontext()\n            with adapter_ctx:\n                output = self.engine.infer_batch(data, loss_function=loss_function)\n        delta_time = timer.last\n\n        if self.engine.is_mp_src_rank_with_outputs():\n            final_output = self._postprocess_output(\n                output,\n                global_token_num=global_token_num,\n                delta_time=delta_time,\n                forward_only=True,\n                images_seqlens=images_seqlens,\n            ).cpu()\n        else:\n            final_output = None\n\n        return final_output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):\n        return self.engine.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False):\n        return self.engine.load_checkpoint(local_path, hdfs_path, del_local_after_load)\n\n\nclass ActorRolloutRefWorker(Worker, DistProfilerExtension):\n    \"\"\"Hybrid worker that includes actor model, rollout and optional ref model.\n    For standalone actor or rollout, use ActorWorker or BaseRollout respectively.\n\n    NOTE: ActorRolloutRefWorker no longer support spmd mode and run native server mode.\n    \"\"\"\n\n    def __init__(self, config: DictConfig, role: str, **kwargs):\n        Worker.__init__(self)\n        self.config = config\n        self.role = role\n        self.actor: TrainingWorker = None\n        self.ref: TrainingWorker = None\n        self.rollout: BaseRollout = None\n        assert self.role in [\"actor\", \"rollout\", \"ref\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_actor = self.role in [\"actor\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_rollout = self.role in [\"rollout\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_ref = self.role in [\"ref\", \"actor_rollout_ref\"]\n\n        if self._is_actor:\n            omega_profiler_config = config.actor.get(\"profiler\", {})\n        elif self._is_rollout:\n            # NOTE: In colocation mode, rollout config may not take effect (follow the actor config)\n            # This is for extendability in AsyncRL cases\n            omega_profiler_config = config.rollout.get(\"profiler\", {})\n        else:\n            omega_profiler_config = config.ref.get(\"profiler\", {})\n\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n\n        self.enable_routing_replay = (\n            self.config.actor.strategy == \"megatron\" and self.config.actor.megatron.router_replay.mode != \"disabled\"\n        )\n\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def set_loss_fn(self, loss_fn):\n        self.actor.set_loss_fn(loss_fn=loss_fn)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def to(self, device, model=True, optimizer=True, grad=True):\n        \"\"\"Manual control of load/offload\"\"\"\n        self.actor.to(device=device, model=model, optimizer=optimizer, grad=grad)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model)\n\n        # 1. build reference model\n        if \"ref\" in self.role:\n            # TODO: align ref config with actor config\n            with open_dict(self.config.ref):\n                self.config.ref.ppo_mini_batch_size = self.config.actor.ppo_mini_batch_size\n                self.config.ref.ppo_micro_batch_size = self.config.ref.pop(\"log_prob_micro_batch_size\", None)\n                self.config.ref.ppo_micro_batch_size_per_gpu = self.config.ref.pop(\n                    \"log_prob_micro_batch_size_per_gpu\", None\n                )\n                self.config.ref.use_dynamic_bsz = self.config.ref.pop(\"log_prob_use_dynamic_bsz\", False)\n                self.config.ref.ppo_max_token_len_per_gpu = self.config.ref.pop(\"log_prob_max_token_len_per_gpu\", None)\n            ref_config: ActorConfig = omega_conf_to_dataclass(self.config.ref)\n\n            # The ref model does not need to enable MTP; force it to false.\n            ref_config.model_config = deepcopy(model_config)\n            ref_config.model_config.mtp = MtpConfig(enable=False)\n\n            # construct TrainingWorkerConfig\n            ref_training_config = TrainingWorkerConfig(\n                model_type=\"language_model\",\n                model_config=ref_config.model_config,\n                engine_config=ref_config.engine,\n                optimizer_config=ref_config.optim,\n                checkpoint_config=ref_config.checkpoint,\n            )\n\n            # assign engine configs\n            ref_training_config.engine_config.use_dynamic_bsz = self.config.ref.use_dynamic_bsz\n            ref_training_config.engine_config.infer_max_token_len_per_gpu = self.config.ref.ppo_max_token_len_per_gpu\n            ref_training_config.engine_config.infer_micro_batch_size_per_gpu = (\n                self.config.ref.ppo_micro_batch_size_per_gpu\n            )\n            ref_training_config.engine_config.use_remove_padding = model_config.use_remove_padding\n\n            self.ref = TrainingWorker(config=ref_training_config)\n            self.ref.reset()\n            self.set_dispatch_collect(mesh_name=\"ref\", **self.ref.get_dispatch_collect())\n\n        # 2. build actor model\n        if \"actor\" in self.role:\n            actor_config: ActorConfig = omega_conf_to_dataclass(self.config.actor)\n            actor_config.model_config = model_config\n            actor_training_config = TrainingWorkerConfig(\n                model_type=\"language_model\",\n                model_config=actor_config.model_config,\n                engine_config=actor_config.engine,\n                optimizer_config=actor_config.optim,\n                checkpoint_config=actor_config.checkpoint,\n            )\n\n            assert self.config.actor.use_dynamic_bsz == self.config.rollout.log_prob_use_dynamic_bsz\n\n            # assign engine configs\n            actor_training_config.engine_config.use_dynamic_bsz = self.config.actor.use_dynamic_bsz\n            actor_training_config.engine_config.infer_max_token_len_per_gpu = (\n                self.config.rollout.log_prob_max_token_len_per_gpu\n            )\n            actor_training_config.engine_config.infer_micro_batch_size_per_gpu = (\n                self.config.rollout.log_prob_micro_batch_size_per_gpu\n            )\n            actor_training_config.engine_config.max_token_len_per_gpu = self.config.actor.ppo_max_token_len_per_gpu\n            actor_training_config.engine_config.micro_batch_size_per_gpu = (\n                self.config.actor.ppo_micro_batch_size_per_gpu\n            )\n            actor_training_config.engine_config.use_remove_padding = model_config.use_remove_padding\n\n            if self.config.actor.use_dynamic_bsz:\n                assert self.config.rollout.log_prob_max_token_len_per_gpu is not None\n                assert self.config.actor.ppo_max_token_len_per_gpu is not None\n            else:\n                assert self.config.rollout.log_prob_micro_batch_size_per_gpu is not None\n                assert self.config.actor.ppo_micro_batch_size_per_gpu is not None\n\n            self.loss_fn = partial(ppo_loss, config=actor_config)\n            self.actor = TrainingWorker(config=actor_training_config)\n            self.actor.reset()\n            self.actor.set_loss_fn(self.loss_fn)\n            self.set_dispatch_collect(mesh_name=\"actor\", **self.actor.get_dispatch_collect())\n\n        # 3. build rollout engine\n        if \"rollout\" in self.role:\n            rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout)\n\n            # TODO: move rollout_device_mesh into ServerAdapter\n            # 3.1 build rollout device mesh (sglang need only)\n            infer_tp = rollout_config.tensor_model_parallel_size * rollout_config.data_parallel_size\n            infer_pp = rollout_config.pipeline_model_parallel_size\n            infer_world_size = infer_tp * infer_pp\n            dp = self.world_size // infer_world_size\n            assert self.world_size % infer_world_size == 0, (\n                f\"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}\"\n            )\n            rollout_device_mesh = init_device_mesh(\n                get_device_name(), mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=[\"dp\", \"infer_tp\", \"infer_pp\"]\n            )\n\n            # 3.2 initialize rollout engine\n            rollout_cls: type[BaseRollout] = get_rollout_class(rollout_config.name, rollout_config.mode)\n            self.rollout = rollout_cls(\n                config=rollout_config, model_config=model_config, device_mesh=rollout_device_mesh\n            )\n\n            # used for LoRA\n            self.base_sync_done: bool = \"dummy\" not in self.config.rollout.load_format\n            self.layered_summon = self.config.rollout.get(\"layered_summon\", False)\n            self.peft_merge: bool = model_config.lora.get(\"merge\", False)\n\n        # 4. build checkpoint engine\n        if \"actor\" in self.role:\n            checkpoint_engine_config = omega_conf_to_dataclass(self.config.rollout.checkpoint_engine)\n            backend = checkpoint_engine_config.backend\n            bucket_size = checkpoint_engine_config.update_weights_bucket_megabytes << 20\n            engine_kwargs = checkpoint_engine_config.engine_kwargs.get(backend, {})\n            self.checkpoint_engine = CheckpointEngineRegistry.new(\n                backend, is_master=(torch.distributed.get_rank() == 0), bucket_size=bucket_size, **engine_kwargs\n            )\n\n        # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo\n        aggressive_empty_cache(force_sync=True)\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"ref\"))\n    @DistProfiler.annotate(color=\"olive\", role=\"ref_compute_log_prob\")\n    @_with_routing_replay_flag(enabled=False)\n    def compute_ref_log_prob(self, data: TensorDict) -> TensorDict:\n        output = self.ref.infer_batch(data=data)\n        return output.cpu() if output is not None else None\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @DistProfiler.annotate(color=\"blue\", role=\"actor_compute_log_prob\")\n    @_with_routing_replay_flag(enabled=True)\n    def compute_log_prob(self, data: TensorDict) -> TensorDict:\n        output = self.actor.infer_batch(data)\n\n        return output.cpu() if output is not None else None\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @DistProfiler.annotate(color=\"red\", role=\"actor_update\")\n    @_with_routing_replay_flag(enabled=True)\n    def update_actor(self, data: TensorDict) -> TensorDict:\n        output = self.actor.train_mini_batch(data=data)\n        return output.cpu() if output is not None else None\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False):\n        assert \"actor\" in self.role, \"load_checkpoint only support actor role\"\n        self.actor.load_checkpoint(local_path, hdfs_path, del_local_after_load)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):\n        assert \"actor\" in self.role, \"save_checkpoint only support actor role\"\n        self.actor.save_checkpoint(local_path, hdfs_path, global_step, max_ckpt_to_keep)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    async def update_weights(self, global_steps: int = None):\n        \"\"\"Update weights from trainer to rollout.\n\n        1. For sync training with colocated trainer and rollout, update rollout directly from model engine.\n           - before update_weights: rollout should be in sleep mode.\n           - after update_weights: rollout should be in wake_up mode.\n        2. For async training with disaggregated trainer and rollout, send_weights only by checkpoint engine.\n        \"\"\"\n\n        # 0. send_weights only for async training with disaggregated trainer and rollout\n        if self.config.rollout.checkpoint_engine.backend != \"naive\":\n            per_tensor_param, _ = self.actor.engine.get_per_tensor_param()\n            await self.checkpoint_engine.send_weights(per_tensor_param)\n            return\n\n        set_expandable_segments(False)\n        log_gpu_memory_usage(\"Before resume weights\", logger=logger)\n\n        # 1. resume weights and update weights\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"weights\"])\n        log_gpu_memory_usage(\"After resume weights\", logger=logger)\n\n        # 2. get per tensor generator from engine, this will load model to gpu\n        per_tensor_param, peft_config = self.actor.engine.get_per_tensor_param(\n            layered_summon=self.layered_summon, base_sync_done=True\n        )\n\n        await self.rollout.update_weights(\n            per_tensor_param, peft_config=peft_config, base_sync_done=True, global_steps=global_steps\n        )\n\n        do_lora_base_sync = False\n        if not self.peft_merge and peft_config is not None:\n            # set sleep level for LoRA adapter weights only sync\n            # TODO: make this configurable so that users with small\n            # main memory can trade sync time to avoid OOM\n            self.rollout.sleep_level = 1\n\n            do_lora_base_sync = (not self.base_sync_done) or (\n                self.rollout.sleep_level != 1 and self.config.rollout.free_cache_engine\n            )\n\n        if do_lora_base_sync:\n            per_tensor_base_params, _ = self.actor.engine.get_per_tensor_param(\n                layered_summon=self.layered_summon, base_sync_done=False\n            )\n            await self.rollout.update_weights(per_tensor_base_params, peft_config=peft_config, base_sync_done=False)\n\n        log_gpu_memory_usage(\"After update_weights\", logger=logger)\n\n        # 3. offload model to cpu\n        self.actor.engine.to(\"cpu\", model=True, optimizer=False, grad=False)\n        aggressive_empty_cache(force_sync=True)\n\n        # 4. resume kv_cache\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"kv_cache\"])\n        log_gpu_memory_usage(\"After resume kv_cache\", logger=logger)\n\n        self.base_sync_done = True\n        set_expandable_segments(True)\n\n    @register(dispatch_mode=Dispatch.DP_COMPUTE, blocking=False)\n    def execute_checkpoint_engine(self, method: str, *args, **kwargs):\n        \"\"\"Execute checkpoint engine method.\n\n        Args:\n            method (str): Checkpoint engine method name.\n            *args: Variable length argument list.\n            **kwargs: Arbitrary keyword arguments.\n\n        \"\"\"\n        return getattr(self.checkpoint_engine, method)(*args, **kwargs)\n"
  },
  {
    "path": "verl/workers/fsdp_workers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe main entry point to run the PPO algorithm\n\"\"\"\n\nimport datetime\nimport json\nimport logging\nimport os\nimport warnings\nfrom dataclasses import asdict\n\nimport psutil\nimport torch\nimport torch.distributed\nimport torch.distributed as dist\nfrom codetiming import Timer\nfrom omegaconf import DictConfig, OmegaConf, open_dict\nfrom omegaconf.errors import ConfigAttributeError\nfrom peft import LoraConfig, TaskType, get_peft_model\nfrom safetensors.torch import save_file\nfrom torch.distributed.device_mesh import init_device_mesh\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom torch.distributed.fsdp.api import FullStateDictConfig, ShardedStateDictConfig, StateDictType\n\ntry:\n    # for torch 2.5+\n    from torch.distributed.tensor import DTensor\nexcept ImportError:\n    from torch.distributed._tensor import DTensor\n\nfrom verl import DataProto\nfrom verl.models.transformers.monkey_patch import apply_monkey_patch\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils import hf_processor, hf_tokenizer\nfrom verl.utils.activation_offload import enable_activation_offloading\nfrom verl.utils.checkpoint.fsdp_checkpoint_manager import FSDPCheckpointManager\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import (\n    get_device_id,\n    get_device_name,\n    get_nccl_backend,\n    get_torch_device,\n    set_expandable_segments,\n)\nfrom verl.utils.flops_counter import FlopsCounter\nfrom verl.utils.fs import copy_to_local\nfrom verl.utils.fsdp_utils import (\n    CPUOffloadPolicy,\n    MixedPrecisionPolicy,\n    apply_fsdp2,\n    collect_lora_params,\n    fsdp2_load_full_state_dict,\n    fsdp_version,\n    get_fsdp_wrap_policy,\n    get_init_weight_context_manager,\n    get_shard_placement_fn,\n    init_fn,\n    layered_summon_lora_params,\n    load_fsdp_model_to_gpu,\n    load_fsdp_optimizer,\n    offload_fsdp_model_to_cpu,\n    offload_fsdp_optimizer,\n    replace_lora_wrapper,\n)\nfrom verl.utils.import_utils import import_external_libs\nfrom verl.utils.memory_utils import aggressive_empty_cache\nfrom verl.utils.model import convert_weight_keys\nfrom verl.utils.profiler import DistProfiler, DistProfilerExtension, ProfilerConfig, log_gpu_memory_usage, simple_timer\nfrom verl.utils.profiler.performance import reduce_timing, topk_reduce_ratio_min_max\nfrom verl.utils.py_functional import convert_to_regular_types\n\n# QAT support\nfrom verl.utils.qat import apply_qat, enable_qat_fuse\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.transformers_compat import get_auto_model_for_vision2seq\nfrom verl.workers.config import FSDPCriticConfig, FSDPEngineConfig, HFModelConfig, RolloutConfig\nfrom verl.workers.config.optimizer import build_optimizer\nfrom verl.workers.rollout import get_rollout_class\nfrom verl.workers.sharding_manager.fsdp_ulysses import FSDPUlyssesShardingManager\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\ndevice_name = get_device_name()\n\n\ndef create_device_mesh(world_size, fsdp_size):\n    if fsdp_size < 0 or fsdp_size >= world_size:\n        device_mesh = init_device_mesh(device_name, mesh_shape=(world_size,), mesh_dim_names=[\"fsdp\"])\n    else:\n        device_mesh = init_device_mesh(\n            device_name, mesh_shape=(world_size // fsdp_size, fsdp_size), mesh_dim_names=[\"ddp\", \"fsdp\"]\n        )\n    return device_mesh\n\n\ndef get_sharding_strategy(device_mesh, zero3_enable=True):\n    from torch.distributed.fsdp import ShardingStrategy\n\n    if zero3_enable:\n        fsdp_strategy = ShardingStrategy.FULL_SHARD\n        hsdp_strategy = ShardingStrategy.HYBRID_SHARD\n    else:\n        fsdp_strategy = ShardingStrategy.SHARD_GRAD_OP\n        hsdp_strategy = ShardingStrategy._HYBRID_SHARD_ZERO2\n\n    if device_mesh.ndim == 1:\n        sharding_strategy = fsdp_strategy\n    elif device_mesh.ndim == 2:\n        sharding_strategy = hsdp_strategy\n    else:\n        raise NotImplementedError(f\"Get device mesh ndim={device_mesh.ndim}, but only support 1 or 2\")\n    return sharding_strategy\n\n\ndef get_vl_model_vision_tower(vl_model_instance):\n    \"\"\"\n    Util to extract Vision Tower from a VL model instance\n    \"\"\"\n    if hasattr(vl_model_instance, \"model\") and hasattr(vl_model_instance.model, \"visual\"):\n        # transformers >= 4.52.0\n        return vl_model_instance.model.visual\n    elif hasattr(vl_model_instance, \"visual\"):\n        # transformers < 4.52.0\n        return vl_model_instance.visual\n    return None\n\n\nclass ActorRolloutRefWorker(Worker, DistProfilerExtension):\n    \"\"\"\n    This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy\n    or a hybrid engine based on the config.rollout\n    \"\"\"\n\n    def __init__(self, config: DictConfig, role: str, **kwargs):\n        Worker.__init__(self)\n\n        self.config = config\n        import torch.distributed\n\n        if not torch.distributed.is_initialized():\n            rank = int(os.environ.get(\"RANK\", 0))\n            world_size = int(os.environ.get(\"WORLD_SIZE\", 1))\n            torch.distributed.init_process_group(\n                backend=f\"cpu:gloo,{get_device_name()}:{get_nccl_backend()}\",\n                rank=rank,\n                world_size=world_size,\n                timeout=datetime.timedelta(seconds=self.config.get(\"nccl_timeout\", 600)),\n                init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n            )\n\n        # Apply NPU patches for FSDP backend\n        from verl.workers.engine.fsdp.utils import apply_npu_fsdp_patches\n\n        apply_npu_fsdp_patches()\n\n        # build device mesh for FSDP\n        world_size = torch.distributed.get_world_size()\n        # TODO(sgm): support FSDP hybrid shard for larger model\n        self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=self.config.actor.fsdp_config.fsdp_size)\n\n        # build device mesh for Ulysses Sequence Parallel\n        self.ulysses_device_mesh = None\n        self.ulysses_sequence_parallel_size = self.config.actor.get(\"ulysses_sequence_parallel_size\", 1)\n        dp = world_size // self.ulysses_sequence_parallel_size\n        if self.ulysses_sequence_parallel_size > 1:\n            self.ulysses_device_mesh = init_device_mesh(\n                device_name, mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=[\"dp\", \"sp\"]\n            )\n\n        # create training dispatch\n        if self.ulysses_device_mesh is not None:\n            is_collect = self.ulysses_device_mesh[\"sp\"].get_local_rank() == 0\n            self._register_dispatch_collect_info(\n                \"actor\", dp_rank=self.ulysses_device_mesh[\"dp\"].get_local_rank(), is_collect=is_collect\n            )\n        else:\n            self._register_dispatch_collect_info(\"actor\", dp_rank=self.rank, is_collect=True)\n\n        self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh)\n        self._lora_rank = self.config.model.get(\"lora_rank\", 0)\n        self._is_lora = self.config.model.get(\"lora_adapter_path\") is not None or self._lora_rank > 0\n\n        self.role = role\n        assert self.role in [\"actor\", \"rollout\", \"ref\", \"actor_rollout\", \"actor_rollout_ref\"]\n\n        self._is_actor = self.role in [\"actor\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_rollout = self.role in [\"rollout\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_ref = self.role in [\"ref\", \"actor_rollout_ref\"]\n        self.use_orig_params = self.config.actor.fsdp_config.get(\"use_orig_params\", False)\n\n        # TODO(haibin.lin):\n        # As of now the type of config is DictConfig, if we assign config.profiler with ProfilerConfig,\n        # it will actually convert the ProfilerConfig dataclass back to a DictConfig.\n        # We can still use ProfilerConfig for testing purpose (tests/utils/test_nvtx_profile.py)\n        # as they provides DictConfig-like interface\n        # The benefit of creating the dataclass config is to perform validation during __post_init__\n        if self._is_actor:\n            omega_profiler_config = config.actor.get(\"profiler\", {})\n        elif self._is_rollout:\n            # NOTE: In colocation mode, rollout config may not take effect (follow the actor config)\n            # This is for extendability in AsyncRL cases\n            omega_profiler_config = config.rollout.get(\"profiler\", {})\n        elif self._is_ref:\n            omega_profiler_config = config.ref.get(\"profiler\", {})\n        else:\n            raise ValueError(\n                f\"Invalid role {self.role}, should be one of \"\n                \"['actor', 'rollout', 'ref', 'actor_rollout', 'actor_rollout_ref']\"\n            )\n        # omega_profiler_config is DictConfig\n        # profiler_config is a ProfilerConfig dataclass\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n\n        self._is_offload_param = False\n        self._is_offload_optimizer = False\n        if self._is_actor:\n            self._is_offload_param = self.config.actor.fsdp_config.get(\"param_offload\", False)\n            self._is_offload_optimizer = self.config.actor.fsdp_config.get(\"optimizer_offload\", False)\n        elif self._is_ref:\n            # TODO: it seems that manual offload is slowly than FSDP offload\n            self._is_offload_param = self.config.ref.fsdp_config.get(\"param_offload\", False)\n\n        # normalize config\n        if self._is_actor:\n            self.config.actor.ppo_mini_batch_size *= self.config.rollout.n\n            self.config.actor.ppo_mini_batch_size //= self.device_mesh.size() // self.ulysses_sequence_parallel_size\n            assert self.config.actor.ppo_mini_batch_size > 0, (\n                f\"ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be larger than 0 after \"\n                f\"normalization\"\n            )\n            # micro bsz\n            if self.config.actor.ppo_micro_batch_size is not None:\n                self.config.actor.ppo_micro_batch_size //= (\n                    self.device_mesh.size() // self.ulysses_sequence_parallel_size\n                )\n                self.config.actor.ppo_micro_batch_size_per_gpu = self.config.actor.ppo_micro_batch_size\n\n            if self.config.actor.ppo_micro_batch_size_per_gpu is not None:\n                assert self.config.actor.ppo_mini_batch_size % self.config.actor.ppo_micro_batch_size_per_gpu == 0, (\n                    f\"normalized ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be divisible by \"\n                    f\"ppo_micro_batch_size_per_gpu {self.config.actor.ppo_micro_batch_size_per_gpu}\"\n                )\n                assert self.config.actor.ppo_mini_batch_size // self.config.actor.ppo_micro_batch_size_per_gpu > 0, (\n                    f\"normalized ppo_mini_batch_size {self.config.actor.ppo_mini_batch_size} should be larger than \"\n                    f\"ppo_micro_batch_size_per_gpu {self.config.actor.ppo_micro_batch_size_per_gpu}\"\n                )\n\n        # normalize rollout config\n        if self._is_rollout and self.config.rollout.log_prob_micro_batch_size is not None:\n            self.config.rollout.log_prob_micro_batch_size //= (\n                self.device_mesh.size() // self.ulysses_sequence_parallel_size\n            )\n            self.config.rollout.log_prob_micro_batch_size_per_gpu = self.config.rollout.log_prob_micro_batch_size\n        # normalize ref config\n        if self._is_ref and self.config.ref.log_prob_micro_batch_size is not None:\n            self.config.ref.log_prob_micro_batch_size //= self.device_mesh.size() // self.ulysses_sequence_parallel_size\n            self.config.ref.log_prob_micro_batch_size_per_gpu = self.config.ref.log_prob_micro_batch_size\n\n    def _init_qat_config(self):\n        \"\"\"Initialize QAT configuration from actor.qat.\"\"\"\n        try:\n            self.qat_config = self.config.actor.qat\n            self._qat_enabled = self.qat_config.enable\n            if self._qat_enabled:\n                logger.info(\n                    f\"QAT enabled: mode={self.qat_config.mode}, config_path={self.qat_config.quantization_config_path}\"\n                )\n        except (AttributeError, KeyError, ConfigAttributeError):\n            # QAT config not provided, disable QAT\n            self._qat_enabled = False\n            self.qat_config = None\n\n    def _restore_w4a4_input_scales(self, model, model_path):\n        \"\"\"Restore input_global_scale and input_amax from checkpoint for W4A4 mode.\"\"\"\n        import glob\n\n        from safetensors import safe_open\n\n        safetensor_files = glob.glob(f\"{model_path}/model*.safetensors\")\n        loaded_count = 0\n\n        for sf_path in safetensor_files:\n            with safe_open(sf_path, framework=\"pt\") as f:\n                for key in f.keys():\n                    if \"input_global_scale\" in key:\n                        module_path = key.replace(\".input_global_scale\", \"\")\n                        amax_key = f\"{module_path}.input_amax\"\n\n                        module = model\n                        for part in module_path.split(\".\"):\n                            module = getattr(module, part)\n\n                        scale_val = f.get_tensor(key)\n                        val = scale_val.item() if scale_val.numel() == 1 else scale_val.max().item()\n                        module.input_global_scale.fill_(val)\n\n                        amax_val = f.get_tensor(amax_key)\n                        amax = amax_val.item() if amax_val.numel() == 1 else amax_val.max().item()\n                        module.input_amax.fill_(amax)\n                        loaded_count += 1\n\n        if self.rank == 0:\n            logger.info(f\"[W4A4] Loaded {loaded_count} input scales from checkpoint\")\n\n    def _build_model_optimizer(\n        self,\n        model_path,\n        fsdp_config: FSDPEngineConfig,\n        optim_config,\n        override_model_config,\n        use_remove_padding=False,\n        use_fused_kernels=False,\n        enable_gradient_checkpointing=False,\n        trust_remote_code=False,\n        use_liger=False,\n        role=\"actor\",\n        enable_activation_offload=False,\n        use_prefix_grouper=False,\n        use_tiled_mlp=False,\n        tiled_mlp_shards=4,\n    ):\n        from torch.distributed.fsdp import CPUOffload, MixedPrecision\n        from transformers import (\n            AutoConfig,\n            AutoModel,\n            AutoModelForCausalLM,\n        )\n\n        try:\n            from transformers import AutoModelForVision2Seq\n        except ImportError:\n            AutoModelForVision2Seq = None\n        try:\n            from transformers import AutoModelForImageTextToText\n        except ImportError:\n            AutoModelForImageTextToText = AutoModelForVision2Seq\n\n        from verl.utils.model import get_generation_config, print_model_size, update_model_config\n        from verl.utils.torch_dtypes import PrecisionType\n\n        AutoModelForVision2Seq = get_auto_model_for_vision2seq()\n\n        assert role in [\"actor\", \"ref\"]\n\n        # TiledMLP requires FSDP2 for correct gradient computation\n        if use_tiled_mlp and self.config.actor.strategy == \"fsdp\":\n            raise ValueError(\"TiledMLP requires FSDP2. Set `actor_rollout_ref.actor.strategy=fsdp2`.\")\n\n        log_gpu_memory_usage(f\"Before init {role} from HF AutoModel\", logger=logger)\n        local_path = model_path\n\n        # note that we have to create model in fp32. Otherwise, the optimizer is in bf16, which is incorrect\n        # TODO(zhangchi.usc1992): 1. support create from random initialized model. 2. Support init with FSDP directly\n        self.tokenizer = hf_tokenizer(local_path, trust_remote_code=trust_remote_code)\n        self.processor = hf_processor(local_path, trust_remote_code=trust_remote_code)\n\n        if self.config.model.get(\"custom_chat_template\", None) is not None:\n            if self.processor is not None:\n                self.processor.chat_template = self.config.model.custom_chat_template\n            else:\n                self.tokenizer.chat_template = self.config.model.custom_chat_template\n\n        torch_dtype = fsdp_config.get(\"model_dtype\", None)\n        if torch_dtype is None:\n            torch_dtype = torch.float32 if self._is_actor else torch.bfloat16\n        else:\n            torch_dtype = PrecisionType.to_dtype(torch_dtype)\n\n        # override model kwargs\n        attn_implementation = override_model_config.get(\"attn_implementation\", \"flash_attention_2\")\n        actor_model_config = AutoConfig.from_pretrained(\n            local_path, trust_remote_code=trust_remote_code, attn_implementation=attn_implementation\n        )\n        # TODO: VL models use VisionAttention, which directly uses flash_attention in transformers>=4.53\n        # which will be patched by _ulysses_flash_attention_forward, but errorly misses position_ids\n        # Maybe support Ulysses in VisionAttention in the future and remove this patch\n        if self.ulysses_sequence_parallel_size > 1 and hasattr(actor_model_config, \"vision_config\"):\n            actor_model_config.vision_config._attn_implementation = \"eager\"\n\n        # patch for qwen2.5-vl: when using flash_attention_3, set vision tower to use flash_attention_2\n        # because the vision tower does not support flash_attention_3\n        if (\n            getattr(actor_model_config, \"model_type\", None) == \"qwen2_5_vl\"\n            and attn_implementation == \"flash_attention_3\"\n            and hasattr(actor_model_config, \"vision_config\")\n        ):\n            actor_model_config.vision_config._attn_implementation = \"flash_attention_2\"\n\n        # patch for kimi-vl\n        if getattr(actor_model_config, \"model_type\", None) == \"kimi_vl\":\n            actor_model_config.text_config.topk_method = \"greedy\"\n\n        self.generation_config = get_generation_config(local_path, trust_remote_code=trust_remote_code)\n\n        override_config_kwargs = {\n            \"bos_token_id\": self.tokenizer.bos_token_id,\n            \"eos_token_id\": self.tokenizer.eos_token_id,\n            \"pad_token_id\": self.tokenizer.pad_token_id,\n        }\n\n        if self.config.model.get(\"mtp\", {}).get(\"enable\", False):\n            raise NotImplementedError(\"Right now,  MTP is not supported in FSDP\")\n        else:\n            if hasattr(actor_model_config, \"num_nextn_predict_layers\"):\n                actor_model_config.num_nextn_predict_layers = 0\n\n        override_config_kwargs.update(override_model_config)\n        update_model_config(actor_model_config, override_config_kwargs=override_config_kwargs)\n        if self.rank == 0:\n            print(f\"Model config after override: {actor_model_config}\")\n\n        # NOTE(fix me): tie_word_embedding causes meta_tensor init to hang\n        init_context = get_init_weight_context_manager(\n            use_meta_tensor=not actor_model_config.tie_word_embeddings, mesh=self.device_mesh\n        )\n\n        with init_context(), warnings.catch_warnings():\n            warnings.simplefilter(\"ignore\")\n            has_remote_code = hasattr(actor_model_config, \"auto_map\") and any(\n                actor_model_config.architectures[0] in val for val in actor_model_config.auto_map.values()\n            )\n            if has_remote_code:\n                auto_class = next(\n                    k for k, v in actor_model_config.auto_map.items() if actor_model_config.architectures[0] in v\n                )\n                match auto_class:\n                    case \"AutoModelForVision2Seq\":\n                        actor_module_class = AutoModelForVision2Seq\n                    case \"AutoModelForCausalLM\":\n                        actor_module_class = AutoModelForCausalLM\n                    case \"AutoModelForImageTextToText\":\n                        actor_module_class = AutoModelForImageTextToText\n                    case _:\n                        actor_module_class = AutoModel\n            else:\n                if type(actor_model_config) in AutoModelForVision2Seq._model_mapping.keys():\n                    actor_module_class = AutoModelForVision2Seq\n                elif type(actor_model_config) in AutoModelForCausalLM._model_mapping.keys():\n                    actor_module_class = AutoModelForCausalLM\n                elif type(actor_model_config) in AutoModelForImageTextToText._model_mapping.keys():\n                    actor_module_class = AutoModelForImageTextToText\n                else:\n                    actor_module_class = AutoModel\n\n            actor_module = actor_module_class.from_pretrained(\n                pretrained_model_name_or_path=local_path,\n                torch_dtype=torch_dtype,\n                config=actor_model_config,\n                trust_remote_code=trust_remote_code,\n                attn_implementation=attn_implementation,\n            )\n\n            # Apply Liger kernel to the model if use_liger is set to True\n            if use_liger:\n                from liger_kernel.transformers.monkey_patch import _apply_liger_kernel_to_instance\n\n                _apply_liger_kernel_to_instance(model=actor_module)\n\n            fused_kernel_options = self.config.model.get(\"fused_kernel_options\", None)\n            fused_kernels_backend = (\n                fused_kernel_options.get(\"impl_backend\", None) if fused_kernel_options is not None else None\n            )\n\n            apply_monkey_patch(\n                model=actor_module,\n                use_remove_padding=use_remove_padding,\n                ulysses_sp_size=self.ulysses_sequence_parallel_size,\n                use_fused_kernels=use_fused_kernels,\n                fused_kernels_backend=fused_kernels_backend,\n                use_prefix_grouper=use_prefix_grouper,\n                use_tiled_mlp=use_tiled_mlp,\n                tiled_mlp_shards=tiled_mlp_shards,\n            )\n\n            # some parameters may not in torch_dtype. TODO(zhangchi.usc1992) remove this after we switch to fsdp2\n            actor_module.to(torch_dtype)\n\n            if enable_gradient_checkpointing:\n                actor_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n\n        if self._is_lora:\n            print(\"Applying LoRA to actor module\")\n            actor_module.enable_input_require_grads()\n\n            lora_adapter_path = self.config.model.get(\"lora_adapter_path\")\n            if lora_adapter_path is not None:\n                from peft import PeftModel\n\n                print(f\"Loading pre-trained LoRA adapter to {role} from: {lora_adapter_path}\")\n\n                # Copy adapter to local if needed\n                local_adapter_path = copy_to_local(lora_adapter_path, use_shm=self.config.model.get(\"use_shm\", False))\n\n                actor_module = PeftModel.from_pretrained(actor_module, local_adapter_path, is_trainable=True)\n                peft_config = actor_module.peft_config[\"default\"]\n                # Ensure task_type is TaskType enum, not string\n                if isinstance(peft_config.task_type, str):\n                    peft_config.task_type = TaskType.CAUSAL_LM\n\n            else:\n                # Convert config to regular Python types before creating PEFT model\n                lora_config = {\n                    \"task_type\": TaskType.CAUSAL_LM,\n                    \"r\": self.config.model.lora_rank,\n                    \"lora_alpha\": self.config.model.lora_alpha,\n                    \"target_modules\": convert_to_regular_types(self.config.model.target_modules),\n                    \"exclude_modules\": convert_to_regular_types(self.config.model.exclude_modules),\n                    \"bias\": \"none\",\n                }\n                actor_module = get_peft_model(actor_module, LoraConfig(**lora_config))\n\n        self.use_orig_params = fsdp_config.get(\"use_orig_params\", False)\n        if self.config.actor.get(\"freeze_vision_tower\", False):\n            vision_tower = get_vl_model_vision_tower(actor_module)\n            if vision_tower is not None:\n                vision_tower.requires_grad_(False)\n                self.use_orig_params = True\n                if self.rank == 0:\n                    print(\"[actor model] Vision tower is set to not trainable.\")\n            else:\n                if self.rank == 0:\n                    print(\"[actor model] No vision tower found.\")\n\n        # Apply QAT before FSDP wrapping (actor only)\n        if role == \"actor\" and self._qat_enabled:\n            actor_module = apply_qat(actor_module, self.qat_config)\n            enable_qat_fuse(actor_module)\n            if self.qat_config.mode == \"w4a4\":\n                self._restore_w4a4_input_scales(actor_module, self.config.model.path)\n\n        torch.distributed.barrier()\n\n        if self.rank == 0:\n            print_model_size(actor_module)\n\n        log_gpu_memory_usage(f\"After init {role} from HF AutoModel\", logger=logger)\n\n        # We wrap FSDP for rollout as well\n        mixed_precision_config = fsdp_config.get(\"mixed_precision\", None)\n        if mixed_precision_config is not None:\n            param_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"param_dtype\", \"bf16\"))\n            reduce_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"reduce_dtype\", \"fp32\"))\n            buffer_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"buffer_dtype\", \"fp32\"))\n        else:\n            param_dtype = PrecisionType.to_dtype(fsdp_config.dtype)\n            reduce_dtype = torch.float32\n            buffer_dtype = torch.float32\n\n        mixed_precision = MixedPrecision(param_dtype=param_dtype, reduce_dtype=reduce_dtype, buffer_dtype=buffer_dtype)\n\n        # Store param_dtype for QAT quantizer\n        self._param_dtype = param_dtype\n\n        auto_wrap_policy = get_fsdp_wrap_policy(\n            module=actor_module,\n            config=fsdp_config.get(\"wrap_policy\", None),\n            is_lora=self._is_lora,\n        )\n\n        # if self._is_rollout and self.config.rollout.name == \"hf\":\n        #     # TODO(zhangchi.usc1992, shengguangming) fix me.\n        #     Current, auto_wrap_policy causes HFRollout to hang in Gemma\n        #     auto_wrap_policy = None\n\n        if self.rank == 0:\n            print(f\"wrap_policy: {auto_wrap_policy}\")\n\n        fsdp_mesh = self.device_mesh\n        fsdp_enable_zero3 = fsdp_config.reshard_after_forward\n        sharding_strategy = get_sharding_strategy(fsdp_mesh, fsdp_enable_zero3)\n\n        # TODO: add transformer policy\n        # We force reference policy to use CPUOffload to save memory.\n        # We force turn off CPUOffload for actor because it causes incorrect results when using grad accumulation\n        cpu_offload = None if role == \"actor\" else CPUOffload(offload_params=True)\n        fsdp_strategy = self.config.actor.strategy\n        if fsdp_strategy == \"fsdp\":\n            actor_module_fsdp = FSDP(\n                actor_module,\n                cpu_offload=cpu_offload,\n                param_init_fn=init_fn,\n                auto_wrap_policy=auto_wrap_policy,\n                device_id=get_device_id(),\n                sharding_strategy=sharding_strategy,  # zero3\n                mixed_precision=mixed_precision,\n                sync_module_states=True,\n                device_mesh=self.device_mesh,\n                use_orig_params=self.use_orig_params,\n                forward_prefetch=fsdp_config.get(\"forward_prefetch\", False),\n            )\n        elif fsdp_strategy == \"fsdp2\":\n            assert CPUOffloadPolicy is not None, \"PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)\"\n            mp_policy = MixedPrecisionPolicy(\n                param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True\n            )\n            if role == \"actor\" and fsdp_config.offload_policy:\n                cpu_offload = CPUOffloadPolicy(pin_memory=True)\n                self._is_offload_param = False\n                self._is_offload_optimizer = False\n            else:\n                cpu_offload = None if role == \"actor\" else CPUOffloadPolicy(pin_memory=True)\n\n            fsdp_kwargs = {\n                \"mesh\": fsdp_mesh,\n                \"mp_policy\": mp_policy,\n                \"offload_policy\": cpu_offload,\n                \"reshard_after_forward\": fsdp_config.reshard_after_forward,\n                \"shard_placement_fn\": get_shard_placement_fn(fsdp_size=self.device_mesh.shape[-1]),\n            }\n            full_state = actor_module.state_dict()\n            apply_fsdp2(actor_module, fsdp_kwargs, fsdp_config)\n            fsdp2_load_full_state_dict(actor_module, full_state, fsdp_mesh, cpu_offload)\n            actor_module_fsdp = actor_module\n        else:\n            raise NotImplementedError(f\"not implement {fsdp_strategy}\")\n\n        if enable_activation_offload:\n            enable_activation_offloading(actor_module_fsdp, fsdp_strategy, enable_gradient_checkpointing)\n\n        log_gpu_memory_usage(f\"After {role} FSDP init\", logger=logger)\n\n        # TODO: add more optimizer args into config\n        if role == \"actor\" and optim_config is not None:\n            from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup\n\n            actor_optimizer = build_optimizer(actor_module_fsdp.parameters(), optim_config)\n\n            total_steps = optim_config.get(\"total_training_steps\", 0)\n            num_warmup_steps = int(optim_config.get(\"lr_warmup_steps\", -1))\n            lr_scheduler_type = optim_config.get(\"lr_scheduler_type\", \"constant\")\n            min_lr_ratio = optim_config.get(\"min_lr_ratio\", 0.0)\n            num_cycles = optim_config.get(\"num_cycles\", 0.5)\n            if num_warmup_steps < 0:\n                num_warmup_steps_ratio = optim_config.get(\"lr_warmup_steps_ratio\", 0.0)\n                num_warmup_steps = int(num_warmup_steps_ratio * total_steps)\n\n            if self.rank == 0:\n                print(f\"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}\")\n\n            if lr_scheduler_type == \"constant\":\n                actor_lr_scheduler = get_constant_schedule_with_warmup(\n                    optimizer=actor_optimizer, num_warmup_steps=num_warmup_steps\n                )\n            elif lr_scheduler_type == \"cosine\":\n                actor_lr_scheduler = get_cosine_schedule_with_warmup(\n                    optimizer=actor_optimizer,\n                    num_warmup_steps=num_warmup_steps,\n                    num_training_steps=total_steps,\n                    min_lr_ratio=min_lr_ratio,\n                    num_cycles=num_cycles,\n                )\n            else:\n                raise NotImplementedError(f\"LR scheduler type {lr_scheduler_type} is not supported\")\n\n            log_gpu_memory_usage(f\"After {role} optimizer init\", logger=logger)\n        else:\n            actor_optimizer = None\n            actor_lr_scheduler = None\n\n        return actor_module_fsdp, actor_optimizer, actor_lr_scheduler, actor_model_config\n\n    def _build_rollout(self, trust_remote_code=False):\n        from torch.distributed.device_mesh import init_device_mesh\n\n        # 1. parse rollout and huggingface model config\n        rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout)\n        model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model, dataclass_type=HFModelConfig)\n        self.model_config = model_config\n\n        # 2. build rollout device mesh\n        infer_tp = self.config.rollout.tensor_model_parallel_size * self.config.rollout.data_parallel_size\n        infer_pp = self.config.rollout.pipeline_model_parallel_size\n        infer_world_size = infer_tp * infer_pp\n        dp = self.world_size // infer_world_size\n        assert self.world_size % infer_world_size == 0, (\n            f\"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}\"\n        )\n        rollout_device_mesh = init_device_mesh(\n            device_name, mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=[\"dp\", \"infer_tp\", \"infer_pp\"]\n        )\n        rollout_name = self.config.rollout.name\n\n        self.rollout_device_mesh = rollout_device_mesh\n\n        if rollout_name == \"hf\":\n            self._register_dispatch_collect_info(\"rollout\", dp_rank=self.rank, is_collect=True)\n        else:\n            is_collect = (\n                rollout_device_mesh[\"infer_tp\"].get_local_rank() == 0\n                and rollout_device_mesh[\"infer_pp\"].get_local_rank() == 0\n            )\n            self._register_dispatch_collect_info(\n                \"rollout\", dp_rank=rollout_device_mesh[\"dp\"].get_local_rank(), is_collect=is_collect\n            )\n\n        # 4. build rollout model\n        log_gpu_memory_usage(f\"Before building {self.config.rollout.name} rollout\", logger=logger)\n        self.rollout = get_rollout_class(rollout_config.name, rollout_config.mode)(\n            config=rollout_config, model_config=model_config, device_mesh=rollout_device_mesh\n        )\n        log_gpu_memory_usage(f\"After building {self.config.rollout.name} rollout\", logger=logger)\n\n        # Full params\n        if torch.distributed.get_world_size() == 1 and fsdp_version(self.actor_module_fsdp) == 1:\n            FSDP.set_state_dict_type(\n                self.actor_module_fsdp,\n                state_dict_type=StateDictType.FULL_STATE_DICT,\n                state_dict_config=FullStateDictConfig(),\n            )\n        elif fsdp_version(self.actor_module_fsdp) == 1:\n            FSDP.set_state_dict_type(\n                self.actor_module_fsdp,\n                state_dict_type=StateDictType.SHARDED_STATE_DICT,\n                state_dict_config=ShardedStateDictConfig(),\n            )\n\n        # used for LoRA\n        self.base_sync_done: bool = \"dummy\" not in self.config.rollout.load_format\n        self.layered_summon = self.config.rollout.get(\"layered_summon\", False)\n\n        # 5. switch to trainer mode\n        # NOTE: It's critical that hybrid engine in trainer mode initially to load checkpoint.\n        # For async mode, we can't call run_until_complete here, so we will switch to trainer mode in AgentLoopManager.\n        # Note: sync mode is deprecated and rejected in RolloutConfig.__post_init__\n\n    async def rollout_mode(self):\n        \"\"\"Context switch hybridengine to rollout mode.\"\"\"\n        aggressive_empty_cache(force_sync=True)\n\n        log_gpu_memory_usage(\"Before load_fsdp_model_to_gpu\", logger=logger)\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.actor_module_fsdp)\n        log_gpu_memory_usage(\"After load_fsdp_model_to_gpu\", logger=logger)\n\n        peft_config = None\n        peft_model = getattr(self.actor_module_fsdp, \"_fsdp_wrapped_module\", self.actor_module_fsdp)\n        if hasattr(peft_model, \"peft_config\"):  # LoRA\n            peft_config = peft_model.peft_config.get(\"default\", None)\n            params = collect_lora_params(\n                module=self.actor_module_fsdp,\n                layered_summon=self.config.rollout.get(\"layered_summon\", False),\n                base_sync_done=self.base_sync_done,\n            )\n            if not self.base_sync_done:\n                params = {replace_lora_wrapper(k, peft_config): v for k, v in params.items()}\n        else:\n            params = self.actor_module_fsdp.state_dict()\n\n        params = convert_weight_keys(\n            params, getattr(self.actor_module_fsdp, \"_fsdp_wrapped_module\", self.actor_module_fsdp)\n        )\n\n        # Special handling for LoRA with sleep_level=2:\n        # When sleep_level=2, base model weights are destroyed during each sleep cycle.\n        # separately collect and update LoRA weights and base model weights through their respective interfaces.\n        # Here: params contains LoRA weights, base_model_params contains base model weights.\n        # Only needed if the rollout engine actually sleeps/frees weights (free_cache_engine=True).\n        if (\n            peft_config is not None\n            and getattr(self.rollout, \"sleep_level\", None) == 2\n            and self.config.rollout.free_cache_engine\n        ):\n            base_model_params = collect_lora_params(\n                module=self.actor_module_fsdp,\n                layered_summon=self.layered_summon,\n                base_sync_done=False,\n            )\n            base_model_params = {replace_lora_wrapper(k, peft_config): v for k, v in base_model_params.items()}\n            base_model_params = convert_weight_keys(\n                base_model_params, getattr(self.actor_module_fsdp, \"_fsdp_wrapped_module\", self.actor_module_fsdp)\n            )\n\n        log_gpu_memory_usage(\"Before offload_fsdp_model_to_cpu\", logger=logger)\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n        log_gpu_memory_usage(\"After offload_fsdp_model_to_cpu\", logger=logger)\n\n        set_expandable_segments(False)\n\n        if peft_config is not None and self.base_sync_done:\n            per_tensor_param = params.items() if isinstance(params, dict) else params  # Fixed: handle dict case\n        else:\n            device = get_device_id()  # used when fsdp2 set cpu_offload_policy\n            per_tensor_param = (\n                (name, param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param)\n                for name, param in params.items()\n            )\n\n        # QAT: quantize weights before sending to vLLM\n        if self._qat_enabled:\n            from verl.utils.qat.quantizer import QATQuantizer\n\n            quantizer = QATQuantizer(\n                mode=self.qat_config.mode,\n                group_size=self.qat_config.group_size,\n                ignore_patterns=self.qat_config.ignore_patterns,\n                device=torch.device(get_device_id()),\n                param_dtype=self._param_dtype,\n            )\n            per_tensor_param = quantizer.quantize_with_fusion(\n                per_tensor_param,\n                target_device=torch.device(\"cpu\"),\n            )\n            aggressive_empty_cache(force_sync=True)\n\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"weights\"])\n        log_gpu_memory_usage(\"After resume weights\", logger=logger)\n\n        if (\n            peft_config is not None\n            and getattr(self.rollout, \"sleep_level\", None) == 2\n            and self.config.rollout.free_cache_engine\n        ):\n            per_tensor_base_params = (\n                (name, param.to(device, non_blocking=True).full_tensor() if isinstance(param, DTensor) else param)\n                for name, param in base_model_params.items()\n            )\n            await self.rollout.update_weights(per_tensor_base_params, base_sync_done=False)\n            del base_model_params, per_tensor_base_params\n\n        await self.rollout.update_weights(per_tensor_param, peft_config=peft_config, base_sync_done=self.base_sync_done)\n        log_gpu_memory_usage(\"After update_weights\", logger=logger)\n        del params, per_tensor_param\n        aggressive_empty_cache(force_sync=True)\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"kv_cache\"])\n        log_gpu_memory_usage(\"After resume kv_cache\", logger=logger)\n\n        self.base_sync_done = True\n        set_expandable_segments(True)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        from verl.workers.actor import DataParallelPPOActor\n\n        # This is used to import external_lib into the huggingface systems\n        import_external_libs(self.config.model.get(\"external_lib\", None))\n\n        # Initialize QAT config before _build_model_optimizer\n        self._init_qat_config()\n\n        override_model_config = OmegaConf.to_container(OmegaConf.create(self.config.model.get(\"override_config\", {})))\n        use_remove_padding = self.config.model.get(\"use_remove_padding\", False)\n        use_shm = self.config.model.get(\"use_shm\", False)\n        use_fused_kernels = self.config.model.get(\"use_fused_kernels\", False)\n\n        if self._is_actor or self._is_rollout:\n            # we need the model for actor and rollout\n            if self._is_actor:\n                optim_config = self.config.actor.optim\n                fsdp_config = omega_conf_to_dataclass(self.config.actor.fsdp_config)\n            else:\n                optim_config = None\n                fsdp_config = FSDPEngineConfig()\n\n            local_path = copy_to_local(self.config.model.path, use_shm=use_shm)\n            # TiledMLP configuration for memory-efficient MLP computation\n            tiled_mlp_config = self.config.model.get(\"tiled_mlp\", {})\n            use_tiled_mlp = tiled_mlp_config.get(\"enabled\", False)\n            tiled_mlp_shards = tiled_mlp_config.get(\"num_shards\", 4)\n\n            (\n                self.actor_module_fsdp,\n                self.actor_optimizer,\n                self.actor_lr_scheduler,\n                self.actor_model_config,\n            ) = self._build_model_optimizer(\n                model_path=local_path,\n                fsdp_config=fsdp_config,\n                optim_config=optim_config,\n                override_model_config=override_model_config,\n                use_remove_padding=use_remove_padding,\n                use_fused_kernels=use_fused_kernels,\n                enable_gradient_checkpointing=self.config.model.get(\"enable_gradient_checkpointing\", False),\n                trust_remote_code=self.config.model.get(\"trust_remote_code\", False),\n                use_liger=self.config.model.get(\"use_liger\", False),\n                role=\"actor\",\n                enable_activation_offload=self.config.model.get(\"enable_activation_offload\", False),\n                use_prefix_grouper=self.config.actor.get(\"use_prefix_grouper\", False),\n                use_tiled_mlp=use_tiled_mlp,\n                tiled_mlp_shards=tiled_mlp_shards,\n            )\n\n            # get the original unwrapped module\n            if fsdp_version(self.actor_module_fsdp) == 1:\n                self.actor_module = self.actor_module_fsdp._fsdp_wrapped_module\n\n            if self._is_offload_param:\n                offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n                log_gpu_memory_usage(\"After offload actor model during init\", logger=logger)\n\n            if self._is_offload_optimizer:\n                offload_fsdp_optimizer(optimizer=self.actor_optimizer)\n                log_gpu_memory_usage(\"After offload actor optimizer during init\", logger=logger)\n\n        if self._is_actor:\n            actor_cfg = omega_conf_to_dataclass(self.config.actor)\n            self.actor = DataParallelPPOActor(\n                config=actor_cfg, actor_module=self.actor_module_fsdp, actor_optimizer=self.actor_optimizer\n            )\n\n        if self._is_rollout:\n            self._build_rollout(trust_remote_code=self.config.model.get(\"trust_remote_code\", False))\n\n        if self._is_ref:\n            ref_model_path = self.config.model.path\n            ref_model = self.config.ref.get(\"model\", None)\n            if ref_model is not None:\n                ref_model_path = ref_model.get(\"path\", self.config.model.path)\n\n            if self.rank == 0:\n                print(\"reference model:\", ref_model_path)\n            local_path = copy_to_local(ref_model_path, use_shm=use_shm)\n            use_prefix_grouper = hasattr(self.config, \"actor\") and self.config.actor.get(\"use_prefix_grouper\", False)\n\n            # TiledMLP for ref model: use ref config if specified, otherwise use actor config\n            ref_tiled_mlp_config = self.config.ref.get(\"tiled_mlp\", None)\n            if ref_tiled_mlp_config is None:\n                ref_tiled_mlp_config = self.config.model.get(\"tiled_mlp\", {})\n            ref_use_tiled_mlp = ref_tiled_mlp_config.get(\"enabled\", False)\n            ref_tiled_mlp_shards = ref_tiled_mlp_config.get(\"num_shards\", 4)\n\n            self.ref_module_fsdp = self._build_model_optimizer(\n                model_path=local_path,\n                fsdp_config=omega_conf_to_dataclass(self.config.ref.fsdp_config),\n                optim_config=None,\n                override_model_config=override_model_config,\n                use_remove_padding=use_remove_padding,\n                use_fused_kernels=use_fused_kernels,\n                trust_remote_code=self.config.model.get(\"trust_remote_code\", False),\n                use_liger=self.config.model.get(\"use_liger\", False),\n                role=\"ref\",\n                use_prefix_grouper=use_prefix_grouper,\n                use_tiled_mlp=ref_use_tiled_mlp,\n                tiled_mlp_shards=ref_tiled_mlp_shards,\n            )[0]\n            OmegaConf.set_struct(self.config.ref, True)\n            with open_dict(self.config.ref):\n                self.config.ref.use_remove_padding = use_remove_padding\n                self.config.ref.use_fused_kernels = use_fused_kernels\n                if use_prefix_grouper:\n                    self.config.ref.use_prefix_grouper = use_prefix_grouper\n            self.ref_policy = DataParallelPPOActor(config=self.config.ref, actor_module=self.ref_module_fsdp)\n\n        if self._is_actor:\n            self.flops_counter = FlopsCounter(self.actor_model_config)\n            self.checkpoint_manager = FSDPCheckpointManager(\n                model=self.actor_module_fsdp,\n                optimizer=self.actor.actor_optimizer,\n                lr_scheduler=self.actor_lr_scheduler,\n                processing_class=self.processor if self.processor is not None else self.tokenizer,\n                checkpoint_config=self.config.actor.checkpoint,\n                trust_remote_code=self.config.model.get(\"trust_remote_code\", False),\n            )\n\n        if not self._is_actor and self._is_rollout:\n            # If ActorRolloutRefWorker is initialized as a standalone rollout,\n            # create a checkpoint manager for FSDP model to allow loading FSDP checkpoints for rollout.\n\n            checkpoint_contents = OmegaConf.create({\"load_contents\": [\"model\"], \"save_contents\": []})\n            self.checkpoint_manager = FSDPCheckpointManager(\n                model=self.actor_module_fsdp,\n                optimizer=None,\n                lr_scheduler=None,\n                processing_class=self.processor if self.processor is not None else self.tokenizer,\n                checkpoint_config=checkpoint_contents,\n            )\n\n        # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo\n        aggressive_empty_cache(force_sync=True)\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @DistProfiler.annotate(color=\"red\", role=\"actor_update\")\n    def update_actor(self, data: DataProto):\n        assert self._is_actor\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.actor_module_fsdp)\n        if self._is_offload_optimizer:\n            load_fsdp_optimizer(optimizer=self.actor_optimizer, device_id=get_device_id())\n\n        with self.ulysses_sharding_manager:\n            data = data.to(\"cpu\")  # data will to device with each micro batch on actor.update_policy\n            data.meta_info.setdefault(\"pad_token_id\", self.tokenizer.pad_token_id)\n            # perform training\n            with Timer(name=\"update_policy\", logger=None) as timer:\n                metrics = self.actor.update_policy(data=data)\n            delta_time = timer.last\n            global_num_tokens = data.meta_info[\"global_token_num\"]\n            images_seqlens = data.meta_info.get(\"images_seqlens\", None)\n            estimated_flops, promised_flops = self.flops_counter.estimate_flops(\n                global_num_tokens, delta_time, images_seqlens=images_seqlens\n            )\n            metrics[\"perf/mfu/actor\"] = (\n                estimated_flops * self.config.actor.ppo_epochs / promised_flops / self.world_size\n            )\n            metrics[\"perf/max_memory_allocated_gb\"] = get_torch_device().max_memory_allocated() / (1024**3)\n            metrics[\"perf/max_memory_reserved_gb\"] = get_torch_device().max_memory_reserved() / (1024**3)\n            metrics[\"perf/cpu_memory_used_gb\"] = psutil.virtual_memory().used / (1024**3)\n\n            lr = self.actor_lr_scheduler.get_last_lr()[0]\n            metrics[\"actor/lr\"] = lr.item() if torch.is_tensor(lr) else lr\n            self.actor_lr_scheduler.step()\n\n            # TODO: here, we should return all metrics\n            output = DataProto(meta_info={\"metrics\": metrics})\n\n            output = output.to(\"cpu\")\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n            log_gpu_memory_usage(\"After offload actor model during update_actor\", logger=logger)\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(optimizer=self.actor_optimizer)\n            log_gpu_memory_usage(\"After offload actor optimizer during update_actor\", logger=logger)\n\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"rollout\"))\n    @DistProfiler.annotate(color=\"red\", role=\"rollout_generate\")\n    def generate_sequences(self, prompts: DataProto):\n        # Support all hardwares\n        assert self._is_rollout\n        prompts = prompts.to(get_device_id())\n\n        meta_info = {\n            \"eos_token_id\": self.generation_config.eos_token_id\n            if self.generation_config is not None\n            else self.tokenizer.eos_token_id,\n            \"pad_token_id\": self.generation_config.pad_token_id\n            if self.generation_config is not None\n            else self.tokenizer.pad_token_id,\n        }\n        prompts.meta_info.update(meta_info)\n\n        timing_generate = {}\n        if self._is_actor:  # For rollout only, we do not switch context.\n            loop = get_event_loop()\n            loop.run_until_complete(self.rollout_mode())\n            log_gpu_memory_usage(\"After switch to rollout mode\", logger=logger)\n\n        with simple_timer(\"generate_sequences\", timing_generate):\n            output = self.rollout.generate_sequences(prompts=prompts)\n\n        if self._is_actor:\n            loop.run_until_complete(self.trainer_mode())\n            log_gpu_memory_usage(\"After switch to trainer mode\", logger=logger)\n\n        # We calculate the average timing across all ranks\n        # to make sure meta_info[\"timing\"] is the same\n        timing_generate_topk_ratio, timing_generate_min, timing_generate_max = topk_reduce_ratio_min_max(\n            timing_generate[\"generate_sequences\"]\n        )\n        timing_generate = reduce_timing(timing_generate)\n        timing_generate.update(\n            {\n                \"generation_timing/max\": timing_generate_max,\n                \"generation_timing/min\": timing_generate_min,\n                \"generation_timing/topk_ratio\": timing_generate_topk_ratio,\n            }\n        )\n        output.meta_info[\"timing\"] = timing_generate\n        output = output.to(\"cpu\")\n\n        # clear kv cache\n        get_torch_device().empty_cache()\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @DistProfiler.annotate(color=\"blue\", role=\"actor_compute_log_prob\")\n    def compute_log_prob(self, data: DataProto):\n        # when is_lora is True, we use the actor without lora applied to calculate the log_prob\n        # which is mostly used for ref log_prob calculation\n        assert self._is_actor\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.actor_module_fsdp)\n\n        # Support all hardwares\n        from contextlib import nullcontext\n\n        is_lora = data.meta_info.pop(\"is_lora\", False)\n        adapter_ctx = self.actor.actor_module.disable_adapter() if is_lora else nullcontext()\n        # we should always recompute old_log_probs when it is HybridEngine\n        config_source = self.config.ref if is_lora else self.config.rollout\n        data.meta_info[\"micro_batch_size\"] = config_source.log_prob_micro_batch_size_per_gpu\n        data.meta_info[\"max_token_len\"] = config_source.log_prob_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = config_source.log_prob_use_dynamic_bsz\n        data.meta_info[\"temperature\"] = self.config.rollout.temperature\n        data.meta_info.setdefault(\"pad_token_id\", self.tokenizer.pad_token_id)\n        # perform recompute log_prob\n        calculate_entropy = not is_lora\n        with self.ulysses_sharding_manager:\n            with adapter_ctx:\n                outputs = self.actor.compute_log_prob(data=data, calculate_entropy=calculate_entropy)\n            if not is_lora:\n                tensors = {\"old_log_probs\": outputs[\"log_probs\"]}\n            else:\n                tensors = {\"ref_log_prob\": outputs[\"log_probs\"]}\n            if calculate_entropy:\n                tensors[\"entropys\"] = outputs[\"entropys\"]\n            if \"sum_pi_squared\" in outputs:\n                tensors[\"sum_pi_squared\"] = outputs[\"sum_pi_squared\"]\n            output = DataProto.from_dict(\n                tensors=tensors,\n                meta_info={\"temperature\": self.config.rollout.temperature},\n            )\n\n        output = output.to(\"cpu\")\n\n        # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes\n        # unshard the root FSDP module\n        if self.world_size > 1 and fsdp_version(self.actor.actor_module) == 1:\n            self.actor.actor_module._handle.reshard(True)\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n            log_gpu_memory_usage(\"After offload actor model during compute_log_prob\", logger=logger)\n\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @DistProfiler.annotate(color=\"olive\", role=\"ref_compute_log_prob\")\n    def compute_ref_log_prob(self, data: DataProto):\n        if self._is_lora:\n            # if _is_lora, actor without lora applied is the ref\n            data.meta_info[\"is_lora\"] = True\n            return self.compute_log_prob(data)\n        assert self._is_ref\n        # else:\n        # otherwise, the class have a standalone ref model\n\n        micro_batch_size = self.config.ref.log_prob_micro_batch_size_per_gpu\n        data.meta_info[\"micro_batch_size\"] = micro_batch_size\n        data.meta_info[\"temperature\"] = self.config.rollout.temperature\n        data.meta_info[\"max_token_len\"] = self.config.ref.log_prob_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = self.config.ref.log_prob_use_dynamic_bsz\n        data.meta_info.setdefault(\"pad_token_id\", self.tokenizer.pad_token_id)\n        with self.ulysses_sharding_manager:\n            data = data.to(\"cpu\")  # data will to device with each micro batch on ref.compute_log_prob\n            outputs = self.ref_policy.compute_log_prob(data=data, calculate_entropy=False)\n            output = DataProto.from_dict(tensors={\"ref_log_prob\": outputs[\"log_probs\"]})\n\n        output = output.to(\"cpu\")\n\n        # https://pytorch.org/docs/stable/notes/fsdp.html#fsdp-notes\n        # unshard the root FSDP module\n        if self.world_size > 1:\n            if fsdp_version(self.ref_policy.actor_module) == 1:\n                self.ref_policy.actor_module._handle.reshard(True)\n            elif fsdp_version(self.ref_policy.actor_module) == 2:\n                self.ref_policy.actor_module.reshard()\n\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):\n        from verl.utils.logger import log_with_rank\n\n        # only support save and load ckpt for actor\n        assert self._is_actor\n\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.actor_module_fsdp)\n\n        self.checkpoint_manager.save_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n        dist.barrier()\n\n        if self._is_lora and hasattr(getattr(self, \"actor_module\", self.actor_module_fsdp), \"peft_config\"):\n            lora_save_path = os.path.join(local_path, \"lora_adapter\")\n            peft_model = getattr(self, \"actor_module\", self.actor_module_fsdp)\n            peft_config = {}\n            if dist.get_rank() == 0:\n                os.makedirs(lora_save_path, exist_ok=True)\n                peft_config = asdict(peft_model.peft_config.get(\"default\", {}))\n                peft_config[\"task_type\"] = peft_config[\"task_type\"].value\n                peft_config[\"peft_type\"] = peft_config[\"peft_type\"].value\n                peft_config[\"target_modules\"] = list(peft_config[\"target_modules\"])\n            try:\n                if fsdp_version(self.actor_module_fsdp) > 0:\n                    self.actor_module_fsdp = self.actor_module_fsdp.to(get_device_name())\n                    lora_params = layered_summon_lora_params(self.actor_module_fsdp)\n                    if dist.get_rank() == 0:\n                        save_file(lora_params, os.path.join(lora_save_path, \"adapter_model.safetensors\"))\n                        with open(os.path.join(lora_save_path, \"adapter_config.json\"), \"w\", encoding=\"utf-8\") as f:\n                            json.dump(peft_config, f, ensure_ascii=False, indent=4)\n            except Exception as e:\n                log_with_rank(\n                    f\"Save LoRA Adapter Error ({e})\", rank=dist.get_rank(), logger=logger, log_only_rank_0=True\n                )\n\n            dist.barrier()\n            log_with_rank(\n                f\"[rank-{self.rank}]: Saved LoRA adapter to: {lora_save_path}\",\n                rank=dist.get_rank(),\n                logger=logger,\n                log_only_rank_0=True,\n            )\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=False):\n        assert self._is_actor or (not self._is_actor and self._is_rollout), (\n            f\"Checkpoint loading is only supported for Actor or standalone Rollout Workers, but got \"\n            f\"{self._is_actor} and {self._is_rollout}\"\n        )\n\n        # No checkpoint to load, just offload the model and optimizer to CPU\n        if local_path is None:\n            if self._is_offload_param:\n                offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n            if self._is_offload_optimizer:\n                offload_fsdp_optimizer(self.actor_optimizer)\n            return\n\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.actor_module_fsdp)\n\n        self.checkpoint_manager.load_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.actor_module_fsdp)\n\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(self.actor_optimizer)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def start_profile(self, **kwargs) -> None:\n        \"\"\"Start profiling for the current rank in the current training step.\"\"\"\n        self.profiler.start(**kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def stop_profile(self) -> None:\n        \"\"\"Stop profiling for the current rank in the current training step.\"\"\"\n        self.profiler.stop()\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def dump_memory_snapshot(self, tag: str = \"manual\", sub_dir: str = None) -> None:\n        \"\"\"Manually trigger a CUDA memory snapshot dump on all ranks.\"\"\"\n        # Memory snapshot is now handled by the profiler system\n        # This method is kept for backward compatibility but delegates to profiler\n        if hasattr(self, \"profiler\") and hasattr(self.profiler, \"_impl\"):\n            try:\n                # Try to use the profiler's memory snapshot functionality\n                if hasattr(self.profiler._impl, \"sampler\"):\n                    out_dir = OmegaConf.select(self.config, \"actor.profiler.save_path\") or \".\"\n                    self.profiler._impl.sampler.dump_memory_snapshot(out_dir=out_dir, tag=tag, sub_dir=sub_dir)\n            except Exception:\n                # silently ignore if profiler doesn't support memory snapshots\n                pass\n\n\nclass CriticWorker(Worker, DistProfilerExtension):\n    def __init__(self, config: FSDPCriticConfig):\n        Worker.__init__(self)\n        omega_profiler_config = config.get(\"profiler\", {})\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n        import torch.distributed\n\n        self.config = config\n        if not torch.distributed.is_initialized():\n            torch.distributed.init_process_group(\n                backend=get_nccl_backend(),\n                timeout=datetime.timedelta(seconds=self.config.get(\"nccl_timeout\", 600)),\n                init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n            )\n        self.config: FSDPCriticConfig = config\n\n        # build device mesh for Ulysses Sequence Parallel\n        world_size = torch.distributed.get_world_size()\n        from torch.distributed.device_mesh import init_device_mesh\n\n        fsdp_size = self.config.model.fsdp_config.fsdp_size\n        self.device_mesh = create_device_mesh(world_size=world_size, fsdp_size=fsdp_size)\n\n        self.ulysses_device_mesh = None\n        self.ulysses_sequence_parallel_size = self.config.get(\"ulysses_sequence_parallel_size\", 1)\n        dp = world_size // self.ulysses_sequence_parallel_size\n        if self.ulysses_sequence_parallel_size > 1:\n            self.ulysses_device_mesh = init_device_mesh(\n                device_name, mesh_shape=(dp, self.ulysses_sequence_parallel_size), mesh_dim_names=[\"dp\", \"sp\"]\n            )\n\n        # create training dispatch\n        if self.ulysses_device_mesh is not None:\n            is_collect = self.ulysses_device_mesh[\"sp\"].get_local_rank() == 0\n            self._register_dispatch_collect_info(\n                \"critic\", dp_rank=self.ulysses_device_mesh[\"dp\"].get_local_rank(), is_collect=is_collect\n            )\n        else:\n            self._register_dispatch_collect_info(\"critic\", dp_rank=self.rank, is_collect=True)\n\n        self.ulysses_sharding_manager = FSDPUlyssesShardingManager(self.ulysses_device_mesh)\n\n        # set FSDP offload params\n        self._is_offload_param = self.config.model.fsdp_config.param_offload\n        self._is_offload_optimizer = self.config.model.fsdp_config.optimizer_offload\n\n        # normalize config\n        self.config.ppo_mini_batch_size *= self.config.rollout_n\n        self.config.ppo_mini_batch_size //= torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size\n        if self.config.ppo_micro_batch_size is not None:\n            self.config.ppo_micro_batch_size //= (\n                torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size\n            )\n            self.config.forward_micro_batch_size //= (\n                torch.distributed.get_world_size() // self.ulysses_sequence_parallel_size\n            )\n            self.config.ppo_micro_batch_size_per_gpu = self.config.ppo_micro_batch_size\n            self.config.forward_micro_batch_size_per_gpu = self.config.forward_micro_batch_size\n\n        if self.config.ppo_micro_batch_size_per_gpu is not None:\n            assert self.config.ppo_mini_batch_size % self.config.ppo_micro_batch_size_per_gpu == 0, (\n                f\"normalized ppo_mini_batch_size {self.config.ppo_mini_batch_size} should be divisible by \"\n                f\"ppo_micro_batch_size_per_gpu {self.config.ppo_micro_batch_size_per_gpu}\"\n            )\n            assert self.config.ppo_mini_batch_size // self.config.ppo_micro_batch_size_per_gpu > 0, (\n                f\"normalized ppo_mini_batch_size {self.config.ppo_mini_batch_size} should be larger than \"\n                f\"ppo_micro_batch_size_per_gpu {self.config.ppo_micro_batch_size_per_gpu}\"\n            )\n        self._is_lora = (\n            self.config.model.get(\"lora_adapter_path\") is not None or self.config.model.get(\"lora_rank\", 0) > 0\n        )\n        self.use_orig_params = self.config.model.fsdp_config.get(\"use_orig_params\", False)\n\n    def _build_critic_model_optimizer(self, config: FSDPCriticConfig):\n        # the following line is necessary\n        from torch.distributed.fsdp import MixedPrecision\n\n        from verl.utils.model import load_valuehead_model, print_model_size\n        from verl.utils.torch_dtypes import PrecisionType\n\n        use_shm = config.model.get(\"use_shm\", False)\n        local_path = copy_to_local(config.model.path, use_shm=use_shm)\n        # note that the tokenizer between actor and critic may be different. So override tokenizer info with actor info\n        # using random initialized model from any architecture. May not be the same as Actor.\n\n        tokenizer_path = copy_to_local(config.model.tokenizer_path, use_shm=use_shm)\n        self.tokenizer = hf_tokenizer(tokenizer_path, trust_remote_code=config.model.get(\"trust_remote_code\", False))\n        self.processor = hf_processor(tokenizer_path, trust_remote_code=config.model.get(\"trust_remote_code\", False))\n\n        if self.config.model.get(\"custom_chat_template\", None) is not None:\n            if self.processor is not None:\n                self.processor.chat_template = self.config.model.custom_chat_template\n            else:\n                self.tokenizer.chat_template = self.config.model.custom_chat_template\n        override_config = OmegaConf.to_container(OmegaConf.create(self.config.model.get(\"override_config\", {})))\n        override_config_kwargs = {\n            \"bos_token_id\": self.tokenizer.bos_token_id,\n            \"eos_token_id\": self.tokenizer.eos_token_id,\n            \"pad_token_id\": self.tokenizer.pad_token_id,\n        }\n        override_config_kwargs.update(override_config)\n        if self.rank == 0:\n            print(f\"Critic overriding config {override_config_kwargs}\")\n\n        torch_dtype = self.config.model.fsdp_config.get(\"model_dtype\", \"fp32\")\n        torch_dtype = PrecisionType.to_dtype(torch_dtype)\n\n        from transformers import AutoConfig\n\n        # override model kwargs\n        attn_implementation = override_config.get(\"attn_implementation\", \"flash_attention_2\")\n        critic_model_config = AutoConfig.from_pretrained(\n            local_path,\n            attn_implementation=attn_implementation,\n            trust_remote_code=config.model.get(\"trust_remote_code\", False),\n        )\n        # TODO: VL models use VisionAttention, which directly uses flash_attention in transformers>=4.53\n        # which will be patched by _ulysses_flash_attention_forward, but errorly misses position_ids\n        # Maybe support Ulysses in VisionAttention in the future and remove this patch\n        if self.ulysses_sequence_parallel_size > 1 and hasattr(critic_model_config, \"vision_config\"):\n            critic_model_config.vision_config._attn_implementation = \"eager\"\n\n        critic_model_config.num_labels = 1\n        # patch for kimi-vl\n        if getattr(critic_model_config, \"model_type\", None) == \"kimi_vl\":\n            critic_model_config.text_config.topk_method = \"greedy\"\n\n        init_context = get_init_weight_context_manager(\n            use_meta_tensor=not critic_model_config.tie_word_embeddings, mesh=self.device_mesh\n        )\n\n        # TiledMLP configuration for memory-efficient MLP computation\n        tiled_mlp_config = config.model.get(\"tiled_mlp\", {})\n        use_tiled_mlp = tiled_mlp_config.get(\"enabled\", False)\n        tiled_mlp_shards = tiled_mlp_config.get(\"num_shards\", 4)\n\n        # TiledMLP requires FSDP2 for correct gradient computation\n        if use_tiled_mlp and config.strategy == \"fsdp\":\n            raise ValueError(\"TiledMLP requires FSDP2. Set `critic.strategy=fsdp2`.\")\n\n        with init_context(), warnings.catch_warnings():\n            warnings.simplefilter(\"ignore\")\n            critic_model_config.classifier_dropout = 0.0\n            critic_model_config.hidden_dropout = \"0\"\n            critic_model_config.summary_dropout_prob = 0.0\n\n            critic_module = load_valuehead_model(\n                local_path,\n                torch_dtype,\n                critic_model_config,\n                config.model.get(\"trust_remote_code\", False),\n            )\n\n            use_remove_padding = config.model.get(\"use_remove_padding\", False)\n\n            apply_monkey_patch(\n                model=critic_module,\n                use_remove_padding=use_remove_padding,\n                ulysses_sp_size=self.ulysses_sequence_parallel_size,\n                use_tiled_mlp=use_tiled_mlp,\n                tiled_mlp_shards=tiled_mlp_shards,\n            )\n\n            # some parameters may not in torch_dtype\n            critic_module.to(torch_dtype)\n\n            if config.model.get(\"enable_gradient_checkpointing\", False):\n                critic_module.gradient_checkpointing_enable(gradient_checkpointing_kwargs={\"use_reentrant\": False})\n\n        if self._is_lora:\n            print(\"Applying LoRA to critic module\")\n            critic_module.enable_input_require_grads()\n\n            # Check if we should load a pre-trained LoRA adapter\n            lora_adapter_path = self.config.model.get(\"lora_adapter_path\")\n            if lora_adapter_path is not None:\n                from peft import PeftModel\n\n                print(f\"Loading pre-trained LoRA adapter to critic from: {lora_adapter_path}\")\n\n                # Copy adapter to local if needed\n                local_adapter_path = copy_to_local(lora_adapter_path, use_shm=self.config.model.get(\"use_shm\", False))\n\n                critic_module = PeftModel.from_pretrained(critic_module, local_adapter_path, is_trainable=True)\n                peft_config = critic_module.peft_config[\"default\"]\n                # Ensure task_type is TaskType enum, not string\n                # Use TOKEN_CLS for Critic since it's loaded as AutoModelForTokenClassification\n                if isinstance(peft_config.task_type, str):\n                    peft_config.task_type = TaskType.TOKEN_CLS\n\n            else:\n                # Convert config to regular Python types before creating PEFT model\n                # Use TOKEN_CLS for Critic since it's loaded as AutoModelForTokenClassification\n                lora_config = {\n                    \"task_type\": TaskType.TOKEN_CLS,\n                    \"r\": self.config.model.lora_rank,\n                    \"lora_alpha\": self.config.model.lora_alpha,\n                    \"target_modules\": convert_to_regular_types(self.config.model.target_modules),\n                    \"bias\": \"none\",\n                }\n                critic_module = get_peft_model(critic_module, LoraConfig(**lora_config))\n\n        if self.rank == 0:\n            print_model_size(critic_module)\n\n        self.critic_model_config = critic_model_config\n\n        fsdp_config = self.config.model.fsdp_config\n        mixed_precision_config = fsdp_config.get(\"mixed_precision\", None)\n        if mixed_precision_config is not None:\n            param_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"param_dtype\", \"bf16\"))\n            reduce_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"reduce_dtype\", \"fp32\"))\n            buffer_dtype = PrecisionType.to_dtype(mixed_precision_config.get(\"buffer_dtype\", \"fp32\"))\n        else:\n            param_dtype = torch.bfloat16\n            reduce_dtype = torch.float32\n            buffer_dtype = torch.float32\n\n        mixed_precision = MixedPrecision(param_dtype=param_dtype, reduce_dtype=reduce_dtype, buffer_dtype=buffer_dtype)\n\n        auto_wrap_policy = get_fsdp_wrap_policy(\n            module=critic_module,\n            config=self.config.model.fsdp_config.wrap_policy,\n            is_lora=self._is_lora,\n        )\n\n        log_gpu_memory_usage(\"Before critic FSDP\", logger=None)\n\n        fsdp_mesh = self.device_mesh\n        sharding_strategy = get_sharding_strategy(fsdp_mesh)\n\n        self.use_orig_params = fsdp_config.get(\"use_orig_params\", False)\n        if self.config.model.get(\"freeze_vision_tower\", False):\n            vision_tower = get_vl_model_vision_tower(critic_module)\n            if vision_tower is not None:\n                vision_tower.requires_grad_(False)\n                self.use_orig_params = True\n                if self.rank == 0:\n                    print(\"[critic model] Vision tower is set to not trainable.\")\n            else:\n                if self.rank == 0:\n                    print(\"[critic model] No vision tower found.\")\n\n        # Note: We force turn off CPUOffload for critic because it causes incorrect results when using grad accumulation\n        if config.strategy == \"fsdp\":\n            critic_module = FSDP(\n                critic_module,\n                param_init_fn=init_fn,\n                use_orig_params=self.use_orig_params,\n                auto_wrap_policy=auto_wrap_policy,\n                device_id=get_device_id(),\n                sharding_strategy=sharding_strategy,\n                mixed_precision=mixed_precision,\n                sync_module_states=True,\n                forward_prefetch=self.config.model.fsdp_config.forward_prefetch,\n                device_mesh=self.device_mesh,\n                cpu_offload=None,\n            )\n        elif config.strategy == \"fsdp2\":\n            assert CPUOffloadPolicy is not None, \"PyTorch version >= 2.4 is required for using fully_shard API (FSDP2)\"\n            mp_policy = MixedPrecisionPolicy(\n                param_dtype=param_dtype, reduce_dtype=reduce_dtype, cast_forward_inputs=True\n            )\n            offload_policy = None\n            if fsdp_config.offload_policy:\n                self._is_offload_param = False\n                self._is_offload_optimizer = False\n                offload_policy = CPUOffloadPolicy(pin_memory=True)\n\n            fsdp_kwargs = {\n                \"mesh\": fsdp_mesh,\n                \"mp_policy\": mp_policy,\n                \"offload_policy\": offload_policy,\n                \"reshard_after_forward\": fsdp_config.reshard_after_forward,\n                \"shard_placement_fn\": get_shard_placement_fn(fsdp_size=self.device_mesh.shape[-1]),\n            }\n            full_state = critic_module.state_dict()\n            apply_fsdp2(critic_module, fsdp_kwargs, fsdp_config)\n            fsdp2_load_full_state_dict(critic_module, full_state, fsdp_mesh, offload_policy)\n        else:\n            raise NotImplementedError(f\"Unknown strategy {config.strategy}\")\n\n        if config.model.get(\"enable_activation_offload\", False):\n            enable_gradient_checkpointing = config.model.get(\"enable_gradient_checkpointing\", False)\n            enable_activation_offloading(critic_module, config.strategy, enable_gradient_checkpointing)\n\n        log_gpu_memory_usage(\"After critic FSDP\", logger=None)\n\n        critic_optimizer = build_optimizer(critic_module.parameters(), config.optim)\n\n        total_steps = config.optim.get(\"total_training_steps\", 0)\n        num_warmup_steps = int(config.optim.get(\"lr_warmup_steps\", -1))\n\n        lr_scheduler_type = config.optim.get(\"lr_scheduler_type\", \"constant\")\n        if num_warmup_steps < 0:\n            num_warmup_steps_ratio = config.optim.get(\"lr_warmup_steps_ratio\", 0.0)\n            num_warmup_steps = int(num_warmup_steps_ratio * total_steps)\n\n        if self.rank == 0:\n            print(f\"Total steps: {total_steps}, num_warmup_steps: {num_warmup_steps}\")\n\n        from verl.utils.torch_functional import get_constant_schedule_with_warmup, get_cosine_schedule_with_warmup\n\n        if lr_scheduler_type == \"constant\":\n            critic_lr_scheduler = get_constant_schedule_with_warmup(\n                optimizer=critic_optimizer, num_warmup_steps=num_warmup_steps\n            )\n        elif lr_scheduler_type == \"cosine\":\n            min_lr_ratio = config.optim.get(\"min_lr_ratio\", 0.0)\n            num_cycles = config.optim.get(\"num_cycles\", 0.5)\n            critic_lr_scheduler = get_cosine_schedule_with_warmup(\n                optimizer=critic_optimizer,\n                num_warmup_steps=num_warmup_steps,\n                num_training_steps=total_steps,\n                min_lr_ratio=min_lr_ratio,\n                num_cycles=num_cycles,\n            )\n        else:\n            raise NotImplementedError(f\"LR scheduler type {lr_scheduler_type} is not supported\")\n\n        return critic_module, critic_optimizer, critic_lr_scheduler\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        # This is used to import external_lib into the huggingface systems\n        import_external_libs(self.config.model.get(\"external_lib\", None))\n\n        from verl.workers.critic import DataParallelPPOCritic\n\n        self.critic_module, self.critic_optimizer, self.critic_lr_scheduler = self._build_critic_model_optimizer(\n            self.config\n        )\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.critic_module)\n            log_gpu_memory_usage(\"After offload critic model during init\", logger=logger)\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(optimizer=self.critic_optimizer)\n            log_gpu_memory_usage(\"After offload critic optimizer during init\", logger=logger)\n\n        self.critic = DataParallelPPOCritic(\n            config=self.config, critic_module=self.critic_module, critic_optimizer=self.critic_optimizer\n        )\n\n        self.flops_counter = FlopsCounter(self.critic_model_config)\n        self.checkpoint_manager = FSDPCheckpointManager(\n            model=self.critic_module,\n            optimizer=self.critic_optimizer,\n            lr_scheduler=self.critic_lr_scheduler,\n            processing_class=self.processor if self.processor is not None else self.tokenizer,\n            checkpoint_config=self.config.checkpoint,\n            trust_remote_code=self.config.model.get(\"trust_remote_code\", False),\n        )\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"critic\"))\n    @DistProfiler.annotate(color=\"cyan\", role=\"compute_values\")\n    def compute_values(self, data: DataProto):\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.critic_module)\n        micro_batch_size = self.config.forward_micro_batch_size_per_gpu\n        data.meta_info[\"micro_batch_size\"] = micro_batch_size\n        data.meta_info[\"max_token_len\"] = self.config.forward_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = self.config.use_dynamic_bsz\n        # perform forward computation\n        with self.ulysses_sharding_manager:\n            data = data.to(\"cpu\")  # data will to device with each micro batch on critic.compute_values\n            values = self.critic.compute_values(data=data)\n            output = DataProto.from_dict(tensors={\"values\": values})\n\n        output = output.to(\"cpu\")\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.critic_module)\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"critic\"))\n    @DistProfiler.annotate(color=\"pink\", role=\"critic_update\")\n    def update_critic(self, data: DataProto):\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.critic_module)\n        if self._is_offload_optimizer:\n            load_fsdp_optimizer(optimizer=self.critic_optimizer, device_id=get_device_id())\n\n        # perform forward computation\n        with self.ulysses_sharding_manager:\n            data = data.to(\"cpu\")  # data will to device with each micro batch on critic.update_critic\n            with Timer(name=\"update_critic\", logger=None) as timer:\n                metrics = self.critic.update_critic(data=data)\n            delta_time = timer.last\n\n            global_num_tokens = data.meta_info[\"global_token_num\"]\n            estimated_flops, promised_flops = self.flops_counter.estimate_flops(global_num_tokens, delta_time)\n            metrics[\"perf/mfu/critic\"] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size\n\n            lr = self.critic_lr_scheduler.get_last_lr()[0]\n            metrics[\"critic/lr\"] = lr\n            self.critic_lr_scheduler.step()\n\n            output = DataProto(batch=None, meta_info={\"metrics\": metrics})\n\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.critic_module)\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(optimizer=self.critic_optimizer)\n\n        output = output.to(\"cpu\")\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, local_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):\n        import torch\n\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.critic_module)\n\n        self.checkpoint_manager.save_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.critic_module)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, local_path, hdfs_path=None, del_local_after_load=True):\n        import torch\n\n        if self._is_offload_param:\n            load_fsdp_model_to_gpu(self.critic_module)\n\n        self.checkpoint_manager.load_checkpoint(\n            local_path=local_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_fsdp_model_to_cpu(self.critic_module)\n\n        if self._is_offload_optimizer:\n            offload_fsdp_optimizer(self.critic_optimizer)\n\n\n# ================================= Async related workers =================================\nclass AsyncActorRolloutRefWorker(ActorRolloutRefWorker):\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    async def update_weights(self, global_steps: int = None):\n        await self.rollout_mode()\n        return True\n"
  },
  {
    "path": "verl/workers/megatron_workers.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe main entry point to run the PPO algorithm\n\"\"\"\n\nimport datetime\nimport logging\nimport os\nimport time\n\nimport psutil\nimport torch\nimport torch.distributed\nfrom codetiming import Timer\nfrom omegaconf import DictConfig, OmegaConf\n\ntry:\n    from verl.workers.engine.mindspeed.transformer_impl import repatch\nexcept ImportError:\n    repatch = None\n\nfrom contextlib import nullcontext\n\nfrom megatron.core import parallel_state as mpu\n\nfrom verl import DataProto\nfrom verl.models.mcore import get_mcore_weight_converter\nfrom verl.single_controller.base import Worker\nfrom verl.single_controller.base.decorator import Dispatch, make_nd_compute_dataproto_dispatch_fn, register\nfrom verl.utils import hf_tokenizer\nfrom verl.utils.checkpoint.megatron_checkpoint_manager import MegatronCheckpointManager\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import (\n    get_device_id,\n    get_device_name,\n    get_nccl_backend,\n    get_torch_device,\n    set_expandable_segments,\n)\nfrom verl.utils.distributed import set_numa_affinity\nfrom verl.utils.flops_counter import FlopsCounter\nfrom verl.utils.fs import copy_to_local\nfrom verl.utils.megatron.router_replay_patch import RouterReplay, RouterReplayAction, apply_router_replay_patch\nfrom verl.utils.megatron_peft_utils import add_base_layer_suffix, build_peft_config_for_vllm\nfrom verl.utils.megatron_utils import (\n    load_megatron_model_to_gpu,\n    load_megatron_optimizer,\n    offload_megatron_model_to_cpu,\n    offload_megatron_optimizer,\n    per_tensor_generator,\n    register_megatron_training_hooks,\n)\nfrom verl.utils.memory_utils import aggressive_empty_cache\nfrom verl.utils.model import get_hf_model_path, load_mcore_dist_weights, load_megatron_gptmodel_weights\nfrom verl.utils.profiler import (\n    DistProfiler,\n    DistProfilerExtension,\n    GPUMemoryLogger,\n    ProfilerConfig,\n    log_gpu_memory_usage,\n    simple_timer,\n)\nfrom verl.utils.profiler.performance import reduce_timing, topk_reduce_ratio_min_max\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.torch_functional import use_original_torch_compile\nfrom verl.workers.actor.megatron_actor import MegatronPPOActor\nfrom verl.workers.config import HFModelConfig, McoreCriticConfig, RolloutConfig\nfrom verl.workers.critic.megatron_critic import MegatronPPOCritic\nfrom verl.workers.rollout import get_rollout_class\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\ndef set_random_seed(seed, only_rollout=False):\n    import random\n\n    import numpy as np\n    import torch\n\n    torch.manual_seed(seed)\n    np.random.seed(seed)\n    random.seed(seed)\n    if not only_rollout and get_torch_device().device_count() > 0:\n        from megatron.core import tensor_parallel\n\n        tensor_parallel.model_parallel_cuda_manual_seed(seed)\n    # FIXME: torch cumsum not support deterministic (used in vllm sampler),\n    # https://github.com/pytorch/pytorch/issues/89492\n    # torch.use_deterministic_algorithms(True, warn_only=True)\n    # os.environ['CUBLAS_WORKSPACE_CONFIG'] = ':4096:8'\n\n\nclass MegatronWorker(Worker):\n    def _init_hf_config_and_tf_config(\n        self,\n        model_path,\n        tokenizer_or_path,\n        dtype,\n        override_model_config,\n        override_transformer_config,\n        trust_remote_code=False,\n        megatron_config=None,\n        enable_mtp=False,\n    ):\n        from transformers import AutoConfig\n\n        from verl.models.mcore import hf_to_mcore_config\n        from verl.utils import hf_processor\n        from verl.utils.model import update_model_config\n\n        # Step 1: initialize the tokenizer\n        self.local_path = copy_to_local(model_path)\n        if tokenizer_or_path is None:\n            self.tokenizer = hf_tokenizer(self.local_path, trust_remote_code=trust_remote_code)\n            self.processor = hf_processor(self.local_path, trust_remote_code=trust_remote_code)\n        elif isinstance(tokenizer_or_path, str):\n            self.tokenizer = hf_tokenizer(copy_to_local(tokenizer_or_path), trust_remote_code=trust_remote_code)\n            self.processor = hf_processor(copy_to_local(tokenizer_or_path), trust_remote_code=trust_remote_code)\n        else:\n            self.tokenizer = tokenizer_or_path\n            self.processor = tokenizer_or_path\n\n        if self.config.model.get(\"custom_chat_template\", None) is not None:\n            if self.processor is not None:\n                self.processor.chat_template = self.config.model.custom_chat_template\n            else:\n                self.tokenizer.chat_template = self.config.model.custom_chat_template\n\n        # Step 2: get the hf\n        hf_config = AutoConfig.from_pretrained(self.local_path, trust_remote_code=trust_remote_code)\n\n        # Step 3: override the hf config\n        override_config_kwargs = {\n            \"bos_token_id\": self.tokenizer.bos_token_id,\n            \"eos_token_id\": self.tokenizer.eos_token_id,\n            \"pad_token_id\": self.tokenizer.pad_token_id,\n        }\n        override_config_kwargs.update(override_model_config.get(\"model_config\", {}))\n        self.share_embeddings_and_output_weights = getattr(hf_config, \"tie_word_embeddings\", False)\n\n        # only actor need enable mtp\n        if enable_mtp:\n            assert hf_config.num_nextn_predict_layers > 0, \"MTP requires at least one nextn_predict_layer\"\n            assert megatron_config.use_mbridge, \"MTP requires use_mbridge to be True\"\n            override_transformer_config[\"mtp_loss_scaling_factor\"] = self.config.model.mtp.mtp_loss_scaling_factor\n        else:\n            if hasattr(hf_config, \"num_nextn_predict_layers\"):\n                hf_config.num_nextn_predict_layers = 0\n\n        self.enable_mtp = enable_mtp\n\n        update_model_config(hf_config, override_config_kwargs=override_config_kwargs)\n        self.architectures = getattr(hf_config, \"architectures\", None)\n        if self.rank == 0:\n            print(f\"Model config after override: {hf_config}\")\n\n        from verl.models.mcore.config_converter import mapping_string_to_attn_backend\n\n        # todo: remove this line after mcore adopt mbridge 0.15, now for compatibility\n        override_transformer_config = mapping_string_to_attn_backend(override_transformer_config)\n        fp16 = dtype == torch.float16\n        bf16 = dtype == torch.bfloat16\n        if fp16:\n            assert megatron_config.use_mbridge, \"fp16 mode requires use_mbridge to be True\"\n\n        self.provider = None\n        self.vanilla_bridge = megatron_config.get(\"vanilla_mbridge\", True)\n        if megatron_config.use_mbridge:\n            if self.vanilla_bridge:\n                from verl.models.mcore.mbridge import AutoBridge\n\n                bridge = AutoBridge.from_config(hf_config, dtype=dtype)\n                bridge.set_extra_args(**override_transformer_config)\n                tf_config = bridge.config\n                tf_config.fp16 = fp16\n                tf_config.bf16 = bf16\n            else:\n                from verl.models.mcore.bridge import AutoBridge\n\n                # Use Megatron-Bridge to convert HF config to Megatron config\n                bridge = AutoBridge.from_hf_pretrained(self.local_path, trust_remote_code=trust_remote_code)\n                # Get Megatron provider and configure it\n                provider = bridge.to_megatron_provider(load_weights=False)\n\n                # In case of invalid overrides, we need to make sure some critical params are set correctly\n                provider.params_dtype = dtype\n\n                # Ensure dtype settings propagate to Megatron-Bridge/TE\n                provider.fp16 = fp16\n                provider.bf16 = bf16\n\n                # Pass distributed info\n                provider.tensor_model_parallel_size = megatron_config.tensor_model_parallel_size\n                provider.pipeline_model_parallel_size = megatron_config.pipeline_model_parallel_size\n                provider.expert_model_parallel_size = megatron_config.expert_model_parallel_size\n                provider.expert_tensor_parallel_size = megatron_config.expert_tensor_parallel_size\n                provider.virtual_pipeline_model_parallel_size = megatron_config.virtual_pipeline_model_parallel_size\n                provider.context_parallel_size = megatron_config.context_parallel_size\n                provider.sequence_parallel = megatron_config.sequence_parallel\n\n                # Match verl implementation (need variable_seq_lengths)\n                from megatron.core.transformer.enums import AttnBackend\n\n                provider.attention_backend = AttnBackend.flash\n                provider.variable_seq_lengths = True\n                provider.moe_token_dispatcher_type = \"alltoall\"\n                provider.moe_router_load_balancing_type = \"none\"\n\n                # Apply transformer config overrides\n                for key, value in override_transformer_config.items():\n                    setattr(provider, key, value)\n\n                provider.finalize()\n                self.provider = provider\n                tf_config = None  # Will be set after model creation\n            self.bridge = bridge\n        else:\n            tf_config = hf_to_mcore_config(hf_config, dtype, **override_transformer_config)\n            self.bridge = None\n\n        if torch.distributed.get_rank() == 0:\n            if tf_config is not None:\n                print(f\"TF config: {tf_config}\")\n        self.hf_config = hf_config\n        self.tf_config = tf_config\n\n        # Get PEFT config from model.lora if specified\n        from verl.workers.config.megatron_peft import get_peft_cls\n\n        self.peft_cls = get_peft_cls(\n            model_config=self.config.model, bridge=self.bridge, provider=self.provider, dtype=dtype\n        )\n\n\nclass ActorRolloutRefWorker(MegatronWorker, DistProfilerExtension):\n    \"\"\"\n    This worker can be instantiated as a standalone actor or a standalone rollout or a standalone reference policy\n    or a hybrid engine based on the config.rollout\n    \"\"\"\n\n    def __init__(self, config: DictConfig, role: str, **kwargs):\n        Worker.__init__(self)\n        self.config = config\n        if repatch is not None:\n            # NPU MindSpeed patch, will be refactored with MindSpeedEngine.\n            repatch(self.config.actor.megatron.get(\"override_transformer_config\", {}))\n\n        self.role = role\n        assert self.role in [\"actor\", \"rollout\", \"ref\", \"actor_rollout\", \"actor_rollout_ref\"]\n\n        self._is_actor = self.role in [\"actor\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_rollout = self.role in [\"rollout\", \"actor_rollout\", \"actor_rollout_ref\"]\n        self._is_ref = self.role in [\"ref\", \"actor_rollout_ref\"]\n\n        # NOTE(sgm): We utilize colocate WorkerGroup by default.\n        # As a result, Workers for different model share the same process.\n        # Therefore, we only require one distribute initialization.\n        # To utilize different parallel strategy in different models:\n        # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models,\n        # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385\n        if not torch.distributed.is_initialized():\n            set_numa_affinity()\n            rank = int(os.environ[\"LOCAL_RANK\"])\n            torch.distributed.init_process_group(\n                backend=f\"cpu:gloo,{get_device_name()}:{get_nccl_backend()}\",\n                timeout=datetime.timedelta(seconds=self.config.get(\"nccl_timeout\", 600)),\n                init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n            )\n            get_torch_device().set_device(rank)\n\n            if self._is_actor or self._is_ref:\n                mpu.initialize_model_parallel(\n                    tensor_model_parallel_size=self.config.actor.megatron.tensor_model_parallel_size,\n                    pipeline_model_parallel_size=self.config.actor.megatron.pipeline_model_parallel_size,\n                    virtual_pipeline_model_parallel_size=self.config.actor.megatron.virtual_pipeline_model_parallel_size,\n                    use_sharp=False,\n                    context_parallel_size=self.config.actor.megatron.context_parallel_size,\n                    expert_model_parallel_size=self.config.actor.megatron.expert_model_parallel_size,\n                    expert_tensor_parallel_size=self.config.actor.megatron.expert_tensor_parallel_size,\n                    nccl_communicator_config_path=None,\n                )\n\n        if self._is_actor or self._is_ref:\n            is_collect = (\n                mpu.get_tensor_model_parallel_rank() == 0\n                and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1\n                and mpu.get_context_parallel_rank() == 0\n            )\n            self._register_dispatch_collect_info(\n                mesh_name=\"actor\", dp_rank=mpu.get_data_parallel_rank(), is_collect=is_collect\n            )\n        only_rollout = self._is_rollout and not self._is_actor\n\n        self.enable_routing_replay = False\n        if self._is_actor:\n            self.router_replay = self.config.actor.router_replay\n            self.enable_routing_replay = self.router_replay.mode != \"disabled\"\n\n        if self.enable_routing_replay:\n            apply_router_replay_patch()\n\n        set_random_seed(seed=self.config.actor.megatron.seed, only_rollout=only_rollout)\n\n        if self._is_actor:\n            omega_profiler_config = config.actor.get(\"profiler\", {})\n        elif self._is_rollout:\n            # NOTE: In colocation mode, rollout config may not take effect (follow the actor config)\n            # This is for extendability in AsyncRL cases\n            omega_profiler_config = config.rollout.get(\"profiler\", {})\n        elif self._is_ref:\n            omega_profiler_config = config.ref.get(\"profiler\", {})\n        else:\n            raise ValueError(\n                f\"Invalid role {self.role}, should be one of \"\n                \"['actor', 'rollout', 'ref', 'actor_rollout', 'actor_rollout_ref']\"\n            )\n        # omega_profiler_config is DictConfig\n        # profiler_config is a ProfilerConfig dataclass\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n\n        # TODO(sgm): Currently, we only support reference model param offload\n        # will support other offload later\n        self._is_offload_param = False\n        self._is_offload_grad = False\n        self._is_offload_optimizer = False\n\n        # Initialize LoRA-related attributes (will be updated in _build_rollout if needed)\n        self.base_sync_done = False\n        self.peft_merge = False\n\n        # normalize config\n        if self._is_actor:\n            self.config.actor.ppo_mini_batch_size *= self.config.rollout.n\n            self.config.actor.ppo_mini_batch_size //= mpu.get_data_parallel_world_size()\n            if self.config.actor.get(\"ppo_micro_batch_size\", None):\n                self.config.actor.ppo_micro_batch_size //= mpu.get_data_parallel_world_size()\n                self.config.rollout.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size()\n                self.config.actor.ppo_micro_batch_size_per_gpu = self.config.actor.ppo_micro_batch_size\n                self.config.rollout.log_prob_micro_batch_size_per_gpu = self.config.rollout.log_prob_micro_batch_size\n\n            self._is_offload_param = self.config.actor.megatron.get(\"param_offload\", False)\n            self._is_offload_grad = self.config.actor.megatron.get(\"grad_offload\", False)\n            self._is_offload_optimizer = self.config.actor.megatron.get(\"optimizer_offload\", False)\n        elif self._is_ref:\n            if self.config.ref.get(\"log_prob_micro_batch_size\", None):\n                self.config.ref.log_prob_micro_batch_size //= mpu.get_data_parallel_world_size()\n                self.config.ref.log_prob_micro_batch_size_per_gpu = self.config.ref.log_prob_micro_batch_size\n            else:\n                assert self.config.ref.get(\"log_prob_micro_batch_size_per_gpu\", None) is not None, (\n                    \"Please note that in the ref policy configuration, `log_prob_micro_batch_size_per_gpu` and \"\n                    \"`log_prob_micro_batch_size` should not be None at the same time.\"\n                )\n            self._ref_is_offload_param = self.config.ref.megatron.get(\"param_offload\", False)\n\n    def _build_model_optimizer(\n        self, model_path, optim_config, override_model_config, override_transformer_config, override_ddp_config=None\n    ):\n        from verl.utils.megatron.optimizer import (\n            get_megatron_optimizer,\n            get_megatron_optimizer_param_scheduler,\n            init_megatron_optim_config,\n        )\n        from verl.utils.megatron_utils import McoreModuleWrapperConfig, make_megatron_module\n        from verl.utils.model import get_generation_config, print_model_size\n\n        self._init_hf_config_and_tf_config(\n            model_path,\n            self.config.model.get(\"tokenizer_path\") or model_path,\n            self.dtype,\n            override_model_config,\n            override_transformer_config,\n            self.config.model.get(\"trust_remote_code\", False),\n            self.config.actor.megatron if not self._is_ref else self.config.ref.megatron,\n            self.config.model.get(\"mtp\", {}).get(\"enable\", False),\n        )\n        self.generation_config = get_generation_config(\n            self.local_path,\n            self.config.model.get(\"trust_remote_code\", False),\n        )\n\n        if self._is_actor or self._is_rollout:\n            wrap_config = McoreModuleWrapperConfig(\n                is_value_model=False,  # actor is not value model\n                share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,\n                wrap_with_ddp=True,\n                use_distributed_optimizer=self.config.actor.megatron.use_distributed_optimizer,\n            )\n            actor_module, updated_tf_config = make_megatron_module(\n                wrap_config=wrap_config,\n                tf_config=self.tf_config,\n                hf_config=self.hf_config,\n                bridge=self.bridge,\n                provider=self.provider,\n                override_model_config=override_model_config,\n                override_ddp_config=override_ddp_config,\n                peft_cls=self.peft_cls,\n                peft_config=self.config.model.get(\"lora\", None),\n            )\n            self.tf_config = updated_tf_config\n            print(f\"actor_module: {len(actor_module)}\")\n            if self.config.actor.load_weight:\n                if self.config.actor.megatron.use_dist_checkpointing:\n                    load_mcore_dist_weights(\n                        actor_module,\n                        self.config.actor.megatron.dist_checkpointing_path,\n                        is_value_model=False,\n                        prefix=self.config.actor.megatron.dist_checkpointing_prefix,\n                    )\n                else:\n                    if self.bridge is not None:\n                        local_model_path = get_hf_model_path(self.config)\n                        if self.vanilla_bridge:\n                            self.bridge.load_weights(actor_module, local_model_path)\n                        else:\n                            self.bridge.load_hf_weights(actor_module, local_model_path)\n                    else:\n                        load_megatron_gptmodel_weights(\n                            self.config, self.hf_config, actor_module, params_dtype=self.dtype, is_value_model=False\n                        )\n\n            if self.rank == 0:\n                print_model_size(actor_module[0])\n            log_gpu_memory_usage(\"After MegatronPPOActor init\", logger=logger)\n        elif self._is_ref:\n            wrap_config = McoreModuleWrapperConfig(\n                is_value_model=False,  # ref is not value model\n                share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,\n                wrap_with_ddp=False,\n                use_distributed_optimizer=self.config.ref.megatron.use_distributed_optimizer,\n            )\n            ref_module, updated_tf_config = make_megatron_module(\n                wrap_config=wrap_config,\n                tf_config=self.tf_config,\n                hf_config=self.hf_config,\n                bridge=self.bridge,\n                provider=self.provider,\n                override_model_config=override_model_config,\n            )\n            self.tf_config = updated_tf_config\n            if self.config.ref.load_weight:  # should align with the actor:\n                assert self.config.actor.load_weight == self.config.ref.load_weight\n                print(\"load ref weight start\")\n                if self.config.ref.megatron.use_dist_checkpointing:\n                    load_mcore_dist_weights(\n                        ref_module,\n                        self.config.ref.megatron.dist_checkpointing_path,\n                        is_value_model=False,\n                        prefix=self.config.ref.megatron.dist_checkpointing_prefix,\n                    )\n                else:\n                    if self.bridge is not None:\n                        local_model_path = get_hf_model_path(self.config)\n                        if self.vanilla_bridge:\n                            self.bridge.load_weights(ref_module, local_model_path)\n                        else:\n                            self.bridge.load_hf_weights(ref_module, local_model_path)\n                    else:\n                        load_megatron_gptmodel_weights(\n                            self.config, self.hf_config, ref_module, params_dtype=self.dtype, is_value_model=False\n                        )\n            log_gpu_memory_usage(\"After ref module init\", logger=logger)\n            return ref_module, self.hf_config\n\n        # TODO: add more optimizer args into config\n        if self._is_actor:\n            optim_config_megatron = init_megatron_optim_config(\n                optim_config,\n                use_distributed_optimizer=wrap_config.use_distributed_optimizer,\n                fp16=self.dtype == torch.float16,\n            )\n            actor_optimizer = get_megatron_optimizer(model=actor_module, config=optim_config_megatron)\n            actor_optimizer_scheduler = get_megatron_optimizer_param_scheduler(\n                optimizer=actor_optimizer, config=optim_config\n            )\n        else:\n            optim_config = None\n            actor_optimizer = None\n            actor_optimizer_scheduler = None\n\n        log_gpu_memory_usage(\"After actor optimizer init\", logger=logger)\n\n        register_megatron_training_hooks(actor_module, actor_optimizer)\n\n        return actor_module, actor_optimizer, actor_optimizer_scheduler, self.hf_config, optim_config\n\n    def _build_rollout(self, trust_remote_code=False):\n        from torch.distributed.device_mesh import init_device_mesh\n\n        # 1. parse rollout and huggingface model config\n        rollout_config: RolloutConfig = omega_conf_to_dataclass(self.config.rollout)\n        model_config: HFModelConfig = omega_conf_to_dataclass(self.config.model)\n\n        # 2. build rollout device mesh\n        infer_tp = self.config.rollout.tensor_model_parallel_size * self.config.rollout.data_parallel_size\n        infer_pp = self.config.rollout.pipeline_model_parallel_size\n        infer_world_size = infer_tp * infer_pp\n        dp = self.world_size // infer_world_size\n        assert self.world_size % infer_world_size == 0, (\n            f\"rollout world_size: {self.world_size} is not divisible by infer_world_size: {infer_world_size}\"\n        )\n        rollout_device_mesh = init_device_mesh(\n            get_device_name(), mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=[\"dp\", \"infer_tp\", \"infer_pp\"]\n        )\n\n        self.rollout_device_mesh = rollout_device_mesh\n\n        is_collect = (\n            rollout_device_mesh[\"infer_tp\"].get_local_rank() == 0\n            and rollout_device_mesh[\"infer_pp\"].get_local_rank() == 0\n        )\n        self._register_dispatch_collect_info(\n            \"rollout\", dp_rank=rollout_device_mesh[\"dp\"].get_local_rank(), is_collect=is_collect\n        )\n\n        # 4. build rollout model\n        log_gpu_memory_usage(f\"Before building {self.config.rollout.name} rollout\", logger=logger)\n        self.rollout = get_rollout_class(rollout_config.name, rollout_config.mode)(\n            config=rollout_config, model_config=model_config, device_mesh=rollout_device_mesh\n        )\n        log_gpu_memory_usage(f\"After building {self.config.rollout.name} rollout\", logger=logger)\n\n        # Initialize base_sync_done for LoRA\n        self.base_sync_done: bool = \"dummy\" not in self.config.rollout.load_format\n        self.peft_merge: bool = model_config.lora.get(\"merge\", False)\n\n        # 5. switch to trainer mode\n        # NOTE: It's critical that hybrid engine in trainer mode initially to load checkpoint.\n        # For async mode, we can't call run_until_complete here, so we will switch to trainer mode in AgentLoopManager.\n        # Note: sync mode is deprecated and rejected in RolloutConfig.__post_init__\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        if self.config.model.get(\"external_lib\", None) is not None:\n            # This is used to import external_lib into the huggingface systems\n            import importlib\n\n            importlib.import_module(self.config.model.external_lib)\n\n        from verl.utils.torch_dtypes import PrecisionType\n\n        override_model_config = OmegaConf.to_container(OmegaConf.create(self.config.model.get(\"override_config\", {})))\n        if self._is_actor:\n            override_transformer_config = OmegaConf.to_container(\n                OmegaConf.create(self.config.actor.megatron.get(\"override_transformer_config\", {}))\n            )\n            if self.enable_routing_replay:\n                override_transformer_config[\"enable_routing_replay\"] = True\n            override_ddp_config = OmegaConf.to_container(\n                OmegaConf.create(self.config.actor.megatron.get(\"override_ddp_config\", {}))\n            )\n        elif self._is_ref:\n            override_transformer_config = OmegaConf.to_container(\n                OmegaConf.create(self.config.ref.megatron.get(\"override_transformer_config\", {}))\n            )\n        else:\n            override_transformer_config = {}\n        self.param_dtype = PrecisionType.to_dtype(self.config.actor.megatron.dtype)\n        log_gpu_memory_usage(\"Before init actor model and optimizer\", logger=logger)\n        self.dtype = PrecisionType.to_dtype(self.param_dtype)\n        if self._is_actor:\n            # we need the model for actor and rollout\n            optim_config = self.config.actor.optim if self._is_actor else None\n            (\n                self.actor_module,\n                self.actor_optimizer,\n                self.actor_optimizer_scheduler,\n                self.actor_model_config,\n                self.actor_optim_config,\n            ) = self._build_model_optimizer(\n                model_path=self.config.model.path,\n                optim_config=optim_config,\n                override_model_config=override_model_config,\n                override_transformer_config=override_transformer_config,\n                override_ddp_config=override_ddp_config,\n            )\n            if self._is_offload_param:\n                offload_megatron_model_to_cpu(self.actor_module)\n                log_gpu_memory_usage(\"After offload actor params and grad during init\", logger=logger)\n            if self._is_offload_optimizer:\n                offload_megatron_optimizer(self.actor_optimizer)\n                log_gpu_memory_usage(\"After offload actor optimizer during init\", logger=logger)\n\n        if self._is_actor:\n            actor_cfg = omega_conf_to_dataclass(self.config.actor)\n            self.actor = MegatronPPOActor(\n                config=actor_cfg,\n                model_config=self.actor_model_config,\n                hf_config=self.hf_config,\n                tf_config=self.tf_config,\n                actor_module=self.actor_module,\n                actor_optimizer=self.actor_optimizer,\n                mtp_config=self.config.model.mtp if self.config.model.mtp.enable else None,\n            )\n            print(f\"routing replay layers: {len(RouterReplay.router_instances)}\")\n            log_gpu_memory_usage(\"After MegatronPPOActor init\", logger=logger)\n\n        if self._is_rollout:\n            with use_original_torch_compile():\n                self._build_rollout(trust_remote_code=self.config.model.get(\"trust_remote_code\", False))\n            log_gpu_memory_usage(\"After rollout init\", logger=logger)\n\n        if self._is_ref:\n            self.ref_module, self.ref_model_config = self._build_model_optimizer(\n                model_path=self.config.model.path,\n                optim_config=None,\n                override_model_config=override_model_config,\n                override_transformer_config=override_transformer_config,\n            )\n            log_gpu_memory_usage(\"After ref model init\", logger=logger)\n            self.ref_policy = MegatronPPOActor(\n                config=self.config.ref,\n                model_config=self.ref_model_config,\n                hf_config=self.hf_config,\n                tf_config=self.tf_config,\n                actor_module=self.ref_module,\n                actor_optimizer=None,\n            )\n            if self._ref_is_offload_param:\n                offload_megatron_model_to_cpu(self.ref_module)\n                log_gpu_memory_usage(\"After offload ref params during init\", logger=logger)\n\n        if self._is_actor:\n            self.flops_counter = FlopsCounter(self.actor_model_config)\n            self.checkpoint_mananager = MegatronCheckpointManager(\n                config=self.config,\n                checkpoint_config=self.config.actor.checkpoint,\n                model_config=self.actor_model_config,\n                transformer_config=self.tf_config,\n                role=\"actor\",\n                model=self.actor_module,\n                arch=self.architectures[0],\n                hf_config=self.hf_config,\n                param_dtype=self.param_dtype,\n                share_embeddings_and_output_weights=self.share_embeddings_and_output_weights,\n                processing_class=self.processor if self.processor is not None else self.tokenizer,\n                optimizer=self.actor_optimizer,\n                optimizer_scheduler=self.actor_optimizer_scheduler,\n                use_distributed_optimizer=self.config.actor.megatron.use_distributed_optimizer,\n                use_checkpoint_opt_param_scheduler=self.config.actor.optim.use_checkpoint_opt_param_scheduler,\n                bridge=self.bridge,\n                provider=self.provider,\n                use_dist_checkpointing=self.config.actor.megatron.use_dist_checkpointing,\n                peft_cls=self.peft_cls,\n            )\n\n            self.layer_name_mapping = {\n                \"qkv_layer_name\": \"self_attention.linear_qkv.\",\n                \"gate_proj_layer_name\": \"linear_fc1.\",\n            }\n            self.weight_converter = None\n            if not self.config.actor.megatron.use_mbridge:\n                self.weight_converter = get_mcore_weight_converter(self.actor_model_config, self.dtype)\n\n        # Free cached GPU memory so colocated vLLM processes can see it via cudaMemGetInfo\n        aggressive_empty_cache(force_sync=True)\n        log_gpu_memory_usage(\"After init_model finish\", logger=logger)\n\n    async def rollout_mode(self):\n        \"\"\"Context switch hybridengine to rollout mode.\"\"\"\n        aggressive_empty_cache(force_sync=True)\n        set_expandable_segments(False)\n\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor.actor_module, load_grad=False)\n            log_gpu_memory_usage(\"After load actor params during rollout_mode\", logger=logger)\n\n        # Build peft_config for vLLM LoRA support\n        peft_config = None\n        do_lora_base_sync = False\n        if not self.peft_merge and self.peft_cls is not None:\n            peft_config = build_peft_config_for_vllm(self.config.model.get(\"lora\", {}))\n            # set sleep level for LoRA adapter weights only sync\n            # TODO: make this configurable so that users with small\n            # main memory can trade sync time to avoid OOM\n            self.rollout.sleep_level = 1\n\n            do_lora_base_sync = (not self.base_sync_done) or (\n                self.rollout.sleep_level != 1 and self.config.rollout.free_cache_engine\n            )\n\n        if self.bridge is not None:\n            if self.vanilla_bridge:\n                per_tensor_param = self.bridge.export_weights(self.actor.actor_module)\n            elif not self.peft_merge and self.peft_cls is not None:\n                # Only export adapter weights\n                per_tensor_param = self.bridge.export_adapter_weights(self.actor.actor_module)\n            else:\n                per_tensor_param = self.bridge.export_hf_weights(self.actor.actor_module)\n        else:\n            per_tensor_param = per_tensor_generator(\n                self.actor.actor_module,\n                self.actor_model_config,\n                self.weight_converter,\n                self.tf_config,\n                self.layer_name_mapping,\n            )\n\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"weights\"])\n        if do_lora_base_sync:\n            # Base layer sync\n            per_tensor_param_lora_base = self.bridge.export_hf_weights(\n                self.actor.actor_module, merge_adapter_weights=False\n            )\n            await self.rollout.update_weights(\n                add_base_layer_suffix(per_tensor_param_lora_base, model_type=self.hf_config.model_type),\n                peft_config=peft_config,\n                base_sync_done=False,\n            )\n\n            # Mark base sync as done after first successful sync\n            self.base_sync_done = True\n\n        await self.rollout.update_weights(per_tensor_param, peft_config=peft_config, base_sync_done=True)\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.actor.actor_module)\n        aggressive_empty_cache(force_sync=True)\n        if self.config.rollout.free_cache_engine:\n            await self.rollout.resume(tags=[\"kv_cache\"])\n\n        set_expandable_segments(True)\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @GPUMemoryLogger(role=\"update_actor\", logger=logger)\n    @DistProfiler.annotate(color=\"red\", role=\"actor_update\")\n    def update_actor(self, data: DataProto):\n        assert self._is_actor\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor_module)\n            log_gpu_memory_usage(\"After load actor params and grad during update_actor\", logger=logger)\n        if self._is_offload_optimizer:\n            load_megatron_optimizer(self.actor_optimizer)\n            log_gpu_memory_usage(\"After load actor optimizer during update_actor\", logger=logger)\n\n        micro_batch_size = self.config.actor.ppo_micro_batch_size_per_gpu\n        data.meta_info[\"micro_batch_size\"] = micro_batch_size\n        dataloader = self.actor.make_minibatch_iterator(data=data)\n        with Timer(name=\"update_policy\", logger=None) as timer:\n            metrics = self.actor.update_policy(dataloader=dataloader)\n        delta_time = timer.last\n        global_num_tokens = data.meta_info[\"global_token_num\"]\n        images_seqlens = data.meta_info.get(\"images_seqlens\", None)\n        estimated_flops, promised_flops = self.flops_counter.estimate_flops(\n            global_num_tokens, delta_time, images_seqlens=images_seqlens\n        )\n        metrics[\"perf/mfu/actor\"] = estimated_flops * self.config.actor.ppo_epochs / promised_flops / self.world_size\n        metrics[\"perf/max_memory_allocated_gb\"] = get_torch_device().max_memory_allocated() / (1024**3)\n        metrics[\"perf/max_memory_reserved_gb\"] = get_torch_device().max_memory_reserved() / (1024**3)\n        metrics[\"perf/cpu_memory_used_gb\"] = psutil.virtual_memory().used / (1024**3)\n        from verl.utils.megatron.optimizer import get_megatron_last_lr\n\n        metrics[\"actor/lr\"] = get_megatron_last_lr(self.actor_optimizer)\n        self.actor_optimizer_scheduler.step(1)\n\n        # TODO: here, we should return all metrics\n        output = DataProto(meta_info={\"metrics\": metrics})\n        output = output.to(\"cpu\")\n\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.actor_module)\n            log_gpu_memory_usage(\"After offload actor params and grad during update_actor\", logger=logger)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.actor_optimizer)\n            log_gpu_memory_usage(\"After offload actor optimizer during update_actor\", logger=logger)\n\n        aggressive_empty_cache(force_sync=True)\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"rollout\"))\n    @GPUMemoryLogger(role=\"generate_sequences\", logger=logger)\n    @DistProfiler.annotate(color=\"red\", role=\"rollout_generate\")\n    def generate_sequences(self, prompts: DataProto):\n        assert self._is_rollout\n        prompts = prompts.to(get_device_name())\n        meta_info = {\n            \"eos_token_id\": self.generation_config.eos_token_id\n            if self.generation_config is not None\n            else self.tokenizer.eos_token_id,\n            \"pad_token_id\": self.generation_config.pad_token_id\n            if self.generation_config is not None\n            else self.tokenizer.pad_token_id,\n        }\n        prompts.meta_info.update(meta_info)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.actor_optimizer)\n\n        timing_generate = {}\n        if self._is_actor:  # For rollout only, we do not switch context.\n            loop = get_event_loop()\n            loop.run_until_complete(self.rollout_mode())\n            log_gpu_memory_usage(\"After switch to rollout mode\", logger=logger)\n\n        with simple_timer(\"generate_sequences\", timing_generate):\n            output = self.rollout.generate_sequences(prompts=prompts)\n\n        if self._is_actor:\n            loop.run_until_complete(self.trainer_mode())\n            log_gpu_memory_usage(\"After switch to trainer mode\", logger=logger)\n\n        # We calculate the average timing across all ranks\n        # to make sure meta_info[\"timing\"] is the same\n        timing_generate_topk_ratio, timing_generate_min, timing_generate_max = topk_reduce_ratio_min_max(\n            timing_generate[\"generate_sequences\"]\n        )\n        timing_generate = reduce_timing(timing_generate)\n        timing_generate.update(\n            {\n                \"generation_timing/max\": timing_generate_max,\n                \"generation_timing/min\": timing_generate_min,\n                \"generation_timing/topk_ratio\": timing_generate_topk_ratio,\n            }\n        )\n        output.meta_info[\"timing\"] = timing_generate\n        output = output.to(\"cpu\")\n        # clear kv cache\n        aggressive_empty_cache(force_sync=True)\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @GPUMemoryLogger(role=\"compute_ref_log_prob\", logger=logger)\n    @DistProfiler.annotate(color=\"olive\", role=\"ref_compute_log_prob\")\n    def compute_ref_log_prob(self, data: DataProto):\n        if self.peft_cls is not None:\n            # if is lora, actor without lora applied is the ref\n            data.meta_info[\"is_lora\"] = True\n            return self.compute_log_prob(data)\n        assert self._is_ref\n        if self._ref_is_offload_param:\n            load_megatron_model_to_gpu(self.ref_module, load_grad=False)\n            log_gpu_memory_usage(\"After load ref params and grad during compute_ref_log_prob\", logger=logger)\n        micro_batch_size = self.config.ref.log_prob_micro_batch_size_per_gpu\n        data.meta_info[\"micro_batch_size\"] = micro_batch_size\n        data.meta_info[\"max_token_len\"] = self.config.ref.log_prob_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = self.config.ref.log_prob_use_dynamic_bsz\n        data.meta_info[\"temperature\"] = self.config.rollout.temperature\n        output, _, _ = self.ref_policy.compute_log_prob(data=data, calculate_entropy=False)\n        output = DataProto.from_dict(tensors={\"ref_log_prob\": output})\n        output = output.to(\"cpu\")\n        if self._ref_is_offload_param:\n            offload_megatron_model_to_cpu(self.ref_module)\n            log_gpu_memory_usage(\"After offload ref params and grad during compute_ref_log_prob\", logger=logger)\n        aggressive_empty_cache(force_sync=True)\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"actor\"))\n    @GPUMemoryLogger(role=\"compute_log_prob\", logger=logger)\n    @DistProfiler.annotate(color=\"blue\", role=\"actor_compute_log_prob\")\n    def compute_log_prob(self, data: DataProto):\n        assert self._is_actor\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor_module, load_grad=False)\n            log_gpu_memory_usage(\"After load actor params and grad during compute_log_prob\", logger=logger)\n        is_lora = data.meta_info.pop(\"is_lora\", False)\n        adapter_ctx = self.peft_cls.disable_adapter(self.actor_module) if is_lora else nullcontext()\n        # we should always recompute old_log_probs when it is HybridEngine\n        config_source = self.config.ref if is_lora else self.config.rollout\n        data.meta_info[\"micro_batch_size\"] = config_source.log_prob_micro_batch_size_per_gpu\n        data.meta_info[\"max_token_len\"] = config_source.log_prob_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = config_source.log_prob_use_dynamic_bsz\n        data.meta_info[\"temperature\"] = self.config.rollout.temperature\n\n        if self.enable_routing_replay and self.config.actor.router_replay.mode == \"R2\":\n            RouterReplay.set_global_router_replay_action(RouterReplayAction.RECORD)\n\n        if self.enable_routing_replay and self.config.actor.router_replay.mode == \"R3\":\n            RouterReplay.set_global_router_replay_action(RouterReplayAction.REPLAY_FORWARD)\n\n        with adapter_ctx:\n            output, entropys, layers_topk_idx = self.actor.compute_log_prob(data=data, calculate_entropy=not is_lora)\n        tensors = {\"ref_log_prob\": output} if is_lora else {\"old_log_probs\": output}\n        if not is_lora:\n            tensors[\"entropys\"] = entropys\n        output = DataProto.from_dict(\n            tensors=tensors,\n            meta_info={\"temperature\": self.config.rollout.temperature},\n        )\n        if self.config.actor.router_replay.mode == \"R2\":\n            output.batch[\"routed_experts\"] = layers_topk_idx\n\n        if self.config.actor.router_replay.mode in [\"R2\", \"R3\"]:\n            RouterReplay.clear_global_indices()\n            RouterReplay.clear_global_router_replay_action()\n\n        output = output.to(\"cpu\")\n        # clear kv cache\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.actor_module)\n            log_gpu_memory_usage(\"After offload actor params and grad during compute_log_prob\", logger=logger)\n        aggressive_empty_cache(force_sync=True)\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load=True):\n        # No checkpoint to load, just offload the model and optimizer to CPU\n        if checkpoint_path is None:\n            if self._is_offload_param:\n                offload_megatron_model_to_cpu(self.actor_module)\n            if self._is_offload_optimizer:\n                offload_megatron_optimizer(self.actor_optimizer)\n            log_gpu_memory_usage(\"After offload actor params and optimizer during load_checkpoint\", logger=logger)\n            return\n\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor_module)\n        self.checkpoint_mananager.load_checkpoint(\n            local_path=checkpoint_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.actor_module)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.actor_optimizer)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_pretrained_model(self, checkpoint_path, del_local_after_load=True):\n        pass\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, checkpoint_path, hdfs_path=None, global_step=0, max_ckpt_to_keep=None):\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.actor_module)\n        if self.checkpoint_mananager.checkpoint_config.async_save and self._is_offload_optimizer:\n            load_megatron_optimizer(self.actor_optimizer)\n        self.checkpoint_mananager.save_checkpoint(\n            local_path=checkpoint_path, hdfs_path=hdfs_path, global_step=global_step, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n        torch.distributed.barrier()\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.actor_module)\n        if self.checkpoint_mananager.checkpoint_config.async_save and self._is_offload_optimizer:\n            offload_megatron_optimizer(self.actor_optimizer)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def async_calls_finalize_fn_exec(self, blocking=False):\n        from megatron.core.dist_checkpointing.strategies.base import async_calls\n\n        async_calls.maybe_finalize_async_calls(blocking=blocking)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def start_profile(self, **kwargs) -> None:\n        \"\"\"Start profiling for the current rank in the current training step.\"\"\"\n        self.profiler.start(**kwargs)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def stop_profile(self) -> None:\n        \"\"\"Stop profiling for the current rank in the current training step.\"\"\"\n        self.profiler.stop()\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def dump_memory_snapshot(self, tag: str = \"manual\", sub_dir: str = None) -> None:\n        \"\"\"Manually trigger a CUDA memory snapshot dump on all ranks.\"\"\"\n        # Memory snapshot is now handled by the profiler system\n        # This method is kept for backward compatibility but delegates to profiler\n        if hasattr(self, \"profiler\") and hasattr(self.profiler, \"_impl\"):\n            try:\n                # Try to use the profiler's memory snapshot functionality\n                if hasattr(self.profiler._impl, \"sampler\"):\n                    out_dir = OmegaConf.select(self.config, \"actor.profiler.save_path\") or \".\"\n                    self.profiler._impl.sampler.dump_memory_snapshot(out_dir=out_dir, tag=tag, sub_dir=sub_dir)\n            except Exception as e:\n                # Log a warning if memory snapshot fails. This might be expected if the profiler doesn't support it.\n                logger.warning(f\"Failed to dump memory snapshot: {e}\")\n\n\nclass AsyncActorRolloutRefWorker(ActorRolloutRefWorker):\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL, blocking=False)\n    async def update_weights(self, global_steps: int = None):\n        await self.rollout_mode()\n        return True\n\n\nclass CriticWorker(MegatronWorker, DistProfilerExtension):\n    def __init__(self, config: McoreCriticConfig):\n        Worker.__init__(self)\n\n        omega_profiler_config = config.get(\"profiler\", {})\n        profiler_config = omega_conf_to_dataclass(omega_profiler_config, dataclass_type=ProfilerConfig)\n        if omega_profiler_config.get(\"tool\", None) in [\"npu\", \"nsys\", \"torch\", \"torch_memory\"]:\n            tool_config = omega_conf_to_dataclass(\n                omega_profiler_config.get(\"tool_config\", {}).get(omega_profiler_config.get(\"tool\"))\n            )\n        else:\n            tool_config = None\n        DistProfilerExtension.__init__(\n            self, DistProfiler(rank=self.rank, config=profiler_config, tool_config=tool_config)\n        )\n        self.config: McoreCriticConfig = config\n\n        # NOTE(sgm): We utilize colocate WorkerGroup by default.\n        # As a result, Workers for different model share the same process.\n        # Therefore, we only require one distribute initialization.\n        # To utilize different parallel strategy in different models:\n        # 1, users should disable WorkerDict; 2.assign different ResourcePool to different models,\n        # 3. and apply the following patch in ray==2.10, https://github.com/ray-project/ray/pull/44385\n        if not torch.distributed.is_initialized():\n            set_numa_affinity()\n            rank = int(os.environ[\"LOCAL_RANK\"])\n            torch.distributed.init_process_group(\n                backend=get_nccl_backend(),\n                timeout=datetime.timedelta(seconds=self.config.get(\"nccl_timeout\", 600)),\n                init_method=os.environ.get(\"DIST_INIT_METHOD\", None),\n            )\n            get_torch_device().set_device(rank)\n\n            mpu.initialize_model_parallel(\n                tensor_model_parallel_size=self.config.megatron.tensor_model_parallel_size,\n                pipeline_model_parallel_size=self.config.megatron.pipeline_model_parallel_size,\n                virtual_pipeline_model_parallel_size=self.config.megatron.virtual_pipeline_model_parallel_size,\n                use_sharp=False,\n                context_parallel_size=self.config.megatron.context_parallel_size,\n                expert_model_parallel_size=self.config.megatron.expert_model_parallel_size,\n                expert_tensor_parallel_size=self.config.megatron.expert_tensor_parallel_size,\n                nccl_communicator_config_path=None,\n            )\n\n        is_collect = (\n            mpu.get_tensor_model_parallel_rank() == 0\n            and mpu.get_pipeline_model_parallel_rank() == mpu.get_pipeline_model_parallel_world_size() - 1\n            and mpu.get_context_parallel_rank() == 0\n        )\n        self._register_dispatch_collect_info(\n            mesh_name=\"critic\", dp_rank=mpu.get_data_parallel_rank(), is_collect=is_collect\n        )\n\n        set_random_seed(seed=self.config.megatron.seed)\n\n        # set FSDP offload params\n        self._is_offload_param = self.config.megatron.param_offload\n        self._is_offload_optimizer = self.config.megatron.optimizer_offload\n\n        # normalize config\n        self.config.ppo_mini_batch_size *= self.config.rollout_n\n        self.config.ppo_mini_batch_size //= mpu.get_data_parallel_world_size()\n        if self.config.get(\"ppo_micro_batch_size\", None):\n            self.config.ppo_micro_batch_size //= mpu.get_data_parallel_world_size()\n            self.config.ppo_micro_batch_size_per_gpu = self.config.ppo_micro_batch_size\n\n        # TODO(sgm): support critic model offload\n\n    def _build_critic_model_optimizer(\n        self, model_path, optim_config, override_model_config, override_transformer_config, override_ddp_config\n    ):\n        from verl.utils.megatron.optimizer import (\n            get_megatron_optimizer,\n            get_megatron_optimizer_param_scheduler,\n            init_megatron_optim_config,\n        )\n        from verl.utils.megatron_utils import McoreModuleWrapperConfig, make_megatron_module\n        from verl.utils.model import print_model_size\n\n        self._init_hf_config_and_tf_config(\n            model_path,\n            self.config.model.get(\"tokenizer_path\") or model_path,\n            self.dtype,\n            override_model_config,\n            override_transformer_config,\n            self.config.model.get(\"trust_remote_code\", False),\n            self.config.megatron,\n        )\n\n        wrap_config = McoreModuleWrapperConfig(\n            is_value_model=True,  # critic is value model\n            share_embeddings_and_output_weights=False,\n            wrap_with_ddp=True,\n            use_distributed_optimizer=self.config.megatron.use_distributed_optimizer,\n        )\n        critic_module, updated_tf_config = make_megatron_module(\n            wrap_config=wrap_config,\n            tf_config=self.tf_config,\n            hf_config=self.hf_config,\n            bridge=self.bridge,\n            provider=self.provider,\n            override_model_config=override_model_config,\n            override_ddp_config=override_ddp_config,\n            peft_cls=self.peft_cls,\n            peft_config=self.config.model.get(\"lora\", None),\n        )\n        self.tf_config = updated_tf_config\n        # note that here critic_module will be a list to be compatible with the construction of interleaved pp (vpp).\n        # but here, we do not use pp (vpp) yet. For simplicity, we remove the list\n        # critic_module = nn.ModuleList(critic_module)\n\n        if self.config.load_weight:\n            t0 = time.time()\n            if self.config.megatron.use_dist_checkpointing:\n                load_mcore_dist_weights(\n                    critic_module,\n                    self.config.megatron.dist_checkpointing_path,\n                    is_value_model=True,\n                    prefix=self.config.megatron.dist_checkpointing_prefix,\n                )\n            else:\n                if self.bridge is not None:\n                    local_model_path = get_hf_model_path(self.config)\n                    if self.vanilla_bridge:\n                        self.bridge.load_weights(critic_module, local_model_path)\n                    else:\n                        self.bridge.load_hf_weights(\n                            critic_module, local_model_path, allowed_mismatched_params=[\"output_layer.weight\"]\n                        )\n                else:\n                    load_megatron_gptmodel_weights(\n                        self.config, self.hf_config, critic_module, params_dtype=self.dtype, is_value_model=True\n                    )\n            t1 = time.time()\n            if torch.distributed.get_rank() == 0:\n                print(f\"critic load_weight time: {t1 - t0}\")\n        if self.rank == 0:\n            print_model_size(critic_module[0])\n\n        # TODO: add more optimizer args into config\n        optim_config_megatron = init_megatron_optim_config(\n            optim_config,\n            use_distributed_optimizer=wrap_config.use_distributed_optimizer,\n            fp16=self.dtype == torch.float16,\n        )\n        critic_optimizer = get_megatron_optimizer(model=critic_module, config=optim_config_megatron)\n        critic_optimizer_scheduler = get_megatron_optimizer_param_scheduler(\n            optimizer=critic_optimizer, config=optim_config\n        )\n        get_torch_device().empty_cache()\n\n        register_megatron_training_hooks(critic_module, critic_optimizer)\n\n        return critic_module, critic_optimizer, critic_optimizer_scheduler, self.hf_config, optim_config\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def init_model(self):\n        # create critic\n\n        from verl.utils.torch_dtypes import PrecisionType\n\n        if self.config.model.get(\"external_lib\", None) is not None:\n            # This is used to import external_lib into the huggingface systems\n            import importlib\n\n            importlib.import_module(self.config.model.external_lib)\n        override_model_config = OmegaConf.to_container(OmegaConf.create(self.config.model.get(\"override_config\", {})))\n        override_transformer_config = OmegaConf.to_container(\n            OmegaConf.create(self.config.megatron.get(\"override_transformer_config\", {}))\n        )\n        override_ddp_config = OmegaConf.to_container(\n            OmegaConf.create(self.config.megatron.get(\"override_ddp_config\", {}))\n        )\n        self.param_dtype = PrecisionType.to_dtype(self.config.megatron.dtype)\n        self.dtype = PrecisionType.to_dtype(self.param_dtype)\n        (\n            self.critic_module,\n            self.critic_optimizer,\n            self.critic_optimizer_scheduler,\n            self.critic_model_config,\n            critic_optimizer_config,\n        ) = self._build_critic_model_optimizer(\n            model_path=self.config.model.path,\n            optim_config=self.config.optim,\n            override_model_config=override_model_config,\n            override_transformer_config=override_transformer_config,\n            override_ddp_config=override_ddp_config,\n        )\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.critic_module)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.critic_optimizer)\n\n        self.critic = MegatronPPOCritic(\n            config=self.config,\n            model_config=self.critic_model_config,\n            hf_config=self.hf_config,\n            tf_config=self.tf_config,\n            critic_module=self.critic_module,\n            critic_optimizer=self.critic_optimizer,\n            critic_optimizer_config=critic_optimizer_config,\n        )\n        self.flops_counter = FlopsCounter(self.critic_model_config)\n        self.checkpoint_mananager = MegatronCheckpointManager(\n            config=self.config,\n            checkpoint_config=self.config.checkpoint,\n            model_config=self.critic_model_config,\n            transformer_config=self.tf_config,\n            role=\"critic\",\n            model=self.critic_module,\n            arch=self.architectures[0],\n            hf_config=self.hf_config,\n            param_dtype=self.param_dtype,\n            share_embeddings_and_output_weights=False,\n            processing_class=self.processor if self.processor is not None else self.tokenizer,\n            optimizer=self.critic_optimizer,\n            optimizer_scheduler=self.critic_optimizer_scheduler,\n            use_distributed_optimizer=self.config.megatron.use_distributed_optimizer,\n            use_checkpoint_opt_param_scheduler=self.config.optim.use_checkpoint_opt_param_scheduler,\n            bridge=self.bridge,\n            provider=self.provider,\n            use_dist_checkpointing=self.config.megatron.use_dist_checkpointing,\n            peft_cls=self.peft_cls,\n        )\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"critic\"))\n    @DistProfiler.annotate(color=\"cyan\", role=\"compute_values\")\n    def compute_values(self, data: DataProto):\n        micro_batch_size = self.config.ppo_micro_batch_size_per_gpu\n        data.meta_info[\"micro_batch_size\"] = micro_batch_size\n        data.meta_info[\"max_token_len\"] = self.config.forward_max_token_len_per_gpu\n        data.meta_info[\"use_dynamic_bsz\"] = self.config.use_dynamic_bsz\n        data = data.to(get_device_id())\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.critic_module)\n        values = self.critic.compute_values(data=data)\n        output = DataProto.from_dict(tensors={\"values\": values})\n        output = output.to(\"cpu\")\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.critic_module)\n        return output\n\n    @register(dispatch_mode=make_nd_compute_dataproto_dispatch_fn(mesh_name=\"critic\"))\n    @DistProfiler.annotate(color=\"pink\", role=\"critic_update\")\n    def update_critic(self, data: DataProto):\n        data = data.to(get_device_id())\n\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.critic_module)\n        if self._is_offload_optimizer:\n            load_megatron_optimizer(self.critic_optimizer)\n\n        dataloader = self.critic.make_minibatch_iterator(data)\n        with Timer(name=\"update_critic\", logger=None) as timer:\n            metrics = self.critic.update_critic(dataloader=dataloader)\n        delta_time = timer.last\n        global_num_tokens = data.meta_info[\"global_token_num\"]\n        estimated_flops, promised_flops = self.flops_counter.estimate_flops(global_num_tokens, delta_time)\n        metrics[\"perf/mfu/critic\"] = estimated_flops * self.config.ppo_epochs / promised_flops / self.world_size\n        from verl.utils.megatron.optimizer import get_megatron_last_lr\n\n        metrics[\"critic/lr\"] = get_megatron_last_lr(self.critic_optimizer)\n        self.critic_optimizer_scheduler.step(1)\n\n        output = DataProto(batch=None, meta_info={\"metrics\": metrics})\n\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.critic_module)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.critic_optimizer)\n        output = output.to(\"cpu\")\n        return output\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def load_checkpoint(self, checkpoint_path, hdfs_path=None, del_local_after_load=True):\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.critic_module)\n        self.checkpoint_mananager.load_checkpoint(\n            local_path=checkpoint_path, hdfs_path=hdfs_path, del_local_after_load=del_local_after_load\n        )\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.critic_module)\n        if self._is_offload_optimizer:\n            offload_megatron_optimizer(self.critic_optimizer)\n\n    @register(dispatch_mode=Dispatch.ONE_TO_ALL)\n    def save_checkpoint(self, checkpoint_path, hdfs_path=None, global_steps=0, max_ckpt_to_keep=None):\n        if self._is_offload_param:\n            load_megatron_model_to_gpu(self.critic_module)\n        self.checkpoint_mananager.save_checkpoint(\n            local_path=checkpoint_path, hdfs_path=hdfs_path, global_step=global_steps, max_ckpt_to_keep=max_ckpt_to_keep\n        )\n        if self._is_offload_param:\n            offload_megatron_model_to_cpu(self.critic_module)\n"
  },
  {
    "path": "verl/workers/reward_manager/__init__.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .registry import get_reward_manager_cls, register  # noqa: I001\nfrom .batch import BatchRewardManager\nfrom .dapo import DAPORewardManager\nfrom .naive import NaiveRewardManager\nfrom .prime import PrimeRewardManager\n\n# Note(haibin.lin): no need to include all reward managers here in case of complicated dependencies\n__all__ = [\n    \"BatchRewardManager\",\n    \"DAPORewardManager\",\n    \"NaiveRewardManager\",\n    \"PrimeRewardManager\",\n    \"register\",\n    \"get_reward_manager_cls\",\n]\n\n# Import experimental reward managers to ensure they are registered\ntry:\n    from verl.experimental.reward_loop.reward_manager.limited import RateLimitedRewardManager  # noqa: F401\n\n    __all__.append(\"RateLimitedRewardManager\")\nexcept ImportError:\n    pass  # Optional dependency, may not be available\n"
  },
  {
    "path": "verl/workers/reward_manager/abstract.py",
    "content": "# Copyright 2023-2025 SGLang Team\n# Copyright Amazon.com, Inc. or its affiliates.\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom abc import ABC, abstractmethod\nfrom typing import Any, Callable\n\nimport torch\n\nfrom verl.protocol import DataProto\n\nRawRewardFn = Callable[..., Any]\n\n\nclass AbstractRewardManager(ABC):\n    @abstractmethod\n    def __init__(\n        self,\n        tokenizer: Any,\n        num_examine: int,\n        compute_score: RawRewardFn | None,\n        reward_fn_key: str = \"data_source\",\n        **kwargs: Any,\n    ):\n        pass\n\n    @abstractmethod\n    def __call__(\n        self,\n        data: DataProto,\n        return_dict: bool = False,\n    ) -> torch.Tensor | dict[str, Any]:\n        pass\n\n    def _extract_reward_from_rm_scores(\n        self, data: DataProto, return_dict: bool = False\n    ) -> torch.Tensor | dict[str, Any] | None:\n        \"\"\"\n        Extract reward from already-computed rm_scores if available.\n        This has been deprecated.\n\n        Args:\n            data: DataProto object containing the batch data\n            return_dict: Whether to return a dictionary with reward_tensor and reward_extra_info\n\n        Returns:\n            If rm_scores exists:\n                - If return_dict=True: dict with \"reward_tensor\" and \"reward_extra_info\"\n                - If return_dict=False: torch.Tensor of rm_scores\n            If rm_scores doesn't exist: None\n        \"\"\"\n        if \"rm_scores\" not in data.batch.keys():\n            return None\n\n        if return_dict:\n            reward_extra_keys = data.meta_info.get(\"reward_extra_keys\", [])\n            reward_extra_info = {key: data.non_tensor_batch[key] for key in reward_extra_keys}\n            return {\"reward_tensor\": data.batch[\"rm_scores\"], \"reward_extra_info\": reward_extra_info}\n        else:\n            return data.batch[\"rm_scores\"]\n"
  },
  {
    "path": "verl/workers/reward_manager/batch.py",
    "content": "# Copyright 2025 Individual Contributor: Mert Unsal\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections import defaultdict\nfrom typing import Any\n\nimport torch\n\nfrom verl import DataProto\nfrom verl.workers.reward_manager import register\nfrom verl.workers.reward_manager.abstract import AbstractRewardManager, RawRewardFn\n\n\n@register(\"batch\")\nclass BatchRewardManager(AbstractRewardManager):\n    \"\"\"\n    A batch reward manager that computes rewards for a batch of data.\n\n    Args:\n        tokenizer (Tokenizer): The tokenizer to use for decoding the responses.\n        num_examine (int): The number of responses to examine.\n        compute_score (callable): The function to compute the rewards.\n        reward_fn_key (str): The key to use for the reward function.\n        reward_kwargs (dict): The keyword arguments to pass to the reward function.\n    \"\"\"\n\n    def __init__(\n        self, tokenizer, num_examine, compute_score: RawRewardFn, reward_fn_key=\"data_source\", **reward_kwargs\n    ):\n        self.tokenizer = tokenizer\n        self.num_examine = num_examine\n        self.compute_score = compute_score\n        self.reward_fn_key = reward_fn_key\n        self.reward_kwargs = reward_kwargs\n\n    def verify(self, data):\n        prompt_ids = data.batch[\"prompts\"]\n        response_ids = data.batch[\"responses\"]\n        attention_mask = data.batch[\"attention_mask\"]\n\n        prompt_len = prompt_ids.shape[-1]\n        valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1)\n\n        responses_str = []\n        for i in range(len(data)):\n            valid_len = valid_response_lengths[i]\n            valid_response_ids = response_ids[i][:valid_len]\n            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n            responses_str.append(response_str)\n\n        ground_truths = [item.non_tensor_batch[\"reward_model\"].get(\"ground_truth\", None) for item in data]\n        data_sources = data.non_tensor_batch[self.reward_fn_key]\n        rollout_reward_scores = data.non_tensor_batch.get(\"reward_scores\", [{} for _ in range(len(data))])\n        extras = data.non_tensor_batch.get(\"extra_info\", [{} for _ in range(len(data))])\n\n        for i in range(len(data)):\n            extras[i][\"rollout_reward_scores\"] = rollout_reward_scores[i]\n\n        scores = self.compute_score(\n            data_sources=data_sources,\n            solution_strs=responses_str,\n            ground_truths=ground_truths,\n            extra_infos=extras,\n            **self.reward_kwargs,\n        )\n\n        return scores\n\n    def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]:\n        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn\n        reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict)\n        if reward_from_rm_scores is not None:\n            return reward_from_rm_scores\n\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n        reward_extra_info = defaultdict(list)\n        prompt_ids = data.batch[\"prompts\"]\n        prompt_len = prompt_ids.shape[-1]\n        attention_mask = data.batch[\"attention_mask\"]\n        valid_response_lengths = attention_mask[:, prompt_len:].sum(dim=-1)\n        data_sources = data.non_tensor_batch[self.reward_fn_key]\n\n        scores = self.verify(data)\n        rewards = []\n        already_printed: dict[str, Any] = {}\n\n        for i in range(len(data)):\n            length = valid_response_lengths[i].item()\n            score = scores[i]\n\n            if isinstance(score, dict):\n                reward = score[\"score\"]\n                for key, value in score.items():\n                    reward_extra_info[key].append(value)\n            else:\n                reward = score\n\n            rewards.append(reward)\n            reward_tensor[i, length - 1] = reward\n\n            data_source = data_sources[i]\n            if already_printed.get(data_source, 0) < self.num_examine:\n                response_str = self.tokenizer.decode(data.batch[\"responses\"][i][:length], skip_special_tokens=True)\n                prompt_str = self.tokenizer.decode(data.batch[\"prompts\"][i], skip_special_tokens=True)\n                ground_truth = data[i].non_tensor_batch[\"reward_model\"].get(\"ground_truth\", None)\n                print(\"[prompt]\", prompt_str)\n                print(\"[response]\", response_str)\n                print(\"[ground_truth]\", ground_truth)\n                print(\"[score]\", scores[i])\n                already_printed[data_source] = already_printed.get(data_source, 0) + 1\n\n        data.batch[\"acc\"] = torch.tensor(rewards, dtype=torch.float32, device=prompt_ids.device)\n\n        if return_dict:\n            return {\"reward_tensor\": reward_tensor, \"reward_extra_info\": reward_extra_info}\n        else:\n            return reward_tensor\n"
  },
  {
    "path": "verl/workers/reward_manager/dapo.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections import defaultdict\n\nimport torch\n\nfrom verl import DataProto\nfrom verl.utils.reward_score import default_compute_score\nfrom verl.workers.reward_manager import register\nfrom verl.workers.reward_manager.abstract import AbstractRewardManager\n\n\n@register(\"dapo\")\nclass DAPORewardManager(AbstractRewardManager):\n    \"\"\"The reward manager.\"\"\"\n\n    def __init__(\n        self,\n        tokenizer,\n        num_examine,\n        compute_score=None,\n        reward_fn_key=\"data_source\",\n        max_resp_len=None,\n        overlong_buffer_cfg=None,\n    ) -> None:\n        self.tokenizer = tokenizer\n        self.num_examine = num_examine  # the number of batches of decoded responses to print to the console\n        self.compute_score = compute_score or default_compute_score\n        self.reward_fn_key = reward_fn_key\n        self.overlong_buffer_cfg = overlong_buffer_cfg\n        self.max_resp_len = max_resp_len\n\n        if self.overlong_buffer_cfg is not None:\n            assert self.max_resp_len is not None, (\n                f\"max_resp_len must be provided if {overlong_buffer_cfg=}, but got None\"\n            )\n            assert self.max_resp_len >= self.overlong_buffer_cfg.len, (\n                \"max_resp_len must be larger than overlong_buffer.len\"\n            )\n            assert not self.overlong_buffer_cfg.enable or self.overlong_buffer_cfg.len > 0, (\n                \"overlong_buffer.len must be positive when overlong penalty is enabled,\"\n                f\"but got {self.overlong_buffer_cfg.len}.\"\n                \"To disable the overlong penalty, set overlong_buffer.enable = False\"\n            )\n\n    def __call__(self, data: DataProto, return_dict: bool = False):\n        \"\"\"We will expand this function gradually based on the available datasets\"\"\"\n\n        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn\n        reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict)\n        if reward_from_rm_scores is not None:\n            return reward_from_rm_scores\n\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n        reward_extra_info = defaultdict(list)\n\n        already_print_data_sources = {}\n\n        for i in range(len(data)):\n            data_item = data[i]  # DataProtoItem\n\n            prompt_ids = data_item.batch[\"prompts\"]\n\n            prompt_length = prompt_ids.shape[-1]\n\n            valid_prompt_length = data_item.batch[\"attention_mask\"][:prompt_length].sum()\n            valid_prompt_ids = prompt_ids[-valid_prompt_length:]\n\n            response_ids = data_item.batch[\"responses\"]\n            valid_response_length = data_item.batch[\"attention_mask\"][prompt_length:].sum()\n            valid_response_ids = response_ids[:valid_response_length]\n\n            # decode\n            prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True)\n            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n            eos_token = self.tokenizer.eos_token\n            if response_str.endswith(eos_token):\n                response_str = response_str[: -len(eos_token)]\n\n            ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n\n            data_source = data_item.non_tensor_batch[self.reward_fn_key]\n\n            extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n\n            rollout_reward_scores = data_item.non_tensor_batch.get(\"reward_scores\", {})\n\n            extra_info[\"rollout_reward_scores\"] = rollout_reward_scores\n\n            result = self.compute_score(\n                data_source=data_source,\n                solution_str=response_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n            )\n\n            score: float\n            if isinstance(result, dict):\n                score = result[\"score\"]\n                # Store the information including original reward\n                for key, value in result.items():\n                    reward_extra_info[key].append(value)\n            else:\n                score = result\n                reward_extra_info[\"acc\"].append(score)\n\n            reward = score\n\n            if self.overlong_buffer_cfg.enable:\n                overlong_buffer_len = self.overlong_buffer_cfg.len\n                expected_len = self.max_resp_len - overlong_buffer_len\n                exceed_len = valid_response_length - expected_len\n                overlong_penalty_factor = self.overlong_buffer_cfg.penalty_factor\n                overlong_reward = min(-exceed_len / overlong_buffer_len * overlong_penalty_factor, 0)\n                reward += overlong_reward\n                if self.overlong_buffer_cfg.log:\n                    reward_extra_info[\"overlong_reward\"].append(overlong_reward)\n                    reward_extra_info[\"overlong\"].append(overlong_reward < 0)\n\n            reward_tensor[i, valid_response_length - 1] = reward\n\n            if data_source not in already_print_data_sources:\n                already_print_data_sources[data_source] = 0\n\n            if already_print_data_sources[data_source] < self.num_examine:\n                already_print_data_sources[data_source] += 1\n                print(\"[prompt]\", prompt_str)\n                print(\"[response]\", response_str)\n                print(\"[ground_truth]\", ground_truth)\n                if isinstance(result, dict):\n                    for key, value in result.items():\n                        print(f\"[{key}]\", value)\n                else:\n                    print(\"[score]\", score)\n\n        if return_dict:\n            return {\n                \"reward_tensor\": reward_tensor,\n                \"reward_extra_info\": reward_extra_info,\n            }\n        else:\n            return reward_tensor\n"
  },
  {
    "path": "verl/workers/reward_manager/naive.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom collections import defaultdict\nfrom typing import Any\n\nimport torch\n\nfrom verl import DataProto\nfrom verl.utils.reward_score import default_compute_score\nfrom verl.workers.reward_manager import register\nfrom verl.workers.reward_manager.abstract import AbstractRewardManager\n\n\n@register(\"naive\")\nclass NaiveRewardManager(AbstractRewardManager):\n    \"\"\"The reward manager.\"\"\"\n\n    def __init__(self, tokenizer, num_examine, compute_score=None, reward_fn_key=\"data_source\") -> None:\n        \"\"\"\n        Initialize the NaiveRewardManager instance.\n\n        Args:\n            tokenizer: The tokenizer used to decode token IDs into text.\n            num_examine: The number of batches of decoded responses to print to the console for debugging purpose.\n            compute_score: A function to compute the reward score. If None, `default_compute_score` will be used.\n            reward_fn_key: The key used to access the data source in the non-tensor batch data. Defaults to\n                \"data_source\".\n        \"\"\"\n        self.tokenizer = tokenizer  # Store the tokenizer for decoding token IDs\n        self.num_examine = num_examine  # the number of batches of decoded responses to print to the console\n        self.compute_score = compute_score or default_compute_score\n        self.reward_fn_key = reward_fn_key  # Store the key for accessing the data source\n\n    def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]:\n        \"\"\"We will expand this function gradually based on the available datasets\"\"\"\n\n        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn\n        reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict)\n        if reward_from_rm_scores is not None:\n            return reward_from_rm_scores\n\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n        reward_extra_info = defaultdict(list)\n\n        already_print_data_sources = {}\n\n        for i in range(len(data)):\n            data_item = data[i]  # DataProtoItem\n\n            prompt_ids = data_item.batch[\"prompts\"]\n\n            prompt_length = prompt_ids.shape[-1]\n\n            valid_prompt_length = data_item.batch[\"attention_mask\"][:prompt_length].sum()\n            valid_prompt_ids = prompt_ids[-valid_prompt_length:]\n\n            response_ids = data_item.batch[\"responses\"]\n            valid_response_length = data_item.batch[\"attention_mask\"][prompt_length:].sum()\n            valid_response_ids = response_ids[:valid_response_length]\n\n            # decode\n            prompt_str = self.tokenizer.decode(valid_prompt_ids, skip_special_tokens=True)\n            response_str = self.tokenizer.decode(valid_response_ids, skip_special_tokens=True)\n\n            ground_truth = data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"]\n            data_source = data_item.non_tensor_batch[self.reward_fn_key]\n            extra_info = data_item.non_tensor_batch.get(\"extra_info\", {})\n            num_turns = data_item.non_tensor_batch.get(\"__num_turns__\", None)\n            rollout_reward_scores = data_item.non_tensor_batch.get(\"reward_scores\", {})\n            extra_info[\"num_turns\"] = num_turns\n            extra_info[\"rollout_reward_scores\"] = rollout_reward_scores\n\n            score = self.compute_score(\n                data_source=data_source,\n                solution_str=response_str,\n                ground_truth=ground_truth,\n                extra_info=extra_info,\n            )\n\n            if isinstance(score, dict):\n                reward = score[\"score\"]\n                # Store the information including original reward\n                for key, value in score.items():\n                    reward_extra_info[key].append(value)\n            else:\n                reward = score\n\n            reward_tensor[i, valid_response_length - 1] = reward\n\n            if data_source not in already_print_data_sources:\n                already_print_data_sources[data_source] = 0\n\n            if already_print_data_sources[data_source] < self.num_examine:\n                already_print_data_sources[data_source] += 1\n                print(\"[prompt]\", prompt_str)\n                print(\"[response]\", response_str)\n                print(\"[ground_truth]\", ground_truth)\n                if isinstance(score, dict):\n                    for key, value in score.items():\n                        print(f\"[{key}]\", value)\n                else:\n                    print(\"[score]\", score)\n\n        if return_dict:\n            return {\n                \"reward_tensor\": reward_tensor,\n                \"reward_extra_info\": reward_extra_info,\n            }\n        else:\n            return reward_tensor\n"
  },
  {
    "path": "verl/workers/reward_manager/prime.py",
    "content": "# Copyright 2024 PRIME team and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport asyncio\nfrom concurrent.futures import ProcessPoolExecutor\nfrom functools import partial\nfrom typing import Any, Callable, Optional\n\nimport psutil\nimport torch\nfrom transformers import PreTrainedTokenizer\n\nfrom verl import DataProto\nfrom verl.utils.ray_utils import get_event_loop\nfrom verl.utils.reward_score import default_compute_score\nfrom verl.workers.reward_manager import register\nfrom verl.workers.reward_manager.abstract import AbstractRewardManager\n\n\nasync def single_compute_score(evaluation_func, completion, reference, task, task_extra_info, executor, timeout=300.0):\n    loop = get_event_loop()\n    try:\n        # Ensure process_completion is called properly\n        future = loop.run_in_executor(executor, partial(evaluation_func, task, completion, reference, task_extra_info))\n        return await asyncio.wait_for(future, timeout=timeout)\n    except asyncio.TimeoutError:\n        print(f\"[Timeout] Task timeout: {completion}\")\n        return None  # Default value for timed-out rows\n    except Exception as e:\n        print(f\"[Error] Task failed: {e}, completion: {completion[:80]}\")\n        return None  # Default value for failed rows\n\n\nasync def parallel_compute_score_async(\n    evaluation_func, completions, references, tasks, extra_info=None, num_processes=64\n):\n    if extra_info is None:\n        extra_info = [None] * len(tasks)\n    scores = []\n    with ProcessPoolExecutor(max_workers=num_processes) as executor:\n        # to prevent very occasional starvation caused by some anomalous programs ( like infinite loop ), the\n        # exceptions in async programs will instantly halt the evaluation, and all summoned processes will be killed.\n        try:\n            # Create tasks for all rows\n            tasks_async = [\n                single_compute_score(evaluation_func, c, r, t, ei, executor, timeout=300.0)\n                for c, r, t, ei in zip(completions, references, tasks, extra_info, strict=True)\n            ]\n            results = await asyncio.gather(*tasks_async, return_exceptions=False)\n        except Exception as e:\n            print(f\"[Exception] async gather failed: {e}\")\n            raise\n        finally:\n            terminated_count = 0\n            for pid, proc in executor._processes.items():\n                try:\n                    p = psutil.Process(pid)\n                    p.terminate()\n                    try:\n                        p.wait(timeout=5)\n                    except psutil.TimeoutExpired:\n                        p.kill()\n                    terminated_count += 1\n                except Exception:\n                    pass\n            print(f\"[Shutdown] {terminated_count} subprocess(es) terminated.\")\n\n    # Process results\n    for result, completion, reference, task in zip(results, completions, references, tasks, strict=True):\n        if isinstance(result, Exception) or result is None:\n            # Handle failed or timed-out tasks\n            scores.append(0.0)\n        elif isinstance(result, int | float | bool):\n            scores.append(float(result))\n        else:\n            scores.append(float(result[0]))\n    return scores\n\n\ndef run_reward_scoring(evaluation_func, completions, references, tasks, extra_info=None, num_processes=64):\n    loop = asyncio.new_event_loop()\n    asyncio.set_event_loop(loop)\n    try:\n        return loop.run_until_complete(\n            parallel_compute_score_async(evaluation_func, completions, references, tasks, extra_info, num_processes)\n        )\n    finally:\n        loop.close()\n\n\n@register(\"prime\")\nclass PrimeRewardManager(AbstractRewardManager):\n    \"\"\"\n    The Reward Manager used in https://github.com/PRIME-RL/PRIME\n    \"\"\"\n\n    def __init__(\n        self,\n        tokenizer: PreTrainedTokenizer,\n        num_examine: int,\n        compute_score: Optional[Callable] = None,\n        reward_fn_key: str = \"data_source\",\n    ) -> None:\n        self.tokenizer = tokenizer\n        self.num_examine = num_examine  # the number of batches of decoded responses to print to the console\n        self.compute_score = compute_score or default_compute_score\n        self.reward_fn_key = reward_fn_key\n\n    def verify(self, data):\n        \"\"\"\n        verify the batch and save as ``acc`` tensor\n        \"\"\"\n        # batched scoring\n        prompt_ids = data.batch[\"prompts\"]\n\n        response_ids = data.batch[\"responses\"]\n        sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True)\n        ground_truth = [data_item.non_tensor_batch[\"reward_model\"][\"ground_truth\"] for data_item in data]\n        data_sources = data.non_tensor_batch[self.reward_fn_key]\n        extra_info = data.non_tensor_batch.get(\"extra_info\", None)\n\n        assert len(sequences_str) == len(ground_truth) == len(data_sources)\n        try:\n            scores = run_reward_scoring(\n                self.compute_score,\n                completions=sequences_str,\n                references=ground_truth,\n                tasks=data_sources,\n                extra_info=extra_info,\n                num_processes=64,\n            )\n        except asyncio.TimeoutError:\n            print(\"[Timeout] Global reward scoring timed out. Setting all as 0.\")\n            scores = [0.0 for _ in range(len(sequences_str))]\n        except Exception as e:\n            print(f\"[Error] Unexpected error during scoring. Setting all as 0. {e}\")\n            scores = [0.0 for _ in range(len(sequences_str))]\n        data.batch[\"acc\"] = torch.tensor(scores, dtype=torch.float32, device=prompt_ids.device)\n        return scores\n\n    def __call__(self, data: DataProto, return_dict: bool = False) -> torch.Tensor | dict[str, Any]:\n        \"\"\"We will expand this function gradually based on the available datasets\"\"\"\n\n        # If there is rm score, we directly return rm score. Otherwise, we compute via rm_score_fn\n        reward_from_rm_scores = self._extract_reward_from_rm_scores(data, return_dict)\n        if reward_from_rm_scores is not None:\n            return reward_from_rm_scores\n\n        reward_tensor = torch.zeros_like(data.batch[\"responses\"], dtype=torch.float32)\n\n        already_print_data_sources = {}\n\n        # batched scoring\n        prompt_ids = data.batch[\"prompts\"]\n        prompt_length = prompt_ids.shape[-1]\n\n        response_ids = data.batch[\"responses\"]\n        valid_response_length = data.batch[\"attention_mask\"][:, prompt_length:].sum(dim=-1)\n        sequences_str = self.tokenizer.batch_decode(response_ids, skip_special_tokens=True)\n        data_sources = data.non_tensor_batch[\"data_source\"]\n\n        scores = self.verify(data)\n\n        for i in range(len(data)):\n            data_source = data_sources[i]\n            reward_tensor[i, valid_response_length[i].item() - 1] = scores[i]\n\n            if data_source not in already_print_data_sources:\n                already_print_data_sources[data_source] = 0\n\n            if already_print_data_sources[data_source] < self.num_examine:\n                already_print_data_sources[data_source] += 1\n                print(sequences_str)\n\n        if return_dict:\n            return {\"reward_tensor\": reward_tensor}\n        else:\n            return reward_tensor\n"
  },
  {
    "path": "verl/workers/reward_manager/registry.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom typing import Callable\n\nfrom verl.workers.reward_manager.abstract import AbstractRewardManager\n\n__all__ = [\"register\", \"get_reward_manager_cls\"]\n\nREWARD_MANAGER_REGISTRY: dict[str, type[AbstractRewardManager]] = {}\n\n\ndef register(name: str) -> Callable[[type[AbstractRewardManager]], type[AbstractRewardManager]]:\n    \"\"\"Decorator to register a reward manager class with a given name.\n\n    Args:\n        name: `(str)`\n            The name of the reward manager.\n    \"\"\"\n\n    def decorator(cls: type[AbstractRewardManager]) -> type[AbstractRewardManager]:\n        if name in REWARD_MANAGER_REGISTRY and REWARD_MANAGER_REGISTRY[name] != cls:\n            raise ValueError(\n                f\"Reward manager {name} has already been registered: {REWARD_MANAGER_REGISTRY[name]} vs {cls}\"\n            )\n        REWARD_MANAGER_REGISTRY[name] = cls\n        return cls\n\n    return decorator\n\n\ndef get_reward_manager_cls(name: str) -> type[AbstractRewardManager]:\n    \"\"\"Get the reward manager class with a given name.\n\n    Args:\n        name: `(str)`\n            The name of the reward manager.\n\n    Returns:\n        `(type)`: The reward manager class.\n    \"\"\"\n    if name not in REWARD_MANAGER_REGISTRY:\n        raise ValueError(f\"Unknown reward manager: {name}\")\n    return REWARD_MANAGER_REGISTRY[name]\n"
  },
  {
    "path": "verl/workers/rollout/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .base import BaseRollout, get_rollout_class\nfrom .hf_rollout import HFRollout\nfrom .naive import NaiveRollout\nfrom .replica import RolloutReplica\n\n__all__ = [\"BaseRollout\", \"NaiveRollout\", \"HFRollout\", \"get_rollout_class\", \"RolloutReplica\"]\n"
  },
  {
    "path": "verl/workers/rollout/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport importlib\nfrom abc import ABC, abstractmethod\nfrom typing import Generator\n\nimport torch\nfrom torch.distributed.device_mesh import DeviceMesh\n\nfrom verl import DataProto\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.workers.config import HFModelConfig, RolloutConfig\n\n__all__ = [\"BaseRollout\"]\n\n\nclass BaseRollout(ABC):\n    \"\"\"Base class for rollout.\"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        device_mesh: DeviceMesh,\n        *args,\n        **kwargs,\n    ):\n        self.config = omega_conf_to_dataclass(config)\n        self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)\n        self.device_mesh = device_mesh\n\n    @abstractmethod\n    async def resume(self, tags: list[str]):\n        \"\"\"Resume rollout weights or kv cache in GPU memory.\n\n        Args:\n            tags: weights or kv_cache.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    async def update_weights(\n        self,\n        weights: Generator[tuple[str, torch.Tensor], None, None],\n        **kwargs,\n    ):\n        \"\"\"Update the weights of the rollout model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    async def release(self):\n        \"\"\"Release weights and kv cache in GPU memory.\"\"\"\n        pass\n\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Batch generate sequences in sync mode.\n\n        Args:\n            prompts: The input prompts.\n\n        Returns:\n            The output sequences.\n        \"\"\"\n        raise NotImplementedError\n\n\n_ROLLOUT_REGISTRY = {\n    (\"vllm\", \"async\"): \"verl.workers.rollout.vllm_rollout.ServerAdapter\",\n    (\"sglang\", \"async\"): \"verl.workers.rollout.sglang_rollout.sglang_rollout.ServerAdapter\",\n    (\"trtllm\", \"async\"): \"verl.workers.rollout.trtllm_rollout.trtllm_rollout.ServerAdapter\",\n}\n\n\ndef get_rollout_class(rollout_name: str, mode: str = \"async\") -> type[BaseRollout]:\n    \"\"\"Get the rollout class by name.\n\n    Args:\n        rollout_name: The name of the rollout.\n        mode: The mode of the rollout, async: server mode.\n\n    Returns:\n        The rollout class.\n    \"\"\"\n    assert (rollout_name, mode) in _ROLLOUT_REGISTRY, f\"Rollout {rollout_name} with mode {mode} not found\"\n    fqdn = _ROLLOUT_REGISTRY[(rollout_name, mode)]\n    module_name, class_name = fqdn.rsplit(\".\", 1)\n    rollout_module = importlib.import_module(module_name)\n    return getattr(rollout_module, class_name)\n"
  },
  {
    "path": "verl/workers/rollout/hf_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nRollout with huggingface models.\nTODO: refactor this class. Currently, it will hang when using FSDP HybridShard. We should actually create a single\nGPU model. Then, get full state_dict and bind the state_dict to the single GPU model. Then, use the single GPU model\nto perform generation.\n\"\"\"\n\nimport contextlib\n\nimport torch\nimport torch.distributed\nfrom tensordict import TensorDict\nfrom torch import nn\nfrom torch.distributed.fsdp import FullyShardedDataParallel as FSDP\nfrom transformers import GenerationConfig\n\nfrom verl import DataProto\nfrom verl.utils.device import get_device_name, get_torch_device\nfrom verl.utils.torch_functional import get_response_mask\n\nfrom .base import BaseRollout\n\n__all__ = [\"HFRollout\"]\n\n\nclass HFRollout(BaseRollout):\n    def __init__(self, module: nn.Module, config):\n        super().__init__()\n        self.config = config\n        self.module = module\n\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        batch_size = prompts.batch.batch_size[0]\n        num_chunks = max(batch_size // self.config.get(\"micro_batch_size\", batch_size), 1)\n        batch_prompts = prompts.chunk(chunks=num_chunks)\n        output = [self._generate_minibatch(p) for p in batch_prompts]\n        output = DataProto.concat(output)\n        return output\n\n    @torch.no_grad()\n    def _generate_minibatch(self, prompts: DataProto) -> DataProto:\n        # make sampling args can be overridden by inputs\n        do_sample = prompts.meta_info.get(\"do_sample\", self.config.do_sample)\n        is_validate = prompts.meta_info.get(\"validate\", False)\n\n        temperature = prompts.meta_info.get(\"temperature\", self.config.temperature)\n        response_length = prompts.meta_info.get(\"response_length\", self.config.response_length)\n        top_p = prompts.meta_info.get(\"top_p\", self.config.get(\"top_p\", 1.0))\n        top_k = max(0, prompts.meta_info.get(\"top_k\", self.config.get(\"top_k\", 0)))  # to be compatible with vllm\n\n        if not do_sample:\n            # do_sample==False -> greedy decoding\n            kwargs = {\n                \"do_sample\": False,\n                \"num_beams\": 1,\n            }\n        elif is_validate:\n            # do validate and do sample -> use val_kwargs\n            kwargs = {\n                \"do_sample\": True,\n                \"num_beams\": 1,\n                \"top_k\": max(0, self.config.val_kwargs.top_k),  # to be compatible with vllm\n                \"top_p\": self.config.val_kwargs.top_p,\n                \"temperature\": self.config.val_kwargs.temperature,\n                \"num_return_sequences\": 1,  # if validate, already repeat in ray_trainer\n            }\n        else:\n            # do_sample -> use rollout config\n            kwargs = {\n                \"do_sample\": True,\n                \"num_beams\": 1,\n                \"top_p\": top_p,\n                \"top_k\": top_k,\n                \"temperature\": temperature,\n                # already repeat in ray_trainer\n                # https://github.com/volcengine/verl/blob/2fdfbdcba6f2e076f64bc47922d8fe6cf7dc7da5/verl/trainer/ppo/ray_trainer.py#L1117\n                \"num_return_sequences\": 1,\n            }\n\n        # make config according to generate mode\n        generation_config = GenerationConfig(**kwargs)\n\n        idx = prompts.batch[\"input_ids\"]  # (bs, prompt_length)\n        prompt_length = idx.size(1)\n        attention_mask = prompts.batch[\"attention_mask\"]  # left-padded attention_mask\n        position_ids = prompts.batch[\"position_ids\"]\n\n        # used to construct attention_mask\n        eos_token_id = prompts.meta_info[\"eos_token_id\"]\n        pad_token_id = prompts.meta_info[\"pad_token_id\"]\n\n        self.module.eval()\n        param_ctx = contextlib.nullcontext()\n\n        if isinstance(self.module, FSDP):\n            # recurse need to set to False according to https://github.com/pytorch/pytorch/issues/100069\n            param_ctx = FSDP.summon_full_params(self.module, writeback=False, recurse=False)\n        with param_ctx, torch.autocast(device_type=get_device_name(), dtype=torch.bfloat16):\n            output = self.module.generate(\n                input_ids=idx,\n                attention_mask=attention_mask,\n                position_ids=position_ids,\n                do_sample=do_sample,\n                max_new_tokens=response_length,\n                eos_token_id=eos_token_id,\n                pad_token_id=pad_token_id,\n                generation_config=generation_config,\n                output_scores=False,  # this is potentially very large\n                return_dict_in_generate=True,\n                use_cache=True,\n            )\n\n        # TODO: filter out the seq with no answers like ds-chat\n        seq = output.sequences\n        generated_batch_size = seq.size(0)  # bs * num_return_sequences\n\n        # huggingface generate will stop generating when all the batch reaches [EOS].\n        # We have to pad to response_length\n        sequence_length = prompt_length + self.config.response_length\n        delta_length = sequence_length - seq.shape[1]\n\n        if delta_length > 0:\n            delta_tokens = torch.ones(size=(generated_batch_size, delta_length), device=seq.device, dtype=seq.dtype)\n            delta_tokens = pad_token_id * delta_tokens\n            seq = torch.cat((seq, delta_tokens), dim=1)\n        assert seq.shape[1] == sequence_length\n\n        # make necessary reputations if num_return_sequences > 1\n        num_return_sequences = kwargs.get(\"num_return_sequences\", 1)\n        if num_return_sequences > 1:\n            position_ids = position_ids.repeat_interleave(num_return_sequences, dim=0)\n            attention_mask = attention_mask.repeat_interleave(num_return_sequences, dim=0)\n\n        prompt = seq[:, :prompt_length]  # (generated_batch_size, prompt_length)\n        response = seq[:, prompt_length:]  # (generated_batch_size, response_length)\n\n        response_length = response.size(1)\n        delta_position_id = torch.arange(1, response_length + 1, device=position_ids.device)\n        delta_position_id = delta_position_id.unsqueeze(0).repeat(generated_batch_size, 1)\n\n        response_position_ids = position_ids[:, -1:] + delta_position_id\n        position_ids = torch.cat([position_ids, response_position_ids], dim=-1)\n\n        response_attention_mask = get_response_mask(\n            response_id=response, eos_token=eos_token_id, dtype=attention_mask.dtype\n        )\n        attention_mask = torch.cat((attention_mask, response_attention_mask), dim=-1)\n\n        batch = TensorDict(\n            {\n                \"prompts\": prompt,\n                \"responses\": response,\n                \"input_ids\": seq,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n            },\n            batch_size=generated_batch_size,\n        )\n\n        # empty cache before compute old_log_prob\n        get_torch_device().empty_cache()\n\n        self.module.train()\n        return DataProto(batch=batch)\n"
  },
  {
    "path": "verl/workers/rollout/naive/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nfrom .naive_rollout import NaiveRollout\n\n__all__ = [\"NaiveRollout\"]\n"
  },
  {
    "path": "verl/workers/rollout/naive/naive_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nIn single GPU rollout, the sequences are generated directly by sampling from the model.\nThe output will contain\n1. output_ids\n2. attention_masks (left padding)\n3. eos_masks\n4. log_probs\n\"\"\"\n\nimport torch\nimport torch.nn.functional as F\nfrom tensordict import TensorDict\nfrom torch import nn\n\nfrom verl import DataProto\nfrom verl.utils.torch_functional import logprobs_from_logits\n\nfrom ..base import BaseRollout\n\n__all__ = [\"NaiveRollout\"]\n\n\nclass NaiveRollout(BaseRollout):\n    def __init__(self, module: nn.Module, config):\n        \"\"\"A naive rollout. It requires the module to be compatible with huggingface APIs. That is:\n        The module should define __call__ to receive input_ids, attention_mask and position_ids.\n        It outputs a structure that contains logits field.\n\n        Args:\n            module: module here follows huggingface APIs\n            config: DictConfig\n        \"\"\"\n        super().__init__()\n        self.config = config\n        self.module = module\n\n    @torch.no_grad()\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Generate sequences\"\"\"\n        idx = prompts.batch[\"input_ids\"]  # (bs, prompt_length)\n        attention_mask = prompts.batch[\"attention_mask\"]  # left-padded attention_mask\n        position_ids = prompts.batch[\"position_ids\"]\n\n        # used to construct attention_mask\n        eos_token_id = prompts.meta_info[\"eos_token_id\"]\n\n        batch_size = idx.size(0)\n        prompt_length = idx.size(1)\n\n        self.module.eval()\n\n        prev_attention_mask = torch.ones(size=(batch_size, 1), dtype=attention_mask.dtype, device=attention_mask.device)\n\n        logits_lst = []\n        for _ in range(self.config.response_length):\n            # if the sequence context is growing too long we must crop it at block_size\n            # idx_cond = idx if idx.size(1) <= self.config.block_size else idx[:, -self.config.block_size:]\n            idx_cond = idx\n            # forward the model to get the logits for the index in the sequence\n            # we use huggingface APIs here\n            output = self.module(input_ids=idx_cond, attention_mask=attention_mask, position_ids=position_ids)\n            logits = output.logits\n            # pluck the logits at the final step and scale by desired temperature\n            logits = logits[:, -1, :] / self.config.temperature  # (bs, vocab_size)\n            # optionally crop the logits to only the top k options\n            if self.config.top_k is not None:\n                v, _ = torch.topk(logits, min(self.config.top_k, logits.size(-1)))\n                logits[logits < v[:, [-1]]] = -float(\"Inf\")\n            # apply softmax to convert logits to (normalized) probabilities\n            probs = F.softmax(logits, dim=-1)\n            # sample from the distribution\n            if self.config.do_sample:\n                idx_next = torch.multinomial(probs, num_samples=1)\n            else:\n                idx_next = torch.argmax(probs, dim=-1, keepdim=True)\n\n            attention_mask = torch.cat((attention_mask, prev_attention_mask), dim=-1)\n\n            for token_id in eos_token_id:\n                prev_attention_mask = torch.logical_and(idx_next != token_id, prev_attention_mask.bool())\n            prev_attention_mask.to(attention_mask.dtype)\n\n            position_ids = torch.cat((position_ids, position_ids[:, -1:] + 1), dim=-1)\n\n            # append sampled index to the running sequence and continue\n            idx = torch.cat((idx, idx_next), dim=1)\n            logits_lst.append(logits)\n\n        logits = torch.stack(logits_lst, dim=1)  # (bs, response_length, vocab_size)\n        prompts = idx[:, :prompt_length]  # (bs, prompt_length)\n        response = idx[:, prompt_length:]  # (bs, response_length)\n        log_probs = logprobs_from_logits(logits=logits, labels=response)\n        batch = TensorDict(\n            {\n                \"input_ids\": prompts,\n                \"responses\": response,\n                \"sequences\": idx,\n                \"old_log_probs\": log_probs,\n                \"attention_mask\": attention_mask,\n                \"position_ids\": position_ids,\n            },\n            batch_size=batch_size,\n        )\n\n        self.module.train()\n\n        return DataProto(batch=batch)\n"
  },
  {
    "path": "verl/workers/rollout/replica.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nfrom abc import ABC, abstractmethod\nfrom enum import Enum\nfrom typing import Any, Callable, Optional\n\nimport ray\nfrom omegaconf import DictConfig\nfrom pydantic import BaseModel\nfrom ray.actor import ActorHandle\n\nfrom verl.single_controller.ray import RayClassWithInitArgs, RayResourcePool, RayWorkerGroup, ResourcePoolManager\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import is_torch_npu_available\nfrom verl.workers.config import HFModelConfig, RolloutConfig\n\nlogger = logging.getLogger(__file__)\n\n\n# Max number of concurrent calls to the methods of Rollout,\n# excluding calls to generate method.\nCONTROL_METHOD_CONCURRENCY = 16\n\n\nclass TokenOutput(BaseModel):\n    token_ids: list[int]\n    \"\"\"response token ids\"\"\"\n    log_probs: Optional[list[float]] = None\n    \"\"\"logprobs of response token ids\"\"\"\n    routed_experts: Optional[Any] = None\n    \"\"\"routed experts of response token ids\"\"\"\n    stop_reason: Optional[str] = None\n    \"\"\"stop reason: 'completed', 'aborted', or None for unknown\"\"\"\n    num_preempted: Optional[int] = None\n    \"\"\"number of preempted times for metric calculation\"\"\"\n    extra_fields: dict[str, Any] = {}\n    \"\"\"Extra fields for dynamic addition.\"\"\"\n\n\nclass RolloutMode(Enum):\n    # Rollout engine and training engine(fsdp/megatron) fused in same process\n    # Rollout and trainer share GPUs, switch context with weight synchronization.\n    # Usage scenarios: on-policy training.\n    HYBRID = \"hybrid\"\n\n    # Rollout engine colocated with hybrid engine in same ray placement group but in separate process.\n    # Rollout and hybrid processes share GPUs, switch context without weight synchronization.\n    # Usage scenarios: GRM (LLM as a judge).\n    COLOCATED = \"colocated\"\n\n    # Standalone rollout server with separate GPU resource, disaggregated architecture.\n    # Usage scenarios: off-policy training.\n    STANDALONE = \"standalone\"\n\n\nclass RolloutReplica(ABC):\n    \"\"\"Rollout replica is an individual server instance, which may be deployed on single or multiple nodes.\n    It is equivalent to launch server in each node with command line:\n\n    SGLang:\n    ```\n    python -m sglang.launch_server --node-rank 0 --nnode 2 ...\n    python -m sglang.launch_server --node-rank 1 --nnode 2 ...\n    ```\n\n    vLLM:\n    ```\n    vllm serve --data-parallel-size 16 --data-parallel-size-local 8 --data-parallel-start-rank 0 ...\n    vllm serve --data-parallel-size 16 --data-parallel-size-local 8 --data-parallel-start-rank 8 ...\n    ```\n\n    Args:\n        replica_rank: int, rank of this rollout replica.\n        config: RolloutConfig, full config.\n        model_config: DictConfig, model config.\n        gpus_per_node: int, number of gpus per node.\n    \"\"\"\n\n    def __init__(\n        self,\n        replica_rank: int,\n        config: RolloutConfig,\n        model_config: DictConfig,\n        gpus_per_node: int = 8,\n        is_reward_model: bool = False,\n    ) -> None:\n        self.replica_rank = replica_rank\n        self.config: RolloutConfig = omega_conf_to_dataclass(config)\n        self.model_config: HFModelConfig = model_config\n\n        self.world_size = (\n            self.config.tensor_model_parallel_size\n            * self.config.data_parallel_size\n            * self.config.pipeline_model_parallel_size\n        )\n        self.gpus_per_node = gpus_per_node\n        self.gpus_per_replica_node = min(gpus_per_node, self.world_size)\n        assert self.world_size % self.gpus_per_replica_node == 0, (\n            f\"world_size {self.world_size} must be divisible by gpus_per_node {self.gpus_per_replica_node}\"\n        )\n        self.nnodes = self.world_size // self.gpus_per_replica_node\n        self.is_reward_model = is_reward_model\n\n        self.rollout_mode: RolloutMode = None\n        self.workers: list[ActorHandle] = []\n        self.resource_pool: RayResourcePool = None\n        self.bundle_indices: list[int] = []\n\n        self.servers: list[ActorHandle] = []\n        self._server_address: str = None\n        self._server_handle: ActorHandle = None\n\n    async def init_hybrid(self, worker_group: RayWorkerGroup):\n        \"\"\"Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process.\n\n        Args:\n            worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized.\n        \"\"\"\n        self.rollout_mode = RolloutMode.HYBRID\n        self.workers = worker_group.workers[\n            self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1)\n        ]\n        await self.launch_servers()\n\n    async def init_hybrid_colocated(self, worker_group: RayWorkerGroup, resource_pool: RayResourcePool):\n        \"\"\"Init hybrid rollout server, rollout engine and training engine(fsdp/megatron) fused in same process.\n\n        Args:\n            worker_group: RayWorkerGroup, fused workers where training engine(fsdp/megatron) have been initialized.\n            resource_pool: RayResourcePool, ray placement group where hybrid engine processes have been launched.\n            bundle_indices: list[int], bundle indices for this rollout replica.\n        \"\"\"\n        self.rollout_mode = RolloutMode.HYBRID\n        self.workers = worker_group.workers[\n            self.world_size * self.replica_rank : self.world_size * (self.replica_rank + 1)\n        ]\n        self.resource_pool = resource_pool\n        self.bundle_indices = [self.replica_rank * self.world_size + idx for idx in range(self.world_size)]\n        await self.launch_servers()\n\n    # TODO(sgm): this should be the default solution, but need to make the RolloutMode more clear.\n    async def init_colocated(self, resource_pool: RayResourcePool):\n        \"\"\"Init colocated rollout server, rollout engine and hybrid engine colocated in same ray placement group\n        but in separate processes.\n\n        Args:\n            resource_pool: RayResourcePool, ray placement group where hybrid engine processes have been launched.\n        \"\"\"\n        self.rollout_mode = RolloutMode.COLOCATED\n        self.resource_pool = resource_pool\n        use_gpu = self.rollout_worker_use_gpu()\n\n        worker_group = RayWorkerGroup(\n            resource_pool=self.resource_pool,\n            ray_cls_with_init=self.get_ray_class_with_init_args(),\n            bin_pack=False,\n            name_prefix=f\"rollout_colocate_{self.replica_rank}\"\n            if not self.is_reward_model\n            else f\"rollout_reward_colocate_{self.replica_rank}\",\n            use_gpu=use_gpu,\n            device_name=\"cuda\" if not is_torch_npu_available(check_device=False) else \"npu\",\n        )\n        self.workers = worker_group.workers\n        await self.launch_servers()\n\n    async def init_standalone(self):\n        \"\"\"Init standalone rollout server, create new resource pool for this rollout.\"\"\"\n        # create resource pool for this rollout\n        self.rollout_mode = RolloutMode.STANDALONE\n        resource_pool_name = (\n            f\"rollout_pool_{self.replica_rank}\"\n            if not self.is_reward_model\n            else f\"rollout_pool_reward_{self.replica_rank}\"\n        )\n        resource_pool_spec = {\n            resource_pool_name: [self.gpus_per_replica_node] * self.nnodes,\n        }\n        resource_pool_manager = ResourcePoolManager(resource_pool_spec=resource_pool_spec, mapping=None)\n        resource_pool_manager.create_resource_pool()\n        self.resource_pool = resource_pool_manager.resource_pool_dict[resource_pool_name]\n\n        # create worker group for this rollout\n        use_gpu = self.rollout_worker_use_gpu()\n        worker_group = RayWorkerGroup(\n            resource_pool=self.resource_pool,\n            ray_cls_with_init=self.get_ray_class_with_init_args(),\n            bin_pack=False,\n            name_prefix=f\"rollout_standalone_{self.replica_rank}\"\n            if not self.is_reward_model\n            else f\"rollout_reward_standalone_{self.replica_rank}\",\n            use_gpu=use_gpu,\n            device_name=\"cuda\" if not is_torch_npu_available(check_device=False) else \"npu\",\n        )\n        self.workers = worker_group.workers\n        await self.launch_servers()\n\n    def get_ray_class_with_init_args(self) -> RayClassWithInitArgs:\n        \"\"\"Get rollout worker actor class for colocated and standalone mode.\"\"\"\n        from verl.checkpoint_engine.base import CheckpointEngineWorker\n\n        rollout_worker_actor_cls = ray.remote(CheckpointEngineWorker)\n\n        return RayClassWithInitArgs(\n            cls=rollout_worker_actor_cls,\n            rollout_config=self.config,\n            model_config=self.model_config,\n            replica_rank=self.replica_rank,\n        )\n\n    @abstractmethod\n    async def launch_servers(self):\n        \"\"\"Launch http server in each node.\"\"\"\n        raise NotImplementedError\n\n    @property\n    def server_address(self) -> str:\n        \"\"\"Get rollout server address for OpenAI chat completion.\"\"\"\n        return self._server_address\n\n    @property\n    def server_handle(self) -> ActorHandle:\n        \"\"\"Get rollout server handle for Token-in-token-out generation.\"\"\"\n        return self._server_handle\n\n    @property\n    def max_concurrency(self) -> int:\n        # 1000 is Ray's default max_concurrency for async execution.\n        # Add some margin to account for control method call.\n        return max(1000, self.config.max_num_seqs + CONTROL_METHOD_CONCURRENCY)\n\n    def rollout_worker_use_gpu(self) -> bool:\n        return True\n\n    async def wake_up(self):\n        \"\"\"Wake up each rollout server.\"\"\"\n        await asyncio.gather(*[server.wake_up.remote() for server in self.servers])\n\n    async def sleep(self):\n        \"\"\"Sleep each rollout server.\"\"\"\n        await asyncio.gather(*[server.sleep.remote() for server in self.servers])\n\n    async def abort_all_requests(self):\n        \"\"\"Partial rollout: abort and save all unfinished requests in each rollout server.\"\"\"\n        await asyncio.gather(*[server.abort_all_requests.remote() for server in self.servers])\n\n    async def resume_generation(self):\n        \"\"\"Resume generation on all servers after abort_all_requests.\"\"\"\n        await asyncio.gather(*[server.resume_generation.remote() for server in self.servers])\n\n    async def clear_kv_cache(self):\n        \"\"\"reset kv cache in each rollout server.\"\"\"\n        await asyncio.gather(*[server.clear_kv_cache.remote() for server in self.servers])\n\n    async def start_profile(self, **kwargs):\n        \"\"\"Start profiling on the replica.\"\"\"\n        await asyncio.gather(*[server.start_profile.remote(**kwargs) for server in self.servers])\n\n    async def stop_profile(self):\n        \"\"\"Stop profiling on the replica.\"\"\"\n        await asyncio.gather(*[server.stop_profile.remote() for server in self.servers])\n\n\nclass RolloutReplicaRegistry:\n    \"\"\"Factory for managing rollout replica implementations.\"\"\"\n\n    _registry: dict[str, Callable[[], type[RolloutReplica]]] = {}\n\n    @classmethod\n    def register(cls, name: str, loader: Callable[[], type[RolloutReplica]]) -> None:\n        \"\"\"Register a new rollout replica type.\"\"\"\n        cls._registry[name] = loader\n\n    @classmethod\n    def get(cls, name: str) -> type[RolloutReplica]:\n        \"\"\"Get a rollout replica class by name.\"\"\"\n        if name not in cls._registry:\n            raise ValueError(f\"Unknown rollout mode: {name}. Available: {list(cls._registry.keys())}\")\n        return cls._registry[name]()\n\n\n# Loader functions for built-in types\ndef _load_vllm():\n    from verl.workers.rollout.vllm_rollout.vllm_async_server import vLLMReplica\n\n    return vLLMReplica\n\n\ndef _load_sglang():\n    os.environ[\"SGLANG_USE_CPU_ENGINE\"] = \"1\"\n\n    try:\n        import vllm  # noqa: F401\n    except ImportError:\n        import sys\n        import types\n        from unittest.mock import Mock\n\n        mock_vllm = types.ModuleType(\"vllm\")\n\n        mock_custom_ops = types.ModuleType(\"vllm._custom_ops\")\n        mock_custom_ops.scaled_fp8_quant = Mock()\n        mock_vllm._custom_ops = mock_custom_ops\n\n        mock_model_executor = types.ModuleType(\"vllm.model_executor\")\n        mock_layers = types.ModuleType(\"vllm.model_executor.layers\")\n        mock_activation = types.ModuleType(\"vllm.model_executor.layers.activation\")\n\n        class GeluAndMul:  # noqa: N801\n            pass\n\n        class SiluAndMul:  # noqa: N801\n            pass\n\n        mock_activation.GeluAndMul = GeluAndMul\n        mock_activation.SiluAndMul = SiluAndMul\n        mock_layers.activation = mock_activation\n        mock_model_executor.layers = mock_layers\n        mock_vllm.model_executor = mock_model_executor\n\n        sys.modules[\"vllm\"] = mock_vllm\n        sys.modules[\"vllm._custom_ops\"] = mock_custom_ops\n        sys.modules[\"vllm.model_executor\"] = mock_model_executor\n        sys.modules[\"vllm.model_executor.layers\"] = mock_layers\n        sys.modules[\"vllm.model_executor.layers.activation\"] = mock_activation\n\n    from verl.workers.rollout.sglang_rollout.async_sglang_server import SGLangReplica\n\n    del os.environ[\"SGLANG_USE_CPU_ENGINE\"]\n    return SGLangReplica\n\n\ndef _load_trtllm():\n    from verl.workers.rollout.trtllm_rollout.trtllm_async_server import TRTLLMReplica\n\n    return TRTLLMReplica\n\n\n# Register built-in types\nRolloutReplicaRegistry.register(\"vllm\", _load_vllm)\nRolloutReplicaRegistry.register(\"sglang\", _load_sglang)\nRolloutReplicaRegistry.register(\"trtllm\", _load_trtllm)\n\n\n# Original function for backward compatibility\ndef get_rollout_replica_class(rollout: str) -> type[RolloutReplica]:\n    return RolloutReplicaRegistry.get(rollout)\n"
  },
  {
    "path": "verl/workers/rollout/schemas.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport difflib\nimport logging\nimport os\nfrom enum import Enum\nfrom typing import Any, Optional\n\nimport torch\nfrom pydantic import BaseModel, ConfigDict, model_validator\nfrom transformers import PreTrainedTokenizer, PreTrainedTokenizerFast, ProcessorMixin\n\nfrom verl.tools.schemas import OpenAIFunctionToolCall, OpenAIFunctionToolSchema, ToolResponse\nfrom verl.utils.model import compute_position_id_with_mask\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\nBASE_CHAT_HISTORY = [\n    {\"role\": \"system\", \"content\": \"You are a helpful assistant.\"},\n    {\"role\": \"user\", \"content\": \"I am a user.\"},\n]\n\n\nclass FinishReasonTypeEnum(str, Enum):\n    \"\"\"The enum for finish reason type.\"\"\"\n\n    LENGTH = \"length\"\n    STOP = \"stop\"\n    TOOL_CALL = \"tool_calls\"\n\n    @classmethod\n    def from_str(cls, value: str) -> \"FinishReasonTypeEnum\":\n        if value == \"stop\":\n            return cls.STOP\n        elif value == \"length\":\n            return cls.LENGTH\n        elif value == \"tool_calls\":\n            return cls.TOOL_CALL\n        else:\n            raise ValueError(f\"Unsupported finish reason type: {value}\")\n\n\nclass Message(BaseModel):\n    role: str\n    content: str | dict[str, Any] | list[dict[str, Any]] | ToolResponse\n    tool_calls: Optional[list[OpenAIFunctionToolCall]] = None\n\n\nclass AsyncRolloutRequestStateEnum(str, Enum):\n    \"\"\"The enum for async rollout request state.\"\"\"\n\n    PENDING = \"pending\"\n    RUNNING = \"running\"\n    COMPLETED = \"completed\"\n    FAILED = \"failed\"\n    TOOL_CALLING = \"tool_calling\"\n    INTERACTING = \"interacting\"\n\n\nclass TokenizationSanityCheckModeEnum(str, Enum):\n    \"\"\"The enum for tokenization sanity check mode.\"\"\"\n\n    DISABLE = \"disable\"\n    STRICT = \"strict\"\n    IGNORE_STRIPPABLE = \"ignore_strippable\"\n\n\nclass AsyncRolloutRequest(BaseModel):\n    \"\"\"The data model for async rollout.\"\"\"\n\n    model_config = ConfigDict(arbitrary_types_allowed=True)\n\n    batch_data_id: int = 0\n    rollout_offset: int = 0\n    request_id: str\n    state: AsyncRolloutRequestStateEnum\n    messages: list[Message]\n    multi_modal_keys: Optional[list[str]] = None\n    multi_modal_data: Optional[dict[str, Any]] = None\n    multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None\n    tool_schemas: Optional[list[OpenAIFunctionToolSchema]] = None\n    tools_kwargs: dict[str, Any] = {}\n    interaction_kwargs: dict[str, Any] = {}\n    input_ids: Optional[torch.Tensor] = None\n    prompt_ids: Optional[torch.Tensor] = None\n    response_ids: Optional[torch.Tensor] = None\n    attention_mask: Optional[torch.Tensor] = None\n    prompt_attention_mask: Optional[torch.Tensor] = None\n    response_attention_mask: Optional[torch.Tensor] = None\n    position_ids: Optional[torch.Tensor] = None\n    prompt_position_ids: Optional[torch.Tensor] = None\n    response_position_ids: Optional[torch.Tensor] = None\n    loss_mask: Optional[torch.Tensor] = None\n    prompt_loss_mask: Optional[torch.Tensor] = None\n    response_loss_mask: Optional[torch.Tensor] = None\n    reward_scores: dict[str, float]\n    max_prompt_len: int\n    max_response_len: int = 8192\n    max_model_len: int = 32768\n    metrics: dict[str, list[Any]] = {}\n    output_token_ids: torch.Tensor | None = None\n    rollout_log_probs: torch.Tensor | None = None\n\n    use_inference_chat_template: bool\n    tokenization_sanity_check_mode: TokenizationSanityCheckModeEnum\n    generation_prompt_ids: Optional[torch.Tensor] = None\n    base_conv_wo_gen_prompt_end_pos: int\n    base_conv_with_gen_prompt_end_pos: int\n\n    @model_validator(mode=\"before\")\n    @classmethod\n    def initialize_request(cls, values):\n        if not (messages := values.get(\"messages\")):\n            raise ValueError(\"messages is required for AsyncRolloutRequest initialization\")\n        if not (max_prompt_len := values.get(\"max_prompt_len\")):\n            raise ValueError(\"max_prompt_len is required for AsyncRolloutRequest initialization\")\n        if not (processing_class := values.pop(\"processing_class\", None)):\n            raise ValueError(\"processing_class is required for AsyncRolloutRequest initialization\")\n\n        values[\"messages\"] = [Message.model_validate(msg) for msg in messages]\n\n        # If there is no multi_modal_keys, we assume the multi-modal data is image and video.\n        if not values.get(\"multi_modal_keys\"):\n            values[\"multi_modal_keys\"] = [\"image\", \"video\"]\n        if not values.get(\"multi_modal_data\"):\n            values[\"multi_modal_data\"] = {key: [] for key in values[\"multi_modal_keys\"]}\n        else:\n            # check if all multi_modal_keys are in multi_modal_data\n            for key in values[\"multi_modal_keys\"]:\n                if key not in values[\"multi_modal_data\"]:\n                    values[\"multi_modal_data\"][key] = []\n        if not values.get(\"multi_modal_inputs\"):\n            values[\"multi_modal_inputs\"] = {}\n\n        tools = (\n            [tool.model_dump() for tool in tool_schemas] if (tool_schemas := values.get(\"tool_schemas\", [])) else None\n        )\n\n        multi_modal_data = values[\"multi_modal_data\"]\n        tokens_without_prompt = cls._handle_apply_chat_template(\n            processing_class,\n            messages,\n            multi_modal_data=multi_modal_data,\n            tools=tools,\n            add_generation_prompt=False,\n            tokenize=True,\n        )\n        if (\n            values.get(\"input_ids\") is None\n            or values.get(\"attention_mask\") is None\n            or values.get(\"position_ids\") is None\n        ):\n            tokenization_dict_with_prompt = cls._handle_apply_chat_template(\n                processing_class,\n                messages,\n                multi_modal_data=multi_modal_data,\n                tools=tools,\n                add_generation_prompt=True,\n                tokenize=True,\n                return_dict=True,\n            )\n\n            values[\"input_ids\"], values[\"attention_mask\"] = (\n                tokenization_dict_with_prompt[\"input_ids\"],\n                tokenization_dict_with_prompt[\"attention_mask\"],\n            )\n            if values[\"input_ids\"].shape[-1] > max_prompt_len:\n                # Only log the warning to avoid truncating in the middle of generation prompt. Consider raising an\n                # error for this case in the future.\n                # Ensure batch_data_id exists with default value if not provided\n                if \"batch_data_id\" not in values:\n                    values[\"batch_data_id\"] = cls.model_fields[\"batch_data_id\"].default\n                logger.warning(\n                    f\"Prompt {values['batch_data_id']} has length {values['input_ids'].shape[-1]} \"\n                    f\"which is greater than max_prompt_len {max_prompt_len} after applied chat template with tools.\"\n                )\n\n            # Process multi_modal_inputs\n            multi_modal_inputs = tokenization_dict_with_prompt.copy()\n            multi_modal_inputs.pop(\"input_ids\", None)\n            multi_modal_inputs.pop(\"attention_mask\", None)\n            values[\"multi_modal_inputs\"] = multi_modal_inputs\n\n            values[\"position_ids\"] = values[\"prompt_position_ids\"] = cls._get_position_ids(\n                processing_class, values[\"input_ids\"], values[\"attention_mask\"], multi_modal_inputs\n            )\n\n        values[\"prompt_ids\"], values[\"prompt_attention_mask\"] = values[\"input_ids\"], values[\"attention_mask\"]\n        values[\"loss_mask\"] = values[\"prompt_loss_mask\"] = torch.zeros_like(values[\"input_ids\"], dtype=torch.bool)\n        values[\"generation_prompt_ids\"] = values[\"input_ids\"][..., tokens_without_prompt.shape[-1] :]\n        values[\"base_conv_wo_gen_prompt_end_pos\"] = cls._handle_apply_chat_template(\n            processing_class,\n            BASE_CHAT_HISTORY,\n            multi_modal_data=multi_modal_data,\n            tools=tools,\n            add_generation_prompt=False,\n            tokenize=True,\n        ).shape[-1]\n\n        values[\"base_conv_with_gen_prompt_end_pos\"] = cls._handle_apply_chat_template(\n            processing_class,\n            BASE_CHAT_HISTORY,\n            multi_modal_data=multi_modal_data,\n            tools=tools,\n            add_generation_prompt=True,\n            tokenize=True,\n        ).shape[-1]\n\n        return values\n\n    @staticmethod\n    def _handle_apply_chat_template(\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        messages: list[Message],\n        multi_modal_data: dict[str, Any],\n        tools: Optional[list[OpenAIFunctionToolSchema]] = None,\n        add_generation_prompt: bool = False,\n        tokenize: bool = False,\n        return_dict: bool = False,\n    ):\n        raw_prompt = processing_class.apply_chat_template(\n            messages, tools=tools, add_generation_prompt=add_generation_prompt, tokenize=False\n        )\n        if not tokenize:\n            return raw_prompt\n\n        if isinstance(processing_class, PreTrainedTokenizer) or isinstance(processing_class, PreTrainedTokenizerFast):\n            if any(len(values) > 0 for values in multi_modal_data.values()):\n                logger.warning(\n                    \"There is multi_modal_data but you are not using a processor. Multi-modal data will be ignored.\"\n                )\n            model_inputs = processing_class(text=[raw_prompt], return_tensors=\"pt\")\n        elif isinstance(processing_class, ProcessorMixin):\n            # When we update multi_model_keys, we also need to update this logic\n            images = images if len(images := multi_modal_data.get(\"image\", [])) > 0 else None\n            videos = videos if len(videos := multi_modal_data.get(\"video\", [])) > 0 else None\n            model_inputs = processing_class(text=[raw_prompt], images=images, videos=videos, return_tensors=\"pt\")\n        else:\n            raise ValueError(f\"Unsupported processing class type: {type(processing_class)}\")\n\n        model_inputs = dict(model_inputs)\n        if return_dict:\n            return model_inputs\n        else:\n            return model_inputs[\"input_ids\"]\n\n    @staticmethod\n    def _get_position_ids(\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        input_ids: torch.Tensor,\n        attention_mask: torch.Tensor,\n        multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None,\n    ) -> torch.Tensor:\n        # special case for qwen2vl\n        is_qwen2vl = (\n            hasattr(processing_class, \"image_processor\")\n            and \"Qwen2VLImageProcessor\" in processing_class.image_processor.__class__.__name__\n        )\n        if is_qwen2vl:\n            from verl.models.transformers.qwen2_vl import get_rope_index\n\n            image_grid_thw = video_grid_thw = second_per_grid_ts = None\n            if multi_modal_inputs:\n                image_grid_thw = multi_modal_inputs.get(\"image_grid_thw\")\n                video_grid_thw = multi_modal_inputs.get(\"video_grid_thw\")\n                second_per_grid_ts = multi_modal_inputs.get(\"second_per_grid_ts\")\n\n            assert input_ids.dim() == 2 and input_ids.shape[0] == 1, (\n                f\"input_ids should be 2D with batch size 1, but got shape {input_ids.shape}\"\n            )\n            assert attention_mask.dim() == 2 and attention_mask.shape[0] == 1, (\n                f\"attention_mask should be 2D with batch size 1, but got shape {attention_mask.shape}\"\n            )\n            new_position_ids = get_rope_index(\n                processing_class,\n                input_ids=input_ids.squeeze(0),\n                image_grid_thw=image_grid_thw,\n                video_grid_thw=video_grid_thw,\n                second_per_grid_ts=second_per_grid_ts,\n                attention_mask=attention_mask.squeeze(0),\n            )\n            return new_position_ids  # (3, seq_len)\n        else:\n            return compute_position_id_with_mask(attention_mask)  # (1, seq_len)\n\n    def _update_input_ids(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        new_input_ids: torch.Tensor,\n        attention_mask: bool,\n        loss_mask: bool,\n        new_multi_modal_inputs: Optional[dict[str, torch.Tensor]] = None,\n    ) -> None:\n        \"\"\"\n        Update the input_ids, attention_mask, position_ids, and loss_mask of the request in additive manner.\n        \"\"\"\n        self.input_ids = torch.cat([self.input_ids, new_input_ids], dim=-1)\n        attention_mask = torch.ones_like(new_input_ids) * int(attention_mask)\n        self.attention_mask = torch.cat([self.attention_mask, attention_mask], dim=-1)\n        loss_mask = torch.ones_like(new_input_ids) * int(loss_mask)\n        self.loss_mask = torch.cat([self.loss_mask, loss_mask], dim=-1)\n\n        if new_multi_modal_inputs:\n            self._update_multi_modal_inputs(new_multi_modal_inputs)\n\n        new_position_ids = self._get_position_ids(\n            processing_class, new_input_ids, attention_mask, new_multi_modal_inputs\n        )\n\n        last_pos = self.position_ids[..., -1:]\n        new_position_ids = new_position_ids + (last_pos + 1)\n\n        self.position_ids = torch.cat([self.position_ids, new_position_ids], dim=-1)\n\n        assert (\n            self.input_ids.shape[-1]\n            == self.attention_mask.shape[-1]\n            == self.position_ids.shape[-1]\n            == self.loss_mask.shape[-1]\n        ), f\"\"\"Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, \n            {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}\"\"\"\n\n    def _update_multi_modal_inputs(self, new_multi_modal_inputs: dict[str, torch.Tensor]) -> None:\n        \"\"\"\n        Update the multi_modal_inputs of the request in additive manner.\n        \"\"\"\n        for key in new_multi_modal_inputs:\n            input_tensor = new_multi_modal_inputs[key]\n            self.multi_modal_inputs[key] = (\n                torch.cat([self.multi_modal_inputs[key], input_tensor], dim=0)\n                if key in self.multi_modal_inputs\n                else input_tensor\n            )\n\n    def get_generation_prompt_ids(\n        self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin\n    ) -> list[int]:\n        \"\"\"\n        Get the generation prompt ids for rollout engine.\n\n        Because rollout engine(SGLang) requires the ids to be a list, we need to convert the tensor to a list.\n        \"\"\"\n        generation_prompt_ids = (\n            None\n            if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all()\n            else self.generation_prompt_ids\n        )\n        if generation_prompt_ids is not None:\n            self._update_input_ids(processing_class, generation_prompt_ids, attention_mask=True, loss_mask=False)\n\n        if self.use_inference_chat_template:\n            messages = [msg.model_dump() for msg in self.messages]\n            tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None\n            generation_prompt_ids = self._handle_apply_chat_template(\n                processing_class,\n                messages,\n                multi_modal_data=self.multi_modal_data,\n                tools=tools,\n                add_generation_prompt=True,\n                tokenize=True,\n            )\n            return generation_prompt_ids.squeeze(0).tolist()\n        else:\n            return self.input_ids.squeeze(0).tolist()\n\n    def add_user_message(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        content: str,\n    ) -> None:\n        self.messages.append(Message(role=\"user\", content=content))\n        messages = [*BASE_CHAT_HISTORY, self.messages[-1]]\n        tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None\n\n        # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine\n        # Inference, it is pure text.\n        content_ids = self._handle_apply_chat_template(\n            processing_class, messages, multi_modal_data={}, tools=tools, add_generation_prompt=False, tokenize=True\n        )[..., self.base_conv_wo_gen_prompt_end_pos :]\n        self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=False)\n\n    def add_assistant_message(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        content: str,\n        content_ids: Optional[torch.Tensor] = None,\n        tool_calls: Optional[list[OpenAIFunctionToolCall]] = None,\n    ) -> None:\n        self.messages.append(Message(role=\"assistant\", content=content, tool_calls=tool_calls))\n        if content_ids is None:\n            messages = [*BASE_CHAT_HISTORY, self.messages[-1]]\n            tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None\n\n            # We don't need to pass multi_modal_data here because we don't have any multi-modal data from Engine\n            # Inference, it is pure text.\n            content_ids = self._handle_apply_chat_template(\n                processing_class, messages, multi_modal_data={}, tools=tools, add_generation_prompt=False, tokenize=True\n            )[..., self.base_conv_with_gen_prompt_end_pos :]\n        self._update_input_ids(processing_class, content_ids, attention_mask=True, loss_mask=True)\n\n    def add_tool_response_messages(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        contents: list[ToolResponse],\n    ) -> None:\n        if not contents or all(content.is_empty() for content in contents):\n            return\n        # We also handle the case when tool returns image\n        # We require the processing of the image and video to be done at tool.execute() level\n        delta_multi_modal_data = {key: [] for key in self.multi_modal_keys}\n        for content in contents:\n            if content.is_text_only():\n                self.messages.append(Message(role=\"tool\", content=content.text))\n            else:\n                content_list = []\n                # When we update multi_model_keys, we also need to update this logic\n                if content.image:\n                    content_list.extend([{\"type\": \"image\"} for _ in content.image])\n                    delta_multi_modal_data[\"image\"].extend(content.image)\n                if content.video:\n                    content_list.extend([{\"type\": \"video\"} for _ in content.video])\n                    delta_multi_modal_data[\"video\"].extend(content.video)\n                if content.text:\n                    content_list.append({\"type\": \"text\", \"text\": content.text})\n                self.messages.append(Message(role=\"tool\", content=content_list))\n\n        messages = [*BASE_CHAT_HISTORY, *self.messages[-len(contents) :]]\n        tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None\n\n        for key in self.multi_modal_keys:\n            if len(delta_multi_modal_data[key]) > 0:\n                self.multi_modal_data[key].extend(delta_multi_modal_data[key])\n\n        # We just passed the new multi-modal data to the chat template to update the input_ids.\n        content_info = self._handle_apply_chat_template(\n            processing_class,\n            messages,\n            multi_modal_data=delta_multi_modal_data,\n            tools=tools,\n            add_generation_prompt=False,\n            tokenize=True,\n            return_dict=True,\n        )\n        content_ids = content_info[\"input_ids\"][..., self.base_conv_wo_gen_prompt_end_pos :]\n\n        # process multi_modal_inputs\n        multi_modal_inputs = content_info.copy()\n        multi_modal_inputs.pop(\"input_ids\", None)\n        multi_modal_inputs.pop(\"attention_mask\", None)\n\n        # chat templates include generation prompt tokens (e.g., \"<im_start>assistant\\n\")\n        # So when tool response is added, we need to explicitly remove these tokens.\n        self._remove_generation_prompt_ids_if_present()\n\n        self._update_input_ids(\n            processing_class,\n            content_ids,\n            attention_mask=True,\n            loss_mask=False,\n            new_multi_modal_inputs=multi_modal_inputs,\n        )\n\n    def update_metrics(self, metrics: Any, tool_id: str) -> None:\n        \"\"\"\n        metrics: should be a dict of tools_name -> Any\n        \"\"\"\n        if self.metrics.get(tool_id) is None:\n            self.metrics[tool_id] = []\n        self.metrics[tool_id].append(metrics)\n\n    def _get_prompt_diffs(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        full_prompt_ids: torch.Tensor,\n        current_prompt_ids: torch.Tensor,\n        diff_surrounding_chars: int = 10,\n    ) -> list[dict[str, Any]]:\n        \"\"\"Get differences between full prompt and current prompt with surrounding context.\n\n        This function helps debug tokenization mismatches by showing the differences between\n        full prompt and current prompt with surrounding context. Instead of just showing\n        the exact diff, it includes additional tokens before and after to help locate\n        the issue in the chat template.\n\n        For example, if the actual diff is a newline change from \"\\n\\n\" to \"\\n\", with\n        diff_surrounding_chars the output might look like:\n\n        full_prompt_chunk:    \"<|im_start|>assistant\\n\\nI think...\"\n        current_prompt_chunk: \"<|im_start|>assistant\\nI think...\"\n\n        This context makes it much easier to identify where in the chat template the\n        mismatch occurs.\n\n        Args:\n            processing_class: The processing class to use for decoding the token IDs\n            full_prompt_ids: Token IDs from applying chat template to all messages at once\n            current_prompt_ids: Token IDs from incremental chat template application\n            diff_surrounding_chars: Number of surrounding characters to include for context (default: 10)\n\n        Returns:\n            List of dicts containing the differing chunks with context and their indices\n        \"\"\"\n        full_prompt_ids = full_prompt_ids.squeeze(0)\n        current_prompt_ids = current_prompt_ids.squeeze(0)\n        full_prompt = processing_class.decode(full_prompt_ids, skip_special_tokens=False)\n        current_prompt = processing_class.decode(current_prompt_ids, skip_special_tokens=False)\n        s = difflib.SequenceMatcher(None, full_prompt, current_prompt, autojunk=False)\n        diffs = []\n        for tag, i1, i2, j1, j2 in s.get_opcodes():\n            if tag == \"equal\":\n                continue\n\n            # Get the surrounding context for better readability\n            start_i = max(0, i1 - diff_surrounding_chars)\n            end_i = min(len(full_prompt), i2 + diff_surrounding_chars)\n            start_j = max(0, j1 - diff_surrounding_chars)\n            end_j = min(len(current_prompt), j2 + diff_surrounding_chars)\n\n            diffs.append(\n                {\n                    \"full_prompt_chunk\": full_prompt[start_i:end_i],\n                    \"current_prompt_chunk\": current_prompt[start_j:end_j],\n                    \"indices\": (start_i, end_i, start_j, end_j),\n                }\n            )\n        return diffs\n\n    def _remove_generation_prompt_ids_if_present(self) -> None:\n        \"\"\"\n        Remove generation prompt IDs from input tensors if they are present at the end.\n        \"\"\"\n        if self.input_ids[..., -self.generation_prompt_ids.shape[-1] :].eq(self.generation_prompt_ids).all():\n            self.input_ids = self.input_ids[..., : -self.generation_prompt_ids.shape[-1]]\n            self.attention_mask = self.attention_mask[..., : -self.generation_prompt_ids.shape[-1]]\n            self.position_ids = self.position_ids[..., : -self.generation_prompt_ids.shape[-1]]\n            self.loss_mask = self.loss_mask[..., : -self.generation_prompt_ids.shape[-1]]\n\n    def finalize(\n        self,\n        processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin,\n        reward_scores: dict[str, list[float]],\n        finish_reason_type: FinishReasonTypeEnum = FinishReasonTypeEnum.STOP,\n    ) -> None:\n        self.state = AsyncRolloutRequestStateEnum.COMPLETED\n        self.reward_scores = reward_scores\n\n        # In case we failed to generate the assistant message and the generation prompt ids were already added to\n        # input_ids, remove them from the end of input_ids\n        self._remove_generation_prompt_ids_if_present()\n\n        self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :]\n\n        if self.tokenization_sanity_check_mode != TokenizationSanityCheckModeEnum.DISABLE:\n            # When there is a diff, we log the diffs with diff_surrounding_chars context\n            diff_surrounding_chars = 10\n\n            messages = [msg.model_dump() for msg in self.messages]\n            tools = [tool.model_dump() for tool in self.tool_schemas] if self.tool_schemas else None\n            full_prompt_info = self._handle_apply_chat_template(\n                processing_class,\n                messages,\n                multi_modal_data=self.multi_modal_data,\n                tools=tools,\n                add_generation_prompt=False,\n                tokenize=True,\n                return_dict=True,\n            )\n            full_prompt_ids = full_prompt_info[\"input_ids\"]\n\n            # We must use dict(full_prompt_info) to convert BatchFeature values to a new dict\n            # because np.array() only keeps the keys for BatchFeature.\n            full_prompt_multi_modal_inputs = full_prompt_info.copy()\n            full_prompt_multi_modal_inputs.pop(\"input_ids\", None)\n            full_prompt_multi_modal_inputs.pop(\"attention_mask\", None)\n\n            for multi_modal_inputs_key in self.multi_modal_inputs:\n                if multi_modal_inputs_key in full_prompt_multi_modal_inputs:\n                    if (\n                        not self.multi_modal_inputs[multi_modal_inputs_key]\n                        .eq(full_prompt_multi_modal_inputs[multi_modal_inputs_key])\n                        .all()\n                    ):\n                        logger.warning(\n                            f\"Multi-modal data {multi_modal_inputs_key} is not consistent. \"\n                            f\"This may lead to unexpected behavior during training. \"\n                            f\"Please review your multi_modal_inputs logic.\"\n                        )\n                else:\n                    logger.warning(\n                        f\"Multi-modal inputs key {multi_modal_inputs_key} is not found in the multi_modal_inputs. \"\n                        f\"This may lead to unexpected behavior during training.\"\n                        f\"Please review your multi_modal_inputs logic.\"\n                    )\n\n            if diffs := self._get_prompt_diffs(\n                processing_class, full_prompt_ids, self.input_ids, diff_surrounding_chars=diff_surrounding_chars\n            ):\n                log_warning = False\n                if self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.STRICT:\n                    log_warning = True\n                elif self.tokenization_sanity_check_mode == TokenizationSanityCheckModeEnum.IGNORE_STRIPPABLE:\n                    non_strippable_diffs_exist = any(\n                        d[\"full_prompt_chunk\"].strip() or d[\"current_prompt_chunk\"].strip() for d in diffs\n                    )\n                    if non_strippable_diffs_exist:\n                        log_warning = True\n\n                if log_warning:\n                    mode_str = f\" ({self.tokenization_sanity_check_mode.value})\"\n                    logger.warning(\n                        f\"Inconsistent training and inference tokenization detected{mode_str}. This may lead to \"\n                        f\"unexpected behavior during training. Please review your chat template to determine if this \"\n                        f\"is intentional. For more information, refer to the multiturn README.md.\"\n                    )\n                    logger.warning(\n                        f\"Showing {diff_surrounding_chars} characters before and after the diffs for context and \"\n                        f\"better readability.\"\n                    )\n                    diff_details_list = []\n                    for d in diffs:\n                        i1, i2, j1, j2 = d[\"indices\"]\n                        diff_details_list.append(\n                            f\"idx {i1}:{i2} -> {j1}:{j2} | full_prompt_chunk: {repr(d['full_prompt_chunk'])} | \"\n                            f\"current_prompt_chunk: {repr(d['current_prompt_chunk'])}\"\n                        )\n                    diff_details = \"\\n\".join(diff_details_list)\n                    logger.warning(f\"Found differences:\\n{diff_details}\")\n\n        if finish_reason_type == FinishReasonTypeEnum.STOP:\n            pass\n        elif finish_reason_type == FinishReasonTypeEnum.LENGTH:\n            pass\n        else:\n            raise ValueError(f\"Unsupported finalize finish reason type: {finish_reason_type}\")\n        self.truncate_output_ids(processing_class)\n\n        assert (\n            self.input_ids.shape[-1]\n            == self.attention_mask.shape[-1]\n            == self.position_ids.shape[-1]\n            == self.loss_mask.shape[-1]\n        ), f\"\"\"Request {self.request_id} has different length of {self.input_ids.shape[-1]=}, \n            {self.attention_mask.shape[-1]=}, {self.position_ids.shape[-1]=}, {self.loss_mask.shape[-1]=}\"\"\"\n\n    def truncate_output_ids(\n        self, processing_class: PreTrainedTokenizer | PreTrainedTokenizerFast | ProcessorMixin\n    ) -> None:\n        self.input_ids = self.input_ids[..., : self.max_model_len]\n        self.attention_mask = self.attention_mask[..., : self.max_model_len]\n        self.position_ids = self.position_ids[..., : self.max_model_len]\n        self.loss_mask = self.loss_mask[..., : self.max_model_len]\n        self.response_ids = self.input_ids[..., self.prompt_ids.shape[-1] :][..., : self.max_response_len]\n        self.response_attention_mask = self.attention_mask[..., self.prompt_attention_mask.shape[-1] :][\n            ..., : self.max_response_len\n        ]\n        self.response_position_ids = self.position_ids[..., self.prompt_position_ids.shape[-1] :][\n            ..., : self.max_response_len\n        ]\n        self.response_loss_mask = self.loss_mask[..., self.prompt_loss_mask.shape[-1] :][..., : self.max_response_len]\n"
  },
  {
    "path": "verl/workers/rollout/sglang_rollout/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n"
  },
  {
    "path": "verl/workers/rollout/sglang_rollout/async_sglang_server.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport dataclasses\nimport json\nimport logging\nimport os\nfrom typing import Any, Optional\n\nimport ray\nimport sglang\nimport sglang.srt.entrypoints.engine\nimport torch\nfrom packaging import version\nfrom ray.actor import ActorHandle\nfrom sglang.srt.entrypoints.http_server import (\n    ServerArgs,\n    _GlobalState,\n    _launch_subprocesses,\n    app,\n    set_global_state,\n)\nfrom sglang.srt.managers.io_struct import (\n    ContinueGenerationReqInput,\n    GenerateReqInput,\n    PauseGenerationReqInput,\n    ReleaseMemoryOccupationReqInput,\n    ResumeMemoryOccupationReqInput,\n)\nfrom sglang.srt.managers.tokenizer_manager import ServerStatus\n\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import get_visible_devices_keyword\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\nfrom verl.utils.profiler import DistProfiler, build_sglang_profiler_args\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput\nfrom verl.workers.rollout.sglang_rollout.sglang_rollout import _set_envs_and_config\nfrom verl.workers.rollout.utils import get_max_position_embeddings, run_uvicorn\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(logging.INFO)\n\nvisible_devices_keyword = get_visible_devices_keyword()\n\n\nclass SGLangHttpServer:\n    \"\"\"SGLang http server in single node, this is equivalent to launch server with command line:\n    ```\n    python -m sglang.launch_server --node-rank 0 --nnode 1 ...\n    ```\n\n    Args:\n        config (DictConfig): full config.\n        rollout_mode (RolloutMode): rollout mode.\n        replica_rank (int): replica rank, a replica may contain multiple nodes.\n        node_rank (int): node rank.\n        nnodes (int): number of nodes.\n        cuda_visible_devices (str): cuda visible devices.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        rollout_mode: RolloutMode,\n        workers: list[ActorHandle],\n        replica_rank: int,\n        node_rank: int,\n        nnodes: int,\n        cuda_visible_devices: str,\n        base_gpu_id: int,\n    ):\n        print(f\"SGLang http server: {rollout_mode=}, {replica_rank=}, {node_rank=}, {nnodes=}, {cuda_visible_devices=}\")\n        os.environ[visible_devices_keyword] = cuda_visible_devices\n\n        self.config: RolloutConfig = omega_conf_to_dataclass(config)\n        self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)\n        max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config)\n        if self.config.max_model_len is None:\n            self.config.max_model_len = max_position_embeddings\n        else:\n            if self.config.max_model_len > max_position_embeddings:\n                raise ValueError(\n                    f\"max_model_len ({self.config.max_model_len}) should be less than or equal to \"\n                    f\"max_position_embeddings ({max_position_embeddings})\"\n                )\n        self.rollout_mode = rollout_mode\n        self.workers = workers\n\n        self.replica_rank = replica_rank\n        self.node_rank = node_rank\n        self.nnodes = nnodes\n        self.base_gpu_id = base_gpu_id\n        # model weights version, set by ServerAdapter when update weights.\n        self.global_steps = None\n\n        if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == \"dummy\":\n            logger.warning(f\"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto\")\n            self.config.load_format = \"auto\"\n\n        # used for http server\n        self._server_address = ray.util.get_node_ip_address().strip(\"[]\")\n        self._server_port = None\n\n        # used for controlling sglang server profiler\n        profiler_config = self.config.profiler\n        tool_config = None\n        if profiler_config is not None:\n            if profiler_config.tool in [\"torch\", \"npu\"]:\n                tool_config = omega_conf_to_dataclass((profiler_config.tool_config or {}).get(profiler_config.tool))\n            else:\n                logger.warning(f\"agent loop only support torch and npu profiler, got {profiler_config.tool}\")\n                profiler_config = None\n        self.profiler_controller = DistProfiler(self.replica_rank, config=profiler_config, tool_config=tool_config)\n\n        # For multi-node, we need dist_init_addr so nodes can coordinate NCCL init.\n        # For single-node, let SGLang handle port selection internally via nccl_port,\n        # which also avoids port conflicts.\n        self._master_address = None\n        self._master_port = None\n        self._master_sock = None\n        if self.nnodes > 1 and self.node_rank == 0:\n            self._master_address = self._server_address\n            self._master_port, self._master_sock = get_free_port(self._server_address, with_alive_sock=True)\n            logger.info(\n                f\"SGLangHttpServer, replica_rank: {self.replica_rank}, \"\n                f\"master address: {self._master_address}, port: {self._master_port}\"\n            )\n\n    def get_master_address(self):\n        \"\"\"Get master address and port for init NCCL process group.\"\"\"\n        return self._master_address, self._master_port\n\n    def get_server_address(self):\n        \"\"\"Get http server address and port.\"\"\"\n        assert self._server_port is not None, \"http server is not launched, port is None\"\n        return self._server_address, self._server_port\n\n    async def launch_server(self, master_address: str = None, master_port: int = None):\n        if self.nnodes > 1:\n            if self.node_rank != 0:\n                assert master_address and master_port, \"non-master node should provide master address and port\"\n                self._master_address = master_address\n                self._master_port = master_port\n            else:\n                self._master_sock.close()\n\n        engine_kwargs = self.config.get(\"engine_kwargs\", {}).get(\"sglang\", {}) or {}\n        attention_backend = engine_kwargs.pop(\"attention_backend\", None)\n        quantization = self.config.get(\"quantization\", None)\n        if quantization is not None:\n            if quantization == \"fp8\":\n                assert version.parse(sglang.__version__) >= version.parse(\"0.5.5\"), (\n                    \"sglang>=0.5.5 is required for FP8 quantization\"\n                )\n                FP8_BLOCK_QUANT_KWARGS = {\n                    \"activation_scheme\": \"dynamic\",\n                    \"fmt\": \"e4m3\",\n                    \"quant_method\": \"fp8\",\n                    \"weight_block_size\": [128, 128],\n                }\n                fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS)\n            else:\n                raise ValueError(f\"Currently only support fp8 quantization, got: {quantization}\")\n        infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size\n        args = {\n            \"model_path\": self.model_config.local_path,\n            \"dtype\": self.config.dtype,\n            \"mem_fraction_static\": self.config.gpu_memory_utilization,\n            \"disable_cuda_graph\": self.config.enforce_eager,\n            \"enable_memory_saver\": True,\n            \"base_gpu_id\": self.base_gpu_id,\n            \"gpu_id_step\": 1,\n            \"tp_size\": infer_tp,\n            \"dp_size\": self.config.data_parallel_size,\n            \"ep_size\": self.config.expert_parallel_size,\n            \"node_rank\": self.node_rank,\n            \"load_format\": self.config.load_format,\n            \"nnodes\": self.nnodes,\n            \"trust_remote_code\": self.model_config.trust_remote_code,\n            \"max_running_requests\": self.config.get(\"max_num_seqs\", None),\n            \"log_level\": \"error\",\n            \"mm_attention_backend\": \"fa3\",\n            \"attention_backend\": attention_backend if attention_backend is not None else \"fa3\",\n            \"skip_tokenizer_init\": self.config.skip_tokenizer_init,\n            \"skip_server_warmup\": True,\n            \"quantization\": quantization,\n            \"json_model_override_args\": json.dumps({\"quantization_config\": fp8_block_quant_kwargs})\n            if quantization == \"fp8\"\n            else json.dumps({}),\n            **engine_kwargs,\n        }\n\n        # Only set dist_init_addr for multi-node; for single-node, let SGLang\n        # handle port selection internally via nccl_port to avoid conflicts.\n        if self.nnodes > 1:\n            dist_init_addr = (\n                f\"[{self._master_address}]:{self._master_port}\"\n                if is_valid_ipv6_address(self._master_address)\n                else f\"{self._master_address}:{self._master_port}\"\n            )\n            args[\"dist_init_addr\"] = dist_init_addr\n\n        if self.config.prometheus.enable:\n            if self.config.prometheus.served_model_name:\n                # Extract model name from path if it's a full path\n                served_model_name = self.config.prometheus.served_model_name\n                if \"/\" in served_model_name:\n                    # If it's a full path, extract the last part as model name\n                    served_model_name = served_model_name.split(\"/\")[-1]\n                args[\"served_model_name\"] = served_model_name\n\n            # start sglang metrics\n            args[\"enable_metrics\"] = True\n\n        # enable_weights_cpu_backup is supported in sglang>=0.5.3\n        if \"enable_weights_cpu_backup\" in [f.name for f in dataclasses.fields(ServerArgs)]:\n            enable_weights_cpu_backup = True if self.rollout_mode == RolloutMode.COLOCATED else False\n            args[\"enable_weights_cpu_backup\"] = enable_weights_cpu_backup\n\n        if self.config.enable_rollout_routing_replay:\n            args.update({\"enable_return_routed_experts\": True})\n\n        # mtp\n        if self.config.mtp.enable and self.config.mtp.enable_rollout:\n            # Enable weights CPU backup for sglang >= 0.5.6\n            if sglang.__version__ < \"0.5.6\":\n                raise ValueError(f\"sglang version {sglang.__version__} is not supported for MTP rollout\")\n\n            args[\"speculative_algorithm\"] = self.config.mtp.speculative_algorithm\n            args[\"speculative_num_steps\"] = self.config.mtp.speculative_num_steps\n            args[\"speculative_eagle_topk\"] = self.config.mtp.speculative_eagle_topk\n            args[\"speculative_num_draft_tokens\"] = self.config.mtp.speculative_num_draft_tokens\n\n            args[\"enable_weights_cpu_backup\"] = True\n            args[\"enable_draft_weights_cpu_backup\"] = True\n\n        # NOTE: We can't directly call SGLang's launch_server since it's not an async function.\n        # https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server.py\n        sglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config\n        os.environ[\"SGLANG_BLOCK_NONZERO_RANK_CHILDREN\"] = \"0\"\n        server_args = ServerArgs(**args)\n        if version.parse(sglang.__version__) >= version.parse(\"0.5.7\"):\n            self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses(\n                server_args=server_args,\n                init_tokenizer_manager_func=sglang.srt.entrypoints.engine.init_tokenizer_manager,\n                run_scheduler_process_func=sglang.srt.entrypoints.engine.run_scheduler_process,\n                run_detokenizer_process_func=sglang.srt.entrypoints.engine.run_detokenizer_process,\n            )\n        else:\n            self.tokenizer_manager, self.template_manager, self.scheduler_info, *_ = _launch_subprocesses(\n                server_args=server_args\n            )\n\n        # In multi-node cases, non-zero rank nodes should not launch http server.\n        if self.node_rank > 0:\n            return\n\n        set_global_state(\n            _GlobalState(\n                tokenizer_manager=self.tokenizer_manager,\n                template_manager=self.template_manager,\n                scheduler_info=self.scheduler_info,\n            )\n        )\n        app.is_single_tokenizer_mode = True\n\n        # Set warmup_thread_{kw}args to avoid AttributeError in lifespan function\n        app.server_args = server_args\n        app.warmup_thread_kwargs = {\"server_args\": server_args}\n        app.warmup_thread_args = (server_args, None, None)\n\n        # Manually add Prometheus middleware before starting server\n        # This ensures /metrics endpoint is available immediately\n        if server_args.enable_metrics:\n            from sglang.srt.utils.common import add_prometheus_middleware\n\n            add_prometheus_middleware(app)\n\n        self._server_port, self._server_task = await run_uvicorn(app, server_args, self._server_address)\n        self.tokenizer_manager.server_status = ServerStatus.Up\n\n    async def wake_up(self):\n        if self.node_rank != 0:\n            return\n\n        if self.rollout_mode == RolloutMode.HYBRID:\n            # In hybrid mode, rollout is wake up in `update_weights`\n            raise ValueError(f\"wake_up not support rollout_mode {self.rollout_mode}\")\n        elif self.rollout_mode == RolloutMode.COLOCATED:\n            # Directly call engine to wake up without sync weights.\n            obj = ResumeMemoryOccupationReqInput(tags=[\"kv_cache\", \"weights\"])\n            await self.tokenizer_manager.resume_memory_occupation(obj, None)\n            await self.tokenizer_manager.flush_cache()\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            # In standalone mode, resume kv_cache if free_cache_engine is enabled\n            obj = ResumeMemoryOccupationReqInput(tags=[\"kv_cache\"])\n            await self.tokenizer_manager.resume_memory_occupation(obj, None)\n            await self.tokenizer_manager.flush_cache()\n\n    async def sleep(self):\n        if self.node_rank != 0 or not self.config.free_cache_engine:\n            return\n\n        if self.rollout_mode == RolloutMode.HYBRID:\n            obj = ReleaseMemoryOccupationReqInput(tags=[\"kv_cache\", \"weights\"])\n            await self.tokenizer_manager.release_memory_occupation(obj, None)\n        elif self.rollout_mode == RolloutMode.COLOCATED:\n            obj = ReleaseMemoryOccupationReqInput(tags=[\"kv_cache\", \"weights\"])\n            await self.tokenizer_manager.release_memory_occupation(obj, None)\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            # In standalone mode, resume kv_cache if free_cache_engine is enabled\n            obj = ReleaseMemoryOccupationReqInput(tags=[\"kv_cache\"])\n            await self.tokenizer_manager.release_memory_occupation(obj, None)\n\n    async def clear_kv_cache(self):\n        if self.node_rank == 0:\n            await self.tokenizer_manager.flush_cache()\n\n    async def generate(\n        self,\n        prompt_ids: torch.Tensor,\n        sampling_params: dict[str, Any],\n        request_id: str,\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> TokenOutput:\n        \"\"\"Generate sequence with token-in-token-out.\"\"\"\n        # TODO(@wuxibin): switch to `/generate` http endpoint once multi-modal support ready.\n        max_possible_tokens = self.config.max_model_len - len(prompt_ids)\n\n        if max_possible_tokens < 0:\n            raise ValueError(\n                f\"Prompt length ({len(prompt_ids)}) exceeds the model's maximum context length \"\n                f\"({self.config.max_model_len}).\"\n            )\n\n        if \"max_new_tokens\" in sampling_params:\n            max_new_tokens = sampling_params.pop(\"max_new_tokens\")\n        elif \"max_tokens\" in sampling_params:\n            # support vllm-style 'max_tokens' param\n            max_new_tokens = sampling_params.pop(\"max_tokens\")\n        else:\n            # Cap max_tokens by response_length to ensure tensor alignment,\n            # and by remaining budget to prevent OOM in multi-turn rollouts.\n            max_new_tokens = min(\n                self.config.response_length, self.config.prompt_length + self.config.response_length - len(prompt_ids)\n            )\n\n        # Clamp max_new_tokens to the valid range [0, max_possible_tokens]\n        max_new_tokens = max(0, min(max_new_tokens, max_possible_tokens))\n\n        assert max_new_tokens <= max_possible_tokens, (\n            f\"max_new_tokens {max_new_tokens} exceeds available context space {max_possible_tokens}\"\n        )\n        sampling_params[\"max_new_tokens\"] = max_new_tokens\n        return_logprob = sampling_params.pop(\"logprobs\", False)\n\n        request = {\n            \"rid\": request_id,\n            \"input_ids\": prompt_ids,\n            \"sampling_params\": sampling_params,\n            \"return_logprob\": return_logprob,\n            \"image_data\": image_data,\n            # TODO: support video input for sglang\n            # video_data=video_data,\n        }\n\n        if self.config.enable_rollout_routing_replay:\n            request.update({\"return_routed_experts\": True})\n\n        generate_request = GenerateReqInput(**request)\n\n        output = await self.tokenizer_manager.generate_request(generate_request, None).__anext__()\n        finish_reason = output[\"meta_info\"][\"finish_reason\"]\n        finish_reason = finish_reason[\"type\"] if finish_reason else None\n        if return_logprob:\n            output_token_logprobs = output[\"meta_info\"][\"output_token_logprobs\"]\n            log_probs, token_ids = zip(\n                *[(log_prob, token_ids) for log_prob, token_ids, _ in output_token_logprobs], strict=True\n            )\n        else:\n            token_ids = output[\"output_ids\"]\n            log_probs = None\n\n        routed_experts = None\n        if self.config.enable_rollout_routing_replay:\n            if self.config.skip_tokenizer_init:\n                routed_experts = output.get(\"meta_info\", {}).get(\"routed_experts\", None)\n            else:\n                from sglang.srt.layers.moe.routed_experts_capturer import extract_routed_experts_from_meta_info\n\n                hf_config = self.model_config.hf_config\n                if not hasattr(hf_config, \"num_hidden_layers\") or not hasattr(hf_config, \"num_experts_per_tok\"):\n                    raise AttributeError(\n                        \"enable_rollout_routing_replay is set, but hf_config is missing \"\n                        \"'num_hidden_layers' or 'num_experts_per_tok'. This feature requires an MoE model \"\n                        \"configuration that defines these attributes.\"\n                    )\n                routed_experts = extract_routed_experts_from_meta_info(output).reshape(\n                    -1, hf_config.num_hidden_layers, hf_config.num_experts_per_tok\n                )\n\n        return TokenOutput(\n            token_ids=token_ids,\n            log_probs=log_probs,\n            routed_experts=routed_experts,\n            stop_reason=finish_reason,\n            extra_fields={\"global_steps\": self.global_steps},\n        )\n\n    async def set_global_steps(self, global_steps: int):\n        \"\"\"Set the global steps of the model weights.\"\"\"\n        self.global_steps = global_steps\n\n    async def abort_all_requests(self):\n        await self.tokenizer_manager.pause_generation(PauseGenerationReqInput(mode=\"abort\"))\n\n    async def resume_generation(self):\n        await self.tokenizer_manager.continue_generation(ContinueGenerationReqInput())\n\n    async def start_profile(self, **kwargs):\n        if (\n            self.profiler_controller.check_enable()\n            and self.profiler_controller.check_this_rank()\n            and self.profiler_controller.is_discrete_mode()\n        ):\n            profile_args = build_sglang_profiler_args(\n                self.profiler_controller.config, self.profiler_controller.tool_config, self.replica_rank\n            )\n            await self.tokenizer_manager.start_profile(**profile_args)\n\n    async def stop_profile(self):\n        if (\n            self.profiler_controller.check_enable()\n            and self.profiler_controller.check_this_rank()\n            and self.profiler_controller.is_discrete_mode()\n        ):\n            await self.tokenizer_manager.stop_profile()\n\n\nclass SGLangReplica(RolloutReplica):\n    def __init__(\n        self,\n        replica_rank: int,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        gpus_per_node: int = 8,\n        is_reward_model: bool = False,\n    ):\n        super().__init__(replica_rank, config, model_config, gpus_per_node, is_reward_model)\n        self.server_class = ray.remote(SGLangHttpServer)\n\n    async def launch_servers(self):\n        \"\"\"Launch http server in each node.\"\"\"\n        assert len(self.workers) == self.world_size, (\n            f\"worker number {len(self.workers)} not equal to world size {self.world_size}\"\n        )\n\n        # get (node_id, CUDA_VISIBLE_DEVICES) of all workers\n        worker_infos = await asyncio.gather(\n            *[\n                worker.__ray_call__.remote(\n                    lambda self: (ray.get_runtime_context().get_node_id(), os.environ[visible_devices_keyword])\n                )\n                for worker in self.workers\n            ]\n        )\n        worker_cuda_visible_devices = [worker_info[1] for worker_info in worker_infos]\n        worker_node_ids = [worker_info[0] for worker_info in worker_infos]\n        base_gpu_id = 0\n        infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size\n        replica_world_size = infer_tp * self.config.pipeline_model_parallel_size\n        if os.environ.get(f\"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}\", None):\n            logger.warning(f\"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword} is set True!\")\n            base_gpu_id = (0 + self.replica_rank * replica_world_size) % self.gpus_per_node\n        # create server actor in each node with node affinity and cuda visible devices\n        for node_rank in range(self.nnodes):\n            workers = self.workers[\n                node_rank * self.gpus_per_replica_node : (node_rank + 1) * self.gpus_per_replica_node\n            ]\n            node_cuda_visible_devices_set = worker_cuda_visible_devices[\n                node_rank * self.gpus_per_replica_node : (node_rank + 1) * self.gpus_per_replica_node\n            ]\n            node_cuda_visible_devices = \",\".join(\n                map(\n                    str,\n                    sorted(\n                        set(\n                            int(device)\n                            for worker_devices_set in node_cuda_visible_devices_set\n                            for device in worker_devices_set.split(\",\")\n                            if device.strip()\n                        )\n                    ),\n                )\n            )\n\n            node_id = worker_node_ids[node_rank * self.gpus_per_replica_node]\n            name = (\n                f\"sglang_server_{self.replica_rank}_{node_rank}\"\n                if not self.is_reward_model\n                else f\"sglang_server_reward_{self.replica_rank}_{node_rank}\"\n            )\n\n            server = self.server_class.options(\n                scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                    node_id=node_id,\n                    soft=False,\n                ),\n                runtime_env={\"env_vars\": {f\"RAY_EXPERIMENTAL_NOSET_{visible_devices_keyword}\": \"1\"}},\n                name=name,\n                max_concurrency=self.max_concurrency,\n            ).remote(\n                config=self.config,\n                model_config=self.model_config,\n                rollout_mode=self.rollout_mode,\n                workers=workers,\n                replica_rank=self.replica_rank,\n                node_rank=node_rank,\n                nnodes=self.nnodes,\n                cuda_visible_devices=node_cuda_visible_devices,\n                base_gpu_id=base_gpu_id,\n            )\n            self.servers.append(server)\n\n        # launch http server in each node\n        master_address, master_port = None, None\n        if self.nnodes > 1:\n            master_address, master_port = await self.servers[0].get_master_address.remote()\n        await asyncio.gather(\n            *[\n                server.launch_server.remote(master_address=master_address, master_port=master_port)\n                for server in self.servers\n            ]\n        )\n\n        # get http server address from first server\n        server_address, server_port = await self.servers[0].get_server_address.remote()\n        self._server_handle = self.servers[0]\n        self._server_address = (\n            f\"[{server_address}]:{server_port}\"\n            if is_valid_ipv6_address(server_address)\n            else f\"{server_address}:{server_port}\"\n        )\n"
  },
  {
    "path": "verl/workers/rollout/sglang_rollout/http_server_engine.py",
    "content": "# Copyright 2025 z.ai\n# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n#\n# This file is adapted from multiple sources:\n# 1. THUDM/slime project\n#    Original source: https://github.com/THUDM/slime/blob/main/slime/backends/sglang_utils/http_server_engine.py\n#    Copyright 2025 z.ai\n#    Licensed under the Apache License, Version 2.0\n# 2. SGLang project\n#    Original source: https://github.com/sgl-project/sglang/blob/main/python/sglang/srt/entrypoints/http_server_engine.py\n#    Copyright 2023-2024 SGLang Team\n#    Licensed under the Apache License, Version 2.0\n#\n# Modifications made by z.ai and ModelBest Inc. include but are not limited to:\n# - Enhanced error handling and retry logic\n# - Added async support with connection pooling\n# - Extended functionality for distributed weight updates\n# - Improved logging and monitoring capabilities\n# - Additional configuration options and optimizations\n\n\"\"\"HTTP Server Engine Adapter for SGLang.\n\nThis module provides HTTP-based adapters for SGLang engines, allowing communication\nwith SGLang servers through HTTP requests instead of direct engine calls.\n\nClasses:\n    HttpServerAdapter: Synchronous HTTP adapter for SGLang engines\n    AsyncHttpServerAdapter: Asynchronous HTTP adapter for SGLang engines\n\nFunctions:\n    launch_server_process: Launch and initialize an SGLang HTTP server process\n\"\"\"\n\nimport asyncio\nimport logging\nimport multiprocessing\nimport os\nimport time\nfrom contextlib import asynccontextmanager\nfrom typing import Any, Callable, Optional\n\nimport aiohttp\nimport requests\nfrom sglang.srt.entrypoints.EngineBase import EngineBase\nfrom sglang.srt.entrypoints.http_server import launch_server\nfrom sglang.srt.managers.io_struct import (\n    UpdateWeightsFromTensorReqInput,\n)\nfrom sglang.srt.server_args import ServerArgs\nfrom sglang.srt.utils import kill_process_tree\n\n# Configure logger\nlogger = logging.getLogger(__name__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n# Default configuration constants\nDEFAULT_TIMEOUT = 60.0\nDEFAULT_MAX_ATTEMPTS = 3\nDEFAULT_RETRY_DELAY = 2.0\nDEFAULT_MAX_CONNECTIONS = 2000\nDEFAULT_MAX_WAIT_TIME = 300.0\n\n\ndef _read_response(response: requests.Response):\n    if response.status_code == 204 or not response.content:\n        return {}\n    try:\n        return response.json()\n    except ValueError:\n        return {\n            \"content_type\": response.headers.get(\"Content-Type\", \"\"),\n            \"text\": response.text,\n        }\n\n\nasync def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]:\n    if resp.status == 204 or (resp.content_length == 0):\n        return {}\n\n    try:\n        return await resp.json(content_type=None)\n    except Exception:\n        try:\n            text = await resp.text()\n        except Exception:\n            return {}\n        return {\n            \"content_type\": (resp.headers.get(\"Content-Type\") or \"\"),\n            \"text\": text,\n        }\n\n\ndef launch_server_process(\n    server_args: ServerArgs,\n    timeout: float = DEFAULT_TIMEOUT,\n    max_wait_time=DEFAULT_MAX_WAIT_TIME,\n    first_rank_in_node=False,\n) -> multiprocessing.Process:\n    \"\"\"Launch an SGLang HTTP server process and wait for it to be ready.\n\n    This function starts a new process running an SGLang HTTP server, then waits\n    for the server to become ready by polling its health endpoints. It ensures\n    the server is fully operational before returning.\n\n    Args:\n        server_args (ServerArgs): Server configuration arguments including host, port, and other settings\n        timeout (float, optional): Timeout for individual HTTP requests during health checks.\n            Defaults to DEFAULT_TIMEOUT.\n\n    Returns:\n        multiprocessing.Process: The launched multiprocessing.Process instance\n\n    Raises:\n        RuntimeError: If the server process terminates unexpectedly during startup or cache flush\n        TimeoutError: If server fails to become ready within reasonable time (300 seconds)\n        requests.RequestException: If health check requests fail repeatedly\n\n    Note:\n        This function will return immediately for non-master nodes (node_rank != 0),\n        but the process will still be started and returned.\n        This is for consistency; except for the process obtained by node_rank = 0,\n        other processes have no actual effect.\n    \"\"\"\n    p = multiprocessing.Process(target=launch_server, args=(server_args,))\n    if server_args.node_rank != 0 or not first_rank_in_node:\n        logger.info(f\"Server process started with PID {p.pid} for node rank {server_args.node_rank}\", flush=True)\n        return p\n\n    p.start()\n\n    base_url = server_args.url()\n    headers = {\n        \"Content-Type\": \"application/json; charset=utf-8\",\n        \"Authorization\": f\"Bearer {server_args.api_key}\",\n    }\n\n    # Health check with overall timeout\n    start_time = time.time()\n\n    with requests.Session() as session:\n        while time.time() - start_time < max_wait_time:\n            if not p.is_alive():\n                raise RuntimeError(\"Server process terminated unexpectedly during startup\")\n\n            try:\n                if server_args.is_embedding:\n                    response = session.get(f\"{base_url}/health\", headers=headers, timeout=timeout)\n                else:\n                    response = session.get(f\"{base_url}/health_generate\", headers=headers, timeout=timeout)\n                if response.status_code == 200:\n                    break\n            except requests.RequestException as e:\n                logger.debug(f\"Health check failed: {e}\")\n\n            time.sleep(2)\n        else:\n            p.terminate()\n            logger.error(f\"Server in {base_url} failed to become healthy within timeout period\")\n            raise TimeoutError(\"Server failed to become healthy within timeout period\")\n\n        # Ensure cache is ready\n        while time.time() - start_time < max_wait_time:\n            if not p.is_alive():\n                raise RuntimeError(\"Server process terminated unexpectedly during cache flush\")\n\n            try:\n                response = session.get(f\"{base_url}/flush_cache\", headers=headers, timeout=timeout)\n                if response.status_code == 200:\n                    break\n            except requests.RequestException as e:\n                logger.debug(f\"Cache flush check failed: {e}\")\n\n            time.sleep(2)\n        else:\n            p.terminate()\n            raise TimeoutError(\"Server cache flush failed within timeout period\")\n\n    return p\n\n\nclass HttpServerAdapter(EngineBase):\n    \"\"\"HTTP-based adapter for SGLang engines.\n\n    This adapter allows interaction with SGLang engines through HTTP requests\n    instead of direct engine calls. It launches an HTTP server process and\n    provides methods to communicate with it via REST API calls.\n\n    You can use this class to launch a server from a HttpServerAdapter instance.\n    We recommend using this class only when you need to use http server.\n    Otherwise, you can use Engine directly.\n\n    Attributes:\n        router_ip (Optional[str]): IP address of the router for worker registration\n        router_port (Optional[int]): Port of the router for worker registration\n        server_args (ServerArgs): Server configuration arguments\n        node_rank (int): Rank of this node in distributed setup\n        process (multiprocessing.Process): The launched server process\n        timeout (float): HTTP request timeout in seconds\n        max_attempts (int): Maximum number of attempts for requests\n        retry_delay (float): Base delay between retries in seconds\n    \"\"\"\n\n    def __init__(\n        self,\n        router_ip: Optional[str] = None,\n        router_port: Optional[int] = None,\n        timeout: float = DEFAULT_TIMEOUT,\n        max_attempts: int = DEFAULT_MAX_ATTEMPTS,\n        retry_delay: float = DEFAULT_RETRY_DELAY,\n        first_rank_in_node: bool = False,\n        max_start_wait_time: float = DEFAULT_MAX_WAIT_TIME,\n        launch_server: bool = True,\n        **kwargs: Any,\n    ) -> None:\n        \"\"\"Initialize the HTTP server engine adapter.\n\n        Args:\n            router_ip (Optional[str], optional): IP address of router for worker registration.\n                Defaults to None.\n            router_port (Optional[int], optional): Port of router for worker registration.\n                Defaults to None.\n            timeout (float, optional): HTTP request timeout in seconds.\n                Defaults to DEFAULT_TIMEOUT.\n            max_attempts (int, optional): Maximum number of retry attempts for failed requests.\n                Defaults to DEFAULT_MAX_ATTEMPTS.\n            retry_delay (float, optional): Base delay between retries in seconds.\n                Defaults to DEFAULT_RETRY_DELAY.\n            launch_server (bool, optional): Whether to launch the server process.\n                Defaults to True.\n            **kwargs (Any): Additional arguments passed to ServerArgs\n\n        Note:\n            TODO: @ChangyiYang Enable SGLang router for this http server engine\n            If both router_ip and router_port are provided and this is the master node\n            (node_rank == 0), the adapter will automatically register with the router.\n        \"\"\"\n        self.router_ip: Optional[str] = router_ip\n        self.router_port: Optional[int] = router_port\n        self.timeout: float = timeout\n        self.max_attempts: int = max_attempts\n        self.retry_delay: float = retry_delay\n        self.server_args: ServerArgs = ServerArgs(**kwargs)\n        self.node_rank: int = self.server_args.node_rank\n        self.max_start_wait_time: float = max_start_wait_time\n\n        logger.info(\n            f\"Launch HttpServerAdapter at: {self.server_args.host}:{self.server_args.port} with {first_rank_in_node}\"\n        )\n        if launch_server:\n            self.process: multiprocessing.Process = launch_server_process(\n                self.server_args, self.timeout, self.max_start_wait_time, first_rank_in_node\n            )\n\n        if self.node_rank == 0 and self.router_ip and self.router_port:\n            self._register_with_router()\n\n    def _register_with_router(self) -> None:\n        \"\"\"Register worker with router with error handling.\n\n        This method attempts to register the current worker with a router service.\n        If registration fails, it logs an error but does not raise an exception,\n        allowing the server to continue operating without router integration.\n\n        Raises:\n            Does not raise exceptions - all errors are logged and handled gracefully.\n        \"\"\"\n        try:\n            url = f\"http://{self.router_ip}:{self.router_port}/add_worker\"\n            params = {\"url\": f\"http://{self.server_args.host}:{self.server_args.port}\"}\n            response = requests.post(url, params=params, timeout=self.timeout)\n            response.raise_for_status()\n            logger.info(\"Successfully registered with router\")\n        except Exception as e:\n            logger.error(f\"Failed to register with router: {e}\")\n            # Don't raise here - server can still work without router\n\n    def _make_request(\n        self,\n        endpoint: str,\n        payload: Optional[dict[str, Any]] = None,\n        method: str = \"POST\",\n        timeout: float = DEFAULT_TIMEOUT,\n        only_master: bool = True,\n    ) -> dict[str, Any]:\n        \"\"\"Make a HTTP request with retry logic and consistent error handling.\n\n        Args:\n            endpoint (str): The API endpoint to call (without leading slash)\n            payload (Optional[Dict[str, Any]], optional): The JSON payload to send.\n                Defaults to empty dict if None.\n            method (str, optional): HTTP method to use. Defaults to \"POST\".\n\n        Returns:\n            Dict[str, Any]: The JSON response from the server\n\n        Raises:\n            requests.HTTPError: If the HTTP request fails with a client/server error\n            RuntimeError: If all retry attempts are exhausted\n\n        Note:\n            - For non-master nodes (node_rank != 0), returns empty dict immediately\n            - Uses exponential backoff for retries\n            - Logs warnings for timeout and connection errors, errors for HTTP errors\n        \"\"\"\n        if only_master and self.node_rank != 0:\n            return {}\n\n        url = f\"http://{self.server_args.host}:{self.server_args.port}/{endpoint}\"\n\n        for attempt in range(self.max_attempts):\n            try:\n                if method.upper() == \"GET\":\n                    response = requests.get(url, timeout=self.timeout)\n                else:\n                    response = requests.post(url, json=payload or {}, timeout=self.timeout)\n\n                response.raise_for_status()\n                return _read_response(response)\n\n            except requests.exceptions.Timeout:\n                logger.warning(f\"Request to {endpoint} timed out (attempt {attempt + 1})\")\n            except requests.exceptions.ConnectionError:\n                logger.warning(f\"Connection error for {endpoint} (attempt {attempt + 1})\")\n            except requests.exceptions.HTTPError as e:\n                logger.error(f\"HTTP error for {endpoint}: {e}\")\n                raise\n            except Exception as e:\n                logger.error(f\"Unexpected error for {endpoint}: {e}\")\n                if attempt == self.max_attempts - 1:\n                    raise\n\n            if attempt < self.max_attempts - 1:\n                time.sleep(self.retry_delay * (2**attempt))\n\n        raise RuntimeError(f\"Failed to complete request to {endpoint} after {self.max_attempts} attempts\")\n\n    def update_weights_from_tensor(self, req: UpdateWeightsFromTensorReqInput) -> dict[str, Any]:\n        \"\"\"Update model weights from tensor data.\n\n        The HTTP server will only post meta data, and the real weights will be\n        copied directly from GPUs.\n\n        Args:\n            serialized_named_tensors (List[str]): List of serialized tensor data\n            load_format (Optional[str], optional): Format specification for loading weights.\n                Defaults to None.\n            flush_cache (bool, optional): Whether to flush cache after updating weights.\n                Defaults to False.\n\n        Returns:\n            Dict[str, Any]: Server response containing update status\n\n        Note:\n            The model should be on GPUs rather than CPU for this functionality to work properly.\n            If you encounter issues, ensure your model is loaded on GPU devices rather than CPU.\n        \"\"\"\n        import base64\n\n        named_tensors = req.serialized_named_tensors\n        load_format = req.load_format\n        flush_cache = req.flush_cache\n\n        if named_tensors:\n            serialized_named_tensors = [\n                base64.b64encode(named_tensor).decode(\"utf-8\") for named_tensor in named_tensors\n            ]\n        else:\n            serialized_named_tensors = []\n\n        return self._make_request(\n            \"update_weights_from_tensor\",\n            {\n                \"serialized_named_tensors\": serialized_named_tensors,\n                \"load_format\": load_format,\n                \"flush_cache\": flush_cache,\n            },\n        )\n\n    def shutdown(self) -> None:\n        \"\"\"Shutdown the HTTP server and clean up resources.\n\n        This method performs the following cleanup operations:\n        1. Unregisters the worker from the router (if configured)\n        2. Terminates the server process tree\n\n        All operations are performed with error handling to ensure graceful shutdown\n        even if individual steps fail.\n\n        Note:\n            This method should be called when the adapter is no longer needed\n            to ensure proper cleanup of resources and processes.\n        \"\"\"\n        # Unregister from router\n        if self.router_ip and self.router_port:\n            try:\n                url = f\"http://{self.router_ip}:{self.router_port}/remove_worker\"\n                params = {\"url\": f\"http://{self.server_args.host}:{self.server_args.port}\"}\n                requests.post(url, params=params, timeout=5.0)  # Short timeout for shutdown\n                logger.info(\"Successfully unregistered from router\")\n            except Exception as e:\n                logger.warning(f\"Failed to unregister from router: {e}\")\n\n        # Kill server process\n        if hasattr(self, \"process\") and self.process is not None:\n            try:\n                kill_process_tree(self.process.pid)\n                logger.info(\"Server process terminated\")\n            except Exception as e:\n                logger.error(f\"Failed to terminate server process: {e}\")\n\n    def generate(\n        self,\n        prompt: Optional[str] = None,\n        sampling_params: Optional[dict[str, Any]] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        return_logprob: bool = False,\n        logprob_start_len: Optional[int] = None,\n        top_logprobs_num: Optional[int] = None,\n        token_ids_logprob: Optional[list[int]] = None,\n        lora_path: Optional[str] = None,\n        custom_logit_processor: Optional[Callable] = None,\n    ) -> dict[str, Any]:\n        \"\"\"Generate text using the SGLang server.\n\n        Args:\n            prompt (Optional[str], optional): Text prompt for generation. Defaults to None.\n            sampling_params (Optional[Dict[str, Any]], optional): Parameters controlling\n                text generation sampling. Defaults to None.\n            input_ids (Optional[List[int]], optional): Alternative to prompt, direct token IDs input.\n                Defaults to None.\n            image_data (Optional[Any], optional): Image data for multimodal generation.\n                Defaults to None.\n            return_logprob (bool, optional): Whether to return log probabilities.\n                Defaults to False.\n            logprob_start_len (Optional[int], optional): Starting length for log probability calculation.\n                Defaults to None.\n            top_logprobs_num (Optional[int], optional): Number of top log probabilities to return.\n                Defaults to None.\n            token_ids_logprob (Optional[List[int]], optional): Specific token IDs for\n                log probability calculation. Defaults to None.\n            lora_path (Optional[str], optional): Path to LoRA adapter weights. Defaults to None.\n            custom_logit_processor (Optional[Callable], optional): Custom logit processing function.\n                Defaults to None.\n\n        Returns:\n            Dict[str, Any]: Generated text and associated metadata from the server\n\n        Note:\n            Either prompt or input_ids should be provided, but not both.\n            The response format depends on the server configuration and parameters.\n        \"\"\"\n        payload = {\n            \"text\": prompt,\n            \"sampling_params\": sampling_params,\n            \"input_ids\": input_ids,\n            \"image_data\": image_data,\n            \"return_logprob\": return_logprob,\n            \"logprob_start_len\": logprob_start_len,\n            \"top_logprobs_num\": top_logprobs_num,\n            \"token_ids_logprob\": token_ids_logprob,\n            \"lora_path\": lora_path,\n            \"custom_logit_processor\": custom_logit_processor,\n        }\n        # Filter out None values\n        payload = {k: v for k, v in payload.items() if v is not None}\n\n        return self._make_request(\"generate\", payload, only_master=False)\n\n    def reward_score(\n        self,\n        prompt: Optional[str] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        lora_path: Optional[str] = None,\n    ) -> dict[str, Any]:\n        assert self.server_args.is_embedding, \"Score is only supported for embedding models\"\n        payload = {\n            \"text\": prompt,\n            \"input_ids\": input_ids,\n            \"image_data\": image_data,\n            \"lora_path\": lora_path,\n        }\n        # Filter out None values\n        payload = {k: v for k, v in payload.items() if v is not None}\n\n        return self._make_request(\"classify\", payload, only_master=False)\n\n    def flush_cache(self) -> dict[str, Any]:\n        \"\"\"Flush the cache of the server.\n\n        This method repeatedly attempts to flush the server cache until successful.\n        The flush operation will not return status 200 when there are pending requests.\n\n        Returns:\n            Dict[str, Any]: Server response indicating cache flush status.\n                For non-master nodes, returns empty dict.\n\n        Note:\n            Uses retry logic with limited attempts (max_attempts * 2) to avoid infinite loops.\n            Each retry includes a delay to allow pending requests to complete.\n        \"\"\"\n        if self.node_rank != 0:\n            return {}\n\n        # Use retry logic with limited attempts to avoid infinite loops\n        for attempt in range(self.max_attempts * 2):  # Allow more retries for cache flush\n            try:\n                response = requests.get(\n                    f\"http://{self.server_args.host}:{self.server_args.port}/flush_cache\", timeout=self.timeout\n                )\n                if response.status_code == 200:\n                    return _read_response(response)\n            except Exception as e:\n                logger.warning(f\"Error flushing cache (attempt {attempt + 1}): {e}\")\n\n            time.sleep(self.retry_delay)\n\n        logger.error(\"Failed to flush cache after maximum attempts\")\n        return {}\n\n    def release_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]:\n        \"\"\"Release GPU memory occupation temporarily.\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to release.\n                If None, releases all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory release status\n        \"\"\"\n        return self._make_request(\"release_memory_occupation\", {\"tags\": tags})\n\n    def resume_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]:\n        \"\"\"Resume GPU memory occupation.\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to resume.\n                If None, resumes all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory resume status\n        \"\"\"\n        return self._make_request(\"resume_memory_occupation\", {\"tags\": tags})\n\n    def abort_request(self, rid: str = \"\", abort_all: bool = False) -> dict[str, Any]:\n        \"\"\"Abort a request.\n\n        Args:\n            rid (str): The ID of the request to abort\n            abort_all (bool, optional): Whether to abort all requests. Defaults to False.\n\n        Returns:\n            Dict[str, Any]: Server response indicating abort status\n        \"\"\"\n        return self._make_request(\"abort_request\", {\"rid\": rid, \"abort_all\": abort_all})\n\n\nclass AsyncHttpServerAdapter(HttpServerAdapter):\n    \"\"\"Asynchronous HTTP-based adapter for SGLang engines.\n\n    This class inherits from HttpServerAdapter and adds async capabilities\n    for non-blocking HTTP requests to the SGLang server. It provides the same\n    functionality as the synchronous version but with async/await support.\n\n    The async adapter is useful when you need to make multiple concurrent requests\n    or integrate with async frameworks. It uses aiohttp for efficient async HTTP\n    communication and maintains connection pooling for better performance.\n\n    Attributes:\n        max_connections (int): Maximum number of connections in the connection pool\n    \"\"\"\n\n    def __init__(\n        self,\n        router_ip: Optional[str] = None,\n        router_port: Optional[int] = None,\n        timeout: float = DEFAULT_TIMEOUT,\n        max_attempts: int = DEFAULT_MAX_ATTEMPTS,\n        retry_delay: float = DEFAULT_RETRY_DELAY,\n        max_connections: int = DEFAULT_MAX_CONNECTIONS,\n        first_rank_in_node: bool = False,\n        launch_server: bool = True,\n        **kwargs: Any,\n    ) -> None:\n        \"\"\"Initialize the async HTTP server engine adapter.\n\n        Args:\n            router_ip (Optional[str], optional): IP address of router for worker registration.\n                Defaults to None.\n            router_port (Optional[int], optional): Port of router for worker registration.\n                Defaults to None.\n            timeout (float, optional): HTTP request timeout in seconds.\n                Defaults to DEFAULT_TIMEOUT.\n            max_attempts (int, optional): Maximum number of retry attempts for failed requests.\n                Defaults to DEFAULT_MAX_ATTEMPTS.\n            retry_delay (float, optional): Base delay between retries in seconds.\n                Defaults to DEFAULT_RETRY_DELAY.\n            max_connections (int, optional): Maximum number of connections in the connection pool.\n                Defaults to DEFAULT_MAX_CONNECTIONS.\n            launch_server (bool, optional): Whether to launch the server process.\n                Defaults to True.\n            **kwargs (Any): Additional arguments passed to ServerArgs\n        \"\"\"\n        super().__init__(\n            router_ip,\n            router_port,\n            timeout,\n            max_attempts,\n            retry_delay,\n            first_rank_in_node,\n            launch_server=launch_server,\n            **kwargs,\n        )\n        self.max_connections: int = max_connections\n\n    @asynccontextmanager\n    async def _get_session(self) -> aiohttp.ClientSession:\n        \"\"\"Context manager for safe session access with proper connection pooling.\n\n        Yields:\n            aiohttp.ClientSession: Session instance for making HTTP requests\n\n        Note:\n            This method creates a new session for each request to avoid resource competition\n            while still maintaining proper connection pooling through the shared connector.\n        \"\"\"\n        # Create a new session for each request to avoid resource competition\n        connector = aiohttp.TCPConnector(\n            limit=self.max_connections,\n            limit_per_host=self.max_connections // 4,\n            ttl_dns_cache=300,\n            use_dns_cache=True,\n        )\n        timeout = aiohttp.ClientTimeout(total=self.timeout)\n        session = aiohttp.ClientSession(connector=connector, timeout=timeout)\n\n        try:\n            yield session\n        finally:\n            # Always close the session to free up resources\n            if not session.closed:\n                await session.close()\n\n    async def _make_async_request(\n        self,\n        endpoint: str,\n        payload: Optional[dict[str, Any]] = None,\n        method: str = \"POST\",\n        timeout: float = DEFAULT_TIMEOUT,\n        only_master: bool = True,\n    ) -> dict[str, Any]:\n        \"\"\"Make an async HTTP request with retry logic and consistent error handling.\n\n        Args:\n            endpoint (str): The API endpoint to call (without leading slash)\n            payload (Optional[Dict[str, Any]], optional): The JSON payload to send.\n                Defaults to empty dict if None.\n            method (str, optional): HTTP method to use. Defaults to \"POST\".\n\n        Returns:\n            Dict[str, Any]: The JSON response from the server\n\n        Raises:\n            aiohttp.ClientResponseError: If the HTTP request fails with a client/server error\n            RuntimeError: If all retry attempts are exhausted\n\n        Note:\n            - For non-master nodes (node_rank != 0), returns empty dict immediately\n            - Uses exponential backoff for retries\n            - Logs warnings for timeout and connection errors, errors for HTTP errors\n        \"\"\"\n        if only_master and self.node_rank != 0:\n            return {}\n\n        url = f\"http://{self.server_args.host}:{self.server_args.port}/{endpoint}\"\n\n        for attempt in range(self.max_attempts):\n            try:\n                async with self._get_session() as session:\n                    if method.upper() == \"GET\":\n                        async with session.get(url, timeout=timeout) as response:\n                            response.raise_for_status()\n                            return await _read_async_response(response)\n                    else:\n                        async with session.post(url, json=payload or {}, timeout=timeout) as response:\n                            response.raise_for_status()\n                            return await _read_async_response(response)\n\n            except asyncio.TimeoutError:\n                logger.warning(f\"Async request to {endpoint} timed out (attempt {attempt + 1})\")\n            except aiohttp.ClientConnectorError:\n                logger.warning(f\"Connection error for {endpoint} (attempt {attempt + 1})\")\n            except aiohttp.ClientResponseError as e:\n                logger.error(f\"HTTP error for {endpoint}: {e}\")\n                raise\n            except Exception as e:\n                logger.error(f\"Unexpected error for {endpoint}: {e}\")\n                if attempt == self.max_attempts - 1:\n                    raise\n\n            if attempt < self.max_attempts - 1:\n                await asyncio.sleep(self.retry_delay * (2**attempt))\n\n        raise RuntimeError(f\"Failed to complete async request to {endpoint} after {self.max_attempts} attempts\")\n\n    async def release_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]:\n        \"\"\"Release GPU memory occupation temporarily (async version).\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to release.\n                If None, releases all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory release status\n        \"\"\"\n        return await self._make_async_request(\"release_memory_occupation\", {\"tags\": tags})\n\n    async def resume_memory_occupation(self, tags: Optional[list[str]] = None) -> dict[str, Any]:\n        \"\"\"Resume GPU memory occupation (async version).\n\n        Similar to AsyncEngine, this method handles first-time weight reloading\n        by calling release_memory_occupation if needed.\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to resume.\n                If None, resumes all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory resume status\n        \"\"\"\n        return await self._make_async_request(\"resume_memory_occupation\", {\"tags\": tags})\n\n    async def update_weights_from_tensor(\n        self,\n        req: UpdateWeightsFromTensorReqInput,\n    ) -> dict[str, Any]:\n        \"\"\"Update model weights from tensor data asynchronously.\n\n        Args:\n            serialized_named_tensors (List[str]): List of serialized tensor data\n            load_format (Optional[str], optional): Format specification for loading weights.\n                Defaults to None.\n            flush_cache (bool, optional): Whether to flush cache after updating weights.\n                Defaults to True.\n\n        Returns:\n            Dict[str, Any]: Server response containing update status\n        \"\"\"\n        import base64\n\n        named_tensors = req.serialized_named_tensors\n        load_format = req.load_format\n        flush_cache = req.flush_cache\n\n        serialized_named_tensors = [base64.b64encode(named_tensor).decode(\"utf-8\") for named_tensor in named_tensors]\n        return await self._make_async_request(\n            \"update_weights_from_tensor\",\n            {\n                \"serialized_named_tensors\": serialized_named_tensors,\n                \"load_format\": load_format,\n                \"flush_cache\": flush_cache,\n            },\n        )\n\n    async def flush_cache(self) -> dict[str, Any]:\n        \"\"\"Flush the cache of the server asynchronously.\n\n        Similar to the sync version, this method retries until the cache\n        is successfully flushed. It uses async sleep between retries.\n\n        Returns:\n            Dict[str, Any]: Server response indicating cache flush status.\n                For non-master nodes, returns empty dict.\n\n        Note:\n            Uses retry logic with limited attempts (max_attempts * 4) to avoid infinite loops.\n            Each retry includes an async delay to allow pending requests to complete.\n        \"\"\"\n        if self.node_rank != 0:\n            return {}\n\n        # Use retry logic with limited attempts to avoid infinite loops\n        for attempt in range(self.max_attempts * 4):  # Allow more retries for cache flush\n            try:\n                async with self._get_session() as session:\n                    url = f\"http://{self.server_args.host}:{self.server_args.port}/flush_cache\"\n                    async with session.get(url) as response:\n                        if response.status == 200:\n                            return await _read_async_response(response)\n            except Exception as e:\n                logger.warning(f\"Error flushing cache (attempt {attempt + 1}): {e}\")\n\n            await asyncio.sleep(self.retry_delay)\n\n        logger.error(\"Failed to flush cache after maximum attempts\")\n        return {}\n\n    async def generate(\n        self,\n        prompt: Optional[str] = None,\n        sampling_params: Optional[dict[str, Any]] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        return_logprob: bool = False,\n        logprob_start_len: Optional[int] = None,\n        top_logprobs_num: Optional[int] = None,\n        token_ids_logprob: Optional[list[int]] = None,\n        lora_path: Optional[str] = None,\n        custom_logit_processor: Optional[Callable] = None,\n    ) -> dict[str, Any]:\n        \"\"\"Generate text using the SGLang server asynchronously.\"\"\"\n        logger.info(\"generate() started\")\n\n        payload = {\n            \"text\": prompt,\n            \"sampling_params\": sampling_params,\n            \"input_ids\": input_ids,\n            \"image_data\": image_data,\n            \"return_logprob\": return_logprob,\n            \"logprob_start_len\": logprob_start_len,\n            \"top_logprobs_num\": top_logprobs_num,\n            \"token_ids_logprob\": token_ids_logprob,\n            \"lora_path\": lora_path,\n            \"custom_logit_processor\": custom_logit_processor,\n        }\n\n        # Filter out None values\n        payload = {k: v for k, v in payload.items() if v is not None}\n\n        # Send request\n        response = await self._make_async_request(\"generate\", payload, timeout=self.timeout, only_master=False)\n\n        return response\n\n    async def async_generate(\n        self,\n        prompt: Optional[str] = None,\n        sampling_params: Optional[dict[str, Any]] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        return_logprob: bool = False,\n        logprob_start_len: Optional[int] = None,\n        top_logprobs_num: Optional[int] = None,\n        token_ids_logprob: Optional[list[int]] = None,\n        lora_path: Optional[str] = None,\n        custom_logit_processor: Optional[Callable] = None,\n    ) -> dict[str, Any]:\n        \"\"\"Async generate method that mirrors AsyncEngine.async_generate interface.\n\n        This method provides compatibility with AsyncEngine's async_generate method\n        by forwarding the call to the generate method. It ensures API consistency\n        between direct engine usage and HTTP-based engine usage.\n\n        Args:\n            prompt (Optional[str], optional): Text prompt for generation. Defaults to None.\n            sampling_params (Optional[Dict[str, Any]], optional): Parameters controlling\n                text generation sampling. Defaults to None.\n            input_ids (Optional[List[int]], optional): Alternative to prompt, direct token IDs input.\n                Defaults to None.\n            image_data (Optional[Any], optional): Image data for multimodal generation.\n                Defaults to None.\n            return_logprob (bool, optional): Whether to return log probabilities.\n                Defaults to False.\n            logprob_start_len (Optional[int], optional): Starting length for log probability calculation.\n                Defaults to None.\n            top_logprobs_num (Optional[int], optional): Number of top log probabilities to return.\n                Defaults to None.\n            token_ids_logprob (Optional[List[int]], optional): Specific token IDs for\n                log probability calculation. Defaults to None.\n            lora_path (Optional[str], optional): Path to LoRA adapter weights. Defaults to None.\n            custom_logit_processor (Optional[Callable], optional): Custom logit processing function.\n                Defaults to None.\n\n        Returns:\n            Dict[str, Any]: Generated text and associated metadata from the server\n\n        Note:\n            This method is provided for API compatibility with AsyncEngine.\n            It forwards all calls to the generate method.\n        \"\"\"\n        return await self.generate(\n            prompt=prompt,\n            sampling_params=sampling_params,\n            input_ids=input_ids,\n            image_data=image_data,\n            return_logprob=return_logprob,\n            logprob_start_len=logprob_start_len,\n            top_logprobs_num=top_logprobs_num,\n            token_ids_logprob=token_ids_logprob,\n            lora_path=lora_path,\n            custom_logit_processor=custom_logit_processor,\n        )\n\n    async def reward_score(\n        self,\n        prompt: Optional[str] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        lora_path: Optional[str] = None,\n    ) -> dict[str, Any]:\n        logger.info(\"reward_score() started\")\n        payload = {\n            \"text\": prompt,\n            \"input_ids\": input_ids,\n            \"image_data\": image_data,\n            \"lora_path\": lora_path,\n        }\n        # Filter out None values\n        payload = {k: v for k, v in payload.items() if v is not None}\n\n        # Send request\n        response = await self._make_async_request(\"classify\", payload, timeout=self.timeout, only_master=False)\n\n        return response\n\n    async def async_reward_score(\n        self,\n        prompt: Optional[str] = None,\n        input_ids: Optional[list[int]] = None,\n        image_data: Optional[Any] = None,\n        lora_path: Optional[str] = None,\n    ) -> dict[str, Any]:\n        return await self.reward_score(\n            prompt=prompt,\n            input_ids=input_ids,\n            image_data=image_data,\n            lora_path=lora_path,\n        )\n\n    async def abort_request(self, rid: str = \"\", abort_all: bool = False) -> dict[str, Any]:\n        \"\"\"Abort a request asynchronously.\n\n        Args:\n            rid (str): The ID of the request to abort\n            abort_all (bool, optional): Whether to abort all requests. Defaults to False.\n\n        Returns:\n            Dict[str, Any]: Server response indicating abort status\n        \"\"\"\n        return await self._make_async_request(\"abort_request\", {\"rid\": rid, \"abort_all\": abort_all})\n"
  },
  {
    "path": "verl/workers/rollout/sglang_rollout/sglang_rollout.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport logging\nimport multiprocessing as mp\nimport os\nfrom typing import Generator\n\nimport ray\nimport sglang.srt.entrypoints.engine\nimport torch\nfrom sglang.srt.server_args import ServerArgs\nfrom sglang.srt.utils import (\n    assert_pkg_version,\n    is_cuda,\n    set_prometheus_multiproc_dir,\n    set_ulimit,\n)\nfrom sglang.srt.weight_sync.utils import update_weights as sgl_update_weights\nfrom torch.distributed.device_mesh import DeviceMesh, init_device_mesh\n\nfrom verl.utils.net_utils import is_valid_ipv6_address\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.base import BaseRollout\nfrom verl.workers.rollout.sglang_rollout.http_server_engine import AsyncHttpServerAdapter\nfrom verl.workers.rollout.sglang_rollout.utils import get_named_tensor_buckets\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n\n# patch to avoid issue https://github.com/sgl-project/sglang/issues/6723\ndef _set_envs_and_config(server_args: ServerArgs):\n    # Set global environments\n    os.environ[\"TF_CPP_MIN_LOG_LEVEL\"] = \"3\"\n    os.environ[\"NCCL_CUMEM_ENABLE\"] = \"0\"\n    os.environ[\"NCCL_NVLS_ENABLE\"] = str(int(server_args.enable_nccl_nvls))\n    os.environ[\"TORCH_NCCL_AVOID_RECORD_STREAMS\"] = \"1\"\n    os.environ[\"CUDA_DEVICE_MAX_CONNECTIONS\"] = \"4\"\n    os.environ[\"CUDA_MODULE_LOADING\"] = \"AUTO\"\n    # Enable faulthandler in subprocesses\n    os.environ[\"PYTHONFAULTHANDLER\"] = \"1\"\n\n    # Set prometheus env vars\n    if server_args.enable_metrics:\n        set_prometheus_multiproc_dir()\n\n    # Set ulimit\n    set_ulimit()\n\n    # Check flashinfer version\n    if server_args.attention_backend == \"flashinfer\":\n        assert_pkg_version(\n            \"flashinfer_python\",\n            \"0.2.5\",\n            \"Please uninstall the old version and reinstall the latest version by following the instructions at https://docs.flashinfer.ai/installation.html.\",\n        )\n    if is_cuda():\n        assert_pkg_version(\n            \"sgl-kernel\",\n            \"0.1.1\",\n            \"Please reinstall the latest version with `pip install sgl-kernel --force-reinstall`\",\n        )\n\n    # Set mp start method\n    mp.set_start_method(\"spawn\", force=True)\n\n\nsglang.srt.entrypoints.engine._set_envs_and_config = _set_envs_and_config\n\n\n# because chatCompletion is an async method, it makes the whole ray actor be an async actor\n# which can not call loop.run_until_complete. So we need to make the engine to be an async class\nclass ServerAdapter(BaseRollout):\n    \"\"\"SGLang server adapter used in native http server mode, serve as http client to request SGLang server\n    to resume/release/update weights and kv_cache.\n\n    - hybrid mode: reside in each hybrid worker to sync weights between training engine and SGLang server.\n    - standalone/colocated mode: just a dummy placeholder to occupy the GPU to prevent ray scheduling new GPU actor.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        device_mesh: DeviceMesh,\n        replica_rank: int = -1,\n    ):\n        if config.get(\"quantization\", None) == \"fp8\":\n            import sglang\n            from packaging import version\n\n            assert version.parse(sglang.__version__) >= version.parse(\"0.5.5\"), (\n                \"sglang>=0.5.5 is required for FP8 quantization\"\n            )\n            FP8_BLOCK_QUANT_KWARGS = {\n                \"activation_scheme\": \"dynamic\",\n                \"fmt\": \"e4m3\",\n                \"quant_method\": \"fp8\",\n                \"weight_block_size\": [128, 128],\n            }\n            fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS)\n            model_config.hf_config.quantization_config = fp8_block_quant_kwargs\n        super().__init__(config, model_config, device_mesh)\n        self._engine: AsyncHttpServerAdapter = None\n\n        rank = int(os.environ[\"RANK\"])\n        local_world_size = int(os.environ[\"RAY_LOCAL_WORLD_SIZE\"])\n        rollout_world_size = self.config.tensor_model_parallel_size * self.config.data_parallel_size\n        if replica_rank == -1:\n            self.replica_rank = rank // rollout_world_size\n        else:\n            self.replica_rank = replica_rank\n        self.rollout_rank = rank % rollout_world_size\n        self.node_rank = self.rollout_rank // local_world_size\n        self.local_rank = self.rollout_rank % local_world_size\n\n    async def _init_server_adapter(self):\n        if self._engine is not None:\n            return\n\n        # device_mesh is needed to gather cuda ipc handle to update weights\n        if self.device_mesh is None:\n            assert torch.distributed.is_initialized(), \"torch distributed must be initialized\"\n            infer_tp = self.config.tensor_model_parallel_size * self.config.data_parallel_size\n            infer_pp = self.config.pipeline_model_parallel_size\n            infer_world_size = infer_tp * infer_pp\n            dp = torch.distributed.get_world_size() // infer_world_size\n            self.device_mesh = init_device_mesh(\n                \"cpu\", mesh_shape=(dp, infer_tp, infer_pp), mesh_dim_names=[\"dp\", \"infer_tp\", \"infer_pp\"]\n            )\n\n        # Only init http server adapter in tp rank 0\n        if self.device_mesh[\"infer_tp\"].get_local_rank() != 0:\n            return\n\n        # Lazy init http server adapter because http server is launched after hybrid engine.\n        self.server_actor = ray.get_actor(f\"sglang_server_{self.replica_rank}_{self.node_rank}\")\n        server_address, server_port = await self.server_actor.get_server_address.remote()\n        logger.debug(\n            f\"replica_rank={self.replica_rank} node_rank={self.node_rank}, \"\n            f\"server address: {server_address}, port: {server_port}\"\n        )\n        host = f\"[{server_address}]\" if is_valid_ipv6_address(server_address) else server_address\n        self._engine = AsyncHttpServerAdapter(\n            model_path=self.model_config.local_path,\n            host=host,\n            port=server_port,\n            launch_server=False,\n            trust_remote_code=self.model_config.trust_remote_code,\n        )\n\n    async def resume(self, tags: list[str]):\n        \"\"\"Resume rollout weights or kv cache in GPU memory.\n\n        Args:\n            tag: weights or kv_cache.\n        \"\"\"\n        await self._init_server_adapter()\n        if self.device_mesh[\"infer_tp\"].get_local_rank() == 0 and self.config.free_cache_engine:\n            await self._engine.resume_memory_occupation(tags=tags)\n\n    async def release(self):\n        \"\"\"Release weights and kv cache in GPU memory.\"\"\"\n        await self._init_server_adapter()\n        if self.device_mesh[\"infer_tp\"].get_local_rank() == 0 and self.config.free_cache_engine:\n            await self._engine.release_memory_occupation(tags=[\"kv_cache\", \"weights\"])\n\n    async def update_weights(\n        self, weights: Generator[tuple[str, torch.Tensor], None, None], global_steps: int = None, **kwargs\n    ):\n        \"\"\"\n        Update model weights using tensor buckets, similar to THUDM/slime's implementation.\n\n        Notes:\n          - For the best performance of `rebuild_cuda_tensor`, it is recommended to:\n              1. Enable `RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES`.\n              2. Manually set `CUDA_VISIBLE_DEVICES=0,1,2,3,4,5,6,7`\n            when using Tensor Parallelism (TP >= 8).\n          - See reference implementations in SLIME:\n            - Main logic: https://github.com/THUDM/slime/blob/fb7605cc5fb09af0f9369d37f7192f12bddee577/slime/ray/ppo_actor.py#L452\n            - runtime envs: https://github.com/THUDM/slime/blob/fb7605cc5fb09af0f9369d37f7192f12bddee577/slime/ray/ppo_actor.py#L39\n        \"\"\"\n        await self._init_server_adapter()\n\n        update_weights_bucket_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) << 20\n        if self.config.get(\"quantization\", None) == \"fp8\":\n            from verl.utils.sglang.sglang_fp8_utils import SGLangFP8QuantizerHelper\n\n            logger.info(\"Convert bf16 weights to fp8 format before loading\")\n            fp8_quantizer_helper = SGLangFP8QuantizerHelper(self.model_config.hf_config.quantization_config)\n            weights = fp8_quantizer_helper.quant_weights_by_name(\n                weights,\n                dtype=self.model_config.hf_config.dtype,\n            )\n        else:\n            weights = weights\n\n        async for params_batch in get_named_tensor_buckets(weights, update_weights_bucket_bytes):\n            await sgl_update_weights(\n                engine=self._engine,\n                params_batch=params_batch,\n                device_mesh_key=\"infer_tp\",\n                device_mesh=self.device_mesh,\n            )\n\n        if self.device_mesh[\"infer_tp\"].get_local_rank() == 0:\n            await self._engine.flush_cache()\n            if global_steps is not None:\n                await self.server_actor.set_global_steps.remote(global_steps)\n"
  },
  {
    "path": "verl/workers/rollout/sglang_rollout/utils.py",
    "content": "# Copyright 2023-2024 SGLang Team\n# Copyright 2025 ModelBest Inc. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport pickle\nfrom typing import Any, Iterator, Optional\n\nimport numpy as np\nimport torch\nimport torch.distributed as dist\n\nfrom verl.utils.device import get_device_name\nfrom verl.workers.rollout.utils import ensure_async_iterator\n\n\ndef broadcast_pyobj(\n    data: list[Any],\n    rank: int,\n    dist_group: Optional[torch.distributed.ProcessGroup] = None,\n    src: int = 0,\n    force_cpu_device: bool = False,\n):\n    \"\"\"from https://github.com/sgl-project/sglang/blob/844e2f227ab0cce6ef818a719170ce37b9eb1e1b/python/sglang/srt/utils.py#L905\n\n    Broadcast inputs from src rank to all other ranks with torch.dist backend.\n    The `rank` here refer to the source rank on global process group (regardless\n    of dist_group argument).\n    \"\"\"\n    device = torch.device(get_device_name() if not force_cpu_device else \"cpu\")\n\n    if rank == src:\n        if len(data) == 0:\n            tensor_size = torch.tensor([0], dtype=torch.long, device=device)\n            dist.broadcast(tensor_size, src=src, group=dist_group)\n        else:\n            serialized_data = pickle.dumps(data)\n            size = len(serialized_data)\n\n            tensor_data = torch.ByteTensor(np.frombuffer(serialized_data, dtype=np.uint8)).to(device)\n            tensor_size = torch.tensor([size], dtype=torch.long, device=device)\n\n            dist.broadcast(tensor_size, src=src, group=dist_group)\n            dist.broadcast(tensor_data, src=src, group=dist_group)\n        return data\n    else:\n        tensor_size = torch.tensor([0], dtype=torch.long, device=device)\n        dist.broadcast(tensor_size, src=src, group=dist_group)\n        size = tensor_size.item()\n\n        if size == 0:\n            return []\n\n        tensor_data = torch.empty(size, dtype=torch.uint8, device=device)\n        dist.broadcast(tensor_data, src=src, group=dist_group)\n\n        serialized_data = bytes(tensor_data.cpu().numpy())\n        data = pickle.loads(serialized_data)\n        return data\n\n\nasync def get_named_tensor_buckets(\n    iterable: Iterator[tuple[str, torch.Tensor]], bucket_bytes: int\n) -> Iterator[list[tuple[str, torch.Tensor]]]:\n    \"\"\"\n    Group tensors into buckets based on a specified size in megabytes.\n\n    Args:\n        iterable: An iterator of tuples containing tensor names and tensors.\n        bucket_bytes: The maximum size of each bucket in bytes.\n\n    Yields:\n        Lists of tuples, where each tuple contains a tensor name and its corresponding tensor.\n\n    Example:\n        >>> tensors = [('tensor1', torch.randn(1000, 1000)), ('tensor2', torch.randn(2000, 2000))]\n        >>> for bucket in get_named_tensor_buckets(tensors, bucket_size_mb=10):\n        ...     print(bucket)\n        [('tensor1', tensor(...)), ('tensor2', tensor(...))]\n\n    \"\"\"\n    if bucket_bytes <= 0:\n        raise ValueError(f\"bucket_bytes must be greater than 0, got {bucket_bytes}\")\n\n    current_bucket = []\n    current_size = 0\n    async for name, tensor in ensure_async_iterator(iterable):\n        tensor_size = tensor.element_size() * tensor.numel()\n        if current_size + tensor_size > bucket_bytes:\n            if current_bucket:\n                yield current_bucket\n            current_bucket = [(name, tensor.clone())]\n            current_size = tensor_size\n        else:\n            current_bucket.append((name, tensor.clone()))\n            current_size += tensor_size\n\n    if current_bucket:\n        yield current_bucket\n"
  },
  {
    "path": "verl/workers/rollout/tokenizer.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe base tokenizer class, required for any hybrid engine based rollout or inference with vLLM.\n\"\"\"\n\nfrom abc import ABC, abstractmethod\n\nimport numpy as np\nimport torch\n\n__all__ = [\"HybridEngineBaseTokenizer\"]\n\n\nclass HybridEngineBaseTokenizer(ABC):\n    \"\"\"the tokenizer property and function name should align with HF's to meet vllm requirement\"\"\"\n\n    @property\n    @abstractmethod\n    def vocab_size(self):\n        \"\"\"\n        `int`: Size of the base vocabulary (without the added tokens).\n        \"\"\"\n        pass\n\n    @property\n    @abstractmethod\n    def pad_token_id(self):\n        \"\"\"\n        `Optional[int]`: Id of the padding token in the vocabulary. Returns `None` if the token has not been set.\n        \"\"\"\n        pass\n\n    @property\n    @abstractmethod\n    def eos_token_id(self):\n        \"\"\"\n        `Optional[int]`: Id of the end of sentence token in the vocabulary. Returns `None` if the token has not been\n        set.\n        \"\"\"\n        pass\n\n    @property\n    @abstractmethod\n    def all_special_ids(self) -> list[int]:\n        \"\"\"\n        `List[int]`: List the ids of the special tokens(`'<unk>'`, `'<cls>'`, etc.) mapped to class attributes.\n        \"\"\"\n        pass\n\n    @property\n    @abstractmethod\n    def all_special_tokens(self) -> list[str]:\n        \"\"\"\n        `List[str]`: A list of the unique special tokens (`'<unk>'`, `'<cls>'`, ..., etc.).\n\n        Convert tokens of `tokenizers.AddedToken` type to string.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def encode(self, text):\n        \"\"\"\n        Converts a string to a sequence of ids (integer), using the tokenizer and vocabulary.\n\n        Args:\n            text (`str`, `List[str]` or `List[int]`):\n                The first sequence to be encoded. This can be a string, a list of strings (tokenized string using the\n                `tokenize` method) or a list of integers.\n\n            text_pair (`str`, `List[str]` or `List[int]`, *optional*):\n                Optional second sequence to be encoded. This can be a string, a list of strings (tokenized string using\n                the `tokenize` method) or a list of integers.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def decode(\n        self,\n        token_ids: int | list[int] | np.ndarray | torch.Tensor,\n        skip_special_tokens: bool = False,\n        clean_up_tokenization_spaces: bool = None,\n        **kwargs,\n    ) -> str:\n        \"\"\"\n        Converts a sequence of ids in a string, using the tokenizer and vocabulary with options to remove special\n        tokens and clean up tokenization spaces.\n\n        Similar to doing `self.convert_tokens_to_string(self.convert_ids_to_tokens(token_ids))`.\n\n        Args:\n            token_ids (`Union[int, List[int], np.ndarray, torch.Tensor]`):\n                List of tokenized input ids. Can be obtained using the `__call__` method.\n            skip_special_tokens (`bool`, *optional*, defaults to `False`):\n                Whether or not to remove special tokens in the decoding.\n            clean_up_tokenization_spaces (`bool`, *optional*):\n                Whether or not to clean up the tokenization spaces. If `None`, will default to\n                `self.clean_up_tokenization_spaces`.\n            kwargs (additional keyword arguments, *optional*):\n                Will be passed to the underlying model specific decode method.\n\n        Returns:\n            `str`: The decoded sentence.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def convert_ids_to_tokens(self, ids: int | list[int], skip_special_tokens: bool = False) -> str | list[str]:\n        \"\"\"\n        Converts a single index or a sequence of indices in a token or a sequence of tokens, using the vocabulary and\n        added tokens.\n\n        Args:\n            ids (`int` or `List[int]`):\n                The token id (or token ids) to convert to tokens.\n            skip_special_tokens (`bool`, *optional*, defaults to `False`):\n                Whether or not to remove special tokens in the decoding.\n\n        Returns:\n            `str` or `List[str]`: The decoded token(s).\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def get_added_vocab(self) -> dict[str, int]:\n        \"\"\"\n        Returns the added tokens in the vocabulary as a dictionary of token to index. Results might be different from\n        the fast call because for now we always add the tokens even if they are already in the vocabulary. This is\n        something we should change.\n\n        Returns:\n            `Dict[str, int]`: The added tokens.\n        \"\"\"\n        pass\n\n    @abstractmethod\n    def convert_tokens_to_string(self, tokens: list[str]) -> str:\n        \"\"\"\n        Converts a sequence of tokens in a single string. The most simple way to do it is `\" \".join(tokens)` but we\n        often want to remove sub-word tokenization artifacts at the same time.\n\n        Args:\n            tokens (`List[str]`): The token to join in a string.\n\n        Returns:\n            `str`: The joined tokens.\n        \"\"\"\n        pass\n\n    @property\n    def is_fast(self):\n        return False\n"
  },
  {
    "path": "verl/workers/rollout/trtllm_rollout/trtllm_async_rollout.md",
    "content": "# Running VeRL with TensorRT-LLM Rollout\n\nWe provide initial support for [TensorRT-LLM](https://github.com/NVIDIA/TensorRT-LLM) as an asynchronous rollout engine in VERL's reinforcement learning pipeline. It covers key features such as distributed inference with Ray-based orchestration, dynamic weight updates via IPC (Inter-Process Communication), and efficient GPU memory management for GRPO training.\n\nTRT-LLM rollout uses hybrid engine colocate mode, where training and inference workers are colocated on the same GPUs. Memory is managed via `resume()`/`release()` APIs to enable GPU sharing between training and inference workloads.\n\nWhile the current design factors in multi-node use cases, more extensive multi-node testing and functionality will be delivered in the near future. Current focus is on FSDP and Megatron backend support for Qwen model variants.\n\n---\n\n## 1. Quick Start\n\n\n```bash\n# GRPO with FSDP training engine and TP1\n>> bash examples/grpo_trainer/run_qwen2-7b_math_trtllm.sh 1\n```\n\nNote that using the TRT-LLM rollout requires setting the following environment variables before launching the Ray cluster, as included in the above script.\n\n```bash\n# Clean all SLURM/MPI/PMIx env to avoid pmix mismatch error.\nfor v in $(env | awk -F= '/^(PMI|PMIX|MPI|OMPI|SLURM)_/{print $1}'); do\n    unset \"$v\"\ndone\n```\n\n## 2. Architecture Design\n\n### 2.1 High-Level Component Diagram\n\n```mermaid\n%%{init: {'theme':'base', 'themeVariables': { 'fontSize':'18px', 'edgeLabelBackground':'#eeeeee'}}}%%\nflowchart TB\n    space1[\" \"]\n    style space1 fill:none,stroke:none\n    \n    subgraph VERL[\"<b>VERL Training Pipeline</b>\"]\n        subgraph Workers[\"<b>Training Workers</b>\"]\n            Actor[\"<b>Actor Worker</b>\"]\n            Critic[\"<b>Critic Worker</b>\"]\n            RefModel[\"<b>Ref Model Worker</b>\"]\n        end\n        \n        Actor -->|<b>Weight Updates<br/>IPC</b>| Rollout[\"<b>TensorRT-LLM Rollout</b>\"]\n        \n        subgraph RayCluster[\"<b>Rollout Workers<br/>(Ray Cluster)</b>\"]\n            space2[\" \"]\n            style space2 fill:none,stroke:none\n            \n    subgraph AsyncRollout[\"<b>ServerAdapter<br/>(per DP rank)</b>\"]\n        DPLeader[\"<b>• DP Leader coordination</b>\"]\n        IPCMgmt[\"<b>• IPC handle management</b>\"]\n        HTTPAdapter[\"<b>• HTTP adapter for server communication</b>\"]\n    end\n            \n            AsyncRollout -->|<b>HTTP/REST API</b>| HTTPServer\n            \n            subgraph HTTPServer[\"<b>TRTLLMHttpServer<br/>(Ray Actor per Replica)</b>\"]\n                OpenAI[\"<b>• OpenAI Server wrapper</b>\"]\n                EngMgmt[\"<b>• AsyncLLM engine management</b>\"]\n                MemMgmt[\"<b>• Memory management (resume/release)</b>\"]\n            end\n            \n            HTTPServer --> AsyncLLM\n            \n            subgraph AsyncLLM[\"<b>TensorRT-LLM<br/>AsyncLLM Engine</b>\"]\n                GPUWorkers[\"<b>• GPU workers (Tensor Parallel)</b>\"]\n                KVCache[\"<b>• KV Cache management</b>\"]\n                CUDAGraph[\"<b>• CUDA Graph optimization</b>\"]\n            end\n        end\n    end\n    \n    space1 ~~~ VERL\n    \n    style VERL fill:#e1f5ff\n    style RayCluster fill:#fff4e6\n    style AsyncRollout fill:#f3e5f5\n    style HTTPServer fill:#e8f5e9\n    style AsyncLLM fill:#fce4ec\n```\n\n### 2.2 Agent Loop Architecture\n\nTRT-LLM rollout follows the same Agent Loop architecture described in the [VERL documentation](https://verl.readthedocs.io/en/latest/advance/agent_loop.html).\n\nWith TensorRT-LLM rollout, the AsyncLLM engine runs in the same process as the TRTLLMHttpServer (Ray actor). The engine spawns Ray workers as ModelRunner through Ray's native orchestration with placement groups.\n\nAsyncLLM engine communicates with Ray workers through TensorRT-LLM's internal communication layer. When the server receives a request, it directly calls the AsyncLLM engine to generate response_ids. The Ray workers are separate processes from FSDP/Megatron-LM workers but are co-located on the same GPUs in hybrid engine mode.\n\nThe diagram below illustrates TRT-LLM's implementation in hybrid engine mode (Ray Workers and FSDP workers share GPUs):\n\n```mermaid\nflowchart TB\n    generate[generate]\n    \n    generate --> Server\n    \n    Server[TRTLLMHttpServer<br/>AsyncLLM Engine]\n    \n    Server --> Workers\n    \n    subgraph Workers[\"TRT-LLM group (TP4)\"]\n        direction LR\n        subgraph W0[ ]\n            RW0[Ray Worker-0]\n            F0[FSDP-0]\n        end\n        subgraph W1[ ]\n            RW1[Ray Worker-1]\n            F1[FSDP-1]\n        end\n        subgraph W2[ ]\n            RW2[Ray Worker-2]\n            F2[FSDP-2]\n        end\n        subgraph W3[ ]\n            RW3[Ray Worker-3]\n            F3[FSDP-3]\n        end\n    end\n    \n    style Server fill:#ffb6c1\n    style RW0 fill:#ffffe0\n    style RW1 fill:#ffffe0\n    style RW2 fill:#ffffe0\n    style RW3 fill:#ffffe0\n    style F0 fill:#ffb6c1\n    style F1 fill:#ffb6c1\n    style F2 fill:#ffb6c1\n    style F3 fill:#ffb6c1\n    style W0 fill:#d3d3d3\n    style W1 fill:#d3d3d3\n    style W2 fill:#d3d3d3\n    style W3 fill:#d3d3d3\n    style Workers fill:#f5f5f5\n```\n\n\n### 2.3 Ray Placement Group Architecture\n\n1. **Placement APIs & GPU Assignment**: TRT-LLM rollout leverages TRT-LLM's Ray-based APIs (`placement_groups`, `placement_bundle_indices`, `per_worker_gpu_share`) to control GPU placement. Each replica (corresponding to one `TRTLLMHttpServer`) is assigned GPU bundles from placement groups based on its replica rank and TP size.\n\n2. **Server Placement**: `TRTLLMHttpServer` is pinned to the same node as its first bundle using `NodeAffinitySchedulingStrategy`, ensuring efficient communication between the HTTP server and its Ray workers.\n\n3. **GPU Sharing**: In hybrid engine mode, training and inference workers share GPUs. Memory is managed via `resume()`/`release()` APIs. The resource pool uses `max_colocate_count=3` internally to support colocation of ActorRollout, RewardModel, and Critic workers.\n\n4. **Multi-Node Design**: The placement group slicing algorithm supports spanning multiple placement groups for multi-node deployments. **Note**: Formal multi-node testing and functionality will be delivered in subsequent MRs.\n\nThe following diagram shows an example of TP=4 and DP=2. Replica 0 takes bundles 0-3 and Replica 1 takes bundles 4-7 from the same placement group, with each replica managing TP workers across its assigned bundles:\n\n```mermaid\nflowchart TB\n    subgraph RayCluster[\"Ray Cluster Resource Pool\"]\n        subgraph PG0[\"Placement Group 0 (Node 0)\"]\n            B0_0[\"Bundle 0: GPU 0\"]\n            B0_1[\"Bundle 1: GPU 1\"]\n            B0_2[\"Bundle 2: GPU 2\"]\n            B0_3[\"Bundle 3: GPU 3\"]\n            B0_4[\"Bundle 4: GPU 4\"]\n            B0_5[\"Bundle 5: GPU 5\"]\n            B0_6[\"Bundle 6: GPU 6\"]\n            B0_7[\"Bundle 7: GPU 7\"]\n        end\n        \n        subgraph PG1[\"Placement Group 1 (Node 1)\"]\n            B1_0[\"Bundle 0: GPU 0\"]\n            B1_1[\"Bundle 1: GPU 1\"]\n            B1_2[\"Bundle 2: GPU 2\"]\n            B1_3[\"Bundle 3: GPU 3\"]\n            B1_4[\"Bundle 4: GPU 4\"]\n            B1_5[\"Bundle 5: GPU 5\"]\n            B1_6[\"Bundle 6: GPU 6\"]\n            B1_7[\"Bundle 7: GPU 7\"]\n        end\n        \n        PG0 --> Assignment\n        PG1 --> Assignment\n        \n        Assignment[\"Assigned to TRTLLMReplica\"]\n        \n        Assignment --> Replica0\n        Assignment --> Replica1\n        \n        Replica0[\"Replica 0<br/>(bundles 0-3 from PG0)<br/>TP=4, DP=2\"]\n        Replica1[\"Replica 1<br/>(bundles 4-7 from PG0)<br/>TP=4, DP=2\"]\n    end\n    \n    style PG0 fill:#e3f2fd\n    style PG1 fill:#e3f2fd\n    style Replica0 fill:#c8e6c9\n    style Replica1 fill:#c8e6c9\n```\n\n---\n\n## 3. Core Components\n\n### 3.1 `TRTLLMHttpServer`\n\n**Purpose**: Ray actor that wraps TensorRT-LLM's AsyncLLM engine and exposes an OpenAI-compatible HTTP API.\n\n**Key Responsibilities**:\n- Initialize and manage AsyncLLM engine with placement group constraints\n- Wrap AsyncLLM with OpenAIServer to expose HTTP endpoints\n- Handle HTTP server lifecycle (launch, shutdown)\n- Process generation requests with sampling parameters\n- Coordinate memory management (wake_up/sleep) for GPU sharing with training workers\n\n\n### 3.2 `TRTLLMReplica`\n\n**Purpose**: Manages the mapping between replicas and Ray placement groups, orchestrating server deployment.\n\n**Key Responsibilities**:\n- Calculate placement group and bundle index assignments per replica\n- Pin TRTLLMHttpServer to specific nodes using NodeAffinitySchedulingStrategy\n- Launch and coordinate HTTP servers across distributed nodes\n- Validate placement group configurations\n\n\n### 3.3 `ServerAdapter`\n\n**Purpose**: Rollout worker that handles weight updates, memory management, and generation via HTTP adapter.\n\nEach DP rank has one leader (the first TP rank within that DP group), and that leader coordinates weight updates to the corresponding TRTLLMHttpServer replica.\n\n**Key Responsibilities**:\n- Act as DP leader for weight synchronization across exclude_dp mesh\n- Convert PyTorch tensors to IPC handles for zero-copy weight updates\n- Stream weight updates in chunks to avoid memory exhaustion\n- Coordinate resume/release operations for memory management\n- Initialize HTTP adapter for server communication\n\n\n### 3.4 `AsyncTRTLLMHttpAdapter`\n\n**Purpose**: HTTP client for communicating with TRTLLMHttpServer.\n\n**Key Features**:\n- Async request handling with retry logic\n- Connection pooling for high throughput\n- Exponential backoff on failures\n- Timeout management\n\n---\n\n## 4. Data Flow Diagrams\n\n### 4.1 Generation Request Flow\n\n```mermaid\nsequenceDiagram\n    participant Client as Client/Actor\n    participant Rollout as ServerAdapter\n    participant Adapter as AsyncHttpAdapter\n    participant Server as TRTLLMHttpServer\n    participant AsyncLLM as AsyncLLM Engine\n    \n    Client->>Rollout: generate(prompts)\n    \n    rect rgb(240, 248, 255)\n        Note over Rollout: Init adapter if needed\n    end\n    \n    Rollout->>Adapter: POST /v1/completions<br/>{prompt_ids, sampling_params}\n    \n    rect rgb(255, 250, 240)\n        Note over Adapter: Retry loop with backoff\n    end\n    \n    Adapter->>Server: HTTP POST\n    \n    rect rgb(245, 255, 245)\n        Note over Server: Parse request<br/>Validate params\n    end\n    \n    Server->>AsyncLLM: generate_async()\n    \n    rect rgb(255, 245, 245)\n        Note over AsyncLLM: Schedule to execution queue\n        Note over AsyncLLM: Run inference (TP workers)<br/>- Forward pass<br/>- Sample tokens<br/>- Update KV cache\n    end\n    \n    AsyncLLM-->>Server: Output (token_ids, log_probs)\n    \n    Server-->>Adapter: JSON response\n    Adapter-->>Rollout: TokenOutput\n    Rollout-->>Client: Results\n```\n"
  },
  {
    "path": "verl/workers/rollout/trtllm_rollout/trtllm_async_server.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\nimport os\nfrom typing import Any, Optional\n\nimport ray\nimport torch\nfrom omegaconf import DictConfig\nfrom ray.actor import ActorHandle\nfrom ray.util import placement_group_table\nfrom ray.util.placement_group import PlacementGroup\n\nfrom verl.single_controller.ray import SubRayResourcePool\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.net_utils import is_valid_ipv6_address\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput\nfrom verl.workers.rollout.trtllm_rollout.trtllm_rollout import ServerAdapter\nfrom verl.workers.rollout.utils import get_max_position_embeddings, qwen2_5_vl_dedup_image_tokens, run_uvicorn\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(logging.INFO)\n\n\n@ray.remote\nclass TRTLLMHttpServer:\n    \"\"\"TensorRT LLM HTTP server in single node.\n\n    Args:\n        config (DictConfig): full config.\n        model_config (HFModelConfig): model config.\n        is_reward_model (bool): whether this is a reward model.\n        rollout_mode (RolloutMode): rollout mode.\n        workers (list[ActorHandle]): list of rollout workers.\n        replica_rank (int): replica rank, a replica may contain multiple nodes.\n        max_colocate_count (int): max colocate count.\n        pgs (list[PlacementGroup]): placement groups.\n        bundle_indices (list[list[int]]): bundle indices.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        is_reward_model: bool,\n        rollout_mode: RolloutMode,\n        workers: list[ActorHandle],\n        replica_rank: int,\n        max_colocate_count: int,\n        pgs: list[PlacementGroup] = None,\n        bundle_indices: list[list[int]] = None,\n    ):\n        os.environ[\"TRT_LLM_DISABLE_LOAD_WEIGHTS_IN_PARALLEL\"] = \"1\"\n        assert torch.cuda.is_available(), \"TRTLLM http server should run on GPU node\"\n\n        self.config: RolloutConfig = omega_conf_to_dataclass(config)\n        self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)\n        self.is_reward_model = is_reward_model\n        max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config)\n        if self.config.max_model_len is None:\n            self.config.max_model_len = max_position_embeddings\n        else:\n            if self.config.max_model_len > max_position_embeddings:\n                raise ValueError(\n                    f\"max_model_len ({self.config.max_model_len}) should be less than or equal to \"\n                    f\"max_position_embeddings ({max_position_embeddings})\"\n                )\n        self.rollout_mode = rollout_mode\n        self.workers = workers\n        self.replica_rank = replica_rank\n        self.max_colocate_count = max_colocate_count\n        self.pgs = pgs\n        self.bundle_indices = bundle_indices\n        # model weights version, set by ServerAdapter when update weights.\n        self.global_steps = None\n\n        if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == \"dummy\":\n            logger.warning(f\"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto\")\n            self.config.load_format = \"auto\"\n\n        self.is_vlm_model = (\n            self.model_config.hf_config is not None and hasattr(self.model_config.hf_config, \"vision_config\")\n        ) or hasattr(self.model_config, \"vision_config\")\n\n        # used for http server\n        self._server_address = ray.util.get_node_ip_address().strip(\"[]\")\n        self._server_port = None\n\n        logger.info(f\"TRTLLMHttpServer, replica_rank: {self.replica_rank}\")\n\n        self.sampling_args = {\n            \"detokenize\": False,\n            \"end_id\": -1,\n            \"pad_id\": self.model_config.hf_config.pad_token_id,\n            \"stop_token_ids\": [self.model_config.hf_config.eos_token_id],\n            \"include_stop_str_in_output\": True,\n        }\n\n    def get_server_address(self):\n        \"\"\"Get http server address and port.\"\"\"\n        assert self._server_port is not None, \"http server is not launched, port is None\"\n        return self._server_address, self._server_port\n\n    async def launch_server(self):\n        from tensorrt_llm import AsyncLLM\n        from tensorrt_llm.llmapi import CapacitySchedulerPolicy, CudaGraphConfig, KvCacheConfig, SchedulerConfig\n        from tensorrt_llm.serve import OpenAIServer\n\n        assert self.config.pipeline_model_parallel_size == 1, \"pipeline_model_parallel_size > 1 is not supported yet\"\n\n        engine_kwargs = self.config.get(\"engine_kwargs\", {}).get(\"trtllm\", {}) or {}\n        kv_cache_config = KvCacheConfig(\n            enable_block_reuse=self.config.enable_prefix_caching,\n            free_gpu_memory_fraction=self.config.gpu_memory_utilization,\n        )\n\n        per_worker_gpu_share = 1.0 / self.max_colocate_count\n\n        quantization = self.config.quantization\n        if quantization is not None:\n            if quantization == \"fp8\":\n                FP8_BLOCK_QUANT_KWARGS = {\n                    \"activation_scheme\": \"dynamic\",\n                    \"fmt\": \"e4m3\",\n                    \"quant_method\": \"fp8\",\n                    \"weight_block_size\": [128, 128],\n                }\n                engine_kwargs[\"model_kwargs\"] = {\"quantization_config\": FP8_BLOCK_QUANT_KWARGS}\n                if self.config.load_format != \"dummy\":\n                    raise ValueError(\"FP8 quantization is only supported for dummy load format\")\n            else:\n                raise ValueError(f\"Currently only support fp8 quantization, got: {quantization}\")\n\n        llm_kwargs = {\n            \"model\": self.model_config.local_path,\n            \"backend\": \"pytorch\",\n            \"dtype\": self.config.dtype,\n            \"enable_chunked_prefill\": self.config.enable_chunked_prefill,\n            \"skip_tokenizer_init\": self.config.skip_tokenizer_init,\n            \"orchestrator_type\": \"ray\",\n            \"kv_cache_config\": kv_cache_config,\n            \"max_seq_len\": self.config.max_model_len,\n            \"max_batch_size\": self.config.max_num_seqs,\n            \"max_num_tokens\": self.config.max_num_batched_tokens,\n            \"tensor_parallel_size\": self.config.tensor_model_parallel_size,\n            \"pipeline_parallel_size\": self.config.pipeline_model_parallel_size,\n            \"moe_expert_parallel_size\": self.config.expert_parallel_size,\n            \"moe_tensor_parallel_size\": self.config.moe_tensor_parallel_size,\n            \"load_format\": self.config.load_format,\n            \"trust_remote_code\": self.model_config.trust_remote_code,\n            \"placement_groups\": self.pgs,\n            \"placement_bundle_indices\": self.bundle_indices,\n            \"per_worker_gpu_share\": per_worker_gpu_share,\n            \"enable_sleep\": self.config.enable_sleep_mode,\n            \"allreduce_strategy\": \"NCCL\",\n            \"sampler_type\": \"TRTLLMSampler\",\n            **engine_kwargs,\n        }\n\n        self_defined_extension = {\n            \"ray_worker_extension_cls\": \"verl.workers.rollout.trtllm_rollout.trtllm_worker_extension.WorkerExtension\",\n        }\n        if self.is_vlm_model:\n            llm_kwargs.update(self_defined_extension)\n        else:\n            llm_kwargs.update(\n                {\n                    \"ray_worker_extension_cls\": \"tensorrt_llm.llmapi.rlhf_utils.WorkerExtension\",\n                }\n            )\n\n        if self.is_reward_model:\n            llm_kwargs.update(\n                {\n                    \"cuda_graph_config\": None,\n                    \"disable_overlap_scheduler\": True,\n                }\n            )\n        else:\n            llm_kwargs.update(\n                {\n                    \"cuda_graph_config\": CudaGraphConfig(\n                        enable_padding=True,\n                        batch_sizes=self.config.cudagraph_capture_sizes,\n                        max_batch_size=0 if self.config.cudagraph_capture_sizes else self.config.max_num_seqs,\n                    ),\n                    \"scheduler_config\": SchedulerConfig(\n                        capacity_scheduler_policy=CapacitySchedulerPolicy.MAX_UTILIZATION,\n                    ),\n                }\n            )\n\n        self.llm = await AsyncLLM(**llm_kwargs)\n        import inspect\n\n        init_params = inspect.signature(OpenAIServer.__init__).parameters\n        if \"generator\" in init_params:\n            trtllm_server = OpenAIServer(\n                generator=self.llm,\n                model=self.model_config.local_path,\n                tool_parser=None,\n                server_role=None,\n                metadata_server_cfg=None,\n            )\n        else:\n            trtllm_server = OpenAIServer(\n                llm=self.llm,\n                model=self.model_config.local_path,\n                tool_parser=None,\n                server_role=None,\n                metadata_server_cfg=None,\n            )\n\n        app = trtllm_server.app\n        self._server_port, self._server_task = await run_uvicorn(app, None, self._server_address)\n\n    async def generate(\n        self,\n        prompt_ids: str | list[int],\n        sampling_params: dict[str, Any],\n        request_id: str,\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n    ) -> TokenOutput:\n        from tensorrt_llm.llmapi import SamplingParams\n\n        max_tokens = min(self.config.response_length, self.config.max_model_len - len(prompt_ids))\n        sampling_params[\"max_tokens\"] = max_tokens\n        sampling_params[\"logprobs\"] = 1 if sampling_params.pop(\"logprobs\", False) else None\n        if sampling_params[\"top_k\"] == -1:\n            sampling_params[\"top_k\"] = 0\n        sampling_params.update(self.sampling_args)\n\n        trt_llm_sampling_params = SamplingParams(**sampling_params)\n        if self.is_vlm_model and (image_data or video_data):\n            deduped_ids = qwen2_5_vl_dedup_image_tokens(prompt_ids, self.model_config.processor)\n            org_prompt = self.llm.tokenizer.decode(deduped_ids)\n            input_dict = {\n                \"prompt\": org_prompt,\n                \"multi_modal_data\": {},\n                \"mm_processor_kwargs\": {},\n            }\n            if image_data:\n                input_dict[\"multi_modal_data\"][\"image\"] = image_data\n            if video_data:\n                input_dict[\"multi_modal_data\"][\"video\"] = video_data\n\n            outputs = await self.llm.generate_async(\n                inputs=input_dict,\n                sampling_params=trt_llm_sampling_params,\n            )\n        else:\n            outputs = await self.llm.generate_async(\n                inputs=prompt_ids,\n                sampling_params=trt_llm_sampling_params,\n            )\n        token_ids = outputs.outputs[0].token_ids\n        log_probs = None\n        if outputs.outputs[0].logprobs is not None:\n            # When logprobs=1, TRT-LLM returns only the sampled token's logprob at each position\n            log_probs = [list(d.values())[0].logprob for d in outputs.outputs[0].logprobs]\n        return TokenOutput(token_ids=token_ids, log_probs=log_probs, extra_fields={\"global_steps\": self.global_steps})\n\n    async def set_global_steps(self, global_steps: int):\n        \"\"\"Set the global steps of the model weights.\"\"\"\n        self.global_steps = global_steps\n\n    async def abort_all_requests(self):\n        raise NotImplementedError\n\n    async def resume_generation(self):\n        raise NotImplementedError\n\n    async def wake_up(self):\n        if self.rollout_mode == RolloutMode.HYBRID:\n            # In hybrid mode, rollout is wake up in `update_weights`\n            raise ValueError(f\"wake_up not support rollout_mode {self.rollout_mode}\")\n        if self.rollout_mode == RolloutMode.COLOCATED:\n            await self.llm.resume(tags=ServerAdapter.get_full_tags())\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            logger.info(\"skip wake_up in standalone mode\")\n\n    async def sleep(self):\n        if not self.config.free_cache_engine:\n            return\n\n        if self.rollout_mode == RolloutMode.HYBRID:\n            await self.llm.release(tags=ServerAdapter.get_full_tags())\n        elif self.rollout_mode == RolloutMode.COLOCATED:\n            await self.llm.release(tags=ServerAdapter.get_full_tags())\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            logger.info(\"skip sleep in standalone mode\")\n\n    async def report_device_ids(self) -> list[str]:\n        \"\"\"Report GPU device UUIDs from TRT-LLM workers.\"\"\"\n        return await self.llm.collective_rpc(\n            \"report_device_id\",\n            unique_reply_rank=0,\n        )\n\n\nclass TRTLLMReplica(RolloutReplica):\n    def __init__(\n        self,\n        replica_rank: int,\n        config: RolloutConfig,\n        model_config: DictConfig,\n        gpus_per_node: int = 8,\n        is_reward_model: bool = False,\n    ) -> None:\n        super().__init__(replica_rank, config, model_config, gpus_per_node, is_reward_model)\n        self.node_ip = ray.util.get_node_ip_address().strip(\"[]\")\n\n    def rollout_worker_use_gpu(self) -> bool:\n        return False\n\n    def get_pgs_and_bundle_indices(self) -> tuple[list[PlacementGroup], list[list[int]]]:\n        \"\"\"Get placement groups and bundle indices for the replica.\"\"\"\n\n        start_pg_index = 0\n        local_bundle_index = 0\n\n        # For SubRayResourcePool, the replica is assigned sub pool specific for this replica.\n        if isinstance(self.resource_pool, SubRayResourcePool):\n            assert self.resource_pool.subgroup_world_size == self.world_size, (\n                \"Subgroup world size must be equal to world size\"\n            )\n            local_bundle_index = self.resource_pool.start_bundle_index\n        # For RayResourcePool, the replica is assigned to entire resource pool.\n        # We need to find start pg index and local bundle index based on replica rank.\n        else:\n            local_bundle_index = self.world_size * self.replica_rank\n\n        while local_bundle_index >= self.resource_pool.pgs[start_pg_index].bundle_count:\n            start_pg_index += 1\n            local_bundle_index -= self.resource_pool.pgs[start_pg_index].bundle_count\n        assert (\n            start_pg_index < len(self.resource_pool.pgs)\n            and local_bundle_index < self.resource_pool.pgs[start_pg_index].bundle_count\n        ), \"Start pg index or local bundle index out of range\"\n\n        # Global Bundle View for Replica x 2 & TP=4:\n        # ┌───────────────────┬───────────────────┐\n        # │ Placement Group 0 │ Placement Group 1 │\n        # ├────┬────┬────┬────┼────┬────┬────┬────┤\n        # │ 0  │ 1  │ 2  │ 3  │ 0  │ 1  │ 2  │ 3  │\n        # └────┴────┴────┴────┴────┴────┴────┴────┘\n        #   └───────────────┘   └───────────────┘\n        #       Replica 0           Replica 1\n        #       (4 GPUs)            (4 GPUs)\n\n        left_bundle_count = self.world_size\n\n        pgs = []\n        bundle_indices = []\n\n        for pg in self.resource_pool.pgs[start_pg_index:]:\n            if left_bundle_count == 0:\n                break\n\n            left_bundle_count_in_pg = min(left_bundle_count, pg.bundle_count - local_bundle_index)\n            pg_bundle_indices = [local_bundle_index + idx for idx in range(left_bundle_count_in_pg)]\n            pgs.append(pg)\n            bundle_indices.append(pg_bundle_indices)\n            left_bundle_count -= left_bundle_count_in_pg\n            local_bundle_index = 0\n\n        assert left_bundle_count == 0, \"all bundle indices should be assigned\"\n\n        return pgs, bundle_indices\n\n    async def launch_servers(self):\n        assert self.nnodes == 1, \"TRTLLMReplica doesn't support multiple nodes for single replica yet.\"\n        assert self.resource_pool.pgs is not None, \"placement groups are not initialized\"\n\n        pgs, bundle_indices = self.get_pgs_and_bundle_indices()\n\n        # Check server process should be launched on the same node as first bundle of first pg.\n        first_pg_data = placement_group_table(pgs[0])\n        node_id = first_pg_data[\"bundles_to_node_id\"][bundle_indices[0][0]]\n        print(f\"TRTLLMReplica: {self.replica_rank}\")\n        print(f\"pg node_id: {node_id}\")\n        print(f\"pgs: {pgs}\")\n        print(f\"bundle_indices: {bundle_indices}\")\n\n        # TRTLLMReplica is a 1:1 map from replica to TRTLLMHttpServer.\n        name = (\n            f\"trtllm_server_{self.replica_rank}\"\n            if not self.is_reward_model\n            else f\"trtllm_server_reward_{self.replica_rank}\"\n        )\n\n        server = TRTLLMHttpServer.options(\n            scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                node_id=node_id,\n                soft=False,\n            ),\n            runtime_env={\"env_vars\": {\"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES\": \"1\", \"NCCL_CUMEM_ENABLE\": \"0\"}},\n            name=name,\n            max_concurrency=self.max_concurrency,\n        ).remote(\n            config=self.config,\n            model_config=self.model_config,\n            is_reward_model=self.is_reward_model,\n            rollout_mode=self.rollout_mode,\n            workers=self.workers,\n            replica_rank=self.replica_rank,\n            max_colocate_count=self.resource_pool.max_colocate_count,\n            pgs=pgs,\n            bundle_indices=bundle_indices,\n        )\n        self.servers.append(server)\n\n        # launch http server in each node\n        await asyncio.gather(*[server.launch_server.remote() for server in self.servers])\n\n        # get http server address from first server\n        server_address, server_port = await self.servers[0].get_server_address.remote()\n        self._server_handle = self.servers[0]\n        self._server_address = (\n            f\"[{server_address}]:{server_port}\"\n            if is_valid_ipv6_address(server_address)\n            else f\"{server_address}:{server_port}\"\n        )\n"
  },
  {
    "path": "verl/workers/rollout/trtllm_rollout/trtllm_rollout.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nfrom __future__ import annotations\n\nimport asyncio\nimport base64\nimport contextlib\nimport gc\nimport logging\nimport os\nimport pickle\nimport threading\nfrom contextlib import asynccontextmanager\nfrom typing import Any, Generator, Optional\n\nimport aiohttp\nimport pynvml\nimport ray\nimport torch\nimport torch.distributed as dist\nfrom torch.distributed.device_mesh import DeviceMesh, init_device_mesh\nfrom torch.multiprocessing.reductions import reduce_tensor\n\nfrom verl.utils.device import get_torch_device\nfrom verl.utils.net_utils import is_valid_ipv6_address\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.base import BaseRollout\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n# Default configuration constants\nDEFAULT_TIMEOUT = 60.0\nDEFAULT_MAX_ATTEMPTS = 3\nDEFAULT_RETRY_DELAY = 2.0\nDEFAULT_MAX_CONNECTIONS = 2000\nDEFAULT_MAX_WAIT_TIME = 300.0\n\n\n@contextlib.contextmanager\ndef nvml_context():\n    \"\"\"Context manager for NVML initialization and shutdown.\n\n    Raises:\n        RuntimeError: If NVML initialization fails\n    \"\"\"\n    try:\n        pynvml.nvmlInit()\n        yield\n    except pynvml.NVMLError as e:\n        raise RuntimeError(f\"Failed to initialize NVML: {e}\") from e\n    finally:\n        try:\n            pynvml.nvmlShutdown()\n        except pynvml.NVMLError:\n            pass\n\n\n_NVML_INITIALIZED = False\n_NVML_LOCK = threading.Lock()\n\n\ndef get_device_uuid(id: str | int) -> str:\n    \"\"\"Get the UUID of a CUDA device using NVML.\"\"\"\n    id = int(id)  # pynvml expects int; ray.get_gpu_ids() may return str\n    global _NVML_INITIALIZED\n    with _NVML_LOCK:\n        if not _NVML_INITIALIZED:\n            try:\n                pynvml.nvmlInit()\n                _NVML_INITIALIZED = True\n            except pynvml.NVMLError as e:\n                raise RuntimeError(f\"Failed to initialize NVML: {e}\") from e\n\n    # Get the device handle and UUID\n    try:\n        handle = pynvml.nvmlDeviceGetHandleByIndex(id)\n        uuid = pynvml.nvmlDeviceGetUUID(handle)\n        # Ensure the UUID is returned as a string, not bytes\n        if isinstance(uuid, bytes):\n            return uuid.decode(\"utf-8\")\n        elif isinstance(uuid, str):\n            return uuid\n        else:\n            raise RuntimeError(f\"Unexpected UUID type: {type(uuid)} for device {id} (global index: {id})\")\n    except pynvml.NVMLError as e:\n        raise RuntimeError(f\"Failed to get device UUID for device {id} (global index: {id}): {e}\") from e\n\n\nasync def _read_async_response(resp: aiohttp.ClientResponse) -> dict[str, Any]:\n    if resp.status == 204 or (resp.content_length == 0):\n        return {}\n\n    try:\n        return await resp.json(content_type=None)\n    except Exception:\n        try:\n            text = await resp.text()\n        except Exception:\n            return {}\n        return {\n            \"content_type\": (resp.headers.get(\"Content-Type\") or \"\"),\n            \"text\": text,\n        }\n\n\nclass AsyncTRTLLMHttpAdapter:\n    def __init__(\n        self,\n        host: str,\n        port: int,\n        timeout: float = DEFAULT_TIMEOUT,\n        max_attempts: int = DEFAULT_MAX_ATTEMPTS,\n        retry_delay: float = DEFAULT_RETRY_DELAY,\n        max_connections: int = DEFAULT_MAX_CONNECTIONS,\n    ):\n        self.host = host\n        self.port = port\n        self.timeout = timeout\n        self.max_attempts = max_attempts\n        self.retry_delay = retry_delay\n        self.max_connections = max_connections\n\n    @asynccontextmanager\n    async def _get_session(self) -> aiohttp.ClientSession:\n        \"\"\"Context manager for safe session access with proper connection pooling.\n\n        Yields:\n            aiohttp.ClientSession: Session instance for making HTTP requests\n\n        Note:\n            This method creates a new session for each request to avoid resource competition\n            while still maintaining proper connection pooling through the shared connector.\n        \"\"\"\n        # Create a new session for each request to avoid resource competition\n        connector = aiohttp.TCPConnector(\n            limit=self.max_connections,\n            limit_per_host=self.max_connections // 4,\n            ttl_dns_cache=300,\n            use_dns_cache=True,\n        )\n        timeout = aiohttp.ClientTimeout(total=self.timeout)\n        session = aiohttp.ClientSession(connector=connector, timeout=timeout)\n\n        try:\n            yield session\n        finally:\n            # Always close the session to free up resources\n            if not session.closed:\n                await session.close()\n\n    async def _make_async_request(\n        self,\n        endpoint: str,\n        payload: Optional[dict[str, Any]] = None,\n        timeout: float = DEFAULT_TIMEOUT,\n        method: str = \"POST\",\n        return_status: bool = False,\n    ) -> dict[str, Any] | int:\n        \"\"\"Make an async HTTP request with retry logic and consistent error handling.\n\n        Args:\n            endpoint (str): The API endpoint to call (without leading slash)\n            payload (Optional[Dict[str, Any]], optional): The JSON payload to send.\n                Defaults to empty dict if None.\n            method (str, optional): HTTP method to use. Defaults to \"POST\".\n\n        Returns:\n            Dict[str, Any]: The JSON response from the server\n\n        Raises:\n            aiohttp.ClientResponseError: If the HTTP request fails with a client/server error\n            RuntimeError: If all retry attempts are exhausted\n\n        Note:\n            - Uses exponential backoff for retries\n            - Logs warnings for timeout and connection errors, errors for HTTP errors\n        \"\"\"\n\n        url = f\"http://{self.host}:{self.port}/{endpoint}\"\n\n        for attempt in range(self.max_attempts):\n            try:\n                async with self._get_session() as session:\n                    if method.upper() == \"GET\":\n                        async with session.get(url, timeout=timeout) as response:\n                            response.raise_for_status()\n                            return response.status if return_status else await _read_async_response(response)\n                    else:\n                        async with session.post(url, json=payload or {}, timeout=timeout) as response:\n                            response.raise_for_status()\n                            return response.status if return_status else await _read_async_response(response)\n\n            except asyncio.TimeoutError:\n                logger.warning(f\"Async request to {endpoint} timed out (attempt {attempt + 1})\")\n            except aiohttp.ClientConnectorError:\n                logger.warning(f\"Connection error for {endpoint} (attempt {attempt + 1})\")\n            except aiohttp.ClientResponseError as e:\n                logger.error(f\"HTTP error for {endpoint}: {e}\")\n                raise\n            except Exception as e:\n                logger.error(f\"Unexpected error for {endpoint}: {e}\")\n                if attempt == self.max_attempts - 1:\n                    raise\n\n            if attempt < self.max_attempts - 1:\n                await asyncio.sleep(self.retry_delay * (2**attempt))\n\n        raise RuntimeError(f\"Failed to complete async request to {endpoint} after {self.max_attempts} attempts\")\n\n    async def resume_memory_occupation(self, tags: list[str]):\n        \"\"\"Resume GPU memory occupation (async version).\n\n        Similar to AsyncEngine, this method handles first-time weight reloading\n        by calling release_memory_occupation if needed.\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to resume.\n                If None, resumes all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory resume status\n        \"\"\"\n        return await self._make_async_request(\"resume_memory\", {\"tags\": tags})\n\n    async def release_memory_occupation(self, tags: list[str]):\n        \"\"\"Release GPU memory occupation temporarily (async version).\n\n        Args:\n            tags (Optional[List[str]], optional): List of tags to specify which memory to release.\n                If None, releases all memory. Defaults to None. [\"weights\", \"kv_cache\"]\n\n        Returns:\n            Dict[str, Any]: Server response indicating memory release status\n        \"\"\"\n        return await self._make_async_request(\"release_memory\", {\"tags\": tags})\n\n    async def update_weights(self, weights: dict[str, str]):\n        \"\"\"Update model weights from tensor data asynchronously.\n\n        Args:\n            weights: A dictionary that maps the device uuid of the weight handles.\n\n        Returns:\n            Dict[str, Any]: Server response containing update status\n        \"\"\"\n        return await self._make_async_request(\"update_weights\", {\"weights\": weights})\n\n\nclass ServerAdapter(BaseRollout):\n    _WEIGHTS_TAGS = [\n        \"sampler\",\n        \"drafter\",\n        \"guided_decoder\",\n        \"spec_resource_manager\",\n        \"model_extra\",\n        \"executor_extra\",\n        \"model\",\n        \"draft_model\",\n    ]\n\n    @staticmethod\n    def get_full_tags() -> list[str]:\n        return ServerAdapter._WEIGHTS_TAGS + [\"kv_cache\"]\n\n    def __init__(\n        self, config: RolloutConfig, model_config: HFModelConfig, device_mesh: DeviceMesh, replica_rank: int = -1\n    ):\n        if config.get(\"quantization\", None) == \"fp8\":\n            FP8_BLOCK_QUANT_KWARGS = {\n                \"activation_scheme\": \"dynamic\",\n                \"fmt\": \"e4m3\",\n                \"quant_method\": \"fp8\",\n                \"weight_block_size\": [128, 128],\n            }\n            fp8_block_quant_kwargs = dict(FP8_BLOCK_QUANT_KWARGS)\n            model_config.hf_config.quantization_config = fp8_block_quant_kwargs\n        super().__init__(config, model_config, device_mesh)\n        self._adapter = None\n        self.hybrid_device_mesh = None\n        self.gpu_id = None\n        self.is_leader_rank = None\n        self.replica_rank = None\n        self.is_dp_rank = None\n        self._supports_partial_loading = None\n\n        # hybrid mode\n        if self.device_mesh is not None:\n            assert device_mesh.mesh_dim_names.index(\"dp\") == 0, \"DP dim should always be the first dimension\"\n\n            # Clone a new device mesh for CPU backend only (used for internal ranks communication)\n            device_mesh_kwargs = dict(\n                mesh_shape=device_mesh.mesh.shape,\n                mesh_dim_names=device_mesh.mesh_dim_names,\n            )\n            self.hybrid_device_mesh = init_device_mesh(\"cpu\", **device_mesh_kwargs)\n\n            self.hybrid_device_mesh[self.hybrid_device_mesh.mesh_dim_names[1:]]._flatten(mesh_dim_name=\"exclude_dp\")\n            self.is_leader_rank = self.hybrid_device_mesh[\"exclude_dp\"].get_local_rank() == 0\n            logger.info(f\"is_dp_leader: {self.is_leader_rank}\")\n            logger.info(f\"exclude_dp_rank = {self.hybrid_device_mesh['exclude_dp'].get_local_rank()}\")\n            logger.info(f\"exclude_dp_size = {self.hybrid_device_mesh['exclude_dp'].size()}\")\n            self.gpu_id = ray.get_gpu_ids()[0]\n            self.replica_rank = self.hybrid_device_mesh[\"dp\"].get_local_rank()\n            assert len(ray.get_gpu_ids()) == 1, \"ServerAdapter should run on a single GPU node\"\n        else:\n            rank = int(os.environ[\"RANK\"])\n            self.replica_rank = replica_rank\n            self.is_leader_rank = rank == 0\n\n        # Below is required for all modes.\n        assert self.replica_rank >= 0, \"replica_rank is not set\"\n        assert self.is_leader_rank is not None, \"is_leader_rank is not set\"\n\n        self.node_ip = ray.util.get_node_ip_address().strip(\"[]\")\n\n    async def get_supports_partial_loading(self) -> bool:\n        \"\"\"Query and cache whether the model supports partial weight loading.\"\"\"\n        if self._supports_partial_loading is not None:\n            return self._supports_partial_loading\n\n        await self._init_server_adapter()\n        try:\n            self._supports_partial_loading = await self.server_actor.supports_partial_loading.remote()\n        except Exception as e:\n            logger.warning(f\"Failed to query partial loading support: {e}, defaulting to False\")\n            self._supports_partial_loading = False\n\n        logger.info(f\"Model supports partial loading: {self._supports_partial_loading}\")\n        return self._supports_partial_loading\n\n    async def _init_server_adapter(self):\n        if self._adapter is not None:\n            return\n\n        # Lazy init http server adapter because http server is launched after hybrid engine.\n        self.server_actor = ray.get_actor(f\"trtllm_server_{self.replica_rank}\")\n        server_address, server_port = await self.server_actor.get_server_address.remote()\n        assert server_address == self.node_ip, f\"server address: {server_address} != node_ip: {self.node_ip}\"\n\n        logger.debug(f\"replica_rank={self.replica_rank}, server address: {server_address}, port: {server_port}\")\n        host = f\"[{server_address}]\" if is_valid_ipv6_address(server_address) else server_address\n        self._adapter = AsyncTRTLLMHttpAdapter(\n            host=host,\n            port=server_port,\n            timeout=self.config.server.timeout,\n            max_attempts=self.config.server.max_attempts,\n            retry_delay=self.config.server.retry_delay,\n            max_connections=self.config.server.max_connections,\n        )\n\n    async def resume(self, tags: list[str]):\n        \"\"\"Resume rollout weights or kv cache in GPU memory.\n\n        Args:\n            tag: weights or kv_cache.\n        \"\"\"\n        # Synchronize all ranks before resuming KV cache to ensure non-leader ranks\n        # have completed actor offloading to CPU, preventing OOM issue.\n        if \"kv_cache\" in tags and self.config.free_cache_engine:\n            await asyncio.to_thread(dist.barrier, group=self.hybrid_device_mesh[\"exclude_dp\"].get_group())\n        if self.is_leader_rank and self.config.free_cache_engine:\n            if \"weights\" in tags:\n                tags = self._WEIGHTS_TAGS\n            elif \"kv_cache\" in tags:\n                tags = [\"kv_cache\"]\n            else:\n                raise ValueError(f\"Invalid tag: {tags}\")\n            await self._init_server_adapter()\n            await self._adapter.resume_memory_occupation(tags=tags)\n\n    async def release(self):\n        \"\"\"Release weights and kv cache in GPU memory.\"\"\"\n        if self.is_leader_rank and self.config.free_cache_engine:\n            await self._init_server_adapter()\n            tags = self._WEIGHTS_TAGS + [\"kv_cache\"]\n            await self._adapter.release_memory_occupation(tags=tags)\n\n    async def update_weights_from_ipc_handles(self, device_handles):\n        assert self.hybrid_device_mesh is not None, \"hybrid_device_mesh is not set\"\n\n        \"\"\"Update weights from IPC handles.\"\"\"\n        if self.is_leader_rank:\n            gathered_handles = [None for _ in range(self.hybrid_device_mesh[\"exclude_dp\"].size())]\n        else:\n            gathered_handles = None\n\n        await asyncio.to_thread(\n            dist.gather_object,\n            obj=device_handles,\n            object_gather_list=gathered_handles,\n            group_dst=0,\n            group=self.hybrid_device_mesh[\"exclude_dp\"].get_group(),\n        )\n\n        if self.is_leader_rank:\n            all_handles = {k: v for d in gathered_handles for k, v in d.items()}\n            await self._adapter.update_weights(all_handles)\n\n        await asyncio.to_thread(dist.barrier, group=self.hybrid_device_mesh[\"exclude_dp\"].get_group())\n\n    async def update_weights(\n        self, weights: Generator[tuple[str, torch.Tensor], None, None], global_steps: int = None, **kwargs\n    ):\n        assert self.hybrid_device_mesh is not None, \"hybrid_device_mesh is not set\"\n\n        \"\"\"Update the weights of the rollout model.\n\n        Args:\n            weights: A generator that yields the name of the weight tensor and the tensor itself.\n        \"\"\"\n        if self.is_leader_rank:\n            await self._init_server_adapter()\n\n        total_available_bytes = int(self.config.checkpoint_engine.update_weights_bucket_megabytes) * 1024 * 1024\n\n        if self.config.get(\"quantization\", None) == \"fp8\":\n            from verl.utils.trtllm.trtllm_fp8_utils import TRTLLMFP8QuantizerHelper\n\n            fp8_quantizer_helper = TRTLLMFP8QuantizerHelper(self.model_config.hf_config.quantization_config)\n            weights = fp8_quantizer_helper.quant_weights_by_name(\n                weights,\n                dtype=self.model_config.hf_config.dtype,\n            )\n\n        try:\n            device_uuid = get_device_uuid(int(self.gpu_id))\n        except Exception as e:\n            logger.error(f\"Failed to get device UUID in update_weights(): {e}\")\n            device_uuid = None\n            raise e\n\n        cur_available_bytes = total_available_bytes\n        cur_handles = []\n\n        async def flush():\n            nonlocal cur_available_bytes, cur_handles\n            if not cur_handles:\n                return\n            serialized_device_handles = {device_uuid: base64.b64encode(pickle.dumps(cur_handles)).decode(\"utf-8\")}\n            await self.update_weights_from_ipc_handles(serialized_device_handles)\n            cur_available_bytes = total_available_bytes\n            cur_handles = []\n\n        # Query if model supports partial loading\n        supports_partial_loading = await self.get_supports_partial_loading()\n\n        for name, param in weights:\n            if supports_partial_loading:\n                size_in_bytes = param.element_size() * param.numel()\n                if size_in_bytes > cur_available_bytes:\n                    await flush()\n\n                assert cur_available_bytes >= size_in_bytes, (\n                    f\"cur_available_bytes: {cur_available_bytes:,} size_in_bytes: {size_in_bytes:,} name: {name}\"\n                )\n                cur_available_bytes -= size_in_bytes\n\n            handle = reduce_tensor(param.detach())\n            cur_handles.append((name, handle))\n\n        await flush()\n\n        if self.is_leader_rank:\n            # Finalize update weights\n            await self._adapter.update_weights(None)\n            if global_steps is not None:\n                await self.server_actor.set_global_steps.remote(global_steps)\n        await asyncio.to_thread(dist.barrier, group=self.hybrid_device_mesh[\"exclude_dp\"].get_group())\n\n        del weights\n        gc.collect()\n        get_torch_device().empty_cache()\n\n    def _get_attribute(self, name: str):\n        return getattr(self, name)\n"
  },
  {
    "path": "verl/workers/rollout/trtllm_rollout/trtllm_worker_extension.py",
    "content": "# Copyright 2026 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport base64\nimport inspect\nfrom typing import Optional\n\nfrom tensorrt_llm import serialization\nfrom tensorrt_llm._ray_utils import control_action_decorator\nfrom tensorrt_llm._torch.modules.fused_moe.moe_load_balancer import MoeLoadBalancer\nfrom tensorrt_llm._torch.utils import get_device_uuid\nfrom tensorrt_llm.llmapi.rlhf_utils import WorkerExtension as TrtllmWorkerExtension\nfrom tensorrt_llm.logger import logger\n\n\nclass WorkerExtension(TrtllmWorkerExtension):\n    def __init__(self):\n        pass\n\n    @control_action_decorator\n    def supports_partial_loading(self) -> bool:\n        \"\"\"Check if the model supports partial weight loading.\"\"\"\n        try:\n            model = self.engine.model_engine.model\n            load_weights_args = inspect.getfullargspec(model.load_weights).args\n            return \"allow_partial_loading\" in load_weights_args\n        except Exception as e:\n            logger.warning(f\"Failed to check partial loading support: {e}\")\n            return False\n\n    @control_action_decorator\n    def update_weights(self, ipc_handles: Optional[dict] = None):\n        try:\n            if not hasattr(self.engine.model_engine.model, \"first_pre_reload_weights\"):\n                for module in self.engine.model_engine.model.modules():\n                    if hasattr(module, \"pre_reload_weights\") and not getattr(module, \"_weights_removed\", False):\n                        module.pre_reload_weights()\n                self.engine.model_engine.model.first_pre_reload_weights = True\n\n            if ipc_handles is not None:\n                logger.info(\"Update weights from IPC handles\")\n                device_uuid = get_device_uuid(self.device_id)\n\n                if device_uuid not in ipc_handles:\n                    raise ValueError(f\"Device UUID {device_uuid} not found in ipc_handles\")\n\n                weights = {}\n\n                serialized_handles = ipc_handles[device_uuid]\n                if isinstance(serialized_handles, str):\n                    # Data is base64-encoded pickled bytes - deserialize it\n                    # using restricted unpickler from tensorrt_llm.serialization\n                    logger.info(\"Deserializing base64-encoded weight handles\")\n                    decoded_data = base64.b64decode(serialized_handles)\n                    # Allow basic builtins and torch tensor reconstruction classes\n                    approved_imports = {\n                        \"builtins\": [\n                            \"list\",\n                            \"tuple\",\n                            \"str\",\n                            \"int\",\n                            \"float\",\n                            \"bool\",\n                            \"bytes\",\n                            \"dict\",\n                            \"NoneType\",\n                            \"type\",\n                        ],\n                        \"torch\": [\n                            \"Tensor\",\n                            \"FloatTensor\",\n                            \"DoubleTensor\",\n                            \"HalfTensor\",\n                            \"BFloat16Tensor\",\n                            \"IntTensor\",\n                            \"LongTensor\",\n                            \"ShortTensor\",\n                            \"CharTensor\",\n                            \"ByteTensor\",\n                            \"BoolTensor\",\n                            \"Size\",\n                            \"dtype\",\n                            \"device\",\n                            \"float32\",\n                            \"float16\",\n                            \"int32\",\n                            \"int64\",\n                            \"int16\",\n                            \"int8\",\n                            \"uint8\",\n                            \"bool\",\n                        ],\n                        \"torch.multiprocessing.reductions\": [\n                            \"rebuild_cuda_tensor\",\n                            \"rebuild_tensor\",\n                        ],\n                        \"torch._utils\": [\n                            \"_rebuild_tensor_v2\",\n                        ],\n                        \"torch.storage\": [\n                            \"_load_from_bytes\",\n                            \"_TypedStorage\",\n                            \"UntypedStorage\",\n                            \"TypedStorage\",\n                        ],\n                    }\n                    all_handles = serialization.loads(\n                        decoded_data,\n                        approved_imports=approved_imports,\n                    )\n\n                    # Verify the result is a list as expected\n                    if not isinstance(all_handles, list):\n                        raise ValueError(f\"Deserialized data must be a list, got {type(all_handles).__name__} instead\")\n                else:\n                    # Data is already in the correct format (backward compatibility)\n                    all_handles = serialized_handles\n\n                for param_name, tensor_handle in all_handles:\n                    func, args = tensor_handle\n                    list_args = list(args)\n                    list_args[6] = self.device_id\n                    tensor = func(*list_args)\n                    weights[param_name] = tensor\n\n                logger.info(f\"weights key size: {len(weights.keys())}\")\n\n                # Check if model supports partial loading and use appropriate strategy\n                model = self.engine.model_engine.model\n                load_weights_args = inspect.getfullargspec(model.load_weights).args\n                supports_partial_loading = \"allow_partial_loading\" in load_weights_args\n\n                if supports_partial_loading:\n                    self.engine.model_engine.model_loader.reload(model, weights, allow_partial_loading=True)\n                else:\n                    self.engine.model_engine.model_loader.reload(model, weights, allow_partial_loading=False)\n            else:\n                logger.info(\"Finalize update weights\")\n                for module in self.engine.model_engine.model.modules():\n                    if hasattr(module, \"process_weights_after_loading\") and not getattr(\n                        module, \"_weights_removed\", False\n                    ):\n                        module.process_weights_after_loading()\n                    if hasattr(module, \"post_load_weights\") and not getattr(module, \"_weights_removed\", False):\n                        module.post_load_weights()\n                moe_load_balancer = getattr(self.engine.model_engine, \"moe_load_balancer\", None)\n                if isinstance(moe_load_balancer, MoeLoadBalancer):\n                    moe_load_balancer.register_weight_slots_after_to_cuda()\n                    logger.info(\"moe_load_balancer finalizing model...\")\n                    moe_load_balancer.finalize_model()\n                    logger.info(\"moe_load_balancer finalize model done\")\n                self.engine.reset_prefix_cache()\n                delattr(self.engine.model_engine.model, \"first_pre_reload_weights\")\n\n        except Exception as e:\n            logger.error(\"Encountered an error in update_weights\")\n            raise e\n"
  },
  {
    "path": "verl/workers/rollout/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport asyncio\nimport logging\n\nimport numpy as np\nimport uvicorn\nfrom fastapi import FastAPI\n\nlogger = logging.getLogger(__file__)\n\n\ndef get_max_position_embeddings(hf_config) -> int:\n    max_len = getattr(hf_config, \"max_position_embeddings\", None)\n    if max_len is None:\n        text_config = getattr(hf_config, \"text_config\", None)\n        if text_config is not None:\n            max_len = getattr(text_config, \"max_position_embeddings\", None)\n\n    if max_len is None:\n        raise ValueError(\"max_position_embeddings not found in HFModelConfig!\")\n    return int(max_len)\n\n\nclass _UvicornServerAutoPort(uvicorn.Server):\n    \"\"\"Uvicorn Server that reports the system-assigned port when port=0.\"\"\"\n\n    def __init__(self, config: uvicorn.Config) -> None:\n        super().__init__(config)\n        self.actual_port: int | None = None\n        self._startup_done: asyncio.Event = asyncio.Event()\n\n    async def startup(self, sockets=None) -> None:\n        try:\n            await super().startup(sockets=sockets)\n            if self.servers and self.config.port == 0:\n                sock = self.servers[0].sockets[0]\n                self.actual_port = sock.getsockname()[1]\n            else:\n                self.actual_port = self.config.port\n        finally:\n            self._startup_done.set()\n\n    async def get_port(self) -> int | None:\n        await self._startup_done.wait()\n        return self.actual_port\n\n\nasync def run_uvicorn(app: FastAPI, server_args, server_address) -> tuple[int, asyncio.Task]:\n    app.server_args = server_args\n    config = uvicorn.Config(app, host=server_address, port=0, log_level=\"warning\")\n    server = _UvicornServerAutoPort(config)\n    server_task = asyncio.create_task(server.serve())\n    server_port = await server.get_port()\n    if server_port is None:\n        # server.startup() failed. await the task to re-raise exception from server.serve()\n        await server_task\n\n        # Fails on unexpected situation.\n        raise RuntimeError(\"Unexpected: HTTP server started without reporting listened port\")\n    logger.info(f\"HTTP server started on port {server_port}\")\n    return server_port, server_task\n\n\nasync def ensure_async_iterator(iterable):\n    \"\"\"Convert an iterable to an async iterator.\"\"\"\n    if hasattr(iterable, \"__aiter__\"):\n        async for item in iterable:\n            yield item\n    else:\n        for item in iterable:\n            yield item\n\n\ndef qwen2_5_vl_dedup_image_tokens(prompt_ids: list[int], processor):\n    \"\"\"Deduplicate consecutive image tokens in prompt_ids for Qwen2.5-VL, since vLLM will replicate the\n    <|image_pad|> and <|video_pad|> token by image_data.\n    For example,\n    ```\n    <|vision_start|><|image_pad|><|image_pad|>...<|image_pad|><|vision_end|>\n    =>\n    <|vision_start|><|image_pad|><|vision_end|>\n    ```\n    \"\"\"\n    if processor is not None and \"Qwen2VLImageProcessor\" in processor.image_processor.__class__.__name__:\n        prompt_ids = np.array(prompt_ids)\n        mask = np.ones(len(prompt_ids), dtype=bool)\n        is_value = (prompt_ids == processor.image_token_id) | (prompt_ids == processor.video_token_id)\n        mask[1:] &= ~(is_value[1:] & is_value[:-1])\n        return prompt_ids[mask].tolist()\n    else:\n        return prompt_ids\n"
  },
  {
    "path": "verl/workers/rollout/vllm_rollout/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport os\nfrom importlib.metadata import PackageNotFoundError, version\n\nfrom .vllm_rollout import ServerAdapter  # noqa: F401\n\n\ndef get_version(pkg):\n    try:\n        return version(pkg)\n    except PackageNotFoundError:\n        return None\n\n\nvllm_package_name = \"vllm\"\nvllm_package_version = get_version(vllm_package_name)\nif vllm_package_version is None:\n    raise PackageNotFoundError(\n        \"To use vllm rollout, please ensure the 'vllm' package is properly installed. See \"\n        \"https://verl.readthedocs.io/en/latest/start/install.html for more details\"\n    )\n\nif \"ROCM_PATH\" in os.environ:\n    import re\n\n    match = re.match(r\"(\\d+\\.\\d+\\.?\\d*)\", vllm_package_version)\n    if match:\n        vllm_package_version = match.group(1)\n    else:\n        raise ValueError(f\"Warning: Could not parse version format: {vllm_package_version}\")\n"
  },
  {
    "path": "verl/workers/rollout/vllm_rollout/bucketed_weight_transfer.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nBucketed weight transfer via ZMQ + IPC (or shared memory fallback).\n\nNot recommended depending on vllm for this file.\n\"\"\"\n\nimport gc\nimport logging\nimport os\nfrom multiprocessing import shared_memory\nfrom typing import Callable, TypedDict\n\nimport torch\nimport zmq\nfrom torch.multiprocessing.reductions import reduce_tensor\n\nfrom verl.utils.device import get_device_id, get_device_name, get_torch_device\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\n\n\nclass TensorMetadata(TypedDict):\n    name: str\n    shape: torch.Size\n    dtype: torch.dtype\n    offset: int\n\n\n# copy from https://github.com/vllm-project/vllm/blob/main/examples/offline_inference/rlhf_utils.py\ndef rebuild_ipc(handle: tuple[Callable, tuple], device_id: int | None = None) -> torch.Tensor:\n    func, args = handle\n    list_args = list(args)\n    if device_id is not None:\n        # the key is to change device id to the current device id\n        # in case two processes have different CUDA_VISIBLE_DEVICES\n        list_args[6] = device_id\n    buffer = func(*list_args)\n    return buffer\n\n\ndef create_shared_memory(size: int, name: str):\n    \"\"\"Create shared memory for weight transfer. If already exists, attach to it.\"\"\"\n    try:\n        shm = shared_memory.SharedMemory(name=name, create=True, size=size)\n    except FileExistsError:\n        shm = shared_memory.SharedMemory(name=name)\n        assert shm.size >= size, f\"Stale shm segment '{name}': expected {size} bytes, got {shm.size}\"\n    return shm\n\n\ndef rebuild_shared_memory(name: str, size: int, dtype=torch.uint8):\n    \"\"\"Rebuild tensor from shared memory.\"\"\"\n    shm = shared_memory.SharedMemory(name=name)\n    tensor = torch.frombuffer(shm.buf[:size], dtype=dtype)\n\n    return tensor, shm\n\n\nclass BucketedWeightSender:\n    \"\"\"\n    Send model weights via bucketed IPC transfer over ZMQ.\n\n    Packs weight tensors into a fixed-size communication buffer and sends them\n    in buckets to the receiver. Supports CUDA IPC and shared memory fallback.\n\n    Args:\n        zmq_handle: ZMQ IPC socket path (e.g., \"ipc:///tmp/rl-colocate-zmq-<uuid>.sock\")\n        bucket_size_mb: Communication buffer size in MB\n        use_shm: Use shared memory instead of CUDA IPC (for NPU compatibility)\n    \"\"\"\n\n    def __init__(\n        self,\n        zmq_handle: str,\n        bucket_size_mb: int = 512,\n        use_shm: bool = False,\n    ):\n        self.zmq_handle = zmq_handle\n        self.bucket_size_mb = bucket_size_mb\n        self.bucket_size = int(bucket_size_mb) << 20\n        self.use_shm = use_shm\n\n        self.zmq_context = zmq.Context.instance()\n        self.socket = None\n        self.buffer = None\n        self.shm = None\n\n    async def async_send_weights(self, weights):\n        \"\"\"\n        Send weights to the receiver. Accepts a sync generator or async iterator.\n\n        Args:\n            weights: Generator or async iterator yielding (name, tensor) pairs\n        \"\"\"\n        from verl.workers.rollout.utils import ensure_async_iterator\n\n        try:\n            self._init_socket()\n            self._init_buffer()\n\n            # send bucket weights\n            offset = 0\n            bucket_meta: dict[str, TensorMetadata] = {}\n            # dtype = PrecisionType.to_dtype(self.config.dtype)\n            async for name, weight in ensure_async_iterator(weights):\n                # model parameters are in fp32 full precision\n                # (vermouth1992) we should not force cast weight here because some parameters\n                # (such as moe gate) have to keep fp32 precision. If a weight is bf16 in the rollout side,\n                # the rollout should automatically cast on demand. However, this would incur a higher weight\n                # transfer volume.\n                # weight = weight.to(dtype, non_blocking=True)\n\n                # fill the tensor bucket\n                if offset + weight.nbytes > self.bucket_size:\n                    get_torch_device().synchronize()\n                    self.socket.send_pyobj({\"bucket_meta\": bucket_meta, \"is_last\": False})\n                    self.socket.recv()\n                    bucket_meta = {}\n                    offset = 0\n\n                # TODO: slice embedding layer weight into chunks\n                assert offset + weight.nbytes <= self.bucket_size, (\n                    f\"Weight {name}({weight.shape}, {weight.dtype}) is too large to fit in the bucket.\"\n                    f\"Please increase rollout.update_weights_bucket_megabytes({self.bucket_size_mb} MB).\"\n                )\n                bucket_meta[name] = {\n                    \"name\": name,\n                    \"shape\": weight.shape,\n                    \"dtype\": weight.dtype,\n                    \"offset\": offset,\n                }\n                self.buffer[offset : offset + weight.nbytes].copy_(weight.view(-1).view(torch.uint8), non_blocking=True)\n                offset += weight.nbytes\n\n            # send the last bucket\n            get_torch_device().synchronize()\n            self.socket.send_pyobj({\"bucket_meta\": bucket_meta, \"is_last\": True})\n            self.socket.recv()\n        finally:\n            self._cleanup()\n\n    def _init_socket(self):\n        \"\"\"Initialize ZMQ REQ socket and bind.\"\"\"\n        self.socket = self.zmq_context.socket(zmq.REQ)\n        self.socket.bind(self.zmq_handle)\n\n    def _init_buffer(self):\n        \"\"\"build communication buffer\"\"\"\n        buffer, shm = None, None\n        if not self.use_shm:\n            buffer = torch.empty(self.bucket_size, dtype=torch.uint8, device=f\"{get_device_name()}:{get_device_id()}\")\n            handle = reduce_tensor(buffer)\n            self.socket.send_pyobj(handle)\n        else:\n            import uuid\n\n            # Create unique name for shared memory\n            shm_name = f\"verl_weights_{uuid.uuid4().hex}\"\n            shm = create_shared_memory(self.bucket_size, shm_name)\n            buffer = torch.frombuffer(shm.buf, dtype=torch.uint8)\n\n            comm_metadata = {\"name\": shm_name, \"size\": self.bucket_size}\n            self.socket.send_pyobj(comm_metadata)\n\n        self.socket.recv()\n        self.buffer = buffer\n        self.shm = shm\n\n    def _cleanup(self):\n        \"\"\"clean up\"\"\"\n        if self.socket is not None:\n            self.socket.close()\n            self.socket = None\n        del self.buffer\n        self.buffer = None\n        if self.shm is not None:\n            self.shm.close()\n            self.shm.unlink()\n            del self.shm\n            self.shm = None\n        gc.collect()\n        get_torch_device().ipc_collect()\n        get_torch_device().empty_cache()\n\n\nclass BucketedWeightReceiver:\n    \"\"\"\n    Receive model weights via bucketed IPC transfer over ZMQ.\n\n    Receives weight tensors from BucketedWeightSender and passes each\n    bucket to a callback for processing (e.g., loading into the model).\n\n    Args:\n        zmq_handle: ZMQ IPC socket path (must match sender)\n        device: Target device for received tensors\n        use_shm: Use shared memory instead of CUDA IPC\n    \"\"\"\n\n    def __init__(\n        self,\n        zmq_handle: str,\n        device: torch.device,\n        use_shm: bool = False,\n    ):\n        self.zmq_handle = zmq_handle\n        self.device = device\n        self.use_shm = use_shm\n\n        self.zmq_context = zmq.Context.instance()\n        self.socket = None\n        self.buffer = None\n        self.shm = None\n\n    def receive_weights(self, on_bucket_received: callable):\n        \"\"\"\n        Receive weights from sender and process each bucket via callback.\n\n        Args:\n            on_bucket_received: Callback function(weights: list[(name, tensor)]) called per bucket.\n        \"\"\"\n        try:\n            self._init_socket()\n            self._init_buffer()\n\n            # receive bucket and update weights\n            while True:\n                metadata = self.socket.recv_pyobj()\n                weights, tensor = [], None\n                for name, meta in metadata[\"bucket_meta\"].items():\n                    shape, dtype, offset = meta[\"shape\"], meta[\"dtype\"], meta[\"offset\"]\n                    size = dtype.itemsize * shape.numel()\n                    # NOTE: we need to clone the tensor to release CUDA IPC memory\n                    # but for shared memory, it's not necessary and if we do clone,\n                    # it will cause extra memory copy overhead and slow down the process.\n                    tensor = self.buffer[offset : offset + size].view(dtype=dtype).view(shape)\n                    if not self.use_shm:\n                        tensor = tensor.clone()\n                    else:\n                        tensor = tensor.to(self.device)\n                    weights.append((name, tensor))\n                get_torch_device().synchronize()\n                self.socket.send(b\"\")\n                on_bucket_received(weights)\n                del weights, tensor\n                if metadata[\"is_last\"]:\n                    break\n        finally:\n            self._cleanup()\n\n    def _init_socket(self):\n        \"\"\"Initialize ZMQ REP socket and connect.\"\"\"\n        self.socket = self.zmq_context.socket(zmq.REP)\n        self.socket.connect(self.zmq_handle)\n\n    def _init_buffer(self):\n        \"\"\"Receive and rebuild communication buffer from sender.\"\"\"\n        comm_metadata = self.socket.recv_pyobj()\n        buffer, shm = None, None\n        if not self.use_shm:\n            handle = comm_metadata\n            buffer = rebuild_ipc(handle, self.device.index)\n            assert buffer.dtype == torch.uint8\n        else:\n            shm_name = comm_metadata[\"name\"]\n            shm_size = comm_metadata[\"size\"]\n            buffer, shm = rebuild_shared_memory(shm_name, shm_size, dtype=torch.uint8)\n        self.socket.send(b\"\")\n        self.buffer = buffer\n        self.shm = shm\n\n    def _cleanup(self):\n        \"\"\"clean up\"\"\"\n        if self.socket is not None:\n            self.socket.close()\n            self.socket = None\n        # Synchronize before releasing the buffer to ensure all async ops\n        # referencing it (e.g. clone, .to()) have completed.\n        get_torch_device().synchronize()\n        del self.buffer\n        self.buffer = None\n        if self.shm is not None:\n            self.shm.close()\n            del self.shm\n            self.shm = None\n        gc.collect()\n        get_torch_device().ipc_collect()\n        get_torch_device().empty_cache()\n"
  },
  {
    "path": "verl/workers/rollout/vllm_rollout/utils.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport ctypes\nimport json\nimport logging\nimport os\nimport platform\nimport signal\nimport threading\nfrom types import MethodType\nfrom typing import Any, Literal, get_args\n\nimport torch\n\nfrom verl.utils.device import is_npu_available\nfrom verl.utils.vllm import TensorLoRARequest, VLLMHijack\nfrom verl.utils.vllm.patch import patch_vllm_moe_model_weight_loader\nfrom verl.utils.vllm.vllm_fp8_utils import apply_vllm_fp8_patches, is_fp8_model, load_quanted_weights\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"WARN\"))\n\n# magic numbers that ensure we are using the same LoRA adapter during the rollout and training process\nVLLM_LORA_INT_ID = 123\nVLLM_LORA_NAME = \"123\"\nVLLM_LORA_PATH = \"simon_lora_path\"\n\nVLLM_ASCEND_REQUIRED_ENV_VARS = {\"VLLM_ALL2ALL_BACKEND\": \"flashinfer_all2allv\", \"VLLM_ASCEND_ENABLE_NZ\": \"0\"}\n\n\ndef set_death_signal():\n    \"\"\"Kill the current process when the parent process exits.\"\"\"\n    if platform.system() != \"Linux\":\n        return\n    libc = ctypes.CDLL(\"libc.so.6\")\n    libc.prctl(1, signal.SIGKILL)\n    if os.getppid() == 1:\n        os.kill(os.getpid(), signal.SIGKILL)\n\n\ndef get_device_uuid(device_id: int) -> str:\n    from vllm.platforms import current_platform\n\n    # Convert torch.npu.current_device to its corresponding ASCEND_RT_VISIBLE_DEVICES.\n    if is_npu_available:\n        if os.getenv(\"ASCEND_RT_VISIBLE_DEVICES\") is not None:\n            npu_visible_devices = os.environ[\"ASCEND_RT_VISIBLE_DEVICES\"].split(\",\")\n            assert device_id < len(npu_visible_devices), f\"device_id {device_id} must less than {npu_visible_devices}\"\n            return \"NPU-\" + npu_visible_devices[device_id]\n        else:\n            return f\"NPU-{device_id}\"\n    else:\n        return current_platform.get_device_uuid(device_id)\n\n\ndef get_vllm_max_lora_rank(lora_rank: int):\n    \"\"\"\n    For vLLM, automatically adjusts the `max_lora_rank` to the nearest allowed value.\n    The allowed values are retrieved from vLLM's MaxLoRARanks type definition.\n    \"\"\"\n    assert lora_rank > 0, f\"lora_rank must be greater than 0, get {lora_rank}\"\n\n    try:\n        from vllm.config.lora import MaxLoRARanks\n    except Exception:\n        # FIXME: migrate vllm version https://github.com/vllm-project/vllm/blob/main/vllm/config/lora.py#L25\n        MaxLoRARanks = Literal[1, 8, 16, 32, 64, 128, 256, 320, 512]\n\n    vllm_max_lora_ranks = sorted(get_args(MaxLoRARanks))\n    if lora_rank > vllm_max_lora_ranks[-1]:\n        raise ValueError(f\"lora_rank must be less than or equal to {vllm_max_lora_ranks[-1]}, but got {lora_rank}\")\n\n    for rank in vllm_max_lora_ranks:\n        if lora_rank <= rank:\n            return rank\n\n\n# https://github.com/vllm-project/vllm/issues/13175\ndef monkey_patch_compute_logits(model, vocab_size: int):\n    original_compute_logits = model.compute_logits\n\n    def compute_logits(\n        self,\n        *args,\n        **kwargs,\n    ) -> torch.Tensor:\n        logits = original_compute_logits(*args, **kwargs)\n        logits[..., vocab_size:] = float(\"-inf\")\n        return logits\n\n    model.compute_logits = MethodType(compute_logits, model)\n\n\nclass vLLMColocateWorkerExtension:\n    \"\"\"\n    The class for vLLM's worker to inherit from, in the colocate setting.\n    By defining an extension class, the code can work no matter what is\n    the underlying worker class. This way, the code can be compatible\n    with both vLLM V0 and V1.\n    NOTE: we define this class in a separate module, and the main module\n    should pass the full qualified name as `worker_extension_cls` argument.\n\n    Feature support:\n    1. LoRA\n    2. Online FP8 quantization\n    \"\"\"\n\n    def __new__(cls, **kwargs):\n        set_death_signal()\n\n        # 1. patch for Lora\n        VLLMHijack.hijack()\n        # 2. patch online fp8 quant\n        if os.environ.get(\"VERL_VLLM_FP8_QUANT_ENABLED\", \"0\") == \"1\":\n            apply_vllm_fp8_patches()\n        # 3. patch QAT (compressed-tensors NVFP4) for dynamic weight loading\n        vllm_config = kwargs.get(\"vllm_config\")\n        quant_config = getattr(vllm_config, \"quant_config\", None) if vllm_config else None\n        _is_qat_model = getattr(quant_config, \"quant_format\", None) == \"nvfp4-pack-quantized\"\n        if _is_qat_model:\n            from verl.utils.qat import apply_qat_patches\n\n            apply_qat_patches()\n            logger.info(\"Applied QAT patches in vLLM worker subprocess\")\n\n        # TODO: For ascend NPU, when the corresponding vllm-ascend version is upgraded to v0.13.0,\n        # please remove the VLLM_ASCEND_REQUIRED_ENV_VARS variable replacement action.\n        # This is only a fix for vllm version < v0.13.0.\n        if is_npu_available:\n            for k in VLLM_ASCEND_REQUIRED_ENV_VARS:\n                if k not in os.environ:\n                    os.environ[k] = VLLM_ASCEND_REQUIRED_ENV_VARS[k]\n\n        instance = super().__new__(cls)\n        instance._is_qat_model = _is_qat_model\n        return instance\n\n    def monkey_patch_model(self, vocab_size: int):\n        # patch compute_logits to avoid sampling OOV token\n        monkey_patch_compute_logits(self.model_runner.model, vocab_size)\n        # patch weight loader to support MoE model\n        patch_vllm_moe_model_weight_loader(self.model_runner.model)\n\n    def update_weights_from_ipc(self, peft_config: dict = None, base_sync_done=False, use_shm: bool = False):\n        \"\"\"Update the weights of the rollout model.\"\"\"\n        from vllm.platforms import current_platform\n\n        from verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightReceiver\n\n        if current_platform.device_type == \"npu\" and self.device is None:\n            self.device = torch.device(f\"npu:{self.local_rank}\")\n\n        # In async mode, make sure the old lora is removed before adding the new one\n        if peft_config and base_sync_done:\n            self.remove_lora(VLLM_LORA_INT_ID)\n\n        use_standard_weight_load = not (peft_config and base_sync_done) and not is_fp8_model(\n            self.model_runner.vllm_config\n        )\n\n        if self._is_qat_model:\n            # QAT: Prepare for weight loading BEFORE receiving any buckets\n            from verl.utils.qat import prepare_qat_for_load_weights\n\n            prepare_qat_for_load_weights(self.model_runner.model, device=self.device)\n            logger.info(\"QAT: prepare_qat_for_load_weights completed\")\n        elif use_standard_weight_load:\n            # Re-apply here because async IPC weight sync can happen long after init and lose MoE weight_loader attrs.\n            patch_vllm_moe_model_weight_loader(self.model_runner.model)\n\n        assert self.device is not None\n        receiver = BucketedWeightReceiver(\n            zmq_handle=self._get_zmq_handle(),\n            device=self.device,\n            use_shm=use_shm,\n        )\n        receiver.receive_weights(\n            on_bucket_received=lambda weights: self._update_weights(\n                weights, peft_config=peft_config, base_sync_done=base_sync_done\n            )\n        )\n\n        if self._is_qat_model:\n            # QAT: call process_weights_after_loading AFTER all buckets are received\n            from verl.utils.qat import manual_process_weights_after_loading\n\n            manual_process_weights_after_loading(self.model_runner.model)\n            logger.info(\"QAT: process_weights_after_loading completed\")\n        elif use_standard_weight_load:\n            # Some post-load transforms are non-idempotent; run once after all buckets.\n            from vllm.model_executor.model_loader.utils import process_weights_after_loading\n\n            model = self.model_runner.model\n            model_config = self.model_runner.vllm_config.model_config\n            process_weights_after_loading(model, model_config, self.device)\n\n    def _update_weights(self, weights: list[tuple[str, torch.Tensor]], peft_config: dict, base_sync_done: bool):\n        if peft_config and base_sync_done:\n            weights = dict(weights)\n            lora_request = TensorLoRARequest(\n                lora_name=VLLM_LORA_NAME,\n                lora_int_id=VLLM_LORA_INT_ID,\n                lora_path=VLLM_LORA_PATH,\n                peft_config=peft_config,\n                lora_tensors=weights,\n            )\n            self.add_lora(lora_request)\n            logger.info(f\"vLLM load weights, loaded_params: {len(weights)}\")\n        else:\n            # Add the FP8 related logic here as sharding manager has been deprecated.\n            # Check if FP8 quantization is enabled and apply appropriate weight loading\n            if is_fp8_model(self.model_runner.vllm_config):\n                logger.info(f\"FP8 model detected (async): {self.model_runner.vllm_config.quant_config}\")\n                # Convert bf16 weights to fp8 format before loading\n                loaded_params = load_quanted_weights(weights, self.model_runner)\n                logger.info(f\"FP8 weights loaded (async), loaded_params: {len(loaded_params)}\")\n            else:\n                logger.info(\"Loading standard weights (non-FP8, async)\")\n                self.model_runner.model.load_weights(weights)\n\n    def _get_zmq_handle(self) -> str:\n        \"\"\"Get ZMQ handle for communication.\"\"\"\n        if not hasattr(self, \"device_uuid\") or not self.device_uuid:\n            self.device_uuid = get_device_uuid(self.device.index)\n        return f\"ipc:///tmp/rl-colocate-zmq-{self.device_uuid}.sock\"\n\n\nclass SuppressSignalInThread:\n    def __enter__(self):\n        self.original_signal = signal.signal\n\n        def no_op_signal(sig, action):\n            if threading.current_thread() is not threading.main_thread():\n                print(f\"Ignored signal {sig} in thread {threading.current_thread().name}\")\n                return\n            return self.original_signal(sig, action)\n\n        signal.signal = no_op_signal\n        return self\n\n    def __exit__(self, exc_type, exc_val, exc_tb):\n        signal.signal = self.original_signal\n\n\ndef build_cli_args_from_config(config: dict[str, Any]) -> list[str]:\n    \"\"\"\n    Convert a config dictionary to CLI arguments for vLLM server.\n\n    Handles different value types appropriately:\n    - None: skipped\n    - bool True: adds '--key'\n    - bool False: skipped\n    - list: expands to '--key item1 item2 ...'\n    - empty list: skipped (vLLM uses nargs=\"+\" which requires at least one value)\n    - dict: JSON serialized\n    - other: string converted\n\n    Args:\n        config: Dictionary of configuration key-value pairs\n\n    Returns:\n        List of CLI argument strings\n    \"\"\"\n    cli_args = []\n    for k, v in config.items():\n        if v is None:\n            continue\n        if isinstance(v, bool):\n            if v:\n                cli_args.append(f\"--{k}\")\n        elif isinstance(v, list):\n            if not v:\n                # Skip empty lists - vLLM uses nargs=\"+\" which requires at least one value\n                continue\n            # Lists need to be expanded as multiple separate arguments\n            # e.g., --cuda-graph-sizes 1 2 4 8 becomes ['--cuda-graph-sizes', '1', '2', '4', '8']\n            cli_args.append(f\"--{k}\")\n            cli_args.extend([str(item) for item in v])\n        else:\n            cli_args.append(f\"--{k}\")\n            # Use json.dumps for dict to ensure valid JSON format\n            cli_args.append(json.dumps(v) if isinstance(v, dict) else str(v))\n    return cli_args\n"
  },
  {
    "path": "verl/workers/rollout/vllm_rollout/vllm_async_server.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\nimport argparse\nimport asyncio\nimport inspect\nimport json\nimport logging\nimport os\nfrom pprint import pprint\nfrom typing import Any, Callable, Optional\n\nimport ray\nimport vllm.entrypoints.cli.serve\nfrom packaging import version\nfrom ray.actor import ActorHandle\nfrom vllm import SamplingParams\nfrom vllm.engine.arg_utils import AsyncEngineArgs\nfrom vllm.entrypoints.cli.serve import run_headless\nfrom vllm.entrypoints.openai.api_server import build_app, init_app_state\nfrom vllm.inputs import TokensPrompt\nfrom vllm.lora.request import LoRARequest\nfrom vllm.outputs import RequestOutput\nfrom vllm.usage.usage_lib import UsageContext\nfrom vllm.v1.engine.async_llm import AsyncLLM\n\nfrom verl.utils.config import omega_conf_to_dataclass\nfrom verl.utils.device import get_resource_name, get_visible_devices_keyword, is_npu_available, is_torch_npu_available\nfrom verl.utils.net_utils import get_free_port, is_valid_ipv6_address\nfrom verl.utils.profiler import DistProfiler, build_vllm_profiler_args\nfrom verl.utils.tokenizer import normalize_token_ids\nfrom verl.utils.vllm.vllm_fp8_utils import apply_vllm_fp8_patches\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.replica import RolloutMode, RolloutReplica, TokenOutput\nfrom verl.workers.rollout.utils import get_max_position_embeddings, qwen2_5_vl_dedup_image_tokens, run_uvicorn\nfrom verl.workers.rollout.vllm_rollout.utils import (\n    VLLM_LORA_INT_ID,\n    VLLM_LORA_NAME,\n    VLLM_LORA_PATH,\n    SuppressSignalInThread,\n    build_cli_args_from_config,\n    get_vllm_max_lora_rank,\n)\n\n_VLLM_VERSION = version.parse(vllm.__version__)\n\nif _VLLM_VERSION > version.parse(\"0.11.0\"):\n    from vllm.utils.argparse_utils import FlexibleArgumentParser\n\n    if _VLLM_VERSION == version.parse(\"0.12.0\"):\n        from vllm.entrypoints.harmony_utils import get_encoding\n\n    elif _VLLM_VERSION >= version.parse(\"0.13.0\"):\n        from vllm.entrypoints.openai.parser.harmony_utils import get_encoding\n\n    else:\n        get_encoding = None\n\n    if get_encoding is not None and os.getenv(\"VERL_USE_GPT_OSS\", \"0\") == \"1\":\n        get_encoding()\nelse:\n    from vllm.utils import FlexibleArgumentParser\n\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(logging.INFO)\n\n\nclass vLLMHttpServer:\n    \"\"\"vLLM http server in single node, this is equivalent to launch server with command line:\n    ```\n    vllm serve --tensor-parallel-size=8 ...\n    ```\n    \"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        rollout_mode: RolloutMode,\n        workers: list[ActorHandle],\n        replica_rank: int,\n        node_rank: int,\n        gpus_per_node: int,\n        nnodes: int,\n        cuda_visible_devices: str,\n    ):\n        \"\"\"\n        Args:\n            config (RolloutConfig): full config.\n            model_config (HFModelConfig): model config.\n            rollout_mode (RolloutMode): rollout mode.\n            replica_rank (int): replica rank, a replica may contain multiple nodes.\n            node_rank (int): node rank.\n            gpus_per_node (int): number of gpus per node.\n            nnodes (int): number of nodes.\n            cuda_visible_devices (str): cuda visible devices.\n        \"\"\"\n        os.environ[get_visible_devices_keyword()] = cuda_visible_devices\n\n        self.config: RolloutConfig = omega_conf_to_dataclass(config)\n        self.model_config: HFModelConfig = omega_conf_to_dataclass(model_config, dataclass_type=HFModelConfig)\n        max_position_embeddings = get_max_position_embeddings(self.model_config.hf_config)\n        if self.config.max_model_len is None:\n            self.config.max_model_len = max_position_embeddings\n        else:\n            if self.config.max_model_len > max_position_embeddings:\n                raise ValueError(\n                    f\"max_model_len ({self.config.max_model_len}) should be less than or equal to \"\n                    f\"max_position_embeddings ({max_position_embeddings})\"\n                )\n\n        self.rollout_mode = rollout_mode\n        self.workers = workers\n\n        self.replica_rank = replica_rank\n        self.node_rank = node_rank\n        self.gpus_per_node = gpus_per_node\n        self.nnodes = nnodes\n        # model weights version, set by ServerAdapter when update weights.\n        self.global_steps = None\n\n        if self.rollout_mode != RolloutMode.HYBRID and self.config.load_format == \"dummy\":\n            logger.warning(f\"rollout mode is {self.rollout_mode}, load_format is dummy, set to auto\")\n            self.config.load_format = \"auto\"\n\n        # used for http server\n        self._server_address = ray.util.get_node_ip_address().strip(\"[]\")\n        self._server_port = None\n\n        # used for controlling vllm server profiler\n        profiler_config = self.config.profiler\n        tool_config = None\n        if profiler_config is not None:\n            if profiler_config.tool in [\"torch\", \"npu\"]:\n                tool_config = omega_conf_to_dataclass((profiler_config.tool_config or {}).get(profiler_config.tool))\n            else:\n                logger.warning(f\"agent loop only support torch and npu profiler, got {profiler_config.tool}\")\n                profiler_config = None\n        self.profiler_controller = DistProfiler(self.replica_rank, config=profiler_config, tool_config=tool_config)\n\n        # used for data parallel: --data-parallel-address, --data-parallel-rpc-port\n        if self.node_rank == 0:\n            self._master_address = self._server_address\n            # used for torch.distributed.init_process_group\n            self._master_port, self._master_sock = get_free_port(self._server_address, with_alive_sock=True)\n            # used for data parallel: --data-parallel-address, --data-parallel-rpc-port\n            self._dp_rpc_port, self._dp_rpc_sock = get_free_port(self._server_address, with_alive_sock=True)\n            self._dp_master_port, self._dp_master_sock = get_free_port(self._server_address, with_alive_sock=True)\n        else:\n            self._master_address = None\n            self._master_port = None\n            self._dp_rpc_port = None\n            self._dp_master_port = None\n\n        logger.info(\n            f\"vLLMHttpServer, replica_rank: {self.replica_rank}, node_rank: {self.node_rank}, \"\n            f\"{get_visible_devices_keyword()}: {cuda_visible_devices}, \"\n            f\"master_address: {self._master_address}, master_port: {self._master_port}, \"\n            f\"data_parallel_rpc_port: {self._dp_rpc_port}, data_parallel_master_port: {self._dp_master_port}\"\n        )\n\n    def get_master_address(self):\n        \"\"\"Get master address and port for data parallel.\n        Returns:\n            tuple: (master_address, master_port, dp_rpc_port)\n        \"\"\"\n        return self._master_address, self._master_port, self._dp_rpc_port\n\n    def get_server_address(self):\n        \"\"\"Get http server address and port.\"\"\"\n        assert self._server_port is not None, \"http server is not launched, port is None\"\n        return self._server_address, self._server_port\n\n    @property\n    def lora_as_adapter(self) -> bool:\n        return (\n            self.model_config.lora_rank > 0 or self.model_config.lora.get(\"rank\", 0) > 0\n        ) and not self.model_config.lora.get(\"merge\", False)\n\n    async def collective_rpc(\n        self,\n        method: str | Callable,\n        timeout: float | None = None,\n        args: tuple = (),\n        kwargs: dict[str, Any] | None = None,\n    ):\n        await self.engine.collective_rpc(\n            method=method,\n            timeout=timeout,\n            args=args,\n            kwargs=kwargs,\n        )\n\n    async def launch_server(self, master_address: str = None, master_port: int = None, dp_rpc_port: int = None):\n        if self.node_rank != 0:\n            assert master_address and master_port and dp_rpc_port, (\n                \"non-master node should provide master_address, master_port and dp_rpc_port\"\n            )\n            self._master_address = master_address\n            self._master_port = master_port\n            self._dp_rpc_port = dp_rpc_port\n\n        # 1. setup vllm serve cli args\n        engine_kwargs = self.config.get(\"engine_kwargs\", {}).get(\"vllm\", {}) or {}\n        engine_kwargs = {key: val for key, val in engine_kwargs.items() if val is not None}\n        if self.config.get(\"limit_images\", None):  # support for multi-image data\n            engine_kwargs[\"limit_mm_per_prompt\"] = {\"image\": self.config.get(\"limit_images\")}\n        if self.config.cudagraph_capture_sizes:\n            engine_kwargs[\"cuda_graph_sizes\"] = self.config.cudagraph_capture_sizes\n\n        # Override default generation config from hugging face model config,\n        # user can still override them by passing kwargs in each request.\n        override_generation_config = dict(\n            temperature=self.config.temperature,\n            top_k=self.config.top_k,\n            top_p=self.config.top_p,\n            repetition_penalty=1.0,\n            max_new_tokens=self.config.response_length,\n        )\n        logger.info(f\"override_generation_config: {override_generation_config}\")\n\n        logger.info(f\"enable_sleep_mode: {self.config.enable_sleep_mode}\")\n        if not self.config.enable_sleep_mode:\n            from verl.utils.device import set_expandable_segments\n\n            set_expandable_segments(True)\n\n        quantization = self.config.quantization\n        hf_overrides = {}\n        if is_torch_npu_available(check_device=False):\n            from verl.utils.vllm.npu_vllm_patch import check_vllm_ascend_before_server_launch\n\n            check_vllm_ascend_before_server_launch()\n        # Handle QAT (Quantization-Aware Training) configuration\n        qat_config_dict = getattr(self.config, \"qat\", {}) or {}\n        if qat_config_dict.get(\"enable\", False):\n            # QAT uses compressed-tensors quantization, apply patches for dynamic weight loading\n            from verl.utils.qat import QATConfig, apply_qat_patches, load_quantization_config\n\n            apply_qat_patches()\n\n            # Load quantization config from JSON file\n            qat_config = QATConfig(**qat_config_dict)\n            quantization_config_dict = load_quantization_config(qat_config)\n            hf_overrides[\"quantization_config\"] = quantization_config_dict\n            quantization = \"compressed-tensors\"\n\n            logger.info(\"QAT quantization config injected to vLLM async server\")\n        elif quantization is not None:\n            # Handle other quantization methods (fp8, torchao)\n            _SUPPORTED_QUANTIZATION = [\"fp8\", \"torchao\"]\n            if quantization not in _SUPPORTED_QUANTIZATION:\n                raise ValueError(f\"Currently only support {_SUPPORTED_QUANTIZATION} quantization, got: {quantization}\")\n\n            if quantization == \"fp8\":\n                # Ignore MoE router layers for FP8 quantization\n                all_mlp_gate_layers = []\n                for layer in range(self.model_config.hf_config.num_hidden_layers):\n                    all_mlp_gate_layers.append(f\"model.layers.{layer}.mlp.gate\")\n\n                FP8_BLOCK_QUANT_KWARGS = {\n                    \"activation_scheme\": \"dynamic\",\n                    \"fmt\": \"e4m3\",\n                    \"quant_method\": \"fp8\",\n                    \"weight_block_size\": [128, 128],\n                    \"ignored_layers\": all_mlp_gate_layers,\n                }\n                hf_overrides[\"quantization_config\"] = dict(FP8_BLOCK_QUANT_KWARGS)\n                # Apply vllm fp8 patches\n                # Will remove the patch after vllm support on-the-fly quant for rollout natively.\n                apply_vllm_fp8_patches()\n                # for subprocesses patching\n                os.environ[\"VERL_VLLM_FP8_QUANT_ENABLED\"] = \"1\"\n\n        if quantization is not None and self.config.quantization_config_file is not None:\n            hf_overrides[\"quantization_config_file\"] = self.config.quantization_config_file\n\n        compilation_config = engine_kwargs.pop(\"compilation_config\", None) or {}\n        if isinstance(compilation_config, str):\n            compilation_config = json.loads(compilation_config)\n        compilation_config.setdefault(\"cudagraph_mode\", \"FULL_AND_PIECEWISE\")\n\n        # FULL cuda graph is not yet supported with DCP, downgrade to PIECEWISE\n        dcp_size = engine_kwargs.get(\"decode_context_parallel_size\", 1) or 1\n        if dcp_size > 1 and compilation_config[\"cudagraph_mode\"] == \"FULL_AND_PIECEWISE\":\n            logger.warning(\n                \"FULL cuda graph is not supported with DCP (decode_context_parallel_size=%d), \"\n                \"downgrading cudagraph_mode to PIECEWISE.\",\n                dcp_size,\n            )\n            compilation_config[\"cudagraph_mode\"] = \"PIECEWISE\"\n\n        compilation_config = json.dumps(compilation_config)\n        args = {\n            \"dtype\": self.config.dtype,\n            \"load_format\": self.config.load_format,\n            \"skip_tokenizer_init\": False,\n            \"distributed_executor_backend\": \"mp\",\n            \"worker_extension_cls\": \"verl.workers.rollout.vllm_rollout.utils.vLLMColocateWorkerExtension\",\n            \"trust_remote_code\": self.model_config.trust_remote_code,\n            \"max_model_len\": self.config.max_model_len,\n            \"max_num_seqs\": self.config.max_num_seqs,\n            \"enable_chunked_prefill\": self.config.enable_chunked_prefill,\n            \"max_num_batched_tokens\": self.config.max_num_batched_tokens,\n            \"enable_prefix_caching\": self.config.enable_prefix_caching,\n            \"enable_sleep_mode\": self.config.enable_sleep_mode,\n            \"logprobs_mode\": self.config.logprobs_mode,\n            \"enforce_eager\": self.config.enforce_eager,\n            \"gpu_memory_utilization\": self.config.gpu_memory_utilization,\n            \"disable_log_stats\": self.config.disable_log_stats,\n            \"tensor_parallel_size\": self.config.tensor_model_parallel_size,\n            \"seed\": self.replica_rank + self.config.get(\"seed\", 0),\n            \"override_generation_config\": json.dumps(override_generation_config),\n            \"quantization\": quantization,\n            \"hf_overrides\": hf_overrides,\n            \"scheduling_policy\": self.config.scheduling_policy,\n            \"compilation_config\": compilation_config,\n            **engine_kwargs,\n        }\n\n        # update profiler args\n        profiler_args = build_vllm_profiler_args(\n            self.profiler_controller.config, self.profiler_controller.tool_config, self.replica_rank\n        )\n        if _VLLM_VERSION >= version.parse(\"0.13.0\"):\n            # vLLM >= 0.13.0 supports profiler config via CLI args; env vars still work but will be deprecated\n            args.update(profiler_args)\n\n        if self.config.prometheus.enable:\n            if self.config.prometheus.served_model_name:\n                # Extract model name from path if it's a full path\n                served_model_name = self.config.prometheus.served_model_name\n                if \"/\" in served_model_name:\n                    # If it's a full path, extract the last part as model name\n                    served_model_name = served_model_name.split(\"/\")[-1]\n                args[\"served_model_name\"] = served_model_name\n\n        # mtp\n        if self.config.mtp.enable and self.config.mtp.enable_rollout:\n            speculative_config = {\n                \"method\": self.config.mtp.method,\n                \"num_speculative_tokens\": self.config.mtp.num_speculative_tokens,\n            }\n            args[\"speculative_config\"] = speculative_config\n\n        if self.config.data_parallel_size > 1:\n            assert self.gpus_per_node % self.config.tensor_model_parallel_size == 0, (\n                \"gpus_per_node should be divisible by tensor_model_parallel_size\"\n            )\n            data_parallel_size_local = self.gpus_per_node // self.config.tensor_model_parallel_size\n            assert len(self.workers) == data_parallel_size_local * self.config.tensor_model_parallel_size, (\n                f\"num workers ({len(self.workers)}) should be equal to \"\n                f\"dp_size_local ({data_parallel_size_local}) * tp_size ({self.config.tensor_model_parallel_size})\"\n            )\n            dp_args = {\n                \"data_parallel_size\": self.config.data_parallel_size,\n                \"data_parallel_size_local\": data_parallel_size_local,\n                \"data_parallel_start_rank\": self.node_rank * data_parallel_size_local,\n                \"data_parallel_address\": self._master_address,\n                \"data_parallel_rpc_port\": self._dp_rpc_port,\n            }\n            args.update(dp_args)\n\n        args.update({\"enable_expert_parallel\": self.config.expert_parallel_size > 1})\n\n        # used for torch.distributed.init_process_group\n        if self.nnodes > 1:\n            args.update(\n                {\n                    \"master_addr\": self._master_address,\n                    \"master_port\": self._master_port,\n                    \"node_rank\": self.node_rank,\n                    \"nnodes\": self.nnodes,\n                    \"data_parallel_address\": self._master_address,\n                    \"data_parallel_rpc_port\": self._dp_rpc_port,\n                }\n            )\n\n        # update lora-related args\n        lora_rank = self.model_config.lora.get(\"rank\", 0)\n        if lora_rank <= 0:\n            lora_rank = (\n                self.model_config.lora_rank\n            )  # FIXME: fallback to lora_rank for now, we should unify lora settings.\n\n        if self.model_config.lora.get(\"merge\", False):\n            lora_rank = 0\n\n        if lora_rank > 0:\n            lora_args = {\n                \"enable_lora\": True,\n                \"max_loras\": 1,\n                \"max_lora_rank\": get_vllm_max_lora_rank(lora_rank),\n            }\n            if self.model_config.lora.get(\"fully_sharded_loras\", False):\n                lora_args[\"fully_sharded_loras\"] = True\n            args.update(lora_args)\n\n        if self.config.enable_rollout_routing_replay:\n            args.update({\"enable_return_routed_experts\": True})\n\n        server_args = [\"serve\", self.model_config.local_path] + build_cli_args_from_config(args)\n\n        if self.replica_rank == 0:\n            pprint(server_args)\n\n        CMD_MODULES = [vllm.entrypoints.cli.serve]\n        parser = FlexibleArgumentParser(description=\"vLLM CLI\")\n        subparsers = parser.add_subparsers(required=False, dest=\"subparser\")\n        cmds = {}\n        for cmd_module in CMD_MODULES:\n            new_cmds = cmd_module.cmd_init()\n            for cmd in new_cmds:\n                cmd.subparser_init(subparsers).set_defaults(dispatch_function=cmd.cmd)\n                cmds[cmd.name] = cmd\n        server_args = parser.parse_args(args=server_args)\n        server_args.model = server_args.model_tag\n        if server_args.subparser in cmds:\n            cmds[server_args.subparser].validate(server_args)\n\n        # 3. launch server\n        if self.node_rank == 0:\n            self._master_sock.close()\n            self._dp_rpc_sock.close()\n            self._dp_master_sock.close()\n            await self.run_server(server_args)\n        else:\n            # TODO: avoid connect before master_sock close\n            await asyncio.sleep(3)\n            await self.run_headless(server_args)\n\n    async def run_server(self, args: argparse.Namespace):\n        engine_args = AsyncEngineArgs.from_cli_args(args)\n        usage_context = UsageContext.OPENAI_API_SERVER\n        vllm_config = engine_args.create_engine_config(usage_context=usage_context)\n        vllm_config.parallel_config.data_parallel_master_port = self._dp_master_port\n\n        fn_args = set(dict(inspect.signature(AsyncLLM.from_vllm_config).parameters).keys())\n        kwargs = {}\n        if \"enable_log_requests\" in fn_args:\n            kwargs[\"enable_log_requests\"] = engine_args.enable_log_requests\n        if \"disable_log_stats\" in fn_args:\n            kwargs[\"disable_log_stats\"] = engine_args.disable_log_stats\n\n        engine_client = AsyncLLM.from_vllm_config(vllm_config=vllm_config, usage_context=usage_context, **kwargs)\n\n        # Don't keep the dummy data in memory\n        await engine_client.reset_mm_cache()\n        await engine_client.collective_rpc(\n            method=\"monkey_patch_model\", kwargs={\"vocab_size\": len(self.model_config.tokenizer)}\n        )\n\n        build_app_sig = inspect.signature(build_app)\n        supported_tasks: tuple[Any, ...] = ()\n        if \"supported_tasks\" in build_app_sig.parameters:\n            supported_tasks = await engine_client.get_supported_tasks()\n            app = build_app(args, supported_tasks)\n        else:\n            app = build_app(args)\n\n        init_app_sig = inspect.signature(init_app_state)\n        if \"vllm_config\" in init_app_sig.parameters:\n            await init_app_state(engine_client, vllm_config, app.state, args)\n        elif \"supported_tasks\" in init_app_sig.parameters:\n            await init_app_state(engine_client, app.state, args, supported_tasks)\n        else:\n            await init_app_state(engine_client, app.state, args)\n        if self.replica_rank == 0 and self.node_rank == 0:\n            logger.info(f\"Initializing a V1 LLM engine with config: {vllm_config}\")\n\n        self.engine = engine_client\n        self._server_port, self._server_task = await run_uvicorn(app, args, self._server_address)\n\n    async def run_headless(self, args: argparse.Namespace):\n        \"\"\"Run headless server in a separate thread.\"\"\"\n        args.api_server_count = 0\n\n        def run_headless_wrapper():\n            with SuppressSignalInThread():\n                run_headless(args)\n\n        def on_run_headless_done(future: asyncio.Future):\n            try:\n                exc = future.exception()\n                if exc:\n                    logger.exception(f\"run_headless failed with exception: {exc}\")\n                else:\n                    logger.warning(\"run_headless completed successfully, but it's not expected.\")\n            except Exception as e:\n                logger.exception(f\"get result from run_headless failed: {e}\")\n            finally:\n                os._exit(1)\n\n        self.task = asyncio.create_task(asyncio.to_thread(run_headless_wrapper))\n        self.task.add_done_callback(on_run_headless_done)\n\n    async def generate(\n        self,\n        prompt_ids: list[int],\n        sampling_params: dict[str, Any],\n        request_id: str,\n        image_data: Optional[list[Any]] = None,\n        video_data: Optional[list[Any]] = None,\n        priority: int = 0,\n    ) -> TokenOutput:\n        \"\"\"Generate sequence with token-in-token-out.\"\"\"\n        prompt_ids = normalize_token_ids(prompt_ids)\n\n        # Calculate the maximum possible new tokens based on available context space\n        # This serves as a safety upper bound\n        max_possible_tokens = self.config.max_model_len - len(prompt_ids)\n        if max_possible_tokens < 0:\n            raise ValueError(\n                f\"Prompt length ({len(prompt_ids)}) exceeds the model's maximum context length \"\n                f\"({self.config.max_model_len}).\"\n            )\n\n        # Determine max_tokens from sampling_params or use configured response_length as default\n        if \"max_tokens\" in sampling_params:\n            max_tokens = sampling_params.pop(\"max_tokens\")\n        elif \"max_new_tokens\" in sampling_params:\n            # support sglang-style 'max_new_tokens' param\n            max_tokens = sampling_params.pop(\"max_new_tokens\")\n        else:\n            # Default to a calculation that considers configured lengths\n            # Cap max_tokens by response_length to ensure tensor alignment,\n            # and by remaining budget to prevent OOM in multi-turn rollouts.\n            max_tokens = min(\n                self.config.response_length, self.config.prompt_length + self.config.response_length - len(prompt_ids)\n            )\n\n        # Clamp max_tokens to the valid range [0, max_possible_tokens]\n        max_tokens = max(0, min(max_tokens, max_possible_tokens))\n\n        assert max_tokens <= max_possible_tokens, (\n            f\"max_tokens {max_tokens} exceeds available context space {max_possible_tokens}\"\n        )\n        sampling_params[\"logprobs\"] = 0 if sampling_params.pop(\"logprobs\", False) else None\n        sampling_params.setdefault(\"repetition_penalty\", self.config.get(\"repetition_penalty\", 1.0))\n        sampling_params = SamplingParams(max_tokens=max_tokens, **sampling_params)\n        prompt_ids = qwen2_5_vl_dedup_image_tokens(prompt_ids, self.model_config.processor)\n        multi_modal_data = {}\n        if image_data is not None:\n            multi_modal_data[\"image\"] = image_data\n        if video_data is not None:\n            multi_modal_data[\"video\"] = video_data\n\n        prompt = TokensPrompt(prompt_token_ids=prompt_ids, multi_modal_data=multi_modal_data)\n\n        # Add lora request\n        lora_request = None\n        if self.lora_as_adapter:\n            # Make sure we also check that the lora is already loaded in the engine\n            lora_loaded = VLLM_LORA_INT_ID in await self.engine.list_loras()\n            if lora_loaded:\n                lora_request = LoRARequest(\n                    lora_name=VLLM_LORA_NAME, lora_int_id=VLLM_LORA_INT_ID, lora_path=VLLM_LORA_PATH\n                )\n\n        generator = self.engine.generate(\n            prompt=prompt,\n            sampling_params=sampling_params,\n            request_id=request_id,\n            lora_request=lora_request,\n            priority=priority,\n        )\n\n        # Get final response\n        final_res: Optional[RequestOutput] = None\n        async for output in generator:\n            final_res = output\n        assert final_res is not None\n\n        token_ids = final_res.outputs[0].token_ids\n        log_probs = None\n        if sampling_params.logprobs is not None:\n            log_probs = [logprobs[token_ids[i]].logprob for i, logprobs in enumerate(final_res.outputs[0].logprobs)]\n\n        routed_experts = None\n        if self.config.enable_rollout_routing_replay:\n            routed_experts = final_res.outputs[0].routed_experts\n\n        # Determine stop reason from finish_reason\n        finish_reason = final_res.outputs[0].finish_reason\n        if finish_reason == \"abort\":\n            stop_reason = \"aborted\"\n        elif finish_reason in (\"stop\", \"length\"):\n            stop_reason = \"completed\"\n        else:\n            stop_reason = finish_reason  # for more stop reason in the future\n\n        num_preempted = None\n\n        if hasattr(final_res.outputs[0], \"num_preempted\"):\n            num_preempted = final_res.outputs[0].num_preempted\n\n        return TokenOutput(\n            token_ids=token_ids,\n            log_probs=log_probs,\n            routed_experts=routed_experts,\n            stop_reason=stop_reason,\n            num_preempted=num_preempted,\n            extra_fields={\"global_steps\": self.global_steps},\n        )\n\n    async def wake_up(self):\n        if self.node_rank != 0:\n            return\n\n        if self.rollout_mode == RolloutMode.HYBRID:\n            # In hybrid mode, rollout is wake up in `update_weights`\n            raise ValueError(f\"wake_up not support rollout_mode {self.rollout_mode}\")\n        elif self.rollout_mode == RolloutMode.COLOCATED:\n            # Directly call engine to wake up without sync weights.\n            await self.engine.wake_up(tags=[\"kv_cache\", \"weights\"])\n            await self.engine.reset_prefix_cache()\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            logger.info(\"skip wake_up in standalone mode\")\n\n    async def sleep(self):\n        if self.node_rank != 0 or not self.config.free_cache_engine:\n            return\n\n        if self.rollout_mode == RolloutMode.HYBRID:\n            # Don't use engine.sleep(level=2) here\n            # lora only update adapter weights, so set sleep level to 1\n            if self.lora_as_adapter or is_npu_available:\n                sleep_level = 1\n            else:\n                sleep_level = 2\n            await self.engine.collective_rpc(\"sleep\", kwargs={\"level\": sleep_level})\n\n            # clear encoder cache: https://github.com/vllm-project/vllm/pull/33452\n            # await self.engine.reset_encoder_cache()\n        elif self.rollout_mode == RolloutMode.COLOCATED:\n            await self.engine.sleep(level=1)\n        elif self.rollout_mode == RolloutMode.STANDALONE:\n            logger.info(\"skip sleep in standalone mode\")\n\n    async def start_profile(self, **kwargs):\n        if (\n            self.profiler_controller.check_enable()\n            and self.profiler_controller.check_this_rank()\n            and self.profiler_controller.is_discrete_mode()\n        ):\n            await self.engine.start_profile(**kwargs)\n\n    async def stop_profile(self):\n        if (\n            self.profiler_controller.check_enable()\n            and self.profiler_controller.check_this_rank()\n            and self.profiler_controller.is_discrete_mode()\n        ):\n            await self.engine.stop_profile()\n\n    async def clear_kv_cache(self):\n        if self.node_rank == 0:\n            await self.engine.reset_prefix_cache()\n\n    async def set_global_steps(self, global_steps: int):\n        \"\"\"Set the global steps of the model weights.\"\"\"\n        self.global_steps = global_steps\n\n    async def wait_for_requests_to_drain(self):\n        await self.engine.wait_for_requests_to_drain()\n\n    async def abort_all_requests(self, reset_prefix_cache: bool = True) -> dict[str, Any]:\n        \"\"\"Abort all ongoing generation requests.\n\n        On vLLM >= 0.12.0, uses AsyncLLM.pause_generation() to abort in-flight\n        requests, drain, and clear caches. The engine remains paused after this\n        call — use resume_generation() to accept new requests (e.g. before\n        validation).\n\n        On vLLM < 0.12.0, manually aborts each request and resets prefix cache.\n\n        Returns:\n            dict[str, Any]: Dictionary containing:\n                - aborted_count: Number of requests aborted\n                - request_ids: List of aborted request IDs\n        \"\"\"\n        try:\n            if _VLLM_VERSION >= version.parse(\"0.12.0\"):\n                # Snapshot request IDs before pausing for reporting\n                request_ids = list(self.engine.output_processor.request_states.keys())\n\n                # pause_generation with wait_for_inflight_requests=False will:\n                # 1. Set engine to paused state (blocks new generate calls)\n                # 2. Abort all in-flight requests\n                # 3. Wait for requests to drain\n                # 4. Clear prefix and mm caches if clear_cache=True\n                await self.engine.pause_generation(\n                    wait_for_inflight_requests=False,\n                    clear_cache=reset_prefix_cache,\n                )\n            else:\n                # Take an atomic snapshot to avoid race conditions with the vLLM engine thread\n                request_states_snapshot = list(self.engine.output_processor.request_states.items())\n                request_ids = [req_id for req_id, _ in request_states_snapshot]\n\n                if not request_ids:\n                    return {\"aborted_count\": 0, \"request_ids\": []}\n\n                # For each request, create an abort output and put it to its queue\n                # This allows the generator to receive the aborted result\n                from vllm.v1.engine import FinishReason\n\n                for _, req_state in request_states_snapshot:\n                    request_output = req_state.make_request_output(\n                        [], pooling_output=None, finish_reason=FinishReason.ABORT, stop_reason=None\n                    )\n                    req_state.queue.put(request_output)\n\n                # Abort requests in the output processor and engine core\n                self.engine.output_processor.abort_requests(request_ids)\n                await self.engine.engine_core.abort_requests_async(request_ids)\n\n                # Try to reset prefix cache to ensure clean state\n                if reset_prefix_cache:\n                    await self.clear_kv_cache()\n                    logger.info(\"Prefix cache reset after abort\")\n\n            logger.info(f\"Aborted {len(request_ids)} requests: {request_ids}\")\n            return {\"aborted_count\": len(request_ids), \"request_ids\": request_ids}\n\n        except Exception as e:\n            logger.error(f\"Error aborting requests: {e}\")\n            return {\"aborted_count\": 0, \"request_ids\": [], \"error\": str(e)}\n\n    async def resume_generation(self):\n        \"\"\"Resume generation after abort_all_requests (pause_generation).\n\n        Only effective on vLLM >= 0.12.0 where pause_generation is used.\n        No-op on older versions.\n        \"\"\"\n        if self.node_rank != 0:\n            return\n        if _VLLM_VERSION >= version.parse(\"0.12.0\"):\n            await self.engine.resume_generation()\n\n    async def abort_request(self, request_id: str, reset_prefix_cache: bool = True) -> dict[str, Any]:\n        \"\"\"Abort a specific generation request.\n\n        Args:\n            request_id: The ID of the request to abort.\n\n        Returns:\n            dict[str, Any]: Dictionary containing abort result.\n        \"\"\"\n        try:\n            request_states = self.engine.output_processor.request_states\n            req_state = request_states.get(request_id)\n\n            if req_state is None:\n                return {\"aborted\": False, \"error\": f\"Request {request_id} not found\"}\n\n            # Create abort output and put it to the queue\n            from vllm.v1.engine import FinishReason\n\n            request_output = req_state.make_request_output(\n                [], pooling_output=None, finish_reason=FinishReason.ABORT, stop_reason=None\n            )\n            req_state.queue.put(request_output)\n\n            # Abort in output processor and engine core\n            self.engine.output_processor.abort_requests([request_id])\n            await self.engine.engine_core.abort_requests_async([request_id])\n\n            # Try to reset prefix cache to ensure clean state\n            if reset_prefix_cache:\n                await self.clear_kv_cache()\n                logger.info(f\"Prefix cache reset after abort request {request_id}\")\n\n            logger.info(f\"Aborted request: {request_id}\")\n            return {\"aborted\": True, \"request_id\": request_id}\n\n        except Exception as e:\n            logger.error(f\"Error aborting request {request_id}: {e}\")\n            return {\"aborted\": False, \"request_id\": request_id, \"error\": str(e)}\n\n\nclass vLLMReplica(RolloutReplica):\n    def __init__(\n        self,\n        replica_rank: int,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        gpus_per_node: int = 8,\n        is_reward_model: bool = False,\n    ):\n        super().__init__(replica_rank, config, model_config, gpus_per_node, is_reward_model)\n        self.server_class = ray.remote(vLLMHttpServer)\n\n    async def launch_servers(self):\n        \"\"\"Launch http server in each node.\"\"\"\n        assert len(self.workers) == self.world_size, (\n            f\"worker number {len(self.workers)} not equal to world size {self.world_size}\"\n        )\n\n        # NOTE: We always use MP Executor backend whether it's single-node or multi-node.\n        # For multi-node without DP (e.g TP=16), need vllm>=0.11.1, https://github.com/vllm-project/vllm/pull/23691\n        if self.config.data_parallel_size == 1 and self.nnodes > 1:\n            assert _VLLM_VERSION >= version.parse(\"0.11.1\"), (\n                \"For multi-node MP Executor, either (1) set data_parallel_size > 1 or (2) upgrade vLLM to >= 0.11.1\"\n            )\n\n        # get (node_id, CUDA_VISIBLE_DEVICES) of all workers\n        worker_infos = await asyncio.gather(\n            *[\n                worker.__ray_call__.remote(\n                    lambda self: (\n                        ray.get_runtime_context().get_node_id(),\n                        ray.get_runtime_context().get_accelerator_ids()[get_resource_name()][0],\n                    )\n                )\n                for worker in self.workers\n            ]\n        )\n        worker_cuda_visible_devices = [worker_info[1] for worker_info in worker_infos]\n        worker_node_ids = [worker_info[0] for worker_info in worker_infos]\n\n        # create server actor in each node with node affinity and cuda visible devices\n        nnodes, gpus_per_replica_node = self.nnodes, self.gpus_per_replica_node\n        for node_rank in range(nnodes):\n            workers = self.workers[node_rank * gpus_per_replica_node : (node_rank + 1) * gpus_per_replica_node]\n            node_cuda_visible_devices = \",\".join(\n                worker_cuda_visible_devices[node_rank * gpus_per_replica_node : (node_rank + 1) * gpus_per_replica_node]\n            )\n            node_id = worker_node_ids[node_rank * gpus_per_replica_node]\n            name = (\n                f\"vllm_server_{self.replica_rank}_{node_rank}\"\n                if not self.is_reward_model\n                else f\"vllm_server_reward_{self.replica_rank}_{node_rank}\"\n            )\n\n            server = self.server_class.options(\n                scheduling_strategy=ray.util.scheduling_strategies.NodeAffinitySchedulingStrategy(\n                    node_id=node_id,\n                    soft=False,\n                ),\n                runtime_env={\n                    \"env_vars\": {\n                        \"RAY_EXPERIMENTAL_NOSET_CUDA_VISIBLE_DEVICES\": \"1\",\n                        \"RAY_EXPERIMENTAL_NOSET_ASCEND_RT_VISIBLE_DEVICES\": \"1\",\n                        # To prevent hanging or crash during synchronization of weights between actor and rollout\n                        # in disaggregated mode. See:\n                        # https://docs.vllm.ai/en/latest/usage/troubleshooting.html?h=nccl_cumem_enable#known-issues\n                        # https://github.com/vllm-project/vllm/blob/c6b0a7d3ba03ca414be1174e9bd86a97191b7090/vllm/worker/worker_base.py#L445\n                        \"NCCL_CUMEM_ENABLE\": \"0\",\n                    }\n                },\n                name=name,\n                max_concurrency=self.max_concurrency,\n            ).remote(\n                config=self.config,\n                model_config=self.model_config,\n                rollout_mode=self.rollout_mode,\n                workers=workers,\n                replica_rank=self.replica_rank,\n                node_rank=node_rank,\n                gpus_per_node=gpus_per_replica_node,\n                nnodes=nnodes,\n                cuda_visible_devices=node_cuda_visible_devices,\n            )\n            self.servers.append(server)\n\n        # launch http server in each node\n        master_address, master_port, dp_rpc_port = await self.servers[0].get_master_address.remote()\n        await asyncio.gather(\n            *[\n                server.launch_server.remote(\n                    master_address=master_address, master_port=master_port, dp_rpc_port=dp_rpc_port\n                )\n                for server in self.servers\n            ]\n        )\n\n        # get http server address from first server\n        server_address, server_port = await self.servers[0].get_server_address.remote()\n        self._server_handle = self.servers[0]\n        self._server_address = (\n            f\"[{server_address}]:{server_port}\"\n            if is_valid_ipv6_address(server_address)\n            else f\"{server_address}:{server_port}\"\n        )\n\n    async def sleep(self):\n        \"\"\"Sleep each rollout server.\"\"\"\n        # Drain DP engines for safe sleep.\n        await self.servers[0].wait_for_requests_to_drain.remote()\n        await asyncio.gather(*[server.sleep.remote() for server in self.servers])\n\n    async def abort_all_requests(self) -> dict[str, Any]:\n        \"\"\"Abort all ongoing generation requests across all servers.\n\n        Returns:\n            dict[str, Any]: Combined abort results from all servers.\n        \"\"\"\n        results = await asyncio.gather(*[server.abort_all_requests.remote() for server in self.servers])\n\n        total_aborted = sum(r.get(\"aborted_count\", 0) for r in results)\n        all_request_ids = []\n        for r in results:\n            all_request_ids.extend(r.get(\"request_ids\", []))\n\n        return {\n            \"aborted_count\": total_aborted,\n            \"request_ids\": all_request_ids,\n            \"server_results\": results,\n        }\n\n    async def resume_generation(self):\n        \"\"\"Resume generation on all servers after abort_all_requests.\"\"\"\n        await asyncio.gather(*[server.resume_generation.remote() for server in self.servers])\n\n    async def abort_request(self, request_id: str) -> dict[str, Any]:\n        \"\"\"Abort a specific request. Tries all servers since we don't know which one has it.\n\n        Args:\n            request_id: The ID of the request to abort.\n\n        Returns:\n            dict[str, Any]: Abort result.\n        \"\"\"\n        # TODO(petersh6): we should only abort on the server that has the request.\n        results = await asyncio.gather(*[server.abort_request.remote(request_id) for server in self.servers])\n\n        for r in results:\n            if r.get(\"aborted\", False):\n                return r\n\n        return {\"aborted\": False, \"request_id\": request_id, \"error\": \"Request not found on any server\"}\n"
  },
  {
    "path": "verl/workers/rollout/vllm_rollout/vllm_rollout.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nThe vllm_rollout that can be applied in different backend\nWhen working with FSDP:\n- Use DTensor weight loader (recommended) or HF weight loader\n- Utilize state_dict from the FSDP to synchronize the weights among tp ranks in vLLM\nWhen working with Megatron:\n- Use Megatron weight loader\n- During training, only the current pp stage holds the parameters\n- Before inference, broadcast the parameters of the current pp rank\n  to all other pp ranks (all pp ranks holds all the parameters)\n- Bind the parameters to the inference engine\n- Do inference in tp. pp is treated as additional dp\n- After inference, all the parameters that doesn't belong to this pp rank is freed.\n\"\"\"\n\nimport logging\nimport os\nimport time\nfrom typing import Any, Generator, Optional\n\nimport ray\nimport torch\nfrom packaging import version as vs\nfrom torch.distributed.device_mesh import DeviceMesh\n\nfrom verl import DataProto\nfrom verl.third_party.vllm import VLLM_SLEEP_LEVEL, get_version\nfrom verl.utils.device import get_device_id, is_support_ipc\nfrom verl.workers.config import HFModelConfig, RolloutConfig\nfrom verl.workers.rollout.base import BaseRollout\nfrom verl.workers.rollout.vllm_rollout.bucketed_weight_transfer import BucketedWeightSender\nfrom verl.workers.rollout.vllm_rollout.utils import get_device_uuid\n\nlogger = logging.getLogger(__file__)\nlogger.setLevel(os.getenv(\"VERL_LOGGING_LEVEL\", \"INFO\"))\n\n\ndef _check_vllm_version_for_sleep_level():\n    # https://github.com/vllm-project/vllm/issues/25171\n    minver = \"0.11.0\"\n    current_version = get_version(\"vllm\")\n    if not current_version:\n        logger.warning(\"Could not determine vLLM version, assuming an older version for sleep_level configuration.\")\n        return False\n    return vs.parse(current_version) >= vs.parse(minver)\n\n\nclass ServerAdapter(BaseRollout):\n    \"\"\"\n    vLLM server adapter used in native async mode, serve as a client to request vLLM server\n    to resume/release/update weights and kv_cache.\n    \"\"\"\n\n    def __init__(\n        self,\n        config: RolloutConfig,\n        model_config: HFModelConfig,\n        device_mesh: DeviceMesh,\n        replica_rank: int = -1,\n    ):\n        super().__init__(config, model_config, device_mesh)\n        self.server_handle: ray.actor.ActorHandle = None\n\n        rank = int(os.environ[\"RANK\"])\n        local_world_size = int(os.environ[\"RAY_LOCAL_WORLD_SIZE\"])\n        rollout_world_size = (\n            self.config.tensor_model_parallel_size\n            * self.config.data_parallel_size\n            * self.config.pipeline_model_parallel_size\n        )\n        if replica_rank == -1:\n            self.replica_rank = rank // rollout_world_size\n        else:\n            self.replica_rank = replica_rank\n        self.rollout_rank = rank % rollout_world_size\n        self.node_rank = self.rollout_rank // local_world_size\n\n        if config.layered_summon or (config.expert_parallel_size > 1 and not _check_vllm_version_for_sleep_level()):\n            logger.warning(\"Setting the sleep level to 1 may cause a memory overflow.\")\n            self.sleep_level = 1\n        else:\n            self.sleep_level = VLLM_SLEEP_LEVEL\n\n        self.device_uuid = get_device_uuid(get_device_id())\n        self.zmq_handle = f\"ipc:///tmp/rl-colocate-zmq-{self.device_uuid}.sock\"\n\n        self.use_shm = not is_support_ipc()\n        if self.use_shm:\n            logger.warning(\n                \"IPC is not supported on your devices. Falling back to shared memory for weight transfer, \"\n                \"which may cause performance degradation. If you are using Ascend NPUs, please ensure that \"\n                \"your software and CANN toolkit versions meet the requirements for IPC support. (Ascend HDK version \"\n                \">= 25.3.rc1 and CANN toolkit version >= 8.3.RC1)\"\n            )\n\n    async def _execute_method(\n        self,\n        method: str,\n        non_block: bool = False,\n        timeout: Optional[float] = None,\n        args: tuple = (),\n        kwargs: Optional[dict] = None,\n    ) -> Any:\n        \"\"\"Execute method on inference engine via ray.\n\n        Args:\n            method: The method name to execute on the server.\n            non_block: If True, execute the method asynchronously and return immediately.\n            timeout: Timeout for the collective_rpc call.\n            args: Positional arguments for the method.\n            kwargs: Keyword arguments for the method.\n\n        Returns:\n            The result of the method execution, or None if non_block=True.\n        \"\"\"\n        if self.rollout_rank != 0:\n            return None\n\n        # Lazy init http server adapter because http server is launched after hybrid engine.\n        if self.server_handle is None:\n            self.server_handle = ray.get_actor(f\"vllm_server_{self.replica_rank}_{self.node_rank}\")\n\n        future = self.server_handle.collective_rpc.remote(method, timeout=timeout, args=args, kwargs=kwargs)\n        return future if non_block else await future\n\n    async def resume(self, tags: list[str]):\n        \"\"\"Resume rollout weights or kv cache in GPU memory.\n\n        Args:\n            tags: weights or kv_cache.\n        \"\"\"\n        if self.config.free_cache_engine:\n            await self._execute_method(\"wake_up\", kwargs={\"tags\": tags})\n\n    async def release(self):\n        \"\"\"Release weights and kv cache in GPU memory.\"\"\"\n        if self.config.free_cache_engine:\n            await self._execute_method(\"sleep\", kwargs={\"level\": self.sleep_level})\n\n    @torch.no_grad()\n    async def update_weights(\n        self, weights: Generator[tuple[str, torch.Tensor], None, None], global_steps: int = None, **kwargs\n    ):\n        \"\"\"Update model weights via CUDA IPC (fallback to shared memory if IPC not supported) to inference workers.\"\"\"\n        start_time = time.time()\n\n        future = await self._execute_method(\n            \"update_weights_from_ipc\",\n            non_block=True,\n            kwargs={**kwargs, \"use_shm\": self.use_shm},\n        )\n\n        bucket_size_mb = self.config.checkpoint_engine.update_weights_bucket_megabytes\n        sender = BucketedWeightSender(\n            zmq_handle=self.zmq_handle,\n            bucket_size_mb=bucket_size_mb,\n            use_shm=self.use_shm,\n        )\n        await sender.async_send_weights(weights)\n\n        if future is not None:\n            await future\n\n        # reset prefix cache after updating weights\n        if self.rollout_rank == 0:\n            await self.server_handle.clear_kv_cache.remote()\n            if global_steps is not None:\n                await self.server_handle.set_global_steps.remote(global_steps)\n\n        if self.replica_rank == 0 and self.rollout_rank == 0:\n            logger.info(f\"update_weights done, time cost: {time.time() - start_time:.2f}s\")\n\n    def generate_sequences(self, prompts: DataProto) -> DataProto:\n        \"\"\"Batch generate sequences in sync mode.\n\n        Note: ServerAdapter uses async server mode and does not support synchronous\n        generation. Since SPMD mode was retired (PR #4411), the generation workflow\n        should use the async server interface instead.\n\n        Raises:\n            NotImplementedError: Always raised as sync generation is not supported.\n        \"\"\"\n        raise NotImplementedError(\n            \"ServerAdapter does not support synchronous generate_sequences(). \"\n            \"The vLLM SPMD mode was retired in PR #4411. For batch generation, \"\n            \"please use the async server interface via vLLMReplica and AsyncLLMServerManager, \"\n            \"or use HFRollout for synchronous generation. \"\n            \"See https://github.com/volcengine/verl/issues/4682 for more details.\"\n        )\n"
  },
  {
    "path": "verl/workers/sharding_manager/__init__.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/workers/sharding_manager/base.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nSharding manager to implement HybridEngine\n\"\"\"\n\nfrom verl import DataProto\n\n\nclass BaseShardingManager:\n    def __init__(self):\n        self.timing = {}\n\n    def __enter__(self):\n        pass\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        pass\n\n    def preprocess_data(self, data: DataProto) -> DataProto:\n        return data\n\n    def postprocess_data(self, data: DataProto) -> DataProto:\n        return data\n"
  },
  {
    "path": "verl/workers/sharding_manager/fsdp_ulysses.py",
    "content": "# Copyright 2024 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\"\"\"\nContains a resharding manager that binds weights from FSDP zero3 to XPerfGPT\n\"\"\"\n\nfrom torch.distributed.device_mesh import DeviceMesh\n\nfrom verl import DataProto\nfrom verl.protocol import all_gather_data_proto\nfrom verl.utils.ulysses import get_ulysses_sequence_parallel_group, set_ulysses_sequence_parallel_group\n\nfrom .base import BaseShardingManager\n\n\nclass FSDPUlyssesShardingManager(BaseShardingManager):\n    \"\"\"\n    Sharding manager to support data resharding when using FSDP + Ulysses\n    \"\"\"\n\n    def __init__(self, device_mesh: DeviceMesh):\n        super().__init__()\n        self.device_mesh = device_mesh\n        self.seed_offset = 12345\n\n    def __enter__(self):\n        if self.device_mesh is not None:\n            # We have a global SP group\n            # so we have to change to use model-specific sp group\n            self.prev_sp_group = get_ulysses_sequence_parallel_group()\n            set_ulysses_sequence_parallel_group(self.device_mesh[\"sp\"].get_group())\n            # TODO: check how to set seed for each model\n\n    def __exit__(self, exc_type, exc_value, traceback):\n        # restore random states\n        if self.device_mesh is not None:\n            # revert to previous sp group\n            set_ulysses_sequence_parallel_group(self.prev_sp_group)\n            # TODO: check how to set seed for each model\n\n    def preprocess_data(self, data: DataProto) -> DataProto:\n        \"\"\"\n        AllGather data from sp region\n        This is because the data is first sharded along the FSDP dimension as we utilize the DP_COMPUTE\n        In Ulysses, we need to make sure the same data is used across a SP group\n        \"\"\"\n        if self.device_mesh is not None:\n            group = self.device_mesh[\"sp\"].get_group()\n\n            all_gather_data_proto(data=data, process_group=group)\n        return data\n\n    def postprocess_data(self, data: DataProto) -> DataProto:\n        \"\"\"\n        Split the data to follow FSDP partition\n        \"\"\"\n        if self.device_mesh is not None:\n            sp_size = self.device_mesh[\"sp\"].size()\n            sp_rank = self.device_mesh[\"sp\"].get_local_rank()\n            data = data.chunk(chunks=sp_size)[sp_rank]\n        return data\n"
  },
  {
    "path": "verl/workers/utils/__init__.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n"
  },
  {
    "path": "verl/workers/utils/losses.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\n\nimport torch\nimport torch.nn.functional as F\nfrom tensordict import TensorDict\n\nfrom verl.trainer.ppo.core_algos import agg_loss, compute_value_loss, get_policy_loss_fn, kl_penalty\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.dataset.dataset_utils import DatasetPadMode\nfrom verl.utils.metric import AggregationType, Metric\nfrom verl.utils.torch_functional import masked_mean, masked_sum\nfrom verl.workers.config import ActorConfig, CriticConfig\nfrom verl.workers.utils.padding import no_padding_2_padding\n\n\ndef sft_loss(config: ActorConfig, model_output, data: TensorDict, dp_group=None):\n    pad_mode = tu.get_non_tensor_data(data=data, key=\"pad_mode\", default=DatasetPadMode.NO_PADDING)\n    dp_size = data[\"dp_size\"]\n    batch_num_tokens = data[\"batch_num_tokens\"]\n\n    log_prob = model_output[\"log_probs\"]\n\n    if pad_mode == DatasetPadMode.NO_PADDING:\n        # log_prob and loss mask are nested tensors of shape [bsz, j1]\n        # for each sample, loss mask shape is [1, prompt_length + response_length]\n        loss_mask = data[\"loss_mask\"]\n\n        log_prob_flatten = log_prob.values()\n        loss_mask_flatten = loss_mask.values()\n\n        # left-shift the loss mask by one token to align with log_prob\n        loss_mask_flatten = torch.roll(loss_mask_flatten, shifts=-1, dims=0)\n\n        # NOTE: loss is averaged over all tokens in the batch across all data parallel groups,\n        # For FSDP backend, the loss is directly used for backward; while for Megatron backend,\n        # the loss should be scaled by `num_microbatches` for pp schedule.\n        loss = -masked_sum(log_prob_flatten, loss_mask_flatten) / batch_num_tokens * dp_size\n    else:\n        response_mask = data[\"response_mask\"].to(bool)\n        loss = -masked_sum(log_prob, response_mask) / batch_num_tokens * dp_size\n\n    return loss, {}\n\n\ndef _slice_response_from_unpad_output(tensor: torch.Tensor, data: TensorDict) -> torch.Tensor:\n    \"\"\"Slice response from unpad model output.\n\n    Args:\n        tensor: model output tensor of shape [bsz, 1]\n        data: TensorDict with \"prompt_ids\", \"response_ids\", \"attention_mask\"\n\n    Returns:\n        tensor: sliced response tensor of shape [bsz, max_response_len]\n    \"\"\"\n    values = tensor.values() if tensor.is_nested else tensor\n    prompt_ids = data[\"prompts\"]\n    response_ids = data[\"responses\"]\n    attention_mask = data[\"attention_mask\"]\n\n    if prompt_ids.is_nested:\n        prompt_lens = prompt_ids.offsets().diff()\n        response_lens = response_ids.offsets().diff()\n        max_response_len = response_ids.offsets().max().item()\n    else:\n        assert not attention_mask.is_nested\n        prompt_lens = attention_mask[:, : prompt_ids.shape[1]].sum(dim=1)\n        response_lens = attention_mask[:, prompt_ids.shape[1] :].sum(dim=1)\n        max_response_len = response_ids.shape[1]\n\n    sequence_lens = prompt_lens + response_lens\n    sequence_offsets = sequence_lens.cumsum(dim=0)\n    assert sequence_offsets[-1].item() == values.shape[0]\n\n    response_list = []\n    for resp_len, seq_offset in zip(response_lens, sequence_offsets, strict=True):\n        pad_size = max_response_len - resp_len\n        # left-shift model output by one token for log_probs/values\n        response_list.append(F.pad(values[seq_offset - resp_len - 1 : seq_offset - 1], (0, pad_size)))\n\n    output = torch.stack(response_list, dim=0)\n    return output\n\n\ndef ppo_loss(config: ActorConfig, model_output, data: TensorDict, dp_group=None):\n    \"\"\"Computes ppo loss from model output (log_prob, entropy, values, etc. ) and old_log_probs from data.\"\"\"\n    log_prob = no_padding_2_padding(model_output[\"log_probs\"], data)\n    entropy = model_output.get(\"entropy\", None)\n    if entropy is not None:\n        entropy = no_padding_2_padding(entropy, data)\n\n    # global batch info for loss aggregation\n    config.global_batch_info[\"dp_size\"] = data[\"dp_size\"]\n    config.global_batch_info[\"batch_num_tokens\"] = data[\"batch_num_tokens\"]\n    config.global_batch_info[\"global_batch_size\"] = data[\"global_batch_size\"]\n    config.global_batch_info[\"loss_scale_factor\"] = config.loss_scale_factor\n\n    # assumes that if any of the global batch info is set, the policy_loss_fn will\n    # normalize using dp_size/global_bsz/global_token; in this case, metric aggregation should be SUM\n    # to reflect the mean loss over the global batch\n    if (\n        data[\"dp_size\"] > 1\n        or data[\"batch_num_tokens\"] is not None\n        or data[\"global_batch_size\"] is not None\n        or config.loss_scale_factor is not None\n    ):\n        metric_aggregation = AggregationType.SUM\n    else:\n        metric_aggregation = AggregationType.MEAN\n\n    metrics = {}\n\n    response_mask = data[\"response_mask\"].to(bool)\n    # compute policy loss\n    old_log_prob = data[\"old_log_probs\"]\n    advantages = data[\"advantages\"]\n    rollout_is_weights = data.get(\"rollout_is_weights\", None)\n\n    loss_agg_mode = config.loss_agg_mode\n\n    loss_mode = config.policy_loss.get(\"loss_mode\", \"vanilla\")\n\n    policy_loss_fn = get_policy_loss_fn(loss_mode)\n    pg_loss, pg_metrics = policy_loss_fn(\n        old_log_prob=old_log_prob,\n        log_prob=log_prob,\n        advantages=advantages,\n        response_mask=response_mask,\n        loss_agg_mode=loss_agg_mode,\n        config=config,\n        rollout_is_weights=rollout_is_weights,\n    )\n\n    # AggregationType.MEAN for pg metrics: assumes policy_loss_fn normalizes by local_bsz/local_tokens\n    # Ex: in compute_policy_loss_vanilla, pg_metrics are pg_clipfrac, ppo_kl, pg_clipfrac_lower\n    pg_metrics = Metric.from_dict(pg_metrics, aggregation=AggregationType.MEAN)\n\n    metrics.update(pg_metrics)\n    metrics[\"actor/pg_loss\"] = Metric(value=pg_loss, aggregation=metric_aggregation)\n    policy_loss = pg_loss\n\n    # add entropy loss\n    if entropy is not None:\n        entropy_loss = agg_loss(\n            loss_mat=entropy, loss_mask=response_mask, loss_agg_mode=loss_agg_mode, **config.global_batch_info\n        )\n        entropy_coeff = config.entropy_coeff\n        policy_loss -= entropy_coeff * entropy_loss\n        metrics[\"actor/entropy_loss\"] = Metric(value=entropy_loss, aggregation=metric_aggregation)\n\n    # add kl loss\n    if config.use_kl_loss:\n        ref_log_prob = data[\"ref_log_prob\"]\n        # compute kl loss\n        kld = kl_penalty(logprob=log_prob, ref_logprob=ref_log_prob, kl_penalty=config.kl_loss_type)\n        kl_loss = agg_loss(\n            loss_mat=kld, loss_mask=response_mask, loss_agg_mode=config.loss_agg_mode, **config.global_batch_info\n        )\n\n        policy_loss += kl_loss * config.kl_loss_coef\n        metrics[\"kl_loss\"] = Metric(value=kl_loss, aggregation=metric_aggregation)\n        metrics[\"kl_coef\"] = config.kl_loss_coef\n\n    return policy_loss, metrics\n\n\ndef value_loss(config: CriticConfig, model_output, data: TensorDict, dp_group=None):\n    \"\"\"value loss\n\n    Args:\n        config: CriticConfig\n        model_output: model output from the model\n        data: the input to the model\n        dp_group: data paralle group\n\n    Returns:\n        value loss\n    \"\"\"\n    vpreds = _slice_response_from_unpad_output(model_output[\"values\"], data)  # (bsz, response_length)\n\n    values = data[\"values\"]\n    returns = data[\"returns\"]\n    response_mask = data[\"response_mask\"].to(bool)\n\n    vf_loss, vf_clipfrac = compute_value_loss(\n        vpreds=vpreds,\n        values=values,\n        returns=returns,\n        response_mask=response_mask,\n        cliprange_value=config.cliprange_value,\n        loss_agg_mode=config.loss_agg_mode,\n    )\n\n    metrics = {}\n\n    metrics.update(\n        {\n            \"critic/vf_loss\": vf_loss.detach().item(),\n            \"critic/vf_clipfrac\": vf_clipfrac.detach().item(),\n            \"critic/vpred_mean\": masked_mean(vpreds, response_mask).detach().item(),\n        }\n    )\n\n    return vf_loss, metrics\n"
  },
  {
    "path": "verl/workers/utils/padding.py",
    "content": "# Copyright 2025 Bytedance Ltd. and/or its affiliates\n#\n# Licensed under the Apache License, Version 2.0 (the \"License\");\n# you may not use this file except in compliance with the License.\n# You may obtain a copy of the License at\n#\n#     http://www.apache.org/licenses/LICENSE-2.0\n#\n# Unless required by applicable law or agreed to in writing, software\n# distributed under the License is distributed on an \"AS IS\" BASIS,\n# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.\n# See the License for the specific language governing permissions and\n# limitations under the License.\n\nimport torch\nimport torch.nn.functional as F\nfrom tensordict import TensorDict\n\nfrom verl.utils import tensordict_utils as tu\nfrom verl.utils.attention_utils import index_first_axis, unpad_input\n\n\ndef left_right_2_no_padding(data: TensorDict) -> TensorDict:\n    \"\"\"\n    Convert TensorDict from left-right padding to no-padding format.\n\n    Args:\n        data: TensorDict with \"input_ids\", \"attention_mask\", \"response_mask\", \"position_ids\"\n\n    Returns:\n        data: TensorDict with\n        - Tensor includes NestedTensors like \"input_ids\", \"loss_mask\", \"position_ids\"\n        - NonTensorData includes \"max_seq_len\", \"max_response_len\", \"indices\"\n\n    Note:\n    1. the return input_ids/position_ids/loss_mask are nested tensor.\n    2. we will remove \"attention_mask\", \"response\" in the return data, but \"response_mask\" is kept.\n    \"\"\"\n    assert \"input_ids\" in data, \"input_ids is required in left-right padding data\"\n    assert \"attention_mask\" in data, \"attention_mask is required in left-right padding data\"\n    assert \"response_mask\" in data, \"response_mask is required in left-right padding data\"\n    assert \"position_ids\" in data, \"position_ids is required in left-right padding data\"\n\n    input_ids = data.pop(\"input_ids\")\n    attention_mask = data[\"attention_mask\"]\n    response_mask = data[\"response_mask\"]\n    position_ids = data[\"position_ids\"]  # (bs, seq_len) or # (bs, 4, seq_len)\n\n    max_seq_len, max_response_len = input_ids.shape[1], response_mask.shape[1]\n    tu.assign_non_tensor_data(data, \"max_seq_len\", max_seq_len)\n    tu.assign_non_tensor_data(data, \"max_response_len\", max_response_len)\n\n    input_ids_rmpad, indices, cu_seqlens, *_ = unpad_input(input_ids.unsqueeze(-1), attention_mask)\n    tu.assign_non_tensor_data(data, \"indices\", indices)\n\n    input_ids_nested = torch.nested.nested_tensor_from_jagged(input_ids_rmpad.squeeze(-1), offsets=cu_seqlens)\n\n    position_ids_list = []\n    for i in range(attention_mask.shape[0]):\n        curr_mask = attention_mask[i].bool()\n        curr_pos_ids = position_ids[i]\n        if curr_pos_ids.dim() == 1:  # (seq_len,)\n            valid_ids = curr_pos_ids[curr_mask]\n        else:  # (4, seq_len)\n            valid_ids = curr_pos_ids[:, curr_mask]\n        position_ids_list.append(valid_ids)\n    position_ids_nested = torch.nested.as_nested_tensor(position_ids_list, layout=torch.jagged)\n\n    data[\"input_ids\"] = input_ids_nested\n    data[\"position_ids\"] = position_ids_nested\n    data[\"loss_mask\"] = data[\"response_mask\"]\n\n    routed_experts = data.get(\"routed_experts\", None)\n    if routed_experts is not None and not routed_experts.is_nested:\n        if routed_experts.max() <= 255:\n            routed_experts = routed_experts.to(torch.uint8)\n        routed_experts_rmpad = index_first_axis(routed_experts.unsqueeze(-1).flatten(0, 1), indices)\n        routed_experts_nested = torch.nested.nested_tensor_from_jagged(\n            routed_experts_rmpad.squeeze(-1), offsets=cu_seqlens\n        )\n        data[\"routed_experts\"] = routed_experts_nested\n\n    return data\n\n\ndef no_padding_2_padding(tensor: torch.Tensor, data: TensorDict) -> torch.Tensor:\n    \"\"\"Slice response from unpad model output.\n\n    Args:\n        tensor: a nested tensor or a 1D tensor in shape (total_nnz,),\n            total_nnz is the total number of tokens across all sequences in the batch\n        data: TensorDict with \"prompts\", \"responses\", \"attention_mask\"\n\n    Returns:\n        tensor: sliced response tensor of shape [bsz, max_response_len]\n    \"\"\"\n    values = tensor.values() if tensor.is_nested else tensor\n    prompt_ids = data[\"prompts\"]\n    response_ids = data[\"responses\"]\n    attention_mask = data[\"attention_mask\"]\n\n    max_response_len = tu.get_non_tensor_data(data=data, key=\"max_response_len\", default=-1)\n\n    if prompt_ids.is_nested:\n        prompt_lens = prompt_ids.offsets().diff()\n        response_lens = response_ids.offsets().diff()\n        if max_response_len < 0:\n            max_response_len = response_lens.max().item()\n    else:\n        assert not attention_mask.is_nested\n        prompt_lens = attention_mask[:, : prompt_ids.shape[1]].sum(dim=1)\n        response_lens = attention_mask[:, prompt_ids.shape[1] :].sum(dim=1)\n        max_response_len = response_ids.shape[1]\n\n    sequence_lens = prompt_lens + response_lens\n    sequence_offsets = sequence_lens.cumsum(dim=0)\n    assert sequence_offsets[-1].item() == values.shape[0]\n\n    response_list = []\n    for resp_len, seq_offset in zip(response_lens, sequence_offsets, strict=True):\n        pad_size = max_response_len - resp_len\n        # left-shift model output by one token for log_probs/values\n        response_list.append(F.pad(values[seq_offset - resp_len - 1 : seq_offset - 1], (0, pad_size)))\n\n    output = torch.stack(response_list, dim=0)\n    return output\n"
  }
]